fix: connections stats were reset at each request
Massimo Melina committed
Mar 2, 2022 at 11:43 UTC
70510579dc4bca3427118bcfdb6ed9a26a7a96e1
3 files changed
+25
-7
src/ThrottledStream.ts
+6
-1
@@ -9,8 +9,13 @@ export class ThrottledStream extends Transform {
9
private lastSpeedTime = Date.now()
10
private totalSent: number = 0
11
12
- constructor(private group: ThrottleGroup) {
12
+ constructor(private group: ThrottleGroup, copyStats?: ThrottledStream) {
13
super()
14
+ if (!copyStats) return
15
+ this.sent = copyStats.sent
16
+ this.totalSent = copyStats.totalSent
17
+ this.lastSpeedTime = copyStats.lastSpeedTime
18
+ this.lastSpeed = copyStats.lastSpeed
19
}
20
21
async _transform(chunk: any, encoding: BufferEncoding, done: TransformCallback) {
src/connections.ts
+1
@@ -8,6 +8,7 @@ export interface Connection {
8
got: number
9
sent: number
10
outSpeed?: number
11
+ [rest:symbol]: any, // let other modules add extra data, but using symbols to avoid name collision
12
}
13
14
const all: Connection[] = []
src/throttler.ts
+18
-6
@@ -17,6 +17,8 @@ const ip2group: Record<string, {
17
destroy: () => void
18
}> = {}
19
20
+const SymThrStr = Symbol('stream')
21
+
22
export function throttler(): Koa.Middleware {
23
return async (ctx, next) => {
24
await next()
@@ -29,19 +31,29 @@ export function throttler(): Koa.Middleware {
31
tg.updateLimit(v ?? Infinity))
32
return { group:tg, count:0, destroy: unsub }
33
})
32
- const ts = new ThrottledStream(ipGroup.group)
34
+ const conn = socket2connection(ctx.socket)
35
+ if (!conn) throw 'assert throttler connection'
36
+
37
+ const ts = conn[SymThrStr] = new ThrottledStream(ipGroup.group, conn[SymThrStr])
38
+
39
+ const DELAY = 1000
40
+ const update = _.debounce(() => {
41
+ updateConnection(conn, {
42
+ sent: ts.getBytesSent(),
43
+ outSpeed: _.round(ts.getSpeed(), 1)
44
+ })
45
+ },
46
+ DELAY, { maxWait:DELAY })
47
+ ts.on('sent', update)
48
+
49
++ipGroup.count
50
ts.on('close', ()=> {
51
+ update.flush()
52
if (--ipGroup.count) return // any left?
53
ipGroup.destroy?.()
54
delete ip2group[ctx.ip]
55
})
56
40
- const conn = socket2connection(ctx.socket)
41
- if (conn)
42
- ts.on('sent', _.debounce(() =>
43
- updateConnection(conn, { sent: ts.getBytesSent(), outSpeed: _.round(ts.getSpeed(),1) }), 1000, { maxWait:1000 }))
44
-
57
const bak = ctx.response.length // preserve
58
ctx.body = ctx.body.pipe(ts)
59