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  }