Skip to content

Fix acquire timeout release lock #108

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

Open
jessicaochen wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix acquire timeout release lock #108

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
9f128b2
support releasing lock if user times our on aquire
jessicaochen Jul 15, 2026
d1195c5
removed timing from tests
jessicaochen Jul 15, 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
28 changes: 28 additions & 0 deletions pkg/accelerator-orchestrator/server/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ func (s *Server) Acquire(ctx context.Context, req *pb.AcquireRequest) (*pb.Acqui
select {
case <-ctx.Done():
slog.InfoContext(ctx, "Acquire context cancelled", "error", ctx.Err())
s.cancelLockRequest(ctx, groupID, jobID)
return nil, status.FromContextError(ctx.Err()).Err()
case <-ticker.C:
resp, err, done := s.checkAcquire(ctx, groupID, jobID, startTime)
Expand All @@ -99,6 +100,33 @@ func (s *Server) Acquire(ctx context.Context, req *pb.AcquireRequest) (*pb.Acqui
}
}

// cancelLockRequest undoes the lock request made by Acquire when the caller
// stops waiting, so the group is not left locked (or the job queued) for a
// caller that believes the acquire failed.
func (s *Server) cancelLockRequest(ctx context.Context, groupID, jobID string) {
// The request context is already cancelled; detach so the store updates can proceed.
ctx = context.WithoutCancel(ctx)

// Re-read group to get the latest status and spec from the store
group, err := s.groupStore.Get(ctx, groupID)
if err != nil {
slog.ErrorContext(ctx, "Failed to get group to cancel lock request", "error", err)
return
}

released, err := group.Spec().CancelLockRequest(ctx, jobID)
if err != nil {
slog.ErrorContext(ctx, "Failed to cancel lock request", "error", err)
return
}
Comment on lines +103 to +121

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win

Detached work needs a timeout

context.WithoutCancel removes both cancellation and deadline, so groupStore.Get and CancelLockRequest can block indefinitely on a slow backing store while holding the group mutex. Wrap the detached context in WithTimeout before calling into the store.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/accelerator-orchestrator/server/server.go` around lines 103 - 121, Add a
bounded timeout to the detached context in Server.cancelLockRequest before
calling groupStore.Get and group.Spec().CancelLockRequest. Derive the timeout
context from context.WithoutCancel(ctx), use the established cancellation/defer
pattern, and ensure both store operations receive the timed context so cleanup
cannot block indefinitely.

if released {
slog.InfoContext(ctx, "Released lock held by cancelled acquire")
if s.ctrl != nil {
s.ctrl.EnqueueWork(groupID)
}
}
Comment on lines +117 to +127

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== Files mentioning CancelLockRequest / lockingJob / fault ==\n'
rg -n --hidden --glob '!**/.git/**' 'CancelLockRequest|lockingJob|isGroupFaulted|faulted|lease|heartbeat|reconcile|reconciler' pkg

printf '\n== Candidate files ==\n'
git ls-files 'pkg/**' | rg 'accelerator-orchestrator|group\.go|server\.go|recon|fault|lease|heartbeat'

printf '\n== Outline of server.go and likely group file(s) ==\n'
ast-grep outline pkg/accelerator-orchestrator/server/server.go --view expanded || true
fd -a 'group.go' pkg || true

Repository: llm-d-incubation/llm-d-rl-time-slicing

Length of output: 20710

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== store/group.go relevant slice ==\n'
sed -n '150,320p' pkg/accelerator-orchestrator/store/group.go | cat -n

printf '\n== controller/controller.go relevant slice around reconcile and group state ==\n'
sed -n '200,520p' pkg/accelerator-orchestrator/controller/controller.go | cat -n

printf '\n== server/server.go fault handling slice ==\n'
sed -n '130,190p' pkg/accelerator-orchestrator/server/server.go | cat -n

Repository: llm-d-incubation/llm-d-rl-time-slicing

Length of output: 21953

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== all unlock / CancelLockRequest call sites ==\n'
rg -n --hidden --glob '!**/.git/**' 'CancelLockRequest\(|\.unlock\(|Yield\(' pkg/accelerator-orchestrator

printf '\n== lock store implementations ==\n'
sed -n '1,220p' pkg/accelerator-orchestrator/store/configmap_lockstore.go | cat -n
printf '\n---\n'
sed -n '1,140p' pkg/accelerator-orchestrator/store/lock_memstore.go | cat -n

printf '\n== server Acquire/cancel path ==\n'
sed -n '59,130p' pkg/accelerator-orchestrator/server/server.go | cat -n

Repository: llm-d-incubation/llm-d-rl-time-slicing

Length of output: 12468

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== controller queue / enqueue sources ==\n'
rg -n --hidden --glob '!**/.git/**' 'EnqueueWork|Run\(|NewController|workqueue|ticker|resync|periodic|reconcileAll|List\(\)' pkg/accelerator-orchestrator/controller

printf '\n== controller top-level structure ==\n'
sed -n '1,220p' pkg/accelerator-orchestrator/controller/controller.go | cat -n

Repository: llm-d-incubation/llm-d-rl-time-slicing

Length of output: 24740

Retry CancelLockRequest before giving up. The periodic resync only requeues groups; it does not retry the unlock, so a transient lockStore.Unlock error can leave lockingJob set and block the queue with no automatic recovery path.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/accelerator-orchestrator/server/server.go` around lines 117 - 127, Update
the cancellation flow around CancelLockRequest to retry transient unlock
failures before returning, ensuring a failed lockStore.Unlock does not leave
lockingJob set without an automatic recovery path. Reuse the existing retry or
resync mechanism where available, and preserve the current logging and
EnqueueWork behavior after a successful release.

}

func (s *Server) defaultCheckAcquire(
ctx context.Context,
groupID, jobID string,
Expand Down
156 changes: 111 additions & 45 deletions pkg/accelerator-orchestrator/server/server_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ import (
"github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/accelerator-orchestrator/store"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/status"
"google.golang.org/grpc/test/bufconn"
)
Expand Down Expand Up @@ -139,14 +138,15 @@ func (m *MockGroupLockStore) Unlock(ctx context.Context, jobID string) error {

func TestServer_Acquire_Whitebox(t *testing.T) {
tests := []struct {
name string
setupStores func(t *testing.T, ctx context.Context) (GroupStore, JobStore)
groupID string
jobID string
expectedCode codes.Code
verify func(t *testing.T, resp *pb.AcquireResponse, err error)
expectEnqueue bool
hook func(t *testing.T, srv *Server, gs GroupStore, js JobStore, cancel context.CancelFunc)
name string
setupStores func(t *testing.T, ctx context.Context) (GroupStore, JobStore)
groupID string
jobID string
expectedCode codes.Code
verify func(t *testing.T, resp *pb.AcquireResponse, err error)
wantEnqueues int
hook func(t *testing.T, srv *Server, gs GroupStore, js JobStore, cancel context.CancelFunc)
verifyAfter func(t *testing.T, gs GroupStore)
}{
{
name: "block when it has not yet reconciled but another job requesting lock",
Expand All @@ -168,10 +168,10 @@ func TestServer_Acquire_Whitebox(t *testing.T) {
g.Status().SetLoadedJob("job-1")
return gs, store.NewJobStore()
},
groupID: "group-1",
jobID: "job-1", // job-1 tries to acquire, should block because lock is for job-2
expectedCode: codes.Canceled, // We expect Canceled because we will cancel it manually
expectEnqueue: true,
groupID: "group-1",
jobID: "job-1", // job-1 tries to acquire, should block because lock is for job-2
expectedCode: codes.Canceled, // We expect Canceled because we will cancel it manually
wantEnqueues: 1,
hook: func(t *testing.T, srv *Server, gs GroupStore, js JobStore, cancel context.CancelFunc) {
t.Helper()
tickCalled := make(chan struct{}, 1)
Expand All @@ -192,6 +192,77 @@ func TestServer_Acquire_Whitebox(t *testing.T) {
cancel()
}()
},
verifyAfter: func(t *testing.T, gs GroupStore) {
t.Helper()
g, err := gs.Get(context.Background(), "group-1")
if err != nil {
t.Fatalf("failed to get group: %v", err)
}
// Cancelled acquire must remove job-1 from the waiting queue
// without touching job-2's lock.
if g.Spec().GetWaitingJobQueue().Exists("job-1") {
t.Errorf("expected job-1 to be removed from waiting queue")
}
if got := g.Spec().LockingJob(); got != "job-2" {
t.Errorf("LockingJob() = %q, want %q", got, "job-2")
}
},
},
{
name: "release lock when acquire is cancelled after lock granted",
setupStores: func(t *testing.T, ctx context.Context) (GroupStore, JobStore) {
t.Helper()
lockStore := store.NewMemLockStore()
// job-1 already holds the lock (promoted while its Acquire was
// waiting), but its context has not been loaded yet, so the
// acquire keeps blocking.
if err := lockStore.Lock(ctx, "group-1", "job-1"); err != nil {
t.Fatalf("failed to lock: %v", err)
}
gs := store.NewGroupStore(lockStore)
g, _, err := gs.GetOrCreate(ctx, "group-1")
if err != nil {
t.Fatalf("failed to create group: %v", err)
}
g.Status().SetLoadedJob("job-2")
return gs, store.NewJobStore()
},
groupID: "group-1",
jobID: "job-1",
expectedCode: codes.Canceled,
// One enqueue from the lock request, one from releasing the lock
// so the controller can promote the next waiter.
wantEnqueues: 2,
hook: func(t *testing.T, srv *Server, gs GroupStore, js JobStore, cancel context.CancelFunc) {
t.Helper()
tickCalled := make(chan struct{}, 1)
origCheck := srv.checkAcquire
srv.checkAcquire = func(ctx context.Context, groupID, jobID string, startTime time.Time) (*pb.AcquireResponse, error, bool) {
resp, err, done := origCheck(ctx, groupID, jobID, startTime)
select {
case tickCalled <- struct{}{}:
default:
}
return resp, err, done
}
// Wait for 5 ticks to be sure it is blocked, then cancel
go func() {
for i := 0; i < 5; i++ {
<-tickCalled
}
cancel()
}()
},
verifyAfter: func(t *testing.T, gs GroupStore) {
t.Helper()
g, err := gs.Get(context.Background(), "group-1")
if err != nil {
t.Fatalf("failed to get group: %v", err)
}
if got := g.Spec().LockingJob(); got != "" {
t.Errorf("LockingJob() = %q, want empty (lock released)", got)
}
},
},
{
name: "succeed when job loads later",
Expand Down Expand Up @@ -222,10 +293,10 @@ func TestServer_Acquire_Whitebox(t *testing.T) {
}
return gs, store.NewJobStore()
},
groupID: "group-1",
jobID: "job-2", // job-2 tries to acquire, should succeed after job-2 is loaded
expectedCode: codes.OK,
expectEnqueue: true,
groupID: "group-1",
jobID: "job-2", // job-2 tries to acquire, should succeed after job-2 is loaded
expectedCode: codes.OK,
wantEnqueues: 1,
hook: func(t *testing.T, srv *Server, gs GroupStore, js JobStore, cancel context.CancelFunc) {
t.Helper()
mGS, ok := gs.(*MockGroupStore)
Expand Down Expand Up @@ -274,40 +345,24 @@ func TestServer_Acquire_Whitebox(t *testing.T) {
defer cancel()
gs, js := tc.setupStores(t, serverCtx)

srv, mq, cleanup := InitGRPCServer(gs, js)
defer cleanup()
mq := &MockWorkQueue{}
ctrl := controller.NewController(nil, nil, mq, nil, nil)
srv := NewServer(ctrl, gs, js)
srv.acquirePollInterval = 1 * time.Millisecond

if tc.hook != nil {
tc.hook(t, srv, gs, js, cancel)
}

conn, err := grpc.NewClient(
"passthrough:///bufnet",
grpc.WithContextDialer(BufDialer),
grpc.WithTransportCredentials(insecure.NewCredentials()),
)
if err != nil {
t.Fatalf("Failed to dial bufnet: %v", err)
}
defer conn.Close()
client := pb.NewAcceleratorOrchestratorServiceClient(conn)

resp, err := client.Acquire(clientCtx, &pb.AcquireRequest{
// Call the handler directly (not through gRPC) so that when it
// returns, all of its side effects — including lock-request
// cleanup on cancellation — have completed and can be asserted
// synchronously.
resp, err := srv.Acquire(clientCtx, &pb.AcquireRequest{
GroupId: tc.groupID,
JobId: tc.jobID,
})

added := mq.GetAdded()
if tc.expectEnqueue {
if len(added) != 1 || added[0] != tc.groupID {
t.Errorf("expected group %s to be enqueued, got %v", tc.groupID, added)
}
} else {
if len(added) != 0 {
t.Errorf("expected no enqueue, got %v", added)
}
}

if tc.expectedCode != codes.OK {
if err == nil {
t.Fatalf("Expected error, got nil")
Expand All @@ -319,11 +374,22 @@ func TestServer_Acquire_Whitebox(t *testing.T) {
if st.Code() != tc.expectedCode {
t.Errorf("Expected code %v, got %v", tc.expectedCode, st.Code())
}
return
} else if tc.verify != nil {
tc.verify(t, resp, err)
}

if tc.verify != nil {
tc.verify(t, resp, err)
if tc.verifyAfter != nil {
tc.verifyAfter(t, gs)
}

added := mq.GetAdded()
if len(added) != tc.wantEnqueues {
t.Errorf("expected %d enqueues, got %v", tc.wantEnqueues, added)
}
for _, id := range added {
if id != tc.groupID {
t.Errorf("expected only group %s to be enqueued, got %v", tc.groupID, added)
}
}
})
}
Expand Down
21 changes: 21 additions & 0 deletions pkg/accelerator-orchestrator/store/group.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,27 @@ func (s *GroupSpec) Yield(ctx context.Context, jobID string) error {
// to yield their job if there is a temporary issue locking for next job.
}

// CancelLockRequest undoes a previous RequestLock for the given job when the
// caller stops waiting for the lock (e.g. the acquire was cancelled or timed out).
// If the job is still waiting in the queue it is removed. If the job has already
// been promoted to be the locking job, the lock is released so the group is not
// left locked for a job that believes it failed to acquire.
// Returns true if the lock was released, meaning the next waiter can be promoted.
func (s *GroupSpec) CancelLockRequest(ctx context.Context, jobID string) (bool, error) {
s.mu.Lock()
defer s.mu.Unlock()

if s.lockingJob == jobID {
if err := s.unlock(ctx, jobID); err != nil {
return false, err
}
return true, nil
}

s.queue.Remove(jobID)
return false, nil
}

// RequestLock requests a lock for the given job.
// If the job already holds the lock, it returns immediately.
// Otherwise, it enqueues the job in the waiting queue.
Expand Down
80 changes: 80 additions & 0 deletions pkg/accelerator-orchestrator/store/group_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,86 @@ func TestGroupSpec_Yield(t *testing.T) {
}
}

func TestGroupSpec_CancelLockRequest(t *testing.T) {
tests := []struct {
name string
initialLock string
queuedJobs []string
cancelJob string
wantReleased bool
wantLock string
wantQueue []string
}{
{
name: "cancel by lock holder releases lock",
initialLock: "job-1",
cancelJob: "job-1",
wantReleased: true,
wantLock: "",
},
{
name: "cancel by queued job removes it from queue",
initialLock: "job-1",
queuedJobs: []string{"job-2", "job-3"},
cancelJob: "job-2",
wantReleased: false,
wantLock: "job-1",
wantQueue: []string{"job-3"},
},
{
name: "cancel by unknown job is a no-op",
initialLock: "job-1",
queuedJobs: []string{"job-2"},
cancelJob: "job-4",
wantReleased: false,
wantLock: "job-1",
wantQueue: []string{"job-2"},
},
}

for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
ctx := context.Background()
lockStore := store.NewMemLockStore()
if tc.initialLock != "" {
if err := lockStore.Lock(ctx, "group-1", tc.initialLock); err != nil {
t.Fatalf("failed to lock: %v", err)
}
}
wrapped := store.NewGroupLockStoreWrapper(lockStore, "group-1")
group, err := store.NewGroup(ctx, "group-1", wrapped)
if err != nil {
t.Fatalf("failed to create group: %v", err)
}
spec := group.Spec()
for _, jobID := range tc.queuedJobs {
spec.RequestLock(jobID)
}

released, err := spec.CancelLockRequest(ctx, tc.cancelJob)
if err != nil {
t.Fatalf("CancelLockRequest() error = %v", err)
}
if released != tc.wantReleased {
t.Errorf("CancelLockRequest() released = %v, want %v", released, tc.wantReleased)
}
if spec.LockingJob() != tc.wantLock {
t.Errorf("LockingJob() = %q, want %q", spec.LockingJob(), tc.wantLock)
}

queue := spec.GetWaitingJobQueue()
if queue.Len() != len(tc.wantQueue) {
t.Fatalf("queue len = %d, want %d", queue.Len(), len(tc.wantQueue))
}
for i, waiting := range queue.List() {
if waiting.JobID != tc.wantQueue[i] {
t.Errorf("queue[%d] = %q, want %q", i, waiting.JobID, tc.wantQueue[i])
}
}
})
}
}

func TestGroupSpec_TryPromote(t *testing.T) {
tests := []struct {
name string
Expand Down
15 changes: 15 additions & 0 deletions pkg/accelerator-orchestrator/store/waiting_job_queue.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,21 @@ func (q *WaitingJobQueue) Dequeue() (string, bool) {
return job.JobID, true
}

// Remove removes the given job from the queue regardless of its position.
// Returns true if the job was removed, false if it was not in the queue.
func (q *WaitingJobQueue) Remove(jobID string) bool {
q.mu.Lock()
defer q.mu.Unlock()

elem, ok := q.exist[jobID]
if !ok {
return false
}
q.jobs.Remove(elem)
delete(q.exist, jobID)
return true
}

// Peek returns the next job from the front of the queue without removing it.
// Returns the jobID and true if successful, or empty string and false if the queue is empty.
func (q *WaitingJobQueue) Peek() (string, bool) {
Expand Down
Loading
Loading