diff options
Diffstat (limited to 'vendor/google.golang.org/grpc/transport/flowcontrol.go')
-rw-r--r-- | vendor/google.golang.org/grpc/transport/flowcontrol.go | 10 |
1 files changed, 8 insertions, 2 deletions
diff --git a/vendor/google.golang.org/grpc/transport/flowcontrol.go b/vendor/google.golang.org/grpc/transport/flowcontrol.go index 378f5c450..bbf98b6f5 100644 --- a/vendor/google.golang.org/grpc/transport/flowcontrol.go +++ b/vendor/google.golang.org/grpc/transport/flowcontrol.go @@ -58,14 +58,20 @@ type writeQuota struct { ch chan struct{} // done is triggered in error case. done <-chan struct{} + // replenish is called by loopyWriter to give quota back to. + // It is implemented as a field so that it can be updated + // by tests. + replenish func(n int) } func newWriteQuota(sz int32, done <-chan struct{}) *writeQuota { - return &writeQuota{ + w := &writeQuota{ quota: sz, ch: make(chan struct{}, 1), done: done, } + w.replenish = w.realReplenish + return w } func (w *writeQuota) get(sz int32) error { @@ -83,7 +89,7 @@ func (w *writeQuota) get(sz int32) error { } } -func (w *writeQuota) replenish(n int) { +func (w *writeQuota) realReplenish(n int) { sz := int32(n) a := atomic.AddInt32(&w.quota, sz) b := a - sz |