pkg/utils/threadsafemultimap.go (77 lines of code) (raw):

// ------------------------------------------------------------------------------------------- // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. See License.txt in the project root for license information. // -------------------------------------------------------------------------------------------- package utils import ( "sync" mapset "github.com/deckarep/golang-set" ) // ThreadsafeMultiMap is a thread safe implementation of a multimap. type ThreadsafeMultiMap interface { Insert(key interface{}, value interface{}) Clear(key interface{}) Erase(key interface{}) bool EraseValue(value interface{}) bool ContainsPair(key interface{}, value interface{}) bool ContainsValue(value interface{}) bool } type threadsafeMultiMap struct { // TODO: Evaluate if sync.Map is needed here // for our usage pattern. sync.RWMutex v map[interface{}]mapset.Set } // NewThreadsafeMultimap creates a ThreadsafeMultiMap. func NewThreadsafeMultimap() ThreadsafeMultiMap { return &threadsafeMultiMap{ v: make(map[interface{}]mapset.Set), } } // Insert inserts a key value pair into the multimap. func (m *threadsafeMultiMap) Insert(key interface{}, value interface{}) { m.Lock() defer m.Unlock() if m.v[key] == nil { m.v[key] = mapset.NewSet() } m.v[key].Add(value) } // Clear removes all the values associated with a key. func (m *threadsafeMultiMap) Clear(key interface{}) { m.Lock() defer m.Unlock() if m.v[key] != nil { m.v[key].Clear() } } // Erase removes a key from the multimap. func (m *threadsafeMultiMap) Erase(key interface{}) bool { m.Lock() defer m.Unlock() _, exists := m.v[key] if exists { delete(m.v, key) return true } return false } // EraseValue removes all the occurrences of a particular value from the multimap. func (m *threadsafeMultiMap) EraseValue(value interface{}) bool { m.Lock() defer m.Unlock() erased := false for i := range m.v { if m.v[i] != nil && m.v[i].Contains(value) { erased = true m.v[i].Remove(value) } } return erased } // ContainsPair checks if a particular key value pair exists in the multimap. func (m *threadsafeMultiMap) ContainsPair(key interface{}, value interface{}) bool { m.RLock() defer m.RUnlock() if m.v[key] != nil && m.v[key].Contains(value) { return true } return false } // ContainsValue checks if a particular value exists in the multimap. func (m *threadsafeMultiMap) ContainsValue(value interface{}) bool { m.RLock() defer m.RUnlock() for i := range m.v { if m.v[i] != nil && m.v[i].Contains(value) { return true } } return false }