sdk/db: do not hold the lock on Close

sdk/database/dbplugin/v5/grpc_server.go

@@ -290,27 +290,37 @@ func (g *gRPCServer) Type(ctx context.Context, _ *proto.Empty) (*proto.TypeRespo
 
 func (g *gRPCServer) Close(ctx context.Context, _ *proto.Empty) (*proto.Empty, error) {
 	g.Lock()
-	defer g.Unlock()
 
 	impl, err := g.getDatabaseInternal(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	err = impl.Close()
-	if err != nil {
-		return &proto.Empty{}, status.Errorf(codes.Internal, "unable to close database plugin: %s", err)
-	}
-
+	var id string
 	if g.singleImpl == nil {
 		// only cleanup instances map when multiplexing is supported
-		id, err := pluginutil.GetMultiplexIDFromContext(ctx)
+		id, err = pluginutil.GetMultiplexIDFromContext(ctx)
 		if err != nil {
 			return nil, err
 		}
 		delete(g.instances, id)
 	}
 
+	// unlock here so that the subsequent call to Close() does not hold the
+	// lock in case the DB is slow to respond
+	g.Unlock()
+
+	err = impl.Close()
+	if err != nil {
+		// call to Close failed, so we will put the DB instance back in the map
+		g.Lock()
+		defer g.Unlock()
+		if g.singleImpl == nil {
+			g.instances[id] = impl
+		}
+		return &proto.Empty{}, status.Errorf(codes.Internal, "unable to close database plugin: %s", err)
+	}
+
 	return &proto.Empty{}, nil
 }
 

sdk/database/dbplugin/v5/grpc_server_test.go

@@ -551,6 +551,19 @@ func TestGRPCServer_Close(t *testing.T) {
 				}
 			},
 		},
+		"error path for multiplexed plugin": {
+			db: fakeDatabase{
+				closeErr: errors.New("close error"),
+			},
+			expectErr:     true,
+			expectCode:    codes.Internal,
+			grpcSetupFunc: testGrpcServer,
+			assertFunc: func(t *testing.T, g gRPCServer) {
+				if len(g.instances) != 1 {
+					t.Fatalf("err expected instances map to contain exactly 1 element")
+				}
+			},
+		},
 		"happy path for non-multiplexed plugin": {
 			db:            fakeDatabase{},
 			expectErr:     false,