<commit_before>package rkeworkerupgrader

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

	"github.com/docker/docker/pkg/locker"
	"github.com/rancher/rancher/pkg/controllers/management/clusterprovisioner"
	nodeserver "github.com/rancher/rancher/pkg/rkenodeconfigserver"
	"github.com/rancher/rancher/pkg/systemaccount"
	"github.com/rancher/rke/util"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/errgroup"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
)

type upgradeHandler struct {
	clusters             v3.ClusterInterface
	nodes                v3.NodeInterface
	nodeLister           v3.NodeLister
	clusterLister        v3.ClusterLister
	lookup               *nodeserver.BundleLookup
	systemAccountManager *systemaccount.Manager
	serviceOptionsLister v3.RKEK8sServiceOptionLister
	serviceOptions       v3.RKEK8sServiceOptionInterface
	sysImagesLister      v3.RKEK8sSystemImageLister
	sysImages            v3.RKEK8sSystemImageInterface
	clusterLock          *locker.Locker
	ctx                  context.Context
}

func Register(ctx context.Context, mgmt *config.ManagementContext, scaledContext *config.ScaledContext) {

	uh := &upgradeHandler{
		clusters:             mgmt.Management.Clusters(""),
		clusterLister:        mgmt.Management.Clusters("").Controller().Lister(),
		nodes:                mgmt.Management.Nodes(""),
		nodeLister:           mgmt.Management.Nodes("").Controller().Lister(),
		lookup:               nodeserver.NewLookup(scaledContext.Core.Namespaces(""), scaledContext.Core),
		systemAccountManager: systemaccount.NewManagerFromScale(scaledContext),
		serviceOptionsLister: mgmt.Management.RKEK8sServiceOptions("").Controller().Lister(),
		serviceOptions:       mgmt.Management.RKEK8sServiceOptions(""),
		sysImagesLister:      mgmt.Management.RKEK8sSystemImages("").Controller().Lister(),
		sysImages:            mgmt.Management.RKEK8sSystemImages(""),
		clusterLock:          locker.New(),
		ctx:                  ctx,
	}

	mgmt.Management.Nodes("").Controller().AddHandler(ctx, "rke-worker-upgrader", uh.Sync)
}

func (uh *upgradeHandler) Sync(key string, node *v3.Node) (runtime.Object, error) {
	if strings.HasSuffix(key, "upgrade_") {

		cName := strings.Split(key, "/")[0]
		cluster, err := uh.clusterLister.Get("", cName)
		if err != nil {
			return nil, err
		}
		if cluster.DeletionTimestamp != nil || cluster.Status.AppliedSpec.RancherKubernetesEngineConfig == nil {
			return nil, nil
		}

		if cluster.Annotations[clusterprovisioner.RkeRestoreAnnotation] == "true" {
			// rke completed etcd snapshot restore, need to update plan for nodes so worker nodes reset to old rkeConfig
			return node, uh.restore(cluster)
		}

		logrus.Infof("checking cluster [%s] for worker nodes upgrade", cluster.Name)

		if toUpgrade, planChanged, err := uh.toUpgradeCluster(cluster); err != nil {
			return nil, err
		} else if toUpgrade {
			if err := uh.upgradeCluster(cluster, key, planChanged); err != nil {
				return nil, err
			}
		}

		return nil, nil
	}

	if node == nil || node.DeletionTimestamp != nil || !v3.NodeConditionProvisioned.IsTrue(node) {
		return node, nil
	}

	cluster, err := uh.clusterLister.Get("", node.Namespace)
	if err != nil {
		return nil, err
	}

	if cluster.DeletionTimestamp != nil || cluster.Status.AppliedSpec.RancherKubernetesEngineConfig == nil {
		return node, nil
	}

	if node.Status.NodePlan == nil {
		return uh.updateNodePlan(node, cluster, true)
	}

	if v3.ClusterConditionUpdated.IsUnknown(cluster) || cluster.Annotations[clusterprovisioner.RkeRestoreAnnotation] == "true" {
		return node, nil
	}

	if v3.ClusterConditionUpgraded.IsUnknown(cluster) {
		// if sync is for a node that was just updated, do nothing

		// node is already updated to cordon/drain/uncordon, do nothing
		if node.Spec.DesiredNodeUnschedulable != "" {
			logrus.Debugf("cluster [%s] worker-upgrade: return node [%s], unschedulable field [%v]", cluster.Name,
				node.Name, node.Spec.DesiredNodeUnschedulable)
			return node, nil
		}
		// node is upgrading and already updated with new node plan, do nothing
		if v3.NodeConditionUpgraded.IsUnknown(node) && node.Status.AppliedNodeVersion != cluster.Status.NodeVersion {
			if node.Status.NodePlan.Version == cluster.Status.NodeVersion {
				logrus.Debugf("cluster [%s] worker-upgrade: return node [%s], plan's updated [%v]", cluster.Name,
					node.Name, cluster.Status.NodeVersion)
				return node, nil
			}
		}

		logrus.Infof("cluster [%s] worker-upgrade: call upgrade to reconcile for node [%s]", cluster.Name, node.Name)
		if err := uh.upgradeCluster(cluster, node.Name, false); err != nil {
			return nil, err
		}
		return node, nil
	}

	// proceed only if node and cluster's versions mismatch
	if cluster.Status.NodeVersion == node.Status.AppliedNodeVersion {
		if v3.NodeConditionUpgraded.IsUnknown(node) {
			if err := uh.updateNodeActive(node); err != nil {
				return nil, err
			}
		}
		return node, nil
	}

	nodePlan, err := uh.getNodePlan(node, cluster)
	if err != nil {
		return nil, err
	}

	if node.Status.AppliedNodeVersion == 0 {
		// node never received appliedNodeVersion
		if planChangedForUpgrade(nodePlan, node.Status.NodePlan.Plan) || planChangedForUpdate(nodePlan, node.Status.NodePlan.Plan) {
			return uh.updateNodePlan(node, cluster, false)
		}
	} else {
		if planChangedForUpgrade(nodePlan, node.Status.NodePlan.Plan) {
			logrus.Infof("cluster [%s] worker-upgrade: plan changed for node [%s], call upgrade to reconcile cluster", cluster.Name, node.Name)
			if err := uh.upgradeCluster(cluster, node.Name, false); err != nil {
				return nil, err
			}
			return node, nil
		}
		if planChangedForUpdate(nodePlan, node.Status.NodePlan.Plan) {
			logrus.Infof("cluster [%s] worker-upgrade: plan changed for update [%s]", cluster.Name, node.Name)
			return uh.updateNodePlan(node, cluster, false)
		}
	}

	return node, nil
}

func (uh *upgradeHandler) updateNodePlan(node *v3.Node, cluster *v3.Cluster, create bool) (*v3.Node, error) {
	if node.Status.NodeConfig == nil {
		return node, nil
	}
	nodePlan, err := uh.getNodePlan(node, cluster)
	if err != nil {
		return nil, fmt.Errorf("getNodePlan error for node [%s]: %v", node.Name, err)
	}

	nodeCopy := node.DeepCopy()
	np := &v3.NodePlan{
		Plan:    nodePlan,
		Version: cluster.Status.NodeVersion,
	}
	if node.Status.NodePlan == nil || node.Status.NodePlan.AgentCheckInterval == 0 {
		// default<commit_msg>remove unused code<commit_after>package rkeworkerupgrader

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

	"github.com/docker/docker/pkg/locker"
	"github.com/rancher/rancher/pkg/controllers/management/clusterprovisioner"
	nodeserver "github.com/rancher/rancher/pkg/rkenodeconfigserver"
	"github.com/rancher/rancher/pkg/systemaccount"
	"github.com/rancher/rke/util"
	v3 "github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/errgroup"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
)

type upgradeHandler struct {
	clusters             v3.ClusterInterface
	nodes                v3.NodeInterface
	nodeLister           v3.NodeLister
	clusterLister        v3.ClusterLister
	lookup               *nodeserver.BundleLookup
	systemAccountManager *systemaccount.Manager
	serviceOptionsLister v3.RKEK8sServiceOptionLister
	serviceOptions       v3.RKEK8sServiceOptionInterface
	sysImagesLister      v3.RKEK8sSystemImageLister
	sysImages            v3.RKEK8sSystemImageInterface
	clusterLock          *locker.Locker
	ctx                  context.Context
}

func Register(ctx context.Context, mgmt *config.ManagementContext, scaledContext *config.ScaledContext) {

	uh := &upgradeHandler{
		clusters:             mgmt.Management.Clusters(""),
		clusterLister:        mgmt.Management.Clusters("").Controller().Lister(),
		nodes:                mgmt.Management.Nodes(""),
		nodeLister:           mgmt.Management.Nodes("").Controller().Lister(),
		lookup:               nodeserver.NewLookup(scaledContext.Core.Namespaces(""), scaledContext.Core),
		systemAccountManager: systemaccount.NewManagerFromScale(scaledContext),
		serviceOptionsLister: mgmt.Management.RKEK8sServiceOptions("").Controller().Lister(),
		serviceOptions:       mgmt.Management.RKEK8sServiceOptions(""),
		sysImagesLister:      mgmt.Management.RKEK8sSystemImages("").Controller().Lister(),
		sysImages:            mgmt.Management.RKEK8sSystemImages(""),
		clusterLock:          locker.New(),
		ctx:                  ctx,
	}

	mgmt.Management.Nodes("").Controller().AddHandler(ctx, "rke-worker-upgrader", uh.Sync)
}

func (uh *upgradeHandler) Sync(key string, node *v3.Node) (runtime.Object, error) {
	if strings.HasSuffix(key, "upgrade_") {

		cName := strings.Split(key, "/")[0]
		cluster, err := uh.clusterLister.Get("", cName)
		if err != nil {
			return nil, err
		}
		if cluster.DeletionTimestamp != nil || cluster.Status.AppliedSpec.RancherKubernetesEngineConfig == nil {
			return nil, nil
		}

		if cluster.Annotations[clusterprovisioner.RkeRestoreAnnotation] == "true" {
			// rke completed etcd snapshot restore, need to update plan for nodes so worker nodes reset to old rkeConfig
			return node, uh.restore(cluster)
		}

		logrus.Infof("checking cluster [%s] for worker nodes upgrade", cluster.Name)

		if toUpgrade, planChanged, err := uh.toUpgradeCluster(cluster); err != nil {
			return nil, err
		} else if toUpgrade {
			if err := uh.upgradeCluster(cluster, key, planChanged); err != nil {
				return nil, err
			}
		}

		return nil, nil
	}

	if node == nil || node.DeletionTimestamp != nil || !v3.NodeConditionProvisioned.IsTrue(node) {
		return node, nil
	}

	cluster, err := uh.clusterLister.Get("", node.Namespace)
	if err != nil {
		return nil, err
	}

	if cluster.DeletionTimestamp != nil || cluster.Status.AppliedSpec.RancherKubernetesEngineConfig == nil {
		return node, nil
	}

	if node.Status.NodePlan == nil {
		return uh.updateNodePlan(node, cluster, true)
	}

	if v3.ClusterConditionUpdated.IsUnknown(cluster) || cluster.Annotations[clusterprovisioner.RkeRestoreAnnotation] == "true" {
		return node, nil
	}

	if v3.ClusterConditionUpgraded.IsUnknown(cluster) {
		// if sync is for a node that was just updated, do nothing

		// node is already updated to cordon/drain/uncordon, do nothing
		if node.Spec.DesiredNodeUnschedulable != "" {
			logrus.Debugf("cluster [%s] worker-upgrade: return node [%s], unschedulable field [%v]", cluster.Name,
				node.Name, node.Spec.DesiredNodeUnschedulable)
			return node, nil
		}
		// node is upgrading and already updated with new node plan, do nothing
		if v3.NodeConditionUpgraded.IsUnknown(node) && node.Status.AppliedNodeVersion != cluster.Status.NodeVersion {
			if node.Status.NodePlan.Version == cluster.Status.NodeVersion {
				logrus.Debugf("cluster [%s] worker-upgrade: return node [%s], plan's updated [%v]", cluster.Name,
					node.Name, cluster.Status.NodeVersion)
				return node, nil
			}
		}

		logrus.Infof("cluster [%s] worker-upgrade: call upgrade to reconcile for node [%s]", cluster.Name, node.Name)
		if err := uh.upgradeCluster(cluster, node.Name, false); err != nil {
			return nil, err
		}
		return node, nil
	}

	// proceed only if node and cluster's versions mismatch
	if cluster.Status.NodeVersion == node.Status.NodePlan.Version {
		return node, nil
	}

	nodePlan, err := uh.getNodePlan(node, cluster)
	if err != nil {
		return nil, err
	}

	if node.Status.AppliedNodeVersion == 0 {
		// node never received appliedNodeVersion
		if planChangedForUpgrade(nodePlan, node.Status.NodePlan.Plan) || planChangedForUpdate(nodePlan, node.Status.NodePlan.Plan) {
			return uh.updateNodePlan(node, cluster, false)
		}
	} else {
		if planChangedForUpdate(nodePlan, node.Status.NodePlan.Plan) {
			logrus.Infof("cluster [%s] worker-upgrade: plan changed for update [%s]", cluster.Name, node.Name)
			return uh.updateNodePlan(node, cluster, false)
		}
	}

	return node, nil
}

func (uh *upgradeHandler) updateNodePlan(node *v3.Node, cluster *v3.Cluster, create bool) (*v3.Node, error) {
	if node.Status.NodeConfig == nil {
		return node, nil
	}
	nodePlan, err := uh.getNodePlan(node, cluster)
	if err != nil {
		return nil, fmt.Errorf("getNodePlan error for node [%s]: %v", node.Name, err)
	}

	nodeCopy := node.DeepCopy()
	np := &v3.NodePlan{
		Plan:    nodePlan,
		Version: cluster.Status.NodeVersion,
	}
	if node.Status.NodePlan == nil || node.Status.NodePlan.AgentCheckInterval == 0 {
		// default
{'fraction_non_alphanumeric': 0.12423277134643773, 'fraction_numerical': 0.005982441146764043, 'mean_word_length': 3.870223231176693, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5827761', 'n_tokens_mistral': 4601, 'n_tokens_neox': 4238, 'n_words': 898}