node-labeler/main.go
okozachenko 977f62ba56 Fix getting node name
Change-Id: Id34db7684a132a3e5eb108b4d5e925714e3b1fec
2020-07-23 20:01:31 +03:00

149 lines
2.8 KiB
Go

package main
import (
"encoding/json"
"os"
"strings"
"time"
"github.com/avast/retry-go"
jsonpatch "github.com/evanphx/json-patch"
"github.com/gosimple/slug"
"github.com/jaypipes/ghw"
"github.com/prometheus/common/log"
"go.uber.org/zap"
v1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apimachinery/pkg/types"
"k8s.io/client-go/kubernetes"
"k8s.io/client-go/rest"
)
func main() {
// NOTE(mnaser): Due to the fact that we don't want to run ghw in
// a privileged context, we simply disable all warnings.
os.Setenv("GHW_DISABLE_WARNINGS", "1")
logger, _ := zap.NewProduction()
log := logger.Named("node-labeler")
nodeName, ok := os.LookupEnv("NODE")
if !ok {
log.Fatal("Environment variable is not defined",
zap.String("var", "NODE"),
)
}
log.Info("Starting service",
zap.String("node", nodeName),
)
config, err := rest.InClusterConfig()
if err != nil {
log.Fatal(err.Error())
}
clientset, err := kubernetes.NewForConfig(config)
if err != nil {
log.Fatal(err.Error())
}
for {
chassis, err := ghw.Chassis()
if err != nil {
log.Fatal(err.Error())
}
product, err := ghw.Product()
if err != nil {
log.Fatal(err.Error())
}
labels := map[string]string{
"node.vexxhost.com/vendor": slug.Make(chassis.Vendor),
"node.vexxhost.com/product": slug.Make(product.Name),
}
if CheckComputeNode(nodeName) {
labels["node-role.openstack.org"] = "compute"
}
var node *v1.Node
err = retry.Do(
func() error {
node, err = clientset.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
return err
},
)
if err != nil {
log.Fatal(err.Error())
}
for label, value := range labels {
err = addLabelToNode(clientset, node, label, value)
if err != nil {
log.Fatal(err.Error())
}
}
time.Sleep(600 * time.Second)
}
}
// checkComputeSubstring checks if the node name includes compute substring.
func CheckComputeNode(nodeName string) bool {
compute_tag := []string{
"kvm",
"compute",
}
for _, sub := range compute_tag {
if strings.Contains(nodeName, sub) {
return true
}
}
return false
}
func addLabelToNode(clientset *kubernetes.Clientset, node *v1.Node, key string, value string) error {
log.Info("Applying node label",
zap.String(key, value),
)
originalNode, err := json.Marshal(node)
if err != nil {
return err
}
node.ObjectMeta.Labels[key] = value
newNode, err := json.Marshal(node)
if err != nil {
return err
}
patch, err := jsonpatch.CreateMergePatch(originalNode, newNode)
if err != nil {
return err
}
log.Info("Patching Node resource",
zap.String("node", node.ObjectMeta.Name),
zap.String("patch", string(patch)),
)
err = retry.Do(
func() error {
_, err = clientset.CoreV1().Nodes().Patch(node.Name, types.MergePatchType, patch)
return err
},
)
if err != nil {
return err
}
return nil
}