github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/internal/datatest/datatest.go (about)

     1  // Copyright 2022 The LevelDB-Go and Pebble 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 datatest provides common datadriven test commands for use outside of
     6  // the root Pebble package.
     7  package datatest
     8  
     9  import (
    10  	"strings"
    11  	"sync"
    12  
    13  	"github.com/cockroachdb/datadriven"
    14  	"github.com/cockroachdb/errors"
    15  	"github.com/cockroachdb/pebble"
    16  )
    17  
    18  // TODO(jackson): Consider a refactoring that can consolidate this package and
    19  // the datadriven commands defined in pebble/data_test.go.
    20  
    21  // DefineBatch interprets the provided datadriven command as a sequence of write
    22  // operations, one-per-line, to apply to the provided batch.
    23  func DefineBatch(d *datadriven.TestData, b *pebble.Batch) error {
    24  	for _, line := range strings.Split(d.Input, "\n") {
    25  		parts := strings.Fields(line)
    26  		if len(parts) == 0 {
    27  			continue
    28  		}
    29  		if parts[1] == `<nil>` {
    30  			parts[1] = ""
    31  		}
    32  		var err error
    33  		switch parts[0] {
    34  		case "set":
    35  			if len(parts) != 3 {
    36  				return errors.Errorf("%s expects 2 arguments", parts[0])
    37  			}
    38  			err = b.Set([]byte(parts[1]), []byte(parts[2]), nil)
    39  		case "del":
    40  			if len(parts) != 2 {
    41  				return errors.Errorf("%s expects 1 argument", parts[0])
    42  			}
    43  			err = b.Delete([]byte(parts[1]), nil)
    44  		case "singledel":
    45  			if len(parts) != 2 {
    46  				return errors.Errorf("%s expects 1 argument", parts[0])
    47  			}
    48  			err = b.SingleDelete([]byte(parts[1]), nil)
    49  		case "del-range":
    50  			if len(parts) != 3 {
    51  				return errors.Errorf("%s expects 2 arguments", parts[0])
    52  			}
    53  			err = b.DeleteRange([]byte(parts[1]), []byte(parts[2]), nil)
    54  		case "merge":
    55  			if len(parts) != 3 {
    56  				return errors.Errorf("%s expects 2 arguments", parts[0])
    57  			}
    58  			err = b.Merge([]byte(parts[1]), []byte(parts[2]), nil)
    59  		case "range-key-set":
    60  			if len(parts) != 5 {
    61  				return errors.Errorf("%s expects 4 arguments", parts[0])
    62  			}
    63  			err = b.RangeKeySet(
    64  				[]byte(parts[1]),
    65  				[]byte(parts[2]),
    66  				[]byte(parts[3]),
    67  				[]byte(parts[4]),
    68  				nil)
    69  		case "range-key-unset":
    70  			if len(parts) != 4 {
    71  				return errors.Errorf("%s expects 3 arguments", parts[0])
    72  			}
    73  			err = b.RangeKeyUnset(
    74  				[]byte(parts[1]),
    75  				[]byte(parts[2]),
    76  				[]byte(parts[3]),
    77  				nil)
    78  		case "range-key-del":
    79  			if len(parts) != 3 {
    80  				return errors.Errorf("%s expects 2 arguments", parts[0])
    81  			}
    82  			err = b.RangeKeyDelete(
    83  				[]byte(parts[1]),
    84  				[]byte(parts[2]),
    85  				nil)
    86  		default:
    87  			return errors.Errorf("unknown op: %s", parts[0])
    88  		}
    89  		if err != nil {
    90  			return err
    91  		}
    92  	}
    93  	return nil
    94  }
    95  
    96  // CompactionTracker is a listener that tracks the number of compactions.
    97  type CompactionTracker struct {
    98  	sync.Cond
    99  	count    int
   100  	attached bool
   101  }
   102  
   103  // NewCompactionTracker setups the necessary options to keep track of the
   104  // compactions that are in flight.
   105  func NewCompactionTracker(options *pebble.Options) *CompactionTracker {
   106  	ct := CompactionTracker{}
   107  	ct.Cond = sync.Cond{
   108  		L: &sync.Mutex{},
   109  	}
   110  	ct.attached = true
   111  	el := pebble.EventListener{
   112  		CompactionEnd: func(info pebble.CompactionInfo) {
   113  			ct.L.Lock()
   114  			ct.count--
   115  			ct.Broadcast()
   116  			ct.L.Unlock()
   117  		},
   118  		CompactionBegin: func(info pebble.CompactionInfo) {
   119  			ct.L.Lock()
   120  			ct.count++
   121  			ct.Broadcast()
   122  			ct.L.Unlock()
   123  		},
   124  	}
   125  
   126  	options.AddEventListener(el)
   127  	return &ct
   128  }
   129  
   130  // WaitForInflightCompactionsToEqual waits until compactions meet the specified target.
   131  func (cql *CompactionTracker) WaitForInflightCompactionsToEqual(target int) {
   132  	cql.L.Lock()
   133  	if !cql.attached {
   134  		panic("Cannot wait for compactions if listener has not been attached")
   135  	}
   136  	for cql.count != target {
   137  		cql.Wait()
   138  	}
   139  	cql.L.Unlock()
   140  }