objectrestarter.go

Documentation: edge-infra.dev/pkg/k8s/runtime/objectrestarter

     1  // Package objectrestarter provides a utility for restarting workloads using the
     2  // kubectl annotation used by K8s for graceful restarting.
     3  // Inspired by https://github.com/kubernetes/kubernetes/blob/master/staging/src/k8s.io/kubectl/pkg/polymorphichelpers/objectrestarter.go.
     4  package objectrestarter
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"fmt"
    10  	"time"
    11  
    12  	appsv1 "k8s.io/api/apps/v1"
    13  	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
    14  	"k8s.io/apimachinery/pkg/types"
    15  	"sigs.k8s.io/controller-runtime/pkg/client"
    16  )
    17  
    18  var ErrWorkloadIsNotRestartable = errors.New("workload kind is not restartable")
    19  
    20  // restartAnnotation is the annotation added by kubectl when you run
    21  // `kubectl rollout restart`
    22  const restartAnnotation = "kubectl.kubernetes.io/restartedAt"
    23  
    24  // Restart performs a graceful restart of the provided workload object, if it
    25  // is a restartable workload.
    26  func Restart(ctx context.Context, c client.Client, workload client.Object) error {
    27  	patch := []byte(fmt.Sprintf(
    28  		`{"spec": { "template": {"metadata":{"annotations":{"%s": "%s"}}}}}`,
    29  		restartAnnotation, time.Now().Format(time.RFC3339)),
    30  	)
    31  	return c.Patch(ctx, workload, client.RawPatch(types.MergePatchType, patch))
    32  }
    33  
    34  // IsKindRestartable returns true if the provided Kind string refers to a workload
    35  // object that can be restarted.
    36  func IsKindRestartable(kind string) bool {
    37  	return kind == "StatefulSet" || kind == "DaemonSet" || kind == "Deployment"
    38  }
    39  
    40  // IsObjectRestartable determines if the provided Object interface represents
    41  // a workload object that can be restarted.  If an Unstructured object is provided,
    42  // it must have had its Kind information set explicitly.
    43  func IsObjectRestartable(obj client.Object) bool {
    44  	switch obj.(type) {
    45  	case *appsv1.Deployment:
    46  		return true
    47  	case *appsv1.DaemonSet:
    48  		return true
    49  	case *appsv1.StatefulSet:
    50  		return true
    51  	case *unstructured.Unstructured:
    52  		return IsKindRestartable(obj.GetObjectKind().GroupVersionKind().Kind)
    53  	default:
    54  		return false
    55  	}
    56  }
    57
View as plain text