1// SPDX-License-Identifier: MIT
2
3package cachesvc
4
5import (
6 "net"
7 "sync"
8 "time"
9)
10
11// bandwidth is the aggregate rate limit of one channel: a token bucket over
12// time, shared by every connection of the sandbox and by both directions, and
13// applied to the bytes read from the guest (request bodies) and written to it
14// (replies, archives).
15// The cache stream of a micro-VM does not cross its network interface, so the
16// interface's rate limit does not see it; this one stands in for it.
17//
18// A transfer is slowed, never cut: each Read or Write moves at most chunk bytes
19// and then sleeps for what those bytes cost, so a throttled transfer keeps
20// making progress in small steps (the progress deadlines of the handlers and the
21// relay's idle cut of the guest see activity). Sleeping ends when the connection
22// is closed.
23type bandwidth struct {
24 rate float64 // bytes per second
25 chunk int // largest Read or Write step, in bytes
26
27 mu sync.Mutex
28 free time.Time // when the bucket is empty: the cost of the bytes taken so far
29}
30
31const maxChunk = 32 << 10
32
33// newBandwidth returns a limiter of bytesPerSec, or nil (no limit) when it is
34// not positive.
35func newBandwidth(bytesPerSec int64) *bandwidth {
36 if bytesPerSec <= 0 {
37 return nil
38 }
39 chunk := int(min(max(bytesPerSec/10, 512), maxChunk))
40 // No burst: idle time earns no credit, so a steady transfer stays at the
41 // rate and an idle channel does not bank one.
42 return &bandwidth{rate: float64(bytesPerSec), chunk: chunk}
43}
44
45// mbitToBytes converts the Mbit/s of the sandbox configuration (same unit and
46// meaning as the VM's eth0 limit) to bytes per second; 0 or less is no limit.
47func mbitToBytes(mbit int) int64 {
48 if mbit <= 0 {
49 return 0
50 }
51 return int64(mbit) * 1_000_000 / 8
52}
53
54// reserve takes n bytes of the bucket and returns how long to wait before
55// moving them.
56func (b *bandwidth) reserve(n int) time.Duration {
57 now := time.Now()
58 b.mu.Lock()
59 defer b.mu.Unlock()
60 if b.free.Before(now) {
61 b.free = now
62 }
63 b.free = b.free.Add(time.Duration(float64(n) / b.rate * float64(time.Second)))
64 return b.free.Sub(now)
65}
66
67// wait pays for n bytes: it returns when they are due, or with net.ErrClosed
68// as soon as done is closed.
69func (b *bandwidth) wait(done <-chan struct{}, n int) error {
70 d := b.reserve(n)
71 if d <= 0 {
72 return nil
73 }
74 t := time.NewTimer(d)
75 defer t.Stop()
76 select {
77 case <-t.C:
78 return nil
79 case <-done:
80 return net.ErrClosed
81 }
82}
83
84// throttleListener hands out connections that pay into b.
85type throttleListener struct {
86 net.Listener
87 b *bandwidth
88}
89
90func (l throttleListener) Accept() (net.Conn, error) {
91 c, err := l.Listener.Accept()
92 if err != nil {
93 return nil, err
94 }
95 return &throttledConn{Conn: c, b: l.b, done: make(chan struct{})}, nil
96}
97
98// throttledConn moves its bytes in steps of at most b.chunk and pays for each
99// step once it has moved, after the Read or Write itself: the deadlines of the
100// handlers, set before a call, then cover the transfer and not the wait, and a
101// step costs at most chunk/rate (about 0.1 s) per connection ahead in line. It
102// does not implement ReadFrom/WriteTo, so nothing bypasses the steps (no
103// sendfile).
104type throttledConn struct {
105 net.Conn
106 b *bandwidth
107 once sync.Once
108 done chan struct{}
109}
110
111func (c *throttledConn) Read(p []byte) (int, error) {
112 if len(p) > c.b.chunk {
113 p = p[:c.b.chunk]
114 }
115 n, err := c.Conn.Read(p)
116 if n > 0 {
117 if werr := c.b.wait(c.done, n); werr != nil {
118 return 0, werr
119 }
120 }
121 return n, err
122}
123
124func (c *throttledConn) Write(p []byte) (int, error) {
125 var sent int
126 for len(p) > 0 {
127 step := p
128 if len(step) > c.b.chunk {
129 step = step[:c.b.chunk]
130 }
131 n, err := c.Conn.Write(step)
132 sent += n
133 if err != nil {
134 return sent, err
135 }
136 if err := c.b.wait(c.done, n); err != nil {
137 return sent, err
138 }
139 p = p[len(step):]
140 }
141 return sent, nil
142}
143
144func (c *throttledConn) Close() error {
145 c.once.Do(func() { close(c.done) })
146 return c.Conn.Close()
147}
148
149// ServeOption configures a channel (Serve).
150type ServeOption func(*serveConfig)
151
152type serveConfig struct {
153 bytesPerSec int64
154}
155
156// WithBandwidth limits the channel to mbit Mbit/s, summed over all its
157// connections and both directions. Zero or less means no limit, as for the VM
158// interface.
159//
160// The cache is the one stream of a micro-VM that does not cross its network
161// interface, so this is its share of the sandbox's NetMbit.
162func WithBandwidth(mbit int) ServeOption {
163 return func(c *serveConfig) { c.bytesPerSec = mbitToBytes(mbit) }
164}
165
166// withRate is WithBandwidth in bytes per second (tests use small rates).
167func withRate(bytesPerSec int64) ServeOption {
168 return func(c *serveConfig) { c.bytesPerSec = bytesPerSec }
169}