Skip to content

perf: Use per-output-base lifecycle locking for the Bazel Remote Output Service #47

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
EdSchouten merged 2 commits into from
Oct 8, 2026
Merged

perf: Use per-output-base lifecycle locking for the Bazel Remote Output Service #47

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
792ec0c
Serialize output-base lifecycle work without holding the global mutex
joeljeske Oct 7, 2026
ed49a31
Address review feedback on output-base lifecycle locking
joeljeske Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 78 additions & 8 deletions pkg/filesystem/virtual/bazel_output_service_directory.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,9 @@ type BazelOutputServiceDirectory struct {
outputBaseIDs map[path.Component]*outputPathState
buildIDs map[string]*outputPathState
outputPaths outputPathState

busyOutputBaseIDsLock sync.Mutex
busyOutputBaseIDs map[path.Component]<-chan struct{}
}

var (
Expand All @@ -100,20 +103,56 @@ func NewBazelOutputServiceDirectory(handleAllocator virtual.StatefulHandleAlloca

outputBaseIDs: map[path.Component]*outputPathState{},
buildIDs: map[string]*outputPathState{},

busyOutputBaseIDs: map[path.Component]<-chan struct{}{},
}
d.handle = handleAllocator.New().AsStatefulDirectory(d)
d.outputPaths.previous = &d.outputPaths
d.outputPaths.next = &d.outputPaths
return d
}

// acquireOutputBase serializes lifecycle calls, including calls for bases
// that do not have an output path yet. Waiting does not hold the registry lock.
func (d *BazelOutputServiceDirectory) acquireOutputBase(ctx context.Context, outputBaseID path.Component) (func(), error) {
for {
d.busyOutputBaseIDsLock.Lock()
if busy, ok := d.busyOutputBaseIDs[outputBaseID]; ok {
d.busyOutputBaseIDsLock.Unlock()
select {
case <-busy:
continue
case <-ctx.Done():
return nil, util.StatusFromContext(ctx)
}
}

busy := make(chan struct{})
d.busyOutputBaseIDs[outputBaseID] = busy
d.busyOutputBaseIDsLock.Unlock()

return func() {
d.busyOutputBaseIDsLock.Lock()
delete(d.busyOutputBaseIDs, outputBaseID)
close(busy)
d.busyOutputBaseIDsLock.Unlock()
}, nil
}
}

// Clean all build outputs associated with a single output base.
func (d *BazelOutputServiceDirectory) Clean(ctx context.Context, request *bazeloutputservice.CleanRequest) (*bazeloutputservice.CleanResponse, error) {
outputBaseID, ok := path.NewComponent(request.OutputBaseId)
if !ok {
return nil, status.Error(codes.InvalidArgument, "Output base ID is not a valid filename")
}

release, err := d.acquireOutputBase(ctx, outputBaseID)
if err != nil {
return nil, err
}
defer release()

d.lock.Lock()
outputPathState, ok := d.outputBaseIDs[outputBaseID]
d.lock.Unlock()
Expand Down Expand Up @@ -291,6 +330,12 @@ func (d *BazelOutputServiceDirectory) StartBuild(ctx context.Context, request *b
return nil, err
}

release, err := d.acquireOutputBase(ctx, outputBaseID)
if err != nil {
return nil, err
}
defer release()

d.lock.Lock()
state, ok := d.buildIDs[request.BuildId]
if !ok {
Expand All @@ -303,6 +348,7 @@ func (d *BazelOutputServiceDirectory) StartBuild(ctx context.Context, request *b
state.buildState = nil
}
} else {
d.lock.Unlock()
// No previous builds have been run for this
// output base. Create a new output path.
//
Expand All @@ -321,8 +367,11 @@ func (d *BazelOutputServiceDirectory) StartBuild(ctx context.Context, request *b
),
d.handleAllocator.New(),
)
rootDirectory := d.outputPathFactory.StartInitialBuild(outputBaseID, casFileFactory, digestFunction, errorLogger)

d.lock.Lock()
state = &outputPathState{
rootDirectory: d.outputPathFactory.StartInitialBuild(outputBaseID, casFileFactory, digestFunction, errorLogger),
rootDirectory: rootDirectory,
casFileFactory: casFileFactory,

previous: d.outputPaths.previous,
Expand All @@ -336,8 +385,8 @@ func (d *BazelOutputServiceDirectory) StartBuild(ctx context.Context, request *b
d.changeID++
}

// Allow StageArtifacts() and BatchStat() requests for
// the new build ID.
// Register the new build ID. StageArtifacts() and BatchStat()
// requests are admitted once this lifecycle call releases the base.
state.buildState = &buildState{
id: request.BuildId,
digestFunction: digestFunction,
Expand Down Expand Up @@ -367,9 +416,8 @@ func (d *BazelOutputServiceDirectory) StartBuild(ctx context.Context, request *b
}

// getOutputPathAndBuildState returns the state objects associated with
// a given build ID. This function is used by all gRPC methods that can
// only be invoked as part of a build (e.g., StageArtifacts(),
// BatchStat()).
// a given build ID, rejecting StageArtifacts() and BatchStat() requests
// while a lifecycle call owns the output base.
func (d *BazelOutputServiceDirectory) getOutputPathAndBuildState(buildID string) (*outputPathState, *buildState, error) {
d.lock.Lock()
defer d.lock.Unlock()
Expand All @@ -378,6 +426,12 @@ func (d *BazelOutputServiceDirectory) getOutputPathAndBuildState(buildID string)
if !ok {
return nil, nil, status.Error(codes.FailedPrecondition, "Build ID is not associated with any running build")
}
d.busyOutputBaseIDsLock.Lock()
_, busy := d.busyOutputBaseIDs[outputPathState.outputBaseID]
d.busyOutputBaseIDsLock.Unlock()
if busy {
return nil, nil, status.Error(codes.FailedPrecondition, "Output base is currently finalizing a build or getting cleaned")
}
return outputPathState, outputPathState.buildState, nil
}

Expand Down Expand Up @@ -665,16 +719,32 @@ func (d *BazelOutputServiceDirectory) FinalizeArtifacts(ctx context.Context, req
// build has completed. This prevents successive StageArtifacts() and
// BatchStat() calls from being processed.
func (d *BazelOutputServiceDirectory) FinalizeBuild(ctx context.Context, request *bazeloutputservice.FinalizeBuildRequest) (*bazeloutputservice.FinalizeBuildResponse, error) {
// Look up the output base ID associated with the build ID,
// so that we may mark that output base ID busy.
d.lock.Lock()
defer d.lock.Unlock()
outputPathState, ok := d.buildIDs[request.BuildId]
d.lock.Unlock()
if !ok {
return &bazeloutputservice.FinalizeBuildResponse{}, nil
}

release, err := d.acquireOutputBase(ctx, outputPathState.outputBaseID)
if err != nil {
return nil, err
}
defer release()

// Silently ignore requests for unknown build IDs. This ensures
// that FinalizeBuild() remains idempotent.
d.lock.Lock()
if outputPathState, ok := d.buildIDs[request.BuildId]; ok {
buildState := outputPathState.buildState
outputPathState.rootDirectory.FinalizeBuild(ctx, buildState.digestFunction)
delete(d.buildIDs, buildState.id)
outputPathState.buildState = nil
d.lock.Unlock()
outputPathState.rootDirectory.FinalizeBuild(ctx, buildState.digestFunction)
} else {
d.lock.Unlock()
}

@EdSchouten EdSchouten Oct 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What are your thoughts on writing this as follows instead?

	// Look up the output base ID associated with the build ID,
	// so that we may mark that output base ID busy.
	d.lock.Lock()
	outputPathState, ok := d.buildIDs[request.BuildId]
	d.lock.Unlock()
	if !ok {
		return &bazeloutputservice.FinalizeBuildResponse{}, nil
	}

	release, err := d.acquireOutputBase(ctx, outputPathState.outputBaseID)
	if err != nil {
		return nil, err
	}
	defer release()

	// Silently ignore requests for unknown build IDs. This ensures
	// that FinalizeBuild() remains idempotent.
	d.lock.Lock()
	if outputPathState, ok := d.buildIDs[request.BuildId]; ok {
		buildState := outputPathState.buildState	
		delete(d.buildIDs, buildState.id)
		outputPathState.buildState = nil
		d.lock.Unlock()
		outputPathState.rootDirectory.FinalizeBuild(ctx, buildState.digestFunction)
	} else {
		d.lock.Unlock()
	}

That does add a second map lookup, but it does keep the code more readable.

@joeljeske joeljeske Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, the tradeoff seems worth it. thanks! Updated.

return &bazeloutputservice.FinalizeBuildResponse{}, nil
}
Expand Down
195 changes: 195 additions & 0 deletions pkg/filesystem/virtual/bazel_output_service_directory_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ import (
"context"
"syscall"
"testing"
"testing/synctest"
"time"

remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
"github.com/buildbarn/bb-clientd/internal/mock"
Expand All @@ -14,6 +16,7 @@ import (
"github.com/buildbarn/bb-storage/pkg/digest"
"github.com/buildbarn/bb-storage/pkg/filesystem/path"
"github.com/buildbarn/bb-storage/pkg/testutil"
"github.com/buildbarn/bb-storage/pkg/util"
"github.com/stretchr/testify/require"

"google.golang.org/grpc/codes"
Expand Down Expand Up @@ -1315,3 +1318,195 @@ func TestBazelOutputServiceDirectoryVirtualReadDir(t *testing.T) {
)
})
}

func TestBazelOutputServiceDirectoryLifecycleContention(t *testing.T) {
for _, phase := range []string{"Restore", "Filter", "Finalize", "Clean", "UnopenedClean"} {
t.Run(phase, func(t *testing.T) {
synctest.Test(t, func(t *testing.T) {
ctrl, ctx := gomock.WithContext(context.Background(), t)
handleAllocator := mock.NewMockStatefulHandleAllocator(ctrl)
outputPathFactory := mock.NewMockOutputPathFactory(ctrl)
dHandleAllocation := mock.NewMockStatefulHandleAllocation(ctrl)
handleAllocator.EXPECT().New().Return(dHandleAllocation)
dHandle := mock.NewMockStatefulDirectoryHandle(ctrl)
dHandleAllocation.EXPECT().AsStatefulDirectory(gomock.Any()).Return(dHandle)
d := cd_vfs.NewBazelOutputServiceDirectory(
handleAllocator,
outputPathFactory,
mock.NewMockBlobAccess(ctrl),
mock.NewMockBlobAccess(ctrl),
mock.NewMockDirectoryFetcher(ctrl),
mock.NewMockSymlinkFactory(ctrl),
10000,
)

casFileHandleAllocation := mock.NewMockStatefulHandleAllocation(ctrl)
handleAllocator.EXPECT().New().Return(casFileHandleAllocation).AnyTimes()
casFileHandleAllocation.EXPECT().AsStatelessAllocator().Return(mock.NewMockStatelessHandleAllocator(ctrl)).AnyTimes()
digestFunction := digest.MustNewFunction("", remoteexecution.DigestFunction_SHA256)
args, err := anypb.New(&bazeloutputservicerev2.StartBuildArgs{
DigestFunction: remoteexecution.DigestFunction_SHA256,
})
require.NoError(t, err)
startBuild := func(ctx context.Context, baseID, buildID string) error {
_, err := d.StartBuild(ctx, &bazeloutputservice.StartBuildRequest{
OutputBaseId: baseID,
BuildId: buildID,
Args: args,
OutputPathPrefix: "/outputs",
})
return err
}

entered, resume := make(chan struct{}), make(chan struct{})
block := func() {
close(entered)
<-resume
}
outputPath := mock.NewMockOutputPath(ctrl)
if phase != "UnopenedClean" {
restore := outputPathFactory.EXPECT().StartInitialBuild(path.MustNewComponent("base"), gomock.Any(), digestFunction, gomock.Any()).Return(outputPath)
filter := outputPath.EXPECT().FilterChildren(gomock.Any())
if phase == "Restore" {
restore.Do(func(path.Component, re_vfs.CASFileFactory, digest.Function, util.ErrorLogger) { block() })
} else if phase == "Filter" {
filter.Do(func(re_vfs.ChildFilter) { block() })
} else {
require.NoError(t, startBuild(ctx, "base", "build"))
}
}

var operation func() error
switch phase {
case "Restore", "Filter":
operation = func() error { return startBuild(ctx, "base", "build") }
case "Finalize":
outputPath.EXPECT().FinalizeBuild(ctx, digestFunction).Do(func(context.Context, digest.Function) { block() })
operation = func() error {
_, err := d.FinalizeBuild(ctx, &bazeloutputservice.FinalizeBuildRequest{BuildId: "build"})
return err
}
case "Clean", "UnopenedClean":
if phase == "Clean" {
outputPath.EXPECT().RemoveAllChildren(true).Do(func(bool) { block() })
dHandle.EXPECT().NotifyRemoval(path.MustNewComponent("base"))
} else {
outputPathFactory.EXPECT().Clean(path.MustNewComponent("base")).Do(func(path.Component) { block() })
}
operation = func() error {
_, err := d.Clean(ctx, &bazeloutputservice.CleanRequest{OutputBaseId: "base"})
return err
}
}
done := make(chan error, 1)
go func() { done <- operation() }()
<-entered

// New stage/stat calls fail immediately while a lifecycle call
// owns the base. Finalization unregisters its build first.
if phase != "Restore" && phase != "UnopenedClean" {
expected := status.Error(codes.FailedPrecondition, "Output base is currently finalizing a build or getting cleaned")
if phase == "Finalize" {
expected = status.Error(codes.FailedPrecondition, "Build ID is not associated with any running build")
}
_, err = d.StageArtifacts(ctx, &bazeloutputservice.StageArtifactsRequest{BuildId: "build"})
testutil.RequireEqualStatus(t, expected, err)
_, err = d.BatchStat(ctx, &bazeloutputservice.BatchStatRequest{BuildId: "build"})
testutil.RequireEqualStatus(t, expected, err)
}

// Every RPC on another base can finish while this one is busy.
otherPath := mock.NewMockOutputPath(ctrl)
outputPathFactory.EXPECT().StartInitialBuild(path.MustNewComponent("other"), gomock.Any(), digestFunction, gomock.Any()).Return(otherPath)
otherPath.EXPECT().FilterChildren(gomock.Any())
require.NoError(t, startBuild(ctx, "other", "other-build"))
_, err = d.StageArtifacts(ctx, &bazeloutputservice.StageArtifactsRequest{BuildId: "other-build"})
require.NoError(t, err)
_, err = d.BatchStat(ctx, &bazeloutputservice.BatchStatRequest{BuildId: "other-build"})
require.NoError(t, err)
otherPath.EXPECT().FinalizeBuild(ctx, digestFunction)
_, err = d.FinalizeBuild(ctx, &bazeloutputservice.FinalizeBuildRequest{BuildId: "other-build"})
require.NoError(t, err)
otherPath.EXPECT().RemoveAllChildren(true)
dHandle.EXPECT().NotifyRemoval(path.MustNewComponent("other"))
_, err = d.Clean(ctx, &bazeloutputservice.CleanRequest{OutputBaseId: "other"})
require.NoError(t, err)

// Same-base lifecycle calls wait, and cancellation releases only
// the waiter, not the operation that owns the base.
waitCtx, cancel := context.WithCancel(ctx)
defer cancel()
canceled := make(chan error, 3)
go func() { canceled <- startBuild(waitCtx, "base", "canceled-build") }()
go func() {
_, err := d.Clean(waitCtx, &bazeloutputservice.CleanRequest{OutputBaseId: "base"})
canceled <- err
}()
waiters := 2
if phase == "Filter" || phase == "Clean" {
waiters++
go func() {
_, err := d.FinalizeBuild(waitCtx, &bazeloutputservice.FinalizeBuildRequest{BuildId: "build"})
canceled <- err
}()
}
synctest.Wait()
require.Len(t, canceled, 0)
cancel()
for range waiters {
require.Equal(t, codes.Canceled, status.Code(<-canceled))
}
require.Len(t, done, 0)

if phase == "Finalize" {
deadlineCtx, cancel := context.WithTimeout(ctx, time.Second)
defer cancel()
_, err := d.Clean(deadlineCtx, &bazeloutputservice.CleanRequest{OutputBaseId: "base"})
require.Equal(t, codes.DeadlineExceeded, status.Code(err))
}

// A finalizer already waiting on a removed build must not
// finalize the next build, even if that build wins the race.
// Once finalization has unregistered the build, a duplicate
// FinalizeBuild() succeeds without waiting.
var finalized chan error
if phase == "Finalize" || phase == "Clean" {
finalized = make(chan error, 1)
go func() {
_, err := d.FinalizeBuild(ctx, &bazeloutputservice.FinalizeBuildRequest{BuildId: "build"})
finalized <- err
}()
}
if phase == "Clean" || phase == "UnopenedClean" {
outputPathFactory.EXPECT().StartInitialBuild(path.MustNewComponent("base"), gomock.Any(), digestFunction, gomock.Any()).Return(outputPath)
}
outputPath.EXPECT().FilterChildren(gomock.Any())
started := make(chan error, 1)
go func() { started <- startBuild(ctx, "base", "next-build") }()
synctest.Wait()
require.Len(t, started, 0)
if phase == "Finalize" {
require.NoError(t, <-finalized)
finalized = nil
} else {
require.Len(t, finalized, 0, "FinalizeBuild must wait for the busy base")
}

close(resume)
require.NoError(t, <-done)
require.NoError(t, <-started)
if finalized != nil {
require.NoError(t, <-finalized)
}
_, err = d.BatchStat(ctx, &bazeloutputservice.BatchStatRequest{BuildId: "build"})
require.Equal(t, codes.FailedPrecondition, status.Code(err))
_, err = d.StageArtifacts(ctx, &bazeloutputservice.StageArtifactsRequest{BuildId: "next-build"})
require.NoError(t, err)
_, err = d.BatchStat(ctx, &bazeloutputservice.BatchStatRequest{BuildId: "next-build"})
require.NoError(t, err)
_, err = d.FinalizeBuild(ctx, &bazeloutputservice.FinalizeBuildRequest{BuildId: "build"})
require.NoError(t, err)
})
})
}
}
Loading