Make callback handle lookups lock-free

callback.go

@@ -29,6 +29,7 @@ import (
 	"math"
 	"reflect"
 	"sync"
+	"sync/atomic"
 	"unsafe"
 )
 
@@ -104,24 +105,26 @@ type handleVal struct {
 }
 
 var handleLock sync.Mutex
-var handleVals = make(map[unsafe.Pointer]handleVal)
+var handleVals atomic.Value // stores map[unsafe.Pointer]handleVal
 
 func newHandle(db *SQLiteConn, v any) unsafe.Pointer {
-	handleLock.Lock()
-	defer handleLock.Unlock()
 	val := handleVal{db: db, val: v}
 	var p unsafe.Pointer = C.malloc(C.size_t(1))
 	if p == nil {
 		panic("can't allocate 'cgo-pointer hack index pointer': ptr == nil")
 	}
-	handleVals[p] = val
+
+	handleLock.Lock()
+	defer handleLock.Unlock()
+
+	next := cloneHandleVals(len(loadHandleVals()) + 1)
+	next[p] = val
+	handleVals.Store(next)
 	return p
 }
 
 func lookupHandleVal(handle unsafe.Pointer) handleVal {
-	handleLock.Lock()
+	return loadHandleVals()[handle]
-	defer handleLock.Unlock()
-	return handleVals[handle]
 }
 
 func lookupHandle(handle unsafe.Pointer) any {
@@ -131,12 +134,34 @@ func lookupHandle(handle unsafe.Pointer) any {
 func deleteHandles(db *SQLiteConn) {
 	handleLock.Lock()
 	defer handleLock.Unlock()
-	for handle, val := range handleVals {
+
+	current := loadHandleVals()
+	if len(current) == 0 {
+		return
+	}
+
+	next := make(map[unsafe.Pointer]handleVal, len(current))
+	for handle, val := range current {
 		if val.db == db {
-			delete(handleVals, handle)
 			C.free(handle)
+			continue
 		}
+		next[handle] = val
 	}
+	handleVals.Store(next)
+}
+
+func loadHandleVals() map[unsafe.Pointer]handleVal {
+	m, _ := handleVals.Load().(map[unsafe.Pointer]handleVal)
+	return m
+}
+
+func cloneHandleVals(size int) map[unsafe.Pointer]handleVal {
+	next := make(map[unsafe.Pointer]handleVal, size)
+	for handle, val := range loadHandleVals() {
+		next[handle] = val
+	}
+	return next
 }
 
 // This is only here so that tests can refer to it.

callback_bench_test.go

@@ -0,0 +1,85 @@
+// Copyright (C) 2019 Yasuhiro Matsumoto <mattn.jp@gmail.com>.
+//
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+//go:build cgo
+// +build cgo
+
+package sqlite3
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+	"unsafe"
+)
+
+func BenchmarkHandleLookupParallel(b *testing.B) {
+	d := SQLiteDriver{}
+	conn, err := d.Open(":memory:")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer conn.Close()
+	c := conn.(*SQLiteConn)
+
+	handle := newHandle(c, func() {})
+
+	benchmarkHandleLookupParallel(b, func() any {
+		return lookupHandle(handle)
+	})
+}
+
+func BenchmarkHandleLookupBeforeAfter(b *testing.B) {
+	value := handleVal{val: func() {}}
+	handle := unsafe.Pointer(&value)
+
+	before := mutexHandleTable{vals: map[unsafe.Pointer]handleVal{handle: value}}
+	after := atomicHandleTable{}
+	after.vals.Store(map[unsafe.Pointer]handleVal{handle: value})
+
+	b.Run("before_mutex", func(b *testing.B) {
+		benchmarkHandleLookupParallel(b, func() any {
+			return before.lookup(handle).val
+		})
+	})
+	b.Run("after_atomic", func(b *testing.B) {
+		benchmarkHandleLookupParallel(b, func() any {
+			return after.lookup(handle).val
+		})
+	})
+}
+
+func benchmarkHandleLookupParallel(b *testing.B, lookup func() any) {
+	b.Helper()
+	b.ReportAllocs()
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			if lookup() == nil {
+				b.Fatal("lookup returned nil")
+			}
+		}
+	})
+}
+
+type mutexHandleTable struct {
+	mu   sync.Mutex
+	vals map[unsafe.Pointer]handleVal
+}
+
+func (t *mutexHandleTable) lookup(handle unsafe.Pointer) handleVal {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.vals[handle]
+}
+
+type atomicHandleTable struct {
+	vals atomic.Value
+}
+
+func (t *atomicHandleTable) lookup(handle unsafe.Pointer) handleVal {
+	m, _ := t.vals.Load().(map[unsafe.Pointer]handleVal)
+	return m[handle]
+}