github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/performance/kvbench/prolly_store.go (about)

     1  // Copyright 2021 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package kvbench
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  
    21  	"github.com/dolthub/dolt/go/store/chunks"
    22  	"github.com/dolthub/dolt/go/store/datas"
    23  	"github.com/dolthub/dolt/go/store/nbs"
    24  
    25  	"github.com/dolthub/dolt/go/store/prolly/tree"
    26  	"github.com/dolthub/dolt/go/store/types"
    27  )
    28  
    29  const (
    30  	defaultMemTableSize = 256 * 1024 * 1024
    31  )
    32  
    33  func newMemoryProllyStore() keyValStore {
    34  	ctx := context.Background()
    35  	cs := &chunks.TestStorage{}
    36  	return newProllyStore(ctx, cs.NewViewWithDefaultFormat())
    37  }
    38  
    39  func newNBSProllyStore(dir string) keyValStore {
    40  	ctx := context.Background()
    41  	verStr := types.Format_Default.VersionString()
    42  	cs, err := nbs.NewLocalStore(ctx, verStr, dir, defaultMemTableSize, nbs.NewUnlimitedMemQuotaProvider())
    43  	if err != nil {
    44  		panic(err)
    45  	}
    46  	return newProllyStore(ctx, cs)
    47  }
    48  
    49  func newProllyStore(ctx context.Context, cs chunks.ChunkStore) keyValStore {
    50  	vrw := types.NewValueStore(cs)
    51  	ns := tree.NewNodeStore(cs)
    52  	db := datas.NewTypesDatabase(vrw, ns)
    53  	m, err := types.NewMap(ctx, vrw)
    54  	if err != nil {
    55  		panic(err)
    56  	}
    57  	return &prollyStore{
    58  		store:  m,
    59  		editor: types.NewMapEditor(m),
    60  		db:     db,
    61  		vrw:    vrw,
    62  	}
    63  }
    64  
    65  type prollyStore struct {
    66  	store  types.Map
    67  	editor *types.MapEditor
    68  	vrw    types.ValueReadWriter
    69  	db     datas.Database
    70  	mu     sync.RWMutex
    71  }
    72  
    73  var _ keyValStore = &prollyStore{}
    74  
    75  func (m *prollyStore) get(key []byte) (val []byte, ok bool) {
    76  	m.mu.RLock()
    77  	defer m.mu.RUnlock()
    78  
    79  	m.flush()
    80  
    81  	ctx := context.Background()
    82  	v, ok, err := m.store.MaybeGet(ctx, types.String(key))
    83  	if err != nil {
    84  		panic(err)
    85  	}
    86  
    87  	val = []byte(v.(types.String))
    88  	return val, ok
    89  }
    90  
    91  func (m *prollyStore) put(key, val []byte) {
    92  	m.mu.Lock()
    93  	defer m.mu.Unlock()
    94  
    95  	m.set(key, val)
    96  	m.flush()
    97  }
    98  
    99  func (m *prollyStore) delete(key []byte) {
   100  	m.mu.Lock()
   101  	defer m.mu.Unlock()
   102  
   103  	m.set(key, nil)
   104  	m.flush()
   105  }
   106  
   107  func (m *prollyStore) set(key, val []byte) {
   108  	k := types.String(key)
   109  	v := types.Value(nil)
   110  	if val != nil {
   111  		v = types.String(val)
   112  	}
   113  	m.editor.Set(k, v)
   114  }
   115  
   116  func (m *prollyStore) putMany(keys, vals [][]byte) {
   117  	m.mu.Lock()
   118  	defer m.mu.Unlock()
   119  
   120  	for i := range keys {
   121  		k := types.String(keys[i])
   122  		v := types.String(vals[i])
   123  		m.editor.Set(k, v)
   124  	}
   125  	m.flush()
   126  }
   127  
   128  func (m *prollyStore) flush() {
   129  	if m.editor.NumEdits() == 0 {
   130  		return
   131  	}
   132  
   133  	var err error
   134  	ctx := context.Background()
   135  
   136  	m.store, err = m.editor.Map(ctx)
   137  	if err != nil {
   138  		panic(err)
   139  	}
   140  
   141  	// persist
   142  	_, err = m.vrw.WriteValue(ctx, m.store)
   143  	if err != nil {
   144  		panic(err)
   145  	}
   146  	m.editor = types.NewMapEditor(m.store)
   147  }