Source File
	synchronize.go

Belonging Package
	github.com/polarsignals/frostdb/query/physicalplan

package physicalplan

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"

	"github.com/apache/arrow-go/v18/arrow"
)

// Synchronizer is used to combine the results of multiple parallel streams
// into a single stream concurrent stream. It also forms a barrier on the
// finishers, by waiting to call next plan's finish until all previous parallel
// stages have finished.
type Synchronizer struct {
	next    PhysicalPlan
	nextMtx sync.Mutex
	running *atomic.Int64
	open    *atomic.Int64
}

func Synchronize(concurrency int) *Synchronizer {
	running := &atomic.Int64{}
	running.Add(int64(concurrency))
	open := &atomic.Int64{}
	open.Add(int64(concurrency))
	return &Synchronizer{running: running, open: open}
}

func (m *Synchronizer) Callback(ctx context.Context, r arrow.Record) error {
	// multiple threads can emit the results to the next step, but they will do
	// it synchronously
	m.nextMtx.Lock()
	defer m.nextMtx.Unlock()

	err := m.next.Callback(ctx, r)
	if err != nil {
		return err
	}
	return nil
}

func (m *Synchronizer) Finish(ctx context.Context) error {
	running := m.running.Add(-1)
	if running < 0 {
		return errors.New("too many Synchronizer Finish calls")
	}
	if running > 0 {
		return nil
	}
	return m.next.Finish(ctx)
}

func (m *Synchronizer) SetNext(next PhysicalPlan) {
	m.next = next
}

func (m *Synchronizer) SetNextPlan(nextPlan PhysicalPlan) {
	m.next = nextPlan
}

func (m *Synchronizer) Draw() *Diagram {
	return &Diagram{Details: "Synchronizer", Child: m.next.Draw()}
}

func (m *Synchronizer) Close() {
	open := m.open.Add(-1)
	if open < 0 {
		panic("too many Synchronizer Close calls")
	}
	if open > 0 {
		return
	}
	m.next.Close()
}

The pages are generated with Golds v0.8.2. (GOOS=linux GOARCH=amd64)
Golds is a Go 101 project developed by Tapir Liu.
PR and bug reports are welcome and can be submitted to the issue list.
Please follow @zigo_101 (reachable from the left QR code) to get the latest news of Golds.