|
| 1 | +// Copyright 2015 The Cockroach Authors. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or |
| 12 | +// implied. See the License for the specific language governing |
| 13 | +// permissions and limitations under the License. See the AUTHORS file |
| 14 | +// for names of contributors. |
| 15 | +// |
| 16 | +// Author: Spencer Kimball (spencer.kimball@gmail.com) |
| 17 | +// Author: Veteran Lu (23907238@qq.com) |
| 18 | + |
| 19 | +package roachpb |
| 20 | + |
| 21 | +import ( |
| 22 | + "encoding/hex" |
| 23 | + "reflect" |
| 24 | + "strings" |
| 25 | + "testing" |
| 26 | +) |
| 27 | + |
| 28 | +func TestBatchSplit(t *testing.T) { |
| 29 | + get := &GetRequest{} |
| 30 | + scan := &ScanRequest{} |
| 31 | + put := &PutRequest{} |
| 32 | + spl := &AdminSplitRequest{} |
| 33 | + dr := &DeleteRangeRequest{} |
| 34 | + bt := &BeginTransactionRequest{} |
| 35 | + et := &EndTransactionRequest{} |
| 36 | + rv := &ReverseScanRequest{} |
| 37 | + np := &NoopRequest{} |
| 38 | + testCases := []struct { |
| 39 | + reqs []Request |
| 40 | + sizes []int |
| 41 | + canSplitET bool |
| 42 | + }{ |
| 43 | + {[]Request{get, put}, []int{1, 1}, true}, |
| 44 | + {[]Request{get, get, get, put, put, get, get}, []int{3, 2, 2}, true}, |
| 45 | + {[]Request{spl, get, scan, spl, get}, []int{1, 2, 1, 1}, true}, |
| 46 | + {[]Request{spl, spl, get, spl}, []int{1, 1, 1, 1}, true}, |
| 47 | + {[]Request{bt, put, et}, []int{2, 1}, true}, |
| 48 | + {[]Request{get, scan, get, dr, rv, put, et}, []int{3, 1, 1, 1, 1}, true}, |
| 49 | + // Same one again, but this time don't allow EndTransaction to be split. |
| 50 | + {[]Request{get, scan, get, dr, rv, put, et}, []int{3, 1, 1, 2}, false}, |
| 51 | + // An invalid request in real life, but it demonstrates that we'll always |
| 52 | + // split **after** an EndTransaction (because either the next request |
| 53 | + // wants to be alone, or its flags can't match the current flags, which |
| 54 | + // have isAlone set). Could be useful if we ever want to allow executing |
| 55 | + // multiple batches back-to-back. |
| 56 | + {[]Request{et, scan, et}, []int{1, 2}, false}, |
| 57 | + // Check that Noop can mix with other requests regardless of flags. |
| 58 | + {[]Request{np, put, np}, []int{3}, true}, |
| 59 | + {[]Request{np, spl, np}, []int{3}, true}, |
| 60 | + {[]Request{np, rv, np}, []int{3}, true}, |
| 61 | + {[]Request{np, np, et}, []int{3}, true}, // et does not split off |
| 62 | + } |
| 63 | + |
| 64 | + for i, test := range testCases { |
| 65 | + ba := BatchRequest{} |
| 66 | + for _, args := range test.reqs { |
| 67 | + ba.Add(args) |
| 68 | + } |
| 69 | + var partLen []int |
| 70 | + var recombined []RequestUnion |
| 71 | + for _, part := range ba.Split(test.canSplitET) { |
| 72 | + recombined = append(recombined, part...) |
| 73 | + partLen = append(partLen, len(part)) |
| 74 | + } |
| 75 | + if !reflect.DeepEqual(partLen, test.sizes) { |
| 76 | + t.Errorf("%d: expected chunks %v, got %v", i, test.sizes, partLen) |
| 77 | + } |
| 78 | + if !reflect.DeepEqual(recombined, ba.Requests) { |
| 79 | + t.Errorf("%d: started with:\n%+v\ngot back:\n%+v", i, ba.Requests, recombined) |
| 80 | + } |
| 81 | + } |
| 82 | +} |
| 83 | + |
| 84 | +func TestFlagsToStr(t *testing.T) { |
| 85 | + var ba BatchRequest |
| 86 | + ba.Add(&PutRequest{}) |
| 87 | + ba.Add(&AdminSplitRequest{}) |
| 88 | + exp := "AdWrAl" |
| 89 | + if act := flagsToStr(ba.flags()); act != exp { |
| 90 | + t.Fatalf("expected %s, got %s", exp, act) |
| 91 | + } |
| 92 | +} |
| 93 | + |
| 94 | +func TestBatchTraceID(t *testing.T) { |
| 95 | + var ba BatchRequest |
| 96 | + ba.Add(&ReverseScanRequest{}) |
| 97 | + ba.Add(&IncrementRequest{}) |
| 98 | + |
| 99 | + expID := "cRdWrRgRv@0.000000000,0" |
| 100 | + expName := expID[1:] |
| 101 | + if actID, actName := ba.TraceID(), ba.TraceName(); expID != actID || expName != actName { |
| 102 | + t.Fatalf("expected (%s,%s), got (%s, %s)", expID, expName, actID, actName) |
| 103 | + } |
| 104 | + |
| 105 | + uuidStr := "a53eef22-35f3-4c69-8e98-c563100e028c" |
| 106 | + bytes, err := hex.DecodeString(strings.Replace(uuidStr, "-", "", -1)) |
| 107 | + if err != nil { |
| 108 | + t.Fatal(err) |
| 109 | + } |
| 110 | + ba.Txn = &Transaction{ID: bytes} |
| 111 | + expID = "t" + uuidStr |
| 112 | + expName = expID[:9] |
| 113 | + if actID, actName := ba.TraceID(), ba.TraceName(); expID != actID || expName != actName { |
| 114 | + t.Fatalf("expected (%s,%s), got (%s, %s)", expID, expName, actID, actName) |
| 115 | + } |
| 116 | +} |
| 117 | + |
| 118 | +func TestBatchRequestGetArg(t *testing.T) { |
| 119 | + testCases := []struct { |
| 120 | + bu []RequestUnion |
| 121 | + expB, expG bool |
| 122 | + }{ |
| 123 | + {[]RequestUnion{}, false, false}, |
| 124 | + {[]RequestUnion{{Get: &GetRequest{}}}, false, true}, |
| 125 | + {[]RequestUnion{{EndTransaction: &EndTransactionRequest{}}, {Get: &GetRequest{}}}, false, true}, |
| 126 | + {[]RequestUnion{{EndTransaction: &EndTransactionRequest{}}}, true, false}, |
| 127 | + {[]RequestUnion{{Get: &GetRequest{}}, {EndTransaction: &EndTransactionRequest{}}}, true, true}, |
| 128 | + } |
| 129 | + |
| 130 | + for i, c := range testCases { |
| 131 | + br := BatchRequest{Requests: c.bu} |
| 132 | + if _, r := br.GetArg(EndTransaction); r != c.expB { |
| 133 | + t.Errorf("%d: unexpected batch request for %v: %v", i, c.bu, r) |
| 134 | + } |
| 135 | + if _, r := br.GetArg(Get); r != c.expG { |
| 136 | + t.Errorf("%d: unexpected get match for %v: %v", i, c.bu, r) |
| 137 | + } |
| 138 | + |
| 139 | + } |
| 140 | +} |
0 commit comments