vendor/github.com/syndtr/goleveldb/leveldb/memdb/memdb.go

   1 // Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
   2 // All rights reserved.
   3 //
   4 // Use of this source code is governed by a BSD-style license that can be
   5 // found in the LICENSE file.
   6
   7 // Package memdb provides in-memory key/value database implementation.
   8 package memdb
   9
  10 import (
  11         "math/rand"
  12         "sync"
  13
  14         "github.com/syndtr/goleveldb/leveldb/comparer"
  15         "github.com/syndtr/goleveldb/leveldb/errors"
  16         "github.com/syndtr/goleveldb/leveldb/iterator"
  17         "github.com/syndtr/goleveldb/leveldb/util"
  18 )
  19
  20 // Common errors.
  21 var (
  22         ErrNotFound     = errors.ErrNotFound
  23         ErrIterReleased = errors.New("leveldb/memdb: iterator released")
  24 )
  25
  26 const tMaxHeight = 12
  27
  28 type dbIter struct {
  29         util.BasicReleaser
  30         p          *DB
  31         slice      *util.Range
  32         node       int
  33         forward    bool
  34         key, value []byte
  35         err        error
  36 }
  37
  38 func (i *dbIter) fill(checkStart, checkLimit bool) bool {
  39         if i.node != 0 {
  40                 n := i.p.nodeData[i.node]
  41                 m := n + i.p.nodeData[i.node+nKey]
  42                 i.key = i.p.kvData[n:m]
  43                 if i.slice != nil {
  44                         switch {
  45                         case checkLimit && i.slice.Limit != nil && i.p.cmp.Compare(i.key, i.slice.Limit) >= 0:
  46                                 fallthrough
  47                         case checkStart && i.slice.Start != nil && i.p.cmp.Compare(i.key, i.slice.Start) < 0:
  48                                 i.node = 0
  49                                 goto bail
  50                         }
  51                 }
  52                 i.value = i.p.kvData[m : m+i.p.nodeData[i.node+nVal]]
  53                 return true
  54         }
  55 bail:
  56         i.key = nil
  57         i.value = nil
  58         return false
  59 }
  60
  61 func (i *dbIter) Valid() bool {
  62         return i.node != 0
  63 }
  64
  65 func (i *dbIter) First() bool {
  66         if i.Released() {
  67                 i.err = ErrIterReleased
  68                 return false
  69         }
  70
  71         i.forward = true
  72         i.p.mu.RLock()
  73         defer i.p.mu.RUnlock()
  74         if i.slice != nil && i.slice.Start != nil {
  75                 i.node, _ = i.p.findGE(i.slice.Start, false)
  76         } else {
  77                 i.node = i.p.nodeData[nNext]
  78         }
  79         return i.fill(false, true)
  80 }
  81
  82 func (i *dbIter) Last() bool {
  83         if i.Released() {
  84                 i.err = ErrIterReleased
  85                 return false
  86         }
  87
  88         i.forward = false
  89         i.p.mu.RLock()
  90         defer i.p.mu.RUnlock()
  91         if i.slice != nil && i.slice.Limit != nil {
  92                 i.node = i.p.findLT(i.slice.Limit)
  93         } else {
  94                 i.node = i.p.findLast()
  95         }
  96         return i.fill(true, false)
  97 }
  98
  99 func (i *dbIter) Seek(key []byte) bool {
 100         if i.Released() {
 101                 i.err = ErrIterReleased
 102                 return false
 103         }
 104
 105         i.forward = true
 106         i.p.mu.RLock()
 107         defer i.p.mu.RUnlock()
 108         if i.slice != nil && i.slice.Start != nil && i.p.cmp.Compare(key, i.slice.Start) < 0 {
 109                 key = i.slice.Start
 110         }
 111         i.node, _ = i.p.findGE(key, false)
 112         return i.fill(false, true)
 113 }
 114
 115 func (i *dbIter) Next() bool {
 116         if i.Released() {
 117                 i.err = ErrIterReleased
 118                 return false
 119         }
 120
 121         if i.node == 0 {
 122                 if !i.forward {
 123                         return i.First()
 124                 }
 125                 return false
 126         }
 127         i.forward = true
 128         i.p.mu.RLock()
 129         defer i.p.mu.RUnlock()
 130         i.node = i.p.nodeData[i.node+nNext]
 131         return i.fill(false, true)
 132 }
 133
 134 func (i *dbIter) Prev() bool {
 135         if i.Released() {
 136                 i.err = ErrIterReleased
 137                 return false
 138         }
 139
 140         if i.node == 0 {
 141                 if i.forward {
 142                         return i.Last()
 143                 }
 144                 return false
 145         }
 146         i.forward = false
 147         i.p.mu.RLock()
 148         defer i.p.mu.RUnlock()
 149         i.node = i.p.findLT(i.key)
 150         return i.fill(true, false)
 151 }
 152
 153 func (i *dbIter) Key() []byte {
 154         return i.key
 155 }
 156
 157 func (i *dbIter) Value() []byte {
 158         return i.value
 159 }
 160
 161 func (i *dbIter) Error() error { return i.err }
 162
 163 func (i *dbIter) Release() {
 164         if !i.Released() {
 165                 i.p = nil
 166                 i.node = 0
 167                 i.key = nil
 168                 i.value = nil
 169                 i.BasicReleaser.Release()
 170         }
 171 }
 172
 173 const (
 174         nKV = iota
 175         nKey
 176         nVal
 177         nHeight
 178         nNext
 179 )
 180
 181 // DB is an in-memory key/value database.
 182 type DB struct {
 183         cmp comparer.BasicComparer
 184         rnd *rand.Rand
 185
 186         mu     sync.RWMutex
 187         kvData []byte
 188         // Node data:
 189         // [0]         : KV offset
 190         // [1]         : Key length
 191         // [2]         : Value length
 192         // [3]         : Height
 193         // [3..height] : Next nodes
 194         nodeData  []int
 195         prevNode  [tMaxHeight]int
 196         maxHeight int
 197         n         int
 198         kvSize    int
 199 }
 200
 201 func (p *DB) randHeight() (h int) {
 202         const branching = 4
 203         h = 1
 204         for h < tMaxHeight && p.rnd.Int()%branching == 0 {
 205                 h++
 206         }
 207         return
 208 }
 209
 210 // Must hold RW-lock if prev == true, as it use shared prevNode slice.
 211 func (p *DB) findGE(key []byte, prev bool) (int, bool) {
 212         node := 0
 213         h := p.maxHeight - 1
 214         for {
 215                 next := p.nodeData[node+nNext+h]
 216                 cmp := 1
 217                 if next != 0 {
 218                         o := p.nodeData[next]
 219                         cmp = p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key)
 220                 }
 221                 if cmp < 0 {
 222                         // Keep searching in this list
 223                         node = next
 224                 } else {
 225                         if prev {
 226                                 p.prevNode[h] = node
 227                         } else if cmp == 0 {
 228                                 return next, true
 229                         }
 230                         if h == 0 {
 231                                 return next, cmp == 0
 232                         }
 233                         h--
 234                 }
 235         }
 236 }
 237
 238 func (p *DB) findLT(key []byte) int {
 239         node := 0
 240         h := p.maxHeight - 1
 241         for {
 242                 next := p.nodeData[node+nNext+h]
 243                 o := p.nodeData[next]
 244                 if next == 0 || p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key) >= 0 {
 245                         if h == 0 {
 246                                 break
 247                         }
 248                         h--
 249                 } else {
 250                         node = next
 251                 }
 252         }
 253         return node
 254 }
 255
 256 func (p *DB) findLast() int {
 257         node := 0
 258         h := p.maxHeight - 1
 259         for {
 260                 next := p.nodeData[node+nNext+h]
 261                 if next == 0 {
 262                         if h == 0 {
 263                                 break
 264                         }
 265                         h--
 266                 } else {
 267                         node = next
 268                 }
 269         }
 270         return node
 271 }
 272
 273 // Put sets the value for the given key. It overwrites any previous value
 274 // for that key; a DB is not a multi-map.
 275 //
 276 // It is safe to modify the contents of the arguments after Put returns.
 277 func (p *DB) Put(key []byte, value []byte) error {
 278         p.mu.Lock()
 279         defer p.mu.Unlock()
 280
 281         if node, exact := p.findGE(key, true); exact {
 282                 kvOffset := len(p.kvData)
 283                 p.kvData = append(p.kvData, key...)
 284                 p.kvData = append(p.kvData, value...)
 285                 p.nodeData[node] = kvOffset
 286                 m := p.nodeData[node+nVal]
 287                 p.nodeData[node+nVal] = len(value)
 288                 p.kvSize += len(value) - m
 289                 return nil
 290         }
 291
 292         h := p.randHeight()
 293         if h > p.maxHeight {
 294                 for i := p.maxHeight; i < h; i++ {
 295                         p.prevNode[i] = 0
 296                 }
 297                 p.maxHeight = h
 298         }
 299
 300         kvOffset := len(p.kvData)
 301         p.kvData = append(p.kvData, key...)
 302         p.kvData = append(p.kvData, value...)
 303         // Node
 304         node := len(p.nodeData)
 305         p.nodeData = append(p.nodeData, kvOffset, len(key), len(value), h)
 306         for i, n := range p.prevNode[:h] {
 307                 m := n + nNext + i
 308                 p.nodeData = append(p.nodeData, p.nodeData[m])
 309                 p.nodeData[m] = node
 310         }
 311
 312         p.kvSize += len(key) + len(value)
 313         p.n++
 314         return nil
 315 }
 316
 317 // Delete deletes the value for the given key. It returns ErrNotFound if
 318 // the DB does not contain the key.
 319 //
 320 // It is safe to modify the contents of the arguments after Delete returns.
 321 func (p *DB) Delete(key []byte) error {
 322         p.mu.Lock()
 323         defer p.mu.Unlock()
 324
 325         node, exact := p.findGE(key, true)
 326         if !exact {
 327                 return ErrNotFound
 328         }
 329
 330         h := p.nodeData[node+nHeight]
 331         for i, n := range p.prevNode[:h] {
 332                 m := n + 4 + i
 333                 p.nodeData[m] = p.nodeData[p.nodeData[m]+nNext+i]
 334         }
 335
 336         p.kvSize -= p.nodeData[node+nKey] + p.nodeData[node+nVal]
 337         p.n--
 338         return nil
 339 }
 340
 341 // Contains returns true if the given key are in the DB.
 342 //
 343 // It is safe to modify the contents of the arguments after Contains returns.
 344 func (p *DB) Contains(key []byte) bool {
 345         p.mu.RLock()
 346         _, exact := p.findGE(key, false)
 347         p.mu.RUnlock()
 348         return exact
 349 }
 350
 351 // Get gets the value for the given key. It returns error.ErrNotFound if the
 352 // DB does not contain the key.
 353 //
 354 // The caller should not modify the contents of the returned slice, but
 355 // it is safe to modify the contents of the argument after Get returns.
 356 func (p *DB) Get(key []byte) (value []byte, err error) {
 357         p.mu.RLock()
 358         if node, exact := p.findGE(key, false); exact {
 359                 o := p.nodeData[node] + p.nodeData[node+nKey]
 360                 value = p.kvData[o : o+p.nodeData[node+nVal]]
 361         } else {
 362                 err = ErrNotFound
 363         }
 364         p.mu.RUnlock()
 365         return
 366 }
 367
 368 // Find finds key/value pair whose key is greater than or equal to the
 369 // given key. It returns ErrNotFound if the table doesn't contain
 370 // such pair.
 371 //
 372 // The caller should not modify the contents of the returned slice, but
 373 // it is safe to modify the contents of the argument after Find returns.
 374 func (p *DB) Find(key []byte) (rkey, value []byte, err error) {
 375         p.mu.RLock()
 376         if node, _ := p.findGE(key, false); node != 0 {
 377                 n := p.nodeData[node]
 378                 m := n + p.nodeData[node+nKey]
 379                 rkey = p.kvData[n:m]
 380                 value = p.kvData[m : m+p.nodeData[node+nVal]]
 381         } else {
 382                 err = ErrNotFound
 383         }
 384         p.mu.RUnlock()
 385         return
 386 }
 387
 388 // NewIterator returns an iterator of the DB.
 389 // The returned iterator is not safe for concurrent use, but it is safe to use
 390 // multiple iterators concurrently, with each in a dedicated goroutine.
 391 // It is also safe to use an iterator concurrently with modifying its
 392 // underlying DB. However, the resultant key/value pairs are not guaranteed
 393 // to be a consistent snapshot of the DB at a particular point in time.
 394 //
 395 // Slice allows slicing the iterator to only contains keys in the given
 396 // range. A nil Range.Start is treated as a key before all keys in the
 397 // DB. And a nil Range.Limit is treated as a key after all keys in
 398 // the DB.
 399 //
 400 // The iterator must be released after use, by calling Release method.
 401 //
 402 // Also read Iterator documentation of the leveldb/iterator package.
 403 func (p *DB) NewIterator(slice *util.Range) iterator.Iterator {
 404         return &dbIter{p: p, slice: slice}
 405 }
 406
 407 // Capacity returns keys/values buffer capacity.
 408 func (p *DB) Capacity() int {
 409         p.mu.RLock()
 410         defer p.mu.RUnlock()
 411         return cap(p.kvData)
 412 }
 413
 414 // Size returns sum of keys and values length. Note that deleted
 415 // key/value will not be accounted for, but it will still consume
 416 // the buffer, since the buffer is append only.
 417 func (p *DB) Size() int {
 418         p.mu.RLock()
 419         defer p.mu.RUnlock()
 420         return p.kvSize
 421 }
 422
 423 // Free returns keys/values free buffer before need to grow.
 424 func (p *DB) Free() int {
 425         p.mu.RLock()
 426         defer p.mu.RUnlock()
 427         return cap(p.kvData) - len(p.kvData)
 428 }
 429
 430 // Len returns the number of entries in the DB.
 431 func (p *DB) Len() int {
 432         p.mu.RLock()
 433         defer p.mu.RUnlock()
 434         return p.n
 435 }
 436
 437 // Reset resets the DB to initial empty state. Allows reuse the buffer.
 438 func (p *DB) Reset() {
 439         p.mu.Lock()
 440         p.rnd = rand.New(rand.NewSource(0xdeadbeef))
 441         p.maxHeight = 1
 442         p.n = 0
 443         p.kvSize = 0
 444         p.kvData = p.kvData[:0]
 445         p.nodeData = p.nodeData[:nNext+tMaxHeight]
 446         p.nodeData[nKV] = 0
 447         p.nodeData[nKey] = 0
 448         p.nodeData[nVal] = 0
 449         p.nodeData[nHeight] = tMaxHeight
 450         for n := 0; n < tMaxHeight; n++ {
 451                 p.nodeData[nNext+n] = 0
 452                 p.prevNode[n] = 0
 453         }
 454         p.mu.Unlock()
 455 }
 456
 457 // New creates a new initialized in-memory key/value DB. The capacity
 458 // is the initial key/value buffer capacity. The capacity is advisory,
 459 // not enforced.
 460 //
 461 // This DB is append-only, deleting an entry would remove entry node but not
 462 // reclaim KV buffer.
 463 //
 464 // The returned DB instance is safe for concurrent use.
 465 func New(cmp comparer.BasicComparer, capacity int) *DB {
 466         p := &DB{
 467                 cmp:       cmp,
 468                 rnd:       rand.New(rand.NewSource(0xdeadbeef)),
 469                 maxHeight: 1,
 470                 kvData:    make([]byte, 0, capacity),
 471                 nodeData:  make([]int, 4+tMaxHeight),
 472         }
 473         p.nodeData[nHeight] = tMaxHeight
 474         return p
 475 }