protocol/triple/triple_protocol/handler_stream_compat.go

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 triple_protocol

import (
	"context"
)

import (
	"github.com/dubbogo/grpc-go"
	"github.com/dubbogo/grpc-go/metadata"
)

import (
	"dubbo.apache.org/dubbo-go/v3/common/constant"
)

type compatHandlerStream struct {
	ctx  context.Context
	conn StreamingHandlerConn
}

func (c *compatHandlerStream) SetHeader(md metadata.MD) error {
	// todo(DMwangnima): add header method for streaming
	return nil
}

func (c *compatHandlerStream) SendHeader(md metadata.MD) error {
	// todo(DMwangnima): add header method for streaming
	return nil
}

func (c *compatHandlerStream) SetTrailer(md metadata.MD) {
	// todo(DMwangnima): add trailer method for streaming
}

func (c *compatHandlerStream) Context() context.Context {
	return c.ctx
}

func (c *compatHandlerStream) SetContext(ctx context.Context) {
	c.ctx = ctx
}

func (c *compatHandlerStream) SendMsg(m any) error {
	return c.conn.Send(m)
}

func (c *compatHandlerStream) RecvMsg(m any) error {
	return c.conn.Receive(m)
}

func NewCompatStreamHandler(
	procedure string,
	srv any,
	typ StreamType,
	streamFunc func(srv any, stream grpc.ServerStream) error,
	options ...HandlerOption,
) *Handler {
	config := newHandlerConfig(procedure, options)
	implementation := generateCompatStreamHandlerFunc(procedure, srv, streamFunc, config.Interceptor)
	protocolHandlers := config.newProtocolHandlers(typ)

	hdl := &Handler{
		spec:             config.newSpec(typ),
		implementations:  make(map[string]StreamingHandlerFunc, defaultImplementationsSize),
		protocolHandlers: protocolHandlers,
		allowMethod:      sortedAllowMethodValue(protocolHandlers),
		acceptPost:       sortedAcceptPostValue(protocolHandlers),
	}
	hdl.processImplementation(getIdentifier(config.Group, config.Version), implementation)

	return hdl
}

func generateCompatStreamHandlerFunc(
	procedure string,
	srv any,
	streamFunc func(any, grpc.ServerStream) error,
	interceptor Interceptor,
) StreamingHandlerFunc {
	implementation := func(ctx context.Context, conn StreamingHandlerConn) error {
		ctx = metadata.NewIncomingContext(ctx, metadata.MD(conn.ExportableHeader()))
		// staticcheck error: SA1029. Stub code generated by protoc-gen-go-triple makes use of "XXX_TRIPLE_GO_INTERFACE_NAME" directly
		//nolint:staticcheck
		ctx = context.WithValue(ctx, constant.TripleGoInterfaceName, procedure)
		stream := &compatHandlerStream{
			ctx:  ctx,
			conn: conn,
		}
		return streamFunc(srv, stream)
	}
	if interceptor != nil {
		implementation = interceptor.WrapStreamingHandler(implementation)
	}

	return implementation
}