Speed up RefcountMap

Signed-off-by: Vladimir Popov <vladimir.popov@xored.com>
networkservicemesh · Nov 25, 2020 · 583c8ab · 583c8ab
1 parent fb1a08f
commit 583c8ab
Show file tree

Hide file tree

Showing 3 changed files with 288 additions and 73 deletions.
diff --git a/pkg/tools/clientmap/refcount.go b/pkg/tools/clientmap/refcount.go
@@ -25,22 +25,19 @@ import (
 
 type refcountClient struct {
 	count int32
+	lock  sync.RWMutex
 	networkservice.NetworkServiceClient
 }
 
 // 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 {
-	m    Map
-	lock sync.RWMutex
+	m Map
 }
 
 // Store sets the value for a key.
 // Increments the refcount
 func (r *RefcountMap) Store(key string, value networkservice.NetworkServiceClient) {
-	r.lock.RLock()
-	defer r.lock.RUnlock()
-
 	client := &refcountClient{
 		count:                1,
 		NetworkServiceClient: value,
@@ -52,12 +49,21 @@ func (r *RefcountMap) Store(key string, value networkservice.NetworkServiceClien
 // 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) {
-	r.lock.RLock()
-	defer r.lock.RUnlock()
-
-	raw, loaded := r.m.LoadOrStore(key, &refcountClient{NetworkServiceClient: value})
+	raw, loaded := r.m.LoadOrStore(key, &refcountClient{
+		count:                1,
+		NetworkServiceClient: value,
+	})
 	client := raw.(*refcountClient)
+	if !loaded {
+		return client.NetworkServiceClient, false
+	}
 
+	client.lock.RLock()
+	defer client.lock.RUnlock()
+
+	if atomic.LoadInt32(&client.count) <= 0 {
+		return r.LoadOrStore(key, value)
+	}
 	atomic.AddInt32(&client.count, 1)
 
 	return client.NetworkServiceClient, loaded
@@ -67,15 +73,18 @@ func (r *RefcountMap) LoadOrStore(key string, value networkservice.NetworkServic
 // value was found in the map.
 // Increments refcount.
 func (r *RefcountMap) Load(key string) (networkservice.NetworkServiceClient, bool) {
-	r.lock.RLock()
-	defer r.lock.RUnlock()
-
 	raw, loaded := r.m.Load(key)
 	if !loaded {
 		return nil, false
 	}
 	client := raw.(*refcountClient)
 
+	client.lock.RLock()
+	defer client.lock.RUnlock()
+
+	if atomic.LoadInt32(&client.count) <= 0 {
+		return nil, false
+	}
 	atomic.AddInt32(&client.count, 1)
 
 	return client.NetworkServiceClient, loaded
@@ -101,16 +110,17 @@ func (r *RefcountMap) LoadAndDelete(key string) (networkservice.NetworkServiceCl
 	}
 	client := raw.(*refcountClient)
 
-	r.lock.Lock()
-	defer r.lock.Unlock()
+	client.lock.Lock()
+	defer client.lock.Unlock()
 
-	switch count := atomic.AddInt32(&client.count, -1); {
-	case count < 0:
+	switch client.count {
+	case 0:
 		return nil, false
-	case count == 0:
+	case 1:
 		r.m.Delete(key)
 		fallthrough
 	default:
+		client.count--
 		return client.NetworkServiceClient, true
 	}
 }

diff --git a/pkg/tools/clientmap/refcount_bench_test.go b/pkg/tools/clientmap/refcount_bench_test.go
@@ -0,0 +1,177 @@
+// 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 (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/networkservicemesh/api/pkg/api/networkservice"
+
+	"github.com/networkservicemesh/sdk/pkg/networkservice/common/null"
+	"github.com/networkservicemesh/sdk/pkg/tools/clientmap"
+)
+
+type refcountClient struct {
+	count int32
+	networkservice.NetworkServiceClient
+}
+
+type syncRefcountMap struct {
+	m    map[string]*refcountClient
+	lock sync.Mutex
+}
+
+func (r *syncRefcountMap) Store(key string, value networkservice.NetworkServiceClient) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	client := &refcountClient{
+		count:                1,
+		NetworkServiceClient: value,
+	}
+	r.m[key] = client
+}
+
+func (r *syncRefcountMap) LoadOrStore(key string, value networkservice.NetworkServiceClient) (networkservice.NetworkServiceClient, bool) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	client, loaded := r.m[key]
+	if !loaded {
+		client = &refcountClient{NetworkServiceClient: value}
+		r.m[key] = client
+	}
+
+	atomic.AddInt32(&client.count, 1)
+
+	return client.NetworkServiceClient, loaded
+}
+
+// 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 *syncRefcountMap) Load(key string) (networkservice.NetworkServiceClient, bool) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	client, loaded := r.m[key]
+	if !loaded {
+		return nil, false
+	}
+
+	atomic.AddInt32(&client.count, 1)
+
+	return client.NetworkServiceClient, loaded
+}
+
+// LoadUnsafe is same as Load but doesn't modify refcount.
+func (r *syncRefcountMap) LoadUnsafe(key string) (networkservice.NetworkServiceClient, bool) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	client, loaded := r.m[key]
+	if !loaded {
+		return nil, false
+	}
+
+	return client.NetworkServiceClient, loaded
+}
+
+// LoadAndDelete returns the value stored in the map for a key, and decrements the refcount. If no value found or the
+// refcount < 0, returns (nil, false). If refcount == 0, deletes the value.
+func (r *syncRefcountMap) LoadAndDelete(key string) (networkservice.NetworkServiceClient, bool) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	client, loaded := r.m[key]
+	if !loaded {
+		return nil, false
+	}
+
+	switch count := atomic.AddInt32(&client.count, -1); {
+	case count < 0:
+		return nil, false
+	case count == 0:
+		delete(r.m, key)
+		fallthrough
+	default:
+		return client.NetworkServiceClient, true
+	}
+}
+
+var ids = []string{
+	"00", "01", "02", "03", "04", "05", "06", "07", "08", "09",
+	"10", "11", "12", "13", "14", "15", "16", "17", "18", "19",
+	"20", "21", "22", "23", "24", "25", "26", "27", "28", "29",
+	"30", "31", "32", "33", "34", "35", "36", "37", "38", "39",
+	"40", "41", "42", "43", "44", "45", "46", "47", "48", "49",
+	"50", "51", "52", "53", "54", "55", "56", "57", "58", "59",
+	"60", "61", "62", "63", "64", "65", "66", "67", "68", "69",
+	"70", "71", "72", "73", "74", "75", "76", "77", "78", "79",
+	"80", "81", "82", "83", "84", "85", "86", "87", "88", "89",
+	"90", "91", "92", "93", "94", "95", "96", "97", "98", "99", "100",
+}
+
+func BenchmarkRefcountMap(b *testing.B) {
+	var m clientmap.RefcountMap
+	b.SetParallelism(1000)
+	b.RunParallel(func(pb *testing.PB) {
+		client := null.NewClient()
+		for i := 0; pb.Next(); i++ {
+			id := ids[i%len(ids)]
+			switch i % 8 {
+			case 0:
+				m.Store(id, client)
+			case 1:
+				client, _ = m.LoadOrStore(id, client)
+			case 2:
+				client, _ = m.Load(id)
+			case 3:
+				client, _ = m.LoadUnsafe(id)
+			case 4, 5, 6, 7:
+				client, _ = m.LoadAndDelete(id)
+			}
+		}
+	})
+}
+
+func BenchmarkSyncRefcountMap(b *testing.B) {
+	m := syncRefcountMap{
+		m: map[string]*refcountClient{},
+	}
+	b.SetParallelism(1000)
+	b.RunParallel(func(pb *testing.PB) {
+		client := null.NewClient()
+		for i := 0; pb.Next(); i++ {
+			id := ids[i%len(ids)]
+			switch i % 8 {
+			case 0:
+				m.Store(id, client)
+			case 1:
+				client, _ = m.LoadOrStore(id, client)
+			case 2:
+				client, _ = m.Load(id)
+			case 3:
+				client, _ = m.LoadUnsafe(id)
+			case 4, 5, 6, 7:
+				client, _ = m.LoadAndDelete(id)
+			}
+		}
+	})
+}