cluster.go

package utils

import (
	"sync"

	"gitlab.com/tensorsecurity-rd/waf-console/pkg/generated/clientset/versioned"

	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

type KubeClient struct {
	*kubernetes.Clientset
	Versioned *versioned.Clientset
}

func NewKubeClient(apiServer string, caData []byte, clientCertData []byte, clientKeyData []byte, insecure bool) *KubeClient {
	versionedClient := versioned.NewForConfigOrDie(&rest.Config{
		Host: apiServer,
		TLSClientConfig: rest.TLSClientConfig{
			Insecure: insecure,
			CAData:   caData,
			CertData: clientCertData,
			KeyData:  clientKeyData,
		},
	})
	clientset := kubernetes.NewForConfigOrDie(&rest.Config{
		Host: apiServer,
		TLSClientConfig: rest.TLSClientConfig{
			Insecure: insecure,
			CAData:   caData,
			CertData: clientCertData,
			KeyData:  clientKeyData,
		},
	})

	return &KubeClient{
		Clientset: clientset,
		Versioned: versionedClient,
	}
}

type ClusterClientManager struct {
	clients map[string]*KubeClient
}

func NewClusterClientManager() *ClusterClientManager {
	return &ClusterClientManager{
		clients: make(map[string]*KubeClient),
	}
}

func (c *ClusterClientManager) GetClient(regionCode string) *KubeClient {
	return c.clients[regionCode]
}

func (c *ClusterClientManager) AddClient(regionCode string, client *KubeClient) {
	c.clients[regionCode] = client
}

func (c *ClusterClientManager) RemoveClient(regionCode string) {
	delete(c.clients, regionCode)
}

func (c *ClusterClientManager) GetAllClients() map[string]*KubeClient {
	return c.clients
}

func (c *ClusterClientManager) GetClientByRegionCode(regionCode string) *KubeClient {
	return c.clients[regionCode]
}

func (c *ClusterClientManager) ForEach(fn func(regionCode string, client *KubeClient)) {
	wg := sync.WaitGroup{}
	for regionCode, client := range c.clients {
		wg.Add(1)
		go func(regionCode string, client *KubeClient) {
			defer wg.Done()
			fn(regionCode, client)
		}(regionCode, client)
	}
	wg.Wait()
}