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

common/data/treapstore/store_test.go

+// 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())
+                }
+        })
+}