github.com/cloud-foundations/dominator@v0.0.0-20221004181915-6e4fee580046/cmd/mdbd/loadFleetManager.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/Cloud-Foundations/Dominator/lib/constants"
     9  	"github.com/Cloud-Foundations/Dominator/lib/log"
    10  	"github.com/Cloud-Foundations/Dominator/lib/mdb"
    11  	"github.com/Cloud-Foundations/Dominator/lib/srpc"
    12  	fm_proto "github.com/Cloud-Foundations/Dominator/proto/fleetmanager"
    13  	hyper_proto "github.com/Cloud-Foundations/Dominator/proto/hypervisor"
    14  )
    15  
    16  type fleetManagerGeneratorType struct {
    17  	fleetManager string
    18  	location     string
    19  	logger       log.DebugLogger
    20  	eventChannel chan<- struct{}
    21  	mutex        sync.Mutex
    22  	machines     map[string]*fm_proto.Machine
    23  	vms          map[string]*hyper_proto.VmInfo
    24  }
    25  
    26  func newFleetManagerGenerator(args []string,
    27  	logger log.DebugLogger) (generator, error) {
    28  	g := &fleetManagerGeneratorType{
    29  		fleetManager: fmt.Sprintf("%s:%d",
    30  			args[0], constants.FleetManagerPortNumber),
    31  		logger:   logger,
    32  		machines: make(map[string]*fm_proto.Machine),
    33  		vms:      make(map[string]*hyper_proto.VmInfo),
    34  	}
    35  	if len(args) > 1 {
    36  		g.location = args[1]
    37  	}
    38  	go g.daemon()
    39  	return g, nil
    40  }
    41  
    42  func (g *fleetManagerGeneratorType) daemon() {
    43  	for {
    44  		if err := g.getUpdates(g.fleetManager); err != nil {
    45  			g.logger.Println(err)
    46  			time.Sleep(time.Second)
    47  		}
    48  	}
    49  }
    50  
    51  func (g *fleetManagerGeneratorType) getUpdates(fleetManager string) error {
    52  	client, err := srpc.DialHTTP("tcp", g.fleetManager, 0)
    53  	if err != nil {
    54  		return err
    55  	}
    56  	defer client.Close()
    57  	conn, err := client.Call("FleetManager.GetUpdates")
    58  	if err != nil {
    59  		return err
    60  	}
    61  	defer conn.Close()
    62  	request := fm_proto.GetUpdatesRequest{Location: g.location}
    63  	if err := conn.Encode(request); err != nil {
    64  		return err
    65  	}
    66  	if err := conn.Flush(); err != nil {
    67  		return err
    68  	}
    69  	initialUpdate := true
    70  	for {
    71  		var update fm_proto.Update
    72  		if err := conn.Decode(&update); err != nil {
    73  			return err
    74  		}
    75  		g.update(update, initialUpdate)
    76  		initialUpdate = false
    77  		select {
    78  		case g.eventChannel <- struct{}{}:
    79  		default:
    80  		}
    81  	}
    82  }
    83  
    84  func (g *fleetManagerGeneratorType) Generate(unused_datacentre string,
    85  	logger log.Logger) (*mdb.Mdb, error) {
    86  	var newMdb mdb.Mdb
    87  	g.mutex.Lock()
    88  	defer g.mutex.Unlock()
    89  	for _, machine := range g.machines {
    90  		var ipAddr string
    91  		if len(machine.HostIpAddress) > 0 {
    92  			ipAddr = machine.HostIpAddress.String()
    93  		}
    94  		tags := machine.Tags
    95  		if tags == nil {
    96  			tags = emptyTags
    97  		}
    98  		_, disableUpdates := tags["DisableUpdates"]
    99  		newMdb.Machines = append(newMdb.Machines, mdb.Machine{
   100  			Hostname:       machine.Hostname,
   101  			IpAddress:      ipAddr,
   102  			RequiredImage:  tags["RequiredImage"],
   103  			PlannedImage:   tags["PlannedImage"],
   104  			DisableUpdates: disableUpdates,
   105  			Tags:           machine.Tags,
   106  		})
   107  	}
   108  	for ipAddr, vm := range g.vms {
   109  		if vm.State == hyper_proto.StateRunning {
   110  			tags := vm.Tags
   111  			if tags == nil {
   112  				tags = emptyTags
   113  			}
   114  			_, disableUpdates := tags["DisableUpdates"]
   115  			var ownerGroup string
   116  			if len(vm.OwnerGroups) > 0 {
   117  				ownerGroup = vm.OwnerGroups[0]
   118  			}
   119  			newMdb.Machines = append(newMdb.Machines, mdb.Machine{
   120  				Hostname:       ipAddr,
   121  				IpAddress:      ipAddr,
   122  				RequiredImage:  tags["RequiredImage"],
   123  				PlannedImage:   tags["PlannedImage"],
   124  				DisableUpdates: disableUpdates,
   125  				OwnerGroup:     ownerGroup,
   126  				Tags:           vm.Tags,
   127  			})
   128  		}
   129  	}
   130  	return &newMdb, nil
   131  }
   132  
   133  func (g *fleetManagerGeneratorType) RegisterEventChannel(
   134  	events chan<- struct{}) {
   135  	g.eventChannel = events
   136  }
   137  
   138  func (g *fleetManagerGeneratorType) update(update fm_proto.Update,
   139  	initialUpdate bool) {
   140  	machinesToDelete := make(map[string]struct{}, len(g.machines))
   141  	vmsToDelete := make(map[string]struct{}, len(g.vms))
   142  	if initialUpdate {
   143  		for hostname := range g.machines {
   144  			machinesToDelete[hostname] = struct{}{}
   145  		}
   146  		for ipAddr := range g.vms {
   147  			vmsToDelete[ipAddr] = struct{}{}
   148  		}
   149  	}
   150  	g.mutex.Lock()
   151  	defer g.mutex.Unlock()
   152  	for _, machine := range update.ChangedMachines {
   153  		g.machines[machine.Hostname] = machine
   154  		delete(machinesToDelete, machine.Hostname)
   155  	}
   156  	for _, hostname := range update.DeletedMachines {
   157  		delete(g.machines, hostname)
   158  	}
   159  	for hostname := range machinesToDelete {
   160  		delete(g.machines, hostname)
   161  	}
   162  	for ipAddr, vm := range update.ChangedVMs {
   163  		g.vms[ipAddr] = vm
   164  		delete(vmsToDelete, ipAddr)
   165  	}
   166  	for _, ipAddr := range update.DeletedVMs {
   167  		delete(g.vms, ipAddr)
   168  	}
   169  	for ipAddr := range vmsToDelete {
   170  		delete(g.vms, ipAddr)
   171  	}
   172  }