Merge pull request #413 from SiaFoundation/matt/graceful-shutdown

Matt/graceful shutdown

Author: n8mgr

.changeset/rhp4_graceful_shutdown.md

@@ -0,0 +1,5 @@
+---
+default: minor
+---
+
+# RHP4 graceful shutdown

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/quic-go/quic-go v0.59.0
 	github.com/quic-go/webtransport-go v0.10.0
 	go.etcd.io/bbolt v1.4.3
-	go.sia.tech/core v0.19.0
+	go.sia.tech/core v0.19.1-0.20260323200052-1bfa3facc408
 	go.sia.tech/mux v1.4.0
 	go.uber.org/zap v1.27.1
 	golang.org/x/crypto v0.49.0

go.sum

@@ -14,8 +14,8 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 go.etcd.io/bbolt v1.4.3 h1:dEadXpI6G79deX5prL3QRNP6JB8UxVkqo4UPnHaNXJo=
 go.etcd.io/bbolt v1.4.3/go.mod h1:tKQlpPaYCVFctUIgFKFnAlvbmB3tpy1vkTnDWohtc0E=
-go.sia.tech/core v0.19.0 h1:mj/lsixiI25hNTq1FzLHs94BCewTABulkqq2pHSHmdo=
-go.sia.tech/core v0.19.0/go.mod h1:Gge/hpiE9m1ugPLz8RR1ZMoYZTPWLEdRWviHr/4rVeA=
+go.sia.tech/core v0.19.1-0.20260323200052-1bfa3facc408 h1:euiK0PJ9ChJCK/moH3mqDuhnZq8Oaz42QB0PLE+HLb8=
+go.sia.tech/core v0.19.1-0.20260323200052-1bfa3facc408/go.mod h1:HR5UCFehifHhlOc9DadUVFlMLJYRBtqwXX4UbOX8E14=
 go.sia.tech/mux v1.4.0 h1:LgsLHtn7l+25MwrgaPaUCaS8f2W2/tfvHIdXps04sVo=
 go.sia.tech/mux v1.4.0/go.mod h1:iNFi9ifFb2XhuD+LF4t2HBb4Mvgq/zIPKqwXU/NlqHA=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=

rhp/v4/rpc_test.go

@@ -5,6 +5,7 @@ import (
 	"cmp"
 	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"maps"
 	"math"
@@ -2952,3 +2953,29 @@ func BenchmarkContractUpload(b *testing.B) {
 		b.Fatalf("expected %v sectors, got %v", b.N, appendResult.Revision.Filesize/proto4.SectorSize)
 	}
 }
+
+func TestGracefulShutdown(t *testing.T) {
+	n, genesis := testutil.V2Network()
+	cm, w := startTestNode(t, n, genesis)
+
+	ss := testutil.NewEphemeralSectorStore()
+	c := testutil.NewEphemeralContractor(cm)
+	sr := testutil.NewEphemeralSettingsReporter()
+
+	hostKey := types.GeneratePrivateKey()
+	rs := rhp4.NewServer(hostKey, cm, c, w, sr, ss, rhp4.WithPriceTableValidity(2*time.Minute))
+	hostAddr := testutil.ServeSiaMux(t, rs, zap.NewNop())
+
+	transport, err := siamux.Dial(context.Background(), hostAddr, hostKey.PublicKey())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer transport.Close()
+
+	rs.Close()
+
+	_, err = rhp4.RPCSettings(context.Background(), transport)
+	if !errors.Is(err, proto4.ErrHostShuttingDown) {
+		t.Fatalf("expected ErrHostShuttingDown, got %v", err)
+	}
+}

rhp/v4/server.go

@@ -12,6 +12,7 @@ import (
 	"go.sia.tech/core/consensus"
 	rhp4 "go.sia.tech/core/rhp/v4"
 	"go.sia.tech/core/types"
+	"go.sia.tech/coreutils/threadgroup"
 	"go.sia.tech/coreutils/wallet"
 	"go.sia.tech/mux/v2"
 	"go.uber.org/zap"
@@ -128,6 +129,8 @@ type (
 
 	// A Server handles incoming RHP4 RPC.
 	Server struct {
+		tg *threadgroup.ThreadGroup
+
 		hostKey            types.PrivateKey
 		priceTableValidity time.Duration
 		rpcTimeout         time.Duration
@@ -1216,6 +1219,11 @@ func (s *Server) HostKey() types.PrivateKey {
 	return s.hostKey
 }
 
+// Close stops accepting new RPCs and waits for in-flight handlers to complete.
+func (s *Server) Close() {
+	s.tg.Stop()
+}
+
 // Serve accepts incoming streams on the provided multiplexer and handles them
 func (s *Server) Serve(t TransportMux, log *zap.Logger) error {
 	defer t.Close()
@@ -1230,6 +1238,14 @@ func (s *Server) Serve(t TransportMux, log *zap.Logger) error {
 		log := log.With(zap.String("streamID", hex.EncodeToString(frand.Bytes(4))))
 		log.Debug("accepted stream")
 		go func() {
+			done, err := s.tg.Add()
+			if err != nil {
+				log.Debug("rejected stream, server is shutting down")
+				rhp4.WriteResponse(stream, rhp4.ErrHostShuttingDown.(*rhp4.RPCError))
+				stream.Close()
+				return
+			}
+			defer done()
 			defer func() {
 				if err := stream.Close(); err != nil {
 					log.Debug("failed to close stream", zap.Error(err))
@@ -1255,6 +1271,8 @@ func errorDecodingError(f string, p ...any) error {
 // NewServer creates a new RHP4 server
 func NewServer(pk types.PrivateKey, cm ChainManager, contracts Contractor, wallet Wallet, settings Settings, sectors Sectors, opts ...ServerOption) *Server {
 	s := &Server{
+		tg: threadgroup.New(),
+
 		hostKey:            pk,
 		priceTableValidity: 30 * time.Minute,
 		rpcTimeout:         10 * time.Minute,