Add support for systemd (#110)

* Add support for systemd

Tweek systemd

* Add exec start args

* Fix data race

build/systemd/wayback.service

@@ -3,9 +3,9 @@ Description=Wayback webpage
 After=network.target
 
 [Service]
-Type=simple
+Type=notify
 User=wayback
-ExecStart=/usr/bin/wayback
+ExecStart=/usr/bin/wayback -d web
 Restart=always
 
 # https://www.freedesktop.org/software/systemd/man/systemd.exec.html#NoNewPrivileges=

cmd/wayback/serve.go

@@ -20,6 +20,7 @@ import (
 	"github.com/wabarc/wayback/service/telegram"
 	"github.com/wabarc/wayback/service/twitter"
 	"github.com/wabarc/wayback/storage"
+	"github.com/wabarc/wayback/systemd"
 )
 
 type target struct {
@@ -45,9 +46,17 @@ func serve(_ *cobra.Command, _ []string) {
 	srv := &service{}
 	_ = srv.run(ctx, store, pool)
 
-	go srv.stop(cancel)
+	if systemd.HasNotifySocket() {
+		logger.Info("sending readiness notification to Systemd")
+
+		if err := systemd.SdNotify(systemd.SdNotifyReady); err != nil {
+			logger.Error("unable to send readiness notification to systemd: %v", err)
+		}
+	}
 
+	go srv.stop(cancel)
 	<-ctx.Done()
+
 	logger.Info("wayback service stopped.")
 }
 
@@ -170,8 +179,10 @@ func (srv *service) stop(cancel context.CancelFunc) {
 		sig := <-signalChan
 		if sig == os.Interrupt {
 			logger.Info("Signal SIGINT is received, probably due to `Ctrl-C`, exiting...")
-			once.Do(func() { srv.shutdown() })
-			cancel()
+			once.Do(func() {
+				srv.shutdown()
+				cancel()
+			})
 			return
 		}
 	}

service/discord/discord_test.go

@@ -184,6 +184,7 @@ func TestServe(t *testing.T) {
 	d := &Discord{ctx: ctx, bot: bot, pool: pool}
 	time.AfterFunc(3*time.Second, func() {
 		d.Shutdown()
+		time.Sleep(3)
 		cancel()
 	})
 	got := d.Serve()

service/telegram/telegram.go

@@ -33,6 +33,8 @@ import (
 // ErrServiceClosed is returned by the Service's Serve method after a call to Shutdown.
 var ErrServiceClosed = errors.New("telegram: Service closed")
 
+var pollTick = 3 * time.Second
+
 // Telegram represents a Telegram service in the application.
 type Telegram struct {
 	ctx context.Context
@@ -57,7 +59,7 @@ func New(ctx context.Context, store *storage.Storage, pool pooling.Pool) *Telegr
 		Token: config.Opts.TelegramToken(),
 		// Verbose:   config.Opts.HasDebugMode(),
 		ParseMode: telegram.ModeHTML,
-		Poller:    &telegram.LongPoller{Timeout: 3 * time.Second},
+		Poller:    &telegram.LongPoller{Timeout: pollTick},
 		Reporter: func(err error) {
 			if err != nil {
 				logger.Warn(err.Error())

service/telegram/telegram_test.go

@@ -218,24 +218,21 @@ func TestServe(t *testing.T) {
 	defer server.Close()
 	handle(mux, `{"ok":true, "result":[]}`)
 
-	bot, err := telegram.NewBot(telegram.Settings{
-		URL:    server.URL,
-		Token:  token,
-		Client: httpClient,
-		Poller: &telegram.LongPoller{Timeout: time.Second},
-	})
+	pool := pooling.New(config.Opts.PoolingSize())
+	defer pool.Close()
+
+	tg, cancel, err := newTelegram(httpClient, server.URL)
 	if err != nil {
 		t.Fatal(err)
 	}
+	if tg.store == nil {
+		t.Fatalf("Open storage failed: %v", err)
+	}
+	defer tg.store.Close()
 
-	ctx, cancel := context.WithCancel(context.Background())
-
-	pool := pooling.New(config.Opts.PoolingSize())
-	defer pool.Close()
-
-	tg := &Telegram{ctx: ctx, bot: bot, pool: pool}
-	time.AfterFunc(3*time.Second, func() {
+	time.AfterFunc(pollTick, func() {
 		tg.Shutdown()
+		time.Sleep(time.Second)
 		cancel()
 	})
 
@@ -261,8 +258,6 @@ func TestProcess(t *testing.T) {
 		t.Fatalf("Parse environment variables or flags failed, error: %v", err)
 	}
 
-	done := make(chan bool, 1)
-
 	httpClient, mux, server := helper.MockServer()
 	defer server.Close()
 	handle(mux, getUpdatesJSON)
@@ -276,6 +271,7 @@ func TestProcess(t *testing.T) {
 	}
 	defer tg.store.Close()
 
+	done := make(chan bool, 1)
 	tg.bot.Poller = telegram.NewMiddlewarePoller(tg.bot.Poller, func(update *telegram.Update) bool {
 		switch {
 		// case update.Callback != nil:
@@ -323,8 +319,6 @@ func TestProcessPlayback(t *testing.T) {
 		t.Fatalf("Parse environment variables or flags failed, error: %v", err)
 	}
 
-	done := make(chan bool, 1)
-
 	getUpdatesJSON = `{
   "ok": true,
   "result": [
@@ -367,6 +361,7 @@ func TestProcessPlayback(t *testing.T) {
 	}
 	defer tg.store.Close()
 
+	done := make(chan bool, 1)
 	tg.bot.Poller = telegram.NewMiddlewarePoller(tg.bot.Poller, func(update *telegram.Update) bool {
 		switch {
 		// case update.Callback != nil:

systemd/doc.go

@@ -0,0 +1,9 @@
+// Copyright 2021 Wayback Archiver. All rights reserved.
+// Use of this source code is governed by the GNU GPL v3
+// license that can be found in the LICENSE file.
+
+/*
+Package systemd provides a Go implementation of the sd_notify protocol.
+It can be used to inform systemd of service start-up completion.
+*/
+package systemd // import "github.com/wabarc/wayback/systemd"

systemd/systemd.go

@@ -0,0 +1,46 @@
+// Copyright 2021 Wayback Archiver. All rights reserved.
+// Use of this source code is governed by the GNU GPL v3
+// license that can be found in the LICENSE file.
+
+package systemd // import "github.com/wabarc/wayback/systemd"
+
+import (
+	"net"
+	"os"
+)
+
+// SdNotifyReady tells the service manager that service startup is
+// finished, or the service finished loading its configuration.
+// https://www.freedesktop.org/software/systemd/man/sd_notify.html#READY=1
+const SdNotifyReady = "READY=1"
+
+// HasNotifySocket checks if the process is supervised by Systemd and has the notify socket.
+func HasNotifySocket() bool {
+	return os.Getenv("NOTIFY_SOCKET") != ""
+}
+
+// SdNotify sends a message to systemd using the sd_notify protocol.
+// See https://www.freedesktop.org/software/systemd/man/sd_notify.html.
+func SdNotify(state string) error {
+	addr := &net.UnixAddr{
+		Net:  "unixgram",
+		Name: os.Getenv("NOTIFY_SOCKET"),
+	}
+
+	// We're not running under systemd (NOTIFY_SOCKET is not set).
+	if addr.Name == "" {
+		return nil
+	}
+
+	conn, err := net.DialUnix(addr.Net, nil, addr)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	if _, err = conn.Write([]byte(state)); err != nil {
+		return err
+	}
+
+	return nil
+}

systemd/systemd_test.go

@@ -0,0 +1,58 @@
+// Copyright 2021 Wayback Archiver. All rights reserved.
+// Use of this source code is governed by the GNU GPL v3
+// license that can be found in the LICENSE file.
+
+package systemd // import "github.com/wabarc/wayback/systemd"
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net"
+	"os"
+	"testing"
+)
+
+func TestSdNotify(t *testing.T) {
+	testDir, err := ioutil.TempDir("", "test-")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(testDir)
+
+	notifySocket := testDir + "/notify-socket.sock"
+	laddr := net.UnixAddr{
+		Name: notifySocket,
+		Net:  "unixgram",
+	}
+	if _, err := net.ListenUnixgram("unixgram", &laddr); err != nil {
+		t.Fatal(err)
+	}
+
+	tests := []struct {
+		envSocket string
+
+		werr bool
+	}{
+		// (nil) - notification supported, data has been sent
+		{notifySocket, false},
+		// (err) - notification supported, but failure happened
+		{testDir + "/missing.sock", true},
+		// (nil) - notification not supported
+		{"", false},
+	}
+
+	for i, tt := range tests {
+		os.Unsetenv("NOTIFY_SOCKET")
+
+		if tt.envSocket != "" {
+			os.Setenv("NOTIFY_SOCKET", tt.envSocket)
+		}
+		err := SdNotify(fmt.Sprintf("TestSdNotify test message #%d", i))
+
+		if tt.werr && err == nil {
+			t.Errorf("#%d: want non-nil err, got nil", i)
+		} else if !tt.werr && err != nil {
+			t.Errorf("#%d: want nil err, got %v", i, err)
+		}
+	}
+}