common/ordered_set.go

   1 package common
   2
   3 import (
   4         "errors"
   5         "sync"
   6 )
   7
   8 // OrderedSet is a set with limited capacity.
   9 // Items are evicted according to their insertion order.
  10 type OrderedSet struct {
  11         capacity int
  12         set      map[interface{}]struct{}
  13         queue    []interface{}
  14         start    int
  15         end      int
  16
  17         lock sync.RWMutex
  18 }
  19
  20 // NewOrderedSet creates an ordered set with given capacity
  21 func NewOrderedSet(capacity int) (*OrderedSet, error) {
  22         if capacity < 1 {
  23                 return nil, errors.New("capacity must be a positive integer")
  24         }
  25
  26         return &OrderedSet{
  27                 capacity: capacity,
  28                 set:      map[interface{}]struct{}{},
  29                 queue:    make([]interface{}, capacity),
  30                 end:      -1,
  31         }, nil
  32 }
  33
  34 // Add inserts items into the set.
  35 // If capacity is reached, oldest items are evicted
  36 func (os *OrderedSet) Add(items ...interface{}) {
  37         os.lock.Lock()
  38         defer os.lock.Unlock()
  39
  40         for _, item := range items {
  41                 if _, ok := os.set[item]; ok {
  42                         continue
  43                 }
  44
  45                 next := (os.end + 1) % os.capacity
  46                 if os.end != -1 && next == os.start {
  47                         delete(os.set, os.queue[os.start])
  48                         os.start = (os.start + 1) % os.capacity
  49                 }
  50                 os.end = next
  51                 os.queue[os.end] = item
  52                 os.set[item] = struct{}{}
  53         }
  54 }
  55
  56 // Has checks if certain items exists in the set
  57 func (os *OrderedSet) Has(item interface{}) bool {
  58         os.lock.RLock()
  59         defer os.lock.RUnlock()
  60
  61         _, ok := os.set[item]
  62         return ok
  63 }
  64
  65 // Size returns the size of the set
  66 func (os *OrderedSet) Size() int {
  67         os.lock.RLock()
  68         defer os.lock.RUnlock()
  69
  70         return len(os.set)
  71 }