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

common/data/treapstore/store_test.go

Index: common/data/treapstore/store_test.go
diff --git a/common/data/treapstore/store_test.go b/common/data/treapstore/store_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..b5219d92e0f4d87a00a5269ee9fbc89333c04777
--- /dev/null
+++ b/common/data/treapstore/store_test.go
@@ -0,0 +1,180 @@
+// 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
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/luci/gtreap"
+
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+func newStringStore() *Store {
+	return New(func(a, b interface{}) int {
+		return strings.Compare(a.(string), b.(string))
+	})
+}
+
+func iterAll(it *gtreap.Iterator) []string {
+	all := []string{}
+	for {
+		v, ok := it.Next()
+		if !ok {
+			return all
+		}
+		all = append(all, v.(string))
+	}
+}
+
+func TestStore(t *testing.T) {
+	t.Parallel()
+
+	Convey(`Testing a string Store`, t, func() {
+		st := newStringStore()
+		st.GetOrCreateCollection("")
+
+		Convey(`When empty`, func() {
+			checkEmpty := func(c *Collection) {
+				So(c.Get("foo"), ShouldBeNil)
+				So(iterAll(c.Iterator("")), ShouldResemble, []string{})
+			}
+
+			// Check the basic Store.
+			checkEmpty(st.GetCollection(""))
+
+			// Take a snapshot, then mutate the base Store. Assert that the snapshot
+			// is still empty.
+			snap := st.Snapshot()
+			st.GetCollection("").Put("foo")
+			checkEmpty(snap.GetCollection(""))
+		})
+
+		Convey(`With keys`, func() {
+			coll := st.GetCollection("")
+			coll.Put("a")
+			coll.Put("b")
+			coll.Put("z")
+
+			checkKeys := func(coll *Collection, keys ...string) {
+				for _, k := range keys {
+					So(coll.Get(k), ShouldEqual, k)
+				}
+
+				So(iterAll(coll.Iterator("")), ShouldResemble, keys)
+				for i, k := range keys {
+					So(iterAll(coll.Iterator(k)), ShouldResemble, keys[i:])
+					So(iterAll(coll.Iterator(k+"1")), ShouldResemble, keys[i+1:])
+				}
+			}
+			checkKeys(coll, "a", "b", "z")
+
+			// Take a snapshot, then mutate the base Store. Assert that the snapshot
+			// is still empty.
+			snap := st.Snapshot()
+			snapColl := snap.GetCollection("")
+			coll.Put("foo")
+			checkKeys(snapColl, "a", "b", "z")
+			checkKeys(coll, "a", "b", "foo", "z")
+		})
+	})
+}
+
+// TestStoreParallel performs several rounds of parallel accesses. Each round
+// takes a snapshot of the "head" Store, then simultaneusly dispatches a round
+// of parallel writes agains the "head" store, reads against the snapshot, and
+// reads against the "head" store.
+//
+// This is meant to be run with "-race" to trigger on race conditions.
+func TestStoreParallel(t *testing.T) {
+	t.Parallel()
+
+	Convey(`Testing a string Store for parallel access.`, t, func() {
+		const (
+			readers = 128
+			writers = 16
+			rounds  = 8
+		)
+
+		head := newStringStore()
+		var snaps []*Store
+
+		// Dispatch readers.
+		doReads := func() int {
+			readDoneC := make(chan int, readers)
+			for i := 0; i < readers; i++ {
+				go func() {
+					var (
+						doneC = make(chan int, 1+len(snaps))
+					)
+
+					// "head"
+					go func() {
+						doneC <- len(iterAll(head.GetCollection("").Iterator("")))
+					}()
+
+					// "snap"
+					for _, snap := range snaps {
+						go func(snap *Store) {
+							doneC <- len(iterAll(snap.GetCollection("").Iterator("")))
+						}(snap)
+					}
+
+					total := 0
+					for i := 0; i < 1+len(snaps); i++ {
+						total += <-doneC
+					}
+					readDoneC <- total
+				}()
+			}
+
+			total := 0
+			for i := 0; i < readers; i++ {
+				total += <-readDoneC
+			}
+			return total
+		}
+
+		// Dispatch writers.
+		doWrite := func(base string) {
+			writeDoneC := make(chan struct{}, writers)
+			for i := 0; i < writers; i++ {
+				go func(idx int) {
+					head.GetOrCreateCollection("").Put(fmt.Sprintf("%s.%d", base, idx))
+					writeDoneC <- struct{}{}
+				}(i)
+			}
+
+			for i := 0; i < writers; i++ {
+				<-writeDoneC
+			}
+		}
+
+		// Main loop.
+		for i := 0; i < rounds; i++ {
+			writeDoneC := make(chan struct{})
+			readDoneC := make(chan int)
+			go func() {
+				doWrite(strconv.Itoa(i))
+				close(writeDoneC)
+			}()
+			if i > 0 {
+				// The first round has to actually create the Collection.
+				go func() {
+					readDoneC <- doReads()
+				}()
+			}
+
+			<-writeDoneC
+			if i > 0 {
+				So(<-readDoneC, ShouldBeGreaterThan, 0)
+			}
+			snaps = append(snaps, head.Snapshot())
+		}
+	})
+}