vitess.io/vitess@v0.16.2/go/vt/vtgate/buffer/buffer_helper_test.go (about) 1 package buffer 2 3 import ( 4 "context" 5 "errors" 6 "fmt" 7 "strings" 8 "testing" 9 "time" 10 11 "vitess.io/vitess/go/vt/discovery" 12 "vitess.io/vitess/go/vt/proto/query" 13 topodatapb "vitess.io/vitess/go/vt/proto/topodata" 14 vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc" 15 "vitess.io/vitess/go/vt/vterrors" 16 ) 17 18 type failover func(buf *Buffer, tablet *topodatapb.Tablet, keyspace, shard string, now time.Time) 19 20 func testAllImplementations(t *testing.T, runTest func(t *testing.T, fail failover)) { 21 t.Helper() 22 23 t.Run("HealthCheck", func(t *testing.T) { 24 t.Helper() 25 runTest(t, func(buf *Buffer, tablet *topodatapb.Tablet, keyspace, shard string, now time.Time) { 26 buf.ProcessPrimaryHealth(&discovery.TabletHealth{ 27 Tablet: tablet, 28 Target: &query.Target{Keyspace: keyspace, Shard: shard, TabletType: topodatapb.TabletType_PRIMARY}, 29 PrimaryTermStartTime: now.Unix(), 30 }) 31 }) 32 }) 33 34 t.Run("KeyspaceEvent", func(t *testing.T) { 35 t.Helper() 36 runTest(t, func(buf *Buffer, tablet *topodatapb.Tablet, keyspace, shard string, now time.Time) { 37 buf.HandleKeyspaceEvent(&discovery.KeyspaceEvent{ 38 Keyspace: keyspace, 39 Shards: []discovery.ShardEvent{ 40 { 41 Tablet: tablet.Alias, 42 Target: &query.Target{Keyspace: keyspace, Shard: shard, TabletType: topodatapb.TabletType_PRIMARY}, 43 Serving: true, 44 }, 45 }, 46 }) 47 }) 48 }) 49 } 50 51 // issueRequest simulates executing a request which goes through the buffer. 52 // If the buffering returned an error, it will be sent on the returned channel. 53 func issueRequest(ctx context.Context, t *testing.T, b *Buffer, err error) chan error { 54 return issueRequestAndBlockRetry(ctx, t, b, err, nil /* markRetryDone */) 55 } 56 57 // issueRequestAndBlockRetry is the same as issueRequest() but allows to signal 58 // when the buffer should be informed that the retry is done. (For that, 59 // the channel "markRetryDone" must be closed.) 60 func issueRequestAndBlockRetry(ctx context.Context, t *testing.T, b *Buffer, err error, markRetryDone chan struct{}) chan error { 61 bufferingStopped := make(chan error) 62 63 go func() { 64 retryDone, err := b.WaitForFailoverEnd(ctx, keyspace, shard, failoverErr) 65 if err != nil { 66 bufferingStopped <- err 67 } 68 if markRetryDone != nil { 69 // Wait for the test's signal before we tell the buffer that we retried. 70 <-markRetryDone 71 } 72 if retryDone != nil { 73 defer retryDone() 74 } 75 defer close(bufferingStopped) 76 }() 77 78 return bufferingStopped 79 } 80 81 // waitForRequestsInFlight blocks until the buffer queue has reached "count". 82 // This check is potentially racy and therefore retried up to a timeout of 10s. 83 func waitForRequestsInFlight(b *Buffer, count int) error { 84 start := time.Now() 85 sb := b.getOrCreateBuffer(keyspace, shard) 86 for { 87 got, want := sb.testGetSize(), count 88 if got == want { 89 return nil 90 } 91 92 if time.Since(start) > 10*time.Second { 93 return fmt.Errorf("wrong buffered requests in flight: got = %v, want = %v", got, want) 94 } 95 time.Sleep(1 * time.Millisecond) 96 } 97 } 98 99 // waitForState polls the buffer data for up to 10 seconds and returns an error 100 // if shardBuffer doesn't have the wanted state by then. 101 func waitForState(b *Buffer, want bufferState) error { 102 sb := b.getOrCreateBuffer(keyspace, shard) 103 start := time.Now() 104 for { 105 got := sb.testGetState() 106 if got == want { 107 return nil 108 } 109 110 if time.Since(start) > 10*time.Second { 111 return fmt.Errorf("wrong buffer state: got = %v, want = %v", got, want) 112 } 113 time.Sleep(1 * time.Millisecond) 114 } 115 } 116 117 // waitForPoolSlots waits up to 10s that all buffer pool slots have been 118 // returned. The wait is necessary because in some cases the buffer code 119 // does not block itself on the wait. But in any case, the slot should be 120 // returned when the request has finished. See also shardBuffer.unblockAndWait(). 121 func waitForPoolSlots(b *Buffer, want int) error { 122 start := time.Now() 123 for { 124 got := b.bufferSizeSema.Size() 125 if got == want { 126 return nil 127 } 128 129 if time.Since(start) > 10*time.Second { 130 return fmt.Errorf("not all pool slots were returned: got = %v, want = %v", got, want) 131 } 132 time.Sleep(1 * time.Millisecond) 133 } 134 } 135 136 func waitForRequestsExceededWindow(count int) error { 137 start := time.Now() 138 for { 139 got, want := requestsEvicted.Counts()[statsKeyJoinedWindowExceeded], int64(count) 140 if got == want { 141 return nil 142 } 143 144 if time.Since(start) > 10*time.Second { 145 return fmt.Errorf("wrong number of requests which exceeded their buffering window: got = %v, want = %v", got, want) 146 } 147 time.Sleep(1 * time.Millisecond) 148 } 149 } 150 151 func isCanceledError(err error) error { 152 if err == nil { 153 return fmt.Errorf("buffering should have stopped early and returned an error because the request was canceled from the outside") 154 } 155 if got, want := vterrors.Code(err), vtrpcpb.Code_UNAVAILABLE; got != want { 156 return fmt.Errorf("wrong error code for canceled buffered request. got = %v, want = %v", got, want) 157 } 158 if got, want := err.Error(), "context was canceled before failover finished"; !strings.Contains(got, want) { 159 return fmt.Errorf("canceled buffered request should return a different error message. got = %v, want = %v", got, want) 160 } 161 return nil 162 } 163 164 // isEvictedError returns nil if "err" is or contains "entryEvictedError". 165 func isEvictedError(err error) error { 166 if err == nil { 167 return errors.New("request should have been evicted because the buffer was full") 168 } 169 if got, want := vterrors.Code(err), vtrpcpb.Code_UNAVAILABLE; got != want { 170 return fmt.Errorf("wrong error code for evicted buffered request. got = %v, want = %v full error: %v", got, want, err) 171 } 172 if got, want := err.Error(), entryEvictedError.Error(); !strings.Contains(got, want) { 173 return fmt.Errorf("evicted buffered request should return a different error message. got = %v, want substring = %v", got, want) 174 } 175 return nil 176 } 177 178 // resetVariables resets the task level variables. The code does not reset these 179 // with very failover. 180 func resetVariables() { 181 starts.ResetAll() 182 stops.ResetAll() 183 184 utilizationSum.ResetAll() 185 utilizationDryRunSum.ResetAll() 186 187 requestsBuffered.ResetAll() 188 requestsBufferedDryRun.ResetAll() 189 requestsDrained.ResetAll() 190 requestsEvicted.ResetAll() 191 requestsSkipped.ResetAll() 192 } 193 194 // checkVariables makes sure that the invariants described in variables.go 195 // hold up. 196 func checkVariables(t *testing.T) { 197 for k, buffered := range requestsBuffered.Counts() { 198 evicted := int64(0) 199 // The evicted count is grouped by Reason i.e. the entries are named 200 // "<Keyspace>.<Shard>.<Reason>". Match all reasons for this shard. 201 for withReason, v := range requestsEvicted.Counts() { 202 if strings.HasPrefix(withReason, fmt.Sprintf("%s.", k)) { 203 evicted += v 204 } 205 } 206 drained := requestsDrained.Counts()[k] 207 208 if buffered != evicted+drained { 209 t.Fatalf("buffered == evicted + drained is violated: %v != %v + %v", buffered, evicted, drained) 210 } else { 211 t.Logf("buffered == evicted + drained: %v == %v + %v", buffered, evicted, drained) 212 } 213 } 214 }