Minor flipcall refactoring.

PiperOrigin-RevId: 342768550

3 files changed, 73 insertions, 70 deletions

diff --git a/pkg/flipcall/ctrl_futex.go b/pkg/flipcall/ctrl_futex.go
index e7c3a3a0b..2e8452a02 100644
--- a/pkg/flipcall/ctrl_futex.go
+++ b/pkg/flipcall/ctrl_futex.go
@@ -40,17 +40,41 @@ func (ep *Endpoint) ctrlInit(opts ...EndpointOption) error {
 	return nil
 }
 
-type ctrlHandshakeRequest struct{}
-
-type ctrlHandshakeResponse struct{}
-
 func (ep *Endpoint) ctrlConnect() error {
 	if err := ep.enterFutexWait(); err != nil {
 		return err
 	}
-	_, err := ep.futexConnect(&ctrlHandshakeRequest{})
-	ep.exitFutexWait()
-	return err
+	defer ep.exitFutexWait()
+
+	// Write the connection request.
+	w := ep.NewWriter()
+	if err := json.NewEncoder(w).Encode(struct{}{}); err != nil {
+		return fmt.Errorf("error writing connection request: %v", err)
+	}
+	*ep.dataLen() = w.Len()
+
+	// Exchange control with the server.
+	if err := ep.futexSetPeerActive(); err != nil {
+		return err
+	}
+	if err := ep.futexWakePeer(); err != nil {
+		return err
+	}
+	if err := ep.futexWaitUntilActive(); err != nil {
+		return err
+	}
+
+	// Read the connection response.
+	var resp struct{}
+	respLen := atomic.LoadUint32(ep.dataLen())
+	if respLen > ep.dataCap {
+		return fmt.Errorf("invalid connection response length %d (maximum %d)", respLen, ep.dataCap)
+	}
+	if err := json.NewDecoder(ep.NewReader(respLen)).Decode(&resp); err != nil {
+		return fmt.Errorf("error reading connection response: %v", err)
+	}
+
+	return nil
 }
 
 func (ep *Endpoint) ctrlWaitFirst() error {
@@ -59,52 +83,61 @@ func (ep *Endpoint) ctrlWaitFirst() error {
 	}
 	defer ep.exitFutexWait()
 
-	// Wait for the handshake request.
-	if err := ep.futexSwitchFromPeer(); err != nil {
+	// Wait for the connection request.
+	if err := ep.futexWaitUntilActive(); err != nil {
 		return err
 	}
 
-	// Read the handshake request.
+	// Read the connection request.
 	reqLen := atomic.LoadUint32(ep.dataLen())
 	if reqLen > ep.dataCap {
-		return fmt.Errorf("invalid handshake request length %d (maximum %d)", reqLen, ep.dataCap)
+		return fmt.Errorf("invalid connection request length %d (maximum %d)", reqLen, ep.dataCap)
 	}
-	var req ctrlHandshakeRequest
+	var req struct{}
 	if err := json.NewDecoder(ep.NewReader(reqLen)).Decode(&req); err != nil {
-		return fmt.Errorf("error reading handshake request: %v", err)
+		return fmt.Errorf("error reading connection request: %v", err)
 	}
 
-	// Write the handshake response.
+	// Write the connection response.
 	w := ep.NewWriter()
-	if err := json.NewEncoder(w).Encode(ctrlHandshakeResponse{}); err != nil {
-		return fmt.Errorf("error writing handshake response: %v", err)
+	if err := json.NewEncoder(w).Encode(struct{}{}); err != nil {
+		return fmt.Errorf("error writing connection response: %v", err)
 	}
 	*ep.dataLen() = w.Len()
 
 	// Return control to the client.
 	raceBecomeInactive()
-	if err := ep.futexSwitchToPeer(); err != nil {
+	if err := ep.futexSetPeerActive(); err != nil {
+		return err
+	}
+	if err := ep.futexWakePeer(); err != nil {
 		return err
 	}
 
-	// Wait for the first non-handshake message.
-	return ep.futexSwitchFromPeer()
+	// Wait for the first non-connection message.
+	return ep.futexWaitUntilActive()
 }
 
 func (ep *Endpoint) ctrlRoundTrip() error {
-	if err := ep.futexSwitchToPeer(); err != nil {
+	if err := ep.enterFutexWait(); err != nil {
 		return err
 	}
-	if err := ep.enterFutexWait(); err != nil {
+	defer ep.exitFutexWait()
+
+	if err := ep.futexSetPeerActive(); err != nil {
 		return err
 	}
-	err := ep.futexSwitchFromPeer()
-	ep.exitFutexWait()
-	return err
+	if err := ep.futexWakePeer(); err != nil {
+		return err
+	}
+	return ep.futexWaitUntilActive()
 }
 
 func (ep *Endpoint) ctrlWakeLast() error {
-	return ep.futexSwitchToPeer()
+	if err := ep.futexSetPeerActive(); err != nil {
+		return err
+	}
+	return ep.futexWakePeer()
 }
 
 func (ep *Endpoint) enterFutexWait() error {
diff --git a/pkg/flipcall/flipcall_unsafe.go b/pkg/flipcall/flipcall_unsafe.go
index ac974b232..580bf23a4 100644
--- a/pkg/flipcall/flipcall_unsafe.go
+++ b/pkg/flipcall/flipcall_unsafe.go
@@ -41,11 +41,11 @@ const (
 )
 
 func (ep *Endpoint) connState() *uint32 {
-	return (*uint32)((unsafe.Pointer)(ep.packet))
+	return (*uint32)(unsafe.Pointer(ep.packet))
 }
 
 func (ep *Endpoint) dataLen() *uint32 {
-	return (*uint32)((unsafe.Pointer)(ep.packet + 4))
+	return (*uint32)(unsafe.Pointer(ep.packet + 4))
 }
 
 // Data returns the datagram part of ep's packet window as a byte slice.
@@ -63,7 +63,7 @@ func (ep *Endpoint) dataLen() *uint32 {
 // all.
 func (ep *Endpoint) Data() []byte {
 	var bs []byte
-	bsReflect := (*reflect.SliceHeader)((unsafe.Pointer)(&bs))
+	bsReflect := (*reflect.SliceHeader)(unsafe.Pointer(&bs))
 	bsReflect.Data = ep.packet + PacketHeaderBytes
 	bsReflect.Len = int(ep.dataCap)
 	bsReflect.Cap = int(ep.dataCap)
@@ -76,12 +76,12 @@ var ioSync int64
 
 func raceBecomeActive() {
 	if sync.RaceEnabled {
-		sync.RaceAcquire((unsafe.Pointer)(&ioSync))
+		sync.RaceAcquire(unsafe.Pointer(&ioSync))
 	}
 }
 
 func raceBecomeInactive() {
 	if sync.RaceEnabled {
-		sync.RaceReleaseMerge((unsafe.Pointer)(&ioSync))
+		sync.RaceReleaseMerge(unsafe.Pointer(&ioSync))
 	}
 }
diff --git a/pkg/flipcall/futex_linux.go b/pkg/flipcall/futex_linux.go
index 168c1ccff..0e559ee16 100644
--- a/pkg/flipcall/futex_linux.go
+++ b/pkg/flipcall/futex_linux.go
@@ -17,7 +17,6 @@
 package flipcall
 
 import (
-	"encoding/json"
 	"fmt"
 	"runtime"
 	"sync/atomic"
@@ -26,55 +25,26 @@ import (
 	"gvisor.dev/gvisor/pkg/abi/linux"
 )
 
-func (ep *Endpoint) futexConnect(req *ctrlHandshakeRequest) (ctrlHandshakeResponse, error) {
-	var resp ctrlHandshakeResponse
-
-	// Write the handshake request.
-	w := ep.NewWriter()
-	if err := json.NewEncoder(w).Encode(req); err != nil {
-		return resp, fmt.Errorf("error writing handshake request: %v", err)
-	}
-	*ep.dataLen() = w.Len()
-
-	// Exchange control with the server.
-	if err := ep.futexSwitchToPeer(); err != nil {
-		return resp, err
+func (ep *Endpoint) futexSetPeerActive() error {
+	if atomic.CompareAndSwapUint32(ep.connState(), ep.activeState, ep.inactiveState) {
+		return nil
 	}
-	if err := ep.futexSwitchFromPeer(); err != nil {
-		return resp, err
+	switch cs := atomic.LoadUint32(ep.connState()); cs {
+	case csShutdown:
+		return ShutdownError{}
+	default:
+		return fmt.Errorf("unexpected connection state before FUTEX_WAKE: %v", cs)
 	}
-
-	// Read the handshake response.
-	respLen := atomic.LoadUint32(ep.dataLen())
-	if respLen > ep.dataCap {
-		return resp, fmt.Errorf("invalid handshake response length %d (maximum %d)", respLen, ep.dataCap)
-	}
-	if err := json.NewDecoder(ep.NewReader(respLen)).Decode(&resp); err != nil {
-		return resp, fmt.Errorf("error reading handshake response: %v", err)
-	}
-
-	return resp, nil
 }
 
-func (ep *Endpoint) futexSwitchToPeer() error {
-	// Update connection state to indicate that the peer should be active.
-	if !atomic.CompareAndSwapUint32(ep.connState(), ep.activeState, ep.inactiveState) {
-		switch cs := atomic.LoadUint32(ep.connState()); cs {
-		case csShutdown:
-			return ShutdownError{}
-		default:
-			return fmt.Errorf("unexpected connection state before FUTEX_WAKE: %v", cs)
-		}
-	}
-
-	// Wake the peer's Endpoint.futexSwitchFromPeer().
+func (ep *Endpoint) futexWakePeer() error {
 	if err := ep.futexWakeConnState(1); err != nil {
 		return fmt.Errorf("failed to FUTEX_WAKE peer Endpoint: %v", err)
 	}
 	return nil
 }
 
-func (ep *Endpoint) futexSwitchFromPeer() error {
+func (ep *Endpoint) futexWaitUntilActive() error {
 	for {
 		switch cs := atomic.LoadUint32(ep.connState()); cs {
 		case ep.activeState: