feat: graceful shutdown server (#1016)

bin/server/main.go

 package main
 
 import (
+	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 
 	_ "github.com/mattn/go-sqlite3"
-	"github.com/pkg/errors"
 
 	"context"
 	"fmt"
@@ -24,11 +26,11 @@ const (
 `
 )
 
-func run() error {
-	ctx := context.Background()
+func main() {
 	profile, err := profile.GetProfile()
 	if err != nil {
-		return err
+		fmt.Printf("failed to get profile, error: %+v\n", err)
+		return
 	}
 	println("---")
 	println("profile")
@@ -38,19 +40,35 @@ func run() error {
 	println("version:", profile.Version)
 	println("---")
 
-	serverInstance, err := server.NewServer(ctx, profile)
+	ctx, cancel := context.WithCancel(context.Background())
+	s, err := server.NewServer(ctx, profile)
 	if err != nil {
-		return errors.Wrap(err, "failed to start server")
+		cancel()
+		fmt.Printf("failed to create server, error: %+v\n", err)
+		return
 	}
 
+	c := make(chan os.Signal, 1)
+	// Trigger graceful shutdown on SIGINT or SIGTERM.
+	// The default signal sent by the `kill` command is SIGTERM,
+	// which is taken as the graceful shutdown signal for many systems, eg., Kubernetes, Gunicorn.
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		sig := <-c
+		fmt.Printf("%s received.\n", sig.String())
+		s.Shutdown(ctx)
+		cancel()
+	}()
+
 	println(greetingBanner)
 	fmt.Printf("Version %s has started at :%d\n", profile.Version, profile.Port)
-	return serverInstance.Start(ctx)
-}
-
-func main() {
-	if err := run(); err != nil {
-		fmt.Printf("error: %+v\n", err)
-		os.Exit(1)
+	if err := s.Start(ctx); err != nil {
+		if err != http.ErrServerClosed {
+			fmt.Printf("failed to start server, error: %+v\n", err)
+			cancel()
+		}
 	}
+
+	// Wait for CTRL-C.
+	<-ctx.Done()
 }

common/log/logger.go

+// Package log implements a simple logging package.
+package log
+
+import (
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+)
+
+var (
+	// `gl` is the global logger.
+	// Other packages should use public methods such as Info/Error to do the logging.
+	// For other types of logging, e.g. logging to a separate file, they should use their own loggers.
+	gl     *zap.Logger
+	gLevel zap.AtomicLevel
+)
+
+// Initializes the global console logger.
+func init() {
+	gLevel = zap.NewAtomicLevelAt(zap.InfoLevel)
+	gl, _ = zap.Config{
+		Level:       gLevel,
+		Development: true,
+		// Use "console" to print readable stacktrace.
+		Encoding:         "console",
+		EncoderConfig:    zap.NewDevelopmentEncoderConfig(),
+		OutputPaths:      []string{"stderr"},
+		ErrorOutputPaths: []string{"stderr"},
+	}.Build(
+		// Skip one caller stack to locate the correct caller.
+		zap.AddCallerSkip(1),
+	)
+}
+
+// SetLevel wraps the zap Level's SetLevel method.
+func SetLevel(level zapcore.Level) {
+	gLevel.SetLevel(level)
+}
+
+// EnabledLevel wraps the zap Level's Enabled method.
+func EnabledLevel(level zapcore.Level) bool {
+	return gLevel.Enabled(level)
+}
+
+// Debug wraps the zap Logger's Debug method.
+func Debug(msg string, fields ...zap.Field) {
+	gl.Debug(msg, fields...)
+}
+
+// Info wraps the zap Logger's Info method.
+func Info(msg string, fields ...zap.Field) {
+	gl.Info(msg, fields...)
+}
+
+// Warn wraps the zap Logger's Warn method.
+func Warn(msg string, fields ...zap.Field) {
+	gl.Warn(msg, fields...)
+}
+
+// Error wraps the zap Logger's Error method.
+func Error(msg string, fields ...zap.Field) {
+	gl.Error(msg, fields...)
+}
+
+// Sync wraps the zap Logger's Sync method.
+func Sync() {
+	_ = gl.Sync()
+}

go.mod

@@ -37,17 +37,20 @@ require (
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
 	github.com/valyala/fasttemplate v1.2.1 // indirect
 	github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c // indirect
+	go.uber.org/atomic v1.9.0 // indirect
+	go.uber.org/multierr v1.6.0 // indirect
 	golang.org/x/sys v0.1.0 // indirect
 	golang.org/x/text v0.4.0 // indirect
 	golang.org/x/time v0.0.0-20220722155302-e5dcc9cfc0b9 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
-	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
 require (
 	github.com/pkg/errors v0.9.1
 	github.com/segmentio/analytics-go v3.1.0+incompatible
-	github.com/stretchr/testify v1.7.0
+	github.com/stretchr/testify v1.8.0
+	go.uber.org/zap v1.24.0
 	golang.org/x/exp v0.0.0-20230111222715-75897c7a292a
 	golang.org/x/mod v0.6.0
 )

go.sum

@@ -2,6 +2,7 @@ github.com/VictoriaMetrics/fastcache v1.10.0 h1:5hDJnLsKLpnUEToub7ETuRu8RCkb40wo
 github.com/VictoriaMetrics/fastcache v1.10.0/go.mod h1:tjiYeEfYXCqacuvYw/7UoDIeJaNxq6132xHICNP77w8=
 github.com/allegro/bigcache v1.2.1-0.20190218064605-e24eb225f156 h1:eMwmnE/GDgah4HI848JfFxHt+iPb26b4zyfspmqY0/8=
 github.com/allegro/bigcache v1.2.1-0.20190218064605-e24eb225f156/go.mod h1:Cb/ax3seSYIx7SuZdm2G2xzfwmv3TPSk2ucNfQESPXM=
+github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4YnC6+E63dPcxHo2sUxDIu8g3QgEJdRY=
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE=
@@ -56,15 +57,26 @@ github.com/segmentio/analytics-go v3.1.0+incompatible/go.mod h1:C7CYBtQWk4vRk2Ry
 github.com/segmentio/backo-go v1.0.1 h1:68RQccglxZeyURy93ASB/2kc9QudzgIDexJ927N++y4=
 github.com/segmentio/backo-go v1.0.1/go.mod h1:9/Rh6yILuLysoQnZ2oNooD2g7aBnvM7r/fNVxRNWfBc=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
-github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/valyala/fasttemplate v1.2.1 h1:TVEnxayobAdVkhQfrfes2IzOB6o+z4roRkPF52WA1u4=
 github.com/valyala/fasttemplate v1.2.1/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ=
 github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c h1:3lbZUMbMiGUW/LMkfsEABsc5zNT9+b1CvsJx47JzJ8g=
 github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c/go.mod h1:UrdRz5enIKZ63MEE3IF9l2/ebyx59GyGgPi+tICQdmM=
+go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
+go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/goleak v1.1.11 h1:wy28qYRKZgnJTxGxvye5/wgWr1EKjmUDGYox5mGlRlI=
+go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4=
+go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
+go.uber.org/zap v1.24.0 h1:FiJd5l1UOLj0wCgbSE0rwwXHzEdAZS6hiiSnxJN/D60=
+go.uber.org/zap v1.24.0/go.mod h1:2kMP+WWQ8aoFoedH3T2sq6iJ2yDWpHbP0f6MQbS9Gkg=
 golang.org/x/crypto v0.1.0 h1:MDRAIl0xIo9Io2xV565hzXHw3zVseKrJKodhohM5CjU=
 golang.org/x/crypto v0.1.0/go.mod h1:RecgLatLF4+eUMCP1PoPZQb+cVrJcOPbHkTkbkB9sbw=
 golang.org/x/exp v0.0.0-20230111222715-75897c7a292a h1:/YWeLOBWYV5WAQORVPkZF3Pq9IppkcT72GKnWjNf5W8=
@@ -87,5 +99,6 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

scripts/.air.toml

@@ -11,3 +11,5 @@ tmp_dir = ".air"
   exclude_unchanged = false
   follow_symlink = false
   full_bin = ""
+  send_interrupt = true
+  kill_delay = 2000

server/server.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"database/sql"
 	"encoding/json"
 	"fmt"
 	"time"
@@ -20,7 +21,8 @@ import (
 )
 
 type Server struct {
-	e *echo.Echo
+	e  *echo.Echo
+	db *sql.DB
 
 	ID        string
 	Profile   *profile.Profile
@@ -34,16 +36,16 @@ func NewServer(ctx context.Context, profile *profile.Profile) (*Server, error) {
 	e.HideBanner = true
 	e.HidePort = true
 
-	s := &Server{
-		e:       e,
-		Profile: profile,
-	}
-
 	db := db.NewDB(profile)
 	if err := db.Open(ctx); err != nil {
 		return nil, errors.Wrap(err, "cannot open db")
 	}
 
+	s := &Server{
+		e:       e,
+		db:      db.DBInstance,
+		Profile: profile,
+	}
 	storeInstance := store.New(db.DBInstance, profile)
 	s.Store = storeInstance
 
@@ -125,6 +127,23 @@ func (s *Server) Start(ctx context.Context) error {
 	return s.e.Start(fmt.Sprintf(":%d", s.Profile.Port))
 }
 
+func (s *Server) Shutdown(ctx context.Context) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
+	// Shutdown echo server
+	if err := s.e.Shutdown(ctx); err != nil {
+		fmt.Printf("failed to shutdown server, error: %v\n", err)
+	}
+
+	// Close database connection
+	if err := s.db.Close(); err != nil {
+		fmt.Printf("failed to close database, error: %v\n", err)
+	}
+
+	fmt.Printf("memos stopped properly\n")
+}
+
 func (s *Server) createServerStartActivity(ctx context.Context) error {
 	payload := api.ActivityServerStartPayload{
 		ServerID: s.ID,

store/db/db.go

@@ -43,7 +43,7 @@ func (db *DB) Open(ctx context.Context) (err error) {
 	}
 
 	// Connect to the database without foreign_key.
-	sqliteDB, err := sql.Open("sqlite3", db.profile.DSN+"?_foreign_keys=0")
+	sqliteDB, err := sql.Open("sqlite3", db.profile.DSN+"?cache=shared&_foreign_keys=0&_journal_mode=WAL")
 	if err != nil {
 		return fmt.Errorf("failed to open db with dsn: %s, err: %w", db.profile.DSN, err)
 	}