github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/storage/sst_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  	"io/ioutil"
    15  	"path/filepath"
    16  	"reflect"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/testutils"
    20  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    21  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    22  )
    23  
    24  func runTestSSTIterator(t *testing.T, iter SimpleIterator, allKVs []MVCCKeyValue) {
    25  	// Drop the first kv so we can test Seek.
    26  	expected := allKVs[1:]
    27  
    28  	// Run the test multiple times to check re-Seeking.
    29  	for i := 0; i < 3; i++ {
    30  		var kvs []MVCCKeyValue
    31  		for iter.SeekGE(expected[0].Key); ; iter.Next() {
    32  			ok, err := iter.Valid()
    33  			if err != nil {
    34  				t.Fatalf("%+v", err)
    35  			}
    36  			if !ok {
    37  				break
    38  			}
    39  			kv := MVCCKeyValue{
    40  				Key: MVCCKey{
    41  					Key:       append([]byte(nil), iter.UnsafeKey().Key...),
    42  					Timestamp: iter.UnsafeKey().Timestamp,
    43  				},
    44  				Value: append([]byte(nil), iter.UnsafeValue()...),
    45  			}
    46  			kvs = append(kvs, kv)
    47  		}
    48  		if ok, err := iter.Valid(); err != nil {
    49  			t.Fatalf("%+v", err)
    50  		} else if ok {
    51  			t.Fatalf("expected !ok")
    52  		}
    53  
    54  		lastElemKey := expected[len(expected)-1].Key
    55  		seekTo := MVCCKey{Key: lastElemKey.Key.Next()}
    56  
    57  		iter.SeekGE(seekTo)
    58  		if ok, err := iter.Valid(); err != nil {
    59  			t.Fatalf("%+v", err)
    60  		} else if ok {
    61  			foundKey := iter.UnsafeKey()
    62  			t.Fatalf("expected !ok seeking to lastEmem.Next(). foundKey %s < seekTo %s: %t",
    63  				foundKey, seekTo, foundKey.Less(seekTo))
    64  		}
    65  
    66  		if !reflect.DeepEqual(kvs, expected) {
    67  			t.Fatalf("got %+v but expected %+v", kvs, expected)
    68  		}
    69  	}
    70  }
    71  
    72  func TestSSTIterator(t *testing.T) {
    73  	defer leaktest.AfterTest(t)()
    74  
    75  	sstFile := &MemFile{}
    76  	sst := MakeIngestionSSTWriter(sstFile)
    77  	defer sst.Close()
    78  	var allKVs []MVCCKeyValue
    79  	for i := 0; i < 10; i++ {
    80  		kv := MVCCKeyValue{
    81  			Key: MVCCKey{
    82  				Key:       []byte{'A' + byte(i)},
    83  				Timestamp: hlc.Timestamp{WallTime: int64(i)},
    84  			},
    85  			Value: []byte{'a' + byte(i)},
    86  		}
    87  		if err := sst.Put(kv.Key, kv.Value); err != nil {
    88  			t.Fatalf("%+v", err)
    89  		}
    90  		allKVs = append(allKVs, kv)
    91  	}
    92  
    93  	if err := sst.Finish(); err != nil {
    94  		t.Fatalf("%+v", err)
    95  	}
    96  
    97  	t.Run("Disk", func(t *testing.T) {
    98  		tempDir, cleanup := testutils.TempDir(t)
    99  		defer cleanup()
   100  
   101  		path := filepath.Join(tempDir, "data.sst")
   102  		if err := ioutil.WriteFile(path, sstFile.Data(), 0600); err != nil {
   103  			t.Fatalf("%+v", err)
   104  		}
   105  
   106  		iter, err := NewSSTIterator(path)
   107  		if err != nil {
   108  			t.Fatalf("%+v", err)
   109  		}
   110  		defer iter.Close()
   111  		runTestSSTIterator(t, iter, allKVs)
   112  	})
   113  	t.Run("Mem", func(t *testing.T) {
   114  		iter, err := NewMemSSTIterator(sstFile.Data(), false)
   115  		if err != nil {
   116  			t.Fatalf("%+v", err)
   117  		}
   118  		defer iter.Close()
   119  		runTestSSTIterator(t, iter, allKVs)
   120  	})
   121  }
   122  
   123  func TestCockroachComparer(t *testing.T) {
   124  	defer leaktest.AfterTest(t)()
   125  
   126  	keyAMetadata := MVCCKey{
   127  		Key: []byte("a"),
   128  	}
   129  	keyA2 := MVCCKey{
   130  		Key:       []byte("a"),
   131  		Timestamp: hlc.Timestamp{WallTime: 2},
   132  	}
   133  	keyA1 := MVCCKey{
   134  		Key:       []byte("a"),
   135  		Timestamp: hlc.Timestamp{WallTime: 1},
   136  	}
   137  	keyB2 := MVCCKey{
   138  		Key:       []byte("b"),
   139  		Timestamp: hlc.Timestamp{WallTime: 2},
   140  	}
   141  
   142  	if x := MVCCComparer.Compare(EncodeKey(keyAMetadata), EncodeKey(keyA1)); x != -1 {
   143  		t.Errorf("expected key metadata to sort first got: %d", x)
   144  	}
   145  	if x := MVCCComparer.Compare(EncodeKey(keyA2), EncodeKey(keyA1)); x != -1 {
   146  		t.Errorf("expected higher timestamp to sort first got: %d", x)
   147  	}
   148  	if x := MVCCComparer.Compare(EncodeKey(keyA2), EncodeKey(keyB2)); x != -1 {
   149  		t.Errorf("expected lower key to sort first got: %d", x)
   150  	}
   151  }