When database plugin Clean is called, close connections in goroutines

builtin/logical/database/backend.go

@@ -6,8 +6,10 @@ import (
 	"net/rpc"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
+	"github.com/armon/go-metrics"
 	log "github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/go-secure-stdlib/strutil"
 	"github.com/hashicorp/go-uuid"
@@ -18,6 +20,7 @@ import (
 	"github.com/hashicorp/vault/sdk/helper/locksutil"
 	"github.com/hashicorp/vault/sdk/logical"
 	"github.com/hashicorp/vault/sdk/queue"
+	"golang.org/x/sync/semaphore"
 )
 
 const (
@@ -27,6 +30,32 @@ const (
 	minRootCredRollbackAge = 1 * time.Minute
 )
 
+// mutable for tests only
+var cleanupMaxWaitTime = 500 * time.Millisecond
+
+// metrics collection
+var (
+	gaugeSync = sync.Mutex{}
+	gauges    = map[string]*int32{}
+	gaugeKey  = []string{"secrets", "database", "backend", "connections", "count"}
+)
+
+func createGauge(name string) {
+	gaugeSync.Lock()
+	defer gaugeSync.Unlock()
+	gauges[name] = new(int32)
+}
+
+func addToGauge(dbType, version string, amount int32) {
+	labels := []metrics.Label{{"type", dbType}, {"version", version}}
+	gaugeName := fmt.Sprintf("%s-%s", dbType, version)
+	if _, ok := gauges[gaugeName]; !ok {
+		createGauge(gaugeName)
+	}
+	val := atomic.AddInt32(gauges[gaugeName], amount)
+	metrics.SetGaugeWithLabels(gaugeKey, float32(val), labels)
+}
+
 type dbPluginInstance struct {
 	sync.RWMutex
 	database databaseVersionWrapper
@@ -280,10 +309,26 @@ func (b *databaseBackend) GetConnectionWithConfig(ctx context.Context, name stri
 		id:       id,
 		name:     name,
 	}
+	b.addConnectionsCounter(dbw, 1)
 	b.connections[name] = dbi
 	return dbi, nil
 }
 
+// addConnectionsCounter is used to keep metrics on how many and what types of databases we have open
+func (b *databaseBackend) addConnectionsCounter(dbw databaseVersionWrapper, amount int) {
+	// keep track of what databases we open
+	dbType, err := dbw.Type()
+	if err != nil {
+		b.Logger().Debug("Error getting database type", "err", err)
+		dbType = "unknown"
+	}
+	version := "5"
+	if dbw.isV4() {
+		version = "4"
+	}
+	addToGauge(dbType, version, int32(amount))
+}
+
 // invalidateQueue cancels any background queue loading and destroys the queue.
 func (b *databaseBackend) invalidateQueue() {
 	// cancel context before grabbing lock to start closing any open connections
@@ -311,6 +356,7 @@ func (b *databaseBackend) clearConnection(name string) error {
 	if ok {
 		// Ignore error here since the database client is always killed
 		db.Close()
+		b.addConnectionsCounter(db.database, -1)
 		delete(b.connections, name)
 	}
 	return nil
@@ -331,6 +377,7 @@ func (b *databaseBackend) CloseIfShutdown(db *dbPluginInstance, err error) {
 			// Ensure we are deleting the correct connection
 			mapDB, ok := b.connections[db.name]
 			if ok && db.id == mapDB.id {
+				b.addConnectionsCounter(db.database, -1)
 				delete(b.connections, db.name)
 			}
 		}()
@@ -339,18 +386,45 @@ func (b *databaseBackend) CloseIfShutdown(db *dbPluginInstance, err error) {
 
 // clean closes all connections from all database types
 // and cancels any rotation queue loading operation.
+// It spawns goroutines to close the databases since these
+// are not guaranteed to finish quickly.
 func (b *databaseBackend) clean(ctx context.Context) {
-	// invalidateQueue acquires it's own lock on the backend, removes queue, and
+	cleanupCtx, cancel := context.WithDeadline(ctx, time.Now().Add(cleanupMaxWaitTime))
+	defer cancel()
+
+	// invalidateQueue acquires its own lock on the backend, removes queue, and
 	// terminates the background ticker
 	b.invalidateQueue()
 
 	b.Lock()
 	defer b.Unlock()
 
-	for _, db := range b.connections {
-		db.Close()
-	}
+	// copy all connections so we can close asynchronously
+	connectionsCopy := b.connections
 	b.connections = make(map[string]*dbPluginInstance)
+
+	// we will try to wait for all the connections to close
+	// ... but not too long
+	sem := semaphore.NewWeighted(int64(len(connectionsCopy)))
+	err := sem.Acquire(cleanupCtx, int64(len(connectionsCopy)))
+	if err != nil {
+		b.Logger().Debug("Error acquiring semaphore; ignoring", "error", err)
+	}
+
+	for _, db := range connectionsCopy {
+		go func(db *dbPluginInstance) {
+			defer sem.Release(1)
+			b.addConnectionsCounter(db.database, -1)
+			err := db.Close()
+			if err != nil {
+				b.Logger().Debug("Error closing database while cleaning up plugin; ignoring", "error", err)
+			}
+		}(db)
+	}
+	err = sem.Acquire(cleanupCtx, int64(len(connectionsCopy)))
+	if err != nil {
+		b.Logger().Debug("Error in cleanup semaphore; ignoring", "error", err)
+	}
 }
 
 const backendHelp = `

builtin/logical/database/backend_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/hashicorp/vault/plugins/database/mongodb"
 	"github.com/hashicorp/vault/plugins/database/postgresql"
 	v4 "github.com/hashicorp/vault/sdk/database/dbplugin"
+	"github.com/hashicorp/vault/sdk/database/dbplugin/v5"
 	v5 "github.com/hashicorp/vault/sdk/database/dbplugin/v5"
 	"github.com/hashicorp/vault/sdk/database/helper/dbutil"
 	"github.com/hashicorp/vault/sdk/framework"
@@ -1461,6 +1462,88 @@ func TestBackend_ConnectionURL_redacted(t *testing.T) {
 	}
 }
 
+type hangingPlugin struct{}
+
+func (h hangingPlugin) Initialize(ctx context.Context, req v5.InitializeRequest) (v5.InitializeResponse, error) {
+	return v5.InitializeResponse{
+		Config: req.Config,
+	}, nil
+}
+
+func (h hangingPlugin) NewUser(ctx context.Context, req v5.NewUserRequest) (v5.NewUserResponse, error) {
+	return v5.NewUserResponse{}, nil
+}
+
+func (h hangingPlugin) UpdateUser(ctx context.Context, req v5.UpdateUserRequest) (v5.UpdateUserResponse, error) {
+	return v5.UpdateUserResponse{}, nil
+}
+
+func (h hangingPlugin) DeleteUser(ctx context.Context, req v5.DeleteUserRequest) (v5.DeleteUserResponse, error) {
+	return v5.DeleteUserResponse{}, nil
+}
+
+func (h hangingPlugin) Type() (string, error) {
+	return "hanging", nil
+}
+
+func (h hangingPlugin) Close() error {
+	time.Sleep(1000 * time.Second)
+	return nil
+}
+
+var _ dbplugin.Database = (*hangingPlugin)(nil)
+
+func TestBackend_PluginMain_Hanging(t *testing.T) {
+	if os.Getenv(pluginutil.PluginVaultVersionEnv) == "" {
+		return
+	}
+	v5.Serve(&hangingPlugin{})
+}
+
+func TestBackend_Closes_Cleanly_Even_If_Plugin_Hangs(t *testing.T) {
+	cleanupMaxWaitTime = 100 * time.Millisecond
+	cluster, sys := getCluster(t)
+	vault.TestAddTestPlugin(t, cluster.Cores[0].Core, "hanging-plugin", consts.PluginTypeDatabase, "TestBackend_PluginMain_Hanging", []string{}, "")
+	t.Cleanup(cluster.Cleanup)
+
+	config := logical.TestBackendConfig()
+	config.StorageView = &logical.InmemStorage{}
+	config.System = sys
+
+	b, err := Factory(context.Background(), config)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Configure a connection
+	data := map[string]interface{}{
+		"connection_url": "doesn't matter",
+		"plugin_name":    "hanging-plugin",
+		"allowed_roles":  []string{"plugin-role-test"},
+	}
+	req := &logical.Request{
+		Operation: logical.UpdateOperation,
+		Path:      "config/hang",
+		Storage:   config.StorageView,
+		Data:      data,
+	}
+	_, err = b.HandleRequest(namespace.RootContext(nil), req)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	timeout := time.NewTimer(750 * time.Millisecond)
+	done := make(chan bool)
+	go func() {
+		b.Cleanup(context.Background())
+		done <- true
+	}()
+	select {
+	case <-timeout.C:
+		t.Error("Hanging plugin caused Close() to time out")
+	case <-done:
+	}
+}
+
 func testCredsExist(t *testing.T, resp *logical.Response, connURL string) bool {
 	t.Helper()
 	var d struct {

builtin/logical/database/mocks_test.go

@@ -37,7 +37,7 @@ func (m *mockNewDatabase) DeleteUser(ctx context.Context, req v5.DeleteUserReque
 
 func (m *mockNewDatabase) Type() (string, error) {
 	args := m.Called()
-	return args.String(0), args.Error(1)
+	return args.Get(0).(string), args.Error(1)
 }
 
 func (m *mockNewDatabase) Close() error {

builtin/logical/database/rotation_test.go

@@ -1379,6 +1379,7 @@ func setupMockDB(b *databaseBackend) *mockNewDatabase {
 	mockDB := &mockNewDatabase{}
 	mockDB.On("Initialize", mock.Anything, mock.Anything).Return(v5.InitializeResponse{}, nil)
 	mockDB.On("Close").Return(nil)
+	mockDB.On("Type").Return("mock", nil)
 	dbw := databaseVersionWrapper{
 		v5: mockDB,
 	}