Use grace instead of lifetime

internal/cleanup.go

@@ -0,0 +1,46 @@
+package internal
+
+import (
+	"context"
+	"github.com/tomwright/grace"
+	"log"
+	"time"
+)
+
+// NewCleanupRunner returns a runner that can be used cleanup old diagrams.
+func NewCleanupRunner(generator Generator) grace.Runner {
+	return &cleanupService{
+		generator:   generator,
+		runEvery:    time.Minute * 5,
+		cleanupLast: time.Hour,
+	}
+}
+
+// cleanupService is a runner that is used cleanup old diagrams.
+type cleanupService struct {
+	generator   Generator
+	runEvery    time.Duration
+	cleanupLast time.Duration
+}
+
+// Run starts the cleanup process.
+func (s *cleanupService) Run(ctx context.Context) error {
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+		}
+
+		if err := s.generator.CleanUp(s.cleanupLast); err != nil {
+			log.Printf("error when cleaning up: %s", err.Error())
+		}
+
+		select {
+		case <-time.After(s.runEvery):
+			continue
+		case <-ctx.Done():
+			return nil
+		}
+	}
+}

internal/cleanup_service.go

@@ -1,41 +0,0 @@
-package internal
-
-import (
-	"log"
-	"time"
-)
-
-// NewCleanupService returns a service that can be used cleanup old diagrams.
-func NewCleanupService(generator Generator) *cleanupService {
-	return &cleanupService{
-		generator: generator,
-		stopCh:    make(chan struct{}),
-	}
-}
-
-// cleanupService is a service that can be used cleanup old diagrams.
-type cleanupService struct {
-	generator Generator
-	stopCh    chan struct{}
-}
-
-// Start starts the cleanup service.
-func (s *cleanupService) Start() error {
-	for {
-		if err := s.generator.CleanUp(time.Hour); err != nil {
-			log.Printf("error when cleaning up: %s", err.Error())
-		}
-
-		select {
-		case <-time.After(time.Minute * 5):
-			continue
-		case <-s.stopCh:
-			return nil
-		}
-	}
-}
-
-// Stop stops the cleanup service.
-func (s *cleanupService) Stop() {
-	close(s.stopCh)
-}

internal/http.go

@@ -3,13 +3,32 @@ package internal
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/tomwright/grace"
+	"github.com/tomwright/gracehttpserverrunner"
 	"io/ioutil"
 	"log"
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 )
 
+// NewHTTPRunner returns a grace runner that runs a HTTP server.
+func NewHTTPRunner(generator Generator) grace.Runner {
+	httpHandler := generateHTTPHandler(generator)
+
+	r := http.NewServeMux()
+	r.Handle("/generate", http.HandlerFunc(httpHandler))
+
+	return &gracehttpserverrunner.HTTPServerRunner{
+		Server: &http.Server{
+			Addr:    ":80",
+			Handler: r,
+		},
+		ShutdownTimeout: time.Second * 5,
+	}
+}
+
 func writeJSON(rw http.ResponseWriter, value interface{}, status int) {
 	bytes, err := json.Marshal(value)
 	if err != nil {

internal/http_service.go

@@ -1,47 +0,0 @@
-package internal
-
-import (
-	"net/http"
-)
-
-// NewHTTPService returns a service that can be used to start a http server
-// that will generate diagrams.
-func NewHTTPService(generator Generator) *httpService {
-	return &httpService{
-		generator: generator,
-	}
-}
-
-// httpService is a service that can be used to start a http server
-// that will generate diagrams.
-type httpService struct {
-	httpServer *http.Server
-	generator  Generator
-}
-
-// Start starts the HTTP server.
-func (s *httpService) Start() error {
-	httpHandler := generateHTTPHandler(s.generator)
-
-	r := http.NewServeMux()
-	r.Handle("/generate", http.HandlerFunc(httpHandler))
-
-	s.httpServer = &http.Server{
-		Addr:    ":80",
-		Handler: r,
-	}
-
-	if err := s.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
-		if err != http.ErrServerClosed {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func (s *httpService) Stop() {
-	if s != nil {
-		_ = s.httpServer.Close()
-	}
-}