Retry fetch grant token

Author: ivov

internal/auth/handshake.go

@@ -4,9 +4,9 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
-	"task-runner-launcher/internal/logs"
 	"net/url"
 	"strings"
+	"task-runner-launcher/internal/logs"
 
 	"github.com/gorilla/websocket"
 )
@@ -201,6 +201,7 @@ func Handshake(cfg HandshakeConfig) error {
 	case err := <-errReceived:
 		return err
 	case <-handshakeComplete:
+		logs.Info("Runner's task offer was accepted")
 		return nil
 	}
 }

internal/auth/token_exchange.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	launcherHttp "task-runner-launcher/internal/http"
 )
 
 type grantTokenResponse struct {
@@ -13,10 +14,7 @@ type grantTokenResponse struct {
 	} `json:"data"`
 }
 
-// FetchGrantToken exchanges the launcher's auth token for a less privileged
-// grant token returned by the n8n main instance. The launcher will later pass
-// this grant token to the task runner.
-func FetchGrantToken(n8nURI, authToken string) (string, error) {
+func fetchGrantToken(n8nURI, authToken string) (string, error) {
 	url := fmt.Sprintf("http://%s/runners/auth", n8nURI)
 
 	payload := map[string]string{"token": authToken}
@@ -44,8 +42,29 @@ func FetchGrantToken(n8nURI, authToken string) (string, error) {
 
 	var tokenResp grantTokenResponse
 	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
-		return "", err
+		return "", fmt.Errorf("failed to decode grant token response: %w", err)
 	}
 
 	return tokenResp.Data.Token, nil
 }
+
+// FetchGrantTokenWithRetry exchanges the launcher's auth token for a less
+// privileged grant token returned by the n8n main instance. This exchange is
+// retried in case the n8n main instance is temporarily unavailable.
+func FetchGrantTokenWithRetry(n8nURI, authToken string) (string, error) {
+	grantTokenFetch := func() (string, error) {
+		token, err := fetchGrantToken(n8nURI, authToken)
+		if err != nil {
+			return "", fmt.Errorf("failed to fetch grant token: %w", err)
+		}
+		return token, nil
+	}
+
+	token, err := launcherHttp.LimitedRetry("grant-token-fetch", grantTokenFetch)
+
+	if err != nil {
+		return "", fmt.Errorf("exhausted retries to fetch grant token: %w", err)
+	}
+
+	return token, nil
+}

internal/commands/launch.go

@@ -92,14 +92,14 @@ func (l *LaunchCommand) Execute() error {
 
 	// 4. wait for n8n instance to be ready
 
-	if err := http.WaitForN8n(n8nMainServerURI); err != nil {
+	if err := http.WaitForN8nReadiness(n8nMainServerURI); err != nil {
 		return fmt.Errorf("exhausted retries while waiting for n8n instance to be ready: %w", err)
 	}
 
 	for {
 		// 5. fetch grant token for launcher
 
-		launcherGrantToken, err := auth.FetchGrantToken(n8nRunnersServerURI, token)
+		launcherGrantToken, err := auth.FetchGrantTokenWithRetry(n8nRunnersServerURI, token)
 		if err != nil {
 			return fmt.Errorf("failed to fetch grant token for launcher: %w", err)
 		}
@@ -120,7 +120,7 @@ func (l *LaunchCommand) Execute() error {
 
 		// 7. fetch grant token for runner
 
-		runnerGrantToken, err := auth.FetchGrantToken(n8nRunnersServerURI, token)
+		runnerGrantToken, err := auth.FetchGrantTokenWithRetry(n8nRunnersServerURI, token)
 		if err != nil {
 			return fmt.Errorf("failed to fetch grant token for runner: %w", err)
 		}

internal/http/retry.go

@@ -25,14 +25,15 @@ type retryConfig struct {
 	WaitTimeBetweenRetries time.Duration
 }
 
-func retry(operationName string, operationFn func() error, cfg retryConfig) error {
+func retry[T any](operationName string, operationFn func() (T, error), cfg retryConfig) (T, error) {
 	var lastErr error
+	var zero T
 	startTime := time.Now()
 	attempt := 1
 
 	for {
 		if cfg.MaxRetryTime > 0 && time.Since(startTime) > cfg.MaxRetryTime {
-			return fmt.Errorf(
+			return zero, fmt.Errorf(
 				"gave up retrying operation `%s` on reaching max retry time %v, last error: %w",
 				operationName,
 				cfg.MaxRetryTime,
@@ -41,17 +42,17 @@ func retry(operationName string, operationFn func() error, cfg retryConfig) erro
 		}
 
 		if cfg.MaxAttempts > 0 && attempt > cfg.MaxAttempts {
-			return fmt.Errorf(
+			return zero, fmt.Errorf(
 				"gave up retrying operation `%s` on reaching max retry attempts %d, last error: %w",
 				operationName,
 				cfg.MaxAttempts,
 				lastErr,
 			)
 		}
 
-		err := operationFn()
+		str, err := operationFn()
 		if err == nil {
-			return nil
+			return str, nil
 		}
 
 		lastErr = err
@@ -63,7 +64,7 @@ func retry(operationName string, operationFn func() error, cfg retryConfig) erro
 }
 
 // UnlimitedRetry retries an operation forever.
-func UnlimitedRetry(operationName string, operation func() error) error {
+func UnlimitedRetry[T any](operationName string, operation func() (T, error)) (T, error) {
 	return retry(operationName, operation, retryConfig{
 		MaxRetryTime:           0,
 		MaxAttempts:            0,
@@ -72,7 +73,7 @@ func UnlimitedRetry(operationName string, operation func() error) error {
 }
 
 // LimitedRetry retries an operation until max retry time or until max attempts.
-func LimitedRetry(operationName string, operation func() error) error {
+func LimitedRetry[T any](operationName string, operation func() (T, error)) (T, error) {
 	return retry(operationName, operation, retryConfig{
 		MaxRetryTime:           defaultMaxRetryTime,
 		MaxAttempts:            defaultMaxRetries,

internal/http/wait_for_n8n.go

@@ -28,27 +28,27 @@ func sendReadinessRequest(n8nMainServerURI string) (*http.Response, error) {
 	return client.Do(req)
 }
 
-// WaitForN8n retries indefinitely until the n8n main instance is ready, i.e.
-// until its database is connected and migrated. In case of long-running
-// migrations, n8n instance readiness may take a long time.
-func WaitForN8n(n8nMainServerURI string) error {
+// WaitForN8nReadiness checks forever until the n8n main instance is ready, i.e.
+// until its DB is connected and migrated. In case of long-running migrations,
+// readiness may take a long time. Returns nil when ready.
+func WaitForN8nReadiness(n8nMainServerURI string) error {
 	logs.Info("Waiting for n8n to be ready...")
 
-	readinessCheck := func() error {
+	readinessCheck := func() (string, error) {
 		resp, err := sendReadinessRequest(n8nMainServerURI)
 		if err != nil {
-			return fmt.Errorf("failed to send readiness check request to n8n main server: %w", err)
+			return "", fmt.Errorf("n8n readiness check failed with error: %w", err)
 		}
 		defer resp.Body.Close()
 
 		if resp.StatusCode != http.StatusOK {
-			return fmt.Errorf("readiness check failed with status code: %d", resp.StatusCode)
+			return "", fmt.Errorf("readiness check failed with status code: %d", resp.StatusCode)
 		}
 
-		return nil
+		return "", nil
 	}
 
-	if err := UnlimitedRetry("readiness-check", readinessCheck); err != nil {
+	if _, err := UnlimitedRetry("readiness-check", readinessCheck); err != nil {
 		return fmt.Errorf("encountered error while waiting for n8n to be ready: %w", err)
 	}