<commit_before>package exportertools

import (
    "github.com/prometheus/client_golang/prometheus"
    "github.com/prometheus/common/log"

    "strings"
    "math"
    "strconv"
)

type Metric struct {
    Name        string
    Description string
    Type        MetricType
    Value       interface{}
    Labels      map[string]string
}

type MetricType int

type MetricCollector interface {
    Collect() ([]*Metric, error)
}

const (
    Counter MetricType = iota
    Gauge   MetricType = iota
    Untyped MetricType = iota
)

func (m *Metric) PromDescription(exporterName string) *prometheus.Desc {
    return prometheus.NewDesc(
        prometheus.BuildFQName("", exporterName, m.Name),
        m.Description,
        nil, prometheus.Labels(m.Labels),
    )
}

func (m *Metric) PromType() prometheus.ValueType {
    switch m.Type {
    case Counter:
        return prometheus.CounterValue
    case Gauge:
        return prometheus.GaugeValue
    default:
        return prometheus.UntypedValue
    }
}

func (m *Metric) PromValue() float64 {
    switch v := m.Value.(type) {
    case int64:
        return float64(v)
    case float64:
        return v
    case time.Time:
        return float64(v.Unix())
    case []byte:
        // Try and convert to string and then parse to a float64
        strV := string(v)
        result, err := strconv.ParseFloat(strV, 64)
        if err != nil {
            return math.NaN()
        }
        return result
    case string:
        result, err := strconv.ParseFloat(v, 64)
        if err != nil {
            return math.NaN()
        }
        return result
    case nil:
        return math.NaN()
    default:
        return math.NaN()
    }
}

func StringToType(s string) MetricType {
    switch strings.ToLower(s) {
    case "gauge":
        return Gauge
    case "counter":
        return Counter
    case "untyped":
        return Untyped
    default:
        log.Errorf("Undefined metric type: %v", s)
        return Untyped
    }
}
<commit_msg>fix import<commit_after>package exportertools

import (
    "github.com/prometheus/client_golang/prometheus"
    "github.com/prometheus/common/log"

    "strings"
    "math"
    "strconv"
    "time"
)

type Metric struct {
    Name        string
    Description string
    Type        MetricType
    Value       interface{}
    Labels      map[string]string
}

type MetricType int

type MetricCollector interface {
    Collect() ([]*Metric, error)
}

const (
    Counter MetricType = iota
    Gauge   MetricType = iota
    Untyped MetricType = iota
)

func (m *Metric) PromDescription(exporterName string) *prometheus.Desc {
    return prometheus.NewDesc(
        prometheus.BuildFQName("", exporterName, m.Name),
        m.Description,
        nil, prometheus.Labels(m.Labels),
    )
}

func (m *Metric) PromType() prometheus.ValueType {
    switch m.Type {
    case Counter:
        return prometheus.CounterValue
    case Gauge:
        return prometheus.GaugeValue
    default:
        return prometheus.UntypedValue
    }
}

func (m *Metric) PromValue() float64 {
    switch v := m.Value.(type) {
    case int64:
        return float64(v)
    case float64:
        return v
    case time.Time:
        return float64(v.Unix())
    case []byte:
        // Try and convert to string and then parse to a float64
        strV := string(v)
        result, err := strconv.ParseFloat(strV, 64)
        if err != nil {
            return math.NaN()
        }
        return result
    case string:
        result, err := strconv.ParseFloat(v, 64)
        if err != nil {
            return math.NaN()
        }
        return result
    case nil:
        return math.NaN()
    default:
        return math.NaN()
    }
}

func StringToType(s string) MetricType {
    switch strings.ToLower(s) {
    case "gauge":
        return Gauge
    case "counter":
        return Counter
    case "untyped":
        return Untyped
    default:
        log.Errorf("Undefined metric type: %v", s)
        return Untyped
    }
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/google/gopacket/layers"
	"bytes"
	"time"
	"github.com/google/gopacket"
	"sort"
	"io"
	"sync"
	"strconv"
)

// gopacket provide a tcp connection, however it split one tcp connection into two stream.
// So it is hard to match http request and response. we make our own connection here

const MAX_TCP_SEQ uint32 = 0xFFFFFFFF
const TCP_SEQ_WINDOW = 1 << 10

type TcpAssembler struct {
	connectionDict    map[string]*TcpConnection
	lock              sync.Mutex
	connectionHandler ConnectionHandler
	filterIp          string
	filterPort        uint16
}

func newTcpAssembler(connectionHandler ConnectionHandler) *TcpAssembler {
	return &TcpAssembler{connectionDict:map[string]*TcpConnection{}, connectionHandler:connectionHandler}
}

func (assembler *TcpAssembler) assemble(flow gopacket.Flow, tcp *layers.TCP, timestamp time.Time) {
	src := EndPoint{ip:flow.Src().String(), port:uint16(tcp.SrcPort)}
	dst := EndPoint{ip:flow.Dst().String(), port:uint16(tcp.DstPort)}
	dropped := false
	if assembler.filterIp != "" {
		if src.ip != assembler.filterIp && dst.ip != assembler.filterIp {
			dropped = true
		}
	}
	if assembler.filterPort != 0 {
		if src.port != assembler.filterPort && dst.port != assembler.filterPort {
			dropped = true
		}
	}
	if dropped {
		return
	}

	srcString := src.String()
	dstString := dst.String()
	var key string
	if srcString < dstString {
		key = srcString + "-" + dstString
	} else {
		key = dstString + "-" + srcString
	}
	connection := assembler.newConnection(src, dst, key)

	connection.onReceive(src, dst, tcp, timestamp)

	if connection.closed() {
		assembler.deleteConnection(key)
		connection.finish()
	}

	//TODO: cleanup timeout connections
}

func (assembler *TcpAssembler) newConnection(src, dst EndPoint, key string) *TcpConnection {
	assembler.lock.Lock()
	defer assembler.lock.Unlock()
	connection := assembler.connectionDict[key]
	if (connection == nil) {
		connection = newTcpConnection()
		assembler.connectionDict[key] = connection
		assembler.connectionHandler.handle(src, dst, connection)
	}
	return connection
}

func (assembler *TcpAssembler) deleteConnection(key string) {
	assembler.lock.Lock()
	defer assembler.lock.Unlock()
	delete(assembler.connectionDict, key)
}

func (assembler *TcpAssembler) flushOlderThan(time time.Time) {
	//
}

func (assembler *TcpAssembler) finishAll() {
	assembler.lock.Lock()
	defer assembler.lock.Unlock()
	for _, connection := range assembler.connectionDict {
		connection.finish()
	}
	assembler.connectionDict = nil
	assembler.connectionHandler.finish()
}

type ConnectionHandler interface {
	handle(src EndPoint, dst EndPoint, connection *TcpConnection)
	finish()
}

// one tcp connection
type  TcpConnection struct {
	upStream      *NetworkStream // stream from client to server
	downStream    *NetworkStream // stream from server to client
	clientId      EndPoint       // the client key(by ip and port)
	lastTimestamp time.Time      // timestamp receive last packet
	isHttp        bool
}

type EndPoint struct {
	ip   string
	port uint16
}

func (p EndPoint) equals(p2 EndPoint) bool {
	return p.ip == p2.ip && p.port == p2.port
}

func (p EndPoint) String() string {
	return p.ip + ":" + strconv.Itoa(int(p.port))
}

type ConnectionId struct {
	src EndPoint
	dst EndPoint
}

// create tcp connection, by the first tcp packet. this packet should from client to server
func newTcpConnection() *TcpConnection {
	connection := &TcpConnection{
		upStream:newNetworkStream(),
		downStream:newNetworkStream(),
	}
	return connection
}

// when receive tcp packet
func (connection *TcpConnection) onReceive(src, dst EndPoint, tcp *layers.TCP, timestamp time.Time) {
	connection.lastTimestamp = timestamp
	payload := tcp.Payload

	if !connection.isHttp {
		// skip no-http data
		if !isHttpRequestData(payload) {
			return
		}
		// receive first valid http data packet
		connection.clientId = src
		connection.isHttp = true
	}

	var sendStream, confirmStream *NetworkStream
	//var up bool
	if connection.clientId.equals(src) {
		sendStream = connection.upStream
		confirmStream = connection.downStream
		//up = true
	} else {
		sendStream = connection.downStream
		confirmStream = connection.upStream
		//up = false
	}

	if len(payload) > 0 {
		sendStream.appendPacket(tcp)
	}

	if tcp.SYN {
		// do nothing
	}

	if tcp.ACK {
		// confirm
		confirmStream.confirmPacket(tcp.Ack)
	}

	// terminate connection
	if tcp.FIN || tcp.RST {
		sendStream.closed = true
	}
}

func (connection *TcpConnection) closed() bool {
	return connection.upStream.closed && connection.downStream.closed
}

func (connection *TcpConnection) finish() {
	connection.upStream.finish()
	connection.downStream.finish()
}

// tread one-direction tcp data as stream. impl reader closer
type NetworkStream struct {
	buffer []*layers.TCP
	c      chan []byte
	remain []byte
	ignore bool
	closed bool
}

func newNetworkStream() *NetworkStream {
	return &NetworkStream{c:make(chan []byte, 1000)}
}

func (stream *NetworkStream) appendPacket(tcp *layers.TCP) {
	if stream.ignore {
		return
	}
	stream.buffer = append(stream.buffer, tcp)
}

func (stream *NetworkStream) confirmPacket(ack uint32) {
	if stream.ignore {
		return
	}
	var confirmedBuffer, remainedBuffer Buffer
	for _, tcp := range stream.buffer {
		if compareTcpSeq(tcp.Seq, ack) <= 0 {
			confirmedBuffer = append(confirmedBuffer, tcp)
		} else {
			remainedBuffer = append(remainedBuffer, tcp)
		}
	}

	if len(confirmedBuffer) > 0 {
		sort.Sort(confirmedBuffer)
	}
	var lastSeq uint32
	for _, tcp := range confirmedBuffer {
		seq := uint32(tcp.Seq)
		if (seq == lastSeq) {
			continue
		}
		lastSeq = seq
		stream.c <- tcp.Payload
	}

	stream.buffer = remainedBuffer
}

func (stream *NetworkStream) finish() {
	close(stream.c)
}

func (stream *NetworkStream) Read(p []byte) (n int, err error) {
	for len(stream.remain) == 0 {
		data, ok := <-stream.c
		if !ok {
			err = io.EOF
			return
		}
		stream.remain = data
	}

	if len(stream.remain) > len(p) {
		n = copy(p, stream.remain[:len(p)])
		stream.remain = stream.remain[len(p):]
	} else {
		n = copy(p, stream.remain)
		stream.remain = nil
	}
	return
}

func (stream *NetworkStream) Close() error {
	stream.ignore = true
	return nil
}

type Buffer []*layers.TCP
// impl sort.Interface
// Len is the number of elements in the collection.
func (buffer Buffer) Len() int {
	return len(buffer)
}
// Less reports whether the element with
// index i should sort before the element with index j.
func (buffer Buffer) Less(i, j int) bool {
	return compareTcpSeq(buffer[i].Seq, buffer[j].Seq) < 0
}
// Swap swaps the elements with indexes i and j.
func (buffer Buffer) Swap(i, j int) {
	buffer[i], buffer[j] = buffer[j], buffer[i]
}

// compare two tcp sequences, if seq1 is earlier, return num < 0, if seq1 == seq2, return 0, else return num > 0
func compareTcpSeq(seq1, seq2 uint32) int {
	if seq1 < TCP_SEQ_WINDOW && seq2 > MAX_TCP_SEQ - TCP_SEQ_WINDOW {
		return int(seq1 + MAX_TCP_SEQ - seq2)
	} else if seq2 < TCP_SEQ_WINDOW && seq1 > MAX_TCP_SEQ - TCP_SEQ_WINDOW {
		return int(seq1 - (MAX_TCP_SEQ + seq2))
	}
	return int(int32(seq1 - seq2))
}

var HTTP_METHODS = map[string]bool{"GET":true, "POST":true, "PUT":true, "DELETE":true, "HEAD":true, "TRACE":true,
	"OPTIONS":true, "PATCH":true}

// if is first http request packet
func isHttpRequestData(body []byte) bool {
	if len(body) < 8 {
		return false
	}
	data := body[0:8]
	idx := bytes.IndexByte(data, byte(' '))
	if (idx < 0) {
		return false
	}

	method := string(data[:idx])
	return HTTP_METHODS[method]
}<commit_msg>use ringbuffer as receive window<commit_after>package main

import (
	"github.com/google/gopacket/layers"
	"bytes"
	"time"
	"github.com/google/gopacket"
	"io"
	"sync"
	"strconv"
)

// gopacket provide a tcp connection, however it split one tcp connection into two stream.
// So it is hard to match http request and response. we make our own connection here

const MAX_TCP_SEQ uint32 = 0xFFFFFFFF
const TCP_SEQ_WINDOW = 1 << 10

type TcpAssembler struct {
	connectionDict    map[string]*TcpConnection
	lock              sync.Mutex
	connectionHandler ConnectionHandler
	filterIp          string
	filterPort        uint16
}

func newTcpAssembler(connectionHandler ConnectionHandler) *TcpAssembler {
	return &TcpAssembler{connectionDict:map[string]*TcpConnection{}, connectionHandler:connectionHandler}
}

func (assembler *TcpAssembler) assemble(flow gopacket.Flow, tcp *layers.TCP, timestamp time.Time) {
	src := EndPoint{ip:flow.Src().String(), port:uint16(tcp.SrcPort)}
	dst := EndPoint{ip:flow.Dst().String(), port:uint16(tcp.DstPort)}
	dropped := false
	if assembler.filterIp != "" {
		if src.ip != assembler.filterIp && dst.ip != assembler.filterIp {
			dropped = true
		}
	}
	if assembler.filterPort != 0 {
		if src.port != assembler.filterPort && dst.port != assembler.filterPort {
			dropped = true
		}
	}
	if dropped {
		return
	}

	srcString := src.String()
	dstString := dst.String()
	var key string
	if srcString < dstString {
		key = srcString + "-" + dstString
	} else {
		key = dstString + "-" + srcString
	}
	connection := assembler.retrieveConnection(src, dst, key)

	connection.onReceive(src, dst, tcp, timestamp)

	if connection.closed() {
		assembler.deleteConnection(key)
		connection.finish()
	}

	//TODO: cleanup timeout connections
}

// get connection this packet belong to; create new one if is new connection
func (assembler *TcpAssembler) retrieveConnection(src, dst EndPoint, key string) *TcpConnection {
	assembler.lock.Lock()
	defer assembler.lock.Unlock()
	connection := assembler.connectionDict[key]
	if (connection == nil) {
		connection = newTcpConnection(key)
		assembler.connectionDict[key] = connection
		assembler.connectionHandler.handle(src, dst, connection)
	}
	return connection
}

// remove connection (when is closed or timeout)
func (assembler *TcpAssembler) deleteConnection(key string) {
	assembler.lock.Lock()
	defer assembler.lock.Unlock()
	delete(assembler.connectionDict, key)
}

// flush older packets
func (assembler *TcpAssembler) flushOlderThan(time time.Time) {
	var connections []*TcpConnection
	assembler.lock.Lock()
	for _, connection := range assembler.connectionDict {
		connections = append(connections, connection)
	}
	assembler.lock.Unlock()

	//for _, connection := range connections {
	//	connection.flushOlderThan(time)
	//}
}

func (assembler *TcpAssembler) finishAll() {
	assembler.lock.Lock()
	defer assembler.lock.Unlock()
	for _, connection := range assembler.connectionDict {
		connection.finish()
	}
	assembler.connectionDict = nil
	assembler.connectionHandler.finish()
}

type ConnectionHandler interface {
	handle(src EndPoint, dst EndPoint, connection *TcpConnection)
	finish()
}

// one tcp connection
type  TcpConnection struct {
	upStream      *NetworkStream // stream from client to server
	downStream    *NetworkStream // stream from server to client
	clientId      EndPoint       // the client key(by ip and port)
	lastTimestamp time.Time      // timestamp receive last packet
	isHttp        bool
	key           string
}

type EndPoint struct {
	ip   string
	port uint16
}

func (p EndPoint) equals(p2 EndPoint) bool {
	return p.ip == p2.ip && p.port == p2.port
}

func (p EndPoint) String() string {
	return p.ip + ":" + strconv.Itoa(int(p.port))
}

type ConnectionId struct {
	src EndPoint
	dst EndPoint
}

// create tcp connection, by the first tcp packet. this packet should from client to server
func newTcpConnection(key string) *TcpConnection {
	connection := &TcpConnection{
		upStream:newNetworkStream(),
		downStream:newNetworkStream(),
		key: key,
	}
	return connection
}

// when receive tcp packet
func (connection *TcpConnection) onReceive(src, dst EndPoint, tcp *layers.TCP, timestamp time.Time) {
	connection.lastTimestamp = timestamp
	payload := tcp.Payload

	if !connection.isHttp {
		// skip no-http data
		if !isHttpRequestData(payload) {
			return
		}
		// receive first valid http data packet
		connection.clientId = src
		connection.isHttp = true
	}

	var sendStream, confirmStream *NetworkStream
	//var up bool
	if connection.clientId.equals(src) {
		sendStream = connection.upStream
		confirmStream = connection.downStream
		//up = true
	} else {
		sendStream = connection.downStream
		confirmStream = connection.upStream
		//up = false
	}

	if len(payload) > 0 {
		sendStream.appendPacket(tcp)
	}

	if tcp.SYN {
		// do nothing
	}

	if tcp.ACK {
		// confirm
		confirmStream.confirmPacket(tcp.Ack)
	}

	// terminate connection
	if tcp.FIN || tcp.RST {
		sendStream.closed = true
	}
}

func (connection *TcpConnection) flushOlderThan(time time.Time) {
	// flush all data
	//connection.upStream.window
	//connection.downStream.window
	// remove and close connection
	connection.upStream.closed = true
	connection.downStream.closed = true
	connection.finish()

}

func (connection *TcpConnection) closed() bool {
	return connection.upStream.closed && connection.downStream.closed
}

func (connection *TcpConnection) finish() {
	connection.upStream.finish()
	connection.downStream.finish()
}

// tread one-direction tcp data as stream. impl reader closer
type NetworkStream struct {
	window *ReceiveWindow
	c      chan []byte
	remain []byte
	ignore bool
	closed bool
}

func newNetworkStream() *NetworkStream {
	return &NetworkStream{window:newReceiveWindow(), c:make(chan []byte, 1000)}
}

func (stream *NetworkStream) appendPacket(tcp *layers.TCP) {
	if stream.ignore {
		return
	}
	stream.window.insert(tcp)
}

func (stream *NetworkStream) confirmPacket(ack uint32) {
	if stream.ignore {
		return
	}
	stream.window.confirm(ack, stream.c)
}

func (stream *NetworkStream) finish() {
	close(stream.c)
}

func (stream *NetworkStream) Read(p []byte) (n int, err error) {
	for len(stream.remain) == 0 {
		data, ok := <-stream.c
		if !ok {
			err = io.EOF
			return
		}
		stream.remain = data
	}

	if len(stream.remain) > len(p) {
		n = copy(p, stream.remain[:len(p)])
		stream.remain = stream.remain[len(p):]
	} else {
		n = copy(p, stream.remain)
		stream.remain = nil
	}
	return
}

func (stream *NetworkStream) Close() error {
	stream.ignore = true
	return nil
}

type ReceiveWindow struct {
	size   int
	start  int
	buffer []*layers.TCP
}

func newReceiveWindow() *ReceiveWindow {
	buffer := make([]*layers.TCP, 32)
	return &ReceiveWindow{buffer:buffer}
}

func (window *ReceiveWindow) destroy() {
	window.size = 0
	window.start = 0
	window.buffer = nil
}

func (window *ReceiveWindow) insert(packet *layers.TCP) {
	idx := 0
	for ; idx < window.size; idx ++ {
		index := (idx + window.start) % len(window.buffer)
		current := window.buffer[index]
		result := compareTcpSeq(current.Seq, packet.Seq)
		if result == 0 {
			// duplicated
			return
		}
		if result > 0 {
			// insert at index
			break
		}
	}

	if window.size == len(window.buffer) {
		window.expand()
	}

	if idx == window.size {
		// append at last
		index := (idx + window.start) % len(window.buffer)
		window.buffer[index] = packet
	} else {
		// insert at index
		for i := window.size; i >= idx; i-- {
			next := (i + window.start + 1) % len(window.buffer)
			current := (i + window.start) % len(window.buffer)
			window.buffer[next] = window.buffer[current]
		}
		index := (idx + window.start) % len(window.buffer)
		window.buffer[index] = packet

	}

	window.size++
}

func (window *ReceiveWindow) confirm(ack uint32, c chan []byte) {
	idx := 0
	for ; idx < window.size; idx ++ {
		index := (idx + window.start) % len(window.buffer)
		current := window.buffer[index]
		result := compareTcpSeq(current.Seq, ack)
		if result > 0 {
			break
		}
		window.buffer[index] = nil
		c <- current.Payload
	}
	window.start = (window.start + idx) % len(window.buffer)
	window.size = window.size - idx
}

func (window *ReceiveWindow) expand() {
	buffer := make([]*layers.TCP, len(window.buffer) * 2)
	end := window.start + window.size
	if end < len(window.buffer) {
		copy(buffer, window.buffer[window.start:window.start + window.size])
	} else {
		copy(buffer, window.buffer[window.start:])
		copy(buffer[len(window.buffer) - window.start:], window.buffer[:end - len(window.buffer)])
	}
	window.start = 0
	window.buffer = buffer
}

type Buffer []*layers.TCP
// impl sort.Interface
// Len is the number of elements in the collection.
func (buffer Buffer) Len() int {
	return len(buffer)
}
// Less reports whether the element with
// index i should sort before the element with index j.
func (buffer Buffer) Less(i, j int) bool {
	return compareTcpSeq(buffer[i].Seq, buffer[j].Seq) < 0
}
// Swap swaps the elements with indexes i and j.
func (buffer Buffer) Swap(i, j int) {
	buffer[i], buffer[j] = buffer[j], buffer[i]
}

// compare two tcp sequences, if seq1 is earlier, return num < 0, if seq1 == seq2, return 0, else return num > 0
func compareTcpSeq(seq1, seq2 uint32) int {
	if seq1 < TCP_SEQ_WINDOW && seq2 > MAX_TCP_SEQ - TCP_SEQ_WINDOW {
		return int(seq1 + MAX_TCP_SEQ - seq2)
	} else if seq2 < TCP_SEQ_WINDOW && seq1 > MAX_TCP_SEQ - TCP_SEQ_WINDOW {
		return int(seq1 - (MAX_TCP_SEQ + seq2))
	}
	return int(int32(seq1 - seq2))
}

var HTTP_METHODS = map[string]bool{"GET":true, "POST":true, "PUT":true, "DELETE":true, "HEAD":true, "TRACE":true,
	"OPTIONS":true, "PATCH":true}

// if is first http request packet
func isHttpRequestData(body []byte) bool {
	if len(body) < 8 {
		return false
	}
	data := body[0:8]
	idx := bytes.IndexByte(data, byte(' '))
	if (idx < 0) {
		return false
	}

	method := string(data[:idx])
	return HTTP_METHODS[method]
}<|endoftext|>


<commit_before>// Package connhelper provides helpers for connecting to a remote daemon host with custom logic.
package connhelper

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/docker/cli/cli/connhelper/ssh"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// ConnectionHelper allows to connect to a remote host with custom stream provider binary.
type ConnectionHelper struct {
	Dialer func(ctx context.Context, network, addr string) (net.Conn, error)
	Host   string // dummy URL used for HTTP requests. e.g. "http://docker"
}

// GetConnectionHelper returns Docker-specific connection helper for the given URL.
// GetConnectionHelper returns nil without error when no helper is registered for the scheme.
// URL is like "ssh://me@server01".
func GetConnectionHelper(daemonURL string) (*ConnectionHelper, error) {
	u, err := url.Parse(daemonURL)
	if err != nil {
		return nil, err
	}
	switch scheme := u.Scheme; scheme {
	case "ssh":
		sshCmd, sshArgs, err := ssh.New(daemonURL)
		if err != nil {
			return nil, err
		}
		return &ConnectionHelper{
			Dialer: func(ctx context.Context, network, addr string) (net.Conn, error) {
				return newCommandConn(ctx, sshCmd, sshArgs...)
			},
			Host: "http://docker",
		}, nil
	}
	// Future version may support plugins via ~/.docker/config.json. e.g. "dind"
	// See docker/cli#889 for the previous discussion.
	return nil, err
}

func newCommandConn(ctx context.Context, cmd string, args ...string) (net.Conn, error) {
	var (
		c   commandConn
		err error
	)
	c.cmd = exec.CommandContext(ctx, cmd, args...)
	// we assume that args never contains sensitive information
	logrus.Debugf("connhelper: starting %s with %v", cmd, args)
	c.cmd.Env = os.Environ()
	setPdeathsig(c.cmd)
	c.stdin, err = c.cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	c.stdout, err = c.cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	c.cmd.Stderr = &stderrWriter{
		stderrMu:    &c.stderrMu,
		stderr:      &c.stderr,
		debugPrefix: fmt.Sprintf("connhelper (%s):", cmd),
	}
	c.localAddr = dummyAddr{network: "dummy", s: "dummy-0"}
	c.remoteAddr = dummyAddr{network: "dummy", s: "dummy-1"}
	return &c, c.cmd.Start()
}

// commandConn implements net.Conn
type commandConn struct {
	cmd           *exec.Cmd
	cmdExited     bool
	cmdWaitErr    error
	cmdMutex      sync.Mutex
	stdin         io.WriteCloser
	stdout        io.ReadCloser
	stderrMu      sync.Mutex
	stderr        bytes.Buffer
	stdioClosedMu sync.Mutex // for stdinClosed and stdoutClosed
	stdinClosed   bool
	stdoutClosed  bool
	localAddr     net.Addr
	remoteAddr    net.Addr
}

// killIfStdioClosed kills the cmd if both stdin and stdout are closed.
func (c *commandConn) killIfStdioClosed() error {
	c.stdioClosedMu.Lock()
	stdioClosed := c.stdoutClosed && c.stdinClosed
	c.stdioClosedMu.Unlock()
	if !stdioClosed {
		return nil
	}
	return c.kill()
}

// killAndWait tries sending SIGTERM to the process before sending SIGKILL.
func killAndWait(cmd *exec.Cmd) error {
	var werr error
	if runtime.GOOS != "windows" {
		werrCh := make(chan error)
		go func() { werrCh <- cmd.Wait() }()
		cmd.Process.Signal(syscall.SIGTERM)
		select {
		case werr = <-werrCh:
		case <-time.After(3 * time.Second):
			cmd.Process.Kill()
			werr = <-werrCh
		}
	} else {
		cmd.Process.Kill()
		werr = cmd.Wait()
	}
	return werr
}

// kill returns nil if the command terminated, regardless to the exit status.
func (c *commandConn) kill() error {
	var werr error
	c.cmdMutex.Lock()
	if c.cmdExited {
		werr = c.cmdWaitErr
	} else {
		werr = killAndWait(c.cmd)
		c.cmdWaitErr = werr
		c.cmdExited = true
	}
	c.cmdMutex.Unlock()
	if werr == nil {
		return nil
	}
	wExitErr, ok := werr.(*exec.ExitError)
	if ok {
		if wExitErr.ProcessState.Exited() {
			return nil
		}
	}
	return errors.Wrapf(werr, "connhelper: failed to wait")
}

func (c *commandConn) onEOF(eof error) error {
	// when we got EOF, the command is going to be terminated
	var werr error
	c.cmdMutex.Lock()
	if c.cmdExited {
		werr = c.cmdWaitErr
	} else {
		werrCh := make(chan error)
		go func() { werrCh <- c.cmd.Wait() }()
		select {
		case werr = <-werrCh:
			c.cmdWaitErr = werr
			c.cmdExited = true
		case <-time.After(10 * time.Second):
			c.cmdMutex.Unlock()
			c.stderrMu.Lock()
			stderr := c.stderr.String()
			c.stderrMu.Unlock()
			return errors.Errorf("command %v did not exit after %v: stderr=%q", c.cmd.Args, eof, stderr)
		}
	}
	c.cmdMutex.Unlock()
	if werr == nil {
		return eof
	}
	c.stderrMu.Lock()
	stderr := c.stderr.String()
	c.stderrMu.Unlock()
	return errors.Errorf("command %v has exited with %v, please make sure the URL is valid, and Docker 18.09 or later is installed on the remote host: stderr=%q", c.cmd.Args, werr, stderr)
}

func ignorableCloseError(err error) bool {
	errS := err.Error()
	ss := []string{
		os.ErrClosed.Error(),
	}
	for _, s := range ss {
		if strings.Contains(errS, s) {
			return true
		}
	}
	return false
}

func (c *commandConn) CloseRead() error {
	// NOTE: maybe already closed here
	if err := c.stdout.Close(); err != nil && !ignorableCloseError(err) {
		logrus.Warnf("commandConn.CloseRead: %v", err)
	}
	c.stdioClosedMu.Lock()
	c.stdoutClosed = true
	c.stdioClosedMu.Unlock()
	if err := c.killIfStdioClosed(); err != nil {
		logrus.Warnf("commandConn.CloseRead: %v", err)
	}
	return nil
}

func (c *commandConn) Read(p []byte) (int, error) {
	n, err := c.stdout.Read(p)
	if err == io.EOF {
		err = c.onEOF(err)
	}
	return n, err
}

func (c *commandConn) CloseWrite() error {
	// NOTE: maybe already closed here
	if err := c.stdin.Close(); err != nil && !ignorableCloseError(err) {
		logrus.Warnf("commandConn.CloseWrite: %v", err)
	}
	c.stdioClosedMu.Lock()
	c.stdinClosed = true
	c.stdioClosedMu.Unlock()
	if err := c.killIfStdioClosed(); err != nil {
		logrus.Warnf("commandConn.CloseWrite: %v", err)
	}
	return nil
}

func (c *commandConn) Write(p []byte) (int, error) {
	n, err := c.stdin.Write(p)
	if err == io.EOF {
		err = c.onEOF(err)
	}
	return n, err
}

func (c *commandConn) Close() error {
	var err error
	if err = c.CloseRead(); err != nil {
		logrus.Warnf("commandConn.Close: CloseRead: %v", err)
	}
	if err = c.CloseWrite(); err != nil {
		logrus.Warnf("commandConn.Close: CloseWrite: %v", err)
	}
	return err
}

func (c *commandConn) LocalAddr() net.Addr {
	return c.localAddr
}
func (c *commandConn) RemoteAddr() net.Addr {
	return c.remoteAddr
}
func (c *commandConn) SetDeadline(t time.Time) error {
	logrus.Debugf("unimplemented call: SetDeadline(%v)", t)
	return nil
}
func (c *commandConn) SetReadDeadline(t time.Time) error {
	logrus.Debugf("unimplemented call: SetReadDeadline(%v)", t)
	return nil
}
func (c *commandConn) SetWriteDeadline(t time.Time) error {
	logrus.Debugf("unimplemented call: SetWriteDeadline(%v)", t)
	return nil
}

type dummyAddr struct {
	network string
	s       string
}

func (d dummyAddr) Network() string {
	return d.network
}

func (d dummyAddr) String() string {
	return d.s
}

type stderrWriter struct {
	stderrMu    *sync.Mutex
	stderr      *bytes.Buffer
	debugPrefix string
}

func (w *stderrWriter) Write(p []byte) (int, error) {
	logrus.Debugf("%s%s", w.debugPrefix, string(p))
	w.stderrMu.Lock()
	if w.stderr.Len() > 4096 {
		w.stderr.Reset()
	}
	n, err := w.stderr.Write(p)
	w.stderrMu.Unlock()
	return n, err
}
<commit_msg>err message improve when ssh fail<commit_after>// Package connhelper provides helpers for connecting to a remote daemon host with custom logic.
package connhelper

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/docker/cli/cli/connhelper/ssh"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// ConnectionHelper allows to connect to a remote host with custom stream provider binary.
type ConnectionHelper struct {
	Dialer func(ctx context.Context, network, addr string) (net.Conn, error)
	Host   string // dummy URL used for HTTP requests. e.g. "http://docker"
}

// GetConnectionHelper returns Docker-specific connection helper for the given URL.
// GetConnectionHelper returns nil without error when no helper is registered for the scheme.
// URL is like "ssh://me@server01".
func GetConnectionHelper(daemonURL string) (*ConnectionHelper, error) {
	u, err := url.Parse(daemonURL)
	if err != nil {
		return nil, err
	}
	switch scheme := u.Scheme; scheme {
	case "ssh":
		sshCmd, sshArgs, err := ssh.New(daemonURL)
		if err != nil {
			return nil, err
		}
		return &ConnectionHelper{
			Dialer: func(ctx context.Context, network, addr string) (net.Conn, error) {
				return newCommandConn(ctx, sshCmd, sshArgs...)
			},
			Host: "http://docker",
		}, nil
	}
	// Future version may support plugins via ~/.docker/config.json. e.g. "dind"
	// See docker/cli#889 for the previous discussion.
	return nil, err
}

func newCommandConn(ctx context.Context, cmd string, args ...string) (net.Conn, error) {
	var (
		c   commandConn
		err error
	)
	c.cmd = exec.CommandContext(ctx, cmd, args...)
	// we assume that args never contains sensitive information
	logrus.Debugf("connhelper: starting %s with %v", cmd, args)
	c.cmd.Env = os.Environ()
	setPdeathsig(c.cmd)
	c.stdin, err = c.cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	c.stdout, err = c.cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	c.cmd.Stderr = &stderrWriter{
		stderrMu:    &c.stderrMu,
		stderr:      &c.stderr,
		debugPrefix: fmt.Sprintf("connhelper (%s):", cmd),
	}
	c.localAddr = dummyAddr{network: "dummy", s: "dummy-0"}
	c.remoteAddr = dummyAddr{network: "dummy", s: "dummy-1"}
	return &c, c.cmd.Start()
}

// commandConn implements net.Conn
type commandConn struct {
	cmd           *exec.Cmd
	cmdExited     bool
	cmdWaitErr    error
	cmdMutex      sync.Mutex
	stdin         io.WriteCloser
	stdout        io.ReadCloser
	stderrMu      sync.Mutex
	stderr        bytes.Buffer
	stdioClosedMu sync.Mutex // for stdinClosed and stdoutClosed
	stdinClosed   bool
	stdoutClosed  bool
	localAddr     net.Addr
	remoteAddr    net.Addr
}

// killIfStdioClosed kills the cmd if both stdin and stdout are closed.
func (c *commandConn) killIfStdioClosed() error {
	c.stdioClosedMu.Lock()
	stdioClosed := c.stdoutClosed && c.stdinClosed
	c.stdioClosedMu.Unlock()
	if !stdioClosed {
		return nil
	}
	return c.kill()
}

// killAndWait tries sending SIGTERM to the process before sending SIGKILL.
func killAndWait(cmd *exec.Cmd) error {
	var werr error
	if runtime.GOOS != "windows" {
		werrCh := make(chan error)
		go func() { werrCh <- cmd.Wait() }()
		cmd.Process.Signal(syscall.SIGTERM)
		select {
		case werr = <-werrCh:
		case <-time.After(3 * time.Second):
			cmd.Process.Kill()
			werr = <-werrCh
		}
	} else {
		cmd.Process.Kill()
		werr = cmd.Wait()
	}
	return werr
}

// kill returns nil if the command terminated, regardless to the exit status.
func (c *commandConn) kill() error {
	var werr error
	c.cmdMutex.Lock()
	if c.cmdExited {
		werr = c.cmdWaitErr
	} else {
		werr = killAndWait(c.cmd)
		c.cmdWaitErr = werr
		c.cmdExited = true
	}
	c.cmdMutex.Unlock()
	if werr == nil {
		return nil
	}
	wExitErr, ok := werr.(*exec.ExitError)
	if ok {
		if wExitErr.ProcessState.Exited() {
			return nil
		}
	}
	return errors.Wrapf(werr, "connhelper: failed to wait")
}

func (c *commandConn) onEOF(eof error) error {
	// when we got EOF, the command is going to be terminated
	var werr error
	c.cmdMutex.Lock()
	if c.cmdExited {
		werr = c.cmdWaitErr
	} else {
		werrCh := make(chan error)
		go func() { werrCh <- c.cmd.Wait() }()
		select {
		case werr = <-werrCh:
			c.cmdWaitErr = werr
			c.cmdExited = true
		case <-time.After(10 * time.Second):
			c.cmdMutex.Unlock()
			c.stderrMu.Lock()
			stderr := c.stderr.String()
			c.stderrMu.Unlock()
			return errors.Errorf("command %v did not exit after %v: stderr=%q", c.cmd.Args, eof, stderr)
		}
	}
	c.cmdMutex.Unlock()
	if werr == nil {
		return eof
	}
	c.stderrMu.Lock()
	stderr := c.stderr.String()
	c.stderrMu.Unlock()
	return errors.Errorf("command %v has exited with %v, please make sure the URL is valid, and Docker 18.09 or later is installed on the remote host: stderr=%s", c.cmd.Args, werr, stderr)
}

func ignorableCloseError(err error) bool {
	errS := err.Error()
	ss := []string{
		os.ErrClosed.Error(),
	}
	for _, s := range ss {
		if strings.Contains(errS, s) {
			return true
		}
	}
	return false
}

func (c *commandConn) CloseRead() error {
	// NOTE: maybe already closed here
	if err := c.stdout.Close(); err != nil && !ignorableCloseError(err) {
		logrus.Warnf("commandConn.CloseRead: %v", err)
	}
	c.stdioClosedMu.Lock()
	c.stdoutClosed = true
	c.stdioClosedMu.Unlock()
	if err := c.killIfStdioClosed(); err != nil {
		logrus.Warnf("commandConn.CloseRead: %v", err)
	}
	return nil
}

func (c *commandConn) Read(p []byte) (int, error) {
	n, err := c.stdout.Read(p)
	if err == io.EOF {
		err = c.onEOF(err)
	}
	return n, err
}

func (c *commandConn) CloseWrite() error {
	// NOTE: maybe already closed here
	if err := c.stdin.Close(); err != nil && !ignorableCloseError(err) {
		logrus.Warnf("commandConn.CloseWrite: %v", err)
	}
	c.stdioClosedMu.Lock()
	c.stdinClosed = true
	c.stdioClosedMu.Unlock()
	if err := c.killIfStdioClosed(); err != nil {
		logrus.Warnf("commandConn.CloseWrite: %v", err)
	}
	return nil
}

func (c *commandConn) Write(p []byte) (int, error) {
	n, err := c.stdin.Write(p)
	if err == io.EOF {
		err = c.onEOF(err)
	}
	return n, err
}

func (c *commandConn) Close() error {
	var err error
	if err = c.CloseRead(); err != nil {
		logrus.Warnf("commandConn.Close: CloseRead: %v", err)
	}
	if err = c.CloseWrite(); err != nil {
		logrus.Warnf("commandConn.Close: CloseWrite: %v", err)
	}
	return err
}

func (c *commandConn) LocalAddr() net.Addr {
	return c.localAddr
}
func (c *commandConn) RemoteAddr() net.Addr {
	return c.remoteAddr
}
func (c *commandConn) SetDeadline(t time.Time) error {
	logrus.Debugf("unimplemented call: SetDeadline(%v)", t)
	return nil
}
func (c *commandConn) SetReadDeadline(t time.Time) error {
	logrus.Debugf("unimplemented call: SetReadDeadline(%v)", t)
	return nil
}
func (c *commandConn) SetWriteDeadline(t time.Time) error {
	logrus.Debugf("unimplemented call: SetWriteDeadline(%v)", t)
	return nil
}

type dummyAddr struct {
	network string
	s       string
}

func (d dummyAddr) Network() string {
	return d.network
}

func (d dummyAddr) String() string {
	return d.s
}

type stderrWriter struct {
	stderrMu    *sync.Mutex
	stderr      *bytes.Buffer
	debugPrefix string
}

func (w *stderrWriter) Write(p []byte) (int, error) {
	logrus.Debugf("%s%s", w.debugPrefix, string(p))
	w.stderrMu.Lock()
	if w.stderr.Len() > 4096 {
		w.stderr.Reset()
	}
	n, err := w.stderr.Write(p)
	w.stderrMu.Unlock()
	return n, err
}
<|endoftext|>


<commit_before>package spoon

import (
	"crypto/tls"
	"log"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

// TCPListener is spoon's Listener interface with extra helper methods.
type TCPListener interface {
	net.Listener
	SetKeepAlive(d time.Duration)
	SetForceTimeout(d time.Duration)
	ListenAndServe(addr string, handler http.Handler) error
	ListenAndServeTLS(addr string, certFile string, keyFile string, handler http.Handler) error
	RunServer(server *http.Server) error
}

func newtcpListener(l *net.TCPListener) *tcpListener {
	return &tcpListener{
		TCPListener:          l,
		keepaliveDuration:    3 * time.Minute,
		forceTimeoutDuration: 5 * time.Minute,
		wg:                   new(sync.WaitGroup),
		conns:                map[*net.TCPConn]*net.TCPConn{},
		m:                    new(sync.Mutex),
	}
}

type tcpListener struct {
	*net.TCPListener
	keepaliveDuration    time.Duration
	forceTimeoutDuration time.Duration
	wg                   *sync.WaitGroup
	conns                map[*net.TCPConn]*net.TCPConn
	m                    *sync.Mutex
}

var _ net.Listener = new(tcpListener)
var _ TCPListener = new(tcpListener)

func (l *tcpListener) Accept() (net.Conn, error) {

	conn, err := l.TCPListener.AcceptTCP()
	if err != nil {
		return nil, err
	}

	conn.SetKeepAlive(true)                      // see http.tcpKeepAliveListener
	conn.SetKeepAlivePeriod(l.keepaliveDuration) // see http.tcpKeepAliveListener
	// conn.SetLinger(0) // is the default already accoring to the docs https://golang.org/pkg/net/#TCPConn.SetLinger

	zconn := zeroTCPConn{
		TCPConn: conn,
		wg:      l.wg,
		l:       l,
	}

	l.m.Lock()
	l.conns[conn] = conn
	l.m.Unlock()

	l.wg.Add(1)

	return zconn, nil
}

// blocking wait for close
func (l *tcpListener) Close() error {

	//stop accepting connections - release fd
	err := l.TCPListener.Close()
	c := make(chan struct{})

	go func() {

		log.Println("Closing Remaining Keepalives")

		l.m.Lock()
		for _, v := range l.conns {
			v.Close() // this is OK to close, see (*TCPConn) SetLinger, just can't reduce waitgroup until it's actually closed!
		}
		l.m.Unlock()
	}()

	go func() {
		l.wg.Wait()
		close(c)
	}()

	select {
	case <-c:
	// closed gracefully
	case <-time.After(l.forceTimeoutDuration):
		log.Println("timeout reached, force shutdown")
		// not waiting any longer, letting this go.
		// spoon will think it's been closed and when
		// the process dies, connections will get cut
	}

	return err
}

func (l *tcpListener) File() *os.File {

	// returns a dup(2) - FD_CLOEXEC flag *not* set
	tl := l.TCPListener
	fl, _ := tl.File()

	return fl
}

//notifying on close net.Conn
type zeroTCPConn struct {
	*net.TCPConn
	wg *sync.WaitGroup
	l  *tcpListener
}

func (conn zeroTCPConn) Close() (err error) {

	if err = conn.TCPConn.Close(); err != nil {
		log.Println("ERROR CLOSING CONNECTION, OK if connection already closed, we must have triggered a restart: ", err)
	}

	conn.l.m.Lock()
	delete(conn.l.conns, conn.TCPConn)
	conn.l.m.Unlock()

	conn.wg.Done()
	return
}

//
// HTTP Section for tcpListener
//

// ListenAndServe mimics the std libraries http.ListenAndServe but uses our custom listener
// for graceful restarts.
// NOTE: addr is ignored, the address of the listener is used, only reason it is a param is for
// easier conversion from stdlib http.ListenAndServe
func (l *tcpListener) ListenAndServe(addr string, handler http.Handler) error {

	server := &http.Server{Addr: l.Addr().String(), Handler: handler}

	go server.Serve(l)

	return nil
}

// ListenAndServeTLS mimics the std libraries http.ListenAndServeTLS but uses out custom listener
// for graceful restarts.
// NOTE: addr is ignored, the address of the listener is used, only reason it is a param is for
// easier conversion from stdlib http.ListenAndServeTLS
func (l *tcpListener) ListenAndServeTLS(addr string, certFile string, keyFile string, handler http.Handler) error {

	var err error

	tlsConfig := &tls.Config{
		NextProtos:   []string{"http/1.1"},
		Certificates: make([]tls.Certificate, 1),
	}

	tlsConfig.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(l, tlsConfig)

	server := &http.Server{Addr: tlsListener.Addr().String(), Handler: handler, TLSConfig: tlsConfig}

	go server.Serve(tlsListener)

	return nil
}

// RunServer runs the provided http.Server, if using TLS, TLSConfig must be setup prior to
// calling this function
func (l *tcpListener) RunServer(server *http.Server) error {

	var lis net.Listener

	if server.TLSConfig != nil {
		if server.TLSConfig.NextProtos == nil {
			server.TLSConfig.NextProtos = make([]string, 0)
		}

		if len(server.TLSConfig.NextProtos) == 0 || !strSliceContains(server.TLSConfig.NextProtos, "http/1.1") {
			server.TLSConfig.NextProtos = append(server.TLSConfig.NextProtos, "http/1.1")
		}

		lis = tls.NewListener(l, server.TLSConfig)
	} else {
		lis = l
	}

	go server.Serve(lis)

	return nil
}

// SetKeepAlive sets the listener's connection keep alive timeout.
// NOTE: method is NOT thread safe, must set prior to sp.Run()
// DEFAULT: time.Minute * 3
func (l *tcpListener) SetKeepAlive(d time.Duration) {
	l.keepaliveDuration = d
}

// SetKeepAlive sets the listener's connection keep alive timeout.
// NOTE: method is NOT thread safe, must set prior to sp.Run()
// DEFAULT: time.Minute * 5
func (l *tcpListener) SetForceTimeout(d time.Duration) {
	l.forceTimeoutDuration = d
}

func strSliceContains(ss []string, s string) bool {
	for _, v := range ss {
		if v == s {
			return true
		}
	}
	return false
}
<commit_msg>remove debug.<commit_after>package spoon

import (
	"crypto/tls"
	"log"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

// TCPListener is spoon's Listener interface with extra helper methods.
type TCPListener interface {
	net.Listener
	SetKeepAlive(d time.Duration)
	SetForceTimeout(d time.Duration)
	ListenAndServe(addr string, handler http.Handler) error
	ListenAndServeTLS(addr string, certFile string, keyFile string, handler http.Handler) error
	RunServer(server *http.Server) error
}

func newtcpListener(l *net.TCPListener) *tcpListener {
	return &tcpListener{
		TCPListener:          l,
		keepaliveDuration:    3 * time.Minute,
		forceTimeoutDuration: 5 * time.Minute,
		wg:                   new(sync.WaitGroup),
		conns:                map[*net.TCPConn]*net.TCPConn{},
		m:                    new(sync.Mutex),
	}
}

type tcpListener struct {
	*net.TCPListener
	keepaliveDuration    time.Duration
	forceTimeoutDuration time.Duration
	wg                   *sync.WaitGroup
	conns                map[*net.TCPConn]*net.TCPConn
	m                    *sync.Mutex
}

var _ net.Listener = new(tcpListener)
var _ TCPListener = new(tcpListener)

func (l *tcpListener) Accept() (net.Conn, error) {

	conn, err := l.TCPListener.AcceptTCP()
	if err != nil {
		return nil, err
	}

	conn.SetKeepAlive(true)                      // see http.tcpKeepAliveListener
	conn.SetKeepAlivePeriod(l.keepaliveDuration) // see http.tcpKeepAliveListener
	// conn.SetLinger(0) // is the default already accoring to the docs https://golang.org/pkg/net/#TCPConn.SetLinger

	zconn := zeroTCPConn{
		TCPConn: conn,
		wg:      l.wg,
		l:       l,
	}

	l.m.Lock()
	l.conns[conn] = conn
	l.m.Unlock()

	l.wg.Add(1)

	return zconn, nil
}

// blocking wait for close
func (l *tcpListener) Close() error {

	//stop accepting connections - release fd
	err := l.TCPListener.Close()
	c := make(chan struct{})

	go func() {
		l.m.Lock()
		for _, v := range l.conns {
			v.Close() // this is OK to close, see (*TCPConn) SetLinger, just can't reduce waitgroup until it's actually closed!
		}
		l.m.Unlock()
	}()

	go func() {
		l.wg.Wait()
		close(c)
	}()

	select {
	case <-c:
	// closed gracefully
	case <-time.After(l.forceTimeoutDuration):
		log.Println("timeout reached, force shutdown")
		// not waiting any longer, letting this go.
		// spoon will think it's been closed and when
		// the process dies, connections will get cut
	}

	return err
}

func (l *tcpListener) File() *os.File {

	// returns a dup(2) - FD_CLOEXEC flag *not* set
	tl := l.TCPListener
	fl, _ := tl.File()

	return fl
}

//notifying on close net.Conn
type zeroTCPConn struct {
	*net.TCPConn
	wg *sync.WaitGroup
	l  *tcpListener
}

func (conn zeroTCPConn) Close() (err error) {

	if err = conn.TCPConn.Close(); err != nil {
		log.Println("ERROR CLOSING CONNECTION, OK if connection already closed, we must have triggered a restart: ", err)
	}

	conn.l.m.Lock()
	delete(conn.l.conns, conn.TCPConn)
	conn.l.m.Unlock()

	conn.wg.Done()
	return
}

//
// HTTP Section for tcpListener
//

// ListenAndServe mimics the std libraries http.ListenAndServe but uses our custom listener
// for graceful restarts.
// NOTE: addr is ignored, the address of the listener is used, only reason it is a param is for
// easier conversion from stdlib http.ListenAndServe
func (l *tcpListener) ListenAndServe(addr string, handler http.Handler) error {

	server := &http.Server{Addr: l.Addr().String(), Handler: handler}

	go server.Serve(l)

	return nil
}

// ListenAndServeTLS mimics the std libraries http.ListenAndServeTLS but uses out custom listener
// for graceful restarts.
// NOTE: addr is ignored, the address of the listener is used, only reason it is a param is for
// easier conversion from stdlib http.ListenAndServeTLS
func (l *tcpListener) ListenAndServeTLS(addr string, certFile string, keyFile string, handler http.Handler) error {

	var err error

	tlsConfig := &tls.Config{
		NextProtos:   []string{"http/1.1"},
		Certificates: make([]tls.Certificate, 1),
	}

	tlsConfig.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	tlsListener := tls.NewListener(l, tlsConfig)

	server := &http.Server{Addr: tlsListener.Addr().String(), Handler: handler, TLSConfig: tlsConfig}

	go server.Serve(tlsListener)

	return nil
}

// RunServer runs the provided http.Server, if using TLS, TLSConfig must be setup prior to
// calling this function
func (l *tcpListener) RunServer(server *http.Server) error {

	var lis net.Listener

	if server.TLSConfig != nil {
		if server.TLSConfig.NextProtos == nil {
			server.TLSConfig.NextProtos = make([]string, 0)
		}

		if len(server.TLSConfig.NextProtos) == 0 || !strSliceContains(server.TLSConfig.NextProtos, "http/1.1") {
			server.TLSConfig.NextProtos = append(server.TLSConfig.NextProtos, "http/1.1")
		}

		lis = tls.NewListener(l, server.TLSConfig)
	} else {
		lis = l
	}

	go server.Serve(lis)

	return nil
}

// SetKeepAlive sets the listener's connection keep alive timeout.
// NOTE: method is NOT thread safe, must set prior to sp.Run()
// DEFAULT: time.Minute * 3
func (l *tcpListener) SetKeepAlive(d time.Duration) {
	l.keepaliveDuration = d
}

// SetKeepAlive sets the listener's connection keep alive timeout.
// NOTE: method is NOT thread safe, must set prior to sp.Run()
// DEFAULT: time.Minute * 5
func (l *tcpListener) SetForceTimeout(d time.Duration) {
	l.forceTimeoutDuration = d
}

func strSliceContains(ss []string, s string) bool {
	for _, v := range ss {
		if v == s {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package elastic_load_balancing

import (
	"github.com/jagregory/cfval/constraints"
	. "github.com/jagregory/cfval/schema"
)

// see: http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-elb-listener.html
var listener = NestedResource{
	Description: "ElasticLoadBalancing Listener",
	Properties: Properties{
		"InstancePort": Schema{
			Type:     ValueString,
			Required: constraints.Always,
		},

		"InstanceProtocol": Schema{
			Type: instanceProtocol,
			// TODO:
			// * If the front-end protocol is HTTP or HTTPS, InstanceProtocol has to
			//   be at the same protocol layer, i.e., HTTP or HTTPS. Likewise, if the
			//   front-end protocol is TCP or SSL, InstanceProtocol has to be TCP
			//   or SSL.
			// * If there is another listener with the same InstancePort whose
			//   InstanceProtocol is secure, i.e., HTTPS or SSL, the listener's
			//   InstanceProtocol has to be secure, i.e., HTTPS or SSL. If there is
			//   another listener with the same InstancePort whose InstanceProtocol is
			//   HTTP or TCP, the listener's InstanceProtocol must be either HTTP
			//   or TCP.
		},

		"LoadBalancerPort": Schema{
			Type:     ValueString,
			Required: constraints.Always,
		},

		"PolicyNames": Schema{
			Type: Multiple(policy),
		},

		"Protocol": Schema{
			Required: constraints.Always,
			Type:     instanceProtocol,
		},

		"SSLCertificateId": Schema{
			Type: ValueString,
		},
	},
}
<commit_msg>Fix for LoadBalancer Listeners PolicyNames type<commit_after>package elastic_load_balancing

import (
	"github.com/jagregory/cfval/constraints"
	. "github.com/jagregory/cfval/schema"
)

// see: http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-elb-listener.html
var listener = NestedResource{
	Description: "ElasticLoadBalancing Listener",
	Properties: Properties{
		"InstancePort": Schema{
			Type:     ValueString,
			Required: constraints.Always,
		},

		"InstanceProtocol": Schema{
			Type: instanceProtocol,
			// TODO:
			// * If the front-end protocol is HTTP or HTTPS, InstanceProtocol has to
			//   be at the same protocol layer, i.e., HTTP or HTTPS. Likewise, if the
			//   front-end protocol is TCP or SSL, InstanceProtocol has to be TCP
			//   or SSL.
			// * If there is another listener with the same InstancePort whose
			//   InstanceProtocol is secure, i.e., HTTPS or SSL, the listener's
			//   InstanceProtocol has to be secure, i.e., HTTPS or SSL. If there is
			//   another listener with the same InstancePort whose InstanceProtocol is
			//   HTTP or TCP, the listener's InstanceProtocol must be either HTTP
			//   or TCP.
		},

		"LoadBalancerPort": Schema{
			Type:     ValueString,
			Required: constraints.Always,
		},

		"PolicyNames": Schema{
			Type: Multiple(ValueString),
		},

		"Protocol": Schema{
			Required: constraints.Always,
			Type:     instanceProtocol,
		},

		"SSLCertificateId": Schema{
			Type: ValueString,
		},
	},
}
<|endoftext|>


<commit_before>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by the license
// that can be found in the LICENSE file.

package batchfossilizer

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/stratumn/go/fossilizer"
	"github.com/stratumn/goprivate/merkle"
	"github.com/stratumn/goprivate/testutil"
	"github.com/stratumn/goprivate/types"
)

const interval = 10 * time.Millisecond

var (
	pathA0     merkle.Path
	pathAB0    merkle.Path
	pathAB1    merkle.Path
	pathABC0   merkle.Path
	pathABC1   merkle.Path
	pathABC2   merkle.Path
	pathABCD0  merkle.Path
	pathABCD1  merkle.Path
	pathABCD2  merkle.Path
	pathABCD3  merkle.Path
	pathABCDE0 merkle.Path
	pathABCDE1 merkle.Path
	pathABCDE2 merkle.Path
	pathABCDE3 merkle.Path
	pathABCDE4 merkle.Path
)

func TestMain(m *testing.M) {
	seed := int64(time.Now().Nanosecond())
	fmt.Printf("using seed %d\n", seed)
	rand.Seed(seed)

	loadPath("testdata/path-a-0.json", &pathA0)
	loadPath("testdata/path-ab-0.json", &pathAB0)
	loadPath("testdata/path-ab-1.json", &pathAB1)
	loadPath("testdata/path-abc-0.json", &pathABC0)
	loadPath("testdata/path-abc-1.json", &pathABC1)
	loadPath("testdata/path-abc-2.json", &pathABC2)
	loadPath("testdata/path-abcd-0.json", &pathABCD0)
	loadPath("testdata/path-abcd-1.json", &pathABCD1)
	loadPath("testdata/path-abcd-2.json", &pathABCD2)
	loadPath("testdata/path-abcd-3.json", &pathABCD3)
	loadPath("testdata/path-abcde-0.json", &pathABCDE0)
	loadPath("testdata/path-abcde-1.json", &pathABCDE1)
	loadPath("testdata/path-abcde-2.json", &pathABCDE2)
	loadPath("testdata/path-abcde-3.json", &pathABCDE3)
	loadPath("testdata/path-abcde-4.json", &pathABCDE4)

	flag.Parse()
	os.Exit(m.Run())
}

func TestGetInfo(t *testing.T) {
	a := New(&Config{})
	info, err := a.GetInfo()
	if err != nil {
		t.Fatal(err)
	}
	if info == nil {
		t.Fatal("info is nil")
	}
}

func loadPath(filename string, path *merkle.Path) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}
	if err = json.Unmarshal(data, path); err != nil {
		panic(err)
	}
}

func TestFossilize(t *testing.T) {
	a := New(&Config{Interval: interval})
	tests := []fossilizeTest{
		{atos(sha256.Sum256([]byte("a"))), []byte("test a"), pathABCDE0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b"), pathABCDE1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c"), pathABCDE2, 0, false},
		{atos(sha256.Sum256([]byte("d"))), []byte("test d"), pathABCDE3, 0, false},
		{atos(sha256.Sum256([]byte("e"))), []byte("test e"), pathABCDE4, 0, false},
	}
	testFossilizeMultiple(t, a, tests)
}

func TestFossilizeMaxLeaves(t *testing.T) {
	a := New(&Config{Interval: interval, MaxLeaves: 4})
	tests := []fossilizeTest{
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 1"), pathABCD0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 1"), pathABCD1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 1"), pathABCD2, 0, false},
		{atos(sha256.Sum256([]byte("d"))), []byte("test d 1"), pathABCD3, 0, false},
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 2"), pathABC0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 2"), pathABC1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 2"), pathABC2, 0, false},
	}
	testFossilizeMultiple(t, a, tests)
}

func TestFossilizeInterval(t *testing.T) {
	a := New(&Config{Interval: interval})
	tests := []fossilizeTest{
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 1"), pathABC0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 1"), pathABC1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 1"), pathABC2, interval * 2, false},
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 2"), pathABCD0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 2"), pathABCD1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 2"), pathABCD2, 0, false},
		{atos(sha256.Sum256([]byte("d"))), []byte("test d 2"), pathABCD3, interval * 2, false},
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 3"), pathABC0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 3"), pathABC1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 3"), pathABC2, 0, false},
	}
	testFossilizeMultiple(t, a, tests)
}

func BenchmarkFossilizeMaxLeaves100(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 100})
}

func BenchmarkFossilizeMaxLeaves1000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 1000})
}

func BenchmarkFossilizeMaxLeaves10000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 10000})
}

func BenchmarkFossilizeMaxLeaves100000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 100000})
}

func BenchmarkFossilizeMaxLeaves1000000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 1000000})
}

type fossilizeTest struct {
	data       []byte
	meta       []byte
	path       merkle.Path
	sleep      time.Duration
	fossilized bool
}

func testFossilizeMultiple(t *testing.T, a *Fossilizer, tests []fossilizeTest) {
	go a.Start()
	defer a.Stop()
	rc := make(chan *fossilizer.Result)
	a.AddResultChan(rc)

	for _, test := range tests {
		if err := a.Fossilize(test.data, test.meta); err != nil {
			t.Fatal(err)
		}
		if test.sleep > 0 {
			time.Sleep(test.sleep)
		}
	}

RESULT_LOOP:
	for _ = range tests {
		r := <-rc
		for i := range tests {
			test := &tests[i]
			if string(test.meta) == string(r.Meta) {
				test.fossilized = true
				if !reflect.DeepEqual(r.Data, test.data) {
					a := hex.EncodeToString(r.Data)
					e := hex.EncodeToString(test.data)
					t.Logf("actual: %s; expected %s\n", a, e)
					t.Error("unexpected result data")
				}
				evidence := r.Evidence.(*EvidenceWrapper).Evidence
				if !reflect.DeepEqual(evidence.Path, test.path) {
					ajs, _ := json.MarshalIndent(evidence.Path, "", "  ")
					ejs, _ := json.MarshalIndent(test.path, "", "  ")
					t.Logf("actual: %s; expected %s\n", string(ajs), string(ejs))
					t.Error("unexpected merkle path")
				}
				continue RESULT_LOOP
			}
		}
		t.Errorf("unexpected result meta: %s", r.Meta)
	}

	for _, test := range tests {
		if !test.fossilized {
			t.Errorf("not fossilized: %s\n", test.meta)
		}
	}
}

func benchmarkFossilize(b *testing.B, config *Config) {
	a := New(config)
	go a.Start()
	defer a.Stop()
	rc := make(chan *fossilizer.Result)
	a.AddResultChan(rc)

	data := make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		data[i] = atos(testutil.RandomHash())
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if err := a.Fossilize(data[i], data[i]); err != nil {
			b.Fatal(err)
		}
	}

	for i := 0; i < b.N; i++ {
		<-rc
	}
}

func atos(a types.Bytes32) []byte {
	return a[:]
}
<commit_msg>batchfossilizer: Use hash pointers<commit_after>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by the license
// that can be found in the LICENSE file.

package batchfossilizer

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/stratumn/go/fossilizer"
	"github.com/stratumn/goprivate/merkle"
	"github.com/stratumn/goprivate/testutil"
	"github.com/stratumn/goprivate/types"
)

const interval = 10 * time.Millisecond

var (
	pathA0     merkle.Path
	pathAB0    merkle.Path
	pathAB1    merkle.Path
	pathABC0   merkle.Path
	pathABC1   merkle.Path
	pathABC2   merkle.Path
	pathABCD0  merkle.Path
	pathABCD1  merkle.Path
	pathABCD2  merkle.Path
	pathABCD3  merkle.Path
	pathABCDE0 merkle.Path
	pathABCDE1 merkle.Path
	pathABCDE2 merkle.Path
	pathABCDE3 merkle.Path
	pathABCDE4 merkle.Path
)

func TestMain(m *testing.M) {
	seed := int64(time.Now().Nanosecond())
	fmt.Printf("using seed %d\n", seed)
	rand.Seed(seed)

	loadPath("testdata/path-a-0.json", &pathA0)
	loadPath("testdata/path-ab-0.json", &pathAB0)
	loadPath("testdata/path-ab-1.json", &pathAB1)
	loadPath("testdata/path-abc-0.json", &pathABC0)
	loadPath("testdata/path-abc-1.json", &pathABC1)
	loadPath("testdata/path-abc-2.json", &pathABC2)
	loadPath("testdata/path-abcd-0.json", &pathABCD0)
	loadPath("testdata/path-abcd-1.json", &pathABCD1)
	loadPath("testdata/path-abcd-2.json", &pathABCD2)
	loadPath("testdata/path-abcd-3.json", &pathABCD3)
	loadPath("testdata/path-abcde-0.json", &pathABCDE0)
	loadPath("testdata/path-abcde-1.json", &pathABCDE1)
	loadPath("testdata/path-abcde-2.json", &pathABCDE2)
	loadPath("testdata/path-abcde-3.json", &pathABCDE3)
	loadPath("testdata/path-abcde-4.json", &pathABCDE4)

	flag.Parse()
	os.Exit(m.Run())
}

func TestGetInfo(t *testing.T) {
	a := New(&Config{})
	info, err := a.GetInfo()
	if err != nil {
		t.Fatal(err)
	}
	if info == nil {
		t.Fatal("info is nil")
	}
}

func loadPath(filename string, path *merkle.Path) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}
	if err = json.Unmarshal(data, path); err != nil {
		panic(err)
	}
}

func TestFossilize(t *testing.T) {
	a := New(&Config{Interval: interval})
	tests := []fossilizeTest{
		{atos(sha256.Sum256([]byte("a"))), []byte("test a"), pathABCDE0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b"), pathABCDE1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c"), pathABCDE2, 0, false},
		{atos(sha256.Sum256([]byte("d"))), []byte("test d"), pathABCDE3, 0, false},
		{atos(sha256.Sum256([]byte("e"))), []byte("test e"), pathABCDE4, 0, false},
	}
	testFossilizeMultiple(t, a, tests)
}

func TestFossilizeMaxLeaves(t *testing.T) {
	a := New(&Config{Interval: interval, MaxLeaves: 4})
	tests := []fossilizeTest{
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 1"), pathABCD0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 1"), pathABCD1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 1"), pathABCD2, 0, false},
		{atos(sha256.Sum256([]byte("d"))), []byte("test d 1"), pathABCD3, 0, false},
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 2"), pathABC0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 2"), pathABC1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 2"), pathABC2, 0, false},
	}
	testFossilizeMultiple(t, a, tests)
}

func TestFossilizeInterval(t *testing.T) {
	a := New(&Config{Interval: interval})
	tests := []fossilizeTest{
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 1"), pathABC0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 1"), pathABC1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 1"), pathABC2, interval * 2, false},
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 2"), pathABCD0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 2"), pathABCD1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 2"), pathABCD2, 0, false},
		{atos(sha256.Sum256([]byte("d"))), []byte("test d 2"), pathABCD3, interval * 2, false},
		{atos(sha256.Sum256([]byte("a"))), []byte("test a 3"), pathABC0, 0, false},
		{atos(sha256.Sum256([]byte("b"))), []byte("test b 3"), pathABC1, 0, false},
		{atos(sha256.Sum256([]byte("c"))), []byte("test c 3"), pathABC2, 0, false},
	}
	testFossilizeMultiple(t, a, tests)
}

func BenchmarkFossilizeMaxLeaves100(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 100})
}

func BenchmarkFossilizeMaxLeaves1000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 1000})
}

func BenchmarkFossilizeMaxLeaves10000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 10000})
}

func BenchmarkFossilizeMaxLeaves100000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 100000})
}

func BenchmarkFossilizeMaxLeaves1000000(b *testing.B) {
	benchmarkFossilize(b, &Config{Interval: interval, MaxLeaves: 1000000})
}

type fossilizeTest struct {
	data       []byte
	meta       []byte
	path       merkle.Path
	sleep      time.Duration
	fossilized bool
}

func testFossilizeMultiple(t *testing.T, a *Fossilizer, tests []fossilizeTest) {
	go a.Start()
	defer a.Stop()
	rc := make(chan *fossilizer.Result)
	a.AddResultChan(rc)

	for _, test := range tests {
		if err := a.Fossilize(test.data, test.meta); err != nil {
			t.Fatal(err)
		}
		if test.sleep > 0 {
			time.Sleep(test.sleep)
		}
	}

RESULT_LOOP:
	for _ = range tests {
		r := <-rc
		for i := range tests {
			test := &tests[i]
			if string(test.meta) == string(r.Meta) {
				test.fossilized = true
				if !reflect.DeepEqual(r.Data, test.data) {
					a := hex.EncodeToString(r.Data)
					e := hex.EncodeToString(test.data)
					t.Logf("actual: %s; expected %s\n", a, e)
					t.Error("unexpected result data")
				}
				evidence := r.Evidence.(*EvidenceWrapper).Evidence
				if !reflect.DeepEqual(evidence.Path, test.path) {
					ajs, _ := json.MarshalIndent(evidence.Path, "", "  ")
					ejs, _ := json.MarshalIndent(test.path, "", "  ")
					t.Logf("actual: %s; expected %s\n", string(ajs), string(ejs))
					t.Error("unexpected merkle path")
				}
				continue RESULT_LOOP
			}
		}
		t.Errorf("unexpected result meta: %s", r.Meta)
	}

	for _, test := range tests {
		if !test.fossilized {
			t.Errorf("not fossilized: %s\n", test.meta)
		}
	}
}

func benchmarkFossilize(b *testing.B, config *Config) {
	a := New(config)
	go a.Start()
	defer a.Stop()
	rc := make(chan *fossilizer.Result)
	a.AddResultChan(rc)

	data := make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		data[i] = atos(*testutil.RandomHash())
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if err := a.Fossilize(data[i], data[i]); err != nil {
			b.Fatal(err)
		}
	}

	for i := 0; i < b.N; i++ {
		<-rc
	}
}

func atos(a types.Bytes32) []byte {
	return a[:]
}
<|endoftext|>


<commit_before>package beanspike

import (
	"errors"
	"fmt"
	"strings"
	"sync"

	as "github.com/aerospike/aerospike-client-go"
	ast "github.com/aerospike/aerospike-client-go/types"
)

var rwMutex = sync.RWMutex{}

var tubesMap = map[string]*Tube{}

func (conn *Conn) newJobID() (int64, error) {
	key, err := as.NewKey(AerospikeNamespace, AerospikeMetadataSet, "seq")
	if err != nil {
		return 0, err
	}

	bin := as.NewBin("last", as.NewLongValue(1))
	record, err := conn.aerospike.Operate(as.NewWritePolicy(0, 0), key, as.AddOp(bin), as.GetOp())
	if err != nil {
		return 0, err
	}

	// This type conversion seem required as the value appears to be a
	// int instead of an int64
	id := int64(record.Bins[bin.Name].(int))
	return id, nil
}

// Note, index limits may cause this to fail
func (conn *Conn) Use(name string) (*Tube, error) {
	if conn == nil {
		return nil, errors.New("Aerospike connection not established")
	}

	name = strings.Trim(name, " ")
	if name == "" {
		return nil, errors.New("Tube name must not be blank")
	}

	if name == AerospikeMetadataSet {
		return nil, fmt.Errorf("Tube name %v is reserved", name)
	}

	rwMutex.RLock()
	t := tubesMap[name]
	if t != nil {
		rwMutex.RUnlock()
		return t, nil
	}
	rwMutex.RUnlock()

	rwMutex.Lock()
	defer func() {
		rwMutex.Unlock()
	}()

	t = tubesMap[name]
	if t != nil {
		return t, nil
	}

	task, err := conn.aerospike.CreateIndex(nil, AerospikeNamespace, name, "idx_tube_"+name+"_"+AerospikeNameStatus, AerospikeNameStatus, as.STRING)
	if err != nil {
		if ae, ok := err.(ast.AerospikeError); ok && ae.ResultCode() == ast.INDEX_FOUND {
			// skipping index creation
			// println("Skipping index creation")
		} else {
			return nil, err
		}
	}

	if task == nil {
		//TODO: Check that this is ok
	} else {
		for ierr := range task.OnComplete() {
			if ierr != nil {
				return nil, ierr
			}
		}
	}
	tube := &Tube{Conn: conn, Name: name, first: true}

	tubesMap[name] = tube
	return tube, nil
}

// Any Tubes that reference this name
// should be discarded after this operation
func (conn *Conn) Delete(name string) error {
	if conn == nil {
		return errors.New("Aerospike connection not established")
	}

	client := conn.aerospike

	// No Truncate for sets yet, scan and purge
	recordset, err := client.ScanAll(nil, AerospikeNamespace, name)

	if err != nil {
		return err
	}

	defer recordset.Close()

	for res := range recordset.Results() {
		if res.Err != nil {
			return res.Err
		}
		key := res.Record.Key

		// nil out body before deleting record to address aerospike limitations.
		// Ref: https://discuss.aerospike.com/t/expired-deleted-data-reappears-after-server-is-restarted/470
		policy := as.NewWritePolicy(res.Record.Generation, 0)
		policy.RecordExistsAction = as.UPDATE_ONLY
		policy.SendKey = true
		policy.CommitLevel = as.COMMIT_MASTER
		policy.GenerationPolicy = as.EXPECT_GEN_EQUAL

		binBody := as.NewBin(AerospikeNameBody, nil)
		binCSize := as.NewBin(AerospikeNameCompressedSize, nil)
		binSize := as.NewBin(AerospikeNameSize, 0)
		binStatus := as.NewBin(AerospikeNameStatus, AerospikeSymDeleted)

		err = client.PutBins(policy, key, binBody, binCSize, binSize, binStatus)
		if err != nil {
			return err
		}

		_, err = client.Delete(nil, key)
		if err != nil {
			return err
		}

		conn.stats("tube.delete.count", name, float64(1))
	}

	tk, _ := as.NewKey(AerospikeNamespace, AerospikeMetadataSet, name+":"+AerospikeKeySuffixTtr)
	client.Delete(nil, tk)

	dk, _ := as.NewKey(AerospikeNamespace, AerospikeMetadataSet, name+":"+AerospikeKeySuffixDelayed)
	client.Delete(nil, dk)

	err = client.DropIndex(nil, AerospikeNamespace, name, "idx_tube_"+name+"_"+AerospikeNameStatus)

	if err != nil {
		return err
	}

	return nil
}
<commit_msg>code review<commit_after>package beanspike

import (
	"errors"
	"fmt"
	"strings"
	"sync"

	as "github.com/aerospike/aerospike-client-go"
	ast "github.com/aerospike/aerospike-client-go/types"
)

var tubesMap = struct {
	sync.RWMutex
	m map[string]*Tube
}{m: make(map[string]*Tube)}

func (conn *Conn) newJobID() (int64, error) {
	key, err := as.NewKey(AerospikeNamespace, AerospikeMetadataSet, "seq")
	if err != nil {
		return 0, err
	}

	bin := as.NewBin("last", as.NewLongValue(1))
	record, err := conn.aerospike.Operate(as.NewWritePolicy(0, 0), key, as.AddOp(bin), as.GetOp())
	if err != nil {
		return 0, err
	}

	// This type conversion seem required as the value appears to be a
	// int instead of an int64
	id := int64(record.Bins[bin.Name].(int))
	return id, nil
}

// Note, index limits may cause this to fail
func (conn *Conn) Use(name string) (*Tube, error) {
	if conn == nil {
		return nil, errors.New("Aerospike connection not established")
	}

	name = strings.Trim(name, " ")
	if name == "" {
		return nil, errors.New("Tube name must not be blank")
	}

	if name == AerospikeMetadataSet {
		return nil, fmt.Errorf("Tube name %v is reserved", name)
	}

	tubesMap.RLock()
	t := tubesMap.m[name]
	if t != nil {
		tubesMap.RUnlock()
		return t, nil
	}
	tubesMap.RUnlock()

	tubesMap.Lock()
	defer tubesMap.Unlock()

	t = tubesMap.m[name]
	if t != nil {
		return t, nil
	}

	task, err := conn.aerospike.CreateIndex(nil, AerospikeNamespace, name, "idx_tube_"+name+"_"+AerospikeNameStatus, AerospikeNameStatus, as.STRING)
	if err != nil {
		if ae, ok := err.(ast.AerospikeError); ok && ae.ResultCode() == ast.INDEX_FOUND {
			// skipping index creation
			// println("Skipping index creation")
		} else {
			return nil, err
		}
	}

	if task == nil {
		//TODO: Check that this is ok
	} else {
		for ierr := range task.OnComplete() {
			if ierr != nil {
				return nil, ierr
			}
		}
	}
	tube := &Tube{Conn: conn, Name: name, first: true}

	tubesMap.m[name] = tube
	return tube, nil
}

// Any Tubes that reference this name
// should be discarded after this operation
func (conn *Conn) Delete(name string) error {
	if conn == nil {
		return errors.New("Aerospike connection not established")
	}

	client := conn.aerospike

	// No Truncate for sets yet, scan and purge
	recordset, err := client.ScanAll(nil, AerospikeNamespace, name)

	if err != nil {
		return err
	}

	defer recordset.Close()

	for res := range recordset.Results() {
		if res.Err != nil {
			return res.Err
		}
		key := res.Record.Key

		// nil out body before deleting record to address aerospike limitations.
		// Ref: https://discuss.aerospike.com/t/expired-deleted-data-reappears-after-server-is-restarted/470
		policy := as.NewWritePolicy(res.Record.Generation, 0)
		policy.RecordExistsAction = as.UPDATE_ONLY
		policy.SendKey = true
		policy.CommitLevel = as.COMMIT_MASTER
		policy.GenerationPolicy = as.EXPECT_GEN_EQUAL

		binBody := as.NewBin(AerospikeNameBody, nil)
		binCSize := as.NewBin(AerospikeNameCompressedSize, nil)
		binSize := as.NewBin(AerospikeNameSize, 0)
		binStatus := as.NewBin(AerospikeNameStatus, AerospikeSymDeleted)

		err = client.PutBins(policy, key, binBody, binCSize, binSize, binStatus)
		if err != nil {
			return err
		}

		_, err = client.Delete(nil, key)
		if err != nil {
			return err
		}

		conn.stats("tube.delete.count", name, float64(1))
	}

	tk, _ := as.NewKey(AerospikeNamespace, AerospikeMetadataSet, name+":"+AerospikeKeySuffixTtr)
	client.Delete(nil, tk)

	dk, _ := as.NewKey(AerospikeNamespace, AerospikeMetadataSet, name+":"+AerospikeKeySuffixDelayed)
	client.Delete(nil, dk)

	err = client.DropIndex(nil, AerospikeNamespace, name, "idx_tube_"+name+"_"+AerospikeNameStatus)

	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package pgconn_test

import (
	"context"
	"os"
	"testing"

	"github.com/jackc/pgx/pgconn"
	"github.com/stretchr/testify/require"
)

func BenchmarkConnect(b *testing.B) {
	benchmarks := []struct {
		name string
		env  string
	}{
		{"Unix socket", "PGX_TEST_UNIX_SOCKET_CONN_STRING"},
		{"TCP", "PGX_TEST_TCP_CONN_STRING"},
	}

	for _, bm := range benchmarks {
		b.Run(bm.name, func(b *testing.B) {
			connString := os.Getenv(bm.env)
			if connString == "" {
				b.Skipf("Skipping due to missing environment variable %v", bm.env)
			}

			for i := 0; i < b.N; i++ {
				conn, err := pgconn.Connect(context.Background(), connString)
				require.Nil(b, err)

				err = conn.Close(context.Background())
				require.Nil(b, err)
			}
		})
	}
}

func BenchmarkExecPrepared(b *testing.B) {
	conn, err := pgconn.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
	require.Nil(b, err)
	defer closeConn(b, conn)

	err = conn.Prepare(context.Background(), "ps1", "select 'hello'::text as a, 42::int4 as b, '2019-01-01'::date", nil)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		_, err := conn.ExecPrepared(context.Background(), "ps1", nil, nil, nil)
		require.Nil(b, err)
	}
}
<commit_msg>Add non-buffered benchmark<commit_after>package pgconn_test

import (
	"context"
	"os"
	"testing"

	"github.com/jackc/pgx/pgconn"
	"github.com/stretchr/testify/require"
)

func BenchmarkConnect(b *testing.B) {
	benchmarks := []struct {
		name string
		env  string
	}{
		{"Unix socket", "PGX_TEST_UNIX_SOCKET_CONN_STRING"},
		{"TCP", "PGX_TEST_TCP_CONN_STRING"},
	}

	for _, bm := range benchmarks {
		b.Run(bm.name, func(b *testing.B) {
			connString := os.Getenv(bm.env)
			if connString == "" {
				b.Skipf("Skipping due to missing environment variable %v", bm.env)
			}

			for i := 0; i < b.N; i++ {
				conn, err := pgconn.Connect(context.Background(), connString)
				require.Nil(b, err)

				err = conn.Close(context.Background())
				require.Nil(b, err)
			}
		})
	}
}

func BenchmarkExecPrepared(b *testing.B) {
	conn, err := pgconn.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
	require.Nil(b, err)
	defer closeConn(b, conn)

	err = conn.Prepare(context.Background(), "ps1", "select 'hello'::text as a, 42::int4 as b, '2019-01-01'::date", nil)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		_, err := conn.ExecPrepared(context.Background(), "ps1", nil, nil, nil)
		require.Nil(b, err)
	}
}

func BenchmarkSendExecPrepared(b *testing.B) {
	conn, err := pgconn.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
	require.Nil(b, err)
	defer closeConn(b, conn)

	err = conn.Prepare(context.Background(), "ps1", "select 'hello'::text as a, 42::int4 as b, '2019-01-01'::date", nil)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		conn.SendExecPrepared("ps1", nil, nil, nil)
		err := conn.Flush(context.Background())
		require.Nil(b, err)

		for conn.NextResult(context.Background()) {
			_, err := conn.ResultReader().Close()
			require.Nil(b, err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package snap

import (
	"errors"
	"fmt"
	"hash/crc32"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

	pioutil "go.etcd.io/etcd/pkg/v3/ioutil"
	"go.etcd.io/etcd/pkg/v3/pbutil"
	"go.etcd.io/etcd/raft/v3"
	"go.etcd.io/etcd/raft/v3/raftpb"
	"go.etcd.io/etcd/server/v3/etcdserver/api/snap/snappb"
	"go.etcd.io/etcd/server/v3/storage/wal/walpb"

	"go.uber.org/zap"
)

const snapSuffix = ".snap"

var (
	ErrNoSnapshot    = errors.New("snap: no available snapshot")
	ErrEmptySnapshot = errors.New("snap: empty snapshot")
	ErrCRCMismatch   = errors.New("snap: crc mismatch")
	crcTable         = crc32.MakeTable(crc32.Castagnoli)

	// A map of valid files that can be present in the snap folder.
	validFiles = map[string]bool{
		"db": true,
	}
)

type Snapshotter struct {
	lg  *zap.Logger
	dir string
}

func New(lg *zap.Logger, dir string) *Snapshotter {
	if lg == nil {
		lg = zap.NewNop()
	}
	return &Snapshotter{
		lg:  lg,
		dir: dir,
	}
}

func (s *Snapshotter) SaveSnap(snapshot raftpb.Snapshot) error {
	if raft.IsEmptySnap(snapshot) {
		return nil
	}
	return s.save(&snapshot)
}

func (s *Snapshotter) save(snapshot *raftpb.Snapshot) error {
	start := time.Now()

	fname := fmt.Sprintf("%016x-%016x%s", snapshot.Metadata.Term, snapshot.Metadata.Index, snapSuffix)
	b := pbutil.MustMarshal(snapshot)
	crc := crc32.Update(0, crcTable, b)
	snap := snappb.Snapshot{Crc: crc, Data: b}
	d, err := snap.Marshal()
	if err != nil {
		return err
	}
	snapMarshallingSec.Observe(time.Since(start).Seconds())

	spath := filepath.Join(s.dir, fname)

	fsyncStart := time.Now()
	err = pioutil.WriteAndSyncFile(spath, d, 0666)
	snapFsyncSec.Observe(time.Since(fsyncStart).Seconds())

	if err != nil {
		s.lg.Warn("failed to write a snap file", zap.String("path", spath), zap.Error(err))
		rerr := os.Remove(spath)
		if rerr != nil {
			s.lg.Warn("failed to remove a broken snap file", zap.String("path", spath), zap.Error(err))
		}
		return err
	}

	snapSaveSec.Observe(time.Since(start).Seconds())
	return nil
}

// Load returns the newest snapshot.
func (s *Snapshotter) Load() (*raftpb.Snapshot, error) {
	return s.loadMatching(func(*raftpb.Snapshot) bool { return true })
}

// LoadNewestAvailable loads the newest snapshot available that is in walSnaps.
func (s *Snapshotter) LoadNewestAvailable(walSnaps []walpb.Snapshot) (*raftpb.Snapshot, error) {
	return s.loadMatching(func(snapshot *raftpb.Snapshot) bool {
		m := snapshot.Metadata
		for i := len(walSnaps) - 1; i >= 0; i-- {
			if m.Term == walSnaps[i].Term && m.Index == walSnaps[i].Index {
				return true
			}
		}
		return false
	})
}

// loadMatching returns the newest snapshot where matchFn returns true.
func (s *Snapshotter) loadMatching(matchFn func(*raftpb.Snapshot) bool) (*raftpb.Snapshot, error) {
	names, err := s.snapNames()
	if err != nil {
		return nil, err
	}
	var snap *raftpb.Snapshot
	for _, name := range names {
		if snap, err = loadSnap(s.lg, s.dir, name); err == nil && matchFn(snap) {
			return snap, nil
		}
	}
	return nil, ErrNoSnapshot
}

func loadSnap(lg *zap.Logger, dir, name string) (*raftpb.Snapshot, error) {
	fpath := filepath.Join(dir, name)
	snap, err := Read(lg, fpath)
	if err != nil {
		brokenPath := fpath + ".broken"
		if lg != nil {
			lg.Warn("failed to read a snap file", zap.String("path", fpath), zap.Error(err))
		}
		if rerr := os.Rename(fpath, brokenPath); rerr != nil {
			if lg != nil {
				lg.Warn("failed to rename a broken snap file", zap.String("path", fpath), zap.String("broken-path", brokenPath), zap.Error(rerr))
			}
		} else {
			if lg != nil {
				lg.Warn("renamed to a broken snap file", zap.String("path", fpath), zap.String("broken-path", brokenPath))
			}
		}
	}
	return snap, err
}

// Read reads the snapshot named by snapname and returns the snapshot.
func Read(lg *zap.Logger, snapname string) (*raftpb.Snapshot, error) {
	b, err := os.ReadFile(snapname)
	if err != nil {
		if lg != nil {
			lg.Warn("failed to read a snap file", zap.String("path", snapname), zap.Error(err))
		}
		return nil, err
	}

	if len(b) == 0 {
		if lg != nil {
			lg.Warn("failed to read empty snapshot file", zap.String("path", snapname))
		}
		return nil, ErrEmptySnapshot
	}

	var serializedSnap snappb.Snapshot
	if err = serializedSnap.Unmarshal(b); err != nil {
		if lg != nil {
			lg.Warn("failed to unmarshal snappb.Snapshot", zap.String("path", snapname), zap.Error(err))
		}
		return nil, err
	}

	if len(serializedSnap.Data) == 0 || serializedSnap.Crc == 0 {
		if lg != nil {
			lg.Warn("failed to read empty snapshot data", zap.String("path", snapname))
		}
		return nil, ErrEmptySnapshot
	}

	crc := crc32.Update(0, crcTable, serializedSnap.Data)
	if crc != serializedSnap.Crc {
		if lg != nil {
			lg.Warn("snap file is corrupt",
				zap.String("path", snapname),
				zap.Uint32("prev-crc", serializedSnap.Crc),
				zap.Uint32("new-crc", crc),
			)
		}
		return nil, ErrCRCMismatch
	}

	var snap raftpb.Snapshot
	if err = snap.Unmarshal(serializedSnap.Data); err != nil {
		if lg != nil {
			lg.Warn("failed to unmarshal raftpb.Snapshot", zap.String("path", snapname), zap.Error(err))
		}
		return nil, err
	}
	return &snap, nil
}

// snapNames returns the filename of the snapshots in logical time order (from newest to oldest).
// If there is no available snapshots, an ErrNoSnapshot will be returned.
func (s *Snapshotter) snapNames() ([]string, error) {
	dir, err := os.Open(s.dir)
	if err != nil {
		return nil, err
	}
	defer dir.Close()
	names, err := dir.Readdirnames(-1)
	if err != nil {
		return nil, err
	}
	filenames, err := s.cleanupSnapdir(names)
	if err != nil {
		return nil, err
	}
	snaps := checkSuffix(s.lg, filenames)
	if len(snaps) == 0 {
		return nil, ErrNoSnapshot
	}
	sort.Sort(sort.Reverse(sort.StringSlice(snaps)))
	return snaps, nil
}

func checkSuffix(lg *zap.Logger, names []string) []string {
	snaps := []string{}
	for i := range names {
		if strings.HasSuffix(names[i], snapSuffix) {
			snaps = append(snaps, names[i])
		} else {
			// If we find a file which is not a snapshot then check if it's
			// a vaild file. If not throw out a warning.
			if _, ok := validFiles[names[i]]; !ok {
				if lg != nil {
					lg.Warn("found unexpected non-snap file; skipping", zap.String("path", names[i]))
				}
			}
		}
	}
	return snaps
}

// cleanupSnapdir removes any files that should not be in the snapshot directory:
// - db.tmp prefixed files that can be orphaned by defragmentation
func (s *Snapshotter) cleanupSnapdir(filenames []string) (names []string, err error) {
	names = make([]string, 0, len(filenames))
	for _, filename := range filenames {
		if strings.HasPrefix(filename, "db.tmp") {
			s.lg.Info("found orphaned defragmentation file; deleting", zap.String("path", filename))
			if rmErr := os.Remove(filepath.Join(s.dir, filename)); rmErr != nil && !os.IsNotExist(rmErr) {
				return names, fmt.Errorf("failed to remove orphaned .snap.db file %s: %v", filename, rmErr)
			}
		} else {
			names = append(names, filename)
		}
	}
	return names, nil
}

func (s *Snapshotter) ReleaseSnapDBs(snap raftpb.Snapshot) error {
	dir, err := os.Open(s.dir)
	if err != nil {
		return err
	}
	defer dir.Close()
	filenames, err := dir.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, filename := range filenames {
		if strings.HasSuffix(filename, ".snap.db") {
			hexIndex := strings.TrimSuffix(filepath.Base(filename), ".snap.db")
			index, err := strconv.ParseUint(hexIndex, 16, 64)
			if err != nil {
				s.lg.Error("failed to parse index from filename", zap.String("path", filename), zap.String("error", err.Error()))
				continue
			}
			if index < snap.Metadata.Index {
				s.lg.Info("found orphaned .snap.db file; deleting", zap.String("path", filename))
				if rmErr := os.Remove(filepath.Join(s.dir, filename)); rmErr != nil && !os.IsNotExist(rmErr) {
					s.lg.Error("failed to remove orphaned .snap.db file", zap.String("path", filename), zap.String("error", rmErr.Error()))
				}
			}
		}
	}
	return nil
}
<commit_msg>fix a typo: print the correct error info<commit_after>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package snap

import (
	"errors"
	"fmt"
	"hash/crc32"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

	pioutil "go.etcd.io/etcd/pkg/v3/ioutil"
	"go.etcd.io/etcd/pkg/v3/pbutil"
	"go.etcd.io/etcd/raft/v3"
	"go.etcd.io/etcd/raft/v3/raftpb"
	"go.etcd.io/etcd/server/v3/etcdserver/api/snap/snappb"
	"go.etcd.io/etcd/server/v3/storage/wal/walpb"

	"go.uber.org/zap"
)

const snapSuffix = ".snap"

var (
	ErrNoSnapshot    = errors.New("snap: no available snapshot")
	ErrEmptySnapshot = errors.New("snap: empty snapshot")
	ErrCRCMismatch   = errors.New("snap: crc mismatch")
	crcTable         = crc32.MakeTable(crc32.Castagnoli)

	// A map of valid files that can be present in the snap folder.
	validFiles = map[string]bool{
		"db": true,
	}
)

type Snapshotter struct {
	lg  *zap.Logger
	dir string
}

func New(lg *zap.Logger, dir string) *Snapshotter {
	if lg == nil {
		lg = zap.NewNop()
	}
	return &Snapshotter{
		lg:  lg,
		dir: dir,
	}
}

func (s *Snapshotter) SaveSnap(snapshot raftpb.Snapshot) error {
	if raft.IsEmptySnap(snapshot) {
		return nil
	}
	return s.save(&snapshot)
}

func (s *Snapshotter) save(snapshot *raftpb.Snapshot) error {
	start := time.Now()

	fname := fmt.Sprintf("%016x-%016x%s", snapshot.Metadata.Term, snapshot.Metadata.Index, snapSuffix)
	b := pbutil.MustMarshal(snapshot)
	crc := crc32.Update(0, crcTable, b)
	snap := snappb.Snapshot{Crc: crc, Data: b}
	d, err := snap.Marshal()
	if err != nil {
		return err
	}
	snapMarshallingSec.Observe(time.Since(start).Seconds())

	spath := filepath.Join(s.dir, fname)

	fsyncStart := time.Now()
	err = pioutil.WriteAndSyncFile(spath, d, 0666)
	snapFsyncSec.Observe(time.Since(fsyncStart).Seconds())

	if err != nil {
		s.lg.Warn("failed to write a snap file", zap.String("path", spath), zap.Error(err))
		rerr := os.Remove(spath)
		if rerr != nil {
			s.lg.Warn("failed to remove a broken snap file", zap.String("path", spath), zap.Error(rerr))
		}
		return err
	}

	snapSaveSec.Observe(time.Since(start).Seconds())
	return nil
}

// Load returns the newest snapshot.
func (s *Snapshotter) Load() (*raftpb.Snapshot, error) {
	return s.loadMatching(func(*raftpb.Snapshot) bool { return true })
}

// LoadNewestAvailable loads the newest snapshot available that is in walSnaps.
func (s *Snapshotter) LoadNewestAvailable(walSnaps []walpb.Snapshot) (*raftpb.Snapshot, error) {
	return s.loadMatching(func(snapshot *raftpb.Snapshot) bool {
		m := snapshot.Metadata
		for i := len(walSnaps) - 1; i >= 0; i-- {
			if m.Term == walSnaps[i].Term && m.Index == walSnaps[i].Index {
				return true
			}
		}
		return false
	})
}

// loadMatching returns the newest snapshot where matchFn returns true.
func (s *Snapshotter) loadMatching(matchFn func(*raftpb.Snapshot) bool) (*raftpb.Snapshot, error) {
	names, err := s.snapNames()
	if err != nil {
		return nil, err
	}
	var snap *raftpb.Snapshot
	for _, name := range names {
		if snap, err = loadSnap(s.lg, s.dir, name); err == nil && matchFn(snap) {
			return snap, nil
		}
	}
	return nil, ErrNoSnapshot
}

func loadSnap(lg *zap.Logger, dir, name string) (*raftpb.Snapshot, error) {
	fpath := filepath.Join(dir, name)
	snap, err := Read(lg, fpath)
	if err != nil {
		brokenPath := fpath + ".broken"
		if lg != nil {
			lg.Warn("failed to read a snap file", zap.String("path", fpath), zap.Error(err))
		}
		if rerr := os.Rename(fpath, brokenPath); rerr != nil {
			if lg != nil {
				lg.Warn("failed to rename a broken snap file", zap.String("path", fpath), zap.String("broken-path", brokenPath), zap.Error(rerr))
			}
		} else {
			if lg != nil {
				lg.Warn("renamed to a broken snap file", zap.String("path", fpath), zap.String("broken-path", brokenPath))
			}
		}
	}
	return snap, err
}

// Read reads the snapshot named by snapname and returns the snapshot.
func Read(lg *zap.Logger, snapname string) (*raftpb.Snapshot, error) {
	b, err := os.ReadFile(snapname)
	if err != nil {
		if lg != nil {
			lg.Warn("failed to read a snap file", zap.String("path", snapname), zap.Error(err))
		}
		return nil, err
	}

	if len(b) == 0 {
		if lg != nil {
			lg.Warn("failed to read empty snapshot file", zap.String("path", snapname))
		}
		return nil, ErrEmptySnapshot
	}

	var serializedSnap snappb.Snapshot
	if err = serializedSnap.Unmarshal(b); err != nil {
		if lg != nil {
			lg.Warn("failed to unmarshal snappb.Snapshot", zap.String("path", snapname), zap.Error(err))
		}
		return nil, err
	}

	if len(serializedSnap.Data) == 0 || serializedSnap.Crc == 0 {
		if lg != nil {
			lg.Warn("failed to read empty snapshot data", zap.String("path", snapname))
		}
		return nil, ErrEmptySnapshot
	}

	crc := crc32.Update(0, crcTable, serializedSnap.Data)
	if crc != serializedSnap.Crc {
		if lg != nil {
			lg.Warn("snap file is corrupt",
				zap.String("path", snapname),
				zap.Uint32("prev-crc", serializedSnap.Crc),
				zap.Uint32("new-crc", crc),
			)
		}
		return nil, ErrCRCMismatch
	}

	var snap raftpb.Snapshot
	if err = snap.Unmarshal(serializedSnap.Data); err != nil {
		if lg != nil {
			lg.Warn("failed to unmarshal raftpb.Snapshot", zap.String("path", snapname), zap.Error(err))
		}
		return nil, err
	}
	return &snap, nil
}

// snapNames returns the filename of the snapshots in logical time order (from newest to oldest).
// If there is no available snapshots, an ErrNoSnapshot will be returned.
func (s *Snapshotter) snapNames() ([]string, error) {
	dir, err := os.Open(s.dir)
	if err != nil {
		return nil, err
	}
	defer dir.Close()
	names, err := dir.Readdirnames(-1)
	if err != nil {
		return nil, err
	}
	filenames, err := s.cleanupSnapdir(names)
	if err != nil {
		return nil, err
	}
	snaps := checkSuffix(s.lg, filenames)
	if len(snaps) == 0 {
		return nil, ErrNoSnapshot
	}
	sort.Sort(sort.Reverse(sort.StringSlice(snaps)))
	return snaps, nil
}

func checkSuffix(lg *zap.Logger, names []string) []string {
	snaps := []string{}
	for i := range names {
		if strings.HasSuffix(names[i], snapSuffix) {
			snaps = append(snaps, names[i])
		} else {
			// If we find a file which is not a snapshot then check if it's
			// a vaild file. If not throw out a warning.
			if _, ok := validFiles[names[i]]; !ok {
				if lg != nil {
					lg.Warn("found unexpected non-snap file; skipping", zap.String("path", names[i]))
				}
			}
		}
	}
	return snaps
}

// cleanupSnapdir removes any files that should not be in the snapshot directory:
// - db.tmp prefixed files that can be orphaned by defragmentation
func (s *Snapshotter) cleanupSnapdir(filenames []string) (names []string, err error) {
	names = make([]string, 0, len(filenames))
	for _, filename := range filenames {
		if strings.HasPrefix(filename, "db.tmp") {
			s.lg.Info("found orphaned defragmentation file; deleting", zap.String("path", filename))
			if rmErr := os.Remove(filepath.Join(s.dir, filename)); rmErr != nil && !os.IsNotExist(rmErr) {
				return names, fmt.Errorf("failed to remove orphaned .snap.db file %s: %v", filename, rmErr)
			}
		} else {
			names = append(names, filename)
		}
	}
	return names, nil
}

func (s *Snapshotter) ReleaseSnapDBs(snap raftpb.Snapshot) error {
	dir, err := os.Open(s.dir)
	if err != nil {
		return err
	}
	defer dir.Close()
	filenames, err := dir.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, filename := range filenames {
		if strings.HasSuffix(filename, ".snap.db") {
			hexIndex := strings.TrimSuffix(filepath.Base(filename), ".snap.db")
			index, err := strconv.ParseUint(hexIndex, 16, 64)
			if err != nil {
				s.lg.Error("failed to parse index from filename", zap.String("path", filename), zap.String("error", err.Error()))
				continue
			}
			if index < snap.Metadata.Index {
				s.lg.Info("found orphaned .snap.db file; deleting", zap.String("path", filename))
				if rmErr := os.Remove(filepath.Join(s.dir, filename)); rmErr != nil && !os.IsNotExist(rmErr) {
					s.lg.Error("failed to remove orphaned .snap.db file", zap.String("path", filename), zap.String("error", rmErr.Error()))
				}
			}
		}
	}
	return nil
}
<|endoftext|>