109 lines
2.1 KiB
Go
109 lines
2.1 KiB
Go
package httpd_test
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/influxdata/influxdb/services/httpd"
|
|
)
|
|
|
|
type fakeListener struct {
|
|
AcceptFn func() (net.Conn, error)
|
|
}
|
|
|
|
func (l *fakeListener) Accept() (net.Conn, error) {
|
|
if l.AcceptFn != nil {
|
|
return l.AcceptFn()
|
|
}
|
|
return &fakeConn{}, nil
|
|
}
|
|
|
|
func (*fakeListener) Close() error { return nil }
|
|
func (*fakeListener) Addr() net.Addr { return nil }
|
|
|
|
type fakeConn struct {
|
|
closed bool
|
|
}
|
|
|
|
func (*fakeConn) Read([]byte) (int, error) { return 0, io.EOF }
|
|
func (*fakeConn) Write(b []byte) (int, error) { return len(b), nil }
|
|
func (c *fakeConn) Close() error {
|
|
c.closed = true
|
|
return nil
|
|
}
|
|
func (*fakeConn) LocalAddr() net.Addr { return nil }
|
|
func (*fakeConn) RemoteAddr() net.Addr { return nil }
|
|
func (*fakeConn) SetDeadline(time.Time) error { return nil }
|
|
func (*fakeConn) SetReadDeadline(time.Time) error { return nil }
|
|
func (*fakeConn) SetWriteDeadline(time.Time) error { return nil }
|
|
|
|
func TestLimitListener(t *testing.T) {
|
|
conns := make(chan net.Conn, 2)
|
|
l := httpd.LimitListener(&fakeListener{
|
|
AcceptFn: func() (net.Conn, error) {
|
|
select {
|
|
case c := <-conns:
|
|
if c != nil {
|
|
return c, nil
|
|
}
|
|
default:
|
|
}
|
|
return nil, io.EOF
|
|
},
|
|
}, 1)
|
|
c1, c2 := &fakeConn{}, &fakeConn{}
|
|
conns <- c1
|
|
conns <- c2
|
|
|
|
var c net.Conn
|
|
var err error
|
|
if c, err = l.Accept(); err != nil {
|
|
t.Fatalf("expected accept to succeed: %s", err)
|
|
}
|
|
|
|
if _, err = l.Accept(); err != io.EOF {
|
|
t.Fatalf("expected eof, got %s", err)
|
|
} else if !c2.closed {
|
|
t.Fatalf("expected connection to be automatically closed")
|
|
}
|
|
c.Close()
|
|
|
|
conns <- &fakeConn{}
|
|
if _, err = l.Accept(); err != nil {
|
|
t.Fatalf("expeced accept to succeed: %s", err)
|
|
}
|
|
}
|
|
|
|
func BenchmarkLimitListener(b *testing.B) {
|
|
var wg sync.WaitGroup
|
|
wg.Add(b.N)
|
|
|
|
l := httpd.LimitListener(&fakeListener{}, b.N)
|
|
errC := make(chan error)
|
|
for i := 0; i < b.N; i++ {
|
|
go func() {
|
|
defer wg.Done()
|
|
c, err := l.Accept()
|
|
if err != nil {
|
|
errC <- err
|
|
return
|
|
}
|
|
c.Close()
|
|
}()
|
|
}
|
|
|
|
go func() {
|
|
wg.Wait()
|
|
close(errC)
|
|
}()
|
|
|
|
for err := range errC {
|
|
if err != nil {
|
|
b.Error(err)
|
|
}
|
|
}
|
|
}
|