This is an automated email from the ASF dual-hosted git repository.
lizhimins pushed a commit to branch rocketmq-studio
in repository https://gitbox.apache.org/repos/asf/rocketmq-dashboard.git
The following commit(s) were added to refs/heads/rocketmq-studio by this push:
new c01908beb fix(rmqctl): respect cancellation while waiting for MCP
session access (#6604)
c01908beb is described below
commit c01908bebdc58c1341e4308220971f7a2441b606
Author: pilichoumao <[email protected]>
AuthorDate: Thu Oct 15 18:16:00 2026 +0800
fix(rmqctl): respect cancellation while waiting for MCP session access
(#6604)
---
docs/ai-agent-architecture.md | 8 +
.../internal/studio/mcp_gate_integration_test.go | 227 +++++++++++++++++++++
rmqctl/internal/studio/mcp_http.go | 6 +-
rmqctl/internal/studio/mcp_message.go | 14 +-
rmqctl/internal/studio/send_gate.go | 134 ++++++++++++
rmqctl/internal/studio/send_gate_test.go | 209 +++++++++++++++++++
6 files changed, 590 insertions(+), 8 deletions(-)
diff --git a/docs/ai-agent-architecture.md b/docs/ai-agent-architecture.md
index ab01cd757..caa28fa17 100644
--- a/docs/ai-agent-architecture.md
+++ b/docs/ai-agent-architecture.md
@@ -62,6 +62,14 @@ stdio 桥最多转发 64 个在途调用；满载时暂停读取后续输入，
- **独立性**：Studio 挂了 `rmqctl` 就无用。
- **隔离**：工具最终仍在拉起 Agent 的同一个 JVM 与同一个容器里执行，没有租户或网络隔离。
+### MCP 请求取消与会话重连
+
+`rmqctl` 的单次消息超时包含等待发送和重连的时间。普通请求可以并发发送；重连在已有发送完成后
+独占会话，完成 `initialize` 和 `notifications/initialized` 才允许后续请求进入新会话。
+等待队列按到达顺序放行，连续的普通请求可一起进入，后来的请求不能越过已排队的重连。
+等待中的请求取消或超时后会退出队列，不会迟到发送；关闭会话也会唤醒等待者。
+取消一条等待中的请求不会取消另一条正常调用或正在执行的重连，也不改变既有的重试次数和代际检查。
+
### 已否决的替代方案
把 MCP server 直接以 `{"type":"http","url":"http://127.0.0.1:8888/api/mcp"}` 注册进
diff --git a/rmqctl/internal/studio/mcp_gate_integration_test.go
b/rmqctl/internal/studio/mcp_gate_integration_test.go
new file mode 100644
index 000000000..f1aa20238
--- /dev/null
+++ b/rmqctl/internal/studio/mcp_gate_integration_test.go
@@ -0,0 +1,227 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 studio
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "net/http"
+ "sync/atomic"
+ "testing"
+ "time"
+)
+import mcptransport "github.com/mark3labs/mcp-go/client/transport"
+
+func TestSessionDeadlineWhileReconnectOwnsSendLock(t *testing.T) {
+ var initCount atomic.Int32
+ entered := make(chan struct{})
+ release := make(chan struct{})
+ client := NewClient(&http.Client{Transport: roundTripFunc(func(req
*http.Request) (*http.Response, error) {
+ if req.Method == http.MethodDelete {
+ return mcpHTTPResponse(204, "", ""), nil
+ }
+ payload := readMCPPayload(t, req)
+ switch payload.Method {
+ case "initialize":
+ n := initCount.Add(1)
+ if n == 2 {
+ close(entered)
+ select {
+ case <-release:
+ case <-req.Context().Done():
+ return nil, req.Context().Err()
+ }
+ }
+ resp := mcpJSONResultResponse(payload.ID,
map[string]any{"protocolVersion": "2025-11-25", "capabilities":
map[string]any{}, "serverInfo": map[string]any{"name": "test", "version": "1"}})
+ resp.Header.Set(mcptransport.HeaderKeySessionID,
fmt.Sprintf("session-%d", n))
+ return resp, nil
+ case "notifications/initialized":
+ return mcpHTTPResponse(202, "", ""), nil
+ case "tools/list":
+ if req.Header.Get(mcptransport.HeaderKeySessionID) ==
"session-1" {
+ return mcpHTTPResponse(404, "text/plain",
"expired"), nil
+ }
+ return mcpJSONResultResponse(payload.ID,
map[string]any{"tools": []any{}}), nil
+ }
+ return nil, fmt.Errorf("unexpected method %s", payload.Method)
+ })})
+ session := newMCPTestSession(t, client, Target{Server:
"http://localhost";, InstanceID: "dev", Credential: Credential{AccessKey:
"test", SecretKey: "test"}, Timeout: 5 * time.Second})
+ defer session.Close()
+ initializeClientSession(t, context.Background(), session)
+ firstDone := make(chan error, 1)
+ go func() {
+ _, _, err := session.SendMessage(context.Background(),
json.RawMessage(`{"jsonrpc":"2.0","id":"first","method":"tools/list"}`))
+ firstDone <- err
+ }()
+ select {
+ case <-entered:
+ case <-time.After(time.Second):
+ t.Fatal("reconnect did not start")
+ }
+ ctx, cancel := context.WithTimeout(context.Background(),
30*time.Millisecond)
+ defer cancel()
+ done := make(chan error, 1)
+ go func() {
+ _, _, err := session.SendMessage(ctx,
json.RawMessage(`{"jsonrpc":"2.0","id":"second","method":"tools/list"}`))
+ done <- err
+ }()
+ gotDone := false
+ select {
+ case err := <-done:
+ gotDone = true
+ if !errors.Is(err, context.DeadlineExceeded) {
+ t.Errorf("unexpected cancellation error %v", err)
+ }
+ case <-time.After(time.Second):
+ t.Error("cancelled request remained blocked by another
request's reconnect")
+ }
+ close(release)
+ select {
+ case err := <-firstDone:
+ if err != nil {
+ t.Errorf("healthy reconnect failed: %v", err)
+ }
+ case <-time.After(time.Second):
+ t.Error("healthy request stalled")
+ }
+ if !gotDone {
+ select {
+ case <-done:
+ case <-time.After(time.Second):
+ t.Error("cancelled call never returned")
+ }
+ }
+}
+
+func TestSessionDeadlineWhileAnotherCallOwnsReadLock(t *testing.T) {
+ entered := make(chan struct{})
+ release := make(chan struct{})
+ client := NewClient(&http.Client{Transport: roundTripFunc(func(req
*http.Request) (*http.Response, error) {
+ if req.Method == http.MethodDelete {
+ return mcpHTTPResponse(204, "", ""), nil
+ }
+ payload := readMCPPayload(t, req)
+ switch payload.Method {
+ case "initialize":
+ resp := mcpJSONResultResponse(payload.ID,
map[string]any{"protocolVersion": "2025-11-25", "capabilities":
map[string]any{}, "serverInfo": map[string]any{"name": "test", "version": "1"}})
+ resp.Header.Set(mcptransport.HeaderKeySessionID,
"session-1")
+ return resp, nil
+ case "notifications/initialized":
+ return mcpHTTPResponse(202, "", ""), nil
+ case "tools/list":
+ if string(payload.ID) == `"long"` {
+ close(entered)
+ select {
+ case <-release:
+ case <-req.Context().Done():
+ return nil, req.Context().Err()
+ }
+ return mcpJSONResultResponse(payload.ID,
map[string]any{"tools": []any{}}), nil
+ }
+ return mcpHTTPResponse(404, "text/plain", "expired"),
nil
+ }
+ return nil, fmt.Errorf("unexpected method %s", payload.Method)
+ })})
+ session := newMCPTestSession(t, client, Target{Server:
"http://localhost";, InstanceID: "dev", Credential: Credential{AccessKey:
"test", SecretKey: "test"}, Timeout: 5 * time.Second})
+ defer session.Close()
+ initializeClientSession(t, context.Background(), session)
+ longDone := make(chan error, 1)
+ go func() {
+ _, _, err := session.SendMessage(context.Background(),
json.RawMessage(`{"jsonrpc":"2.0","id":"long","method":"tools/list"}`))
+ longDone <- err
+ }()
+ select {
+ case <-entered:
+ case <-time.After(time.Second):
+ t.Fatal("long request did not start")
+ }
+ ctx, cancel := context.WithTimeout(context.Background(),
30*time.Millisecond)
+ defer cancel()
+ done := make(chan error, 1)
+ go func() {
+ _, _, err := session.SendMessage(ctx,
json.RawMessage(`{"jsonrpc":"2.0","id":"short","method":"tools/list"}`))
+ done <- err
+ }()
+ gotDone := false
+ select {
+ case err := <-done:
+ gotDone = true
+ if !errors.Is(err, context.DeadlineExceeded) {
+ t.Errorf("unexpected cancellation error %v", err)
+ }
+ case <-time.After(time.Second):
+ t.Error("reconnect lock acquisition ignored the caller
deadline")
+ }
+ close(release)
+ select {
+ case err := <-longDone:
+ if err != nil {
+ t.Errorf("healthy long request failed: %v", err)
+ }
+ case <-time.After(time.Second):
+ t.Error("long request stalled")
+ }
+ if !gotDone {
+ select {
+ case <-done:
+ case <-time.After(time.Second):
+ t.Error("cancelled call never returned")
+ }
+ }
+}
+
+func TestSessionConfiguredTimeoutAndCloseWhileQueued(t *testing.T) {
+ for _, method := range []string{"request", "notification"} {
+ for _, termination := range []string{"timeout", "close"} {
+ t.Run(method+"/"+termination, func(t *testing.T) {
+ transport := &reconnectFailureTransport{}
+ session := &MCPClientSession{transport:
transport, closed: make(chan struct{})}
+ if termination == "timeout" {
+ session.timeout = 30 * time.Millisecond
+ }
+ if err :=
session.sendGate.acquire(context.Background(), session.closed, true); err !=
nil {
+ t.Fatal(err)
+ }
+ defer session.sendGate.release(true)
+ defer session.Close()
+ payload :=
json.RawMessage(`{"jsonrpc":"2.0","id":"queued","method":"tools/list"}`)
+ if method == "notification" {
+ payload =
json.RawMessage(`{"jsonrpc":"2.0","method":"notifications/test"}`)
+ }
+ done := make(chan error, 1)
+ go func() { _, _, err :=
session.SendMessage(context.Background(), payload); done <- err }()
+ waitSendGateQueue(t, &session.sendGate, 1)
+ want := error(context.DeadlineExceeded)
+ if termination == "close" {
+ want = errMCPClientSessionClosed
+ if err := session.Close(); err != nil {
+ t.Fatal(err)
+ }
+ }
+ receiveSendGateResult(t, done, want)
+ transport.mu.Lock()
+ ordinary, notifications :=
transport.ordinaryCalls, transport.notificationCalls
+ transport.mu.Unlock()
+ if ordinary != 0 || notifications != 0 {
+ t.Fatalf("queued message was sent:
calls=%d notifications=%d", ordinary, notifications)
+ }
+ })
+ }
+ }
+}
diff --git a/rmqctl/internal/studio/mcp_http.go
b/rmqctl/internal/studio/mcp_http.go
index f8a1bb3af..78307a619 100644
--- a/rmqctl/internal/studio/mcp_http.go
+++ b/rmqctl/internal/studio/mcp_http.go
@@ -53,11 +53,11 @@ type MCPClientSession struct {
closed chan struct{}
notifications chan json.RawMessage
- // sendMu allows ordinary messages to be concurrent while preventing any
+ // sendGate allows ordinary messages to be concurrent while preventing
any
// message from entering a replacement session before reinitialization
and
// its initialized notification have completed.
- sendMu sync.RWMutex
- state struct {
+ sendGate sendGate
+ state struct {
sync.RWMutex
initialize *mcptransport.JSONRPCRequest
generation uint64
diff --git a/rmqctl/internal/studio/mcp_message.go
b/rmqctl/internal/studio/mcp_message.go
index 8babade41..7c4d7a6b7 100644
--- a/rmqctl/internal/studio/mcp_message.go
+++ b/rmqctl/internal/studio/mcp_message.go
@@ -158,7 +158,9 @@ func (session *MCPClientSession) sendWithReconnect(ctx
context.Context, skipReco
// A call may recover once and send at most twice. Only a definitive
session
// termination permits replay; an incomplete handshake has not sent the
call.
for attempt := 0; attempt < 2; attempt++ {
- session.sendMu.RLock()
+ if err := session.sendGate.acquire(ctx, session.closed, false);
err != nil {
+ return err
+ }
generation, hadSession := session.sendSnapshot()
session.state.RLock()
pending := session.state.reconnectPending
@@ -168,7 +170,7 @@ func (session *MCPClientSession) sendWithReconnect(ctx
context.Context, skipReco
} else {
err = send()
}
- session.sendMu.RUnlock()
+ session.sendGate.release(false)
if err == nil || skipReconnect || attempt == 1 {
return err
}
@@ -185,7 +187,7 @@ func (session *MCPClientSession) sendWithReconnect(ctx
context.Context, skipReco
}
// sendSnapshot returns the current session generation and whether a session
-// id is already established. Callers must hold sendMu (at least RLock) so the
+// id is already established. Callers must hold a sendGate permit so the
// snapshot is consistent with the transport state used for the send.
func (session *MCPClientSession) sendSnapshot() (uint64, bool) {
session.state.RLock()
@@ -235,8 +237,10 @@ func initializationProtocolVersion(response
*mcptransport.JSONRPCResponse) (stri
func (session *MCPClientSession) reinitialize(ctx context.Context,
expectedGeneration uint64) error {
// A 404 for an established session requires a fresh initialize
exchange.
// MCP 2025-11-25:
https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#session-management
- session.sendMu.Lock()
- defer session.sendMu.Unlock()
+ if err := session.sendGate.acquire(ctx, session.closed, true); err !=
nil {
+ return err
+ }
+ defer session.sendGate.release(true)
session.state.RLock()
if session.state.generation != expectedGeneration &&
!session.state.reconnectPending {
session.state.RUnlock()
diff --git a/rmqctl/internal/studio/send_gate.go
b/rmqctl/internal/studio/send_gate.go
new file mode 100644
index 000000000..b1e0f3e63
--- /dev/null
+++ b/rmqctl/internal/studio/send_gate.go
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 studio
+
+import (
+ "container/list"
+ "context"
+ "errors"
+ "sync"
+)
+
+var errMCPClientSessionClosed = errors.New("MCP client session is closed")
+
+// sendGate admits concurrent ordinary sends, but grants reconnects
exclusively.
+// A FIFO queue prevents new sends from starving a reconnect. Its zero value is
+// ready to use; waiting needs neither a polling timer nor a helper goroutine.
+type sendGate struct {
+ mu sync.Mutex
+ readers int
+ writer bool
+ waiters list.List
+}
+
+type sendGateWaiter struct {
+ exclusive bool
+ ready chan struct{}
+ granted bool
+}
+
+func sendGateWaitError(ctx context.Context, closed <-chan struct{}) error {
+ if err := ctx.Err(); err != nil {
+ return err
+ }
+ select {
+ case <-closed:
+ return errMCPClientSessionClosed
+ default:
+ return nil
+ }
+}
+
+func (gate *sendGate) acquire(ctx context.Context, closed <-chan struct{},
exclusive bool) error {
+ gate.mu.Lock()
+ if err := sendGateWaitError(ctx, closed); err != nil {
+ gate.mu.Unlock()
+ return err
+ }
+ if gate.waiters.Len() == 0 && !gate.writer && (!exclusive ||
gate.readers == 0) {
+ if exclusive {
+ gate.writer = true
+ } else {
+ gate.readers++
+ }
+ gate.mu.Unlock()
+ return nil
+ }
+ waiter := &sendGateWaiter{exclusive: exclusive, ready: make(chan
struct{})}
+ element := gate.waiters.PushBack(waiter)
+ gate.mu.Unlock()
+
+ select {
+ case <-waiter.ready:
+ case <-ctx.Done():
+ case <-closed:
+ }
+ gate.mu.Lock()
+ defer gate.mu.Unlock()
+ // Cancellation can race with a grant. Return the permit in that case
rather
+ // than leaving a reader or writer behind that blocks every subsequent
call.
+ if err := sendGateWaitError(ctx, closed); err != nil {
+ if waiter.granted {
+ gate.releaseLocked(exclusive)
+ } else {
+ gate.waiters.Remove(element)
+ gate.grantLocked()
+ }
+ return err
+ }
+ return nil
+}
+
+func (gate *sendGate) release(exclusive bool) {
+ gate.mu.Lock()
+ defer gate.mu.Unlock()
+ gate.releaseLocked(exclusive)
+}
+
+func (gate *sendGate) releaseLocked(exclusive bool) {
+ if exclusive {
+ gate.writer = false
+ } else {
+ gate.readers--
+ }
+ gate.grantLocked()
+}
+
+// Grant an exclusive waiter only after existing readers finish, or admit the
+// consecutive readers at the head of the queue as one concurrent batch.
+func (gate *sendGate) grantLocked() {
+ if gate.writer {
+ return
+ }
+ for element := gate.waiters.Front(); element != nil; element =
gate.waiters.Front() {
+ waiter := element.Value.(*sendGateWaiter)
+ if waiter.exclusive && gate.readers != 0 {
+ return
+ }
+ gate.waiters.Remove(element)
+ waiter.granted = true
+ if waiter.exclusive {
+ gate.writer = true
+ } else {
+ gate.readers++
+ }
+ close(waiter.ready)
+ if waiter.exclusive {
+ return
+ }
+ }
+}
diff --git a/rmqctl/internal/studio/send_gate_test.go
b/rmqctl/internal/studio/send_gate_test.go
new file mode 100644
index 000000000..438e1b2a2
--- /dev/null
+++ b/rmqctl/internal/studio/send_gate_test.go
@@ -0,0 +1,209 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 studio
+
+import (
+ "context"
+ "errors"
+ "testing"
+ "time"
+)
+
+func waitSendGateQueue(t *testing.T, gate *sendGate, count int) {
+ t.Helper()
+ deadline := time.After(2 * time.Second)
+ ticker := time.NewTicker(time.Millisecond)
+ defer ticker.Stop()
+ for {
+ gate.mu.Lock()
+ size := gate.waiters.Len()
+ gate.mu.Unlock()
+ if size == count {
+ return
+ }
+ select {
+ case <-deadline:
+ t.Fatalf("queued waiters = %d, want %d", size, count)
+ case <-ticker.C:
+ }
+ }
+}
+
+func acquireSendGateAsync(gate *sendGate, ctx context.Context, closed <-chan
struct{}, exclusive bool) <-chan error {
+ result := make(chan error, 1)
+ go func() { result <- gate.acquire(ctx, closed, exclusive) }()
+ return result
+}
+
+func receiveSendGateResult(t *testing.T, result <-chan error, want error) {
+ t.Helper()
+ select {
+ case err := <-result:
+ if !errors.Is(err, want) {
+ t.Fatalf("acquire error = %v, want %v", err, want)
+ }
+ case <-time.After(2 * time.Second):
+ t.Fatal("gate waiter did not return")
+ }
+}
+
+func TestSendGateConcurrentReadersAndFIFOReconnects(t *testing.T) {
+ var gate sendGate
+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+ defer cancel()
+ closed := make(chan struct{})
+ if err := gate.acquire(ctx, closed, false); err != nil {
+ t.Fatal(err)
+ }
+ // A second ordinary send can overlap the first one.
+ if err := gate.acquire(ctx, closed, false); err != nil {
+ t.Fatal(err)
+ }
+ writer1 := acquireSendGateAsync(&gate, ctx, closed, true)
+ waitSendGateQueue(t, &gate, 1)
+ reader1 := acquireSendGateAsync(&gate, ctx, closed, false)
+ waitSendGateQueue(t, &gate, 2)
+ reader2 := acquireSendGateAsync(&gate, ctx, closed, false)
+ waitSendGateQueue(t, &gate, 3)
+ writer2 := acquireSendGateAsync(&gate, ctx, closed, true)
+ waitSendGateQueue(t, &gate, 4)
+ lateReader := acquireSendGateAsync(&gate, ctx, closed, false)
+ waitSendGateQueue(t, &gate, 5)
+ gate.release(false)
+ select {
+ case <-writer1:
+ t.Fatal("writer overlapped a remaining reader")
+ default:
+ }
+ gate.release(false)
+ receiveSendGateResult(t, writer1, nil)
+ select {
+ case <-reader1:
+ t.Fatal("reader overtook reconnect")
+ default:
+ }
+ gate.release(true)
+ receiveSendGateResult(t, reader1, nil)
+ receiveSendGateResult(t, reader2, nil)
+ gate.release(false)
+ select {
+ case <-writer2:
+ t.Fatal("writer overlapped reader batch")
+ default:
+ }
+ gate.release(false)
+ receiveSendGateResult(t, writer2, nil)
+ select {
+ case <-lateReader:
+ t.Fatal("late reader overtook queued reconnect")
+ default:
+ }
+ gate.release(true)
+ receiveSendGateResult(t, lateReader, nil)
+ gate.release(false)
+}
+
+func TestSendGateCancelledWriterUnblocksQueuedReaders(t *testing.T) {
+ var gate sendGate
+ closed := make(chan struct{})
+ if err := gate.acquire(context.Background(), closed, false); err != nil
{
+ t.Fatal(err)
+ }
+ writerCtx, cancel := context.WithCancel(context.Background())
+ writer := acquireSendGateAsync(&gate, writerCtx, closed, true)
+ waitSendGateQueue(t, &gate, 1)
+ readerCtx, readerCancel := context.WithTimeout(context.Background(),
2*time.Second)
+ defer readerCancel()
+ reader := acquireSendGateAsync(&gate, readerCtx, closed, false)
+ waitSendGateQueue(t, &gate, 2)
+ cancel()
+ receiveSendGateResult(t, writer, context.Canceled)
+ // The original reader is still active: removal must wake compatible
waiters.
+ receiveSendGateResult(t, reader, nil)
+ gate.release(false)
+ gate.release(false)
+ waitSendGateQueue(t, &gate, 0)
+}
+
+func TestSendGateCancelledGrantReturnsPermit(t *testing.T) {
+ for _, exclusive := range []bool{false, true} {
+ name := "reader"
+ if exclusive {
+ name = "writer"
+ }
+ t.Run(name, func(t *testing.T) {
+ for i := 0; i < 100; i++ {
+ var gate sendGate
+ closed := make(chan struct{})
+ if err := gate.acquire(context.Background(),
closed, true); err != nil {
+ t.Fatal(err)
+ }
+ ctx, cancel :=
context.WithCancel(context.Background())
+ result := acquireSendGateAsync(&gate, ctx,
closed, exclusive)
+ waitSendGateQueue(t, &gate, 1)
+ // Grant the permit, but cancel before the
waiter can finish acquiring it.
+ gate.mu.Lock()
+ gate.releaseLocked(true)
+ cancel()
+ gate.mu.Unlock()
+ receiveSendGateResult(t, result,
context.Canceled)
+ nextCtx, nextCancel :=
context.WithTimeout(context.Background(), time.Second)
+ err := gate.acquire(nextCtx, closed, true)
+ nextCancel()
+ if err != nil {
+ t.Fatalf("iteration %d leaked a granted
permit: %v", i, err)
+ }
+ gate.release(true)
+ }
+ })
+ }
+}
+
+func TestSendGateCloseWakesAllWaiters(t *testing.T) {
+ var gate sendGate
+ closed := make(chan struct{})
+ if err := gate.acquire(context.Background(), closed, true); err != nil {
+ t.Fatal(err)
+ }
+ reader := acquireSendGateAsync(&gate, context.Background(), closed,
false)
+ waitSendGateQueue(t, &gate, 1)
+ writer := acquireSendGateAsync(&gate, context.Background(), closed,
true)
+ waitSendGateQueue(t, &gate, 2)
+ close(closed)
+ receiveSendGateResult(t, reader, errMCPClientSessionClosed)
+ receiveSendGateResult(t, writer, errMCPClientSessionClosed)
+ waitSendGateQueue(t, &gate, 0)
+ gate.release(true)
+ if err := gate.acquire(context.Background(), closed, false);
!errors.Is(err, errMCPClientSessionClosed) {
+ t.Fatalf("post-close acquire = %v", err)
+ }
+}
+
+func TestSendGatePreCancelledContextNeverAcquires(t *testing.T) {
+ var gate sendGate
+ closed := make(chan struct{})
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+ for _, exclusive := range []bool{false, true} {
+ if err := gate.acquire(ctx, closed, exclusive); !errors.Is(err,
context.Canceled) {
+ t.Fatalf("acquire = %v", err)
+ }
+ }
+ if gate.readers != 0 || gate.writer || gate.waiters.Len() != 0 {
+ t.Fatal("cancelled acquisition changed gate state")
+ }
+}
