sstable: fix nil pointer with DisableValueBlocks set

Previously, the EstimatedSize method of the columnar sstable writer would panic
when the writer was configured with DisableValueBlocks enabled.

sstable/colblk_writer.go

@@ -203,11 +203,13 @@ func (w *RawColumnWriter) EstimatedSize() uint64 {
 	if w.rangeKeyBlock.KeyCount() > 0 {
 		sz += uint64(w.rangeKeyBlock.Size())
 	}
-	for _, blk := range w.valueBlock.blocks {
-		sz += uint64(blk.block.LengthWithTrailer())
-	}
-	if w.valueBlock.buf != nil {
-		sz += uint64(len(w.valueBlock.buf.b))
+	if w.valueBlock != nil {
+		for _, blk := range w.valueBlock.blocks {
+			sz += uint64(blk.block.LengthWithTrailer())
+		}
+		if w.valueBlock.buf != nil {
+			sz += uint64(len(w.valueBlock.buf.b))
+		}
 	}
 	// TODO(jackson): Include an estimate of the properties, filter and meta
 	// index blocks sizes.

sstable/data_test.go

@@ -103,39 +103,8 @@ func runBuildMemObjCmd(
 			_ = w.Close()
 		}
 	}()
-	for _, data := range strings.Split(td.Input, "\n") {
-		err := func() (err error) {
-			defer func() {
-				if r := recover(); r != nil {
-					err = errors.Errorf("%v", r)
-				}
-			}()
-
-			switch {
-			case strings.HasPrefix(data, "EncodeSpan:"):
-				return w.EncodeSpan(keyspan.ParseSpan(strings.TrimPrefix(data, "EncodeSpan:")))
-			default:
-				forceObsolete := strings.HasPrefix(data, "force-obsolete:")
-				if forceObsolete {
-					data = strings.TrimSpace(strings.TrimPrefix(data, "force-obsolete:"))
-				}
-				j := strings.Index(data, ":")
-				key := base.ParseInternalKey(data[:j])
-				value := []byte(data[j+1:])
-				switch key.Kind() {
-				case InternalKeyKindRangeDelete:
-					if forceObsolete {
-						return errors.Errorf("force-obsolete is not allowed for RANGEDEL")
-					}
-					return w.AddWithForceObsolete(key, value, false /* forceObsolete */)
-				default:
-					return w.AddWithForceObsolete(key, value, forceObsolete)
-				}
-			}
-		}()
-		if err != nil {
-			return nil, nil, err
-		}
+	if err := writeKVs(w, td.Input); err != nil {
+		return nil, nil, err
 	}
 	if err := w.Close(); err != nil {
 		return nil, nil, err
@@ -148,6 +117,41 @@ func runBuildMemObjCmd(
 	return meta, obj, nil
 }
 
+func writeKVs(w RawWriter, input string) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = errors.Errorf("%v", r)
+		}
+	}()
+	for _, data := range strings.Split(input, "\n") {
+		switch {
+		case strings.HasPrefix(data, "EncodeSpan:"):
+			err = w.EncodeSpan(keyspan.ParseSpan(strings.TrimPrefix(data, "EncodeSpan:")))
+		default:
+			forceObsolete := strings.HasPrefix(data, "force-obsolete:")
+			if forceObsolete {
+				data = strings.TrimSpace(strings.TrimPrefix(data, "force-obsolete:"))
+			}
+			j := strings.Index(data, ":")
+			key := base.ParseInternalKey(data[:j])
+			value := []byte(data[j+1:])
+			switch key.Kind() {
+			case InternalKeyKindRangeDelete:
+				if forceObsolete {
+					return errors.Errorf("force-obsolete is not allowed for RANGEDEL")
+				}
+				err = w.AddWithForceObsolete(key, value, false /* forceObsolete */)
+			default:
+				err = w.AddWithForceObsolete(key, value, forceObsolete)
+			}
+		}
+		if err != nil {
+			return err
+		}
+	}
+	return err
+}
+
 func runBuildCmd(
 	td *datadriven.TestData, writerOpts *WriterOptions, cacheSize int,
 ) (*WriterMetadata, *Reader, error) {

sstable/testdata/writer_v5

@@ -327,3 +327,18 @@ rocksdb.property.collectors: [obsolete-key]
 pebble.raw.range-key.key.size: 6
 pebble.raw.range-key.value.size: 9
 obsolete-key: hex:0074
+
+open-writer disable-value-blocks
+----
+
+write-kvs
+[email protected]:a2
+[email protected]:a1
+[email protected]:b2
+----
+EstimatedSize()=53
+
+close
+----
+point:    [a@2#1,SET-b@2#1,SET]
+seqnums:  [1-1]

sstable/writer_test.go

@@ -120,20 +120,24 @@ func formatWriterMetadata(td *datadriven.TestData, m *WriterMetadata) string {
 
 func runDataDriven(t *testing.T, file string, tableFormat TableFormat, parallelism bool) {
 	var r *Reader
+	var w RawWriter
 	var obj *objstorage.MemObj
-	defer func() {
+	closeReaderAndWriter := func() {
 		if r != nil {
 			require.NoError(t, r.Close())
 		}
-	}()
+		if w != nil {
+			require.NoError(t, w.Close())
+		}
+		r = nil
+		w = nil
+	}
+	defer closeReaderAndWriter()
 
 	datadriven.RunTest(t, file, func(t *testing.T, td *datadriven.TestData) string {
 		switch td.Cmd {
 		case "build":
-			if r != nil {
-				_ = r.Close()
-				r = nil
-			}
+			closeReaderAndWriter()
 			var meta *WriterMetadata
 			var err error
 			wopts := &WriterOptions{
@@ -152,6 +156,7 @@ func runDataDriven(t *testing.T, file string, tableFormat TableFormat, paralleli
 			return formatWriterMetadata(td, meta)
 
 		case "build-raw":
+			closeReaderAndWriter()
 			if r != nil {
 				_ = r.Close()
 				r = nil
@@ -167,6 +172,41 @@ func runDataDriven(t *testing.T, file string, tableFormat TableFormat, paralleli
 			}
 			return formatWriterMetadata(td, meta)
 
+		case "open-writer":
+			if r != nil {
+				_ = r.Close()
+				r = nil
+			}
+			wopts := &WriterOptions{
+				Comparer:           testkeys.Comparer,
+				DisableValueBlocks: td.HasArg("disable-value-blocks"),
+				TableFormat:        tableFormat,
+				Parallelism:        parallelism,
+			}
+			obj := &objstorage.MemObj{}
+			if err := optsFromArgs(td, wopts); err != nil {
+				return err.Error()
+			}
+			w = NewRawWriter(obj, *wopts)
+			return ""
+
+		case "write-kvs":
+			if err := writeKVs(w, td.Input); err != nil {
+				return err.Error()
+			}
+			return fmt.Sprintf("EstimatedSize()=%d", w.EstimatedSize())
+
+		case "close":
+			if err := w.Close(); err != nil {
+				return err.Error()
+			}
+			meta, err := w.Metadata()
+			if err != nil {
+				return err.Error()
+			}
+			w = nil
+			return formatWriterMetadata(td, meta)
+
 		case "scan":
 			iter, err := r.NewIter(NoTransforms, nil /* lower */, nil /* upper */)
 			if err != nil {