From 461a15be35059bd1e184121d5776f2dea15c89a8 Mon Sep 17 00:00:00 2001
From: Andrei Vagin <avagin@gmail.com>
Date: Thu, 1 Dec 2022 15:15:09 -0800
Subject: [PATCH] tcpip/stack: use atomic refs for addressState

In this case, we don't need to take a read-write lock in IncRef and DecRef.
---
 pkg/tcpip/stack/BUILD                         | 12 +++++
 pkg/tcpip/stack/addressable_endpoint_state.go | 48 +++++++------------
 2 files changed, 29 insertions(+), 31 deletions(-)

diff --git a/pkg/tcpip/stack/BUILD b/pkg/tcpip/stack/BUILD
index 6248ff237..8cbce3b67 100644
--- a/pkg/tcpip/stack/BUILD
+++ b/pkg/tcpip/stack/BUILD
@@ -50,9 +50,21 @@ go_template_instance(
     },
 )
 
+go_template_instance(
+    name = "address_state_refs",
+    out = "address_state_refs.go",
+    package = "stack",
+    prefix = "addressState",
+    template = "//pkg/refs:refs_template",
+    types = {
+        "T": "addressState",
+    },
+)
+
 go_library(
     name = "stack",
     srcs = [
+        "address_state_refs.go",
         "addressable_endpoint_state.go",
         "conntrack.go",
         "gro.go",
diff --git a/pkg/tcpip/stack/addressable_endpoint_state.go b/pkg/tcpip/stack/addressable_endpoint_state.go
index df9953375..dc1027326 100644
--- a/pkg/tcpip/stack/addressable_endpoint_state.go
+++ b/pkg/tcpip/stack/addressable_endpoint_state.go
@@ -223,7 +223,7 @@ func (a *AddressableEndpointState) addAndAcquireAddressLocked(addr tcpip.Address
 		}
 
 		addrState.mu.RLock()
-		if addrState.refs == 0 {
+		if addrState.refs.ReadRefs() == 0 {
 			panic(fmt.Sprintf("found an address that should have been released (ref count == 0); address = %s", addrState.addr))
 		}
 		isPermanent := addrState.kind.IsPermanent()
@@ -257,6 +257,7 @@ func (a *AddressableEndpointState) addAndAcquireAddressLocked(addr tcpip.Address
 				break
 			}
 		}
+		addrState.refs.IncRef()
 	} else {
 		addrState = &addressState{
 			addressableEndpointState: a,
@@ -266,6 +267,7 @@ func (a *AddressableEndpointState) addAndAcquireAddressLocked(addr tcpip.Address
 			// results in allocations on every call.
 			subnet: addr.Subnet(),
 		}
+		addrState.refs.InitRefs()
 		a.endpoints[addr.Address] = addrState
 		// We never promote an address to temporary - it can only be added as such.
 		// If we are actually adding a permanent address, it is promoted below.
@@ -287,11 +289,10 @@ func (a *AddressableEndpointState) addAndAcquireAddressLocked(addr tcpip.Address
 		}
 
 		// Primary addresses are biased by 1.
-		addrState.refs++
+		addrState.refs.IncRef()
 		addrState.kind = kind
 	}
 	// Acquire the address before returning it.
-	addrState.refs++
 	addrState.configType = properties.ConfigType
 	lifetimes := properties.Lifetimes
 	lifetimes.sanitize()
@@ -384,19 +385,16 @@ func (a *AddressableEndpointState) decAddressRef(addrState *addressState) {
 //
 // +checklocks:a.mu
 func (a *AddressableEndpointState) decAddressRefLocked(addrState *addressState) {
-	addrState.mu.Lock()
-	defer addrState.mu.Unlock()
+	destroy := false
+	addrState.refs.DecRef(func() {
+		destroy = true
+	})
 
-	if addrState.refs == 0 {
-		panic(fmt.Sprintf("attempted to decrease ref count for AddressEndpoint w/ addr = %s when it is already released", addrState.addr))
-	}
-
-	addrState.refs--
-
-	if addrState.refs != 0 {
+	if !destroy {
 		return
 	}
-
+	addrState.mu.Lock()
+	defer addrState.mu.Unlock()
 	// A non-expired permanent address must not have its reference count dropped
 	// to 0.
 	if addrState.kind.IsPermanent() {
@@ -701,9 +699,8 @@ type addressState struct {
 	//
 	// AddressableEndpointState.mu
 	//   addressState.mu
-	mu sync.RWMutex
-	// checklocks:mu
-	refs uint32
+	mu   sync.RWMutex
+	refs addressStateRefs
 	// checklocks:mu
 	kind AddressKind
 	// checklocks:mu
@@ -788,14 +785,7 @@ func (a *addressState) IsAssigned(allowExpired bool) bool {
 
 // IncRef implements AddressEndpoint.
 func (a *addressState) IncRef() bool {
-	a.mu.Lock()
-	defer a.mu.Unlock()
-	if a.refs == 0 {
-		return false
-	}
-
-	a.refs++
-	return true
+	return a.refs.TryIncRef()
 }
 
 // DecRef implements AddressEndpoint.
@@ -809,13 +799,9 @@ func (a *addressState) DecRef() {
 // Panics if the ref count is less than 2 after acquiring the lock in this
 // function.
 func (a *addressState) decRefMustNotFree() {
-	a.mu.Lock()
-	defer a.mu.Unlock()
-
-	if a.refs < 2 {
-		panic(fmt.Sprintf("cannot decrease addressState %s ref count %d without freeing the endpoint", a.addr, a.refs))
-	}
-	a.refs--
+	a.refs.DecRef(func() {
+		panic(fmt.Sprintf("cannot decrease addressState %s without freeing the endpoint", a.addr))
+	})
 }
 
 // ConfigType implements AddressEndpoint.