notify/server_stream.go

package notify

import "context"

func (s *ServerCommon) SetStreamHandler(fn func(StreamAcceptInfo) error) {
	runtime := s.getStreamRuntime()
	if runtime == nil {
		return
	}
	runtime.setHandler(fn)
}

func (s *ServerCommon) OpenStreamLogical(ctx context.Context, logical *LogicalConn, opt StreamOpenOptions) (Stream, error) {
	if s == nil {
		return nil, errStreamServerNil
	}
	if logical == nil {
		return nil, errStreamLogicalConnNil
	}
	return s.openStreamTransport(ctx, logical.CurrentTransportConn(), opt)
}

func (s *ServerCommon) OpenStreamTransport(ctx context.Context, transport *TransportConn, opt StreamOpenOptions) (Stream, error) {
	if s == nil {
		return nil, errStreamServerNil
	}
	if transport == nil {
		return nil, errStreamTransportNil
	}
	return s.openStreamTransport(ctx, transport, opt)
}

func (s *ServerCommon) openStreamTransport(ctx context.Context, transport *TransportConn, opt StreamOpenOptions) (Stream, error) {
	if s == nil {
		return nil, errStreamServerNil
	}
	if transport == nil {
		return nil, errStreamTransportNil
	}
	if err := s.ensureServerTransportSendReady(transport); err != nil {
		return nil, err
	}
	logical := transport.LogicalConn()
	if logical == nil {
		return nil, errStreamLogicalConnNil
	}
	runtime := s.getStreamRuntime()
	if runtime == nil {
		return nil, errStreamRuntimeNil
	}
	req := serverStreamRequest(runtime, opt)
	scope := serverFileScope(logical)
	if _, exists := runtime.lookup(scope, req.StreamID); exists {
		return nil, errStreamAlreadyExists
	}
	dataID, err := runtime.reserveDataID(scope)
	if err != nil {
		return nil, err
	}
	req.DataID = dataID
	stream := newStreamHandle(logical.stopContextSnapshot(), runtime, scope, req, 0, logical, transport, transport.TransportGeneration(), serverStreamCloseSender(s, logical, transport), serverStreamResetSender(s, logical, transport), serverStreamDataSender(s, transport), runtime.configSnapshot())
	if err := runtime.adoptReserved(scope, stream); err != nil {
		runtime.releaseDataID(scope, req.DataID)
		return nil, err
	}
	resp, err := sendStreamOpenServerTransport(ctx, s, transport, req)
	if err != nil {
		s.bestEffortStreamResetTransport(transport, StreamResetRequest{StreamID: req.StreamID, DataID: req.DataID, Error: err.Error()})
		stream.markReset(err)
		return nil, err
	}
	if resp.DataID != 0 && resp.DataID != req.DataID {
		err = errStreamAlreadyExists
		s.bestEffortStreamResetTransport(transport, StreamResetRequest{StreamID: req.StreamID, Error: "stream data id mismatch"})
		stream.markReset(err)
		return nil, err
	}
	if resp.FastPathVersion != 0 {
		stream.setFastPathVersion(resp.FastPathVersion)
	} else {
		stream.setFastPathVersion(streamFastPathVersionV1)
	}
	stream.metadata = mergeStreamMetadata(req.Metadata, resp.Metadata)
	if resp.TransportGeneration != 0 {
		stream.setTransportGeneration(resp.TransportGeneration)
	}
	return stream, nil
}

func serverStreamRequest(runtime *streamRuntime, opt StreamOpenOptions) StreamOpenRequest {
	id := opt.ID
	if id == "" && runtime != nil {
		id = runtime.nextID()
	}
	return normalizeStreamOpenRequest(StreamOpenRequest{
		StreamID:        id,
		FastPathVersion: streamFastPathVersionCurrent,
		Channel:         opt.Channel,
		Metadata:        cloneStreamMetadata(opt.Metadata),
		ReadTimeout:     opt.ReadTimeout,
		WriteTimeout:    opt.WriteTimeout,
	})
}

func serverStreamCloseSender(s *ServerCommon, logical *LogicalConn, transport *TransportConn) streamCloseSender {
	return func(ctx context.Context, stream *streamHandle, full bool) error {
		req := StreamCloseRequest{
			StreamID: stream.ID(),
			DataID:   stream.dataIDSnapshot(),
			Full:     full,
		}
		if transport != nil {
			_, err := sendStreamCloseServerTransport(ctx, s, transport, req)
			return err
		}
		_, err := sendStreamCloseServerLogical(ctx, s, logical, req)
		return err
	}
}

func serverStreamResetSender(s *ServerCommon, logical *LogicalConn, transport *TransportConn) streamResetSender {
	return func(ctx context.Context, stream *streamHandle, message string) error {
		req := StreamResetRequest{
			StreamID:      stream.ID(),
			DataID:        stream.dataIDSnapshot(),
			Error:         message,
			RecordFailure: stream.recordResetFailure(),
		}
		if transport != nil {
			_, err := sendStreamResetServerTransport(ctx, s, transport, req)
			return err
		}
		_, err := sendStreamResetServerLogical(ctx, s, logical, req)
		return err
	}
}

func (s *ServerCommon) bestEffortStreamResetTransport(transport *TransportConn, req StreamResetRequest) {
	if s == nil || transport == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), streamDispatchRejectTimeout)
	defer cancel()
	_, _ = sendStreamResetServerTransport(ctx, s, transport, req)
}

func serverStreamDataSender(s *ServerCommon, transport *TransportConn) streamDataSender {
	return func(ctx context.Context, stream *streamHandle, chunk []byte) error {
		if s == nil {
			return errStreamServerNil
		}
		if transport == nil {
			return errStreamTransportNil
		}
		if !transport.IsCurrent() {
			return errTransportDetached
		}
		if ctx != nil {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}
		}
		if dataID := stream.dataIDSnapshot(); dataID != 0 {
			return s.sendFastStreamDataTransport(ctx, stream.LogicalConn(), transport, stream, chunk)
		}
		return s.sendEnvelopeTransport(transport, newStreamDataEnvelope(stream.ID(), chunk))
	}
}