limit.go RawHistoryFolder

1// SPDX-License-Identifier: MIT
2
3package cachesvc
4
5import (
6	"net"
7	"sync"
8)
9
10// limitListener bounds the open connections of a listener (local) and of the
11// service (global, shared by all listeners). A slot is taken only once a
12// connection has been accepted, so an idle listener holds none; a connection
13// that finds no slot is closed at once (never queued or waited for), which makes
14// the client retry. Idle and stalled connections are cut by the server's
15// timeouts and the progress deadlines of the handlers, which frees their slots.
16type limitListener struct {
17	net.Listener
18	local  chan struct{}
19	global chan struct{}
20}
21
22func newLimitListener(l net.Listener, perListener int, global chan struct{}) *limitListener {
23	return &limitListener{Listener: l, local: make(chan struct{}, perListener), global: global}
24}
25
26// take takes a slot of c without waiting.
27func take(c chan struct{}) bool {
28	select {
29	case c <- struct{}{}:
30		return true
31	default:
32		return false
33	}
34}
35
36func (l *limitListener) Accept() (net.Conn, error) {
37	for {
38		c, err := l.Listener.Accept()
39		if err != nil {
40			return nil, err
41		}
42		if !take(l.local) {
43			_ = c.Close()
44			continue
45		}
46		if !take(l.global) {
47			<-l.local
48			_ = c.Close()
49			continue
50		}
51		return &limitConn{Conn: c, release: func() { <-l.global; <-l.local }}, nil
52	}
53}
54
55type limitConn struct {
56	net.Conn
57	once    sync.Once
58	release func()
59}
60
61func (c *limitConn) Close() error {
62	err := c.Conn.Close()
63	c.once.Do(c.release)
64	return err
65}