libroach,engine: support pagination of ExportToSst

c-deps/libroach/db.cc

@@ -1075,9 +1075,9 @@ DBStatus DBUnlockFile(DBFileLock lock) {
   return ToDBStatus(rocksdb::Env::Default()->UnlockFile((rocksdb::FileLock*)lock));
 }
 
-DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, DBIterOptions iter_opts,
-                       DBEngine* engine, DBString* data, DBString* write_intent,
-                       DBString* summary) {
+DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, uint64_t target_size,
+                       DBIterOptions iter_opts, DBEngine* engine, DBString* data,
+                       DBString* write_intent, DBString* summary, DBString* resume) {
   DBSstFileWriter* writer = DBSstFileWriterNew();
   DBStatus status = DBSstFileWriterOpen(writer);
   if (status.data != NULL) {
@@ -1092,14 +1092,23 @@ DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, DBIter
   bool skip_current_key_versions = !export_all_revisions;
   DBIterState state;
   const std::string end_key = EncodeKey(end);
-  for (state = iter.seek(start);; state = iter.next(skip_current_key_versions)) {
+  // cur_key is used when paginated is true and export_all_revisions is
+  // true. If we're exporting all revisions and we're returning a paginated
+  // SST then we need to keep track of when we've finished adding all of the
+  // versions of a key to the writer.
+  const bool paginated = target_size > 0;
+  std::string cur_key;
+  std::string resume_key;
+  // Seek to the MVCC metadata key for the provided start key and let the
+  // incremental iterator find the appropriate version.
+  const DBKey seek_key = { .key = start.key };
+  for (state = iter.seek(seek_key);; state = iter.next(skip_current_key_versions)) {
     if (state.status.data != NULL) {
       DBSstFileWriterClose(writer);
       return state.status;
     } else if (!state.valid || kComparator.Compare(iter.key(), end_key) >= 0) {
       break;
     }
-
     rocksdb::Slice decoded_key;
     int64_t wall_time = 0;
     int32_t logical_time = 0;
@@ -1109,13 +1118,32 @@ DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, DBIter
       return ToDBString("Unable to decode key");
     }
 
+    const bool is_new_key = !export_all_revisions || decoded_key.compare(cur_key) != 0;
+    if (paginated && export_all_revisions && is_new_key) {
+      // Reuse the underlying buffer in cur_key.
+      cur_key.clear();
+      cur_key.reserve(decoded_key.size());
+      cur_key.assign(decoded_key.data(), decoded_key.size());
+    }
+
     // Skip tombstone (len=0) records when start time is zero (non-incremental)
     // and we are not exporting all versions.
-    bool is_skipping_deletes = start.wall_time == 0 && start.logical == 0 && !export_all_revisions;
+    const bool is_skipping_deletes = start.wall_time == 0 && start.logical == 0 && !export_all_revisions;
     if (is_skipping_deletes && iter.value().size() == 0) {
       continue;
     }
 
+    // Check to see if this is the first version of key and adding it would
+    // put us over the limit (we might already be over the limit).
+    const int64_t cur_size = bulkop_summary.data_size();
+    const int64_t new_size = cur_size + decoded_key.size() + iter.value().size();
+    const bool is_over_target = cur_size > 0 && new_size > target_size;
+    if (paginated && is_new_key && is_over_target) {
+      resume_key.reserve(decoded_key.size());
+      resume_key.assign(decoded_key.data(), decoded_key.size());
+      break;
+    }
+
     // Insert key into sst and update statistics.
     status = DBSstFileWriterAddRaw(writer, iter.key(), iter.value());
     if (status.data != NULL) {
@@ -1126,8 +1154,7 @@ DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, DBIter
     if (!row_counter.Count((iter.key()), &bulkop_summary)) {
       return ToDBString("Error in row counter");
     }
-    bulkop_summary.set_data_size(bulkop_summary.data_size() + decoded_key.size() +
-                                 iter.value().size());
+    bulkop_summary.set_data_size(new_size);
   }
   *summary = ToDBString(bulkop_summary.SerializeAsString());
 
@@ -1139,6 +1166,11 @@ DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, DBIter
   auto res = DBSstFileWriterFinish(writer, data);
   DBSstFileWriterClose(writer);
 
+  // If we're not returning an error, check to see if we need to return the resume key.
+  if (res.data == NULL && resume_key.length() > 0) {
+    *resume = ToDBString(resume_key);
+  }
+
   return res;
 }
 

c-deps/libroach/include/libroach.h

@@ -549,14 +549,23 @@ typedef void* DBFileLock;
 // DBLockFile sets a lock on the specified file using RocksDB's file locking interface.
 DBStatus DBLockFile(DBSlice filename, DBFileLock* lock);
 
-// DBUnlockFile unlocks the file asscoiated with the specified lock and GCs any allocated memory for
+// DBUnlockFile unlocks the file associated with the specified lock and GCs any allocated memory for
 // the lock.
 DBStatus DBUnlockFile(DBFileLock lock);
 
 // DBExportToSst exports changes over the keyrange and time interval between the
 // start and end DBKeys to an SSTable using an IncrementalIterator.
-DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, DBIterOptions iter_opts,
-                       DBEngine* engine, DBString* data, DBString* write_intent, DBString* summary);
+// If target_size is positive, it indicates that the export should produce SSTs
+// which are roughly target size. Specifically, it will produce SSTs which contain
+// all relevant versions of a key and will not add the first version of a new
+// key if it would lead to the SST exceeding the target_size. If export_all_revisions
+// is false, the returned SST will be smaller than target_size so long as the first
+// kv pair is smaller than target_size. If export_all_revisions is true then
+// target_size may be exceeded. If the SST construction stops due to the target_size,
+// then resume will be set to the value of the resume key.
+DBStatus DBExportToSst(DBKey start, DBKey end, bool export_all_revisions, uint64_t target_size,
+                       DBIterOptions iter_opts, DBEngine* engine, DBString* data,
+                       DBString* write_intent, DBString* summary, DBString* resume);
 
 // DBEnvOpenReadableFile opens a DBReadableFile in the given engine.
 DBStatus DBEnvOpenReadableFile(DBEngine* db, DBSlice path, DBReadableFile* file);

pkg/ccl/storageccl/export.go

@@ -24,6 +24,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/spanset"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	crdberrors "github.com/cockroachdb/errors"
 	"github.com/pkg/errors"
 )
 
@@ -130,9 +131,14 @@ func evalExport(
 	}
 
 	e := spanset.GetDBEngine(batch, roachpb.Span{Key: args.Key, EndKey: args.EndKey})
-
-	data, summary, err := e.ExportToSst(args.Key, args.EndKey, args.StartTime, h.Timestamp, exportAllRevisions, io)
-
+	// TODO(ajwerner): Add a constant or internal cluster setting to control the
+	// target size for files and then paginate the actual export. The external
+	// API may need to be modified to deal with the case where ReturnSST is true.
+	const targetSize = 0 // unlimited
+	data, summary, resume, err := e.ExportToSst(args.Key, args.EndKey, args.StartTime, h.Timestamp, exportAllRevisions, targetSize, io)
+	if resume != nil {
+		return result.Result{}, crdberrors.AssertionFailedf("expected nil resume key with unlimited target size")
+	}
 	if err != nil {
 		return result.Result{}, err
 	}

pkg/ccl/storageccl/export_test.go

@@ -31,6 +31,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+	"github.com/stretchr/testify/require"
 )
 
 func TestExportCmd(t *testing.T) {
@@ -302,6 +303,7 @@ func assertEqualKVs(
 	startTime, endTime hlc.Timestamp,
 	exportAllRevisions bool,
 	enableTimeBoundIteratorOptimization bool,
+	targetSize uint64,
 ) func(*testing.T) {
 	return func(t *testing.T) {
 		t.Helper()
@@ -328,17 +330,19 @@ func assertEqualKVs(
 			io.MaxTimestampHint = endTime
 			io.MinTimestampHint = startTime.Next()
 		}
-		sst, _, err := e.ExportToSst(startKey, endKey, startTime, endTime, exportAllRevisions, io)
-		if err != nil {
-			t.Fatal(err)
+		var kvs []engine.MVCCKeyValue
+		for start := startKey; start != nil; {
+			var sst []byte
+			sst, _, start, err = e.ExportToSst(start, endKey, startTime, endTime, exportAllRevisions, targetSize, io)
+			require.NoError(t, err)
+			loaded := loadSST(t, sst, startKey, endKey)
+			kvs = append(kvs, loaded...)
 		}
 
 		// Compare new C++ implementation against the oracle.
 		expectedKVS := loadSST(t, expected, startKey, endKey)
-		kvs := loadSST(t, sst, startKey, endKey)
-
 		if len(kvs) != len(expectedKVS) {
-			t.Fatalf("got %d kvs (%+v) but expected %d (%+v)", len(kvs), kvs, len(expected), expected)
+			t.Fatalf("got %d kvs but expected %d:\n%v\n%v", len(kvs), len(expectedKVS), kvs, expectedKVS)
 		}
 
 		for i := range kvs {
@@ -430,27 +434,40 @@ func TestRandomKeyAndTimestampExport(t *testing.T) {
 				timestamps[i].Logical < timestamps[j].Logical)
 	})
 
+	testWithTargetSize := func(t *testing.T, targetSize uint64) {
+		if testing.Short() && targetSize > 0 && targetSize < 1<<15 {
+			t.Skipf("testing with size %d is slow", targetSize)
+		}
+		t.Run("ts (0-∞], latest, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, false, false, targetSize))
+		t.Run("ts (0-∞], all, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, true, false, targetSize))
+		t.Run("ts (0-∞], latest, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, false, true, targetSize))
+		t.Run("ts (0-∞], all, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, true, true, targetSize))
+
+		upperBound := randutil.RandIntInRange(rnd, 1, numKeys)
+		lowerBound := rnd.Intn(upperBound)
+
+		// Exercise random key ranges.
+		t.Run("kv [randLower, randUpper), latest, nontimebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, false, false, targetSize))
+		t.Run("kv [randLower, randUpper), all, nontimebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, true, false, targetSize))
+		t.Run("kv [randLower, randUpper), latest, timebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, false, true, targetSize))
+		t.Run("kv [randLower, randUpper), all, timebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, true, true, targetSize))
+
+		upperBound = randutil.RandIntInRange(rnd, 1, numKeys)
+		lowerBound = rnd.Intn(upperBound)
+
+		// Exercise random timestamps.
+		t.Run("kv (randLowerTime, randUpperTime], latest, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], false, false, targetSize))
+		t.Run("kv (randLowerTime, randUpperTime], all, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], true, false, targetSize))
+		t.Run("kv (randLowerTime, randUpperTime], latest, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], false, true, targetSize))
+		t.Run("kv (randLowerTime, randUpperTime], all, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], true, true, targetSize))
+	}
 	// Exercise min to max time and key ranges.
-	t.Run("ts (0-∞], latest, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, false, false))
-	t.Run("ts (0-∞], all, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, true, false))
-	t.Run("ts (0-∞], latest, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, false, true))
-	t.Run("ts (0-∞], all, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, tsMin, tsMax, true, true))
-
-	upperBound := randutil.RandIntInRange(rnd, 1, numKeys)
-	lowerBound := rnd.Intn(upperBound)
-
-	// Exercise random key ranges.
-	t.Run("kv [randLower, randUpper), latest, nontimebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, false, false))
-	t.Run("kv [randLower, randUpper), all, nontimebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, true, false))
-	t.Run("kv [randLower, randUpper), latest, timebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, false, true))
-	t.Run("kv [randLower, randUpper), all, timebound", assertEqualKVs(ctx, e, keys[lowerBound], keys[upperBound], tsMin, tsMax, true, true))
-
-	upperBound = randutil.RandIntInRange(rnd, 1, numKeys)
-	lowerBound = rnd.Intn(upperBound)
-
-	// Exercise random timestamps.
-	t.Run("kv (randLowerTime, randUpperTime], latest, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], false, false))
-	t.Run("kv (randLowerTime, randUpperTime], all, nontimebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], true, false))
-	t.Run("kv (randLowerTime, randUpperTime], latest, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], false, true))
-	t.Run("kv (randLowerTime, randUpperTime], all, timebound", assertEqualKVs(ctx, e, keyMin, keyMax, timestamps[lowerBound], timestamps[upperBound], true, true))
+	for _, targetSize := range []uint64{
+		0 /* unlimited */, 1 << 10, 1 << 16, 1 << 20,
+	} {
+		t.Run(fmt.Sprintf("targetSize=%d", targetSize), func(t *testing.T) {
+			testWithTargetSize(t, targetSize)
+		})
+	}
+
 }

pkg/storage/engine/engine.go

@@ -200,9 +200,19 @@ type Reader interface {
 	// within the interval is exported. Deletions are included if all revisions are
 	// requested or if the start.Timestamp is non-zero. Returns the bytes of an
 	// SSTable containing the exported keys, the size of exported data, or an error.
+	// If targetSize is positive, it indicates that the export should produce SSTs
+	// which are roughly target size. Specifically, it will produce SSTs which contain
+	// all relevant versions of a key and will not add the first version of a new
+	// key if it would lead to the SST exceeding the targetSize. If exportAllRevisions
+	// is false, the returned SST will be smaller than target_size so long as the first
+	// kv pair is smaller than targetSize. If exportAllRevisions is true then
+	// targetSize may be exceeded by as much as the size of all of the versions of
+	// the last key. If the SST construction stops due to the targetSize,
+	// then a non-nil resumeKey will be returned.
 	ExportToSst(
-		startKey, endKey roachpb.Key, startTS, endTS hlc.Timestamp, exportAllRevisions bool, io IterOptions,
-	) ([]byte, roachpb.BulkOpSummary, error)
+		startKey, endKey roachpb.Key, startTS, endTS hlc.Timestamp,
+		exportAllRevisions bool, targetSize uint64, io IterOptions,
+	) (sst []byte, _ roachpb.BulkOpSummary, resumeKey roachpb.Key, _ error)
 	// Get returns the value for the given key, nil otherwise.
 	//
 	// Deprecated: use MVCCGet instead.

pkg/storage/engine/pebble.go

@@ -492,9 +492,10 @@ func (p *Pebble) ExportToSst(
 	startKey, endKey roachpb.Key,
 	startTS, endTS hlc.Timestamp,
 	exportAllRevisions bool,
+	targetSize uint64,
 	io IterOptions,
-) ([]byte, roachpb.BulkOpSummary, error) {
-	return pebbleExportToSst(p, startKey, endKey, startTS, endTS, exportAllRevisions, io)
+) ([]byte, roachpb.BulkOpSummary, roachpb.Key, error) {
+	return pebbleExportToSst(p, startKey, endKey, startTS, endTS, exportAllRevisions, targetSize, io)
 }
 
 // Get implements the Engine interface.
@@ -938,9 +939,10 @@ func (p *pebbleReadOnly) ExportToSst(
 	startKey, endKey roachpb.Key,
 	startTS, endTS hlc.Timestamp,
 	exportAllRevisions bool,
+	targetSize uint64,
 	io IterOptions,
-) ([]byte, roachpb.BulkOpSummary, error) {
-	return pebbleExportToSst(p, startKey, endKey, startTS, endTS, exportAllRevisions, io)
+) ([]byte, roachpb.BulkOpSummary, roachpb.Key, error) {
+	return pebbleExportToSst(p, startKey, endKey, startTS, endTS, exportAllRevisions, targetSize, io)
 }
 
 func (p *pebbleReadOnly) Get(key MVCCKey) ([]byte, error) {
@@ -1059,9 +1061,10 @@ func (p *pebbleSnapshot) ExportToSst(
 	startKey, endKey roachpb.Key,
 	startTS, endTS hlc.Timestamp,
 	exportAllRevisions bool,
+	targetSize uint64,
 	io IterOptions,
-) ([]byte, roachpb.BulkOpSummary, error) {
-	return pebbleExportToSst(p, startKey, endKey, startTS, endTS, exportAllRevisions, io)
+) ([]byte, roachpb.BulkOpSummary, roachpb.Key, error) {
+	return pebbleExportToSst(p, startKey, endKey, startTS, endTS, exportAllRevisions, targetSize, io)
 }
 
 // Get implements the Reader interface.
@@ -1113,8 +1116,9 @@ func pebbleExportToSst(
 	startKey, endKey roachpb.Key,
 	startTS, endTS hlc.Timestamp,
 	exportAllRevisions bool,
+	targetSize uint64,
 	io IterOptions,
-) ([]byte, roachpb.BulkOpSummary, error) {
+) ([]byte, roachpb.BulkOpSummary, roachpb.Key, error) {
 	sstFile := &MemFile{}
 	sstWriter := MakeBackupSSTWriter(sstFile)
 	defer sstWriter.Close()
@@ -1128,12 +1132,15 @@ func pebbleExportToSst(
 			EndTime:     endTS,
 		})
 	defer iter.Close()
+	var curKey roachpb.Key // only used if exportAllRevisions
+	var resumeKey roachpb.Key
+	paginated := targetSize > 0
 	for iter.SeekGE(MakeMVCCMetadataKey(startKey)); ; {
 		ok, err := iter.Valid()
 		if err != nil {
 			// The error may be a WriteIntentError. In which case, returning it will
 			// cause this command to be retried.
-			return nil, roachpb.BulkOpSummary{}, err
+			return nil, roachpb.BulkOpSummary{}, nil, err
 		}
 		if !ok {
 			break
@@ -1143,18 +1150,30 @@ func pebbleExportToSst(
 			break
 		}
 		unsafeValue := iter.UnsafeValue()
+		isNewKey := !exportAllRevisions || !unsafeKey.Key.Equal(curKey)
+		if paginated && exportAllRevisions && isNewKey {
+			curKey = append(curKey[:0], unsafeKey.Key...)
+		}
 
 		// Skip tombstone (len=0) records when start time is zero (non-incremental)
 		// and we are not exporting all versions.
 		skipTombstones := !exportAllRevisions && startTS.IsEmpty()
 		if len(unsafeValue) > 0 || !skipTombstones {
 			if err := rows.Count(unsafeKey.Key); err != nil {
-				return nil, roachpb.BulkOpSummary{}, errors.Wrapf(err, "decoding %s", unsafeKey)
+				return nil, roachpb.BulkOpSummary{}, nil, errors.Wrapf(err, "decoding %s", unsafeKey)
+			}
+			curSize := rows.BulkOpSummary.DataSize
+			newSize := curSize + int64(len(unsafeKey.Key)+len(unsafeValue))
+			isOverTarget := paginated && curSize > 0 && uint64(newSize) > targetSize
+			if isNewKey && isOverTarget {
+				// Allocate the right size for resumeKey rather than using curKey.
+				resumeKey = append(make(roachpb.Key, 0, len(unsafeKey.Key)), unsafeKey.Key...)
+				break
 			}
-			rows.BulkOpSummary.DataSize += int64(len(unsafeKey.Key) + len(unsafeValue))
 			if err := sstWriter.Put(unsafeKey, unsafeValue); err != nil {
-				return nil, roachpb.BulkOpSummary{}, errors.Wrapf(err, "adding key %s", unsafeKey)
+				return nil, roachpb.BulkOpSummary{}, nil, errors.Wrapf(err, "adding key %s", unsafeKey)
 			}
+			rows.BulkOpSummary.DataSize = newSize
 		}
 
 		if exportAllRevisions {
@@ -1167,12 +1186,12 @@ func pebbleExportToSst(
 	if rows.BulkOpSummary.DataSize == 0 {
 		// If no records were added to the sstable, skip completing it and return a
 		// nil slice – the export code will discard it anyway (based on 0 DataSize).
-		return nil, roachpb.BulkOpSummary{}, nil
+		return nil, roachpb.BulkOpSummary{}, nil, nil
 	}
 
 	if err := sstWriter.Finish(); err != nil {
-		return nil, roachpb.BulkOpSummary{}, err
+		return nil, roachpb.BulkOpSummary{}, nil, err
 	}
 
-	return sstFile.Data(), rows.BulkOpSummary, nil
+	return sstFile.Data(), rows.BulkOpSummary, resumeKey, nil
 }