Add support for timeouts and servers closing connections

This commit is contained in:
Leonid Bugaev
2015-07-02 18:33:38 +05:00
parent decb9483a5
commit a08ecc79fc
3 changed files with 129 additions and 18 deletions
+10 -8
View File
@@ -1,31 +1,33 @@
SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go settings_header_filters.go settings_header_hash_filters.go settings_headers.go settings_methods.go settings_option.go settings_url_regexp.go test_input.go elasticsearch.go settings_url_map.go
SOURCE_PATH = /gopath/src/github.com/buger/gor/
release: release-x86 release-x64
release-x64:
docker run -v `pwd`:/gopath/src/gor -t --env GOOS=linux --env GOARCH=amd64 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x64.tar.gz gor && rm gor
docker run -v `pwd`:$(SOURCE_PATH) -t --env GOOS=linux --env GOARCH=amd64 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x64.tar.gz gor && rm gor
release-x86:
docker run -v `pwd`:/gopath/src/gor -t --env GOOS=linux --env GOARCH=386 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x86.tar.gz gor && rm gor
docker run -v `pwd`:$(SOURCE_PATH) -t --env GOOS=linux --env GOARCH=386 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x86.tar.gz gor && rm gor
dbuild:
docker build -t gor .
dtest:
docker run -v `pwd`:/gopath/src/gor -t -i --env GORACE="halt_on_error=1" gor go test $(ARGS) -race -v
docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test $(ARGS) -race -v
dfmt:
docker run -v `pwd`:/gopath/src/gor -t -i gor go fmt
docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go fmt
dvet:
docker run -v `pwd`:/gopath/src/gor -t -i gor go vet
docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go vet
dbench:
docker run -v `pwd`:/gopath/src/gor -t -i gor go test -v -run NOT_EXISTING -bench HTTP
docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go test -v -run NOT_EXISTING -bench HTTP
# Used mainly for debugging, because docker container do not have access to parent machine ports
drun:
docker run -v `pwd`:/gopath/src/gor -t -i gor go run $(SOURCE) --input-dummy=0 --input-http=:9000 --output-http="http://localhost:9000" --verbose
docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go run $(SOURCE) --input-dummy=0 --input-http=:9000 --output-http="http://localhost:9000" --verbose
dbash:
docker run -v `pwd`:/gopath/src/gor -t -i gor /bin/bash
docker run -v `pwd`:$(SOURCE_PATH) -t -i gor /bin/bash
+49 -10
View File
@@ -2,9 +2,11 @@ package main
import (
"crypto/tls"
"io"
"net"
"net/url"
"strings"
"time"
)
var defaultPorts = map[string]string{
@@ -15,13 +17,17 @@ var defaultPorts = map[string]string{
type HTTPClient struct {
baseURL *url.URL
conn net.Conn
buf []byte
respBuf []byte
}
func NewHTTPClient(baseURL string) *HTTPClient {
if !strings.HasPrefix(baseURL, "http") {
baseURL = "http://" + baseURL
}
client := new(HTTPClient)
client.baseURL, _ = url.Parse(baseURL)
client.buf = make([]byte, 4096*10)
client.respBuf = make([]byte, 4096*10)
if !strings.Contains(client.baseURL.Host, ":") {
client.baseURL.Host += ":" + defaultPorts[client.baseURL.Scheme]
@@ -31,11 +37,17 @@ func NewHTTPClient(baseURL string) *HTTPClient {
}
func (c *HTTPClient) Connect() (err error) {
c.Disconnect()
c.conn, err = net.Dial("tcp", c.baseURL.Host)
if c.baseURL.Scheme == "https" {
tlsConn := tls.Client(c.conn, &tls.Config{InsecureSkipVerify: true})
err = tlsConn.Handshake()
if err = tlsConn.Handshake(); err != nil {
return
}
c.conn = tlsConn
}
@@ -43,19 +55,46 @@ func (c *HTTPClient) Connect() (err error) {
}
func (c *HTTPClient) Disconnect() {
c.conn.Close()
c.conn = nil
if c.conn != nil {
c.conn.Close()
c.conn = nil
Debug("Disconnected: ", c.baseURL)
}
}
func (c *HTTPClient) isAlive() bool {
one := make([]byte, 1)
// Ready 1 byte from socket without timeout to check if it not closed
c.conn.SetReadDeadline(time.Now().Add(time.Millisecond))
if _, err := c.conn.Read(one); err == io.EOF {
return false
}
return true
}
func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
if c.conn == nil {
if c.conn == nil || !c.isAlive() {
Debug("Connecting:", c.baseURL)
c.Connect()
}
_, err = c.conn.Write(data)
n, err := c.conn.Read(c.buf)
timeout := time.Now().Add(5 * time.Second)
Debug(string(c.buf[:n]))
c.conn.SetWriteDeadline(timeout)
if _, err = c.conn.Write(data); err != nil {
Debug("Write error:", err, c.baseURL)
return
}
return c.buf[:n], err
c.conn.SetReadDeadline(timeout)
n, err := c.conn.Read(c.respBuf)
if err != nil {
Debug("READ ERRORR!", err, c.conn)
return
}
return c.respBuf[:n], err
}
+70
View File
@@ -2,11 +2,13 @@ package main
import (
"io/ioutil"
"net"
"net/http"
"net/http/httptest"
"net/http/httputil"
"sync"
"testing"
_ "time"
)
func TestHTTPClientURLPort(t *testing.T) {
@@ -24,6 +26,11 @@ func TestHTTPClientURLPort(t *testing.T) {
if c3.baseURL.String() != "https://example.com:1" {
t.Error("Sould use specified port:", c3.baseURL.String())
}
c4 := NewHTTPClient("example.com")
if c4.baseURL.String() != "http://example.com:80" {
t.Error("Sould add default protocol:", c4.baseURL.String())
}
}
func TestHTTPClientSend(t *testing.T) {
@@ -111,3 +118,66 @@ func TestHTTPClientHTTPSSend(t *testing.T) {
wg.Wait()
}
func TestHTTPClientServerInstantDisconnect(t *testing.T) {
wg := new(sync.WaitGroup)
GET_payload := []byte("GET / HTTP/1.1\r\n\r\n")
ln, _ := net.Listen("tcp", ":0")
go func() {
for {
conn, _ := ln.Accept()
conn.Close()
wg.Done()
}
}()
client := NewHTTPClient(ln.Addr().String())
wg.Add(2)
client.Send(GET_payload)
client.Send(GET_payload)
wg.Wait()
}
func TestHTTPClientServerNoKeepAlive(t *testing.T) {
wg := new(sync.WaitGroup)
GET_payload := []byte("GET / HTTP/1.1\r\n\r\n")
ln, _ := net.Listen("tcp", ":0")
go func() {
for {
conn, err := ln.Accept()
if err != nil {
// handle error
}
buf := make([]byte, 4096)
reqLen, err := conn.Read(buf)
if err != nil {
t.Error("Error reading:", err.Error())
}
Debug("Received: ", string(buf[0:reqLen]))
conn.Write([]byte("OK"))
// No keep-alive connections
conn.Close()
wg.Done()
}
}()
client := NewHTTPClient(ln.Addr().String())
wg.Add(2)
client.Send(GET_payload)
client.Send(GET_payload)
wg.Wait()
}