github.com/Cloud-Foundations/Dominator@v0.3.4/lib/filegen/connect.go (about)

     1  package filegen
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"io/ioutil"
     7  	"time"
     8  
     9  	"github.com/Cloud-Foundations/Dominator/lib/mdb"
    10  	"github.com/Cloud-Foundations/Dominator/lib/srpc"
    11  	proto "github.com/Cloud-Foundations/Dominator/proto/filegenerator"
    12  )
    13  
    14  func (t *rpcType) Connect(conn *srpc.Conn) error {
    15  	return t.manager.connect(conn) // Long-lived.
    16  }
    17  
    18  func (m *Manager) connect(conn *srpc.Conn) error {
    19  	defer conn.Flush()
    20  	clientChannel := make(chan *proto.ServerMessage, 4096)
    21  	m.rwMutex.Lock()
    22  	m.clients[clientChannel] = clientChannel
    23  	m.rwMutex.Unlock()
    24  	defer func() {
    25  		m.rwMutex.Lock()
    26  		delete(m.clients, clientChannel)
    27  		m.rwMutex.Unlock()
    28  	}()
    29  	closeNotifyChannel := make(chan struct{})
    30  	go m.handleClientRequests(conn, clientChannel, closeNotifyChannel)
    31  	for {
    32  		select {
    33  		case serverMessage := <-clientChannel:
    34  			if err := conn.Encode(serverMessage); err != nil {
    35  				m.logger.Printf("error encoding ServerMessage: %s\n", err)
    36  				return err
    37  			}
    38  			if len(clientChannel) < 1 {
    39  				if err := conn.Flush(); err != nil {
    40  					m.logger.Printf("error flushing: %s\n", err)
    41  					return err
    42  				}
    43  			}
    44  		case <-closeNotifyChannel:
    45  			return nil
    46  		}
    47  	}
    48  }
    49  
    50  func (m *Manager) handleClientRequests(decoder srpc.Decoder,
    51  	messageChan chan<- *proto.ServerMessage,
    52  	closeNotifyChannel chan<- struct{}) {
    53  	for {
    54  		if err := m.handleRequest(decoder, messageChan); err != nil {
    55  			if err != io.EOF {
    56  				m.logger.Println(err)
    57  			}
    58  			closeNotifyChannel <- struct{}{}
    59  		}
    60  	}
    61  }
    62  
    63  func (m *Manager) handleRequest(decoder srpc.Decoder,
    64  	messageChan chan<- *proto.ServerMessage) error {
    65  	var request proto.ClientRequest
    66  	if err := decoder.Decode(&request); err != nil {
    67  		if err == io.EOF {
    68  			return err
    69  		}
    70  		return errors.New("error decoding ClientRequest: " + err.Error())
    71  	}
    72  	serverMessage := new(proto.ServerMessage)
    73  	if request := request.YieldRequest; request != nil {
    74  		m.updateMachineData(request.Machine)
    75  		fileInfos := make([]proto.FileInfo, 0, len(request.Pathnames))
    76  		for _, pathname := range request.Pathnames {
    77  			if fileInfo, ok := m.computeFile(request.Machine, pathname); ok {
    78  				fileInfos = append(fileInfos, fileInfo)
    79  			}
    80  		}
    81  		serverMessage.YieldResponse = &proto.YieldResponse{
    82  			Hostname: request.Machine.Hostname,
    83  			Files:    fileInfos}
    84  	}
    85  	if request := request.GetObjectRequest; request != nil {
    86  		_, reader, err := m.objectServer.GetObject(request.Hash)
    87  		if err != nil {
    88  			return err
    89  		} else {
    90  			data, _ := ioutil.ReadAll(reader)
    91  			serverMessage.GetObjectResponse = &proto.GetObjectResponse{
    92  				Hash: request.Hash,
    93  				Data: data}
    94  			reader.Close()
    95  		}
    96  	}
    97  	messageChan <- serverMessage
    98  	return nil
    99  }
   100  
   101  func (m *Manager) updateMachineData(machine mdb.Machine) {
   102  	m.rwMutex.Lock()
   103  	defer m.rwMutex.Unlock()
   104  	if oldMachine, ok := m.machineData[machine.Hostname]; !ok {
   105  		m.machineData[machine.Hostname] = machine
   106  	} else if !oldMachine.Compare(machine) {
   107  		m.machineData[machine.Hostname] = machine
   108  		for _, pathMgr := range m.pathManagers {
   109  			delete(pathMgr.machineHashes, machine.Hostname)
   110  		}
   111  	}
   112  }
   113  
   114  func (m *Manager) computeFile(machine mdb.Machine, pathname string) (
   115  	proto.FileInfo, bool) {
   116  	fileInfo := proto.FileInfo{Pathname: pathname}
   117  	m.rwMutex.RLock()
   118  	pathMgr, ok := m.pathManagers[pathname]
   119  	if !ok {
   120  		m.rwMutex.RUnlock()
   121  		m.logger.Println("no generator for: " + pathname)
   122  		return fileInfo, false
   123  	}
   124  	if fi, ok := pathMgr.machineHashes[machine.Hostname]; ok {
   125  		if fi.validUntil.IsZero() || time.Now().Before(fi.validUntil) {
   126  			m.rwMutex.RUnlock()
   127  			fileInfo.Hash = fi.hash
   128  			fileInfo.Length = fi.length
   129  			return fileInfo, true
   130  		}
   131  	}
   132  	m.rwMutex.RUnlock()
   133  	hashVal, length, validUntil, err := pathMgr.generator.generate(machine,
   134  		m.logger)
   135  	if err != nil {
   136  		return fileInfo, false
   137  	}
   138  	fileInfo.Hash = hashVal
   139  	fileInfo.Length = length
   140  	m.rwMutex.Lock()
   141  	defer m.rwMutex.Unlock()
   142  	pathMgr.machineHashes[machine.Hostname] = expiringHash{
   143  		hashVal, length, validUntil}
   144  	m.scheduleTimer(pathname, machine.Hostname, validUntil)
   145  	return fileInfo, true
   146  }