github.com/grafana/pyroscope@v1.18.0/pkg/util/delayhandler/connect.go (about)

     1  package delayhandler
     2  
     3  import (
     4  	"context"
     5  
     6  	"connectrpc.com/connect"
     7  )
     8  
     9  type delayInterceptor struct {
    10  	limits Limits
    11  }
    12  
    13  func NewConnect(limits Limits) connect.Interceptor {
    14  	return &delayInterceptor{limits: limits}
    15  }
    16  
    17  func (i *delayInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
    18  	return func(ctx context.Context, req connect.AnyRequest) (resp connect.AnyResponse, err error) {
    19  		start := timeNow()
    20  		delay := getDelay(ctx, i.limits)
    21  		delayCtx := context.Background()
    22  		if delay > 0 {
    23  			var cancel context.CancelFunc
    24  			delayCtx, cancel = context.WithCancel(context.Background())
    25  			defer cancel()
    26  			ctx = WithDelayCancel(ctx, cancel)
    27  		}
    28  
    29  		// now run the chain after me
    30  		resp, err = next(ctx, req)
    31  
    32  		// if there is an error, return it immediately
    33  		if err != nil {
    34  			return resp, err
    35  		}
    36  
    37  		// The delay has been cancelled down the chain.
    38  		if delayCtx.Err() != nil {
    39  			return resp, err
    40  		}
    41  
    42  		// no delay, return immediately
    43  		if delay <= 0 {
    44  			return resp, err
    45  		}
    46  
    47  		delayLeft := delay - timeNow().Sub(start)
    48  
    49  		// if the delay is already expired, return immediately
    50  		if delayLeft <= 0 {
    51  			return resp, err
    52  		}
    53  
    54  		// add delay header
    55  		addDelayHeader(resp.Header(), delayLeft)
    56  
    57  		// if the delay is not expired, sleep for the remaining time
    58  		<-timeAfter(delayLeft)
    59  
    60  		return resp, nil
    61  	}
    62  }
    63  
    64  // do nothing for streaming handlers
    65  func (delayInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
    66  	return func(ctx context.Context, conn connect.StreamingHandlerConn) (err error) {
    67  		panic("delayInterceptor not implemented")
    68  	}
    69  }
    70  
    71  // do nothing for streaming clients
    72  func (delayInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
    73  	panic("delayInterceptor not implemented")
    74  }