Upgrade versions of vendored libs

vendor/github.com/cockroachdb/pebble/batch.go

@@ -5,6 +5,7 @@
 package pebble
 
 import (
+	"context"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -12,6 +13,7 @@ import (
 	"sort"
 	"sync"
 	"sync/atomic"
+	"time"
 	"unsafe"
 
 	"github.com/cockroachdb/errors"
@@ -82,7 +84,7 @@ func (d DeferredBatchOp) Finish() error {
 // consumers should use a batch per goroutine or provide their own
 // synchronization.
 //
-// Indexing
+// # Indexing
 //
 // Batches can be optionally indexed (see DB.NewIndexedBatch). An indexed batch
 // allows iteration via an Iterator (see Batch.NewIter). The iterator provides
@@ -104,7 +106,7 @@ func (d DeferredBatchOp) Finish() error {
 // significantly slower than inserting into a non-indexed batch. Only use an
 // indexed batch if you require reading from it.
 //
-// Atomic commit
+// # Atomic commit
 //
 // The operations in a batch are persisted by calling Batch.Commit which is
 // equivalent to calling DB.Apply(batch). A batch is committed atomically by
@@ -115,7 +117,7 @@ func (d DeferredBatchOp) Finish() error {
 // Batch.Commit will guarantee that the batch is persisted to disk before
 // returning. See commitPipeline for more on the implementation details.
 //
-// Large batches
+// # Large batches
 //
 // The size of a batch is limited only by available memory (be aware that
 // indexed batches require considerably additional memory for the skiplist
@@ -140,34 +142,35 @@ func (d DeferredBatchOp) Finish() error {
 // memtable have the same big-O time, but the constant factor dominates
 // here. Sorting is significantly faster and uses significantly less memory.
 //
-// Internal representation
+// # Internal representation
 //
 // The internal batch representation is a contiguous byte buffer with a fixed
 // 12-byte header, followed by a series of records.
 //
-//   +-------------+------------+--- ... ---+
-//   | SeqNum (8B) | Count (4B) |  Entries  |
-//   +-------------+------------+--- ... ---+
+//	+-------------+------------+--- ... ---+
+//	| SeqNum (8B) | Count (4B) |  Entries  |
+//	+-------------+------------+--- ... ---+
 //
 // Each record has a 1-byte kind tag prefix, followed by 1 or 2 length prefixed
 // strings (varstring):
 //
-//   +-----------+-----------------+-------------------+
-//   | Kind (1B) | Key (varstring) | Value (varstring) |
-//   +-----------+-----------------+-------------------+
+//	+-----------+-----------------+-------------------+
+//	| Kind (1B) | Key (varstring) | Value (varstring) |
+//	+-----------+-----------------+-------------------+
 //
 // A varstring is a varint32 followed by N bytes of data. The Kind tags are
 // exactly those specified by InternalKeyKind. The following table shows the
 // format for records of each kind:
 //
-//   InternalKeyKindDelete         varstring
-//   InternalKeyKindLogData        varstring
-//   InternalKeyKindSet            varstring varstring
-//   InternalKeyKindMerge          varstring varstring
-//   InternalKeyKindRangeDelete    varstring varstring
-//   InternalKeyKindRangeKeySet    varstring varstring
-//   InternalKeyKindRangeKeyUnset  varstring varstring
-//   InternalKeyKindRangeKeyDelete varstring varstring
+//	InternalKeyKindDelete         varstring
+//	InternalKeyKindLogData        varstring
+//	InternalKeyKindIngestSST      varstring
+//	InternalKeyKindSet            varstring varstring
+//	InternalKeyKindMerge          varstring varstring
+//	InternalKeyKindRangeDelete    varstring varstring
+//	InternalKeyKindRangeKeySet    varstring varstring
+//	InternalKeyKindRangeKeyUnset  varstring varstring
+//	InternalKeyKindRangeKeyDelete varstring varstring
 //
 // The intuitive understanding here are that the arguments to Delete, Set,
 // Merge, DeleteRange and RangeKeyDelete are encoded into the batch. The
@@ -210,6 +213,8 @@ type Batch struct {
 
 	// The db to which the batch will be committed. Do not change this field
 	// after the batch has been created as it might invalidate internal state.
+	// Batch.memTableSize is only refreshed if Batch.db is set. Setting db to
+	// nil once it has been set implies that the Batch has encountered an error.
 	db *DB
 
 	// The count of records in the batch. This count will be stored in the batch
@@ -255,9 +260,74 @@ type Batch struct {
 	// memtable.
 	flushable *flushableBatch
 
+	// ingestedSSTBatch indicates that the batch contains one or more key kinds
+	// of InternalKeyKindIngestSST. If the batch contains key kinds of IngestSST
+	// then it will only contain key kinds of IngestSST.
+	ingestedSSTBatch bool
+
+	// Synchronous Apply uses the commit WaitGroup for both publishing the
+	// seqnum and waiting for the WAL fsync (if needed). Asynchronous
+	// ApplyNoSyncWait, which implies WriteOptions.Sync is true, uses the commit
+	// WaitGroup for publishing the seqnum and the fsyncWait WaitGroup for
+	// waiting for the WAL fsync.
+	//
+	// TODO(sumeer): if we find that ApplyNoSyncWait in conjunction with
+	// SyncWait is causing higher memory usage because of the time duration
+	// between when the sync is already done, and a goroutine calls SyncWait
+	// (followed by Batch.Close), we could separate out {fsyncWait, commitErr}
+	// into a separate struct that is allocated separately (using another
+	// sync.Pool), and only that struct needs to outlive Batch.Close (which
+	// could then be called immediately after ApplyNoSyncWait). commitStats
+	// will also need to be in this separate struct.
 	commit    sync.WaitGroup
+	fsyncWait sync.WaitGroup
+
+	commitStats BatchCommitStats
+
 	commitErr error
-	applied   uint32 // updated atomically
+	applied   atomic.Bool
+}
+
+// BatchCommitStats exposes stats related to committing a batch.
+//
+// NB: there is no Pebble internal tracing (using LoggerAndTracer) of slow
+// batch commits. The caller can use these stats to do their own tracing as
+// needed.
+type BatchCommitStats struct {
+	// TotalDuration is the time spent in DB.{Apply,ApplyNoSyncWait} or
+	// Batch.Commit, plus the time waiting in Batch.SyncWait. If there is a gap
+	// between calling ApplyNoSyncWait and calling SyncWait, that gap could
+	// include some duration in which real work was being done for the commit
+	// and will not be included here. This missing time is considered acceptable
+	// since the goal of these stats is to understand user-facing latency.
+	//
+	// TotalDuration includes time spent in various queues both inside Pebble
+	// and outside Pebble (I/O queues, goroutine scheduler queue, mutex wait
+	// etc.). For some of these queues (which we consider important) the wait
+	// times are included below -- these expose low-level implementation detail
+	// and are meant for expert diagnosis and subject to change. There may be
+	// unaccounted time after subtracting those values from TotalDuration.
+	TotalDuration time.Duration
+	// SemaphoreWaitDuration is the wait time for semaphores in
+	// commitPipeline.Commit.
+	SemaphoreWaitDuration time.Duration
+	// WALQueueWaitDuration is the wait time for allocating memory blocks in the
+	// LogWriter (due to the LogWriter not writing fast enough).
+	WALQueueWaitDuration time.Duration
+	// MemTableWriteStallDuration is the wait caused by a write stall due to too
+	// many memtables (due to not flushing fast enough).
+	MemTableWriteStallDuration time.Duration
+	// L0ReadAmpWriteStallDuration is the wait caused by a write stall due to
+	// high read amplification in L0 (due to not compacting fast enough out of
+	// L0).
+	L0ReadAmpWriteStallDuration time.Duration
+	// WALRotationDuration is the wait time for WAL rotation, which includes
+	// syncing and closing the old WAL and creating (or reusing) a new one.
+	WALRotationDuration time.Duration
+	// CommitWaitDuration is the wait for publishing the seqnum plus the
+	// duration for the WAL sync (if requested). The former should be tiny and
+	// one can assume that this is all due to the WAL sync.
+	CommitWaitDuration time.Duration
 }
 
 var _ Reader = (*Batch)(nil)
@@ -348,13 +418,16 @@ func (b *Batch) refreshMemTableSize() {
 		if !ok {
 			break
 		}
-		b.memTableSize += memTableEntrySize(len(key), len(value))
 		switch kind {
 		case InternalKeyKindRangeDelete:
 			b.countRangeDels++
 		case InternalKeyKindRangeKeySet, InternalKeyKindRangeKeyUnset, InternalKeyKindRangeKeyDelete:
 			b.countRangeKeys++
+		case InternalKeyKindIngestSST:
+			// This key kind doesn't contribute to the memtable size.
+			continue
 		}
+		b.memTableSize += memTableEntrySize(len(key), len(value))
 	}
 }
 
@@ -362,6 +435,9 @@ func (b *Batch) refreshMemTableSize() {
 //
 // It is safe to modify the contents of the arguments after Apply returns.
 func (b *Batch) Apply(batch *Batch, _ *WriteOptions) error {
+	if b.ingestedSSTBatch {
+		panic("pebble: invalid batch application")
+	}
 	if len(batch.data) == 0 {
 		return nil
 	}
@@ -392,6 +468,8 @@ func (b *Batch) Apply(batch *Batch, _ *WriteOptions) error {
 				b.countRangeDels++
 			case InternalKeyKindRangeKeySet, InternalKeyKindRangeKeyUnset, InternalKeyKindRangeKeyDelete:
 				b.countRangeKeys++
+			case InternalKeyKindIngestSST:
+				panic("pebble: invalid key kind for batch")
 			}
 			if b.index != nil {
 				var err error
@@ -797,6 +875,28 @@ func (b *Batch) LogData(data []byte, _ *WriteOptions) error {
 	return nil
 }
 
+// IngestSST adds the FileNum for an sstable to the batch. The data will only be
+// written to the WAL (not added to memtables or sstables).
+func (b *Batch) ingestSST(fileNum base.FileNum) {
+	if b.Empty() {
+		b.ingestedSSTBatch = true
+	} else if !b.ingestedSSTBatch {
+		// Batch contains other key kinds.
+		panic("pebble: invalid call to ingestSST")
+	}
+
+	origMemTableSize := b.memTableSize
+	var buf [binary.MaxVarintLen64]byte
+	length := binary.PutUvarint(buf[:], uint64(fileNum))
+	b.prepareDeferredKeyRecord(length, InternalKeyKindIngestSST)
+	copy(b.deferredOp.Key, buf[:length])
+	// Since IngestSST writes only to the WAL and does not affect the memtable,
+	// we restore b.memTableSize to its original value. Note that Batch.count
+	// is not reset because for the InternalKeyKindIngestSST the count is the
+	// number of sstable paths which have been added to the batch.
+	b.memTableSize = origMemTableSize
+}
+
 // Empty returns true if the batch is empty, and false otherwise.
 func (b *Batch) Empty() bool {
 	return len(b.data) <= batchHeaderLen
@@ -845,10 +945,16 @@ func (b *Batch) SetRepr(data []byte) error {
 // later mutations. Its view can be refreshed via RefreshBatchSnapshot or
 // SetOptions().
 func (b *Batch) NewIter(o *IterOptions) *Iterator {
+	return b.NewIterWithContext(context.Background(), o)
+}
+
+// NewIterWithContext is like NewIter, and additionally accepts a context for
+// tracing.
+func (b *Batch) NewIterWithContext(ctx context.Context, o *IterOptions) *Iterator {
 	if b.index == nil {
 		return &Iterator{err: ErrNotIndexed}
 	}
-	return b.db.newIterInternal(b, nil /* snapshot */, o)
+	return b.db.newIter(ctx, b, nil /* snapshot */, o)
 }
 
 // newInternalIter creates a new internalIterator that iterates over the
@@ -1089,8 +1195,10 @@ func (b *Batch) Reset() {
 	b.rangeKeysSeqNum = 0
 	b.flushable = nil
 	b.commit = sync.WaitGroup{}
+	b.fsyncWait = sync.WaitGroup{}
+	b.commitStats = BatchCommitStats{}
 	b.commitErr = nil
-	atomic.StoreUint32(&b.applied, 0)
+	b.applied.Store(false)
 	if b.data != nil {
 		if cap(b.data) > batchMaxRetainedSize {
 			// If the capacity of the buffer is larger than our maximum
@@ -1159,7 +1267,9 @@ func (b *Batch) setCount(v uint32) {
 }
 
 // Count returns the count of memtable-modifying operations in this batch. All
-// operations with the except of LogData increment this count.
+// operations with the except of LogData increment this count. For IngestSSTs,
+// count is only used to indicate the number of SSTs ingested in the record, the
+// batch isn't applied to the memtable.
 func (b *Batch) Count() uint32 {
 	if b.count > math.MaxUint32 {
 		panic(ErrInvalidBatch)
@@ -1205,6 +1315,26 @@ func batchDecodeStr(data []byte) (odata []byte, s []byte, ok bool) {
 	return data[v:], data[:v], true
 }
 
+// SyncWait is to be used in conjunction with DB.ApplyNoSyncWait.
+func (b *Batch) SyncWait() error {
+	now := time.Now()
+	b.fsyncWait.Wait()
+	if b.commitErr != nil {
+		b.db = nil // prevent batch reuse on error
+	}
+	waitDuration := time.Since(now)
+	b.commitStats.CommitWaitDuration += waitDuration
+	b.commitStats.TotalDuration += waitDuration
+	return b.commitErr
+}
+
+// CommitStats returns stats related to committing the batch. Should be called
+// after Batch.Commit, DB.Apply. If DB.ApplyNoSyncWait is used, should be
+// called after Batch.SyncWait.
+func (b *Batch) CommitStats() BatchCommitStats {
+	return b.commitStats
+}
+
 // BatchReader iterates over the entries contained in a batch.
 type BatchReader []byte
 
@@ -1336,6 +1466,19 @@ func (i *batchIter) Next() (*InternalKey, base.LazyValue) {
 	return ikey, base.MakeInPlaceValue(i.value())
 }
 
+func (i *batchIter) NextPrefix(succKey []byte) (*InternalKey, LazyValue) {
+	// Because NextPrefix was invoked `succKey` must be ≥ the key at i's current
+	// position. Seek the arena iterator using TrySeekUsingNext.
+	ikey := i.iter.SeekGE(succKey, base.SeekGEFlagsNone.EnableTrySeekUsingNext())
+	for ikey != nil && ikey.SeqNum() >= i.snapshot {
+		ikey = i.iter.Next()
+	}
+	if ikey == nil {
+		return nil, base.LazyValue{}
+	}
+	return ikey, base.MakeInPlaceValue(i.value())
+}
+
 func (i *batchIter) Prev() (*InternalKey, base.LazyValue) {
 	ikey := i.iter.Prev()
 	for ikey != nil && ikey.SeqNum() >= i.snapshot {
@@ -1773,6 +1916,12 @@ func (i *flushableBatchIter) Prev() (*InternalKey, base.LazyValue) {
 	return &i.key, i.value()
 }
 
+// Note: flushFlushableBatchIter.NextPrefix mirrors the implementation of
+// flushableBatchIter.NextPrefix due to performance. Keep the two in sync.
+func (i *flushableBatchIter) NextPrefix(succKey []byte) (*InternalKey, LazyValue) {
+	return i.SeekGE(succKey, base.SeekGEFlagsNone.EnableTrySeekUsingNext())
+}
+
 func (i *flushableBatchIter) getKey(index int) InternalKey {
 	e := &i.offsets[index]
 	kind := InternalKeyKind(i.data[e.offset])
@@ -1866,6 +2015,10 @@ func (i *flushFlushableBatchIter) First() (*InternalKey, base.LazyValue) {
 	return key, val
 }
 
+func (i *flushFlushableBatchIter) NextPrefix(succKey []byte) (*InternalKey, base.LazyValue) {
+	panic("pebble: Prev unimplemented")
+}
+
 // Note: flushFlushableBatchIter.Next mirrors the implementation of
 // flushableBatchIter.Next due to performance. Keep the two in sync.
 func (i *flushFlushableBatchIter) Next() (*InternalKey, base.LazyValue) {