Rework clientmap.RefcountMap, add tests, add bench tests.

pkg/tools/clientmap/README.md

@@ -0,0 +1,23 @@
+# clientmap.Map
+
+It is a `sync.Map` typed for the `string` keys and `networkservice.NetworkServiceClient` values.
+
+# clientmap.RefcountMap
+
+It is a `clientmap.Map` wrapped with refcounting:
+```
+Store, LoadOrStore (store)  -->  count = 1
+Load, LoadOrStore (load)    -->  count += 1
+LoadAndDelete, Delete       -->  count -= 1
+```
+if count becomes 0, value deletes.
+
+## Performance
+
+`clientmap.RefcountMap` is a very thin wrapper, it doesn't hardly affect the `clientmap.Map` performance:
+```
+BenchmarkMap
+BenchmarkMap-4           	20850049	        54.7 ns/op
+BenchmarkRefcountMap
+BenchmarkRefcountMap-4   	14408266	        76.3 ns/op
+```

pkg/tools/clientmap/bench_test.go

@@ -0,0 +1,80 @@
+// Copyright (c) 2020 Doc.ai and/or its affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// 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 clientmap_test
+
+import (
+	"testing"
+
+	"github.com/networkservicemesh/sdk/pkg/networkservice/common/null"
+	"github.com/networkservicemesh/sdk/pkg/tools/clientmap"
+)
+
+const (
+	parallelCount = 20
+)
+
+var ids = []string{
+	"00", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10",
+}
+
+func BenchmarkMap(b *testing.B) {
+	var m clientmap.Map
+	client := null.NewClient()
+	b.SetParallelism(parallelCount)
+
+	b.ResetTimer()
+
+	b.RunParallel(func(pb *testing.PB) {
+		for i := 0; pb.Next(); i++ {
+			id := ids[i%len(ids)]
+			switch i % 6 {
+			case 0:
+				m.Store(id, client)
+			case 1:
+				_, _ = m.LoadOrStore(id, client)
+			case 2:
+				_, _ = m.Load(id)
+			case 3, 4, 5:
+				_, _ = m.LoadAndDelete(id)
+			}
+		}
+	})
+}
+
+func BenchmarkRefcountMap(b *testing.B) {
+	var m clientmap.RefcountMap
+	client := null.NewClient()
+	b.SetParallelism(parallelCount)
+
+	b.ResetTimer()
+
+	b.RunParallel(func(pb *testing.PB) {
+		for i := 0; pb.Next(); i++ {
+			id := ids[i%len(ids)]
+			switch i % 6 {
+			case 0:
+				m.Store(id, client)
+			case 1:
+				_, _ = m.LoadOrStore(id, client)
+			case 2:
+				_, _ = m.Load(id)
+			case 3, 4, 5:
+				_, _, _ = m.LoadAndDelete(id)
+			}
+		}
+	})
+}

pkg/tools/clientmap/refcount.go

@@ -1,3 +1,5 @@
+// Copyright (c) 2020 Doc.ai and/or its affiliates.
+//
 // Copyright (c) 2020 Cisco and/or its affiliates.
 //
 // SPDX-License-Identifier: Apache-2.0
@@ -17,73 +19,114 @@
 package clientmap
 
 import (
-	"context"
-	"sync/atomic"
+	"sync"
 
-	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/networkservicemesh/api/pkg/api/networkservice"
-	"google.golang.org/grpc"
 )
 
-type refcountClient struct {
-	count  int32
-	client networkservice.NetworkServiceClient
+type entry struct {
+	count int
+	lock  sync.Mutex
+	value networkservice.NetworkServiceClient
 }
 
-func (r *refcountClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
-	return r.client.Request(ctx, request)
+// RefcountMap is a sync.Map wrapped with refcounting. Each Store, Load, LoadOrStore increments the count, each
+// LoadAndDelete, Delete decrements the count. When the count == 1, LoadAndDelete, Delete deletes.
+type RefcountMap struct {
+	m sync.Map
 }
 
-func (r *refcountClient) Close(ctx context.Context, conn *networkservice.Connection, opts ...grpc.CallOption) (*empty.Empty, error) {
-	return r.client.Close(ctx, conn)
+// Store sets the value for a key
+// count = 1.
+func (m *RefcountMap) Store(key string, newValue networkservice.NetworkServiceClient) {
+	m.m.Store(key, &entry{
+		count: 1,
+		value: newValue,
+	})
 }
 
-// RefcountMap - clientmap.Map wrapped with refcounting.  Each Store,Load,LoadOrStore increments the count,
-//               each Delete decrements the count.  When the count is zero, Delete deletes.
-type RefcountMap struct {
-	Map
-}
+// LoadOrStore returns the existing value for the key if present. Otherwise, it stores and returns the given value. The
+// loaded result is true if the value was loaded, false if stored.
+// store  -->  count = 1
+// load   -->  count += 1
+func (m *RefcountMap) LoadOrStore(key string, newValue networkservice.NetworkServiceClient) (value networkservice.NetworkServiceClient, loaded bool) {
+	var raw interface{}
+	raw, loaded = m.m.LoadOrStore(key, &entry{
+		count: 1,
+		value: newValue,
+	})
+	entry := raw.(*entry)
+	if !loaded {
+		return entry.value, false
+	}
 
-// Store sets the value for a key.
-// Increments the refcount
-func (r *RefcountMap) Store(key string, value networkservice.NetworkServiceClient) {
-	client := &refcountClient{client: value}
-	r.Map.Store(key, client)
-	atomic.AddInt32(&client.count, 1)
-}
+	entry.lock.Lock()
 
-// LoadOrStore returns the existing value for the key if present. Otherwise, it stores and returns the given value. The loaded result is true if the value was loaded, false if stored.
-// Increments the refcount.
-func (r *RefcountMap) LoadOrStore(key string, value networkservice.NetworkServiceClient) (networkservice.NetworkServiceClient, bool) {
-	client := &refcountClient{client: value}
-	rv, loaded := r.Map.LoadOrStore(key, client)
-	if client, ok := rv.(*refcountClient); ok {
-		atomic.AddInt32(&client.count, 1)
-		return client.client, loaded
+	if entry.count == 0 {
+		entry.lock.Unlock()
+		return m.LoadOrStore(key, newValue)
 	}
-	return rv, loaded
+	entry.count++
+
+	entry.lock.Unlock()
+
+	return entry.value, true
 }
 
-// Load returns the value stored in the map for a key, or nil if no value is present. The ok result indicates whether value was found in the map.
-// Increments refcount.
-func (r *RefcountMap) Load(key string) (networkservice.NetworkServiceClient, bool) {
-	rv, loaded := r.Map.Load(key)
-	if client, ok := rv.(*refcountClient); ok {
-		atomic.AddInt32(&client.count, 1)
-		return client.client, loaded
+// Load returns the value stored in the map for a key, or nil if no value is present. The loaded result indicates
+// whether value was found in the map.
+// count += 1
+func (m *RefcountMap) Load(key string) (value networkservice.NetworkServiceClient, loaded bool) {
+	var raw interface{}
+	raw, loaded = m.m.Load(key)
+	if !loaded {
+		return nil, false
 	}
-	return rv, loaded
+	entry := raw.(*entry)
+
+	entry.lock.Lock()
+	defer entry.lock.Unlock()
+
+	if entry.count == 0 {
+		return nil, false
+	}
+	entry.count++
+
+	return entry.value, true
 }
 
-// Delete decrements the refcount and deletes the value for a key if refcount is zero. Returns true if value is no (more) present.
-func (r *RefcountMap) Delete(key string) bool {
-	rv, loaded := r.Map.Load(key)
+// LoadAndDelete is trying to delete the value for a key, returning the previous value if any. The loaded result
+// reports whether the key was present.
+// count == 1  -->  delete
+// count > 1   -->  count -= 1
+func (m *RefcountMap) LoadAndDelete(key string) (value networkservice.NetworkServiceClient, loaded, deleted bool) {
+	var raw interface{}
+	raw, loaded = m.m.Load(key)
 	if !loaded {
-		return true
+		return nil, false, true
+	}
+	entry := raw.(*entry)
+
+	entry.lock.Lock()
+	defer entry.lock.Unlock()
+
+	if entry.count == 0 {
+		return nil, false, true
 	}
-	if client, ok := rv.(*refcountClient); ok && atomic.AddInt32(&client.count, -1) == 0 {
-		r.Map.Delete(key)
-		return true
+	entry.count--
+
+	if entry.count == 0 {
+		m.m.Delete(key)
+		deleted = true
 	}
-	return false
+
+	return entry.value, true, deleted
+}
+
+// Delete is trying to delete the value for a key.
+// count == 1  -->  delete
+// count > 1   -->  count -= 1
+func (m *RefcountMap) Delete(key string) (deleted bool) {
+	_, _, deleted = m.LoadAndDelete(key)
+	return deleted
 }