github.com/cloud-foundations/dominator@v0.0.0-20221004181915-6e4fee580046/lib/mdb/mdbd/watchd.go (about)

     1  package mdbd
     2  
     3  import (
     4  	"bufio"
     5  	"encoding/gob"
     6  	"encoding/json"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  	"path"
    11  	"reflect"
    12  	"sort"
    13  	"time"
    14  
    15  	"github.com/Cloud-Foundations/Dominator/lib/fsutil"
    16  	jsonwriter "github.com/Cloud-Foundations/Dominator/lib/json"
    17  	"github.com/Cloud-Foundations/Dominator/lib/log"
    18  	"github.com/Cloud-Foundations/Dominator/lib/mdb"
    19  	"github.com/Cloud-Foundations/Dominator/lib/srpc"
    20  	"github.com/Cloud-Foundations/Dominator/proto/mdbserver"
    21  )
    22  
    23  func startMdbDaemon(mdbFileName string, logger log.Logger) <-chan *mdb.Mdb {
    24  	mdbChannel := make(chan *mdb.Mdb, 1)
    25  	if *mdbServerHostname != "" && *mdbServerPortNum > 0 {
    26  		go serverWatchDaemon(*mdbServerHostname, *mdbServerPortNum, mdbFileName,
    27  			mdbChannel, logger)
    28  	} else {
    29  		go fileWatchDaemon(mdbFileName, mdbChannel, logger)
    30  	}
    31  	return mdbChannel
    32  }
    33  
    34  type genericDecoder interface {
    35  	Decode(v interface{}) error
    36  }
    37  
    38  func fileWatchDaemon(mdbFileName string, mdbChannel chan<- *mdb.Mdb,
    39  	logger log.Logger) {
    40  	var lastMdb *mdb.Mdb
    41  	for readCloser := range fsutil.WatchFile(mdbFileName, logger) {
    42  		mdb := loadFile(readCloser, mdbFileName, logger)
    43  		readCloser.Close()
    44  		if mdb == nil {
    45  			continue
    46  		}
    47  		compareStartTime := time.Now()
    48  		if lastMdb == nil || !reflect.DeepEqual(lastMdb, mdb) {
    49  			if lastMdb != nil {
    50  				mdbCompareTimeDistribution.Add(time.Since(compareStartTime))
    51  			}
    52  			mdbChannel <- mdb
    53  			lastMdb = mdb
    54  		}
    55  	}
    56  }
    57  
    58  func serverWatchDaemon(mdbServerHostname string, mdbServerPortNum uint,
    59  	mdbFileName string, mdbChannel chan<- *mdb.Mdb, logger log.Logger) {
    60  	if file, err := os.Open(mdbFileName); err == nil {
    61  		fileMdb := loadFile(file, mdbFileName, logger)
    62  		file.Close()
    63  		if fileMdb != nil {
    64  			sort.Sort(fileMdb)
    65  			mdbChannel <- fileMdb
    66  		}
    67  	}
    68  	address := fmt.Sprintf("%s:%d", mdbServerHostname, mdbServerPortNum)
    69  	for ; ; time.Sleep(time.Second) {
    70  		client, err := srpc.DialHTTP("tcp", address, time.Second*15)
    71  		if err != nil {
    72  			logger.Println(err)
    73  			continue
    74  		}
    75  		conn, err := client.Call("MdbServer.GetMdbUpdates")
    76  		if err != nil {
    77  			logger.Println(err)
    78  			client.Close()
    79  			continue
    80  		}
    81  		lastMdb := &mdb.Mdb{}
    82  		for {
    83  			var mdbUpdate mdbserver.MdbUpdate
    84  			if err := conn.Decode(&mdbUpdate); err != nil {
    85  				logger.Println(err)
    86  				break
    87  			} else {
    88  				lastMdb = processUpdate(lastMdb, mdbUpdate)
    89  				sort.Sort(lastMdb)
    90  				mdbChannel <- lastMdb
    91  				if file, err := os.Create(mdbFileName + "~"); err != nil {
    92  					logger.Println(err)
    93  				} else {
    94  					writer := bufio.NewWriter(file)
    95  					var err error
    96  					if isGob(mdbFileName) {
    97  						encoder := gob.NewEncoder(writer)
    98  						err = encoder.Encode(lastMdb.Machines)
    99  					} else {
   100  						err = jsonwriter.WriteWithIndent(writer, "    ",
   101  							lastMdb.Machines)
   102  					}
   103  					if err != nil {
   104  						logger.Println(err)
   105  						os.Remove(mdbFileName + "~")
   106  					} else {
   107  						writer.Flush()
   108  						file.Close()
   109  						os.Rename(mdbFileName+"~", mdbFileName)
   110  					}
   111  				}
   112  			}
   113  		}
   114  		conn.Close()
   115  		client.Close()
   116  	}
   117  }
   118  
   119  func loadFile(reader io.Reader, filename string, logger log.Logger) *mdb.Mdb {
   120  	decoder := getDecoder(reader, filename)
   121  	var mdb mdb.Mdb
   122  	decodeStartTime := time.Now()
   123  	if err := decoder.Decode(&mdb.Machines); err != nil {
   124  		logger.Printf("Error decoding MDB data: %s\n", err)
   125  		return nil
   126  	}
   127  	sortStartTime := time.Now()
   128  	mdbDecodeTimeDistribution.Add(sortStartTime.Sub(decodeStartTime))
   129  	sort.Sort(&mdb)
   130  	mdbSortTimeDistribution.Add(time.Since(sortStartTime))
   131  	return &mdb
   132  }
   133  
   134  func isGob(filename string) bool {
   135  	switch path.Ext(filename) {
   136  	case ".gob":
   137  		return true
   138  	default:
   139  		return false
   140  	}
   141  }
   142  
   143  func getDecoder(reader io.Reader, filename string) genericDecoder {
   144  	if isGob(filename) {
   145  		return gob.NewDecoder(reader)
   146  	} else {
   147  		return json.NewDecoder(reader)
   148  	}
   149  }
   150  
   151  func processUpdate(oldMdb *mdb.Mdb, mdbUpdate mdbserver.MdbUpdate) *mdb.Mdb {
   152  	newMdb := &mdb.Mdb{}
   153  	if len(oldMdb.Machines) < 1 {
   154  		newMdb.Machines = mdbUpdate.MachinesToAdd
   155  		return newMdb
   156  	}
   157  	newMachines := make(map[string]mdb.Machine)
   158  	for _, machine := range oldMdb.Machines {
   159  		newMachines[machine.Hostname] = machine
   160  	}
   161  	for _, machine := range mdbUpdate.MachinesToAdd {
   162  		newMachines[machine.Hostname] = machine
   163  	}
   164  	for _, machine := range mdbUpdate.MachinesToUpdate {
   165  		newMachines[machine.Hostname] = machine
   166  	}
   167  	for _, name := range mdbUpdate.MachinesToDelete {
   168  		delete(newMachines, name)
   169  	}
   170  	newMdb.Machines = make([]mdb.Machine, 0, len(newMachines))
   171  	for _, machine := range newMachines {
   172  		newMdb.Machines = append(newMdb.Machines, machine)
   173  	}
   174  	return newMdb
   175  }