Block any database interaction when shutting down

database/controller.go

@@ -50,13 +50,13 @@ func (c *Controller) Injected() bool {
 
 // Get return the record with the given key.
 func (c *Controller) Get(key string) (record.Record, error) {
+	c.readLock.RLock()
+	defer c.readLock.RUnlock()
+
 	if shuttingDown.IsSet() {
 		return nil, ErrShuttingDown
 	}
 
-	c.readLock.RLock()
-	defer c.readLock.RUnlock()
-
 	// process hooks
 	for _, hook := range c.hooks {
 		if hook.h.UsesPreGet() && hook.q.MatchesKey(key) {
@@ -98,6 +98,9 @@ func (c *Controller) Get(key string) (record.Record, error) {
 
 // Put saves a record in the database.
 func (c *Controller) Put(r record.Record) (err error) {
+	c.writeLock.RLock()
+	defer c.writeLock.RUnlock()
+
 	if shuttingDown.IsSet() {
 		return ErrShuttingDown
 	}
@@ -116,9 +119,6 @@ func (c *Controller) Put(r record.Record) (err error) {
 		}
 	}
 
-	c.writeLock.RLock()
-	defer c.writeLock.RUnlock()
-
 	err = c.storage.Put(r)
 	if err != nil {
 		return err
@@ -139,11 +139,13 @@ func (c *Controller) Put(r record.Record) (err error) {
 
 // Query executes the given query on the database.
 func (c *Controller) Query(q *query.Query, local, internal bool) (*iterator.Iterator, error) {
+	c.readLock.RLock()
+
 	if shuttingDown.IsSet() {
+		c.readLock.RUnlock()
 		return nil, ErrShuttingDown
 	}
 
-	c.readLock.RLock()
 	it, err := c.storage.Query(q, local, internal)
 	if err != nil {
 		c.readLock.RUnlock()
@@ -160,6 +162,10 @@ func (c *Controller) PushUpdate(r record.Record) {
 		c.readLock.RLock()
 		defer c.readLock.RUnlock()
 
+		if shuttingDown.IsSet() {
+			return
+		}
+
 		for _, sub := range c.subscriptions {
 			if r.Meta().CheckPermission(sub.local, sub.internal) && sub.q.Matches(r) {
 				select {
@@ -177,6 +183,10 @@ func (c *Controller) addSubscription(sub *Subscription) {
 	c.writeLock.Lock()
 	defer c.writeLock.Unlock()
 
+	if shuttingDown.IsSet() {
+		return
+	}
+
 	c.subscriptions = append(c.subscriptions, sub)
 }
 
@@ -189,6 +199,11 @@ func (c *Controller) readUnlockerAfterQuery(it *iterator.Iterator) {
 func (c *Controller) Maintain() error {
 	c.writeLock.RLock()
 	defer c.writeLock.RUnlock()
+
+	if shuttingDown.IsSet() {
+		return nil
+	}
+
 	return c.storage.Maintain()
 }
 
@@ -196,11 +211,21 @@ func (c *Controller) Maintain() error {
 func (c *Controller) MaintainThorough() error {
 	c.writeLock.RLock()
 	defer c.writeLock.RUnlock()
+
+	if shuttingDown.IsSet() {
+		return nil
+	}
+
 	return c.storage.MaintainThorough()
 }
 
 // Shutdown shuts down the storage.
 func (c *Controller) Shutdown() error {
-	// TODO: should we wait for gets/puts/queries to complete?
+	// aquire full locks
+	c.readLock.Lock()
+	defer c.readLock.Unlock()
+	c.writeLock.Lock()
+	defer c.writeLock.Unlock()
+
 	return c.storage.Shutdown()
 }

database/controllers.go

@@ -2,15 +2,15 @@ package database
 
 import (
 	"errors"
-	"sync"
 	"fmt"
+	"sync"
 
 	"github.com/Safing/portbase/database/storage"
 )
 
 var (
 	controllers     = make(map[string]*Controller)
-	controllersLock sync.Mutex
+	controllersLock sync.RWMutex
 )
 
 func getController(name string) (*Controller, error) {
@@ -18,13 +18,19 @@ func getController(name string) (*Controller, error) {
 		return nil, errors.New("database not initialized")
 	}
 
+	// return database if already started
+	controllersLock.RLock()
+	controller, ok := controllers[name]
+	controllersLock.RUnlock()
+	if ok {
+		return controller, nil
+	}
+
 	controllersLock.Lock()
 	defer controllersLock.Unlock()
 
-  // return database if already started
-	controller, ok := controllers[name]
-	if ok {
-    return controller, nil
+	if shuttingDown.IsSet() {
+		return nil, ErrShuttingDown
 	}
 
 	// get db registration
@@ -60,6 +66,10 @@ func InjectDatabase(name string, storageInt storage.Interface) (*Controller, err
 	controllersLock.Lock()
 	defer controllersLock.Unlock()
 
+	if shuttingDown.IsSet() {
+		return nil, ErrShuttingDown
+	}
+
 	_, ok := controllers[name]
 	if ok {
 		return nil, errors.New(`database "%s" already loaded`)

database/main.go

@@ -50,10 +50,14 @@ func Initialize() error {
 func Shutdown() (err error) {
 	if shuttingDown.SetToIf(false, true) {
 		close(shutdownSignal)
+	} else {
+		return
 	}
 
-	all := duplicateControllers()
-	for _, c := range all {
+	controllersLock.RLock()
+	defer controllersLock.RUnlock()
+
+	for _, c := range controllers {
 		err = c.Shutdown()
 		if err != nil {
 			return