[Resource Pool] Shorten critical section for register/unregister

go/pools/numbered.go

@@ -64,31 +64,47 @@ func NewNumbered() *Numbered {
 // It does not lock the object.
 // It returns an error if the id already exists.
 func (nu *Numbered) Register(id int64, val interface{}, enforceTimeout bool) error {
-	nu.mu.Lock()
-	defer nu.mu.Unlock()
-	if _, ok := nu.resources[id]; ok {
-		return fmt.Errorf("already present")
-	}
+	// Optimistically assume we're not double registering.
 	now := time.Now()
-	nu.resources[id] = &numberedWrapper{
+	resource := &numberedWrapper{
 		val:            val,
 		timeCreated:    now,
 		timeUsed:       now,
 		enforceTimeout: enforceTimeout,
 	}
+
+	nu.mu.Lock()
+	defer nu.mu.Unlock()
+
+	_, ok := nu.resources[id]
+	if ok {
+		return fmt.Errorf("already present")
+	}
+	nu.resources[id] = resource
 	return nil
 }
 
-// Unregiester forgets the specified resource.
-// If the resource is not present, it's ignored.
+// Unregister forgets the specified resource.  If the resource is not present, it's ignored.
 func (nu *Numbered) Unregister(id int64, reason string) {
+	success := nu.unregister(id, reason)
+	if success {
+		nu.recentlyUnregistered.Set(
+			fmt.Sprintf("%v", id), &unregistered{reason: reason, timeUnregistered: time.Now()})
+	}
+}
+
+// unregister forgets the resource, if it exists. Returns whether or not the resource existed at
+// time of Unregister.
+func (nu *Numbered) unregister(id int64, reason string) bool {
 	nu.mu.Lock()
 	defer nu.mu.Unlock()
+
+	_, ok := nu.resources[id]
 	delete(nu.resources, id)
 	if len(nu.resources) == 0 {
 		nu.empty.Broadcast()
 	}
-	nu.recentlyUnregistered.Set(fmt.Sprintf("%v", id), &unregistered{reason: reason, timeUnregistered: time.Now()})
+	return ok
 }
 
 // Get locks the resource for use. It accepts a purpose as a string.

go/pools/numbered_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package pools
 
 import (
+	"math/rand"
 	"strings"
 	"testing"
 	"time"
@@ -103,3 +104,35 @@ func TestNumbered(t *testing.T) {
 	}()
 	p.WaitForEmpty()
 }
+
+/*
+go test --test.run=XXX --test.bench=. --test.benchtime=10s
+
+golang.org/x/tools/cmd/benchcmp /tmp/bad.out /tmp/good.out
+
+benchmark                                 old ns/op     new ns/op     delta
+BenchmarkRegisterUnregister-8             667           596           -10.64%
+BenchmarkRegisterUnregisterParallel-8     2430          1752          -27.90%
+*/
+func BenchmarkRegisterUnregister(b *testing.B) {
+	p := NewNumbered()
+	id := int64(1)
+	val := "foobarbazdummyval"
+	for i := 0; i < b.N; i++ {
+		p.Register(id, val, false)
+		p.Unregister(id, "some reason")
+	}
+}
+
+func BenchmarkRegisterUnregisterParallel(b *testing.B) {
+	p := NewNumbered()
+	val := "foobarbazdummyval"
+	b.SetParallelism(200)
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			id := rand.Int63()
+			p.Register(id, val, false)
+			p.Unregister(id, "some reason")
+		}
+	})
+}