merging_processor.go

package main

import (
	"bytes"
	"sort"
	"strings"
	"sync"
	"text/template"
)

type MergingProcessor struct {
	queues        map[string][]*Block
	done          map[string]bool
	headerColumns []string
	Emitter
	sync.Mutex
}

func MergingProcessorCreate(id string) *MergingProcessor {
	return &MergingProcessor{
		queues:        make(map[string][]*Block, 0),
		done:          make(map[string]bool, 0),
		headerColumns: make([]string, 0),
		Emitter: Emitter{
			id: id,
		},
	}
}

func (m *MergingProcessor) Subscribed(id string) {
	m.queues[id] = make([]*Block, 0)
}

func (m *MergingProcessor) Header(id string, header Header) error {
	m.Lock()
	defer m.Unlock()

	if len(header.columns) > 0 {
		m.headerColumns = append(m.headerColumns, header.columns)
	}
	return nil
}

func (m *MergingProcessor) sendHeader() error {
	tmpl, err := template.New("merging").Parse(`#
# IP to {{.Id}} mapping
#
# Merges
{{range $id := .Ids -}}
#   {{$id}}
{{end -}}
#
# Generated by haproxy-mapper at {{.Now}}
#
`,
	)
	if err != nil {
		return err
	}

	ids := make([]string, 0)
	for id, _ := range m.queues {
		ids = append(ids, id)
	}

	sort.Strings(ids)

	context := struct {
		Id  string
		Ids []string
		Now string
	}{
		Id:  m.id,
		Ids: ids,
		Now: NowIso8610(),
	}

	var rendered bytes.Buffer
	if err = tmpl.Execute(&rendered, context); err != nil {
		return err
	}

	// Make sure order is consistent and not just when we received them
	sort.Strings(m.headerColumns)
	header := Header{
		general: rendered.String(),
		columns: strings.Join(m.headerColumns, ""),
	}
	for _, receiver := range m.receivers {
		if err := receiver.Header(m.id, header); err != nil {
			return err
		}
	}

	// clear out our header columns now
	m.headerColumns = make([]string, 0)

	return nil
}

func (m *MergingProcessor) Receive(id string, block *Block) error {
	if len(m.headerColumns) > 0 {
		// We have header columns (on our first receive) we're need to emit the
		// header which will clean them and we'll be done with it.
		if err := m.sendHeader(); err != nil {
			return err
		}
	}
	// Assuming each source (id) is single threaded, thus no locking
	m.queues[id] = append(m.queues[id], block)

	// We can get in checkEmit from multiple emitters and pull from any queue
	// that's ready thus we need locking. This will also provide safety or
	// amything downstream of us so they don't have to worry about it.
	m.Lock()
	defer m.Unlock()

	return m.checkEmit()
}

func (m *MergingProcessor) checkEmit() error {
	// This will run until it finds an empty queue that's not done or fails to find anything
	for {
		var min *Block
		var minId string
		for id, queue := range m.queues {
			if len(queue) == 0 {
				if m.done[id] {
					// It's empty and done so that's Ok. We should get rid of
					// it and keep checking others
					delete(m.queues, id)
					continue
				}
				// We have an empty queue and therefore can't know what to emit
				// next, maybe next time through
				return nil
			}

			// If we don't yet have a min or this one is less than what we have
			if min == nil || queue[0].Less(min) {
				min = queue[0]
				minId = id
			}
		}
		if min == nil {
			return nil
		}

		// Pop the first item off
		m.queues[minId] = m.queues[minId][1:len(m.queues[minId])]

		// Emit it (min)
		err := m.Emit(min)
		if err != nil {
			return err
		}
	}
}

func (m *MergingProcessor) Done(id string) error {
	m.done[id] = true

	// We can get in here from multiple emitters and thus need safety while in
	// checkEmit and when we look to see if we're done, to avoid multiple Done
	// events.
	m.Lock()
	defer m.Unlock()

	err := m.checkEmit()
	if err != nil {
		return err
	}

	if len(m.queues) == 0 {
		return m.Emitter.Done()
	}

	return nil
}