client/stream.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package client

import (
	"context"

	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/remotecli"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/streaming"
)

// Streaming client streaming interface for code generate
type Streaming interface {
	Stream(ctx context.Context, method string, request, response interface{}) error
}

// Stream implements the Streaming interface
func (kc *kClient) Stream(ctx context.Context, method string, request, response interface{}) error {
	if !kc.inited {
		panic("client not initialized")
	}
	if kc.closed {
		panic("client is already closed")
	}
	if ctx == nil {
		panic("ctx is nil")
	}
	var ri rpcinfo.RPCInfo
	ctx, ri = kc.initRPCInfo(ctx, method)

	rpcinfo.AsMutableRPCConfig(ri.Config()).SetInteractionMode(rpcinfo.Streaming)
	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)

	ctx = kc.opt.TracerCtl.DoStart(ctx, ri)
	return kc.sEps(ctx, request, response)
}

func (kc *kClient) invokeStreamingEndpoint() (endpoint.Endpoint, error) {
	handler, err := kc.opt.RemoteOpt.CliHandlerFactory.NewTransHandler(kc.opt.RemoteOpt)
	if err != nil {
		return nil, err
	}
	for _, h := range kc.opt.MetaHandlers {
		if shdlr, ok := h.(remote.StreamingMetaHandler); ok {
			kc.opt.RemoteOpt.StreamingMetaHandlers = append(kc.opt.RemoteOpt.StreamingMetaHandlers, shdlr)
		}
	}
	return func(ctx context.Context, req, resp interface{}) (err error) {
		// req and resp as &streaming.Stream
		ri := rpcinfo.GetRPCInfo(ctx)
		st, err := remotecli.NewStream(ctx, ri, handler, kc.opt.RemoteOpt)
		if err != nil {
			return
		}
		st = &stream{stream: st, kc: kc}
		resp.(*streaming.Result).Stream = st
		return
	}, nil
}

type stream struct {
	stream streaming.Stream
	kc     *kClient
}

func (s *stream) Context() context.Context {
	return s.stream.Context()
}

func (s *stream) RecvMsg(m interface{}) error {
	return s.stream.RecvMsg(m)
}

func (s *stream) SendMsg(m interface{}) error {
	return s.stream.SendMsg(m)
}

func (s *stream) Close() error {
	ctx := s.stream.Context()
	ri := rpcinfo.GetRPCInfo(ctx)
	s.kc.opt.TracerCtl.DoFinish(ctx, ri, nil)
	return s.stream.Close()
}