feat(kbagent): implement request-level timeout override and retry for action execution (#10158)

Co-authored-by: gongna-au <gongna1@xiaomi.com>

Author: gongna-au

pkg/kbagent/proto/proto.go

@@ -74,8 +74,8 @@ type ActionRequest struct {
 	Action         string            `json:"action"`
 	Parameters     map[string]string `json:"parameters,omitempty"`
 	NonBlocking    *bool             `json:"nonBlocking,omitempty"`
-	TimeoutSeconds *int32            `json:"timeoutSeconds,omitempty"` // TODO: not implemented
-	RetryPolicy    *RetryPolicy      `json:"retryPolicy,omitempty"`    // TODO: not implemented
+	TimeoutSeconds *int32            `json:"timeoutSeconds,omitempty"`
+	RetryPolicy    *RetryPolicy      `json:"retryPolicy,omitempty"`
 }
 
 type ActionResponse struct {

pkg/kbagent/service/action.go

@@ -26,6 +26,7 @@ import (
 	"net"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/go-logr/logr"
 	"github.com/pkg/errors"
@@ -124,19 +125,20 @@ func (s *actionService) handleRequest(ctx context.Context, req *proto.ActionRequ
 	if err := checkReconfigure(ctx, req); err != nil {
 		return nil, err
 	}
+	timeout := resolveTimeout(&action.TimeoutSeconds, req.TimeoutSeconds)
 	if req.NonBlocking == nil || !*req.NonBlocking {
-		return blockingCallAction(ctx, action, req.Parameters, &action.TimeoutSeconds)
+		return callActionWithRetry(ctx, action, req.Parameters, timeout, req.RetryPolicy)
 	}
-	return s.handleRequestNonBlocking(ctx, req, action)
+	return s.handleRequestNonBlocking(ctx, req, action, timeout)
 }
 
-func (s *actionService) handleRequestNonBlocking(ctx context.Context, req *proto.ActionRequest, action *proto.Action) ([]byte, error) {
+func (s *actionService) handleRequestNonBlocking(ctx context.Context, req *proto.ActionRequest, action *proto.Action, timeout *int32) ([]byte, error) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
 	running, ok := s.runningActions[req.Action]
 	if !ok {
-		resultChan, err := nonBlockingCallAction(ctx, action, req.Parameters, &action.TimeoutSeconds)
+		resultChan, err := nonBlockingCallAction(ctx, action, req.Parameters, timeout)
 		if err != nil {
 			return nil, err
 		}
@@ -155,3 +157,33 @@ func (s *actionService) handleRequestNonBlocking(ctx context.Context, req *proto
 	}
 	return (*result).stdout.Bytes(), nil
 }
+
+func resolveTimeout(actionTimeout *int32, requestTimeout *int32) *int32 {
+	if requestTimeout != nil {
+		return requestTimeout
+	}
+	return actionTimeout
+}
+
+func callActionWithRetry(ctx context.Context, action *proto.Action, parameters map[string]string, timeout *int32, retryPolicy *proto.RetryPolicy) ([]byte, error) {
+	output, err := blockingCallAction(ctx, action, parameters, timeout)
+	if err == nil || retryPolicy == nil || retryPolicy.MaxRetries <= 0 {
+		return output, err
+	}
+
+	interval := retryPolicy.RetryInterval
+	for i := 0; i < retryPolicy.MaxRetries; i++ {
+		if interval > 0 {
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			case <-time.After(interval):
+			}
+		}
+		output, err = blockingCallAction(ctx, action, parameters, timeout)
+		if err == nil {
+			return output, nil
+		}
+	}
+	return output, err
+}