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