set: Generalized Set type

set/item.go

@@ -0,0 +1,42 @@
+package set
+
+import "time"
+
+// Interface for an item storeable in the set
+type Item interface {
+	Key() string
+	Value() interface{}
+}
+
+type StringItem string
+
+func (item StringItem) Key() string {
+	return string(item)
+}
+
+func (item StringItem) Value() interface{} {
+	return string(item)
+}
+
+func Expire(item Item, d time.Duration) Item {
+	return &ExpiringItem{
+		Item: item,
+		Time: time.Now().Add(d),
+	}
+}
+
+type ExpiringItem struct {
+	Item
+	time.Time
+}
+
+func (item *ExpiringItem) Expired() bool {
+	return time.Now().After(item.Time)
+}
+
+func (item *ExpiringItem) Value() interface{} {
+	if item.Expired() {
+		return nil
+	}
+	return item.Item.Value()
+}

set/set.go

@@ -0,0 +1,191 @@
+package set
+
+import (
+	"errors"
+	"strings"
+	"sync"
+)
+
+// Returned when an added key already exists in the set.
+var ErrCollision = errors.New("key already exists")
+
+// Returned when a requested item does not exist in the set.
+var ErrMissing = errors.New("item does not exist")
+
+// Returned when a nil item is added. Nil values are considered expired and invalid.
+var ErrNil = errors.New("item value must not be nil")
+
+type Set struct {
+	sync.RWMutex
+	lookup    map[string]Item
+	normalize func(string) string
+}
+
+// New creates a new set with case-insensitive keys
+func New() *Set {
+	return &Set{
+		lookup:    map[string]Item{},
+		normalize: normalize,
+	}
+}
+
+// Clear removes all items and returns the number removed.
+func (s *Set) Clear() int {
+	s.Lock()
+	n := len(s.lookup)
+	s.lookup = map[string]Item{}
+	s.Unlock()
+	return n
+}
+
+// Len returns the size of the set right now.
+func (s *Set) Len() int {
+	s.RLock()
+	defer s.RUnlock()
+	return len(s.lookup)
+}
+
+// In checks if an item exists in this set.
+func (s *Set) In(key string) bool {
+	key = s.normalize(key)
+	s.RLock()
+	item, ok := s.lookup[key]
+	s.RUnlock()
+	if ok && item.Value() == nil {
+		s.cleanup(key)
+	}
+	return ok
+}
+
+// Get returns an item with the given key.
+func (s *Set) Get(key string) (Item, error) {
+	key = s.normalize(key)
+	s.RLock()
+	item, ok := s.lookup[key]
+	s.RUnlock()
+
+	if !ok {
+		return nil, ErrMissing
+	}
+	if item.Value() == nil {
+		s.cleanup(key)
+	}
+
+	return item, nil
+}
+
+// Remove potentially expired key (normalized).
+func (s *Set) cleanup(key string) {
+	s.Lock()
+	item, ok := s.lookup[key]
+	if ok && item == nil {
+		delete(s.lookup, key)
+	}
+	s.Unlock()
+}
+
+// Add item to this set if it does not exist already.
+func (s *Set) Add(item Item) error {
+	if item.Value() == nil {
+		return ErrNil
+	}
+	key := s.normalize(item.Key())
+
+	s.Lock()
+	defer s.Unlock()
+
+	oldItem, found := s.lookup[key]
+	if found && oldItem.Value() != nil {
+		return ErrCollision
+	}
+
+	s.lookup[key] = item
+	return nil
+}
+
+// Remove item from this set.
+func (s *Set) Remove(key string) error {
+	key = s.normalize(key)
+
+	s.Lock()
+	defer s.Unlock()
+
+	_, found := s.lookup[key]
+	if !found {
+		return ErrMissing
+	}
+	delete(s.lookup, key)
+	return nil
+}
+
+// Replace oldKey with a new item, which might be a new key.
+// Can be used to rename items.
+func (s *Set) Replace(oldKey string, item Item) error {
+	if item.Value() == nil {
+		return ErrNil
+	}
+	newKey := s.normalize(item.Key())
+	oldKey = s.normalize(oldKey)
+
+	s.Lock()
+	defer s.Unlock()
+
+	if newKey != oldKey {
+		// Check if it already exists
+		_, found := s.lookup[newKey]
+		if found {
+			return ErrCollision
+		}
+
+		// Remove oldKey
+		_, found = s.lookup[oldKey]
+		if !found {
+			return ErrMissing
+		}
+		delete(s.lookup, oldKey)
+	}
+
+	// Add new item
+	s.lookup[newKey] = item
+
+	return nil
+}
+
+// Each loops over every item while holding a read lock and applies fn to each
+// element.
+func (s *Set) Each(fn func(item Item)) {
+	cleanup := []string{}
+	s.RLock()
+	for key, item := range s.lookup {
+		if item.Value() == nil {
+			cleanup = append(cleanup, key)
+			continue
+		}
+		fn(item)
+	}
+	s.RUnlock()
+
+	if len(cleanup) == 0 {
+		return
+	}
+	for _, key := range cleanup {
+		s.cleanup(key)
+	}
+}
+
+// ListPrefix returns a list of items with a prefix, normalized.
+// TODO: Add trie for efficient prefix lookup
+func (s *Set) ListPrefix(prefix string) []Item {
+	r := []Item{}
+	prefix = s.normalize(prefix)
+
+	s.Each(func(item Item) {
+		r = append(r, item)
+	})
+
+	return r
+}
+
+func normalize(key string) string {
+	return strings.ToLower(key)
+}

set/set_test.go

@@ -0,0 +1,64 @@
+package set
+
+import (
+	"testing"
+	"time"
+)
+
+func TestSetExpiring(t *testing.T) {
+	s := New()
+	if s.In("foo") {
+		t.Error("matched before set.")
+	}
+
+	if err := s.Add(StringItem("foo")); err != nil {
+		t.Fatalf("failed to add foo: %s", err)
+	}
+	if !s.In("foo") {
+		t.Errorf("not matched after set")
+	}
+	if s.Len() != 1 {
+		t.Error("not len 1 after set")
+	}
+
+	item := &ExpiringItem{nil, time.Now().Add(-time.Nanosecond * 1)}
+	if !item.Expired() {
+		t.Errorf("ExpiringItem a nanosec ago is not expiring")
+	}
+
+	item = &ExpiringItem{nil, time.Now().Add(time.Minute * 2)}
+	if item.Expired() {
+		t.Errorf("ExpiringItem in 2 minutes is expiring now")
+	}
+
+	item = Expire(StringItem("bar"), time.Minute*2).(*ExpiringItem)
+	until := item.Time
+	if !until.After(time.Now().Add(time.Minute*1)) || !until.Before(time.Now().Add(time.Minute*3)) {
+		t.Errorf("until is not a minute after %s: %s", time.Now(), until)
+	}
+	if item.Value() == nil {
+		t.Errorf("bar expired immediately")
+	}
+	if err := s.Add(item); err != nil {
+		t.Fatalf("failed to add item: %s", err)
+	}
+	_, ok := s.lookup["bar"]
+	if !ok {
+		t.Fatalf("expired bar added to lookup")
+	}
+	s.lookup["bar"] = item
+
+	if !s.In("bar") {
+		t.Errorf("not matched after timed set")
+	}
+	if s.Len() != 2 {
+		t.Error("not len 2 after set")
+	}
+
+	if err := s.Replace("bar", Expire(StringItem("bar"), time.Minute*5)); err != nil {
+		t.Fatalf("failed to add bar: %s", err)
+	}
+	if !s.In("bar") {
+		t.Error("failed to match before expiry")
+	}
+}