Add treapstore, a treap-based in-memory data store

common/data/treapstore/store.go

Index: common/data/treapstore/store.go
diff --git a/common/data/treapstore/store.go b/common/data/treapstore/store.go
new file mode 100644
index 0000000000000000000000000000000000000000..ebc3babba0946e1e97ba234376b2fa7836979723
--- /dev/null
+++ b/common/data/treapstore/store.go
@@ -0,0 +1,213 @@
+// Copyright 2016 The LUCI Authors. All rights reserved.
+// Use of this source code is governed under the Apache License, Version 2.0
+// that can be found in the LICENSE file.
+
+// Package treapstore is a lightweight append-only in-memory key-value store
+// built on top a treap (tree + heap) implementation.
+//
+// treapstore is specifically focused on supporting the in-memory datastore
+// implementation at "github.com/luci/gae/impl/memory".
+package treapstore
+
+import (
+	"math/rand"
+	"sort"
+	"sync"
+	"sync/atomic"
+
+	"github.com/luci/gtreap"

[Vadim Sh., 2016/12/28 01:33:40]

do we need to fork it? IIRC, we forked gkvlite to fix some bugs that upstream
ignored. luci/gtreap looks to be identical to steveyen/gtreap.

[Vadim Sh., 2016/12/28 01:34:27]

nvm, saw another CL :)

[dnj, 2016/12/28 03:01:06]

Acknowledged. Yeah will actually put that forward as a pull request post-review
and if steveyen merges it back in will un-fork.

+)
+
+// Store is a generaly-purpose concurrently-accessible copy-on-write store built
+// on top of the github.com/steveyen/gtreap treap (tree + heap) implementation.
+//
+// A Store is composed of a series of named Collection instances, each of which
+// can hold data elements.
+//
+// A Store is created by calling New, and is initially in a read/write mode.
+// Derivative Stores can be created by calling a Store's Snapshot capability.
+// These Stores will be read-only, meaning their collections and those
+// collections' data data may only be accessed, not modified.
+type Store struct {
+	compare  gtreap.Compare
+	colls    map[string]*Collection
+	collLock *sync.RWMutex
+}
+
+// New creates a new read/write Store.
+func New(c gtreap.Compare) *Store {
+	return &Store{
+		compare:  c,
+		collLock: &sync.RWMutex{},
+	}
+}
+
+// IsReadOnly returns true if this Store is read-only.
+func (s *Store) IsReadOnly() bool { return s.collLock == nil }
+
+func (s *Store) assertNotReadOnly() {
+	if s.IsReadOnly() {
+		panic("store is read-only")
+	}
+}
+
+// Snapshot creates a snapshot of the Store and all of its Collection instances.

[Vadim Sh., 2016/12/28 01:33:40]

nit: "Snapshot creates a read-only copy of the Store ..."

This is more informative.

[dnj, 2016/12/28 03:01:06]

Done.

+// Because a Store is copy-on-write, this is a cheap operation.
+func (s *Store) Snapshot() *Store {
+	if s.IsReadOnly() {
+		return s
+	}
+
+	s.collLock.RLock()
+	defer s.collLock.RUnlock()
+
+	// Return a read-only Store with the new Collection set.
+	snap := &Store{
+		compare:  s.compare,
+		colls:    make(map[string]*Collection, len(s.colls)),
+		collLock: nil,
+	}
+
+	// Create a read-only Collection for each coll.
+	for k, coll := range s.colls {
+		newColl := &Collection{
+			s:    snap,
+			name: coll.name,
+		}
+		newColl.setRoot(coll.currentRoot())
+		snap.colls[k] = newColl
+	}
+	return snap
+}
+
+// GetCollectionNames returns the names of the Collections in this Store, sorted
+// alphabetically.
+func (s *Store) GetCollectionNames() []string {
+	if !s.IsReadOnly() {
+		s.collLock.RLock()
+		defer s.collLock.RUnlock()
+	}
+
+	names := make([]string, 0, len(s.colls))
+	for k := range s.colls {
+		names = append(names, k)
+	}
+	sort.Strings(names)
+	return names
+}
+
+// GetCollection returns the Collection with the specified name. If no such
+// Collection exists, GetCollection will return nil.
+func (s *Store) GetCollection(name string) *Collection {
+	if !s.IsReadOnly() {
+		s.collLock.RLock()
+		defer s.collLock.RUnlock()
+	}
+	return s.colls[name]
+}
+
+// GetOrCreateCollection returns a Collection with the specified name, creating
+// it if it does not exist. If s is read-only, this will panic regardless of
+// whether the named Collection exists.
+func (s *Store) GetOrCreateCollection(name string) *Collection {
+	s.assertNotReadOnly()
+
+	// Fast path: get the Collection using read lock.
+	coll := s.GetCollection(name)
+	if coll != nil {
+		return coll
+	}
+
+	// Might have to create the Collection. Grab the write lock.
+	s.collLock.Lock()
+	defer s.collLock.Unlock()
+
+	// Try to get the existing Collection again now that we hold the write lock.
+	coll = s.colls[name]
+	if coll != nil {
+		return coll
+	}
+
+	// Create a new read/write Collection.
+	coll = &Collection{
+		s:       s,
+		name:    name,
+		putLock: &sync.Mutex{},
+	}
+	coll.setRoot(gtreap.NewTreap(s.compare))
+
+	if s.colls == nil {
+		s.colls = make(map[string]*Collection)
+	}
+	s.colls[name] = coll
+	return coll
+}
+
+// Collection is a collection of Items.
+//
+// Collections belonging to read/write Store instances are, themselves,
+// read/write. Collections belonging to Snapshot instances are read-only.
+type Collection struct {
+	s    *Store

[Vadim Sh., 2016/12/28 01:33:39]

nit: this is unnecessary, just use 'putLock == nil' as a check for the read-only
bit, as in Store.

[dnj, 2016/12/28 03:01:06]

Done.

+	name string
+
+	root    *atomic.Value

[Vadim Sh., 2016/12/28 01:33:39]

I'm not sure atomic.Value is necessary since you use putLock too. Is there a
reason to use atomic.Value in addition to the lock?

[dnj, 2016/12/28 03:01:06]

Readers will grab this value w/out taking out a lock. I suppose I could use
RWLock. Probably cleaner, better performance. I'll do that.

+	putLock *sync.Mutex
+}
+
+// Name returns this Collection's name.
+func (c *Collection) Name() string { return c.name }
+
+// IsReadOnly returns true if this Collection is read-only.
+func (c *Collection) IsReadOnly() bool { return c.s.IsReadOnly() }
+
+// Min returns the smallest item in the collection.
+func (c *Collection) Min() gtreap.Item { return c.currentRoot().Min() }
+
+// Max returns the largest item in the collection.
+func (c *Collection) Max() gtreap.Item { return c.currentRoot().Max() }
+
+// Get returns the item in the Store that matches i, or nil if no such item
+// exists.
+func (c *Collection) Get(i gtreap.Item) gtreap.Item { return c.currentRoot().Get(i) }
+
+// Put adds an item to the Store.
+//
+// If the Store is read-only, Put will panic.
+func (c *Collection) Put(i gtreap.Item) {
+	c.s.assertNotReadOnly()
+
+	// Lock around the entire Upsert operation to serialize Puts.
+	c.putLock.Lock()
+	c.setRoot(c.currentRoot().Upsert(i, rand.Int()))

[Vadim Sh., 2016/12/28 01:33:40]

nit: do rand.Int() outside the lock

[dnj, 2016/12/28 03:01:06]

um ... yes done. /ashamed

+	c.putLock.Unlock()
+}
+
+// Delete deletes an item from the Collection, if such an item exists.
+func (c *Collection) Delete(i gtreap.Item) {
+	c.s.assertNotReadOnly()
+
+	c.putLock.Lock()
+	c.setRoot(c.currentRoot().Delete(i))
+	c.putLock.Unlock()
+}
+
+func (c *Collection) currentRoot() *gtreap.Treap {
+	if c.root == nil {
+		panic("collection has no root")
+	}
+	return c.root.Load().(*gtreap.Treap)
+}
+
+func (c *Collection) setRoot(root *gtreap.Treap) {
+	if c.root == nil {
+		c.root = &atomic.Value{}
+	}
+	c.root.Store(root)
+}
+
+// Iterator returns an iterator over the Store, starting at the supplied pivot
+// item.
+func (c *Collection) Iterator(pivot gtreap.Item) *gtreap.Iterator {
+	return c.currentRoot().Iterator(pivot)
+}