diff --git a/set/item.go b/set/item.go new file mode 100644 index 0000000..bcb2a1b --- /dev/null +++ b/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() +} diff --git a/set/set.go b/set/set.go new file mode 100644 index 0000000..06ccec3 --- /dev/null +++ b/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) +} diff --git a/set/set_test.go b/set/set_test.go new file mode 100644 index 0000000..7b55dc8 --- /dev/null +++ b/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") + } +}