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 }