github.com/psiphon-labs/psiphon-tunnel-core@v2.0.28+incompatible/psiphon/dataStore_badger.go (about) 1 //go:build PSIPHON_USE_BADGER_DB 2 // +build PSIPHON_USE_BADGER_DB 3 4 /* 5 * Copyright (c) 2018, Psiphon Inc. 6 * All rights reserved. 7 * 8 * This program is free software: you can redistribute it and/or modify 9 * it under the terms of the GNU General Public License as published by 10 * the Free Software Foundation, either version 3 of the License, or 11 * (at your option) any later version. 12 * 13 * This program is distributed in the hope that it will be useful, 14 * but WITHOUT ANY WARRANTY; without even the implied warranty of 15 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 16 * GNU General Public License for more details. 17 * 18 * You should have received a copy of the GNU General Public License 19 * along with this program. If not, see <http://www.gnu.org/licenses/>. 20 * 21 */ 22 23 package psiphon 24 25 import ( 26 "os" 27 "path/filepath" 28 29 "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/errors" 30 "github.com/dgraph-io/badger" 31 "github.com/dgraph-io/badger/options" 32 ) 33 34 const ( 35 DATA_STORE_DIRECTORY = "psiphon.badgerdb" 36 ) 37 38 type datastoreDB struct { 39 badgerDB *badger.DB 40 } 41 42 type datastoreTx struct { 43 badgerTx *badger.Txn 44 } 45 46 type datastoreBucket struct { 47 name []byte 48 tx *datastoreTx 49 } 50 51 type datastoreCursor struct { 52 badgerIterator *badger.Iterator 53 prefix []byte 54 } 55 56 func datastoreOpenDB( 57 rootDataDirectory string, _ bool) (*datastoreDB, error) { 58 59 dbDirectory := filepath.Join(rootDataDirectory, "psiphon.badgerdb") 60 61 err := os.MkdirAll(dbDirectory, 0700) 62 if err != nil { 63 return nil, errors.Trace(err) 64 } 65 66 opts := badger.DefaultOptions 67 68 opts.Dir = dbDirectory 69 opts.ValueDir = dbDirectory 70 71 opts.TableLoadingMode = options.FileIO 72 opts.ValueLogLoadingMode = options.FileIO 73 opts.MaxTableSize = 1 << 16 74 opts.ValueLogFileSize = 1 << 20 75 opts.NumMemtables = 1 76 opts.NumLevelZeroTables = 1 77 opts.NumLevelZeroTablesStall = 2 78 opts.NumCompactors = 1 79 80 db, err := badger.Open(opts) 81 if err != nil { 82 return nil, errors.Trace(err) 83 } 84 85 for { 86 if db.RunValueLogGC(0.5) != nil { 87 break 88 } 89 } 90 91 return &datastoreDB{badgerDB: db}, nil 92 } 93 94 func (db *datastoreDB) close() error { 95 return db.badgerDB.Close() 96 } 97 98 func (db *datastoreDB) getDataStoreMetrics() string { 99 // TODO: report metrics 100 return "" 101 } 102 103 func (db *datastoreDB) view(fn func(tx *datastoreTx) error) error { 104 return db.badgerDB.View( 105 func(tx *badger.Txn) error { 106 err := fn(&datastoreTx{badgerTx: tx}) 107 if err != nil { 108 return errors.Trace(err) 109 } 110 return nil 111 }) 112 } 113 114 func (db *datastoreDB) update(fn func(tx *datastoreTx) error) error { 115 return db.badgerDB.Update( 116 func(tx *badger.Txn) error { 117 err := fn(&datastoreTx{badgerTx: tx}) 118 if err != nil { 119 return errors.Trace(err) 120 } 121 return nil 122 }) 123 } 124 125 func (tx *datastoreTx) bucket(name []byte) *datastoreBucket { 126 return &datastoreBucket{ 127 name: name, 128 tx: tx, 129 } 130 } 131 132 func (tx *datastoreTx) clearBucket(name []byte) error { 133 b := tx.bucket(name) 134 c := b.cursor() 135 for key := c.firstKey(); key != nil; key = c.nextKey() { 136 err := tx.badgerTx.Delete(key) 137 if err != nil { 138 return errors.Trace(err) 139 } 140 } 141 return nil 142 } 143 144 func (b *datastoreBucket) get(key []byte) []byte { 145 keyWithPrefix := append(b.name, key...) 146 item, err := b.tx.badgerTx.Get(keyWithPrefix) 147 if err != nil { 148 if err != badger.ErrKeyNotFound { 149 // The original datastore interface does not return an error from 150 // Get, so emit notice. 151 NoticeWarning("get failed: %s: %s", 152 string(keyWithPrefix), errors.Trace(err)) 153 } 154 return nil 155 } 156 value, err := item.Value() 157 if err != nil { 158 NoticeWarning("get failed: %s: %s", 159 string(keyWithPrefix), errors.Trace(err)) 160 return nil 161 } 162 return value 163 } 164 165 func (b *datastoreBucket) put(key, value []byte) error { 166 keyWithPrefix := append(b.name, key...) 167 err := b.tx.badgerTx.Set(keyWithPrefix, value) 168 if err != nil { 169 return errors.Trace(err) 170 } 171 return nil 172 } 173 174 func (b *datastoreBucket) delete(key []byte) error { 175 keyWithPrefix := append(b.name, key...) 176 err := b.tx.badgerTx.Delete(keyWithPrefix) 177 if err != nil { 178 return errors.Trace(err) 179 } 180 return nil 181 } 182 183 func (b *datastoreBucket) cursor() *datastoreCursor { 184 opts := badger.DefaultIteratorOptions 185 opts.PrefetchValues = false 186 iterator := b.tx.badgerTx.NewIterator(opts) 187 return &datastoreCursor{badgerIterator: iterator, prefix: b.name} 188 } 189 190 func (c *datastoreCursor) firstKey() []byte { 191 c.badgerIterator.Seek(c.prefix) 192 return c.currentKey() 193 } 194 195 func (c *datastoreCursor) currentKey() []byte { 196 if !c.badgerIterator.ValidForPrefix(c.prefix) { 197 return nil 198 } 199 item := c.badgerIterator.Item() 200 return item.Key()[len(c.prefix):] 201 } 202 203 func (c *datastoreCursor) nextKey() []byte { 204 c.badgerIterator.Next() 205 return c.currentKey() 206 } 207 208 func (c *datastoreCursor) first() ([]byte, []byte) { 209 c.badgerIterator.Seek(c.prefix) 210 return c.current() 211 } 212 213 func (c *datastoreCursor) current() ([]byte, []byte) { 214 if !c.badgerIterator.ValidForPrefix(c.prefix) { 215 return nil, nil 216 } 217 item := c.badgerIterator.Item() 218 value, err := item.Value() 219 if err != nil { 220 return nil, nil 221 } 222 return item.Key()[len(c.prefix):], value 223 } 224 225 func (c *datastoreCursor) next() ([]byte, []byte) { 226 c.badgerIterator.Next() 227 return c.current() 228 } 229 230 func (c *datastoreCursor) close() { 231 c.badgerIterator.Close() 232 }