Add shutdown context. (#20)

README.md

@@ -20,8 +20,7 @@
   - Request fee address (`POST /feeaddress`)
   - Pay fee (`POST /payFee`)
   - Ticket status (`POST /ticketstatus`)
-  - Set voting preferences (TBD)
+  - Set voting preferences (`POST /setvotebits`)
-
 - A minimal, static, web front-end providing pool stats and basic connection instructions.
 - Fees have an expiry period. If the fee is not paid within this period, the
   client must request a new fee. This enables the VSP to alter its fee rate.
@@ -33,7 +32,6 @@
 - Status check API call as described in [dcrstakepool #628](https://github.com/decred/dcrstakepool/issues/628).
 - Accountability for both client and server changes to voting preferences.
 - Consistency checking across connected wallets.
-- Validate votebits provided in PayFee request are valid per current agendas.
 
 ## Issue Tracker
 

database/database.go

@@ -79,6 +79,5 @@ func New(dbFile string) (*VspDatabase, error) {
 // Close releases all database resources. It will block waiting for any open
 // transactions to finish before closing the database and returning.
 func (vdb *VspDatabase) Close() error {
-	log.Debug("Closing database")
 	return vdb.db.Close()
 }

main.go

@@ -1,8 +1,13 @@
 package main
 
 import (
+	"context"
+	"errors"
 	"fmt"
+	"net"
+	"net/http"
 	"os"
+	"time"
 
 	"github.com/jholdstock/dcrvsp/database"
 	"github.com/jrick/wsrpc/v2"
@@ -15,28 +20,94 @@ var (
 )
 
 func main() {
+	// Create a context that is cancelled when a shutdown request is received
+	// through an interrupt signal.
+	ctx := withShutdownCancel(context.Background())
+	go shutdownListener()
+
+	// Run until error is returned, or shutdown is requested.
+	if err := run(ctx); err != nil && !errors.Is(err, context.Canceled) {
+		os.Exit(1)
+	}
+}
+
+// run is the main startup and teardown logic performed by the main package. It
+// is responsible for parsing the config, creating a dcrwallet RPC client,
+// opening the database, starting the webserver, and stopping all started
+// services when the context is cancelled.
+func run(ctx context.Context) error {
 	var err error
+
+	// Load config file and parse CLI args.
 	cfg, err = loadConfig()
 	if err != nil {
-		// Don't use logger here because it may not be initialised yet.
+		// Don't use logger here because it may not be initialised.
-		fmt.Fprintf(os.Stderr, "config error: %v", err)
+		fmt.Fprintf(os.Stderr, "Config error: %v", err)
-		os.Exit(1)
+		return err
 	}
 
+	// Open database.
 	db, err = database.New(cfg.dbPath)
 	if err != nil {
-		log.Errorf("database error: %v", err)
+		log.Errorf("Database error: %v", err)
-		os.Exit(1)
+		return err
 	}
-	defer db.Close()
+	// Close database.
-
+	defer func() {
-	// Start HTTP server
+		log.Debug("Closing database...")
-	log.Infof("Listening on %s", cfg.Listen)
+		err := db.Close()
-	// TODO: Make releaseMode properly configurable.
-	releaseMode := false
-	err = newRouter(releaseMode).Run(cfg.Listen)
 		if err != nil {
-		log.Errorf("web server terminated with error: %v", err)
+			log.Errorf("Error closing database: %v", err)
-		os.Exit(1)
+		} else {
+			log.Debug("Database closed")
 		}
+	}()
+
+	// Create TCP listener for webserver.
+	var listenConfig net.ListenConfig
+	listener, err := listenConfig.Listen(ctx, "tcp", cfg.Listen)
+	if err != nil {
+		log.Errorf("Failed to create tcp listener: %v", err)
+		return err
+	}
+	log.Infof("Listening on %s", cfg.Listen)
+
+	// Create webserver.
+	// TODO: Make releaseMode properly configurable.
+	releaseMode := true
+	srv := http.Server{
+		Handler:      newRouter(releaseMode),
+		ReadTimeout:  5 * time.Second,  // slow requests should not hold connections opened
+		WriteTimeout: 60 * time.Second, // hung responses must die
+	}
+
+	// Start webserver.
+	go func() {
+		err = srv.Serve(listener)
+		// If the server dies for any reason other than ErrServerClosed (from
+		// graceful server.Shutdown), log the error and request dcrvsp be
+		// shutdown.
+		if err != nil && err != http.ErrServerClosed {
+			log.Errorf("Unexpected webserver error: %v", err)
+			requestShutdown()
+		}
+	}()
+
+	// Stop webserver.
+	defer func() {
+		log.Debug("Stopping webserver...")
+		// Give the webserver 5 seconds to finish what it is doing.
+		timeoutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := srv.Shutdown(timeoutCtx); err != nil {
+			log.Errorf("Failed to stop webserver cleanly: %v", err)
+		}
+		log.Debug("Webserver stopped")
+	}()
+
+	// Wait until shutdown is signaled before returning and running deferred
+	// shutdown tasks.
+	<-ctx.Done()
+
+	return ctx.Err()
 }

signal.go

@@ -0,0 +1,71 @@
+// Copyright (c) 2013-2014 The btcsuite developers
+// Use of this source code is governed by an ISC
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"context"
+	"os"
+	"os/signal"
+)
+
+// shutdownRequestChannel is used to initiate shutdown from one of the
+// subsystems using the same code paths as when an interrupt signal is received.
+var shutdownRequestChannel = make(chan struct{})
+
+// shutdownSignaled is closed whenever shutdown is invoked through an interrupt
+// signal or from an JSON-RPC stop request.  Any contexts created using
+// withShutdownChannel are cancelled when this is closed.
+var shutdownSignaled = make(chan struct{})
+
+// signals defines the signals that are handled to do a clean shutdown.
+// Conditional compilation is used to also include SIGTERM on Unix.
+var signals = []os.Signal{os.Interrupt}
+
+// withShutdownCancel creates a copy of a context that is cancelled whenever
+// shutdown is invoked through an interrupt signal or from an JSON-RPC stop
+// request.
+func withShutdownCancel(ctx context.Context) context.Context {
+	ctx, cancel := context.WithCancel(ctx)
+	go func() {
+		<-shutdownSignaled
+		cancel()
+	}()
+	return ctx
+}
+
+// requestShutdown signals for starting the clean shutdown of the process
+// through an internal component (such as through the JSON-RPC stop request).
+func requestShutdown() {
+	shutdownRequestChannel <- struct{}{}
+}
+
+// shutdownListener listens for shutdown requests and cancels all contexts
+// created from withShutdownCancel.  This function never returns and is intended
+// to be spawned in a new goroutine.
+func shutdownListener() {
+	interruptChannel := make(chan os.Signal, 1)
+	signal.Notify(interruptChannel, signals...)
+
+	// Listen for the initial shutdown signal
+	select {
+	case sig := <-interruptChannel:
+		log.Infof("Received signal (%s).  Shutting down...", sig)
+	case <-shutdownRequestChannel:
+		log.Info("Shutdown requested.  Shutting down...")
+	}
+
+	// Cancel all contexts created from withShutdownCancel.
+	close(shutdownSignaled)
+
+	// Listen for any more shutdown signals and log that shutdown has already
+	// been signaled.
+	for {
+		select {
+		case <-interruptChannel:
+		case <-shutdownRequestChannel:
+		}
+		log.Info("Shutdown signaled.  Already shutting down...")
+	}
+}

signalsigterm.go

@@ -0,0 +1,16 @@
+// Copyright (c) 2016 The btcsuite developers
+// Use of this source code is governed by an ISC
+// license that can be found in the LICENSE file.
+
+// +build darwin dragonfly freebsd linux netbsd openbsd solaris
+
+package main
+
+import (
+	"os"
+	"syscall"
+)
+
+func init() {
+	signals = []os.Signal{os.Interrupt, syscall.SIGTERM}
+}