github.com/tachunwu/scale@v0.0.0-20230215135019-20d7de63a560/pkg/database/pebble.go (about)

     1  package database
     2  
     3  import (
     4  	"log"
     5  
     6  	"github.com/cockroachdb/pebble"
     7  	"github.com/cockroachdb/pebble/bloom"
     8  	"github.com/tachunwu/scale/pkg/database/bytealloc"
     9  )
    10  
    11  type DB interface {
    12  	NewIter(*pebble.IterOptions) iterator
    13  	NewBatch() batch
    14  	Scan(iter iterator, key []byte, count int64, reverse bool) error
    15  	Metrics() *pebble.Metrics
    16  	Flush() error
    17  	CommitBatch(*pebble.Batch) error
    18  }
    19  
    20  type iterator interface {
    21  	SeekLT(key []byte) bool
    22  	SeekGE(key []byte) bool
    23  	Valid() bool
    24  	Key() []byte
    25  	Value() []byte
    26  	First() bool
    27  	Next() bool
    28  	Last() bool
    29  	Prev() bool
    30  	Close() error
    31  }
    32  
    33  type batch interface {
    34  	Close() error
    35  	Commit(opts *pebble.WriteOptions) error
    36  	Set(key, value []byte, opts *pebble.WriteOptions) error
    37  	Delete(key []byte, opts *pebble.WriteOptions) error
    38  	LogData(data []byte, opts *pebble.WriteOptions) error
    39  }
    40  
    41  // Adapters for Pebble. Since the interfaces above are based on Pebble's
    42  // interfaces, it can simply forward calls for everything.
    43  type pebbleDB struct {
    44  	d       *pebble.DB
    45  	ballast []byte
    46  }
    47  
    48  func NewPebbleDB(dir string) DB {
    49  	verbose := true
    50  	cache := pebble.NewCache(1 << 30)
    51  	defer cache.Unref()
    52  	opts := &pebble.Options{
    53  		Cache:                       cache,
    54  		DisableWAL:                  false,
    55  		FormatMajorVersion:          pebble.FormatNewest,
    56  		L0CompactionThreshold:       2,
    57  		L0StopWritesThreshold:       1000,
    58  		LBaseMaxBytes:               64 << 20, // 64 MB
    59  		Levels:                      make([]pebble.LevelOptions, 7),
    60  		MaxOpenFiles:                16384,
    61  		MemTableSize:                64 << 20,
    62  		MemTableStopWritesThreshold: 4,
    63  		Merger: &pebble.Merger{
    64  			Name: "cockroach_merge_operator",
    65  		},
    66  		MaxConcurrentCompactions: func() int {
    67  			return 3
    68  		},
    69  	}
    70  
    71  	for i := 0; i < len(opts.Levels); i++ {
    72  		l := &opts.Levels[i]
    73  		l.BlockSize = 32 << 10       // 32 KB
    74  		l.IndexBlockSize = 256 << 10 // 256 KB
    75  		l.FilterPolicy = bloom.FilterPolicy(10)
    76  		l.FilterType = pebble.TableFilter
    77  		if i > 0 {
    78  			l.TargetFileSize = opts.Levels[i-1].TargetFileSize * 2
    79  		}
    80  		l.EnsureDefaults()
    81  	}
    82  	opts.Levels[6].FilterPolicy = nil
    83  	opts.FlushSplitBytes = opts.Levels[0].TargetFileSize
    84  
    85  	opts.EnsureDefaults()
    86  
    87  	if verbose {
    88  		lel := pebble.MakeLoggingEventListener(nil)
    89  		opts.EventListener = &lel
    90  		opts.EventListener.TableDeleted = nil
    91  		opts.EventListener.TableIngested = nil
    92  		opts.EventListener.WALCreated = nil
    93  		opts.EventListener.WALDeleted = nil
    94  	}
    95  
    96  	p, err := pebble.Open(dir, opts)
    97  	if err != nil {
    98  		log.Fatal(err)
    99  	}
   100  	return pebbleDB{
   101  		d:       p,
   102  		ballast: make([]byte, 1<<30),
   103  	}
   104  }
   105  
   106  func (p pebbleDB) Flush() error {
   107  	return p.d.Flush()
   108  }
   109  
   110  func (p pebbleDB) NewIter(opts *pebble.IterOptions) iterator {
   111  	return p.d.NewIter(opts)
   112  }
   113  
   114  func (p pebbleDB) NewBatch() batch {
   115  	return p.d.NewBatch()
   116  }
   117  
   118  // Add by testing
   119  func (p pebbleDB) CommitBatch(b *pebble.Batch) error {
   120  	return p.d.Apply(b, pebble.Sync)
   121  }
   122  
   123  func (p pebbleDB) Scan(iter iterator, key []byte, count int64, reverse bool) error {
   124  	var data bytealloc.A
   125  	if reverse {
   126  		for i, valid := 0, iter.SeekLT(key); valid; valid = iter.Prev() {
   127  			data, _ = data.Copy(iter.Key())
   128  			data, _ = data.Copy(iter.Value())
   129  			i++
   130  			if i >= int(count) {
   131  				break
   132  			}
   133  		}
   134  	} else {
   135  		for i, valid := 0, iter.SeekGE(key); valid; valid = iter.Next() {
   136  			data, _ = data.Copy(iter.Key())
   137  			data, _ = data.Copy(iter.Value())
   138  			i++
   139  			if i >= int(count) {
   140  				break
   141  			}
   142  		}
   143  	}
   144  	return nil
   145  }
   146  
   147  func (p pebbleDB) Metrics() *pebble.Metrics {
   148  	return p.d.Metrics()
   149  }