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 }