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 }