github.com/cloudwego/kitex@v0.9.0/internal/stream/stream_option.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package stream
    18  
    19  import (
    20  	"context"
    21  
    22  	"github.com/cloudwego/kitex/pkg/endpoint"
    23  	"github.com/cloudwego/kitex/pkg/stats"
    24  )
    25  
    26  // StreamEventHandler is used to handle stream events
    27  type StreamEventHandler func(ctx context.Context, evt stats.Event, err error)
    28  
    29  type StreamingConfig struct {
    30  	RecvMiddlewareBuilders []endpoint.RecvMiddlewareBuilder
    31  	RecvMiddlewares        []endpoint.RecvMiddleware
    32  
    33  	SendMiddlewareBuilders []endpoint.SendMiddlewareBuilder
    34  	SendMiddlewares        []endpoint.SendMiddleware
    35  
    36  	EventHandler StreamEventHandler
    37  }
    38  
    39  func (c *StreamingConfig) InitMiddlewares(ctx context.Context) {
    40  	if c.EventHandler != nil {
    41  		c.RecvMiddlewares = append(c.RecvMiddlewares, c.streamRecvTraceMW(ctx))
    42  		c.SendMiddlewares = append(c.SendMiddlewares, c.streamSendTraceMW(ctx))
    43  	}
    44  
    45  	if len(c.RecvMiddlewareBuilders) > 0 {
    46  		recvMiddlewares := make([]endpoint.RecvMiddleware, 0, len(c.RecvMiddlewareBuilders))
    47  		for _, mwb := range c.RecvMiddlewareBuilders {
    48  			recvMiddlewares = append(recvMiddlewares, mwb(ctx))
    49  		}
    50  		c.RecvMiddlewares = append(c.RecvMiddlewares, recvMiddlewares...)
    51  	}
    52  
    53  	if len(c.SendMiddlewareBuilders) > 0 {
    54  		sendMiddlewares := make([]endpoint.SendMiddleware, 0, len(c.SendMiddlewareBuilders))
    55  		for _, mwb := range c.SendMiddlewareBuilders {
    56  			sendMiddlewares = append(sendMiddlewares, mwb(ctx))
    57  		}
    58  		c.SendMiddlewares = append(c.SendMiddlewares, sendMiddlewares...)
    59  	}
    60  }
    61  
    62  func (c *StreamingConfig) BuildRecvInvokeChain(ep endpoint.RecvEndpoint) endpoint.RecvEndpoint {
    63  	return endpoint.RecvChain(c.RecvMiddlewares...)(ep)
    64  }
    65  
    66  func (c *StreamingConfig) BuildSendInvokeChain(ep endpoint.SendEndpoint) endpoint.SendEndpoint {
    67  	return endpoint.SendChain(c.SendMiddlewares...)(ep)
    68  }