github.com/aminovpavel/nomad@v0.11.8/nomad/drainer/drain_heap_test.go (about)

     1  package drainer
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/hashicorp/nomad/testutil"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestDeadlineHeap_Interface(t *testing.T) {
    13  	t.Parallel()
    14  	require := require.New(t)
    15  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    16  	require.Implements((*DrainDeadlineNotifier)(nil), h)
    17  }
    18  
    19  func TestDeadlineHeap_WatchAndGet(t *testing.T) {
    20  	t.Parallel()
    21  	require := require.New(t)
    22  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    23  
    24  	now := time.Now()
    25  	nodeID := "1"
    26  	wait := 10 * time.Millisecond
    27  	deadline := now.Add(wait)
    28  	h.Watch(nodeID, deadline)
    29  
    30  	var batch []string
    31  	select {
    32  	case batch = <-h.NextBatch():
    33  	case <-time.After(testutil.Timeout(3 * wait)):
    34  		t.Fatal("timeout")
    35  	}
    36  
    37  	require.Len(batch, 1)
    38  	require.Equal(nodeID, batch[0])
    39  }
    40  
    41  func TestDeadlineHeap_WatchThenUpdateAndGet(t *testing.T) {
    42  	t.Parallel()
    43  	require := require.New(t)
    44  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    45  
    46  	now := time.Now()
    47  	nodeID := "1"
    48  	wait := 10 * time.Millisecond
    49  	deadline := now.Add(wait)
    50  
    51  	// Initially watch way in the future
    52  	h.Watch(nodeID, now.Add(24*time.Hour))
    53  
    54  	// Rewatch
    55  	h.Watch(nodeID, deadline)
    56  
    57  	var batch []string
    58  	select {
    59  	case batch = <-h.NextBatch():
    60  	case <-time.After(testutil.Timeout(2 * wait)):
    61  		t.Fatal("timeout")
    62  	}
    63  
    64  	require.Len(batch, 1)
    65  	require.Equal(nodeID, batch[0])
    66  }
    67  
    68  func TestDeadlineHeap_MultiwatchAndDelete(t *testing.T) {
    69  	t.Parallel()
    70  	require := require.New(t)
    71  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    72  
    73  	now := time.Now()
    74  	wait := 50 * time.Millisecond
    75  	deadline := now.Add(wait)
    76  
    77  	nodeID1 := "1"
    78  	nodeID2 := "2"
    79  	h.Watch(nodeID1, deadline)
    80  	h.Watch(nodeID2, deadline)
    81  
    82  	time.Sleep(1 * time.Millisecond)
    83  	h.Remove(nodeID2)
    84  
    85  	var batch []string
    86  	select {
    87  	case batch = <-h.NextBatch():
    88  	case <-time.After(testutil.Timeout(2 * wait)):
    89  		t.Fatal("timeout")
    90  	}
    91  
    92  	require.Len(batch, 1)
    93  	require.Equal(nodeID1, batch[0])
    94  }
    95  
    96  func TestDeadlineHeap_WatchCoalesce(t *testing.T) {
    97  	t.Parallel()
    98  	require := require.New(t)
    99  	h := NewDeadlineHeap(context.Background(), 100*time.Millisecond)
   100  
   101  	now := time.Now()
   102  
   103  	group1 := map[string]time.Time{
   104  		"1": now.Add(5 * time.Millisecond),
   105  		"2": now.Add(10 * time.Millisecond),
   106  		"3": now.Add(20 * time.Millisecond),
   107  		"4": now.Add(100 * time.Millisecond),
   108  	}
   109  
   110  	group2 := map[string]time.Time{
   111  		"10": now.Add(350 * time.Millisecond),
   112  		"11": now.Add(360 * time.Millisecond),
   113  	}
   114  
   115  	for _, g := range []map[string]time.Time{group1, group2} {
   116  		for n, d := range g {
   117  			h.Watch(n, d)
   118  		}
   119  	}
   120  
   121  	var batch []string
   122  	select {
   123  	case batch = <-h.NextBatch():
   124  	case <-time.After(testutil.Timeout(time.Second)):
   125  		t.Fatal("timeout")
   126  	}
   127  
   128  	require.Len(batch, len(group1))
   129  	for nodeID := range group1 {
   130  		require.Contains(batch, nodeID)
   131  	}
   132  	batch = nil
   133  
   134  	select {
   135  	case batch = <-h.NextBatch():
   136  	case <-time.After(testutil.Timeout(2 * time.Second)):
   137  		t.Fatal("timeout")
   138  	}
   139  
   140  	require.Len(batch, len(group2))
   141  	for nodeID := range group2 {
   142  		require.Contains(batch, nodeID)
   143  	}
   144  
   145  	select {
   146  	case <-h.NextBatch():
   147  		t.Fatal("unexpected batch")
   148  	case <-time.After(testutil.Timeout(100 * time.Millisecond)):
   149  	}
   150  }
   151  
   152  func TestDeadlineHeap_MultipleForce(t *testing.T) {
   153  	t.Parallel()
   154  	require := require.New(t)
   155  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
   156  
   157  	nodeID := "1"
   158  	deadline := time.Time{}
   159  	h.Watch(nodeID, deadline)
   160  
   161  	var batch []string
   162  	select {
   163  	case batch = <-h.NextBatch():
   164  	case <-time.After(testutil.Timeout(10 * time.Millisecond)):
   165  		t.Fatal("timeout")
   166  	}
   167  
   168  	require.Len(batch, 1)
   169  	require.Equal(nodeID, batch[0])
   170  
   171  	nodeID = "2"
   172  	h.Watch(nodeID, deadline)
   173  	select {
   174  	case batch = <-h.NextBatch():
   175  	case <-time.After(testutil.Timeout(10 * time.Millisecond)):
   176  		t.Fatal("timeout")
   177  	}
   178  
   179  	require.Len(batch, 1)
   180  	require.Equal(nodeID, batch[0])
   181  }