[utils] Implement ShardedMap #248
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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{} { | ||
|
There was a problem hiding this comment. I observe that nothing about this package uses the resulting This means you'd also have to parameterize the constructor to construct caller-specified container shards rather than assuming they're maps. One nice benefit of this approach is that (apart from the need to type-assert the result of There was a problem hiding this comment.
With the current API, the only thing that we would probably want to be asserted is that
Because the caller (me in There was a problem hiding this comment. I think this is a change we could make in a future PR once we have a clear use case in mind (likely just wrapping a |
||
| 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() | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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()) | ||
|
There was a problem hiding this comment. It looks like none of the goroutines in There was a problem hiding this comment. I think adding There was a problem hiding this comment. Appreciate the careful review! |
||
|
|
||
| // 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") | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Increased this to stop a flaky test 😢