From 090d0aee0b3739794db34af069d199e06e36b6ee Mon Sep 17 00:00:00 2001
From: Cameron Sparr <cameronsparr@gmail.com>
Date: Fri, 22 Jul 2016 13:20:48 +0100
Subject: [PATCH] Allow for configurable number of shards

closes #17
---
 concurrent_map.go            | 24 +++++++++++++++---------
 concurrent_map_bench_test.go |  6 +++---
 concurrent_map_test.go       |  6 +-----
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/concurrent_map.go b/concurrent_map.go
index 2993038..d822985 100644
--- a/concurrent_map.go
+++ b/concurrent_map.go
@@ -5,10 +5,10 @@ import (
 	"sync"
 )
 
-var SHARD_COUNT = 32
+const SHARD_COUNT = 32
 
 // A "thread" safe map of type string:Anything.
-// To avoid lock bottlenecks this map is dived to several (SHARD_COUNT) map shards.
+// To avoid lock bottlenecks this map is dived to several (len(m)) map shards.
 type ConcurrentMap []*ConcurrentMapShared
 
 // A "thread" safe string to anything map.
@@ -18,9 +18,15 @@ type ConcurrentMapShared struct {
 }
 
 // Creates a new concurrent map.
-func New() ConcurrentMap {
-	m := make(ConcurrentMap, SHARD_COUNT)
-	for i := 0; i < SHARD_COUNT; i++ {
+func New(shardCount ...int) ConcurrentMap {
+	var nShards int
+	if len(shardCount) > 0 {
+		nShards = shardCount[0]
+	} else {
+		nShards = SHARD_COUNT
+	}
+	m := make(ConcurrentMap, nShards)
+	for i := 0; i < nShards; i++ {
 		m[i] = &ConcurrentMapShared{items: make(map[string]interface{})}
 	}
 	return m
@@ -94,7 +100,7 @@ func (m ConcurrentMap) Get(key string) (interface{}, bool) {
 // Returns the number of elements within the map.
 func (m ConcurrentMap) Count() int {
 	count := 0
-	for i := 0; i < SHARD_COUNT; i++ {
+	for i := 0; i < len(m); i++ {
 		shard := m[i]
 		shard.RLock()
 		count += len(shard.items)
@@ -141,7 +147,7 @@ func (m ConcurrentMap) Iter() <-chan Tuple {
 	ch := make(chan Tuple)
 	go func() {
 		wg := sync.WaitGroup{}
-		wg.Add(SHARD_COUNT)
+		wg.Add(len(m))
 		// Foreach shard.
 		for _, shard := range m {
 			go func(shard *ConcurrentMapShared) {
@@ -165,7 +171,7 @@ func (m ConcurrentMap) IterBuffered() <-chan Tuple {
 	ch := make(chan Tuple, m.Count())
 	go func() {
 		wg := sync.WaitGroup{}
-		wg.Add(SHARD_COUNT)
+		wg.Add(len(m))
 		// Foreach shard.
 		for _, shard := range m {
 			go func(shard *ConcurrentMapShared) {
@@ -203,7 +209,7 @@ func (m ConcurrentMap) Keys() []string {
 	go func() {
 		// Foreach shard.
 		wg := sync.WaitGroup{}
-		wg.Add(SHARD_COUNT)
+		wg.Add(len(m))
 		for _, shard := range m {
 			go func(shard *ConcurrentMapShared) {
 				// Foreach key, value pair.
diff --git a/concurrent_map_bench_test.go b/concurrent_map_bench_test.go
index 47cb8d8..7b8c0e6 100644
--- a/concurrent_map_bench_test.go
+++ b/concurrent_map_bench_test.go
@@ -3,6 +3,8 @@ package cmap
 import "testing"
 import "strconv"
 
+var m ConcurrentMap
+
 func BenchmarkItems(b *testing.B) {
 	m := New()
 
@@ -177,10 +179,8 @@ func GetSet(m ConcurrentMap, finished chan struct{}) (set func(key, value string
 }
 
 func runWithShards(bench func(b *testing.B), b *testing.B, shardsCount int) {
-	oldShardsCount := SHARD_COUNT
-	SHARD_COUNT = shardsCount
+	m = New(shardsCount)
 	bench(b)
-	SHARD_COUNT = oldShardsCount
 }
 
 func BenchmarkKeys(b *testing.B) {
diff --git a/concurrent_map_test.go b/concurrent_map_test.go
index 3f0411e..cc48a3f 100644
--- a/concurrent_map_test.go
+++ b/concurrent_map_test.go
@@ -272,12 +272,8 @@ func TestConcurrent(t *testing.T) {
 }
 
 func TestJsonMarshal(t *testing.T) {
-	SHARD_COUNT = 2
-	defer func() {
-		SHARD_COUNT = 32
-	}()
 	expected := "{\"a\":1,\"b\":2}"
-	m := New()
+	m := New(2)
 	m.Set("a", 1)
 	m.Set("b", 2)
 	j, err := json.Marshal(m)