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
	}
	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
	}
	resp, err := sendStreamOpenServerLogical(ctx, s, logical, req)
	if err != nil {
		return nil, err
	}
	if resp.DataID != 0 {
		req.DataID = resp.DataID
	}
	if resp.FastPathVersion != 0 {
		req.FastPathVersion = resp.FastPathVersion
	} else {
		req.FastPathVersion = streamFastPathVersionV1
	}
	req.Metadata = mergeStreamMetadata(req.Metadata, resp.Metadata)
	transport := logical.CurrentTransportConn()
	stream := newStreamHandle(logical.stopContextSnapshot(), runtime, scope, req, 0, logical, transport, resp.TransportGeneration, serverStreamCloseSender(s, logical, nil), serverStreamResetSender(s, logical, nil), serverStreamDataSender(s, transport), runtime.configSnapshot())
	if err := runtime.register(scope, stream); err != nil {
		_, _ = sendStreamResetServerLogical(context.Background(), s, logical, StreamResetRequest{
			StreamID: req.StreamID,
			Error:    err.Error(),
		})
		return nil, err
	}
	return stream, nil
}

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
	}
	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
	}
	resp, err := sendStreamOpenServerTransport(ctx, s, transport, req)
	if err != nil {
		return nil, err
	}
	if resp.DataID != 0 {
		req.DataID = resp.DataID
	}
	if resp.FastPathVersion != 0 {
		req.FastPathVersion = resp.FastPathVersion
	} else {
		req.FastPathVersion = streamFastPathVersionV1
	}
	req.Metadata = mergeStreamMetadata(req.Metadata, resp.Metadata)
	stream := newStreamHandle(logical.stopContextSnapshot(), runtime, scope, req, 0, logical, transport, resp.TransportGeneration, serverStreamCloseSender(s, logical, transport), serverStreamResetSender(s, logical, transport), serverStreamDataSender(s, transport), runtime.configSnapshot())
	if err := runtime.register(scope, stream); err != nil {
		_, _ = sendStreamResetServerTransport(context.Background(), s, transport, StreamResetRequest{
			StreamID: req.StreamID,
			Error:    err.Error(),
		})
		return nil, err
	}
	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(),
			Full:     full,
		}
		if logical != nil {
			_, err := sendStreamCloseServerLogical(ctx, s, logical, req)
			return err
		}
		_, err := sendStreamCloseServerTransport(ctx, s, transport, 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(),
			Error:    message,
		}
		if logical != nil {
			_, err := sendStreamResetServerLogical(ctx, s, logical, req)
			return err
		}
		_, err := sendStreamResetServerTransport(ctx, s, transport, req)
		return err
	}
}

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))
	}
}