Line data    Source code

       1             : // Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
       2             : // of this source code is governed by a BSD-style license that can be found in
       3             : // the LICENSE file.
       4             : 
       5             : package rowblk
       6             : 
       7             : import (
       8             :         "bytes"
       9             : 
      10             :         "github.com/cockroachdb/errors"
      11             :         "github.com/cockroachdb/pebble/internal/base"
      12             :         "github.com/cockroachdb/pebble/internal/bytealloc"
      13             :         "github.com/cockroachdb/pebble/sstable/block"
      14             : )
      15             : 
      16             : // NewRewriter constructs a new rewriter.
      17           1 : func NewRewriter(comparer *base.Comparer, restartInterval int) *Rewriter {
      18           1 :         rw := &Rewriter{comparer: comparer}
      19           1 :         rw.writer.RestartInterval = restartInterval
      20           1 :         return rw
      21           1 : }
      22             : 
      23             : // Rewriter may be used to rewrite row-based blocks.
      24             : type Rewriter struct {
      25             :         comparer *base.Comparer
      26             :         // bufs resued each call to Rewrite.
      27             :         writer     Writer
      28             :         iter       Iter
      29             :         scratchKey base.InternalKey
      30             :         // alloc grown throughout the lifetime of the rewriter.
      31             :         keyAlloc bytealloc.A
      32             : }
      33             : 
      34             : // RewriteSuffixes rewrites the input block. It expects the input block to only
      35             : // contain keys with the suffix `from`. It rewrites the block to contain the
      36             : // same keys with the suffix `to`.
      37             : //
      38             : // RewriteSuffixes returns the start and end keys of the rewritten block, and the
      39             : // finished rewritten block.
      40             : func (r *Rewriter) RewriteSuffixes(
      41             :         input []byte, from []byte, to []byte,
      42           1 : ) (start, end base.InternalKey, rewritten []byte, err error) {
      43           1 :         if err := r.iter.Init(r.comparer.Compare, r.comparer.Split, input, block.NoTransforms); err != nil {
      44           0 :                 return base.InternalKey{}, base.InternalKey{}, nil, err
      45           0 :         }
      46           1 :         if cap(r.writer.restarts) < int(r.iter.restarts) {
      47           1 :                 r.writer.restarts = make([]uint32, 0, r.iter.restarts)
      48           1 :         }
      49           1 :         if cap(r.writer.buf) == 0 {
      50           1 :                 r.writer.buf = make([]byte, 0, len(input))
      51           1 :         }
      52           1 :         if cap(r.writer.restarts) < int(r.iter.numRestarts) {
      53           0 :                 r.writer.restarts = make([]uint32, 0, r.iter.numRestarts)
      54           0 :         }
      55           1 :         for kv := r.iter.First(); kv != nil; kv = r.iter.Next() {
      56           1 :                 if kv.Kind() != base.InternalKeyKindSet {
      57           0 :                         return base.InternalKey{}, base.InternalKey{}, nil,
      58           0 :                                 errors.New("key does not have expected kind (set)")
      59           0 :                 }
      60           1 :                 si := r.comparer.Split(kv.K.UserKey)
      61           1 :                 oldSuffix := kv.K.UserKey[si:]
      62           1 :                 if !bytes.Equal(oldSuffix, from) {
      63           1 :                         return base.InternalKey{}, base.InternalKey{}, nil,
      64           1 :                                 errors.Errorf("key has suffix %q, expected %q", oldSuffix, from)
      65           1 :                 }
      66           1 :                 newLen := si + len(to)
      67           1 :                 if cap(r.scratchKey.UserKey) < newLen {
      68           1 :                         r.scratchKey.UserKey = make([]byte, 0, len(kv.K.UserKey)*2+len(to)-len(from))
      69           1 :                 }
      70             : 
      71           1 :                 r.scratchKey.Trailer = kv.K.Trailer
      72           1 :                 r.scratchKey.UserKey = r.scratchKey.UserKey[:newLen]
      73           1 :                 copy(r.scratchKey.UserKey, kv.K.UserKey[:si])
      74           1 :                 copy(r.scratchKey.UserKey[si:], to)
      75           1 : 
      76           1 :                 // NB: for TableFormatPebblev3 and higher, since
      77           1 :                 // !iter.lazyValueHandling.hasValuePrefix, it will return the raw value
      78           1 :                 // in the block, which includes the 1-byte prefix. This is fine since bw
      79           1 :                 // also does not know about the prefix and will preserve it in bw.add.
      80           1 :                 v := kv.InPlaceValue()
      81           1 :                 r.writer.Add(r.scratchKey, v)
      82           1 :                 if start.UserKey == nil {
      83           1 :                         // Copy the first key.
      84           1 :                         start.Trailer = r.scratchKey.Trailer
      85           1 :                         r.keyAlloc, start.UserKey = r.keyAlloc.Copy(r.scratchKey.UserKey)
      86           1 :                 }
      87             :         }
      88             :         // Copy the last key.
      89           1 :         end.Trailer = r.scratchKey.Trailer
      90           1 :         r.keyAlloc, end.UserKey = r.keyAlloc.Copy(r.scratchKey.UserKey)
      91           1 : 
      92           1 :         r.iter = r.iter.ResetForReuse()
      93           1 :         return start, end, r.writer.Finish(), nil
      94             : }