wenshu/k8s-orchestrator/internal/controller/nodeinventory_controller.go

package controller

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	meta "k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	orcv1 "github.com/cyrus/spark-orchestrator/api/v1alpha1"
	"github.com/cyrus/spark-orchestrator/internal/k8s"
	"github.com/cyrus/spark-orchestrator/internal/onboarding"
)

type NodeInventoryReconciler struct {
	client.Client
	Scheme             *runtime.Scheme
	Onboarder          onboarding.Onboarder
	DefaultJoinCommand string
}

func (r *NodeInventoryReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	logger := log.FromContext(ctx).WithValues("nodeInventory", req.Name)

	var nodeInv orcv1.NodeInventory
	if err := r.Get(ctx, req.NamespacedName, &nodeInv); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	if nodeInv.Status.Phase == "" {
		nodeInv.Status.Phase = orcv1.NodePhaseNew
		nodeInv.Status.ObservedAt = metav1.Now()
		nodeInv.Status.Conditions = []metav1.Condition{orcv1.NewCondition(orcv1.ConditionRunning, "Initialized", "Node registration initialized", metav1.ConditionTrue, nodeInv.Generation)}
		if err := r.Status().Update(ctx, &nodeInv); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{RequeueAfter: 1 * time.Second}, nil
	}

	switch nodeInv.Status.Phase {
	case orcv1.NodePhaseNew:
		nodeInv.Status.Phase = orcv1.NodePhaseProvisioning
		nodeInv.Status.ObservedAt = metav1.Now()
		setCondition(&nodeInv.Status.Conditions, orcv1.NewCondition(orcv1.ConditionRunning, "Provisioning", "Node provisioning started", metav1.ConditionTrue, nodeInv.Generation))
		if err := r.Status().Update(ctx, &nodeInv); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{Requeue: true}, nil
	case orcv1.NodePhaseProvisioning:
		if err := r.ensureNodeJoined(ctx, &nodeInv); err != nil {
			logger.Error(err, "node onboarding failed")
			nodeInv.Status.Phase = orcv1.NodePhaseFailed
			nodeInv.Status.LastError = err.Error()
			nodeInv.Status.ObservedAt = metav1.Now()
			setCondition(&nodeInv.Status.Conditions, orcv1.NewCondition(orcv1.ConditionFailed, "OnboardingFailed", err.Error(), metav1.ConditionTrue, nodeInv.Generation))
			_ = r.Status().Update(ctx, &nodeInv)
			return ctrl.Result{RequeueAfter: 30 * time.Second}, nil
		}

		// Prefer the name already resolved by ensureNodeJoined (set via --node-name).
		// Fall back to IP-based lookup for nodes joined without an explicit name.
		nodeName := nodeInv.Status.KubernetesNodeName
		if nodeName == "" {
			name, err := k8s.FindNodeByInternalIP(ctx, r.Client, nodeInv.Spec.Address)
			if err != nil {
				if apierrors.IsNotFound(err) {
					return ctrl.Result{RequeueAfter: 10 * time.Second}, nil
				}
				return ctrl.Result{}, err
			}
			nodeName = name
		}
		var k8sNode corev1.Node
		if err := r.Get(ctx, types.NamespacedName{Name: nodeName}, &k8sNode); err != nil {
			if apierrors.IsNotFound(err) {
				return ctrl.Result{RequeueAfter: 5 * time.Second}, nil
			}
			return ctrl.Result{}, err
		}
		nodeInv.Status.KubernetesNodeName = nodeName
		if !k8s.IsNodeReady(&k8sNode) {
			nodeInv.Status.ObservedAt = metav1.Now()
			if err := r.Status().Update(ctx, &nodeInv); err != nil {
				return ctrl.Result{}, err
			}
			return ctrl.Result{RequeueAfter: 5 * time.Second}, nil
		}

		nodeInv.Status.Phase = orcv1.NodePhaseIdle
		nodeInv.Status.LastError = ""
		nodeInv.Status.ObservedAt = metav1.Now()
		setCondition(&nodeInv.Status.Conditions, orcv1.NewCondition(orcv1.ConditionReady, "NodeReady", "Node is ready for allocation", metav1.ConditionTrue, nodeInv.Generation))
		if err := r.Status().Update(ctx, &nodeInv); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	case orcv1.NodePhaseFailed:
		// Retry failed nodes in controller loop.
		nodeInv.Status.Phase = orcv1.NodePhaseProvisioning
		nodeInv.Status.ObservedAt = metav1.Now()
		if err := r.Status().Update(ctx, &nodeInv); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{RequeueAfter: 5 * time.Second}, nil
	default:
		return ctrl.Result{}, nil
	}
}

func (r *NodeInventoryReconciler) ensureNodeJoined(ctx context.Context, nodeInv *orcv1.NodeInventory) error {
	if nodeInv.Status.KubernetesNodeName != "" {
		return nil
	}

	// Check by predictable node name (set via --node-name in kubeadm join).
	// This handles the case where join succeeded but the node name wasn't yet persisted.
	var candidate corev1.Node
	if err := r.Get(ctx, types.NamespacedName{Name: nodeInv.Name}, &candidate); err == nil {
		nodeInv.Status.KubernetesNodeName = candidate.Name
		return nil
	}

	// Fallback: search by internal IP (covers nodes joined without --node-name).
	if existing, err := k8s.FindNodeByInternalIP(ctx, r.Client, nodeInv.Spec.Address); err == nil && existing != "" {
		nodeInv.Status.KubernetesNodeName = existing
		return nil
	}

	if r.Onboarder == nil {
		return fmt.Errorf("onboarder is not configured")
	}
	secretNs, secretName := parseSecretRef(nodeInv.Spec.SSHSecretRef)
	if secretNs == "" || secretName == "" {
		return fmt.Errorf("invalid sshSecretRef %q", nodeInv.Spec.SSHSecretRef)
	}
	var secret corev1.Secret
	if err := r.Get(ctx, types.NamespacedName{Namespace: secretNs, Name: secretName}, &secret); err != nil {
		return fmt.Errorf("load ssh secret: %w", err)
	}
	user := string(secret.Data["sshUser"])
	if user == "" {
		user = "root"
	}
	password := string(secret.Data["sshPassword"])
	if password == "" {
		return fmt.Errorf("sshPassword key missing in secret")
	}
	port := int32(22)
	if rawPort := string(secret.Data["sshPort"]); rawPort != "" {
		parsed, err := strconv.Atoi(rawPort)
		if err != nil {
			return fmt.Errorf("invalid sshPort in secret: %w", err)
		}
		port = int32(parsed)
	}
	joinCmd := nodeInv.Spec.Bootstrap.JoinCommandOverride
	if joinCmd == "" {
		joinCmd = r.DefaultJoinCommand
	}
	if joinCmd == "" {
		return fmt.Errorf("join command is empty")
	}
	return r.Onboarder.Onboard(ctx, onboarding.NodeJoinRequest{
		Address:      nodeInv.Spec.Address,
		Port:         port,
		User:         user,
		Password:     password,
		JoinCommand:  joinCmd,
		NodeNameHint: nodeInv.Name,
	})
}

func (r *NodeInventoryReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&orcv1.NodeInventory{}).
		Complete(r)
}

func setCondition(conditions *[]metav1.Condition, cond metav1.Condition) {
	meta.SetStatusCondition(conditions, cond)
}

func parseSecretRef(ref string) (string, string) {
	parts := strings.Split(ref, "/")
	if len(parts) != 2 {
		return "", ""
	}
	return parts[0], parts[1]
}