Add Workload DnsController

Signed-off-by: Tom <[email protected]>

Author: Flying-Tom

pkg/controller/controller.go

@@ -148,6 +148,7 @@ func (c *Controller) Start(stopCh <-chan struct{}) error {
 	c.client = NewXdsClient(c.mode, c.bpfAdsObj, c.bpfWorkloadObj, c.bpfConfig.EnableMonitoring, c.bpfConfig.EnableProfiling)
 
 	if c.client.WorkloadController != nil {
+		c.client.WorkloadController.StartDnsController(stopCh)
 		c.client.WorkloadController.Run(ctx)
 	} else {
 		c.client.AdsController.StartDnsController(stopCh)

pkg/controller/workload/dns.go

@@ -0,0 +1,230 @@
+package workload
+
+import (
+	"net"
+	"net/netip"
+	"sync"
+	"time"
+
+	"google.golang.org/protobuf/proto"
+	"kmesh.net/kmesh/api/v2/workloadapi"
+	"kmesh.net/kmesh/pkg/controller/workload/cache"
+	"kmesh.net/kmesh/pkg/dns"
+)
+
+type dnsController struct {
+	workloadsChan chan []*workloadapi.Workload
+	cache         cache.WorkloadCache
+	dnsResolver   *dns.DNSResolver
+	// store the copy of pendingResolveWorkload.
+	workloadCache map[string]*pendingResolveDomain
+	// store all pending hostnames in the workloads
+	pendingHostnames map[string][]string
+	sync.RWMutex
+}
+
+// pending resolve domain info of Dual-Engine Mode,
+// workload is used for create the apiworkload
+type pendingResolveDomain struct {
+	Workloads   []*workloadapi.Workload
+	RefreshRate time.Duration
+}
+
+func NewDnsController(cache cache.WorkloadCache) (*dnsController, error) {
+	resolver, err := dns.NewDNSResolver()
+	if err != nil {
+		return nil, err
+	}
+	return &dnsController{
+		workloadsChan:    make(chan []*workloadapi.Workload),
+		cache:            cache,
+		dnsResolver:      resolver,
+		workloadCache:    make(map[string]*pendingResolveDomain),
+		pendingHostnames: make(map[string][]string),
+	}, nil
+}
+
+func (r *dnsController) Run(stopCh <-chan struct{}) {
+	go r.dnsResolver.StartDnsResolver(stopCh)
+	go r.refreshWorker(stopCh)
+	go r.processWorkloads()
+	go func() {
+		<-stopCh
+		close(r.workloadsChan)
+	}()
+}
+
+func (r *dnsController) processWorkloads() {
+	for workloads := range r.workloadsChan {
+		r.processDomains(workloads)
+	}
+}
+
+func (r *dnsController) processDomains(workload []*workloadapi.Workload) {
+	domains := getPendingResolveDomain(workload)
+
+	// store all pending hostnames of clusters in pendingHostnames
+	for _, workload := range workload {
+		workloadName := workload.GetName()
+		hostname := workload.GetHostname()
+		r.pendingHostnames[workloadName] = []string{hostname}
+		r.workloadCache[hostname] = &pendingResolveDomain{
+			Workloads:   []*workloadapi.Workload{workload},
+			RefreshRate: 15 * time.Second,
+		}
+	}
+
+	// delete any scheduled re-resolve for domains we no longer care about
+	r.dnsResolver.RemoveUnwatchDomain(domains)
+
+	// update workloadCache with pendingResolveWorkload
+	for k, v := range domains {
+		addresses := r.dnsResolver.GetDNSAddresses(k)
+		if addresses != nil {
+			go r.updateWorkloads(v.(*pendingResolveDomain), k, addresses)
+		} else {
+			// Initialize the newly added hostname
+			// and add it to the dns queue to be resolved.
+			domainInfo := &dns.DomainInfo{
+				Domain:      k,
+				RefreshRate: v.(*pendingResolveDomain).RefreshRate,
+			}
+			r.dnsResolver.AddDomainInQueue(domainInfo, 0)
+		}
+	}
+}
+
+func (r *dnsController) refreshWorker(stop <-chan struct{}) {
+	for {
+		select {
+		case <-stop:
+			return
+		case domain := <-r.dnsResolver.DnsChan:
+			pendingDomain := r.getWorkloadsByDomain(domain)
+			// log.Infof("dnsController refreshWorker: domain %s, pendingDomain %v", domain, pendingDomain)
+			addrs := r.dnsResolver.GetDNSAddresses(domain)
+			maxRetry := 3
+			for range maxRetry {
+				if len(addrs) > 0 {
+					r.updateWorkloads(pendingDomain, domain, addrs)
+				}
+				time.Sleep(1 * time.Second)
+			}
+		}
+	}
+}
+
+func (r *dnsController) updateWorkloads(pendingDomain *pendingResolveDomain, domain string, addrs []string) {
+	isWorkerUpdate := false
+	if pendingDomain == nil || addrs == nil{
+		return
+	}
+	log.Infof("dnsController updateWorkloads: pendingDomain %v, domain %s, addrs %v", pendingDomain, domain, addrs)
+
+	for _, workload := range pendingDomain.Workloads {
+		if ready, newWorkload := r.overwriteDnsWorkload(workload, domain, addrs); ready {
+			// log.Infof("dnsController update cache for workload %s with addresses %v", workload.ResourceName(), addrs)
+			if r.cache.GetWorkloadByUid(workload.GetUid()) != nil {
+				r.cache.AddOrUpdateWorkload(newWorkload)
+				delete(r.workloadCache, domain)
+				isWorkerUpdate = true
+			}
+		}
+	}
+
+	if isWorkerUpdate {
+		// w.cache.Flush()
+	}
+}
+
+func (r *dnsController) overwriteDnsWorkload(workload *workloadapi.Workload, domain string, addrs []string) (bool, *workloadapi.Workload) {
+	ready := true
+	hostNames := r.pendingHostnames[workload.GetName()]
+	addressesOfHostname := make(map[string][]string)
+
+	for _, hostName := range hostNames {
+		// log.Infof("overwriteDnsWorkload: checking hostname %s for workload %s with domain %s", hostName, workload.ResourceName(), domain)
+		addresses := r.dnsResolver.GetDNSAddresses(hostName)
+		// There are hostnames in this Cluster that are not resolved.
+		if addresses != nil {
+			addressesOfHostname[hostName] = addresses
+		} else {
+			ready = false
+		}
+	}
+
+	if ready {
+		// log.Infof("overwriteDnsWorkload ready for workload %s with domain %s", workload.ResourceName(), domain)
+		newWorkload := cloneWorkload(workload)
+		for _, addr := range addrs {
+			if ip := net.ParseIP(addr); ip != nil && ip.To4() != nil {
+				newWorkload.Addresses = append(newWorkload.Addresses, netip.MustParseAddr(addr).AsSlice())
+			}
+		}
+
+		return ready, newWorkload
+	}
+
+	return ready, nil
+}
+
+func getPendingResolveDomain(workloads []*workloadapi.Workload) map[string]any {
+	domains := make(map[string]any)
+
+	for _, workload := range workloads {
+		hostname := workload.GetHostname()
+		if hostname == "" {
+			continue
+		}
+
+		if _, err := netip.ParseAddr(hostname); err == nil {
+			// This is an ip address
+			continue
+		}
+
+		// log.Infof("getPendingResolveDomain: processing workload %s with hostname %s", workload.ResourceName(), hostname)
+		if v, ok := domains[hostname]; ok {
+			v.(*pendingResolveDomain).Workloads = append(v.(*pendingResolveDomain).Workloads, workload)
+		} else {
+
+			domainWithRefreshRate := &pendingResolveDomain{
+				Workloads:   []*workloadapi.Workload{workload},
+				RefreshRate: 15 * time.Second,
+			}
+			domains[hostname] = domainWithRefreshRate
+		}
+	}
+
+	return domains
+}
+
+func (r *dnsController) newWorkloadCache() {
+	r.Lock()
+	defer r.Unlock()
+
+	if r.workloadCache != nil {
+		log.Debug("clean up dns workloads")
+		r.workloadCache = map[string]*pendingResolveDomain{}
+		return
+	}
+}
+
+func (r *dnsController) getWorkloadsByDomain(domain string) *pendingResolveDomain {
+	r.RLock()
+	defer r.RUnlock()
+
+	if r.workloadCache != nil {
+		if v, ok := r.workloadCache[domain]; ok {
+			return v
+		}
+	}
+	return nil
+}
+
+func cloneWorkload(workload *workloadapi.Workload) *workloadapi.Workload {
+	if workload == nil {
+		return nil
+	}
+	workloadCopy := proto.Clone(workload).(*workloadapi.Workload)
+	return workloadCopy
+}

pkg/controller/workload/workload_controller.go

@@ -40,6 +40,7 @@ var log = logger.NewLoggerScope("workload_controller")
 type Controller struct {
 	Stream                    discoveryv3.AggregatedDiscoveryService_DeltaAggregatedResourcesClient
 	Processor                 *Processor
+	dnsResolverController     *dnsController
 	Rbac                      *auth.Rbac
 	MetricController          *telemetry.MetricController
 	MapMetricController       *telemetry.MapMetricController
@@ -48,9 +49,17 @@ type Controller struct {
 }
 
 func NewController(bpfWorkload *bpfwl.BpfWorkload, enableMonitoring, enablePerfMonitor bool) *Controller {
+	processor := NewProcessor(bpfWorkload.SockConn.KmeshCgroupSockWorkloadObjects.KmeshCgroupSockWorkloadMaps)
+	dnsResolverController, err := NewDnsController(processor.WorkloadCache)
+	if err != nil {
+		log.Errorf("dns resolver of Dual-Engine mode create failed: %v", err)
+		return nil
+	}
+	processor.DnsResolverChan = dnsResolverController.workloadsChan
 	c := &Controller{
-		Processor:      NewProcessor(bpfWorkload.SockConn.KmeshCgroupSockWorkloadObjects.KmeshCgroupSockWorkloadMaps),
-		bpfWorkloadObj: bpfWorkload,
+		dnsResolverController: dnsResolverController,
+		Processor:             processor,
+		bpfWorkloadObj:        bpfWorkload,
 	}
 	// do some initialization when restart
 	// restore endpoint index, otherwise endpoint number can double
@@ -142,6 +151,7 @@ func (c *Controller) HandleWorkloadStream() error {
 		return fmt.Errorf("stream recv failed, %s", err)
 	}
 
+	// c.dnsResolverController.newWorkloadCache()
 	c.Processor.processWorkloadResponse(rspDelta, c.Rbac)
 
 	if err = c.Stream.Send(c.Processor.ack); err != nil {
@@ -188,3 +198,9 @@ func (c *Controller) SetConnectionMetricTrigger(enable bool) {
 func (c *Controller) GetConnectionMetricTrigger() bool {
 	return c.MetricController.EnableConnectionMetric.Load()
 }
+
+func (c *Controller) StartDnsController(stopCh <-chan struct{}) {
+	if c.dnsResolverController != nil {
+		c.dnsResolverController.Run(stopCh)
+	}
+}