// TestPool tests the pool is functional.
func TestPool(t *testing.T) {
tests.ResetLog()
defer tests.DisplayLog()
t.Log("Given the need to validate the work pool functions.")
{
cfg := pool.Config{
MinRoutines: func() int { return 100 },
MaxRoutines: func() int { return 5000 },
}
p, err := pool.New("TestPool", "Pool1", cfg)
if err != nil {
t.Fatal("\tShould not get error creating pool.", tests.Failed, err)
}
t.Log("\tShould not get error creating pool.", tests.Success)
for i := 0; i < 100; i++ {
p.Do("TestPool", &theWork{privateID: i})
}
time.Sleep(100 * time.Millisecond)
p.Shutdown("TestPool")
}
}
// ExampleNew provides a basic example for using a pool.
func ExampleNew() {
tests.ResetLog()
defer tests.DisplayLog()
// Create a configuration.
cfg := pool.Config{
MinRoutines: func() int { return 3 },
MaxRoutines: func() int { return 4 },
}
// Create a new pool.
p, err := pool.New("TEST", "TheWork", cfg)
if err != nil {
fmt.Println(err)
return
}
// Pass in some work to be performed.
p.Do("TEST", &theWork{})
p.Do("TEST", &theWork{})
p.Do("TEST", &theWork{})
// Wait to the work to be processed.
time.Sleep(100 * time.Millisecond)
// Shutdown the pool.
p.Shutdown("TEST")
}
func main() {
const context = "main"
const totalWork = 100
wg.Add(totalWork)
// Create the configuration.
cfg := pool.Config{
MinRoutines: func() int { return cfg.MustInt(cfgMinRoutines) },
MaxRoutines: func() int { return cfg.MustInt(cfgMaxRoutines) },
}
// Create a pool.
p, err := pool.New(context, "test", cfg)
if err != nil {
log.Error(context, "main", err, "Creating pool")
return
}
// Look at stats for the work.
go func() {
for {
time.Sleep(250 * time.Millisecond)
log.User(context, "Stats", "%#v", p.Stats())
}
}()
// Perform some work.
for i := 0; i < totalWork; i++ {
p.Do(context, &Task{Name: strconv.Itoa(i)})
}
// Wait until all the work is complete.
wg.Wait()
// Shutdown the pool.
p.Shutdown(context)
}
// TestRateLimit tests we can drop connections when they come in too fast.
func TestRateLimit(t *testing.T) {
tests.ResetLog()
defer tests.DisplayLog()
const ratelimit = 1 * time.Second
t.Log("Given the need to drop TCP connections.")
{
recvCfg := pool.Config{
MinRoutines: func() int { return 2 },
MaxRoutines: func() int { return 1000 },
}
recv, err := pool.New(tests.Context, "Test-Recv", recvCfg)
if err != nil {
t.Fatal("\tShould be able to create a work pool for the recv.", tests.Failed, err)
}
sendCfg := pool.Config{
MinRoutines: func() int { return 2 },
MaxRoutines: func() int { return 1000 },
}
send, err := pool.New(tests.Context, "Test-Send", sendCfg)
if err != nil {
t.Fatal("\tShould be able to create a work pool for the send.", tests.Failed, err)
}
// Create a configuration.
cfg := tcp.Config{
NetType: "tcp4",
Addr: ":0",
ConnHandler: tcpConnHandler{},
ReqHandler: tcpReqHandler{},
RespHandler: tcpRespHandler{},
OptUserPool: tcp.OptUserPool{
RecvPool: recv,
SendPool: send,
},
OptRateLimit: tcp.OptRateLimit{
RateLimit: func() time.Duration { return ratelimit },
},
}
// Create a new TCP value.
u, err := tcp.New(tests.Context, "TEST", cfg)
if err != nil {
t.Fatal("\tShould be able to create a new TCP listener.", tests.Failed, err)
}
t.Log("\tShould be able to create a new TCP listener.", tests.Success)
// Start accepting client data.
if err := u.Start(tests.Context); err != nil {
t.Fatal("\tShould be able to start the TCP listener.", tests.Failed, err)
}
t.Log("\tShould be able to start the TCP listener.", tests.Success)
defer u.Stop(tests.Context)
newconn := func() (*bufio.Writer, *bufio.Reader, net.Conn, error) {
// Let's connect to the host:port.
conn, err := net.Dial("tcp4", u.Addr().String())
if err != nil {
return nil, nil, nil, err
}
return bufio.NewWriter(conn), bufio.NewReader(conn), conn, nil
}
// Make a successful connection
successfulTest := func(ctx interface{}) {
w, r, c, err := newconn()
if err != nil {
t.Fatal("\tShould be able to dial a new TCP connection.", ctx, tests.Failed, err)
}
t.Log("\tShould be able to dial a new TCP connection.", ctx, tests.Success)
defer c.Close()
if _, err := w.WriteString("Hello\n"); err != nil {
t.Fatal("\tShould be able to send data to the connection.", ctx, tests.Failed, err)
}
t.Log("\tShould be able to send data to the connection.", ctx, tests.Success)
if err := w.Flush(); err != nil {
t.Fatal("\tShould be able to flush the writer.", ctx, tests.Failed, err)
}
t.Log("\tShould be able to flush the writer.", ctx, tests.Success)
// Let's read the response.
response, err := r.ReadString('\n')
if err != nil {
t.Fatal("\tShould be able to read the response from the connection.", ctx, tests.Failed, err)
}
t.Log("\tShould be able to read the response from the connection.", ctx, tests.Success)
t.Log(response)
}
successfulTest("PRE-LIMIT")
// The next 100 connections should fail (assuming it's all under rateLimit amount of time).
for i := 0; i < 100; i++ {
// Apparently, even though the connection should not exist, we are still allowed
// to connect to the remote socket and write to it. The error is exhibited
// only when it's time to perform a read on that connection.
w, r, c, err := newconn()
if err != nil {
t.Fatal("\tShould be able to dial a non-first TCP connection.", tests.Failed, err)
}
t.Log("\tShould be able to dial a non-first TCP connection", c.LocalAddr(), tests.Success)
defer c.Close()
if _, err := w.WriteString("Hello\n"); err != nil {
t.Fatal("\tShould be able to send data to the connection.", tests.Failed, err)
}
t.Log("\tShould be able to send data to the connection.", tests.Success)
if err := w.Flush(); err != nil {
t.Fatal("\tShould be able to flush the writer.", tests.Failed, err)
}
t.Log("\tShould be able to flush the writer.", tests.Success)
// Let's read the response.
_, err = r.ReadString('\n')
if err == nil {
t.Fatal("\tShould have tests.Failed to read from the connection.", tests.Failed)
}
t.Log("\tShould have tests.Failed to read from the connection", tests.Success, err)
}
// Sleep for rateLimit to perform another successful test.
time.Sleep(ratelimit)
successfulTest("POST-LIMIT")
// NOTE If you call another 'successfulTest' here, we will fail because we expect
// the test to fail due to the limit.
}
}
// TestDropConnections tests we can drop connections when configured.
func TestDropConnections(t *testing.T) {
tests.ResetLog()
defer tests.DisplayLog()
t.Log("Given the need to drop TCP connections.")
{
recvCfg := pool.Config{
MinRoutines: func() int { return 2 },
MaxRoutines: func() int { return 1000 },
}
recv, err := pool.New(tests.Context, "Test-Recv", recvCfg)
if err != nil {
t.Fatal("\tShould be able to create a work pool for the recv.", tests.Failed, err)
}
sendCfg := pool.Config{
MinRoutines: func() int { return 2 },
MaxRoutines: func() int { return 1000 },
}
send, err := pool.New(tests.Context, "Test-Send", sendCfg)
if err != nil {
t.Fatal("\tShould be able to create a work pool for the send.", tests.Failed, err)
}
// Create a configuration.
cfg := tcp.Config{
NetType: "tcp4",
Addr: ":0",
ConnHandler: tcpConnHandler{},
ReqHandler: tcpReqHandler{},
RespHandler: tcpRespHandler{},
OptUserPool: tcp.OptUserPool{
RecvPool: recv,
SendPool: send,
},
}
// Create a new TCP value.
u, err := tcp.New(tests.Context, "TEST", cfg)
if err != nil {
t.Fatal("\tShould be able to create a new TCP listener.", tests.Failed, err)
}
t.Log("\tShould be able to create a new TCP listener.", tests.Success)
// Set the drop connection flag to true.
t.Log("\tSet the drop connections flag to TRUE.", tests.Success)
u.DropConnections(tests.Context, true)
// Start accepting client data.
if err := u.Start(tests.Context); err != nil {
t.Fatal("\tShould be able to start the TCP listener.", tests.Failed, err)
}
t.Log("\tShould be able to start the TCP listener.", tests.Success)
defer u.Stop(tests.Context)
// Let's connect to the host:port.
conn, err := net.Dial("tcp4", u.Addr().String())
if err != nil {
t.Fatal("\tShould be able to dial a new TCP connection.", tests.Failed, err)
}
t.Log("\tShould be able to dial a new TCP connection.", tests.Success)
// An attempt to read should result in an EOF.
b := make([]byte, 1)
if _, err = conn.Read(b); err == nil {
t.Fatal("\tShould not be able to read the response from the connection.", tests.Failed, err)
}
t.Log("\tShould not be able to read the response from the connection.", tests.Success)
}
}
// New creates a new manager to service clients.
func New(context interface{}, name string, cfg Config) (*TCP, error) {
log.Dev(context, "New", "Started : Name[%s] NetType[%s] Addr[%s] RecvMaxPoolSize[%d] SendMaxPoolSize[%d]", name, cfg.NetType, cfg.Addr, cfg.RecvMaxPoolSize, cfg.SendMaxPoolSize)
// Validate the configuration.
if err := cfg.Validate(); err != nil {
log.Error(context, "New", err, "Completed")
return nil, err
}
// Resolve the addr that is provided.
tcpAddr, err := net.ResolveTCPAddr(cfg.NetType, cfg.Addr)
if err != nil {
log.Error(context, "New", err, "Completed")
return nil, err
}
log.Dev(context, "New", "Address[ %s ] Zone[%s]", join(tcpAddr.IP.String(), tcpAddr.Port), tcpAddr.Zone)
// Need a work pool to handle the received messages.
var recv *pool.Pool
if cfg.RecvPool != nil {
recv = cfg.RecvPool
} else {
recvCfg := pool.Config{
MinRoutines: cfg.RecvMinPoolSize,
MaxRoutines: cfg.RecvMaxPoolSize,
}
var err error
if recv, err = pool.New(context, name+"-Recv", recvCfg); err != nil {
log.Error(context, "New", err, "Completed")
return nil, err
}
}
// Need a work pool to handle the messages to send.
var send *pool.Pool
if cfg.SendPool != nil {
send = cfg.SendPool
} else {
sendCfg := pool.Config{
MinRoutines: cfg.SendMinPoolSize,
MaxRoutines: cfg.SendMaxPoolSize,
}
var err error
if send, err = pool.New(context, name+"-Send", sendCfg); err != nil {
log.Error(context, "New", err, "Completed")
return nil, err
}
}
// Are we using user provided work pools. Validation is helping us
// only have to check one of the two configuration options for this.
var userPools bool
if cfg.RecvPool != nil {
log.Dev(context, "New", "Using User Pools")
userPools = true
}
// Create a TCP for this ipaddress and port.
t := TCP{
Config: cfg,
Name: name,
ipAddress: tcpAddr.IP.String(),
port: tcpAddr.Port,
tcpAddr: tcpAddr,
clients: make(map[string]*client),
recv: recv,
send: send,
userPools: userPools,
}
log.Dev(context, "New", "Completed")
return &t, nil
}