github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/storage/multi_iterator_test.go (about)

     1  // Copyright 2017 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 storage
    12  
    13  import (
    14  	"bytes"
    15  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/keys"
    19  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    20  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    21  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    22  )
    23  
    24  func TestMultiIterator(t *testing.T) {
    25  	defer leaktest.AfterTest(t)()
    26  
    27  	rocksDB := newRocksDBInMem(roachpb.Attributes{}, 1<<20)
    28  	defer rocksDB.Close()
    29  
    30  	// Each `input` is turned into an iterator and these are passed to a new
    31  	// MultiIterator, which is fully iterated (using either NextKey or Next) and
    32  	// turned back into a string in the same format as `input`. This is compared
    33  	// to expectedNextKey or expectedNext.
    34  	//
    35  	// Input is a string containing key, timestamp, value tuples: first a single
    36  	// character key, then a single character timestamp walltime. If the
    37  	// character after the timestamp is an M, this entry is a "metadata" key
    38  	// (timestamp=0, sorts before any non-0 timestamp, and no value). If the
    39  	// character after the timestamp is an X, this entry is a deletion
    40  	// tombstone. Otherwise the value is the same as the timestamp.
    41  	tests := []struct {
    42  		inputs          []string
    43  		expectedNextKey string
    44  		expectedNext    string
    45  	}{
    46  		{[]string{}, "", ""},
    47  
    48  		{[]string{"a1"}, "a1", "a1"},
    49  		{[]string{"a1b1"}, "a1b1", "a1b1"},
    50  		{[]string{"a2a1"}, "a2", "a2a1"},
    51  		{[]string{"a2a1b1"}, "a2b1", "a2a1b1"},
    52  
    53  		{[]string{"a1", "a2"}, "a2", "a2a1"},
    54  		{[]string{"a2", "a1"}, "a2", "a2a1"},
    55  		{[]string{"a1", "b2"}, "a1b2", "a1b2"},
    56  		{[]string{"b2", "a1"}, "a1b2", "a1b2"},
    57  		{[]string{"a1b2", "b3"}, "a1b3", "a1b3b2"},
    58  		{[]string{"a1c2", "b3"}, "a1b3c2", "a1b3c2"},
    59  
    60  		{[]string{"aM", "a1"}, "aM", "aMa1"},
    61  		{[]string{"a1", "aM"}, "aM", "aMa1"},
    62  		{[]string{"aMa2", "a1"}, "aM", "aMa2a1"},
    63  		{[]string{"aMa1", "a2"}, "aM", "aMa2a1"},
    64  
    65  		{[]string{"a1", "a2X"}, "a2X", "a2Xa1"},
    66  		{[]string{"a1", "a2X", "a3"}, "a3", "a3a2Xa1"},
    67  		{[]string{"a1", "a2Xb2"}, "a2Xb2", "a2Xa1b2"},
    68  		{[]string{"a1b2", "a2X"}, "a2Xb2", "a2Xa1b2"},
    69  
    70  		{[]string{"a1", "a1"}, "a1", "a1"},
    71  		{[]string{"a4a2a1", "a4a3a1"}, "a4", "a4a3a2a1"},
    72  		{[]string{"a1b1", "a1b2"}, "a1b2", "a1b2b1"},
    73  		{[]string{"a1b2", "a1b1"}, "a1b2", "a1b2b1"},
    74  		{[]string{"a1b1", "a1b1"}, "a1b1", "a1b1"},
    75  	}
    76  	for _, test := range tests {
    77  		name := fmt.Sprintf("%q", test.inputs)
    78  		t.Run(name, func(t *testing.T) {
    79  			var iters []SimpleIterator
    80  			for _, input := range test.inputs {
    81  				batch := rocksDB.NewBatch()
    82  				defer batch.Close()
    83  				for i := 0; ; {
    84  					if i == len(input) {
    85  						break
    86  					}
    87  					k := []byte{input[i]}
    88  					ts := hlc.Timestamp{WallTime: int64(input[i+1])}
    89  					var v []byte
    90  					if i+1 < len(input) && input[i+1] == 'M' {
    91  						ts = hlc.Timestamp{}
    92  						v = nil
    93  					} else if i+2 < len(input) && input[i+2] == 'X' {
    94  						v = nil
    95  						i++
    96  					} else {
    97  						v = []byte{input[i+1]}
    98  					}
    99  					i += 2
   100  					if err := batch.Put(MVCCKey{Key: k, Timestamp: ts}, v); err != nil {
   101  						t.Fatalf("%+v", err)
   102  					}
   103  				}
   104  				iter := batch.NewIterator(IterOptions{UpperBound: roachpb.KeyMax})
   105  				defer iter.Close()
   106  				iters = append(iters, iter)
   107  			}
   108  
   109  			subtests := []struct {
   110  				name     string
   111  				expected string
   112  				fn       func(SimpleIterator)
   113  			}{
   114  				{"NextKey", test.expectedNextKey, (SimpleIterator).NextKey},
   115  				{"Next", test.expectedNext, (SimpleIterator).Next},
   116  			}
   117  			for _, subtest := range subtests {
   118  				t.Run(subtest.name, func(t *testing.T) {
   119  					var output bytes.Buffer
   120  					it := MakeMultiIterator(iters)
   121  					for it.SeekGE(MVCCKey{Key: keys.MinKey}); ; subtest.fn(it) {
   122  						ok, err := it.Valid()
   123  						if err != nil {
   124  							t.Fatalf("unexpected error: %+v", err)
   125  						}
   126  						if !ok {
   127  							break
   128  						}
   129  						output.Write(it.UnsafeKey().Key)
   130  						if it.UnsafeKey().Timestamp == (hlc.Timestamp{}) {
   131  							output.WriteRune('M')
   132  						} else {
   133  							output.WriteByte(byte(it.UnsafeKey().Timestamp.WallTime))
   134  							if len(it.UnsafeValue()) == 0 {
   135  								output.WriteRune('X')
   136  							}
   137  						}
   138  					}
   139  					if actual := output.String(); actual != subtest.expected {
   140  						t.Errorf("got %q expected %q", actual, subtest.expected)
   141  					}
   142  				})
   143  			}
   144  		})
   145  	}
   146  }