Merge pull request #1 from mhilton/001-service-implementation

go-service implementation

README

@@ -0,0 +1,7 @@
+Go-service provides a helper for long-running services that require
+a mechanism for graceful shutdown. The service starts shutting down
+either when a component stops and returns an error, or a chosen signal
+is received. A service can call registered functions when shutting down.
+
+The package documentation for go-service can be found at
+https://pkg.go.dev/github.com/canonical/go-service.

go.mod

@@ -0,0 +1,5 @@
+module github.com/canonical/go-service
+
+go 1.16
+
+require golang.org/x/sync v0.0.0-20210220032951-036812b2e83c

go.sum

@@ -0,0 +1,2 @@
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=

service.go

@@ -0,0 +1,103 @@
+// Copyright 2021 Canonical Ltd.
+
+// Package service provides helpers for long-running service applications.
+package service
+
+import (
+	"context"
+	"os"
+	"os/signal"
+
+	"golang.org/x/sync/errgroup"
+)
+
+// A Service is a service provided by a number of goroutines which will
+// initiate a graceful shutdown when either one of those goroutines errors,
+// or on the receipt of chosen signals.
+type Service struct {
+	g *errgroup.Group
+
+	doneC     <-chan struct{}
+	shutdownC chan<- func()
+}
+
+// NewService creates a new service instance using the given context. If
+// any signals are specified the service will start a shutdown upon
+// receiving that signal.
+func NewService(ctx context.Context, sig ...os.Signal) (context.Context, *Service) {
+	g, ctx := errgroup.WithContext(ctx)
+
+	if len(sig) > 0 {
+		sigC := make(chan os.Signal, 1)
+		g.Go(func() error {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case sig := <-sigC:
+				return &SignalError{
+					Signal: sig,
+				}
+			}
+		})
+		signal.Notify(sigC, sig...)
+	}
+
+	shutdownC := make(chan func())
+	g.Go(func() error {
+		var funcs []func()
+		for {
+			select {
+			case f := <-shutdownC:
+				funcs = append(funcs, f)
+			case <-ctx.Done():
+				for i := len(funcs) - 1; i >= 0; i-- {
+					funcs[i]()
+				}
+				return ctx.Err()
+			}
+		}
+	})
+
+	return ctx, &Service{
+		g:         g,
+		doneC:     ctx.Done(),
+		shutdownC: shutdownC,
+	}
+}
+
+// Go calls the given function in a new goroutine.
+//
+// The first call to return a non-nil error cancels the service; its error
+// will be returned by Wait.
+func (s *Service) Go(f func() error) {
+	s.g.Go(f)
+}
+
+// Wait waits for all goroutines started by this service and all functions
+// registered with OnShutdown to complete. The error returned will be the
+// error that caused the service to be canceled, if any.
+func (s *Service) Wait() error {
+	return s.g.Wait()
+}
+
+// OnShutdown registers a function to be called when the service determines
+// it is shutting down. The Wait function will wait for all functions
+// provided to OnShutdown to complete before returning.
+func (s *Service) OnShutdown(f func()) {
+	select {
+	case s.shutdownC <- f:
+	case <-s.doneC:
+		f()
+	}
+}
+
+// A SignalError is the type of error returned when a Service has shutdown
+// due to receiving a signal.
+type SignalError struct {
+	Signal os.Signal
+}
+
+// Error implements the error interface.
+func (e *SignalError) Error() string {
+	return "received " + e.Signal.String()
+}

service_test.go

@@ -0,0 +1,81 @@
+// Copyright 2021 Canonical Ltd.
+
+package service
+
+import (
+	"context"
+	"errors"
+	"os"
+	"sync"
+	"syscall"
+	"testing"
+)
+
+func TestSignal(t *testing.T) {
+	ctx, svc := NewService(context.Background(), syscall.SIGUSR1)
+	svc.Go(func() error {
+		p, err := os.FindProcess(os.Getpid())
+		if err != nil {
+			return err
+		}
+		if err := p.Signal(syscall.SIGUSR1); err != nil {
+			return err
+		}
+		<-ctx.Done()
+		return nil
+	})
+	err := svc.Wait()
+	if err.Error() != "received user defined signal 1" {
+		t.Error("unexpected error:", err)
+	}
+}
+
+func TestServiceError(t *testing.T) {
+	_, svc := NewService(context.Background(), syscall.SIGUSR2)
+	svc.Go(func() error {
+		return errors.New("test error")
+	})
+	err := svc.Wait()
+	if err.Error() != "test error" {
+		t.Error("unexpected error:", err)
+	}
+}
+
+func TestOnShutdown(t *testing.T) {
+	_, svc := NewService(context.Background())
+	var mu sync.Mutex
+	var ops []string
+	svc.OnShutdown(func() {
+		mu.Lock()
+		defer mu.Unlock()
+		ops = append(ops, "shutdown-1")
+	})
+	svc.OnShutdown(func() {
+		mu.Lock()
+		defer mu.Unlock()
+		ops = append(ops, "shutdown-2")
+	})
+	svc.Go(func() error {
+		mu.Lock()
+		defer mu.Unlock()
+		ops = append(ops, "go-1")
+		return errors.New("test error")
+	})
+	err := svc.Wait()
+	if err.Error() != "test error" {
+		t.Error("unexpected error:", err)
+	}
+	svc.OnShutdown(func() {
+		mu.Lock()
+		defer mu.Unlock()
+		ops = append(ops, "shutdown-3")
+	})
+	mu.Lock()
+	defer mu.Unlock()
+	if len(ops) != 4 {
+		t.Fatal("unexpected operations occured:", ops)
+	}
+	if ops[0] != "go-1" {
+		t.Fatal("shutdown operations happened too early", ops)
+	}
+}