feat: add gorums.System and proper io.Closer support

This adds the System helper for registration, starting and stopping
servers and closing connections. This also implements proper io.Closer
support in the Manager's Close() method, and prevents duplicate calls to
the channel's close().

Author: meling

channel.go

@@ -81,6 +81,7 @@ type channel struct {
 	// to the caller.
 	responseRouters map[uint64]request
 	responseMut     sync.Mutex
+	closeOnceFunc   func() error
 }
 
 // newChannel creates a new channel for the given node and starts the sender
@@ -102,19 +103,22 @@ func newChannel(parentCtx context.Context, conn *grpc.ClientConn, id uint32, sen
 		responseRouters: make(map[uint64]request),
 		streamReady:     make(chan struct{}, 1),
 	}
+	c.closeOnceFunc = sync.OnceValue(func() error {
+		// important to cancel first to stop goroutines
+		c.connCancel()
+		if c.conn != nil {
+			return c.conn.Close()
+		}
+		return nil
+	})
 	go c.sender()
 	go c.receiver()
 	return c
 }
 
-// close closes the channel and the underlying connection.
+// close closes the channel and the underlying connection exactly once.
 func (c *channel) close() error {
-	// important to cancel first to stop goroutines
-	c.connCancel()
-	if c.conn != nil {
-		return c.conn.Close()
-	}
-	return nil
+	return c.closeOnceFunc()
 }
 
 // newNodeStream creates a new stream for this channel.

channel_test.go

@@ -82,7 +82,7 @@ func testNodeWithoutServer(t testing.TB, opts ...ManagerOption) *Node {
 	t.Helper()
 	mgrOpts := append([]ManagerOption{InsecureDialOptions(t)}, opts...)
 	mgr := NewManager(mgrOpts...)
-	t.Cleanup(mgr.Close)
+	t.Cleanup(Closer(t, mgr))
 	// Use a high port number that's unlikely to have anything listening.
 	// We use a fixed ID for simplicity.
 	node, err := mgr.newNode("127.0.0.1:59999", 1)

config_test.go

@@ -26,7 +26,7 @@ var (
 func TestNewConfigurationEmptyNodeList(t *testing.T) {
 	wantErr := errors.New("config: missing required node addresses")
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	_, err := gorums.NewConfiguration(mgr, gorums.WithNodeList([]string{}))
 	if err == nil {
@@ -39,7 +39,7 @@ func TestNewConfigurationEmptyNodeList(t *testing.T) {
 
 func TestNewConfigurationNodeList(t *testing.T) {
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	cfg, err := gorums.NewConfiguration(mgr, gorums.WithNodeList(nodes))
 	if err != nil {
@@ -77,7 +77,7 @@ func TestNewConfigurationNodeList(t *testing.T) {
 
 func TestNewConfigurationNodeMap(t *testing.T) {
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	cfg, err := gorums.NewConfiguration(mgr, gorums.WithNodeMap(nodeMap))
 	if err != nil {
@@ -103,7 +103,7 @@ func TestNewConfigurationNodeMap(t *testing.T) {
 
 func TestNewConfigurationNodeIDs(t *testing.T) {
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	c1, err := gorums.NewConfiguration(mgr, gorums.WithNodeList(nodes))
 	if err != nil {
@@ -141,7 +141,7 @@ func TestNewConfigurationNodeIDs(t *testing.T) {
 
 func TestNewConfigurationAnd(t *testing.T) {
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	c1, err := gorums.NewConfiguration(mgr, gorums.WithNodeList(nodes))
 	if err != nil {
@@ -195,7 +195,7 @@ func TestNewConfigurationAnd(t *testing.T) {
 
 func TestNewConfigurationExcept(t *testing.T) {
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	c1, err := gorums.NewConfiguration(mgr, gorums.WithNodeList(nodes))
 	if err != nil {
@@ -254,7 +254,7 @@ func TestConfigConcurrentAccess(t *testing.T) {
 
 func TestConfigurationWithoutErrors(t *testing.T) {
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	cfg, err := gorums.NewConfiguration(mgr, gorums.WithNodeMap(nodeMap))
 	if err != nil {

mgr.go

@@ -1,6 +1,7 @@
 package gorums
 
 import (
+	"errors"
 	"fmt"
 	"log"
 	"sync"
@@ -50,23 +51,15 @@ func NewManager(opts ...ManagerOption) *Manager {
 	return m
 }
 
-func (m *Manager) closeNodeConns() {
-	for _, node := range m.nodes {
-		err := node.close()
-		if err != nil && m.logger != nil {
-			m.logger.Printf("error closing: %v", err)
-		}
-	}
-}
-
 // Close closes all node connections and any client streams.
-func (m *Manager) Close() {
+func (m *Manager) Close() error {
+	var err error
 	m.closeOnce.Do(func() {
-		if m.logger != nil {
-			m.logger.Printf("closing")
+		for _, node := range m.nodes {
+			err = errors.Join(err, node.close())
 		}
-		m.closeNodeConns()
 	})
+	return err
 }
 
 // NodeIDs returns the identifier of each available node. IDs are returned in

mgr_test.go

@@ -23,17 +23,17 @@ func TestManagerLogging(t *testing.T) {
 		logger = log.New(&buf, "logger: ", log.Lshortfile)
 	)
 	mgr := NewManager(InsecureDialOptions(t), WithLogger(logger))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(Closer(t, mgr))
 
-	want := "logger: mgr.go:48: ready"
+	want := "logger: mgr.go:49: ready"
 	if strings.TrimSpace(buf.String()) != want {
 		t.Errorf("logger: got %q, want %q", buf.String(), want)
 	}
 }
 
 func TestManagerNewNode(t *testing.T) {
 	mgr := NewManager(InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(Closer(t, mgr))
 
 	_, err := NewConfiguration(mgr, WithNodeMap(nodeMap))
 	if err != nil {
@@ -62,7 +62,7 @@ func TestManagerNewNode(t *testing.T) {
 func TestManagerNewNodeWithConn(t *testing.T) {
 	addrs := TestServers(t, 3, nil)
 	mgr := NewManager(InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(Closer(t, mgr))
 
 	// Create configuration with only first 2 nodes
 	_, err := NewConfiguration(mgr, WithNodeList(addrs[:2]))

server_test.go

@@ -121,7 +121,7 @@ func TestTCPReconnection(t *testing.T) {
 	}()
 
 	mgr := gorums.NewManager(gorums.InsecureDialOptions(t))
-	t.Cleanup(mgr.Close)
+	t.Cleanup(gorums.Closer(t, mgr))
 
 	cfg, err := gorums.NewConfiguration(mgr, gorums.WithNodeList([]string{addr}))
 	if err != nil {

system.go

@@ -0,0 +1,61 @@
+package gorums
+
+import (
+	"errors"
+	"io"
+	"net"
+)
+
+// System encapsulates the state of a Gorums system, including the server,
+// listener, and any registered closers (e.g. managers).
+type System struct {
+	closers []io.Closer
+	srv     *Server
+	lis     net.Listener
+}
+
+// NewSystem creates a new Gorums System with the provided listener and server options.
+func NewSystem(lis net.Listener, opts ...ServerOption) *System {
+	return &System{
+		srv: NewServer(opts...),
+		lis: lis,
+	}
+}
+
+// RegisterService registers the service with the server using the provided register function.
+// The closer is added to the list of closers to be closed when the system is stopped.
+//
+// Example usage:
+//
+//	gs := NewSystem(lis)
+//	mgr := NewManager(...)
+//	impl := &srvImpl{}
+//	gs.RegisterService(mgr, func(srv *Server) {
+//		pb.RegisterMultiPaxosServer(srv, impl)
+//	})
+func (s *System) RegisterService(closer io.Closer, registerFunc func(*Server)) {
+	if closer != nil {
+		s.closers = append(s.closers, closer)
+	}
+	registerFunc(s.srv)
+}
+
+// Serve starts the server.
+func (s *System) Serve() error {
+	return s.srv.Serve(s.lis)
+}
+
+// Stop stops the Gorums server and closes all registered closers.
+// It immediately closes all open connections and listeners. It cancels
+// all active RPCs on the server side and the corresponding pending RPCs
+// on the client side will get notified by connection errors.
+func (s *System) Stop() (errs error) {
+	// We cannot use graceful stop here since multicast methods does not
+	// respond to the client, and thus would block indefinitely.
+	// The server's listener is closed by s.srv.Stop().
+	s.srv.Stop()
+	for _, closer := range s.closers {
+		errs = errors.Join(errs, closer.Close())
+	}
+	return errs
+}

system_test.go

@@ -0,0 +1,87 @@
+package gorums_test
+
+import (
+	"errors"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/relab/gorums"
+)
+
+type mockCloser struct {
+	closed bool
+	err    error
+}
+
+func (m *mockCloser) Close() error {
+	m.closed = true
+	return m.err
+}
+
+func TestSystemLifecycle(t *testing.T) {
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to listen: %v", err)
+	}
+
+	sys := gorums.NewSystem(lis)
+	if sys == nil {
+		t.Fatal("NewSystem returned nil")
+	}
+
+	closer1 := &mockCloser{}
+	closer2 := &mockCloser{}
+
+	sys.RegisterService(closer1, func(srv *gorums.Server) {
+		// In a real scenario, we would register a Gorums service here.
+	})
+	sys.RegisterService(closer2, func(srv *gorums.Server) {
+		// Register another service or just use the callback.
+	})
+
+	go func() {
+		// Serve acts as a blocking call, so run in goroutine
+		if err := sys.Serve(); err != nil {
+			// Serve returns error on Stop usually (or net closed)
+			t.Logf("Serve returned: %v", err)
+		}
+	}()
+
+	// Give it a moment to start
+	time.Sleep(10 * time.Millisecond)
+
+	// Stop the system
+	if err := sys.Stop(); err != nil {
+		t.Errorf("Stop returned error: %v", err)
+	}
+
+	if !closer1.closed {
+		t.Error("closer1 was not closed")
+	}
+	if !closer2.closed {
+		t.Error("closer2 was not closed")
+	}
+}
+
+func TestSystemStopError(t *testing.T) {
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to listen: %v", err)
+	}
+
+	sys := gorums.NewSystem(lis)
+	errCloser := &mockCloser{err: errors.New("closer error")}
+
+	sys.RegisterService(errCloser, func(srv *gorums.Server) {})
+
+	go func() {
+		_ = sys.Serve()
+	}()
+	time.Sleep(10 * time.Millisecond)
+
+	err = sys.Stop()
+	if err == nil {
+		t.Error("expected error from Stop, got nil")
+	}
+}

testing_gorums.go

@@ -3,6 +3,7 @@ package gorums
 import (
 	"context"
 	"fmt"
+	"io"
 	"iter"
 	"maps"
 	"net"
@@ -148,6 +149,16 @@ type ServerIface interface {
 	Stop()
 }
 
+// Closer returns a cleanup function that closes the given io.Closer.
+func Closer(t testing.TB, c io.Closer) func() {
+	t.Helper()
+	return func() {
+		if err := c.Close(); err != nil {
+			t.Errorf("c.Close() = %q, expected no error", err.Error())
+		}
+	}
+}
+
 type serverState struct {
 	srv     ServerIface
 	lis     net.Listener

testopts.go

@@ -53,7 +53,7 @@ func (to *testOptions) getOrCreateManager(t testing.TB) *Manager {
 	// Create manager and register its cleanup LAST so it runs FIRST (LIFO)
 	mgrOpts := append([]ManagerOption{InsecureDialOptions(t)}, to.managerOpts...)
 	mgr := NewManager(mgrOpts...)
-	t.Cleanup(mgr.Close)
+	t.Cleanup(Closer(t, mgr))
 	return mgr
 }