service_inline.go

/*
 * Copyright 2023 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"
	"errors"
	"runtime/debug"
	"unsafe"

	"github.com/cloudwego/kitex/client/callopt"
	"github.com/cloudwego/kitex/internal/client"
	"github.com/cloudwego/kitex/pkg/consts"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

type ContextServiceInlineHandler interface {
	WriteMeta(cliCtx context.Context, req interface{}) (newCliCtx context.Context, err error)
	ReadMeta(cliCtx context.Context, resp interface{}) (err error)
}

type serviceInlineClient struct {
	svcInfo *serviceinfo.ServiceInfo
	mws     []endpoint.Middleware
	eps     endpoint.Endpoint
	opt     *client.Options

	inited bool
	closed bool

	// server info
	serverEps endpoint.Endpoint

	contextServiceInlineHandler ContextServiceInlineHandler
}

type ServerInitialInfo interface {
	BuildServiceInlineInvokeChain() endpoint.Endpoint
}

// NewServiceInlineClient creates a kitex.Client with the given ServiceInfo, it is from generated code.
func NewServiceInlineClient(svcInfo *serviceinfo.ServiceInfo, s ServerInitialInfo, opts ...Option) (Client, error) {
	if svcInfo == nil {
		return nil, errors.New("NewClient: no service info")
	}
	kc := &serviceInlineClient{}
	kc.svcInfo = svcInfo
	kc.opt = client.NewOptions(opts)
	kc.serverEps = s.BuildServiceInlineInvokeChain()
	if err := kc.init(); err != nil {
		_ = kc.Close()
		return nil, err
	}
	return kc, nil
}

func (kc *serviceInlineClient) SetContextServiceInlineHandler(simh ContextServiceInlineHandler) {
	kc.contextServiceInlineHandler = simh
}

func (kc *serviceInlineClient) init() (err error) {
	if err = kc.checkOptions(); err != nil {
		return err
	}
	ctx := kc.initContext()
	kc.initMiddlewares(ctx)
	kc.richRemoteOption()
	if err = kc.buildInvokeChain(); err != nil {
		return err
	}
	kc.inited = true
	return nil
}

func (kc *serviceInlineClient) checkOptions() (err error) {
	if kc.opt.Svr.ServiceName == "" {
		return errors.New("service name is required")
	}
	return nil
}

func (kc *serviceInlineClient) initContext() context.Context {
	ctx := context.Background()
	ctx = context.WithValue(ctx, endpoint.CtxEventBusKey, kc.opt.Bus)
	ctx = context.WithValue(ctx, endpoint.CtxEventQueueKey, kc.opt.Events)
	return ctx
}

func (kc *serviceInlineClient) initMiddlewares(ctx context.Context) {
	builderMWs := richMWsWithBuilder(ctx, kc.opt.MWBs)
	kc.mws = append(kc.mws, contextMW)
	kc.mws = append(kc.mws, builderMWs...)
}

// initRPCInfo initializes the RPCInfo structure and attaches it to context.
func (kc *serviceInlineClient) initRPCInfo(ctx context.Context, method string) (context.Context, rpcinfo.RPCInfo, *callopt.CallOptions) {
	return initRPCInfo(ctx, method, kc.opt, kc.svcInfo, 0, nil)
}

// Call implements the Client interface .
func (kc *serviceInlineClient) Call(ctx context.Context, method string, request, response interface{}) (err error) {
	validateForCall(ctx, kc.inited, kc.closed)
	var ri rpcinfo.RPCInfo
	var callOpts *callopt.CallOptions
	ctx, ri, callOpts = kc.initRPCInfo(ctx, method)

	ctx = kc.opt.TracerCtl.DoStart(ctx, ri)
	var reportErr error
	defer func() {
		if panicInfo := recover(); panicInfo != nil {
			reportErr = rpcinfo.ClientPanicToErr(ctx, panicInfo, ri, true)
		}
		kc.opt.TracerCtl.DoFinish(ctx, ri, reportErr)
		// If the user start a new goroutine and return before endpoint finished, it may cause panic.
		// For example,, if the user writes a timeout Middleware and times out, rpcinfo will be recycled,
		// but in fact, rpcinfo is still being used when it is executed inside
		// So if endpoint returns err, client won't recycle rpcinfo.
		if reportErr == nil {
			rpcinfo.PutRPCInfo(ri)
		}
		callOpts.Recycle()
	}()
	reportErr = kc.eps(ctx, request, response)

	if reportErr == nil {
		err = ri.Invocation().BizStatusErr()
	} else {
		err = reportErr
	}
	return err
}

func (kc *serviceInlineClient) richRemoteOption() {
	kc.opt.RemoteOpt.SvcInfo = kc.svcInfo
}

func (kc *serviceInlineClient) buildInvokeChain() error {
	innerHandlerEp, err := kc.invokeHandleEndpoint()
	if err != nil {
		return err
	}
	kc.eps = endpoint.Chain(kc.mws...)(innerHandlerEp)
	return nil
}

func (kc *serviceInlineClient) invokeHandleEndpoint() (endpoint.Endpoint, error) {
	return func(ctx context.Context, req, resp interface{}) (err error) {
		cliRpcInfo := rpcinfo.GetRPCInfo(ctx)
		if v, ok := cliRpcInfo.Invocation().(rpcinfo.InvocationSetter); ok {
			v.SetExtra(consts.SERVICE_INLINE_SERVICE_NAME, kc.svcInfo.ServiceName)
		}
		ctx = context.WithValue(ctx, consts.SERVICE_INLINE_RPCINFO_KEY, unsafe.Pointer(&cliRpcInfo))
		if kc.contextServiceInlineHandler != nil {
			ctx, err = kc.contextServiceInlineHandler.WriteMeta(ctx, req)
			if err != nil {
				return err
			}
		}

		// server logic
		err = kc.serverEps(ctx, req, resp)

		if kc.contextServiceInlineHandler != nil {
			err1 := kc.contextServiceInlineHandler.ReadMeta(ctx, resp)
			if err1 != nil {
				return err1
			}
		}
		return err
	}, nil
}

// Close is not concurrency safe.
func (kc *serviceInlineClient) Close() error {
	defer func() {
		if err := recover(); err != nil {
			klog.Warnf("KITEX: panic when close client, error=%s, stack=%s", err, string(debug.Stack()))
		}
	}()
	if kc.closed {
		return nil
	}
	kc.closed = true
	var errs utils.ErrChain
	for _, cb := range kc.opt.CloseCallbacks {
		if err := cb(); err != nil {
			errs.Append(err)
		}
	}
	if errs.HasError() {
		return errs
	}
	return nil
}