<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2018 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

// Package testutils is indended only for use in tests, do not import in production code!
package testutils

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/andybalholm/brotli"
	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
	"github.com/loadimpact/k6/lib/netext"
	"github.com/loadimpact/k6/lib/netext/httpext"
	"github.com/mccutchen/go-httpbin/httpbin"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/http2"
)

// GetTLSClientConfig returns a TLS config that trusts the supplied
// httptest.Server certificate as well as all the system root certificates
func GetTLSClientConfig(t testing.TB, srv *httptest.Server) *tls.Config {
	var err error

	certs := x509.NewCertPool()

	if runtime.GOOS != "windows" {
		certs, err = x509.SystemCertPool()
		require.NoError(t, err)
	}

	for _, c := range srv.TLS.Certificates {
		roots, err := x509.ParseCertificates(c.Certificate[len(c.Certificate)-1])
		require.NoError(t, err)
		for _, root := range roots {
			certs.AddCert(root)
		}
	}
	return &tls.Config{
		RootCAs:            certs,
		InsecureSkipVerify: false,
		Renegotiation:      tls.RenegotiateFreelyAsClient,
	}
}

const httpDomain = "httpbin.local"

// We have to use example.com if we want a real HTTPS domain with a valid
// certificate because the default httptest certificate is for example.com:
// https://golang.org/src/net/http/internal/testcert.go?s=399:410#L10
const httpsDomain = "example.com"

// HTTPMultiBin can be used as a local alternative of httpbin.org. It offers both http and https servers, as well as real domains
type HTTPMultiBin struct {
	Mux             *http.ServeMux
	ServerHTTP      *httptest.Server
	ServerHTTPS     *httptest.Server
	Replacer        *strings.Replacer
	TLSClientConfig *tls.Config
	Dialer          *netext.Dialer
	HTTPTransport   *http.Transport
	Context         context.Context
	Cleanup         func()
}

type jsonBody struct {
	Header      http.Header `json:"headers"`
	Compression string      `json:"compression"`
}

func getWebsocketEchoHandler(t testing.TB) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		t.Logf("[%p %s] Upgrading to websocket connection...", req, req.URL)
		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
		require.NoError(t, err)
		t.Logf("[%p %s] Upgraded...", req, req.URL)

		for {
			mt, message, err := conn.ReadMessage()
			select {
			case <-req.Context().Done():
				return
			default:
			}
			t.Logf("[%p %s] Read message '%s' of type %d (error '%v')", req, req.URL, message, mt, err)
			if err != nil {
				break
			}
			err = conn.WriteMessage(mt, message)
			select {
			case <-req.Context().Done():
				return
			default:
			}

			t.Logf("[%p %s] Wrote back message '%s' of type %d and closed the connection", req, req.URL, message, mt)

			if err != nil {
				break
			}
		}
	})
}

func getWebsocketCloserHandler(t testing.TB) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
		if !assert.NoError(t, err) {
			return
		}
		assert.NoError(t, conn.Close())
	})
}

func writeJSON(w io.Writer, v interface{}) error {
	e := json.NewEncoder(w)
	e.SetIndent("", "  ")
	return errors.Wrap(e.Encode(v), "failed to encode JSON")
}

func getEncodedHandler(t testing.TB, compressionType httpext.CompressionType) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		var (
			encoding string
			err      error
			encw     io.WriteCloser
		)

		switch compressionType {
		case httpext.CompressionTypeBr:
			encw = brotli.NewWriter(rw)
			encoding = "br"
		case httpext.CompressionTypeZstd:
			encw, _ = zstd.NewWriter(rw)
			encoding = "zstd"
		}

		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}
		err = writeJSON(encw, data)
		_ = encw.Close()
		if !assert.NoError(t, err) {
			return
		}
	})
}

func getZstdBrHandler(t testing.TB) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		encoding := "zstd, br"
		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}

		bw := brotli.NewWriter(rw)
		zw, _ := zstd.NewWriter(bw)
		defer func() {
			_ = zw.Close()
			_ = bw.Close()
		}()

		require.NoError(t, writeJSON(zw, data))
	})
}

// NewHTTPMultiBin returns a fully configured and running HTTPMultiBin
func NewHTTPMultiBin(t testing.TB) *HTTPMultiBin {
	// Create a http.ServeMux and set the httpbin handler as the default
	mux := http.NewServeMux()
	mux.Handle("/brotli", getEncodedHandler(t, httpext.CompressionTypeBr))
	mux.Handle("/ws-echo", getWebsocketEchoHandler(t))
	mux.Handle("/ws-close", getWebsocketCloserHandler(t))
	mux.Handle("/zstd", getEncodedHandler(t, httpext.CompressionTypeZstd))
	mux.Handle("/zstd-br", getZstdBrHandler(t))
	mux.Handle("/", httpbin.New().Handler())

	// Initialize the HTTP server and get its details
	httpSrv := httptest.NewServer(mux)
	httpURL, err := url.Parse(httpSrv.URL)
	require.NoError(t, err)
	httpIP := net.ParseIP(httpURL.Hostname())
	require.NotNil(t, httpIP)

	// Initialize the HTTPS server and get its details and tls config
	httpsSrv := httptest.NewTLSServer(mux)
	httpsURL, err := url.Parse(httpsSrv.URL)
	require.NoError(t, err)
	httpsIP := net.ParseIP(httpsURL.Hostname())
	require.NotNil(t, httpsIP)
	tlsConfig := GetTLSClientConfig(t, httpsSrv)

	// Set up the dialer with shorter timeouts and the custom domains
	dialer := netext.NewDialer(net.Dialer{
		Timeout:   2 * time.Second,
		KeepAlive: 10 * time.Second,
		DualStack: true,
	})
	dialer.Hosts = map[string]net.IP{
		httpDomain:  httpIP,
		httpsDomain: httpsIP,
	}

	// Pre-configure the HTTP client transport with the dialer and TLS config (incl. HTTP2 support)
	transport := &http.Transport{
		DialContext:     dialer.DialContext,
		TLSClientConfig: tlsConfig,
	}
	require.NoError(t, http2.ConfigureTransport(transport))

	ctx, ctxCancel := context.WithCancel(context.Background())
	return &HTTPMultiBin{
		Mux:         mux,
		ServerHTTP:  httpSrv,
		ServerHTTPS: httpsSrv,
		Replacer: strings.NewReplacer(
			"HTTPBIN_IP_URL", httpSrv.URL,
			"HTTPBIN_DOMAIN", httpDomain,
			"HTTPBIN_URL", fmt.Sprintf("http://%s:%s", httpDomain, httpURL.Port()),
			"WSBIN_URL", fmt.Sprintf("ws://%s:%s", httpDomain, httpURL.Port()),
			"HTTPBIN_IP", httpIP.String(),
			"HTTPBIN_PORT", httpURL.Port(),
			"HTTPSBIN_IP_URL", httpsSrv.URL,
			"HTTPSBIN_DOMAIN", httpsDomain,
			"HTTPSBIN_URL", fmt.Sprintf("https://%s:%s", httpsDomain, httpsURL.Port()),
			"WSSBIN_URL", fmt.Sprintf("wss://%s:%s", httpsDomain, httpsURL.Port()),
			"HTTPSBIN_IP", httpsIP.String(),
			"HTTPSBIN_PORT", httpsURL.Port(),
		),
		TLSClientConfig: tlsConfig,
		Dialer:          dialer,
		HTTPTransport:   transport,
		Context:         ctx,
		Cleanup: func() {
			httpsSrv.Close()
			httpSrv.Close()
			ctxCancel()
		},
	}
}
<commit_msg>Don't use testing.TB in ws handlers<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2018 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

// Package testutils is indended only for use in tests, do not import in production code!
package testutils

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/andybalholm/brotli"
	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
	"github.com/loadimpact/k6/lib/netext"
	"github.com/loadimpact/k6/lib/netext/httpext"
	"github.com/mccutchen/go-httpbin/httpbin"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/http2"
)

// GetTLSClientConfig returns a TLS config that trusts the supplied
// httptest.Server certificate as well as all the system root certificates
func GetTLSClientConfig(t testing.TB, srv *httptest.Server) *tls.Config {
	var err error

	certs := x509.NewCertPool()

	if runtime.GOOS != "windows" {
		certs, err = x509.SystemCertPool()
		require.NoError(t, err)
	}

	for _, c := range srv.TLS.Certificates {
		roots, err := x509.ParseCertificates(c.Certificate[len(c.Certificate)-1])
		require.NoError(t, err)
		for _, root := range roots {
			certs.AddCert(root)
		}
	}
	return &tls.Config{
		RootCAs:            certs,
		InsecureSkipVerify: false,
		Renegotiation:      tls.RenegotiateFreelyAsClient,
	}
}

const httpDomain = "httpbin.local"

// We have to use example.com if we want a real HTTPS domain with a valid
// certificate because the default httptest certificate is for example.com:
// https://golang.org/src/net/http/internal/testcert.go?s=399:410#L10
const httpsDomain = "example.com"

// HTTPMultiBin can be used as a local alternative of httpbin.org. It offers both http and https servers, as well as real domains
type HTTPMultiBin struct {
	Mux             *http.ServeMux
	ServerHTTP      *httptest.Server
	ServerHTTPS     *httptest.Server
	Replacer        *strings.Replacer
	TLSClientConfig *tls.Config
	Dialer          *netext.Dialer
	HTTPTransport   *http.Transport
	Context         context.Context
	Cleanup         func()
}

type jsonBody struct {
	Header      http.Header `json:"headers"`
	Compression string      `json:"compression"`
}

func websocketEchoHandler(w http.ResponseWriter, req *http.Request) {
	conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
	if err != nil {
		return
	}

	for {
		mt, message, err := conn.ReadMessage()
		if err != nil {
			break
		}
		err = conn.WriteMessage(mt, message)
		if err != nil {
			break
		}
	}
}

func websocketCloserHandler(w http.ResponseWriter, req *http.Request) {
	conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
	if err != nil {
		return
	}
	_ = conn.Close()
}

func writeJSON(w io.Writer, v interface{}) error {
	e := json.NewEncoder(w)
	e.SetIndent("", "  ")
	return errors.Wrap(e.Encode(v), "failed to encode JSON")
}

func getEncodedHandler(t testing.TB, compressionType httpext.CompressionType) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		var (
			encoding string
			err      error
			encw     io.WriteCloser
		)

		switch compressionType {
		case httpext.CompressionTypeBr:
			encw = brotli.NewWriter(rw)
			encoding = "br"
		case httpext.CompressionTypeZstd:
			encw, _ = zstd.NewWriter(rw)
			encoding = "zstd"
		}

		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}
		err = writeJSON(encw, data)
		_ = encw.Close()
		if !assert.NoError(t, err) {
			return
		}
	})
}

func getZstdBrHandler(t testing.TB) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		encoding := "zstd, br"
		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}

		bw := brotli.NewWriter(rw)
		zw, _ := zstd.NewWriter(bw)
		defer func() {
			_ = zw.Close()
			_ = bw.Close()
		}()

		require.NoError(t, writeJSON(zw, data))
	})
}

// NewHTTPMultiBin returns a fully configured and running HTTPMultiBin
func NewHTTPMultiBin(t testing.TB) *HTTPMultiBin {
	// Create a http.ServeMux and set the httpbin handler as the default
	mux := http.NewServeMux()
	mux.Handle("/brotli", getEncodedHandler(t, httpext.CompressionTypeBr))
	mux.HandleFunc("/ws-echo", websocketEchoHandler)
	mux.HandleFunc("/ws-close", websocketCloserHandler)
	mux.Handle("/zstd", getEncodedHandler(t, httpext.CompressionTypeZstd))
	mux.Handle("/zstd-br", getZstdBrHandler(t))
	mux.Handle("/", httpbin.New().Handler())

	// Initialize the HTTP server and get its details
	httpSrv := httptest.NewServer(mux)
	httpURL, err := url.Parse(httpSrv.URL)
	require.NoError(t, err)
	httpIP := net.ParseIP(httpURL.Hostname())
	require.NotNil(t, httpIP)

	// Initialize the HTTPS server and get its details and tls config
	httpsSrv := httptest.NewTLSServer(mux)
	httpsURL, err := url.Parse(httpsSrv.URL)
	require.NoError(t, err)
	httpsIP := net.ParseIP(httpsURL.Hostname())
	require.NotNil(t, httpsIP)
	tlsConfig := GetTLSClientConfig(t, httpsSrv)

	// Set up the dialer with shorter timeouts and the custom domains
	dialer := netext.NewDialer(net.Dialer{
		Timeout:   2 * time.Second,
		KeepAlive: 10 * time.Second,
		DualStack: true,
	})
	dialer.Hosts = map[string]net.IP{
		httpDomain:  httpIP,
		httpsDomain: httpsIP,
	}

	// Pre-configure the HTTP client transport with the dialer and TLS config (incl. HTTP2 support)
	transport := &http.Transport{
		DialContext:     dialer.DialContext,
		TLSClientConfig: tlsConfig,
	}
	require.NoError(t, http2.ConfigureTransport(transport))

	ctx, ctxCancel := context.WithCancel(context.Background())
	return &HTTPMultiBin{
		Mux:         mux,
		ServerHTTP:  httpSrv,
		ServerHTTPS: httpsSrv,
		Replacer: strings.NewReplacer(
			"HTTPBIN_IP_URL", httpSrv.URL,
			"HTTPBIN_DOMAIN", httpDomain,
			"HTTPBIN_URL", fmt.Sprintf("http://%s:%s", httpDomain, httpURL.Port()),
			"WSBIN_URL", fmt.Sprintf("ws://%s:%s", httpDomain, httpURL.Port()),
			"HTTPBIN_IP", httpIP.String(),
			"HTTPBIN_PORT", httpURL.Port(),
			"HTTPSBIN_IP_URL", httpsSrv.URL,
			"HTTPSBIN_DOMAIN", httpsDomain,
			"HTTPSBIN_URL", fmt.Sprintf("https://%s:%s", httpsDomain, httpsURL.Port()),
			"WSSBIN_URL", fmt.Sprintf("wss://%s:%s", httpsDomain, httpsURL.Port()),
			"HTTPSBIN_IP", httpsIP.String(),
			"HTTPSBIN_PORT", httpsURL.Port(),
		),
		TLSClientConfig: tlsConfig,
		Dialer:          dialer,
		HTTPTransport:   transport,
		Context:         ctx,
		Cleanup: func() {
			httpsSrv.Close()
			httpSrv.Close()
			ctxCancel()
		},
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/pschanely/scoredb"
	"log"
	"os"
	"path"
	"runtime"
	"strings"
	"time"
)

func MakeStandardDb(dataDir string, numShards int) (*scoredb.BaseDb, error) {
	var shards []scoredb.StreamingDb

	if scoredb.Exists(dataDir) && scoredb.Exists(path.Join(dataDir, "shard.0")) {
		i := 0
		shards = make([]scoredb.StreamingDb, 0, numShards)
		for {
			shardDir := path.Join(dataDir, fmt.Sprintf("shard.%d", i))
			if scoredb.Exists(shardDir) {
				shards = append(shards, scoredb.BaseStreamingDb{scoredb.NewFsScoreDb(shardDir)})
			} else {
				break
			}
			i += 1
		}
	} else {
		shards = make([]scoredb.StreamingDb, numShards)
		for i := range shards {
			shardDir := path.Join(dataDir, fmt.Sprintf("shard.%d", i))
			shards[i] = scoredb.BaseStreamingDb{scoredb.NewFsScoreDb(shardDir)}
		}
	}
	idDb, err := scoredb.NewBoltIdDb(path.Join(dataDir, "iddb"))
	if err != nil {
		return nil, err
	}
	return &scoredb.BaseDb{
		StreamingDb: scoredb.ShardedDb{
			Shards: shards,
		},
		IdDb: idDb,
	}, nil
}

func watchDir(db *scoredb.MigratableDb, baseDir string, namePrefix string) {
	log.Printf("Watching for databases at %s%s*\n", baseDir, namePrefix)
	var lastName = ""
	for {
		dir, err := os.Open(baseDir)
		var fileInfos []os.FileInfo
		if err == nil {
			fileInfos, err = dir.Readdir(0)
		}
		if err != nil {
			log.Printf("Unable to read %v: %v\n", dir, err)
			time.Sleep(55 * time.Second)
		} else {
			var newDbName = ""
			for _, fileInfo := range fileInfos {
				name := fileInfo.Name()
				if strings.HasPrefix(name, namePrefix) {
					if name > newDbName {
						newDbName = name
					}
				}
			}
			if newDbName > lastName {
				fmt.Printf("Detected database at %s%s\n", baseDir, newDbName)
				fullDbName := path.Join(baseDir, newDbName)
				newDb, err := MakeStandardDb(fullDbName, 1)
				if err != nil {
					log.Printf("Unable to load database at %s%s (%v); ignoring\n", dir, fullDbName, err)
				} else {
					fmt.Printf("The database at %s%s is live\n", baseDir, fullDbName)
					db.Current = newDb
					lastName = newDbName
				}
			}
		}
		time.Sleep(5 * time.Second)
	}
}

func SetupDirLoading(databaseDir string) *scoredb.MigratableDb {
	migratable := scoredb.MigratableDb{Current: nil}
	baseDir, namePrefix := path.Split(databaseDir)
	fmt.Printf("Watching for new databases named %s* in %s\n", namePrefix, baseDir)
	go watchDir(&migratable, baseDir, namePrefix)
	return &migratable
}

func main() {

	serveCommand := flag.NewFlagSet("serve", flag.ExitOnError)
	servePort := serveCommand.Int("port", 11625, "listening port in http mode, defaults to 11625")
	serveIntf := serveCommand.String("interface", "", "network interface to listen on in http mode, defaults to empty string (any interface)")
	serveDataDir := serveCommand.String("datadir", "./data", "Storage directory for database")
	serveNumShards := serveCommand.Int("numshards", 4, "Number of shards")
	serveReadOnly := serveCommand.Bool("readonly", false, "Only allow GET requests")
	serveAutoMigrate := serveCommand.Bool("automigrate", false, "When new directories appear matching <datadir>*, atomically swap in the database at that directory. (lexigraphically last)")

	loadCommand := flag.NewFlagSet("load", flag.ExitOnError)
	loadDataDir := loadCommand.String("datadir", "./data", "Storage directory for database")
	loadNumShards := loadCommand.Int("numshards", 4, "Number of shards (ignored if db already exists)")

	benchCommand := flag.NewFlagSet("benchmark", flag.ExitOnError)
	benchCsvFilename := benchCommand.String("csv", "", "csv filename of census data")
	benchMaxRecords := benchCommand.Int64("maxrecords", 1000*1000, "Maximum size of database to benchmark (in # of records)")
	benchCsvOutput := benchCommand.String("out", "output.csv", "csv of performance data to output")
	benchEsUrl := benchCommand.String("esurl", "http://localhost:9200/", "URL of elasticsearch instance")
	benchEsIndex := benchCommand.String("esindex", "benchmark_scoredb", "Index name to use for elasticsearch")
	benchFsDataDir := benchCommand.String("fsdatadir", "./benchmark_data", "Storage directory for native scoredb database")

	/*
		for cmd := range([]*flag.FlagSet{serveCommand, benchCommand}) {
			// common args here
		}
	*/

	if len(os.Args) <= 1 {
		fmt.Println("usage: scoredb <command> [<args>]")
		fmt.Println("Commands:")
		fmt.Println(" serve      Run a scoredb server")
		fmt.Println(" load       Load json lines from stdin")
		fmt.Println(" benchmark  Run performance benchmarks")
		fmt.Println("For more help, run scoredb <command> -h")
		os.Exit(1)
	}
	var db scoredb.Db
	var err error
	switch os.Args[1] {
	case "serve":
		serveCommand.Parse(os.Args[2:])
		if *serveAutoMigrate {
			db = SetupDirLoading(*serveDataDir)
		} else {
			db, err = MakeStandardDb(*serveDataDir, *serveNumShards)
			if err != nil {
				log.Fatalf("Failed to initialize database at %v: %v\n", *serveDataDir, err)
			}
		}
		addr := fmt.Sprintf("%s:%d", *serveIntf, *servePort)
		fmt.Printf("Serving on %s\n", addr)
		log.Fatal(scoredb.ServeHttp(addr, db, *serveReadOnly))
	case "load":
		loadCommand.Parse(os.Args[2:])
		db, err := MakeStandardDb(*loadDataDir, *loadNumShards)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to initialize database at %v: %v\n", *loadDataDir, err))
		}
		scanner := bufio.NewScanner(os.Stdin)
		batchSize := 200
		batchIndex := 0
		var batch = make([]scoredb.Record, batchSize)
		for scanner.Scan() {
			record := scoredb.Record{}
			line := scanner.Bytes()
			json.Unmarshal(line, &record)
			batch[batchIndex] = record
			batchIndex += 1
			if batchIndex >= batchSize {
				db.BulkIndex(batch)
				batchIndex = 0
				batch = make([]scoredb.Record, batchSize)
			}
		}
		if batchIndex > 0 {
			db.BulkIndex(batch[:batchIndex])
		}
	case "benchmark":
		outputFd, err := os.Create(*benchCsvOutput)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to output output csv file at %v: %v\n", *benchCsvOutput, err))
		}

		runtime.GOMAXPROCS(runtime.NumCPU())
		benchCommand.Parse(os.Args[2:])
		esDb := &scoredb.EsScoreDb{BaseURL: *benchEsUrl, Index: *benchEsIndex}
		fsDb, err := MakeStandardDb(*benchFsDataDir, 4)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to initialize database at %v: %v\n", *benchFsDataDir, err))
		}
		if !scoredb.Exists(*benchCsvFilename) {
			log.Fatal(fmt.Sprintf("Cannot find source csv data file at '%s'", *benchCsvFilename))
		}

		fmt.Printf("Running es benchmarks\n")
		esDb.DeleteIndex()
		esDb.CreateIndex()
		counts, esIndexTimes, esQueryTimes, err := scoredb.RunBenchmark(esDb, *benchCsvFilename, *benchMaxRecords)
		//esDb.DeleteIndex()
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to run es benchmark: %v\n", err))
		}

		fmt.Printf("Running native benchmarks\n")
		_, fsIndexTimes, fsQueryTimes, err := scoredb.RunBenchmark(fsDb, *benchCsvFilename, *benchMaxRecords)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to run native benchmark: %v\n", err))
		}

		fmt.Fprintf(outputFd, "records,es_index,native_index,es_query_1,native_query_1,es_query_2,native_query_2\n")
		for idx := 0; idx < len(esIndexTimes); idx++ {
			fmt.Fprintf(outputFd, "%v,%v,%v", counts[idx], esIndexTimes[idx], fsIndexTimes[idx])
			for idx2 := 0; idx2 < len(esQueryTimes[idx]); idx2++ {
				fmt.Fprintf(outputFd, ",%v,%v", esQueryTimes[idx][idx2], fsQueryTimes[idx][idx2])
			}
			fmt.Fprintf(outputFd, "\n")
		}
		outputFd.Close()
	default:
		fmt.Printf("%q is not valid command.\n", os.Args[1])
		os.Exit(2)
	}
}
<commit_msg>close watch directory fd when done<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/pschanely/scoredb"
	"log"
	"os"
	"path"
	"runtime"
	"strings"
	"time"
)

func MakeStandardDb(dataDir string, numShards int) (*scoredb.BaseDb, error) {
	var shards []scoredb.StreamingDb

	if scoredb.Exists(dataDir) && scoredb.Exists(path.Join(dataDir, "shard.0")) {
		i := 0
		shards = make([]scoredb.StreamingDb, 0, numShards)
		for {
			shardDir := path.Join(dataDir, fmt.Sprintf("shard.%d", i))
			if scoredb.Exists(shardDir) {
				shards = append(shards, scoredb.BaseStreamingDb{scoredb.NewFsScoreDb(shardDir)})
			} else {
				break
			}
			i += 1
		}
	} else {
		shards = make([]scoredb.StreamingDb, numShards)
		for i := range shards {
			shardDir := path.Join(dataDir, fmt.Sprintf("shard.%d", i))
			shards[i] = scoredb.BaseStreamingDb{scoredb.NewFsScoreDb(shardDir)}
		}
	}
	idDb, err := scoredb.NewBoltIdDb(path.Join(dataDir, "iddb"))
	if err != nil {
		return nil, err
	}
	return &scoredb.BaseDb{
		StreamingDb: scoredb.ShardedDb{
			Shards: shards,
		},
		IdDb: idDb,
	}, nil
}

func watchDir(db *scoredb.MigratableDb, baseDir string, namePrefix string) {
	log.Printf("Watching for databases at %s%s*\n", baseDir, namePrefix)
	var lastName = ""
	for {
		dir, err := os.Open(baseDir)
		var fileInfos []os.FileInfo
		if err == nil {
			fileInfos, err = dir.Readdir(0)
			dir.Close()
		}
		if err != nil {
			log.Printf("Unable to read %v: %v\n", dir, err)
			time.Sleep(55 * time.Second)
		} else {
			var newDbName = ""
			for _, fileInfo := range fileInfos {
				name := fileInfo.Name()
				if strings.HasPrefix(name, namePrefix) {
					if name > newDbName {
						newDbName = name
					}
				}
			}
			if newDbName > lastName {
				fmt.Printf("Detected database at %s%s\n", baseDir, newDbName)
				fullDbName := path.Join(baseDir, newDbName)
				newDb, err := MakeStandardDb(fullDbName, 1)
				if err != nil {
					log.Printf("Unable to load database at %s%s (%v); ignoring\n", dir, fullDbName, err)
				} else {
					fmt.Printf("The database at %s%s is live\n", baseDir, fullDbName)
					db.Current = newDb
					lastName = newDbName
				}
			}
		}
		time.Sleep(10 * time.Second)
	}
}

func SetupDirLoading(databaseDir string) *scoredb.MigratableDb {
	migratable := scoredb.MigratableDb{Current: nil}
	baseDir, namePrefix := path.Split(databaseDir)
	fmt.Printf("Watching for new databases named %s* in %s\n", namePrefix, baseDir)
	go watchDir(&migratable, baseDir, namePrefix)
	return &migratable
}

func main() {

	serveCommand := flag.NewFlagSet("serve", flag.ExitOnError)
	servePort := serveCommand.Int("port", 11625, "listening port in http mode, defaults to 11625")
	serveIntf := serveCommand.String("interface", "", "network interface to listen on in http mode, defaults to empty string (any interface)")
	serveDataDir := serveCommand.String("datadir", "./data", "Storage directory for database")
	serveNumShards := serveCommand.Int("numshards", 4, "Number of shards")
	serveReadOnly := serveCommand.Bool("readonly", false, "Only allow GET requests")
	serveAutoMigrate := serveCommand.Bool("automigrate", false, "When new directories appear matching <datadir>*, atomically swap in the database at that directory. (lexigraphically last)")

	loadCommand := flag.NewFlagSet("load", flag.ExitOnError)
	loadDataDir := loadCommand.String("datadir", "./data", "Storage directory for database")
	loadNumShards := loadCommand.Int("numshards", 4, "Number of shards (ignored if db already exists)")

	benchCommand := flag.NewFlagSet("benchmark", flag.ExitOnError)
	benchCsvFilename := benchCommand.String("csv", "", "csv filename of census data")
	benchMaxRecords := benchCommand.Int64("maxrecords", 1000*1000, "Maximum size of database to benchmark (in # of records)")
	benchCsvOutput := benchCommand.String("out", "output.csv", "csv of performance data to output")
	benchEsUrl := benchCommand.String("esurl", "http://localhost:9200/", "URL of elasticsearch instance")
	benchEsIndex := benchCommand.String("esindex", "benchmark_scoredb", "Index name to use for elasticsearch")
	benchFsDataDir := benchCommand.String("fsdatadir", "./benchmark_data", "Storage directory for native scoredb database")

	/*
		for cmd := range([]*flag.FlagSet{serveCommand, benchCommand}) {
			// common args here
		}
	*/

	if len(os.Args) <= 1 {
		fmt.Println("usage: scoredb <command> [<args>]")
		fmt.Println("Commands:")
		fmt.Println(" serve      Run a scoredb server")
		fmt.Println(" load       Load json lines from stdin")
		fmt.Println(" benchmark  Run performance benchmarks")
		fmt.Println("For more help, run scoredb <command> -h")
		os.Exit(1)
	}
	var db scoredb.Db
	var err error
	switch os.Args[1] {
	case "serve":
		serveCommand.Parse(os.Args[2:])
		if *serveAutoMigrate {
			db = SetupDirLoading(*serveDataDir)
		} else {
			db, err = MakeStandardDb(*serveDataDir, *serveNumShards)
			if err != nil {
				log.Fatalf("Failed to initialize database at %v: %v\n", *serveDataDir, err)
			}
		}
		addr := fmt.Sprintf("%s:%d", *serveIntf, *servePort)
		fmt.Printf("Serving on %s\n", addr)
		log.Fatal(scoredb.ServeHttp(addr, db, *serveReadOnly))
	case "load":
		loadCommand.Parse(os.Args[2:])
		db, err := MakeStandardDb(*loadDataDir, *loadNumShards)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to initialize database at %v: %v\n", *loadDataDir, err))
		}
		scanner := bufio.NewScanner(os.Stdin)
		batchSize := 200
		batchIndex := 0
		var batch = make([]scoredb.Record, batchSize)
		for scanner.Scan() {
			record := scoredb.Record{}
			line := scanner.Bytes()
			json.Unmarshal(line, &record)
			batch[batchIndex] = record
			batchIndex += 1
			if batchIndex >= batchSize {
				db.BulkIndex(batch)
				batchIndex = 0
				batch = make([]scoredb.Record, batchSize)
			}
		}
		if batchIndex > 0 {
			db.BulkIndex(batch[:batchIndex])
		}
	case "benchmark":
		outputFd, err := os.Create(*benchCsvOutput)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to output output csv file at %v: %v\n", *benchCsvOutput, err))
		}

		runtime.GOMAXPROCS(runtime.NumCPU())
		benchCommand.Parse(os.Args[2:])
		esDb := &scoredb.EsScoreDb{BaseURL: *benchEsUrl, Index: *benchEsIndex}
		fsDb, err := MakeStandardDb(*benchFsDataDir, 4)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to initialize database at %v: %v\n", *benchFsDataDir, err))
		}
		if !scoredb.Exists(*benchCsvFilename) {
			log.Fatal(fmt.Sprintf("Cannot find source csv data file at '%s'", *benchCsvFilename))
		}

		fmt.Printf("Running es benchmarks\n")
		esDb.DeleteIndex()
		esDb.CreateIndex()
		counts, esIndexTimes, esQueryTimes, err := scoredb.RunBenchmark(esDb, *benchCsvFilename, *benchMaxRecords)
		//esDb.DeleteIndex()
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to run es benchmark: %v\n", err))
		}

		fmt.Printf("Running native benchmarks\n")
		_, fsIndexTimes, fsQueryTimes, err := scoredb.RunBenchmark(fsDb, *benchCsvFilename, *benchMaxRecords)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to run native benchmark: %v\n", err))
		}

		fmt.Fprintf(outputFd, "records,es_index,native_index,es_query_1,native_query_1,es_query_2,native_query_2\n")
		for idx := 0; idx < len(esIndexTimes); idx++ {
			fmt.Fprintf(outputFd, "%v,%v,%v", counts[idx], esIndexTimes[idx], fsIndexTimes[idx])
			for idx2 := 0; idx2 < len(esQueryTimes[idx]); idx2++ {
				fmt.Fprintf(outputFd, ",%v,%v", esQueryTimes[idx][idx2], fsQueryTimes[idx][idx2])
			}
			fmt.Fprintf(outputFd, "\n")
		}
		outputFd.Close()
	default:
		fmt.Printf("%q is not valid command.\n", os.Args[1])
		os.Exit(2)
	}
}
<|endoftext|>


<commit_before>package rabbitpubsub

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/go-cloud/internal/pubsub"
	"github.com/google/go-cloud/internal/pubsub/driver"
	"github.com/streadway/amqp"
)

type topic struct {
	exchange string // the AMQP exchange
	conn     *amqp.Connection

	mu   sync.Mutex
	ch   *amqp.Channel            // AMQP channel used for all communication.
	pubc <-chan amqp.Confirmation // Go channel for server acks of publishes
}

// Values for the amqp client.
// See https://www.rabbitmq.com/amqp-0-9-1-reference.html.
const (
	// Many methods of the amqp client take a "no-wait" parameter, which
	// if true causes the client to return without waiting for a server
	// response. We always want to wait.
	wait = false

	// Always use the empty routing key. This driver expects to be used with topic
	// exchanges, which disregard the routing key.
	routingKey = ""

	// If the message can't be enqueued, return it to the sender rather than silently dropping it.
	mandatory = true

	// If there are no waiting consumers, enqueue the message instead of dropping it
	notImmediate = false
)

// OpenTopic returns a *pubsub.Topic corresponding to the named exchange. The
// exchange must have been previously created (for instance, by using
// amqp.Channel.ExchangeDeclare). For the model of Go Cloud Pub/Sub to make sense,
// the exchange should be a fanout exchange, although nothing in this package
// enforces that.
//
// OpenTopic uses the supplied amqp.Connection for all communication. It is the caller's
// responsibility to establish this connection before calling OpenTopic, and to close
// it when Close has been called on all Topics opened with it.
//
// The documentation of the amqp package recommends using separate connections for
// publishing and subscribing.
func OpenTopic(conn *amqp.Connection, name string) (*pubsub.Topic, error) {
	// TODO(jba): support context.Context
	t, err := newTopic(conn, name)
	if err != nil {
		return nil, err
	}
	return pubsub.NewTopic(t), nil

}

func newTopic(conn *amqp.Connection, name string) (*topic, error) {
	t := &topic{
		exchange: name,
		conn:     conn,
	}
	t.mu.Lock()
	defer t.mu.Unlock()

	if err := t.establishChannel(); err != nil {
		return nil, err
	}
	return t, nil
}

// establishChannel creates an AMQP channel if necessary. According to the amqp
// package docs, once an error is returned from the channel, it must be discarded and
// a new one created.
//
// Must be called with the t.mu held.
func (t *topic) establishChannel() error {
	// TODO(jba): support context.Context
	if t.ch != nil {
		// We already have a channel; nothing to do.
		return nil
	}
	// Create a new channel.
	ch, err := t.conn.Channel()
	if err != nil {
		return err
	}
	// Put the channel into a mode where confirmations are delivered for each publish.
	if err := ch.Confirm(wait); err != nil {
		return err
	}
	t.ch = ch
	// Get a Go channel which will hold acks from the server. The server
	// will send an ack for each published message to confirm that it was received.
	t.pubc = ch.NotifyPublish(make(chan amqp.Confirmation)) // NotifyPublish returns its arg
	return nil
}

// SendBatch implements driver.SendBatch.
func (t *topic) SendBatch(ctx context.Context, ms []*driver.Message) error {
	// It is simplest to allow only one SendBatch at a time. Allowing concurrent
	// calls to SendBatch would complicate the logic of receiving publish
	// confirmations and returns. We can go that route if performance warrants it.
	t.mu.Lock()
	defer t.mu.Unlock()

	if err := t.establishChannel(); err != nil {
		return err
	}

	// Read from the channel established with NotifyPublish.
	// Do so concurrently or we will deadlock with the Publish RPC.
	// (The amqp package docs recommend setting the capacity of the channel
	// to the number of messages to be published, but we can't do that because
	// we want to reuse the channel for all calls to SendBatch.)
	errc := make(chan error, 1)
	go func() {
		// Consume all the acknowledgments for the messages we are publishing.
		// Since this method holds the lock, we expect exactly as many acks as messages.
		//
		// This goroutine can safely access the mutex-protected fields t.ch and t.pubc
		// because its lifetime is within the lifetime of the lock held by SendBatch.
		//
		// TODO(jba): look at AMQP "returns" for more information about publish errors.
		// See https://godoc.org/github.com/streadway/amqp#Channel.NotifyReturn.
		ok := true
		for range ms {
			select {
			case <-ctx.Done():
				errc <- ctx.Err()
				return
			case conf, ok := <-t.pubc:
				if !ok {
					// t.pubc was closed
					errc <- errors.New("rabbitpubsub: publish listener closed unexpectedly")
					t.ch = nil // re-create the channel on next use
					return
				}
				if !conf.Ack {
					ok = false
				}
			}
		}
		if !ok {
			errc <- errors.New("rabbitpubsub: ack failed on publish")
		} else {
			errc <- nil
		}
	}()

	for _, m := range ms {
		pub := toPublishing(m)
		if err := t.ch.Publish(t.exchange, routingKey, mandatory, notImmediate, pub); err != nil {
			t.ch = nil // AMQP channel is broken after error
			return err
		}
	}
	return <-errc
}

// toPublishing converts a driver.Message to an amqp.Publishing.
func toPublishing(m *driver.Message) amqp.Publishing {
	h := amqp.Table{}
	for k, v := range m.Metadata {
		h[k] = v
	}
	return amqp.Publishing{
		Headers: h,
		Body:    m.Body,
	}
}

// IsRetryable implements driver.Topic.IsRetryable.
func (*topic) IsRetryable(error) bool {
	// TODO(jba): figure out what errors can be retried.
	return false
}

// OpenSubscription returns a *pubsub.Subscription corresponding to the named queue. The
// queue must have been previously created (for instance, by using
// amqp.Channel.QueueDeclare) and bound to an exchange.
//
// OpenSubscription uses the supplied amqp.Connection for all communication. It is
// the caller's responsibility to establish this connection before calling OpenSubscription
// and to close it when Close has been called on all Subscriptions opened with it.
//
// The documentation of the amqp package recommends using separate connections for
// publishing and subscribing.
func OpenSubscription(conn *amqp.Connection, name string) (*pubsub.Subscription, error) {
	// TODO(jba): support context.Context
	s, err := newSubscription(conn, name)
	if err != nil {
		return nil, err
	}
	return pubsub.NewSubscription(s), nil
}

type subscription struct {
	conn     *amqp.Connection
	queue    string // the AMQP queue name
	consumer string // the client-generated name for this particular subscriber

	mu   sync.Mutex
	ch   *amqp.Channel // AMQP channel used for all communication.
	delc <-chan amqp.Delivery
}

var nextConsumer int64 // atomic

func newSubscription(conn *amqp.Connection, name string) (*subscription, error) {
	s := &subscription{
		conn:     conn,
		queue:    name,
		consumer: fmt.Sprintf("c%d", atomic.AddInt64(&nextConsumer, 1)),
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := s.establishChannel(); err != nil {
		return nil, err
	}
	return s, nil
}

// Must be called with s.mu held.
func (s *subscription) establishChannel() error {
	// TODO(jba): support context.Context
	if s.ch != nil {
		// We already have a channel; nothing to do.
		return nil
	}
	// Create a new channel.
	ch, err := s.conn.Channel()
	if err != nil {
		return err
	}
	// Subscribe to messages from the queue.
	s.delc, err = ch.Consume(s.queue, s.consumer,
		false, // autoAck
		false, // exclusive
		false, // noLocal
		wait,
		nil) // args
	if err != nil {
		return err
	}
	s.ch = ch
	return nil
}

// ReceiveBatch implements driver.Subscription.ReceiveBatch.
func (s *subscription) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if err := s.establishChannel(); err != nil {
		return nil, err
	}

	// Get up to maxMessages waiting messages, but don't take too long.
	var ms []*driver.Message
	maxTime := time.After(50 * time.Millisecond)
	for {
		select {
		case <-ctx.Done():
			// Cancel the Consume.
			_ = s.ch.Cancel(s.consumer, wait) // ignore the error
			s.ch = nil
			return nil, ctx.Err()

		case d, ok := <-s.delc:
			if !ok { // channel closed
				s.ch = nil // re-establish the channel next time
				if len(ms) > 0 {
					return ms, nil
				}
				// TODO(jba): get more information for a better message
				return nil, errors.New("delivery channel closed")
			}
			ms = append(ms, toMessage(d))
			if len(ms) >= maxMessages {
				return ms, nil
			}

		case <-maxTime:
			// Timed out. Return whatever we have.
			return ms, nil
		}
	}
}

// toMessage converts an amqp.Delivery (a received message) to a driver.Message.
func toMessage(d amqp.Delivery) *driver.Message {
	// Delivery.Headers is a map[string]interface{}, so we have to
	// convert each value to a string.
	md := map[string]string{}
	for k, v := range d.Headers {
		md[k] = fmt.Sprint(v)
	}
	return &driver.Message{
		Body:     d.Body,
		AckID:    d.DeliveryTag,
		Metadata: md,
	}
}

// SendAcks implements driver.Subscription.SendAcks.
func (s *subscription) SendAcks(ctx context.Context, ackIDs []driver.AckID) error {
	// TODO(#853): consider a separate channel for acks, so ReceiveBatch and SendAcks don't
	// block each other.
	s.mu.Lock()
	defer s.mu.Unlock()

	if err := s.establishChannel(); err != nil {
		return err
	}

	// The Ack call doesn't wait for a response, so this loop should execute relatively
	// quickly.
	// It wouldn't help to make it concurrent, because Channel.Ack grabs a channel-wide mutex.
	// (We could consider using multiple channels if performance becomes an issue.)
	for _, id := range ackIDs {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		err := s.ch.Ack(id.(uint64), false) // multiple=false: acking only this ID
		if err != nil {
			s.ch = nil // re-establish channel after an error
			return err
		}
	}
	return nil
}

// IsRetryable implements driver.Subscription.IsRetryable.
func (*subscription) IsRetryable(error) bool {
	// TODO(jba): figure out what errors can be retried.
	return false
}
<commit_msg>s/notImmediate/immediate/ (#859)<commit_after>package rabbitpubsub

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/go-cloud/internal/pubsub"
	"github.com/google/go-cloud/internal/pubsub/driver"
	"github.com/streadway/amqp"
)

type topic struct {
	exchange string // the AMQP exchange
	conn     *amqp.Connection

	mu   sync.Mutex
	ch   *amqp.Channel            // AMQP channel used for all communication.
	pubc <-chan amqp.Confirmation // Go channel for server acks of publishes
}

// Values for the amqp client.
// See https://www.rabbitmq.com/amqp-0-9-1-reference.html.
const (
	// Many methods of the amqp client take a "no-wait" parameter, which
	// if true causes the client to return without waiting for a server
	// response. We always want to wait.
	wait = false

	// Always use the empty routing key. This driver expects to be used with topic
	// exchanges, which disregard the routing key.
	routingKey = ""

	// If the message can't be enqueued, return it to the sender rather than silently dropping it.
	mandatory = true

	// If there are no waiting consumers, enqueue the message instead of dropping it
	immediate = false
)

// OpenTopic returns a *pubsub.Topic corresponding to the named exchange. The
// exchange must have been previously created (for instance, by using
// amqp.Channel.ExchangeDeclare). For the model of Go Cloud Pub/Sub to make sense,
// the exchange should be a fanout exchange, although nothing in this package
// enforces that.
//
// OpenTopic uses the supplied amqp.Connection for all communication. It is the caller's
// responsibility to establish this connection before calling OpenTopic, and to close
// it when Close has been called on all Topics opened with it.
//
// The documentation of the amqp package recommends using separate connections for
// publishing and subscribing.
func OpenTopic(conn *amqp.Connection, name string) (*pubsub.Topic, error) {
	// TODO(jba): support context.Context
	t, err := newTopic(conn, name)
	if err != nil {
		return nil, err
	}
	return pubsub.NewTopic(t), nil

}

func newTopic(conn *amqp.Connection, name string) (*topic, error) {
	t := &topic{
		exchange: name,
		conn:     conn,
	}
	t.mu.Lock()
	defer t.mu.Unlock()

	if err := t.establishChannel(); err != nil {
		return nil, err
	}
	return t, nil
}

// establishChannel creates an AMQP channel if necessary. According to the amqp
// package docs, once an error is returned from the channel, it must be discarded and
// a new one created.
//
// Must be called with the t.mu held.
func (t *topic) establishChannel() error {
	// TODO(jba): support context.Context
	if t.ch != nil {
		// We already have a channel; nothing to do.
		return nil
	}
	// Create a new channel.
	ch, err := t.conn.Channel()
	if err != nil {
		return err
	}
	// Put the channel into a mode where confirmations are delivered for each publish.
	if err := ch.Confirm(wait); err != nil {
		return err
	}
	t.ch = ch
	// Get a Go channel which will hold acks from the server. The server
	// will send an ack for each published message to confirm that it was received.
	t.pubc = ch.NotifyPublish(make(chan amqp.Confirmation)) // NotifyPublish returns its arg
	return nil
}

// SendBatch implements driver.SendBatch.
func (t *topic) SendBatch(ctx context.Context, ms []*driver.Message) error {
	// It is simplest to allow only one SendBatch at a time. Allowing concurrent
	// calls to SendBatch would complicate the logic of receiving publish
	// confirmations and returns. We can go that route if performance warrants it.
	t.mu.Lock()
	defer t.mu.Unlock()

	if err := t.establishChannel(); err != nil {
		return err
	}

	// Read from the channel established with NotifyPublish.
	// Do so concurrently or we will deadlock with the Publish RPC.
	// (The amqp package docs recommend setting the capacity of the channel
	// to the number of messages to be published, but we can't do that because
	// we want to reuse the channel for all calls to SendBatch.)
	errc := make(chan error, 1)
	go func() {
		// Consume all the acknowledgments for the messages we are publishing.
		// Since this method holds the lock, we expect exactly as many acks as messages.
		//
		// This goroutine can safely access the mutex-protected fields t.ch and t.pubc
		// because its lifetime is within the lifetime of the lock held by SendBatch.
		//
		// TODO(jba): look at AMQP "returns" for more information about publish errors.
		// See https://godoc.org/github.com/streadway/amqp#Channel.NotifyReturn.
		ok := true
		for range ms {
			select {
			case <-ctx.Done():
				errc <- ctx.Err()
				return
			case conf, ok := <-t.pubc:
				if !ok {
					// t.pubc was closed
					errc <- errors.New("rabbitpubsub: publish listener closed unexpectedly")
					t.ch = nil // re-create the channel on next use
					return
				}
				if !conf.Ack {
					ok = false
				}
			}
		}
		if !ok {
			errc <- errors.New("rabbitpubsub: ack failed on publish")
		} else {
			errc <- nil
		}
	}()

	for _, m := range ms {
		pub := toPublishing(m)
		if err := t.ch.Publish(t.exchange, routingKey, mandatory, immediate, pub); err != nil {
			t.ch = nil // AMQP channel is broken after error
			return err
		}
	}
	return <-errc
}

// toPublishing converts a driver.Message to an amqp.Publishing.
func toPublishing(m *driver.Message) amqp.Publishing {
	h := amqp.Table{}
	for k, v := range m.Metadata {
		h[k] = v
	}
	return amqp.Publishing{
		Headers: h,
		Body:    m.Body,
	}
}

// IsRetryable implements driver.Topic.IsRetryable.
func (*topic) IsRetryable(error) bool {
	// TODO(jba): figure out what errors can be retried.
	return false
}

// OpenSubscription returns a *pubsub.Subscription corresponding to the named queue. The
// queue must have been previously created (for instance, by using
// amqp.Channel.QueueDeclare) and bound to an exchange.
//
// OpenSubscription uses the supplied amqp.Connection for all communication. It is
// the caller's responsibility to establish this connection before calling OpenSubscription
// and to close it when Close has been called on all Subscriptions opened with it.
//
// The documentation of the amqp package recommends using separate connections for
// publishing and subscribing.
func OpenSubscription(conn *amqp.Connection, name string) (*pubsub.Subscription, error) {
	// TODO(jba): support context.Context
	s, err := newSubscription(conn, name)
	if err != nil {
		return nil, err
	}
	return pubsub.NewSubscription(s), nil
}

type subscription struct {
	conn     *amqp.Connection
	queue    string // the AMQP queue name
	consumer string // the client-generated name for this particular subscriber

	mu   sync.Mutex
	ch   *amqp.Channel // AMQP channel used for all communication.
	delc <-chan amqp.Delivery
}

var nextConsumer int64 // atomic

func newSubscription(conn *amqp.Connection, name string) (*subscription, error) {
	s := &subscription{
		conn:     conn,
		queue:    name,
		consumer: fmt.Sprintf("c%d", atomic.AddInt64(&nextConsumer, 1)),
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := s.establishChannel(); err != nil {
		return nil, err
	}
	return s, nil
}

// Must be called with s.mu held.
func (s *subscription) establishChannel() error {
	// TODO(jba): support context.Context
	if s.ch != nil {
		// We already have a channel; nothing to do.
		return nil
	}
	// Create a new channel.
	ch, err := s.conn.Channel()
	if err != nil {
		return err
	}
	// Subscribe to messages from the queue.
	s.delc, err = ch.Consume(s.queue, s.consumer,
		false, // autoAck
		false, // exclusive
		false, // noLocal
		wait,
		nil) // args
	if err != nil {
		return err
	}
	s.ch = ch
	return nil
}

// ReceiveBatch implements driver.Subscription.ReceiveBatch.
func (s *subscription) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if err := s.establishChannel(); err != nil {
		return nil, err
	}

	// Get up to maxMessages waiting messages, but don't take too long.
	var ms []*driver.Message
	maxTime := time.After(50 * time.Millisecond)
	for {
		select {
		case <-ctx.Done():
			// Cancel the Consume.
			_ = s.ch.Cancel(s.consumer, wait) // ignore the error
			s.ch = nil
			return nil, ctx.Err()

		case d, ok := <-s.delc:
			if !ok { // channel closed
				s.ch = nil // re-establish the channel next time
				if len(ms) > 0 {
					return ms, nil
				}
				// TODO(jba): get more information for a better message
				return nil, errors.New("delivery channel closed")
			}
			ms = append(ms, toMessage(d))
			if len(ms) >= maxMessages {
				return ms, nil
			}

		case <-maxTime:
			// Timed out. Return whatever we have.
			return ms, nil
		}
	}
}

// toMessage converts an amqp.Delivery (a received message) to a driver.Message.
func toMessage(d amqp.Delivery) *driver.Message {
	// Delivery.Headers is a map[string]interface{}, so we have to
	// convert each value to a string.
	md := map[string]string{}
	for k, v := range d.Headers {
		md[k] = fmt.Sprint(v)
	}
	return &driver.Message{
		Body:     d.Body,
		AckID:    d.DeliveryTag,
		Metadata: md,
	}
}

// SendAcks implements driver.Subscription.SendAcks.
func (s *subscription) SendAcks(ctx context.Context, ackIDs []driver.AckID) error {
	// TODO(#853): consider a separate channel for acks, so ReceiveBatch and SendAcks don't
	// block each other.
	s.mu.Lock()
	defer s.mu.Unlock()

	if err := s.establishChannel(); err != nil {
		return err
	}

	// The Ack call doesn't wait for a response, so this loop should execute relatively
	// quickly.
	// It wouldn't help to make it concurrent, because Channel.Ack grabs a channel-wide mutex.
	// (We could consider using multiple channels if performance becomes an issue.)
	for _, id := range ackIDs {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		err := s.ch.Ack(id.(uint64), false) // multiple=false: acking only this ID
		if err != nil {
			s.ch = nil // re-establish channel after an error
			return err
		}
	}
	return nil
}

// IsRetryable implements driver.Subscription.IsRetryable.
func (*subscription) IsRetryable(error) bool {
	// TODO(jba): figure out what errors can be retried.
	return false
}
<|endoftext|>


<commit_before>package filters

import (
	"strings"

	"regexp"

	"fmt"

	d "github.com/appPlant/alpinepass/data"
	"github.com/appPlant/alpinepass/util"
	"github.com/fatih/structs"
)

//PropertyFilter allows filtering a Config for the content of its properties.
type PropertyFilter struct {
	Slices []string
}

func (p PropertyFilter) filter(data d.Config) d.Config {
	verifyFlags(p.Slices) //TODO verify only once and not for every Config
	for i := 0; i < len(p.Slices); i++ {
		data = filterProperty(p.Slices[i], data)
	}
	return data
}

func filterProperty(property string, data d.Config) d.Config {
	var split = strings.Split(property, ":")
	var key = split[0]
	var value = split[1]

	var regex = regexp.MustCompile(value)
	fmt.Println(regex.MatchString(data.Location))

	//TODO rather use reflection?
	switch key {
	case "id":
		if !strings.Contains(data.ID, value) {
			data.IsValid = false
		}
	case "title":
		if !strings.Contains(data.Title, value) {
			data.IsValid = false
		}
	case "location":
		if !regex.MatchString(data.Location) {
			data.IsValid = false
		}
	case "environment":
		if !strings.Contains(data.Environment, value) {
			data.IsValid = false
		}
	case "user":
		if !strings.Contains(data.User, value) {
			data.IsValid = false
		}
	case "host":
		if !strings.Contains(data.Host, value) {
			data.IsValid = false
		}
	}

	return data
}

//verifyFlags checks that the input flags are valid.
func verifyFlags(flags []string) {
	for i := 0; i < len(flags); i++ {
		flag := flags[i]

		//Do some simple checks.
		if !strings.Contains(flag, ":") {
			util.ThrowError("The filter does not contain ':'! Flag: " + flag)
		}
		if strings.Count(flag, ":") > 1 {
			util.ThrowError("The filter contains too many ':'! Flag: " + flag)
		}

		//check that the key part of a flag matches the fields available in struct Config.
		data := d.Config{}
		fieldNames := structs.Names(data)
		key := strings.ToLower(strings.Split(flag, ":")[0])
		isContained := false
		for j := 0; j < len(fieldNames); j++ {
			if key == strings.ToLower(fieldNames[j]) {
				isContained = true
			}
		}
		if !isContained {
			util.ThrowError("The filter type is not valid! Flag: " + flag)
		}
	}
}
<commit_msg>Use regex matching for all fields<commit_after>package filters

import (
	"strings"

	"regexp"

	d "github.com/appPlant/alpinepass/data"
	"github.com/appPlant/alpinepass/util"
	"github.com/fatih/structs"
)

//PropertyFilter allows filtering a Config for the content of its properties.
type PropertyFilter struct {
	Slices []string
}

func (p PropertyFilter) filter(data d.Config) d.Config {
	verifyFlags(p.Slices) //TODO verify only once and not for every Config
	for i := 0; i < len(p.Slices); i++ {
		data = filterProperty(p.Slices[i], data)
	}
	return data
}

func filterProperty(property string, data d.Config) d.Config {
	var split = strings.Split(property, ":")
	var key = split[0]
	var value = split[1]
	var regex = regexp.MustCompile(value)

	//TODO rather use reflection?
	switch key {
	case "id":
		if !regex.MatchString(data.ID) {
			data.IsValid = false
		}
	case "title":
		if !regex.MatchString(data.Title) {
			data.IsValid = false
		}
	case "location":
		if !regex.MatchString(data.Location) {
			data.IsValid = false
		}
	case "environment":
		if !regex.MatchString(data.Environment) {
			data.IsValid = false
		}
	case "user":
		if !regex.MatchString(data.User) {
			data.IsValid = false
		}
	case "host":
		if !regex.MatchString(data.Host) {
			data.IsValid = false
		}
	}

	return data
}

//verifyFlags checks that the input flags are valid.
func verifyFlags(flags []string) {
	for i := 0; i < len(flags); i++ {
		flag := flags[i]

		//Do some simple checks.
		if !strings.Contains(flag, ":") {
			util.ThrowError("The filter does not contain ':'! Flag: " + flag)
		}
		if strings.Count(flag, ":") > 1 {
			util.ThrowError("The filter contains too many ':'! Flag: " + flag)
		}

		//check that the key part of a flag matches the fields available in struct Config.
		data := d.Config{}
		fieldNames := structs.Names(data)
		key := strings.ToLower(strings.Split(flag, ":")[0])
		isContained := false
		for j := 0; j < len(fieldNames); j++ {
			if key == strings.ToLower(fieldNames[j]) {
				isContained = true
			}
		}
		if !isContained {
			util.ThrowError("The filter type is not valid! Flag: " + flag)
		}
	}
}
<|endoftext|>


<commit_before>package updown

import (
	"errors"
	"fmt"
	"net/http"
)

// SSL represents the SSL section of a check
type SSL struct {
	TestedAt string `json:"tested_at,omitempty"`
	Valid    bool   `json:"valid,omitempty"`
	Error    string `json:"error,omitempty"`
}

// Check represents a check performed by Updown on a regular basis
type Check struct {
	Token             string            `json:"token,omitempty"`
	URL               string            `json:"url,omitempty"`
	Alias             string            `json:"alias,omitempty"`
	LastStatus        int               `json:"last_status,omitempty"`
	Uptime            float64           `json:"uptime,omitempty"`
	Down              bool              `json:"down,omitempty"`
	DownSince         string            `json:"down_since,omitempty"`
	Error             string            `json:"error,omitempty"`
	Period            int               `json:"period,omitempty"`
	Apdex             float64           `json:"apdex_t,omitempty"`
	Enabled           bool              `json:"enabled,omitempty"`
	Published         bool              `json:"published,omitempty"`
	LastCheckAt       string            `json:"last_check_at,omitempty"`
	NextCheckAt       string            `json:"next_check_at,omitempty"`
	FaviconURL        string            `json:"favicon_url,omitempty"`
	SSL               SSL               `json:"ssl,omitempty"`
	StringMatch       string            `json:"string_match,omitempty"`
	MuteUntil         string            `json:"mute_until,omitempty"`
	DisabledLocations []string          `json:"disabled_locations,omitempty"`
	CustomHeaders     map[string]string `json:"custom_headers,omitempty"`
}

// CheckItem represents a new check you want to be performed by Updown
type CheckItem struct {
	// The URL you want to monitor
	URL string `json:"url,omitempty"`
	// Interval in seconds (30, 60, 120, 300 or 600)
	Period int `json:"period,omitempty"`
	// APDEX threshold in seconds (0.125, 0.25, 0.5 or 1.0)
	Apdex float64 `json:"apdex_t,omitempty"`
	// Is the check enabled
	Enabled bool `json:"enabled,omitempty"`
	// Shall the status page be public
	Published bool `json:"published,omitempty"`
	// Human readable name
	Alias string `json:"alias,omitempty"`
	// Search for this string in the page
	StringMatch string `json:"string_match,omitempty"`
	// Mute notifications until given time, accepts a time, 'recovery' or 'forever'
	MuteUntil string `json:"mute_until,omitempty"`
	// Disabled monitoring locations. It's an array of abbreviated location names
	DisabledLocations []string `json:"disabled_locations,omitempty"`
	// The HTTP headers you want in updown requests
	CustomHeaders map[string]string `json:"custom_headers,omitempty"`
}

// CheckService interacts with the checks section of the API
type CheckService struct {
	client *Client
	cache  Cache
}

type removeResponse struct {
	Deleted bool `json:"deleted,omitempty"`
}

// ErrTokenNotFound indicates that we cannot find a token for the given name
var ErrTokenNotFound = errors.New("Could not determine a token for the given name")

// TokenForAlias finds the Updown token for a check's alias
func (s *CheckService) TokenForAlias(name string) (string, error) {
	// Retrieve from cache
	if has, val := s.cache.Get(name); has {
		return val, nil
	}

	// List all checks
	checks, _, err := s.List()
	if err != nil {
		return "", err
	}

	// And try to find the appropriate name
	token, found := "", false
	for _, check := range checks {
		s.cache.Put(check.Alias, check.Token)
		if check.Alias == name {
			found, token = true, check.Token
		}
	}

	if found {
		return token, nil
	}

	// Could not find a match
	return "", ErrTokenNotFound
}

// List lists all the checks
func (s *CheckService) List() ([]Check, *http.Response, error) {
	req, err := s.client.NewRequest("GET", "checks", nil)
	if err != nil {
		return nil, nil, err
	}

	var res []Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return nil, resp, err
	}

	return res, resp, err
}

// Get gets a single check by its token
func (s *CheckService) Get(token string) (Check, *http.Response, error) {
	req, err := s.client.NewRequest("GET", pathForToken(token), nil)
	if err != nil {
		return Check{}, nil, err
	}

	var res Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return Check{}, resp, err
	}

	return res, resp, err
}

// Add adds a new check you want to be performed
func (s *CheckService) Add(data CheckItem) (Check, *http.Response, error) {
	req, err := s.client.NewRequest("POST", "checks", data)
	if err != nil {
		return Check{}, nil, err
	}

	var res Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return Check{}, resp, err
	}

	return res, resp, err
}

// Update updates a check performed by Updown
func (s *CheckService) Update(token string, data CheckItem) (Check, *http.Response, error) {
	req, err := s.client.NewRequest("PUT", pathForToken(token), data)
	if err != nil {
		return Check{}, nil, err
	}

	var res Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return Check{}, resp, err
	}

	return res, resp, err
}

// Remove removes a check from Updown by its token
func (s *CheckService) Remove(token string) (bool, *http.Response, error) {
	req, err := s.client.NewRequest("DELETE", pathForToken(token), nil)
	if err != nil {
		return false, nil, err
	}

	var res removeResponse
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return false, resp, err
	}

	return res.Deleted, resp, err
}

func pathForToken(token string) string {
	return fmt.Sprintf("checks/%s", token)
}
<commit_msg>Handle empty booleans in CheckItem structure #9<commit_after>package updown

import (
	"errors"
	"fmt"
	"net/http"
)

// SSL represents the SSL section of a check
type SSL struct {
	TestedAt string `json:"tested_at,omitempty"`
	Valid    bool   `json:"valid,omitempty"`
	Error    string `json:"error,omitempty"`
}

// Check represents a check performed by Updown on a regular basis
type Check struct {
	Token             string            `json:"token,omitempty"`
	URL               string            `json:"url,omitempty"`
	Alias             string            `json:"alias,omitempty"`
	LastStatus        int               `json:"last_status,omitempty"`
	Uptime            float64           `json:"uptime,omitempty"`
	Down              bool              `json:"down"`
	DownSince         string            `json:"down_since,omitempty"`
	Error             string            `json:"error,omitempty"`
	Period            int               `json:"period,omitempty"`
	Apdex             float64           `json:"apdex_t,omitempty"`
	Enabled           bool              `json:"enabled"`
	Published         bool              `json:"published"`
	LastCheckAt       string            `json:"last_check_at,omitempty"`
	NextCheckAt       string            `json:"next_check_at,omitempty"`
	FaviconURL        string            `json:"favicon_url,omitempty"`
	SSL               SSL               `json:"ssl,omitempty"`
	StringMatch       string            `json:"string_match,omitempty"`
	MuteUntil         string            `json:"mute_until,omitempty"`
	DisabledLocations []string          `json:"disabled_locations,omitempty"`
	CustomHeaders     map[string]string `json:"custom_headers,omitempty"`
}

// CheckItem represents a new check you want to be performed by Updown
type CheckItem struct {
	// The URL you want to monitor
	URL string `json:"url,omitempty"`
	// Interval in seconds (30, 60, 120, 300 or 600)
	Period int `json:"period,omitempty"`
	// APDEX threshold in seconds (0.125, 0.25, 0.5 or 1.0)
	Apdex float64 `json:"apdex_t,omitempty"`
	// Is the check enabled
	Enabled bool `json:"enabled"`
	// Shall the status page be public
	Published bool `json:"published"`
	// Human readable name
	Alias string `json:"alias,omitempty"`
	// Search for this string in the page
	StringMatch string `json:"string_match,omitempty"`
	// Mute notifications until given time, accepts a time, 'recovery' or 'forever'
	MuteUntil string `json:"mute_until,omitempty"`
	// Disabled monitoring locations. It's an array of abbreviated location names
	DisabledLocations []string `json:"disabled_locations,omitempty"`
	// The HTTP headers you want in updown requests
	CustomHeaders map[string]string `json:"custom_headers,omitempty"`
}

// CheckService interacts with the checks section of the API
type CheckService struct {
	client *Client
	cache  Cache
}

type removeResponse struct {
	Deleted bool `json:"deleted,omitempty"`
}

// ErrTokenNotFound indicates that we cannot find a token for the given name
var ErrTokenNotFound = errors.New("Could not determine a token for the given name")

// TokenForAlias finds the Updown token for a check's alias
func (s *CheckService) TokenForAlias(name string) (string, error) {
	// Retrieve from cache
	if has, val := s.cache.Get(name); has {
		return val, nil
	}

	// List all checks
	checks, _, err := s.List()
	if err != nil {
		return "", err
	}

	// And try to find the appropriate name
	token, found := "", false
	for _, check := range checks {
		s.cache.Put(check.Alias, check.Token)
		if check.Alias == name {
			found, token = true, check.Token
		}
	}

	if found {
		return token, nil
	}

	// Could not find a match
	return "", ErrTokenNotFound
}

// List lists all the checks
func (s *CheckService) List() ([]Check, *http.Response, error) {
	req, err := s.client.NewRequest("GET", "checks", nil)
	if err != nil {
		return nil, nil, err
	}

	var res []Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return nil, resp, err
	}

	return res, resp, err
}

// Get gets a single check by its token
func (s *CheckService) Get(token string) (Check, *http.Response, error) {
	req, err := s.client.NewRequest("GET", pathForToken(token), nil)
	if err != nil {
		return Check{}, nil, err
	}

	var res Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return Check{}, resp, err
	}

	return res, resp, err
}

// Add adds a new check you want to be performed
func (s *CheckService) Add(data CheckItem) (Check, *http.Response, error) {
	req, err := s.client.NewRequest("POST", "checks", data)
	if err != nil {
		return Check{}, nil, err
	}

	var res Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return Check{}, resp, err
	}

	return res, resp, err
}

// Update updates a check performed by Updown
func (s *CheckService) Update(token string, data CheckItem) (Check, *http.Response, error) {
	req, err := s.client.NewRequest("PUT", pathForToken(token), data)
	if err != nil {
		return Check{}, nil, err
	}

	var res Check
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return Check{}, resp, err
	}

	return res, resp, err
}

// Remove removes a check from Updown by its token
func (s *CheckService) Remove(token string) (bool, *http.Response, error) {
	req, err := s.client.NewRequest("DELETE", pathForToken(token), nil)
	if err != nil {
		return false, nil, err
	}

	var res removeResponse
	resp, err := s.client.Do(req, &res)
	if err != nil {
		return false, resp, err
	}

	return res.Deleted, resp, err
}

func pathForToken(token string) string {
	return fmt.Sprintf("checks/%s", token)
}
<|endoftext|>


<commit_before>package circle

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

var client http.Client

func init() {
	client = http.Client{
		Timeout: 10 * time.Second,
	}
}

const VERSION = "0.5"
const baseUri = "https://circleci.com/api/v1/project"

type TreeBuild struct {
	BuildNum    int            `json:"build_num"`
	BuildURL    string         `json:"build_url"`
	Status      string         `json:"status"`
	StartTime   CircleNullTime `json:"start_time"`
	StopTime    CircleNullTime `json:"stop_time"`
	VCSRevision string         `json:"vcs_revision"`
}

func (tb *TreeBuild) Passed() bool {
	return tb.Status == "success" || tb.Status == "fixed"
}

func (tb *TreeBuild) Failed() bool {
	return tb.Status == "failed" || tb.Status == "timedout"
}

type CircleBuild struct {
	Parallel uint8  `json:"parallel"`
	Steps    []Step `json:"steps"`
}

type Step struct {
	Name    string   `json:"name"`
	Actions []Action `json:"actions"`
}

type Action struct {
	Name      string         `json:"name"`
	OutputURL URL            `json:"output_url"`
	Runtime   CircleDuration `json:"run_time_millis"`
	Status    string         `json:"status"`
}

func (a *Action) Failed() bool {
	return a.Status == "failed" || a.Status == "timedout"
}

func getTreeUri(org string, project string, branch string, token string) string {
	return fmt.Sprintf("%s/%s/%s/tree/%s?circle-token=%s", baseUri, org, project, branch, token)
}

func getBuildUri(org string, project string, build int, token string) string {
	return fmt.Sprintf("%s/%s/%s/%d?circle-token=%s", baseUri, org, project, build, token)
}

type CircleTreeResponse []TreeBuild

func makeRequest(method, uri string) (io.ReadCloser, error) {
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Accept", "application/json")
	req.Header.Add("User-Agent", fmt.Sprintf("circle-command-line-client/%s", VERSION))
	resp, err := client.Do(req)
	if err != nil {
		fmt.Println("request error")
		return nil, err
	}
	return resp.Body, nil
}

func GetTree(org string, project string, branch string) (*CircleTreeResponse, error) {
	token, err := getToken(org)
	if err != nil {
		return nil, err
	}
	uri := getTreeUri(org, project, branch, token)
	body, err := makeRequest("GET", uri)
	defer body.Close()
	var cr CircleTreeResponse
	d := json.NewDecoder(body)
	err = d.Decode(&cr)
	return &cr, err
}

func GetBuild(org string, project string, buildNum int) (*CircleBuild, error) {
	token, err := getToken(org)
	if err != nil {
		return nil, err
	}
	uri := getBuildUri(org, project, buildNum, token)
	body, err := makeRequest("GET", uri)
	if err != nil {
		return nil, err
	}
	defer body.Close()
	var cb CircleBuild
	d := json.NewDecoder(body)
	err = d.Decode(&cb)
	return &cb, err
}
<commit_msg>Check error returned from HTTP request<commit_after>package circle

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

var client http.Client

func init() {
	client = http.Client{
		Timeout: 10 * time.Second,
	}
}

const VERSION = "0.5"
const baseUri = "https://circleci.com/api/v1/project"

type TreeBuild struct {
	BuildNum    int            `json:"build_num"`
	BuildURL    string         `json:"build_url"`
	Status      string         `json:"status"`
	StartTime   CircleNullTime `json:"start_time"`
	StopTime    CircleNullTime `json:"stop_time"`
	VCSRevision string         `json:"vcs_revision"`
}

func (tb *TreeBuild) Passed() bool {
	return tb.Status == "success" || tb.Status == "fixed"
}

func (tb *TreeBuild) Failed() bool {
	return tb.Status == "failed" || tb.Status == "timedout"
}

type CircleBuild struct {
	Parallel uint8  `json:"parallel"`
	Steps    []Step `json:"steps"`
}

type Step struct {
	Name    string   `json:"name"`
	Actions []Action `json:"actions"`
}

type Action struct {
	Name      string         `json:"name"`
	OutputURL URL            `json:"output_url"`
	Runtime   CircleDuration `json:"run_time_millis"`
	Status    string         `json:"status"`
}

func (a *Action) Failed() bool {
	return a.Status == "failed" || a.Status == "timedout"
}

func getTreeUri(org string, project string, branch string, token string) string {
	return fmt.Sprintf("%s/%s/%s/tree/%s?circle-token=%s", baseUri, org, project, branch, token)
}

func getBuildUri(org string, project string, build int, token string) string {
	return fmt.Sprintf("%s/%s/%s/%d?circle-token=%s", baseUri, org, project, build, token)
}

type CircleTreeResponse []TreeBuild

func makeRequest(method, uri string) (io.ReadCloser, error) {
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Accept", "application/json")
	req.Header.Add("User-Agent", fmt.Sprintf("circle-command-line-client/%s", VERSION))
	resp, err := client.Do(req)
	if err != nil {
		fmt.Println("request error")
		return nil, err
	}
	return resp.Body, nil
}

func GetTree(org string, project string, branch string) (*CircleTreeResponse, error) {
	token, err := getToken(org)
	if err != nil {
		return nil, err
	}
	uri := getTreeUri(org, project, branch, token)
	body, err := makeRequest("GET", uri)
	if err != nil {
		return nil, err
	}
	defer body.Close()
	var cr CircleTreeResponse
	d := json.NewDecoder(body)
	err = d.Decode(&cr)
	return &cr, err
}

func GetBuild(org string, project string, buildNum int) (*CircleBuild, error) {
	token, err := getToken(org)
	if err != nil {
		return nil, err
	}
	uri := getBuildUri(org, project, buildNum, token)
	body, err := makeRequest("GET", uri)
	if err != nil {
		return nil, err
	}
	defer body.Close()
	var cb CircleBuild
	d := json.NewDecoder(body)
	err = d.Decode(&cb)
	return &cb, err
}
<|endoftext|>


<commit_before>// circle.go

package circle

import (
	"errors"
	"hash/adler32"

	"github.com/charlesetc/dive"
)

var hash func([]byte) uint32 = adler32.Checksum

type Circle struct {
	address []byte
	hash    uint32
	next    *Circle
}

const (
	ReplicationDepth int = 1
)

func NewCircleHead() *Circle {
	circle := new(Circle)
	circle.hash = 0 // 0 means it's the head
	// circle.address is undefined
	circle.next = circle
	return circle
}

func NewCircle(address []byte) *Circle {
	circle := new(Circle)
	circle.address = address
	circle.hash = hash(address)
	return circle
}

func NewCircleString(address string) *Circle {
	return NewCircle([]byte(address))
}

func (c *Circle) Add(incoming *Circle) *Circle {
	var current *Circle
	for current = c; current.next.hash < incoming.hash; current = current.next {
		if current.next.hash == 0 {
			break
		}
	}
	incoming.next = current.next
	current.next = incoming
	return incoming
}

func (c *Circle) AddString(address string) *Circle {
	return c.Add(NewCircleString(address))
}

func CircleFromNode(node *dive.Node) *Circle {
	circle := NewCircleHead()
	for _, rec := range dive.GetAliveFromMap(node.Members) {
		circle.AddString(rec.Address)
	}
	return circle
}

// Will loop forever with an empty node...
func (c *Circle) KeyAddress(key []byte) func() ([]byte, error) {
	hashed := hash(key)

	var current *Circle
	for current = c; current.next.hash < hashed; current = current.next {
		if current.next.hash == 0 {
			break
		}
	}

	i := 0
	return func() ([]byte, error) {
		output := current.address
		i++

		if i > ReplicationDepth {
			return []byte{}, errors.New("No more replications.")
		}

		current = current.next
		if current.hash == 0 {
			current = current.next
		}

		return output, nil
	}
}
<commit_msg>Add method to create circle from list of strings<commit_after>// circle.go

package circle

import (
	"errors"
	"hash/adler32"

	"github.com/charlesetc/dive"
)

var hash func([]byte) uint32 = adler32.Checksum

type Circle struct {
	address []byte
	hash    uint32
	next    *Circle
}

const (
	ReplicationDepth int = 1
)

func NewCircleHead() *Circle {
	circle := new(Circle)
	circle.hash = 0 // 0 means it's the head
	// circle.address is undefined
	circle.next = circle
	return circle
}

func NewCircle(address []byte) *Circle {
	circle := new(Circle)
	circle.address = address
	circle.hash = hash(address)
	return circle
}

func NewCircleString(address string) *Circle {
	return NewCircle([]byte(address))
}

func (c *Circle) Add(incoming *Circle) *Circle {
	var current *Circle
	for current = c; current.next.hash < incoming.hash; current = current.next {
		if current.next.hash == 0 {
			break
		}
	}
	incoming.next = current.next
	current.next = incoming
	return incoming
}

func (c *Circle) AddString(address string) *Circle {
	return c.Add(NewCircleString(address))
}

func CircleFromList(nodes []string) *Circle {
	circle := NewCircleHead()

	for _, node := range nodes {
		circle.AddString(node)
	}

	return circle
}

func CircleFromNode(node *dive.Node) *Circle {
	circle := NewCircleHead()
	for _, rec := range dive.GetAliveFromMap(node.Members) {
		circle.AddString(rec.Address)
	}
	return circle
}

// Will loop forever with an empty node...
func (c *Circle) KeyAddress(key []byte) func() ([]byte, error) {
	hashed := hash(key)

	var current *Circle
	for current = c; current.next.hash < hashed; current = current.next {
		if current.next.hash == 0 {
			break
		}
	}

	i := 0
	return func() ([]byte, error) {
		output := current.address
		i++

		if i > ReplicationDepth {
			return []byte{}, errors.New("No more replications.")
		}

		current = current.next
		if current.hash == 0 {
			current = current.next
		}

		return output, nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Loadcat Authors. All rights reserved.

package ui

import (
	"net/http"

	"github.com/gorilla/mux"
	"github.com/gorilla/schema"
	"gopkg.in/mgo.v2/bson"

	"github.com/hjr265/loadcat/data"
	"github.com/hjr265/loadcat/feline"
)

func ServeBalancerList(w http.ResponseWriter, r *http.Request) {
	bals, err := data.ListBalancers()
	if err != nil {
		panic(err)
	}

	err = TplBalancerList.Execute(w, struct {
		Balancers []data.Balancer
	}{
		Balancers: bals,
	})
	if err != nil {
		panic(err)
	}
}

func ServeBalancerNewForm(w http.ResponseWriter, r *http.Request) {
	err := TplBalancerNewForm.Execute(w, struct {
	}{})
	if err != nil {
		panic(err)
	}
}

func HandleBalancerCreate(w http.ResponseWriter, r *http.Request) {
	err := r.ParseForm()
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}
	body := struct {
		Label string `schema:"label"`
	}{}
	err = schema.NewDecoder().Decode(&body, r.PostForm)
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	bal := data.Balancer{}
	bal.Label = body.Label
	err = bal.Put()
	if err != nil {
		panic(err)
	}

	err = feline.Commit(&bal)
	if err != nil {
		panic(err)
	}

	http.Redirect(w, r, "/balancers/"+bal.Id.Hex()+"/edit", http.StatusSeeOther)
}

func ServeBalancer(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	if !bson.IsObjectIdHex(vars["id"]) {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}
	bal, err := data.GetBalancer(bson.ObjectIdHex(vars["id"]))
	if err != nil {
		panic(err)
	}

	err = TplBalancerView.Execute(w, struct {
		Balancer *data.Balancer
	}{
		Balancer: bal,
	})
	if err != nil {
		panic(err)
	}
}

func ServeBalancerEditForm(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	if !bson.IsObjectIdHex(vars["id"]) {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}
	bal, err := data.GetBalancer(bson.ObjectIdHex(vars["id"]))
	if err != nil {
		panic(err)
	}

	err = TplBalancerEditForm.Execute(w, struct {
		Balancer     *data.Balancer
		Protocols    []data.Protocol
		Algorithms   []data.Algorithm
		CipherSuites []data.CipherSuite
	}{
		Balancer:     bal,
		Protocols:    data.Protocols,
		Algorithms:   data.Algorithms,
		CipherSuites: data.CipherSuites,
	})
	if err != nil {
		panic(err)
	}
}

func HandleBalancerUpdate(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	if !bson.IsObjectIdHex(vars["id"]) {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}
	bal, err := data.GetBalancer(bson.ObjectIdHex(vars["id"]))
	if err != nil {
		panic(err)
	}

	err = r.ParseForm()
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}
	body := struct {
		Label    string `schema:"label"`
		Settings struct {
			Hostname   string `schema:"hostname"`
			Port       int    `schema:"port"`
			Protocol   string `schema:"protocol"`
			Algorithm  string `schema:"algorithm"`
			SSLOptions struct {
				CipherSuite string  `schema:"cipher_suite"`
				Certificate *string `schema:"certificate"`
				PrivateKey  *string `schema:"private_key"`
			} `schema:"ssl_options"`
		} `schema:"settings"`
	}{}
	err = schema.NewDecoder().Decode(&body, r.PostForm)
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	bal.Label = body.Label
	bal.Settings.Hostname = body.Settings.Hostname
	bal.Settings.Port = body.Settings.Port
	bal.Settings.Protocol = data.Protocol(body.Settings.Protocol)
	bal.Settings.Algorithm = data.Algorithm(body.Settings.Algorithm)
	if body.Settings.Protocol == "https" {
		bal.Settings.SSLOptions.CipherSuite = "recommended"
		if body.Settings.SSLOptions.Certificate != nil {
			bal.Settings.SSLOptions.Certificate = []byte(*body.Settings.SSLOptions.Certificate)
		}
		if body.Settings.SSLOptions.PrivateKey != nil {
			bal.Settings.SSLOptions.PrivateKey = []byte(*body.Settings.SSLOptions.PrivateKey)
		}
	}
	err = bal.Put()
	if err != nil {
		panic(err)
	}

	err = feline.Commit(bal)
	if err != nil {
		panic(err)
	}

	http.Redirect(w, r, "/balancers/"+bal.Id.Hex(), http.StatusSeeOther)
}

func init() {
	Router.NewRoute().
		Methods("GET").
		Path("/balancers").
		Handler(http.HandlerFunc(ServeBalancerList))
	Router.NewRoute().
		Methods("GET").
		Path("/balancers/new").
		Handler(http.HandlerFunc(ServeBalancerNewForm))
	Router.NewRoute().
		Methods("POST").
		Path("/balancers/new").
		Handler(http.HandlerFunc(HandleBalancerCreate))
	Router.NewRoute().
		Methods("GET").
		Path("/balancers/{id}").
		Handler(http.HandlerFunc(ServeBalancer))
	Router.NewRoute().
		Methods("GET").
		Path("/balancers/{id}/edit").
		Handler(http.HandlerFunc(ServeBalancerEditForm))
	Router.NewRoute().
		Methods("POST").
		Path("/balancers/{id}/edit").
		Handler(http.HandlerFunc(HandleBalancerUpdate))
}
<commit_msg>Define balancer settings defaults<commit_after>// Copyright 2015 The Loadcat Authors. All rights reserved.

package ui

import (
	"net/http"

	"github.com/gorilla/mux"
	"github.com/gorilla/schema"
	"gopkg.in/mgo.v2/bson"

	"github.com/hjr265/loadcat/data"
	"github.com/hjr265/loadcat/feline"
)

func ServeBalancerList(w http.ResponseWriter, r *http.Request) {
	bals, err := data.ListBalancers()
	if err != nil {
		panic(err)
	}

	err = TplBalancerList.Execute(w, struct {
		Balancers []data.Balancer
	}{
		Balancers: bals,
	})
	if err != nil {
		panic(err)
	}
}

func ServeBalancerNewForm(w http.ResponseWriter, r *http.Request) {
	err := TplBalancerNewForm.Execute(w, struct {
	}{})
	if err != nil {
		panic(err)
	}
}

func HandleBalancerCreate(w http.ResponseWriter, r *http.Request) {
	err := r.ParseForm()
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}
	body := struct {
		Label string `schema:"label"`
	}{}
	err = schema.NewDecoder().Decode(&body, r.PostForm)
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	bal := data.Balancer{}
	bal.Label = body.Label
	bal.Settings.Hostname = "example.com"
	bal.Settings.Port = 80
	err = bal.Put()
	if err != nil {
		panic(err)
	}

	err = feline.Commit(&bal)
	if err != nil {
		panic(err)
	}

	http.Redirect(w, r, "/balancers/"+bal.Id.Hex()+"/edit", http.StatusSeeOther)
}

func ServeBalancer(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	if !bson.IsObjectIdHex(vars["id"]) {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}
	bal, err := data.GetBalancer(bson.ObjectIdHex(vars["id"]))
	if err != nil {
		panic(err)
	}

	err = TplBalancerView.Execute(w, struct {
		Balancer *data.Balancer
	}{
		Balancer: bal,
	})
	if err != nil {
		panic(err)
	}
}

func ServeBalancerEditForm(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	if !bson.IsObjectIdHex(vars["id"]) {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}
	bal, err := data.GetBalancer(bson.ObjectIdHex(vars["id"]))
	if err != nil {
		panic(err)
	}

	err = TplBalancerEditForm.Execute(w, struct {
		Balancer     *data.Balancer
		Protocols    []data.Protocol
		Algorithms   []data.Algorithm
		CipherSuites []data.CipherSuite
	}{
		Balancer:     bal,
		Protocols:    data.Protocols,
		Algorithms:   data.Algorithms,
		CipherSuites: data.CipherSuites,
	})
	if err != nil {
		panic(err)
	}
}

func HandleBalancerUpdate(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	if !bson.IsObjectIdHex(vars["id"]) {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}
	bal, err := data.GetBalancer(bson.ObjectIdHex(vars["id"]))
	if err != nil {
		panic(err)
	}

	err = r.ParseForm()
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}
	body := struct {
		Label    string `schema:"label"`
		Settings struct {
			Hostname   string `schema:"hostname"`
			Port       int    `schema:"port"`
			Protocol   string `schema:"protocol"`
			Algorithm  string `schema:"algorithm"`
			SSLOptions struct {
				CipherSuite string  `schema:"cipher_suite"`
				Certificate *string `schema:"certificate"`
				PrivateKey  *string `schema:"private_key"`
			} `schema:"ssl_options"`
		} `schema:"settings"`
	}{}
	err = schema.NewDecoder().Decode(&body, r.PostForm)
	if err != nil {
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	bal.Label = body.Label
	bal.Settings.Hostname = body.Settings.Hostname
	bal.Settings.Port = body.Settings.Port
	bal.Settings.Protocol = data.Protocol(body.Settings.Protocol)
	bal.Settings.Algorithm = data.Algorithm(body.Settings.Algorithm)
	if body.Settings.Protocol == "https" {
		bal.Settings.SSLOptions.CipherSuite = "recommended"
		if body.Settings.SSLOptions.Certificate != nil {
			bal.Settings.SSLOptions.Certificate = []byte(*body.Settings.SSLOptions.Certificate)
		}
		if body.Settings.SSLOptions.PrivateKey != nil {
			bal.Settings.SSLOptions.PrivateKey = []byte(*body.Settings.SSLOptions.PrivateKey)
		}
	}
	err = bal.Put()
	if err != nil {
		panic(err)
	}

	err = feline.Commit(bal)
	if err != nil {
		panic(err)
	}

	http.Redirect(w, r, "/balancers/"+bal.Id.Hex(), http.StatusSeeOther)
}

func init() {
	Router.NewRoute().
		Methods("GET").
		Path("/balancers").
		Handler(http.HandlerFunc(ServeBalancerList))
	Router.NewRoute().
		Methods("GET").
		Path("/balancers/new").
		Handler(http.HandlerFunc(ServeBalancerNewForm))
	Router.NewRoute().
		Methods("POST").
		Path("/balancers/new").
		Handler(http.HandlerFunc(HandleBalancerCreate))
	Router.NewRoute().
		Methods("GET").
		Path("/balancers/{id}").
		Handler(http.HandlerFunc(ServeBalancer))
	Router.NewRoute().
		Methods("GET").
		Path("/balancers/{id}/edit").
		Handler(http.HandlerFunc(ServeBalancerEditForm))
	Router.NewRoute().
		Methods("POST").
		Path("/balancers/{id}/edit").
		Handler(http.HandlerFunc(HandleBalancerUpdate))
}
<|endoftext|>


<commit_before>package pubsub

import (
	"math"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/golang/glog"

	"gopkg.in/vmihailenco/msgpack.v2"
)

// ClientConfig configures a Client.
type ClientConfig struct {
	// Dial: function for opening a connection to the broker
	Dial func() (net.Conn, error)

	// BackoffBase: in case client needs to redial, it will back-off
	// exponentially using the formula BackoffBase^(numFailures). Defaults to 1s.
	BackoffBase time.Duration

	// MaxBackoff: when backing off during redialing, this caps the wait time.
	// Defaults to 1m.
	MaxBackoff time.Duration

	// KeepalivePeriod: the client will send a message after being idle for
	// keepalivePeriod. Set this to the largest value you can that still keeps the
	// underlying connections from timing out due to idle timeouts on intervening
	// servers. Defaults to 30s.
	KeepalivePeriod time.Duration

	// AuthenticationKey: if supplied, client will immediately authenticate with
	// this key.
	AuthenticationKey string
}

// Client is a client that is able to communicate with a Broker over the
// network.
type Client struct {
	cfg                *ClientConfig
	subscriptions      map[string][]byte
	subscriptionsMutex sync.RWMutex
	keepaliveTimer     *time.Timer
	out                chan *Sendable
	in                 chan *Message
	forceReconnect     int32
	connect            chan interface{}
	conn               net.Conn
	enc                *msgpack.Encoder
}

// Sendable represents a Message that can be published to a specific client.
type Sendable struct {
	c   *Client
	msg *Message
}

// Connect creates a Client
func Connect(cfg *ClientConfig) *Client {
	// Set sensible defaults
	if cfg.BackoffBase == 0 {
		cfg.BackoffBase = 1 * time.Second
	}

	if cfg.MaxBackoff == 0 {
		cfg.MaxBackoff = 1 * time.Minute
	}

	if cfg.KeepalivePeriod == 0 {
		cfg.KeepalivePeriod = 30 * time.Second
	}

	client := &Client{
		cfg:            cfg,
		subscriptions:  make(map[string][]byte, 0),
		keepaliveTimer: time.NewTimer(cfg.KeepalivePeriod),
		out:            make(chan *Sendable, 0),
		in:             make(chan *Message, 0),
		connect:        make(chan interface{}, 1),
	}
	go client.process()

	return client
}

// Read reads the next Message from the Broker.
func (c *Client) Read() *Message {
	return <-c.in
}

// ReadTimeout reads the next message from the Broker or returns nil, false if
// no message is read within the given timeout.
func (c *Client) ReadTimeout(timeout time.Duration) (msg *Message, ok bool) {
	select {
	case msg := <-c.in:
		return msg, true
	case <-time.After(timeout):
		return nil, false
	}
}

// Subscribe subscribes this client to the given topic.
func (c *Client) Subscribe(topic []byte) {
	c.subscriptionsMutex.Lock()
	c.subscriptions[string(topic)] = topic
	c.subscriptionsMutex.Unlock()
	c.subscribe(topic).Send()
}

func (c *Client) subscribe(topic []byte) *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: Subscribe, Topic: topic},
	}
}

// Unsubscribe unsubscribes this client from the given topic.
func (c *Client) Unsubscribe(topic []byte) {
	c.subscriptionsMutex.Lock()
	delete(c.subscriptions, string(topic))
	c.subscriptionsMutex.Unlock()
	c.unsubscribe(topic).Send()
}

func (c *Client) unsubscribe(topic []byte) *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: Unsubscribe, Topic: topic},
	}
}

// Publish publishes a message with the given body to the given topic.
func (c *Client) Publish(topic, body []byte) {
	s := &Sendable{
		c:   c,
		msg: &Message{Type: Publish, Topic: topic, Body: body},
	}
	s.Send()
}

// Send sends the Sendable eventually (queues behind other Sends and ops).
func (s *Sendable) Send() {
	s.c.out <- s
}

func (s *Sendable) sendImmediate() error {
	s.c.resetKeepaliveTimer()
	err := s.c.enc.Encode(s.msg)
	if err == nil {
		glog.Infof("Sent message %v", s.msg)
	}
	return err
}

func (c *Client) process() {
	c.forceConnect()
	for {
		select {
		case s := <-c.out:
			glog.Info("Sending message")
			c.doWithConnection(s.sendImmediate)
		case <-c.keepaliveTimer.C:
			glog.Info("Sending KeepAlive")
			c.doWithConnection(c.keepAlive().sendImmediate)
		case <-c.connect:
			glog.Info("Ensuring we have a connection")
			c.doWithConnection(func() error { return nil })
		}
	}
}

func (c *Client) doWithConnection(op func() error) {
	for numFailures := 0; numFailures < math.MaxInt32; numFailures++ {
		// Back off if necessary
		if numFailures > 0 {
			backoff := time.Duration(math.Pow(float64(c.cfg.BackoffBase), float64(numFailures)))
			if backoff > c.cfg.MaxBackoff {
				backoff = c.cfg.MaxBackoff
			}
			glog.Infof("%d consecutive failures, waiting %v before retrying", numFailures, backoff)
			time.Sleep(backoff)
		}

		var err error
		if atomic.CompareAndSwapInt32(&c.forceReconnect, 1, 0) || c.conn == nil {
			c.close()
			glog.Info("Dialing new conn")
			c.conn, err = c.cfg.Dial()
			if err != nil {
				c.conn = nil
				glog.Warningf("Unable to dial broker: %v", err)
				continue
			}

			// Set up msgpack
			c.enc = msgpack.NewEncoder(c.conn)

			// Send initial messages
			err = c.sendInitialMessages()
			if err != nil {
				glog.Warningf("Error sending initial messages: %v", err)
				c.close()
				continue
			}

			// Successfully connected, start reading
			go c.readLoop(msgpack.NewDecoder(c.conn))
		}

		err = op()
		if err != nil {
			glog.Warningf("Unable to process op: %v", err)
			c.close()
			continue
		}

		glog.Info("Op succesful")
		return
	}
}

func (c *Client) sendInitialMessages() error {
	glog.Info("Sending initial messages")
	if c.cfg.AuthenticationKey != "" {
		err := c.authenticate(c.cfg.AuthenticationKey).sendImmediate()
		if err != nil {
			return err
		}
	}

	c.subscriptionsMutex.RLock()
	initialTopics := make([][]byte, 0, len(c.subscriptions))
	for _, topic := range c.subscriptions {
		initialTopics = append(initialTopics, topic)
	}
	c.subscriptionsMutex.RUnlock()

	for _, topic := range initialTopics {
		err := c.subscribe(topic).sendImmediate()
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *Client) readLoop(dec *msgpack.Decoder) {
	glog.Info("Start reading")
	for {
		msg := &Message{}
		err := dec.Decode(msg)
		if err != nil {
			glog.Warningf("Unable to read, stopping readLoop: %v", err)
			atomic.StoreInt32(&c.forceReconnect, 1)
			c.forceConnect()
			return
		}
		c.resetKeepaliveTimer()
		if msg.Type != KeepAlive {
			glog.Info("Read message")
			c.in <- msg
		} else {
			glog.Info("Ignoring KeepAlive message")
		}
	}
}

func (c *Client) forceConnect() {
	select {
	case c.connect <- nil:
	default:
		// Already attempting to force connection, don't bother doing it again
	}
}

func (c *Client) authenticate(authenticationKey string) *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: Authenticate, Body: []byte(authenticationKey)},
	}
}

func (c *Client) keepAlive() *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: KeepAlive},
	}
}

func (c *Client) resetKeepaliveTimer() {
	c.keepaliveTimer.Reset(c.cfg.KeepalivePeriod)
}

func (c *Client) close() {
	if c.conn != nil {
		glog.Info("Closing connection")
		c.conn.Close()
		c.conn = nil
	}
}
<commit_msg>Reduced keepalive timeout to 20 seconds<commit_after>package pubsub

import (
	"math"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/golang/glog"

	"gopkg.in/vmihailenco/msgpack.v2"
)

// ClientConfig configures a Client.
type ClientConfig struct {
	// Dial: function for opening a connection to the broker
	Dial func() (net.Conn, error)

	// BackoffBase: in case client needs to redial, it will back-off
	// exponentially using the formula BackoffBase^(numFailures). Defaults to 1s.
	BackoffBase time.Duration

	// MaxBackoff: when backing off during redialing, this caps the wait time.
	// Defaults to 1m.
	MaxBackoff time.Duration

	// KeepalivePeriod: the client will send a message after being idle for
	// keepalivePeriod. Set this to the largest value you can that still keeps the
	// underlying connections from timing out due to idle timeouts on intervening
	// servers. Defaults to 20s.
	KeepalivePeriod time.Duration

	// AuthenticationKey: if supplied, client will immediately authenticate with
	// this key.
	AuthenticationKey string
}

// Client is a client that is able to communicate with a Broker over the
// network.
type Client struct {
	cfg                *ClientConfig
	subscriptions      map[string][]byte
	subscriptionsMutex sync.RWMutex
	keepaliveTimer     *time.Timer
	out                chan *Sendable
	in                 chan *Message
	forceReconnect     int32
	connect            chan interface{}
	conn               net.Conn
	enc                *msgpack.Encoder
}

// Sendable represents a Message that can be published to a specific client.
type Sendable struct {
	c   *Client
	msg *Message
}

// Connect creates a Client
func Connect(cfg *ClientConfig) *Client {
	// Set sensible defaults
	if cfg.BackoffBase == 0 {
		cfg.BackoffBase = 1 * time.Second
	}

	if cfg.MaxBackoff == 0 {
		cfg.MaxBackoff = 1 * time.Minute
	}

	if cfg.KeepalivePeriod == 0 {
		cfg.KeepalivePeriod = 20 * time.Second
	}

	client := &Client{
		cfg:            cfg,
		subscriptions:  make(map[string][]byte, 0),
		keepaliveTimer: time.NewTimer(cfg.KeepalivePeriod),
		out:            make(chan *Sendable, 0),
		in:             make(chan *Message, 0),
		connect:        make(chan interface{}, 1),
	}
	go client.process()

	return client
}

// Read reads the next Message from the Broker.
func (c *Client) Read() *Message {
	return <-c.in
}

// ReadTimeout reads the next message from the Broker or returns nil, false if
// no message is read within the given timeout.
func (c *Client) ReadTimeout(timeout time.Duration) (msg *Message, ok bool) {
	select {
	case msg := <-c.in:
		return msg, true
	case <-time.After(timeout):
		return nil, false
	}
}

// Subscribe subscribes this client to the given topic.
func (c *Client) Subscribe(topic []byte) {
	c.subscriptionsMutex.Lock()
	c.subscriptions[string(topic)] = topic
	c.subscriptionsMutex.Unlock()
	c.subscribe(topic).Send()
}

func (c *Client) subscribe(topic []byte) *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: Subscribe, Topic: topic},
	}
}

// Unsubscribe unsubscribes this client from the given topic.
func (c *Client) Unsubscribe(topic []byte) {
	c.subscriptionsMutex.Lock()
	delete(c.subscriptions, string(topic))
	c.subscriptionsMutex.Unlock()
	c.unsubscribe(topic).Send()
}

func (c *Client) unsubscribe(topic []byte) *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: Unsubscribe, Topic: topic},
	}
}

// Publish publishes a message with the given body to the given topic.
func (c *Client) Publish(topic, body []byte) {
	s := &Sendable{
		c:   c,
		msg: &Message{Type: Publish, Topic: topic, Body: body},
	}
	s.Send()
}

// Send sends the Sendable eventually (queues behind other Sends and ops).
func (s *Sendable) Send() {
	s.c.out <- s
}

func (s *Sendable) sendImmediate() error {
	s.c.resetKeepaliveTimer()
	err := s.c.enc.Encode(s.msg)
	if err == nil {
		glog.Infof("Sent message %v", s.msg)
	}
	return err
}

func (c *Client) process() {
	c.forceConnect()
	for {
		select {
		case s := <-c.out:
			glog.Info("Sending message")
			c.doWithConnection(s.sendImmediate)
		case <-c.keepaliveTimer.C:
			glog.Info("Sending KeepAlive")
			c.doWithConnection(c.keepAlive().sendImmediate)
		case <-c.connect:
			glog.Info("Ensuring we have a connection")
			c.doWithConnection(func() error { return nil })
		}
	}
}

func (c *Client) doWithConnection(op func() error) {
	for numFailures := 0; numFailures < math.MaxInt32; numFailures++ {
		// Back off if necessary
		if numFailures > 0 {
			backoff := time.Duration(math.Pow(float64(c.cfg.BackoffBase), float64(numFailures)))
			if backoff > c.cfg.MaxBackoff {
				backoff = c.cfg.MaxBackoff
			}
			glog.Infof("%d consecutive failures, waiting %v before retrying", numFailures, backoff)
			time.Sleep(backoff)
		}

		var err error
		if atomic.CompareAndSwapInt32(&c.forceReconnect, 1, 0) || c.conn == nil {
			c.close()
			glog.Info("Dialing new conn")
			c.conn, err = c.cfg.Dial()
			if err != nil {
				c.conn = nil
				glog.Warningf("Unable to dial broker: %v", err)
				continue
			}

			// Set up msgpack
			c.enc = msgpack.NewEncoder(c.conn)

			// Send initial messages
			err = c.sendInitialMessages()
			if err != nil {
				glog.Warningf("Error sending initial messages: %v", err)
				c.close()
				continue
			}

			// Successfully connected, start reading
			go c.readLoop(msgpack.NewDecoder(c.conn))
		}

		err = op()
		if err != nil {
			glog.Warningf("Unable to process op: %v", err)
			c.close()
			continue
		}

		glog.Info("Op succesful")
		return
	}
}

func (c *Client) sendInitialMessages() error {
	glog.Info("Sending initial messages")
	if c.cfg.AuthenticationKey != "" {
		err := c.authenticate(c.cfg.AuthenticationKey).sendImmediate()
		if err != nil {
			return err
		}
	}

	c.subscriptionsMutex.RLock()
	initialTopics := make([][]byte, 0, len(c.subscriptions))
	for _, topic := range c.subscriptions {
		initialTopics = append(initialTopics, topic)
	}
	c.subscriptionsMutex.RUnlock()

	for _, topic := range initialTopics {
		err := c.subscribe(topic).sendImmediate()
		if err != nil {
			return err
		}
	}

	return nil
}

func (c *Client) readLoop(dec *msgpack.Decoder) {
	glog.Info("Start reading")
	for {
		msg := &Message{}
		err := dec.Decode(msg)
		if err != nil {
			glog.Warningf("Unable to read, stopping readLoop: %v", err)
			atomic.StoreInt32(&c.forceReconnect, 1)
			c.forceConnect()
			return
		}
		c.resetKeepaliveTimer()
		if msg.Type != KeepAlive {
			glog.Info("Read message")
			c.in <- msg
		} else {
			glog.Info("Ignoring KeepAlive message")
		}
	}
}

func (c *Client) forceConnect() {
	select {
	case c.connect <- nil:
	default:
		// Already attempting to force connection, don't bother doing it again
	}
}

func (c *Client) authenticate(authenticationKey string) *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: Authenticate, Body: []byte(authenticationKey)},
	}
}

func (c *Client) keepAlive() *Sendable {
	return &Sendable{
		c:   c,
		msg: &Message{Type: KeepAlive},
	}
}

func (c *Client) resetKeepaliveTimer() {
	c.keepaliveTimer.Reset(c.cfg.KeepalivePeriod)
}

func (c *Client) close() {
	if c.conn != nil {
		glog.Info("Closing connection")
		c.conn.Close()
		c.conn = nil
	}
}
<|endoftext|>