pixie/flowcontrol.go at main · soap-phia/pixie · GitHub

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
package pixie

import (
	"context"
	"runtime"
	"sync/atomic"
)

type FlowControlMode uint8

type FlowControl struct {
	Mode         FlowControlMode
	Buffer       atomic.Uint32
	TargetBuffer uint32
	Wakeup       chan struct{}
}

const (
	fcDisabled FlowControlMode = iota
	fcTrackAmt
	fcSendMessages
)

func NewFlowControl(mode FlowControlMode, bufferSize uint32) *FlowControl {
	fc := &FlowControl{
		Mode:         mode,
		TargetBuffer: (bufferSize * 9) / 10,
		Wakeup:       make(chan struct{}, 1),
	}
	fc.Buffer.Store(bufferSize)
	return fc
}

func (fc *FlowControl) CanSend() bool {
	if fc.Mode != fcTrackAmt {
		return true
	}
	return fc.Buffer.Load() > 0
}

func (fc *FlowControl) Subtract() {
	if fc.Mode == fcDisabled {
		return
	}

	for {
		old := fc.Buffer.Load()
		if old == 0 {
			return
		}

		new := old - 1
		if fc.Buffer.CompareAndSwap(old, new) {
			return
		}

		runtime.Gosched()
	}
}

func (fc *FlowControl) Add(amount uint32) uint32 {
	return fc.Buffer.Add(amount)
}

func (fc *FlowControl) Set(amount uint32) {
	fc.Buffer.Store(amount)
	select {
	case fc.Wakeup <- struct{}{}:
	default:
	}
}

func (fc *FlowControl) Get() uint32 {
	return fc.Buffer.Load()
}

func (fc *FlowControl) ShouldContinue(readCount uint32) bool {
	if fc.Mode != fcSendMessages {
		return false
	}
	return readCount >= fc.TargetBuffer
}

func (fc *FlowControl) WaitBuffer(ctx context.Context) error {
	if fc.CanSend() {
		return nil
	}

	select {
	case <-fc.Wakeup:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (fc *FlowControl) Awaken() {
	select {
	case fc.Wakeup <- struct{}{}:
	default:
	}
}