address-haiyang-review

Signed-off-by: Ning Wang <n.wang.chn@hotmail.com>

Author: nwangfw

apps/console/api/planner/client/submission.go

@@ -25,13 +25,22 @@ import (
 // because its shape will evolve as the RM populates per-group allocation
 // details (currently always nil — RM doesn't return them yet).
 type PlannerDecision struct {
-	ProvisionID               string `json:"provision_id,omitempty"`
-	ProvisionResourceDeadline int64  `json:"provision_resource_deadline,omitempty"`
+	// ProvisionID mirrors rmtypes.ProvisionResult.ProvisionID returned
+	// by Provisioner.Provision.
+	ProvisionID string `json:"provision_id,omitempty"`
+	// ProvisionResourceDeadline is the unix-seconds deadline.
+	ProvisionResourceDeadline int64 `json:"provision_resource_deadline,omitempty"`
 	ResourceDetails           []struct {
-		ResourceType    string `json:"resource_type"`
+		// ResourceType maps to rmtypes.ResourceProvisionType
+		// (kubernetes / aws / lambdaCloud); may grow finer-grained.
+		ResourceType string `json:"resource_type"`
+		// EndpointCluster identifies the cluster serving this group;
 		EndpointCluster string `json:"endpoint_cluster,omitempty"`
-		GPUType         string `json:"gpu_type,omitempty"`
-		WorkerNum       int    `json:"worker_num,omitempty"`
+		// GPUType identifies the GPU model of the provisioned nodes.
+		// The GPU count is not defined here; it comes from the ModelTemplate.
+		GPUType string `json:"gpu_type,omitempty"`
+		// WorkerNum identifies the number of replicas to provision
+		WorkerNum int `json:"worker_num,omitempty"`
 	} `json:"resource_details,omitempty"`
 }
 

apps/console/api/planner/impl/planner.go

@@ -41,7 +41,7 @@ type Planner struct {
 	bc   plannerclient.BatchClient
 	prov provisioner.Provisioner
 
-	submit chan string // buffered FIFO of pending JobIDs
+	queue pendingQueue
 
 	baseCtx    context.Context
 	baseCancel context.CancelFunc
@@ -60,7 +60,7 @@ type Planner struct {
 // it to "pending" or "provisioning". After submission, status comes from
 // MDS.
 //
-//	Pending      → buffered in submit channel, no worker yet.
+//	Pending      → buffered in the queue, no worker yet.
 //	Provisioning → worker holds the job; Provision + CreateBatch in flight.
 //	Submitted    → CreateBatch returned; MDS owns the lifecycle from here.
 //	Failed       → Provision or CreateBatch errored.
@@ -79,13 +79,14 @@ const (
 )
 
 type queuedJob struct {
-	req        *plannerapi.EnqueueRequest
-	state      jobState
-	batchID    string // populated when state == jobStateSubmitted
-	err        error  // populated when state == jobStateFailed
-	enqueuedAt time.Time
-	failedAt   time.Time // populated when state == jobStateFailed
-	canceledAt time.Time // populated when state == jobStateCanceled
+	req         *plannerapi.EnqueueRequest
+	state       jobState
+	provisionID string // populated once Provision returns accepted
+	batchID     string // populated when state == jobStateSubmitted
+	err         error  // populated when state == jobStateFailed
+	enqueuedAt  time.Time
+	failedAt    time.Time // populated when state == jobStateFailed
+	canceledAt  time.Time // populated when state == jobStateCanceled
 }
 
 // terminalTime returns the timestamp at which the job transitioned into a
@@ -101,10 +102,6 @@ func terminalTime(j *queuedJob) time.Time {
 	return time.Time{}
 }
 
-// queueCapacity caps the submit channel. When full, Enqueue blocks on the
-// caller's context.
-const queueCapacity = 256
-
 // DefaultWorkerCount sizes the worker pool.
 const DefaultWorkerCount = 8
 
@@ -127,7 +124,7 @@ func NewPlanner(bc plannerclient.BatchClient, prov provisioner.Provisioner, work
 	q := &Planner{
 		bc:               bc,
 		prov:             prov,
-		submit:           make(chan string, queueCapacity),
+		queue:            newFIFOPendingQueue(queueCapacity),
 		baseCtx:          ctx,
 		baseCancel:       cancel,
 		jobs:             make(map[string]*queuedJob),
@@ -146,6 +143,7 @@ var _ plannerapi.Planner = (*Planner)(nil)
 
 // Close cancels in-flight work and waits for workers to exit.
 func (q *Planner) Close() error {
+	q.queue.Close()
 	q.baseCancel()
 	q.wg.Wait()
 	return nil
@@ -154,12 +152,11 @@ func (q *Planner) Close() error {
 func (q *Planner) run() {
 	defer q.wg.Done()
 	for {
-		select {
-		case <-q.baseCtx.Done():
+		jobID, err := q.queue.Pop(q.baseCtx)
+		if err != nil {
 			return
-		case jobID := <-q.submit:
-			q.process(jobID)
 		}
+		q.process(jobID)
 	}
 }
 
@@ -171,11 +168,12 @@ func (q *Planner) process(jobID string) {
 		q.mu.Unlock()
 		return
 	}
+	// Example: a pending job is canceled before a worker picks it up, so the
+	// worker later observes a non-pending state here and skips provisioning.
 	if job.state != jobStatePending {
-		// Cancel raced ahead; drop without provisioning.
 		state := job.state
 		q.mu.Unlock()
-		klog.Infof("[planner] skip job_id=%q state=%d", jobID, state)
+		klog.Infof("[planner] invalid state before provisioning job_id=%q state=%d", jobID, state)
 		return
 	}
 	job.state = jobStateProvisioning
@@ -193,6 +191,9 @@ func (q *Planner) process(jobID string) {
 		q.markFailed(jobID, errors.Join(plannerapi.ErrInsufficientResources, err))
 		return
 	}
+	q.mu.Lock()
+	q.jobs[jobID].provisionID = provResult.ProvisionID
+	q.mu.Unlock()
 
 	// Provision returns when the request is accepted, not when the resource
 	// is ready. Wait for Running before submitting to MDS, which rejects
@@ -284,8 +285,9 @@ func (q *Planner) waitForProvisionReady(provisionID string) error {
 }
 
 // releaseAfter performs a best-effort RM release and logs failures. The
-// reason string ("wait failure", "CreateBatch failure", "cancel-race")
-// appears in the log line so each call site is self-identifying.
+// reason string ("wait failure", "CreateBatch failure", "cancel-race",
+// "cancel submitted") appears in the log line so each call site is
+// self-identifying.
 func (q *Planner) releaseAfter(jobID, provisionID, reason string) {
 	if err := q.prov.Release(q.baseCtx, provisionID); err != nil {
 		klog.Warningf("[planner] release after %s job_id=%q provision_id=%q: %v",
@@ -295,16 +297,15 @@ func (q *Planner) releaseAfter(jobID, provisionID, reason string) {
 
 func (q *Planner) markFailed(jobID string, err error) {
 	q.mu.Lock()
-	if job, ok := q.jobs[jobID]; ok {
-		job.state = jobStateFailed
-		job.err = err
-		job.failedAt = time.Now()
-	}
+	job := q.jobs[jobID]
+	job.state = jobStateFailed
+	job.err = err
+	job.failedAt = time.Now()
 	q.mu.Unlock()
 	klog.Warningf("[planner] job_id=%q failed: %v", jobID, err)
 }
 
-// Enqueue records the job, pushes it onto the worker channel, and returns
+// Enqueue records the job, pushes it onto the queue, and returns
 // a placeholder batch in "pending" status.
 func (q *Planner) Enqueue(ctx context.Context, req *plannerapi.EnqueueRequest) (*plannerapi.Job, error) {
 	if req == nil {
@@ -322,6 +323,9 @@ func (q *Planner) Enqueue(ctx context.Context, req *plannerapi.EnqueueRequest) (
 	if q.prov == nil {
 		return nil, fmt.Errorf("%w: missing provisioner", plannerapi.ErrInsufficientResources)
 	}
+	if err := q.baseCtx.Err(); err != nil {
+		return nil, fmt.Errorf("planner closed: %w", err)
+	}
 
 	now := time.Now()
 	q.mu.Lock()
@@ -336,17 +340,14 @@ func (q *Planner) Enqueue(ctx context.Context, req *plannerapi.EnqueueRequest) (
 	}
 	q.mu.Unlock()
 
-	select {
-	case q.submit <- req.JobID:
-		// Happy path: a worker will dequeue and drive the entry; nothing to roll back.
-	case <-ctx.Done():
-		// Caller gave up while q.submit was full; the bookkeeping insert is orphaned.
+	if err := q.queue.Push(ctx, req.JobID); err != nil {
 		q.rollbackEnqueue(req.JobID)
-		return nil, ctx.Err()
-	case <-q.baseCtx.Done():
-		// Planner shutting down while q.submit was full; roll back the orphaned insert.
-		q.rollbackEnqueue(req.JobID)
-		return nil, fmt.Errorf("planner closed: %w", q.baseCtx.Err())
+		if errors.Is(err, errQueueClosed) {
+			// Planner shutting down while the queue was full; roll back the orphaned insert.
+			return nil, fmt.Errorf("planner closed: %w", q.baseCtx.Err())
+		}
+		// Caller gave up while the queue was full; the bookkeeping insert is orphaned.
+		return nil, err
 	}
 
 	klog.Infof("[planner] enqueue job_id=%q", req.JobID)
@@ -405,6 +406,7 @@ func (q *Planner) Cancel(ctx context.Context, jobID string) (*plannerapi.Job, er
 	}
 	state := job.state
 	batchID := job.batchID
+	provisionID := job.provisionID
 	req := job.req
 	enqueuedAt := job.enqueuedAt
 	var terminalAt time.Time
@@ -428,6 +430,7 @@ func (q *Planner) Cancel(ctx context.Context, jobID string) (*plannerapi.Job, er
 		if err != nil {
 			return nil, err
 		}
+		q.releaseAfter(jobID, provisionID, "cancel submitted")
 		return &plannerapi.Job{JobID: jobID, Batch: batch}, nil
 	}
 	// Already terminal (failed/canceled) — return current view, no double-cancel side effects.

apps/console/api/planner/impl/planner_test.go

@@ -711,7 +711,7 @@ func TestCancelQueuedJobBeforeWorkerPicksUp(t *testing.T) {
 	}
 }
 
-func TestCancelSubmittedJobForwardsToMDS(t *testing.T) {
+func TestCancelSubmittedJobForwardsToMDSAndReleasesProvision(t *testing.T) {
 	prov := &fakeProvisioner{}
 	bc := &fakeBatchClient{}
 	q := newTestPlanner(t, bc, prov, 1)
@@ -737,6 +737,11 @@ func TestCancelSubmittedJobForwardsToMDS(t *testing.T) {
 	if len(cancels) != 1 || cancels[0] != "batch-j-sub" {
 		t.Errorf("bc.CancelBatch calls = %v; want [batch-j-sub]", cancels)
 	}
+
+	_, releases, _ := prov.snapshot()
+	if len(releases) != 1 || releases[0] != "prov-j-sub" {
+		t.Errorf("prov.Release calls = %v; want [prov-j-sub]", releases)
+	}
 }
 
 func TestCancelUnknownJobReturnsNotFound(t *testing.T) {
@@ -910,6 +915,27 @@ func TestCloseIsIdempotent(t *testing.T) {
 	}
 }
 
+// TestEnqueueAfterCloseReturnsClosed locks down planner shutdown semantics:
+// once Close returns, new Enqueue calls must fail immediately instead of
+// slipping into the buffered pending queue.
+func TestEnqueueAfterCloseReturnsClosed(t *testing.T) {
+	q := NewPlanner(&fakeBatchClient{}, &fakeProvisioner{}, 1)
+	if err := q.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+
+	// This guards the refactor from queue-owned shutdown to planner-owned
+	// shutdown checks. Without the Enqueue-side closed check, the buffered
+	// queue can still accept a job after Close.
+	_, err := q.Enqueue(context.Background(), validReq("j-closed"))
+	if err == nil {
+		t.Fatal("Enqueue after Close unexpectedly succeeded")
+	}
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("want Close error to wrap context.Canceled; got %v", err)
+	}
+}
+
 // TestWorkerCountFloor: a non-positive workerCount must be floored to 1
 // rather than starting zero goroutines (which would silently hang every
 // Enqueue forever).

apps/console/api/planner/impl/queue.go

@@ -0,0 +1,81 @@
+/*
+Copyright 2026 The Aibrix Team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package impl
+
+import (
+	"context"
+	"errors"
+	"sync"
+)
+
+var errQueueClosed = errors.New("queue closed")
+
+const queueCapacity = 256
+
+// pendingQueue is the planner's buffer of pending job IDs.
+type pendingQueue interface {
+	// Push returns ctx.Err() on caller cancel, errQueueClosed on shutdown.
+	Push(ctx context.Context, jobID string) error
+	// Pop returns ctx.Err() on caller cancel, errQueueClosed on shutdown.
+	Pop(ctx context.Context) (string, error)
+	Close()
+}
+
+// fifoPendingQueue preserves enqueue order for pending job IDs.
+type fifoPendingQueue struct {
+	ch        chan string
+	done      chan struct{}
+	closeOnce sync.Once
+}
+
+func newFIFOPendingQueue(capacity int) *fifoPendingQueue {
+	if capacity < 1 {
+		capacity = 1
+	}
+	return &fifoPendingQueue{
+		ch:   make(chan string, capacity),
+		done: make(chan struct{}),
+	}
+}
+
+func (q *fifoPendingQueue) Push(ctx context.Context, jobID string) error {
+	select {
+	case q.ch <- jobID:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-q.done:
+		return errQueueClosed
+	}
+}
+
+func (q *fifoPendingQueue) Pop(ctx context.Context) (string, error) {
+	select {
+	case jobID := <-q.ch:
+		return jobID, nil
+	case <-ctx.Done():
+		return "", ctx.Err()
+	case <-q.done:
+		return "", errQueueClosed
+	}
+}
+
+func (q *fifoPendingQueue) Close() {
+	q.closeOnce.Do(func() {
+		close(q.done)
+	})
+}