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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
package sftp
/*
Imported from: https://github.com/pkg/sftp
*/
import (
"sync"
)
type allocator struct {
sync.Mutex
available [][]byte
// map key is the request order
used map[uint32][][]byte
}
func newAllocator() *allocator {
return &allocator{
// micro optimization: initialize available pages with an initial capacity
available: make([][]byte, 0, SftpServerWorkerCount*2),
used: make(map[uint32][][]byte),
}
}
// GetPage returns a previously allocated and unused []byte or create a new one.
// The slice have a fixed size = maxMsgLength, this value is suitable for both
// receiving new packets and reading the files to serve
func (a *allocator) GetPage(requestOrderID uint32) []byte {
a.Lock()
defer a.Unlock()
var result []byte
// get an available page and remove it from the available ones.
if len(a.available) > 0 {
truncLength := len(a.available) - 1
result = a.available[truncLength]
a.available[truncLength] = nil // clear out the internal pointer
a.available = a.available[:truncLength] // truncate the slice
}
// no preallocated slice found, just allocate a new one
if result == nil {
result = make([]byte, maxMsgLength)
}
// put result in used pages
a.used[requestOrderID] = append(a.used[requestOrderID], result)
return result
}
// ReleasePages marks unused all pages in use for the given requestID
func (a *allocator) ReleasePages(requestOrderID uint32) {
a.Lock()
defer a.Unlock()
if used := a.used[requestOrderID]; len(used) > 0 {
a.available = append(a.available, used...)
}
delete(a.used, requestOrderID)
}
// Free removes all the used and available pages.
// Call this method when the allocator is not needed anymore
func (a *allocator) Free() {
a.Lock()
defer a.Unlock()
a.available = nil
a.used = make(map[uint32][][]byte)
}
func (a *allocator) countUsedPages() int {
a.Lock()
defer a.Unlock()
num := 0
for _, p := range a.used {
num += len(p)
}
return num
}
func (a *allocator) countAvailablePages() int {
a.Lock()
defer a.Unlock()
return len(a.available)
}
func (a *allocator) isRequestOrderIDUsed(requestOrderID uint32) bool {
a.Lock()
defer a.Unlock()
_, ok := a.used[requestOrderID]
return ok
}
|
