Add sharding support for ops-manager (#12)

Signed-off-by: souravbiswassanto <saurov@appscode.com>
Signed-off-by: Tamal Saha <tamal@appscode.com>
Signed-off-by: Arnob Kumar Saha <arnob@appscode.com>
Co-authored-by: Tamal Saha <tamal@appscode.com>
Co-authored-by: Arnob Kumar Saha <arnob@appscode.com>

Author: 3 people

Makefile

@@ -55,7 +55,7 @@ endif
 ### These variables should not need tweaking.
 ###
 
-SRC_PKGS := cmd crds pkg
+SRC_PKGS := cmd crds pkg api
 SRC_DIRS := $(SRC_PKGS) # directories which hold app source (not vendored)
 
 DOCKER_PLATFORMS := linux/amd64 linux/arm64

api/v1alpha1/shardconfiguration_types.go

@@ -32,7 +32,13 @@ type ShardConfigurationSpec struct {
 	// +kubebuilder:validation:MinItems=1
 	Controllers []kmapi.TypedObjectReference `json:"controllers,omitempty"`
 	// +kubebuilder:validation:MinItems=1
-	Resources []kmapi.TypeReference `json:"resources,omitempty"`
+	Resources []ResourceInfo `json:"resources,omitempty"`
+}
+
+type ResourceInfo struct {
+	kmapi.TypeReference          `json:",inline"`
+	ShardKey                     *string `json:"shardKey,omitempty"`
+	UseCooperativeShardMigration bool    `json:"useCooperativeShardMigration,omitempty"`
 }
 
 type ControllerAllocation struct {

api/v1alpha1/zz_generated.deepcopy.go

@@ -66,12 +66,15 @@ spec:
                 type: array
               resources:
                 items:
-                  description: TypeReference represents an object type.
                   properties:
                     apiGroup:
                       type: string
                     kind:
                       type: string
+                    shardKey:
+                      type: string
+                    useCooperativeShardMigration:
+                      type: boolean
                   type: object
                 minItems: 1
                 type: array

crds/operator.k8s.appscode.com_shardconfigurations.yaml

@@ -72,7 +72,11 @@ func getBetterPartitionCount(members int, load float64) int {
 	return candidate
 }
 
-func newConsistentConfig(members []consistent.Member, shardCount int) *consistent.Consistent {
+func newConsistentConfig(shardCount int) *consistent.Consistent {
+	members := make([]consistent.Member, 0, shardCount)
+	for i := 0; i < shardCount; i++ {
+		members = append(members, Member{ID: i})
+	}
 	return consistent.New(members, consistent.Config{
 		PartitionCount:    getBetterPartitionCount(shardCount, 1.0),
 		ReplicationFactor: 1,

pkg/controller/hashing.go

@@ -22,7 +22,6 @@ import (
 	"fmt"
 	"sort"
 	"strconv"
-	"strings"
 	"sync"
 	"time"
 
@@ -33,6 +32,7 @@ import (
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
@@ -155,11 +155,13 @@ func (r *ShardConfigurationReconciler) Reconcile(ctx context.Context, req ctrl.R
 	if shardCount <= 0 {
 		return ctrl.Result{RequeueAfter: time.Second * 2}, nil
 	}
-	members := make([]consistent.Member, 0, shardCount)
-	for i := 0; i < shardCount; i++ {
-		members = append(members, Member{ID: i})
+
+	cc := newConsistentConfig(shardCount)
+	resourceLists, err := r.d.ServerPreferredResources()
+	if err != nil && !discovery.IsGroupDiscoveryFailedError(err) {
+		return ctrl.Result{}, err
 	}
-	cc := newConsistentConfig(members, shardCount)
+
 	for _, resource := range cfg.Spec.Resources {
 		if resource.Kind != "" {
 			mapping, err := r.mapper.RESTMapping(schema.GroupKind{
@@ -175,14 +177,10 @@ func (r *ShardConfigurationReconciler) Reconcile(ctx context.Context, req ctrl.R
 				return ctrl.Result{}, err
 			}
 
-			if err := r.UpdateShardLabel(ctx, cc, gvk, &cfg); err != nil {
+			if err := r.UpdateShardLabel(ctx, cc, gvk, &cfg, resource); err != nil {
 				return ctrl.Result{}, err
 			}
 		} else {
-			resourceLists, err := r.d.ServerPreferredResources()
-			if err != nil && !discovery.IsGroupDiscoveryFailedError(err) {
-				return ctrl.Result{}, err
-			}
 			for _, resourceList := range resourceLists {
 				if gv, err := schema.ParseGroupVersion(resourceList.GroupVersion); err == nil && gv.Group == resource.APIGroup {
 					for _, apiResource := range resourceList.APIResources {
@@ -191,7 +189,7 @@ func (r *ShardConfigurationReconciler) Reconcile(ctx context.Context, req ctrl.R
 							if err := r.RegisterResourceWatcher(gvk); err != nil {
 								return ctrl.Result{}, err
 							}
-							if err := r.UpdateShardLabel(ctx, cc, gvk, &cfg); err != nil {
+							if err := r.UpdateShardLabel(ctx, cc, gvk, &cfg, resource); err != nil {
 								return ctrl.Result{}, err
 							}
 						}
@@ -204,30 +202,71 @@ func (r *ShardConfigurationReconciler) Reconcile(ctx context.Context, req ctrl.R
 	return ctrl.Result{}, nil
 }
 
-func (r *ShardConfigurationReconciler) UpdateShardLabel(ctx context.Context, cc *consistent.Consistent, gvk schema.GroupVersionKind, cfg *shardapi.ShardConfiguration) error {
-	log := log.FromContext(ctx)
+func (r *ShardConfigurationReconciler) UpdateShardLabel(ctx context.Context, cc *consistent.Consistent, gvk schema.GroupVersionKind, cfg *shardapi.ShardConfiguration, ri shardapi.ResourceInfo) error {
+	logger := log.FromContext(ctx)
 	shardKey := fmt.Sprintf("shard.%s/%s", shardapi.SchemeGroupVersion.Group, cfg.Name)
-	var list metav1.PartialObjectMetadataList
+	nextShardKey := fmt.Sprintf("next.%s/%s", shardapi.SchemeGroupVersion.Group, cfg.Name)
+
+	var list unstructured.UnstructuredList
 	list.SetGroupVersionKind(gvk)
 	err := r.List(ctx, &list)
 	if err != nil {
 		return err
 	}
+
+	ifShardKeyLabelNeedsToBeChanged := func(labels map[string]string, shardKey string, member consistent.Member) bool {
+		return labels[shardKey] != "" && labels[shardKey] != member.String()
+	}
+
 	for _, obj := range list.Items {
-		m := cc.LocateKey([]byte(fmt.Sprintf("%s/%s", obj.GetNamespace(), obj.GetName())))
+		var key []byte
+		if ri.ShardKey != nil {
+			val, found := EvaluateJSONPath(obj.Object, *ri.ShardKey)
+			if !found {
+				return fmt.Errorf("failed to extract shard key from %s/%s %s/%s using jsonPath %s", obj.GroupVersionKind().Group, obj.GroupVersionKind().Kind, obj.GetNamespace(), obj.GetName(), *ri.ShardKey)
+			}
+			key = []byte(fmt.Sprintf("%s/%s", obj.GetNamespace(), val))
+		} else {
+			key = []byte(fmt.Sprintf("%s/%s", obj.GetNamespace(), obj.GetName()))
+		}
+		m := cc.LocateKey(key)
+		changed := false
+		labels := obj.GetLabels()
+		if labels == nil {
+			labels = make(map[string]string)
+		}
+		if labels[shardKey] == "" {
+			labels[shardKey] = m.String()
+			changed = true
+		} else if ifShardKeyLabelNeedsToBeChanged(labels, shardKey, m) {
+			switch ri.UseCooperativeShardMigration {
+			case true:
+				if len(cfg.Status.Controllers) > 0 && len(cfg.Status.Controllers[0].Pods) > 0 {
+					id, err := strconv.Atoi(labels[shardKey])
+					if err != nil {
+						return err
+					}
+					if id >= len(cfg.Status.Controllers[0].Pods) {
+						labels[shardKey] = m.String()
+					} else {
+						labels[nextShardKey] = m.String()
+					}
+				}
+			case false:
+				labels[shardKey] = m.String()
+			}
+			changed = true
+		}
 
-		if obj.Labels[shardKey] != m.String() {
+		if changed {
 			opr, err := controllerutil.CreateOrPatch(ctx, r.Client, &obj, func() error {
-				if obj.Labels == nil {
-					obj.Labels = map[string]string{}
-				}
-				obj.Labels[shardKey] = m.String()
+				obj.SetLabels(labels)
 				return nil
 			})
 			if err != nil {
-				log.Error(err, fmt.Sprintf("failed to update labels for %s/%s %s/%s", obj.GroupVersionKind().Group, obj.GroupVersionKind().Kind, obj.GetNamespace(), obj.GetName()))
+				logger.Error(err, fmt.Sprintf("failed to update labels for %s/%s %s/%s", obj.GroupVersionKind().Group, obj.GroupVersionKind().Kind, obj.GetNamespace(), obj.GetName()))
 			} else {
-				log.Info(fmt.Sprintf("%s/%s %s/%s %s", obj.GroupVersionKind().Group, obj.GroupVersionKind().Kind, obj.GetNamespace(), obj.GetName(), opr))
+				logger.Info(fmt.Sprintf("%s/%s %s/%s %s", obj.GroupVersionKind().Group, obj.GroupVersionKind().Kind, obj.GetNamespace(), obj.GetName(), opr))
 			}
 		}
 	}
@@ -360,40 +399,8 @@ func ListPods(ctx context.Context, kc client.Client, ref kmapi.TypedObjectRefere
 		if err != nil {
 			return nil, err
 		}
-		sel, err := metav1.LabelSelectorAsSelector(obj.Spec.Selector)
-		if err != nil {
-			return nil, err
-		}
-		var list metav1.PartialObjectMetadataList
-		list.SetGroupVersionKind(schema.GroupVersionKind{
-			Group:   "",
-			Kind:    "Pod",
-			Version: "v1",
-		})
-		err = kc.List(ctx, &list, client.MatchingLabelsSelector{Selector: sel})
-		if err != nil {
-			return nil, err
-		}
-		pods := make([]string, 0, len(list.Items))
-		for _, pod := range list.Items {
-			if metav1.IsControlledBy(&pod, &obj) {
-				pods = append(pods, pod.Name)
-			}
-		}
-		sort.Slice(pods, func(i, j int) bool {
-			idx_i := strings.LastIndexByte(pods[i], '-')
-			idx_j := strings.LastIndexByte(pods[j], '-')
-			if idx_i == -1 || idx_j == -1 {
-				return pods[i] < pods[j]
-			}
-			oi, err_i := strconv.Atoi(pods[i][idx_i+1:])
-			oj, err_j := strconv.Atoi(pods[j][idx_j+1:])
-			if err_i != nil || err_j != nil {
-				return pods[i] < pods[j]
-			}
-			return oi < oj
-		})
-		return pods, nil
+		// Important: Resharding feature is not available for stateful sets
+		return buildPodList(obj.GetName(), *obj.Spec.Replicas), nil
 	case "DaemonSet":
 		var obj apps.DaemonSet
 		err := kc.Get(ctx, client.ObjectKey{Name: ref.Name, Namespace: ref.Namespace}, &obj)

pkg/controller/shardconfiguration_controller.go

@@ -16,7 +16,13 @@ limitations under the License.
 
 package controller
 
-import "gomodules.xyz/sets"
+import (
+	"fmt"
+	"strconv"
+	"strings"
+
+	"gomodules.xyz/sets"
+)
 
 func getUpdatedPodLists(existing, podLists []string) []string {
 	if len(existing) > len(podLists) {
@@ -25,6 +31,9 @@ func getUpdatedPodLists(existing, podLists []string) []string {
 	return handleUpdateOrUpScaling(existing, podLists)
 }
 
+// Both the handleDownScaling & handleUpdateOrUpScaling func is to sort the new 'podLists' according to the sequence of 'existing'.
+
+// scaleDown example: existing=[2, 3, 0, 1], podLists=[1, 3]. output will be = [3, 1]. Because 3 appears before 1 in the 'existing' array.
 func handleDownScaling(existing, podLists []string) []string {
 	newPods := make([]string, 0)
 
@@ -48,6 +57,7 @@ func handleDownScaling(existing, podLists []string) []string {
 	return newPods
 }
 
+// scaleUp example: existing=[1, 3], podLists=[2, 3, 0, 1]. output will be = [1, 3, 2, 0]. Keeping [1,3] as it is. Then appending the remaining ones from 'podLists'.
 func handleUpdateOrUpScaling(existing, podLists []string) []string {
 	newPods := make([]string, len(podLists))
 
@@ -87,3 +97,68 @@ func getNextAvailableIndex(next int, pods []string) int {
 func isReadable(verbs []string) bool {
 	return sets.NewString(verbs...).HasAll("get", "list", "watch")
 }
+
+// EvaluateJSONPath evaluates a simple JSONPath expression on an unstructured object
+// Supports paths like ".spec.databaseRef.name"
+// Returns the value as a string and true if found, empty string and false otherwise
+func EvaluateJSONPath(obj map[string]any, jsonPath string) (string, bool) {
+	if obj == nil || jsonPath == "" {
+		return "", false
+	}
+
+	// Remove leading dot if present
+	jsonPath = strings.TrimPrefix(jsonPath, ".")
+
+	// Split the path into parts
+	parts := strings.Split(jsonPath, ".")
+
+	// Navigate through the object structure
+	current := any(obj)
+	for i, part := range parts {
+		if part == "" {
+			continue
+		}
+
+		// Try to cast current to map[string]any
+		currentMap, ok := current.(map[string]any)
+		if !ok {
+			return "", false
+		}
+
+		// Get the next value
+		next, exists := currentMap[part]
+		if !exists {
+			return "", false
+		}
+
+		// If this is the last part, try to convert to string
+		if i == len(parts)-1 {
+			switch v := next.(type) {
+			case string:
+				return v, true
+			case int:
+				return strconv.Itoa(v), true
+			case int64:
+				return strconv.FormatInt(v, 10), true
+			case float64:
+				return strconv.FormatFloat(v, 'f', -1, 64), true
+			case bool:
+				return strconv.FormatBool(v), true
+			default:
+				return fmt.Sprintf("%v", v), true
+			}
+		}
+
+		current = next
+	}
+
+	return "", false
+}
+
+func buildPodList(ctrlName string, replCount int32) []string {
+	pods := make([]string, replCount)
+	for c := int32(0); c < replCount; c++ {
+		pods[c] = fmt.Sprintf("%s-%d", ctrlName, c)
+	}
+	return pods
+}