nezha_domains/service/rpc/io_stream_lifecycle.go

package rpc

import (
	"context"
	"errors"
	"io"
	"time"
)

type ioStreamDetach struct {
	stream    *ioStreamContext
	endpoints []io.ReadWriteCloser
}

func detachStreamLocked(streamID string, retainedStream *ioStreamContext, streams map[string]*ioStreamContext) (ioStreamDetach, bool) {
	current, live := streams[streamID]
	if streamID == "" || !live || current != retainedStream {
		return ioStreamDetach{}, false
	}
	retainedStream.revoke()
	endpoints := make([]io.ReadWriteCloser, 0, 2)
	if retainedStream.userIo != nil {
		endpoints = append(endpoints, retainedStream.userIo)
	}
	if retainedStream.agentIo != nil {
		endpoints = append(endpoints, retainedStream.agentIo)
	}
	delete(streams, streamID)
	return ioStreamDetach{stream: retainedStream, endpoints: endpoints}, true
}

func (s *NezhaHandler) detachExactStream(streamID string, retainedStream *ioStreamContext) error {
	s.ioStreamMutex.Lock()
	detached, ok := detachStreamLocked(streamID, retainedStream, s.ioStreams)
	if !ok {
		s.ioStreamMutex.Unlock()
		return nil
	}
	s.publishIOStreamStateChangeLocked()
	s.ioStreamMutex.Unlock()

	var closeErrors []error
	for _, endpoint := range detached.endpoints {
		if err := endpoint.Close(); err != nil {
			closeErrors = append(closeErrors, err)
		}
	}
	return errors.Join(closeErrors...)
}

func (s *NezhaHandler) detachStreams(shouldDetach func(*ioStreamContext) bool) (int, error) {
	s.ioStreamMutex.Lock()
	detached := make([]ioStreamDetach, 0)
	for streamID, stream := range s.ioStreams {
		if !shouldDetach(stream) {
			continue
		}
		item, ok := detachStreamLocked(streamID, stream, s.ioStreams)
		if ok {
			detached = append(detached, item)
		}
	}
	if len(detached) > 0 {
		s.publishIOStreamStateChangeLocked()
	}
	s.ioStreamMutex.Unlock()

	var closeErrors []error
	for _, item := range detached {
		// Registry publication must precede endpoint Close so Close implementations may reenter safely.
		for _, endpoint := range item.endpoints {
			if err := endpoint.Close(); err != nil {
				closeErrors = append(closeErrors, err)
			}
		}
	}
	return len(detached), errors.Join(closeErrors...)
}

func (s *NezhaHandler) CloseStream(streamID string) error {
	_, err := s.detachStreams(func(stream *ioStreamContext) bool {
		return stream != nil && streamID != "" && stream == s.ioStreams[streamID]
	})
	return err
}

func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamID string, timeout time.Duration) (io.ReadWriteCloser, bool) {
	deadline := time.NewTimer(timeout)
	defer deadline.Stop()
	for {
		s.ioStreamMutex.RLock()
		stream, ok := s.ioStreams[streamID]
		if ok && stream.agentIo != nil {
			agentIo := stream.agentIo
			s.ioStreamMutex.RUnlock()
			return agentIo, true
		}
		if !ok {
			s.ioStreamMutex.RUnlock()
			return nil, false
		}
		revokedCh := stream.revokedCh
		agentIoConnectCh := stream.agentIoConnectCh
		stream.waitStartedOnce.Do(func() { close(stream.waitStartedCh) })
		s.ioStreamMutex.RUnlock()
		select {
		case <-ctx.Done():
			return nil, false
		case <-deadline.C:
			return nil, false
		case <-revokedCh:
			return nil, false
		case <-agentIoConnectCh:
		}
	}
}

func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) {
	if serverID == 0 {
		return
	}
	_, _ = s.detachStreams(func(stream *ioStreamContext) bool { return stream.targetServerID == serverID })
}

func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int {
	revoked, _ := s.detachStreams(func(stream *ioStreamContext) bool { return stream.purpose == purpose })
	return revoked
}