Repository: vitessio/vitess (go)

Issue:
Feature Request: ability to remove a registered vtgateconn dialer
### Feature Description

As of now, the Go `vitessdriver` package will register in a new dialer each time you specify a new Protocol name in the configuration. This means that ultimately an unbounded number of dialers will continue to accumulate throughout the lifetime of an application, unless you reuse the same name(s) repeatedly.

Being able to remove a dialer from the list means that it is possible for many unique dial configurations to be used throughout the lifetime of a program, while also keeping memory use in check.

### Use Case(s)

Programs which open a large number of vtgate dialer connections with differing gRPC dial option configurations over their lifetime.

Pull request:
vtgateconn: add DeregisterDialer hook
<!--
  Thank you for your contribution to the Vitess project.
  How to contribute: https://vitess.io/docs/contributing/
  Please first make sure there is an open Issue to discuss the feature/fix suggested in this PR.
  If this is a new feature, please mark the Issue as "RFC".
 -->

<!-- if this PR is Work in Progress please create it as a Draft Pull Request -->

## Description

Allows removing a named dialer to clean up when it is no longer needed. Inspired by similar APIs in go-sql-driver/mysql: https://pkg.go.dev/github.com/go-sql-driver/mysql

<!-- A few sentences describing the overall goals of the pull request's commits. -->
<!-- If this is a bug fix and you think the fix should be backported, please write so. -->

## Related Issue(s)

Fixes #12212.

<!-- List related issues and pull requests. If this PR fixes an issue, please add it using Fixes #????  -->

## Checklist

-   [ ] "Backport to:" labels have been added if this change should be back-ported
-   [x] Tests were added or are not required
-   [x] Documentation was added or is not required

## Deployment Notes

<!-- Notes regarding deployment of the contained body of work. These should note any db migrations, etc. -->


Patch:
diff --git a/go/vt/vtgate/vtgateconn/vtgateconn.go b/go/vt/vtgate/vtgateconn/vtgateconn.go
index a47a5e5181..eca75c7d86 100644
--- a/go/vt/vtgate/vtgateconn/vtgateconn.go
+++ b/go/vt/vtgate/vtgateconn/vtgateconn.go
@@ -207,6 +207,17 @@ func RegisterDialer(name string, dialer DialerFunc) {
 	dialers[name] = dialer
 }
 
+// DeregisterDialer removes the named DialerFunc from the registered list of
+// dialers. If the named DialerFunc does not exist, it is a noop.
+//
+// This is useful to avoid unbounded memory use if many different dialer
+// implementations are used throughout the lifetime of a program.
+func DeregisterDialer(name string) {
+	dialersM.Lock()
+	defer dialersM.Unlock()
+	delete(dialers, name)
+}
+
 // DialProtocol dials a specific protocol, and returns the *VTGateConn
 func DialProtocol(ctx context.Context, protocol string, address string) (*VTGateConn, error) {
 	dialersM.Lock()
diff --git a/go/vt/vttablet/onlineddl/executor.go b/go/vt/vttablet/onlineddl/executor.go
index b35f3c0f1a..b3771dbe91 100644
--- a/go/vt/vttablet/onlineddl/executor.go
+++ b/go/vt/vttablet/onlineddl/executor.go
@@ -4577,33 +4577,43 @@ func (e *Executor) submitCallbackIfNonConflicting(
 	}
 	// This is either singleton or singleton-context
 
-	e.migrationMutex.Lock()
-	defer e.migrationMutex.Unlock()
-
-	pendingUUIDs, err := e.readPendingMigrationsUUIDs(ctx)
-	if err != nil {
-		return nil, err
-	}
-	switch {
-	case onlineDDL.StrategySetting().IsSingleton():
-		// We will reject this migration if there's any pending migration
-		if len(pendingUUIDs) > 0 {
-			return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "singleton migration rejected: found pending migrations [%s]", strings.Join(pendingUUIDs, ", "))
+	// This entire next logic is wrapped in an anonymous func just to get the migrationMutex released
+	// before calling the callback function. Reason is: the callback function itself may need to acquire
+	// the mutex. And specifically, one of the callback functions used is e.RetryMigration(), which does
+	// lock the mutex...
+	err = func() error {
+		e.migrationMutex.Lock()
+		defer e.migrationMutex.Unlock()
+
+		pendingUUIDs, err := e.readPendingMigrationsUUIDs(ctx)
+		if err != nil {
+			return err
 		}
-	case onlineDDL.StrategySetting().IsSingletonContext():
-		// We will reject this migration if there's any pending migration within a different context
-		for _, pendingUUID := range pendingUUIDs {
-			pendingOnlineDDL, _, err := e.readMigration(ctx, pendingUUID)
-			if err != nil {
-				return nil, vterrors.Wrapf(err, "validateSingleton() migration: %s", pendingUUID)
+		switch {
+		case onlineDDL.StrategySetting().IsSingleton():
+			// We will reject this migration if there's any pending migration
+			if len(pendingUUIDs) > 0 {
+				return vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "singleton migration rejected: found pending migrations [%s]", strings.Join(pendingUUIDs, ", "))
 			}
-			if e.submittedMigrationConflictsWithPendingMigrationInSingletonContext(ctx, onlineDDL, pendingOnlineDDL) {
-				return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "singleton-context migration rejected: found pending migration: %s in different context: %s", pendingUUID, pendingOnlineDDL.MigrationContext)
+		case onlineDDL.StrategySetting().IsSingletonContext():
+			// We will reject this migration if there's any pending migration within a different context
+			for _, pendingUUID := range pendingUUIDs {
+				pendingOnlineDDL, _, err := e.readMigration(ctx, pendingUUID)
+				if err != nil {
+					return vterrors.Wrapf(err, "validateSingleton() migration: %s", pendingUUID)
+				}
+				if e.submittedMigrationConflictsWithPendingMigrationInSingletonContext(ctx, onlineDDL, pendingOnlineDDL) {
+					return vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "singleton-context migration rejected: found pending migration: %s in different context: %s", pendingUUID, pendingOnlineDDL.MigrationContext)
+				}
+				// no conflict? continue looking for other pending migrations
 			}
-			// no conflict? continue looking for other pending migrations
 		}
+		return nil
+	}()
+	if err != nil {
+		return nil, err
 	}
-	// OK to go! We can go
+	// OK to go!
 	return callback()
 }
 


Test patch:
diff --git a/go/test/endtoend/onlineddl/scheduler/onlineddl_scheduler_test.go b/go/test/endtoend/onlineddl/scheduler/onlineddl_scheduler_test.go
index 2c0e563571..875ac646d2 100644
--- a/go/test/endtoend/onlineddl/scheduler/onlineddl_scheduler_test.go
+++ b/go/test/endtoend/onlineddl/scheduler/onlineddl_scheduler_test.go
@@ -803,6 +803,43 @@ func testScheduler(t *testing.T) {
 			}
 		})
 	})
+
+	t.Run("Idempotent submission, retry failed migration in singleton context", func(t *testing.T) {
+		uuid := "00000000_1111_3333_3333_444444444444"
+		ddlStrategy := ddlStrategy + " --singleton-context"
+		overrideVtctlParams = &cluster.VtctlClientParams{DDLStrategy: ddlStrategy, SkipPreflight: true, UUIDList: uuid, MigrationContext: "idempotent:1111-3333-3333"}
+		defer func() { overrideVtctlParams = nil }()
+		// create a migration and cancel it. We don't let it complete. We want it in "failed" state
+		t.Run("start and fail migration", func(t *testing.T) {
+			executedUUID := testOnlineDDLStatement(t, createParams(trivialAlterT1Statement, ddlStrategy+" --postpone-completion", "vtctl", "", "", true)) // skip wait
+			require.Equal(t, uuid, executedUUID)
+			onlineddl.WaitForMigrationStatus(t, &vtParams, shards, uuid, normalWaitTime, schema.OnlineDDLStatusRunning)
+			// let's cancel it
+			onlineddl.CheckCancelMigration(t, &vtParams, shards, uuid, true)
+			status := onlineddl.WaitForMigrationStatus(t, &vtParams, shards, uuid, normalWaitTime, schema.OnlineDDLStatusFailed, schema.OnlineDDLStatusCancelled)
+			fmt.Printf("# Migration status (for debug purposes): <%s>\n", status)
+			onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusCancelled)
+		})
+
+		// now, we submit the exact same migratoin again: same UUID, same migration context.
+		t.Run("resubmit migration", func(t *testing.T) {
+			executedUUID := testOnlineDDLStatement(t, createParams(trivialAlterT1Statement, ddlStrategy, "vtctl", "", "", true)) // skip wait
+			require.Equal(t, uuid, executedUUID)
+
+			// expect it to complete
+			status := onlineddl.WaitForMigrationStatus(t, &vtParams, shards, uuid, normalWaitTime, schema.OnlineDDLStatusComplete, schema.OnlineDDLStatusFailed)
+			fmt.Printf("# Migration status (for debug purposes): <%s>\n", status)
+			onlineddl.CheckMigrationStatus(t, &vtParams, shards, uuid, schema.OnlineDDLStatusComplete)
+
+			rs := onlineddl.ReadMigrations(t, &vtParams, uuid)
+			require.NotNil(t, rs)
+			for _, row := range rs.Named().Rows {
+				retries := row.AsInt64("retries", 0)
+				assert.Greater(t, retries, int64(0))
+			}
+		})
+	})
+
 	// INSTANT DDL
 	instantDDLCapable, err := capableOf(mysql.InstantAddLastColumnFlavorCapability)
 	require.NoError(t, err)
diff --git a/go/vt/vtgate/vtgateconn/vtgateconn_test.go b/go/vt/vtgate/vtgateconn/vtgateconn_test.go
index 7f4969e15c..8bada5b406 100644
--- a/go/vt/vtgate/vtgateconn/vtgateconn_test.go
+++ b/go/vt/vtgate/vtgateconn/vtgateconn_test.go
@@ -44,3 +44,18 @@ func TestGetDialerWithProtocol(t *testing.T) {
 		t.Fatalf("dialerFunc has been registered, should not get nil: %v %v", err, c)
 	}
 }
+
+func TestDeregisterDialer(t *testing.T) {
+	const protocol = "test3"
+
+	RegisterDialer(protocol, func(context.Context, string) (Impl, error) {
+		return nil, nil
+	})
+
+	DeregisterDialer(protocol)
+
+	_, err := DialProtocol(context.Background(), protocol, "")
+	if err == nil || err.Error() != "no dialer registered for VTGate protocol "+protocol {
+		t.Fatalf("protocol: %s is not registered, should return error: %v", protocol, err)
+	}
+}