depot/go/nix/nixbuild/coordinator.go

// SPDX-FileCopyrightText: 2023 Luke Granger-Brown <depot@lukegb.com>
//
// SPDX-License-Identifier: Apache-2.0

package nixbuild

import (
	"context"
	"sort"
	"sync"

	"git.lukegb.com/lukegb/depot/go/nix/nixdrv"
	"git.lukegb.com/lukegb/depot/go/nix/nixstore"
)

type Coordinator struct {
	cfg Config

	mu   sync.Mutex
	work map[string]*WorkItem

	at *nixstore.ActivityTracker
}

func NewCoordinator(cfg Config, at *nixstore.ActivityTracker) *Coordinator {
	return &Coordinator{
		cfg:  cfg,
		work: make(map[string]*WorkItem),
		at:   at,
	}
}

func (c *Coordinator) AddDerivationWork(ctx context.Context, drv *nixdrv.BasicDerivation, drvPath string) *WorkItem {
	var outputNames []string
	for o := range drv.Outputs {
		outputNames = append(outputNames, o)
	}
	sort.Strings(outputNames)

	path := drv.Outputs[outputNames[0]].Path

	c.mu.Lock()
	defer c.mu.Unlock()
	if i := c.work[path]; i != nil {
		return i
	}

	i := &WorkItem{
		doneCh: make(chan struct{}),
		path:   path,

		coord: c,
		at:    c.at,
	}
	c.work[path] = i
	go i.run(ctx)
	return i
}

func (c *Coordinator) AddWork(ctx context.Context, path string) *WorkItem {
	c.mu.Lock()
	defer c.mu.Unlock()

	if i := c.work[path]; i != nil {
		return i
	}

	i := &WorkItem{
		doneCh: make(chan struct{}),
		path:   path,

		coord: c,
		at:    c.at,
	}
	c.work[path] = i
	go i.run(ctx)
	return i
}