mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-22 11:10:12 +00:00
fix(rpc): make IO stream lifecycle race-safe
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
naiba/CloudCode
@@ -0,0 +1,259 @@
|
||||
package rpc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type lifecycleRWC struct {
|
||||
closed chan struct{}
|
||||
}
|
||||
|
||||
type reenteringErrorRWC struct {
|
||||
handler *NezhaHandler
|
||||
streamID string
|
||||
err error
|
||||
}
|
||||
|
||||
func (stream *reenteringErrorRWC) Read([]byte) (int, error) { return 0, io.EOF }
|
||||
func (stream *reenteringErrorRWC) Write(data []byte) (int, error) { return len(data), nil }
|
||||
func (stream *reenteringErrorRWC) Close() error {
|
||||
if _, ok := stream.handler.StreamOwnership(stream.streamID); ok {
|
||||
return errors.Join(stream.err, errors.New("stream remained registered during endpoint close"))
|
||||
}
|
||||
return stream.err
|
||||
}
|
||||
|
||||
func newLifecycleRWC() *lifecycleRWC {
|
||||
return &lifecycleRWC{closed: make(chan struct{})}
|
||||
}
|
||||
|
||||
func (stream *lifecycleRWC) Read([]byte) (int, error) { return 0, io.EOF }
|
||||
func (stream *lifecycleRWC) Write(data []byte) (int, error) { return len(data), nil }
|
||||
func (stream *lifecycleRWC) Close() error {
|
||||
select {
|
||||
case <-stream.closed:
|
||||
default:
|
||||
close(stream.closed)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestIOStreamValidCreateAttachCloseLifecycle(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
user := newLifecycleRWC()
|
||||
agent := newLifecycleRWC()
|
||||
if err := handler.CreateStream("valid-lifecycle", 11, 22); err != nil {
|
||||
t.Fatalf("Given a new stream, CreateStream failed: %v", err)
|
||||
}
|
||||
if err := handler.UserConnected("valid-lifecycle", user); err != nil {
|
||||
t.Fatalf("Given a tracked stream, UserConnected failed: %v", err)
|
||||
}
|
||||
if err := handler.AgentConnected("valid-lifecycle", agent); err != nil {
|
||||
t.Fatalf("Given a tracked stream, AgentConnected failed: %v", err)
|
||||
}
|
||||
if _, ok := handler.StreamOwnership("valid-lifecycle"); !ok {
|
||||
t.Fatal("Then a valid attached stream must remain tracked")
|
||||
}
|
||||
if err := handler.CloseStream("valid-lifecycle"); err != nil {
|
||||
t.Fatalf("When closing the valid stream, CloseStream failed: %v", err)
|
||||
}
|
||||
if _, ok := handler.StreamOwnership("valid-lifecycle"); ok {
|
||||
t.Fatal("Then CloseStream must remove the tracked stream")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateStreamKeepsExistingStreamWhenIDIsDuplicated(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
original := newLifecycleRWC()
|
||||
if err := handler.CreateStream("duplicate-id", 11, 22); err != nil {
|
||||
t.Fatalf("Given a new stream ID, CreateStream failed: %v", err)
|
||||
}
|
||||
if err := handler.AgentConnected("duplicate-id", original); err != nil {
|
||||
t.Fatalf("Given a live stream, AgentConnected failed: %v", err)
|
||||
}
|
||||
|
||||
err := handler.CreateStream("duplicate-id", 33, 44)
|
||||
if !errors.Is(err, ErrStreamAlreadyExists) {
|
||||
t.Fatalf("When reusing a live ID, expected ErrStreamAlreadyExists, got %v", err)
|
||||
}
|
||||
owner, found := handler.StreamOwnership("duplicate-id")
|
||||
if !found || owner != 11 {
|
||||
t.Fatalf("Then the original stream ownership must remain, found=%v owner=%d", found, owner)
|
||||
}
|
||||
select {
|
||||
case <-original.closed:
|
||||
t.Fatal("Then duplicate creation must not close the original endpoint")
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentConnectedRejectsDuplicateEndpointWithoutReplacingLiveRelay(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
first := newLifecycleRWC()
|
||||
second := newLifecycleRWC()
|
||||
if err := handler.CreateStream("agent-once", 11, 22); err != nil {
|
||||
t.Fatalf("Given a new stream, CreateStream failed: %v", err)
|
||||
}
|
||||
if err := handler.AgentConnected("agent-once", first); err != nil {
|
||||
t.Fatalf("Given no agent endpoint, AgentConnected failed: %v", err)
|
||||
}
|
||||
if err := handler.AgentConnected("agent-once", second); !errors.Is(err, ErrAgentStreamAlreadyConnected) {
|
||||
t.Fatalf("When attaching a second agent endpoint, expected ErrAgentStreamAlreadyConnected, got %v", err)
|
||||
}
|
||||
endpoints, err := handler.GetStream("agent-once")
|
||||
if err != nil || endpoints.agentIo != first {
|
||||
t.Fatalf("Then the first endpoint must remain attached, err=%v", err)
|
||||
}
|
||||
select {
|
||||
case <-second.closed:
|
||||
default:
|
||||
t.Fatal("Then the rejected duplicate endpoint must be closed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseStreamWakesWaitForAgentAndAllowsSlotReuse(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
if err := handler.CreateStream("wait-close", 11, 22); err != nil {
|
||||
t.Fatalf("Given a pending stream, CreateStream failed: %v", err)
|
||||
}
|
||||
stream, err := handler.GetStream("wait-close")
|
||||
if err != nil {
|
||||
t.Fatalf("Given a created stream, GetStream failed: %v", err)
|
||||
}
|
||||
result := make(chan bool, 1)
|
||||
go func() {
|
||||
_, ok := handler.WaitForAgent(context.Background(), "wait-close", time.Minute)
|
||||
result <- ok
|
||||
}()
|
||||
select {
|
||||
case <-stream.waitStartedCh:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("WaitForAgent did not enter its blocking select")
|
||||
}
|
||||
|
||||
if err := handler.CloseStream("wait-close"); err != nil {
|
||||
t.Fatalf("When closing a pending stream, CloseStream failed: %v", err)
|
||||
}
|
||||
select {
|
||||
case ok := <-result:
|
||||
if ok {
|
||||
t.Fatal("Then WaitForAgent must report no attached agent")
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Then CloseStream must wake WaitForAgent")
|
||||
}
|
||||
if err := handler.CreateStream("wait-close-reused", 11, 22); err != nil {
|
||||
t.Fatalf("Then the released user/server slot must be reusable: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRevokeStreamsForPurposeWakesWaitForAgentAndIsRepeatable(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
if err := handler.CreateStreamWithPurpose("revoke-wait", 0, 22, PurposeMCPTransfer); err != nil {
|
||||
t.Fatalf("Given a pending MCP stream, CreateStream failed: %v", err)
|
||||
}
|
||||
stream, err := handler.GetStream("revoke-wait")
|
||||
if err != nil {
|
||||
t.Fatalf("Given a created stream, GetStream failed: %v", err)
|
||||
}
|
||||
result := make(chan bool, 1)
|
||||
go func() {
|
||||
_, ok := handler.WaitForAgent(context.Background(), "revoke-wait", time.Minute)
|
||||
result <- ok
|
||||
}()
|
||||
select {
|
||||
case <-stream.waitStartedCh:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("WaitForAgent did not enter its blocking select")
|
||||
}
|
||||
|
||||
if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 {
|
||||
t.Fatalf("When revoking the purpose, expected one stream, got %d", revoked)
|
||||
}
|
||||
if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 {
|
||||
t.Fatalf("When repeating revocation, expected zero streams, got %d", revoked)
|
||||
}
|
||||
select {
|
||||
case ok := <-result:
|
||||
if ok {
|
||||
t.Fatal("Then WaitForAgent must report no attached agent")
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Then revocation must wake WaitForAgent")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseStreamDetachesBeforeReenteringEndpointCloseAndJoinsErrors(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
if err := handler.CreateStream("close-errors", 1, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
firstErr := errors.New("first close error")
|
||||
secondErr := errors.New("second close error")
|
||||
if err := handler.UserConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: firstErr}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := handler.AgentConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: secondErr}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
err := handler.CloseStream("close-errors")
|
||||
if !errors.Is(err, firstErr) || !errors.Is(err, secondErr) {
|
||||
t.Fatalf("close errors were not joined: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartStreamReturnsImmediatelyWhenRevoked(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
if err := handler.CreateStream("start-revoked", 1, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result := make(chan error, 1)
|
||||
go func() { result <- handler.StartStream("start-revoked", time.Minute) }()
|
||||
if revoked := handler.RevokeStreamsForPurpose(PurposeLegacy); revoked != 1 {
|
||||
t.Fatalf("revoked streams: %d", revoked)
|
||||
}
|
||||
select {
|
||||
case err := <-result:
|
||||
if err == nil {
|
||||
t.Fatal("revoked StartStream must return an error")
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("StartStream did not wake on revoke")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentCloseAndRevokePublishOneGeneration(t *testing.T) {
|
||||
handler := NewNezhaHandler()
|
||||
if err := handler.CreateStream("single-generation", 1, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
start := handler.SnapshotIOStreamState()
|
||||
closeDone := make(chan struct{})
|
||||
revokeDone := make(chan struct{})
|
||||
go func() {
|
||||
_ = handler.CloseStream("single-generation")
|
||||
close(closeDone)
|
||||
}()
|
||||
go func() {
|
||||
handler.RevokeStreamsForPurpose(PurposeLegacy)
|
||||
close(revokeDone)
|
||||
}()
|
||||
select {
|
||||
case <-closeDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("CloseStream did not complete")
|
||||
}
|
||||
select {
|
||||
case <-revokeDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("RevokeStreamsForPurpose did not complete")
|
||||
}
|
||||
state := handler.SnapshotIOStreamState()
|
||||
if state.Count != 0 || state.Generation != start.Generation+1 {
|
||||
t.Fatalf("single detach publication: start=%+v final=%+v", start, state)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user