fix(proxy): retry HAKeeper errors during bootstrap

pkg/proxy/bootstrap.go

@@ -16,50 +16,83 @@ package proxy
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"github.com/matrixorigin/matrixone/pkg/common/moerr"
-	pb "github.com/matrixorigin/matrixone/pkg/pb/logservice"
 	"github.com/matrixorigin/matrixone/pkg/util"
 	db_holder "github.com/matrixorigin/matrixone/pkg/util/export/etl/db"
 )
 
 const (
-	BootstrapInterval = time.Millisecond * 200
-	BootstrapTimeout  = time.Minute * 5
+	BootstrapInterval       = time.Millisecond * 200
+	BootstrapTimeout        = time.Minute * 5
+	bootstrapRequestTimeout = time.Second * 3
 )
 
-func (h *handler) bootstrap(ctx context.Context) {
-	ticker := time.NewTicker(time.Millisecond * 200)
+// bootstrap retries until the task-table credentials are available. It returns
+// nil on success, the parent cancellation on cancellation, or a deadline error
+// together with the last HAKeeper error when the bootstrap window expires.
+func (h *handler) bootstrap(ctx context.Context) error {
+	return h.bootstrapWithTimeout(ctx, BootstrapInterval, BootstrapTimeout)
+}
+
+func (h *handler) bootstrapWithTimeout(
+	ctx context.Context,
+	interval time.Duration,
+	timeout time.Duration,
+) error {
+	ctx, cancel := context.WithTimeoutCause(ctx, timeout, moerr.CauseProxyBootstrap)
+	defer cancel()
+
+	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
-	retry := 0
 	getClient := func() util.HAKeeperClient {
 		return h.haKeeperClient
 	}
-	var state pb.CheckerState
-	var err error
-	for retry < int(BootstrapTimeout/BootstrapInterval) {
+	var lastErr error
+	for {
 		select {
 		case <-ticker.C:
-			func(ctx context.Context) {
-				ctx, cancel := context.WithTimeoutCause(ctx, time.Second*3, moerr.CauseProxyBootstrap)
-				defer cancel()
-				state, err = h.haKeeperClient.GetClusterState(ctx)
-				if err != nil {
-					panic(moerr.AttachCause(ctx, err))
-				}
-			}(ctx)
+			requestCtx, requestCancel := context.WithTimeoutCause(
+				ctx,
+				bootstrapRequestTimeout,
+				moerr.CauseProxyBootstrap,
+			)
+			state, err := h.haKeeperClient.GetClusterState(requestCtx)
+			if err != nil {
+				lastErr = moerr.AttachCause(requestCtx, err)
+			} else {
+				lastErr = nil
+			}
+			requestCancel()
+
+			if ctx.Err() != nil {
+				return bootstrapContextError(ctx, lastErr)
+			}
+			if err != nil {
+				continue
+			}
 			if state.TaskTableUser.GetUsername() != "" && state.TaskTableUser.GetPassword() != "" {
 				db_holder.SetSQLWriterDBUser(db_holder.MOLoggerUser, state.TaskTableUser.GetPassword())
 				db_holder.SetSQLWriterDBAddressFunc(util.AddressFunc(h.config.UUID, getClient))
 				h.sqlWorker.SetSQLUser(SQLUsername, state.TaskTableUser.GetPassword())
 				h.sqlWorker.SetAddressFn(util.AddressFunc(h.config.UUID, getClient))
-				return
+				return nil
 			}
 		case <-ctx.Done():
-			return
+			return bootstrapContextError(ctx, lastErr)
 		}
-		retry += 1
 	}
-	panic("proxy bootstrap failed")
+}
+
+func bootstrapContextError(ctx context.Context, lastErr error) error {
+	err := ctx.Err()
+	if cause := context.Cause(ctx); cause != nil && !errors.Is(err, cause) {
+		err = errors.Join(err, cause)
+	}
+	if errors.Is(err, context.DeadlineExceeded) && lastErr != nil {
+		err = errors.Join(err, lastErr)
+	}
+	return err
 }

pkg/proxy/bootstrap_test.go

@@ -16,16 +16,175 @@ package proxy
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
 	"github.com/matrixorigin/matrixone/pkg/common/runtime"
 	"github.com/matrixorigin/matrixone/pkg/common/stopper"
+	logpb "github.com/matrixorigin/matrixone/pkg/pb/logservice"
 	db_holder "github.com/matrixorigin/matrixone/pkg/util/export/etl/db"
 	"github.com/matrixorigin/matrixone/pkg/util/toml"
 	"github.com/stretchr/testify/require"
 )
 
+type bootstrapHAKeeperClient struct {
+	*mockHAKeeperClient
+	getClusterState func(context.Context) (logpb.CheckerState, error)
+}
+
+const bootstrapTestTimeout = 10 * time.Second
+
+func (c *bootstrapHAKeeperClient) GetClusterState(ctx context.Context) (logpb.CheckerState, error) {
+	return c.getClusterState(ctx)
+}
+
+func TestBootstrapRetriesTransientHAKeeperError(t *testing.T) {
+	calls := 0
+	c := &bootstrapHAKeeperClient{
+		mockHAKeeperClient: &mockHAKeeperClient{},
+		getClusterState: func(context.Context) (logpb.CheckerState, error) {
+			calls++
+			if calls == 1 {
+				return logpb.CheckerState{}, errors.New("temporary HAKeeper error")
+			}
+			return logpb.CheckerState{
+				TaskTableUser: logpb.TaskTableUser{
+					Username: "u1",
+					Password: "p1",
+				},
+			}, nil
+		},
+	}
+	h := &handler{
+		haKeeperClient: c,
+		sqlWorker:      newSQLWorker(),
+	}
+
+	require.NoError(t, h.bootstrapWithTimeout(
+		context.Background(),
+		time.Millisecond,
+		time.Second,
+	))
+	require.Equal(t, 2, calls)
+}
+
+func TestBootstrapReturnsContextCancellation(t *testing.T) {
+	called := make(chan struct{})
+	c := &bootstrapHAKeeperClient{
+		mockHAKeeperClient: &mockHAKeeperClient{},
+		getClusterState: func(ctx context.Context) (logpb.CheckerState, error) {
+			close(called)
+			<-ctx.Done()
+			return logpb.CheckerState{}, ctx.Err()
+		},
+	}
+	h := &handler{
+		haKeeperClient: c,
+		sqlWorker:      newSQLWorker(),
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	errC := make(chan error, 1)
+	go func() {
+		errC <- h.bootstrapWithTimeout(ctx, time.Millisecond, time.Second)
+	}()
+
+	select {
+	case <-called:
+	case <-time.After(bootstrapTestTimeout):
+		t.Fatal("HAKeeper request was not started")
+	}
+	cancel()
+	select {
+	case err := <-errC:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(bootstrapTestTimeout):
+		t.Fatal("bootstrap did not return after context cancellation")
+	}
+}
+
+func TestBootstrapStopsWithStopper(t *testing.T) {
+	requestStarted := make(chan struct{})
+	requestCanceled := make(chan error, 1)
+	c := &bootstrapHAKeeperClient{
+		mockHAKeeperClient: &mockHAKeeperClient{},
+		getClusterState: func(ctx context.Context) (logpb.CheckerState, error) {
+			close(requestStarted)
+			<-ctx.Done()
+			requestCanceled <- ctx.Err()
+			return logpb.CheckerState{}, ctx.Err()
+		},
+	}
+	h := &handler{
+		haKeeperClient: c,
+		sqlWorker:      newSQLWorker(),
+	}
+	st := stopper.NewStopper("test-proxy-bootstrap")
+	defer st.Stop()
+	require.NoError(t, runBootstrapTask(context.Background(), st, h))
+
+	select {
+	case <-requestStarted:
+	case <-time.After(bootstrapTestTimeout):
+		t.Fatal("HAKeeper request was not started")
+	}
+
+	stopDone := make(chan struct{})
+	go func() {
+		st.Stop()
+		close(stopDone)
+	}()
+
+	select {
+	case err := <-requestCanceled:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(bootstrapTestTimeout):
+		t.Fatal("HAKeeper request context was not canceled")
+	}
+	select {
+	case <-stopDone:
+	case <-time.After(bootstrapTestTimeout):
+		t.Fatal("stopper did not wait for bootstrap task termination")
+	}
+}
+
+func TestBootstrapReturnsTimeoutAfterPersistentHAKeeperError(t *testing.T) {
+	permanentErr := errors.New("HAKeeper unavailable")
+	calls := 0
+	c := &bootstrapHAKeeperClient{
+		mockHAKeeperClient: &mockHAKeeperClient{},
+		getClusterState: func(context.Context) (logpb.CheckerState, error) {
+			calls++
+			return logpb.CheckerState{}, permanentErr
+		},
+	}
+	h := &handler{
+		haKeeperClient: c,
+		sqlWorker:      newSQLWorker(),
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	errC := make(chan error, 1)
+	go func() {
+		errC <- h.bootstrapWithTimeout(
+			ctx,
+			time.Millisecond,
+			50*time.Millisecond,
+		)
+	}()
+	var err error
+	select {
+	case err = <-errC:
+	case <-time.After(bootstrapTestTimeout):
+		t.Fatal("bootstrap did not return after its timeout")
+	}
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.ErrorIs(t, err, permanentErr)
+	require.Positive(t, calls)
+}
+
 func TestBootstrap(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -39,7 +198,7 @@ func TestBootstrap(t *testing.T) {
 	cfg.Cluster.RefreshInterval = toml.Duration{Duration: defaultRefreshInterval}
 	h, err := newProxyHandler(ctx, rt, cfg, st, nil, &c, true)
 	require.NoError(t, err)
-	h.bootstrap(ctx)
+	require.NoError(t, h.bootstrap(ctx))
 
 	u, err := db_holder.GetSQLWriterDBUser()
 	require.NoError(t, err)

pkg/proxy/server.go

@@ -105,7 +105,9 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err
 		return nil, err
 	}
 
-	go h.bootstrap(ctx)
+	if err := runBootstrapTask(ctx, s.stopper, h); err != nil {
+		return nil, err
+	}
 
 	if err := s.stopper.RunNamedTask("proxy heartbeat", s.heartbeat); err != nil {
 		return nil, err
@@ -127,6 +129,23 @@ func NewServer(ctx context.Context, config Config, opts ...Option) (*Server, err
 	return s, nil
 }
 
+func runBootstrapTask(ctx context.Context, st *stopper.Stopper, h *handler) error {
+	return st.RunNamedTask("proxy bootstrap", func(taskCtx context.Context) {
+		bootstrapCtx, cancel := context.WithCancelCause(ctx)
+		stopCancelPropagation := context.AfterFunc(taskCtx, func() {
+			cancel(context.Cause(taskCtx))
+		})
+		defer func() {
+			stopCancelPropagation()
+			cancel(nil)
+		}()
+
+		if err := h.bootstrap(bootstrapCtx); err != nil && ctx.Err() == nil && taskCtx.Err() == nil {
+			h.logger.Error("proxy bootstrap failed", zap.Error(err))
+		}
+	})
+}
+
 // Start starts the proxy server.
 func (s *Server) Start() error {
 	err := s.app.Start()