controller.go

package mserver

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"time"
	
	"github.com/enfabrica/enkit/lib/knetwork/kdns"
	"github.com/enfabrica/enkit/lib/logger"
	mpb "github.com/enfabrica/enkit/machinist/rpc"
	"github.com/enfabrica/enkit/machinist/state"

	"github.com/miekg/dns"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type Controller struct {
	Log logger.Logger

	startUpFunc []func()

	allRecordsRefreshRate time.Duration

	State         *state.MachineController
	stateFile     string
	stateWriteTTL time.Duration

	dnsServer *kdns.DnsServer
	domains   []string
}

// Init is designed to run after all components have been started up before running itself as a server
func (en *Controller) Init() {
	for _, m := range en.State.Machines {
		en.addNodeToDns(m.Name, m.Ips, m.Tags)
	}
}

func (en *Controller) Nodes() []*state.Machine {
	en.State.Lock()
	defer en.State.Unlock()
	nodes := make([]*state.Machine, len(en.State.Machines))
	copy(nodes, en.State.Machines)
	return nodes
}

func (en *Controller) Download(*mpb.DownloadRequest, mpb.Controller_DownloadServer) error {
	return nil
}

func (en *Controller) Upload(mpb.Controller_UploadServer) error {
	return nil
}

func (en *Controller) HandlePing(stream mpb.Controller_PollServer, ping *mpb.ClientPing) error {
	return stream.Send(
		&mpb.PollResponse{
			Resp: &mpb.PollResponse_Pong{
				Pong: &mpb.ActionPong{
					Payload: ping.Payload,
				},
			},
		})
}

func (en *Controller) HandleRegister(stream mpb.Controller_PollServer, ping *mpb.ClientRegister) error {
	var parsedIps []net.IP
	for _, p := range ping.Ips {
		i := net.ParseIP(p)
		if i != nil {
			parsedIps = append(parsedIps, i)
		}
	}
	if len(parsedIps) == 0 {
		return errors.New("no valid ip sent")
	}
	newMachine := &state.Machine{
		Name: ping.Name,
		Ips:  parsedIps,
		Tags: ping.Tag,
	}
	if err := state.AddMachine(en.State, newMachine); err != nil {
		return status.Errorf(codes.AlreadyExists, err.Error())
	}
	en.addNodeToDns(ping.Name, newMachine.Ips, newMachine.Tags)
	return stream.Send(
		&mpb.PollResponse{
			Resp: &mpb.PollResponse_Result{
				Result: &mpb.ActionResult{},
			},
		})

}

func (en *Controller) Poll(stream mpb.Controller_PollServer) error {
	for {
		in, err := stream.Recv()
		if err != nil {
			return err
		}

		switch r := in.Req.(type) {
		case *mpb.PollRequest_Ping:
			en.HandlePing(stream, r.Ping)

		case *mpb.PollRequest_Register:
			if err = en.HandleRegister(stream, r.Register); err != nil {
				return err
			}
		}
	}
}

func (en *Controller) addNodeToDns(name string, ips []net.IP, tags []string) {
	for _, d := range en.dnsServer.Domains {
		dnsName := dns.CanonicalName(fmt.Sprintf("%s.%s", name, d))
		var recordTags []dns.RR
		for _, t := range tags {
			entry, err := dns.NewRR(fmt.Sprintf("%s %s %s", dnsName, "TXT", t))
			if err != nil {
				continue
			}
			recordTags = append(recordTags, entry)
		}
		for _, i := range ips {
			var recordType string
			if i.To4() != nil {
				recordType = "A"
			}
			if i.To16() != nil && recordType == "" {
				recordType = "AAAA"
			}
			if recordType != "" {
				entry, err := dns.NewRR(fmt.Sprintf("%s %s %s", dnsName, recordType, i.String()))
				if err != nil {
					continue
				}
				en.Log.Infof("Adding %s to the dns ControlPlane %s", dnsName, entry)
				en.dnsServer.SetEntry(dnsName, []dns.RR{entry})
				en.dnsServer.SetEntry(dnsName, recordTags)
			}
		}
	}
}

// ServeAllAndInfoRecords will continuously poll Nodes() and create multiple _all.<domain> records containing the ip addresses
// of all machines attached.
// It also serves the current state of the dns server via the _info record.
// TODO(adam): be able to pass in a wrapped ticker for testing intervals
func (en *Controller) ServeAllAndInfoRecords(killChannel chan struct{}, killChannelAck chan struct{}) {
	for {
		select {
		case <-time.After(en.allRecordsRefreshRate):
			ns := en.Nodes()
			for _, d := range en.dnsServer.Domains {
				dnsName := dns.CanonicalName(fmt.Sprintf("%s.%s", "_all", d))
				infoDnsName := dns.CanonicalName(fmt.Sprintf("%s.%s", "_info", d))
				var allDnsRecords []dns.RR
				var infoDnsRecords []dns.RR
				for _, v := range ns {
					for _, i := range v.Ips {
						rr, err := dns.NewRR(fmt.Sprintf("%s %s %s", dnsName, "A", i.String()))
						if err != nil {
							en.Log.Errorf("err: %v", err)
						}
						infoRR, err := dns.NewRR(fmt.Sprintf("%s %s { name: %s, ip: %s }", infoDnsName, "TXT", v.Name, i.String()))
						if err != nil {
							en.Log.Errorf("err: %v", err)
						}
						allDnsRecords = append(allDnsRecords, rr)
						infoDnsRecords = append(infoDnsRecords, infoRR)
					}
				}
				en.dnsServer.SetEntry(dnsName, allDnsRecords)
				en.dnsServer.SetEntry(infoDnsName, infoDnsRecords)
			}
		case <-killChannel:
			killChannelAck <- struct{}{}
			return
		}
	}
}

// scrapeConfigForHost returns a config object to instruct Prometheus to scrape
// this host. See:
// https://prometheus.io/docs/prometheus/latest/configuration/configuration/#http_sd_config.
// Set a "hostname" label on all metrics from this host.
func scrapeConfigForHost(hostname string, addrs []string) map[string]interface{} {
	return map[string]interface{} {
		"targets": addrs,
		"labels": map[string]string{
			"hostname": hostname,
		},
	}
}

// MetricsTargets is an HTTP handler that satisfies
// https://prometheus.io/docs/prometheus/latest/configuration/configuration/#http_sd_config.
func (en *Controller) MetricsTargets(w http.ResponseWriter, r *http.Request) {
	scrapeConfig := []map[string]interface{}{}
	for _, node := range en.Nodes() {
		var ips []string
		for _, ip := range node.Ips {
			ips = append(ips, ip.String())
		}
		scrapeConfig = append(scrapeConfig, scrapeConfigForHost(node.Name, ips))
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(scrapeConfig)
}

// WriteState writes state to the specified state file every 2 seconds. Will not exit or error out unless no statefile is
// provided.
func (en *Controller) WriteState() {
	if en.stateFile == "" {
		en.Log.Warnf("No path to state provided, state is fully in memory")
		return
	}
	for {
		<-time.After(en.stateWriteTTL)
		if err := state.WriteController(en.State, en.stateFile); err != nil {
			en.Log.Errorf("machinist: writing to state failed with err: %v", err)
		}
	}
}