memdb.go

Documentation: github.com/syndtr/goleveldb/leveldb/memdb

     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 + nNext + 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  // WARNING: Any slice returned by interator (e.g. slice returned by calling
   401  // Iterator.Key() or Iterator.Key() methods), its content should not be modified
   402  // unless noted otherwise.
   403  //
   404  // The iterator must be released after use, by calling Release method.
   405  //
   406  // Also read Iterator documentation of the leveldb/iterator package.
   407  func (p *DB) NewIterator(slice *util.Range) iterator.Iterator {
   408  	return &dbIter{p: p, slice: slice}
   409  }
   410  
   411  // Capacity returns keys/values buffer capacity.
   412  func (p *DB) Capacity() int {
   413  	p.mu.RLock()
   414  	defer p.mu.RUnlock()
   415  	return cap(p.kvData)
   416  }
   417  
   418  // Size returns sum of keys and values length. Note that deleted
   419  // key/value will not be accounted for, but it will still consume
   420  // the buffer, since the buffer is append only.
   421  func (p *DB) Size() int {
   422  	p.mu.RLock()
   423  	defer p.mu.RUnlock()
   424  	return p.kvSize
   425  }
   426  
   427  // Free returns keys/values free buffer before need to grow.
   428  func (p *DB) Free() int {
   429  	p.mu.RLock()
   430  	defer p.mu.RUnlock()
   431  	return cap(p.kvData) - len(p.kvData)
   432  }
   433  
   434  // Len returns the number of entries in the DB.
   435  func (p *DB) Len() int {
   436  	p.mu.RLock()
   437  	defer p.mu.RUnlock()
   438  	return p.n
   439  }
   440  
   441  // Reset resets the DB to initial empty state. Allows reuse the buffer.
   442  func (p *DB) Reset() {
   443  	p.mu.Lock()
   444  	p.rnd = rand.New(rand.NewSource(0xdeadbeef))
   445  	p.maxHeight = 1
   446  	p.n = 0
   447  	p.kvSize = 0
   448  	p.kvData = p.kvData[:0]
   449  	p.nodeData = p.nodeData[:nNext+tMaxHeight]
   450  	p.nodeData[nKV] = 0
   451  	p.nodeData[nKey] = 0
   452  	p.nodeData[nVal] = 0
   453  	p.nodeData[nHeight] = tMaxHeight
   454  	for n := 0; n < tMaxHeight; n++ {
   455  		p.nodeData[nNext+n] = 0
   456  		p.prevNode[n] = 0
   457  	}
   458  	p.mu.Unlock()
   459  }
   460  
   461  // New creates a new initialized in-memory key/value DB. The capacity
   462  // is the initial key/value buffer capacity. The capacity is advisory,
   463  // not enforced.
   464  //
   465  // This DB is append-only, deleting an entry would remove entry node but not
   466  // reclaim KV buffer.
   467  //
   468  // The returned DB instance is safe for concurrent use.
   469  func New(cmp comparer.BasicComparer, capacity int) *DB {
   470  	p := &DB{
   471  		cmp:       cmp,
   472  		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
   473  		maxHeight: 1,
   474  		kvData:    make([]byte, 0, capacity),
   475  		nodeData:  make([]int, 4+tMaxHeight),
   476  	}
   477  	p.nodeData[nHeight] = tMaxHeight
   478  	return p
   479  }
   480
View as plain text