imposm3/parser/pbf/process.go

package pbf

import (
	"runtime"
	"sync"

	"github.com/omniscale/imposm3/element"
)

type parser struct {
	pbf       *Pbf
	coords    chan []element.Node
	nodes     chan []element.Node
	ways      chan []element.Way
	relations chan []element.Relation
	nParser   int
	wg        sync.WaitGroup
	waySync   *barrier
	relSync   *barrier
}

func NewParser(pbf *Pbf, coords chan []element.Node, nodes chan []element.Node, ways chan []element.Way, relations chan []element.Relation) *parser {
	return &parser{
		pbf:       pbf,
		coords:    coords,
		nodes:     nodes,
		ways:      ways,
		relations: relations,
		nParser:   runtime.NumCPU(),
		wg:        sync.WaitGroup{},
	}
}

func (p *parser) Parse() {
	blocks := p.pbf.BlockPositions()
	for i := 0; i < p.nParser; i++ {
		p.wg.Add(1)
		go func() {
			for block := range blocks {
				p.parseBlock(block)
			}
			if p.waySync != nil {
				p.waySync.doneWait()
			}
			if p.relSync != nil {
				p.relSync.doneWait()
			}
			p.wg.Done()
		}()
	}
	p.wg.Wait()

	if p.nodes != nil {
		close(p.nodes)
	}
	if p.coords != nil {
		close(p.coords)
	}
	if p.ways != nil {
		close(p.ways)
	}
	if p.relations != nil {
		close(p.relations)
	}
}

// FinishedCoords registers a single function that gets called when all
// nodes and coords are parsed. The callback should block until it is
// safe to continue with parsing of all ways.
// This only works when the PBF file is ordered by type (nodes before ways before relations).
func (p *parser) FinishedCoords(cb func()) {
	p.waySync = newBarrier(cb)
	p.waySync.add(p.nParser)
}

// FinishedWays registers a single function that gets called when all
// nodes and coords are parsed. The callback should block until it is
// safe to continue with parsing of all ways.
// This only works when the PBF file is ordered by type (nodes before ways before relations).
func (p *parser) FinishedWays(cb func()) {
	p.relSync = newBarrier(cb)
	p.relSync.add(p.nParser)
}

func (p *parser) parseBlock(pos block) {
	block := readPrimitiveBlock(pos)
	stringtable := newStringTable(block.GetStringtable())

	for _, group := range block.Primitivegroup {
		dense := group.GetDense()
		if dense != nil {
			parsedCoords, parsedNodes := readDenseNodes(dense, block, stringtable)
			if len(parsedCoords) > 0 && p.coords != nil {
				p.coords <- parsedCoords
			}
			if len(parsedNodes) > 0 && p.nodes != nil {
				p.nodes <- parsedNodes
			}
		}
		parsedCoords, parsedNodes := readNodes(group.Nodes, block, stringtable)
		if len(parsedCoords) > 0 && p.coords != nil {
			p.coords <- parsedCoords
		}
		if len(parsedNodes) > 0 && p.nodes != nil {
			p.nodes <- parsedNodes
		}
		parsedWays := readWays(group.Ways, block, stringtable)
		if len(parsedWays) > 0 && p.ways != nil {
			if p.waySync != nil {
				p.waySync.doneWait()
			}
			p.ways <- parsedWays
		}
		parsedRelations := readRelations(group.Relations, block, stringtable)
		if len(parsedRelations) > 0 && p.relations != nil {
			if p.waySync != nil {
				p.waySync.doneWait()
			}
			if p.relSync != nil {
				p.relSync.doneWait()
			}
			p.relations <- parsedRelations
		}
	}
}

// barrier is a struct to synchronize multiple goroutines.
// Works similar to a WaitGroup. Except:
// Calls callback function once all goroutines called doneWait().
// doneWait() blocks until the callback returns. doneWait() does not
// block after all goroutines were blocked once.
type barrier struct {
	synced     bool
	wg         sync.WaitGroup
	once       sync.Once
	callbackWg sync.WaitGroup
	callback   func()
}

func newBarrier(callback func()) *barrier {
	s := &barrier{callback: callback}
	s.callbackWg.Add(1)
	return s
}

func (s *barrier) add(delta int) {
	s.wg.Add(delta)
}

func (s *barrier) doneWait() {
	if s.synced {
		return
	}
	s.wg.Done()
	s.wg.Wait()
	s.once.Do(s.call)
	s.callbackWg.Wait()
}

func (s *barrier) call() {
	s.callback()
	s.synced = true
	s.callbackWg.Done()
}