<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/net/proxy"

	"github.com/yahoo/athenz/clients/go/zms"
	"github.com/yahoo/athenz/libs/go/zmscli"
)

var (
	// VERSION gets set by the build script via the LDFLAGS.
	VERSION string

	// BUILD_DATE gets set by the build script via the LDFLAGS.
	BUILD_DATE string
)

func defaultZmsURL() string {
	s := os.Getenv("ZMS")
	if s != "" {
		return s
	}
	return "https://localhost:4443/zms/v1"
}

func defaultIdentity() string {
	return "user." + os.Getenv("USER")
}

func defaultSocksProxy() string {
	return os.Getenv("SOCKS5_PROXY")
}

func defaultDebug() bool {
	sDebug := os.Getenv("ZMS_DEBUG")
	return sDebug == "true"
}

func debugAuthNToken(identity string) string {
	i := strings.LastIndex(identity, ".")
	domain := identity[0:i]
	name := identity[i+1:]
	var buf bytes.Buffer
	buf.WriteString("v=U1;d=" + domain + ";n=" + name + ";s=fakesignature")
	return "v=U1;d=" + domain + ";n=" + name + ";s=fakesignature"
}

// isFreshFile checks the file's last modification time
// and returns true the file was updated within maxAge
// (file is "fresh"), false otherwise (file is "stale").
func isFreshFile(filename string, maxAge float64) bool {
	info, err := os.Stat(filename)
	if err != nil {
		return false
	}
	delta := time.Since(info.ModTime())
	// return false if duration exceeds maxAge
	tooOld := delta.Minutes() > maxAge
	return !tooOld
}

func getCachedNToken() string {
	ntokenFile := os.Getenv("HOME") + "/.ntoken"
	if isFreshFile(ntokenFile, 45) {
		data, err := ioutil.ReadFile(ntokenFile)
		if err == nil {
			return strings.TrimSpace(string(data))
		}
		fmt.Printf("Couldn't read the file, error: %v\n", err)
	}
	return ""
}

const identityPrefix = "user."

func getAuthNToken(identity, authorizedServices, zmsUrl string, tr *http.Transport) (string, error) {
	// our identity must be user
	if !strings.HasPrefix(identity, identityPrefix) {
		return "", errors.New("identity must start with " + identityPrefix)
	}
	user := identity[5:]
	ntoken := getCachedNToken()
	if ntoken != "" {
		return ntoken, nil
	}

	fmt.Fprintf(os.Stderr, "Enter password for "+user+": ")
	pass, err := terminal.ReadPassword(syscall.Stdin)
	if err != nil {
		return "", err
	}
	password := string(pass)
	// The user types <return> during terminal.ReadPassword() however
	// since echo is turned off the cursor doesn't make it to the next line.
	fmt.Fprint(os.Stderr, "\n")

	data := []byte(user + ":" + password)
	str := base64.StdEncoding.EncodeToString(data)

	var authHeader = "Authorization"
	var authCreds = "Basic " + str
	zmsClient := zms.ZMSClient{
		URL:         zmsUrl,
		Transport:   tr,
		CredsHeader: &authHeader,
		CredsToken:  &authCreds,
		Timeout:     0,
	}
	tok, err := zmsClient.GetUserToken(zms.SimpleName(user), authorizedServices, nil)
	if err != nil {
		return "", fmt.Errorf("Cannot get user token for user: %s error: %v", user, err)
	}
	if tok.Token != "" {
		ntokenFile := os.Getenv("HOME") + "/.ntoken"
		data := []byte(tok.Token)
		ioutil.WriteFile(ntokenFile, data, 0600)
	}
	return string(tok.Token), nil
}

func usage() string {
	var buf bytes.Buffer
	buf.WriteString("usage: zms-cli [flags] command [params]\n")
	buf.WriteString(" flags:\n")
	buf.WriteString("   -a auditRef         Audit Reference Token if auditing is required for domain\n")
	buf.WriteString("   -b                  Bulk import/update mode. Do not read/display updated role/policy/service objects (default=false)\n")
	buf.WriteString("   -c cacert_file      CA Certificate file path\n")
	buf.WriteString("   -d domain           The domain used for every command that takes a domain argument\n")
	buf.WriteString("   -f ntoken_file      Principal Authority NToken file used for authentication\n")
	buf.WriteString("   -i identity         User identity to authenticate as if NToken file is not specified\n")
	buf.WriteString("                       (default=" + defaultIdentity() + ")\n")
	buf.WriteString("   -k                  Disable peer verification of SSL certificates.\n")
	buf.WriteString("   -s host:port        The SOCKS5 proxy to route requests through\n")
	buf.WriteString("   -v                  Verbose mode. Full resource names are included in output (default=false)\n")
	buf.WriteString("   -x                  For user token output, exclude the header name (default=false)\n")
	buf.WriteString("   -z zms_url          Base URL of the ZMS server to use\n")
	buf.WriteString("                       (default ZMS=" + defaultZmsURL() + ")\n")
	buf.WriteString("   -debug              Debug mode. Generates debug NTokens (default=false)\n")
	buf.WriteString("\n")
	buf.WriteString(" type 'zms-cli help' to see all available commands\n")
	buf.WriteString(" type 'zms-cli help [command]' for usage of the specified command\n")
	return buf.String()
}

func loadNtokenFromFile(fileName string) (string, error) {
	buf, err := ioutil.ReadFile(fileName)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(buf)), nil
}

func main() {
	pZMS := flag.String("z", defaultZmsURL(), "Base URL of the ZMS server to use")
	pIdentity := flag.String("i", defaultIdentity(), "the identity to authenticate as")
	pNtokenFile := flag.String("f", "", "ntoken file path")
	pCACert := flag.String("c", "", "CA Certificate file path")
	pVerbose := flag.Bool("v", false, "verbose mode. Full resource names are included in output")
	pBulkmode := flag.Bool("b", false, "bulk mode. Do not display updated role/policy/service in output")
	pProductIDSupport := flag.Bool("p", false, "Top Level Domain add operations require product ids")
	pAddSelf := flag.Bool("addself", true, "Add self to domain admin role")
	pDomain := flag.String("d", "", "The domain for the command to execute in. If not specified, only certain commands are available")
	pUserDomain := flag.String("u", "user", "User domain name as configured in Athenz systems")
	pHomeDomain := flag.String("h", "user", "Home domain name as configured in Athenz systems")
	pSocks := flag.String("s", defaultSocksProxy(), "The SOCKS5 proxy to route requests through, i.e. 127.0.0.1:1080")
	pSkipVerify := flag.Bool("k", false, "Disable peer verification of SSL certificates")
	pDebug := flag.Bool("debug", defaultDebug(), "debug mode (for authentication, mainly)")
	pAuditRef := flag.String("a", "", "Audit Reference Token if auditing is enabled for the domain")
	pExcludeHeader := flag.Bool("x", false, "Exclude header in user-token output")
	flag.Usage = func() {
		fmt.Println(usage())
	}

	// first we need to parse our arguments based
	// on the flags we defined above

	flag.Parse()

	if *pZMS == "" {
		fmt.Println("No ZMS Url specified")
		return
	}

	// before processing our arguments verify that the zms
	// url is of expected format and return failure right
	// away otherwise the error message is somewhat
	// confusing when the hostname/port is valid but the
	// uri part is invalid resulting in 404 responses

	if !strings.HasSuffix(*pZMS, "/zms/v1") {
		fmt.Println("Invalid ZMS Url specified: " + *pZMS)
		fmt.Println("Valid ZMS Url has the following form: https://<zms-hostname>:<zms-port>/zms/v1")
		return
	}

	// now process our request

	args := flag.Args()
	if len(args) == 0 {
		fmt.Println(usage())
		return
	} else if args[0] == "help" {
		cli := zmscli.Zms{}
		cli.UserDomain = *pUserDomain
		cli.HomeDomain = *pHomeDomain
		if len(args) == 2 {
			fmt.Println(cli.HelpSpecificCommand(false, args[1]))
		} else {
			fmt.Println(cli.HelpListCommand())
		}
		return
	} else if args[0] == "version" {
		if VERSION == "" {
			fmt.Println("zms-cli (development version)")
		} else {
			fmt.Println("zms-cli " + VERSION + " " + BUILD_DATE)
		}
		return
	}

	identity := *pIdentity
	if !strings.Contains(identity, ".") {
		identity = "user." + identity
	}
	if *pSocks == "" {
		pSocks = nil
	}
	if *pCACert == "" {
		pCACert = nil
	}
	tr := getHttpTransport(pSocks, pCACert, *pSkipVerify)
	var ntoken string
	var err error
	if *pNtokenFile == "" {
		if *pDebug {
			ntoken = debugAuthNToken(identity)
		} else {
			ntoken, err = getAuthNToken(identity, "", *pZMS, tr)
			if err != nil {
				log.Fatalf("Unable to get NToken: %v", err)
			}
		}
	} else {
		ntoken, err = loadNtokenFromFile(*pNtokenFile)
		if err != nil {
			log.Fatalf("Unable to load ntoken from file: '%s' err: %v", *pNtokenFile, err)
		}
	}
	var authHeader = "Athenz-Principal-Auth"

	cli := zmscli.Zms{
		ZmsUrl:           *pZMS,
		Identity:         *pIdentity,
		Verbose:          *pVerbose,
		Bulkmode:         *pBulkmode,
		Interactive:      false,
		Domain:           *pDomain,
		AuditRef:         *pAuditRef,
		UserDomain:       *pUserDomain,
		HomeDomain:       *pHomeDomain,
		ProductIdSupport: *pProductIDSupport,
		Debug:            *pDebug,
		AddSelf:          *pAddSelf,
	}
	cli.SetClient(tr, &authHeader, &ntoken)

	if len(args) > 0 {
		switch args[0] {
		case "get-user-token":
			if len(args) == 2 {
				ntoken, err = getAuthNToken(identity, args[1], *pZMS, tr)
				if err != nil {
					log.Fatalf("Unable to get NToken for service: %s error: %v", args[1], err)
				}
			}
			if !(*pExcludeHeader) {
				fmt.Print(authHeader + ": ")
			}
			fmt.Println(ntoken)
			return
		case "repl":
			cli.Interactive = true
			cli.Repl()
			return
		default:
			break
		}
	}
	msg, err := cli.EvalCommand(args)
	if err != nil {
		fmt.Println("***", err)
		os.Exit(1)
	} else if msg != nil {
		fmt.Println(*msg)
		os.Exit(0)
	}
}

func getHttpTransport(socksProxy, caCertFile *string, skipVerify bool) *http.Transport {
	tr := http.Transport{}
	if socksProxy != nil {
		dialer := &net.Dialer{}
		dialSocksProxy, err := proxy.SOCKS5("tcp", *socksProxy, nil, dialer)
		if err == nil {
			tr.Dial = dialSocksProxy.Dial
		}
	}
	if caCertFile != nil || skipVerify {
		config := &tls.Config{}
		if caCertFile != nil {
			capem, err := ioutil.ReadFile(*caCertFile)
			if err != nil {
				log.Fatalf("Unable to read CA Certificate file %s, error: %v", *caCertFile, err)
			}
			certPool := x509.NewCertPool()
			if !certPool.AppendCertsFromPEM(capem) {
				log.Fatalf("Unable to append CA Certificate to pool")
			}
			config.RootCAs = certPool
		}
		if skipVerify {
			config.InsecureSkipVerify = skipVerify
		}
		tr.TLSClientConfig = config
	}
	return &tr
}
<commit_msg>add key/cert option for zms-cli (#791)<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/net/proxy"

	"github.com/yahoo/athenz/clients/go/zms"
	"github.com/yahoo/athenz/libs/go/zmscli"
)

var (
	// VERSION gets set by the build script via the LDFLAGS.
	VERSION string

	// BUILD_DATE gets set by the build script via the LDFLAGS.
	BUILD_DATE string
)

func defaultZmsURL() string {
	s := os.Getenv("ZMS")
	if s != "" {
		return s
	}
	return "https://localhost:4443/zms/v1"
}

func defaultIdentity() string {
	return "user." + os.Getenv("USER")
}

func defaultSocksProxy() string {
	return os.Getenv("SOCKS5_PROXY")
}

func defaultDebug() bool {
	sDebug := os.Getenv("ZMS_DEBUG")
	return sDebug == "true"
}

func debugAuthNToken(identity string) string {
	i := strings.LastIndex(identity, ".")
	domain := identity[0:i]
	name := identity[i+1:]
	var buf bytes.Buffer
	buf.WriteString("v=U1;d=" + domain + ";n=" + name + ";s=fakesignature")
	return "v=U1;d=" + domain + ";n=" + name + ";s=fakesignature"
}

// isFreshFile checks the file's last modification time
// and returns true the file was updated within maxAge
// (file is "fresh"), false otherwise (file is "stale").
func isFreshFile(filename string, maxAge float64) bool {
	info, err := os.Stat(filename)
	if err != nil {
		return false
	}
	delta := time.Since(info.ModTime())
	// return false if duration exceeds maxAge
	tooOld := delta.Minutes() > maxAge
	return !tooOld
}

func getCachedNToken() string {
	ntokenFile := os.Getenv("HOME") + "/.ntoken"
	if isFreshFile(ntokenFile, 45) {
		data, err := ioutil.ReadFile(ntokenFile)
		if err == nil {
			return strings.TrimSpace(string(data))
		}
		fmt.Printf("Couldn't read the file, error: %v\n", err)
	}
	return ""
}

const identityPrefix = "user."

func getAuthNToken(identity, authorizedServices, zmsUrl string, tr *http.Transport) (string, error) {
	// our identity must be user
	if !strings.HasPrefix(identity, identityPrefix) {
		return "", errors.New("identity must start with " + identityPrefix)
	}
	user := identity[5:]
	ntoken := getCachedNToken()
	if ntoken != "" {
		return ntoken, nil
	}

	fmt.Fprintf(os.Stderr, "Enter password for "+user+": ")
	pass, err := terminal.ReadPassword(syscall.Stdin)
	if err != nil {
		return "", err
	}
	password := string(pass)
	// The user types <return> during terminal.ReadPassword() however
	// since echo is turned off the cursor doesn't make it to the next line.
	fmt.Fprint(os.Stderr, "\n")

	data := []byte(user + ":" + password)
	str := base64.StdEncoding.EncodeToString(data)

	var authHeader = "Authorization"
	var authCreds = "Basic " + str
	zmsClient := zms.ZMSClient{
		URL:         zmsUrl,
		Transport:   tr,
		CredsHeader: &authHeader,
		CredsToken:  &authCreds,
		Timeout:     0,
	}
	tok, err := zmsClient.GetUserToken(zms.SimpleName(user), authorizedServices, nil)
	if err != nil {
		return "", fmt.Errorf("Cannot get user token for user: %s error: %v", user, err)
	}
	if tok.Token != "" {
		ntokenFile := os.Getenv("HOME") + "/.ntoken"
		data := []byte(tok.Token)
		ioutil.WriteFile(ntokenFile, data, 0600)
	}
	return string(tok.Token), nil
}

func usage() string {
	var buf bytes.Buffer
	buf.WriteString("usage: zms-cli [flags] command [params]\n")
	buf.WriteString(" flags:\n")
	buf.WriteString("   -a auditRef         Audit Reference Token if auditing is required for domain\n")
	buf.WriteString("   -b                  Bulk import/update mode. Do not read/display updated role/policy/service objects (default=false)\n")
	buf.WriteString("   -c cacert_file      CA Certificate file path\n")
	buf.WriteString("   -cert x509_cert     Athenz X.509 Certificate file for authentication\n")
	buf.WriteString("   -d domain           The domain used for every command that takes a domain argument\n")
	buf.WriteString("   -f ntoken_file      Principal Authority NToken file used for authentication\n")
	buf.WriteString("   -i identity         User identity to authenticate as if NToken file is not specified\n")
	buf.WriteString("                       (default=" + defaultIdentity() + ")\n")
	buf.WriteString("   -k                  Disable peer verification of SSL certificates.\n")
	buf.WriteString("   -key x509_key       Athenz X.509 Key file for authentication\n")
	buf.WriteString("   -s host:port        The SOCKS5 proxy to route requests through\n")
	buf.WriteString("   -v                  Verbose mode. Full resource names are included in output (default=false)\n")
	buf.WriteString("   -x                  For user token output, exclude the header name (default=false)\n")
	buf.WriteString("   -z zms_url          Base URL of the ZMS server to use\n")
	buf.WriteString("                       (default ZMS=" + defaultZmsURL() + ")\n")
	buf.WriteString("   -debug              Debug mode. Generates debug NTokens (default=false)\n")
	buf.WriteString("\n")
	buf.WriteString(" type 'zms-cli help' to see all available commands\n")
	buf.WriteString(" type 'zms-cli help [command]' for usage of the specified command\n")
	return buf.String()
}

func loadNtokenFromFile(fileName string) (string, error) {
	buf, err := ioutil.ReadFile(fileName)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(buf)), nil
}

func main() {
	pZMS := flag.String("z", defaultZmsURL(), "Base URL of the ZMS server to use")
	pIdentity := flag.String("i", defaultIdentity(), "the identity to authenticate as")
	pNtokenFile := flag.String("f", "", "ntoken file path")
	pCACert := flag.String("c", "", "CA Certificate file path")
	pVerbose := flag.Bool("v", false, "verbose mode. Full resource names are included in output")
	pBulkmode := flag.Bool("b", false, "bulk mode. Do not display updated role/policy/service in output")
	pProductIDSupport := flag.Bool("p", false, "Top Level Domain add operations require product ids")
	pAddSelf := flag.Bool("addself", true, "Add self to domain admin role")
	pDomain := flag.String("d", "", "The domain for the command to execute in. If not specified, only certain commands are available")
	pUserDomain := flag.String("u", "user", "User domain name as configured in Athenz systems")
	pHomeDomain := flag.String("h", "home", "Home domain name as configured in Athenz systems")
	pSocks := flag.String("s", defaultSocksProxy(), "The SOCKS5 proxy to route requests through, i.e. 127.0.0.1:1080")
	pSkipVerify := flag.Bool("k", false, "Disable peer verification of SSL certificates")
	pDebug := flag.Bool("debug", defaultDebug(), "debug mode (for authentication, mainly)")
	pAuditRef := flag.String("a", "", "Audit Reference Token if auditing is enabled for the domain")
	pExcludeHeader := flag.Bool("x", false, "Exclude header in user-token output")
	pX509KeyFile := flag.String("key", "", "x.509 private key file for authentication")
	pX509CertFile := flag.String("cert", "", "x.509 certificate key file for authentication")
	flag.Usage = func() {
		fmt.Println(usage())
	}

	// first we need to parse our arguments based
	// on the flags we defined above

	flag.Parse()

	if *pZMS == "" {
		fmt.Println("No ZMS Url specified")
		return
	}

	// before processing our arguments verify that the zms
	// url is of expected format and return failure right
	// away otherwise the error message is somewhat
	// confusing when the hostname/port is valid but the
	// uri part is invalid resulting in 404 responses

	if !strings.HasSuffix(*pZMS, "/zms/v1") {
		fmt.Println("Invalid ZMS Url specified: " + *pZMS)
		fmt.Println("Valid ZMS Url has the following form: https://<zms-hostname>:<zms-port>/zms/v1")
		return
	}

	// now process our request

	args := flag.Args()
	if len(args) == 0 {
		fmt.Println(usage())
		return
	} else if args[0] == "help" {
		cli := zmscli.Zms{}
		cli.UserDomain = *pUserDomain
		cli.HomeDomain = *pHomeDomain
		if len(args) == 2 {
			fmt.Println(cli.HelpSpecificCommand(false, args[1]))
		} else {
			fmt.Println(cli.HelpListCommand())
		}
		return
	} else if args[0] == "version" {
		if VERSION == "" {
			fmt.Println("zms-cli (development version)")
		} else {
			fmt.Println("zms-cli " + VERSION + " " + BUILD_DATE)
		}
		return
	}

	cli := zmscli.Zms{
		ZmsUrl:           *pZMS,
		Identity:         *pIdentity,
		Verbose:          *pVerbose,
		Bulkmode:         *pBulkmode,
		Interactive:      false,
		Domain:           *pDomain,
		AuditRef:         *pAuditRef,
		UserDomain:       *pUserDomain,
		HomeDomain:       *pHomeDomain,
		ProductIdSupport: *pProductIDSupport,
		Debug:            *pDebug,
		AddSelf:          *pAddSelf,
	}

	if *pX509KeyFile != "" && *pX509CertFile != "" {
		err := cli.SetX509CertClient(*pX509KeyFile, *pX509CertFile, *pCACert, *pSocks, false, *pSkipVerify)
		if err != nil {
			log.Fatalf("Unable to set ZMS x.509 Client: %v\n", err)
		}
	} else {
		identity := *pIdentity
		if !strings.Contains(identity, ".") {
			identity = "user." + identity
		}
		if *pSocks == "" {
			pSocks = nil
		}
		if *pCACert == "" {
			pCACert = nil
		}
		tr := getHttpTransport(pSocks, pCACert, *pSkipVerify)
		var ntoken string
		var err error
		if *pNtokenFile == "" {
			if *pDebug {
				ntoken = debugAuthNToken(identity)
			} else {
				ntoken, err = getAuthNToken(identity, "", *pZMS, tr)
				if err != nil {
					log.Fatalf("Unable to get NToken: %v", err)
				}
			}
		} else {
			ntoken, err = loadNtokenFromFile(*pNtokenFile)
			if err != nil {
				log.Fatalf("Unable to load ntoken from file: '%s' err: %v", *pNtokenFile, err)
			}
		}
		var authHeader = "Athenz-Principal-Auth"
		cli.SetClient(tr, &authHeader, &ntoken)

		if len(args) > 0 && args[0] == "get-user-token" {
			if len(args) == 2 {
				ntoken, err = getAuthNToken(identity, args[1], *pZMS, tr)
				if err != nil {
					log.Fatalf("Unable to get NToken for service: %s error: %v", args[1], err)
				}
			}
			if !(*pExcludeHeader) {
				fmt.Print(authHeader + ": ")
			}
			fmt.Println(ntoken)
			return
		}
	}

	if len(args) > 0 && args[0] == "repl" {
		cli.Interactive = true
		cli.Repl()
		return
	}

	msg, err := cli.EvalCommand(args)
	if err != nil {
		fmt.Println("***", err)
		os.Exit(1)
	} else if msg != nil {
		fmt.Println(*msg)
		os.Exit(0)
	}
}

func getHttpTransport(socksProxy, caCertFile *string, skipVerify bool) *http.Transport {
	tr := http.Transport{}
	if socksProxy != nil {
		dialer := &net.Dialer{}
		dialSocksProxy, err := proxy.SOCKS5("tcp", *socksProxy, nil, dialer)
		if err == nil {
			tr.Dial = dialSocksProxy.Dial
		}
	}
	if caCertFile != nil || skipVerify {
		config := &tls.Config{}
		if caCertFile != nil {
			capem, err := ioutil.ReadFile(*caCertFile)
			if err != nil {
				log.Fatalf("Unable to read CA Certificate file %s, error: %v", *caCertFile, err)
			}
			certPool := x509.NewCertPool()
			if !certPool.AppendCertsFromPEM(capem) {
				log.Fatalf("Unable to append CA Certificate to pool")
			}
			config.RootCAs = certPool
		}
		if skipVerify {
			config.InsecureSkipVerify = skipVerify
		}
		tr.TLSClientConfig = config
	}
	return &tr
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"github.com/siddontang/go/arena"
	"github.com/siddontang/ledisdb/config"
	"github.com/siddontang/ledisdb/ledis"
	"github.com/siddontang/ledisdb/server"
	"net"
	"os"
	"runtime"
	"time"
)

var KB = config.KB
var MB = config.MB
var GB = config.GB

var addr = flag.String("addr", ":6380", "listen addr")
var name = flag.String("db_name", "", "db name")

var ldb *ledis.Ledis
var db *ledis.DB

func setRocksDB(cfg *config.RocksDBConfig) {
	cfg.BlockSize = 64 * KB
	cfg.WriteBufferSize = 64 * MB
	cfg.MaxWriteBufferNum = 2
	cfg.MaxBytesForLevelBase = 512 * MB
	cfg.TargetFileSizeBase = 64 * MB
	cfg.BackgroundThreads = 4
	cfg.HighPriorityBackgroundThreads = 1
	cfg.MaxBackgroundCompactions = 3
	cfg.MaxBackgroundFlushes = 1
	cfg.CacheSize = 512 * MB
	cfg.EnableStatistics = true
	cfg.StatsDumpPeriodSec = 5
	cfg.Level0FileNumCompactionTrigger = 8
	cfg.MaxBytesForLevelMultiplier = 8
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()
	l, err := net.Listen("tcp", *addr)

	println("listen", *addr)

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

	if len(*name) > 0 {
		cfg := config.NewConfigDefault()
		cfg.DataDir = "./var/ledis_respbench"
		cfg.DBName = *name
		os.RemoveAll(cfg.DBPath)
		defer os.RemoveAll(cfg.DBPath)

		os.MkdirAll(cfg.DBPath, 0755)

		cfg.LevelDB.BlockSize = 32 * KB
		cfg.LevelDB.CacheSize = 512 * MB
		cfg.LevelDB.WriteBufferSize = 64 * MB
		cfg.LevelDB.MaxOpenFiles = 1000

		setRocksDB(&cfg.RocksDB)

		ldb, err = ledis.Open(cfg)
		if err != nil {
			println(err.Error())
			return
		}

		db, _ = ldb.Select(0)
	}

	for {
		c, err := l.Accept()
		if err != nil {
			println(err.Error())
			continue
		}
		go run(c)
	}
}

func run(c net.Conn) {
	//buf := make([]byte, 10240)
	ok := []byte("+OK\r\n")
	data := []byte("$4096\r\n")
	data = append(data, make([]byte, 4096)...)
	data = append(data, "\r\n"...)

	var rt time.Duration
	var wt time.Duration

	rb := bufio.NewReaderSize(c, 10240)
	wb := bufio.NewWriterSize(c, 10240)

	a := arena.NewArena(10240)

	for {
		t1 := time.Now()

		a.Reset()

		req, err := server.ReadRequest(rb, a)

		if err != nil {
			break
		}
		t2 := time.Now()

		rt += t2.Sub(t1)

		cmd := string(bytes.ToUpper(req[0]))
		switch cmd {
		case "SET":
			if db != nil {
				db.Set(req[1], req[2])
			}
			wb.Write(ok)
		case "GET":
			if db != nil {
				d, _ := db.GetSlice(req[1])
				if d == nil {
					wb.Write(data)
				} else {
					wb.WriteString(fmt.Sprintf("$%d\r\n", d.Size()))
					wb.Write(d.Data())
					wb.WriteString("\r\n")
					d.Free()
				}
			} else {
				wb.Write(data)
			}
		default:
			wb.WriteString(fmt.Sprintf("-Err %s Not Supported Now\r\n", req[0]))
		}

		wb.Flush()

		t3 := time.Now()
		wt += t3.Sub(t2)
	}

	fmt.Printf("rt:%s wt:%s\n", rt.String(), wt.String())
}
<commit_msg>update bench  tool<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"github.com/siddontang/go/arena"
	"github.com/siddontang/ledisdb/config"
	"github.com/siddontang/ledisdb/ledis"
	"github.com/siddontang/ledisdb/server"
	"net"
	"os"
	"runtime"
	"time"
)

var KB = config.KB
var MB = config.MB
var GB = config.GB

var addr = flag.String("addr", ":6380", "listen addr")
var name = flag.String("db_name", "", "db name")

var ldb *ledis.Ledis
var db *ledis.DB

func setRocksDB(cfg *config.RocksDBConfig) {
	cfg.BlockSize = 64 * KB
	cfg.WriteBufferSize = 64 * MB
	cfg.MaxWriteBufferNum = 2
	cfg.MaxBytesForLevelBase = 512 * MB
	cfg.TargetFileSizeBase = 64 * MB
	cfg.BackgroundThreads = 4
	cfg.HighPriorityBackgroundThreads = 1
	cfg.MaxBackgroundCompactions = 3
	cfg.MaxBackgroundFlushes = 1
	cfg.CacheSize = 512 * MB
	cfg.EnableStatistics = true
	cfg.StatsDumpPeriodSec = 5
	cfg.Level0FileNumCompactionTrigger = 8
	cfg.MaxBytesForLevelMultiplier = 8
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	flag.Parse()
	l, err := net.Listen("tcp", *addr)

	println("listen", *addr)

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

	if len(*name) > 0 {
		cfg := config.NewConfigDefault()
		cfg.DataDir = "./var/ledis_respbench"
		cfg.DBName = *name
		os.RemoveAll(cfg.DBPath)
		defer os.RemoveAll(cfg.DBPath)

		os.MkdirAll(cfg.DBPath, 0755)

		cfg.LevelDB.BlockSize = 32 * KB
		cfg.LevelDB.CacheSize = 512 * MB
		cfg.LevelDB.WriteBufferSize = 64 * MB
		cfg.LevelDB.MaxOpenFiles = 1000

		setRocksDB(&cfg.RocksDB)

		ldb, err = ledis.Open(cfg)
		if err != nil {
			println(err.Error())
			return
		}

		db, _ = ldb.Select(0)
	}

	for {
		c, err := l.Accept()
		if err != nil {
			println(err.Error())
			continue
		}
		go run(c)
	}
}

func run(c net.Conn) {
	//buf := make([]byte, 10240)
	ok := []byte("+OK\r\n")
	data := []byte("$4096\r\n")
	data = append(data, make([]byte, 4096)...)
	data = append(data, "\r\n"...)

	var rt time.Duration
	var wt time.Duration
	var st time.Duration
	var gt time.Duration

	rb := bufio.NewReaderSize(c, 10240)
	wb := bufio.NewWriterSize(c, 10240)

	a := arena.NewArena(10240)

	for {
		t1 := time.Now()

		a.Reset()

		req, err := server.ReadRequest(rb, a)

		if err != nil {
			break
		}
		t2 := time.Now()

		rt += t2.Sub(t1)

		cmd := string(bytes.ToUpper(req[0]))
		switch cmd {
		case "SET":
			if db != nil {
				db.Set(req[1], req[2])
				st += time.Now().Sub(t2)
			}
			wb.Write(ok)
		case "GET":
			if db != nil {
				d, _ := db.GetSlice(req[1])
				gt += time.Now().Sub(t2)
				if d == nil {
					wb.Write(data)
				} else {
					wb.WriteString(fmt.Sprintf("$%d\r\n", d.Size()))
					wb.Write(d.Data())
					wb.WriteString("\r\n")
					d.Free()
				}
			} else {
				wb.Write(data)
			}
		default:
			wb.WriteString(fmt.Sprintf("-Err %s Not Supported Now\r\n", req[0]))
		}

		wb.Flush()

		t3 := time.Now()
		wt += t3.Sub(t2)
	}

	fmt.Printf("rt:%s wt %s, gt:%s, st:%s\n", rt.String(), wt.String(), gt.String(), st.String())
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The OpenEBS Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Common types in the maya-agent
package types
<commit_msg>chore(comment): fixed golint issue in doc.go #1470<commit_after>/*
Copyright 2017 The OpenEBS Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package types contains common types in the maya-agent.
package types
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/datawire/teleproxy/pkg/consulwatch"

	"github.com/datawire/teleproxy/pkg/watt"

	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/supervisor"
)

type aggIsolator struct {
	snapshots     chan string
	k8sWatches    chan []KubernetesWatchSpec
	consulWatches chan []ConsulWatchSpec
	aggregator    *aggregator
	sup           *supervisor.Supervisor
	done          chan struct{}
	t             *testing.T
	cancel        context.CancelFunc
}

func newAggIsolator(t *testing.T, requiredKinds []string, watchHook WatchHook) *aggIsolator {
	// aggregator uses zero length channels for its inputs so we can
	// control the total ordering of all inputs and therefore
	// intentionally trigger any order of events we want to test
	iso := &aggIsolator{
		// we need to create buffered channels for outputs
		// because nothing is asynchronously reading them in
		// the test
		k8sWatches:    make(chan []KubernetesWatchSpec, 100),
		consulWatches: make(chan []ConsulWatchSpec, 100),
		snapshots:     make(chan string, 100),
		// for signaling when the isolator is done
		done: make(chan struct{}),
	}
	iso.aggregator = NewAggregator(iso.snapshots, iso.k8sWatches, iso.consulWatches, requiredKinds, watchHook, nil)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	iso.cancel = cancel
	iso.sup = supervisor.WithContext(ctx)
	iso.sup.Supervise(&supervisor.Worker{
		Name: "aggregator",
		Work: iso.aggregator.Work,
	})
	iso.t = t
	return iso
}

func startAggIsolator(t *testing.T, requiredKinds []string, watchHook WatchHook) *aggIsolator {
	iso := newAggIsolator(t, requiredKinds, watchHook)
	iso.Start()
	return iso
}

func (iso *aggIsolator) Start() {
	go func() {
		errs := iso.sup.Run()
		if len(errs) > 0 {
			iso.t.Errorf("unexpected errors: %v", errs)
		}
		close(iso.done)
	}()
}

func (iso *aggIsolator) Stop() {
	iso.sup.Shutdown()
	iso.cancel()
	<-iso.done
}

func resources(input string) []k8s.Resource {
	result, err := k8s.ParseResources("aggregator-test", input)
	if err != nil {
		panic(err)
	}
	return result
}

var (
	SERVICES = resources(`
---
kind: Service
apiVersion: v1
metadata:
  name: foo
spec:
  selector:
    pod: foo
  ports:
  - protocol: TCP
    port: 80
    targetPort: 80
`)
	RESOLVER = resources(`
---
kind: ConfigMap
apiVersion: v1
metadata:
  name: bar
  annotations:
    "getambassador.io/consul-resolver": "true"
data:
  consulAddress: "127.0.0.1:8500"
  datacenter: "dc1"
  service: "bar"
`)
)

// make sure we shutdown even before achieving a bootstrapped state
func TestAggregatorShutdown(t *testing.T) {
	iso := startAggIsolator(t, nil, nil)
	defer iso.Stop()
}

var WATCH = ConsulWatchSpec{
	ConsulAddress: "127.0.0.1:8500",
	Datacenter:    "dc1",
	ServiceName:   "bar",
}

// Check that we bootstrap properly... this means *not* emitting a
// snapshot until we have:
//
//   a) achieved synchronization with the kubernetes API server
//
//   b) received (possibly empty) endpoint info about all referenced
//      consul services...
func TestAggregatorBootstrap(t *testing.T) {
	watchHook := func(p *supervisor.Process, snapshot string) WatchSet {
		if strings.Contains(snapshot, "configmap") {
			return WatchSet{
				ConsulWatches: []ConsulWatchSpec{WATCH},
			}
		} else {
			return WatchSet{}
		}
	}
	iso := startAggIsolator(t, []string{"service", "configmap"}, watchHook)
	defer iso.Stop()

	// initial kubernetes state is just services
	iso.aggregator.KubernetesEvents <- k8sEvent{"", "service", SERVICES}
	// whenever the aggregator sees updated k8s state, it should
	// send an update to the consul watch manager, in this case it
	// will be empty because there are no resolvers yet
	expect(t, iso.consulWatches, []ConsulWatchSpec(nil))

	// we should not generate a snapshot yet because we specified
	// configmaps are required
	expect(t, iso.snapshots, Timeout(100*time.Millisecond))

	// the configmap references a consul service, so we shouldn't
	// get a snapshot yet, but we should get watches
	iso.aggregator.KubernetesEvents <- k8sEvent{"", "configmap", RESOLVER}
	fmt.Println("1")
	expect(t, iso.snapshots, Timeout(100*time.Millisecond))
	fmt.Println("2")
	expect(t, iso.consulWatches, func(watches []ConsulWatchSpec) bool {
		if len(watches) != 1 {
			t.Logf("expected 1 watch, got %d watches", len(watches))
			return false
		}

		if watches[0].ServiceName != "bar" {
			return false
		}

		return true
	})

	// now lets send in the first endpoints, and we should get a
	// snapshot
	iso.aggregator.ConsulEvents <- consulEvent{
		WATCH.Hash(),
		consulwatch.Endpoints{
			Service: "bar",
			Endpoints: []consulwatch.Endpoint{
				{
					Service: "bar",
					Address: "1.2.3.4",
					Port:    80,
				},
			},
		},
	}

	expect(t, iso.snapshots, func(snapshot string) bool {
		s := &watt.Snapshot{}
		err := json.Unmarshal([]byte(snapshot), s)
		if err != nil {
			return false
		}
		_, ok := s.Consul.Endpoints["bar"]
		return ok
	})
}
<commit_msg>remove printlns<commit_after>package main

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"github.com/datawire/teleproxy/pkg/consulwatch"

	"github.com/datawire/teleproxy/pkg/watt"

	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/supervisor"
)

type aggIsolator struct {
	snapshots     chan string
	k8sWatches    chan []KubernetesWatchSpec
	consulWatches chan []ConsulWatchSpec
	aggregator    *aggregator
	sup           *supervisor.Supervisor
	done          chan struct{}
	t             *testing.T
	cancel        context.CancelFunc
}

func newAggIsolator(t *testing.T, requiredKinds []string, watchHook WatchHook) *aggIsolator {
	// aggregator uses zero length channels for its inputs so we can
	// control the total ordering of all inputs and therefore
	// intentionally trigger any order of events we want to test
	iso := &aggIsolator{
		// we need to create buffered channels for outputs
		// because nothing is asynchronously reading them in
		// the test
		k8sWatches:    make(chan []KubernetesWatchSpec, 100),
		consulWatches: make(chan []ConsulWatchSpec, 100),
		snapshots:     make(chan string, 100),
		// for signaling when the isolator is done
		done: make(chan struct{}),
	}
	iso.aggregator = NewAggregator(iso.snapshots, iso.k8sWatches, iso.consulWatches, requiredKinds, watchHook, nil)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	iso.cancel = cancel
	iso.sup = supervisor.WithContext(ctx)
	iso.sup.Supervise(&supervisor.Worker{
		Name: "aggregator",
		Work: iso.aggregator.Work,
	})
	iso.t = t
	return iso
}

func startAggIsolator(t *testing.T, requiredKinds []string, watchHook WatchHook) *aggIsolator {
	iso := newAggIsolator(t, requiredKinds, watchHook)
	iso.Start()
	return iso
}

func (iso *aggIsolator) Start() {
	go func() {
		errs := iso.sup.Run()
		if len(errs) > 0 {
			iso.t.Errorf("unexpected errors: %v", errs)
		}
		close(iso.done)
	}()
}

func (iso *aggIsolator) Stop() {
	iso.sup.Shutdown()
	iso.cancel()
	<-iso.done
}

func resources(input string) []k8s.Resource {
	result, err := k8s.ParseResources("aggregator-test", input)
	if err != nil {
		panic(err)
	}
	return result
}

var (
	SERVICES = resources(`
---
kind: Service
apiVersion: v1
metadata:
  name: foo
spec:
  selector:
    pod: foo
  ports:
  - protocol: TCP
    port: 80
    targetPort: 80
`)
	RESOLVER = resources(`
---
kind: ConfigMap
apiVersion: v1
metadata:
  name: bar
  annotations:
    "getambassador.io/consul-resolver": "true"
data:
  consulAddress: "127.0.0.1:8500"
  datacenter: "dc1"
  service: "bar"
`)
)

// make sure we shutdown even before achieving a bootstrapped state
func TestAggregatorShutdown(t *testing.T) {
	iso := startAggIsolator(t, nil, nil)
	defer iso.Stop()
}

var WATCH = ConsulWatchSpec{
	ConsulAddress: "127.0.0.1:8500",
	Datacenter:    "dc1",
	ServiceName:   "bar",
}

// Check that we bootstrap properly... this means *not* emitting a
// snapshot until we have:
//
//   a) achieved synchronization with the kubernetes API server
//
//   b) received (possibly empty) endpoint info about all referenced
//      consul services...
func TestAggregatorBootstrap(t *testing.T) {
	watchHook := func(p *supervisor.Process, snapshot string) WatchSet {
		if strings.Contains(snapshot, "configmap") {
			return WatchSet{
				ConsulWatches: []ConsulWatchSpec{WATCH},
			}
		} else {
			return WatchSet{}
		}
	}
	iso := startAggIsolator(t, []string{"service", "configmap"}, watchHook)
	defer iso.Stop()

	// initial kubernetes state is just services
	iso.aggregator.KubernetesEvents <- k8sEvent{"", "service", SERVICES}
	// whenever the aggregator sees updated k8s state, it should
	// send an update to the consul watch manager, in this case it
	// will be empty because there are no resolvers yet
	expect(t, iso.consulWatches, []ConsulWatchSpec(nil))

	// we should not generate a snapshot yet because we specified
	// configmaps are required
	expect(t, iso.snapshots, Timeout(100*time.Millisecond))

	// the configmap references a consul service, so we shouldn't
	// get a snapshot yet, but we should get watches
	iso.aggregator.KubernetesEvents <- k8sEvent{"", "configmap", RESOLVER}
	expect(t, iso.snapshots, Timeout(100*time.Millisecond))
	expect(t, iso.consulWatches, func(watches []ConsulWatchSpec) bool {
		if len(watches) != 1 {
			t.Logf("expected 1 watch, got %d watches", len(watches))
			return false
		}

		if watches[0].ServiceName != "bar" {
			return false
		}

		return true
	})

	// now lets send in the first endpoints, and we should get a
	// snapshot
	iso.aggregator.ConsulEvents <- consulEvent{
		WATCH.Hash(),
		consulwatch.Endpoints{
			Service: "bar",
			Endpoints: []consulwatch.Endpoint{
				{
					Service: "bar",
					Address: "1.2.3.4",
					Port:    80,
				},
			},
		},
	}

	expect(t, iso.snapshots, func(snapshot string) bool {
		s := &watt.Snapshot{}
		err := json.Unmarshal([]byte(snapshot), s)
		if err != nil {
			return false
		}
		_, ok := s.Consul.Endpoints["bar"]
		return ok
	})
}
<|endoftext|>


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

package main

import "shifter/cmd"

func main() {
	cmd.Execute()
}
<commit_msg>Fully qualifying import paths<commit_after>/*
Copyright 2019 Google LLC
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import "github.com/google/shifter/cmd"

func main() {
	cmd.Execute()
}
<|endoftext|>


<commit_before>package nano

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
)

//Stores the path and the handler
type Route struct {
	method  string
	path    string
	handler http.Handler
}

//Matches the url against its own path, returns true when there is a match
func (r *Route) match(req *http.Request) bool {
	if !r.matchURL(req.URL.String()) {
		return false
	}
	//See if http method matches the route
	if req.Method != r.method {
		return false
	}

	return true
}

func (r *Route) matchURL(rawURL string) bool {
	url := strings.Split(rawURL, "/")
	path := strings.Split(r.path, "/")

	reg, _ := regexp.Compile(":[a-zA-Z0-9]")
	for i, c := range url {
		fmt.Println("Comparing: " + c + " to: " + path[i])

		if len(url) != len(path) {
			return false
		}

		if c != path[i] && !(reg.MatchString(path[i])) {
			return false
		}
	}
	return true
}

//Handle executes the handler function attached to this route
func (r *Route) Handle(w http.ResponseWriter, req *http.Request) {
	r.handler.ServeHTTP(w, req)
}

func (r *Route) parseVars(req *http.Request) {
	Vars = make(map[string]string)
	for k, v := range req.URL.Query() {
		Vars[k] = v[0]
	}

	reg, _ := regexp.Compile(":[a-zA-Z0-9]")
	url := strings.Split(req.URL.String(), "/")
	path := strings.Split(r.path, "/")

	for i, p := range path {
		if reg.MatchString(p) {
			fmt.Println("test: " + p[1:len(p)])
			Vars[p[1:len(p)]] = url[i]
		}
	}
}
<commit_msg>Removed obsolute print<commit_after>package nano

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
)

//Stores the path and the handler
type Route struct {
	method  string
	path    string
	handler http.Handler
}

//Matches the url against its own path, returns true when there is a match
func (r *Route) match(req *http.Request) bool {
	if !r.matchURL(req.URL.String()) {
		return false
	}
	//See if http method matches the route
	if req.Method != r.method {
		return false
	}

	return true
}

func (r *Route) matchURL(rawURL string) bool {
	url := strings.Split(rawURL, "/")
	path := strings.Split(r.path, "/")
	reg, _ := regexp.Compile(":[a-zA-Z0-9]")

	for i, c := range url {
		fmt.Println("Comparing: " + c + " to: " + path[i])

		if len(url) != len(path) {
			return false
		}

		if c != path[i] && !(reg.MatchString(path[i])) {
			return false
		}
	}
	return true
}

//Handle executes the handler function attached to this route
func (r *Route) Handle(w http.ResponseWriter, req *http.Request) {
	r.handler.ServeHTTP(w, req)
}

func (r *Route) parseVars(req *http.Request) {
	Vars = make(map[string]string)
	for k, v := range req.URL.Query() {
		Vars[k] = v[0]
	}

	reg, _ := regexp.Compile(":[a-zA-Z0-9]")
	url := strings.Split(req.URL.String(), "/")
	path := strings.Split(r.path, "/")

	for i, p := range path {
		if reg.MatchString(p) {
			Vars[p[1:len(p)]] = url[i]
		}
	}
}
<|endoftext|>


<commit_before>// +build darwin

package cpu

import (
	"context"
	"runtime"
	"strconv"
	"strings"

	"golang.org/x/sys/unix"
)

// sys/resource.h
const (
	CPUser    = 0
	CPNice    = 1
	CPSys     = 2
	CPIntr    = 3
	CPIdle    = 4
	CPUStates = 5
)

// default value. from time.h
var ClocksPerSec = float64(128)

func Times(percpu bool) ([]TimesStat, error) {
	return TimesWithContext(context.Background(), percpu)
}

func TimesWithContext(ctx context.Context, percpu bool) ([]TimesStat, error) {
	if percpu {
		return perCPUTimes()
	}

	return allCPUTimes()
}

// Returns only one CPUInfoStat on FreeBSD
func Info() ([]InfoStat, error) {
	return InfoWithContext(context.Background())
}

func InfoWithContext(ctx context.Context) ([]InfoStat, error) {
	var ret []InfoStat

	c := InfoStat{}
	c.ModelName, _ = unix.Sysctl("machdep.cpu.brand_string")
	family, _ := unix.SysctlUint32("machdep.cpu.family")
	c.Family = strconv.FormatUint(uint64(family), 10)
	model, _ := unix.SysctlUint32("machdep.cpu.model")
	c.Model = strconv.FormatUint(uint64(model), 10)
	stepping, _ := unix.SysctlUint32("machdep.cpu.stepping")
	c.Stepping = int32(stepping)
	features, err := unix.Sysctl("machdep.cpu.features")
	if err == nil {
		for _, v := range strings.Fields(features) {
			c.Flags = append(c.Flags, strings.ToLower(v))
		}
	}
	leaf7Features, err := unix.Sysctl("machdep.cpu.leaf7_features")
	if err == nil {
		for _, v := range strings.Fields(leaf7Features) {
			c.Flags = append(c.Flags, strings.ToLower(v))
		}
	}
	extfeatures, err := unix.Sysctl("machdep.cpu.extfeatures")
	if err == nil {
		for _, v := range strings.Fields(extfeatures) {
			c.Flags = append(c.Flags, strings.ToLower(v))
		}
	}
	cores, _ := unix.SysctlUint32("machdep.cpu.core_count")
	c.Cores = int32(cores)
	cacheSize, _ := unix.SysctlUint32("machdep.cpu.cache.size")
	c.CacheSize = int32(cacheSize)
	c.VendorID, _ = unix.Sysctl("machdep.cpu.vendor")

	// Use the rated frequency of the CPU. This is a static value and does not
	// account for low power or Turbo Boost modes.
	cpuFrequency, err := unix.SysctlUint64("hw.cpufrequency")
	if err != nil {
		return ret, err
	}
	c.Mhz = float64(cpuFrequency) / 1000000.0

	return append(ret, c), nil
}

func CountsWithContext(ctx context.Context, logical bool) (int, error) {
	return runtime.NumCPU(), nil
}
<commit_msg>retrieve a cpu count depends on an boolean argument<commit_after>// +build darwin

package cpu

import (
	"context"
	"runtime"
	"strconv"
	"strings"

	"golang.org/x/sys/unix"
)

// sys/resource.h
const (
	CPUser    = 0
	CPNice    = 1
	CPSys     = 2
	CPIntr    = 3
	CPIdle    = 4
	CPUStates = 5
)

// default value. from time.h
var ClocksPerSec = float64(128)

func Times(percpu bool) ([]TimesStat, error) {
	return TimesWithContext(context.Background(), percpu)
}

func TimesWithContext(ctx context.Context, percpu bool) ([]TimesStat, error) {
	if percpu {
		return perCPUTimes()
	}

	return allCPUTimes()
}

// Returns only one CPUInfoStat on FreeBSD
func Info() ([]InfoStat, error) {
	return InfoWithContext(context.Background())
}

func InfoWithContext(ctx context.Context) ([]InfoStat, error) {
	var ret []InfoStat

	c := InfoStat{}
	c.ModelName, _ = unix.Sysctl("machdep.cpu.brand_string")
	family, _ := unix.SysctlUint32("machdep.cpu.family")
	c.Family = strconv.FormatUint(uint64(family), 10)
	model, _ := unix.SysctlUint32("machdep.cpu.model")
	c.Model = strconv.FormatUint(uint64(model), 10)
	stepping, _ := unix.SysctlUint32("machdep.cpu.stepping")
	c.Stepping = int32(stepping)
	features, err := unix.Sysctl("machdep.cpu.features")
	if err == nil {
		for _, v := range strings.Fields(features) {
			c.Flags = append(c.Flags, strings.ToLower(v))
		}
	}
	leaf7Features, err := unix.Sysctl("machdep.cpu.leaf7_features")
	if err == nil {
		for _, v := range strings.Fields(leaf7Features) {
			c.Flags = append(c.Flags, strings.ToLower(v))
		}
	}
	extfeatures, err := unix.Sysctl("machdep.cpu.extfeatures")
	if err == nil {
		for _, v := range strings.Fields(extfeatures) {
			c.Flags = append(c.Flags, strings.ToLower(v))
		}
	}
	cores, _ := unix.SysctlUint32("machdep.cpu.core_count")
	c.Cores = int32(cores)
	cacheSize, _ := unix.SysctlUint32("machdep.cpu.cache.size")
	c.CacheSize = int32(cacheSize)
	c.VendorID, _ = unix.Sysctl("machdep.cpu.vendor")

	// Use the rated frequency of the CPU. This is a static value and does not
	// account for low power or Turbo Boost modes.
	cpuFrequency, err := unix.SysctlUint64("hw.cpufrequency")
	if err != nil {
		return ret, err
	}
	c.Mhz = float64(cpuFrequency) / 1000000.0

	return append(ret, c), nil
}

func CountsWithContext(ctx context.Context, logical bool) (int, error) {
	var cpuArgument string
	if logical {
		cpuArgument = "hw.logicalcpu"
	} else {
		cpuArgument = "hw.physicalcpu"
	}

	count, err := unix.SysctlUint32(cpuArgument)

	if err != nil {
		return runtime.NumCPU(), nil
	}

	return int(count), nil
}
<|endoftext|>


<commit_before>package index

import (
	"context"
	"database/sql"
	"time"

	"github.com/lib/pq"
	log "github.com/sirupsen/logrus"
)

const UpdateBatchSize = 1000

type FingerprintStore struct {
	db *sql.DB
}

func NewFingerprintStore(db *sql.DB) *FingerprintStore {
	return &FingerprintStore{
		db: db,
	}
}

func (s *FingerprintStore) GetMaxID(ctx context.Context) (int, error) {
	row := s.db.QueryRowContext(ctx, "SELECT max(id) FROM fingerprint")
	var id int
	err := row.Scan(&id)
	if err != nil {
		if err == sql.ErrNoRows {
			return 0, nil
		}
		return 0, err
	}
	return id, nil
}

func (s *FingerprintStore) GetNextFingerprints(ctx context.Context, lastID uint32, limit int) ([]FingerprintInfo, error) {
	rows, err := s.db.QueryContext(ctx, "SELECT id, acoustid_extract_query(fingerprint) FROM fingerprint WHERE id > $1 ORDER BY id LIMIT $2", lastID, limit)
	if err != nil {
		return nil, err
	}
	var fingerprints []FingerprintInfo
	for rows.Next() {
		var id uint32
		var signedHashes []int32
		err = rows.Scan(&id, pq.Array(&signedHashes))
		if err != nil {
			return nil, err
		}
		hashes := make([]uint32, len(signedHashes))
		for i, hash := range signedHashes {
			hashes[i] = uint32(hash)
		}
		fingerprints = append(fingerprints, FingerprintInfo{ID: id, Hashes: hashes})
	}
	err = rows.Close()
	if err != nil {
		return nil, err
	}
	return fingerprints, nil
}

type UpdaterConfig struct {
	Database *DatabaseConfig
	Index    *IndexConfig
	Debug    bool
}

func NewUpdaterConfig() *UpdaterConfig {
	return &UpdaterConfig{
		Database: NewDatabaseConfig(),
		Index:    NewIndexConfig(),
	}
}

func RunUpdater(cfg *UpdaterConfig) {
	if cfg.Debug {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.InfoLevel)
	}

	db, err := sql.Open("postgres", cfg.Database.URL().String())
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}
	defer db.Close()

	err = db.Ping()
	if err != nil {
		log.Fatalf("Can't ping the database: %v", err)
	}

	idx, err := ConnectWithConfig(context.Background(), cfg.Index)
	if err != nil {
		log.Fatalf("Failed to connect to index: %s", err)
	}
	defer idx.Close(context.Background())

	fp := NewFingerprintStore(db)

	const NoDelay = 0 * time.Millisecond
	const MinDelay = 10 * time.Millisecond
	const MaxDelay = time.Minute

	var delay time.Duration

	for {
		if delay > NoDelay {
			if delay > MaxDelay {
				delay = MaxDelay
			}
			log.Debugf("Sleeping for %v", delay)
			time.Sleep(delay)
		}

		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()

		lastID, err := GetLastFingerprintID(ctx, idx)
		if err != nil {
			log.Fatalf("Failed to get the last fingerprint ID in index: %s", err)
			delay = MaxDelay
			continue
		}

		fingerprints, err := fp.GetNextFingerprints(ctx, lastID, UpdateBatchSize)
		if err != nil {
			log.Fatalf("Failed to get the next fingerprints to import: %s", err)
			delay = MaxDelay
			continue
		}

		err = MultiInsert(ctx, idx, fingerprints)
		if err != nil {
			log.Fatalf("Failed to import the fingerprints: %s", err)
			delay = MaxDelay
			continue
		}

		fingerprintCount := len(fingerprints)
		if fingerprintCount > 0 {
			lastID = fingerprints[fingerprintCount-1].ID
		}
		log.Infof("Added %d fingerprints up to ID %d", fingerprintCount, lastID)

		if fingerprintCount == 0 {
			if delay > NoDelay {
				delay += (delay * 10) / 100
			} else {
				delay = MinDelay
			}
		} else {
			delay = NoDelay
		}
	}
}
<commit_msg>Read data from PG as Int64Array<commit_after>package index

import (
	"context"
	"database/sql"
	"time"

	"github.com/lib/pq"
	log "github.com/sirupsen/logrus"
)

const UpdateBatchSize = 1000

type FingerprintStore struct {
	db *sql.DB
}

func NewFingerprintStore(db *sql.DB) *FingerprintStore {
	return &FingerprintStore{
		db: db,
	}
}

func (s *FingerprintStore) GetMaxID(ctx context.Context) (int, error) {
	row := s.db.QueryRowContext(ctx, "SELECT max(id) FROM fingerprint")
	var id int
	err := row.Scan(&id)
	if err != nil {
		if err == sql.ErrNoRows {
			return 0, nil
		}
		return 0, err
	}
	return id, nil
}

func (s *FingerprintStore) GetNextFingerprints(ctx context.Context, lastID uint32, limit int) ([]FingerprintInfo, error) {
	rows, err := s.db.QueryContext(ctx, "SELECT id, acoustid_extract_query(fingerprint) FROM fingerprint WHERE id > $1 ORDER BY id LIMIT $2", lastID, limit)
	if err != nil {
		return nil, err
	}
	var fingerprints []FingerprintInfo
	for rows.Next() {
		var id uint32
		var signedHashes pq.Int64Array
		err = rows.Scan(&id, signedHashes)
		if err != nil {
			return nil, err
		}
		hashes := make([]uint32, len(signedHashes))
		for i, hash := range signedHashes {
			hashes[i] = uint32(hash)
		}
		fingerprints = append(fingerprints, FingerprintInfo{ID: id, Hashes: hashes})
	}
	err = rows.Close()
	if err != nil {
		return nil, err
	}
	return fingerprints, nil
}

type UpdaterConfig struct {
	Database *DatabaseConfig
	Index    *IndexConfig
	Debug    bool
}

func NewUpdaterConfig() *UpdaterConfig {
	return &UpdaterConfig{
		Database: NewDatabaseConfig(),
		Index:    NewIndexConfig(),
	}
}

func RunUpdater(cfg *UpdaterConfig) {
	if cfg.Debug {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.InfoLevel)
	}

	db, err := sql.Open("postgres", cfg.Database.URL().String())
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}
	defer db.Close()

	err = db.Ping()
	if err != nil {
		log.Fatalf("Can't ping the database: %v", err)
	}

	idx, err := ConnectWithConfig(context.Background(), cfg.Index)
	if err != nil {
		log.Fatalf("Failed to connect to index: %s", err)
	}
	defer idx.Close(context.Background())

	fp := NewFingerprintStore(db)

	const NoDelay = 0 * time.Millisecond
	const MinDelay = 10 * time.Millisecond
	const MaxDelay = time.Minute

	var delay time.Duration

	for {
		if delay > NoDelay {
			if delay > MaxDelay {
				delay = MaxDelay
			}
			log.Debugf("Sleeping for %v", delay)
			time.Sleep(delay)
		}

		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()

		lastID, err := GetLastFingerprintID(ctx, idx)
		if err != nil {
			log.Fatalf("Failed to get the last fingerprint ID in index: %s", err)
			delay = MaxDelay
			continue
		}

		fingerprints, err := fp.GetNextFingerprints(ctx, lastID, UpdateBatchSize)
		if err != nil {
			log.Fatalf("Failed to get the next fingerprints to import: %s", err)
			delay = MaxDelay
			continue
		}

		err = MultiInsert(ctx, idx, fingerprints)
		if err != nil {
			log.Fatalf("Failed to import the fingerprints: %s", err)
			delay = MaxDelay
			continue
		}

		fingerprintCount := len(fingerprints)
		if fingerprintCount > 0 {
			lastID = fingerprints[fingerprintCount-1].ID
		}
		log.Infof("Added %d fingerprints up to ID %d", fingerprintCount, lastID)

		if fingerprintCount == 0 {
			if delay > NoDelay {
				delay += (delay * 10) / 100
			} else {
				delay = MinDelay
			}
		} else {
			delay = NoDelay
		}
	}
}
<|endoftext|>


<commit_before>package video

import "github.com/32bitkid/bitreader"
import "errors"

var (
	ErrUnexpectedNonZeroByte = errors.New("unexpected non-zero byte")
)

func (self *VideoSequence) next_start_code() (err error) {
	return next_start_code(self)
}

func next_start_code(br bitreader.BitReader) error {
	if !br.IsByteAligned() {
		if _, err := br.ByteAlign(); err != nil {
			return err
		}
	}

	for {
		if val, err := br.Peek32(24); err != nil {
			return err
		} else if val == StartCodePrefix {
			return nil
		}

		if val, err := br.Read32(8); err != nil {
			return err
		} else if val != StuffingByte {
			return ErrUnexpectedNonZeroByte
		}
	}
}

func marker_bit(br bitreader.BitReader) error {
	marker, err := br.ReadBit()
	if err != nil {
		return err
	}
	if marker == false {
		return ErrMissingMarkerBit
	}
	return nil
}
<commit_msg>implementing trash until to skip until a specific start code.<commit_after>package video

import "github.com/32bitkid/bitreader"
import "errors"

var (
	ErrUnexpectedNonZeroByte = errors.New("unexpected non-zero byte")
)

func (self *VideoSequence) next_start_code() (err error) {
	return next_start_code(self)
}

func next_start_code(br bitreader.BitReader) error {
	if !br.IsByteAligned() {
		if _, err := br.ByteAlign(); err != nil {
			return err
		}
	}

	for {
		if val, err := br.Peek32(24); err != nil {
			return err
		} else if val == StartCodePrefix {
			return nil
		}

		if val, err := br.Read32(8); err != nil {
			return err
		} else if val != StuffingByte {
			return ErrUnexpectedNonZeroByte
		}
	}
}

func marker_bit(br bitreader.BitReader) error {
	marker, err := br.ReadBit()
	if err != nil {
		return err
	}
	if marker == false {
		return ErrMissingMarkerBit
	}
	return nil
}

func (br *VideoSequence) TrashUntil(startCode StartCode) error {
	if !br.IsByteAligned() {
		if _, err := br.ByteAlign(); err != nil {
			return err
		}
	}

	for {
		if val, err := br.Peek32(32); err != nil {
			return err
		} else if StartCode(val) == startCode {
			return nil
		} else if err := br.Trash(8); err != nil {
			return err
		}
	}
}
<|endoftext|>


<commit_before>// +build linux

package sysstats

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// CpuRawStats represents *one* CPU raw statistics of a linux system.
type CpuRawStats map[string]uint64

// CpuStats represents *one* CPU statistics of a linux system.
type CpuStats map[string]float64

// CpusRawStats represents *all* the CPU raw statistics of a linux system.
type CpusRawStats map[string]CpuRawStats

// CpusStats represents *all* the CPU statistics of a linux system.
type CpusStats map[string]CpuStats

// getCpuRawStats gets the Cpu raw stats of a linux system from the
// file /proc/stat
// It retuns a CpusRawStats var. CpusRawStats is a map which keys are the
// 'kernel name' of the CPUs (cpu, cpu0, etc). The value of each key is a
// CpuRawStats var with the statistics for that CPU.
// CpusRawStats has the following key:
//  Name      - Name of the Cpu (as it is on /proc/stat: cpu, cpu0,...).
// CpuRawStats has the following keys:
//  User      - Time spent in user mode.
//  Nice      - Time spent in user mode with low priority (nice).
//  System    - Time spent in system mode.
//  Idle      - Time spent in the idle task.
//  Iowait    - Time spent waiting for I/O to complete (since 2.5.41).
//  Irq       - Time servicing interrupts (since 2.6.0-test4).
//  Softirq   - Time servicing softirqs (since 2.6.0-test4).
//  Steal     - Stolen time, which is the time spent in other operating
//              systems when running a virtualized environment (since 2.6.11).
//  Guest     - Time spent running a virtual Cpu for guest operating
//              systems under the control of the Linux kernel (since 2.6.24).
//  GuestNice - Time spent running a niced guest (virtual Cpu for guest
//              operating systems under the control of the Linux kernel)
//              (since 2.6.33).
//  Total     - Total time (not idle).
func getCpuRawStats() (cpusRawStats CpusRawStats, err error) {
	file, err := os.Open("/proc/stat")
	if err != nil {
		return nil, err
	}
	defer file.Close()

	cpusRawStats = CpusRawStats{}

	re := regexp.MustCompile(`^cpu.*$`)

	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		stats := re.FindString(line)
		if stats == "" {
			// No match so no more cpu 'lines'
			break
		}
		fmt.Println(stats)
		cpuName, rawStats, err := parseCpuRawStats(stats)
		if err != nil {
			return nil, err
		}
		cpusRawStats[cpuName] = rawStats
	}

	return cpusRawStats, nil
}

// parseCpuRawStats parses the CPU stats as they are in the file /proc/stat
// stats has the following format:
//   cpu  294 0 309 10612 71 30 0 0 0 0
// cpuName is the name of the CPU (cpu, cpu0, cpu1, etc) and rawStats has
// the following format:
//   map[User:9366 Nice:0 System:5692 Iowait:114 Steal:0 GuestNice:0
//       Idle:1458880 Irq:806 Softirq:0 Guest:0]
func parseCpuRawStats(stats string) (cpuName string, rawStats CpuRawStats,
	err error) {
	rawStats = CpuRawStats{}

	fields := strings.Fields(stats)
	cpuName = fields[0]
	for i := 1; i < len(fields); i++ {
		stat, err := strconv.ParseUint(fields[i], 10, 64)
		if err != nil {
			return "", nil, err
		}
		switch i {
		case 1:
			rawStats[`User`] = stat
		case 2:
			rawStats[`Nice`] = stat
		case 3:
			rawStats[`System`] = stat
		case 4:
			rawStats[`Idle`] = stat
		case 5:
			rawStats[`Iowait`] = stat
		case 6:
			rawStats[`Irq`] = stat
		case 7:
			rawStats[`Softirq`] = stat
		case 8:
			rawStats[`Steal`] = stat
		case 9:
			rawStats[`Guest`] = stat
		case 10:
			rawStats[`GuestNice`] = stat
		}
	}

	return cpuName, rawStats, nil
}

func getCpuStats() (cpuStats CpusRawStats, err error) {
	return getCpuRawStats()
}
<commit_msg>Added getCpuStats with the average CPU stats (%)<commit_after>// +build linux

package sysstats

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// CpuRawStats represents *one* CPU raw statistics of a linux system.
type CpuRawStats map[string]uint64

// CpuStats represents *one* CPU statistics of a linux system.
type CpuStats map[string]float64

// CpusRawStats represents *all* the CPU raw statistics of a linux system.
type CpusRawStats map[string]CpuRawStats

// CpusStats represents *all* the CPU statistics of a linux system.
type CpusStats map[string]CpuStats

// getCpuRawStats gets the Cpu raw stats of a linux system from the
// file /proc/stat
// It retuns a CpusRawStats var. CpusRawStats is a map which keys are the
// 'kernel name' of the CPUs (cpu, cpu0, etc). The value of each key is a
// CpuRawStats var with the statistics for that CPU.
// CpusRawStats has the following key:
//  Name      - Name of the Cpu (as it is on /proc/stat: cpu, cpu0,...).
// CpuRawStats has the following keys:
//  User      - Time spent in user mode.
//  Nice      - Time spent in user mode with low priority (nice).
//  System    - Time spent in system mode.
//  Idle      - Time spent in the idle task.
//  Iowait    - Time spent waiting for I/O to complete (since 2.5.41).
//  Irq       - Time servicing interrupts (since 2.6.0-test4).
//  Softirq   - Time servicing softirqs (since 2.6.0-test4).
//  Steal     - Stolen time, which is the time spent in other operating
//              systems when running a virtualized environment (since 2.6.11).
//  Guest     - Time spent running a virtual Cpu for guest operating
//              systems under the control of the Linux kernel (since 2.6.24).
//  GuestNice - Time spent running a niced guest (virtual Cpu for guest
//              operating systems under the control of the Linux kernel)
//              (since 2.6.33).
//  Total     - Total time.
func getCpuRawStats() (cpusRawStats CpusRawStats, err error) {
	file, err := os.Open("/proc/stat")
	if err != nil {
		return nil, err
	}
	defer file.Close()

	cpusRawStats = CpusRawStats{}

	re := regexp.MustCompile(`^cpu.*$`)

	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		stats := re.FindString(line)
		if stats == "" {
			// No match so no more cpu 'lines'
			break
		}
		fmt.Println(stats)
		cpuName, rawStats, err := parseCpuRawStats(stats)
		if err != nil {
			return nil, err
		}
		cpusRawStats[cpuName] = rawStats
	}

	return cpusRawStats, nil
}

// parseCpuRawStats parses the CPU stats as they are in the file /proc/stat
// stats has the following format:
//   cpu  294 0 309 10612 71 30 0 0 0 0
// cpuName is the name of the CPU (cpu, cpu0, cpu1, etc) and rawStats has
// the following format:
//   map[User:9366 Nice:0 System:5692 Iowait:114 Steal:0 GuestNice:0
//       Idle:1458880 Irq:806 Softirq:0 Guest:0]
func parseCpuRawStats(stats string) (cpuName string, rawStats CpuRawStats,
	err error) {
	rawStats = CpuRawStats{}

	fields := strings.Fields(stats)
	cpuName = fields[0]
	for i := 1; i < len(fields); i++ {
		stat, err := strconv.ParseUint(fields[i], 10, 64)
		if err != nil {
			return "", nil, err
		}
		rawStats[`Total`] += stat
		switch i {
		case 1:
			rawStats[`User`] = stat
		case 2:
			rawStats[`Nice`] = stat
		case 3:
			rawStats[`System`] = stat
		case 4:
			rawStats[`Idle`] = stat
		case 5:
			rawStats[`Iowait`] = stat
		case 6:
			rawStats[`Irq`] = stat
		case 7:
			rawStats[`Softirq`] = stat
		case 8:
			rawStats[`Steal`] = stat
		case 9:
			rawStats[`Guest`] = stat
		case 10:
			rawStats[`GuestNice`] = stat
		}
	}

	return cpuName, rawStats, nil
}

// getCpuStats calculates the average between 2 CpusRawStats samples and returns
// the % CPU utilization
func getCpuStats(firstSample CpusRawStats, secondSample CpusRawStats) (cpusStats CpusStats, err error) {
	fmt.Println(firstSample)
	fmt.Println(secondSample)

	cpusStats = CpusStats{}

	for cpuName, secondRawStats := range secondSample {
		matched, err := regexp.MatchString(`^cpu.*$`, cpuName)
		if err != nil {
			return nil, err
		}
		if !matched {
			return nil, errors.New("cpuName doesn't match the pattern")
		}

		firstRawStats, ok := firstSample[cpuName]
		if !ok {
			return nil, errors.New("The key " + cpuName + " doesn't exist in the first sample of CpusRawStats")
		}

		cpuStats := CpuStats{}
		totalTime := float64(secondRawStats[`Total`] - firstRawStats[`Total`])
		// Calculate the average between the two samples
		for key, secondValue := range secondRawStats {
			// Don't calculate the average if the key is 'Total'
			if key == `Total` {
				continue
			}
			avg := float64(secondValue-firstRawStats[key]) * 100.0 / totalTime
			avgStr := fmt.Sprintf("%3.2f", avg)
			cpuStats[key], err = strconv.ParseFloat(avgStr, 64)
			if err != nil {
				return nil, err
			}

		}
		cpusStats[cpuName] = cpuStats
	}

	return cpusStats, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"log"
	"net/http"
	"net/url"
	"strings"

	"github.com/PuerkitoBio/goquery"
	"github.com/streadway/amqp"
)

type CrawlerMessageItem struct {
	amqp.Delivery
	HTMLBody []byte

	host           string
	blacklistPaths []string
}

func NewCrawlerMessageItem(delivery amqp.Delivery, host string, blacklistPaths []string) *CrawlerMessageItem {
	return &CrawlerMessageItem{
		Delivery:       delivery,
		host:           host,
		blacklistPaths: blacklistPaths,
	}
}

func (c *CrawlerMessageItem) IsHTML() bool {
	return http.DetectContentType(c.HTMLBody) == "text/html; charset=utf-8"
}

func (c *CrawlerMessageItem) ExtractURLs() ([]string, error) {
	urls := []string{}

	document, err := goquery.NewDocumentFromReader(bytes.NewBuffer(c.HTMLBody))
	if err != nil {
		return urls, err
	}

	urlElementMatches := [][]string{
		[]string{"a", "href"},
		[]string{"img", "src"},
		[]string{"link", "href"},
		[]string{"script", "src"},
	}

	for _, attr := range urlElementMatches {
		element, attr := attr[0], attr[1]
		urls = append(urls, findByElementAttribute(document, c.host, c.blacklistPaths, element, attr)...)
	}

	return urls, err
}

func findByElementAttribute(
	document *goquery.Document,
	host string,
	blacklistPaths []string,
	element string,
	attr string) []string {

	urls := []string{}

	document.Find(element).Each(func(_ int, element *goquery.Selection) {
		href, exists := element.Attr(attr)
		unescapedHref, _ := url.QueryUnescape(href)
		trimmedHref := strings.TrimSpace(unescapedHref)

		u, err := url.Parse(trimmedHref)
		if err != nil {
			log.Fatal(err)
		}

		if exists && len(trimmedHref) > 0 && !isBlacklistedPath(trimmedHref, blacklistPaths) {
			if u.Host == host {
				urls = append(urls, trimmedHref)
			}

			if u.Host == "" && trimmedHref[0] == '/' {
				urls = append(urls, "http://"+host+trimmedHref)
			}
		}
	})

	return urls
}

func isBlacklistedPath(url string, blacklistedPaths []string) bool {
	for _, path := range blacklistedPaths {
		if strings.Contains(url, path) {
			return true
		}
	}

	return false
}
<commit_msg>Provide a way of getting at the URL() easily for a message item<commit_after>package main

import (
	"bytes"
	"log"
	"net/http"
	"net/url"
	"strings"

	"github.com/PuerkitoBio/goquery"
	"github.com/streadway/amqp"
)

type CrawlerMessageItem struct {
	amqp.Delivery
	HTMLBody []byte

	host           string
	blacklistPaths []string
}

func NewCrawlerMessageItem(delivery amqp.Delivery, host string, blacklistPaths []string) *CrawlerMessageItem {
	return &CrawlerMessageItem{
		Delivery:       delivery,
		host:           host,
		blacklistPaths: blacklistPaths,
	}
}

func (c *CrawlerMessageItem) IsHTML() bool {
	return http.DetectContentType(c.HTMLBody) == "text/html; charset=utf-8"
}

func (c *CrawlerMessageItem) URL() string {
	return string(c.Body)
}

func (c *CrawlerMessageItem) ExtractURLs() ([]string, error) {
	urls := []string{}

	document, err := goquery.NewDocumentFromReader(bytes.NewBuffer(c.HTMLBody))
	if err != nil {
		return urls, err
	}

	urlElementMatches := [][]string{
		[]string{"a", "href"},
		[]string{"img", "src"},
		[]string{"link", "href"},
		[]string{"script", "src"},
	}

	for _, attr := range urlElementMatches {
		element, attr := attr[0], attr[1]
		urls = append(urls, findByElementAttribute(document, c.host, c.blacklistPaths, element, attr)...)
	}

	return urls, err
}

func findByElementAttribute(
	document *goquery.Document,
	host string,
	blacklistPaths []string,
	element string,
	attr string) []string {

	urls := []string{}

	document.Find(element).Each(func(_ int, element *goquery.Selection) {
		href, exists := element.Attr(attr)
		unescapedHref, _ := url.QueryUnescape(href)
		trimmedHref := strings.TrimSpace(unescapedHref)

		u, err := url.Parse(trimmedHref)
		if err != nil {
			log.Fatal(err)
		}

		if exists && len(trimmedHref) > 0 && !isBlacklistedPath(trimmedHref, blacklistPaths) {
			if u.Host == host {
				urls = append(urls, trimmedHref)
			}

			if u.Host == "" && trimmedHref[0] == '/' {
				urls = append(urls, "http://"+host+trimmedHref)
			}
		}
	})

	return urls
}

func isBlacklistedPath(url string, blacklistedPaths []string) bool {
	for _, path := range blacklistedPaths {
		if strings.Contains(url, path) {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>package swarm

import (
	"common"
	"crypto/sha256"
	"log"
	"time"
)

/* SwarmInformed

   This State initializes the arbitrary swarm. It uses a simple consensus
   algorithm.

   The first stage is the learning phase. Each node sends a NodeAlive message.
   They then wait to see at least a majority of the swarms NodeAlive messages
   then they resend their NodeAlive message, to make sure that the majority of
   the swarm knows they are alive.

   Once the learning period has expired, rendezvous hashing is used to select
   the host which should compile the block. The selected host creates a block
   and announces it. The hosts who agree that the selected host is the block
   compiler release a heartbeat transaction.

   Once the selected host has compiled heartbeats from the majority of swarms,
   the block compiler makes another block which indicates that the swarm is
   transitioning out of this state.

   If this process fails, then the host should be marked as an invalid compiler
   and the next block compiler attempted, until this process fails or some other
   timeout expires.

   TODO: This implementation handles failure after the first block has appeared
   badly. Also the learning phase -> block compiler transition should probably
   have a pause in it and we should have a constant start time, not just 1 second
   after we started a given structure.

   TODO: We can actually do stage1 hashes in the nodealive message and stage1
   and 2 in the heartbeat message. This would initialize us faster
*/
type StateSwarmInformed struct {
	//Map of hosts seen to number of times they have failed to generate a block
	//Used for both host alive tracking & host block generation tracking
	hostsseen map[string]int

	// How many times we have broadcast that we are alive, we use a two stage
	// process where we broadcast, and then broadcast again when we have seen
	// enough nodes up to form a majority
	broadcastcount uint
	stage2         string

	// This state has two phases, the learning phase where it listens for new
	// hosts and then the commit stage where it listens for a block that
	// is correct according to its knowledge and then votes for it.
	learning bool

	heartbeats []*HeartBeatTransaction

	chain *BlockChain

	// Internal channels
	blockgen      <-chan time.Time
	sendBroadcast chan struct{}
	transaction   chan common.Transaction
	block         chan bwrap
	sync          chan struct{}
}

type bwrap struct {
	block *Block
	ret   chan State
}

func NewStateSwarmInformed(chain *BlockChain) (s *StateSwarmInformed) {
	s = new(StateSwarmInformed)
	s.chain = chain

	// When transitions / timeouts happen
	// Should be dynamically set
	s.blockgen = time.Tick(1 * time.Second)

	s.learning = true
	s.hostsseen = make(map[string]int)
	s.sendBroadcast = make(chan struct{})
	s.transaction = make(chan common.Transaction)
	s.block = make(chan bwrap)
	s.sync = make(chan struct{})

	go s.broadcastLife()
	go s.mainloop()
	return
}

func (s *StateSwarmInformed) HandleTransaction(t common.Transaction) {
	log.Print("STATE: Transaction queed to be handled")
	s.transaction <- t
}

func (s *StateSwarmInformed) Sync() {
	<-s.sync
}

func (s *StateSwarmInformed) HandleBlock(b *Block) State {
	log.Print("STATE: Block queed to be handled")
	c := make(chan State)
	s.block <- bwrap{b, c}
	r := <-c
	return r
}

func (s *StateSwarmInformed) blockCompiler() (compiler string) {

	hosts := make([]string, 0, len(s.hostsseen))

	//Pull all hosts who we haven't seen skipping a block
	for host, skipped := range s.hostsseen {
		if skipped != 0 {
			continue
		}
		hosts = append(hosts, host)
	}

	//Check if we should be the block generator
	compiler = common.RendezvousHash(sha256.New(), hosts, s.chain.Host)
	return
}

func (s *StateSwarmInformed) mainloop() {

	var compiler string

	for {
		select {
		case _ = <-s.sendBroadcast:
			log.Print("STATE: NodeAlive Transaction to be Queed")
			s.broadcastcount += 1
			go func() {
				s.chain.outgoingTransactions <- common.TransactionNetworkObject(
					NewNodeAlive(s.chain.Host, s.chain.Id))
				log.Print("STATE: NodeAlive Transaction Queed")
			}()

		case s.sync <- struct{}{}:

		case t := <-s.transaction:
			log.Print("STATE: Transaction Recieved")
			s.handleTransaction(t)

		case o := <-s.block:
			log.Print("STATE: Block Recieved")
			n := s.handleBlock(o.block)
			o.ret <- n

		case _ = <-s.blockgen:
			log.Print("STATE: Blockgen Recieved")

			if s.learning {
				s.learning = false
				log.Print("STATE: Disable Learning")
			} else if len(s.chain.BlockHistory) == 0 {
				if compiler != "" {
					s.hostsseen[compiler] += 1
					log.Print("STATE: Block Compiler not found")
				}
			}

			//Dont't try to generate a block if we don't have a majority of hosts
			if len(s.hostsseen) <= common.SWARMSIZE/2 {
				continue
				//Should actually switch to state swarmdied after a while
			}

			compiler = s.blockCompiler()

			log.Print("STATE: Block Compiler ", compiler, " Me ", s.chain.Host,
				" chain len ", len(s.chain.BlockHistory),
				" good ", compiler == s.chain.Host)

			if len(s.chain.BlockHistory) == 0 && compiler == s.chain.Host {

				log.Print("STATE: Generating Block type 1")

				id, err := common.RandomString(8)
				if err != nil {
					panic("checkBlockGenRandom" + err.Error())
				}
				b := &Block{id, s.chain.Id, s.chain.Host, nil, nil, nil}
				b.StorageMapping = make(map[string]interface{})
				for host, _ := range s.hostsseen {
					b.StorageMapping[host] = nil
				}

				s.heartbeats = s.heartbeats[0:0]
				time.Sleep(100 * time.Millisecond)
				s.chain.outgoingTransactions <- common.BlockNetworkObject(b)
			}

			if len(s.chain.BlockHistory) == 1 && len(s.heartbeats) > 2 {

				log.Print("STATE: Generating Block type 2")

				id, err := common.RandomString(8)
				if err != nil {
					panic(err)
				}

				b := &Block{id, s.chain.Id, s.chain.Host, nil, nil, nil}
				b.StorageMapping = s.chain.BlockHistory[0].StorageMapping
				b.EntropyStage1 = make(map[string]string)
				for _, h := range s.heartbeats {
					b.EntropyStage1[h.Host] = h.Stage1
				}

				// Arbitrary hard coded constant to make the testcases pass
				time.Sleep(500 * time.Millisecond)
				s.chain.outgoingTransactions <- common.BlockNetworkObject(b)
			}
		}
		log.Print("STATE: Signal Handling Finished")
	}
}

func (s *StateSwarmInformed) broadcastLife() {
	s.sendBroadcast <- struct{}{}
}

func (s *StateSwarmInformed) handleTransaction(t common.Transaction) {
	switch n := t.(type) {
	case *NodeAlive:
		if !s.learning {
			return
		}

		s.hostsseen[n.Node] = 0
		// Resend hostsseen count once you have seen a majority of hosts
		if len(s.hostsseen) > 2 && s.broadcastcount < 2 {
			go s.broadcastLife()
		}

		log.Print("STATE: Node added")

	case *HeartBeatTransaction:

		//If we're learning this is too early
		if s.learning {
			return
		}

		// If we're not trying to compile we don't care
		if s.blockCompiler() != s.chain.Host {
			return
		}

		if len(s.chain.BlockHistory) != 0 && n.Prevblock == s.chain.BlockHistory[0].Id {
			s.heartbeats = append(s.heartbeats, n)
		}

	default:
		return
	}
}

func (s *StateSwarmInformed) handleBlock(b *Block) State {

	// If the learning timeout hasn't expired, don't accept blocks
	if s.learning {
		log.Print("STATE: Block rejected b/c learning")
		return s
	}

	// All blocks in this state should be generated by the ideal host
	if b.Compiler != s.blockCompiler() {
		log.Print("STATE: Block rejected b/c wrong compiler")
		return s
	}

	// We are looking for a block to generate a heartbeat for
	if len(s.chain.BlockHistory) == 0 {
		log.Print("STATE: Block accepted type 1")
		s.chain.AddBlock(b)

		stage1, stage2 := common.HashedRandomData(sha256.New(), 8)
		s.stage2 = stage2

		if _, ok := b.StorageMapping[s.chain.Host]; ok {
			h := NewHeartBeat(s.chain.BlockHistory[0], s.chain.Host, stage1, "")
			go func() {
				s.chain.outgoingTransactions <- common.TransactionNetworkObject(h)
			}()
		}
		return s
	}

	// We're looking for the block with heartbeats to figure out if we're in
	// it
	if len(s.chain.BlockHistory) == 1 {
		log.Print("STATE: Block accepted type 2")

		s.chain.AddBlock(b)

		if _, ok := b.StorageMapping[s.chain.Host]; ok {
			//If we're in the block switch to signal mode
			log.Print("STATE: Switching to connected")
			return NewStateSwarmConnected()
		} else {
			//Join the swarm
			log.Print("STATE: Switching to Join")
			return NewStateSwarmJoin(s.chain)

		}
	}

	return s

}
<commit_msg>Clean up after swarminformed<commit_after>package swarm

import (
	"common"
	"crypto/sha256"
	"log"
	"time"
)

/* SwarmInformed

   This State initializes the arbitrary swarm. It uses a simple consensus
   algorithm.

   The first stage is the learning phase. Each node sends a NodeAlive message.
   They then wait to see at least a majority of the swarms NodeAlive messages
   then they resend their NodeAlive message, to make sure that the majority of
   the swarm knows they are alive.

   Once the learning period has expired, rendezvous hashing is used to select
   the host which should compile the block. The selected host creates a block
   and announces it. The hosts who agree that the selected host is the block
   compiler release a heartbeat transaction.

   Once the selected host has compiled heartbeats from the majority of swarms,
   the block compiler makes another block which indicates that the swarm is
   transitioning out of this state.

   If this process fails, then the host should be marked as an invalid compiler
   and the next block compiler attempted, until this process fails or some other
   timeout expires.

   TODO: This implementation handles failure after the first block has appeared
   badly. Also the learning phase -> block compiler transition should probably
   have a pause in it and we should have a constant start time, not just 1 second
   after we started a given structure.

   TODO: We can actually do stage1 hashes in the nodealive message and stage1
   and 2 in the heartbeat message. This would initialize us faster
*/
type StateSwarmInformed struct {
	//Map of hosts seen to number of times they have failed to generate a block
	//Used for both host alive tracking & host block generation tracking
	hostsseen map[string]int

	// How many times we have broadcast that we are alive, we use a two stage
	// process where we broadcast, and then broadcast again when we have seen
	// enough nodes up to form a majority
	broadcastcount uint
	stage2         string

	// This state has two phases, the learning phase where it listens for new
	// hosts and then the commit stage where it listens for a block that
	// is correct according to its knowledge and then votes for it.
	learning bool

	heartbeats []*HeartBeatTransaction

	chain *BlockChain

	// Internal channels
	blockgen      <-chan time.Time
	sendBroadcast chan struct{}
	transaction   chan common.Transaction
	block         chan bwrap
	sync          chan struct{}
	die           chan struct{}
}

type bwrap struct {
	block *Block
	ret   chan State
}

func NewStateSwarmInformed(chain *BlockChain) (s *StateSwarmInformed) {
	s = new(StateSwarmInformed)
	s.chain = chain

	// When transitions / timeouts happen
	// Should be dynamically set
	s.blockgen = time.Tick(1 * time.Second)

	s.learning = true
	s.hostsseen = make(map[string]int)
	s.sendBroadcast = make(chan struct{})
	s.transaction = make(chan common.Transaction)
	s.block = make(chan bwrap)
	s.sync = make(chan struct{})

	go s.broadcastLife()
	go s.mainloop()
	return
}

func (s *StateSwarmInformed) HandleTransaction(t common.Transaction) {
	log.Print("STATE: Transaction queed to be handled")
	s.transaction <- t
}

func (s *StateSwarmInformed) Sync() {
	<-s.sync
}

func (s *StateSwarmInformed) Die() {
	s.die <- struct{}{}
}

func (s *StateSwarmInformed) HandleBlock(b *Block) State {
	log.Print("STATE: Block queed to be handled")
	c := make(chan State)
	s.block <- bwrap{b, c}
	r := <-c
	return r
}

func (s *StateSwarmInformed) blockCompiler() (compiler string) {

	hosts := make([]string, 0, len(s.hostsseen))

	//Pull all hosts who we haven't seen skipping a block
	for host, skipped := range s.hostsseen {
		if skipped != 0 {
			continue
		}
		hosts = append(hosts, host)
	}

	//Check if we should be the block generator
	compiler = common.RendezvousHash(sha256.New(), hosts, s.chain.Host)
	return
}

func (s *StateSwarmInformed) mainloop() {

	var compiler string

	for {
		select {
		case _ = <-s.die:
			return
		case _ = <-s.sendBroadcast:
			log.Print("STATE: NodeAlive Transaction to be Queed")
			s.broadcastcount += 1
			go func() {
				s.chain.outgoingTransactions <- common.TransactionNetworkObject(
					NewNodeAlive(s.chain.Host, s.chain.Id))
				log.Print("STATE: NodeAlive Transaction Queed")
			}()

		case s.sync <- struct{}{}:

		case t := <-s.transaction:
			log.Print("STATE: Transaction Recieved")
			s.handleTransaction(t)

		case o := <-s.block:
			log.Print("STATE: Block Recieved")
			n := s.handleBlock(o.block)
			o.ret <- n

		case _ = <-s.blockgen:
			log.Print("STATE: Blockgen Recieved")

			if s.learning {
				s.learning = false
				log.Print("STATE: Disable Learning")
			} else if len(s.chain.BlockHistory) == 0 {
				if compiler != "" {
					s.hostsseen[compiler] += 1
					log.Print("STATE: Block Compiler not found")
				}
			}

			//Dont't try to generate a block if we don't have a majority of hosts
			if len(s.hostsseen) <= common.SWARMSIZE/2 {
				continue
				//Should actually switch to state swarmdied after a while
			}

			compiler = s.blockCompiler()

			log.Print("STATE: Block Compiler ", compiler, " Me ", s.chain.Host,
				" chain len ", len(s.chain.BlockHistory),
				" good ", compiler == s.chain.Host)

			if len(s.chain.BlockHistory) == 0 && compiler == s.chain.Host {

				log.Print("STATE: Generating Block type 1")

				id, err := common.RandomString(8)
				if err != nil {
					panic("checkBlockGenRandom" + err.Error())
				}
				b := &Block{id, s.chain.Id, s.chain.Host, nil, nil, nil}
				b.StorageMapping = make(map[string]interface{})
				for host, _ := range s.hostsseen {
					b.StorageMapping[host] = nil
				}

				s.heartbeats = s.heartbeats[0:0]
				time.Sleep(100 * time.Millisecond)
				s.chain.outgoingTransactions <- common.BlockNetworkObject(b)
			}

			if len(s.chain.BlockHistory) == 1 && len(s.heartbeats) > 2 {

				log.Print("STATE: Generating Block type 2")

				id, err := common.RandomString(8)
				if err != nil {
					panic(err)
				}

				b := &Block{id, s.chain.Id, s.chain.Host, nil, nil, nil}
				b.StorageMapping = s.chain.BlockHistory[0].StorageMapping
				b.EntropyStage1 = make(map[string]string)
				for _, h := range s.heartbeats {
					b.EntropyStage1[h.Host] = h.Stage1
				}

				// Arbitrary hard coded constant to make the testcases pass
				time.Sleep(500 * time.Millisecond)
				s.chain.outgoingTransactions <- common.BlockNetworkObject(b)
			}
		}
		log.Print("STATE: Signal Handling Finished")
	}
}

func (s *StateSwarmInformed) broadcastLife() {
	s.sendBroadcast <- struct{}{}
}

func (s *StateSwarmInformed) handleTransaction(t common.Transaction) {
	switch n := t.(type) {
	case *NodeAlive:
		if !s.learning {
			return
		}

		s.hostsseen[n.Node] = 0
		// Resend hostsseen count once you have seen a majority of hosts
		if len(s.hostsseen) > 2 && s.broadcastcount < 2 {
			go s.broadcastLife()
		}

		log.Print("STATE: Node added")

	case *HeartBeatTransaction:

		//If we're learning this is too early
		if s.learning {
			return
		}

		// If we're not trying to compile we don't care
		if s.blockCompiler() != s.chain.Host {
			return
		}

		if len(s.chain.BlockHistory) != 0 && n.Prevblock == s.chain.BlockHistory[0].Id {
			s.heartbeats = append(s.heartbeats, n)
		}

	default:
		return
	}
}

func (s *StateSwarmInformed) handleBlock(b *Block) State {

	// If the learning timeout hasn't expired, don't accept blocks
	if s.learning {
		log.Print("STATE: Block rejected b/c learning")
		return s
	}

	// All blocks in this state should be generated by the ideal host
	if b.Compiler != s.blockCompiler() {
		log.Print("STATE: Block rejected b/c wrong compiler")
		return s
	}

	// We are looking for a block to generate a heartbeat for
	if len(s.chain.BlockHistory) == 0 {
		log.Print("STATE: Block accepted type 1")
		s.chain.AddBlock(b)

		stage1, stage2 := common.HashedRandomData(sha256.New(), 8)
		s.stage2 = stage2

		if _, ok := b.StorageMapping[s.chain.Host]; ok {
			h := NewHeartBeat(s.chain.BlockHistory[0], s.chain.Host, stage1, "")
			go func() {
				s.chain.outgoingTransactions <- common.TransactionNetworkObject(h)
			}()
		}
		return s
	}

	// We're looking for the block with heartbeats to figure out if we're in
	// it
	if len(s.chain.BlockHistory) == 1 {
		log.Print("STATE: Block accepted type 2")

		s.chain.AddBlock(b)

		if _, ok := b.StorageMapping[s.chain.Host]; ok {
			//If we're in the block switch to signal mode
			log.Print("STATE: Switching to connected")
			go s.Die()
			return NewStateSwarmConnected()
		} else {
			//Join the swarm
			log.Print("STATE: Switching to Join")
			go s.Die()
			return NewStateSwarmJoin(s.chain)

		}
	}

	return s

}
<|endoftext|>