graceful stop

fs_store.go

@@ -53,6 +53,12 @@ func NewFsStore(root string, opts ...fsOption) (*fsStore, error) {
 		return nil, fmt.Errorf("create work directory: %w", err)
 	}
 
+	if p, err := os.MkdirTemp(s.work, ""); err != nil {
+		return nil, fmt.Errorf("create work context at work directory: %w", err)
+	} else {
+		s.work = p
+	}
+
 	test_src := filepath.Join(s.work, ".test")
 	test_dst := filepath.Join(s.root, ".test")
 
@@ -131,3 +137,7 @@ func (s *fsStore) Put(ctx context.Context, desc Description, r io.Reader) error
 
 	return nil
 }
+
+func (s *fsStore) Close() error {
+	return os.Remove(s.work)
+}

fs_store_test.go

@@ -1,9 +1,12 @@
 package main_test
 
 import (
+	"context"
+	"io"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	main "github.com/lesomnus/vcpkg-cache-http"
 	"github.com/stretchr/testify/require"
@@ -63,7 +66,7 @@ func TestFsStoreFail(t *testing.T) {
 		require.NoError(err)
 
 		_, err = main.NewFsStore(t.TempDir(), main.WithWorkDir(work))
-		require.ErrorContains(err, "create file at work")
+		require.ErrorIs(err, os.ErrPermission)
 	})
 
 	t.Run("file cannot be renamed from work to store directory", func(t *testing.T) {
@@ -77,4 +80,24 @@ func TestFsStoreFail(t *testing.T) {
 		_, err = main.NewFsStore(root, main.WithWorkDir(t.TempDir()))
 		require.ErrorContains(err, "rename file")
 	})
+
+	t.Run("store cannot be closed if the work left", func(t *testing.T) {
+		t.Parallel()
+		require := require.New(t)
+
+		store, err := main.NewFsStore(t.TempDir())
+		require.NoError(err)
+
+		r, w := io.Pipe()
+		go store.Put(context.Background(), DescriptionFoo, r)
+
+		time.Sleep(time.Millisecond * 10)
+		err = store.Close()
+		require.ErrorContains(err, "not empty")
+
+		w.Close()
+		time.Sleep(time.Millisecond * 10)
+		err = store.Close()
+		require.NoError(err)
+	})
 }

main.go

@@ -1,9 +1,12 @@
 package main
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"net/http"
 	"os"
+	"os/signal"
 	"time"
 
 	"github.com/rs/zerolog"
@@ -38,6 +41,12 @@ func main() {
 		l.Fatal().Err(err).Msg("failed to initialize a store")
 		return
 	}
+	defer func() {
+		err := store.Close()
+		if err != nil {
+			l.Error().Err(err).Msg("failed to close the store")
+		}
+	}()
 
 	handler := &Handler{
 		Store: store,
@@ -61,8 +70,33 @@ func main() {
 		Handler: handler,
 	}
 
+	server_closed := make(chan struct{})
+	defer close(server_closed)
+
+	go func() {
+		signals := make(chan os.Signal, 1)
+		signal.Notify(signals, os.Interrupt)
+
+		for {
+			select {
+			case <-server_closed:
+				return
+
+			case sig := <-signals:
+				l.Warn().Str("signal", sig.String()).Msg("shutdown the server")
+				if err := server.Shutdown(context.Background()); err != nil {
+					l.Error().Err(err).Msg("failed to shutdown the server")
+				}
+			}
+		}
+	}()
+
 	l.Info().Str("addr", addr).Msg("start server")
 	if err := server.ListenAndServe(); err != nil {
-		fmt.Println(err)
+		if errors.Is(err, http.ErrServerClosed) {
+			l.Info().Msg("server closed gracefully")
+		} else {
+			l.Error().Err(err).Msg("unexpected server close")
+		}
 	}
 }

store.go

@@ -20,4 +20,6 @@ type Store interface {
 	Get(ctx context.Context, desc Description, w io.Writer) error
 	Head(ctx context.Context, desc Description) error
 	Put(ctx context.Context, desc Description, r io.Reader) error
+
+	Close() error
 }