Add shutdown method per component

In the current implementation, we provide the user the ability to
initialize things by implementing an `Init` method. However, they don't
have the ability to do a clean shutdown easily.

The user filled multiple github issues (#275 and #765) asking for a `Shutdown`
method.

This PR adds a `Shutdown` method to each component. Note that there is
no guarantee that this method can run - e.g., the app receives a
SIGKILL because the machine suddenly crashes. However, a `Shutdown`
method can enable the user to achieve graceful shutdown in normal
scenarios.

In tests, if someone wants to test the `Shutdown` method, they can get a
pointer to the implementation of the component and explicitly call the
`Shutdown` method.

Author: rgrandl

examples/chat/sqlstore.go

@@ -96,7 +96,7 @@ func (cfg *config) Validate() error {
 	return nil
 }
 
-func (s *sqlStore) Init(ctx context.Context) error {
+func (s *sqlStore) Init(context.Context) error {
 	cfg := s.Config()
 	if cfg.Driver == "" {
 		return fmt.Errorf("missing database driver in config")

internal/tool/generate/generator.go

@@ -602,12 +602,12 @@ func extractComponent(opt Options, pkg *packages.Package, file *ast.File, tset *
 		seenLis[lis] = struct{}{}
 	}
 
-	// Warn the user if the component has a mistyped Init method. Init methods
-	// are supposed to have type "func(context.Context) error", but it's easy
+	// Warn the user if the component has a mistyped Init or Shutdown method. These
+	// methods are supposed to have type "func(context.Context) error", but it's easy
 	// to forget to add a context.Context argument or error return. Without
-	// this warning, the component's Init method will be silently ignored. This
-	// can be very frustrating to debug.
-	if err := checkMistypedInit(pkg, tset, impl); err != nil {
+	// this warning, the component's Init or Shutdown method will be silently ignored.
+	// This can be very frustrating to debug.
+	if err := checkMistypedInitOrShutdown(pkg, tset, impl); err != nil {
 		opt.Warn(err)
 	}
 
@@ -789,27 +789,27 @@ func validateMethods(pkg *packages.Package, tset *typeSet, intf *types.Named) er
 	return errors.Join(errs...)
 }
 
-// checkMistypedInit returns an error if the provided component implementation
-// has an Init method that does not have type "func(context.Context) error".
-func checkMistypedInit(pkg *packages.Package, tset *typeSet, impl *types.Named) error {
+// checkMistypedInitOrShutdown returns an error if the provided component implementation
+// has an Init or a Shutdown method that does not have type "func(context.Context) error".
+func checkMistypedInitOrShutdown(pkg *packages.Package, tset *typeSet, impl *types.Named) error {
 	for i := 0; i < impl.NumMethods(); i++ {
 		m := impl.Method(i)
-		if m.Name() != "Init" {
+		if m.Name() != "Init" && m.Name() != "Shutdown" {
 			continue
 		}
 
 		// TODO(mwhittaker): Highlight the warning yellow instead of red.
 		sig := m.Type().(*types.Signature)
 		err := errorf(pkg.Fset, m.Pos(),
-			`WARNING: Component %v's Init method has type "%v", not type "func(context.Context) error". It will be ignored. See https://serviceweaver.dev/docs.html#components-implementation for more information.`,
-			impl.Obj().Name(), sig)
+			`WARNING: Component %v's %s method has type "%v", not type "func(context.Context) error". It will be ignored. See https://serviceweaver.dev/docs.html#components-implementation for more information.`,
+			impl.Obj().Name(), m.Name(), sig)
 
-		// Check Init's parameters.
+		// Check parameters.
 		if sig.Params().Len() != 1 || !isContext(sig.Params().At(0).Type()) {
 			return err
 		}
 
-		// Check Init's returns.
+		// Check returns.
 		if sig.Results().Len() != 1 || sig.Results().At(0).Type().String() != "error" {
 			return err
 		}

internal/weaver/remoteweavelet.go

@@ -23,10 +23,12 @@ import (
 	"log/slog"
 	"net"
 	"os"
+	"os/signal"
 	"reflect"
 	"strings"
 	"sync"
 	"sync/atomic"
+	"syscall"
 
 	"github.com/ServiceWeaver/weaver/internal/config"
 	"github.com/ServiceWeaver/weaver/internal/control"
@@ -286,6 +288,26 @@ func NewRemoteWeavelet(ctx context.Context, regs []*codegen.Registration, bootst
 		return nil
 	})
 
+	// Start a signal handler to detect when the process is killed. This isn't
+	// perfect, as we can't catch a SIGKILL, but it's good in the common case.
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-done
+		for _, c := range w.componentsByName {
+			if !c.implReady.Load() {
+				continue
+			}
+			// Call Shutdown method if available.
+			if i, ok := c.impl.(interface{ Shutdown(context.Context) error }); ok {
+				if err := i.Shutdown(ctx); err != nil {
+					w.syslogger.Error("Component shutdown failed", "component", c.reg.Name, "err", err)
+				}
+			}
+		}
+		os.Exit(1)
+	}()
+
 	w.syslogger.Debug("🧶 weavelet started", "addr", dialAddr)
 	return w, nil
 }
@@ -537,7 +559,7 @@ func (w *RemoteWeavelet) GetLoad(context.Context, *protos.GetLoadRequest) (*prot
 }
 
 // UpdateComponents implements weaver.controller and conn.WeaverHandler interfaces.
-func (w *RemoteWeavelet) UpdateComponents(ctx context.Context, req *protos.UpdateComponentsRequest) (*protos.UpdateComponentsReply, error) {
+func (w *RemoteWeavelet) UpdateComponents(_ context.Context, req *protos.UpdateComponentsRequest) (*protos.UpdateComponentsReply, error) {
 	var errs []error
 	var components []*component
 	var shortened []string
@@ -578,7 +600,7 @@ func (w *RemoteWeavelet) UpdateComponents(ctx context.Context, req *protos.Updat
 }
 
 // UpdateRoutingInfo implements controller.UpdateRoutingInfo.
-func (w *RemoteWeavelet) UpdateRoutingInfo(ctx context.Context, req *protos.UpdateRoutingInfoRequest) (reply *protos.UpdateRoutingInfoReply, err error) {
+func (w *RemoteWeavelet) UpdateRoutingInfo(_ context.Context, req *protos.UpdateRoutingInfoRequest) (reply *protos.UpdateRoutingInfoReply, err error) {
 	if req.RoutingInfo == nil {
 		w.syslogger.Error("Failed to update nil routing info")
 		return nil, fmt.Errorf("nil RoutingInfo")
@@ -643,7 +665,7 @@ func (w *RemoteWeavelet) UpdateRoutingInfo(ctx context.Context, req *protos.Upda
 }
 
 // GetHealth implements controller.GetHealth.
-func (w *RemoteWeavelet) GetHealth(ctx context.Context, req *protos.GetHealthRequest) (*protos.GetHealthReply, error) {
+func (w *RemoteWeavelet) GetHealth(context.Context, *protos.GetHealthRequest) (*protos.GetHealthReply, error) {
 	// Get the health status for all components. For now, we consider a component
 	// healthy iff it has been successfully initialized. In the future, we will
 	// maintain a real-time health for each component.
@@ -657,7 +679,7 @@ func (w *RemoteWeavelet) GetHealth(ctx context.Context, req *protos.GetHealthReq
 }
 
 // GetMetrics implements controller.GetMetrics.
-func (w *RemoteWeavelet) GetMetrics(ctx context.Context, req *protos.GetMetricsRequest) (*protos.GetMetricsReply, error) {
+func (w *RemoteWeavelet) GetMetrics(context.Context, *protos.GetMetricsRequest) (*protos.GetMetricsReply, error) {
 	// TODO(sanjay): The protocol is currently brittle; if we ever lose a set of
 	// updates, they will be lost forever. Fix by versioning the "last" map in
 	// metrics.Exporter. The reader echoes back the version of the last set of

internal/weaver/singleweavelet.go

@@ -132,7 +132,7 @@ func NewSingleWeavelet(ctx context.Context, regs []*codegen.Registration, opts S
 		fmt.Fprint(os.Stderr, reg.Rolodex())
 	}
 
-	return &SingleWeavelet{
+	w := &SingleWeavelet{
 		ctx:          ctx,
 		regs:         regs,
 		regsByName:   regsByName,
@@ -149,7 +149,29 @@ func NewSingleWeavelet(ctx context.Context, regs []*codegen.Registration, opts S
 		stats:        imetrics.NewStatsProcessor(),
 		components:   map[string]any{},
 		listeners:    map[string]net.Listener{},
-	}, nil
+	}
+
+	// Start a signal handler to detect when the process is killed. This isn't
+	// perfect, as we can't catch a SIGKILL, but it's good in the common case.
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-done
+
+		w.mu.Lock()
+		defer w.mu.Unlock()
+		for c, impl := range w.components {
+			// Call Shutdown method if available.
+			if i, ok := impl.(interface{ Shutdown(context.Context) error }); ok {
+				if err := i.Shutdown(ctx); err != nil {
+					fmt.Printf("Component %s failed to shutdown: %v\n", c, err)
+				}
+			}
+		}
+		os.Exit(1)
+	}()
+
+	return w, nil
 }
 
 // parseSingleConfig parses the "[single]" section of a config file.

website/docs.md

@@ -636,6 +636,20 @@ func (f *foo) Init(context.Context) error {
 }
 ```
 
+If a component implementation implements an `Shutdown(context.Context) error`
+method, it will be called when an instance of the component is destroyed.
+
+```go
+func (f *foo) Shutdown(context.Context) error {
+    // ...
+}
+```
+
+**Note**: There is no guarantee that the `Shutdown` method will always
+be called. `Shutdown` is called **iff** your application receives a
+`SIGINT` or a `SIGTERM` signal. However, if the machine where your application runs
+crashes unexpectedly or becomes unresponsive, the `Shutdown` method is never called.
+
 ## Semantics
 
 When implementing a component, there are a few semantic details to keep in mind: