db: flushable ingested SSTables

batch.go

@@ -160,6 +160,7 @@ func (d DeferredBatchOp) Finish() error {
 //
 //   InternalKeyKindDelete         varstring
 //   InternalKeyKindLogData        varstring
+//   InternalKeyKindIngestSST      varstring
 //   InternalKeyKindSet            varstring varstring
 //   InternalKeyKindMerge          varstring varstring
 //   InternalKeyKindRangeDelete    varstring varstring
@@ -245,9 +246,10 @@ type Batch struct {
 	// memtable.
 	flushable *flushableBatch
 
-	commit    sync.WaitGroup
-	commitErr error
-	applied   uint32 // updated atomically
+	commit       sync.WaitGroup
+	commitErr    error
+	applied      uint32 // updated atomically
+	skipMemtable bool
 }
 
 var _ Reader = (*Batch)(nil)
@@ -328,13 +330,17 @@ 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:
+			// InternalKeyKindIngestSST does not contribute to the size of the
+			// memtable.
+			continue
 		}
+		b.memTableSize += memTableEntrySize(len(key), len(value))
 	}
 }
 
@@ -754,13 +760,33 @@ func (b *Batch) RangeKeyDeleteDeferred(startLen, endLen int) *DeferredBatchOp {
 //
 // It is safe to modify the contents of the argument after LogData returns.
 func (b *Batch) LogData(data []byte, _ *WriteOptions) error {
+	return b.writeRecordToWALOnly(data, InternalKeyKindLogData, false)
+}
+
+// IngestSSTs adds the list of sstable paths to the batch. The data will only be
+// written to the WAL (not added to memtables or sstables).
+func (b *Batch) IngestSSTs(data []byte, _ *WriteOptions) error {
+	return b.writeRecordToWALOnly(data, InternalKeyKindIngestSST, true)
+}
+
+// writeRecordToWALOnly writes a record to a batch that will not be applied to
+// the memtable. `updateCount` dictates whether or not we update batch.Count.
+func (b *Batch) writeRecordToWALOnly(data []byte, kind InternalKeyKind, updateCount bool) error {
 	origCount, origMemTableSize := b.count, b.memTableSize
-	b.prepareDeferredKeyRecord(len(data), InternalKeyKindLogData)
+	defer func() {
+		// Restore b.memTableSize (and b.count when updateCount is true) to the
+		// original value(s) since this batch is not written to the memtable. Note
+		// that Batch.count only refers to records that are added to the memtable.
+		b.memTableSize = origMemTableSize
+		if !updateCount {
+			b.count = origCount
+		}
+	}()
+
+	b.prepareDeferredKeyRecord(len(data), kind)
 	copy(b.deferredOp.Key, data)
-	// Since LogData only writes to the WAL and does not affect the memtable, we
-	// restore b.count and b.memTableSize to their origin values. Note that
-	// Batch.count only refers to records that are added to the memtable.
-	b.count, b.memTableSize = origCount, origMemTableSize
+
+	b.memTableSize = origMemTableSize
 	return nil
 }
 
@@ -1040,7 +1066,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 except for 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)
@@ -1486,6 +1514,10 @@ func (b *flushableBatch) readyForFlush() bool {
 	return true
 }
 
+func (b *flushableBatch) ref() {}
+
+func (b *flushableBatch) unref() {}
+
 // Note: flushableBatchIter mirrors the implementation of batchIter. Keep the
 // two in sync.
 type flushableBatchIter struct {

batch_test.go

@@ -76,6 +76,7 @@ func TestBatch(t *testing.T) {
 		{InternalKeyKindLogData, "", ""},
 		{InternalKeyKindRangeKeyDelete, "grass", "green"},
 		{InternalKeyKindRangeKeyDelete, "", ""},
+		{InternalKeyKindIngestSST, "1.sst,2.sst", ""},
 	}
 	var b Batch
 	for _, tc := range testCases {
@@ -94,6 +95,8 @@ func TestBatch(t *testing.T) {
 			_ = b.LogData([]byte(tc.key), nil)
 		case InternalKeyKindRangeKeyDelete:
 			_ = b.Experimental().RangeKeyDelete([]byte(tc.key), []byte(tc.value), nil)
+		case InternalKeyKindIngestSST:
+			_ = b.IngestSSTs([]byte(tc.key), nil)
 		}
 	}
 	verifyTestCases(&b, testCases)
@@ -137,6 +140,8 @@ func TestBatch(t *testing.T) {
 			copy(d.Key, key)
 			copy(d.Value, value)
 			d.Finish()
+		case InternalKeyKindIngestSST:
+			_ = b.IngestSSTs([]byte(tc.key), nil)
 		}
 	}
 	verifyTestCases(&b, testCases)

commit.go

@@ -280,7 +280,9 @@ func (p *commitPipeline) Commit(b *Batch, syncWAL bool) error {
 // sstable ingestion within the commit pipeline. The prepare callback is
 // invoked with commitPipeline.mu held, but note that DB.mu is not held and
 // must be locked if necessary.
-func (p *commitPipeline) AllocateSeqNum(count int, prepare func(), apply func(seqNum uint64)) {
+func (p *commitPipeline) AllocateSeqNum(
+	count int, prepare func(seqNum uint64), apply func(seqNum uint64),
+) {
 	// This method is similar to Commit and prepare. Be careful about trying to
 	// share additional code with those methods because Commit and prepare are
 	// performance critical code paths.
@@ -331,7 +333,7 @@ func (p *commitPipeline) AllocateSeqNum(count int, prepare func(), apply func(se
 	// Invoke the prepare callback. Note the lack of error reporting. Even if the
 	// callback internally fails, the sequence number needs to be published in
 	// order to allow the commit pipeline to proceed.
-	prepare()
+	prepare(seqNum)
 
 	p.mu.Unlock()
 

commit_test.go

@@ -131,7 +131,7 @@ func TestCommitPipelineAllocateSeqNum(t *testing.T) {
 	for i := 1; i <= n; i++ {
 		go func(i int) {
 			defer wg.Done()
-			p.AllocateSeqNum(i, func() {
+			p.AllocateSeqNum(i, func(seqNum uint64) {
 				atomic.AddUint64(&prepareCount, uint64(1))
 			}, func(seqNum uint64) {
 				atomic.AddUint64(&applyCount, uint64(1))