Internal change

PiperOrigin-RevId: 893956786

Author: aCombray

src/devtools/mobileharness/shared/util/comm/dualconduit/conduit/pump.go

@@ -72,7 +72,7 @@ func handleIngressConnection(c *Conduit, conn net.Conn) {
 
 // AcceptStream handles an incoming RSocket channel, dials the target service, and bridges the streams.
 // This implements the Responder Role.
-func AcceptStream(c *Conduit, destEndpoint string, incoming flux.Flux) flux.Flux {
+func AcceptStream(c *Conduit, destEndpoint string, downstream flux.Flux) flux.Flux {
 	conn, err := net.Dial("tcp", destEndpoint)
 	if err != nil {
 		return flux.Error(err)
@@ -86,8 +86,11 @@ func AcceptStream(c *Conduit, destEndpoint string, incoming flux.Flux) flux.Flux
 
 	streamCtx, cancel := context.WithCancel(c.Context())
 
-	// Read from RSocket and write to backend TCP.
-	rc.FromFlux(streamCtx, incoming)
+	// 1. Read from RSocket and write to backend TCP.
+	rc.FromFlux(streamCtx, downstream)
+
+	// 2. Read from backend TCP and send to RSocket.
+	upstream := rc.ToFlux(streamCtx)
 
 	// Clean up TCP connection when streaming finishes.
 	go func() {
@@ -96,6 +99,5 @@ func AcceptStream(c *Conduit, destEndpoint string, incoming flux.Flux) flux.Flux
 		conn.Close()
 	}()
 
-	// Read from backend TCP and send to RSocket.
-	return rc.ToFlux(streamCtx)
+	return upstream
 }

src/devtools/mobileharness/shared/util/comm/dualconduit/conduit/pump_test.go

@@ -308,33 +308,36 @@ func TestCombinedLoopAndAccept(t *testing.T) {
 	defer conn.Close()
 
 	// 3. Simultaneously run an AcceptStream toward another local server
-	l2, err := net.Listen("tcp", "localhost:0")
+	lis2, err := net.Listen("tcp", "localhost:0")
 	if err != nil {
 		t.Fatalf("Listen failed: %v", err)
 	}
-	defer l2.Close()
-	port2 := l2.Addr().(*net.TCPAddr).Port
+	defer lis2.Close()
+	port2 := lis2.Addr().(*net.TCPAddr).Port
 
 	go func() {
-		c, err := l2.Accept()
+		conn, err := lis2.Accept()
 		if err != nil {
 			if errors.Is(err, net.ErrClosed) {
 				return
 			}
-			t.Errorf("l2.Accept() got error: %v", err)
+			t.Errorf("lis2.Accept() got error: %v", err)
 			return
 		}
-		if c != nil {
-			c.Close()
-		}
+		defer conn.Close()
+
+		io.Copy(io.Discard, conn) // Drain unread data to avoid RST
 	}()
 
 	f := AcceptStream(c, fmt.Sprintf("localhost:%d", port2), flux.Just(payload.New([]byte("hello"), nil)))
 	if f == nil {
 		t.Fatalf("AcceptStream got nil, want non-nil flux")
 	}
 
-	// Wait, no panic = success
+	_, err = f.BlockLast(ctx)
+	if err != nil && !errors.Is(err, io.EOF) {
+		t.Fatalf("BlockLast failed: %v", err)
+	}
 }
 
 func TestHalfClose_Ingress_ClientCloseWrite(t *testing.T) {
@@ -344,19 +347,9 @@ func TestHalfClose_Ingress_ClientCloseWrite(t *testing.T) {
 	completed := make(chan struct{})
 	fakeRSocket := rsocket.NewAbstractSocket(
 		rsocket.RequestChannel(func(msgs flux.Flux) flux.Flux {
-			// Track when upstream is completed (client CloseWrite)
-			go func() {
-				_, err := msgs.BlockLast(ctx)
-				if err == nil {
-					close(completed)
-				}
-			}()
-			// Return downstream data slowly
-			f := flux.Create(func(ctx context.Context, sink flux.Sink) {
-				sink.Next(payload.New([]byte("delayed-response"), nil))
-				sink.Complete()
+			return msgs.DoOnComplete(func() {
+				close(completed)
 			})
-			return f
 		}),
 	)
 
@@ -390,8 +383,8 @@ func TestHalfClose_Ingress_ClientCloseWrite(t *testing.T) {
 		t.Fatalf("Unexpected read error: %v", err)
 	}
 
-	if string(buf[:n]) != "delayed-response" {
-		t.Fatalf("During half-close read got %q, want %q", string(buf[:n]), "delayed-response")
+	if string(buf[:n]) != "req" {
+		t.Fatalf("During half-close read got %q, want %q", string(buf[:n]), "req")
 	}
 
 	select {
@@ -411,8 +404,9 @@ func TestHalfClose_Egress_ServerCloseWrite(t *testing.T) {
 	defer l.Close()
 	port := l.Addr().(*net.TCPAddr).Port
 
-	var testPassed atomic.Int32
+	serverDone := make(chan struct{})
 	go func() {
+		defer close(serverDone)
 		conn, err := l.Accept()
 		if err != nil {
 			return
@@ -425,21 +419,21 @@ func TestHalfClose_Egress_ServerCloseWrite(t *testing.T) {
 
 		buf := make([]byte, 1024)
 		n, err := tcpConn.Read(buf) // Should get "req"
-		if err != nil {
+		if err != nil && err != io.EOF {
 			return
 		}
-		if string(buf[:n]) == "req" {
+		if n > 0 && string(buf[:n]) == "req" {
 			tcpConn.Write([]byte("response"))
 			// Trigger half-close from server end
 			tcpConn.CloseWrite()
 		}
 
-		// Wait to see if we get more data? Actually in Acceptance stream,
-		// remote already pushed all.
-		// Wait to observe EOF from client because RSocket complete propagated?
-		_, err = tcpConn.Read(buf)
-		if err == io.EOF {
-			testPassed.Store(1)
+		// Read remaining data until EOF if not already reached
+		if err != io.EOF {
+			_, err = tcpConn.Read(buf)
+		}
+		if err != io.EOF {
+			t.Errorf("Server side did not receive EOF, got %v", err)
 		}
 	}()
 
@@ -452,7 +446,7 @@ func TestHalfClose_Egress_ServerCloseWrite(t *testing.T) {
 
 	outFlux := AcceptStream(c, fmt.Sprintf("localhost:%d", port), f)
 
-	res, err := outFlux.BlockFirst(ctx)
+	res, err := outFlux.BlockLast(ctx)
 	if err != nil {
 		t.Fatalf("outFlux blockFirst error: %v", err)
 	}
@@ -462,7 +456,5 @@ func TestHalfClose_Egress_ServerCloseWrite(t *testing.T) {
 
 	// Since f completed, and rxconn mirrors EOF locally (CloseWrite on the dialed connection)
 	// the server should have read an EOF.
-	if testPassed.Load() != 1 {
-		t.Fatalf("Server side did not receive EOF (rxconn did not mirror Complete)")
-	}
+	<-serverDone // Wait for server to finish reading EOF and close
 }

src/devtools/mobileharness/shared/util/comm/dualconduit/rxconn/rxconn.go

@@ -3,9 +3,11 @@ package rxconn
 
 import (
 	"context"
+	"errors"
 	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 
 	"github.com/rsocket/rsocket-go/payload"
 	"github.com/rsocket/rsocket-go/rx/flux"
@@ -14,8 +16,9 @@ import (
 
 // Conn bridges a standard network connection with reactive streams.
 type Conn struct {
-	conn net.Conn
-	wg   sync.WaitGroup
+	conn       net.Conn
+	wg         sync.WaitGroup
+	subscribed atomic.Bool // Guard to enforce the single-subscriber policy
 }
 
 // New creates a new Conn bridging the active network connection to reactive fluxes.
@@ -29,12 +32,26 @@ func New(conn net.Conn) (*Conn, error) {
 // ToFlux returns a Flux that reads from the underlying active socket.
 // The context controls pump lifecycle: cancellation tears down background pumps.
 // It may be used after ToFlux returns, as pumps are run in background goroutines.
+//
+// ToFlux enforces a single-subscriber policy using an atomic guard to prevent
+// multiple subscribers from corrupting the stream (net.Conn is stateful and
+// cannot be shared by multiple Read Pumps).
 func (c *Conn) ToFlux(ctx context.Context) flux.Flux {
-	upstreamChan := make(chan []byte, 32)
-	errChan := make(chan error, 1)
-
+	// Eager Tracking: Protect c.wg.Wait() immediately.
+	// This ensures that Wait() does not exit prematurely if the RSocket
+	// subscription is slightly delayed or happens after Wait() is called.
 	c.wg.Add(1)
 	return flux.Create(func(subscriberCtx context.Context, sink flux.Sink) {
+		if c.subscribed.Swap(true) {
+			sink.Error(errors.New("ToFlux allows only one subscriber"))
+			return
+		}
+
+		// Lazy Pumping: Initialize channels and start goroutines ONLY when
+		// a subscriber is ready.
+		upstreamChan := make(chan []byte, 32)
+		errChan := make(chan error, 1)
+
 		// 1. The TCP Read Pump
 		go func() {
 			defer close(upstreamChan)

src/devtools/mobileharness/shared/util/comm/dualconduit/rxconn/rxconn_test.go

@@ -16,14 +16,14 @@ func TestRxConn_ToFlux_EOF(t *testing.T) {
 	clientSide, serverSide := net.Pipe()
 	defer clientSide.Close()
 
-	rxConnObj, err := New(serverSide)
+	conn, err := New(serverSide)
 	if err != nil {
 		t.Fatalf("Failed to create rxConn: %v", err)
 	}
 	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
 	defer cancel()
 
-	f := rxConnObj.ToFlux(ctx)
+	f := conn.ToFlux(ctx)
 	resultChan := make(chan payload.Payload, 10)
 	errChan := make(chan error, 1)
 	doneChan := make(chan struct{})
@@ -74,13 +74,13 @@ func TestRxConn_ToFlux_ContextCancellation(t *testing.T) {
 	clientSide, serverSide := net.Pipe()
 	defer clientSide.Close()
 
-	rxConnObj, err := New(serverSide)
+	conn, err := New(serverSide)
 	if err != nil {
 		t.Fatalf("Failed to create rxConn: %v", err)
 	}
 	ctx, cancel := context.WithCancel(t.Context())
 
-	f := rxConnObj.ToFlux(ctx)
+	f := conn.ToFlux(ctx)
 	errChan := make(chan error, 1)
 
 	f.Subscribe(t.Context(),
@@ -107,7 +107,7 @@ func TestRxConn_FromFlux_DataAndEOFPropagation(t *testing.T) {
 	clientSide, serverSide := net.Pipe()
 	defer clientSide.Close()
 
-	rxConnObj, err := New(serverSide)
+	conn, err := New(serverSide)
 	if err != nil {
 		t.Fatalf("Failed to create rxConn: %v", err)
 	}
@@ -121,7 +121,7 @@ func TestRxConn_FromFlux_DataAndEOFPropagation(t *testing.T) {
 	// Subscribe in a goroutine because FromFlux will synchronously write to the pipe,
 	// which blocks until the test reads from it below.
 	go func() {
-		rxConnObj.FromFlux(t.Context(), f)
+		conn.FromFlux(t.Context(), f)
 	}()
 
 	// Read from pipe to verify
@@ -149,14 +149,14 @@ func TestRxConn_ToFlux_Backpressure(t *testing.T) {
 	defer clientSide.Close()
 	defer serverSide.Close()
 
-	rxConnObj, err := New(serverSide)
+	conn, err := New(serverSide)
 	if err != nil {
 		t.Fatalf("Failed to create rxConn: %v", err)
 	}
 	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
 	defer cancel()
 
-	f := rxConnObj.ToFlux(ctx)
+	f := conn.ToFlux(ctx)
 
 	// We DON'T consume the flux immediately, forcing the upstream channel to fill up.
 	// The channel buffer is 32. We can send 32 bursts without blocking, then it should block.
@@ -205,13 +205,13 @@ func TestRxConn_ToFlux_Backpressure(t *testing.T) {
 func TestRxConn_Wait_CloseFromFluxFirst(t *testing.T) {
 	clientSide, serverSide := net.Pipe()
 
-	rxConnObj, err := New(serverSide)
+	conn, err := New(serverSide)
 	if err != nil {
 		t.Fatalf("Failed to create rxConn: %v", err)
 	}
 	ctx := t.Context()
 
-	fTo := rxConnObj.ToFlux(ctx)
+	fTo := conn.ToFlux(ctx)
 
 	stopFrom := make(chan struct{})
 	fFrom := flux.Create(func(subCtx context.Context, sink flux.Sink) {
@@ -223,11 +223,11 @@ func TestRxConn_Wait_CloseFromFluxFirst(t *testing.T) {
 	})
 
 	fTo.Subscribe(ctx)
-	go rxConnObj.FromFlux(ctx, fFrom)
+	go conn.FromFlux(ctx, fFrom)
 
 	waitDone := make(chan struct{})
 	go func() {
-		rxConnObj.Wait()
+		conn.Wait()
 		close(waitDone)
 	}()
 
@@ -262,13 +262,13 @@ func TestRxConn_Wait_CloseFromFluxFirst(t *testing.T) {
 func TestRxConn_Wait_CloseToFluxFirst(t *testing.T) {
 	clientSide, serverSide := net.Pipe()
 
-	rxConnObj, err := New(serverSide)
+	conn, err := New(serverSide)
 	if err != nil {
 		t.Fatalf("Failed to create rxConn: %v", err)
 	}
 	ctx := t.Context()
 
-	fTo := rxConnObj.ToFlux(ctx)
+	fTo := conn.ToFlux(ctx)
 
 	stopFrom := make(chan struct{})
 	fFrom := flux.Create(func(subCtx context.Context, sink flux.Sink) {
@@ -280,11 +280,11 @@ func TestRxConn_Wait_CloseToFluxFirst(t *testing.T) {
 	})
 
 	fTo.Subscribe(ctx)
-	go rxConnObj.FromFlux(ctx, fFrom)
+	go conn.FromFlux(ctx, fFrom)
 
 	waitDone := make(chan struct{})
 	go func() {
-		rxConnObj.Wait()
+		conn.Wait()
 		close(waitDone)
 	}()
 
@@ -315,3 +315,34 @@ func TestRxConn_Wait_CloseToFluxFirst(t *testing.T) {
 		t.Fatal("Wait() blocked even after both pumps stopped")
 	}
 }
+
+func TestRxConn_ToFlux_SingleSubscriberPolicy(t *testing.T) {
+	_, serverSide := net.Pipe()
+	defer serverSide.Close()
+
+	conn, err := New(serverSide)
+	if err != nil {
+		t.Fatalf("Failed to create rxConn: %v", err)
+	}
+	ctx := t.Context()
+
+	f := conn.ToFlux(ctx)
+
+	// Subscribe once
+	f.Subscribe(ctx)
+
+	// Subscribe second time - should fail
+	errChan := make(chan error, 1)
+	f.Subscribe(ctx, rx.OnError(func(err error) {
+		errChan <- err
+	}))
+
+	select {
+	case err := <-errChan:
+		if err.Error() != "ToFlux allows only one subscriber" {
+			t.Errorf("ToFlux() error = %q, want %q", err.Error(), "ToFlux allows only one subscriber")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for error on second subscription")
+	}
+}