github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/data_test.go (about)

     1  // Copyright 2018 The LevelDB-Go and Pebble and Bitalostored 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 bitalostable
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"io"
    11  	"math"
    12  	"math/rand"
    13  	"strconv"
    14  	"strings"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/errors"
    18  	"github.com/stretchr/testify/require"
    19  	"github.com/zuoyebang/bitalostable/internal/base"
    20  	"github.com/zuoyebang/bitalostable/internal/datadriven"
    21  	"github.com/zuoyebang/bitalostable/internal/keyspan"
    22  	"github.com/zuoyebang/bitalostable/internal/rangedel"
    23  	"github.com/zuoyebang/bitalostable/internal/rangekey"
    24  	"github.com/zuoyebang/bitalostable/internal/testkeys"
    25  	"github.com/zuoyebang/bitalostable/internal/testkeys/blockprop"
    26  	"github.com/zuoyebang/bitalostable/sstable"
    27  	"github.com/zuoyebang/bitalostable/vfs"
    28  )
    29  
    30  type iterCmdOpts struct {
    31  	verboseKey bool
    32  	stats      *base.InternalIteratorStats
    33  }
    34  
    35  type iterCmdOpt func(*iterCmdOpts)
    36  
    37  func iterCmdVerboseKey(opts *iterCmdOpts) { opts.verboseKey = true }
    38  
    39  func iterCmdStats(stats *base.InternalIteratorStats) iterCmdOpt {
    40  	return func(opts *iterCmdOpts) {
    41  		opts.stats = stats
    42  	}
    43  }
    44  
    45  func runGetCmd(td *datadriven.TestData, d *DB) string {
    46  	snap := Snapshot{
    47  		db:     d,
    48  		seqNum: InternalKeySeqNumMax,
    49  	}
    50  
    51  	for _, arg := range td.CmdArgs {
    52  		if len(arg.Vals) != 1 {
    53  			return fmt.Sprintf("%s: %s=<value>", td.Cmd, arg.Key)
    54  		}
    55  		switch arg.Key {
    56  		case "seq":
    57  			var err error
    58  			snap.seqNum, err = strconv.ParseUint(arg.Vals[0], 10, 64)
    59  			if err != nil {
    60  				return err.Error()
    61  			}
    62  		default:
    63  			return fmt.Sprintf("%s: unknown arg: %s", td.Cmd, arg.Key)
    64  		}
    65  	}
    66  
    67  	var buf bytes.Buffer
    68  	for _, data := range strings.Split(td.Input, "\n") {
    69  		v, closer, err := snap.Get([]byte(data))
    70  		if err != nil {
    71  			fmt.Fprintf(&buf, "%s: %s\n", data, err)
    72  		} else {
    73  			fmt.Fprintf(&buf, "%s:%s\n", data, v)
    74  			closer.Close()
    75  		}
    76  	}
    77  	return buf.String()
    78  }
    79  
    80  func runIterCmd(d *datadriven.TestData, iter *Iterator, closeIter bool) string {
    81  	if closeIter {
    82  		defer func() {
    83  			if iter != nil {
    84  				iter.Close()
    85  			}
    86  		}()
    87  	}
    88  	var b bytes.Buffer
    89  	for _, line := range strings.Split(d.Input, "\n") {
    90  		parts := strings.Fields(line)
    91  		if len(parts) == 0 {
    92  			continue
    93  		}
    94  		printValidityState := false
    95  		var valid bool
    96  		var validityState IterValidityState
    97  		switch parts[0] {
    98  		case "seek-ge":
    99  			if len(parts) != 2 {
   100  				return "seek-ge <key>\n"
   101  			}
   102  			valid = iter.SeekGE([]byte(parts[1]))
   103  		case "seek-prefix-ge":
   104  			if len(parts) != 2 {
   105  				return "seek-prefix-ge <key>\n"
   106  			}
   107  			valid = iter.SeekPrefixGE([]byte(parts[1]))
   108  		case "seek-lt":
   109  			if len(parts) != 2 {
   110  				return "seek-lt <key>\n"
   111  			}
   112  			valid = iter.SeekLT([]byte(parts[1]))
   113  		case "seek-ge-limit":
   114  			if len(parts) != 3 {
   115  				return "seek-ge-limit <key> <limit>\n"
   116  			}
   117  			validityState = iter.SeekGEWithLimit(
   118  				[]byte(parts[1]), []byte(parts[2]))
   119  			printValidityState = true
   120  		case "seek-lt-limit":
   121  			if len(parts) != 3 {
   122  				return "seek-lt-limit <key> <limit>\n"
   123  			}
   124  			validityState = iter.SeekLTWithLimit(
   125  				[]byte(parts[1]), []byte(parts[2]))
   126  			printValidityState = true
   127  		case "next-limit":
   128  			if len(parts) != 2 {
   129  				return "next-limit <limit>\n"
   130  			}
   131  			validityState = iter.NextWithLimit([]byte(parts[1]))
   132  			printValidityState = true
   133  		case "prev-limit":
   134  			if len(parts) != 2 {
   135  				return "prev-limit <limit>\n"
   136  			}
   137  			validityState = iter.PrevWithLimit([]byte(parts[1]))
   138  			printValidityState = true
   139  		case "first":
   140  			valid = iter.First()
   141  		case "last":
   142  			valid = iter.Last()
   143  		case "next":
   144  			valid = iter.Next()
   145  		case "prev":
   146  			valid = iter.Prev()
   147  		case "set-bounds":
   148  			if len(parts) <= 1 || len(parts) > 3 {
   149  				return "set-bounds lower=<lower> upper=<upper>\n"
   150  			}
   151  			var lower []byte
   152  			var upper []byte
   153  			for _, part := range parts[1:] {
   154  				arg := strings.Split(part, "=")
   155  				switch arg[0] {
   156  				case "lower":
   157  					lower = []byte(arg[1])
   158  				case "upper":
   159  					upper = []byte(arg[1])
   160  				default:
   161  					return fmt.Sprintf("set-bounds: unknown arg: %s", arg)
   162  				}
   163  			}
   164  			iter.SetBounds(lower, upper)
   165  			valid = iter.Valid()
   166  		case "set-options":
   167  			opts := iter.opts
   168  			if _, err := parseIterOptions(&opts, &iter.opts, parts[1:]); err != nil {
   169  				return fmt.Sprintf("set-options: %s", err.Error())
   170  			}
   171  			iter.SetOptions(&opts)
   172  			valid = iter.Valid()
   173  		case "stats":
   174  			stats := iter.Stats()
   175  			fmt.Fprintf(&b, "stats: %s\n", stats.String())
   176  			continue
   177  		case "clone":
   178  			var opts CloneOptions
   179  			if len(parts) > 1 {
   180  				var iterOpts IterOptions
   181  				if foundAny, err := parseIterOptions(&iterOpts, &iter.opts, parts[1:]); err != nil {
   182  					return fmt.Sprintf("clone: %s", err.Error())
   183  				} else if foundAny {
   184  					opts.IterOptions = &iterOpts
   185  				}
   186  				for _, part := range parts[1:] {
   187  					if arg := strings.Split(part, "="); len(arg) == 2 && arg[0] == "refresh-batch" {
   188  						var err error
   189  						opts.RefreshBatchView, err = strconv.ParseBool(arg[1])
   190  						if err != nil {
   191  							return fmt.Sprintf("clone: refresh-batch: %s", err.Error())
   192  						}
   193  					}
   194  				}
   195  			}
   196  			clonedIter, err := iter.Clone(opts)
   197  			if err != nil {
   198  				fmt.Fprintf(&b, "error in clone, skipping rest of input: err=%v\n", err)
   199  				return b.String()
   200  			}
   201  			if err = iter.Close(); err != nil {
   202  				fmt.Fprintf(&b, "err=%v\n", err)
   203  			}
   204  			iter = clonedIter
   205  		case "is-using-combined":
   206  			if iter.opts.KeyTypes != IterKeyTypePointsAndRanges {
   207  				fmt.Fprintln(&b, "not configured for combined iteration")
   208  			} else if iter.lazyCombinedIter.combinedIterState.initialized {
   209  				fmt.Fprintln(&b, "using combined (non-lazy) iterator")
   210  			} else {
   211  				fmt.Fprintln(&b, "using lazy iterator")
   212  			}
   213  			continue
   214  		default:
   215  			return fmt.Sprintf("unknown op: %s", parts[0])
   216  		}
   217  
   218  		valid = valid || validityState == IterValid
   219  		if valid != iter.Valid() {
   220  			fmt.Fprintf(&b, "mismatched valid states: %t vs %t\n", valid, iter.Valid())
   221  		}
   222  		hasPoint, hasRange := iter.HasPointAndRange()
   223  		hasEither := hasPoint || hasRange
   224  		if hasEither != valid {
   225  			fmt.Fprintf(&b, "mismatched valid/HasPointAndRange states: valid=%t HasPointAndRange=(%t,%t)\n", valid, hasPoint, hasRange)
   226  		}
   227  
   228  		if valid {
   229  			validityState = IterValid
   230  		}
   231  		printIterState(&b, iter, validityState, printValidityState)
   232  	}
   233  	return b.String()
   234  }
   235  
   236  func parseIterOptions(
   237  	opts *IterOptions, ref *IterOptions, parts []string,
   238  ) (foundAny bool, err error) {
   239  	const usageString = "[lower=<lower>] [upper=<upper>] [key-types=point|range|both] [mask-suffix=<suffix>] [mask-filter=<bool>] [only-durable=<bool>] [table-filter=reuse|none] [point-filters=reuse|none]\n"
   240  	for _, part := range parts {
   241  		arg := strings.SplitN(part, "=", 2)
   242  		if len(arg) != 2 {
   243  			return false, errors.Newf(usageString)
   244  		}
   245  		switch arg[0] {
   246  		case "point-filters":
   247  			switch arg[1] {
   248  			case "reuse":
   249  				opts.PointKeyFilters = ref.PointKeyFilters
   250  			case "none":
   251  				opts.PointKeyFilters = nil
   252  			default:
   253  				return false, errors.Newf("unknown arg point-filter=%q:\n%s", arg[1], usageString)
   254  			}
   255  		case "lower":
   256  			opts.LowerBound = []byte(arg[1])
   257  		case "upper":
   258  			opts.UpperBound = []byte(arg[1])
   259  		case "key-types":
   260  			switch arg[1] {
   261  			case "point":
   262  				opts.KeyTypes = IterKeyTypePointsOnly
   263  			case "range":
   264  				opts.KeyTypes = IterKeyTypeRangesOnly
   265  			case "both":
   266  				opts.KeyTypes = IterKeyTypePointsAndRanges
   267  			default:
   268  				return false, errors.Newf("unknown key-type %q:\n%s", arg[1], usageString)
   269  			}
   270  		case "mask-suffix":
   271  			opts.RangeKeyMasking.Suffix = []byte(arg[1])
   272  		case "mask-filter":
   273  			opts.RangeKeyMasking.Filter = func() BlockPropertyFilterMask {
   274  				return blockprop.NewMaskingFilter()
   275  			}
   276  		case "table-filter":
   277  			switch arg[1] {
   278  			case "reuse":
   279  				opts.TableFilter = ref.TableFilter
   280  			case "none":
   281  				opts.TableFilter = nil
   282  			default:
   283  				return false, errors.Newf("unknown arg table-filter=%q:\n%s", arg[1], usageString)
   284  			}
   285  		case "only-durable":
   286  			var err error
   287  			opts.OnlyReadGuaranteedDurable, err = strconv.ParseBool(arg[1])
   288  			if err != nil {
   289  				return false, errors.Newf("cannot parse only-durable=%q: %s", arg[1], err)
   290  			}
   291  		default:
   292  			continue
   293  		}
   294  		foundAny = true
   295  	}
   296  	return foundAny, nil
   297  }
   298  
   299  func printIterState(
   300  	b io.Writer, iter *Iterator, validity IterValidityState, printValidityState bool,
   301  ) {
   302  	var validityStateStr string
   303  	if printValidityState {
   304  		switch validity {
   305  		case IterExhausted:
   306  			validityStateStr = " exhausted"
   307  		case IterValid:
   308  			validityStateStr = " valid"
   309  		case IterAtLimit:
   310  			validityStateStr = " at-limit"
   311  		}
   312  	}
   313  	if err := iter.Error(); err != nil {
   314  		fmt.Fprintf(b, "err=%v\n", err)
   315  	} else if validity == IterValid {
   316  		switch {
   317  		case iter.opts.rangeKeys() && iter.opts.pointKeys():
   318  			hasPoint, hasRange := iter.HasPointAndRange()
   319  			fmt.Fprintf(b, "%s:%s (", iter.Key(), validityStateStr)
   320  			if hasPoint {
   321  				fmt.Fprintf(b, "%s, ", iter.Value())
   322  			} else {
   323  				fmt.Fprint(b, "., ")
   324  			}
   325  			if hasRange {
   326  				start, end := iter.RangeBounds()
   327  				fmt.Fprintf(b, "[%s-%s)", formatASCIIKey(start), formatASCIIKey(end))
   328  				writeRangeKeys(b, iter)
   329  			} else {
   330  				fmt.Fprint(b, ".")
   331  			}
   332  			if iter.RangeKeyChanged() {
   333  				fmt.Fprint(b, " UPDATED")
   334  			}
   335  			fmt.Fprint(b, ")")
   336  		case iter.opts.rangeKeys():
   337  			if iter.Valid() {
   338  				hasPoint, hasRange := iter.HasPointAndRange()
   339  				if hasPoint || !hasRange {
   340  					panic(fmt.Sprintf("bitalostable: unexpected HasPointAndRange (%t, %t)", hasPoint, hasRange))
   341  				}
   342  				start, end := iter.RangeBounds()
   343  				fmt.Fprintf(b, "%s [%s-%s)", iter.Key(), formatASCIIKey(start), formatASCIIKey(end))
   344  				writeRangeKeys(b, iter)
   345  			} else {
   346  				fmt.Fprint(b, ".")
   347  			}
   348  			if iter.RangeKeyChanged() {
   349  				fmt.Fprint(b, " UPDATED")
   350  			}
   351  		default:
   352  			fmt.Fprintf(b, "%s:%s%s", iter.Key(), iter.Value(), validityStateStr)
   353  		}
   354  		fmt.Fprintln(b)
   355  	} else {
   356  		fmt.Fprintf(b, ".%s\n", validityStateStr)
   357  	}
   358  }
   359  
   360  func formatASCIIKey(b []byte) string {
   361  	if bytes.IndexFunc(b, func(r rune) bool { return r < 'A' || r > 'z' }) != -1 {
   362  		// This key is not just ASCII letters. Quote it.
   363  		return fmt.Sprintf("%q", b)
   364  	}
   365  	return string(b)
   366  }
   367  
   368  func writeRangeKeys(b io.Writer, iter *Iterator) {
   369  	rangeKeys := iter.RangeKeys()
   370  	for j := 0; j < len(rangeKeys); j++ {
   371  		if j > 0 {
   372  			fmt.Fprint(b, ",")
   373  		}
   374  		fmt.Fprintf(b, " %s=%s", rangeKeys[j].Suffix, rangeKeys[j].Value)
   375  	}
   376  }
   377  
   378  func runInternalIterCmd(d *datadriven.TestData, iter internalIterator, opts ...iterCmdOpt) string {
   379  	var o iterCmdOpts
   380  	for _, opt := range opts {
   381  		opt(&o)
   382  	}
   383  
   384  	var b bytes.Buffer
   385  	var prefix []byte
   386  	for _, line := range strings.Split(d.Input, "\n") {
   387  		parts := strings.Fields(line)
   388  		if len(parts) == 0 {
   389  			continue
   390  		}
   391  		var key *InternalKey
   392  		var value []byte
   393  		switch parts[0] {
   394  		case "seek-ge":
   395  			if len(parts) < 2 || len(parts) > 3 {
   396  				return "seek-ge <key> [<try-seek-using-next>]\n"
   397  			}
   398  			prefix = nil
   399  			var flags base.SeekGEFlags
   400  			if len(parts) == 3 {
   401  				if trySeekUsingNext, err := strconv.ParseBool(parts[2]); err != nil {
   402  					return err.Error()
   403  				} else if trySeekUsingNext {
   404  					flags = flags.EnableTrySeekUsingNext()
   405  				}
   406  			}
   407  			key, value = iter.SeekGE([]byte(strings.TrimSpace(parts[1])), flags)
   408  		case "seek-prefix-ge":
   409  			if len(parts) != 2 && len(parts) != 3 {
   410  				return "seek-prefix-ge <key> [<try-seek-using-next>]\n"
   411  			}
   412  			prefix = []byte(strings.TrimSpace(parts[1]))
   413  			var flags base.SeekGEFlags
   414  			if len(parts) == 3 {
   415  				if trySeekUsingNext, err := strconv.ParseBool(parts[2]); err != nil {
   416  					return err.Error()
   417  				} else if trySeekUsingNext {
   418  					flags = flags.EnableTrySeekUsingNext()
   419  				}
   420  			}
   421  			key, value = iter.SeekPrefixGE(prefix, prefix /* key */, flags)
   422  		case "seek-lt":
   423  			if len(parts) != 2 {
   424  				return "seek-lt <key>\n"
   425  			}
   426  			prefix = nil
   427  			key, value = iter.SeekLT([]byte(strings.TrimSpace(parts[1])), base.SeekLTFlagsNone)
   428  		case "first":
   429  			prefix = nil
   430  			key, value = iter.First()
   431  		case "last":
   432  			prefix = nil
   433  			key, value = iter.Last()
   434  		case "next":
   435  			key, value = iter.Next()
   436  		case "prev":
   437  			key, value = iter.Prev()
   438  		case "set-bounds":
   439  			if len(parts) <= 1 || len(parts) > 3 {
   440  				return "set-bounds lower=<lower> upper=<upper>\n"
   441  			}
   442  			var lower []byte
   443  			var upper []byte
   444  			for _, part := range parts[1:] {
   445  				arg := strings.Split(strings.TrimSpace(part), "=")
   446  				switch arg[0] {
   447  				case "lower":
   448  					lower = []byte(arg[1])
   449  				case "upper":
   450  					upper = []byte(arg[1])
   451  				default:
   452  					return fmt.Sprintf("set-bounds: unknown arg: %s", arg)
   453  				}
   454  			}
   455  			iter.SetBounds(lower, upper)
   456  			continue
   457  		case "stats":
   458  			if o.stats != nil {
   459  				fmt.Fprintf(&b, "%+v\n", *o.stats)
   460  			}
   461  			continue
   462  		case "reset-stats":
   463  			if o.stats != nil {
   464  				*o.stats = base.InternalIteratorStats{}
   465  			}
   466  			continue
   467  		default:
   468  			return fmt.Sprintf("unknown op: %s", parts[0])
   469  		}
   470  		if key != nil {
   471  			if o.verboseKey {
   472  				fmt.Fprintf(&b, "%s:%s\n", key, value)
   473  			} else {
   474  				fmt.Fprintf(&b, "%s:%s\n", key.UserKey, value)
   475  			}
   476  		} else if err := iter.Error(); err != nil {
   477  			fmt.Fprintf(&b, "err=%v\n", err)
   478  		} else {
   479  			fmt.Fprintf(&b, ".\n")
   480  		}
   481  	}
   482  	return b.String()
   483  }
   484  
   485  func runBatchDefineCmd(d *datadriven.TestData, b *Batch) error {
   486  	for _, line := range strings.Split(d.Input, "\n") {
   487  		parts := strings.Fields(line)
   488  		if len(parts) == 0 {
   489  			continue
   490  		}
   491  		if parts[1] == `<nil>` {
   492  			parts[1] = ""
   493  		}
   494  		var err error
   495  		switch parts[0] {
   496  		case "set":
   497  			if len(parts) != 3 {
   498  				return errors.Errorf("%s expects 2 arguments", parts[0])
   499  			}
   500  			err = b.Set([]byte(parts[1]), []byte(parts[2]), nil)
   501  		case "del":
   502  			if len(parts) != 2 {
   503  				return errors.Errorf("%s expects 1 argument", parts[0])
   504  			}
   505  			err = b.Delete([]byte(parts[1]), nil)
   506  		case "singledel":
   507  			if len(parts) != 2 {
   508  				return errors.Errorf("%s expects 1 argument", parts[0])
   509  			}
   510  			err = b.SingleDelete([]byte(parts[1]), nil)
   511  		case "del-range":
   512  			if len(parts) != 3 {
   513  				return errors.Errorf("%s expects 2 arguments", parts[0])
   514  			}
   515  			err = b.DeleteRange([]byte(parts[1]), []byte(parts[2]), nil)
   516  		case "merge":
   517  			if len(parts) != 3 {
   518  				return errors.Errorf("%s expects 2 arguments", parts[0])
   519  			}
   520  			err = b.Merge([]byte(parts[1]), []byte(parts[2]), nil)
   521  		case "range-key-set":
   522  			if len(parts) != 5 {
   523  				return errors.Errorf("%s expects 4 arguments", parts[0])
   524  			}
   525  			err = b.RangeKeySet(
   526  				[]byte(parts[1]),
   527  				[]byte(parts[2]),
   528  				[]byte(parts[3]),
   529  				[]byte(parts[4]),
   530  				nil)
   531  		case "range-key-unset":
   532  			if len(parts) != 4 {
   533  				return errors.Errorf("%s expects 3 arguments", parts[0])
   534  			}
   535  			err = b.RangeKeyUnset(
   536  				[]byte(parts[1]),
   537  				[]byte(parts[2]),
   538  				[]byte(parts[3]),
   539  				nil)
   540  		case "range-key-del":
   541  			if len(parts) != 3 {
   542  				return errors.Errorf("%s expects 2 arguments", parts[0])
   543  			}
   544  			err = b.RangeKeyDelete(
   545  				[]byte(parts[1]),
   546  				[]byte(parts[2]),
   547  				nil)
   548  		default:
   549  			return errors.Errorf("unknown op: %s", parts[0])
   550  		}
   551  		if err != nil {
   552  			return err
   553  		}
   554  	}
   555  	return nil
   556  }
   557  
   558  func runBuildCmd(td *datadriven.TestData, d *DB, fs vfs.FS) error {
   559  	b := d.NewIndexedBatch()
   560  	if err := runBatchDefineCmd(td, b); err != nil {
   561  		return err
   562  	}
   563  
   564  	if len(td.CmdArgs) != 1 {
   565  		return errors.New("build <path>: argument missing")
   566  	}
   567  	path := td.CmdArgs[0].String()
   568  
   569  	// Override table format, if provided.
   570  	tableFormat := d.opts.FormatMajorVersion.MaxTableFormat()
   571  	for _, cmdArg := range td.CmdArgs {
   572  		switch cmdArg.Key {
   573  		case "format":
   574  			v, err := strconv.Atoi(cmdArg.Vals[0])
   575  			if err != nil {
   576  				return err
   577  			}
   578  			tableFormat = sstable.TableFormat(v)
   579  		}
   580  	}
   581  
   582  	writeOpts := d.opts.MakeWriterOptions(0 /* level */, tableFormat)
   583  
   584  	f, err := fs.Create(path)
   585  	if err != nil {
   586  		return err
   587  	}
   588  	w := sstable.NewWriter(f, writeOpts)
   589  	iter := b.newInternalIter(nil)
   590  	for key, val := iter.First(); key != nil; key, val = iter.Next() {
   591  		tmp := *key
   592  		tmp.SetSeqNum(0)
   593  		if err := w.Add(tmp, val); err != nil {
   594  			return err
   595  		}
   596  	}
   597  	if err := iter.Close(); err != nil {
   598  		return err
   599  	}
   600  
   601  	if rdi := b.newRangeDelIter(nil, math.MaxUint64); rdi != nil {
   602  		for s := rdi.First(); s != nil; s = rdi.Next() {
   603  			err := rangedel.Encode(s, func(k base.InternalKey, v []byte) error {
   604  				k.SetSeqNum(0)
   605  				return w.Add(k, v)
   606  			})
   607  			if err != nil {
   608  				return err
   609  			}
   610  		}
   611  	}
   612  
   613  	if rki := b.newRangeKeyIter(nil, math.MaxUint64); rki != nil {
   614  		for s := rki.First(); s != nil; s = rki.Next() {
   615  			for _, k := range s.Keys {
   616  				var err error
   617  				switch k.Kind() {
   618  				case base.InternalKeyKindRangeKeySet:
   619  					err = w.RangeKeySet(s.Start, s.End, k.Suffix, k.Value)
   620  				case base.InternalKeyKindRangeKeyUnset:
   621  					err = w.RangeKeyUnset(s.Start, s.End, k.Suffix)
   622  				case base.InternalKeyKindRangeKeyDelete:
   623  					err = w.RangeKeyDelete(s.Start, s.End)
   624  				default:
   625  					panic("not a range key")
   626  				}
   627  				if err != nil {
   628  					return err
   629  				}
   630  			}
   631  		}
   632  	}
   633  
   634  	return w.Close()
   635  }
   636  
   637  func runCompactCmd(td *datadriven.TestData, d *DB) error {
   638  	if len(td.CmdArgs) > 4 {
   639  		return errors.Errorf("%s expects at most four arguments", td.Cmd)
   640  	}
   641  	parts := strings.Split(td.CmdArgs[0].Key, "-")
   642  	if len(parts) != 2 {
   643  		return errors.Errorf("expected <begin>-<end>: %s", td.Input)
   644  	}
   645  	parallelize := td.HasArg("parallel")
   646  	if len(td.CmdArgs) >= 2 && strings.HasPrefix(td.CmdArgs[1].Key, "L") {
   647  		levelString := td.CmdArgs[1].String()
   648  		iStart := base.MakeInternalKey([]byte(parts[0]), InternalKeySeqNumMax, InternalKeyKindMax)
   649  		iEnd := base.MakeInternalKey([]byte(parts[1]), 0, 0)
   650  		if levelString[0] != 'L' {
   651  			return errors.Errorf("expected L<n>: %s", levelString)
   652  		}
   653  		level, err := strconv.Atoi(levelString[1:])
   654  		if err != nil {
   655  			return err
   656  		}
   657  		return d.manualCompact(iStart.UserKey, iEnd.UserKey, level, parallelize)
   658  	}
   659  	return d.Compact([]byte(parts[0]), []byte(parts[1]), parallelize)
   660  }
   661  
   662  func runDBDefineCmd(td *datadriven.TestData, opts *Options) (*DB, error) {
   663  	opts = opts.EnsureDefaults()
   664  	opts.FS = vfs.NewMem()
   665  
   666  	if td.Input == "" {
   667  		// Empty LSM.
   668  		d, err := Open("", opts)
   669  		if err != nil {
   670  			return nil, err
   671  		}
   672  		return d, nil
   673  	}
   674  
   675  	var snapshots []uint64
   676  	var levelMaxBytes map[int]int64
   677  	for _, arg := range td.CmdArgs {
   678  		switch arg.Key {
   679  		case "target-file-sizes":
   680  			opts.Levels = make([]LevelOptions, len(arg.Vals))
   681  			for i := range arg.Vals {
   682  				size, err := strconv.ParseInt(arg.Vals[i], 10, 64)
   683  				if err != nil {
   684  					return nil, err
   685  				}
   686  				opts.Levels[i].TargetFileSize = size
   687  			}
   688  		case "snapshots":
   689  			snapshots = make([]uint64, len(arg.Vals))
   690  			for i := range arg.Vals {
   691  				seqNum, err := strconv.ParseUint(arg.Vals[i], 10, 64)
   692  				if err != nil {
   693  					return nil, err
   694  				}
   695  				snapshots[i] = seqNum
   696  				if i > 0 && snapshots[i] < snapshots[i-1] {
   697  					return nil, errors.New("Snapshots must be in ascending order")
   698  				}
   699  			}
   700  		case "level-max-bytes":
   701  			levelMaxBytes = map[int]int64{}
   702  			for i := range arg.Vals {
   703  				j := strings.Index(arg.Vals[i], ":")
   704  				levelStr := strings.TrimSpace(arg.Vals[i][:j])
   705  				level, err := strconv.Atoi(levelStr[1:])
   706  				if err != nil {
   707  					return nil, err
   708  				}
   709  				size, err := strconv.ParseInt(strings.TrimSpace(arg.Vals[i][j+1:]), 10, 64)
   710  				if err != nil {
   711  					return nil, err
   712  				}
   713  				levelMaxBytes[level] = size
   714  			}
   715  		case "auto-compactions":
   716  			switch arg.Vals[0] {
   717  			case "off":
   718  				opts.DisableAutomaticCompactions = true
   719  			case "on":
   720  				opts.DisableAutomaticCompactions = false
   721  			default:
   722  				return nil, errors.Errorf("Unrecognized %q %q arg value: %q", td.Cmd, arg.Key, arg.Vals[0])
   723  			}
   724  		case "enable-table-stats":
   725  			enable, err := strconv.ParseBool(arg.Vals[0])
   726  			if err != nil {
   727  				return nil, errors.Errorf("%s: could not parse %q as bool: %s", td.Cmd, arg.Vals[0], err)
   728  			}
   729  			opts.private.disableTableStats = !enable
   730  		case "block-size":
   731  			size, err := strconv.Atoi(arg.Vals[0])
   732  			if err != nil {
   733  				return nil, err
   734  			}
   735  			for _, levelOpts := range opts.Levels {
   736  				levelOpts.BlockSize = size
   737  			}
   738  		default:
   739  			return nil, errors.Errorf("%s: unknown arg: %s", td.Cmd, arg.Key)
   740  		}
   741  	}
   742  	d, err := Open("", opts)
   743  	if err != nil {
   744  		return nil, err
   745  	}
   746  	d.mu.Lock()
   747  	d.mu.versions.dynamicBaseLevel = false
   748  	for i := range snapshots {
   749  		s := &Snapshot{db: d}
   750  		s.seqNum = snapshots[i]
   751  		d.mu.snapshots.pushBack(s)
   752  	}
   753  	defer d.mu.Unlock()
   754  
   755  	var mem *memTable
   756  	var start, end *base.InternalKey
   757  	ve := &versionEdit{}
   758  	level := -1
   759  
   760  	maybeFlush := func() error {
   761  		if level < 0 {
   762  			return nil
   763  		}
   764  
   765  		toFlush := flushableList{{
   766  			flushable: mem,
   767  			flushed:   make(chan struct{}),
   768  		}}
   769  		c := newFlush(d.opts, d.mu.versions.currentVersion(),
   770  			d.mu.versions.picker.getBaseLevel(), toFlush)
   771  		c.disableSpanElision = true
   772  		// NB: define allows the test to exactly specify which keys go
   773  		// into which sstables. If the test has a small target file
   774  		// size to test grandparent limits, etc, the maxOutputFileSize
   775  		// can cause splitting /within/ the bounds specified to the
   776  		// test. Ignore the target size here, and split only according
   777  		// to the user-defined boundaries.
   778  		c.maxOutputFileSize = math.MaxUint64
   779  
   780  		newVE, _, err := d.runCompaction(0, c)
   781  		if err != nil {
   782  			return err
   783  		}
   784  		for _, f := range newVE.NewFiles {
   785  			if start != nil {
   786  				f.Meta.SmallestPointKey = *start
   787  				f.Meta.Smallest = *start
   788  			}
   789  			if end != nil {
   790  				f.Meta.LargestPointKey = *end
   791  				f.Meta.Largest = *end
   792  			}
   793  			ve.NewFiles = append(ve.NewFiles, newFileEntry{
   794  				Level: level,
   795  				Meta:  f.Meta,
   796  			})
   797  		}
   798  		level = -1
   799  		return nil
   800  	}
   801  
   802  	// Example, a-c.
   803  	parseMeta := func(s string) (*fileMetadata, error) {
   804  		parts := strings.Split(s, "-")
   805  		if len(parts) != 2 {
   806  			return nil, errors.Errorf("malformed table spec: %s", s)
   807  		}
   808  		m := (&fileMetadata{}).ExtendPointKeyBounds(
   809  			opts.Comparer.Compare,
   810  			InternalKey{UserKey: []byte(parts[0])},
   811  			InternalKey{UserKey: []byte(parts[1])},
   812  		)
   813  		return m, nil
   814  	}
   815  
   816  	// Example, compact: a-c.
   817  	parseCompaction := func(outputLevel int, s string) (*compaction, error) {
   818  		m, err := parseMeta(s[len("compact:"):])
   819  		if err != nil {
   820  			return nil, err
   821  		}
   822  		c := &compaction{
   823  			inputs:   []compactionLevel{{}, {level: outputLevel}},
   824  			smallest: m.Smallest,
   825  			largest:  m.Largest,
   826  		}
   827  		c.startLevel, c.outputLevel = &c.inputs[0], &c.inputs[1]
   828  		return c, nil
   829  	}
   830  
   831  	for _, line := range strings.Split(td.Input, "\n") {
   832  		fields := strings.Fields(line)
   833  		if len(fields) > 0 {
   834  			switch fields[0] {
   835  			case "mem":
   836  				if err := maybeFlush(); err != nil {
   837  					return nil, err
   838  				}
   839  				// Add a memtable layer.
   840  				if !d.mu.mem.mutable.empty() {
   841  					d.mu.mem.mutable = newMemTable(memTableOptions{Options: d.opts})
   842  					entry := d.newFlushableEntry(d.mu.mem.mutable, 0, 0)
   843  					entry.readerRefs++
   844  					d.mu.mem.queue = append(d.mu.mem.queue, entry)
   845  					d.updateReadStateLocked(nil)
   846  				}
   847  				mem = d.mu.mem.mutable
   848  				start, end = nil, nil
   849  				fields = fields[1:]
   850  			case "L0", "L1", "L2", "L3", "L4", "L5", "L6":
   851  				if err := maybeFlush(); err != nil {
   852  					return nil, err
   853  				}
   854  				var err error
   855  				if level, err = strconv.Atoi(fields[0][1:]); err != nil {
   856  					return nil, err
   857  				}
   858  				fields = fields[1:]
   859  				start, end = nil, nil
   860  				boundFields := 0
   861  				for _, field := range fields {
   862  					toBreak := false
   863  					switch {
   864  					case strings.HasPrefix(field, "start="):
   865  						ikey := base.ParseInternalKey(strings.TrimPrefix(field, "start="))
   866  						start = &ikey
   867  						boundFields++
   868  					case strings.HasPrefix(field, "end="):
   869  						ikey := base.ParseInternalKey(strings.TrimPrefix(field, "end="))
   870  						end = &ikey
   871  						boundFields++
   872  					default:
   873  						toBreak = true
   874  					}
   875  					if toBreak {
   876  						break
   877  					}
   878  				}
   879  				fields = fields[boundFields:]
   880  				mem = newMemTable(memTableOptions{Options: d.opts})
   881  			}
   882  		}
   883  
   884  		for _, data := range fields {
   885  			i := strings.Index(data, ":")
   886  			// Define in-progress compactions.
   887  			if data[:i] == "compact" {
   888  				c, err := parseCompaction(level, data)
   889  				if err != nil {
   890  					return nil, err
   891  				}
   892  				d.mu.compact.inProgress[c] = struct{}{}
   893  				continue
   894  			}
   895  			if data[:i] == "rangekey" {
   896  				span := keyspan.ParseSpan(data[i:])
   897  				err := rangekey.Encode(&span, func(k base.InternalKey, v []byte) error {
   898  					return mem.set(k, v)
   899  				})
   900  				if err != nil {
   901  					return nil, err
   902  				}
   903  				continue
   904  			}
   905  			key := base.ParseInternalKey(data[:i])
   906  			valueStr := data[i+1:]
   907  			value := []byte(valueStr)
   908  			if valueStr == "<largeval>" {
   909  				value = make([]byte, 4096)
   910  				rnd := rand.New(rand.NewSource(int64(key.SeqNum())))
   911  				if _, err := rnd.Read(value[:]); err != nil {
   912  					return nil, err
   913  				}
   914  			}
   915  			if err := mem.set(key, value); err != nil {
   916  				return nil, err
   917  			}
   918  		}
   919  	}
   920  
   921  	if err := maybeFlush(); err != nil {
   922  		return nil, err
   923  	}
   924  
   925  	if len(ve.NewFiles) > 0 {
   926  		jobID := d.mu.nextJobID
   927  		d.mu.nextJobID++
   928  		d.mu.versions.logLock()
   929  		if err := d.mu.versions.logAndApply(jobID, ve, newFileMetrics(ve.NewFiles), false, func() []compactionInfo {
   930  			return nil
   931  		}); err != nil {
   932  			return nil, err
   933  		}
   934  		d.updateReadStateLocked(nil)
   935  		d.updateTableStatsLocked(ve.NewFiles)
   936  	}
   937  
   938  	for l, maxBytes := range levelMaxBytes {
   939  		d.mu.versions.picker.(*compactionPickerByScore).levelMaxBytes[l] = maxBytes
   940  	}
   941  
   942  	return d, nil
   943  }
   944  
   945  func runTableStatsCmd(td *datadriven.TestData, d *DB) string {
   946  	u, err := strconv.ParseUint(strings.TrimSpace(td.Input), 10, 64)
   947  	if err != nil {
   948  		return err.Error()
   949  	}
   950  	fileNum := base.FileNum(u)
   951  
   952  	d.mu.Lock()
   953  	defer d.mu.Unlock()
   954  	v := d.mu.versions.currentVersion()
   955  	for _, levelMetadata := range v.Levels {
   956  		iter := levelMetadata.Iter()
   957  		for f := iter.First(); f != nil; f = iter.Next() {
   958  			if f.FileNum != fileNum {
   959  				continue
   960  			}
   961  
   962  			if !f.StatsValidLocked() {
   963  				d.waitTableStats()
   964  			}
   965  
   966  			var b bytes.Buffer
   967  			fmt.Fprintf(&b, "num-entries: %d\n", f.Stats.NumEntries)
   968  			fmt.Fprintf(&b, "num-deletions: %d\n", f.Stats.NumDeletions)
   969  			fmt.Fprintf(&b, "num-range-key-sets: %d\n", f.Stats.NumRangeKeySets)
   970  			fmt.Fprintf(&b, "point-deletions-bytes-estimate: %d\n", f.Stats.PointDeletionsBytesEstimate)
   971  			fmt.Fprintf(&b, "range-deletions-bytes-estimate: %d\n", f.Stats.RangeDeletionsBytesEstimate)
   972  			return b.String()
   973  		}
   974  	}
   975  	return "(not found)"
   976  }
   977  
   978  func runPopulateCmd(t *testing.T, td *datadriven.TestData, b *Batch) {
   979  	var timestamps []int
   980  	var maxKeyLength int
   981  	td.ScanArgs(t, "keylen", &maxKeyLength)
   982  	for _, cmdArg := range td.CmdArgs {
   983  		if cmdArg.Key != "timestamps" {
   984  			continue
   985  		}
   986  		for _, timestampVal := range cmdArg.Vals {
   987  			v, err := strconv.Atoi(timestampVal)
   988  			require.NoError(t, err)
   989  			timestamps = append(timestamps, v)
   990  		}
   991  	}
   992  
   993  	ks := testkeys.Alpha(maxKeyLength)
   994  	buf := make([]byte, ks.MaxLen()+testkeys.MaxSuffixLen)
   995  	for i := 0; i < ks.Count(); i++ {
   996  		for _, ts := range timestamps {
   997  			n := testkeys.WriteKeyAt(buf, ks, i, ts)
   998  			require.NoError(t, b.Set(buf[:n], buf[:n], nil))
   999  		}
  1000  	}
  1001  }
  1002  
  1003  // waitTableStats waits until all new files' statistics have been loaded. It's
  1004  // used in tests. The d.mu mutex must be locked while calling this method.
  1005  func (d *DB) waitTableStats() {
  1006  	for d.mu.tableStats.loading || len(d.mu.tableStats.pending) > 0 {
  1007  		d.mu.tableStats.cond.Wait()
  1008  	}
  1009  }
  1010  
  1011  func runIngestCmd(td *datadriven.TestData, d *DB, fs vfs.FS) error {
  1012  	paths := make([]string, 0, len(td.CmdArgs))
  1013  	for _, arg := range td.CmdArgs {
  1014  		paths = append(paths, arg.String())
  1015  	}
  1016  
  1017  	if err := d.Ingest(paths); err != nil {
  1018  		return err
  1019  	}
  1020  	return nil
  1021  }
  1022  
  1023  func runForceIngestCmd(td *datadriven.TestData, d *DB) error {
  1024  	var paths []string
  1025  	var level int
  1026  	for _, arg := range td.CmdArgs {
  1027  		switch arg.Key {
  1028  		case "paths":
  1029  			paths = append(paths, arg.Vals...)
  1030  		case "level":
  1031  			var err error
  1032  			level, err = strconv.Atoi(arg.Vals[0])
  1033  			if err != nil {
  1034  				return err
  1035  			}
  1036  		}
  1037  	}
  1038  	_, err := d.ingest(paths, func(
  1039  		tableNewIters,
  1040  		IterOptions,
  1041  		Compare,
  1042  		*version,
  1043  		int,
  1044  		map[*compaction]struct{},
  1045  		*fileMetadata,
  1046  	) (int, error) {
  1047  		return level, nil
  1048  	})
  1049  	return err
  1050  }
  1051  
  1052  func runLSMCmd(td *datadriven.TestData, d *DB) string {
  1053  	d.mu.Lock()
  1054  	s := d.mu.versions.currentVersion().String()
  1055  	d.mu.Unlock()
  1056  	return s
  1057  }