Repository: pingcap/tidb
File: br/pkg/utils/register_test.go (Go)

Code before:
	t.Log(currTask.MessageToUser())
	require.Equal(t, "/tidb/brie/import/import-into/test", currTask.Key)

	// then register again, this time will only refresh the lease, and left TTL will be larger.
	err = register.RegisterTaskOnce(ctx)
	require.NoError(t, err)
	list, err = GetImportTasksFrom(ctx, client)
	require.NoError(t, err)
	require.Len(t, list.Tasks, 1)
	thisTask := list.Tasks[0]
	require.Equal(t, currTask.Key, thisTask.Key)
	require.Equal(t, currTask.LeaseID, thisTask.LeaseID)
	require.Greater(t, thisTask.TTL, currTask.TTL)
}

func TestTaskRegisterFailedGrant(t *testing.T) {
	integration.BeforeTestExternal(t)
	testEtcdCluster := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1, GRPCKeepAliveInterval: time.Second, GRPCKeepAliveTimeout: 10 * time.Second})
	defer testEtcdCluster.Terminate(t)

	ctx := context.Background()
	client := testEtcdCluster.RandClient()

	register := NewTaskRegisterWithTTL(client, 3*time.Second, RegisterRestore, "test")
	defer func() {
		err := register.Close(ctx)
		// for flaky test, the lease would expire
		if err != nil && !strings.Contains(err.Error(), "requested lease not found") {
			require.NoError(t, err)
		}
	}()
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval", "return(200)"))
	defer func() {
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval")
	}()
	err := register.RegisterTask(ctx)
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		list, err := GetImportTasksFrom(ctx, client)
		if err != nil {
			// The task lease might be revoked during TTL query.
			return strings.Contains(err.Error(), "requested lease not found")
		}
		return len(list.Tasks) == 0

Review comment (refactor):
PR description says `Close` is tolerant to `requested lease not found`, but `taskRegister.Close` (br/pkg/utils/register.go) still returns/logs the revoke error. The test now special-cases the error string in the deferred `Close`, which is brittle and duplicates the intended behavior. Consider handling the lease-not-found case inside `Close` (e.g., ignore/return nil for `rpctypes.ErrLeaseNotFound`) and simplify this defer to a plain `require.NoError`.

Code after:
	require.Equal(t, "/tidb/brie/import/import-into/test", currTask.Key)

	// then register again, this time will only refresh the lease, and left TTL will be larger.
	err = register.RegisterTaskOnce(ctx)
	require.NoError(t, err)
	list, err = GetImportTasksFrom(ctx, client)
	require.NoError(t, err)
	require.Len(t, list.Tasks, 1)
	thisTask := list.Tasks[0]
	require.Equal(t, currTask.Key, thisTask.Key)
	require.Equal(t, currTask.LeaseID, thisTask.LeaseID)
	require.Greater(t, thisTask.TTL, currTask.TTL)
}

func TestTaskRegisterFailedGrant(t *testing.T) {
	integration.BeforeTestExternal(t)
	testEtcdCluster := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1, GRPCKeepAliveInterval: time.Second, GRPCKeepAliveTimeout: 10 * time.Second})
	defer testEtcdCluster.Terminate(t)

	ctx := context.Background()
	client := testEtcdCluster.RandClient()

	register := NewTaskRegisterWithTTL(client, 3*time.Second, RegisterRestore, "test")
	defer func() {
		require.NoError(t, register.Close(ctx))
	}()
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop", "return(true)"))
	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval", "return(200)"))
	defer func() {
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-always-grant")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop")
		failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-retry-interval")
	}()
	err := register.RegisterTask(ctx)
	require.NoError(t, err)

	require.Eventually(t, func() bool {
		list, err := GetImportTasksFrom(ctx, client)
		return err == nil && len(list.Tasks) == 0
	}, 5*time.Second, 50*time.Millisecond)

	failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-keepalive-stop")
	failpoint.Disable("github.com/pingcap/tidb/br/pkg/utils/brie-task-register-failed-to-grant")
	var list RegisterTasksList
	require.Eventually(t, func() bool {
		var err error
		list, err = GetImportTasksFrom(ctx, client)
		return err == nil && len(list.Tasks) > 0