feat(bigtable): NewServerWithListener (#12408)

bhshkh · web-flow · commit 59ea0de1013a · 2025-06-05T15:00:28.000-04:00
diff --git a/bigtable/bttest/inmem.go b/bigtable/bttest/inmem.go
@@ -29,6 +29,19 @@ To use a Server, create it, and then connect to it with no security:
 	client, err := bigtable.NewClient(ctx, proj, instance,
 	        option.WithGRPCConn(conn))
 	...
+
+To use a Server with an in-memory connection, provide a bufconn listener:
+
+	l := bufconn.Listen(1024 * 1024)
+	srv, err := bttest.NewServerWithListener(l)
+	...
+	conn, err := grpc.Dial(
+		"bufnet",
+		grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
+			return l.Dial()
+		}),
+		grpc.WithTransportCredentials(insecure.NewCredentials()))
+	...
 */
 package bttest // import "cloud.google.com/go/bigtable/bttest"
 
@@ -85,9 +98,10 @@ var validLabelTransformer = regexp.MustCompile(`[a-z0-9\-]{1,15}`)
 type Server struct {
 	Addr string
 
-	l   net.Listener
-	srv *grpc.Server
-	s   *server
+	l            net.Listener
+	srv          *grpc.Server
+	s            *server
+	ownsListener bool
 }
 
 // server is the real implementation of the fake.
@@ -105,6 +119,15 @@ type server struct {
 	btpb.BigtableServer
 }
 
+// noopCloserListener wraps a net.Listener, but its Close method is a no-op.
+// This is used to prevent the gRPC server from closing a listener that was
+// provided by a user.
+type noopCloserListener struct {
+	net.Listener
+}
+
+func (n *noopCloserListener) Close() error { return nil }
+
 // NewServer creates a new Server.
 // The Server will be listening for gRPC connections, without TLS,
 // on the provided address. The resolved address is named by the Addr field.
@@ -121,7 +144,10 @@ func NewServer(laddr string, opt ...grpc.ServerOption) (*Server, error) {
 	if err != nil {
 		return nil, err
 	}
+	return newServer(l, true, opt...), nil
+}
 
+func newServer(l net.Listener, ownsListener bool, opt ...grpc.ServerOption) *Server {
 	s := &Server{
 		Addr: l.Addr().String(),
 		l:    l,
@@ -130,14 +156,29 @@ func NewServer(laddr string, opt ...grpc.ServerOption) (*Server, error) {
 			tables:    make(map[string]*table),
 			instances: make(map[string]*btapb.Instance),
 		},
+		ownsListener: ownsListener,
 	}
 	btapb.RegisterBigtableInstanceAdminServer(s.srv, s.s)
 	btapb.RegisterBigtableTableAdminServer(s.srv, s.s)
 	btpb.RegisterBigtableServer(s.srv, s.s)
 
-	go s.srv.Serve(s.l)
+	listenerForGRPC := l
+	if !ownsListener {
+		// If the user owns the listener, wrap it so srv.Stop() doesn't close it.
+		listenerForGRPC = &noopCloserListener{Listener: l}
+	}
+	go s.srv.Serve(listenerForGRPC)
+	return s
+}
 
-	return s, nil
+// NewServerWithListener creates a new Server using the provided listener.
+// The Addr field of the returned Server will be the listener's address.
+//
+// The caller is responsible for closing the listener. The server's Close method
+// will not close the provided listener, nor will it clean up any underlying
+// resources like unix socket files.
+func NewServerWithListener(l net.Listener, opt ...grpc.ServerOption) (*Server, error) {
+	return newServer(l, false, opt...), nil
 }
 
 // Close shuts down the server.
@@ -149,10 +190,12 @@ func (s *Server) Close() {
 	s.s.mu.Unlock()
 
 	s.srv.Stop()
-	s.l.Close()
+	if s.ownsListener {
+		s.l.Close()
+	}
 
 	// clean up unix socket
-	if strings.Contains(s.Addr, "/") {
+	if s.ownsListener && strings.Contains(s.Addr, "/") {
 		_ = os.Remove(s.Addr)
 	}
 }
diff --git a/bigtable/bttest/inmem_test.go b/bigtable/bttest/inmem_test.go
@@ -19,9 +19,12 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"math"
 	"math/rand"
+	"net"
+	"os"
 	"sort"
 	"strconv"
 	"sync"
@@ -46,6 +49,87 @@ import (
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
+type closeSpyListener struct {
+	net.Listener
+	closed bool
+}
+
+func (s *closeSpyListener) Close() error {
+	s.closed = true
+	return s.Listener.Close()
+}
+
+func TestNewServer(t *testing.T) {
+	t.Run("TCP", func(t *testing.T) {
+		srv, err := NewServer("localhost:0")
+		if err != nil {
+			t.Fatalf("NewServer() error = %v", err)
+		}
+		if srv == nil {
+			t.Fatal("NewServer() returned nil server")
+		}
+		defer srv.Close()
+		if srv.Addr == "" {
+			t.Error("NewServer() returned server with empty Addr")
+		}
+	})
+
+	t.Run("Unix", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		sockPath := tmpDir + "/bttest.sock"
+		srv, err := NewServer(sockPath)
+		if err != nil {
+			t.Fatalf("NewServer() error = %v", err)
+		}
+		if srv == nil {
+			t.Fatal("NewServer() returned nil server")
+		}
+
+		if srv.Addr != sockPath {
+			t.Errorf("srv.Addr got %q, want %q", srv.Addr, sockPath)
+		}
+
+		srv.Close()
+
+		// Check that the unix socket file was removed.
+		if _, err := os.Stat(sockPath); !os.IsNotExist(err) {
+			t.Errorf("socket file %q was not removed after Close()", sockPath)
+		}
+	})
+
+	t.Run("WithListener", func(t *testing.T) {
+		l, err := net.Listen("tcp", "localhost:0")
+		if err != nil {
+			t.Fatalf("net.Listen() error = %v", err)
+		}
+
+		// Create a spy that can tell us if its Close method was called.
+		spy := &closeSpyListener{Listener: l}
+		srv, err := NewServerWithListener(spy) // This should not call spy.Close().
+		if err != nil {
+			l.Close()
+			t.Fatalf("NewServerWithListener() error = %v", err)
+		}
+
+		srv.Close() // This should NOT call the Close method on our spy.
+		if spy.closed {
+			t.Error("Listener was closed by the server, but it should not have been.")
+		}
+
+		// Clean up the real listener now that the test is done.
+		l.Close()
+
+		// Validate that the listener is now actually closed.
+		_, err = l.Accept()
+		if err == nil {
+			t.Fatal("l.Accept() should have failed for a closed listener, but it did not")
+		}
+		if !errors.Is(err, net.ErrClosed) {
+			t.Errorf("Expected net.ErrClosed, but got a different error: %v", err)
+		}
+	})
+}
+
 func TestConcurrentMutationsReadModifyAndGC(t *testing.T) {
 	s := &server{
 		tables: make(map[string]*table),
