diff --git a/apis/apps/v1/shardingdefinition_types.go b/apis/apps/v1/shardingdefinition_types.go index 002459c2417..e87da660d19 100644 --- a/apis/apps/v1/shardingdefinition_types.go +++ b/apis/apps/v1/shardingdefinition_types.go @@ -195,6 +195,10 @@ type ShardingLifecycleActions struct { // // - KB_ADD_SHARD_NAME: The name of the shard being added. // + // This Action supports both blocking and non-blocking modes. KubeBlocks does + // not finish adding the shard until the Action has succeeded on every + // selected target shard and Pod. + // // Note: This field is immutable once it has been set. // // +optional @@ -206,6 +210,10 @@ type ShardingLifecycleActions struct { // // - KB_REMOVE_SHARD_NAME: The name of the shard being removed. // + // This Action supports both blocking and non-blocking modes. KubeBlocks does + // not remove the shard until the Action has succeeded on every selected + // target shard and Pod. + // // Note: This field is immutable once it has been set. // // +optional diff --git a/config/crd/bases/apps.kubeblocks.io_shardingdefinitions.yaml b/config/crd/bases/apps.kubeblocks.io_shardingdefinitions.yaml index e2e02d7a8cd..8ef858a5526 100644 --- a/config/crd/bases/apps.kubeblocks.io_shardingdefinitions.yaml +++ b/config/crd/bases/apps.kubeblocks.io_shardingdefinitions.yaml @@ -1002,6 +1002,10 @@ spec: - KB_ADD_SHARD_NAME: The name of the shard being added. + This Action supports both blocking and non-blocking modes. KubeBlocks does + not finish adding the shard until the Action has succeeded on every + selected target shard and Pod. + Note: This field is immutable once it has been set. properties: exec: @@ -1463,6 +1467,10 @@ spec: - KB_REMOVE_SHARD_NAME: The name of the shard being removed. + This Action supports both blocking and non-blocking modes. KubeBlocks does + not remove the shard until the Action has succeeded on every selected + target shard and Pod. + Note: This field is immutable once it has been set. properties: exec: diff --git a/controllers/apps/cluster/sharding_actions.go b/controllers/apps/cluster/sharding_actions.go new file mode 100644 index 00000000000..16fe65365fe --- /dev/null +++ b/controllers/apps/cluster/sharding_actions.go @@ -0,0 +1,199 @@ +/* +Copyright (C) 2022-2026 ApeCloud Co., Ltd + +This file is part of KubeBlocks project + +KubeBlocks is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as published by +the Free Software Foundation, either version 3 of the License, or +(at your option) any later version. + +KubeBlocks is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with KubeBlocks. If not, see . +*/ + +package cluster + +import ( + "slices" + "time" + + "golang.org/x/exp/maps" + "k8s.io/apimachinery/pkg/util/sets" + + appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1" + ictrlutil "github.com/apecloud/kubeblocks/pkg/controllerutil" +) + +type shardingActions struct { + shardingHandler *clusterShardingHandler + transCtx *clusterTransformContext + shardingName string + actions *appsv1.ShardingLifecycleActions + runningComps map[string]*appsv1.Component + toCreate sets.Set[string] + toDelete sets.Set[string] + toUpdate sets.Set[string] +} + +func (h *clusterShardingHandler) handleShardActions(transCtx *clusterTransformContext, + shardingName string, runningComps map[string]*appsv1.Component, + toCreate, toDelete, toUpdate sets.Set[string]) (sets.Set[string], error) { + actions := &shardingActions{ + shardingHandler: h, + transCtx: transCtx, + shardingName: shardingName, + runningComps: runningComps, + toCreate: toCreate, + toDelete: toDelete, + toUpdate: toUpdate, + } + if shardingDef := h.shardingDef(transCtx, shardingName); shardingDef != nil { + actions.actions = shardingDef.Spec.LifecycleActions + } + return actions.reconcile() +} + +func (h *shardingActions) reconcile() (sets.Set[string], error) { + topologyBlocked, err := h.reconcileActions() + // Restoring a still-desired participant is not a new topology change: the + // outstanding request may need this Component in order to make progress. + for _, comp := range h.runningComps { + for _, annotation := range []string{shardingAddActionTargetsKey, shardingRemoveActionTargetsKey} { + targets, found, parseErr := getShardingActionTargets(comp, annotation) + if parseErr != nil || !found { + continue // reconcileActions reports invalid snapshots. + } + for _, target := range targets.Targets { + if h.toCreate.Has(target.Component) { + topologyBlocked.Delete(target.Component) + } + } + } + } + + return topologyBlocked, err +} + +func (h *shardingActions) reconcileActions() (sets.Set[string], error) { + all := h.toCreate.Union(h.toDelete).Union(h.toUpdate) + names := sets.List(sets.KeySet(h.runningComps)) + // An already-started action takes precedence over the latest topology diff. + for _, name := range names { + comp := h.runningComps[name] + switch { + case comp.Annotations[shardingAddActionTargetsKey] != "": + return h.advanceAction(comp, false, all) + case comp.Annotations[shardingRemoveActionTargetsKey] != "": + return h.advanceAction(comp, true, all) + } + } + + blocked, completedRemoves := sets.New[string](), sets.New[string]() + var result error + for _, name := range append(sets.List(h.toUpdate), sets.List(h.toDelete)...) { + comp := h.runningComps[name] + deleting := h.toDelete.Has(name) + for _, remove := range []bool{false, true} { + if (remove && !deleting) || (!remove && comp.Annotations[shardingAddShardKey] == "") { + continue + } + var action *appsv1.ShardingAction + if h.actions != nil { + action = h.actions.ShardAdd + if remove { + action = h.actions.ShardRemove + } + } + if action != nil && action.NonBlocking { + blocked = all.Difference(completedRemoves) + if result != nil { + return blocked, result + } + // Complete earlier deletes before a new request freezes its targets. + if len(completedRemoves) > 0 { + return blocked, pendingShardingAction("sharding", "waiting for shard deletion") + } + // Finish partial creation before selecting potential participants. + // Pod startup and template variables can require all Components. + if len(h.toCreate) > 0 { + return blocked.Difference(h.toCreate), pendingShardingAction("sharding", "waiting for shard creation") + } + return h.advanceAction(comp, remove, blocked) + } + if err := h.callAction(comp, remove); err != nil { + h.transCtx.Logger.Error(err, "failed to call sharding action", "shard", name) + if result == nil { + result = err + } + if deleting { + blocked.Insert(name) + } + break + } + if remove { + completedRemoves.Insert(name) + } + } + } + return blocked, result +} + +// Complete and persist one source request before selecting the next. The same +// topology guard applies to first selection, polling, retries and completion. +func (h *shardingActions) advanceAction(comp *appsv1.Component, remove bool, + blocked sets.Set[string]) (sets.Set[string], error) { + // A DELETE from a preceding reconciliation may still be in progress. + for _, target := range h.runningComps { + if !target.DeletionTimestamp.IsZero() { + return blocked, pendingShardingAction("sharding", "waiting for shard deletion") + } + } + if err := h.callAction(comp, remove); err != nil { + if !ictrlutil.IsDelayedRequeueError(err) { + h.transCtx.Logger.Error(err, "failed to call sharding action", "shard", comp.Name) + } + return blocked, err + } + if remove { + if h.toDelete.Has(comp.Name) { + // Keep the successful remove snapshot until deletion so that a failed + // DELETE can retry the cached result without starting a new request. + blocked.Delete(comp.Name) + return blocked, pendingShardingAction("sharding", "waiting for shard deletion") + } + delete(comp.Annotations, shardingRemoveActionTargetsKey) + if h.actions != nil && h.actions.ShardAdd != nil { + comp.Annotations[shardingAddShardKey] = time.Now().Format(time.RFC3339Nano) + } + } + return blocked, pendingShardingAction("sharding", "waiting for completed action state to persist") +} + +func (h *shardingActions) callAction(comp *appsv1.Component, remove bool) error { + if remove { + return h.shardingHandler.handleShardRemove(h.transCtx, h.shardingName, maps.Values(h.runningComps), comp) + } + return h.shardingHandler.handleShardAdd(h.transCtx, h.shardingName, maps.Values(h.runningComps), comp) +} + +func hasPendingNonBlockingAction(comp *appsv1.Component) bool { + if comp == nil || comp.Annotations == nil { + return false + } + if comp.Annotations[shardingAddActionTargetsKey] != "" || + comp.Annotations[shardingRemoveActionTargetsKey] != "" { + return true + } + if comp.Annotations[shardingAddShardKey] == "" { + return false + } + return slices.ContainsFunc(comp.Spec.CustomActions, func(action appsv1.CustomAction) bool { + return action.Name == shardingAddShardAction && action.Action != nil && action.Action.NonBlocking + }) +} diff --git a/controllers/apps/cluster/sharding_nonblocking_action.go b/controllers/apps/cluster/sharding_nonblocking_action.go new file mode 100644 index 00000000000..6223278d61b --- /dev/null +++ b/controllers/apps/cluster/sharding_nonblocking_action.go @@ -0,0 +1,358 @@ +/* +Copyright (C) 2022-2026 ApeCloud Co., Ltd + +This file is part of KubeBlocks project + +KubeBlocks is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as published by +the Free Software Foundation, either version 3 of the License, or +(at your option) any later version. + +KubeBlocks is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with KubeBlocks. If not, see . +*/ + +package cluster + +import ( + "encoding/json" + "errors" + "fmt" + "sort" + "time" + + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/util/sets" + + appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1" + "github.com/apecloud/kubeblocks/pkg/constant" + "github.com/apecloud/kubeblocks/pkg/controller/component" + "github.com/apecloud/kubeblocks/pkg/controller/lifecycle" + ictrlutil "github.com/apecloud/kubeblocks/pkg/controllerutil" +) + +const ( + shardingActionTargetsVersion = 1 + shardingAddActionTargetsKey = "kubeblocks.io/sharding-add-action-targets" + shardingRemoveActionTargetsKey = "kubeblocks.io/sharding-remove-action-targets" +) + +type shardingActionTargets struct { + Version int `json:"version"` + Targets []shardingActionTarget `json:"targets"` +} + +func (h *clusterShardingHandler) nonBlockingShardingAction(transCtx *clusterTransformContext, + shardingName, actionName, targetsAnnotation string, action *appsv1.ShardingAction, + args map[string]string, runningComps []*appsv1.Component, sourceComp *appsv1.Component) error { + targets, changed, err := h.resolveShardingActionTargets( + transCtx, action, targetsAnnotation, runningComps, sourceComp) + if err != nil { + return err + } + if changed { + if err := setShardingActionTargets(sourceComp, targetsAnnotation, targets); err != nil { + return err + } + return pendingShardingAction(actionName, "targets selected") + } + + comps := make(map[string]*appsv1.Component, len(runningComps)) + for _, comp := range runningComps { + comps[comp.Name] = comp + } + + var callErrors []error + pending := false + for i := range targets.Targets { + target := &targets.Targets[i] + lfa, err := h.newLifecycle(transCtx, comps[target.Component], target.TemplateVars) + if err != nil { + callErrors = append(callErrors, err) + continue + } + for j := range target.Pods { + pod := &target.Pods[j] + opts := &lifecycle.Options{ + Query: pod.Query, + TargetPodName: pod.Name, + PreConditionObjectSelector: constant.GetClusterLabels(transCtx.Cluster.Name, + map[string]string{constant.KBAppShardingNameLabelKey: shardingName}), + } + err := lfa.UserDefined(transCtx.Context, transCtx.Client, opts, actionName, &action.Action, args) + if err = lifecycle.IgnoreNotDefined(err); err == nil { + pod.Query = true + continue + } + switch { + case errors.Is(err, lifecycle.ErrActionInProgress): + pod.Query = true + pending = true + case errors.Is(err, lifecycle.ErrActionBusy): + pending = true + case errors.Is(err, lifecycle.ErrActionResultNotFound): + // Persist the restart decision first. The next call must pass + // startup preconditions again before it can execute anything. + pod.Query = false + pending = true + case isTerminalShardingActionError(err): + pod.Query = false + callErrors = append(callErrors, err) + default: + callErrors = append(callErrors, err) + } + } + } + if err := setShardingActionTargets(sourceComp, targetsAnnotation, targets); err != nil { + return err + } + if len(callErrors) > 0 { + return errors.Join(callErrors...) + } + if pending { + return pendingShardingAction(actionName, "still running") + } + if targetsAnnotation == shardingAddActionTargetsKey { + delete(sourceComp.Annotations, targetsAnnotation) + } + return nil +} + +func pendingShardingAction(actionName, reason string) error { + return ictrlutil.NewDelayedRequeueError(3*time.Second, + fmt.Sprintf("action %s is %s", actionName, reason)) +} +func isTerminalShardingActionError(err error) bool { + return errors.Is(err, lifecycle.ErrActionFailed) || + errors.Is(err, lifecycle.ErrActionTimedOut) || + errors.Is(err, lifecycle.ErrActionInternalError) +} + +func (h *clusterShardingHandler) resolveShardingActionTargets(transCtx *clusterTransformContext, + action *appsv1.ShardingAction, targetsAnnotation string, runningComps []*appsv1.Component, + sourceComp *appsv1.Component) (*shardingActionTargets, bool, error) { + targets, found, err := getShardingActionTargets(sourceComp, targetsAnnotation) + if err != nil { + return nil, false, err + } + if !found { + targets, err = h.selectShardingActionTargets(transCtx, action, runningComps, sourceComp) + return targets, true, err + } + + comps := make(map[string]*appsv1.Component, len(runningComps)) + for _, comp := range runningComps { + comps[comp.Name] = comp + } + changed := false + for i := range targets.Targets { + target := &targets.Targets[i] + comp := comps[target.Component] + if comp == nil { + return nil, false, pendingShardingAction("sharding", + fmt.Sprintf("waiting for target shard %s", target.Component)) + } + pods, err := component.ListOwnedInstances(transCtx.Context, transCtx.Client, comp) + if err != nil { + return nil, false, err + } + existing := sets.New[string]() + for _, pod := range pods { + existing.Insert(pod.Name) + } + surviving := sets.New[string]() + missing := 0 + for _, pod := range target.Pods { + if existing.Has(pod.Name) { + surviving.Insert(pod.Name) + } else { + missing++ + } + } + if missing == 0 { + continue + } + selected, err := selectShardingActionPods(action, pods, comp.Name) + if err != nil { + return nil, false, pendingShardingAction("sharding", + fmt.Sprintf("waiting for replacement pods on shard %s", comp.Name)) + } + replacements := make([]shardingActionTargetPod, 0, missing) + for _, pod := range selected { + if !surviving.Has(pod.Name) { + replacements = append(replacements, pod) + } + } + if len(replacements) < missing { + return nil, false, pendingShardingAction("sharding", + fmt.Sprintf("waiting for %d replacement pods on shard %s", missing, comp.Name)) + } + next := 0 + for j := range target.Pods { + if !existing.Has(target.Pods[j].Name) { + target.Pods[j] = replacements[next] + next++ + } + } + changed = true + } + return targets, changed, nil +} + +func (h *clusterShardingHandler) selectShardingActionTargets(transCtx *clusterTransformContext, + action *appsv1.ShardingAction, runningComps []*appsv1.Component, + sourceComp *appsv1.Component) (*shardingActionTargets, error) { + shards, err := h.selectTargetShard(action, runningComps, sourceComp) + if err != nil { + return nil, err + } + targets := &shardingActionTargets{Version: shardingActionTargetsVersion} + for _, shard := range shards { + pods, err := component.ListOwnedInstances(transCtx.Context, transCtx.Client, shard) + if err != nil { + return nil, err + } + // Do not freeze a partial scale-out or surplus scale-in replica into + // a request. The topology may still be converging from an earlier update. + if shard.Generation != shard.Status.ObservedGeneration || len(pods) != int(shard.Spec.Replicas) { + return nil, pendingShardingAction("sharding", fmt.Sprintf("waiting for shard %s pod topology", shard.Name)) + } + for _, pod := range pods { + if !pod.DeletionTimestamp.IsZero() { + return nil, pendingShardingAction("sharding", fmt.Sprintf("waiting for shard %s pod deletion", shard.Name)) + } + } + compDef := transCtx.componentDefs[shard.Spec.CompDef] + if compDef == nil { + return nil, fmt.Errorf("component definition not found for shard %s", shard.Name) + } + synthesized, err := component.BuildSynthesizedComponent(transCtx.Context, transCtx.Client, compDef, shard) + if err != nil { + return nil, err + } + // Resolve once for this request. Secret-backed environment references are + // kept as references by the existing template-variable resolver. + vars, _, err := component.ResolveTemplateNEnvVars(transCtx.Context, transCtx.Client, synthesized, compDef.Spec.Vars) + if err != nil { + return nil, err + } + selectedPods, err := selectShardingActionPods(action, pods, shard.Name) + if err != nil { + return nil, err + } + targets.Targets = append(targets.Targets, shardingActionTarget{ + Component: shard.Name, + Pods: selectedPods, + TemplateVars: vars, + }) + } + return targets, nil +} + +func selectShardingActionPods(action *appsv1.ShardingAction, pods []*corev1.Pod, + componentName string) ([]shardingActionTargetPod, error) { + if len(pods) == 0 { + return nil, fmt.Errorf("shard %s has no pods to execute action", componentName) + } + selected, err := lifecycle.SelectTargetPods(pods, pods[0], &action.Action) + if err != nil { + return nil, err + } + if len(selected) == 0 { + return nil, fmt.Errorf("shard %s has no pod matching the action target selector", componentName) + } + targets := make([]shardingActionTargetPod, 0, len(selected)) + for _, pod := range selected { + targets = append(targets, shardingActionTargetPod{Name: pod.Name}) + } + return targets, nil +} + +type shardingActionTarget struct { + Component string `json:"component"` + Pods []shardingActionTargetPod `json:"pods"` + TemplateVars map[string]string `json:"templateVars,omitempty"` +} + +type shardingActionTargetPod struct { + Name string `json:"name"` + Query bool `json:"query,omitempty"` +} + +func getShardingActionTargets(comp *appsv1.Component, annotation string) (*shardingActionTargets, bool, error) { + value, found := comp.Annotations[annotation] + if !found { + return nil, false, nil + } + + targets := &shardingActionTargets{} + if err := json.Unmarshal([]byte(value), targets); err != nil { + return nil, false, fmt.Errorf("invalid %s annotation on component %s: %w", annotation, comp.Name, err) + } + if err := validateShardingActionTargets(targets); err != nil { + return nil, false, fmt.Errorf("invalid %s annotation on component %s: %w", annotation, comp.Name, err) + } + return targets, true, nil +} + +func setShardingActionTargets(comp *appsv1.Component, annotation string, targets *shardingActionTargets) error { + sortShardingActionTargets(targets) + data, err := json.Marshal(targets) + if err != nil { + return err + } + if comp.Annotations == nil { + comp.Annotations = map[string]string{} + } + comp.Annotations[annotation] = string(data) + return nil +} + +func sortShardingActionTargets(targets *shardingActionTargets) { + for i := range targets.Targets { + sort.Slice(targets.Targets[i].Pods, func(j, k int) bool { + return targets.Targets[i].Pods[j].Name < targets.Targets[i].Pods[k].Name + }) + } + sort.Slice(targets.Targets, func(i, j int) bool { + return targets.Targets[i].Component < targets.Targets[j].Component + }) +} + +func validateShardingActionTargets(targets *shardingActionTargets) error { + if targets.Version != shardingActionTargetsVersion { + return fmt.Errorf("unsupported version %d", targets.Version) + } + if len(targets.Targets) == 0 { + return fmt.Errorf("targets must not be empty") + } + components := sets.New[string]() + pods := sets.New[string]() + for _, target := range targets.Targets { + if target.Component == "" { + return fmt.Errorf("target component must not be empty") + } + if components.Has(target.Component) { + return fmt.Errorf("duplicate target component %s", target.Component) + } + components.Insert(target.Component) + if len(target.Pods) == 0 { + return fmt.Errorf("target component %s has no pods", target.Component) + } + for _, pod := range target.Pods { + if pod.Name == "" { + return fmt.Errorf("target pod name must not be empty") + } + if pods.Has(pod.Name) { + return fmt.Errorf("duplicate target pod %s", pod.Name) + } + pods.Insert(pod.Name) + } + } + return nil +} diff --git a/controllers/apps/cluster/transformer_cluster_component.go b/controllers/apps/cluster/transformer_cluster_component.go index 20969124aab..b1b346f5ac8 100644 --- a/controllers/apps/cluster/transformer_cluster_component.go +++ b/controllers/apps/cluster/transformer_cluster_component.go @@ -173,6 +173,9 @@ func checkAllCompsUpToDate(transCtx *clusterTransformContext, cluster *appsv1.Cl return false, nil } for _, comp := range compList.Items { + if hasPendingNonBlockingAction(&comp) { + return false, nil + } generation, ok := comp.Annotations[constant.KubeBlocksGenerationKey] if !ok { return false, nil @@ -626,7 +629,7 @@ func (c *phasePrecondition) shardingMatch(transCtx *clusterTransformContext, dag return false, nil } for _, comp := range comps { - if !c.expected(&comp) { + if hasPendingNonBlockingAction(&comp) || !c.expected(&comp) { transCtx.Logger.Info("waiting for predecessor sharding in expected phase", "shard", comp.Name, "predecessor sharding", name) return false, nil @@ -902,36 +905,26 @@ func (h *clusterShardingHandler) update(transCtx *clusterTransformContext, dag * protoCompsMap[comp.Name] = protoComps[i] } - toCreate, toDelete, toUpdate := mapDiff(runningCompsMap, protoCompsMap) - if err := h.handlePostProvision(transCtx, name, maps.Values(runningCompsMap)); err != nil { return err } - pendingAdds := make(map[string]*appsv1.Component) - for name := range toDelete { - if comp := runningCompsMap[name]; comp.Annotations[shardingAddShardKey] != "" { - pendingAdds[name] = comp.DeepCopy() - } - } - errorSkip, err3 := h.handleShardAddNRemove(transCtx, name, runningCompsMap, protoCompsMap, toCreate, toDelete, toUpdate) + return h.updateShards(transCtx, dag, name, runningCompsMap, protoCompsMap) +} - // Preserve completed adds when a subsequent remove failure keeps the shard alive. - graphCli, _ := transCtx.Client.(model.GraphClient) - for name, original := range pendingAdds { - if errorSkip.Has(name) && runningCompsMap[name].Annotations[shardingAddShardKey] == "" { - updated := original.DeepCopy() - delete(updated.Annotations, shardingAddShardKey) - graphCli.Update(dag, original, updated) - } - } +// updateShards applies the topology diff after its prerequisite actions have +// advanced. Both execution modes use this same Component write path. +func (h *clusterShardingHandler) updateShards(transCtx *clusterTransformContext, dag *graph.DAG, name string, + runningComps, protoComps map[string]*appsv1.Component) error { + toCreate, toDelete, toUpdate := mapDiff(runningComps, protoComps) + blocked, err := h.handleShardAddNRemove(transCtx, dag, name, runningComps, protoComps, toCreate, toDelete, toUpdate) // TODO: update strategy - h.deleteComps(transCtx, dag, runningCompsMap, toDelete.Difference(errorSkip)) - h.updateComps(transCtx, dag, runningCompsMap, protoCompsMap, toUpdate.Difference(errorSkip)) - h.createComps(transCtx, dag, protoCompsMap, toCreate) + h.deleteComps(transCtx, dag, runningComps, toDelete.Difference(blocked)) + h.updateComps(transCtx, dag, runningComps, protoComps, toUpdate.Difference(blocked)) + h.createComps(transCtx, dag, protoComps, toCreate.Difference(blocked)) - return err3 + return err } func (h *clusterShardingHandler) createComps(transCtx *clusterTransformContext, dag *graph.DAG, @@ -1359,65 +1352,48 @@ func (h *clusterShardingHandler) updateActionStatus(transCtx *clusterTransformCo transCtx.Cluster.Status.Shardings[shardingName] = shardingStatus } -func (h *clusterShardingHandler) handleShardAddNRemove(transCtx *clusterTransformContext, shardingName string, - runningCompsMap map[string]*appsv1.Component, protoCompsMap map[string]*appsv1.Component, +// handleShardAddNRemove owns action prerequisites and their persisted state. +// The returned names block topology writes, not just failed deletions. +func (h *clusterShardingHandler) handleShardAddNRemove(transCtx *clusterTransformContext, dag *graph.DAG, + shardingName string, runningCompsMap, protoCompsMap map[string]*appsv1.Component, toCreate, toDelete, toUpdate sets.Set[string]) (sets.Set[string], error) { - var ( - errorSkip = sets.Set[string]{} - - create = func() { - shardingDef := h.shardingDef(transCtx, shardingName) - if shardingDef != nil && shardingDef.Spec.LifecycleActions != nil && shardingDef.Spec.LifecycleActions.ShardAdd != nil { - now := time.Now().Format(time.RFC3339Nano) - for name := range toCreate { - protoComp := protoCompsMap[name] - if protoComp.Annotations == nil { - protoComp.Annotations = make(map[string]string) - } - protoComp.Annotations[shardingAddShardKey] = now - } - } + originals := make(map[string]*appsv1.Component) + for name, comp := range runningCompsMap { + if hasPendingNonBlockingAction(comp) || comp.Annotations[shardingAddShardKey] != "" || toDelete.Has(name) { + originals[name] = comp.DeepCopy() } - - update = func() error { - var err error - for name := range toUpdate { - err1 := h.handleShardAdd(transCtx, shardingName, maps.Values(runningCompsMap), runningCompsMap[name]) - if err1 != nil { - transCtx.Logger.Error(err1, "failed to call the shard add action", "shard", name) - if err == nil { - err = err1 - } - // errorSkip.Insert(name) - } - } - return err - } - - _delete = func() error { - var err error - for name := range toDelete { - err1 := h.handleShardRemove(transCtx, shardingName, maps.Values(runningCompsMap), runningCompsMap[name]) - if err1 != nil { - transCtx.Logger.Error(err1, "failed to call the shard remove action", "shard", name) - if err == nil { - err = err1 - } - errorSkip.Insert(name) - } + } + shardingDef := h.shardingDef(transCtx, shardingName) + if shardingDef != nil && shardingDef.Spec.LifecycleActions != nil && shardingDef.Spec.LifecycleActions.ShardAdd != nil { + now := time.Now().Format(time.RFC3339Nano) + for name := range toCreate { + comp := protoCompsMap[name] + if comp.Annotations == nil { + comp.Annotations = make(map[string]string) } - return err + comp.Annotations[shardingAddShardKey] = now } - ) + } - create() - err1 := update() - err2 := _delete() + blocked, err := h.handleShardActions(transCtx, shardingName, runningCompsMap, toCreate, toDelete, toUpdate) - if err1 != nil { - return errorSkip, err1 + // Preserve completed adds when a subsequent remove failure keeps the shard alive. + // The same path persists non-blocking action progress for blocked Components. + // Unblocked Components use the ordinary update/delete path. + graphCli, _ := transCtx.Client.(model.GraphClient) + for name, original := range originals { + current := runningCompsMap[name] + if blocked.Has(name) && shardingActionStateChanged(original, current) { + graphCli.Update(dag, original, current.DeepCopy(), &model.ReplaceIfExistingOption{}) + } } - return errorSkip, err2 + return blocked, err +} + +func shardingActionStateChanged(original, current *appsv1.Component) bool { + return original.Annotations[shardingAddActionTargetsKey] != current.Annotations[shardingAddActionTargetsKey] || + original.Annotations[shardingRemoveActionTargetsKey] != current.Annotations[shardingRemoveActionTargetsKey] || + original.Annotations[shardingAddShardKey] != current.Annotations[shardingAddShardKey] } func (h *clusterShardingHandler) handleShardAdd(transCtx *clusterTransformContext, @@ -1426,7 +1402,8 @@ func (h *clusterShardingHandler) handleShardAdd(transCtx *clusterTransformContex shardingDef = h.shardingDef(transCtx, shardingName) pending = func() bool { - return runningComp.Annotations[shardingAddShardKey] != "" + return runningComp.Annotations[shardingAddShardKey] != "" || + runningComp.Annotations[shardingAddActionTargetsKey] != "" } succeed = func() error { @@ -1459,7 +1436,7 @@ func (h *clusterShardingHandler) handleShardRemove(transCtx *clusterTransformCon } ) - if runningComp.Annotations[shardingAddShardKey] != "" { + if runningComp.Annotations[shardingRemoveActionTargetsKey] == "" && runningComp.Annotations[shardingAddShardKey] != "" { if err := h.handleShardAdd(transCtx, shardingName, runningComps, runningComp); err != nil { return err } @@ -1490,6 +1467,16 @@ func (h *clusterShardingHandler) shardingDef(transCtx *clusterTransformContext, func (h *clusterShardingHandler) shardingAction(transCtx *clusterTransformContext, shardingName, actionName string, action *appsv1.ShardingAction, args map[string]string, runningComps []*appsv1.Component, comp *appsv1.Component) error { + if action.NonBlocking { + switch actionName { + case shardingAddShardAction: + return h.nonBlockingShardingAction(transCtx, shardingName, actionName, shardingAddActionTargetsKey, + action, args, runningComps, comp) + case shardingRemoveShardAction: + return h.nonBlockingShardingAction(transCtx, shardingName, actionName, shardingRemoveActionTargetsKey, + action, args, runningComps, comp) + } + } shards, err := h.selectTargetShard(action, runningComps, comp) if err != nil { return err @@ -1530,12 +1517,13 @@ func (h *clusterShardingHandler) selectTargetShard(shardingAction *appsv1.Shardi } } -func (h *clusterShardingHandler) newLifecycle(transCtx *clusterTransformContext, comp *appsv1.Component) (lifecycle.Lifecycle, error) { +func (h *clusterShardingHandler) newLifecycle(transCtx *clusterTransformContext, comp *appsv1.Component, + templateVars ...map[string]string) (lifecycle.Lifecycle, error) { compDef := transCtx.componentDefs[comp.Spec.CompDef] if compDef == nil { return nil, fmt.Errorf("component definition not found for shard %s", comp.Name) } - return component.NewLifecycle(transCtx.Context, transCtx.Client, compDef, comp) + return component.NewLifecycle(transCtx.Context, transCtx.Client, compDef, comp, templateVars...) } func clusterRunningCompNShardingSet(ctx context.Context, cli client.Reader, cluster *appsv1.Cluster) (sets.Set[string], error) { diff --git a/controllers/apps/cluster/transformer_cluster_component_status.go b/controllers/apps/cluster/transformer_cluster_component_status.go index 8ffebca36af..837855fb195 100644 --- a/controllers/apps/cluster/transformer_cluster_component_status.go +++ b/controllers/apps/cluster/transformer_cluster_component_status.go @@ -144,6 +144,9 @@ func (t *clusterComponentStatusTransformer) clusterCompStatus(cluster *appsv1.Cl status.ObservedGeneration = ig status.UpToDate = comp.Generation == comp.Status.ObservedGeneration && ig == cluster.Generation } + if hasPendingNonBlockingAction(comp) { + status.UpToDate = false + } return status } diff --git a/controllers/apps/cluster/transformer_cluster_component_status_test.go b/controllers/apps/cluster/transformer_cluster_component_status_test.go index f20818a0134..173ef2bab9c 100644 --- a/controllers/apps/cluster/transformer_cluster_component_status_test.go +++ b/controllers/apps/cluster/transformer_cluster_component_status_test.go @@ -85,6 +85,28 @@ var _ = Describe("cluster component status transformer", func() { }) Context("component", func() { + It("detects pending non-blocking actions", func() { + comp := &appsv1.Component{ + ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{shardingAddShardKey: "pending"}, + }, + Spec: appsv1.ComponentSpec{ + CustomActions: []appsv1.CustomAction{{ + Name: shardingAddShardAction, + Action: &appsv1.Action{}, + }}, + }, + } + + Expect(hasPendingNonBlockingAction(comp)).Should(BeFalse()) + comp.Spec.CustomActions[0].Action.NonBlocking = true + Expect(hasPendingNonBlockingAction(comp)).Should(BeTrue()) + + comp.Spec.CustomActions = nil + comp.Annotations = map[string]string{shardingRemoveActionTargetsKey: `{"version":1}`} + Expect(hasPendingNonBlockingAction(comp)).Should(BeTrue()) + }) + It("empty", func() { transCtx.components = nil diff --git a/controllers/apps/cluster/transformer_cluster_component_test.go b/controllers/apps/cluster/transformer_cluster_component_test.go index a3659f1c047..776563c5db5 100644 --- a/controllers/apps/cluster/transformer_cluster_component_test.go +++ b/controllers/apps/cluster/transformer_cluster_component_test.go @@ -21,9 +21,16 @@ package cluster import ( "context" + "encoding/json" + "errors" "fmt" + "os" + "path/filepath" "reflect" "strings" + "time" + + "github.com/go-logr/logr" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -33,18 +40,22 @@ import ( corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/sets" "k8s.io/utils/ptr" "sigs.k8s.io/controller-runtime/pkg/client" appsv1 "github.com/apecloud/kubeblocks/apis/apps/v1" + workloadsv1 "github.com/apecloud/kubeblocks/apis/workloads/v1" appsutil "github.com/apecloud/kubeblocks/controllers/apps/util" "github.com/apecloud/kubeblocks/pkg/constant" "github.com/apecloud/kubeblocks/pkg/controller/component" "github.com/apecloud/kubeblocks/pkg/controller/graph" + "github.com/apecloud/kubeblocks/pkg/controller/lifecycle" "github.com/apecloud/kubeblocks/pkg/controller/model" ictrlutil "github.com/apecloud/kubeblocks/pkg/controllerutil" kbacli "github.com/apecloud/kubeblocks/pkg/kbagent/client" kbagentproto "github.com/apecloud/kubeblocks/pkg/kbagent/proto" + kbagentservice "github.com/apecloud/kubeblocks/pkg/kbagent/service" testapps "github.com/apecloud/kubeblocks/pkg/testutil/apps" ) @@ -2626,5 +2637,1271 @@ var _ = Describe("cluster component transformer test", func() { Entry("add succeeds without remove defined", false, false), ) }) + + Context("non-blocking shard actions", func() { + buildShard := func(name string, podNames ...string) (*appsv1.Component, []*corev1.Pod) { + spec := transCtx.shardingComps[sharding1aName][0].DeepCopy() + spec.Name = name + comp := newCompObj(transCtx, spec, func(comp *appsv1.Component) { + comp.Status.Phase = appsv1.RunningComponentPhase + comp.Status.ObservedGeneration = comp.Generation + comp.Spec.Replicas = int32(len(podNames)) + comp.Labels[constant.KBAppShardingNameLabelKey] = sharding1aName + comp.Labels[constant.ShardingDefLabelKey] = shardingDefName + if comp.Annotations == nil { + comp.Annotations = map[string]string{} + } + comp.Annotations[constant.KBAppClusterUIDKey] = "test-uid" + }) + shortName, err := component.ShortName(transCtx.Cluster.Name, comp.Name) + Expect(err).Should(BeNil()) + pods := make([]*corev1.Pod, 0, len(podNames)) + for _, podName := range podNames { + pods = append(pods, &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Namespace: testCtx.DefaultNamespace, + Name: podName, + Labels: map[string]string{ + constant.AppManagedByLabelKey: constant.AppName, + constant.AppInstanceLabelKey: transCtx.Cluster.Name, + constant.KBAppComponentLabelKey: shortName, + }, + }}) + } + return comp, pods + } + + action := func() *appsv1.ShardingAction { + result := mockShardingAction("non-blocking") + result.NonBlocking = true + return result + } + reconcileActions := func(runningComps, protoComps map[string]*appsv1.Component, + toCreate, toDelete, toUpdate sets.Set[string]) (sets.Set[string], error) { + graphCli := transCtx.Client.(model.GraphClient) + dag = newDAG(graphCli, transCtx.Cluster) + return (&clusterShardingHandler{}).handleShardAddNRemove(transCtx, dag, + sharding1aName, runningComps, protoComps, toCreate, toDelete, toUpdate) + } + + It("continues blocking adds into non-blocking remove preparation in one call", func() { + retained, retainedPods := buildShard("shard-0", "pod-0") + removed, removedPods := buildShard("shard-1", "pod-1") + retained.Annotations[shardingAddShardKey] = "pending" + removed.Annotations[shardingAddShardKey] = "pending" + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardAdd: mockShardingAction("blocking-add"), ShardRemove: action(), + } + running := map[string]*appsv1.Component{retained.Name: retained, removed.Name: removed} + desired := retained.DeepCopy() + desired.Spec.Replicas = 2 + calls := []string{} + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls = append(calls, req.Action) + if req.Action == "udf-"+shardingRemoveShardAction && !req.Query { + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress)}, nil + } + return kbagentproto.ActionResponse{}, nil + }).Times(4) + }) + for round := 0; round < 3; round++ { + cli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{ + running[retained.Name], running[removed.Name], retainedPods[0], removedPods[0], + }}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + Expect(ictrlutil.IsDelayedRequeueError((&clusterShardingHandler{}).updateShards( + transCtx, dag, sharding1aName, running, map[string]*appsv1.Component{desired.Name: desired}))).Should(BeTrue()) + var deleted []string + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok { + switch *n.Action { + case model.UPDATE: + Expect(comp.Spec.Replicas).Should(Equal(int32(1))) + running[comp.Name] = comp.DeepCopy() + case model.DELETE: + deleted = append(deleted, comp.Name) + } + } + return nil + }, nil)).Should(Succeed()) + Expect(running[retained.Name].Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + Expect(running[removed.Name].Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + Expect(running[removed.Name].Annotations).Should(HaveKey(shardingRemoveActionTargetsKey)) + if round == 0 { + // Both blocking adds and the remove snapshot commit together, before any remove RPC. + Expect(calls).Should(Equal([]string{"udf-" + shardingAddShardAction, "udf-" + shardingAddShardAction})) + } + if round < 2 { + Expect(deleted).Should(BeEmpty()) + } else { + Expect(deleted).Should(ConsistOf(removed.Name)) + } + } + Expect(calls).Should(Equal([]string{"udf-" + shardingAddShardAction, "udf-" + shardingAddShardAction, + "udf-" + shardingRemoveShardAction, "udf-" + shardingRemoveShardAction})) + }) + + DescribeTable("does not defer blocking actions because an unused hook is non-blocking", func(remove bool) { + first, firstPods := buildShard("shard-0", "pod-0") + second, secondPods := buildShard("shard-1", "pod-1") + actions := &appsv1.ShardingLifecycleActions{ShardAdd: mockShardingAction("add"), ShardRemove: action()} + toDelete, toUpdate := sets.New[string](), sets.New(first.Name, second.Name) + name := "udf-" + shardingAddShardAction + if remove { + actions.ShardAdd, actions.ShardRemove = action(), mockShardingAction("remove") + toDelete, toUpdate = toUpdate, toDelete + name = "udf-" + shardingRemoveShardAction + } else { + first.Annotations[shardingAddShardKey], second.Annotations[shardingAddShardKey] = "pending", "pending" + } + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = actions + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{first, second, firstPods[0], secondPods[0]}}) + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + Expect(req.Action).Should(Equal(name)) + return kbagentproto.ActionResponse{}, nil + }).Times(2) + }) + blocked, err := reconcileActions(map[string]*appsv1.Component{first.Name: first, second.Name: second}, nil, + sets.New[string](), toDelete, toUpdate) + Expect(err).ShouldNot(HaveOccurred()) + Expect(blocked).Should(BeEmpty()) + }, Entry("blocking adds", false), Entry("blocking removes", true)) + + DescribeTable("resumes persisted actions before fresh blocking work", func(remove bool) { + fresh, freshPods := buildShard("shard-0", "pod-0") + pending, pendingPods := buildShard("shard-1", "pod-1") + actions := &appsv1.ShardingLifecycleActions{ShardAdd: action(), ShardRemove: mockShardingAction("remove")} + annotation, name := shardingAddActionTargetsKey, "udf-"+shardingAddShardAction + toDelete, toUpdate := sets.New(fresh.Name), sets.New(pending.Name) + if remove { + actions.ShardAdd, actions.ShardRemove = mockShardingAction("add"), action() + annotation, name = shardingRemoveActionTargetsKey, "udf-"+shardingRemoveShardAction + toDelete, toUpdate = sets.New(pending.Name), sets.New(fresh.Name) + fresh.Annotations[shardingAddShardKey] = "pending" + } + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = actions + Expect(setShardingActionTargets(pending, annotation, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{Component: pending.Name, Pods: []shardingActionTargetPod{{Name: pendingPods[0].Name, Query: true}}}}, + })).Should(Succeed()) + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{fresh, pending, freshPods[0], pendingPods[0]}}) + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + Expect(req.Action).Should(Equal(name)) + Expect(req.Query).Should(BeTrue()) + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress)}, nil + }).Times(1) + }) + blocked, err := reconcileActions(map[string]*appsv1.Component{fresh.Name: fresh, pending.Name: pending}, nil, + sets.New[string](), toDelete, toUpdate) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(sets.List(blocked)).Should(ConsistOf(fresh.Name, pending.Name)) + }, Entry("pending add before blocking remove", false), Entry("pending remove before blocking add", true)) + + It("waits for earlier blocking removal before freezing non-blocking add targets", func() { + first, firstPods := buildShard("shard-0", "pod-0") + next, nextPods := buildShard("shard-1", "pod-1") + retained, retainedPods := buildShard("shard-2", "pod-2") + next.Annotations[shardingAddShardKey] = "pending" + add := action() + add.TargetShardSelector = appsv1.AllShards + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardAdd: add, ShardRemove: mockShardingAction("blocking-remove"), + } + running := map[string]*appsv1.Component{first.Name: first, next.Name: next, retained.Name: retained} + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + Expect(req.Action).Should(Equal("udf-" + shardingRemoveShardAction)) + Expect(req.Parameters).Should(HaveKeyWithValue(shardingRemoveShardNameVar, first.Name)) + return kbagentproto.ActionResponse{}, nil + }).Times(2) + }) + for round := 0; round < 4; round++ { + switch round { + case 2: + now := metav1.Now() + first.DeletionTimestamp = &now + case 3: + delete(running, first.Name) + } + cli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{ + first, next, retained, firstPods[0], nextPods[0], retainedPods[0], + }}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + Expect(ictrlutil.IsDelayedRequeueError((&clusterShardingHandler{}).updateShards(transCtx, dag, + sharding1aName, running, map[string]*appsv1.Component{retained.Name: retained.DeepCopy()}))).Should(BeTrue()) + var deleted []string + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok && *n.Action == model.DELETE { + deleted = append(deleted, comp.Name) + } + return nil + }, nil)).Should(Succeed()) + if round < 2 { + Expect(deleted).Should(ConsistOf(first.Name)) + } else { + Expect(deleted).Should(BeEmpty()) + } + if round < 3 { + Expect(next.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + } else { + targets, found, err := getShardingActionTargets(next, shardingAddActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets).Should(HaveLen(2)) + Expect([]string{targets.Targets[0].Component, targets.Targets[1].Component}).Should(ConsistOf(next.Name, retained.Name)) + } + } + }) + + It("finishes partial Component creation before preparing the first add", func() { + source, pods := buildShard("shard-0", "pod-0") + missing, missingPods := buildShard("shard-1", "pod-1") + source.Annotations[shardingAddShardKey] = "pending" + transCtx.Cluster.Spec.ComponentSpecs = nil + transCtx.Cluster.Spec.Shardings = []appsv1.ClusterSharding{{Name: sharding1aName, Shards: 2, + Template: appsv1.ClusterComponentSpec{ComponentDef: source.Spec.CompDef}}} + transCtx.componentDefs[source.Spec.CompDef].Spec.Vars = []appsv1.EnvVar{{ + Name: "ALL_HOSTS", ValueFrom: &appsv1.VarSource{ServiceVarRef: &appsv1.ServiceVarSelector{ + ClusterObjectReference: appsv1.ClusterObjectReference{CompDef: source.Spec.CompDef, Name: "headless", + Optional: ptr.To(false), MultipleClusterObjectOption: &appsv1.MultipleClusterObjectOption{ + RequireAllComponentObjects: ptr.To(true), Strategy: appsv1.MultipleClusterObjectStrategyIndividual}}, + ServiceVars: appsv1.ServiceVars{Host: &appsv1.VarRequired}, + }}, + }} + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ShardAdd: action()} + running := map[string]*appsv1.Component{source.Name: source} + objects := []client.Object{transCtx.Cluster, source, pods[0]} + for round := 0; round < 3; round++ { + cli := model.NewGraphClient(&appsutil.MockReader{Objects: objects}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + err := (&clusterShardingHandler{}).updateShards(transCtx, dag, sharding1aName, running, + map[string]*appsv1.Component{source.Name: source.DeepCopy(), missing.Name: missing.DeepCopy()}) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue(), "round %d: %v", round, err) + var created *appsv1.Component + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok && *n.Action == model.CREATE { + Expect(comp.Name).Should(Equal(missing.Name)) + created = comp.DeepCopy() + } + return nil + }, nil)).Should(Succeed()) + if round < 2 { + Expect(source.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + Expect(created).ShouldNot(BeNil()) + Expect(created.Annotations).Should(HaveKey(shardingAddShardKey)) + // Simulate one failed CREATE, then a successful retry. + if round == 1 { + running[created.Name] = created + objects = append(objects, created, missingPods[0]) + for _, comp := range []*appsv1.Component{source, created} { + objects = append(objects, &corev1.Service{ObjectMeta: metav1.ObjectMeta{ + Namespace: comp.Namespace, Name: comp.Name + "-headless"}}) + } + } + } else { + Expect(created).Should(BeNil()) + Expect(source.Annotations).Should(HaveKey(shardingAddActionTargetsKey)) + } + } + }) + + DescribeTable("waits for existing pod topology changes before selecting targets", + func(podCount int, replicas int32, unobserved, deleting bool) { + shard, pods := buildShard("shard-0", "pod-0", "pod-1") + pods = pods[:podCount] + shard.Spec.Replicas = replicas + shard.Generation = 2 + shard.Status.ObservedGeneration = 2 + // Topology stability must not introduce a business-ready precondition. + shard.Status.Phase = appsv1.UpdatingComponentPhase + if unobserved { + shard.Status.ObservedGeneration = 1 + } + if deleting { + now := metav1.Now() + pods[len(pods)-1].DeletionTimestamp = &now + } + setObjects := func() { + objects := []client.Object{shard} + for _, pod := range pods { + objects = append(objects, pod) + } + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: objects}) + } + setObjects() + shardAction := action() + shardAction.TargetPodSelector = appsv1.AllReplicas + calls := 0 + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls++ + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress)}, nil + }).Times(int(replicas)) + }) + poll := func() error { + return (&clusterShardingHandler{}).nonBlockingShardingAction(transCtx, sharding1aName, + shardingAddShardAction, shardingAddActionTargetsKey, shardAction, + nil, []*appsv1.Component{shard}, shard) + } + for i := 0; i < 2; i++ { + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + Expect(calls).Should(BeZero()) + } + + By("selecting the complete stable topology before invoking any target") + _, pods = buildShard("shard-0", "pod-0", "pod-1") + pods = pods[:replicas] + shard.Status.ObservedGeneration = shard.Generation + setObjects() + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) + Expect(calls).Should(BeZero()) + targets, found, err := getShardingActionTargets(shard, shardingAddActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets[0].Pods).Should(HaveLen(int(replicas))) + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) + Expect(calls).Should(Equal(int(replicas))) + }, + Entry("surplus replica during scale-in", 2, int32(1), false, true), + Entry("missing replica during scale-out", 1, int32(2), false, false), + Entry("component update not yet observed", 1, int32(1), true, false), + Entry("deleting pod despite a matching count", 1, int32(1), false, true), + ) + + It("checks startup preconditions for starts and retries but not accepted request polls", func() { + shard, pods := buildShard("shard-0", "pod-0") + shardAction := action() + shardAction.PreCondition = ptr.To(appsv1.ComponentReadyPreConditionType) + shard.Status.Phase = appsv1.UpdatingComponentPhase + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{shard, pods[0]}}) + calls := 0 + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls++ + Expect(req.Query).Should(Equal(calls != 1 && calls != 3)) + switch calls { + case 1, 3: + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress)}, nil + case 2: + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrFailed)}, nil + default: + return kbagentproto.ActionResponse{}, nil + } + }).Times(4) + }) + poll := func() error { + err := (&clusterShardingHandler{}).nonBlockingShardingAction(transCtx, sharding1aName, + shardingAddShardAction, shardingAddActionTargetsKey, shardAction, + nil, []*appsv1.Component{shard}, shard) + Expect(shardAction.PreCondition).ShouldNot(BeNil()) + Expect(*shardAction.PreCondition).Should(Equal(appsv1.ComponentReadyPreConditionType)) + return err + } + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // save targets + Expect(errors.Is(poll(), lifecycle.ErrPreconditionFailed)).Should(BeTrue()) + Expect(calls).Should(BeZero()) + + shard.Status.Phase = appsv1.RunningComponentPhase + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // accepted + shard.Status.Phase = appsv1.UpdatingComponentPhase + Expect(errors.Is(poll(), lifecycle.ErrActionFailed)).Should(BeTrue()) + Expect(calls).Should(Equal(2)) + Expect(errors.Is(poll(), lifecycle.ErrPreconditionFailed)).Should(BeTrue()) // retry still gated + Expect(calls).Should(Equal(2)) + + shard.Status.Phase = appsv1.RunningComponentPhase + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // retry accepted + shard.Status.Phase = appsv1.UpdatingComponentPhase + Expect(poll()).Should(Succeed()) // successful result remains observable + Expect(calls).Should(Equal(4)) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + }) + + DescribeTable("rechecks startup conditions after losing an agent result", func(precondition appsv1.PreConditionType, replacePod bool) { + shard, pods := buildShard("shard-0", "pod-0") + shard.Annotations[shardingAddShardKey] = "pending" + shardAction := action() + shardAction.PreCondition = &precondition + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ShardAdd: shardAction} + its := &workloadsv1.InstanceSet{ObjectMeta: metav1.ObjectMeta{ + Name: shard.Name, Namespace: shard.Namespace, Labels: shard.Labels}, + Spec: workloadsv1.InstanceSetSpec{Replicas: ptr.To(int32(1))}, + Status: workloadsv1.InstanceSetStatus{Replicas: 1, ReadyReplicas: 1, UpdatedReplicas: 1}} + counter := filepath.Join(GinkgoT().TempDir(), "calls") + newAgent := func() kbagentservice.Service { + services, err := kbagentservice.New(logr.Discard(), []kbagentproto.Action{{ + Name: "udf-" + shardingAddShardAction, NonBlocking: true, + Exec: &kbagentproto.ExecAction{Commands: []string{"/bin/sh", "-c", `echo run >> "$1"`, "sh", counter}}, + }}, nil, nil) + Expect(err).ShouldNot(HaveOccurred()) + return services[0] + } + agent := newAgent() + invoke := func(req kbagentproto.ActionRequest) kbagentproto.ActionResponse { + payload, err := json.Marshal(req) + Expect(err).ShouldNot(HaveOccurred()) + data, err := agent.HandleRequest(transCtx.Context, payload) + Expect(err).ShouldNot(HaveOccurred()) + var rsp kbagentproto.ActionResponse + Expect(json.Unmarshal(data, &rsp)).Should(Succeed()) + return rsp + } + var lastRequest kbagentproto.ActionRequest + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + lastRequest = req + return invoke(req), nil + }).AnyTimes() + }) + poll := func() error { + cli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{shard, pods[0], its}}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + err := (&clusterShardingHandler{}).updateShards(transCtx, dag, sharding1aName, + map[string]*appsv1.Component{shard.Name: shard}, map[string]*appsv1.Component{shard.Name: shard.DeepCopy()}) + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok && *n.Action == model.UPDATE { + shard = comp.DeepCopy() + } + return nil + }, nil)).Should(Succeed()) + return err + } + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // snapshot + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // start + lastRequest.Query = true + Eventually(func() string { return invoke(lastRequest).Error }, 3*time.Second).Should(BeEmpty()) + if replacePod { + pods[0] = pods[0].DeepCopy() + pods[0].UID = "replacement" + } + agent = newAgent() // same-name Pod replacement or kb-agent restart + shard.Status.Phase = appsv1.UpdatingComponentPhase + its.Status.ReadyReplicas = 0 + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // miss, no execution + targets, _, err := getShardingActionTargets(shard, shardingAddActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(targets.Targets[0].Pods[0].Query).Should(BeFalse()) + Expect(errors.Is(poll(), lifecycle.ErrPreconditionFailed)).Should(BeTrue()) + data, err := os.ReadFile(counter) + Expect(err).ShouldNot(HaveOccurred()) + Expect(string(data)).Should(Equal("run\n")) + shard.Status.Phase = appsv1.RunningComponentPhase + its.Status.ReadyReplicas = 1 + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // restart after readiness + lastRequest.Query = true + Eventually(func() string { return invoke(lastRequest).Error }, 3*time.Second).Should(BeEmpty()) + shard.Status.Phase = appsv1.UpdatingComponentPhase + its.Status.ReadyReplicas = 0 + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) // terminal result observable + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + data, err = os.ReadFile(counter) + Expect(err).ShouldNot(HaveOccurred()) + Expect(string(data)).Should(Equal("run\nrun\n")) + }, + Entry("ComponentReady after agent restart", appsv1.ComponentReadyPreConditionType, false), + Entry("ComponentReady after Pod replacement", appsv1.ComponentReadyPreConditionType, true), + Entry("RuntimeReady after agent restart", appsv1.RuntimeReadyPreConditionType, false), + Entry("RuntimeReady after Pod replacement", appsv1.RuntimeReadyPreConditionType, true), + ) + + DescribeTable("persists action progress through the shared update path on target errors", func(terminal bool) { + source, pods := buildShard("shard-0", "pod-0", "pod-1") + fresh, _ := buildShard("shard-1", "pod-2") + source.Annotations[shardingAddShardKey] = "pending" + Expect(setShardingActionTargets(source, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{Component: source.Name, Pods: []shardingActionTargetPod{ + {Name: pods[0].Name}, {Name: pods[1].Name}, + }}}, + })).Should(Succeed()) + desired := source.DeepCopy() + desired.Spec.Replicas = 1 + delete(desired.Annotations, shardingAddShardKey) + delete(desired.Annotations, shardingAddActionTargetsKey) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ShardAdd: action()} + round, calls := 0, 0 + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + target := calls % 2 + calls++ + Expect(req.Query).Should(Equal(round != 0 && (round != 1 || target != 1))) + if target == 0 || round == 2 { + return kbagentproto.ActionResponse{}, nil + } + if round == 1 { + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress)}, nil + } + if terminal { + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrFailed)}, nil + } + return kbagentproto.ActionResponse{}, fmt.Errorf("target transport error") + }).Times(6) + }) + + for ; round < 4; round++ { + original := source.DeepCopy() + cli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{source, pods[0], pods[1]}}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + err := (&clusterShardingHandler{}).updateShards(transCtx, dag, sharding1aName, + map[string]*appsv1.Component{source.Name: source}, + map[string]*appsv1.Component{source.Name: desired, fresh.Name: fresh}) + switch round { + case 0: + Expect(err).Should(HaveOccurred()) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeFalse()) + case 1, 2: + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + default: + Expect(err).ShouldNot(HaveOccurred()) + } + var persisted *appsv1.Component + var created []string + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok { + switch *n.Action { + case model.UPDATE: + Expect(n.OriObj).Should(Equal(original)) + persisted = comp.DeepCopy() + case model.CREATE: + created = append(created, comp.Name) + Expect(comp.Annotations).Should(HaveKey(shardingAddShardKey)) + case model.DELETE: + Fail("no Component should be deleted") + } + } + return nil + }, nil)).Should(Succeed()) + Expect(persisted).ShouldNot(BeNil()) + if round < 3 { + Expect(created).Should(BeEmpty()) + Expect(persisted.Spec.Replicas).Should(Equal(int32(2))) + } else { + Expect(created).Should(ConsistOf(fresh.Name)) + Expect(persisted.Spec.Replicas).Should(Equal(int32(1))) + } + if round < 2 { + targets, found, err := getShardingActionTargets(persisted, shardingAddActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets[0].Pods[0].Query).Should(BeTrue()) + Expect(targets.Targets[0].Pods[1].Query).Should(Equal(round != 0)) + } else { + Expect(persisted.Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + Expect(persisted.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + } + source = persisted + } + }, Entry("transport error", false), Entry("terminal failure", true)) + DescribeTable("persists the first shard-add snapshot during scale-in", func(removeDefined bool) { + shard, pods := buildShard("shard-0", "shard-0-0") + shard.Annotations[shardingAddShardKey] = "pending" + actions := &appsv1.ShardingLifecycleActions{ShardAdd: action()} + if removeDefined { + actions.ShardRemove = mockShardingAction("shard-remove") + } + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = actions + var calls []string + addCompleted := false + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls = append(calls, req.Action) + if req.Action == "udf-"+shardingAddShardAction { + Expect(req.Query).Should(Equal(addCompleted)) + if !addCompleted { + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress)}, nil + } + } + return kbagentproto.ActionResponse{}, nil + }).AnyTimes() + }) + + for round := 0; round < 4; round++ { + By(fmt.Sprintf("reconciling scale-in round %d from the previously persisted Component", round)) + original := shard.DeepCopy() + graphCli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{shard, pods[0]}}) + transCtx.Client = graphCli + dag = newDAG(graphCli, transCtx.Cluster) + addCompleted = round >= 2 + err := (&clusterShardingHandler{}).updateShards(transCtx, dag, sharding1aName, + map[string]*appsv1.Component{shard.Name: shard}, map[string]*appsv1.Component{}) + if round < 3 { + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + } else { + // No asynchronous work remains; ordinary deletion needs no extra action round. + Expect(err).ShouldNot(HaveOccurred()) + } + + var persisted *appsv1.Component + deletes := 0 + Expect(dag.WalkReverseTopoOrder(func(vertex graph.Vertex) error { + node := vertex.(*model.ObjectVertex) + if comp, ok := node.Obj.(*appsv1.Component); ok { + switch *node.Action { + case model.DELETE: + deletes++ + Expect(comp.Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + Expect(comp.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + case model.UPDATE: + if round < 3 { + Expect(node.OriObj).Should(Equal(original)) + } + persisted = comp.DeepCopy() + } + } + return nil + }, nil)).Should(Succeed()) + if round < 2 { + Expect(deletes).Should(BeZero()) + Expect(persisted).ShouldNot(BeNil()) + Expect(persisted.Annotations[shardingAddShardKey]).Should(Equal("pending")) + targets, found, err := getShardingActionTargets(persisted, shardingAddActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets).Should(Equal([]shardingActionTarget{{ + Component: shard.Name, Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: round != 0}}, + }})) + Expect(calls).Should(HaveLen(round)) + shard = persisted + } else if round == 2 { + Expect(deletes).Should(BeZero()) + Expect(persisted).ShouldNot(BeNil()) + Expect(persisted.Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + Expect(persisted.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + Expect(calls).Should(HaveLen(2)) + shard = persisted + } else { + Expect(deletes).Should(Equal(1)) + } + } + expectedCalls := []string{"udf-" + shardingAddShardAction, "udf-" + shardingAddShardAction} + if removeDefined { + expectedCalls = append(expectedCalls, "udf-"+shardingRemoveShardAction) + } + Expect(calls).Should(Equal(expectedCalls)) + }, + Entry("without shard-remove", false), + Entry("with blocking shard-remove", true), + ) + + It("waits for a pending action before releasing ordered dependents", func() { + shard, pods := buildShard("shard-0", "pod-0") + shard.Status.ObservedGeneration = shard.Generation + Expect(setShardingActionTargets(shard, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{Component: shard.Name, Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: true}}}}, + })).Should(Succeed()) + transCtx.shardingComps[sharding1aName] = transCtx.shardingComps[sharding1aName][:1] + cli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{shard, pods[0]}}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + Expect(hasPendingNonBlockingAction(shard)).Should(BeTrue()) + ready, err := (&phasePrecondition{}).shardingMatch(transCtx, dag, sharding1aName) + Expect(err).ShouldNot(HaveOccurred()) + Expect(ready).Should(BeFalse()) + delete(shard.Annotations, shardingAddActionTargetsKey) + ready, err = (&phasePrecondition{}).shardingMatch(transCtx, dag, sharding1aName) + Expect(err).ShouldNot(HaveOccurred()) + Expect(ready).Should(BeTrue()) + }) + + It("preserves request parameters while polling the same targets", func() { + shard, pods := buildShard("shard-0", "pod-0") + cm := &corev1.ConfigMap{ObjectMeta: metav1.ObjectMeta{ + Name: "action-input", Namespace: shard.Namespace, + }, Data: map[string]string{"value": "before"}} + secret := &corev1.Secret{ObjectMeta: metav1.ObjectMeta{Name: "action-secret", Namespace: shard.Namespace}, + Data: map[string][]byte{"password": []byte("not-for-the-annotation")}} + transCtx.componentDefs[shard.Spec.CompDef].Spec.Vars = []appsv1.EnvVar{ + {Name: "ACTION_INPUT", ValueFrom: &appsv1.VarSource{ConfigMapKeyRef: &corev1.ConfigMapKeySelector{ + LocalObjectReference: corev1.LocalObjectReference{Name: cm.Name}, Key: "value", + }}}, + {Name: "ACTION_PASSWORD", ValueFrom: &appsv1.VarSource{SecretKeyRef: &corev1.SecretKeySelector{ + LocalObjectReference: corev1.LocalObjectReference{Name: secret.Name}, Key: "password", + }}}, + {Name: "ACTION_REFERENCE", Value: "$(ACTION_PASSWORD)"}, + } + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{shard, pods[0], cm, secret}}) + poll := func() error { + return (&clusterShardingHandler{}).nonBlockingShardingAction(transCtx, sharding1aName, + shardingAddShardAction, shardingAddActionTargetsKey, action(), + map[string]string{shardingAddShardNameVar: shard.Name}, []*appsv1.Component{shard}, shard) + } + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) + targets, found, err := getShardingActionTargets(shard, shardingAddActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets[0].TemplateVars).Should(HaveKeyWithValue("ACTION_INPUT", "before")) + Expect(targets.Targets[0].TemplateVars).ShouldNot(HaveKey("ACTION_PASSWORD")) + Expect(shard.Annotations[shardingAddActionTargetsKey]).ShouldNot(ContainSubstring(string(secret.Data["password"]))) + var requests []kbagentproto.ActionRequest + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + requests = append(requests, req) + return kbagentproto.ActionResponse{Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress)}, nil + }).Times(2) + }) + for i := 0; i < 2; i++ { + Expect(ictrlutil.IsDelayedRequeueError(poll())).Should(BeTrue()) + cm.Data["value"] = "after" + // A continuing request must not even depend on the referenced + // ConfigMap remaining available. + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{shard, pods[0]}}) + } + Expect(requests[0].Parameters).Should(HaveKeyWithValue("ACTION_INPUT", "before")) + Expect(requests[1].Parameters).Should(Equal(requests[0].Parameters)) + }) + + It("recreates a missing target needed by a pending action", func() { + source, pods := buildShard("shard-0", "pod-0") + missing, _ := buildShard("shard-1", "pod-1") + unrelated, _ := buildShard("shard-2", "pod-2") + source.Annotations[shardingAddShardKey] = "pending" + Expect(setShardingActionTargets(source, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{ + {Component: source.Name, Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: true}}}, + {Component: missing.Name, Pods: []shardingActionTargetPod{{Name: "pod-1", Query: true}}}, + }, + })).Should(Succeed()) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ShardAdd: action()} + cli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{source, pods[0]}}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + err := (&clusterShardingHandler{}).updateShards(transCtx, dag, sharding1aName, + map[string]*appsv1.Component{source.Name: source}, + map[string]*appsv1.Component{source.Name: source.DeepCopy(), missing.Name: missing, unrelated.Name: unrelated}) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + created := []string{} + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok && *n.Action == model.CREATE { + created = append(created, comp.Name) + Expect(comp.Annotations).Should(HaveKey(shardingAddShardKey)) + } + return nil + }, nil)).Should(Succeed()) + Expect(created).Should(ConsistOf(missing.Name)) + }) + + It("waits for actual deletion before selecting the next all-shard remove", func() { + first, firstPods := buildShard("shard-0", "pod-0") + next, nextPods := buildShard("shard-1", "pod-1") + retained, retainedPods := buildShard("shard-2", "pod-2") + remove := action() + remove.TargetShardSelector = appsv1.AllShards + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ShardRemove: remove} + Expect(setShardingActionTargets(first, shardingRemoveActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{ + {Component: first.Name, Pods: []shardingActionTargetPod{{Name: firstPods[0].Name, Query: true}}}, + {Component: next.Name, Pods: []shardingActionTargetPod{{Name: nextPods[0].Name, Query: true}}}, + {Component: retained.Name, Pods: []shardingActionTargetPod{{Name: retainedPods[0].Name, Query: true}}}, + }, + })).Should(Succeed()) + calls := 0 + expectedSource := first.Name + testapps.MockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls++ + Expect(req.Parameters).Should(HaveKeyWithValue(shardingRemoveShardNameVar, expectedSource)) + // A failed DELETE must retry cached results, not rerun A. + Expect(req.Query).Should(Equal(expectedSource != next.Name)) + return kbagentproto.ActionResponse{}, nil + }).Times(8) + }) + running := map[string]*appsv1.Component{first.Name: first, next.Name: next, retained.Name: retained} + handle := func() []string { + objects := []client.Object{firstPods[0], nextPods[0], retainedPods[0]} + for _, comp := range running { + objects = append(objects, comp) + } + cli := model.NewGraphClient(&appsutil.MockReader{Objects: objects}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + err := (&clusterShardingHandler{}).updateShards(transCtx, dag, sharding1aName, + running, map[string]*appsv1.Component{retained.Name: retained.DeepCopy()}) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + var deleted []string + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok { + switch *n.Action { + case model.DELETE: + deleted = append(deleted, comp.Name) + case model.UPDATE: + running[comp.Name] = comp.DeepCopy() + } + } + return nil + }, nil)).Should(Succeed()) + return deleted + } + + By("completing A, and retrying its DELETE without selecting B") + for i := 0; i < 2; i++ { + Expect(handle()).Should(ConsistOf(first.Name)) + Expect(running[first.Name].Annotations).Should(HaveKey(shardingRemoveActionTargetsKey)) + Expect(running[next.Name].Annotations).ShouldNot(HaveKey(shardingRemoveActionTargetsKey)) + } + By("waiting while A still exists with a deletion timestamp") + now := metav1.Now() + running[first.Name].DeletionTimestamp = &now + Expect(handle()).Should(BeEmpty()) + Expect(calls).Should(Equal(6)) + Expect(running[next.Name].Annotations).ShouldNot(HaveKey(shardingRemoveActionTargetsKey)) + + By("selecting B only after A has actually disappeared") + delete(running, first.Name) + Expect(handle()).Should(BeEmpty()) + Expect(calls).Should(Equal(6)) + targets, found, err := getShardingActionTargets(running[next.Name], shardingRemoveActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets).Should(HaveLen(2)) + Expect([]string{targets.Targets[0].Component, targets.Targets[1].Component}).Should(ConsistOf(next.Name, retained.Name)) + expectedSource = next.Name + Expect(handle()).Should(ConsistOf(next.Name)) + }) + + It("does not scale down a snapshotted target on initial selection", func() { + source, sourcePods := buildShard("shard-0", "pod-0") + peer, peerPods := buildShard("shard-1", "pod-1", "pod-2") + source.Annotations[shardingAddShardKey] = "pending" + peer.Spec.Replicas = 2 + peerDesired := peer.DeepCopy() + peerDesired.Spec.Replicas = 1 + add := action() + add.TargetShardSelector = appsv1.AllShards + add.TargetPodSelector = appsv1.AllReplicas + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ShardAdd: add} + cli := model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{source, peer, sourcePods[0], peerPods[0], peerPods[1]}}) + transCtx.Client = cli + dag = newDAG(cli, transCtx.Cluster) + err := (&clusterShardingHandler{}).updateShards(transCtx, dag, sharding1aName, + map[string]*appsv1.Component{source.Name: source, peer.Name: peer}, + map[string]*appsv1.Component{source.Name: source.DeepCopy(), peer.Name: peerDesired}) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + targets, found, err := getShardingActionTargets(source, shardingAddActionTargetsKey) + Expect(err).ShouldNot(HaveOccurred()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets[1].Pods).Should(HaveLen(2)) + Expect(dag.WalkReverseTopoOrder(func(v graph.Vertex) error { + n := v.(*model.ObjectVertex) + if comp, ok := n.Obj.(*appsv1.Component); ok && comp.Name == peer.Name && *n.Action == model.UPDATE { + Expect(comp.Spec.Replicas).Should(Equal(int32(2))) + } + return nil + }, nil)).Should(Succeed()) + }) + + It("persists the selected targets before invoking the action", func() { + shard0, pods0 := buildShard("shard-0", "shard-0-0", "shard-0-1") + shard1, pods1 := buildShard("shard-1", "shard-1-0", "shard-1-1") + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{ + shard0, shard1, pods0[0], pods0[1], pods1[0], pods1[1], + }}) + shardAction := action() + shardAction.TargetShardSelector = appsv1.AllShards + shardAction.TargetPodSelector = appsv1.AllReplicas + + err := (&clusterShardingHandler{}).nonBlockingShardingAction( + transCtx, sharding1aName, shardingAddShardAction, shardingAddActionTargetsKey, + shardAction, nil, []*appsv1.Component{shard1, shard0}, shard0) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + + targets, found, err := getShardingActionTargets(shard0, shardingAddActionTargetsKey) + Expect(err).Should(BeNil()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets).Should(HaveLen(2)) + Expect(targets.Targets[0].Component).Should(Equal(shard0.Name)) + Expect(targets.Targets[0].Pods).Should(HaveLen(2)) + Expect(targets.Targets[0].Pods[0].Query).Should(BeFalse()) + Expect(targets.Targets[1].Component).Should(Equal(shard1.Name)) + }) + + It("repairs only missing pod slots in a persisted snapshot", func() { + shard, pods := buildShard("shard-0", "pod-surviving", "pod-replacement") + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{ + Objects: []client.Object{shard, pods[0], pods[1]}, + }) + Expect(setShardingActionTargets(shard, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: shard.Name, + Pods: []shardingActionTargetPod{ + {Name: pods[0].Name, Query: true}, + {Name: "pod-missing", Query: true}, + }, + }}, + })).Should(Succeed()) + shardAction := action() + shardAction.TargetPodSelector = appsv1.AllReplicas + + targets, changed, err := (&clusterShardingHandler{}).resolveShardingActionTargets( + transCtx, shardAction, shardingAddActionTargetsKey, []*appsv1.Component{shard}, shard) + Expect(err).Should(BeNil()) + Expect(changed).Should(BeTrue()) + Expect(targets.Targets).Should(HaveLen(1)) + Expect(targets.Targets[0].Component).Should(Equal(shard.Name)) + Expect(targets.Targets[0].Pods).Should(ConsistOf( + shardingActionTargetPod{Name: pods[0].Name, Query: true}, + shardingActionTargetPod{Name: pods[1].Name})) + }) + + It("does not replace or shrink persisted all-shard targets", func() { + retained, pods := buildShard("shard-0", "pod-0") + replacement, replacementPods := buildShard("shard-2", "pod-2") + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{ + retained, replacement, pods[0], replacementPods[0], + }}) + Expect(setShardingActionTargets(retained, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{ + {Component: retained.Name, Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: true}}}, + {Component: "shard-1", Pods: []shardingActionTargetPod{{Name: "pod-1", Query: true}}}, + }, + })).Should(Succeed()) + shardAction := action() + shardAction.TargetShardSelector = appsv1.AllShards + + targets, changed, err := (&clusterShardingHandler{}).resolveShardingActionTargets( + transCtx, shardAction, shardingAddActionTargetsKey, + []*appsv1.Component{replacement, retained}, retained) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(targets).Should(BeNil()) + Expect(changed).Should(BeFalse()) + persisted, _, err := getShardingActionTargets(retained, shardingAddActionTargetsKey) + Expect(err).Should(BeNil()) + Expect(persisted.Targets).Should(HaveLen(2)) + Expect([]string{persisted.Targets[0].Component, persisted.Targets[1].Component}).Should( + ConsistOf(retained.Name, "shard-1")) + }) + + It("polls every target and retries only terminal failures", func() { + shard, pods := buildShard("shard-0", "pod-0", "pod-1") + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{ + Objects: []client.Object{shard, pods[0], pods[1]}, + }) + Expect(setShardingActionTargets(shard, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: shard.Name, + Pods: []shardingActionTargetPod{ + {Name: pods[0].Name, Query: true}, + {Name: pods[1].Name, Query: true}, + }, + }}, + })).Should(Succeed()) + call := 0 + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).DoAndReturn( + func(context.Context, kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + call++ + if call == 1 { + return kbagentproto.ActionResponse{ + Error: kbagentproto.Error2Type(kbagentproto.ErrFailed), + }, nil + } + return kbagentproto.ActionResponse{ + Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress), + }, nil + }).Times(2) + }) + + err := (&clusterShardingHandler{}).nonBlockingShardingAction( + transCtx, sharding1aName, shardingAddShardAction, shardingAddActionTargetsKey, + action(), nil, []*appsv1.Component{shard}, shard) + Expect(errors.Is(err, lifecycle.ErrActionFailed)).Should(BeTrue()) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeFalse()) + targets, _, err := getShardingActionTargets(shard, shardingAddActionTargetsKey) + Expect(err).Should(BeNil()) + Expect(targets.Targets[0].Pods[0].Query).Should(BeFalse()) + Expect(targets.Targets[0].Pods[1].Query).Should(BeTrue()) + }) + + It("keeps shard removal blocked while the action is in progress", func() { + shard, pods := buildShard("shard-0", "shard-0-0") + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{ + Objects: []client.Object{shard, pods[0]}, + }) + Expect(setShardingActionTargets(shard, shardingRemoveActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: shard.Name, + Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: true}}, + }}, + })).Should(Succeed()) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardRemove: action(), + } + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).Return(kbagentproto.ActionResponse{ + Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress), + }, nil).Times(1) + }) + + errorSkip, err := reconcileActions( + map[string]*appsv1.Component{shard.Name: shard}, + map[string]*appsv1.Component{}, + sets.New[string](), sets.New(shard.Name), sets.New[string]()) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(errorSkip.Has(shard.Name)).Should(BeTrue()) + }) + + It("polls a persisted source before starting another shard add", func() { + persisted, persistedPods := buildShard("shard-z", "shard-z-0") + fresh, freshPods := buildShard("shard-a", "shard-a-0") + persisted.Annotations[shardingAddShardKey] = "pending" + fresh.Annotations[shardingAddShardKey] = "pending" + Expect(setShardingActionTargets(persisted, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: persisted.Name, + Pods: []shardingActionTargetPod{{Name: persistedPods[0].Name, Query: true}}, + }}, + })).Should(Succeed()) + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{Objects: []client.Object{ + persisted, fresh, persistedPods[0], freshPods[0], + }}) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardAdd: action(), + } + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).Return(kbagentproto.ActionResponse{ + Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress), + }, nil).Times(1) + }) + + _, err := reconcileActions( + map[string]*appsv1.Component{persisted.Name: persisted, fresh.Name: fresh}, + map[string]*appsv1.Component{}, sets.New[string](), sets.New[string](), + sets.New(persisted.Name, fresh.Name)) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(fresh.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + }) + + It("finishes a pending shard add before starting shard remove", func() { + shard, pods := buildShard("shard-0", "shard-0-0") + shard.Annotations[shardingAddShardKey] = "pending" + Expect(setShardingActionTargets(shard, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: shard.Name, + Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: true}}, + }}, + })).Should(Succeed()) + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{ + Objects: []client.Object{shard, pods[0]}, + }) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardAdd: action(), + ShardRemove: action(), + } + calls := 0 + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls++ + Expect(req.Action).Should(Equal("udf-" + shardingAddShardAction)) + if calls == 1 { + return kbagentproto.ActionResponse{ + Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress), + }, nil + } + return kbagentproto.ActionResponse{}, nil + }).Times(2) + }) + + handle := func() (sets.Set[string], error) { + return reconcileActions( + map[string]*appsv1.Component{shard.Name: shard}, map[string]*appsv1.Component{}, + sets.New[string](), sets.New(shard.Name), sets.New[string]()) + } + errorSkip, err := handle() + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(errorSkip.Has(shard.Name)).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingRemoveActionTargetsKey)) + + errorSkip, err = handle() + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(errorSkip.Has(shard.Name)).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingRemoveActionTargetsKey)) + + errorSkip, err = handle() + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(errorSkip.Has(shard.Name)).Should(BeTrue()) + Expect(shard.Annotations).Should(HaveKey(shardingRemoveActionTargetsKey)) + }) + + It("resumes shard add before delete when shard remove is not defined", func() { + shard, pods := buildShard("shard-0", "shard-0-0") + shard.Annotations[shardingAddShardKey] = "pending" + Expect(setShardingActionTargets(shard, shardingAddActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: shard.Name, + Pods: []shardingActionTargetPod{{Name: pods[0].Name}}, + }}, + })).Should(Succeed()) + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{ + Objects: []client.Object{shard, pods[0]}, + }) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardAdd: action(), + } + calls := 0 + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls++ + Expect(req.Action).Should(Equal("udf-" + shardingAddShardAction)) + if calls == 1 { + return kbagentproto.ActionResponse{ + Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress), + }, nil + } + return kbagentproto.ActionResponse{}, nil + }).Times(2) + }) + + handle := func() (sets.Set[string], error) { + return reconcileActions( + map[string]*appsv1.Component{shard.Name: shard}, map[string]*appsv1.Component{}, + sets.New[string](), sets.New(shard.Name), sets.New[string]()) + } + + errorSkip, err := handle() + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(errorSkip.Has(shard.Name)).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingRemoveActionTargetsKey)) + targets, found, err := getShardingActionTargets(shard, shardingAddActionTargetsKey) + Expect(err).Should(BeNil()) + Expect(found).Should(BeTrue()) + Expect(targets.Targets[0].Pods[0].Query).Should(BeTrue()) + + errorSkip, err = handle() + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(errorSkip.Has(shard.Name)).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + errorSkip, err = handle() + Expect(err).ShouldNot(HaveOccurred()) + Expect(errorSkip.Has(shard.Name)).Should(BeFalse()) + Expect(calls).Should(Equal(2)) + }) + + It("finishes reversed shard remove before persisting and dispatching shard add", func() { + shard, pods := buildShard("shard-0", "shard-0-0") + Expect(setShardingActionTargets(shard, shardingRemoveActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: shard.Name, + Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: true}}, + }}, + })).Should(Succeed()) + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{ + Objects: []client.Object{shard, pods[0]}, + }) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardAdd: action(), + ShardRemove: action(), + } + calls := 0 + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + calls++ + Expect(req.Action).Should(Equal("udf-" + shardingRemoveShardAction)) + if calls == 1 { + return kbagentproto.ActionResponse{ + Error: kbagentproto.Error2Type(kbagentproto.ErrInProgress), + }, nil + } + return kbagentproto.ActionResponse{}, nil + }).Times(2) + }) + handle := func() error { + _, err := reconcileActions( + map[string]*appsv1.Component{shard.Name: shard}, + map[string]*appsv1.Component{shard.Name: shard.DeepCopy()}, + sets.New[string](), sets.New[string](), sets.New(shard.Name)) + return err + } + + Expect(ictrlutil.IsDelayedRequeueError(handle())).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + Expect(shard.Annotations).Should(HaveKey(shardingRemoveActionTargetsKey)) + + Expect(ictrlutil.IsDelayedRequeueError(handle())).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingRemoveActionTargetsKey)) + Expect(shard.Annotations).Should(HaveKey(shardingAddShardKey)) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddActionTargetsKey)) + + // The next reconciliation persists fresh add targets without invoking + // shardAdd; its invocation can only happen on a later poll. + Expect(ictrlutil.IsDelayedRequeueError(handle())).Should(BeTrue()) + Expect(shard.Annotations).Should(HaveKey(shardingAddActionTargetsKey)) + Expect(calls).Should(Equal(2)) + }) + + It("resumes shard remove after retain intent when shard add is not defined", func() { + shard, pods := buildShard("shard-0", "shard-0-0") + Expect(setShardingActionTargets(shard, shardingRemoveActionTargetsKey, &shardingActionTargets{ + Version: shardingActionTargetsVersion, + Targets: []shardingActionTarget{{ + Component: shard.Name, + Pods: []shardingActionTargetPod{{Name: pods[0].Name, Query: true}}, + }}, + })).Should(Succeed()) + transCtx.Client = model.NewGraphClient(&appsutil.MockReader{ + Objects: []client.Object{shard, pods[0]}, + }) + transCtx.shardingDefs[shardingDefName].Spec.LifecycleActions = &appsv1.ShardingLifecycleActions{ + ShardRemove: action(), + } + testapps.MockKBAgentClient(func(recorder *kbacli.MockClientMockRecorder) { + recorder.Action(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, req kbagentproto.ActionRequest) (kbagentproto.ActionResponse, error) { + Expect(req.Action).Should(Equal("udf-" + shardingRemoveShardAction)) + return kbagentproto.ActionResponse{}, nil + }).Times(1) + }) + + _, err := reconcileActions( + map[string]*appsv1.Component{shard.Name: shard}, + map[string]*appsv1.Component{shard.Name: shard.DeepCopy()}, + sets.New[string](), sets.New[string](), sets.New(shard.Name)) + Expect(ictrlutil.IsDelayedRequeueError(err)).Should(BeTrue()) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingRemoveActionTargetsKey)) + Expect(shard.Annotations).ShouldNot(HaveKey(shardingAddShardKey)) + _, err = reconcileActions( + map[string]*appsv1.Component{shard.Name: shard}, + map[string]*appsv1.Component{shard.Name: shard.DeepCopy()}, + sets.New[string](), sets.New[string](), sets.New(shard.Name)) + Expect(err).ShouldNot(HaveOccurred()) + }) + + It("rejects malformed persisted targets", func() { + shard, _ := buildShard("shard-0") + shard.Annotations[shardingAddActionTargetsKey] = "{" + _, _, err := getShardingActionTargets(shard, shardingAddActionTargetsKey) + Expect(err).Should(MatchError(ContainSubstring("invalid " + shardingAddActionTargetsKey))) + }) + }) }) }) diff --git a/deploy/helm/crds/apps.kubeblocks.io_shardingdefinitions.yaml b/deploy/helm/crds/apps.kubeblocks.io_shardingdefinitions.yaml index e2e02d7a8cd..8ef858a5526 100644 --- a/deploy/helm/crds/apps.kubeblocks.io_shardingdefinitions.yaml +++ b/deploy/helm/crds/apps.kubeblocks.io_shardingdefinitions.yaml @@ -1002,6 +1002,10 @@ spec: - KB_ADD_SHARD_NAME: The name of the shard being added. + This Action supports both blocking and non-blocking modes. KubeBlocks does + not finish adding the shard until the Action has succeeded on every + selected target shard and Pod. + Note: This field is immutable once it has been set. properties: exec: @@ -1463,6 +1467,10 @@ spec: - KB_REMOVE_SHARD_NAME: The name of the shard being removed. + This Action supports both blocking and non-blocking modes. KubeBlocks does + not remove the shard until the Action has succeeded on every selected + target shard and Pod. + Note: This field is immutable once it has been set. properties: exec: diff --git a/docs/developer_docs/api-reference/cluster.md b/docs/developer_docs/api-reference/cluster.md index 20840593fd2..d79bc948ff2 100644 --- a/docs/developer_docs/api-reference/cluster.md +++ b/docs/developer_docs/api-reference/cluster.md @@ -13040,6 +13040,9 @@ ShardingAction +

This Action supports both blocking and non-blocking modes. KubeBlocks does +not finish adding the shard until the Action has succeeded on every +selected target shard and Pod.

Note: This field is immutable once it has been set.

@@ -13059,6 +13062,9 @@ ShardingAction +

This Action supports both blocking and non-blocking modes. KubeBlocks does +not remove the shard until the Action has succeeded on every selected +target shard and Pod.

Note: This field is immutable once it has been set.

diff --git a/pkg/controller/component/component.go b/pkg/controller/component/component.go index c74a500063a..74536f0bc34 100644 --- a/pkg/controller/component/component.go +++ b/pkg/controller/component/component.go @@ -168,14 +168,21 @@ func GetExporter(componentDef appsv1.ComponentDefinitionSpec) *appsv1.Exporter { return nil } -func NewLifecycle(ctx context.Context, cli client.Reader, compDef *appsv1.ComponentDefinition, comp *appsv1.Component) (lifecycle.Lifecycle, error) { +// NewLifecycle uses the supplied template variables, when present, instead of +// resolving them again. This keeps a continuing Action's request inputs stable. +func NewLifecycle(ctx context.Context, cli client.Reader, compDef *appsv1.ComponentDefinition, comp *appsv1.Component, + templateVars ...map[string]string) (lifecycle.Lifecycle, error) { synthesizedComp, err := BuildSynthesizedComponent(ctx, cli, compDef, comp) if err != nil { return nil, err } - synthesizedComp.TemplateVars, _, err = ResolveTemplateNEnvVars(ctx, cli, synthesizedComp, compDef.Spec.Vars) - if err != nil { - return nil, err + if len(templateVars) > 0 { + synthesizedComp.TemplateVars = templateVars[0] + } else { + synthesizedComp.TemplateVars, _, err = ResolveTemplateNEnvVars(ctx, cli, synthesizedComp, compDef.Spec.Vars) + if err != nil { + return nil, err + } } pods, err := ListOwnedInstances(ctx, cli, comp) diff --git a/pkg/controller/lifecycle/errors.go b/pkg/controller/lifecycle/errors.go index 2b4d08c3c4d..3ec46135613 100644 --- a/pkg/controller/lifecycle/errors.go +++ b/pkg/controller/lifecycle/errors.go @@ -29,6 +29,7 @@ var ( ErrPreconditionFailed = errors.New("action precondition is not met") ErrActionInProgress = errors.New("action is in progress") ErrActionBusy = errors.New("action is busy") + ErrActionResultNotFound = errors.New("action result not found") ErrActionTimedOut = errors.New("action timed-out") ErrActionFailed = errors.New("action failed") ErrActionInternalError = errors.New("action internal error") @@ -61,7 +62,8 @@ func IsActionFailure(err error) bool { return !errors.Is(err, ErrActionNotDefined) && !errors.Is(err, ErrPreconditionFailed) && !errors.Is(err, ErrActionInProgress) && - !errors.Is(err, ErrActionBusy) + !errors.Is(err, ErrActionBusy) && + !errors.Is(err, ErrActionResultNotFound) } type actionAggregateError struct { diff --git a/pkg/controller/lifecycle/errors_test.go b/pkg/controller/lifecycle/errors_test.go index 729c5fb13f5..be77203cff9 100644 --- a/pkg/controller/lifecycle/errors_test.go +++ b/pkg/controller/lifecycle/errors_test.go @@ -36,6 +36,7 @@ func TestIsActionFailure(t *testing.T) { {name: "precondition", err: ErrPreconditionFailed, want: false}, {name: "in progress", err: ErrActionInProgress, want: false}, {name: "busy", err: ErrActionBusy, want: false}, + {name: "result not found", err: ErrActionResultNotFound, want: false}, {name: "wrapped waiting", err: fmt.Errorf("wrapped: %w", ErrActionBusy), want: false}, {name: "failed", err: ErrActionFailed, want: true}, {name: "timed out", err: ErrActionTimedOut, want: true}, diff --git a/pkg/controller/lifecycle/kbagent.go b/pkg/controller/lifecycle/kbagent.go index 432b47abce4..1ef91f4aaca 100644 --- a/pkg/controller/lifecycle/kbagent.go +++ b/pkg/controller/lifecycle/kbagent.go @@ -149,6 +149,11 @@ func (a *kbagent) checkedCallAction(ctx context.Context, cli client.Reader, spec if !spec.Defined() { return nil, errors.Wrap(ErrActionNotDefined, lfa.name()) } + if opts != nil && opts.Query { + // Queries cannot start a request, including after the agent loses its + // cached result, so startup preconditions must not gate observation. + return a.callAction(ctx, cli, spec, lfa, opts) + } if err := a.precondition(ctx, cli, spec, func() client.MatchingLabels { if opts == nil || opts.PreConditionObjectSelector == nil { return nil @@ -300,7 +305,7 @@ func (a *kbagent) buildActionRequest(ctx context.Context, cli client.Reader, lfa Parameters: parameters, } if opts != nil { - req.Rerun = opts.Rerun + req.Query = opts.Query if opts.TimeoutSeconds != nil { req.TimeoutSeconds = opts.TimeoutSeconds } @@ -498,6 +503,8 @@ func (a *kbagent) formatError(lfa lifecycleAction, rsp proto.ActionResponse, pod return wrapError(ErrActionInProgress) case errors.Is(err, proto.ErrBusy): return wrapError(ErrActionBusy) + case errors.Is(err, proto.ErrResultNotFound): + return wrapError(ErrActionResultNotFound) case errors.Is(err, proto.ErrTimedOut): return wrapError(ErrActionTimedOut) case errors.Is(err, proto.ErrFailed): diff --git a/pkg/controller/lifecycle/lifecycle.go b/pkg/controller/lifecycle/lifecycle.go index bd5743d97a7..4fd91613bbd 100644 --- a/pkg/controller/lifecycle/lifecycle.go +++ b/pkg/controller/lifecycle/lifecycle.go @@ -30,7 +30,8 @@ import ( ) type Options struct { - Rerun bool + // Query observes an existing non-blocking request without executing it. + Query bool // TargetPodName, when set, executes the Action only on the named Pod and // overrides the Action's targetPodSelector. The Pod must be present in the diff --git a/pkg/controller/lifecycle/lifecycle_test.go b/pkg/controller/lifecycle/lifecycle_test.go index 74a8324fe91..fe485034cd9 100644 --- a/pkg/controller/lifecycle/lifecycle_test.go +++ b/pkg/controller/lifecycle/lifecycle_test.go @@ -251,7 +251,7 @@ var _ = Describe("lifecycle", func() { recorder.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, req proto.ActionRequest) (proto.ActionResponse, error) { Expect(req.Action).Should(Equal("postProvision")) Expect(req.Parameters).Should(BeEmpty()) - Expect(req.Rerun).Should(BeTrue()) + Expect(req.Query).Should(BeFalse()) Expect(req.TimeoutSeconds).ShouldNot(BeNil()) Expect(*req.TimeoutSeconds).Should(Equal(action.TimeoutSeconds)) Expect(req.RetryPolicy).ShouldNot(BeNil()) @@ -262,7 +262,6 @@ var _ = Describe("lifecycle", func() { }) opts := &Options{ - Rerun: true, TimeoutSeconds: &action.TimeoutSeconds, RetryPolicy: action.RetryPolicy, } @@ -585,6 +584,24 @@ var _ = Describe("lifecycle", func() { Expect(err).Should(BeNil()) }) + It("observes a missing result without bypassing preconditions for a new run", func() { + lifecycleActions.PostProvision.PreCondition = ptr.To(appsv1.ClusterReadyPreConditionType) + lfa, err := New(namespace, clusterName, compName, lifecycleActions, nil, nil, pods) + Expect(err).ShouldNot(HaveOccurred()) + reader := &mockReader{cli: k8sClient, objs: []client.Object{&appsv1.Cluster{ + ObjectMeta: metav1.ObjectMeta{Name: clusterName, Namespace: namespace}, + Status: appsv1.ClusterStatus{Phase: appsv1.FailedClusterPhase}, + }}} + mockKBAgentClient(func(r *kbacli.MockClientMockRecorder) { + r.Action(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, req proto.ActionRequest) (proto.ActionResponse, error) { + Expect(req.Query).Should(BeTrue()) + return proto.ActionResponse{Error: proto.Error2Type(proto.ErrResultNotFound)}, nil + }).Times(1) + }) + Expect(errors.Is(lfa.PostProvision(ctx, reader, &Options{Query: true}), ErrActionResultNotFound)).Should(BeTrue()) + Expect(errors.Is(lfa.PostProvision(ctx, reader, nil), ErrPreconditionFailed)).Should(BeTrue()) + }) + It("precondition - fail", func() { clusterReady := appsv1.ClusterReadyPreConditionType lifecycleActions.PostProvision.PreCondition = &clusterReady diff --git a/pkg/kbagent/client/http_client_test.go b/pkg/kbagent/client/http_client_test.go index 0633652710f..e09f1047ad9 100644 --- a/pkg/kbagent/client/http_client_test.go +++ b/pkg/kbagent/client/http_client_test.go @@ -61,6 +61,10 @@ func newHTTPClientForTest(t *testing.T, handler http.HandlerFunc) (*httpClient, } func TestHTTPClientAction(t *testing.T) { + queries := make(chan bool, 2) + queries <- false + queries <- true + close(queries) cli, closeServer := newHTTPClientForTest(t, func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != proto.ServiceAction.URI || r.Method != http.MethodPost { t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path) @@ -69,7 +73,8 @@ func TestHTTPClientAction(t *testing.T) { if err := json.NewDecoder(r.Body).Decode(request); err != nil { t.Fatalf("decode Action request: %v", err) } - if request.Action != "backup" || !request.Rerun { + query, ok := <-queries + if !ok || request.Action != "backup" || request.Query != query { t.Fatalf("unexpected Action request: %#v", request) } w.WriteHeader(http.StatusOK) @@ -77,13 +82,16 @@ func TestHTTPClientAction(t *testing.T) { }) defer closeServer() - resp, err := cli.Action(context.Background(), proto.ActionRequest{Action: "backup", Rerun: true}) + resp, err := cli.Action(context.Background(), proto.ActionRequest{Action: "backup"}) if err != nil { t.Fatalf("Action() error = %v", err) } if resp.Message != "done" || string(resp.Output) != "ok" { t.Fatalf("unexpected response: %#v", resp) } + if _, err = cli.Action(context.Background(), proto.ActionRequest{Action: "backup", Query: true}); err != nil { + t.Fatalf("query Action() error = %v", err) + } resp, err = cli.Action(context.WithValue(context.Background(), constant.DryRunContextKey, true), proto.ActionRequest{Action: "backup"}) if err != nil || resp.Message != "" || resp.Error != "" || len(resp.Output) != 0 { diff --git a/pkg/kbagent/proto/errors.go b/pkg/kbagent/proto/errors.go index d6e1e664ff7..0ba1a8ae489 100644 --- a/pkg/kbagent/proto/errors.go +++ b/pkg/kbagent/proto/errors.go @@ -30,6 +30,7 @@ var ( ErrBadRequest = errors.New("badRequest") ErrInProgress = errors.New("inProgress") ErrBusy = errors.New("busy") + ErrResultNotFound = errors.New("resultNotFound") ErrTimedOut = errors.New("timedOut") ErrFailed = errors.New("failed") ErrInternalError = errors.New("internalError") @@ -52,6 +53,8 @@ func Error2Type(err error) string { return "inProgress" case errors.Is(err, ErrBusy): return "busy" + case errors.Is(err, ErrResultNotFound): + return "resultNotFound" case errors.Is(err, ErrTimedOut): return "timedOut" case errors.Is(err, ErrFailed): @@ -79,6 +82,8 @@ func Type2Error(errType string) error { return ErrInProgress case "busy": return ErrBusy + case "resultNotFound": + return ErrResultNotFound case "timedOut": return ErrTimedOut case "failed": diff --git a/pkg/kbagent/proto/errors_test.go b/pkg/kbagent/proto/errors_test.go index 0df61064515..183e0577638 100644 --- a/pkg/kbagent/proto/errors_test.go +++ b/pkg/kbagent/proto/errors_test.go @@ -37,6 +37,7 @@ func TestError2Type(t *testing.T) { {name: "bad request", err: ErrBadRequest, want: "badRequest"}, {name: "in progress", err: ErrInProgress, want: "inProgress"}, {name: "busy", err: ErrBusy, want: "busy"}, + {name: "result not found", err: ErrResultNotFound, want: "resultNotFound"}, {name: "timed out", err: ErrTimedOut, want: "timedOut"}, {name: "failed", err: ErrFailed, want: "failed"}, {name: "internal error", err: ErrInternalError, want: "internalError"}, @@ -64,6 +65,7 @@ func TestType2Error(t *testing.T) { {errType: "badRequest", want: ErrBadRequest}, {errType: "inProgress", want: ErrInProgress}, {errType: "busy", want: ErrBusy}, + {errType: "resultNotFound", want: ErrResultNotFound}, {errType: "timedOut", want: ErrTimedOut}, {errType: "failed", want: ErrFailed}, {errType: "internalError", want: ErrInternalError}, diff --git a/pkg/kbagent/proto/proto.go b/pkg/kbagent/proto/proto.go index 096eb651f40..1ebd28e53cb 100644 --- a/pkg/kbagent/proto/proto.go +++ b/pkg/kbagent/proto/proto.go @@ -77,9 +77,10 @@ type ActionRequest struct { Arguments [][]string `json:"arguments,omitempty"` TimeoutSeconds *int32 `json:"timeoutSeconds,omitempty"` RetryPolicy *RetryPolicy `json:"retryPolicy,omitempty"` - // Rerun requests a new run instead of returning the previous terminal result. - // It does not interrupt a running Action. - Rerun bool `json:"rerun,omitempty"` + // Query returns the state or result of an equivalent non-blocking request + // without executing it. A missing result returns ErrResultNotFound. + // When false, the request executes again unless an equivalent request is running. + Query bool `json:"query,omitempty"` } type ActionResponse struct { diff --git a/pkg/kbagent/service/action.go b/pkg/kbagent/service/action.go index 5b818606790..981f42ec832 100644 --- a/pkg/kbagent/service/action.go +++ b/pkg/kbagent/service/action.go @@ -120,9 +120,14 @@ func (s *actionService) handleRequest(ctx context.Context, req *proto.ActionRequ if action.Exec == nil && action.HTTP == nil && action.GRPC == nil { return nil, errors.Wrapf(proto.ErrBadRequest, "%s is invalid", req.Action) } + if req.Query && !action.NonBlocking { + return nil, errors.Wrap(proto.ErrBadRequest, "query requires a non-blocking action") + } // HACK: pre-check for the reconfigure action - if err := checkReconfigure(ctx, req); err != nil { - return nil, err + if !req.Query { + if err := checkReconfigure(ctx, req); err != nil { + return nil, err + } } timeout := resolveTimeout(&action.TimeoutSeconds, req.TimeoutSeconds) retryPolicy := resolveRetryPolicy(action.RetryPolicy, req.RetryPolicy) @@ -144,16 +149,22 @@ func (s *actionService) handleRequestNonBlocking(ctx context.Context, req *proto s.mutex.Lock() defer s.mutex.Unlock() - if call, ok := s.calls[req.Action]; ok { + if req.Query { + call, ok := s.calls[req.Action] + if !ok || call.requestFingerprint != fingerprint { + return nil, proto.ErrResultNotFound + } if call.running { - if call.requestFingerprint != fingerprint || req.Rerun { - return nil, proto.ErrBusy - } return nil, proto.ErrInProgress } - if call.requestFingerprint == fingerprint && !req.Rerun { - return call.result.response() + return call.result.response() + } + + if call, ok := s.calls[req.Action]; ok && call.running { + if call.requestFingerprint != fingerprint { + return nil, proto.ErrBusy } + return nil, proto.ErrInProgress } call := &actionCall{ diff --git a/pkg/kbagent/service/action_test.go b/pkg/kbagent/service/action_test.go index c8dc8ec0602..201246b7575 100644 --- a/pkg/kbagent/service/action_test.go +++ b/pkg/kbagent/service/action_test.go @@ -109,6 +109,7 @@ var _ = Describe("action", func() { Expect(out).Should(BeNil()) Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) + req.Query = true Eventually(func() error { out, err = svc.handleRequest(ctx, req) return err @@ -136,7 +137,7 @@ var _ = Describe("action", func() { cancel() Eventually(func() string { - output, callErr := svc.handleRequest(ctx, &proto.ActionRequest{Action: "async"}) + output, callErr := svc.handleRequest(ctx, &proto.ActionRequest{Action: "async", Query: true}) if callErr != nil { return callErr.Error() } @@ -253,6 +254,7 @@ var _ = Describe("action", func() { } second := &proto.ActionRequest{ Action: "shardAdd", + Query: true, Parameters: map[string]string{"first": "1", "second": "2"}, Arguments: [][]string{}, } @@ -277,8 +279,8 @@ var _ = Describe("action", func() { Expect(differentFingerprint).ShouldNot(Equal(firstFingerprint)) }) - It("serializes a single running request and honors rerun after completion", func() { - dir, err := os.MkdirTemp("", "kbagent-action-rerun-*") + DescribeTable("serializes a running request and executes again after completion", func(firstResult string) { + dir, err := os.MkdirTemp("", "kbagent-action-repeat-*") Expect(err).ShouldNot(HaveOccurred()) DeferCleanup(os.RemoveAll, dir) counterPath := filepath.Join(dir, "counter") @@ -287,8 +289,8 @@ var _ = Describe("action", func() { NonBlocking: true, Exec: &proto.ExecAction{Commands: []string{ "/bin/bash", "-c", - `n=0; [ -f "$0" ] && n=$(cat "$0"); n=$((n+1)); echo "$n" > "$0"; sleep 0.1; printf "$n"`, - counterPath, + `n=0; [ -f "$0" ] && n=$(cat "$0"); n=$((n+1)); echo "$n" > "$0"; sleep 0.1; printf "$n"; if [ "$n" -eq 1 ] && [ "$1" = failed ]; then exit 1; fi`, + counterPath, firstResult, }}, } svc, err := newActionService(logr.Discard(), []proto.Action{action}) @@ -299,29 +301,34 @@ var _ = Describe("action", func() { Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) _, err = svc.handleRequest(ctx, req) Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) - _, err = svc.handleRequest(ctx, &proto.ActionRequest{Action: "async", Rerun: true}) - Expect(errors.Is(err, proto.ErrBusy)).Should(BeTrue()) + query := &proto.ActionRequest{Action: "async", Query: true} + _, err = svc.handleRequest(ctx, query) + Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) _, err = svc.handleRequest(ctx, &proto.ActionRequest{ Action: "async", Parameters: map[string]string{"different": "request"}, }) Expect(errors.Is(err, proto.ErrBusy)).Should(BeTrue()) Eventually(func() string { - output, callErr := svc.handleRequest(ctx, req) + output, callErr := svc.handleRequest(ctx, query) if callErr != nil { - return callErr.Error() + return proto.Error2Type(callErr) } return string(output) - }, 2*time.Second, 10*time.Millisecond).Should(Equal("1")) - - output, err := svc.handleRequest(ctx, req) - Expect(err).ShouldNot(HaveOccurred()) - Expect(string(output)).Should(Equal("1")) + }, 2*time.Second, 10*time.Millisecond).Should(Equal(firstResult)) + + output, err := svc.handleRequest(ctx, query) + if firstResult == "failed" { + Expect(errors.Is(err, proto.ErrFailed)).Should(BeTrue()) + } else { + Expect(err).ShouldNot(HaveOccurred()) + Expect(string(output)).Should(Equal(firstResult)) + } - _, err = svc.handleRequest(ctx, &proto.ActionRequest{Action: "async", Rerun: true}) + _, err = svc.handleRequest(ctx, req) Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) Eventually(func() string { - output, callErr := svc.handleRequest(ctx, req) + output, callErr := svc.handleRequest(ctx, query) if callErr != nil { return callErr.Error() } @@ -334,6 +341,7 @@ var _ = Describe("action", func() { } _, err = svc.handleRequest(ctx, differentReq) Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) + differentReq.Query = true Eventually(func() string { output, callErr := svc.handleRequest(ctx, differentReq) if callErr != nil { @@ -341,8 +349,88 @@ var _ = Describe("action", func() { } return string(output) }, 2*time.Second, 10*time.Millisecond).Should(Equal("3")) + }, Entry("success", "1"), Entry("failure", "failed")) + + It("observes cached requests atomically without starting commands on a miss", func() { + counter := filepath.Join(GinkgoT().TempDir(), "calls") + action := proto.Action{Name: "query", NonBlocking: true, Exec: &proto.ExecAction{ + Commands: []string{"/bin/sh", "-c", `echo run >> "$1"; printf done`, "sh", counter}, + }} + svc, err := newActionService(logr.Discard(), []proto.Action{action}) + Expect(err).ShouldNot(HaveOccurred()) + query := &proto.ActionRequest{Action: action.Name, Query: true} + _, err = svc.handleRequest(ctx, query) + Expect(errors.Is(err, proto.ErrResultNotFound)).Should(BeTrue()) + Expect(svc.calls).Should(BeEmpty()) + svc.actions[action.Name].NonBlocking = false + _, err = svc.handleRequest(ctx, query) + Expect(errors.Is(err, proto.ErrBadRequest)).Should(BeTrue()) + svc.actions[action.Name].NonBlocking = true + _, err = os.Stat(counter) + Expect(os.IsNotExist(err)).Should(BeTrue()) + + // Queries racing the first start may see a miss, progress, or the + // completed result, but must never create another invocation. + var wg sync.WaitGroup + for i := range 16 { + wg.Add(1) + go func(start bool) { + defer GinkgoRecover() + defer wg.Done() + _, callErr := svc.handleRequest(ctx, &proto.ActionRequest{Action: action.Name, Query: !start}) + Expect(callErr == nil || errors.Is(callErr, proto.ErrInProgress) || errors.Is(callErr, proto.ErrResultNotFound)).Should(BeTrue()) + }(i == 0) + } + wg.Wait() + Eventually(func() string { + out, callErr := svc.handleRequest(ctx, query) + if callErr != nil { + return callErr.Error() + } + return string(out) + }).Should(Equal("done")) + _, err = svc.handleRequest(ctx, &proto.ActionRequest{Action: action.Name, Query: true, Parameters: map[string]string{"other": "request"}}) + Expect(errors.Is(err, proto.ErrResultNotFound)).Should(BeTrue()) + out, err := svc.handleRequest(ctx, query) + Expect(err).ShouldNot(HaveOccurred()) + Expect(string(out)).Should(Equal("done")) + svc, err = newActionService(logr.Discard(), []proto.Action{action}) + Expect(err).ShouldNot(HaveOccurred()) + _, err = svc.handleRequest(ctx, query) + Expect(errors.Is(err, proto.ErrResultNotFound)).Should(BeTrue()) + data, err := os.ReadFile(counter) + Expect(err).ShouldNot(HaveOccurred()) + Expect(string(data)).Should(Equal("run\n")) }) + DescribeTable("queries never enter an action backend", func(action proto.Action) { + action.Name, action.NonBlocking = "query", true + svc, err := newActionService(logr.Discard(), []proto.Action{action}) + Expect(err).ShouldNot(HaveOccurred()) + req := &proto.ActionRequest{Action: action.Name, Query: true} + _, err = svc.handleRequest(ctx, req) + Expect(errors.Is(err, proto.ErrResultNotFound)).Should(BeTrue()) + hash, err := fingerprintActionRequest(req, &action.TimeoutSeconds, nil) + Expect(err).ShouldNot(HaveOccurred()) + call := &actionCall{requestFingerprint: hash, running: true} + svc.calls[action.Name] = call + _, err = svc.handleRequest(ctx, req) + Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) + different := *req + different.Parameters = map[string]string{"different": "request"} + _, err = svc.handleRequest(ctx, &different) + Expect(errors.Is(err, proto.ErrResultNotFound)).Should(BeTrue()) + call.running = false + call.result = newActionResult(nil, proto.ErrFailed) + _, err = svc.handleRequest(ctx, req) + Expect(errors.Is(err, proto.ErrFailed)).Should(BeTrue()) + Expect(svc.calls[action.Name]).Should(BeIdenticalTo(call)) + }, + Entry("Exec", proto.Action{Exec: &proto.ExecAction{Commands: []string{"does-not-exist"}}}), + Entry("HTTP", proto.Action{HTTP: &proto.HTTPAction{Port: "invalid"}}), + Entry("gRPC", proto.Action{GRPC: &proto.GRPCAction{Port: "invalid"}}), + ) + It("starts only one process for concurrent equivalent requests", func() { dir, err := os.MkdirTemp("", "kbagent-action-concurrent-*") Expect(err).ShouldNot(HaveOccurred()) @@ -375,7 +463,7 @@ var _ = Describe("action", func() { } Eventually(func() string { - output, callErr := svc.handleRequest(ctx, &proto.ActionRequest{Action: "async"}) + output, callErr := svc.handleRequest(ctx, &proto.ActionRequest{Action: "async", Query: true}) if callErr != nil { return callErr.Error() } @@ -459,6 +547,9 @@ var _ = Describe("action", func() { svc.actions["retry"].NonBlocking = true req := &proto.ActionRequest{Action: "retry"} + _, err = svc.handleRequest(ctx, req) + Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) + req.Query = true Eventually(func() string { output, err := svc.handleRequest(ctx, req) if err != nil { diff --git a/pkg/kbagent/service/action_utils_test.go b/pkg/kbagent/service/action_utils_test.go index e0c55bdb0a2..63cf60855d6 100644 --- a/pkg/kbagent/service/action_utils_test.go +++ b/pkg/kbagent/service/action_utils_test.go @@ -466,6 +466,7 @@ var _ = Describe("action utils", func() { req := &proto.ActionRequest{Action: "http"} _, err = svc.handleRequest(ctx, req) Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) + req.Query = true Eventually(func() string { output, callErr := svc.handleRequest(ctx, req) if callErr != nil { @@ -873,6 +874,7 @@ message EchoResponse { req := &proto.ActionRequest{Action: "grpc"} _, err = svc.handleRequest(ctx, req) Expect(errors.Is(err, proto.ErrInProgress)).Should(BeTrue()) + req.Query = true Eventually(func() string { output, callErr := svc.handleRequest(ctx, req) if callErr != nil {