[utils] Implement ShardedMap

go.mod

@@ -17,6 +17,7 @@ require (
 	github.com/lucasjones/reggen v0.0.0-20180717132126-cdb49ff09d77
 	github.com/mitchellh/mapstructure v1.3.3
 	github.com/pkg/errors v0.9.1 // indirect
+	github.com/segmentio/fasthash v1.0.3
 	github.com/stretchr/objx v0.1.1 // indirect
 	github.com/stretchr/testify v1.6.1
 	github.com/tidwall/gjson v1.6.3

go.sum

@@ -226,6 +226,8 @@ github.com/rs/cors v0.0.0-20160617231935-a62a804a8a00/go.mod h1:gFx+x8UowdsKA9Ac
 github.com/rs/xhandler v0.0.0-20160618193221-ed27b6fd6521/go.mod h1:RvLn4FgxWubrpZHtQLnOf6EwhN2hEMusxZOhcW9H3UQ=
 github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g=
 github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/segmentio/fasthash v1.0.3 h1:EI9+KE1EwvMLBWwjpRDc+fEM+prwxDYbslddQGtrmhM=
+github.com/segmentio/fasthash v1.0.3/go.mod h1:waKX8l2N8yckOgmSsXJi7x1ZfdKZ4x7KRMzBtS3oedY=
 github.com/shirou/gopsutil v2.20.5+incompatible h1:tYH07UPoQt0OCQdgWWMgYHy3/a9bcxNpBIysykNIP7I=
 github.com/shirou/gopsutil v2.20.5+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
 github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc=

utils/sharded_map.go

@@ -0,0 +1,85 @@
+// Copyright 2020 Coinbase, Inc.
+//
+// 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 utils
+
+import (
+	"github.com/segmentio/fasthash/fnv1a"
+)
+
+const (
+	// DefaultShards is the default number of shards
+	// to use in ShardedMap.
+	DefaultShards = 256
+)
+
+// shardMapEntry governs access to the shard of
+// the map contained at a particular index.
+type shardMapEntry struct {
+	mutex   *PriorityMutex
+	entries map[string]interface{}
+}
+
+// ShardedMap allows concurrent writes
+// to a map by sharding the map into some
+// number of independently locked subsections.
+type ShardedMap struct {
+	shards []*shardMapEntry
+}
+
+// NewShardedMap creates a new *ShardedMap
+// with some number of shards. The larger the
+// number provided for shards, the less lock
+// contention there will be.
+//
+// As a rule of thumb, shards should usually
+// be set to the concurrency of the caller.
+func NewShardedMap(shards int) *ShardedMap {
+	m := &ShardedMap{
+		shards: make([]*shardMapEntry, shards),
+	}
+
+	for i := 0; i < shards; i++ {
+		m.shards[i] = &shardMapEntry{
+			entries: map[string]interface{}{},
+			mutex:   new(PriorityMutex),
+		}
+	}
+
+	return m
+}
+
+// shardIndex returns the index of the shard
+// that could contain the key.
+func (m *ShardedMap) shardIndex(key string) int {
+	return int(fnv1a.HashString32(key) % uint32(len(m.shards)))
+}
+
+// Lock acquires the lock for a shard that could contain
+// the key. This syntax allows the caller to perform multiple
+// operations while holding the lock for a single shard.
+func (m *ShardedMap) Lock(key string, priority bool) map[string]interface{} {
+	shardIndex := m.shardIndex(key)
+	shard := m.shards[shardIndex]
+	shard.mutex.Lock(priority)
+	return shard.entries
+}
+
+// Unlock releases the lock for a shard that could contain
+// the key.
+func (m *ShardedMap) Unlock(key string) {
+	shardIndex := m.shardIndex(key)
+	shard := m.shards[shardIndex]
+	shard.mutex.Unlock()
+}

utils/sharded_map_test.go

@@ -0,0 +1,69 @@
+// Copyright 2020 Coinbase, Inc.
+//
+// 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 utils
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/sync/errgroup"
+)
+
+func TestShardedMap(t *testing.T) {
+	m := NewShardedMap(2)
+	g, _ := errgroup.WithContext(context.Background())
+
+	// To test locking, we use channels
+	// that will cause deadlock if not executed
+	// concurrently.
+	a := make(chan struct{})
+	b := make(chan struct{})
+
+	g.Go(func() error {
+		s := m.Lock("a", false)
+		assert.Len(t, s, 0)
+		s["test"] = "a"
+		<-a
+		close(b)
+		m.Unlock("a")
+		return nil
+	})
+
+	g.Go(func() error {
+		s := m.Lock("b", false)
+		assert.Len(t, s, 0)
+		s["test"] = "b"
+		close(a)
+		<-b
+		m.Unlock("b")
+		return nil
+	})
+
+	time.Sleep(1 * time.Second)
+	assert.NoError(t, g.Wait())
+
+	// Ensure keys set correctly
+	s := m.Lock("a", false)
+	assert.Len(t, s, 1)
+	assert.Equal(t, s["test"], "a")
+	m.Unlock("a")
+
+	s = m.Lock("b", false)
+	assert.Len(t, s, 1)
+	assert.Equal(t, s["test"], "b")
+	m.Unlock("b")
+}