istio.io/istio@v0.0.0-20240520182934-d79c90f27776/pilot/test/xdstest/grpc.go (about)

     1  // Copyright Istio Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package xdstest
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	"google.golang.org/grpc"
    22  
    23  	"istio.io/istio/pkg/log"
    24  	"istio.io/istio/pkg/sleep"
    25  )
    26  
    27  type slowClientStream struct {
    28  	grpc.ClientStream
    29  	recv, send time.Duration
    30  }
    31  
    32  func (w *slowClientStream) RecvMsg(m any) error {
    33  	if w.recv > 0 {
    34  		sleep.UntilContext(w.Context(), w.recv)
    35  		log.Infof("delayed recv for %v", w.recv)
    36  	}
    37  	return w.ClientStream.RecvMsg(m)
    38  }
    39  
    40  func (w *slowClientStream) SendMsg(m any) error {
    41  	if w.send > 0 {
    42  		sleep.UntilContext(w.Context(), w.send)
    43  		log.Infof("delayed send for %v", w.send)
    44  	}
    45  	return w.ClientStream.SendMsg(m)
    46  }
    47  
    48  // SlowClientInterceptor is an interceptor that allows injecting delays on Send and Recv
    49  func SlowClientInterceptor(recv, send time.Duration) grpc.StreamClientInterceptor {
    50  	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn,
    51  		method string, streamer grpc.Streamer, opts ...grpc.CallOption,
    52  	) (grpc.ClientStream, error) {
    53  		clientStream, err := streamer(ctx, desc, cc, method, opts...)
    54  		return &slowClientStream{clientStream, recv, send}, err
    55  	}
    56  }
    57  
    58  type slowServerStream struct {
    59  	grpc.ServerStream
    60  	recv, send time.Duration
    61  }
    62  
    63  func (w *slowServerStream) RecvMsg(m any) error {
    64  	if w.recv > 0 {
    65  		sleep.UntilContext(w.Context(), w.recv)
    66  		log.Infof("delayed recv for %v", w.recv)
    67  	}
    68  	return w.ServerStream.RecvMsg(m)
    69  }
    70  
    71  func (w *slowServerStream) SendMsg(m any) error {
    72  	if w.send > 0 {
    73  		sleep.UntilContext(w.Context(), w.send)
    74  		log.Infof("delayed send for %v", w.send)
    75  	}
    76  	return w.ServerStream.SendMsg(m)
    77  }
    78  
    79  // SlowServerInterceptor is an interceptor that allows injecting delays on Send and Recv
    80  func SlowServerInterceptor(recv, send time.Duration) grpc.StreamServerInterceptor {
    81  	return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
    82  		return handler(srv, &slowServerStream{ss, recv, send})
    83  	}
    84  }