Skip to content

Commit b1729fd

Browse files
viliakovviliakov
committed
STAC-24228: Address comment
1 parent 237e139 commit b1729fd

1 file changed

Lines changed: 32 additions & 70 deletions

File tree

‎internal/clients/k8s/client.go‎

Lines changed: 32 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -175,11 +175,11 @@ type RestoreLockInfo struct {
175175

176176
// DeploymentUpdateFunc is a function that modifies a deployment.
177177
// It receives a fresh copy of the deployment and should apply the desired changes.
178-
type DeploymentUpdateFunc func(dep *appsv1.Deployment)
178+
type DeploymentUpdateFunc func(dep *appsv1.Deployment) error
179179

180180
// StatefulSetUpdateFunc is a function that modifies a statefulset.
181181
// It receives a fresh copy of the statefulset and should apply the desired changes.
182-
type StatefulSetUpdateFunc func(sts *appsv1.StatefulSet)
182+
type StatefulSetUpdateFunc func(sts *appsv1.StatefulSet) error
183183

184184
// updateDeploymentWithRetry fetches a fresh copy of the deployment and applies the update function,
185185
// retrying on conflict errors (when resource version has changed).
@@ -222,28 +222,16 @@ func updateStatefulSetWithRetry(ctx context.Context, client kubernetes.Interface
222222
func scaleDownDeployment(ctx context.Context, client kubernetes.Interface, namespace, name string) (int32, error) {
223223
var originalReplicas int32
224224

225-
err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
226-
dep, err := client.AppsV1().Deployments(namespace).Get(ctx, name, metav1.GetOptions{})
227-
if err != nil {
228-
return err
229-
}
230-
225+
err := updateDeploymentWithRetry(ctx, client, namespace, name, func(dep *appsv1.Deployment) error {
231226
if dep.Spec.Replicas != nil {
232227
originalReplicas = *dep.Spec.Replicas
233228
}
234-
235-
// Only update if not already at 0
236-
if originalReplicas > 0 {
237-
if dep.Annotations == nil {
238-
dep.Annotations = make(map[string]string)
239-
}
240-
dep.Annotations[PreRestoreReplicasAnnotation] = fmt.Sprintf("%d", originalReplicas)
241-
zero := int32(0)
242-
dep.Spec.Replicas = &zero
243-
244-
_, err = client.AppsV1().Deployments(namespace).Update(ctx, dep, metav1.UpdateOptions{})
245-
return err
229+
if dep.Annotations == nil {
230+
dep.Annotations = make(map[string]string)
246231
}
232+
dep.Annotations[PreRestoreReplicasAnnotation] = fmt.Sprintf("%d", originalReplicas)
233+
zero := int32(0)
234+
dep.Spec.Replicas = &zero
247235
return nil
248236
})
249237

@@ -257,28 +245,16 @@ func scaleDownDeployment(ctx context.Context, client kubernetes.Interface, names
257245
func scaleDownStatefulSet(ctx context.Context, client kubernetes.Interface, namespace, name string) (int32, error) {
258246
var originalReplicas int32
259247

260-
err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
261-
sts, err := client.AppsV1().StatefulSets(namespace).Get(ctx, name, metav1.GetOptions{})
262-
if err != nil {
263-
return err
264-
}
265-
248+
err := updateStatefulSetWithRetry(ctx, client, namespace, name, func(sts *appsv1.StatefulSet) error {
266249
if sts.Spec.Replicas != nil {
267250
originalReplicas = *sts.Spec.Replicas
268251
}
269-
270-
// Only update if not already at 0
271-
if originalReplicas > 0 {
272-
if sts.Annotations == nil {
273-
sts.Annotations = make(map[string]string)
274-
}
275-
sts.Annotations[PreRestoreReplicasAnnotation] = fmt.Sprintf("%d", originalReplicas)
276-
zero := int32(0)
277-
sts.Spec.Replicas = &zero
278-
279-
_, err = client.AppsV1().StatefulSets(namespace).Update(ctx, sts, metav1.UpdateOptions{})
280-
return err
252+
if sts.Annotations == nil {
253+
sts.Annotations = make(map[string]string)
281254
}
255+
sts.Annotations[PreRestoreReplicasAnnotation] = fmt.Sprintf("%d", originalReplicas)
256+
zero := int32(0)
257+
sts.Spec.Replicas = &zero
282258
return nil
283259
})
284260

@@ -291,14 +267,9 @@ func scaleDownStatefulSet(ctx context.Context, client kubernetes.Interface, name
291267
//nolint:dupl // Deployment and StatefulSet are different K8s types requiring separate implementations
292268
func scaleUpDeploymentFromAnnotation(ctx context.Context, client kubernetes.Interface, namespace, name string) (int32, bool, error) {
293269
var scaledTo int32
294-
var found bool
295-
296-
err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
297-
dep, err := client.AppsV1().Deployments(namespace).Get(ctx, name, metav1.GetOptions{})
298-
if err != nil {
299-
return err
300-
}
270+
found := true
301271

272+
err := updateDeploymentWithRetry(ctx, client, namespace, name, func(dep *appsv1.Deployment) error {
302273
if dep.Annotations == nil {
303274
found = false
304275
return nil
@@ -315,15 +286,11 @@ func scaleUpDeploymentFromAnnotation(ctx context.Context, client kubernetes.Inte
315286
return fmt.Errorf("failed to parse replicas annotation: %w", err)
316287
}
317288

318-
dep.Spec.Replicas = &originalReplicas
319289
delete(dep.Annotations, PreRestoreReplicasAnnotation)
290+
dep.Spec.Replicas = &originalReplicas
291+
scaledTo = originalReplicas
320292

321-
_, err = client.AppsV1().Deployments(namespace).Update(ctx, dep, metav1.UpdateOptions{})
322-
if err == nil {
323-
scaledTo = originalReplicas
324-
found = true
325-
}
326-
return err
293+
return nil
327294
})
328295

329296
return scaledTo, found, err
@@ -335,14 +302,9 @@ func scaleUpDeploymentFromAnnotation(ctx context.Context, client kubernetes.Inte
335302
//nolint:dupl // Deployment and StatefulSet are different K8s types requiring separate implementations
336303
func scaleUpStatefulSetFromAnnotation(ctx context.Context, client kubernetes.Interface, namespace, name string) (int32, bool, error) {
337304
var scaledTo int32
338-
var found bool
339-
340-
err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
341-
sts, err := client.AppsV1().StatefulSets(namespace).Get(ctx, name, metav1.GetOptions{})
342-
if err != nil {
343-
return err
344-
}
305+
found := true
345306

307+
err := updateStatefulSetWithRetry(ctx, client, namespace, name, func(sts *appsv1.StatefulSet) error {
346308
if sts.Annotations == nil {
347309
found = false
348310
return nil
@@ -359,15 +321,11 @@ func scaleUpStatefulSetFromAnnotation(ctx context.Context, client kubernetes.Int
359321
return fmt.Errorf("failed to parse replicas annotation: %w", err)
360322
}
361323

362-
sts.Spec.Replicas = &originalReplicas
363324
delete(sts.Annotations, PreRestoreReplicasAnnotation)
325+
sts.Spec.Replicas = &originalReplicas
326+
scaledTo = originalReplicas
364327

365-
_, err = client.AppsV1().StatefulSets(namespace).Update(ctx, sts, metav1.UpdateOptions{})
366-
if err == nil {
367-
scaledTo = originalReplicas
368-
found = true
369-
}
370-
return err
328+
return nil
371329
})
372330

373331
return scaledTo, found, err
@@ -570,12 +528,13 @@ func (c *Client) SetRestoreLock(namespace, labelSelector, datastore, startedAt s
570528
}
571529

572530
for _, dep := range deployments.Items {
573-
err := updateDeploymentWithRetry(ctx, c.clientset, namespace, dep.Name, func(d *appsv1.Deployment) {
531+
err := updateDeploymentWithRetry(ctx, c.clientset, namespace, dep.Name, func(d *appsv1.Deployment) error {
574532
if d.Annotations == nil {
575533
d.Annotations = make(map[string]string)
576534
}
577535
d.Annotations[RestoreInProgressAnnotation] = datastore
578536
d.Annotations[RestoreStartedAtAnnotation] = startedAt
537+
return nil
579538
})
580539
if err != nil {
581540
return fmt.Errorf("failed to set restore lock on deployment %s: %w", dep.Name, err)
@@ -591,12 +550,13 @@ func (c *Client) SetRestoreLock(namespace, labelSelector, datastore, startedAt s
591550
}
592551

593552
for _, sts := range statefulSets.Items {
594-
err := updateStatefulSetWithRetry(ctx, c.clientset, namespace, sts.Name, func(s *appsv1.StatefulSet) {
553+
err := updateStatefulSetWithRetry(ctx, c.clientset, namespace, sts.Name, func(s *appsv1.StatefulSet) error {
595554
if s.Annotations == nil {
596555
s.Annotations = make(map[string]string)
597556
}
598557
s.Annotations[RestoreInProgressAnnotation] = datastore
599558
s.Annotations[RestoreStartedAtAnnotation] = startedAt
559+
return nil
600560
})
601561
if err != nil {
602562
return fmt.Errorf("failed to set restore lock on statefulset %s: %w", sts.Name, err)
@@ -640,10 +600,11 @@ func (c *Client) ClearRestoreLock(namespace, labelSelector string) error {
640600
continue
641601
}
642602

643-
err := updateDeploymentWithRetry(ctx, c.clientset, namespace, dep.Name, func(d *appsv1.Deployment) {
603+
err := updateDeploymentWithRetry(ctx, c.clientset, namespace, dep.Name, func(d *appsv1.Deployment) error {
644604
if d.Annotations != nil {
645605
removeRestoreLockAnnotations(d.Annotations)
646606
}
607+
return nil
647608
})
648609
if err != nil {
649610
return fmt.Errorf("failed to clear restore lock on deployment %s: %w", dep.Name, err)
@@ -663,10 +624,11 @@ func (c *Client) ClearRestoreLock(namespace, labelSelector string) error {
663624
continue
664625
}
665626

666-
err := updateStatefulSetWithRetry(ctx, c.clientset, namespace, sts.Name, func(s *appsv1.StatefulSet) {
627+
err := updateStatefulSetWithRetry(ctx, c.clientset, namespace, sts.Name, func(s *appsv1.StatefulSet) error {
667628
if s.Annotations != nil {
668629
removeRestoreLockAnnotations(s.Annotations)
669630
}
631+
return nil
670632
})
671633
if err != nil {
672634
return fmt.Errorf("failed to clear restore lock on statefulset %s: %w", sts.Name, err)

0 commit comments

Comments
 (0)