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

     1  package filegen
     2  
     3  import (
     4  	"bytes"
     5  	"sort"
     6  	"time"
     7  
     8  	"github.com/Cloud-Foundations/Dominator/lib/hash"
     9  	"github.com/Cloud-Foundations/Dominator/lib/log"
    10  	"github.com/Cloud-Foundations/Dominator/lib/mdb"
    11  	"github.com/Cloud-Foundations/Dominator/lib/objectserver/memory"
    12  	proto "github.com/Cloud-Foundations/Dominator/proto/filegenerator"
    13  )
    14  
    15  type hashGenerator interface {
    16  	generate(machine mdb.Machine, logger log.Logger) (
    17  		hashVal hash.Hash, length uint64, validUntil time.Time, err error)
    18  }
    19  
    20  type hashGeneratorWrapper struct {
    21  	dataGenerator FileGenerator
    22  	objectServer  *memory.ObjectServer
    23  }
    24  
    25  func (m *Manager) registerDataGeneratorForPath(pathname string,
    26  	gen FileGenerator) chan<- string {
    27  	hashGenerator := &hashGeneratorWrapper{gen, m.objectServer}
    28  	return m.registerHashGeneratorForPath(pathname, hashGenerator)
    29  }
    30  
    31  func (m *Manager) registerHashGeneratorForPath(pathname string,
    32  	gen hashGenerator) chan<- string {
    33  	if _, ok := m.pathManagers[pathname]; ok {
    34  		panic(pathname + " already registered")
    35  	}
    36  	notifyChan := make(chan string, 1)
    37  	pathMgr := &pathManager{
    38  		generator:     gen,
    39  		machineHashes: make(map[string]expiringHash)}
    40  	m.rwMutex.Lock()
    41  	m.pathManagers[pathname] = pathMgr
    42  	m.rwMutex.Unlock()
    43  	go m.processPathDataInvalidations(pathname, notifyChan)
    44  	return notifyChan
    45  }
    46  
    47  func (m *Manager) processPathDataInvalidations(pathname string,
    48  	machineNameChannel <-chan string) {
    49  	m.rwMutex.RLock()
    50  	pathMgr := m.pathManagers[pathname]
    51  	m.rwMutex.RUnlock()
    52  	for machineName := range machineNameChannel {
    53  		pathMgr.rwMutex.Lock()
    54  		if machineName == "" {
    55  			for _, mdbData := range m.machineData {
    56  				hashVal, length, validUntil, err := pathMgr.generator.generate(
    57  					mdbData, m.logger)
    58  				if err != nil {
    59  					continue
    60  				}
    61  				pathMgr.machineHashes[mdbData.Hostname] = expiringHash{
    62  					hashVal, length, validUntil}
    63  				files := make([]proto.FileInfo, 1)
    64  				files[0].Pathname = pathname
    65  				files[0].Hash = hashVal
    66  				files[0].Length = length
    67  				yieldResponse := &proto.YieldResponse{mdbData.Hostname, files}
    68  				message := &proto.ServerMessage{YieldResponse: yieldResponse}
    69  				for _, clientChannel := range m.clients {
    70  					clientChannel <- message
    71  				}
    72  				m.scheduleTimer(pathname, mdbData.Hostname, validUntil)
    73  			}
    74  		} else {
    75  			hashVal, length, validUntil, err := pathMgr.generator.generate(
    76  				m.machineData[machineName], m.logger)
    77  			if err != nil {
    78  				continue
    79  			}
    80  			pathMgr.machineHashes[machineName] = expiringHash{
    81  				hashVal, length, validUntil}
    82  			files := make([]proto.FileInfo, 1)
    83  			files[0].Pathname = pathname
    84  			files[0].Hash = hashVal
    85  			files[0].Length = length
    86  			yieldResponse := &proto.YieldResponse{machineName, files}
    87  			message := &proto.ServerMessage{YieldResponse: yieldResponse}
    88  			for _, clientChannel := range m.clients {
    89  				clientChannel <- message
    90  			}
    91  			m.scheduleTimer(pathname, machineName, validUntil)
    92  		}
    93  		pathMgr.rwMutex.Unlock()
    94  	}
    95  }
    96  
    97  func (m *Manager) scheduleTimer(pathname string, hostname string,
    98  	validUntil time.Time) {
    99  	if validUntil.IsZero() || time.Now().After(validUntil) {
   100  		return // No expiration or already expired.
   101  	}
   102  	pathMgr := m.pathManagers[pathname]
   103  	time.AfterFunc(validUntil.Sub(time.Now()), func() {
   104  		pathMgr.rwMutex.Lock()
   105  		defer pathMgr.rwMutex.Unlock()
   106  		mdbData, ok := m.machineData[hostname]
   107  		if !ok {
   108  			return
   109  		}
   110  		hashVal, length, validUntil, err := pathMgr.generator.generate(
   111  			mdbData, m.logger)
   112  		if err != nil {
   113  			return
   114  		}
   115  		pathMgr.machineHashes[hostname] = expiringHash{
   116  			hashVal, length, validUntil}
   117  		files := make([]proto.FileInfo, 1)
   118  		files[0].Pathname = pathname
   119  		files[0].Hash = hashVal
   120  		files[0].Length = length
   121  		yieldResponse := &proto.YieldResponse{mdbData.Hostname, files}
   122  		message := &proto.ServerMessage{YieldResponse: yieldResponse}
   123  		for _, clientChannel := range m.clients {
   124  			clientChannel <- message
   125  		}
   126  		m.scheduleTimer(pathname, mdbData.Hostname, validUntil)
   127  	})
   128  }
   129  
   130  func (m *Manager) getRegisteredPaths() []string {
   131  	pathnames := make([]string, 0, len(m.pathManagers))
   132  	for pathname := range m.pathManagers {
   133  		pathnames = append(pathnames, pathname)
   134  	}
   135  	sort.Strings(pathnames)
   136  	return pathnames
   137  }
   138  
   139  func (g *hashGeneratorWrapper) generate(machine mdb.Machine,
   140  	logger log.Logger) (
   141  	hash.Hash, uint64, time.Time, error) {
   142  	data, validUntil, err := g.dataGenerator.Generate(machine, logger)
   143  	if err != nil {
   144  		return hash.Hash{}, 0, time.Time{}, err
   145  	}
   146  	length := uint64(len(data))
   147  	hashVal, _, err := g.objectServer.AddObject(bytes.NewReader(data), length,
   148  		nil)
   149  	if err != nil {
   150  		return hash.Hash{}, 0, time.Time{}, err
   151  	}
   152  	return hashVal, length, validUntil, nil
   153  }