github.com/xmidt-org/webpa-common@v1.11.9/middleware/concurrent.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/go-kit/kit/endpoint"
     7  )
     8  
     9  // Concurrent produces a middleware that allows only a set number of concurrent calls via
    10  // a semaphore implemented as a buffered channel.  The context is used for cancellation,
    11  // and if the context is cancelled then timeoutError is returned if it is not nil, ctx.Err() otherwise.
    12  func Concurrent(concurrency int, timeoutError error) endpoint.Middleware {
    13  	return func(next endpoint.Endpoint) endpoint.Endpoint {
    14  		semaphore := make(chan struct{}, concurrency)
    15  		for r := 0; r < concurrency; r++ {
    16  			semaphore <- struct{}{}
    17  		}
    18  
    19  		return func(ctx context.Context, value interface{}) (interface{}, error) {
    20  			select {
    21  			case <-ctx.Done():
    22  				if timeoutError != nil {
    23  					return nil, timeoutError
    24  				} else {
    25  					return nil, ctx.Err()
    26  				}
    27  
    28  			case <-semaphore:
    29  			}
    30  
    31  			defer func() {
    32  				semaphore <- struct{}{}
    33  			}()
    34  
    35  			return next(ctx, value)
    36  		}
    37  	}
    38  }