-
Notifications
You must be signed in to change notification settings - Fork 70
/
pool.go
80 lines (67 loc) · 1.71 KB
/
pool.go
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
package port_allocator
import (
"encoding/json"
"errors"
)
var ErrorPortPoolExhausted = errors.New("port pool exhausted")
type Pool struct {
AcquiredPorts map[int]string
}
func (p *Pool) MarshalJSON() ([]byte, error) {
var jsonData struct {
AcquiredPorts map[string][]int `json:"acquired_ports"`
}
jsonData.AcquiredPorts = make(map[string][]int)
for port, handle := range p.AcquiredPorts {
jsonData.AcquiredPorts[handle] = append(jsonData.AcquiredPorts[handle], port)
}
return json.Marshal(jsonData)
}
func (p *Pool) UnmarshalJSON(bytes []byte) error {
var jsonData struct {
AcquiredPorts map[string][]int `json:"acquired_ports"`
}
err := json.Unmarshal(bytes, &jsonData)
if err != nil {
return err
}
p.AcquiredPorts = make(map[int]string)
for handle, ports := range jsonData.AcquiredPorts {
for _, port := range ports {
p.AcquiredPorts[port] = handle
}
}
return nil
}
type Tracker struct {
StartPort int
Capacity int
}
func (t *Tracker) InRange(port int) bool {
return port >= t.StartPort && port < t.StartPort+t.Capacity
}
func (t *Tracker) AcquireOne(pool *Pool, handler string) (int, error) {
if pool.AcquiredPorts == nil {
pool.AcquiredPorts = make(map[int]string)
}
for i := 0; i < t.Capacity; i++ {
candidatePort := t.StartPort + i
if !contains(pool.AcquiredPorts, candidatePort) {
pool.AcquiredPorts[candidatePort] = handler
return candidatePort, nil
}
}
return -1, ErrorPortPoolExhausted
}
func (t *Tracker) ReleaseAll(pool *Pool, handle string) error {
for port, h := range pool.AcquiredPorts {
if h == handle {
delete(pool.AcquiredPorts, port)
}
}
return nil
}
func contains(list map[int]string, candidate int) bool {
_, ok := list[candidate]
return ok
}