notify/client_stream.go

package notify

import (
	"context"
	"net"
)

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

func (c *ClientCommon) OpenStream(ctx context.Context, opt StreamOpenOptions) (Stream, error) {
	if c == nil {
		return nil, errStreamClientNil
	}
	runtime := c.getStreamRuntime()
	if runtime == nil {
		return nil, errStreamRuntimeNil
	}
	req := clientStreamRequest(runtime, opt)
	if req.StreamID == "" {
		return nil, errStreamIDEmpty
	}
	if _, exists := runtime.lookup(clientFileScope(), req.StreamID); exists {
		return nil, errStreamAlreadyExists
	}
	resp, err := sendStreamOpenClient(ctx, c, 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(c.clientStopContextSnapshot(), runtime, clientFileScope(), req, c.currentClientSessionEpoch(), nil, nil, resp.TransportGeneration, clientStreamCloseSender(c), clientStreamResetSender(c), clientStreamDataSender(c, c.currentClientSessionEpoch()), runtime.configSnapshot())
	stream.setClientSnapshotOwner(c)
	stream.setAddrSnapshot(c.clientStreamAddrSnapshot())
	if err := runtime.register(clientFileScope(), stream); err != nil {
		_, _ = sendStreamResetClient(context.Background(), c, StreamResetRequest{
			StreamID: req.StreamID,
			Error:    err.Error(),
		})
		return nil, err
	}
	return stream, nil
}

func (c *ClientCommon) clientStreamAddrSnapshot() (net.Addr, net.Addr) {
	if c == nil {
		return nil, nil
	}
	conn := c.clientTransportConnSnapshot()
	if conn == nil {
		return nil, nil
	}
	return conn.LocalAddr(), conn.RemoteAddr()
}

func clientStreamRequest(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 clientStreamCloseSender(c *ClientCommon) streamCloseSender {
	return func(ctx context.Context, stream *streamHandle, full bool) error {
		_, err := sendStreamCloseClient(ctx, c, StreamCloseRequest{
			StreamID: stream.ID(),
			Full:     full,
		})
		return err
	}
}

func clientStreamResetSender(c *ClientCommon) streamResetSender {
	return func(ctx context.Context, stream *streamHandle, message string) error {
		_, err := sendStreamResetClient(ctx, c, StreamResetRequest{
			StreamID: stream.ID(),
			Error:    message,
		})
		return err
	}
}

func clientStreamDataSender(c *ClientCommon, epoch uint64) streamDataSender {
	return func(ctx context.Context, stream *streamHandle, chunk []byte) error {
		if c == nil {
			return errStreamClientNil
		}
		if epoch != 0 && !c.isClientSessionEpochCurrent(epoch) {
			return errTransportDetached
		}
		if ctx != nil {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
			}
		}
		if dataID := stream.dataIDSnapshot(); dataID != 0 {
			return c.sendFastStreamData(ctx, stream, chunk)
		}
		return c.sendEnvelope(newStreamDataEnvelope(stream.ID(), chunk))
	}
}