From 4f7ee51e9a73e05e5b2dbc6cfa178ace2f83b155 Mon Sep 17 00:00:00 2001
From: Serge Bazanski <serge@bazanski.pl>
Date: Sun, 14 Oct 2018 18:20:59 +0100
Subject: [PATCH] add background context

---
 README   |  4 +++-
 mirko.go | 46 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/README b/README
index 2147adb..d2c369f 100644
--- a/README
+++ b/README
@@ -32,8 +32,10 @@ Usage (dev)
         }
 
         // start any other background processing...
+        // (you can use m.Context() to get a context that will get
+        // canceled when the service is about to shut down)
 
-        select {}
+        <-m.Done()
     }
 
 Usage (running)
diff --git a/mirko.go b/mirko.go
index 417e45d..6ce91fa 100644
--- a/mirko.go
+++ b/mirko.go
@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"time"
 
 	"code.hackerspace.pl/q3k/hspki"
@@ -35,10 +37,19 @@ type Mirko struct {
 	httpListen net.Listener
 	httpServer *http.Server
 	httpMux    *http.ServeMux
+
+	ctx     context.Context
+	cancel  context.CancelFunc
+	waiters []chan bool
 }
 
 func New() *Mirko {
-	return &Mirko{}
+	ctx, cancel := context.WithCancel(context.Background())
+	return &Mirko{
+		ctx:     ctx,
+		cancel:  cancel,
+		waiters: []chan bool{},
+	}
 }
 
 func authRequest(req *http.Request) (any, sensitive bool) {
@@ -113,6 +124,8 @@ func (m *Mirko) Listen() error {
 	return nil
 }
 
+// Trace logs debug information to either a context trace (if present)
+// or stderr (if not)
 func Trace(ctx context.Context, f string, args ...interface{}) {
 	tr, ok := trace.FromContext(ctx)
 	if !ok {
@@ -123,6 +136,7 @@ func Trace(ctx context.Context, f string, args ...interface{}) {
 	tr.LazyPrintf(f, args...)
 }
 
+// GRPC returns the microservice's grpc.Server object
 func (m *Mirko) GRPC() *grpc.Server {
 	if m.grpcServer == nil {
 		panic("GRPC() called before Listen()")
@@ -130,6 +144,7 @@ func (m *Mirko) GRPC() *grpc.Server {
 	return m.grpcServer
 }
 
+// HTTPMux returns the microservice's debug HTTP mux
 func (m *Mirko) HTTPMux() *http.ServeMux {
 	if m.httpMux == nil {
 		panic("HTTPMux() called before Listen()")
@@ -137,6 +152,22 @@ func (m *Mirko) HTTPMux() *http.ServeMux {
 	return m.httpMux
 }
 
+// Context returns a background microservice context that will be canceled
+// when the service is shut down
+func (m *Mirko) Context() context.Context {
+	return m.ctx
+}
+
+// Done() returns a channel that will emit a value when the service is
+// shut down. This should be used in the main() function instead of a select{}
+// call, to allow the background context to be canceled fully.
+func (m *Mirko) Done() chan bool {
+	c := make(chan bool, 1)
+	m.waiters = append(m.waiters, c)
+	return c
+}
+
+// Serve starts serving HTTP and gRPC requests
 func (m *Mirko) Serve() error {
 	errs := make(chan error, 1)
 	go func() {
@@ -150,6 +181,19 @@ func (m *Mirko) Serve() error {
 		}
 	}()
 
+	signalCh := make(chan os.Signal, 1)
+	signal.Notify(signalCh, os.Interrupt)
+	go func() {
+		select {
+		case <-signalCh:
+			m.cancel()
+			time.Sleep(time.Second)
+			for _, w := range m.waiters {
+				w <- true
+			}
+		}
+	}()
+
 	ticker := time.NewTicker(1 * time.Second)
 	select {
 	case <-ticker.C: