diff --git a/README.md b/README.md index 196af84..3a08cbb 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ Flag | Environment Variable | Type | Default | Required | Description ---- | -------------------- | ---- | ------- | -------- | ----------- `log-level` | `LOG_LEVEL` | `string` | `info` | no | The level of log detail. `bind-address` | `BIND_ADDRESS` | `string` | `:9797` | no | The address for binding listener. -`node-taint` | `NODE_TAINT` | `string` | | yes | The startup taint to put on node. -`daemonset-annotation` | `DAEMONSET_ANNOTATION` | `string` | | yes | The annotation of required daemonset. - - +`node-taint` | `NODE_TAINT` | `string` | | yes | The startup taint to remove from node. +`daemonset-annotation` | `DAEMONSET_ANNOTATION` | `string` | | yes | The annotations of required daemonset, eg. key1:val1,key2:val2,key3. +`daemonset-label` | `DAEMONSET_LABEL` | `string` | | yes | The labels of required daemonset, eg. key1:val1,key2:val2,key3. + diff --git a/config/config.go b/config/config.go index 0d01132..36aa760 100644 --- a/config/config.go +++ b/config/config.go @@ -4,6 +4,7 @@ package config type Ops struct { LogLevel string `long:"log-level" env:"LOG_LEVEL" description:"Log level" default:"info"` NodeTaint string `long:"node-taint" env:"NODE_TAINT" description:"The node taints that's going to remove."` - DaemonSetAnnotation string `long:"daemonset-annotation" env:"DAEMONSET_ANNOTATION" description:"The annotation of the daemonset to watch"` + DaemonSetAnnotation string `long:"daemonset-annotation" env:"DAEMONSET_ANNOTATION" description:"The annotation of the daemonset to watch, key:value pairs, comma delemited"` BindAddr string `long:"bind-address" short:"p" env:"BIND_ADDRESS" default:":9797" description:"address for binding metrics listener"` + DaemonSetLabel string `long:"daemonset-label" env:"DAEMONSET_LABEL" description:"The label of the daemonset to watch, key:value pairs, comma delemited"` } diff --git a/main.go b/main.go index 1abff1d..e24da87 100644 --- a/main.go +++ b/main.go @@ -8,6 +8,7 @@ import ( "nodetaint/config" "os" "os/signal" + "strings" "sync" "syscall" "time" @@ -43,8 +44,22 @@ var ( notReadyTaint = &core_v1.Taint{ Effect: core_v1.TaintEffectNoSchedule, } + + annotationsMap = make(map[string]string) + labelsMap = make(map[string]string) + annotationsKeys = make([]string, 0) + labelsKeys = make([]string, 0) ) +func filter(ss []string, testFnc func(string) bool) string { + for _, str := range ss { + if testFnc(str) { + return str + } + } + return "" +} + func setupLogging(logLevel string) { // Use log level level, err := logrus.ParseLevel(logLevel) @@ -73,7 +88,8 @@ func getClientset() (*kubernetes.Clientset, error) { return clientset, nil } -func checkDSStatus(node *core_v1.Node, opts config.Ops) (bool, error) { +func checkDSStatus(node *core_v1.Node) (bool, error) { + isHandling.Lock() defer isHandling.Unlock() @@ -88,9 +104,37 @@ func checkDSStatus(node *core_v1.Node, opts config.Ops) (bool, error) { readyPods := make(map[string]int) for _, obj := range (*podStore[node.Name]).List() { if pod, ok := obj.(*core_v1.Pod); ok { - // only check the pods with specified annotation - if _, ok := pod.Annotations[opts.DaemonSetAnnotation]; !ok { - continue + // only check the pods with specified annotation or label + if filter(annotationsKeys, func(s string) bool { + _, ok := pod.Annotations[s] + return ok + }) == "" { + if filter(labelsKeys, func(s string) bool { + _, ok := pod.Labels[s] + return ok + }) == "" { + continue + } + } + + if annot := filter(annotationsKeys, func(s string) bool { + _, ok := pod.Annotations[s] + return ok + }); annot != "" { + if annotationsMap[annot] != "" && annotationsMap[annot] != pod.Annotations[annot] { + continue + } + logrus.Debugf("Checking annotated pod name: %s annotation: %s on node: %s", pod.Name, annot, node.Name) + } + + if lbl := filter(labelsKeys, func(s string) bool { + _, ok := pod.Labels[s] + return ok + }); lbl != "" { + if labelsMap[lbl] != "" && labelsMap[lbl] != pod.Labels[lbl] { + continue + } + logrus.Debugf("Checking labeled pod name: %s label: %s on node: %s", pod.Name, lbl, node.Name) } // only check the ready pods @@ -111,6 +155,7 @@ func checkDSStatus(node *core_v1.Node, opts config.Ops) (bool, error) { dsName := pod.OwnerReferences[0].Name if _, ok := dsList[dsName]; ok { readyPods[dsName] = 1 + logrus.Debugf("Ready Pod name: %s", pod.Name) } } } @@ -159,12 +204,12 @@ func RemoveTaintOffNode(ctx context.Context, c *kubernetes.Clientset, node *core if !updated { return nil } - return PatchNodeTaints(c, node.Name, oldNode, newNode, ctx) + return PatchNodeTaints(ctx, c, node.Name, oldNode, newNode) }) } // PatchNodeTaints patches node's taints. -func PatchNodeTaints(c *kubernetes.Clientset, nodeName string, oldNode *core_v1.Node, newNode *core_v1.Node, ctx context.Context) error { +func PatchNodeTaints(ctx context.Context, c *kubernetes.Clientset, nodeName string, oldNode *core_v1.Node, newNode *core_v1.Node) error { oldData, err := json.Marshal(oldNode) if err != nil { return fmt.Errorf("failed to marshal old node %#v for node %q: %v", oldNode, nodeName, err) @@ -221,6 +266,31 @@ func main() { notReadyTaint.Key = opts.NodeTaint + annotationsList := strings.Split(opts.DaemonSetAnnotation, ",") + labelsList := strings.Split(opts.DaemonSetLabel, ",") + for _, annot := range annotationsList { + kv := strings.Split(annot, ":") + annotationsMap[strings.TrimSpace(kv[0])] = "" + if len(kv) == 2 { + annotationsMap[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1]) + } + } + for _, lbl := range labelsList { + kv := strings.Split(lbl, ":") + labelsMap[strings.TrimSpace(kv[0])] = "" + if len(kv) == 2 { + labelsMap[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1]) + } + } + + for key := range annotationsMap { + annotationsKeys = append(annotationsKeys, key) + } + + for key := range labelsMap { + labelsKeys = append(labelsKeys, key) + } + setupLogging(opts.LogLevel) clientSet, err := getClientset() @@ -241,9 +311,37 @@ func main() { isHandling.Lock() defer isHandling.Unlock() - if _, ok := ds.Spec.Template.Annotations[opts.DaemonSetAnnotation]; !ok { - delete(dsList, ds.Name) - return + if filter(annotationsKeys, func(s string) bool { + _, ok := ds.Spec.Template.Annotations[s] + return ok + }) == "" { + if filter(labelsKeys, func(s string) bool { + _, ok := ds.Spec.Template.Labels[s] + return ok + }) == "" { + delete(dsList, ds.Name) + return + } + } + + if annot := filter(annotationsKeys, func(s string) bool { + _, ok := ds.Spec.Template.Annotations[s] + return ok + }); annot != "" { + if annotationsMap[annot] != "" && annotationsMap[annot] != ds.Spec.Template.Annotations[annot] { + delete(dsList, ds.Name) + return + } + } + + if lbl := filter(labelsKeys, func(s string) bool { + _, ok := ds.Spec.Template.Labels[s] + return ok + }); lbl != "" { + if labelsMap[lbl] != "" && labelsMap[lbl] != ds.Spec.Template.Labels[lbl] { + delete(dsList, ds.Name) + return + } } switch ops { @@ -255,6 +353,9 @@ func main() { delete(dsList, ds.Name) } logrus.Debugf("Number of required daemonsets is %v", len(dsList)) + for dsName := range dsList { + logrus.Debugf("Required daemonset name: %s", dsName) + } } podHandler := func(pod *core_v1.Pod, node *core_v1.Node, podStopChan chan struct{}) { @@ -278,7 +379,7 @@ func main() { } // check if all daemonset pods are running - finished, err := checkDSStatus(node, *opts) + finished, err := checkDSStatus(node) if err != nil { logrus.Errorf("Failed to sync daemonsets: %v", err) return @@ -337,12 +438,34 @@ func main() { for _, obj := range c.dsIndexer.List() { if ds, ok := obj.(*v1.DaemonSet); ok { - if _, ok := ds.Spec.Template.Annotations[opts.DaemonSetAnnotation]; ok { - dsList[ds.Name] = ds + if annot := filter(annotationsKeys, func(s string) bool { + _, ok := ds.Spec.Template.Annotations[s] + return ok + }); annot != "" { + if annotationsMap[annot] == "" { + dsList[ds.Name] = ds + } + if annotationsMap[annot] != "" && annotationsMap[annot] == ds.Spec.Template.Annotations[annot] { + dsList[ds.Name] = ds + } + } + if lbl := filter(labelsKeys, func(s string) bool { + _, ok := ds.Spec.Template.Labels[s] + return ok + }); lbl != "" { + if labelsMap[lbl] == "" { + dsList[ds.Name] = ds + } + if labelsMap[lbl] != "" && labelsMap[lbl] == ds.Spec.Template.Labels[lbl] { + dsList[ds.Name] = ds + } } } } logrus.Infof("Number of required daemonsets is %v", len(dsList)) + for dsName := range dsList { + logrus.Debugf("Required daemonset name: %s", dsName) + } http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintf(w, "OK\n")