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