github.com/cloud-foundations/dominator@v0.0.0-20221004181915-6e4fee580046/cmd/filegen-client/main.go (about)

     1  package main
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"time"
     9  
    10  	"github.com/Cloud-Foundations/Dominator/lib/filegen/client"
    11  	"github.com/Cloud-Foundations/Dominator/lib/flags/loadflags"
    12  	"github.com/Cloud-Foundations/Dominator/lib/log/cmdlogger"
    13  	"github.com/Cloud-Foundations/Dominator/lib/mdb"
    14  	"github.com/Cloud-Foundations/Dominator/lib/mdb/mdbd"
    15  	"github.com/Cloud-Foundations/Dominator/lib/objectserver/memory"
    16  	"github.com/Cloud-Foundations/Dominator/lib/srpc/setupclient"
    17  	proto "github.com/Cloud-Foundations/Dominator/proto/filegenerator"
    18  )
    19  
    20  var (
    21  	benchmark = flag.Bool("benchmark", false,
    22  		"If true, perform benchmark timing")
    23  	debug = flag.Bool("debug", false,
    24  		"If true, show debugging output")
    25  	mdbFile = flag.String("mdbFile", "/var/lib/Dominator/mdb",
    26  		"File to read MDB data from (default format is JSON)")
    27  
    28  	numMachines int
    29  )
    30  
    31  func printUsage() {
    32  	fmt.Fprintln(os.Stderr,
    33  		"Usage: filegen-client [flags...] pathname source")
    34  	fmt.Fprintln(os.Stderr, "Common flags:")
    35  	flag.PrintDefaults()
    36  }
    37  
    38  func showMdb(mdb *mdb.Mdb) {
    39  	mdb.DebugWrite(os.Stdout)
    40  	fmt.Println()
    41  }
    42  
    43  type messageType struct {
    44  	hostname  string
    45  	fileInfos []proto.FileInfo
    46  }
    47  
    48  func benchmarkMessageHandler(messageChannel <-chan messageType) {
    49  	numMessages := 0
    50  	startTime := time.Now()
    51  	for message := range messageChannel {
    52  		numMessages++
    53  		if numMessages == numMachines {
    54  			fmt.Printf("Time taken: %s\n", time.Since(startTime))
    55  		} else if numMessages > numMachines {
    56  			fmt.Printf("Extra message for machine: %s\n", message.hostname)
    57  		}
    58  	}
    59  }
    60  
    61  func displayMessageHandler(messageChannel <-chan messageType,
    62  	objSrv *memory.ObjectServer) {
    63  	for message := range messageChannel {
    64  		fmt.Printf("For machine: %s:\n", message.hostname)
    65  		for _, fileInfo := range message.fileInfos {
    66  			fmt.Printf("  pathname: %s\n    hash=%x\n    contents:\n",
    67  				fileInfo.Pathname, fileInfo.Hash)
    68  			if _, reader, err := objSrv.GetObject(fileInfo.Hash); err != nil {
    69  				fmt.Println(err)
    70  			} else {
    71  				io.Copy(os.Stdout, reader)
    72  				fmt.Println("\n-----------------------------------------------")
    73  			}
    74  		}
    75  	}
    76  }
    77  
    78  func handleUpdates(hostname string, updateChannel <-chan []proto.FileInfo,
    79  	messageChannel chan<- messageType) {
    80  	for fileInfos := range updateChannel {
    81  		messageChannel <- messageType{hostname, fileInfos}
    82  	}
    83  }
    84  
    85  func main() {
    86  	if err := loadflags.LoadForCli("filegen-client"); err != nil {
    87  		fmt.Fprintln(os.Stderr, err)
    88  		os.Exit(1)
    89  	}
    90  	flag.Usage = printUsage
    91  	flag.Parse()
    92  	if flag.NArg() != 2 {
    93  		printUsage()
    94  		os.Exit(2)
    95  	}
    96  	if err := setupclient.SetupTls(true); err != nil {
    97  		fmt.Fprintln(os.Stderr, err)
    98  		os.Exit(1)
    99  	}
   100  	objectServer := memory.NewObjectServer()
   101  	logger := cmdlogger.New()
   102  	manager := client.New(objectServer, logger)
   103  	mdbChannel := mdbd.StartMdbDaemon(*mdbFile, logger)
   104  	machines := make(map[string]struct{})
   105  	computedFiles := make([]client.ComputedFile, 1)
   106  	computedFiles[0].Pathname = flag.Arg(0)
   107  	computedFiles[0].Source = flag.Arg(1)
   108  	messageChannel := make(chan messageType, 1)
   109  	if *benchmark {
   110  		go benchmarkMessageHandler(messageChannel)
   111  	} else {
   112  		go displayMessageHandler(messageChannel, objectServer)
   113  	}
   114  	for {
   115  		select {
   116  		case mdb := <-mdbChannel:
   117  			if *debug {
   118  				showMdb(mdb)
   119  			}
   120  			numMachines = len(mdb.Machines)
   121  			machinesToDelete := make(map[string]struct{}, len(machines))
   122  			for hostname := range machines {
   123  				machinesToDelete[hostname] = struct{}{}
   124  			}
   125  			for _, mdbEntry := range mdb.Machines {
   126  				delete(machinesToDelete, mdbEntry.Hostname)
   127  				machine := client.Machine{mdbEntry, computedFiles}
   128  				if _, ok := machines[mdbEntry.Hostname]; !ok {
   129  					machines[mdbEntry.Hostname] = struct{}{}
   130  					go handleUpdates(mdbEntry.Hostname, manager.Add(machine, 1),
   131  						messageChannel)
   132  				} else {
   133  					manager.Update(client.Machine{mdbEntry, computedFiles})
   134  				}
   135  			}
   136  			for hostname := range machinesToDelete {
   137  				manager.Remove(hostname)
   138  				delete(machines, hostname)
   139  			}
   140  		}
   141  	}
   142  }