diff --git a/pkg/tools/clientmap/refcount.go b/pkg/tools/clientmap/refcount.go index e159e54570..513811eac4 100644 --- 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 new file mode 100644 index 0000000000..d8f250a2fd --- /dev/null +++ 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) + } + } + }) +} diff --git a/pkg/tools/clientmap/refcount_test.go b/pkg/tools/clientmap/refcount_test.go index f08d2f0c0b..3328e1843d 100644 --- a/pkg/tools/clientmap/refcount_test.go +++ b/pkg/tools/clientmap/refcount_test.go @@ -95,67 +95,95 @@ type sample struct { validator func(bool, []bool) bool } -func TestRefcountMap_Actions(t *testing.T) { - samples := []*sample{ - { - name: "LoadOrStore + LoadOrStore", - // 0 1 -> loaded { false, true } - // 1 0 -> loaded { true, false } - actions: []action{loadOrStore, loadOrStore}, - validator: func(loaded bool, results []bool) bool { - return loaded && !results[0] && results[1] || - loaded && results[0] && !results[1] - }, +var samples = []*sample{ + { + name: "LoadOrStore + LoadOrStore", + // 0 1 -> loaded { false, true } + // 1 0 -> loaded { true, false } + actions: []action{loadOrStore, loadOrStore}, + validator: func(loaded bool, results []bool) bool { + return loaded && !results[0] && results[1] || + loaded && results[0] && !results[1] }, - { - name: "LoadOrStore + LoadAndDelete", - // 0 1 -> !loaded { false, true } - // 1 0 -> loaded { false, false } - actions: []action{loadOrStore, loadAndDelete}, - validator: func(loaded bool, results []bool) bool { - return !loaded && !results[0] && results[1] || - loaded && !results[0] && !results[1] - }, + }, + { + name: "LoadOrStore + LoadAndDelete", + // 0 1 -> !loaded { false, true } + // 1 0 -> loaded { false, false } + actions: []action{loadOrStore, loadAndDelete}, + validator: func(loaded bool, results []bool) bool { + return !loaded && !results[0] && results[1] || + loaded && !results[0] && !results[1] }, - { - name: "LoadAndDelete + LoadAndDelete", - withValue: true, - // 0 1 -> !loaded { true, false } - // 1 0 -> !loaded { false, true } - actions: []action{loadAndDelete, loadAndDelete}, - validator: func(loaded bool, results []bool) bool { - return !loaded && results[0] && !results[1] || - !loaded && !results[0] && results[1] - }, + }, + { + name: "LoadOrStore + LoadAndDelete with value", + withValue: true, + // 0 1 -> loaded { true, true } + // 1 0 -> loaded { false, true } + actions: []action{loadOrStore, loadAndDelete}, + validator: func(loaded bool, results []bool) bool { + return loaded && results[0] && results[1] || + loaded && !results[0] && results[1] }, - { - name: "LoadAndDelete + Load", - withValue: true, - // 0 1 -> !loaded { true, false } - // 1 0 -> loaded { true, true } - actions: []action{loadAndDelete, load}, - validator: func(loaded bool, results []bool) bool { - return !loaded && results[0] && !results[1] || - loaded && results[0] && results[1] - }, + }, + { + name: "LoadAndDelete + LoadAndDelete", + withValue: true, + // 0 1 -> !loaded { true, false } + // 1 0 -> !loaded { false, true } + actions: []action{loadAndDelete, loadAndDelete}, + validator: func(loaded bool, results []bool) bool { + return !loaded && results[0] && !results[1] || + !loaded && !results[0] && results[1] }, - { - name: "LoadAndDelete + LoadAndDelete + Load", - withValue: true, - // 0 1 2 -> !loaded { true, false, false } - // 0 2 1 -> !loaded { true, false, false } (same as ^) - // 1 0 2 -> !loaded { false, true, false } - // 1 2 0 -> !loaded { true, true, true } - // 2 0 1 -> !loaded { false, true, false } (same as ^^) - // 2 1 0 -> !loaded { true, true, true } (same as ^^) - actions: []action{loadAndDelete, loadAndDelete, load}, - validator: func(loaded bool, results []bool) bool { - return !loaded && results[0] && !results[1] && !results[2] || - !loaded && !results[0] && results[1] && !results[2] || - !loaded && results[0] && results[1] && results[2] - }, + }, + { + name: "LoadAndDelete + Load", + withValue: true, + // 0 1 -> !loaded { true, false } + // 1 0 -> loaded { true, true } + actions: []action{loadAndDelete, load}, + validator: func(loaded bool, results []bool) bool { + return !loaded && results[0] && !results[1] || + loaded && results[0] && results[1] }, - } + }, + { + name: "LoadAndDelete + LoadAndDelete + Load", + withValue: true, + // 0 1 2 -> !loaded { true, false, false } + // 0 2 1 -> !loaded { true, false, false } (same as ^) + // 1 0 2 -> !loaded { false, true, false } + // 1 2 0 -> !loaded { true, true, true } + // 2 0 1 -> !loaded { false, true, false } (same as ^^) + // 2 1 0 -> !loaded { true, true, true } (same as ^^) + actions: []action{loadAndDelete, loadAndDelete, load}, + validator: func(loaded bool, results []bool) bool { + return !loaded && results[0] && !results[1] && !results[2] || + !loaded && !results[0] && results[1] && !results[2] || + !loaded && results[0] && results[1] && results[2] + }, + }, + { + name: "LoadOrStore + LoadOrStore + LoadAndDelete", + withValue: true, + // 0 1 2 -> loaded { true, true, true } + // 0 2 1 -> loaded { true, true, true } (same as ^) + // 1 0 2 -> loaded { true, true, true } (same as ^) + // 1 2 0 -> loaded { false, true, true } + // 2 0 1 -> loaded { true, true, true } (same as ^^) + // 2 1 0 -> loaded { true, false, true } + actions: []action{loadOrStore, loadOrStore, loadAndDelete}, + validator: func(loaded bool, results []bool) bool { + return loaded && results[0] && results[1] && results[2] || + loaded && !results[0] && results[1] && results[2] || + loaded && results[0] && !results[1] && results[2] + }, + }, +} + +func TestRefcountMap_Actions(t *testing.T) { for _, s := range samples { sample := s t.Run(sample.name, func(t *testing.T) {