github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/storage/sst_iterator.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  
    16  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    17  	"github.com/cockroachdb/errors"
    18  	"github.com/cockroachdb/pebble/sstable"
    19  	"github.com/cockroachdb/pebble/vfs"
    20  )
    21  
    22  type sstIterator struct {
    23  	sst  *sstable.Reader
    24  	iter sstable.Iterator
    25  
    26  	mvccKey   MVCCKey
    27  	value     []byte
    28  	iterValid bool
    29  	err       error
    30  
    31  	// For allocation avoidance in NextKey.
    32  	nextKeyStart []byte
    33  
    34  	// roachpb.Verify k/v pairs on each call to Next()
    35  	verify bool
    36  }
    37  
    38  // NewSSTIterator returns a `SimpleIterator` for an in-memory sstable.
    39  // It's compatible with sstables written by `RocksDBSstFileWriter` and
    40  // Pebble's `sstable.Writer`, and assumes the keys use Cockroach's MVCC
    41  // format.
    42  func NewSSTIterator(path string) (SimpleIterator, error) {
    43  	file, err := vfs.Default.Open(path)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  	sst, err := sstable.NewReader(file, sstable.ReaderOptions{
    48  		Comparer: MVCCComparer,
    49  	})
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	return &sstIterator{sst: sst}, nil
    54  }
    55  
    56  // NewMemSSTIterator returns a `SimpleIterator` for an in-memory sstable.
    57  // It's compatible with sstables written by `RocksDBSstFileWriter` and
    58  // Pebble's `sstable.Writer`, and assumes the keys use Cockroach's MVCC
    59  // format.
    60  func NewMemSSTIterator(data []byte, verify bool) (SimpleIterator, error) {
    61  	sst, err := sstable.NewReader(vfs.NewMemFile(data), sstable.ReaderOptions{
    62  		Comparer: MVCCComparer,
    63  	})
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  	return &sstIterator{sst: sst, verify: verify}, nil
    68  }
    69  
    70  // Close implements the SimpleIterator interface.
    71  func (r *sstIterator) Close() {
    72  	if r.iter != nil {
    73  		r.err = errors.Wrap(r.iter.Close(), "closing sstable iterator")
    74  	}
    75  	if err := r.sst.Close(); err != nil && r.err == nil {
    76  		r.err = errors.Wrap(err, "closing sstable")
    77  	}
    78  }
    79  
    80  // SeekGE implements the SimpleIterator interface.
    81  func (r *sstIterator) SeekGE(key MVCCKey) {
    82  	if r.err != nil {
    83  		return
    84  	}
    85  	if r.iter == nil {
    86  		// Iterator creation happens on the first Seek as it involves I/O.
    87  		r.iter, r.err = r.sst.NewIter(nil /* lower */, nil /* upper */)
    88  		if r.err != nil {
    89  			return
    90  		}
    91  	}
    92  	var iKey *sstable.InternalKey
    93  	iKey, r.value = r.iter.SeekGE(EncodeKey(key))
    94  	if iKey != nil {
    95  		r.iterValid = true
    96  		r.mvccKey, r.err = DecodeMVCCKey(iKey.UserKey)
    97  	} else {
    98  		r.iterValid = false
    99  		r.err = r.iter.Error()
   100  	}
   101  	if r.iterValid && r.err == nil && r.verify {
   102  		r.err = roachpb.Value{RawBytes: r.value}.Verify(r.mvccKey.Key)
   103  	}
   104  }
   105  
   106  // Valid implements the SimpleIterator interface.
   107  func (r *sstIterator) Valid() (bool, error) {
   108  	return r.iterValid && r.err == nil, r.err
   109  }
   110  
   111  // Next implements the SimpleIterator interface.
   112  func (r *sstIterator) Next() {
   113  	if !r.iterValid || r.err != nil {
   114  		return
   115  	}
   116  	var iKey *sstable.InternalKey
   117  	iKey, r.value = r.iter.Next()
   118  	if iKey != nil {
   119  		r.mvccKey, r.err = DecodeMVCCKey(iKey.UserKey)
   120  	} else {
   121  		r.iterValid = false
   122  		r.err = r.iter.Error()
   123  	}
   124  	if r.iterValid && r.err == nil && r.verify {
   125  		r.err = roachpb.Value{RawBytes: r.value}.Verify(r.mvccKey.Key)
   126  	}
   127  }
   128  
   129  // NextKey implements the SimpleIterator interface.
   130  func (r *sstIterator) NextKey() {
   131  	if !r.iterValid || r.err != nil {
   132  		return
   133  	}
   134  	r.nextKeyStart = append(r.nextKeyStart[:0], r.mvccKey.Key...)
   135  	for r.Next(); r.iterValid && r.err == nil && bytes.Equal(r.nextKeyStart, r.mvccKey.Key); r.Next() {
   136  	}
   137  }
   138  
   139  // UnsafeKey implements the SimpleIterator interface.
   140  func (r *sstIterator) UnsafeKey() MVCCKey {
   141  	return r.mvccKey
   142  }
   143  
   144  // UnsafeValue implements the SimpleIterator interface.
   145  func (r *sstIterator) UnsafeValue() []byte {
   146  	return r.value
   147  }