/util/fdlimiter.go

  1// Copyright 2009 The Go Authors. All rights reserved.
  2// Use of this source code is governed by a BSD-style
  3// license that can be found in the LICENSE file.
  4
  5package util
  6
  7import (
  8	"errors"
  9	"sync"
 10	"time"
 11)
 12
 13var ErrTimeout = errors.New("timeout")
 14
 15// FDLimiter helps keep track of the number of file descriptors in use.
 16type FDLimiter struct {
 17	limit int
 18	count int
 19	lk    sync.Mutex
 20	ch    chan int
 21	nfych chan<- int
 22}
 23
 24// Init initializes (or resets) an FDLimiter object.
 25func (fdl *FDLimiter) Init(fdlim int) {
 26	fdl.lk.Lock()
 27	if fdlim <= 0 {
 28		panic("FDLimiter, bad limit")
 29	}
 30	fdl.limit = fdlim
 31	fdl.count = 0
 32	fdl.ch = make(chan int)
 33	fdl.lk.Unlock()
 34}
 35
 36// SetNotifyChan instructs the FDLimiter to send the current
 37// number of utilized file descriptors every time that number changes.
 38// Calling this method with a nil argument, removes the notify channel.
 39func (fdl *FDLimiter) SetNotifyChan(c chan<- int) {
 40	fdl.lk.Lock()
 41	fdl.nfych = c
 42	fdl.lk.Unlock()
 43}
 44
 45func (fdl *FDLimiter) notify() {
 46	if fdl.nfych != nil {
 47		fdl.nfych <- fdl.count
 48	}
 49}
 50
 51func (fdl *FDLimiter) LockCount() int {
 52	fdl.lk.Lock()
 53	defer fdl.lk.Unlock()
 54	return fdl.count
 55}
 56
 57func (fdl *FDLimiter) Limit() int { return fdl.limit }
 58
 59// Lock blocks until it can allocate one fd without violating the limit.
 60func (fdl *FDLimiter) Lock() {
 61	for {
 62		fdl.lk.Lock()
 63		if fdl.count < fdl.limit {
 64			fdl.count++
 65			fdl.notify()
 66			fdl.lk.Unlock()
 67			return
 68		}
 69		fdl.lk.Unlock()
 70		<-fdl.ch
 71	}
 72	panic("FDLimiter, unreachable")
 73}
 74
 75// LockOrTimeout proceeds as Lock, except that it returns an ErrTimeout
 76// error, if a lock cannot be obtained within ns nanoseconds.
 77func (fdl *FDLimiter) LockOrTimeout(ns int64) error {
 78	waitsofar := int64(0)
 79	for {
 80		// Try to get an fd
 81		fdl.lk.Lock()
 82		if fdl.count < fdl.limit {
 83			fdl.count++
 84			fdl.notify()
 85			fdl.lk.Unlock()
 86			return nil
 87		}
 88		fdl.lk.Unlock()
 89
 90		// Or, wait for an fd or timeout
 91		if waitsofar >= ns {
 92			return ErrTimeout
 93		}
 94		t0 := time.Now().UnixNano()
 95		alrm := alarmOnce(ns - waitsofar)
 96		select {
 97		case <-alrm:
 98		case <-fdl.ch:
 99		}
100		waitsofar += time.Now().UnixNano() - t0
101	}
102	panic("FDLimiter, unreachable")
103}
104
105func (fdl *FDLimiter) LockOrChan(ch <-chan interface{}) (msg interface{}, err error) {
106	for {
107		fdl.lk.Lock()
108		if fdl.count < fdl.limit {
109			fdl.count++
110			fdl.notify()
111			fdl.lk.Unlock()
112			return nil, nil
113		}
114		fdl.lk.Unlock()
115
116		select {
117		case msg = <-ch:
118			return msg, ErrTimeout
119		case <-fdl.ch:
120		}
121	}
122	panic("FDLimiter, unreachable")
123}
124
125// Call Unlock to indicate that a file descriptor has been released.
126func (fdl *FDLimiter) Unlock() {
127	fdl.lk.Lock()
128	if fdl.count <= 0 {
129		panic("FDLimiter")
130	}
131	fdl.count--
132	fdl.notify()
133	if fdl.count == fdl.limit-1 {
134		fdl.ch <- 1
135	}
136	fdl.lk.Unlock()
137}
138
139// alarmOnce sends "1" to the returned chan after ns nanoseconds
140func alarmOnce(ns int64) <-chan int {
141	backchan := make(chan int)
142	go func() {
143		time.Sleep(time.Duration(ns))
144		backchan <- 1
145		close(backchan)
146	}()
147	return backchan
148}