mtbf/conn.go

package main

import (
	"net"
	"time"
)

// Connection represents a network socket.
type Connection struct {
	dialer         net.Dialer
	socket         net.Conn
	connectTimeout time.Duration
	sendTimeout    time.Duration
	recvTimeout    time.Duration
	protocol       string

	endpoint *Endpoint
}

func init() {
	registerParam("connect-timeout-ms", 3000, "")
	registerParam("send-timeout-ms", 2000, "")
	registerParam("recv-timeout-ms", 1000, "")
}

// NewConnection creates a Connection object and optionally connects to an Endpoint.
func NewConnection(endpoint *Endpoint) (*Connection, error) {
	conn := Connection{protocol: "tcp", endpoint: endpoint}
	conn.connectTimeout = getParamDurationMS("connect-timeout-ms")
	conn.sendTimeout = getParamDurationMS("send-timeout-ms")
	conn.recvTimeout = getParamDurationMS("recv-timeout-ms")

	if endpoint != nil {
		return &conn, conn.Connect()
	} else {
		return &conn, nil
	}
}

func (conn *Connection) Close() {
	if conn.socket != nil {
		conn.socket.Close()
		conn.socket = nil
	}
}

// Connect initiates a connection to an Endpoint.
func (conn *Connection) Connect() (err error) {
	conn.endpoint.TakeMutex()
	address := conn.endpoint.String()
	conn.endpoint.ReleaseMutex()

	conn.dialer = net.Dialer{Timeout: conn.connectTimeout, KeepAlive: -1}

	start := time.Now()
	conn.socket, err = conn.dialer.Dial(conn.protocol, address)
	if err != nil {
		log("conn", 2, "cannot connect to \"%v\": %v", address, err.Error())
	}

	conn.endpoint.TakeMutex()
	conn.endpoint.RegisterRTT(time.Since(start))
	conn.endpoint.ReleaseMutex()

	return
}

// SetConnectTimeout sets a custom connect timeout on a Connection.
func (conn *Connection) SetConnectTimeout(timeout time.Duration) {
	conn.connectTimeout = timeout
}

// Send writes data to a Connection.
func (conn *Connection) Send(data []byte) (err error) {
	if len(data) == 0 {
		log("conn", 1, "tried to send empty buffer to a socket, ignoring")
		return nil
	}

	if conn.endpoint != nil {
		conn.endpoint.lastSentAt = time.Now()
	}

	conn.socket.SetWriteDeadline(time.Now().Add(conn.sendTimeout))
	_, err = conn.socket.Write(data)
	return
}

// Recv receives data from a Connection.
func (conn *Connection) Recv() (data []byte, err error) {
	if conn.endpoint != nil {
		conn.endpoint.lastReceivedAt = time.Now()
	}

	data = make([]byte, 1024)
	conn.socket.SetReadDeadline(time.Now().Add(conn.recvTimeout))

	n, err := conn.socket.Read(data)
	if err != nil {
		return nil, err
	}

	data = data[:n]
	return
}