feat: Add support for systemd notify (#425)

cmd/root.go

@@ -39,6 +39,7 @@ import (
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/healthcheck"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/log"
 	"github.com/GoogleCloudPlatform/alloydb-auth-proxy/internal/proxy"
+	"github.com/coreos/go-systemd/daemon"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
@@ -707,9 +708,23 @@ func runSignalWrapper(cmd *Command) (err error) {
 	select {
 	case err := <-shutdownCh:
 		cmd.logger.Errorf("The proxy has encountered a terminal error: %v", err)
+		// If running under systemd with Type=notify, it will send a message to the
+		// service manager that a failure occurred and it is terminating.
+		go func() {
+			if _, err := daemon.SdNotify(false, daemon.SdNotifyStopping); err != nil {
+				cmd.logger.Errorf("Failed to notify systemd of termination: %v", err)
+			}
+		}()
 		return err
 	case p = <-startCh:
 		cmd.logger.Infof("The proxy has started successfully and is ready for new connections!")
+		// If running under systemd with Type=notify, it will send a message to the
+		// service manager that it is ready to handle connections now.
+		go func() {
+			if _, err := daemon.SdNotify(false, daemon.SdNotifyReady); err != nil {
+				cmd.logger.Errorf("Failed to notify systemd of readiness: %v", err)
+			}
+		}()
 	}
 	defer func() {
 		if cErr := p.Close(); cErr != nil {

cmd/root_linux_test.go

@@ -15,10 +15,14 @@
 package cmd
 
 import (
+	"context"
+	"net"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
+	"github.com/coreos/go-systemd/daemon"
 	"github.com/spf13/cobra"
 )
 
@@ -71,3 +75,72 @@ func TestNewCommandArgumentsOnLinux(t *testing.T) {
 		})
 	}
 }
+
+func TestSdNotifyOnLinux(t *testing.T) {
+	tcs := []struct {
+		desc          string
+		proxyMustFail bool
+		notifyState   string
+	}{
+		{
+			desc:          "System with systemd Type=notify and proxy started successfully",
+			proxyMustFail: false,
+			notifyState:   daemon.SdNotifyReady,
+		},
+		{
+			desc:          "System with systemd Type=notify and proxy failed to start",
+			proxyMustFail: true,
+			notifyState:   daemon.SdNotifyStopping,
+		},
+	}
+
+	// Create a temp dir for the socket file.
+	testDir, err := os.MkdirTemp("/tmp/", "test-")
+	if err != nil {
+		t.Fatalf("Fail to create the temp dir: %v", err)
+	}
+	defer os.RemoveAll(testDir)
+
+	//Set up the socket stream to listen for notifications.
+	socketAddr := filepath.Join(testDir, "notify-socket.sock")
+	conn, err := net.ListenUnixgram("unixgram", &net.UnixAddr{Name: socketAddr, Net: "unixgram"})
+	if err != nil {
+		t.Fatalf("net.ListenUnixgram error: %v", err)
+	}
+
+	// To simulate systemd behavior with Type=notify, set NOTIFY_SOCKET
+	// to the name of the socket that listens for notifications.
+	os.Setenv("NOTIFY_SOCKET", socketAddr)
+	defer os.Unsetenv("NOTIFY_SOCKET")
+
+	s := &spyDialer{}
+	c := NewCommand(WithDialer(s))
+	// Keep the test output quiet
+	c.SilenceUsage = false
+	c.SilenceErrors = false
+	c.SetArgs([]string{"projects/proj/locations/region/clusters/clust/instances/inst"})
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+
+			if tc.proxyMustFail {
+				c.conf.FUSEDir = "invalid"
+			}
+
+			go c.ExecuteContext(ctx)
+
+			stateReceived := make([]byte, 4096)
+			length, _, err := conn.ReadFromUnix(stateReceived)
+			if err != nil {
+				t.Fatalf("conn.ReadFromUnix error: %s\n", err)
+			}
+			if string(stateReceived[0:length]) != tc.notifyState {
+				t.Fatalf("Expected Notify State %v, got %v", tc.notifyState, string(stateReceived))
+			}
+
+		})
+	}
+}

go.mod

@@ -6,6 +6,7 @@ require (
 	cloud.google.com/go/alloydbconn v1.3.4
 	contrib.go.opencensus.io/exporter/prometheus v0.4.2
 	contrib.go.opencensus.io/exporter/stackdriver v0.13.14
+	github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf
 	github.com/google/go-cmp v0.5.9
 	github.com/hanwen/go-fuse/v2 v2.3.0
 	github.com/jackc/pgx/v5 v5.4.3

go.sum

@@ -322,6 +322,8 @@ github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3Ee
 github.com/coreos/go-systemd v0.0.0-20161114122254-48702e0da86b/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/go-systemd v0.0.0-20180511133405-39ca1b05acc7/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
+github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf h1:iW4rZ826su+pqaw19uhpSCzhj44qo35pNgKFGqzDKkU=
+github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/go-systemd/v22 v22.0.0/go.mod h1:xO0FLkIi5MaZafQlIrOotqXZ90ih+1atmu1JpKERPPk=
 github.com/coreos/go-systemd/v22 v22.1.0/go.mod h1:xO0FLkIi5MaZafQlIrOotqXZ90ih+1atmu1JpKERPPk=
 github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=