blob: 68b517d2fb1e859c49b5fb738657be3edafb5804 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
|
package llm
import "sync"
type pool[K comparable] struct {
devices map[K]bool
mutex sync.Mutex
cond *sync.Cond
}
func newPool[K comparable](devices ...K) *pool[K] {
p := &pool[K]{
devices: make(map[K]bool),
}
for _, d := range devices {
p.devices[d] = true
}
p.cond = sync.NewCond(&p.mutex)
return p
}
func (p *pool[K]) Len() int {
return len(p.devices)
}
func (p *pool[K]) Acquire() K {
p.mutex.Lock()
defer p.mutex.Unlock()
for {
for k, v := range p.devices {
if v {
p.devices[k] = false
return k
}
}
p.cond.Wait()
}
}
func (p *pool[K]) Release(device K) {
p.mutex.Lock()
p.devices[device] = true
p.cond.Signal()
p.mutex.Unlock()
}
|
