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

     1  package main
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/Cloud-Foundations/Dominator/lib/log"
     9  	"github.com/Cloud-Foundations/Dominator/lib/mdb"
    10  	"github.com/Cloud-Foundations/Dominator/lib/srpc"
    11  	"github.com/Cloud-Foundations/Dominator/proto/mdbserver"
    12  )
    13  
    14  type rpcType struct {
    15  	currentMdb *mdb.Mdb
    16  	logger     log.Logger
    17  	rwMutex    sync.RWMutex
    18  	// Protected by lock.
    19  	updateChannels map[*srpc.Conn]chan<- mdbserver.MdbUpdate
    20  }
    21  
    22  func startRpcd(logger log.Logger) *rpcType {
    23  	rpcObj := &rpcType{
    24  		logger:         logger,
    25  		updateChannels: make(map[*srpc.Conn]chan<- mdbserver.MdbUpdate),
    26  	}
    27  	srpc.RegisterName("MdbServer", rpcObj)
    28  	return rpcObj
    29  }
    30  
    31  func (t *rpcType) GetMdbUpdates(conn *srpc.Conn) error {
    32  	updateChannel := make(chan mdbserver.MdbUpdate, 10)
    33  	t.rwMutex.Lock()
    34  	t.updateChannels[conn] = updateChannel
    35  	t.rwMutex.Unlock()
    36  	defer func() {
    37  		close(updateChannel)
    38  		t.rwMutex.Lock()
    39  		delete(t.updateChannels, conn)
    40  		t.rwMutex.Unlock()
    41  	}()
    42  	if t.currentMdb != nil {
    43  		mdbUpdate := mdbserver.MdbUpdate{MachinesToAdd: t.currentMdb.Machines}
    44  		if err := conn.Encode(mdbUpdate); err != nil {
    45  			return err
    46  		}
    47  		if err := conn.Flush(); err != nil {
    48  			return err
    49  		}
    50  	}
    51  	closeChannel := conn.GetCloseNotifier()
    52  	for {
    53  		var err error
    54  		select {
    55  		case mdbUpdate := <-updateChannel:
    56  			if isEmptyUpdate(mdbUpdate) {
    57  				t.logger.Printf("Queue for: %s is filling up: dropping client")
    58  				return errors.New("update queue too full")
    59  			}
    60  			if err = conn.Encode(mdbUpdate); err != nil {
    61  				break
    62  			}
    63  			if err = conn.Flush(); err != nil {
    64  				break
    65  			}
    66  		case <-closeChannel:
    67  			break
    68  		}
    69  		if err != nil {
    70  			if err != io.EOF {
    71  				t.logger.Println(err)
    72  				return err
    73  			} else {
    74  				return nil
    75  			}
    76  		}
    77  	}
    78  }
    79  
    80  func (t *rpcType) pushUpdateToAll(old, new *mdb.Mdb) {
    81  	t.currentMdb = new
    82  	updateChannels := t.getUpdateChannels()
    83  	if len(updateChannels) < 1 {
    84  		return
    85  	}
    86  	mdbUpdate := mdbserver.MdbUpdate{}
    87  	if old == nil {
    88  		old = &mdb.Mdb{}
    89  	}
    90  	oldMachines := make(map[string]mdb.Machine, len(old.Machines))
    91  	for _, machine := range old.Machines {
    92  		oldMachines[machine.Hostname] = machine
    93  	}
    94  	for _, newMachine := range new.Machines {
    95  		if oldMachine, ok := oldMachines[newMachine.Hostname]; ok {
    96  			if !newMachine.Compare(oldMachine) {
    97  				mdbUpdate.MachinesToUpdate = append(mdbUpdate.MachinesToUpdate,
    98  					newMachine)
    99  			}
   100  		} else {
   101  			mdbUpdate.MachinesToAdd = append(mdbUpdate.MachinesToAdd,
   102  				newMachine)
   103  		}
   104  	}
   105  	for _, machine := range new.Machines {
   106  		delete(oldMachines, machine.Hostname)
   107  	}
   108  	for name := range oldMachines {
   109  		mdbUpdate.MachinesToDelete = append(mdbUpdate.MachinesToDelete, name)
   110  	}
   111  	if isEmptyUpdate(mdbUpdate) {
   112  		t.logger.Println("Ignoring empty update")
   113  		return
   114  	}
   115  	for _, channel := range updateChannels {
   116  		sendUpdate(channel, mdbUpdate)
   117  	}
   118  }
   119  
   120  func (t *rpcType) getUpdateChannels() []chan<- mdbserver.MdbUpdate {
   121  	t.rwMutex.RLock()
   122  	defer t.rwMutex.RUnlock()
   123  	channels := make([]chan<- mdbserver.MdbUpdate, 0, len(t.updateChannels))
   124  	for _, channel := range t.updateChannels {
   125  		channels = append(channels, channel)
   126  	}
   127  	return channels
   128  }
   129  
   130  func isEmptyUpdate(mdbUpdate mdbserver.MdbUpdate) bool {
   131  	if len(mdbUpdate.MachinesToAdd) > 0 {
   132  		return false
   133  	}
   134  	if len(mdbUpdate.MachinesToUpdate) > 0 {
   135  		return false
   136  	}
   137  	if len(mdbUpdate.MachinesToDelete) > 0 {
   138  		return false
   139  	}
   140  	return true
   141  }
   142  
   143  func sendUpdate(channel chan<- mdbserver.MdbUpdate,
   144  	mdbUpdate mdbserver.MdbUpdate) {
   145  	defer func() { recover() }()
   146  	if cap(channel)-len(channel) < 2 {
   147  		// Not enough room for an update and a possible "too much" message next
   148  		// time around: send a "too much" message now.
   149  		channel <- mdbserver.MdbUpdate{}
   150  		return
   151  	}
   152  	channel <- mdbUpdate
   153  }