stats.go raw

   1  /*
   2   *
   3   * Copyright 2025 gRPC authors.
   4   *
   5   * Licensed under the Apache License, Version 2.0 (the "License");
   6   * you may not use this file except in compliance with the License.
   7   * You may obtain a copy of the License at
   8   *
   9   *     http://www.apache.org/licenses/LICENSE-2.0
  10   *
  11   * Unless required by applicable law or agreed to in writing, software
  12   * distributed under the License is distributed on an "AS IS" BASIS,
  13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14   * See the License for the specific language governing permissions and
  15   * limitations under the License.
  16   *
  17   */
  18  
  19  package stats
  20  
  21  import (
  22  	"context"
  23  
  24  	"google.golang.org/grpc/stats"
  25  )
  26  
  27  type combinedHandler struct {
  28  	handlers []stats.Handler
  29  }
  30  
  31  // NewCombinedHandler combines multiple stats.Handlers into a single handler.
  32  //
  33  // It returns nil if no handlers are provided. If only one handler is
  34  // provided, it is returned directly without wrapping.
  35  func NewCombinedHandler(handlers ...stats.Handler) stats.Handler {
  36  	switch len(handlers) {
  37  	case 0:
  38  		return nil
  39  	case 1:
  40  		return handlers[0]
  41  	default:
  42  		return &combinedHandler{handlers: handlers}
  43  	}
  44  }
  45  
  46  func (ch *combinedHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
  47  	for _, h := range ch.handlers {
  48  		ctx = h.TagRPC(ctx, info)
  49  	}
  50  	return ctx
  51  }
  52  
  53  func (ch *combinedHandler) HandleRPC(ctx context.Context, stats stats.RPCStats) {
  54  	for _, h := range ch.handlers {
  55  		h.HandleRPC(ctx, stats)
  56  	}
  57  }
  58  
  59  func (ch *combinedHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
  60  	for _, h := range ch.handlers {
  61  		ctx = h.TagConn(ctx, info)
  62  	}
  63  	return ctx
  64  }
  65  
  66  func (ch *combinedHandler) HandleConn(ctx context.Context, stats stats.ConnStats) {
  67  	for _, h := range ch.handlers {
  68  		h.HandleConn(ctx, stats)
  69  	}
  70  }
  71