5
5
package netutil
6
6
7
7
import (
8
+ "context"
8
9
"errors"
9
- "fmt"
10
10
"io"
11
- "io/ioutil"
12
11
"net"
13
- "net/http"
14
- "runtime"
15
12
"sync"
16
13
"sync/atomic"
17
14
"testing"
18
15
"time"
19
16
)
20
17
21
- const defaultMaxOpenFiles = 256
22
- const timeout = 5 * time .Second
23
-
24
18
func TestLimitListener (t * testing.T ) {
25
- if runtime .GOOS == "plan9" {
26
- t .Skipf ("skipping test known to be flaky on plan9 (https://golang.org/issue/22926)" )
27
- }
28
-
29
- const max = 5
30
- attempts := (maxOpenFiles () - max ) / 2
31
- if attempts > 256 { // maximum length of accept queue is 128 by default
32
- attempts = 256
33
- }
19
+ const (
20
+ max = 5
21
+ attempts = max * 2
22
+ msg = "bye\n "
23
+ )
34
24
35
25
l , err := net .Listen ("tcp" , "127.0.0.1:0" )
36
26
if err != nil {
37
27
t .Fatal (err )
38
28
}
39
- defer l .Close ()
40
29
l = LimitListener (l , max )
41
30
42
- var open int32
43
- go http .Serve (l , http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
44
- if n := atomic .AddInt32 (& open , 1 ); n > max {
45
- t .Errorf ("%d open connections, want <= %d" , n , max )
31
+ var wg sync.WaitGroup
32
+ wg .Add (1 )
33
+ go func () {
34
+ defer wg .Done ()
35
+
36
+ accepted := 0
37
+ for {
38
+ c , err := l .Accept ()
39
+ if err != nil {
40
+ break
41
+ }
42
+ accepted ++
43
+ io .WriteString (c , msg )
44
+
45
+ defer c .Close () // Leave c open until the listener is closed.
46
46
}
47
- defer atomic . AddInt32 ( & open , - 1 )
48
- time . Sleep ( 10 * time . Millisecond )
49
- fmt . Fprint ( w , "some body" )
50
- }) )
47
+ if accepted > max {
48
+ t . Errorf ( "accepted %d simultaneous connections; want at most %d" , accepted , max )
49
+ }
50
+ }( )
51
51
52
- var wg sync.WaitGroup
53
- var failed int32
54
- for i := 0 ; i < attempts ; i ++ {
52
+ // connc keeps the client end of the dialed connections alive until the
53
+ // test completes.
54
+ connc := make (chan []net.Conn , 1 )
55
+ connc <- nil
56
+
57
+ dialCtx , cancelDial := context .WithCancel (context .Background ())
58
+ defer cancelDial ()
59
+ dialer := & net.Dialer {}
60
+
61
+ var served int32
62
+ for n := attempts ; n > 0 ; n -- {
55
63
wg .Add (1 )
56
64
go func () {
57
65
defer wg .Done ()
58
- c := http. Client { Timeout : 3 * time . Second }
59
- r , err := c . Get ( "http://" + l .Addr ().String ())
66
+
67
+ c , err := dialer . DialContext ( dialCtx , l . Addr (). Network (), l .Addr ().String ())
60
68
if err != nil {
61
69
t .Log (err )
62
- atomic .AddInt32 (& failed , 1 )
63
70
return
64
71
}
65
- defer r .Body .Close ()
66
- io .Copy (ioutil .Discard , r .Body )
72
+ defer c .Close ()
73
+
74
+ // Keep this end of the connection alive until after the Listener
75
+ // finishes.
76
+ conns := append (<- connc , c )
77
+ if len (conns ) == max {
78
+ go func () {
79
+ // Give the server a bit of time to make sure it doesn't exceed its
80
+ // limit after serving this connection, then cancel the remaining
81
+ // Dials (if any).
82
+ time .Sleep (10 * time .Millisecond )
83
+ cancelDial ()
84
+ l .Close ()
85
+ }()
86
+ }
87
+ connc <- conns
88
+
89
+ b := make ([]byte , len (msg ))
90
+ if n , err := c .Read (b ); n < len (b ) {
91
+ t .Log (err )
92
+ return
93
+ }
94
+ atomic .AddInt32 (& served , 1 )
67
95
}()
68
96
}
69
97
wg .Wait ()
70
98
71
- // We expect some Gets to fail as the kernel's accept queue is filled,
72
- // but most should succeed.
73
- if int (failed ) >= attempts / 2 {
74
- t .Errorf ("%d requests failed within %d attempts" , failed , attempts )
99
+ conns := <- connc
100
+ for _ , c := range conns {
101
+ c .Close ()
102
+ }
103
+ t .Logf ("with limit %d, served %d connections (of %d dialed, %d attempted)" , max , served , len (conns ), attempts )
104
+ if served != max {
105
+ t .Errorf ("expected exactly %d served" , max )
75
106
}
76
107
}
77
108
@@ -87,27 +118,13 @@ var errFake = errors.New("fake error from errorListener")
87
118
88
119
// This used to hang.
89
120
func TestLimitListenerError (t * testing.T ) {
90
- errCh := make (chan error , 1 )
91
- go func () {
92
- defer close (errCh )
93
- const n = 2
94
- ll := LimitListener (errorListener {}, n )
95
- for i := 0 ; i < n + 1 ; i ++ {
96
- _ , err := ll .Accept ()
97
- if err != errFake {
98
- errCh <- fmt .Errorf ("Accept error = %v; want errFake" , err )
99
- return
100
- }
101
- }
102
- }()
103
-
104
- select {
105
- case err := <- errCh :
106
- if err != nil {
107
- t .Fatalf ("server: %v" , err )
121
+ const n = 2
122
+ ll := LimitListener (errorListener {}, n )
123
+ for i := 0 ; i < n + 1 ; i ++ {
124
+ _ , err := ll .Accept ()
125
+ if err != errFake {
126
+ t .Fatalf ("Accept error = %v; want errFake" , err )
108
127
}
109
- case <- time .After (timeout ):
110
- t .Fatal ("timeout. deadlock?" )
111
128
}
112
129
}
113
130
@@ -122,7 +139,7 @@ func TestLimitListenerClose(t *testing.T) {
122
139
errCh := make (chan error )
123
140
go func () {
124
141
defer close (errCh )
125
- c , err := net .DialTimeout ( "tcp" , ln .Addr ().String (), timeout )
142
+ c , err := net .Dial ( ln . Addr (). Network () , ln .Addr ().String ())
126
143
if err != nil {
127
144
errCh <- err
128
145
return
@@ -138,26 +155,21 @@ func TestLimitListenerClose(t *testing.T) {
138
155
139
156
err = <- errCh
140
157
if err != nil {
141
- t .Fatalf ("DialTimeout : %v" , err )
158
+ t .Fatalf ("Dial : %v" , err )
142
159
}
143
160
144
- acceptDone := make (chan struct {})
145
- go func () {
146
- c , err := ln .Accept ()
147
- if err == nil {
148
- c .Close ()
149
- t .Errorf ("Unexpected successful Accept()" )
150
- }
151
- close (acceptDone )
152
- }()
153
-
154
- // Wait a tiny bit to ensure the Accept() is blocking.
155
- time .Sleep (10 * time .Millisecond )
156
- ln .Close ()
161
+ // Allow the subsequent Accept to block before closing the listener.
162
+ // (Accept should unblock and return.)
163
+ timer := time .AfterFunc (10 * time .Millisecond , func () {
164
+ ln .Close ()
165
+ })
157
166
158
- select {
159
- case <- acceptDone :
160
- case <- time .After (timeout ):
161
- t .Fatalf ("Accept() still blocking" )
167
+ c , err = ln .Accept ()
168
+ if err == nil {
169
+ c .Close ()
170
+ t .Errorf ("Unexpected successful Accept()" )
171
+ }
172
+ if timer .Stop () {
173
+ t .Errorf ("Accept returned before listener closed: %v" , err )
162
174
}
163
175
}
0 commit comments