Buriburizaemon/nezha_domains
1
0
Fork 0
mirror of https://github.com/Buriburizaem0n/nezha_domains.git synced 2026-09-21 18:50:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rpc): make IO stream lifecycle race-safe

Browse Source 
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
2026-07-20 04:26:19 +00:00
co-authored by naiba/CloudCode
parent 8b47ff141f
commit c756ef9385
23 changed files with 2114 additions and 690 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/rpc/wait_for_agent_revoke_test.go
+55 -3
View File
@@ -16,7 +16,9 @@ import (
func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
h := NewNezhaHandler()
const streamID = "kill-switch-wait"
h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer)
if err := h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer); err != nil {
t.Fatalf("create waiter stream: %v", err)
}
done := make(chan struct {
io any
@@ -35,9 +37,16 @@ func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
dur time.Duration
}{stream, ok, time.Since(start)}
}()
waiter, err := h.GetStream(streamID)
if err != nil {
t.Fatalf("get waiter context: %v", err)
}
select {
case <-waiter.waitStartedCh:
case <-time.After(time.Second):
t.Fatal("WaitForAgent did not enter its blocking select")
}
// 让 WaitForAgent 真的进入 select 等待，再触发 kill switch。
time.Sleep(50 * time.Millisecond)
if revoked := h.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 {
t.Fatalf("expected to revoke exactly 1 MCP stream, got %d", revoked)
}
@@ -54,3 +63,46 @@ func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) {
t.Fatalf("WaitForAgent never returned after revoke; kill switch did not wake the waiter")
}
}
func TestRevokeStreamsForServerWakesWaitForAgentAndPreservesNewGeneration(t *testing.T) {
h := NewNezhaHandler()
const streamID = "server-revoke-generation"
if err := h.CreateStream(streamID, 0, 7); err != nil {
t.Fatalf("create waiter stream: %v", err)
}
done := make(chan bool, 1)
go func() {
_, ok := h.WaitForAgent(context.Background(), streamID, time.Minute)
done <- ok
}()
waiter, err := h.GetStream(streamID)
if err != nil {
t.Fatalf("get waiter stream: %v", err)
}
select {
case <-waiter.waitStartedCh:
case <-time.After(time.Second):
t.Fatal("WaitForAgent did not reach its blocking select")
}
h.RevokeStreamsForServer(7)
select {
case ok := <-done:
if ok {
t.Fatal("WaitForAgent must return false after server revocation")
}
case <-time.After(time.Second):
t.Fatal("server revocation did not wake WaitForAgent")
}
h.RevokeStreamsForServer(7)
if err := h.CreateStream(streamID, 0, 8); err != nil {
t.Fatalf("new generation must reuse released ID: %v", err)
}
h.RevokeStreamsForServer(7)
if h.StreamCount() != 1 {
t.Fatalf("new generation must remain tracked, got %d streams", h.StreamCount())
}
if err := h.CloseStream(streamID); err != nil {
t.Fatalf("cleanup new generation: %v", err)
}
}