-
Notifications
You must be signed in to change notification settings - Fork 21
/
expiring_orders.go
106 lines (92 loc) · 2.4 KB
/
expiring_orders.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
// Copyright (c) 2022 Gobalsky Labs Limited
//
// Use of this software is governed by the Business Source License included
// in the LICENSE.VEGA file and at https://www.mariadb.com/bsl11.
//
// Change Date: 18 months from the later of the date of the first publicly
// available Distribution of this version of the repository, and 25 June 2022.
//
// On the date above, in accordance with the Business Source License, use
// of this software will be governed by version 3 or later of the GNU General
// Public License.
package execution
import (
"github.com/google/btree"
)
type ExpiringOrders struct {
orders *btree.BTree
ordersChanged bool
}
type ordersAtTS struct {
ts int64
// order IDs
orders []string
}
func (a *ordersAtTS) Less(b btree.Item) bool {
return a.ts < b.(*ordersAtTS).ts
}
func NewExpiringOrders() *ExpiringOrders {
return &ExpiringOrders{
orders: btree.New(2),
ordersChanged: true,
}
}
func (a *ExpiringOrders) GetExpiryingOrderCount() int {
result := a.orders.Len()
return result
}
func (a *ExpiringOrders) Insert(
orderID string, ts int64,
) {
item := &ordersAtTS{ts: ts}
if item := a.orders.Get(item); item != nil {
item.(*ordersAtTS).orders = append(item.(*ordersAtTS).orders, orderID)
a.ordersChanged = true
return
}
item.orders = []string{orderID}
a.orders.ReplaceOrInsert(item)
a.ordersChanged = true
}
func (a *ExpiringOrders) RemoveOrder(expiresAt int64, orderID string) bool {
item := &ordersAtTS{ts: expiresAt}
if item := a.orders.Get(item); item != nil {
oat := item.(*ordersAtTS)
for i := 0; i < len(oat.orders); i++ {
if oat.orders[i] == orderID {
oat.orders = oat.orders[:i+copy(oat.orders[i:], oat.orders[i+1:])]
// if the slice is empty, remove the parent container
if len(oat.orders) == 0 {
a.orders.Delete(item)
a.ordersChanged = true
}
return true
}
}
}
return false
}
func (a *ExpiringOrders) Expire(ts int64) []string {
if a.orders.Len() == 0 {
return nil
}
orders := []string{}
toDelete := []int64{}
item := &ordersAtTS{ts: ts + 1}
a.orders.AscendLessThan(item, func(i btree.Item) bool {
if ts < i.(*ordersAtTS).ts {
return false
}
orders = append(orders, i.(*ordersAtTS).orders...)
toDelete = append(toDelete, i.(*ordersAtTS).ts)
return true
})
for _, v := range toDelete {
item.ts = v
a.orders.Delete(item)
}
if len(toDelete) > 0 {
a.ordersChanged = true
}
return orders
}