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