/
cqueue.go
143 lines (129 loc) · 2.88 KB
/
cqueue.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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
// Copyright 2015 Aleksey Blinov. All rights reserved.
package psmtp
import (
"container/heap"
"net/smtp"
)
type cqItem struct {
conn *psmtpConn
index int
}
type _cQueue []*cqItem
type cPool struct {
capacity int
all _cQueue
byAuth map[smtp.Auth]_cQueue
}
func newCPool(capacity int) *cPool {
result := &cPool{
capacity: capacity,
all: make(_cQueue, 0, capacity),
byAuth: make(map[smtp.Auth]_cQueue, 0),
}
heap.Init(&result.all)
return result
}
func (this *cPool) Push(conn *psmtpConn) []*psmtpConn {
// Evict if needed
var result []*psmtpConn
if this == nil {
return result
}
if this.capacity > 0 {
for {
if this.all.Len() < this.capacity {
break
}
if v, ok := this.PopAny(); ok {
result = append(result, v)
}
}
item := cqItem{conn: conn}
heap.Push(&this.all, &item)
sub, ok := this.byAuth[conn.auth]
if !ok {
sub = make(_cQueue, 0, this.capacity)
heap.Init(&sub)
}
heap.Push(&sub, &item)
this.byAuth[conn.auth] = sub
} else {
result = append(result, conn)
}
return result
}
// PopAny pops least recently used connection of all
// and correctly updates any by-auth queues.
// Returns ok = false if the pool is empty.
func (this *cPool) PopAny() (*psmtpConn, bool) {
if this == nil {
return nil, false
}
if this.all.Len() == 0 {
return nil, false
}
result := heap.Pop(&this.all).(*cqItem)
sub, ok := this.byAuth[result.conn.auth]
if !ok {
panic("missign by-auth queue")
}
if sub.Len() == 0 {
panic("by-auth queue is empty")
}
check := heap.Pop(&sub).(*cqItem)
this.byAuth[result.conn.auth] = sub
if result.conn != check.conn {
panic("mismatched items in all and by-auth queues")
}
return result.conn, true
}
// PopAny pops least recently used connection
// with given auth and correctly updates global queue.
// Returns ok = false if no matching connection is available.
func (this *cPool) Pop(auth smtp.Auth) (*psmtpConn, bool) {
if this == nil {
return nil, false
}
sub, ok := this.byAuth[auth]
if !ok || sub.Len() == 0 {
return nil, false
}
result := heap.Pop(&sub).(*cqItem)
this.byAuth[auth] = sub
found := false
for i, v := range this.all {
if v == result {
found = true
heap.Remove(&this.all, i)
break
}
}
if !found {
panic("item missing from all queue")
}
return result.conn, true
}
func (this _cQueue) Len() int { return len(this) }
func (this _cQueue) Less(i, j int) bool {
// We want Pop to give us the least recently used.
return this[j].conn.atime.Before(this[i].conn.atime)
}
func (this _cQueue) Swap(i, j int) {
this[i], this[j] = this[j], this[i]
this[i].index = i
this[j].index = j
}
func (this *_cQueue) Push(x interface{}) {
n := len(*this)
item := x.(*cqItem)
item.index = n
*this = append(*this, item)
}
func (this *_cQueue) Pop() interface{} {
old := *this
n := len(old)
item := old[n-1]
item.index = -1 // for safety
*this = old[0 : n-1]
return item
}