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