<commit_before>package peco

import (
	"encoding/json"
	"testing"

	"github.com/nsf/termbox-go"
)

func TestReadRC(t *testing.T) {
	txt := `
{
  "Keymap": {
   "C-j": "peco.Finish"
  }
}
`
	cfg := NewConfig()
	if err := json.Unmarshal([]byte(txt), cfg); err != nil {
		t.Fatalf("Error unmarshaling json: %s", err)
	}
	t.Logf("%#q", cfg)
}

func TestKeymapStrToKeyValue(t *testing.T) {
	expected := map[string]termbox.Key{
		"Insert": termbox.KeyInsert,
		"MouseLeft": termbox.MouseLeft,
		"C-k": termbox.KeyCtrlK,
		"C-h": termbox.KeyCtrlH,
		"C-i": termbox.KeyCtrlI,
		"C-l": termbox.KeyCtrlL,
		"C-m": termbox.KeyCtrlM,
		"C-[": termbox.KeyCtrlLsqBracket,
		"C-\\": termbox.KeyCtrlBackslash,
		"C-_": termbox.KeyCtrlUnderscore,
		"C-8": termbox.KeyCtrl8,
	}

t.Logf("%#q", stringToKey)

	t.Logf("Checking key name -> actual key value mapping...")
	for n, v := range expected {
		t.Logf("    checking %s...", n)
		e, ok := stringToKey[n]
		if !ok {
			t.Errorf("Key name %s not found", n)
		}
		if e != v {
			t.Errorf("Expected '%s' to be '%d', but got '%d'", n, v, stringToKey[n])
		}
	}
}
<commit_msg>remove unused code<commit_after>package peco

import (
	"encoding/json"
	"testing"

	"github.com/nsf/termbox-go"
)

func TestReadRC(t *testing.T) {
	txt := `
{
  "Keymap": {
   "C-j": "peco.Finish"
  }
}
`
	cfg := NewConfig()
	if err := json.Unmarshal([]byte(txt), cfg); err != nil {
		t.Fatalf("Error unmarshaling json: %s", err)
	}
	t.Logf("%#q", cfg)
}

func TestKeymapStrToKeyValue(t *testing.T) {
	expected := map[string]termbox.Key{
		"Insert": termbox.KeyInsert,
		"MouseLeft": termbox.MouseLeft,
		"C-k": termbox.KeyCtrlK,
		"C-h": termbox.KeyCtrlH,
		"C-i": termbox.KeyCtrlI,
		"C-l": termbox.KeyCtrlL,
		"C-m": termbox.KeyCtrlM,
		"C-[": termbox.KeyCtrlLsqBracket,
		"C-\\": termbox.KeyCtrlBackslash,
		"C-_": termbox.KeyCtrlUnderscore,
		"C-8": termbox.KeyCtrl8,
	}

	t.Logf("Checking key name -> actual key value mapping...")
	for n, v := range expected {
		t.Logf("    checking %s...", n)
		e, ok := stringToKey[n]
		if !ok {
			t.Errorf("Key name %s not found", n)
		}
		if e != v {
			t.Errorf("Expected '%s' to be '%d', but got '%d'", n, v, stringToKey[n])
		}
	}
}
<|endoftext|>


<commit_before>package poslog

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

// ImportXML takes a POSLog XML file as the argument and returns
// a POSLog
func ImportXML(filename string) (p POSLog) {
	byteXML, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println("FUCK", err)
	}

	xml.Unmarshal(byteXML, &p)
	p.appendFilename(filepath.Base(filename))
	p.appendDayID()

	return
}

func importReaderXML(f io.ReadCloser, filename string) (p POSLog) {
	byteXML, err := ioutil.ReadAll(f)
	if err != nil {
		fmt.Println("FUCK", err)
	}
	xml.Unmarshal(byteXML, &p)
	p.appendFilename(filepath.Base(filename))
	p.appendDayID()
	return
}

// WriteXML writes a POSLog object to an XML file given as first argument
func WriteXML(filename string, p POSLog) {
	// Drop any other extension and stick a xml on there
	filename = strings.TrimSuffix(filename, filepath.Ext(filename))
	filename = filename + ".xml"
	XMLString, err := xml.MarshalIndent(p, "", "  ")

	if err != nil {
		fmt.Println(err)
	}
	ioutil.WriteFile(filename, XMLString, 0666)
}

// WriteJSON writes a POSLog object to an json file given as first argument
func WriteJSON(filename string, p POSLog) {
	// Drop any other extension and stick a json on there
	filename = strings.TrimSuffix(filename, filepath.Ext(filename))
	filename = filename + ".json"
	XMLString, err := json.MarshalIndent(p, "", "  ")

	if err != nil {
		fmt.Println(err)
	}

	ioutil.WriteFile(filename, XMLString, 0666)
}

func WriteJSONs(folder string, ps POSLogs) {
	os.Mkdir(folder, 0777)
	fmt.Println("This is a test")
	for _, p := range ps.POSLogs {
		op := filepath.Join(folder, p.Filename)
		WriteJSON(op, p)
	}
}
<commit_msg>remove test print<commit_after>package poslog

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

// ImportXML takes a POSLog XML file as the argument and returns
// a POSLog
func ImportXML(filename string) (p POSLog) {
	byteXML, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println("FUCK", err)
	}

	xml.Unmarshal(byteXML, &p)
	p.appendFilename(filepath.Base(filename))
	p.appendDayID()

	return
}

func importReaderXML(f io.ReadCloser, filename string) (p POSLog) {
	byteXML, err := ioutil.ReadAll(f)
	if err != nil {
		fmt.Println("FUCK", err)
	}
	xml.Unmarshal(byteXML, &p)
	p.appendFilename(filepath.Base(filename))
	p.appendDayID()
	return
}

// WriteXML writes a POSLog object to an XML file given as first argument
func WriteXML(filename string, p POSLog) {
	// Drop any other extension and stick a xml on there
	filename = strings.TrimSuffix(filename, filepath.Ext(filename))
	filename = filename + ".xml"
	XMLString, err := xml.MarshalIndent(p, "", "  ")

	if err != nil {
		fmt.Println(err)
	}
	ioutil.WriteFile(filename, XMLString, 0666)
}

// WriteJSON writes a POSLog object to an json file given as first argument
func WriteJSON(filename string, p POSLog) {
	// Drop any other extension and stick a json on there
	filename = strings.TrimSuffix(filename, filepath.Ext(filename))
	filename = filename + ".json"
	XMLString, err := json.MarshalIndent(p, "", "  ")

	if err != nil {
		fmt.Println(err)
	}

	ioutil.WriteFile(filename, XMLString, 0666)
}

func WriteJSONs(folder string, ps POSLogs) {
	os.Mkdir(folder, 0777)
	for _, p := range ps.POSLogs {
		op := filepath.Join(folder, p.Filename)
		WriteJSON(op, p)
	}
}
<|endoftext|>


<commit_before>package pki

import (
  "crypto/rand"
  "crypto/x509"
  "crypto/x509/pkix"
  "encoding/pem"
  "fmt"
  "math/big"
  "net"
  "time"
)

const PemLabelCertificateRequest = "CERTIFICATE REQUEST"

type (
  CertificateData struct {
    Subject  pkix.Name

    DNSNames       []string
    EmailAddresses []string
    IPAddresses    []net.IP
  }

  Certificate x509.Certificate
  CertificateRequest x509.CertificateRequest

  CertificateOptions struct {
    SerialNumber        *big.Int
    NotBefore           time.Time
    NotAfter            time.Time // Validity bounds.
    KeyUsage            x509.KeyUsage
  }
)

func NewCertificateData() *CertificateData {
  return &CertificateData{Subject: pkix.Name{}}
}

// Create a certificate sign request from the input data and the private key of
// the request creator.
func (c *CertificateData) ToCertificateRequest(private_key PrivateKey) (*CertificateRequest, error) {
  csr := &x509.CertificateRequest{}

  csr.Subject        = c.Subject
  csr.DNSNames       = c.DNSNames
  csr.IPAddresses    = c.IPAddresses
  csr.EmailAddresses = c.EmailAddresses

  csr_asn1, err := x509.CreateCertificateRequest(rand.Reader, csr, private_key.PrivateKey())
  if err != nil { return nil, err }
  return LoadCertificateSignRequest(csr_asn1)
}

// Load a certificate sign request from its asn1 representation.
func LoadCertificateSignRequest(raw []byte) (*CertificateRequest, error) {
  csr, err := x509.ParseCertificateRequest(raw)
  if err != nil { return nil, err }
  return (*CertificateRequest)(csr), nil
}

// Return the certificate sign request as a pem block.
func (c *CertificateRequest) MarshalPem() (marshalledPemBlock, error) {
  block := &pem.Block{Type: PemLabelCertificateRequest, Bytes: c.Raw}
  return pem.EncodeToMemory(block), nil
}

// Convert the certificate sign request to a certificate using the private key
// of the signer and the certificate of the signer.
// If the certificate is null, the sign request will be used to sign itself.
// Please also see the certificate options struct for information on mandatory fields.
// For more information, please read http://golang.org/pkg/crypto/x509/#CreateCertificate
func (c *CertificateRequest) ToCertificate(private_key PrivateKey,
        cert_opts CertificateOptions, ca *Certificate) (*Certificate, error) {

  if err := cert_opts.Valid(); err != nil { return nil, err }

  template := &x509.Certificate{}
  template.Subject        = c.Subject
  template.DNSNames       = c.DNSNames
  template.IPAddresses    = c.IPAddresses
  template.EmailAddresses = c.EmailAddresses

  // if no ca is given, we have to set IsCA to self sign
  if ca == nil {
    template.IsCA = true
  }

  template.NotBefore    = cert_opts.NotBefore
  template.NotAfter     = cert_opts.NotAfter
  template.KeyUsage     = cert_opts.KeyUsage
  template.SerialNumber = cert_opts.SerialNumber

  var cert_asn1 []byte
  var err  error
  // if we have no ca which can sign the cert, a self signed cert is wanted
  // (or isn't it? Maybe we should split creation of the template? But that would be ugly)
  if ca == nil {
    cert_asn1, err = x509.CreateCertificate(rand.Reader, template, template, c.PublicKey, private_key.PrivateKey())
  } else {
    cert_asn1, err = x509.CreateCertificate(rand.Reader, template, (*x509.Certificate)(ca), c.PublicKey, private_key.PrivateKey())
  }
  if err != nil { return nil, err }
  return LoadCertificate(cert_asn1)
}

// Load a certificate from its asn1 representation.
func LoadCertificate(raw []byte) (*Certificate, error) {
  cert, err := x509.ParseCertificate(raw)
  if err != nil { return nil, err }
  return (*Certificate)(cert), nil
}

func (co *CertificateOptions) Valid() error {
  if co.SerialNumber == nil { return fmt.Errorf("No serial number set!") }
  return nil
}
<commit_msg>implement marshal pem interface for certificates<commit_after>package pki

import (
  "crypto/rand"
  "crypto/x509"
  "crypto/x509/pkix"
  "encoding/pem"
  "fmt"
  "math/big"
  "net"
  "time"
)

const (
  PemLabelCertificateRequest = "CERTIFICATE REQUEST"
  PemLabelCertificate        = "CERTIFICATE"
)

type (
  CertificateData struct {
    Subject  pkix.Name

    DNSNames       []string
    EmailAddresses []string
    IPAddresses    []net.IP
  }

  Certificate x509.Certificate
  CertificateRequest x509.CertificateRequest

  CertificateOptions struct {
    SerialNumber        *big.Int
    NotBefore           time.Time
    NotAfter            time.Time // Validity bounds.
    KeyUsage            x509.KeyUsage
  }
)

func NewCertificateData() *CertificateData {
  return &CertificateData{Subject: pkix.Name{}}
}

// Create a certificate sign request from the input data and the private key of
// the request creator.
func (c *CertificateData) ToCertificateRequest(private_key PrivateKey) (*CertificateRequest, error) {
  csr := &x509.CertificateRequest{}

  csr.Subject        = c.Subject
  csr.DNSNames       = c.DNSNames
  csr.IPAddresses    = c.IPAddresses
  csr.EmailAddresses = c.EmailAddresses

  csr_asn1, err := x509.CreateCertificateRequest(rand.Reader, csr, private_key.PrivateKey())
  if err != nil { return nil, err }
  return LoadCertificateSignRequest(csr_asn1)
}

// Load a certificate sign request from its asn1 representation.
func LoadCertificateSignRequest(raw []byte) (*CertificateRequest, error) {
  csr, err := x509.ParseCertificateRequest(raw)
  if err != nil { return nil, err }
  return (*CertificateRequest)(csr), nil
}

// Return the certificate sign request as a pem block.
func (c *CertificateRequest) MarshalPem() (marshalledPemBlock, error) {
  block := &pem.Block{Type: PemLabelCertificateRequest, Bytes: c.Raw}
  return pem.EncodeToMemory(block), nil
}

// Convert the certificate sign request to a certificate using the private key
// of the signer and the certificate of the signer.
// If the certificate is null, the sign request will be used to sign itself.
// Please also see the certificate options struct for information on mandatory fields.
// For more information, please read http://golang.org/pkg/crypto/x509/#CreateCertificate
func (c *CertificateRequest) ToCertificate(private_key PrivateKey,
        cert_opts CertificateOptions, ca *Certificate) (*Certificate, error) {

  if err := cert_opts.Valid(); err != nil { return nil, err }

  template := &x509.Certificate{}
  template.Subject        = c.Subject
  template.DNSNames       = c.DNSNames
  template.IPAddresses    = c.IPAddresses
  template.EmailAddresses = c.EmailAddresses

  // if no ca is given, we have to set IsCA to self sign
  if ca == nil {
    template.IsCA = true
  }

  template.NotBefore    = cert_opts.NotBefore
  template.NotAfter     = cert_opts.NotAfter
  template.KeyUsage     = cert_opts.KeyUsage
  template.SerialNumber = cert_opts.SerialNumber

  var cert_asn1 []byte
  var err  error
  // if we have no ca which can sign the cert, a self signed cert is wanted
  // (or isn't it? Maybe we should split creation of the template? But that would be ugly)
  if ca == nil {
    cert_asn1, err = x509.CreateCertificate(rand.Reader, template, template, c.PublicKey, private_key.PrivateKey())
  } else {
    cert_asn1, err = x509.CreateCertificate(rand.Reader, template, (*x509.Certificate)(ca), c.PublicKey, private_key.PrivateKey())
  }
  if err != nil { return nil, err }
  return LoadCertificate(cert_asn1)
}

// Load a certificate from its asn1 representation.
func LoadCertificate(raw []byte) (*Certificate, error) {
  cert, err := x509.ParseCertificate(raw)
  if err != nil { return nil, err }
  return (*Certificate)(cert), nil
}

// marshal the certificate to a pem block
func (c *Certificate) MarshalPem() (marshalledPemBlock, error) {
  block := &pem.Block{Type: PemLabelCertificate, Bytes: c.Raw}
  return pem.EncodeToMemory(block), nil
}

func (co *CertificateOptions) Valid() error {
  if co.SerialNumber == nil { return fmt.Errorf("No serial number set!") }
  return nil
}
<|endoftext|>


<commit_before>// Challenge 51 - Compression Ratio Side-Channel Attacks
// http://cryptopals.com/sets/7/challenges/51

package cryptopals

import (
	"bytes"
	"compress/flate"
	"io"
	"math"
	"text/template"
)

type challenge51 struct {
}

type encryptor func([]byte) []byte
type decryptor func(compressionOracle) string

type compressionOracle struct {
	cookie  string
	prefix  string
	encrypt encryptor
}

const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="

var t = template.Must(template.New("request").Parse(`POST / HTTP/1.1
Host: hapless.com
Cookie: {{ .Cookie }}
Content-Length: {{ len .Data }}
{{ .Data }}`))

func (oracle compressionOracle) process(data string) int {
	req := new(bytes.Buffer)

	t.Execute(req, struct {
		Cookie string
		Data   string
	}{
		Cookie: oracle.cookie,
		Data:   data,
	})

	b := new(bytes.Buffer)
	w, _ := flate.NewWriter(b, flate.BestCompression)

	io.Copy(w, req)
	w.Close()

	ciphertext := oracle.encrypt(b.Bytes())
	return len(ciphertext)
}

func (oracle compressionOracle) cookieLength() int {
	return len(oracle.cookie)
}

func (x challenge51) getRealLength(oracle compressionOracle, data string) int {
	control := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
	real := oracle.process(data)
	guess := data

	for i := 0; ; i++ {
		guess += string(control[i%len(control)])
		l := oracle.process(guess)

		if real < l {
			return real - i - 1
		}
	}
}

func (x challenge51) DecryptAesCtrCompressed(oracle compressionOracle) string {
	body := oracle.prefix

	for len(body) < oracle.cookieLength() {
		best := int(math.MaxInt32)
		var next rune

		for _, c := range alphabet {
			guess := body + string(c)
			l := oracle.process(guess)

			if l < best {
				best = l
				next = c
			}
		}

		body += string(next)
	}

	return body
}

func (x challenge51) DecryptAesCbcCompressed(oracle compressionOracle) string {
	body := oracle.prefix

	for len(body) < oracle.cookieLength() {
		best := int(math.MaxInt32)
		var next rune

		for _, c := range alphabet {
			guess := body + string(c)
			l := x.getRealLength(oracle, guess)

			if l < best {
				best = l
				next = c
			}
		}

		body += string(next)
	}

	return body
}
<commit_msg>Code simplification<commit_after>// Challenge 51 - Compression Ratio Side-Channel Attacks
// http://cryptopals.com/sets/7/challenges/51

package cryptopals

import (
	"bytes"
	"compress/flate"
	"io"
	"math"
	"text/template"
)

type challenge51 struct {
}

type encryptor func([]byte) []byte
type decryptor func(compressionOracle) string

type compressionOracle struct {
	cookie  string
	prefix  string
	encrypt encryptor
}

const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="

var t = template.Must(template.New("request").Parse(`POST / HTTP/1.1
Host: hapless.com
Cookie: {{ .Cookie }}
Content-Length: {{ len .Data }}
{{ .Data }}`))

func (oracle compressionOracle) process(data string) int {
	req := new(bytes.Buffer)

	t.Execute(req, struct {
		Cookie string
		Data   string
	}{
		Cookie: oracle.cookie,
		Data:   data,
	})

	b := new(bytes.Buffer)
	w, _ := flate.NewWriter(b, flate.BestCompression)

	io.Copy(w, req)
	w.Close()

	ciphertext := oracle.encrypt(b.Bytes())
	return len(ciphertext)
}

func (oracle compressionOracle) cookieLength() int {
	return len(oracle.cookie)
}

func (x challenge51) DecryptAesCtrCompressed(oracle compressionOracle) string {
	body := oracle.prefix

	for len(body) < oracle.cookieLength() {
		best := int(math.MaxInt32)
		var next rune

		for _, c := range alphabet {
			guess := body + string(c)
			l := oracle.process(guess)

			if l < best {
				best = l
				next = c
			}
		}

		body += string(next)
	}

	return body
}

func (x challenge51) DecryptAesCbcCompressed(oracle compressionOracle) string {
	body := oracle.prefix

	for len(body) < oracle.cookieLength() {
		best := int(math.MaxInt32)
		var next rune

		for _, c := range alphabet {
			guess := body + string(c)
			l := oracle.process(guess)

			for i := 0; ; i++ {
				guess += string(i % 10)

				if l < oracle.process(guess) {
					l = l - i - 1
					break
				}
			}

			if l < best {
				best = l
				next = c
			}
		}

		body += string(next)
	}

	return body
}
<|endoftext|>


<commit_before>package client

import (
	"github.com/miolini/bankgo/rpc/proto"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

type BalanceRecord struct {
	UserID int64
	Value  int64
}

type BalanceStorageClient struct {
	conn   *grpc.ClientConn
	client proto.BalanceStorageClient
}

func Connect(addr string) (*BalanceStorageClient, error) {
	conn, err := grpc.Dial(addr, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}
	bsc := &BalanceStorageClient{}
	bsc.conn = conn
	bsc.client = proto.NewBalanceStorageClient(bsc.conn)
	return bsc, nil
}

func (bsc *BalanceStorageClient) Close() {
	bsc.conn.Close()
}

func (bsc *BalanceStorageClient) GetBalance(userID int64) (int64, error) {
	request := proto.GetRequest{UserId: userID}
	response, err := bsc.client.Get(context.Background(), &request)
	return response.Value, err
}

func (bsc *BalanceStorageClient) SetValue(userID int64, value int64) (int64, error) {
	request := proto.SetRequest{UserId: userID, Value: value}
	response, err := bsc.client.Set(context.Background(), &request)
	return response.Value, err
}

func (bsc *BalanceStorageClient) IncrementValue(userID int64, amount int64) (int64, error) {
	request := proto.IncrementRequest{UserId: userID, Amount: amount}
	response, err := bsc.client.Increment(context.Background(), &request)
	if err != nil {
		return 0, err
	}
	return response.Value, nil
}

func (bsc *BalanceStorageClient) AllBalances() ([]BalanceRecord, error) {
	response, err := bsc.client.All(context.Background(), &proto.AllRequest{})
	if err != nil {
		return nil, err
	}
	result := make([]BalanceRecord, len(response.Records))
	for i := range response.Records {
		result[i] = BalanceRecord{response.Records[i].UserId, response.Records[i].Value}
	}
	return result, nil
}
<commit_msg>fix rpc client check err<commit_after>package client

import (
	"github.com/miolini/bankgo/rpc/proto"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

type BalanceRecord struct {
	UserID int64
	Value  int64
}

type BalanceStorageClient struct {
	conn   *grpc.ClientConn
	client proto.BalanceStorageClient
}

func Connect(addr string) (*BalanceStorageClient, error) {
	conn, err := grpc.Dial(addr, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}
	bsc := &BalanceStorageClient{}
	bsc.conn = conn
	bsc.client = proto.NewBalanceStorageClient(bsc.conn)
	return bsc, nil
}

func (bsc *BalanceStorageClient) Close() {
	bsc.conn.Close()
}

func (bsc *BalanceStorageClient) GetBalance(userID int64) (int64, error) {
	request := proto.GetRequest{UserId: userID}
	response, err := bsc.client.Get(context.Background(), &request)
	if err != nil {
		return 0, nil
	}
	return response.Value, nil
}

func (bsc *BalanceStorageClient) SetValue(userID int64, value int64) (int64, error) {
	request := proto.SetRequest{UserId: userID, Value: value}
	response, err := bsc.client.Set(context.Background(), &request)
	if err != nil {
		return 0, err
	}
	return response.Value, nil
}

func (bsc *BalanceStorageClient) IncrementValue(userID int64, amount int64) (int64, error) {
	request := proto.IncrementRequest{UserId: userID, Amount: amount}
	response, err := bsc.client.Increment(context.Background(), &request)
	if err != nil {
		return 0, err
	}
	return response.Value, nil
}

func (bsc *BalanceStorageClient) AllBalances() ([]BalanceRecord, error) {
	response, err := bsc.client.All(context.Background(), &proto.AllRequest{})
	if err != nil {
		return nil, err
	}
	result := make([]BalanceRecord, len(response.Records))
	for i := range response.Records {
		result[i] = BalanceRecord{response.Records[i].UserId, response.Records[i].Value}
	}
	return result, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"golang.org/x/net/websocket"
	"html"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	p "path"
	"strconv"
	"strings"
	"time"
)

func main() {
	http.HandleFunc("/", handleAny)
	log.Fatal(http.ListenAndServe(":"+httpPort(), nil))
}

func handleAny(res http.ResponseWriter, req *http.Request) {
	log.Printf("method=%s path=%q", req.Method, req.URL.Path)

	path, err := os.Open(req.URL.Path)
	if err != nil {
		if os.IsNotExist(err) {
			handleError(res, req, err, http.StatusNotFound, "File not found")
			return
		}

		handleError(res, req, err, http.StatusBadRequest, "Error reading path")
		return
	}
	defer path.Close()

	pathInfo, err := path.Stat()
	if err != nil {
		handleError(res, req, err, http.StatusInternalServerError, "Error reading path info")
		return
	}

	if strings.HasPrefix(req.Header.Get("Origin"), "ws://") { // TODO: how to detect? scheme is null
		handleWs(res, req, path, pathInfo)
		return
	}

	switch req.Method {
	case "GET":
		handleGet(res, req, path, pathInfo)
	case "POST":
		handlePost(res, req, path, pathInfo)
	default:
		http.Error(res, "Only GET and POST supported", http.StatusMethodNotAllowed)
	}
}

func handleGet(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo) {
	if pathInfo.Mode().IsDir() {
		fileInfos, err := path.Readdir(0)
		if err != nil {
			handleError(res, req, err, http.StatusInternalServerError, "Error reading directory")
			return
		}

		if strings.Contains(req.Header.Get("Accept"), "html") {
			renderDirHtml(res, path.Name(), fileInfos)
		} else {
			renderDirJson(res, fileInfos)
		}
	} else if pathInfo.Mode().IsRegular() {
		stat, err := path.Stat()
		if err != nil {
			handleError(res, req, err, http.StatusInternalServerError, "Error reading file stat")
			return
		}

		// explicitly set Content-Length for clients to track download progress
		// except for 0 bytes files which could include special 0-byte unix files
		if stat.Size() > 0 {
			res.Header().Set("Content-Length", strconv.FormatInt(stat.Size(), 10))
		}

		io.Copy(res, path)
	} else {
		handleError(res, req, nil,
			http.StatusBadRequest,
			"Invalid file type for GET. Only directories and regular files are supported.")
	}
}

func handlePost(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo) {
	resFlusherWriter := flushWriterWrapper{res.(flushWriter)}
	execCmd(res, req, path, pathInfo, req.Body, resFlusherWriter, false)
}

func handleWs(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo) {
	handler := func(ws *websocket.Conn) {
		execCmd(res, req, path, pathInfo, ws, ws, true)
	}

	websocket.Handler(handler).ServeHTTP(res, req)
}

func execCmd(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo, in io.Reader, out io.Writer, interactive bool) {
	var cmd *exec.Cmd

	if pathInfo.Mode().IsDir() {
		if interactive {
			// TODO: allow interactive session to have a prompt, support heredocs, handle arrow keys, and generally act like a real terminal
			// cmd = exec.Command("bash", "-i") // double echos (ws + bash -i) and displays arrow character

			// pseudo-interactive-bash worksaround `bash -i` echoing problem, but breaks on heredocs and probably other bash special cases
			dir, err := os.Getwd()
			if err != nil {
				log.Fatal(err)
			}
			log.Println(dir)
			cmd = exec.Command(p.Join(dir, "bin", "pseudo-interactive-bash"))
		} else {
			cmd = exec.Command("bash")
		}
		cmd.Dir = path.Name()
	} else if pathInfo.Mode().IsRegular() && pathInfo.Mode()&0110 != 0 /* is executable for user or group */ {
		cmd = exec.Command(path.Name())
		cmd.Dir = path.Name()[0:strings.LastIndex(path.Name(), string(os.PathSeparator))]
	} else {
		handleError(res, req, nil,
			http.StatusBadRequest,
			"Invalid file type for POST. Only directories and regular executable file are supported")
		return
	}

	cmd.Env = append(os.Environ(), cgiEnv(req)...)
	if interactive {
		cmd.Env = append(cmd.Env, "PS1=\\[\\033[01;34m\\]\\w\\[\\033[00m\\] \\[\\033[01;32m\\]$ \\[\\033[00m\\]")
	}
	cmd.Stdin = in
	cmd.Stdout = out
	cmd.Stderr = out
	if err := cmd.Run(); err != nil {
		// error already sent to client. log only
		log.Printf("method=%s path=%q message=%q", req.Method, req.URL.Path, err)
	}
}

func renderDirHtml(res http.ResponseWriter, pathName string, fileInfos []os.FileInfo) {
	res.Header().Set("Content-Type", "text/html")
	fmt.Fprintf(res, "

")
	for _, fi := range fileInfos {
		if !strings.HasSuffix(pathName, "/") {
			pathName += "/"
		}

		label := fi.Name()
		if fi.IsDir() {
			label = "/" + label
		}

		fmt.Fprintf(res, "
%s
", pathName, fi.Name(), html.EscapeString(label))
	}
	fmt.Fprintf(res, "

")
}

func renderDirJson(res http.ResponseWriter, fileInfos []os.FileInfo) {
	res.Header().Set("Content-Type", "application/json")
	fileResponses := map[string]FileInfoDetails{}
	for _, fi := range fileInfos {
		fileResponses[fi.Name()] = toFileInfoDetails(fi)
	}
	fileResponsesJson, err := json.MarshalIndent(fileResponses, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Fprintln(res, string(fileResponsesJson))
}

type FileInfoDetails struct {
	Size    int64     `json:"size"`
	Type    string    `json:"type"`
	Perm    int       `json:"permission"`
	ModTime time.Time `json:"updated_at"`
}

func toFileInfoDetails(fi os.FileInfo) FileInfoDetails {
	return FileInfoDetails{
		Size:    fi.Size(),
		Type:    string(fi.Mode().String()[0]),
		Perm:    int(fi.Mode().Perm()),
		ModTime: fi.ModTime(),
	}
}

type stdError struct {
	Message string `json:"message"`
	Cause   error  `json:"cause"`
}

func handleError(res http.ResponseWriter, req *http.Request, err error, httpStatus int, message string) {
	stdError := stdError{
		Message: message,
		Cause:   err,
	}

	log.Printf("method=%s path=%q message=%q cause=%q", req.Method, req.URL.Path, stdError.Message, stdError.Cause)

	stdErrorJson, err := json.MarshalIndent(stdError, "", "  ")
	if err != nil {
		panic(err)
	}

	res.Header().Set("Content-Type", "application/json")
	http.Error(res, string(stdErrorJson), httpStatus)
}

type flushWriter interface {
	Flush()
	Write(buf []byte) (int, error)
}

type flushWriterWrapper struct {
	fw flushWriter
}

func (fww flushWriterWrapper) Write(p []byte) (n int, err error) {
	n, err = fww.fw.Write(p)
	fww.fw.Flush()
	return
}

func httpPort() string {
	if port := os.Getenv("PORT"); port != "" {
		return port
	}
	return "5000"
}

// copied from net/http/cgi/host.go
func cgiEnv(req *http.Request) []string {
	env := []string{
		"SERVER_SOFTWARE=go",
		"SERVER_NAME=" + req.Host,
		"SERVER_PROTOCOL=HTTP/1.1",
		"HTTP_HOST=" + req.Host,
		"GATEWAY_INTERFACE=CGI/1.1",
		"REQUEST_METHOD=" + req.Method,
		"QUERY_STRING=" + req.URL.RawQuery,
		"REQUEST_URI=" + req.URL.RequestURI(),
		"PATH_INFO=" + req.URL.Path,
		"SCRIPT_NAME=" + req.URL.Path,
		"SCRIPT_FILENAME=" + req.URL.Path,
		"REMOTE_ADDR=" + req.RemoteAddr,
		"REMOTE_HOST=" + req.RemoteAddr,
		"SERVER_PORT=" + httpPort(),
	}

	if req.TLS != nil {
		env = append(env, "HTTPS=on")
	}

	for k, v := range req.Header {
		k = strings.Map(upperCaseAndUnderscore, k)
		joinStr := ", "
		if k == "COOKIE" {
			joinStr = "; "
		}
		env = append(env, "HTTP_"+k+"="+strings.Join(v, joinStr))
	}

	if req.ContentLength > 0 {
		env = append(env, fmt.Sprintf("CONTENT_LENGTH=%d", req.ContentLength))
	}
	if ctype := req.Header.Get("Content-Type"); ctype != "" {
		env = append(env, "CONTENT_TYPE="+ctype)
	}

	return env
}

func upperCaseAndUnderscore(r rune) rune {
	switch {
	case r >= 'a' && r <= 'z':
		return r - ('a' - 'A')
	case r == '-':
		return '_'
	case r == '=':
		// Maybe not part of the CGI 'spec' but would mess up
		// the environment in any case, as Go represents the
		// environment as a slice of "key=value" strings.
		return '_'
	}
	// TODO: other transformations in spec or practice?
	return r
}
<commit_msg>Move slash label to end<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"golang.org/x/net/websocket"
	"html"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	p "path"
	"strconv"
	"strings"
	"time"
)

func main() {
	http.HandleFunc("/", handleAny)
	log.Fatal(http.ListenAndServe(":"+httpPort(), nil))
}

func handleAny(res http.ResponseWriter, req *http.Request) {
	log.Printf("method=%s path=%q", req.Method, req.URL.Path)

	path, err := os.Open(req.URL.Path)
	if err != nil {
		if os.IsNotExist(err) {
			handleError(res, req, err, http.StatusNotFound, "File not found")
			return
		}

		handleError(res, req, err, http.StatusBadRequest, "Error reading path")
		return
	}
	defer path.Close()

	pathInfo, err := path.Stat()
	if err != nil {
		handleError(res, req, err, http.StatusInternalServerError, "Error reading path info")
		return
	}

	if strings.HasPrefix(req.Header.Get("Origin"), "ws://") { // TODO: how to detect? scheme is null
		handleWs(res, req, path, pathInfo)
		return
	}

	switch req.Method {
	case "GET":
		handleGet(res, req, path, pathInfo)
	case "POST":
		handlePost(res, req, path, pathInfo)
	default:
		http.Error(res, "Only GET and POST supported", http.StatusMethodNotAllowed)
	}
}

func handleGet(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo) {
	if pathInfo.Mode().IsDir() {
		fileInfos, err := path.Readdir(0)
		if err != nil {
			handleError(res, req, err, http.StatusInternalServerError, "Error reading directory")
			return
		}

		if strings.Contains(req.Header.Get("Accept"), "html") {
			renderDirHtml(res, path.Name(), fileInfos)
		} else {
			renderDirJson(res, fileInfos)
		}
	} else if pathInfo.Mode().IsRegular() {
		stat, err := path.Stat()
		if err != nil {
			handleError(res, req, err, http.StatusInternalServerError, "Error reading file stat")
			return
		}

		// explicitly set Content-Length for clients to track download progress
		// except for 0 bytes files which could include special 0-byte unix files
		if stat.Size() > 0 {
			res.Header().Set("Content-Length", strconv.FormatInt(stat.Size(), 10))
		}

		io.Copy(res, path)
	} else {
		handleError(res, req, nil,
			http.StatusBadRequest,
			"Invalid file type for GET. Only directories and regular files are supported.")
	}
}

func handlePost(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo) {
	resFlusherWriter := flushWriterWrapper{res.(flushWriter)}
	execCmd(res, req, path, pathInfo, req.Body, resFlusherWriter, false)
}

func handleWs(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo) {
	handler := func(ws *websocket.Conn) {
		execCmd(res, req, path, pathInfo, ws, ws, true)
	}

	websocket.Handler(handler).ServeHTTP(res, req)
}

func execCmd(res http.ResponseWriter, req *http.Request, path *os.File, pathInfo os.FileInfo, in io.Reader, out io.Writer, interactive bool) {
	var cmd *exec.Cmd

	if pathInfo.Mode().IsDir() {
		if interactive {
			// TODO: allow interactive session to have a prompt, support heredocs, handle arrow keys, and generally act like a real terminal
			// cmd = exec.Command("bash", "-i") // double echos (ws + bash -i) and displays arrow character

			// pseudo-interactive-bash worksaround `bash -i` echoing problem, but breaks on heredocs and probably other bash special cases
			dir, err := os.Getwd()
			if err != nil {
				log.Fatal(err)
			}
			log.Println(dir)
			cmd = exec.Command(p.Join(dir, "bin", "pseudo-interactive-bash"))
		} else {
			cmd = exec.Command("bash")
		}
		cmd.Dir = path.Name()
	} else if pathInfo.Mode().IsRegular() && pathInfo.Mode()&0110 != 0 /* is executable for user or group */ {
		cmd = exec.Command(path.Name())
		cmd.Dir = path.Name()[0:strings.LastIndex(path.Name(), string(os.PathSeparator))]
	} else {
		handleError(res, req, nil,
			http.StatusBadRequest,
			"Invalid file type for POST. Only directories and regular executable file are supported")
		return
	}

	cmd.Env = append(os.Environ(), cgiEnv(req)...)
	if interactive {
		cmd.Env = append(cmd.Env, "PS1=\\[\\033[01;34m\\]\\w\\[\\033[00m\\] \\[\\033[01;32m\\]$ \\[\\033[00m\\]")
	}
	cmd.Stdin = in
	cmd.Stdout = out
	cmd.Stderr = out
	if err := cmd.Run(); err != nil {
		// error already sent to client. log only
		log.Printf("method=%s path=%q message=%q", req.Method, req.URL.Path, err)
	}
}

func renderDirHtml(res http.ResponseWriter, pathName string, fileInfos []os.FileInfo) {
	res.Header().Set("Content-Type", "text/html")
	fmt.Fprintf(res, "

")
	for _, fi := range fileInfos {
		if !strings.HasSuffix(pathName, "/") {
			pathName += "/"
		}

		label := fi.Name()
		if fi.IsDir() {
			label = label + "/"
		}

		fmt.Fprintf(res, "
%s
", pathName, fi.Name(), html.EscapeString(label))
	}
	fmt.Fprintf(res, "

")
}

func renderDirJson(res http.ResponseWriter, fileInfos []os.FileInfo) {
	res.Header().Set("Content-Type", "application/json")
	fileResponses := map[string]FileInfoDetails{}
	for _, fi := range fileInfos {
		fileResponses[fi.Name()] = toFileInfoDetails(fi)
	}
	fileResponsesJson, err := json.MarshalIndent(fileResponses, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Fprintln(res, string(fileResponsesJson))
}

type FileInfoDetails struct {
	Size    int64     `json:"size"`
	Type    string    `json:"type"`
	Perm    int       `json:"permission"`
	ModTime time.Time `json:"updated_at"`
}

func toFileInfoDetails(fi os.FileInfo) FileInfoDetails {
	return FileInfoDetails{
		Size:    fi.Size(),
		Type:    string(fi.Mode().String()[0]),
		Perm:    int(fi.Mode().Perm()),
		ModTime: fi.ModTime(),
	}
}

type stdError struct {
	Message string `json:"message"`
	Cause   error  `json:"cause"`
}

func handleError(res http.ResponseWriter, req *http.Request, err error, httpStatus int, message string) {
	stdError := stdError{
		Message: message,
		Cause:   err,
	}

	log.Printf("method=%s path=%q message=%q cause=%q", req.Method, req.URL.Path, stdError.Message, stdError.Cause)

	stdErrorJson, err := json.MarshalIndent(stdError, "", "  ")
	if err != nil {
		panic(err)
	}

	res.Header().Set("Content-Type", "application/json")
	http.Error(res, string(stdErrorJson), httpStatus)
}

type flushWriter interface {
	Flush()
	Write(buf []byte) (int, error)
}

type flushWriterWrapper struct {
	fw flushWriter
}

func (fww flushWriterWrapper) Write(p []byte) (n int, err error) {
	n, err = fww.fw.Write(p)
	fww.fw.Flush()
	return
}

func httpPort() string {
	if port := os.Getenv("PORT"); port != "" {
		return port
	}
	return "5000"
}

// copied from net/http/cgi/host.go
func cgiEnv(req *http.Request) []string {
	env := []string{
		"SERVER_SOFTWARE=go",
		"SERVER_NAME=" + req.Host,
		"SERVER_PROTOCOL=HTTP/1.1",
		"HTTP_HOST=" + req.Host,
		"GATEWAY_INTERFACE=CGI/1.1",
		"REQUEST_METHOD=" + req.Method,
		"QUERY_STRING=" + req.URL.RawQuery,
		"REQUEST_URI=" + req.URL.RequestURI(),
		"PATH_INFO=" + req.URL.Path,
		"SCRIPT_NAME=" + req.URL.Path,
		"SCRIPT_FILENAME=" + req.URL.Path,
		"REMOTE_ADDR=" + req.RemoteAddr,
		"REMOTE_HOST=" + req.RemoteAddr,
		"SERVER_PORT=" + httpPort(),
	}

	if req.TLS != nil {
		env = append(env, "HTTPS=on")
	}

	for k, v := range req.Header {
		k = strings.Map(upperCaseAndUnderscore, k)
		joinStr := ", "
		if k == "COOKIE" {
			joinStr = "; "
		}
		env = append(env, "HTTP_"+k+"="+strings.Join(v, joinStr))
	}

	if req.ContentLength > 0 {
		env = append(env, fmt.Sprintf("CONTENT_LENGTH=%d", req.ContentLength))
	}
	if ctype := req.Header.Get("Content-Type"); ctype != "" {
		env = append(env, "CONTENT_TYPE="+ctype)
	}

	return env
}

func upperCaseAndUnderscore(r rune) rune {
	switch {
	case r >= 'a' && r <= 'z':
		return r - ('a' - 'A')
	case r == '-':
		return '_'
	case r == '=':
		// Maybe not part of the CGI 'spec' but would mess up
		// the environment in any case, as Go represents the
		// environment as a slice of "key=value" strings.
		return '_'
	}
	// TODO: other transformations in spec or practice?
	return r
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
)

var port = flag.String("p", "8888", "Port to listen on")
var storage = flag.String("s", "/var/lib/shelf", "Path to store files")

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:	%v -p <port> -s <storage-path>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func errorResponse(w http.ResponseWriter, e error) {
	w.WriteHeader(http.StatusInternalServerError)
	w.Write([]byte(e.Error()))
	log.Println("error:", e.Error())
}

func main() {
	flag.Parse()

	var storagepath string
	if flag.NArg() == 1 {
		// deprecated: passing storage path as argument
		storagepath = flag.Arg(0)
	} else {
		storagepath = *storage
	}
	os.MkdirAll(storagepath, 0755)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		filepath := storagepath + r.RequestURI
		switch r.Method {
		case "GET":
			file, err := os.Open(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			fi, err := file.Stat()
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("GET", r.RequestURI)
			http.ServeContent(w, r, filepath, fi.ModTime(), file)
		case "PUT":
			os.MkdirAll(path.Dir(filepath), 0755)
			file, err := os.Create(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			_, err = io.Copy(file, r.Body)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("PUT", r.RequestURI)
		case "DELETE":
			err := os.RemoveAll(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("DELETE", r.RequestURI)
		default:
			w.WriteHeader(http.StatusMethodNotAllowed)
		}
	})
	log.Println("Shelf serving files on " + *port + " from " + storagepath)
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<commit_msg>blobstore: Merge pull request #5 from miracle2k/master<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
)

var port = flag.String("p", "8888", "Port to listen on")
var storage = flag.String("s", "/var/lib/shelf", "Path to store files")

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:	%v -p <port> -s <storage-path>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func errorResponse(w http.ResponseWriter, e error) {
	code := http.StatusInternalServerError
	if os.IsNotExist(e) {
		code = http.StatusNotFound
	}
	w.WriteHeader(code)
	w.Write([]byte(e.Error()))
	log.Println("error:", e.Error())
}

func main() {
	flag.Parse()

	var storagepath string
	if flag.NArg() == 1 {
		// deprecated: passing storage path as argument
		storagepath = flag.Arg(0)
	} else {
		storagepath = *storage
	}
	os.MkdirAll(storagepath, 0755)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		filepath := storagepath + r.RequestURI
		switch r.Method {
		case "HEAD", "GET":
			file, err := os.Open(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			fi, err := file.Stat()
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("GET", r.RequestURI)
			http.ServeContent(w, r, filepath, fi.ModTime(), file)
		case "PUT":
			os.MkdirAll(path.Dir(filepath), 0755)
			file, err := os.Create(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			defer file.Close()
			_, err = io.Copy(file, r.Body)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("PUT", r.RequestURI)
		case "DELETE":
			err := os.RemoveAll(filepath)
			if err != nil {
				errorResponse(w, err)
				return
			}
			log.Println("DELETE", r.RequestURI)
		default:
			w.WriteHeader(http.StatusMethodNotAllowed)
		}
	})
	log.Println("Shelf serving files on " + *port + " from " + storagepath)
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<|endoftext|>


<commit_before>// Flake generates unique identifiers that are roughly sortable by time. Flake can
// run on a cluster of machines and still generate unique IDs without requiring
// worker coordination.
//
// A Flake ID is a 64-bit integer will the following components:
//  - 41 bits is the timestamp with millisecond precision
//  - 10 bits is the host id (uses IP modulo 2^10)
//  - 13 bits is an auto-incrementing sequence for ID requests within the same millisecond
//
// Note: In order to make a millisecond timestamp fit within 41 bits, a custom
// epoch of Jan 1, 2015 00:00:00 is used.

package flake

import (
	"net"
	"os"
	"strconv"
	"sync"
	"time"
)

//-----------------------------------------------------------------------------

const (
	HostBits     = 10
	SequenceBits = 13
)

var (
	// Custom Epoch so the timestamp can fit into 41 bits.
	// Jan 1, 2015 00:00:00 UTC
	Epoch       time.Time = time.Date(2015, 1, 1, 0, 0, 0, 0, time.UTC)
	MaxHostId   uint64    = (1 << HostBits) - 1
	MaxSequence uint64    = (1 << SequenceBits) - 1
)

// Id represents a unique k-ordered Id
type Id uint64

// String formats the Id as a base36 string
func (id Id) String() string {
	return strconv.FormatUint(uint64(id), 36)
}

// Uint64 formats the Id as an unsigned integer
func (id Id) Uint64() uint64 {
	return uint64(id)
}

// Flake is a unique Id generator
type Flake struct {
	prevTime uint64
	hostId   uint64
	sequence uint64
	mu       sync.Mutex
}

// New returns a new Id generator and a possible error condition
func New() (*Flake, error) {
	hostId, err := getHostId()
	if err != nil {
		return nil, err
	}

	return &Flake{
		sequence: 0,
		prevTime: getTimestamp(),
		hostId:   hostId,
	}, nil
}

// NextId returns a new Id from the generator
func (f *Flake) NextId() Id {
	f.mu.Lock()
	defer f.mu.Unlock()

	now := getTimestamp()

	if now < f.prevTime {
		now = f.prevTime
	}

	// Use the sequence number if the id request is in the same millisecond as
	// the previous request.
	if now == f.prevTime {
		f.sequence += 1
	} else {
		f.sequence = 0
	}

	// Bump the timestamp by 1ms if we run out of sequence bits.
	if f.sequence > MaxSequence {
		now += 1
		f.sequence = 0
	}

	f.prevTime = now

	timestamp := now << (HostBits + SequenceBits)
	hostid := f.hostId << SequenceBits

	return Id(timestamp | hostid | f.sequence)
}

// getTimestamp returns the timestamp in milliseconds adjusted for the custom
// epoch
func getTimestamp() uint64 {
	return uint64(time.Since(Epoch).Nanoseconds() / 1e6)
}

// getHostId returns the host id using the IP address of the machine
func getHostId() (uint64, error) {
	h, err := os.Hostname()
	if err != nil {
		return 0, err
	}

	addrs, err := net.LookupIP(h)
	if err != nil {
		return 0, err
	}

	a := addrs[0].To4()
	ip := uint64((a[0] * 1 << 24) + (a[1] * 1 << 16) + (a[2] * 1 << 8) + a[3])

	return ip % MaxHostId, nil
}
<commit_msg>refactor id generation<commit_after>// Flake generates unique identifiers that are roughly sortable by time. Flake can
// run on a cluster of machines and still generate unique IDs without requiring
// worker coordination.
//
// A Flake ID is a 64-bit integer will the following components:
//  - 41 bits is the timestamp with millisecond precision
//  - 10 bits is the host id (uses IP modulo 2^10)
//  - 13 bits is an auto-incrementing sequence for ID requests within the same millisecond
//
// Note: In order to make a millisecond timestamp fit within 41 bits, a custom
// epoch of Jan 1, 2015 00:00:00 is used.

package flake

import (
	"net"
	"os"
	"strconv"
	"sync"
	"time"
)

const (
	HostBits     = 10
	SequenceBits = 13
)

var (
	// Custom Epoch so the timestamp can fit into 41 bits.
	// Jan 1, 2015 00:00:00 UTC
	Epoch       time.Time = time.Date(2015, 1, 1, 0, 0, 0, 0, time.UTC)
	MaxHostId   uint64    = (1 << HostBits) - 1
	MaxSequence uint64    = (1 << SequenceBits) - 1
)

// Id represents a unique k-ordered Id
type Id uint64

// String formats the Id as a base36 string
func (id Id) String() string {
	return strconv.FormatUint(uint64(id), 36)
}

// Uint64 formats the Id as an unsigned integer
func (id Id) Uint64() uint64 {
	return uint64(id)
}

// Flake is a unique Id generator
type Flake struct {
	prevTime uint64
	hostId   uint64
	sequence uint64
	mu       sync.Mutex
}

// New returns a new Id generator and a possible error condition
func New() (*Flake, error) {
	hostId, err := getHostId()
	if err != nil {
		return nil, err
	}

	return &Flake{
		sequence: 0,
		prevTime: getTimestamp(),
		hostId:   hostId,
	}, nil
}

// NextId returns a new Id from the generator
func (f *Flake) NextId() Id {
	now := getTimestamp()

	f.mu.Lock()
	sequence := f.sequence

	// Use the sequence number if the id request is in the same millisecond as
	// the previous request.
	if now <= f.prevTime {
		now = f.prevTime
		sequence++
	} else {
		sequence = 0
	}

	// Bump the timestamp by 1ms if we run out of sequence bits.
	if sequence > MaxSequence {
		now++
		sequence = 0
	}

	f.prevTime = now
	f.sequence = sequence
	f.mu.Unlock()

	timestamp := now << (HostBits + SequenceBits)
	hostid := f.hostId << SequenceBits
	return Id(timestamp | hostid | sequence)
}

// getTimestamp returns the timestamp in milliseconds adjusted for the custom
// epoch
func getTimestamp() uint64 {
	return uint64(time.Since(Epoch).Nanoseconds() / 1e6)
}

// getHostId returns the host id using the IP address of the machine
func getHostId() (uint64, error) {
	h, err := os.Hostname()
	if err != nil {
		return 0, err
	}

	addrs, err := net.LookupIP(h)
	if err != nil {
		return 0, err
	}

	a := addrs[0].To4()
	ip := uint64((a[0] * 1 << 24) + (a[1] * 1 << 16) + (a[2] * 1 << 8) + a[3])
	return ip % MaxHostId, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"sync"

	"github.com/uber/tchannel-go/typed"
	"golang.org/x/net/context"
)

var (
	errDuplicateMex        = errors.New("multiple attempts to use the message id")
	errMexChannelFull      = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
	errUnexpectedFrameType = errors.New("unexpected frame received")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"

	// mexChannelBufferSize is the size of the message exchange channel buffer.
	mexChannelBufferSize = 2
)

// A messageExchange tracks this Connections's side of a message exchange with a
// peer.  Each message exchange has a channel that can be used to receive
// frames from the peer, and a Context that can controls when the exchange has
// timed out or been cancelled.
type messageExchange struct {
	recvCh    chan *Frame
	ctx       context.Context
	msgID     uint32
	msgType   messageType
	mexset    *messageExchangeSet
	framePool FramePool
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where
// it can be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	if err := mex.ctx.Err(); err != nil {
		return GetContextError(err)
	}
	select {
	case mex.recvCh <- frame:
		return nil
	case <-mex.ctx.Done():
		// Note: One slow reader processing a large request could stall the connection.
		// If we see this, we need to increase the recvCh buffer size.
		return GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrame waits for a new frame from the peer, or until the context
// expires or is cancelled
func (mex *messageExchange) recvPeerFrame() (*Frame, error) {
	if err := mex.ctx.Err(); err != nil {
		return nil, GetContextError(err)
	}

	select {
	case frame := <-mex.recvCh:
		if frame.Header.ID != mex.msgID {
			mex.mexset.log.WithFields(
				LogField{"msgId", mex.msgID},
				LogField{"header", frame.Header},
			).Error("recvPeerFrame received msg with unexpected ID.")
			return nil, errUnexpectedFrameType
		}
		return frame, nil
	case <-mex.ctx.Done():
		return nil, GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrameOfType waits for a new frame of a given type from the peer, failing
// if the next frame received is not of that type.
// If an error frame is returned, then the errorMessage is returned as the error.
func (mex *messageExchange) recvPeerFrameOfType(msgType messageType) (*Frame, error) {
	frame, err := mex.recvPeerFrame()
	if err != nil {
		return nil, err
	}

	switch frame.Header.messageType {
	case msgType:
		return frame, nil

	case messageTypeError:
		// If we read an error frame, we can release it once we deserialize it.
		defer mex.framePool.Release(frame)

		errMsg := errorMessage{
			id: frame.Header.ID,
		}
		var rbuf typed.ReadBuffer
		rbuf.Wrap(frame.SizedPayload())
		if err := errMsg.read(&rbuf); err != nil {
			return nil, err
		}
		return nil, errMsg

	default:
		// TODO(mmihic): Should be treated as a protocol error
		mex.mexset.log.WithFields(
			LogField{"header", frame.Header},
			LogField{"expectedType", msgType},
			LogField{"expectedID", mex.msgID},
		).Warn("Received unexpected frame.")
		return nil, errUnexpectedFrameType
	}
}

// shutdown shuts down the message exchange, removing it from the message
// exchange set so  that it cannot receive more messages from the peer.  The
// receive channel remains open, however, in case there are concurrent
// goroutines sending to it.
func (mex *messageExchange) shutdown() {
	mex.mexset.removeExchange(mex.msgID)
}

// inboundTimeout is called when an exchange times out, but a handler may still be
// running in the background. Since the handler may still write to the exchange, we
// cannot shutdown the exchange, but we should remove it from the connection's
// exchange list.
func (mex *messageExchange) inboundTimeout() {
	mex.mexset.timeoutExchange(mex.msgID)
}

// A messageExchangeSet manages a set of active message exchanges.  It is
// mainly used to route frames from a peer to the appropriate messageExchange,
// or to cancel or mark a messageExchange as being in error.  Each Connection
// maintains two messageExchangeSets, one to manage exchanges that it has
// initiated (outbound), and another to manage exchanges that the peer has
// initiated (inbound).  The message-type specific handlers are responsible for
// ensuring that their message exchanges are properly registered and removed
// from the corresponding exchange set.
type messageExchangeSet struct {
	sync.RWMutex

	log        Logger
	name       string
	onRemoved  func()
	onAdded    func()
	sendChRefs sync.WaitGroup

	// exchanges is mutable, and is protected by the mutex.
	exchanges map[uint32]*messageExchange
}

// newMessageExchangeSet creates a new messageExchangeSet with a given name.
func newMessageExchangeSet(log Logger, name string) *messageExchangeSet {
	return &messageExchangeSet{
		name:      name,
		log:       log.WithFields(LogField{"exchange", name}),
		exchanges: make(map[uint32]*messageExchange),
	}
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context, framePool FramePool,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Creating new %s message exchange for [%v:%d]", mexset.name, msgType, msgID)
	}

	mex := &messageExchange{
		msgType:   msgType,
		msgID:     msgID,
		ctx:       ctx,
		recvCh:    make(chan *Frame, bufferSize),
		mexset:    mexset,
		framePool: framePool,
	}

	mexset.Lock()
	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.WithFields(
				LogField{"name", mexset.name},
				LogField{"msgType", mex.msgType},
				LogField{"msgID", mex.msgID},
			).Warn("mex registered multiple times.")
		} else {
			mexset.log.WithFields(
				LogField{"msgID", mex.msgID},
				LogField{"existingType", existingMex.msgType},
				LogField{"newType", mex.msgType},
			).Warn("Duplicate msg ID for active and new mex.")
		}

		mexset.Unlock()
		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex
	mexset.sendChRefs.Add(1)
	mexset.Unlock()

	mexset.onAdded()

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// removeExchange removes a message exchange from the set, if it exists.
// It decrements the sendChRefs wait group, signalling that this exchange no longer has
// any active goroutines that will try to send to sendCh.
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)
	}

	mexset.Lock()
	_, found := mexset.exchanges[msgID]
	if found {
		delete(mexset.exchanges, msgID)
	}
	mexset.Unlock()

	if !found {
		mexset.log.WithFields(
			LogField{"msgID", msgID},
		).Error("Tried to remove exchange multiple times")
		return
	}

	// If the message exchange was found, then we perform clean up actions.
	// These clean up actions can only be run once per exchange.
	mexset.sendChRefs.Done()
	mexset.onRemoved()
}

// timeoutExchange is similar to removeExchange, however it does not decrement
// the sendChRefs wait group.
func (mexset *messageExchangeSet) timeoutExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d due to timeout", mexset.name, msgID)

	mexset.Lock()
	delete(mexset.exchanges, msgID)
	mexset.Unlock()

	mexset.onRemoved()
}

// waitForSendCh waits for all goroutines with references to sendCh to complete.
func (mexset *messageExchangeSet) waitForSendCh() {
	mexset.sendChRefs.Wait()
}

func (mexset *messageExchangeSet) count() int {
	mexset.RLock()
	count := len(mexset.exchanges)
	mexset.RUnlock()

	return count
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message
// exchange
func (mexset *messageExchangeSet) forwardPeerFrame(frame *Frame) error {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)
	}

	mexset.RLock()
	mex := mexset.exchanges[frame.Header.ID]
	mexset.RUnlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Infof("received frame %s for %s message exchange that no longer exists",
			frame.Header, mexset.name)
		return nil
	}

	if err := mex.forwardPeerFrame(frame); err != nil {
		mexset.log.Infof("Unable to forward frame %v length %v to %s: %v",
			frame.Header, frame.Header.FrameSize(), mexset.name, err)
		return err
	}

	return nil
}
<commit_msg>Timeouts should not trigger "multiple remove" log<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"sync"

	"github.com/uber/tchannel-go/typed"
	"golang.org/x/net/context"
)

var (
	errDuplicateMex        = errors.New("multiple attempts to use the message id")
	errMexChannelFull      = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
	errUnexpectedFrameType = errors.New("unexpected frame received")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"

	// mexChannelBufferSize is the size of the message exchange channel buffer.
	mexChannelBufferSize = 2
)

// A messageExchange tracks this Connections's side of a message exchange with a
// peer.  Each message exchange has a channel that can be used to receive
// frames from the peer, and a Context that can controls when the exchange has
// timed out or been cancelled.
type messageExchange struct {
	recvCh    chan *Frame
	ctx       context.Context
	msgID     uint32
	msgType   messageType
	mexset    *messageExchangeSet
	framePool FramePool
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where
// it can be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	if err := mex.ctx.Err(); err != nil {
		return GetContextError(err)
	}
	select {
	case mex.recvCh <- frame:
		return nil
	case <-mex.ctx.Done():
		// Note: One slow reader processing a large request could stall the connection.
		// If we see this, we need to increase the recvCh buffer size.
		return GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrame waits for a new frame from the peer, or until the context
// expires or is cancelled
func (mex *messageExchange) recvPeerFrame() (*Frame, error) {
	if err := mex.ctx.Err(); err != nil {
		return nil, GetContextError(err)
	}

	select {
	case frame := <-mex.recvCh:
		if frame.Header.ID != mex.msgID {
			mex.mexset.log.WithFields(
				LogField{"msgId", mex.msgID},
				LogField{"header", frame.Header},
			).Error("recvPeerFrame received msg with unexpected ID.")
			return nil, errUnexpectedFrameType
		}
		return frame, nil
	case <-mex.ctx.Done():
		return nil, GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrameOfType waits for a new frame of a given type from the peer, failing
// if the next frame received is not of that type.
// If an error frame is returned, then the errorMessage is returned as the error.
func (mex *messageExchange) recvPeerFrameOfType(msgType messageType) (*Frame, error) {
	frame, err := mex.recvPeerFrame()
	if err != nil {
		return nil, err
	}

	switch frame.Header.messageType {
	case msgType:
		return frame, nil

	case messageTypeError:
		// If we read an error frame, we can release it once we deserialize it.
		defer mex.framePool.Release(frame)

		errMsg := errorMessage{
			id: frame.Header.ID,
		}
		var rbuf typed.ReadBuffer
		rbuf.Wrap(frame.SizedPayload())
		if err := errMsg.read(&rbuf); err != nil {
			return nil, err
		}
		return nil, errMsg

	default:
		// TODO(mmihic): Should be treated as a protocol error
		mex.mexset.log.WithFields(
			LogField{"header", frame.Header},
			LogField{"expectedType", msgType},
			LogField{"expectedID", mex.msgID},
		).Warn("Received unexpected frame.")
		return nil, errUnexpectedFrameType
	}
}

// shutdown shuts down the message exchange, removing it from the message
// exchange set so  that it cannot receive more messages from the peer.  The
// receive channel remains open, however, in case there are concurrent
// goroutines sending to it.
func (mex *messageExchange) shutdown() {
	mex.mexset.removeExchange(mex.msgID)
}

// inboundTimeout is called when an exchange times out, but a handler may still be
// running in the background. Since the handler may still write to the exchange, we
// cannot shutdown the exchange, but we should remove it from the connection's
// exchange list.
func (mex *messageExchange) inboundTimeout() {
	mex.mexset.timeoutExchange(mex.msgID)
}

// A messageExchangeSet manages a set of active message exchanges.  It is
// mainly used to route frames from a peer to the appropriate messageExchange,
// or to cancel or mark a messageExchange as being in error.  Each Connection
// maintains two messageExchangeSets, one to manage exchanges that it has
// initiated (outbound), and another to manage exchanges that the peer has
// initiated (inbound).  The message-type specific handlers are responsible for
// ensuring that their message exchanges are properly registered and removed
// from the corresponding exchange set.
type messageExchangeSet struct {
	sync.RWMutex

	log        Logger
	name       string
	onRemoved  func()
	onAdded    func()
	sendChRefs sync.WaitGroup

	// maps are mutable, and are protected by the mutex.
	exchanges        map[uint32]*messageExchange
	timeoutExchanges map[uint32]struct{}
}

// newMessageExchangeSet creates a new messageExchangeSet with a given name.
func newMessageExchangeSet(log Logger, name string) *messageExchangeSet {
	return &messageExchangeSet{
		name:             name,
		log:              log.WithFields(LogField{"exchange", name}),
		exchanges:        make(map[uint32]*messageExchange),
		timeoutExchanges: make(map[uint32]struct{}),
	}
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context, framePool FramePool,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Creating new %s message exchange for [%v:%d]", mexset.name, msgType, msgID)
	}

	mex := &messageExchange{
		msgType:   msgType,
		msgID:     msgID,
		ctx:       ctx,
		recvCh:    make(chan *Frame, bufferSize),
		mexset:    mexset,
		framePool: framePool,
	}

	mexset.Lock()
	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.WithFields(
				LogField{"name", mexset.name},
				LogField{"msgType", mex.msgType},
				LogField{"msgID", mex.msgID},
			).Warn("mex registered multiple times.")
		} else {
			mexset.log.WithFields(
				LogField{"msgID", mex.msgID},
				LogField{"existingType", existingMex.msgType},
				LogField{"newType", mex.msgType},
			).Warn("Duplicate msg ID for active and new mex.")
		}

		mexset.Unlock()
		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex
	mexset.sendChRefs.Add(1)
	mexset.Unlock()

	mexset.onAdded()

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// deleteExchange will delete msgID, and return whether it was found or whether it was
// timed out. This method must be called with the lock.
func (mexset *messageExchangeSet) deleteExchange(msgID uint32) (found, timedOut bool) {
	if _, found := mexset.exchanges[msgID]; found {
		delete(mexset.exchanges, msgID)
		return true, false
	}

	if _, timedOut := mexset.timeoutExchanges[msgID]; timedOut {
		delete(mexset.timeoutExchanges, msgID)
		return false, true
	}

	return false, false
}

// removeExchange removes a message exchange from the set, if it exists.
// It decrements the sendChRefs wait group, signalling that this exchange no longer has
// any active goroutines that will try to send to sendCh.
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)
	}

	mexset.Lock()
	found, timedOut := mexset.deleteExchange(msgID)
	mexset.Unlock()

	if !found && !timedOut {
		mexset.log.WithFields(
			LogField{"msgID", msgID},
		).Error("Tried to remove exchange multiple times")
		return
	}

	// If the message exchange was found, then we perform clean up actions.
	// These clean up actions can only be run once per exchange.
	mexset.sendChRefs.Done()
	mexset.onRemoved()
}

// timeoutExchange is similar to removeExchange, however it does not decrement
// the sendChRefs wait group, since there could still be a handler running that
// will write to the send channel.
func (mexset *messageExchangeSet) timeoutExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d due to timeout", mexset.name, msgID)

	mexset.Lock()
	delete(mexset.exchanges, msgID)
	mexset.timeoutExchanges[msgID] = struct{}{}
	mexset.Unlock()

	mexset.onRemoved()
}

// waitForSendCh waits for all goroutines with references to sendCh to complete.
func (mexset *messageExchangeSet) waitForSendCh() {
	mexset.sendChRefs.Wait()
}

func (mexset *messageExchangeSet) count() int {
	mexset.RLock()
	count := len(mexset.exchanges)
	mexset.RUnlock()

	return count
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message
// exchange
func (mexset *messageExchangeSet) forwardPeerFrame(frame *Frame) error {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)
	}

	mexset.RLock()
	mex := mexset.exchanges[frame.Header.ID]
	mexset.RUnlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Infof("received frame %s for %s message exchange that no longer exists",
			frame.Header, mexset.name)
		return nil
	}

	if err := mex.forwardPeerFrame(frame); err != nil {
		mexset.log.Infof("Unable to forward frame %v length %v to %s: %v",
			frame.Header, frame.Header.FrameSize(), mexset.name, err)
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package matrix

import (
	"fmt"
	"math/rand"
)

type Dense struct {
	v          []float64 // [row, row, ..., row]
	rows, cols int
	stride     int // distance between vertically adjacent values
}

// Creates new matrix that refers to v
func NewDense(rows, cols, stride int, v []float64) *Dense {
	if rows*stride > len(v) {
		panic("v is to small")
	}
	return &Dense{v: v, rows: rows, cols: cols, stride: stride}
}

// Retuns new zero matrix
func DenseZero(rows, cols int) *Dense {
	return &Dense{
		v:      make([]float64, cols*rows),
		rows:   rows,
		cols:   cols,
		stride: cols,
	}
}

// Returns identity matrix
func DenseI(n int) *Dense {
	m := DenseZero(n, n)
	inc := n + 1
	for i := 0; i < len(m.v); i += inc {
		m.v[i] = 1
	}
	return m
}

// Returns dimensions of the matrix (rows, cols)
func (m *Dense) Size() (int, int) {
	return m.rows, m.cols
}

// Returns number of rows
func (m *Dense) Rows() int {
	return m.rows
}

// Returns number of columns
func (m *Dense) Cols() int {
	return m.cols
}

// Returns stride
func (m *Dense) Stride() int {
	return m.stride
}

// Returns internal buffer of values
func (m *Dense) Values() []float64 {
	return m.v
}

// Returns value from row i and column k
func (m *Dense) Get(i, k int) float64 {
	m.checkIndexes(i, k)
	return m.v[i*m.stride+k]
}

// Sets value in row i and column k
func (m *Dense) Set(i, k int, a float64) {
	m.checkIndexes(i, k)
	m.v[i*m.stride+k] = a
}

// Sets all elements to a
func (m *Dense) SetAll(a float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		k := m.cols
		for k >= 2 {
			k--
			row[k] = a
			k--
			row[k] = a
		}
		if k != 0 {
			row[0] = a
		}
	}
}

// Sets all elements to: min + rand.Float64() * (max - min)
func (m *Dense) Rand(min, max float64) {
	s := max - min
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = min + s*rand.Float64()
		}
	}
}

// Sets all elements to: mean + stdDev*rand.NormFloat64()
func (m *Dense) RandNorm(mean, stdDev float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = mean + stdDev*rand.NormFloat64()
		}
	}
}

// Returns a slice of a matrix that contains rows from start to stop - 1
func (m *Dense) Hslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.rows {
		panic("bad indexes for horizontal slice")
	}
	return &Dense{
		v:      m.v[start*m.stride : stop*m.stride],
		rows:   stop - start,
		cols:   m.cols,
		stride: m.stride,
	}
}

// Returns a slice of a matrix that contains cols from start to stop - 1
func (m *Dense) Vslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.cols {
		panic("bad indexes for vertical slice")
	}
	return &Dense{
		v:      m.v[start : (m.rows-1)*m.stride+stop],
		rows:   m.rows,
		cols:   stop - start,
		stride: m.stride,
	}

}

// Returns matrix as horizontal vector. Panics if cols != stride.
func (m *Dense) Hvec() *Dense {
	if m.cols != m.stride {
		panic("cols != stride")
	}
	return &Dense{v: m.v, rows: 1, cols: len(m.v), stride: len(m.v)}
}

// Returns matrix as vertical vector. Panics if cols != stride.
func (m *Dense) Vvec() *Dense {
	if m.cols != m.stride {
		panic("cols != stride")
	}
	return &Dense{v: m.v, rows: len(m.v), cols: 1, stride: 1}
}

// Returns true if matrices are equal
func (m *Dense) Equal(a *Dense) bool {
	if m.rows != a.rows || m.cols != a.cols {
		return false
	}
	for i := 0; i < m.rows; i++ {
		mr := m.v[i*m.stride:]
		ar := a.v[i*a.stride:]
		for k := 0; k < m.cols; k++ {
			if mr[k] != ar[k] {
				return false
			}
		}
	}
	return true
}

func (m *Dense) String() string {
	rows, cols := m.Size()
	s := "["
	for i := 0; i < rows; i++ {
		if i != 0 {
			s += "\n "
		}
		o := ""
		for k := 0; k < cols; k++ {
			s += fmt.Sprintf("%s%-g", o, m.Get(i, k))
			if k == 0 {
				o = " "
			}
		}
	}
	return s + "]"
}

// Utils

func (m *Dense) checkIndexes(i, k int) {
	if i < 0 || i >= m.rows {
		panic(fmt.Sprintf("row index (%d) out of range [0, %d]", i, m.rows))
	}
	if k < 0 || k >= m.cols {
		panic(fmt.Sprintf("column index (%d) out of range [0, %d]", k, m.cols))
	}
}

func (m *Dense) checkEqualDims(a *Dense) {
	if m.Rows() != a.Rows() || m.Cols() != a.Cols() {
		panic("dimensions of matrices are not equal")
	}
}
<commit_msg>Set proper length for v in NewDense function<commit_after>package matrix

import (
	"fmt"
	"math/rand"
)

type Dense struct {
	v          []float64 // [row, row, ..., row]
	rows, cols int
	stride     int // distance between vertically adjacent values
}

// Creates new matrix that refers to v
func NewDense(rows, cols, stride int, v []float64) *Dense {
	n := rows*stride
	if n > len(v) {
		panic("v is to small")
	}
	return &Dense{v: v[:n], rows: rows, cols: cols, stride: stride}
}

// Retuns new zero matrix
func DenseZero(rows, cols int) *Dense {
	return &Dense{
		v:      make([]float64, cols*rows),
		rows:   rows,
		cols:   cols,
		stride: cols,
	}
}

// Returns identity matrix
func DenseI(n int) *Dense {
	m := DenseZero(n, n)
	inc := n + 1
	for i := 0; i < len(m.v); i += inc {
		m.v[i] = 1
	}
	return m
}

// Returns dimensions of the matrix (rows, cols)
func (m *Dense) Size() (int, int) {
	return m.rows, m.cols
}

// Returns number of rows
func (m *Dense) Rows() int {
	return m.rows
}

// Returns number of columns
func (m *Dense) Cols() int {
	return m.cols
}

// Returns stride
func (m *Dense) Stride() int {
	return m.stride
}

// Returns internal buffer of values
func (m *Dense) Values() []float64 {
	return m.v
}

// Returns value from row i and column k
func (m *Dense) Get(i, k int) float64 {
	m.checkIndexes(i, k)
	return m.v[i*m.stride+k]
}

// Sets value in row i and column k
func (m *Dense) Set(i, k int, a float64) {
	m.checkIndexes(i, k)
	m.v[i*m.stride+k] = a
}

// Sets all elements to a
func (m *Dense) SetAll(a float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		k := m.cols
		for k >= 2 {
			k--
			row[k] = a
			k--
			row[k] = a
		}
		if k != 0 {
			row[0] = a
		}
	}
}

// Sets all elements to: min + rand.Float64() * (max - min)
func (m *Dense) Rand(min, max float64) {
	s := max - min
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = min + s*rand.Float64()
		}
	}
}

// Sets all elements to: mean + stdDev*rand.NormFloat64()
func (m *Dense) RandNorm(mean, stdDev float64) {
	for i := 0; i < m.rows; i++ {
		row := m.v[i*m.stride:]
		for k := 0; k < m.cols; k++ {
			row[k] = mean + stdDev*rand.NormFloat64()
		}
	}
}

// Returns a slice of a matrix that contains rows from start to stop - 1
func (m *Dense) Hslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.rows {
		panic("bad indexes for horizontal slice")
	}
	return &Dense{
		v:      m.v[start*m.stride : stop*m.stride],
		rows:   stop - start,
		cols:   m.cols,
		stride: m.stride,
	}
}

// Returns a slice of a matrix that contains cols from start to stop - 1
func (m *Dense) Vslice(start, stop int) *Dense {
	if start > stop || start < 0 || stop > m.cols {
		panic("bad indexes for vertical slice")
	}
	return &Dense{
		v:      m.v[start : (m.rows-1)*m.stride+stop],
		rows:   m.rows,
		cols:   stop - start,
		stride: m.stride,
	}

}

// Returns matrix as horizontal vector. Panics if cols != stride.
func (m *Dense) Hvec() *Dense {
	if m.cols != m.stride {
		panic("cols != stride")
	}
	return &Dense{v: m.v, rows: 1, cols: len(m.v), stride: len(m.v)}
}

// Returns matrix as vertical vector. Panics if cols != stride.
func (m *Dense) Vvec() *Dense {
	if m.cols != m.stride {
		panic("cols != stride")
	}
	return &Dense{v: m.v, rows: len(m.v), cols: 1, stride: 1}
}

// Returns true if matrices are equal
func (m *Dense) Equal(a *Dense) bool {
	if m.rows != a.rows || m.cols != a.cols {
		return false
	}
	for i := 0; i < m.rows; i++ {
		mr := m.v[i*m.stride:]
		ar := a.v[i*a.stride:]
		for k := 0; k < m.cols; k++ {
			if mr[k] != ar[k] {
				return false
			}
		}
	}
	return true
}

func (m *Dense) String() string {
	rows, cols := m.Size()
	s := "["
	for i := 0; i < rows; i++ {
		if i != 0 {
			s += "\n "
		}
		o := ""
		for k := 0; k < cols; k++ {
			s += fmt.Sprintf("%s%-g", o, m.Get(i, k))
			if k == 0 {
				o = " "
			}
		}
	}
	return s + "]"
}

// Utils

func (m *Dense) checkIndexes(i, k int) {
	if i < 0 || i >= m.rows {
		panic(fmt.Sprintf("row index (%d) out of range [0, %d]", i, m.rows))
	}
	if k < 0 || k >= m.cols {
		panic(fmt.Sprintf("column index (%d) out of range [0, %d]", k, m.cols))
	}
}

func (m *Dense) checkEqualDims(a *Dense) {
	if m.Rows() != a.Rows() || m.Cols() != a.Cols() {
		panic("dimensions of matrices are not equal")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"html/template"
)

type FormElement struct {
	Input   string            `json:"input", bson:"input"`
	Name    string            `json:"name", bson:"name"`
	Options map[string]string `json:"options", bson:"options"`
}

var (
	formBase = `<!DOCTYPE html>

			<form method="post">
			{{ range . }}
				{{ printinput . }}
			{{ end }}
			<input type="submit" value="Submit" />
			</form>

`

	selectBase = `<select name="{{ .Name }}">
{{ range $label, $value := .Options }}
	<option value="{{ $value }}">{{ $label }}</option>
{{ end }}
</select>
`

	formTmpl   *template.Template
	selectTmpl *template.Template
)

func init() {
	formTmpl = template.Must(template.New("form").Funcs(template.FuncMap{"printinput": printInput}).Parse(formBase))
	selectTmpl = template.Must(template.New("select").Parse(selectBase))
}

func printInput(f *FormElement) (template.HTML, error) {
	var err error
	buf := new(bytes.Buffer)

	switch f.Input {
	case "select":
		err = selectTmpl.Execute(buf, f)
	default:
		return "", fmt.Errorf("invalid form input type %q", f.Input)
	}

	return template.HTML(buf.String()), err
}
<commit_msg>support input type="text"<commit_after>package main

import (
	"bytes"
	"fmt"
	"html/template"
)

type FormElement struct {
	Input   string            `json:"input", bson:"input"`
	Name    string            `json:"name", bson:"name"`
	Options map[string]string `json:"options", bson:"options"`
}

var (
	formBase = `<!DOCTYPE html>

			<form method="post">
			{{ range . }}
				{{ printinput . }}
			{{ end }}
			<input type="submit" value="Submit" />
			</form>

`

	selectBase = `<select name="{{ .Name }}">
{{ range $label, $value := .Options }}
	<option value="{{ $value }}">{{ $label }}</option>
{{ end }}
</select>
`

	textBase = `<input type="text" name="{{ .Name }}">`

	formTmpl   *template.Template
	selectTmpl *template.Template
	textTmpl   *template.Template
)

func init() {
	formTmpl = template.Must(template.New("form").Funcs(template.FuncMap{"printinput": printInput}).Parse(formBase))
	selectTmpl = template.Must(template.New("select").Parse(selectBase))
	textTmpl = template.Must(template.New("text").Parse(textBase))
}

func printInput(f *FormElement) (template.HTML, error) {
	var err error
	buf := new(bytes.Buffer)

	switch f.Input {
	case "text":
		err = textTmpl.Execute(buf, f)
	case "select":
		err = selectTmpl.Execute(buf, f)
	default:
		return "", fmt.Errorf("invalid form input type %q", f.Input)
	}

	return template.HTML(buf.String()), err
}
<|endoftext|>


<commit_before>/*
fperf allows you to build your performace tools easily

Three steps to create your own testcase

1. Create the "NewClient" function

	package demo

	import (
		"fmt"
		"github.com/shafreeck/fperf/client"
		"time"
	)

	type DemoClient struct{}

	func NewDemoClient(flag *client.FlagSet) client.Client {
		return &DemoClient{}
	}

2. Implement the UnaryClient or StreamClient

	func (c *DemoClient) Dial(addr string) error {
		fmt.Println("Dial to", addr)
		return nil
	}

	func (c *DemoClient) Request() error {
		time.Sleep(100 * time.Millisecond)
		return nil
	}

3. Register to fperf

	func init() {
		client.Register("demo", NewDemoClient, "This is a demo client discription")
	}

Run the buildin testcase

http is a simple builtin testcase to benchmark http servers

	fperf -cpu 8 -connection 10 http http://example.com
*/
package main

import (
	"flag"
	"fmt"
	db "github.com/influxdata/influxdb/client/v2"
	"github.com/shafreeck/fperf/client"
	hist "github.com/shafreeck/fperf/stats"
	"golang.org/x/net/context"
	"log"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"sync"
	"syscall"
	"time"
)

type setting struct {
	Connection int
	Stream     int
	Goroutine  int
	CPU        int
	Burst      int
	N          int //number of requests
	Tick       time.Duration
	Address    string
	Send       bool
	Recv       bool
	Delay      time.Duration
	Async      bool
	Target     string
	CallType   string
	InfluxDB   string
}

type statistics struct {
	latencies []time.Duration
	histogram *hist.Histogram
}

//roundtrip will be used in async mode
//the sender and receiver will be in seperate goroutines
type roundtrip struct {
	start time.Time
	end   time.Time
}

//create the testcase clients, n is the number of clients, set by
//flag -connection
func createClients(n int, addr string) []client.Client {
	clients := make([]client.Client, n)
	for i := 0; i < n; i++ {
		cli := client.NewClient(s.Target)
		if cli == nil {
			log.Fatalf("Can not find client %q for benchmark\n", s.Target)
		}

		if err := cli.Dial(addr); err != nil {
			log.Fatalln(err)
		}
		clients[i] = cli
	}
	return clients
}

//create streams for every client. n is the number of streams per client
func createStreams(n int, clients []client.Client) []client.Stream {
	streams := make([]client.Stream, n*len(clients))
	for cur, cli := range clients {
		for i := 0; i < n; i++ {
			if cli, ok := cli.(client.StreamClient); ok {
				stream, err := cli.CreateStream(context.Background())
				if err != nil {
					log.Fatalf("StreamCall faile to create new stream, %v", err)
				}
				streams[cur*n+i] = stream
			} else {
				log.Fatalln(s.Target, " do not implement the client.StreamClient")
			}
		}
	}
	return streams
}

//run benchmark for stream clients, can be in async or sync mode
func benchmarkStream(n int, streams []client.Stream) {
	var wg sync.WaitGroup
	for _, stream := range streams {
		for i := 0; i < n; i++ {
			//Notice here. we must pass stream as a parameter because the varibale stream
			//would be changed after the goroutine created
			if s.Async {
				wg.Add(2)
				go func(stream client.Stream) { send(nil, stream); wg.Done() }(stream)
				go func(stream client.Stream) { recv(nil, stream); wg.Done() }(stream)
			} else {
				wg.Add(1)
				go func(stream client.Stream) { run(nil, stream); wg.Done() }(stream)
			}
		}
	}
	go statPrint()
	wg.Wait()
}

//run benchmark for unary clients
func benchmarkUnary(n int, clients []client.Client) {
	var wg sync.WaitGroup
	for _, cli := range clients {
		for i := 0; i < n; i++ {
			wg.Add(1)
			if cli, ok := cli.(client.UnaryClient); ok {
				go func(cli client.UnaryClient) { runUnary(nil, cli); wg.Done() }(cli)
			} else {
				log.Fatalln(s.Target, " does not implement the client.UnaryClient")
			}
		}
	}
	go statPrint()
	wg.Wait()
}

func runUnary(done <-chan int, cli client.UnaryClient) {
	for i := 0; s.N == 0 || i < s.N; i++ {
		//select {
		//case <-done:
		//	log.Println("run goroutine exit done")
		//	return
		//default:
		start := time.Now()
		if err := cli.Request(); err != nil {
			log.Println(err)
		}
		eplase := time.Since(start)
		stats.latencies = append(stats.latencies, eplase)
		stats.histogram.Add(int64(eplase))
		if s.Delay > 0 {
			time.Sleep(s.Delay)
		}
		//	}
	}
}
func run(done <-chan int, stream client.Stream) {
	for i := 0; s.N == 0 || i < s.N; i++ {
		select {
		case <-done:
			log.Println("run goroutine exit done")
			return
		default:
			start := time.Now()
			if s.Send {
				stream.DoSend()
			}
			if s.Recv {
				stream.DoRecv()
			}
			eplase := time.Since(start)
			stats.latencies = append(stats.latencies, eplase)
			stats.histogram.Add(int64(eplase))
			if s.Delay > 0 {
				time.Sleep(s.Delay)
			}
		}
	}
}

func send(done <-chan int, stream client.Stream) {
	timer := time.NewTimer(time.Second)
	for i := 0; s.N == 0 || i < s.N; i++ {
		select {
		case <-done:
			log.Println("send goroutine exit, done")
			return
		default:
			select {
			case rtts <- &roundtrip{start: time.Now()}:
				timer.Reset(time.Second)
			case <-timer.C:
				log.Println("blocked on send rtts")
			}

			if burst != nil {
				select {
				case burst <- 0:
					timer.Reset(time.Second)
				case <-timer.C:
					log.Println("blocked on send burst chan")
				}
			}

			stream.DoSend()
			if s.Delay > 0 {
				time.Sleep(s.Delay)
			}
		}
	}
}
func recv(done <-chan int, stream client.Stream) {
	timer := time.NewTimer(time.Second)
	for i := 0; s.N == 0 || i < s.N; i++ {
		select {
		case <-done:
			log.Println("recv goroutine exit, done")
			return
		default:
			if burst != nil {
				select {
				case <-burst:
					timer.Reset(time.Second)
				case <-timer.C:
					log.Println("blocked on recv burst chan")
				}
			}
			err := stream.DoRecv()
			if err != nil {
				log.Println("recv goroutine exit", err)
				return
			}
			select {
			case rtt := <-rtts:
				timer.Reset(time.Second)
				eplase := time.Since(rtt.start)
				stats.latencies = append(stats.latencies, eplase)
				stats.histogram.Add(int64(eplase))
			case <-timer.C:
				log.Println("blocked on recv rtts")
			}
			if s.Delay > 0 {
				time.Sleep(s.Delay)
			}
		}
	}
}

func statPrint() {
	tickc := time.Tick(s.Tick)
	var latencies []time.Duration
	total := int64(0)
	for {
		select {
		case <-tickc:
			latencies = stats.latencies
			stats.latencies = stats.latencies[:0]

			sum := time.Duration(0)
			for _, eplase := range latencies {
				total++
				sum += eplase
			}
			count := len(latencies)
			if count != 0 {
				log.Printf("latency %v qps %d total %v\n", sum/time.Duration(count), int64(float64(count)/float64(s.Tick)*float64(time.Second)), total)
				if influxdb != nil {
					bp, _ := db.NewBatchPoints(db.BatchPointsConfig{
						Database:  "fperf",
						Precision: "s",
					})
					tags := map[string]string{"latency": "latency", "qps": "qps"}
					fields := map[string]interface{}{
						"latency": float64(sum) / float64(count) / 1000.0,
						"qps":     int64(float64(count) / float64(s.Tick) * float64(time.Second)),
					}
					pt, err := db.NewPoint("benchmark", tags, fields, time.Now())
					if err != nil {
						log.Println("Error: ", err.Error())
					}
					bp.AddPoint(pt)

					// Write the batch
					err = influxdb.Write(bp)
					if err != nil {
						log.Println("Error: ", err.Error())
					}
				}
			} else {
				log.Printf("blocking...")
			}
		}
	}
}

var s setting
var stats statistics
var rtts = make(chan *roundtrip, 10*1024*1024)
var mutex sync.RWMutex
var burst chan int
var influxdb db.Client

func usage() {
	fmt.Printf("Usage: %v [options] <client>\noptions:\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Println("clients:")
	for name, desc := range client.AllClients() {
		fmt.Printf(" %s", name)
		if len(desc) > 0 {
			fmt.Printf("\t: %s", desc)
		}
		fmt.Println()
	}
}

func main() {
	flag.IntVar(&s.Connection, "connection", 1, "number of connection")
	flag.IntVar(&s.Stream, "stream", 1, "number of streams per connection")
	flag.IntVar(&s.Goroutine, "goroutine", 1, "number of goroutines per stream")
	flag.IntVar(&s.CPU, "cpu", 0, "set the GOMAXPROCS, use go default if 0")
	flag.IntVar(&s.Burst, "burst", 0, "burst a number of request, use with -async=true")
	flag.IntVar(&s.N, "N", 0, "number of request per goroutine")
	flag.BoolVar(&s.Send, "send", true, "perform send action")
	flag.BoolVar(&s.Recv, "recv", true, "perform recv action")
	flag.DurationVar(&s.Delay, "delay", 0, "wait delay time before send the next request")
	flag.DurationVar(&s.Tick, "tick", 2*time.Second, "interval between statistics")
	flag.StringVar(&s.Address, "server", "127.0.0.1:8804", "address of the target server")
	flag.BoolVar(&s.Async, "async", false, "send and recv in seperate goroutines")
	flag.StringVar(&s.CallType, "type", "auto", "set the call type:unary, stream or auto. default is auto")
	flag.StringVar(&s.InfluxDB, "influxdb", "", "writing stats to influxdb, specify the address in this option")
	flag.Usage = usage
	flag.Parse()

	s.Target = flag.Arg(0)
	if len(s.Target) == 0 {
		flag.Usage()
		return
	}
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		_ = <-c
		stats.histogram.Value().Print(os.Stdout)
		os.Exit(0)
	}()

	runtime.GOMAXPROCS(s.CPU)
	go func() {
		runtime.SetBlockProfileRate(1)
		log.Println(http.ListenAndServe(":6060", nil))
	}()

	if s.Burst > 0 {
		burst = make(chan int, s.Burst)
	}

	if len(s.InfluxDB) > 0 {
		c, err := db.NewHTTPClient(db.HTTPConfig{
			Addr: s.InfluxDB,
		})
		if err != nil {
			log.Fatalf("Error creating InfluxDB Client: %v", err.Error())
		}
		defer c.Close()
		q := db.NewQuery("CREATE DATABASE fperf", "", "")
		if response, err := c.Query(q); err == nil && response.Error() == nil {
			log.Println(response.Results)
		}
		influxdb = c
	}

	stats.latencies = make([]time.Duration, 0, 500000)
	histopt := hist.HistogramOptions{
		NumBuckets:         16,
		GrowthFactor:       1.8,
		SmallestBucketSize: 1000,
		MinValue:           10000,
	}
	stats.histogram = hist.NewHistogram(histopt)
	clients := createClients(s.Connection, s.Address)
	cli := clients[0]
	switch s.CallType {
	case "auto":
		switch cli.(type) {
		case client.StreamClient:
			streams := createStreams(s.Stream, clients)
			benchmarkStream(s.Goroutine, streams)
		case client.UnaryClient:
			benchmarkUnary(s.Goroutine, clients)
		}
	case "stream":
		streams := createStreams(s.Stream, clients)
		benchmarkStream(s.Goroutine, streams)
	case "unary":
		benchmarkUnary(s.Goroutine, clients)
	}
}
<commit_msg>Delay histogram to statPrint<commit_after>/*
fperf allows you to build your performace tools easily

Three steps to create your own testcase

1. Create the "NewClient" function

	package demo

	import (
		"fmt"
		"github.com/shafreeck/fperf/client"
		"time"
	)

	type DemoClient struct{}

	func NewDemoClient(flag *client.FlagSet) client.Client {
		return &DemoClient{}
	}

2. Implement the UnaryClient or StreamClient

	func (c *DemoClient) Dial(addr string) error {
		fmt.Println("Dial to", addr)
		return nil
	}

	func (c *DemoClient) Request() error {
		time.Sleep(100 * time.Millisecond)
		return nil
	}

3. Register to fperf

	func init() {
		client.Register("demo", NewDemoClient, "This is a demo client discription")
	}

Run the buildin testcase

http is a simple builtin testcase to benchmark http servers

	fperf -cpu 8 -connection 10 http http://example.com
*/
package main

import (
	"flag"
	"fmt"
	db "github.com/influxdata/influxdb/client/v2"
	"github.com/shafreeck/fperf/client"
	hist "github.com/shafreeck/fperf/stats"
	"golang.org/x/net/context"
	"log"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"sync"
	"syscall"
	"time"
)

type setting struct {
	Connection int
	Stream     int
	Goroutine  int
	CPU        int
	Burst      int
	N          int //number of requests
	Tick       time.Duration
	Address    string
	Send       bool
	Recv       bool
	Delay      time.Duration
	Async      bool
	Target     string
	CallType   string
	InfluxDB   string
}

type statistics struct {
	latencies []time.Duration
	histogram *hist.Histogram
}

//roundtrip will be used in async mode
//the sender and receiver will be in seperate goroutines
type roundtrip struct {
	start time.Time
	end   time.Time
}

//create the testcase clients, n is the number of clients, set by
//flag -connection
func createClients(n int, addr string) []client.Client {
	clients := make([]client.Client, n)
	for i := 0; i < n; i++ {
		cli := client.NewClient(s.Target)
		if cli == nil {
			log.Fatalf("Can not find client %q for benchmark\n", s.Target)
		}

		if err := cli.Dial(addr); err != nil {
			log.Fatalln(err)
		}
		clients[i] = cli
	}
	return clients
}

//create streams for every client. n is the number of streams per client
func createStreams(n int, clients []client.Client) []client.Stream {
	streams := make([]client.Stream, n*len(clients))
	for cur, cli := range clients {
		for i := 0; i < n; i++ {
			if cli, ok := cli.(client.StreamClient); ok {
				stream, err := cli.CreateStream(context.Background())
				if err != nil {
					log.Fatalf("StreamCall faile to create new stream, %v", err)
				}
				streams[cur*n+i] = stream
			} else {
				log.Fatalln(s.Target, " do not implement the client.StreamClient")
			}
		}
	}
	return streams
}

//run benchmark for stream clients, can be in async or sync mode
func benchmarkStream(n int, streams []client.Stream) {
	var wg sync.WaitGroup
	for _, stream := range streams {
		for i := 0; i < n; i++ {
			//Notice here. we must pass stream as a parameter because the varibale stream
			//would be changed after the goroutine created
			if s.Async {
				wg.Add(2)
				go func(stream client.Stream) { send(nil, stream); wg.Done() }(stream)
				go func(stream client.Stream) { recv(nil, stream); wg.Done() }(stream)
			} else {
				wg.Add(1)
				go func(stream client.Stream) { run(nil, stream); wg.Done() }(stream)
			}
		}
	}
	go statPrint()
	wg.Wait()
}

//run benchmark for unary clients
func benchmarkUnary(n int, clients []client.Client) {
	var wg sync.WaitGroup
	for _, cli := range clients {
		for i := 0; i < n; i++ {
			wg.Add(1)
			if cli, ok := cli.(client.UnaryClient); ok {
				go func(cli client.UnaryClient) { runUnary(nil, cli); wg.Done() }(cli)
			} else {
				log.Fatalln(s.Target, " does not implement the client.UnaryClient")
			}
		}
	}
	go statPrint()
	wg.Wait()
}

func runUnary(done <-chan int, cli client.UnaryClient) {
	for i := 0; s.N == 0 || i < s.N; i++ {
		//select {
		//case <-done:
		//	log.Println("run goroutine exit done")
		//	return
		//default:
		start := time.Now()
		if err := cli.Request(); err != nil {
			log.Println(err)
		}
		eplase := time.Since(start)
		stats.latencies = append(stats.latencies, eplase)
		if s.Delay > 0 {
			time.Sleep(s.Delay)
		}
		//	}
	}
}
func run(done <-chan int, stream client.Stream) {
	for i := 0; s.N == 0 || i < s.N; i++ {
		select {
		case <-done:
			log.Println("run goroutine exit done")
			return
		default:
			start := time.Now()
			if s.Send {
				stream.DoSend()
			}
			if s.Recv {
				stream.DoRecv()
			}
			eplase := time.Since(start)
			stats.latencies = append(stats.latencies, eplase)
			if s.Delay > 0 {
				time.Sleep(s.Delay)
			}
		}
	}
}

func send(done <-chan int, stream client.Stream) {
	timer := time.NewTimer(time.Second)
	for i := 0; s.N == 0 || i < s.N; i++ {
		select {
		case <-done:
			log.Println("send goroutine exit, done")
			return
		default:
			select {
			case rtts <- &roundtrip{start: time.Now()}:
				timer.Reset(time.Second)
			case <-timer.C:
				log.Println("blocked on send rtts")
			}

			if burst != nil {
				select {
				case burst <- 0:
					timer.Reset(time.Second)
				case <-timer.C:
					log.Println("blocked on send burst chan")
				}
			}

			stream.DoSend()
			if s.Delay > 0 {
				time.Sleep(s.Delay)
			}
		}
	}
}
func recv(done <-chan int, stream client.Stream) {
	timer := time.NewTimer(time.Second)
	for i := 0; s.N == 0 || i < s.N; i++ {
		select {
		case <-done:
			log.Println("recv goroutine exit, done")
			return
		default:
			if burst != nil {
				select {
				case <-burst:
					timer.Reset(time.Second)
				case <-timer.C:
					log.Println("blocked on recv burst chan")
				}
			}
			err := stream.DoRecv()
			if err != nil {
				log.Println("recv goroutine exit", err)
				return
			}
			select {
			case rtt := <-rtts:
				timer.Reset(time.Second)
				eplase := time.Since(rtt.start)
				stats.latencies = append(stats.latencies, eplase)
			case <-timer.C:
				log.Println("blocked on recv rtts")
			}
			if s.Delay > 0 {
				time.Sleep(s.Delay)
			}
		}
	}
}

func statPrint() {
	tickc := time.Tick(s.Tick)
	var latencies []time.Duration
	total := int64(0)
	for {
		select {
		case <-tickc:
			latencies = stats.latencies
			stats.latencies = stats.latencies[:0]

			sum := time.Duration(0)
			for _, eplase := range latencies {
				total++
				sum += eplase
				stats.histogram.Add(int64(eplase))
			}
			count := len(latencies)
			if count != 0 {
				log.Printf("latency %v qps %d total %v\n", sum/time.Duration(count), int64(float64(count)/float64(s.Tick)*float64(time.Second)), total)
				if influxdb != nil {
					bp, _ := db.NewBatchPoints(db.BatchPointsConfig{
						Database:  "fperf",
						Precision: "s",
					})
					tags := map[string]string{"latency": "latency", "qps": "qps"}
					fields := map[string]interface{}{
						"latency": float64(sum) / float64(count) / 1000.0,
						"qps":     int64(float64(count) / float64(s.Tick) * float64(time.Second)),
					}
					pt, err := db.NewPoint("benchmark", tags, fields, time.Now())
					if err != nil {
						log.Println("Error: ", err.Error())
					}
					bp.AddPoint(pt)

					// Write the batch
					err = influxdb.Write(bp)
					if err != nil {
						log.Println("Error: ", err.Error())
					}
				}
			} else {
				log.Printf("blocking...")
			}
		}
	}
}

var s setting
var stats statistics
var rtts = make(chan *roundtrip, 10*1024*1024)
var mutex sync.RWMutex
var burst chan int
var influxdb db.Client

func usage() {
	fmt.Printf("Usage: %v [options] <client>\noptions:\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Println("clients:")
	for name, desc := range client.AllClients() {
		fmt.Printf(" %s", name)
		if len(desc) > 0 {
			fmt.Printf("\t: %s", desc)
		}
		fmt.Println()
	}
}

func main() {
	flag.IntVar(&s.Connection, "connection", 1, "number of connection")
	flag.IntVar(&s.Stream, "stream", 1, "number of streams per connection")
	flag.IntVar(&s.Goroutine, "goroutine", 1, "number of goroutines per stream")
	flag.IntVar(&s.CPU, "cpu", 0, "set the GOMAXPROCS, use go default if 0")
	flag.IntVar(&s.Burst, "burst", 0, "burst a number of request, use with -async=true")
	flag.IntVar(&s.N, "N", 0, "number of request per goroutine")
	flag.BoolVar(&s.Send, "send", true, "perform send action")
	flag.BoolVar(&s.Recv, "recv", true, "perform recv action")
	flag.DurationVar(&s.Delay, "delay", 0, "wait delay time before send the next request")
	flag.DurationVar(&s.Tick, "tick", 2*time.Second, "interval between statistics")
	flag.StringVar(&s.Address, "server", "127.0.0.1:8804", "address of the target server")
	flag.BoolVar(&s.Async, "async", false, "send and recv in seperate goroutines")
	flag.StringVar(&s.CallType, "type", "auto", "set the call type:unary, stream or auto. default is auto")
	flag.StringVar(&s.InfluxDB, "influxdb", "", "writing stats to influxdb, specify the address in this option")
	flag.Usage = usage
	flag.Parse()

	s.Target = flag.Arg(0)
	if len(s.Target) == 0 {
		flag.Usage()
		return
	}
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		_ = <-c
		stats.histogram.Value().Print(os.Stdout)
		os.Exit(0)
	}()

	runtime.GOMAXPROCS(s.CPU)
	go func() {
		runtime.SetBlockProfileRate(1)
		log.Println(http.ListenAndServe(":6060", nil))
	}()

	if s.Burst > 0 {
		burst = make(chan int, s.Burst)
	}

	if len(s.InfluxDB) > 0 {
		c, err := db.NewHTTPClient(db.HTTPConfig{
			Addr: s.InfluxDB,
		})
		if err != nil {
			log.Fatalf("Error creating InfluxDB Client: %v", err.Error())
		}
		defer c.Close()
		q := db.NewQuery("CREATE DATABASE fperf", "", "")
		if response, err := c.Query(q); err == nil && response.Error() == nil {
			log.Println(response.Results)
		}
		influxdb = c
	}

	stats.latencies = make([]time.Duration, 0, 500000)
	histopt := hist.HistogramOptions{
		NumBuckets:         16,
		GrowthFactor:       1.8,
		SmallestBucketSize: 1000,
		MinValue:           10000,
	}
	stats.histogram = hist.NewHistogram(histopt)
	clients := createClients(s.Connection, s.Address)
	cli := clients[0]
	switch s.CallType {
	case "auto":
		switch cli.(type) {
		case client.StreamClient:
			streams := createStreams(s.Stream, clients)
			benchmarkStream(s.Goroutine, streams)
		case client.UnaryClient:
			benchmarkUnary(s.Goroutine, clients)
		}
	case "stream":
		streams := createStreams(s.Stream, clients)
		benchmarkStream(s.Goroutine, streams)
	case "unary":
		benchmarkUnary(s.Goroutine, clients)
	}
}
<|endoftext|>