github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/replica_batch_updates_test.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package kvserver
    12  
    13  import (
    14  	"fmt"
    15  	"reflect"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/testutils"
    20  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    21  )
    22  
    23  // TestMaybeStripInFlightWrites verifies that in-flight writes declared on an
    24  // EndTxn request are stripped if the corresponding write or query intent is in
    25  // the same batch as the EndTxn.
    26  func TestMaybeStripInFlightWrites(t *testing.T) {
    27  	defer leaktest.AfterTest(t)()
    28  
    29  	keyA, keyB, keyC := roachpb.Key("a"), roachpb.Key("b"), roachpb.Key("c")
    30  	qi1 := &roachpb.QueryIntentRequest{RequestHeader: roachpb.RequestHeader{Key: keyA}}
    31  	qi1.Txn.Sequence = 1
    32  	put2 := &roachpb.PutRequest{RequestHeader: roachpb.RequestHeader{Key: keyB}}
    33  	put2.Sequence = 2
    34  	put3 := &roachpb.PutRequest{RequestHeader: roachpb.RequestHeader{Key: keyC}}
    35  	put3.Sequence = 3
    36  	delRng3 := &roachpb.DeleteRangeRequest{RequestHeader: roachpb.RequestHeader{Key: keyC}}
    37  	delRng3.Sequence = 3
    38  	scan3 := &roachpb.ScanRequest{RequestHeader: roachpb.RequestHeader{Key: keyC}}
    39  	scan3.Sequence = 3
    40  	et := &roachpb.EndTxnRequest{RequestHeader: roachpb.RequestHeader{Key: keyA}, Commit: true}
    41  	et.Sequence = 4
    42  	et.LockSpans = []roachpb.Span{{Key: keyC}}
    43  	et.InFlightWrites = []roachpb.SequencedWrite{{Key: keyA, Sequence: 1}, {Key: keyB, Sequence: 2}}
    44  	testCases := []struct {
    45  		reqs         []roachpb.Request
    46  		expIFW       []roachpb.SequencedWrite
    47  		expLockSpans []roachpb.Span
    48  		expErr       string
    49  	}{
    50  		{
    51  			reqs:         []roachpb.Request{et},
    52  			expIFW:       []roachpb.SequencedWrite{{Key: keyA, Sequence: 1}, {Key: keyB, Sequence: 2}},
    53  			expLockSpans: []roachpb.Span{{Key: keyC}},
    54  		},
    55  		// QueryIntents aren't stripped from the in-flight writes set on the
    56  		// slow-path of maybeStripInFlightWrites. This is intentional.
    57  		{
    58  			reqs:         []roachpb.Request{qi1, et},
    59  			expIFW:       []roachpb.SequencedWrite{{Key: keyA, Sequence: 1}, {Key: keyB, Sequence: 2}},
    60  			expLockSpans: []roachpb.Span{{Key: keyC}},
    61  		},
    62  		{
    63  			reqs:         []roachpb.Request{put2, et},
    64  			expIFW:       []roachpb.SequencedWrite{{Key: keyA, Sequence: 1}},
    65  			expLockSpans: []roachpb.Span{{Key: keyB}, {Key: keyC}},
    66  		},
    67  		{
    68  			reqs:   []roachpb.Request{put3, et},
    69  			expErr: "write in batch with EndTxn missing from in-flight writes",
    70  		},
    71  		{
    72  			reqs:         []roachpb.Request{qi1, put2, et},
    73  			expIFW:       nil,
    74  			expLockSpans: []roachpb.Span{{Key: keyA}, {Key: keyB}, {Key: keyC}},
    75  		},
    76  		{
    77  			reqs:         []roachpb.Request{qi1, put2, delRng3, et},
    78  			expIFW:       nil,
    79  			expLockSpans: []roachpb.Span{{Key: keyA}, {Key: keyB}, {Key: keyC}},
    80  		},
    81  		{
    82  			reqs:         []roachpb.Request{qi1, put2, scan3, et},
    83  			expIFW:       nil,
    84  			expLockSpans: []roachpb.Span{{Key: keyA}, {Key: keyB}, {Key: keyC}},
    85  		},
    86  		{
    87  			reqs:         []roachpb.Request{qi1, put2, delRng3, scan3, et},
    88  			expIFW:       nil,
    89  			expLockSpans: []roachpb.Span{{Key: keyA}, {Key: keyB}, {Key: keyC}},
    90  		},
    91  	}
    92  	for _, c := range testCases {
    93  		var ba roachpb.BatchRequest
    94  		ba.Add(c.reqs...)
    95  		t.Run(fmt.Sprint(ba), func(t *testing.T) {
    96  			resBa, err := maybeStripInFlightWrites(&ba)
    97  			if c.expErr == "" {
    98  				if err != nil {
    99  					t.Errorf("expected no error, got %v", err)
   100  				}
   101  				resArgs, _ := resBa.GetArg(roachpb.EndTxn)
   102  				resEt := resArgs.(*roachpb.EndTxnRequest)
   103  				if !reflect.DeepEqual(resEt.InFlightWrites, c.expIFW) {
   104  					t.Errorf("expected in-flight writes %v, got %v", c.expIFW, resEt.InFlightWrites)
   105  				}
   106  				if !reflect.DeepEqual(resEt.LockSpans, c.expLockSpans) {
   107  					t.Errorf("expected lock spans %v, got %v", c.expLockSpans, resEt.LockSpans)
   108  				}
   109  			} else {
   110  				if !testutils.IsError(err, c.expErr) {
   111  					t.Errorf("expected error %q, got %v", c.expErr, err)
   112  				}
   113  			}
   114  		})
   115  	}
   116  }