mirror of
https://github.com/Oxalide/vsphere-influxdb-go.git
synced 2023-10-10 11:36:51 +00:00
add vendoring with go dep
This commit is contained in:
Adrian Todorov
237
vendor/github.com/influxdata/influxdb/tcp/mux.go
generated
vendored
Normal file
237
vendor/github.com/influxdata/influxdb/tcp/mux.go
generated
vendored
Normal file
@@ -0,0 +1,237 @@
|
||||
// Package tcp provides a simple multiplexer over TCP.
|
||||
package tcp // import "github.com/influxdata/influxdb/tcp"
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
// DefaultTimeout is the default length of time to wait for first byte.
|
||||
DefaultTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// Mux multiplexes a network connection.
|
||||
type Mux struct {
|
||||
mu sync.RWMutex
|
||||
ln net.Listener
|
||||
m map[byte]*listener
|
||||
|
||||
defaultListener *listener
|
||||
|
||||
wg sync.WaitGroup
|
||||
|
||||
// The amount of time to wait for the first header byte.
|
||||
Timeout time.Duration
|
||||
|
||||
// Out-of-band error logger
|
||||
Logger *log.Logger
|
||||
}
|
||||
|
||||
type replayConn struct {
|
||||
net.Conn
|
||||
firstByte byte
|
||||
readFirstbyte bool
|
||||
}
|
||||
|
||||
func (rc *replayConn) Read(b []byte) (int, error) {
|
||||
if rc.readFirstbyte {
|
||||
return rc.Conn.Read(b)
|
||||
}
|
||||
|
||||
if len(b) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
b[0] = rc.firstByte
|
||||
rc.readFirstbyte = true
|
||||
return 1, nil
|
||||
}
|
||||
|
||||
// NewMux returns a new instance of Mux.
|
||||
func NewMux() *Mux {
|
||||
return &Mux{
|
||||
m: make(map[byte]*listener),
|
||||
Timeout: DefaultTimeout,
|
||||
Logger: log.New(os.Stderr, "[tcp] ", log.LstdFlags),
|
||||
}
|
||||
}
|
||||
|
||||
// Serve handles connections from ln and multiplexes then across registered listeners.
|
||||
func (mux *Mux) Serve(ln net.Listener) error {
|
||||
mux.mu.Lock()
|
||||
mux.ln = ln
|
||||
mux.mu.Unlock()
|
||||
for {
|
||||
// Wait for the next connection.
|
||||
// If it returns a temporary error then simply retry.
|
||||
// If it returns any other error then exit immediately.
|
||||
conn, err := ln.Accept()
|
||||
if err, ok := err.(interface {
|
||||
Temporary() bool
|
||||
}); ok && err.Temporary() {
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
// Wait for all connections to be demux
|
||||
mux.wg.Wait()
|
||||
for _, ln := range mux.m {
|
||||
close(ln.c)
|
||||
}
|
||||
|
||||
if mux.defaultListener != nil {
|
||||
close(mux.defaultListener.c)
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// Demux in a goroutine to
|
||||
mux.wg.Add(1)
|
||||
go mux.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (mux *Mux) handleConn(conn net.Conn) {
|
||||
defer mux.wg.Done()
|
||||
// Set a read deadline so connections with no data don't timeout.
|
||||
if err := conn.SetReadDeadline(time.Now().Add(mux.Timeout)); err != nil {
|
||||
conn.Close()
|
||||
mux.Logger.Printf("tcp.Mux: cannot set read deadline: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Read first byte from connection to determine handler.
|
||||
var typ [1]byte
|
||||
if _, err := io.ReadFull(conn, typ[:]); err != nil {
|
||||
conn.Close()
|
||||
mux.Logger.Printf("tcp.Mux: cannot read header byte: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Reset read deadline and let the listener handle that.
|
||||
if err := conn.SetReadDeadline(time.Time{}); err != nil {
|
||||
conn.Close()
|
||||
mux.Logger.Printf("tcp.Mux: cannot reset set read deadline: %s", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Retrieve handler based on first byte.
|
||||
handler := mux.m[typ[0]]
|
||||
if handler == nil {
|
||||
if mux.defaultListener == nil {
|
||||
conn.Close()
|
||||
mux.Logger.Printf("tcp.Mux: handler not registered: %d. Connection from %s closed", typ[0], conn.RemoteAddr())
|
||||
return
|
||||
}
|
||||
|
||||
conn = &replayConn{
|
||||
Conn: conn,
|
||||
firstByte: typ[0],
|
||||
}
|
||||
handler = mux.defaultListener
|
||||
}
|
||||
|
||||
// Send connection to handler. The handler is responsible for closing the connection.
|
||||
timer := time.NewTimer(mux.Timeout)
|
||||
defer timer.Stop()
|
||||
|
||||
select {
|
||||
case handler.c <- conn:
|
||||
case <-timer.C:
|
||||
conn.Close()
|
||||
mux.Logger.Printf("tcp.Mux: handler not ready: %d. Connection from %s closed", typ[0], conn.RemoteAddr())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Listen returns a listener identified by header.
|
||||
// Any connection accepted by mux is multiplexed based on the initial header byte.
|
||||
func (mux *Mux) Listen(header byte) net.Listener {
|
||||
// Ensure two listeners are not created for the same header byte.
|
||||
if _, ok := mux.m[header]; ok {
|
||||
panic(fmt.Sprintf("listener already registered under header byte: %d", header))
|
||||
}
|
||||
|
||||
// Create a new listener and assign it.
|
||||
ln := &listener{
|
||||
c: make(chan net.Conn),
|
||||
mux: mux,
|
||||
}
|
||||
mux.m[header] = ln
|
||||
|
||||
return ln
|
||||
}
|
||||
|
||||
// DefaultListener will return a net.Listener that will pass-through any
|
||||
// connections with non-registered values for the first byte of the connection.
|
||||
// The connections returned from this listener's Accept() method will replay the
|
||||
// first byte of the connection as a short first Read().
|
||||
//
|
||||
// This can be used to pass to an HTTP server, so long as there are no conflicts
|
||||
// with registered listener bytes and the first character of the HTTP request:
|
||||
// 71 ('G') for GET, etc.
|
||||
func (mux *Mux) DefaultListener() net.Listener {
|
||||
if mux.defaultListener == nil {
|
||||
mux.defaultListener = &listener{
|
||||
c: make(chan net.Conn),
|
||||
mux: mux,
|
||||
}
|
||||
}
|
||||
|
||||
return mux.defaultListener
|
||||
}
|
||||
|
||||
// listener is a receiver for connections received by Mux.
|
||||
type listener struct {
|
||||
c chan net.Conn
|
||||
mux *Mux
|
||||
}
|
||||
|
||||
// Accept waits for and returns the next connection to the listener.
|
||||
func (ln *listener) Accept() (c net.Conn, err error) {
|
||||
conn, ok := <-ln.c
|
||||
if !ok {
|
||||
return nil, errors.New("network connection closed")
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// Close is a no-op. The mux's listener should be closed instead.
|
||||
func (ln *listener) Close() error { return nil }
|
||||
|
||||
// Addr returns the Addr of the listener
|
||||
func (ln *listener) Addr() net.Addr {
|
||||
if ln.mux == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
ln.mux.mu.RLock()
|
||||
defer ln.mux.mu.RUnlock()
|
||||
|
||||
if ln.mux.ln == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
return ln.mux.ln.Addr()
|
||||
}
|
||||
|
||||
// Dial connects to a remote mux listener with a given header byte.
|
||||
func Dial(network, address string, header byte) (net.Conn, error) {
|
||||
conn, err := net.Dial(network, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if _, err := conn.Write([]byte{header}); err != nil {
|
||||
return nil, fmt.Errorf("write mux header: %s", err)
|
||||
}
|
||||
|
||||
return conn, nil
|
||||
}
|
||||
156
vendor/github.com/influxdata/influxdb/tcp/mux_test.go
generated
vendored
Normal file
156
vendor/github.com/influxdata/influxdb/tcp/mux_test.go
generated
vendored
Normal file
@@ -0,0 +1,156 @@
|
||||
package tcp_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"testing/quick"
|
||||
"time"
|
||||
|
||||
"github.com/influxdata/influxdb/tcp"
|
||||
)
|
||||
|
||||
// Ensure the muxer can split a listener's connections across multiple listeners.
|
||||
func TestMux(t *testing.T) {
|
||||
if err := quick.Check(func(n uint8, msg []byte) bool {
|
||||
if testing.Verbose() {
|
||||
if len(msg) == 0 {
|
||||
log.Printf("n=%d, <no message>", n)
|
||||
} else {
|
||||
log.Printf("n=%d, hdr=%d, len=%d", n, msg[0], len(msg))
|
||||
}
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
|
||||
// Open single listener on random port.
|
||||
tcpListener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer tcpListener.Close()
|
||||
|
||||
// Setup muxer & listeners.
|
||||
mux := tcp.NewMux()
|
||||
mux.Timeout = 200 * time.Millisecond
|
||||
if !testing.Verbose() {
|
||||
mux.Logger = log.New(ioutil.Discard, "", 0)
|
||||
}
|
||||
|
||||
errC := make(chan error)
|
||||
for i := uint8(0); i < n; i++ {
|
||||
ln := mux.Listen(byte(i))
|
||||
|
||||
wg.Add(1)
|
||||
go func(i uint8, ln net.Listener) {
|
||||
defer wg.Done()
|
||||
|
||||
// Wait for a connection for this listener.
|
||||
conn, err := ln.Accept()
|
||||
if conn != nil {
|
||||
defer conn.Close()
|
||||
}
|
||||
|
||||
// If there is no message or the header byte
|
||||
// doesn't match then expect close.
|
||||
if len(msg) == 0 || msg[0] != byte(i) {
|
||||
if err == nil || err.Error() != "network connection closed" {
|
||||
errC <- fmt.Errorf("unexpected error: %s", err)
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// If the header byte matches this listener
|
||||
// then expect a connection and read the message.
|
||||
var buf bytes.Buffer
|
||||
if _, err := io.CopyN(&buf, conn, int64(len(msg)-1)); err != nil {
|
||||
errC <- err
|
||||
return
|
||||
} else if !bytes.Equal(msg[1:], buf.Bytes()) {
|
||||
errC <- fmt.Errorf("message mismatch:\n\nexp=%x\n\ngot=%x\n\n", msg[1:], buf.Bytes())
|
||||
return
|
||||
}
|
||||
|
||||
// Write response.
|
||||
if _, err := conn.Write([]byte("OK")); err != nil {
|
||||
errC <- err
|
||||
return
|
||||
}
|
||||
}(i, ln)
|
||||
}
|
||||
|
||||
// Begin serving from the listener.
|
||||
go mux.Serve(tcpListener)
|
||||
|
||||
// Write message to TCP listener and read OK response.
|
||||
conn, err := net.Dial("tcp", tcpListener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
} else if _, err = conn.Write(msg); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Read the response into the buffer.
|
||||
var resp [2]byte
|
||||
_, err = io.ReadFull(conn, resp[:])
|
||||
|
||||
// If the message header is less than n then expect a response.
|
||||
// Otherwise we should get an EOF because the mux closed.
|
||||
if len(msg) > 0 && uint8(msg[0]) < n {
|
||||
if string(resp[:]) != `OK` {
|
||||
t.Fatalf("unexpected response: %s", resp[:])
|
||||
}
|
||||
} else {
|
||||
if err == nil || (err != io.EOF && !(strings.Contains(err.Error(), "connection reset by peer") ||
|
||||
strings.Contains(err.Error(), "closed by the remote host"))) {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Close connection.
|
||||
if err := conn.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Close original TCP listener and wait for all goroutines to close.
|
||||
tcpListener.Close()
|
||||
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(errC)
|
||||
}()
|
||||
|
||||
ok := true
|
||||
for err := range errC {
|
||||
if err != nil {
|
||||
ok = false
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
|
||||
return ok
|
||||
}, nil); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure two handlers cannot be registered for the same header byte.
|
||||
func TestMux_Listen_ErrAlreadyRegistered(t *testing.T) {
|
||||
defer func() {
|
||||
if r := recover(); r != `listener already registered under header byte: 5` {
|
||||
t.Fatalf("unexpected recover: %#v", r)
|
||||
}
|
||||
}()
|
||||
|
||||
// Register two listeners with the same header byte.
|
||||
mux := tcp.NewMux()
|
||||
mux.Listen(5)
|
||||
mux.Listen(5)
|
||||
}
|
||||
Reference in New Issue
Block a user