Close database on shutdown

3 files changed, 29 insertions, 10 deletions

diff --git a/cmd/zdns/main.go b/cmd/zdns/main.go
index af616f9..b0537f9 100644
--- a/cmd/zdns/main.go
+++ b/cmd/zdns/main.go
@@ -84,20 +84,19 @@ func newCli(out io.Writer, args []string, configFile string, sig chan os.Signal)
 
 	// SQL backends
 	var (
+		sqlClient *sql.Client
 		sqlLogger *sql.Logger
 		sqlCache  *sql.Cache
 	)
 	if config.DNS.Database != "" {
-		sqlClient, err := sql.New(config.DNS.Database)
+		sqlClient, err = sql.New(config.DNS.Database)
 		fatal(err)
 
 		// Logger
 		sqlLogger = sql.NewLogger(sqlClient, config.DNS.LogMode, config.DNS.LogTTL)
-		sigHandler.OnClose(sqlLogger)
 
 		// Cache
 		sqlCache = sql.NewCache(sqlClient)
-		sigHandler.OnClose(sqlCache)
 	}
 
 	// DNS client
@@ -115,25 +114,43 @@ func newCli(out io.Writer, args []string, configFile string, sig chan os.Signal)
 	} else {
 		dnsCache = cache.New(config.DNS.CacheSize, cacheDNS)
 	}
-	sigHandler.OnClose(dnsCache)
 
 	// DNS server
 	proxy, err := dns.NewProxy(dnsCache, dnsClient, sqlLogger)
 	fatal(err)
-	sigHandler.OnClose(proxy)
 
 	dnsSrv, err := zdns.NewServer(proxy, config)
 	fatal(err)
 	sigHandler.OnReload(dnsSrv)
-	sigHandler.OnClose(dnsSrv)
 	servers := []server{dnsSrv}
 
 	// HTTP server
+	var httpSrv *http.Server
 	if config.DNS.ListenHTTP != "" {
-		httpSrv := http.NewServer(dnsCache, sqlLogger, sqlCache, config.DNS.ListenHTTP)
-		sigHandler.OnClose(httpSrv)
+		httpSrv = http.NewServer(dnsCache, sqlLogger, sqlCache, config.DNS.ListenHTTP)
 		servers = append(servers, httpSrv)
 	}
+
+	// Close proxy first
+	sigHandler.OnClose(proxy)
+
+	// ... then HTTP server
+	if httpSrv != nil {
+		sigHandler.OnClose(httpSrv)
+	}
+
+	// ... then cache
+	sigHandler.OnClose(dnsCache)
+
+	// ... then database components
+	if config.DNS.Database != "" {
+		sigHandler.OnClose(sqlLogger)
+		sigHandler.OnClose(sqlCache)
+		sigHandler.OnClose(sqlClient)
+	}
+
+	// ... and finally the server itself
+	sigHandler.OnClose(dnsSrv)
 	return &cli{servers: servers}, nil
 }
 
diff --git a/signal/signal.go b/signal/signal.go
index 49e05cc..99b2e6f 100644
--- a/signal/signal.go
+++ b/signal/signal.go
@@ -47,10 +47,9 @@ func (s *Handler) readSignal() {
 			log.Printf("received signal %s: shutting down", sig)
 			for _, c := range s.closers {
 				if err := c.Close(); err != nil {
-					log.Printf("close failed: %s", err)
+					log.Printf("close of %T failed: %s", c, err)
 				}
 			}
-
 		default:
 			log.Printf("received signal %s: ignoring", sig)
 		}
diff --git a/sql/sql.go b/sql/sql.go
index ee7b385..5941a00 100644
--- a/sql/sql.go
+++ b/sql/sql.go
@@ -118,6 +118,9 @@ func New(filename string) (*Client, error) {
 	return &Client{db: db}, nil
 }
 
+// Close waits for all queries to complete and then closes the database.
+func (c *Client) Close() error { return c.db.Close() }
+
 func (c *Client) readLog(n int) ([]logEntry, error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()