go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/data/treapstore/store.go

// Copyright 2016 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// 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 "go.chromium.org/luci/gae/impl/memory".
package treapstore

import (
	"fmt"
	"math/rand"
	"sort"
	"sync"

	"github.com/luci/gtreap"
)

// emptyTreap is an empty gtreap.Treap.
var emptyTreap = gtreap.NewTreap(nil)

// Store is a general-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.
//
// A Store's zero value is a valid read-only empty Store, but is not terribly
// useful.
type Store struct {
	collLock  *sync.RWMutex
	colls     map[string]*Collection
	collNames []string // collNames is copy-on-write while holding collLock.
}

// New creates a new read/write Store.
func New() *Store {
	return &Store{
		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 read-only copy of the Store and all of its Collection
// instances. 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{
		collLock:  nil,
		colls:     make(map[string]*Collection, len(s.colls)),
		collNames: s.collNames,
	}

	// Create a read-only Collection for each coll.
	for k, coll := range s.colls {
		newColl := &Collection{
			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.collLock != nil {
		s.collLock.RLock()
		defer s.collLock.RUnlock()
	}

	// Clone our collection names slice.
	if len(s.collNames) == 0 {
		return nil
	}
	return append([]string(nil), s.collNames...)
}

// 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.collLock != nil {
		s.collLock.RLock()
		defer s.collLock.RUnlock()
	}
	return s.colls[name]
}

// CreateCollection returns a Collection with the specified name. If the
// collection already exists, or if s is read-only, CreateCollection will panic.
func (s *Store) CreateCollection(name string, compare gtreap.Compare) *Collection {
	s.assertNotReadOnly()

	s.collLock.Lock()
	defer s.collLock.Unlock()

	// Try to get the existing Collection again now that we hold the write lock.
	if _, ok := s.colls[name]; ok {
		panic(fmt.Errorf("collection %q already exists", name))
	}

	// Create a new read/write Collection.
	coll := &Collection{
		name:     name,
		rootLock: &sync.RWMutex{},
	}
	coll.setRoot(gtreap.NewTreap(compare))

	if s.colls == nil {
		s.colls = make(map[string]*Collection)
	}
	s.colls[name] = coll
	s.collNames = s.insertCollectionName(name)
	return coll
}

// insertCollectionName returns a copy of s.collNames with name inserted
// in its appropriate sorted position.
func (s *Store) insertCollectionName(name string) []string {
	sidx := sort.SearchStrings(s.collNames, name)
	r := make([]string, 0, len(s.collNames)+1)
	return append(append(append(r, s.collNames[:sidx]...), name), s.collNames[sidx:]...)
}

// 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.
//
// A Collection's zero value is a valid read-only empty Collection, which is not
// terribly useful.
type Collection struct {
	name string

	rootLock *sync.RWMutex
	root     *gtreap.Treap
}

func (c *Collection) assertNotReadOnly() {
	if c.IsReadOnly() {
		panic("collection is read-only")
	}
}

// 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.rootLock == nil }

// Min returns the smallest item in the collection.
func (c *Collection) Min() gtreap.Item {
	root := c.currentRoot()
	if root == nil {
		return nil
	}
	return root.Min()
}

// Max returns the largest item in the collection.
func (c *Collection) Max() gtreap.Item {
	root := c.currentRoot()
	if root == nil {
		return nil
	}
	return root.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 {
	root := c.currentRoot()
	if root == nil {
		return nil
	}
	return root.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.assertNotReadOnly()

	// Lock around the entire Upsert operation to serialize Puts.
	priority := rand.Int()
	c.rootLock.Lock()
	c.root = c.root.Upsert(i, priority)
	c.rootLock.Unlock()
}

// Delete deletes an item from the Collection, if such an item exists.
func (c *Collection) Delete(i gtreap.Item) {
	c.assertNotReadOnly()

	c.rootLock.Lock()
	c.root = c.root.Delete(i)
	c.rootLock.Unlock()
}

func (c *Collection) currentRoot() *gtreap.Treap {
	if c.rootLock != nil {
		c.rootLock.RLock()
		defer c.rootLock.RUnlock()
	}
	return c.root
}

func (c *Collection) setRoot(root *gtreap.Treap) {
	if c.rootLock != nil {
		c.rootLock.Lock()
		defer c.rootLock.Unlock()
	}
	c.root = root
}

// VisitAscend traverses the Collection ascendingly, invoking visitor for each
// visited item.
//
// If visitor returns false, iteration will stop prematurely.
//
// VisitAscend is a more efficient traversal than using an Iterator, and is
// useful in times when entry-by-entry iteration is not required.
func (c *Collection) VisitAscend(pivot gtreap.Item, visitor gtreap.ItemVisitor) {
	c.currentRoot().VisitAscend(pivot, visitor)
}

// Iterator returns an iterator over the Collection, starting at the supplied
// pivot item.
func (c *Collection) Iterator(pivot gtreap.Item) *gtreap.Iterator {
	root := c.currentRoot()
	if root == nil {
		root = emptyTreap
	}
	return root.Iterator(pivot)
}