<commit_before>//extend FlagSet with support for flag aliasMap
package uggo

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"time"
)

type FlagSetWithAliases struct {
	*flag.FlagSet
	AutoGnuify   bool //sets whether it interprets options gnuishly (e.g. -lah being equivalent to -l -a -h)
	name           string
	argUsage       string
	out            io.Writer
	aliasMap       map[string][]string
	isPrintUsage   *bool //optional usage (you can use your own, or none, instead)
	isPrintVersion *bool //optional (you can use your own, or none, instead)
	version        string
}

func NewFlagSet(desc string, errorHandling flag.ErrorHandling) FlagSetWithAliases {
	fs := flag.NewFlagSet(desc, errorHandling)
	fs.SetOutput(ioutil.Discard)
	return FlagSetWithAliases{fs, false, desc, "", os.Stderr, map[string][]string{}, nil, nil, "unknown"}
}

func NewFlagSetDefault(name, argUsage, version string) FlagSetWithAliases {
	fs := flag.NewFlagSet(name+" "+argUsage, flag.ContinueOnError)
	fs.SetOutput(ioutil.Discard)
	tmp := false
	tmp2 := false
	flagSet := FlagSetWithAliases{fs, true, name, argUsage, os.Stderr, map[string][]string{}, &tmp, &tmp2, version}
	flagSet.BoolVar(flagSet.isPrintUsage, "help", false, "Show this help")
	flagSet.BoolVar(flagSet.isPrintVersion, "version", false, "Show version")
	flagSet.version = version
	return flagSet
}

func (flagSet FlagSetWithAliases) ProcessHelpOrVersion() bool {
	if flagSet.IsHelp() {
		flagSet.Usage()
		return true
	} else if flagSet.IsVersion() {
		flagSet.PrintVersion()
		return true
	}
	return false
}

//convenience method for storing 'usage' behaviour
func (flagSet FlagSetWithAliases) IsHelp() bool {
	return *flagSet.isPrintUsage
}

//convenience method for storing 'get version' behaviour
func (flagSet FlagSetWithAliases) IsVersion() bool {
	return *flagSet.isPrintVersion
}

func (flagSet FlagSetWithAliases) PrintVersion() {
	fmt.Fprintf(flagSet.out, "`%s` version: '%s'\n", flagSet.name, flagSet.version)
}

func (flagSet FlagSetWithAliases) Usage() {
	fmt.Fprintf(flagSet.out, "Usage: `%s %s`\n", flagSet.name, flagSet.argUsage)
	flagSet.PrintDefaults()
}

func (flagSet FlagSetWithAliases) SetOutput(out io.Writer) {
	flagSet.out = out
}

func (flagSet FlagSetWithAliases) AliasedBoolVar(p *bool, items []string, def bool, description string) {
	flagSet.RecordAliases(items, "bool")
	for _, item := range items {
		flagSet.BoolVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedDurationVar(p *time.Duration, items []string, def time.Duration, description string) {
	flagSet.RecordAliases(items, "duration")
	for _, item := range items {
		flagSet.DurationVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedFloat64Var(p *float64, items []string, def float64, description string) {
	flagSet.RecordAliases(items, "float64")
	for _, item := range items {
		flagSet.Float64Var(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedIntVar(p *int, items []string, def int, description string) {
	flagSet.RecordAliases(items, "int")
	for _, item := range items {
		flagSet.IntVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedInt64Var(p *int64, items []string, def int64, description string) {
	flagSet.RecordAliases(items, "int64")
	for _, item := range items {
		flagSet.Int64Var(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedStringVar(p *string, items []string, def string, description string) {
	flagSet.RecordAliases(items, "string")
	for _, item := range items {
		flagSet.StringVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) isAlternative(name string) bool {
	for _, altSlice := range flagSet.aliasMap {
		for _, alt := range altSlice {
			if alt == name {
				return true
			}
		}
	}
	return false
}

func (flagSet FlagSetWithAliases) RecordAliases(items []string, typ string) {
	var key string
	for i, item := range items {
		if i == 0 {
			key = item
			if _, ok := flagSet.aliasMap[key]; !ok {
				flagSet.aliasMap[key] = []string{}
			}
		} else {
			//key is same as before
			flagSet.aliasMap[key] = append(flagSet.aliasMap[key], item)
		}
	}
}

func (flagSet FlagSetWithAliases) Parse(call []string) error {
	if flagSet.AutoGnuify {
		call = Gnuify(call)
	}
	return flagSet.FlagSet.Parse(call)
}

func (flagSet FlagSetWithAliases) PrintDefaults() {
	flagSet.PrintDefaultsTo(flagSet.out)
}

func (flagSet FlagSetWithAliases) PrintDefaultsTo(out io.Writer) {
	flagSet.FlagSet.VisitAll(func(fl *flag.Flag) {
		format := "-%s=%s"
		l := 0
		alts, isAliased := flagSet.aliasMap[fl.Name]
		if isAliased {
			li, _ := fmt.Fprintf(out, "  ")
			l += li
			if len(fl.Name) > 1 {
				li, _ := fmt.Fprint(out, "-")
				l += li
			}
			//no known straightforward way to test for boolean types
			if fl.DefValue == "false" {
				li, _ = fmt.Fprintf(out, "-%s", fl.Name)
				l += li
			} else {
				li, _ = fmt.Fprintf(out, format, fl.Name, fl.DefValue)
				l += li
			}
			fmt.Fprint(out, " ")
			l += 1
			for _, alt := range alts {
				if len(alt) > 1 {
					li, _ := fmt.Fprint(out, "-")
					l += li
				}
				li, _ := fmt.Fprintf(out, "-%s ", alt)
				l += li
			}
		} else if !flagSet.isAlternative(fl.Name) {
			li, _ := fmt.Fprint(out, "  ")
			l += li
			if len(fl.Name) > 1 {
				li, _ := fmt.Fprint(out, "-")
				l += li
			}
			if fl.DefValue == "false" {
				li, _ = fmt.Fprintf(out, "-%s", fl.Name)
				l += li
			} else {
				li, _ = fmt.Fprintf(out, format, fl.Name, fl.DefValue)
				l += li
			}
		} else {
			//fmt.Fprintf(out, "alias %s\n", fl.Name)
		}
		if !flagSet.isAlternative(fl.Name) {
			for l < 25 {
				l += 1
				fmt.Fprintf(out, " ")
			}
			fmt.Fprintf(out, ": %s\n", fl.Usage)
		}

	})
	fmt.Fprintln(out, "")
}
<commit_msg>tweaking Usage formatting<commit_after>//extend FlagSet with support for flag aliasMap
package uggo

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"time"
)

type FlagSetWithAliases struct {
	*flag.FlagSet
	AutoGnuify   bool //sets whether it interprets options gnuishly (e.g. -lah being equivalent to -l -a -h)
	name           string
	argUsage       string
	out            io.Writer
	aliasMap       map[string][]string
	isPrintUsage   *bool //optional usage (you can use your own, or none, instead)
	isPrintVersion *bool //optional (you can use your own, or none, instead)
	version        string
}

func NewFlagSet(desc string, errorHandling flag.ErrorHandling) FlagSetWithAliases {
	fs := flag.NewFlagSet(desc, errorHandling)
	fs.SetOutput(ioutil.Discard)
	return FlagSetWithAliases{fs, false, desc, "", os.Stderr, map[string][]string{}, nil, nil, "unknown"}
}

func NewFlagSetDefault(name, argUsage, version string) FlagSetWithAliases {
	fs := flag.NewFlagSet(name+" "+argUsage, flag.ContinueOnError)
	fs.SetOutput(ioutil.Discard)
	tmp := false
	tmp2 := false
	flagSet := FlagSetWithAliases{fs, true, name, argUsage, os.Stderr, map[string][]string{}, &tmp, &tmp2, version}
	flagSet.BoolVar(flagSet.isPrintUsage, "help", false, "Show this help")
	flagSet.BoolVar(flagSet.isPrintVersion, "version", false, "Show version")
	flagSet.version = version
	return flagSet
}

func (flagSet FlagSetWithAliases) ProcessHelpOrVersion() bool {
	if flagSet.IsHelp() {
		flagSet.Usage()
		return true
	} else if flagSet.IsVersion() {
		flagSet.PrintVersion()
		return true
	}
	return false
}

//convenience method for storing 'usage' behaviour
func (flagSet FlagSetWithAliases) IsHelp() bool {
	return *flagSet.isPrintUsage
}

//convenience method for storing 'get version' behaviour
func (flagSet FlagSetWithAliases) IsVersion() bool {
	return *flagSet.isPrintVersion
}

func (flagSet FlagSetWithAliases) PrintVersion() {
	fmt.Fprintf(flagSet.out, "`%s` version: '%s'\n", flagSet.name, flagSet.version)
}

func (flagSet FlagSetWithAliases) Usage() {
	fmt.Fprintf(flagSet.out, "Usage: `%s %s`\n", flagSet.name, flagSet.argUsage)
	flagSet.PrintDefaults()
}

func (flagSet FlagSetWithAliases) SetOutput(out io.Writer) {
	flagSet.out = out
}

func (flagSet FlagSetWithAliases) AliasedBoolVar(p *bool, items []string, def bool, description string) {
	flagSet.RecordAliases(items, "bool")
	for _, item := range items {
		flagSet.BoolVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedDurationVar(p *time.Duration, items []string, def time.Duration, description string) {
	flagSet.RecordAliases(items, "duration")
	for _, item := range items {
		flagSet.DurationVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedFloat64Var(p *float64, items []string, def float64, description string) {
	flagSet.RecordAliases(items, "float64")
	for _, item := range items {
		flagSet.Float64Var(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedIntVar(p *int, items []string, def int, description string) {
	flagSet.RecordAliases(items, "int")
	for _, item := range items {
		flagSet.IntVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedInt64Var(p *int64, items []string, def int64, description string) {
	flagSet.RecordAliases(items, "int64")
	for _, item := range items {
		flagSet.Int64Var(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) AliasedStringVar(p *string, items []string, def string, description string) {
	flagSet.RecordAliases(items, "string")
	for _, item := range items {
		flagSet.StringVar(p, item, def, description)
	}
}

func (flagSet FlagSetWithAliases) isAlternative(name string) bool {
	for _, altSlice := range flagSet.aliasMap {
		for _, alt := range altSlice {
			if alt == name {
				return true
			}
		}
	}
	return false
}

func (flagSet FlagSetWithAliases) RecordAliases(items []string, typ string) {
	var key string
	for i, item := range items {
		if i == 0 {
			key = item
			if _, ok := flagSet.aliasMap[key]; !ok {
				flagSet.aliasMap[key] = []string{}
			}
		} else {
			//key is same as before
			flagSet.aliasMap[key] = append(flagSet.aliasMap[key], item)
		}
	}
}

func (flagSet FlagSetWithAliases) Parse(call []string) error {
	if flagSet.AutoGnuify {
		call = Gnuify(call)
	}
	return flagSet.FlagSet.Parse(call)
}

func (flagSet FlagSetWithAliases) PrintDefaults() {
	flagSet.PrintDefaultsTo(flagSet.out)
}

func (flagSet FlagSetWithAliases) PrintDefaultsTo(out io.Writer) {
	flagSet.FlagSet.VisitAll(func(fl *flag.Flag) {
		l := 0
		alts, isAliased := flagSet.aliasMap[fl.Name]
		if isAliased {
			li, _ := fmt.Fprintf(out, "  ")
			l += li
			if len(fl.Name) > 1 {
				li, _ := fmt.Fprint(out, "-")
				l += li
			}
			li, _ = fmt.Fprintf(out, "-%s", fl.Name)
			l += li
			fmt.Fprint(out, " ")
			l += 1
			for _, alt := range alts {
				if len(alt) > 1 {
					li, _ := fmt.Fprint(out, "-")
					l += li
				}
				li, _ := fmt.Fprintf(out, "-%s ", alt)
				l += li
			}
			//defaults:
			//no known straightforward way to test for boolean types
			if fl.DefValue == "false" {
			} else {
				li, _ = fmt.Fprintf(out, " %s", fl.DefValue)
				l += li
			}
		} else if !flagSet.isAlternative(fl.Name) {
			li, _ := fmt.Fprint(out, "  ")
			l += li
			if len(fl.Name) > 1 {
				li, _ := fmt.Fprint(out, "-")
				l += li
			}
			if fl.DefValue == "false" {
				li, _ = fmt.Fprintf(out, "-%s", fl.Name)
				l += li
			} else {
				format := "-%s %s"
				li, _ = fmt.Fprintf(out, format, fl.Name, fl.DefValue)
				l += li
			}
		} else {
			//fmt.Fprintf(out, "alias %s\n", fl.Name)
		}
		if !flagSet.isAlternative(fl.Name) {
			for l < 25 {
				l += 1
				fmt.Fprintf(out, " ")
			}
			fmt.Fprintf(out, ": %s\n", fl.Usage)
		}

	})
	fmt.Fprintln(out, "")
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase

import (
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"sync"

	"github.com/keybase/client/go/externals"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/pvlsource"
	"github.com/keybase/client/go/service"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	"github.com/keybase/kbfs/fsrpc"
	"github.com/keybase/kbfs/libkbfs"
)

var kbCtx *libkb.GlobalContext
var conn net.Conn
var startOnce sync.Once
var logSendContext libkb.LogSendContext
var kbfsConfig libkbfs.Config

// InitOnce runs the Keybase services (only runs one time)
func InitOnce(homeDir string, logFile string, runModeStr string, accessGroupOverride bool) {
	startOnce.Do(func() {
		if err := Init(homeDir, logFile, runModeStr, accessGroupOverride); err != nil {
			kbCtx.Log.Errorf("Init error: %s", err)
		}
	})
}

// Init runs the Keybase services
func Init(homeDir string, logFile string, runModeStr string, accessGroupOverride bool) error {
	fmt.Println("Go: Initializing")
	if logFile != "" {
		fmt.Printf("Go: Using log: %s\n", logFile)
	}

	kbCtx = libkb.G
	kbCtx.Init()
	kbCtx.SetServices(externals.GetServices())
	pvlsource.NewPvlSourceAndInstall(kbCtx)
	usage := libkb.Usage{
		Config:    true,
		API:       true,
		KbKeyring: true,
	}
	runMode, err := libkb.StringToRunMode(runModeStr)
	if err != nil {
		return err
	}
	config := libkb.AppConfig{
		HomeDir:                     homeDir,
		LogFile:                     logFile,
		RunMode:                     runMode,
		Debug:                       true,
		LocalRPCDebug:               "",
		SecurityAccessGroupOverride: accessGroupOverride,
	}
	err = kbCtx.Configure(config, usage)
	if err != nil {
		return err
	}

	svc := service.NewService(kbCtx, false)
	svc.StartLoopbackServer()
	kbCtx.SetService()
	uir := service.NewUIRouter(kbCtx)
	kbCtx.SetUIRouter(uir)
	svc.RunBackgroundOperations(uir)

	serviceLog := config.GetLogFile()
	logs := libkb.Logs{
		Service: serviceLog,
	}

	logSendContext = libkb.LogSendContext{
		Contextified: libkb.NewContextified(kbCtx),
		Logs:         logs,
	}

	// FIXME (MBG): This is causing RPC responses to sometimes not be recieved
	// on iOS. Repro by hooking up getExtendedStatus to a button in the iOS
	// client and watching JS logs. Disabling until we have a root cause / fix.
	kbfsParams := libkbfs.DefaultInitParams(kbCtx)
	// Avoid lots of background routines.
	go func() {
		kbfsParams.Mode = libkbfs.InitMinimalString
		kbfsConfig, _ = libkbfs.Init(kbCtx, kbfsParams, serviceCn{}, func() {}, kbCtx.Log)
	}()

	return Reset()
}

type serviceCn struct {
	ctx *libkb.GlobalContext
}

func (s serviceCn) NewKeybaseService(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.KeybaseService, error) {
	keybaseService := libkbfs.NewKeybaseDaemonRPC(config, ctx, log, true, nil)
	keybaseService.AddProtocols([]rpc.Protocol{
		keybase1.FsProtocol(fsrpc.NewFS(config, log)),
	})
	return keybaseService, nil
}

func (s serviceCn) NewCrypto(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.Crypto, error) {
	return libkbfs.NewCryptoClientRPC(config, ctx), nil
}

// LogSend sends a log to Keybase
func LogSend(feedback string, sendLogs bool, uiLogPath string) (string, error) {
	logSendContext.Logs.Desktop = uiLogPath
	return logSendContext.LogSend(fmt.Sprintf("{\"feedback\": \"%s\"}", feedback), sendLogs, 5*1024*1024)
}

// WriteB64 sends a base64 encoded msgpack rpc payload
func WriteB64(str string) error {
	data, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return fmt.Errorf("Base64 decode error: %s; %s", err, str)
	}
	n, err := conn.Write(data)
	if err != nil {
		return fmt.Errorf("Write error: %s", err)
	}
	if n != len(data) {
		return errors.New("Did not write all the data")
	}
	return nil
}

const targetBufferSize = 50 * 1024

// bufferSize must be divisible by 3 to ensure that we don't split
// our b64 encode across a payload boundary if we go over our buffer
// size.
const bufferSize = targetBufferSize - (targetBufferSize % 3)

// buffer for the conn.Read
var buffer = make([]byte, bufferSize)

// ReadB64 is a blocking read for base64 encoded msgpack rpc data.
// It is called serially by the mobile run loops.
func ReadB64() (string, error) {
	n, err := conn.Read(buffer)
	if n > 0 && err == nil {
		str := base64.StdEncoding.EncodeToString(buffer[0:n])
		return str, nil
	}

	if err != nil {
		// Attempt to fix the connection
		Reset()
		return "", fmt.Errorf("Read error: %s", err)
	}

	return "", nil
}

// Reset resets the socket connection
func Reset() error {
	if conn != nil {
		conn.Close()
	}

	var err error
	conn, err = kbCtx.LoopbackListener.Dial()
	if err != nil {
		return fmt.Errorf("Socket error: %s", err)
	}
	return nil
}

// Version returns semantic version string
func Version() string {
	return libkb.VersionString()
}
<commit_msg>turn on debug logging in KBFS on mobile (#6385)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase

import (
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"sync"

	"github.com/keybase/client/go/externals"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/pvlsource"
	"github.com/keybase/client/go/service"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	"github.com/keybase/kbfs/fsrpc"
	"github.com/keybase/kbfs/libkbfs"
)

var kbCtx *libkb.GlobalContext
var conn net.Conn
var startOnce sync.Once
var logSendContext libkb.LogSendContext
var kbfsConfig libkbfs.Config

// InitOnce runs the Keybase services (only runs one time)
func InitOnce(homeDir string, logFile string, runModeStr string, accessGroupOverride bool) {
	startOnce.Do(func() {
		if err := Init(homeDir, logFile, runModeStr, accessGroupOverride); err != nil {
			kbCtx.Log.Errorf("Init error: %s", err)
		}
	})
}

// Init runs the Keybase services
func Init(homeDir string, logFile string, runModeStr string, accessGroupOverride bool) error {
	fmt.Println("Go: Initializing")
	if logFile != "" {
		fmt.Printf("Go: Using log: %s\n", logFile)
	}

	kbCtx = libkb.G
	kbCtx.Init()
	kbCtx.SetServices(externals.GetServices())
	pvlsource.NewPvlSourceAndInstall(kbCtx)
	usage := libkb.Usage{
		Config:    true,
		API:       true,
		KbKeyring: true,
	}
	runMode, err := libkb.StringToRunMode(runModeStr)
	if err != nil {
		return err
	}
	config := libkb.AppConfig{
		HomeDir:                     homeDir,
		LogFile:                     logFile,
		RunMode:                     runMode,
		Debug:                       true,
		LocalRPCDebug:               "",
		SecurityAccessGroupOverride: accessGroupOverride,
	}
	err = kbCtx.Configure(config, usage)
	if err != nil {
		return err
	}

	svc := service.NewService(kbCtx, false)
	svc.StartLoopbackServer()
	kbCtx.SetService()
	uir := service.NewUIRouter(kbCtx)
	kbCtx.SetUIRouter(uir)
	svc.RunBackgroundOperations(uir)

	serviceLog := config.GetLogFile()
	logs := libkb.Logs{
		Service: serviceLog,
	}

	logSendContext = libkb.LogSendContext{
		Contextified: libkb.NewContextified(kbCtx),
		Logs:         logs,
	}

	go func() {
		kbfsParams := libkbfs.DefaultInitParams(kbCtx)
		// Setting this flag will enable KBFS debug logging to alway be
		// true in a mobile setting. Kill this setting if too spammy.
		kbfsParams.Debug = true
		kbfsParams.Mode = libkbfs.InitMinimalString
		kbfsConfig, _ = libkbfs.Init(kbCtx, kbfsParams, serviceCn{}, func() {}, kbCtx.Log)
	}()

	return Reset()
}

type serviceCn struct {
	ctx *libkb.GlobalContext
}

func (s serviceCn) NewKeybaseService(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.KeybaseService, error) {
	keybaseService := libkbfs.NewKeybaseDaemonRPC(config, ctx, log, true, nil)
	keybaseService.AddProtocols([]rpc.Protocol{
		keybase1.FsProtocol(fsrpc.NewFS(config, log)),
	})
	return keybaseService, nil
}

func (s serviceCn) NewCrypto(config libkbfs.Config, params libkbfs.InitParams, ctx libkbfs.Context, log logger.Logger) (libkbfs.Crypto, error) {
	return libkbfs.NewCryptoClientRPC(config, ctx), nil
}

// LogSend sends a log to Keybase
func LogSend(feedback string, sendLogs bool, uiLogPath string) (string, error) {
	logSendContext.Logs.Desktop = uiLogPath
	return logSendContext.LogSend(fmt.Sprintf("{\"feedback\": \"%s\"}", feedback), sendLogs, 5*1024*1024)
}

// WriteB64 sends a base64 encoded msgpack rpc payload
func WriteB64(str string) error {
	data, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return fmt.Errorf("Base64 decode error: %s; %s", err, str)
	}
	n, err := conn.Write(data)
	if err != nil {
		return fmt.Errorf("Write error: %s", err)
	}
	if n != len(data) {
		return errors.New("Did not write all the data")
	}
	return nil
}

const targetBufferSize = 50 * 1024

// bufferSize must be divisible by 3 to ensure that we don't split
// our b64 encode across a payload boundary if we go over our buffer
// size.
const bufferSize = targetBufferSize - (targetBufferSize % 3)

// buffer for the conn.Read
var buffer = make([]byte, bufferSize)

// ReadB64 is a blocking read for base64 encoded msgpack rpc data.
// It is called serially by the mobile run loops.
func ReadB64() (string, error) {
	n, err := conn.Read(buffer)
	if n > 0 && err == nil {
		str := base64.StdEncoding.EncodeToString(buffer[0:n])
		return str, nil
	}

	if err != nil {
		// Attempt to fix the connection
		Reset()
		return "", fmt.Errorf("Read error: %s", err)
	}

	return "", nil
}

// Reset resets the socket connection
func Reset() error {
	if conn != nil {
		conn.Close()
	}

	var err error
	conn, err = kbCtx.LoopbackListener.Dial()
	if err != nil {
		return fmt.Errorf("Socket error: %s", err)
	}
	return nil
}

// Version returns semantic version string
func Version() string {
	return libkb.VersionString()
}
<|endoftext|>


<commit_before>package goFlags

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"flag"

	"github.com/crgimenes/goConfig/structTag"
)

var parametersStringMap map[*reflect.Value]*string
var parametersIntMap map[*reflect.Value]*int

func init() {
	parametersStringMap = make(map[*reflect.Value]*string)
	parametersIntMap = make(map[*reflect.Value]*int)

	SetTag("flag")
	SetTagDefault("flagDefault")

	structTag.ParseMap[reflect.Int] = reflectInt
	structTag.ParseMap[reflect.String] = reflectString
}

// SetTag set a new tag
func SetTag(tag string) {
	structTag.Tag = tag
}

// SetTagDefault set a new TagDefault to retorn default values
func SetTagDefault(tag string) {
	structTag.TagDefault = tag
}

// Parse configuration
func Parse(config interface{}) (err error) {

	err = structTag.Parse(config, "")
	if err != nil {
		return
	}

	flag.Parse()

	for k, v := range parametersStringMap {
		fmt.Printf("- \"%v\"\n", *v)
		k.SetString(*v)
	}

	for k, v := range parametersIntMap {
		fmt.Printf("- \"%v\"\n", int64(*v))
		k.SetInt(int64(*v))

	}
	return
}

func reflectInt(field *reflect.StructField, value *reflect.Value, tag string) (err error) {
	var aux int
	var defaltValue string
	var defaltValueInt int

	defaltValue = field.Tag.Get(structTag.TagDefault)

	if defaltValue == "" || defaltValue == "0" {
		defaltValueInt = 0
	} else {
		defaltValueInt, err = strconv.Atoi(defaltValue)
		if err != nil {
			return
		}
	}

	flag.IntVar(&aux, strings.ToLower(tag), defaltValueInt, "")
	parametersIntMap[value] = &aux

	fmt.Println(tag, defaltValue)

	return
}

func reflectString(field *reflect.StructField, value *reflect.Value, tag string) (err error) {

	var aux string
	var defaltValue string
	defaltValue = field.Tag.Get(structTag.TagDefault)

	flag.StringVar(&aux, strings.ToLower(tag), defaltValue, "")
	parametersStringMap[value] = &aux

	fmt.Println(tag, defaltValue)

	return
}
<commit_msg>test visit function<commit_after>package goFlags

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"flag"

	"github.com/crgimenes/goConfig/structTag"
)

var parametersStringMap map[*reflect.Value]*string
var parametersIntMap map[*reflect.Value]*int

func init() {
	parametersStringMap = make(map[*reflect.Value]*string)
	parametersIntMap = make(map[*reflect.Value]*int)

	SetTag("flag")
	SetTagDefault("flagDefault")

	structTag.ParseMap[reflect.Int] = reflectInt
	structTag.ParseMap[reflect.String] = reflectString
}

// SetTag set a new tag
func SetTag(tag string) {
	structTag.Tag = tag
}

// SetTagDefault set a new TagDefault to retorn default values
func SetTagDefault(tag string) {
	structTag.TagDefault = tag
}

// Parse configuration
func Parse(config interface{}) (err error) {

	err = structTag.Parse(config, "")
	if err != nil {
		return
	}

	flag.Parse()

	//fmt.Printf("%v#f", flag.CommandLine)
	flag.Visit(visitTest)

	for k, v := range parametersStringMap {
		fmt.Printf("- \"%v\"\n", *v)
		k.SetString(*v)
	}

	for k, v := range parametersIntMap {
		fmt.Printf("- \"%v\"\n", int64(*v))
		k.SetInt(int64(*v))

	}
	return
}

func visitTest(f *flag.Flag) {
	fmt.Printf("name \"%v\"\n", f.Name)
}

func reflectInt(field *reflect.StructField, value *reflect.Value, tag string) (err error) {
	var aux int
	var defaltValue string
	var defaltValueInt int

	defaltValue = field.Tag.Get(structTag.TagDefault)

	if defaltValue == "" || defaltValue == "0" {
		defaltValueInt = 0
	} else {
		defaltValueInt, err = strconv.Atoi(defaltValue)
		if err != nil {
			return
		}
	}

	flag.IntVar(&aux, strings.ToLower(tag), defaltValueInt, "")
	parametersIntMap[value] = &aux

	fmt.Println(tag, defaltValue)

	return
}

func reflectString(field *reflect.StructField, value *reflect.Value, tag string) (err error) {

	var aux string
	var defaltValue string
	defaltValue = field.Tag.Get(structTag.TagDefault)

	flag.StringVar(&aux, strings.ToLower(tag), defaltValue, "")
	parametersStringMap[value] = &aux

	fmt.Println(tag, defaltValue)

	return
}
<|endoftext|>


<commit_before>package formatters

import (
	"io"
	"os"
	"strconv"

	"github.com/mrtazz/checkmake/rules"
	"github.com/olekukonko/tablewriter"
)

// DefaultFormatter is the formatter used by default for CLI output
type DefaultFormatter struct {
	out io.Writer
}

// NewDefaultFormatter returns a DefaultFormatter struct
func NewDefaultFormatter() *DefaultFormatter {
	return &DefaultFormatter{out: os.Stdout}
}

// Format is the function to call to get the formatted output
func (f *DefaultFormatter) Format(violations rules.RuleViolationList) {

	data := make([][]string, len(violations))

	for idx, val := range violations {
		data[idx] = []string{val.Rule,
			val.Violation,
			strconv.Itoa(val.LineNumber)}
	}

	table := tablewriter.NewWriter(f.out)

	table.SetHeader([]string{"Rule", "Description", "Line Number"})

	table.SetCenterSeparator(" ")
	table.SetColumnSeparator(" ")
	table.SetRowSeparator(" ")
	table.SetAutoWrapText(true)

	table.AppendBulk(data)
	table.Render()
}
<commit_msg>Disable borders to remove start/end newlines<commit_after>package formatters

import (
	"io"
	"os"
	"strconv"

	"github.com/mrtazz/checkmake/rules"
	"github.com/olekukonko/tablewriter"
)

// DefaultFormatter is the formatter used by default for CLI output
type DefaultFormatter struct {
	out io.Writer
}

// NewDefaultFormatter returns a DefaultFormatter struct
func NewDefaultFormatter() *DefaultFormatter {
	return &DefaultFormatter{out: os.Stdout}
}

// Format is the function to call to get the formatted output
func (f *DefaultFormatter) Format(violations rules.RuleViolationList) {
	data := make([][]string, len(violations))

	for idx, val := range violations {
		data[idx] = []string{val.Rule,
			val.Violation,
			strconv.Itoa(val.LineNumber)}
	}

	table := tablewriter.NewWriter(f.out)

	table.SetHeader([]string{"Rule", "Description", "Line Number"})

	table.SetCenterSeparator(" ")
	table.SetColumnSeparator(" ")
	table.SetRowSeparator(" ")
	table.SetBorder(false)
	table.SetAutoWrapText(true)

	table.AppendBulk(data)
	table.Render()
}
<|endoftext|>


<commit_before>package semaphore

import (
	"context"
	"fmt"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func ExampleMutex() {
	const routineCount = 5

	var (
		s     = Mutex()
		wg    = new(sync.WaitGroup)
		value int
	)

	wg.Add(routineCount)
	for i := 0; i < routineCount; i++ {
		go func() {
			defer wg.Done()
			defer s.Release()
			s.Acquire()
			value++
			fmt.Println(value)
		}()
	}

	wg.Wait()

	// Unordered output:
	// 1
	// 2
	// 3
	// 4
	// 5
}

func testNewInvalidCount(t *testing.T) {
	for _, c := range []int{0, -1} {
		t.Run(strconv.Itoa(c), func(t *testing.T) {
			assert.Panics(t, func() {
				New(c)
			})
		})
	}
}

func testNewValidCount(t *testing.T) {
	for _, c := range []int{1, 2, 5} {
		t.Run(strconv.Itoa(c), func(t *testing.T) {
			s := New(c)
			assert.NotNil(t, s)
		})
	}
}

func TestNew(t *testing.T) {
	t.Run("InvalidCount", testNewInvalidCount)
	t.Run("ValidCount", testNewValidCount)
}

func testTryAcquire(t *testing.T, s Interface, totalCount int) {
	assert := assert.New(t)
	for i := 0; i < totalCount; i++ {
		assert.True(s.TryAcquire())
	}

	assert.False(s.TryAcquire())
	s.Release()
	assert.True(s.TryAcquire())
	assert.False(s.TryAcquire())
}

func testAcquire(t *testing.T, s Interface, totalCount int) {
	var (
		assert  = assert.New(t)
		require = require.New(t)
	)

	// acquire all the things!
	for i := 0; i < totalCount; i++ {
		done := make(chan struct{})
		go func() {
			defer close(done)
			s.Acquire()
		}()

		select {
		case <-done:
			// passing
		case <-time.After(time.Second):
			assert.FailNow("Acquire blocked unexpectedly")
		}
	}

	// post condition: no point continuing if this fails
	require.False(s.TryAcquire())

	var (
		ready    = make(chan struct{})
		acquired = make(chan struct{})
	)

	go func() {
		defer close(acquired)
		close(ready)
		s.Acquire() // this should now block
	}()

	select {
	case <-ready:
		// passing
		require.False(s.TryAcquire())
		s.Release()
	case <-time.After(time.Second):
		require.FailNow("Unable to spawn acquire goroutine")
	}

	select {
	case <-acquired:
		require.False(s.TryAcquire())
	case <-time.After(time.Second):
		require.FailNow("Acquire blocked unexpectedly")
	}

	s.Release()
	assert.True(s.TryAcquire())
}

func testAcquireWait(t *testing.T, s Interface, totalCount int) {
	var (
		assert  = assert.New(t)
		require = require.New(t)
		timer   = make(chan time.Time)
	)

	// acquire all the things!
	for i := 0; i < totalCount; i++ {
		result := make(chan error)
		go func() {
			result <- s.AcquireWait(timer)
		}()

		select {
		case err := <-result:
			assert.NoError(err)
		case <-time.After(time.Second):
			assert.FailNow("Acquire blocked unexpectedly")
		}
	}

	// post condition: no point continuing if this fails
	require.False(s.TryAcquire())

	var (
		ready  = make(chan struct{})
		result = make(chan error)
	)

	go func() {
		close(ready)
		result <- s.AcquireWait(timer)
	}()

	select {
	case <-ready:
		timer <- time.Time{}
	case <-time.After(time.Second):
		require.FailNow("Unable to spawn acquire goroutine")
	}

	select {
	case err := <-result:
		assert.Equal(ErrTimeout, err)
	case <-time.After(time.Second):
		require.FailNow("AcquireWait blocked unexpectedly")
	}
}

func testAcquireCtx(t *testing.T, s Interface, totalCount int) {
	var (
		assert      = assert.New(t)
		require     = require.New(t)
		ctx, cancel = context.WithCancel(context.Background())
	)

	defer cancel()

	// acquire all the things!
	for i := 0; i < totalCount; i++ {
		result := make(chan error)
		go func() {
			result <- s.AcquireCtx(ctx)
		}()

		select {
		case err := <-result:
			assert.NoError(err)
		case <-time.After(time.Second):
			assert.FailNow("Acquire blocked unexpectedly")
		}
	}

	// post condition: no point continuing if this fails
	require.False(s.TryAcquire())

	var (
		ready  = make(chan struct{})
		result = make(chan error)
	)

	go func() {
		close(ready)
		result <- s.AcquireCtx(ctx)
	}()

	select {
	case <-ready:
		cancel()
	case <-time.After(time.Second):
		require.FailNow("Unable to spawn acquire goroutine")
	}

	select {
	case err := <-result:
		assert.Equal(ctx.Err(), err)
	case <-time.After(time.Second):
		require.FailNow("AcquireWait blocked unexpectedly")
	}
}

func TestSemaphore(t *testing.T) {
	for _, c := range []int{1, 2, 5} {
		t.Run(fmt.Sprintf("count=%d", c), func(t *testing.T) {
			t.Run("TryAcquire", func(t *testing.T) {
				testTryAcquire(t, New(c), c)
			})

			t.Run("Acquire", func(t *testing.T) {
				testAcquire(t, New(c), c)
			})

			t.Run("AcquireWait", func(t *testing.T) {
				testAcquireWait(t, New(c), c)
			})

			t.Run("AcquireCtx", func(t *testing.T) {
				testAcquireCtx(t, New(c), c)
			})
		})
	}
}

func TestMutex(t *testing.T) {
	t.Run("TryAcquire", func(t *testing.T) {
		testTryAcquire(t, Mutex(), 1)
	})

	t.Run("Acquire", func(t *testing.T) {
		testAcquire(t, Mutex(), 1)
	})

	t.Run("AcquireWait", func(t *testing.T) {
		testAcquireWait(t, Mutex(), 1)
	})

	t.Run("AcquireCtx", func(t *testing.T) {
		testAcquireCtx(t, Mutex(), 1)
	})
}
<commit_msg>Added the acquire wait example back in a way go vet accepts<commit_after>package semaphore

import (
	"context"
	"fmt"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func ExampleMutex() {
	const routineCount = 5

	var (
		s     = Mutex()
		wg    = new(sync.WaitGroup)
		value int
	)

	wg.Add(routineCount)
	for i := 0; i < routineCount; i++ {
		go func() {
			defer wg.Done()
			defer s.Release()
			s.Acquire()
			value++
			fmt.Println(value)
		}()
	}

	wg.Wait()

	// Unordered output:
	// 1
	// 2
	// 3
	// 4
	// 5
}

func ExampleInterface_AcquireWait() {
	var (
		s     = Mutex()
		timer = time.NewTimer(100 * time.Millisecond)
	)

	defer timer.Stop()
	s.Acquire() // force AcquireWait to block
	err := s.AcquireWait(timer.C)
	fmt.Println(err != nil)

	// Output:
	// true
}

func testNewInvalidCount(t *testing.T) {
	for _, c := range []int{0, -1} {
		t.Run(strconv.Itoa(c), func(t *testing.T) {
			assert.Panics(t, func() {
				New(c)
			})
		})
	}
}

func testNewValidCount(t *testing.T) {
	for _, c := range []int{1, 2, 5} {
		t.Run(strconv.Itoa(c), func(t *testing.T) {
			s := New(c)
			assert.NotNil(t, s)
		})
	}
}

func TestNew(t *testing.T) {
	t.Run("InvalidCount", testNewInvalidCount)
	t.Run("ValidCount", testNewValidCount)
}

func testTryAcquire(t *testing.T, s Interface, totalCount int) {
	assert := assert.New(t)
	for i := 0; i < totalCount; i++ {
		assert.True(s.TryAcquire())
	}

	assert.False(s.TryAcquire())
	s.Release()
	assert.True(s.TryAcquire())
	assert.False(s.TryAcquire())
}

func testAcquire(t *testing.T, s Interface, totalCount int) {
	var (
		assert  = assert.New(t)
		require = require.New(t)
	)

	// acquire all the things!
	for i := 0; i < totalCount; i++ {
		done := make(chan struct{})
		go func() {
			defer close(done)
			s.Acquire()
		}()

		select {
		case <-done:
			// passing
		case <-time.After(time.Second):
			assert.FailNow("Acquire blocked unexpectedly")
		}
	}

	// post condition: no point continuing if this fails
	require.False(s.TryAcquire())

	var (
		ready    = make(chan struct{})
		acquired = make(chan struct{})
	)

	go func() {
		defer close(acquired)
		close(ready)
		s.Acquire() // this should now block
	}()

	select {
	case <-ready:
		// passing
		require.False(s.TryAcquire())
		s.Release()
	case <-time.After(time.Second):
		require.FailNow("Unable to spawn acquire goroutine")
	}

	select {
	case <-acquired:
		require.False(s.TryAcquire())
	case <-time.After(time.Second):
		require.FailNow("Acquire blocked unexpectedly")
	}

	s.Release()
	assert.True(s.TryAcquire())
}

func testAcquireWait(t *testing.T, s Interface, totalCount int) {
	var (
		assert  = assert.New(t)
		require = require.New(t)
		timer   = make(chan time.Time)
	)

	// acquire all the things!
	for i := 0; i < totalCount; i++ {
		result := make(chan error)
		go func() {
			result <- s.AcquireWait(timer)
		}()

		select {
		case err := <-result:
			assert.NoError(err)
		case <-time.After(time.Second):
			assert.FailNow("Acquire blocked unexpectedly")
		}
	}

	// post condition: no point continuing if this fails
	require.False(s.TryAcquire())

	var (
		ready  = make(chan struct{})
		result = make(chan error)
	)

	go func() {
		close(ready)
		result <- s.AcquireWait(timer)
	}()

	select {
	case <-ready:
		timer <- time.Time{}
	case <-time.After(time.Second):
		require.FailNow("Unable to spawn acquire goroutine")
	}

	select {
	case err := <-result:
		assert.Equal(ErrTimeout, err)
	case <-time.After(time.Second):
		require.FailNow("AcquireWait blocked unexpectedly")
	}
}

func testAcquireCtx(t *testing.T, s Interface, totalCount int) {
	var (
		assert      = assert.New(t)
		require     = require.New(t)
		ctx, cancel = context.WithCancel(context.Background())
	)

	defer cancel()

	// acquire all the things!
	for i := 0; i < totalCount; i++ {
		result := make(chan error)
		go func() {
			result <- s.AcquireCtx(ctx)
		}()

		select {
		case err := <-result:
			assert.NoError(err)
		case <-time.After(time.Second):
			assert.FailNow("Acquire blocked unexpectedly")
		}
	}

	// post condition: no point continuing if this fails
	require.False(s.TryAcquire())

	var (
		ready  = make(chan struct{})
		result = make(chan error)
	)

	go func() {
		close(ready)
		result <- s.AcquireCtx(ctx)
	}()

	select {
	case <-ready:
		cancel()
	case <-time.After(time.Second):
		require.FailNow("Unable to spawn acquire goroutine")
	}

	select {
	case err := <-result:
		assert.Equal(ctx.Err(), err)
	case <-time.After(time.Second):
		require.FailNow("AcquireWait blocked unexpectedly")
	}
}

func TestSemaphore(t *testing.T) {
	for _, c := range []int{1, 2, 5} {
		t.Run(fmt.Sprintf("count=%d", c), func(t *testing.T) {
			t.Run("TryAcquire", func(t *testing.T) {
				testTryAcquire(t, New(c), c)
			})

			t.Run("Acquire", func(t *testing.T) {
				testAcquire(t, New(c), c)
			})

			t.Run("AcquireWait", func(t *testing.T) {
				testAcquireWait(t, New(c), c)
			})

			t.Run("AcquireCtx", func(t *testing.T) {
				testAcquireCtx(t, New(c), c)
			})
		})
	}
}

func TestMutex(t *testing.T) {
	t.Run("TryAcquire", func(t *testing.T) {
		testTryAcquire(t, Mutex(), 1)
	})

	t.Run("Acquire", func(t *testing.T) {
		testAcquire(t, Mutex(), 1)
	})

	t.Run("AcquireWait", func(t *testing.T) {
		testAcquireWait(t, Mutex(), 1)
	})

	t.Run("AcquireCtx", func(t *testing.T) {
		testAcquireCtx(t, Mutex(), 1)
	})
}
<|endoftext|>


<commit_before>package semaphore

import (
	"fmt"
	"sync"
	"testing"
)

// TODO: Mock out consul-api so these aren't integration tests
// Will be harder since api.Client is a struct rather than an interface.

func TestSemaphore(t *testing.T) {
	const (
		path = "tests/integration/semaphore/TestSemaphore"
	)

	sem, err := New(path, "holder")
	if err != nil {
		t.Error(err)
	}

	oldMax, err := sem.SetMax(2)
	if err != nil {
		t.Error(err)
	}

	oldMax, err = sem.SetMax(1)
	if err != nil {
		t.Error(err)
	}

	if oldMax != 2 {
		t.Error("SetMax did not return expected oldMax")
	}

	err = sem.Acquire(false)
	if err != nil {
		t.Error(err)
	}

	err = sem.Release()
	if err != nil {
		t.Error(err)
	}
}

func TestAcquireWait(t *testing.T) {
	const (
		path = "tests/integration/semaphore/TestAcquireWait"
	)

	sem1, err := New(path, "1")
	if err != nil {
		t.Error(err)
	}

	sem2, err := New(path, "2")
	if err != nil {
		t.Error(err)
	}

	err = sem1.Acquire(false)
	if err != nil {
		t.Error(err)
	}

	go func() {
		err = sem2.Acquire(true)
		if err != nil {
			t.Error(err)
		}

		sem2.Release()
		if err != nil {
			t.Error(err)
		}
	}()

	err = sem1.Release()
	if err != nil {
		t.Error(err)
	}
}

func BenchmarkContention(b *testing.B) {
	const (
		path  = "tests/integration/semaphore/BenchmarkContention"
		count = 9
		max   = 2
	)

	for x := 0; x < b.N; x++ {

		sem, err := New(path, ".")
		if err != nil {
			b.Error(err)
		}

		_, err = sem.SetMax(max)
		if err != nil {
			b.Error(err)
		}

		sems := make([]*Semaphore, 0, count)
		for i := 0; i < count; i++ {
			s, err := New(path, fmt.Sprintf("holder-%v", i))
			if err != nil {
				b.Error(err)
			}
			sems = append(sems, s)
		}

		err = sem.Acquire(false)
		if err != nil {
			b.Error(fmt.Sprintf("Error acquiring for holder %v: %v", sem.Holder, err))
		}

		wg := sync.WaitGroup{}

		for _, s := range sems {
			wg.Add(1)

			go func(s *Semaphore) {
				defer wg.Done()

				ea := s.Acquire(true)
				if ea != nil {
					b.Error(fmt.Sprintf("Error acquiring for holder %v: %v", s.Holder, ea))
					b.FailNow()
				}

				er := s.Release()
				if er != nil {
					b.Error(fmt.Sprintf("Error releasing for holder %v: %v", s.Holder, er))
					b.FailNow()
				}
			}(s)
		}

		wg.Add(1)
		go func() {
			defer wg.Done()

			er := sem.Release()
			if er != nil {
				b.Error(fmt.Sprintf("Error releasing for holder %v: %v", sem.Holder, er))
				b.FailNow()
			}
		}()

		wg.Wait()
	}
}
<commit_msg>Benchmark cleanup.<commit_after>package semaphore

import (
	"fmt"
	"io/ioutil"
	"log"
	"sync"
	"testing"
)

// TODO: Mock out consul-api so these aren't integration tests
// Will be harder since api.Client is a struct rather than an interface.

func TestSemaphore(t *testing.T) {
	const (
		path = "tests/integration/semaphore/TestSemaphore"
	)

	sem, err := New(path, "holder")
	if err != nil {
		t.Error(err)
	}

	oldMax, err := sem.SetMax(2)
	if err != nil {
		t.Error(err)
	}

	oldMax, err = sem.SetMax(1)
	if err != nil {
		t.Error(err)
	}

	if oldMax != 2 {
		t.Error("SetMax did not return expected oldMax")
	}

	err = sem.Acquire(false)
	if err != nil {
		t.Error(err)
	}

	err = sem.Release()
	if err != nil {
		t.Error(err)
	}
}

func TestAcquireWait(t *testing.T) {
	const (
		path = "tests/integration/semaphore/TestAcquireWait"
	)

	sem1, err := New(path, "1")
	if err != nil {
		t.Error(err)
	}

	sem2, err := New(path, "2")
	if err != nil {
		t.Error(err)
	}

	err = sem1.Acquire(false)
	if err != nil {
		t.Error(err)
	}

	go func() {
		err = sem2.Acquire(true)
		if err != nil {
			t.Error(err)
		}

		sem2.Release()
		if err != nil {
			t.Error(err)
		}
	}()

	err = sem1.Release()
	if err != nil {
		t.Error(err)
	}
}

func BenchmarkContention(b *testing.B) {
	const (
		path  = "tests/integration/semaphore/BenchmarkContention"
		count = 20
		max   = 1
	)

	log.SetOutput(ioutil.Discard)

	for x := 0; x < b.N; x++ {

		sem, err := New(path, "INITIAL")
		if err != nil {
			b.Error(err)
		}

		_, err = sem.SetMax(max)
		if err != nil {
			b.Error(err)
		}

		sems := make([]*Semaphore, 0, count)
		for i := 0; i < count; i++ {
			s, err := New(path, fmt.Sprintf("CONTENDER-%v", i))
			if err != nil {
				b.Error(err)
			}
			sems = append(sems, s)
		}

		err = sem.Acquire(false)
		if err != nil {
			b.Error(fmt.Sprintf("Error acquiring for holder %v: %v", sem.Holder, err))
		}

		b.Log(fmt.Sprintf("Holder %v acquired", sem.Holder))

		wg := sync.WaitGroup{}

		for _, s := range sems {
			wg.Add(1)

			go func(s *Semaphore) {
				defer wg.Done()

				ea := s.Acquire(true)
				if ea != nil {
					b.Error(fmt.Sprintf("Error acquiring for holder %v: %v", s.Holder, ea))
					return
				}

				b.Log(fmt.Sprintf("Holder %v acquired, releasing", s.Holder))

				er := s.Release()
				if er != nil {
					b.Error(fmt.Sprintf("Error releasing for holder %v: %v", s.Holder, er))
					return
				}
			}(s)
		}

		b.Log(fmt.Sprintf("Holder %v ...", sem.Holder))

		wg.Add(1)
		go func() {
			defer wg.Done()

			b.Log(fmt.Sprintf("Holder %v releasing", sem.Holder))

			er := sem.Release()
			if er != nil {
				b.Error(fmt.Sprintf("Error releasing for holder %v: %v", sem.Holder, er))
				return
			}
		}()

		/*
			go func() {
				time.Sleep(15 * time.Second)
				panic(fmt.Sprintf("Still waiting %v!", wg))
			}()
		*/

		wg.Wait()
	}
}
<|endoftext|>


<commit_before>package classfile

type AttributeTable struct {
	attributes []AttributeInfo
}

func (self *AttributeTable) CodeAttribute() *CodeAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *CodeAttribute:
			return attrInfo.(*CodeAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) ConstantValueAttribute() *ConstantValueAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *ConstantValueAttribute:
			return attrInfo.(*ConstantValueAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) EnclosingMethodAttribute() *EnclosingMethodAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *EnclosingMethodAttribute:
			return attrInfo.(*EnclosingMethodAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) ExceptionsAttribute() *ExceptionsAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *ExceptionsAttribute:
			return attrInfo.(*ExceptionsAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) LineNumberTableAttribute() *LineNumberTableAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *LineNumberTableAttribute:
			return attrInfo.(*LineNumberTableAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) SignatureAttribute() *SignatureAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *SignatureAttribute:
			return attrInfo.(*SignatureAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) SourceFileAttribute() *SourceFileAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *SourceFileAttribute:
			return attrInfo.(*SourceFileAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) UndefinedAttribute(name string) *UndefinedAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *UndefinedAttribute:
			return attrInfo.(*UndefinedAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) RuntimeVisibleAnnotationsAttribute() *UndefinedAttribute {
	return self.UndefinedAttribute("RuntimeVisibleAnnotations")
}
func (self *AttributeTable) RuntimeVisibleParameterAnnotationsAttribute() *UndefinedAttribute {
	return self.UndefinedAttribute("RuntimeVisibleParameterAnnotationsAttribute")
}
func (self *AttributeTable) AnnotationDefaultAttribute() *UndefinedAttribute {
	return self.UndefinedAttribute("AnnotationDefault")
}
<commit_msg>rename func<commit_after>package classfile

type AttributeTable struct {
	attributes []AttributeInfo
}

func (self *AttributeTable) CodeAttribute() *CodeAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *CodeAttribute:
			return attrInfo.(*CodeAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) ConstantValueAttribute() *ConstantValueAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *ConstantValueAttribute:
			return attrInfo.(*ConstantValueAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) EnclosingMethodAttribute() *EnclosingMethodAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *EnclosingMethodAttribute:
			return attrInfo.(*EnclosingMethodAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) ExceptionsAttribute() *ExceptionsAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *ExceptionsAttribute:
			return attrInfo.(*ExceptionsAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) LineNumberTableAttribute() *LineNumberTableAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *LineNumberTableAttribute:
			return attrInfo.(*LineNumberTableAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) SignatureAttribute() *SignatureAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *SignatureAttribute:
			return attrInfo.(*SignatureAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) SourceFileAttribute() *SourceFileAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *SourceFileAttribute:
			return attrInfo.(*SourceFileAttribute)
		}
	}
	return nil
}

func (self *AttributeTable) RuntimeVisibleAnnotationsAttribute() *UndefinedAttribute {
	return self.getUndefinedAttribute("RuntimeVisibleAnnotations")
}
func (self *AttributeTable) RuntimeVisibleParameterAnnotationsAttribute() *UndefinedAttribute {
	return self.getUndefinedAttribute("RuntimeVisibleParameterAnnotationsAttribute")
}
func (self *AttributeTable) AnnotationDefaultAttribute() *UndefinedAttribute {
	return self.getUndefinedAttribute("AnnotationDefault")
}

func (self *AttributeTable) getUndefinedAttribute(name string) *UndefinedAttribute {
	for _, attrInfo := range self.attributes {
		switch attrInfo.(type) {
		case *UndefinedAttribute:
			return attrInfo.(*UndefinedAttribute)
		}
	}
	return nil
}
<|endoftext|>


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

package integration

import (
	"bufio"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/integration"
	"github.com/coreos/etcd/pkg/testutil"
	"github.com/coreos/etcd/pkg/transport"

	grpcprom "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

func TestV3ClientMetrics(t *testing.T) {
	defer testutil.AfterTest(t)

	var (
		addr string = "localhost:27989"
		ln   net.Listener
		err  error
	)

	// listen for all prometheus metrics
	donec := make(chan struct{})
	go func() {
		defer close(donec)

		srv := &http.Server{Handler: prometheus.Handler()}
		srv.SetKeepAlivesEnabled(false)

		ln, err = transport.NewUnixListener(addr)
		if err != nil {
			t.Fatalf("Error: %v occurred while listening on addr: %v", err, addr)
		}

		err = srv.Serve(ln)
		if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
			t.Fatalf("Err serving http requests: %v", err)
		}
	}()

	url := "unix://" + addr + "/metrics"

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	cfg := clientv3.Config{
		Endpoints: []string{clus.Members[0].GRPCAddr()},
		DialOptions: []grpc.DialOption{
			grpc.WithUnaryInterceptor(grpcprom.UnaryClientInterceptor),
			grpc.WithStreamInterceptor(grpcprom.StreamClientInterceptor),
		},
	}
	cli, cerr := clientv3.New(cfg)
	if cerr != nil {
		t.Fatal(cerr)
	}
	defer cli.Close()

	wc := cli.Watch(context.Background(), "foo")

	wBefore := sumCountersForMetricAndLabels(t, url, "grpc_client_msg_received_total", "Watch", "bidi_stream")

	pBefore := sumCountersForMetricAndLabels(t, url, "grpc_client_started_total", "Put", "unary")

	_, err = cli.Put(context.Background(), "foo", "bar")
	if err != nil {
		t.Errorf("Error putting value in key store")
	}

	pAfter := sumCountersForMetricAndLabels(t, url, "grpc_client_started_total", "Put", "unary")
	if pBefore+1 != pAfter {
		t.Errorf("grpc_client_started_total expected %d, got %d", 1, pAfter-pBefore)
	}

	// consume watch response
	select {
	case <-wc:
	case <-time.After(10 * time.Second):
		t.Error("Timeout occurred for getting watch response")
	}

	wAfter := sumCountersForMetricAndLabels(t, url, "grpc_client_msg_received_total", "Watch", "bidi_stream")
	if wBefore+1 != wAfter {
		t.Errorf("grpc_client_msg_received_total expected %d, got %d", 1, wAfter-wBefore)
	}

	ln.Close()
	<-donec
}

func sumCountersForMetricAndLabels(t *testing.T, url string, metricName string, matchingLabelValues ...string) int {
	count := 0
	for _, line := range getHTTPBodyAsLines(t, url) {
		ok := true
		if !strings.HasPrefix(line, metricName) {
			continue
		}

		for _, labelValue := range matchingLabelValues {
			if !strings.Contains(line, `"`+labelValue+`"`) {
				ok = false
				break
			}
		}

		if !ok {
			continue
		}

		valueString := line[strings.LastIndex(line, " ")+1 : len(line)-1]
		valueFloat, err := strconv.ParseFloat(valueString, 32)
		if err != nil {
			t.Fatalf("failed parsing value for line: %v and matchingLabelValues: %v", line, matchingLabelValues)
		}
		count += int(valueFloat)
	}
	return count
}

func getHTTPBodyAsLines(t *testing.T, url string) []string {
	cfgtls := transport.TLSInfo{}
	tr, err := transport.NewTransport(cfgtls, time.Second)
	if err != nil {
		t.Fatalf("Error getting transport: %v", err)
	}

	tr.MaxIdleConns = -1
	tr.DisableKeepAlives = true

	cli := &http.Client{Transport: tr}

	resp, err := cli.Get(url)
	if err != nil {
		t.Fatalf("Error fetching: %v", err)
	}

	reader := bufio.NewReader(resp.Body)
	lines := []string{}
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			} else {
				t.Fatalf("error reading: %v", err)
			}
		}
		lines = append(lines, line)
	}
	resp.Body.Close()
	return lines
}
<commit_msg>clientv3/integration: use 'transport.IsClosedConnError'<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package integration

import (
	"bufio"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/integration"
	"github.com/coreos/etcd/pkg/testutil"
	"github.com/coreos/etcd/pkg/transport"

	grpcprom "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

func TestV3ClientMetrics(t *testing.T) {
	defer testutil.AfterTest(t)

	var (
		addr string = "localhost:27989"
		ln   net.Listener
		err  error
	)

	// listen for all prometheus metrics
	donec := make(chan struct{})
	go func() {
		defer close(donec)

		srv := &http.Server{Handler: prometheus.Handler()}
		srv.SetKeepAlivesEnabled(false)

		ln, err = transport.NewUnixListener(addr)
		if err != nil {
			t.Fatalf("Error: %v occurred while listening on addr: %v", err, addr)
		}

		err = srv.Serve(ln)
		if err != nil && !transport.IsClosedConnError(err) {
			t.Fatalf("Err serving http requests: %v", err)
		}
	}()

	url := "unix://" + addr + "/metrics"

	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	cfg := clientv3.Config{
		Endpoints: []string{clus.Members[0].GRPCAddr()},
		DialOptions: []grpc.DialOption{
			grpc.WithUnaryInterceptor(grpcprom.UnaryClientInterceptor),
			grpc.WithStreamInterceptor(grpcprom.StreamClientInterceptor),
		},
	}
	cli, cerr := clientv3.New(cfg)
	if cerr != nil {
		t.Fatal(cerr)
	}
	defer cli.Close()

	wc := cli.Watch(context.Background(), "foo")

	wBefore := sumCountersForMetricAndLabels(t, url, "grpc_client_msg_received_total", "Watch", "bidi_stream")

	pBefore := sumCountersForMetricAndLabels(t, url, "grpc_client_started_total", "Put", "unary")

	_, err = cli.Put(context.Background(), "foo", "bar")
	if err != nil {
		t.Errorf("Error putting value in key store")
	}

	pAfter := sumCountersForMetricAndLabels(t, url, "grpc_client_started_total", "Put", "unary")
	if pBefore+1 != pAfter {
		t.Errorf("grpc_client_started_total expected %d, got %d", 1, pAfter-pBefore)
	}

	// consume watch response
	select {
	case <-wc:
	case <-time.After(10 * time.Second):
		t.Error("Timeout occurred for getting watch response")
	}

	wAfter := sumCountersForMetricAndLabels(t, url, "grpc_client_msg_received_total", "Watch", "bidi_stream")
	if wBefore+1 != wAfter {
		t.Errorf("grpc_client_msg_received_total expected %d, got %d", 1, wAfter-wBefore)
	}

	ln.Close()
	<-donec
}

func sumCountersForMetricAndLabels(t *testing.T, url string, metricName string, matchingLabelValues ...string) int {
	count := 0
	for _, line := range getHTTPBodyAsLines(t, url) {
		ok := true
		if !strings.HasPrefix(line, metricName) {
			continue
		}

		for _, labelValue := range matchingLabelValues {
			if !strings.Contains(line, `"`+labelValue+`"`) {
				ok = false
				break
			}
		}

		if !ok {
			continue
		}

		valueString := line[strings.LastIndex(line, " ")+1 : len(line)-1]
		valueFloat, err := strconv.ParseFloat(valueString, 32)
		if err != nil {
			t.Fatalf("failed parsing value for line: %v and matchingLabelValues: %v", line, matchingLabelValues)
		}
		count += int(valueFloat)
	}
	return count
}

func getHTTPBodyAsLines(t *testing.T, url string) []string {
	cfgtls := transport.TLSInfo{}
	tr, err := transport.NewTransport(cfgtls, time.Second)
	if err != nil {
		t.Fatalf("Error getting transport: %v", err)
	}

	tr.MaxIdleConns = -1
	tr.DisableKeepAlives = true

	cli := &http.Client{Transport: tr}

	resp, err := cli.Get(url)
	if err != nil {
		t.Fatalf("Error fetching: %v", err)
	}

	reader := bufio.NewReader(resp.Body)
	lines := []string{}
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			} else {
				t.Fatalf("error reading: %v", err)
			}
		}
		lines = append(lines, line)
	}
	resp.Body.Close()
	return lines
}
<|endoftext|>


<commit_before>// Copyright (c) 2020 Tigera, Inc. All rights reserved.
//
// 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 commands

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/projectcalico/felix/bpf"

	"github.com/projectcalico/felix/bpf/conntrack"

	"github.com/docopt/docopt-go"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

func init() {
	conntrackCmd.AddCommand(newConntrackDumpCmd())
	conntrackCmd.AddCommand(newConntrackRemoveCmd())
	rootCmd.AddCommand(conntrackCmd)
}

// conntrackCmd represents the conntrack command
var conntrackCmd = &cobra.Command{
	Use:   "conntrack",
	Short: "Manipulates connection tracking",
}

type conntrackDumpCmd struct {
	*cobra.Command
}

func newConntrackDumpCmd() *cobra.Command {
	cmd := &conntrackDumpCmd{
		Command: &cobra.Command{
			Use:   "dump",
			Short: "Dumps connection tracking table",
		},
	}

	cmd.Command.Args = cmd.Args
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackDumpCmd) Args(c *cobra.Command, args []string) error {
	a, err := docopt.ParseArgs(makeDocUsage(c), args, "")
	if err != nil {
		return errors.New(err.Error())
	}

	err = a.Bind(cmd)
	if err != nil {
		return errors.New(err.Error())
	}

	return nil
}

func (cmd *conntrackDumpCmd) Run(c *cobra.Command, _ []string) {
	mc := &bpf.MapContext{}
	ctMap := conntrack.Map(mc)
	err := ctMap.Iter(func(k, v []byte) {
		var ctKey conntrack.Key
		if len(k) != len(ctKey) {
			log.Panic("Key has unexpected length")
		}
		copy(ctKey[:], k[:])

		var ctVal conntrack.Value
		if len(v) != len(ctVal) {
			log.Panic("Value has unexpected length")
		}
		copy(ctVal[:], v[:])

		fmt.Printf("%v -> %v", ctKey, ctVal)
		dumpExtra(ctKey, ctVal)
		fmt.Printf("\n")
	})
	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}
}

func dumpExtra(k conntrack.Key, v conntrack.Value) {
	now := bpf.KTimeNanos()

	fmt.Printf(" Age: %s Active ago %s",
		time.Duration(now-v.Created()), time.Duration(now-v.LastSeen()))

	if k.Proto() != conntrack.ProtoTCP {
		return
	}

	if v.Type() == conntrack.TypeNATForward {
		return
	}

	data := v.Data()

	if (v.IsForwardDSR() && data.FINsSeenDSR()) || data.FINsSeen() {
		fmt.Printf(" CLOSED")
		return
	}

	if data.Established() {
		fmt.Printf(" ESTABLISHED")
		return
	}

	fmt.Printf(" SYN-SENT")
}

type conntrackRemoveCmd struct {
	*cobra.Command

	Proto string `docopt:"<proto>"`
	IP1   string `docopt:"<ip1>"`
	IP2   string `docopt:"<ip2>"`

	proto uint8
	ip1   net.IP
	ip2   net.IP
}

func newConntrackRemoveCmd() *cobra.Command {
	cmd := &conntrackRemoveCmd{
		Command: &cobra.Command{
			Use:   "remove <proto> <ip1> <ip2>",
			Short: "removes connection tracking",
		},
	}

	cmd.Command.Args = cmd.Args
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackRemoveCmd) Args(c *cobra.Command, args []string) error {
	a, err := docopt.ParseArgs(makeDocUsage(c), args, "")
	if err != nil {
		return errors.New(err.Error())
	}

	err = a.Bind(cmd)
	if err != nil {
		return errors.New(err.Error())
	}

	switch proto := strings.ToLower(args[0]); proto {
	case "udp":
		cmd.proto = 17
	case "tcp":
		cmd.proto = 6
	default:
		return errors.Errorf("unknown protocol %s", proto)
	}

	cmd.ip1 = net.ParseIP(cmd.IP1)
	if cmd.ip1 == nil {
		return errors.Errorf("ip1: %q is not an ip", cmd.IP1)
	}

	cmd.ip2 = net.ParseIP(cmd.IP2)
	if cmd.ip2 == nil {
		return errors.Errorf("ip2: %q is not an ip", cmd.IP2)
	}

	return nil
}

func (cmd *conntrackRemoveCmd) Run(c *cobra.Command, _ []string) {
	mc := &bpf.MapContext{}
	ctMap := conntrack.Map(mc)
	var keysToRemove []conntrack.Key
	err := ctMap.Iter(func(k, v []byte) {
		var ctKey conntrack.Key
		if len(k) != len(ctKey) {
			log.Panic("Key has unexpected length")
		}
		copy(ctKey[:], k[:])

		log.Infof("Examining conntrack key: %v", ctKey)

		if ctKey.Proto() != cmd.proto {
			return
		}

		if ctKey.AddrA().Equal(cmd.ip1) && ctKey.AddrB().Equal(cmd.ip2) {
			log.Info("Match")
			keysToRemove = append(keysToRemove, ctKey)
		} else if ctKey.AddrB().Equal(cmd.ip1) && ctKey.AddrA().Equal(cmd.ip2) {
			log.Info("Match")
			keysToRemove = append(keysToRemove, ctKey)
		}
	})
	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}

	for _, k := range keysToRemove {
		err := ctMap.Delete(k[:])
		if err != nil {
			log.WithError(err).WithField("key", k).Warning("Failed to delete entry from map")
		}
	}
}
<commit_msg>Changes to fix the issues with conntrack test<commit_after>// Copyright (c) 2020 Tigera, Inc. All rights reserved.
//
// 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 commands

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/projectcalico/felix/bpf"

	"github.com/projectcalico/felix/bpf/conntrack"

	"github.com/docopt/docopt-go"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

func init() {
	conntrackCmd.AddCommand(newConntrackDumpCmd())
	conntrackCmd.AddCommand(newConntrackRemoveCmd())
	rootCmd.AddCommand(conntrackCmd)
}

// conntrackCmd represents the conntrack command
var conntrackCmd = &cobra.Command{
	Use:   "conntrack",
	Short: "Manipulates connection tracking",
}

type conntrackDumpCmd struct {
	*cobra.Command
}

func newConntrackDumpCmd() *cobra.Command {
	cmd := &conntrackDumpCmd{
		Command: &cobra.Command{
			Use:   "dump",
			Short: "Dumps connection tracking table",
		},
	}

	cmd.Command.Args = cmd.Args
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackDumpCmd) Args(c *cobra.Command, args []string) error {
	a, err := docopt.ParseArgs(makeDocUsage(c), args, "")
	if err != nil {
		return errors.New(err.Error())
	}

	err = a.Bind(cmd)
	if err != nil {
		return errors.New(err.Error())
	}

	return nil
}

func (cmd *conntrackDumpCmd) Run(c *cobra.Command, _ []string) {
	mc := &bpf.MapContext{}
	ctMap := conntrack.Map(mc)
	if err := ctMap.EnsureExists(); err != nil {
		log.WithError(err).Error("Failed to access ConntrackMap")
	}
	err := ctMap.Iter(func(k, v []byte) {
		var ctKey conntrack.Key
		if len(k) != len(ctKey) {
			log.Panic("Key has unexpected length")
		}
		copy(ctKey[:], k[:])

		var ctVal conntrack.Value
		if len(v) != len(ctVal) {
			log.Panic("Value has unexpected length")
		}
		copy(ctVal[:], v[:])

		fmt.Printf("%v -> %v", ctKey, ctVal)
		dumpExtra(ctKey, ctVal)
		fmt.Printf("\n")
	})
	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}
}

func dumpExtra(k conntrack.Key, v conntrack.Value) {
	now := bpf.KTimeNanos()

	fmt.Printf(" Age: %s Active ago %s",
		time.Duration(now-v.Created()), time.Duration(now-v.LastSeen()))

	if k.Proto() != conntrack.ProtoTCP {
		return
	}

	if v.Type() == conntrack.TypeNATForward {
		return
	}

	data := v.Data()

	if (v.IsForwardDSR() && data.FINsSeenDSR()) || data.FINsSeen() {
		fmt.Printf(" CLOSED")
		return
	}

	if data.Established() {
		fmt.Printf(" ESTABLISHED")
		return
	}

	fmt.Printf(" SYN-SENT")
}

type conntrackRemoveCmd struct {
	*cobra.Command

	Proto string `docopt:"<proto>"`
	IP1   string `docopt:"<ip1>"`
	IP2   string `docopt:"<ip2>"`

	proto uint8
	ip1   net.IP
	ip2   net.IP
}

func newConntrackRemoveCmd() *cobra.Command {
	cmd := &conntrackRemoveCmd{
		Command: &cobra.Command{
			Use:   "remove <proto> <ip1> <ip2>",
			Short: "removes connection tracking",
		},
	}

	cmd.Command.Args = cmd.Args
	cmd.Command.Run = cmd.Run

	return cmd.Command
}

func (cmd *conntrackRemoveCmd) Args(c *cobra.Command, args []string) error {
	a, err := docopt.ParseArgs(makeDocUsage(c), args, "")
	if err != nil {
		return errors.New(err.Error())
	}

	err = a.Bind(cmd)
	if err != nil {
		return errors.New(err.Error())
	}

	switch proto := strings.ToLower(args[0]); proto {
	case "udp":
		cmd.proto = 17
	case "tcp":
		cmd.proto = 6
	default:
		return errors.Errorf("unknown protocol %s", proto)
	}

	cmd.ip1 = net.ParseIP(cmd.IP1)
	if cmd.ip1 == nil {
		return errors.Errorf("ip1: %q is not an ip", cmd.IP1)
	}

	cmd.ip2 = net.ParseIP(cmd.IP2)
	if cmd.ip2 == nil {
		return errors.Errorf("ip2: %q is not an ip", cmd.IP2)
	}

	return nil
}

func (cmd *conntrackRemoveCmd) Run(c *cobra.Command, _ []string) {
	mc := &bpf.MapContext{}
	ctMap := conntrack.Map(mc)
	if err := ctMap.EnsureExists(); err != nil {
		log.WithError(err).Error("Failed to access ConntrackMap")
	}
	var keysToRemove []conntrack.Key
	err := ctMap.Iter(func(k, v []byte) {
		var ctKey conntrack.Key
		if len(k) != len(ctKey) {
			log.Panic("Key has unexpected length")
		}
		copy(ctKey[:], k[:])

		log.Infof("Examining conntrack key: %v", ctKey)

		if ctKey.Proto() != cmd.proto {
			return
		}

		if ctKey.AddrA().Equal(cmd.ip1) && ctKey.AddrB().Equal(cmd.ip2) {
			log.Info("Match")
			keysToRemove = append(keysToRemove, ctKey)
		} else if ctKey.AddrB().Equal(cmd.ip1) && ctKey.AddrA().Equal(cmd.ip2) {
			log.Info("Match")
			keysToRemove = append(keysToRemove, ctKey)
		}
	})
	if err != nil {
		log.WithError(err).Fatal("Failed to iterate over conntrack entries")
	}

	for _, k := range keysToRemove {
		err := ctMap.Delete(k[:])
		if err != nil {
			log.WithError(err).WithField("key", k).Warning("Failed to delete entry from map")
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rectest

import (
	"testing"
	"time"

	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/db/dbtest"
	"gopkg.in/check.v1"
)

func Test(t *testing.T) {
	check.TestingT(t)
}

type CheckerSuite struct{}

var _ = check.Suite(CheckerSuite{})

func (CheckerSuite) SetUpTest(c *check.C) {
	config.Set("database:url", "localhost:27017")
	config.Set("database:name", "tsuru_rectest_test")
	conn, err := db.Conn()
	c.Assert(err, check.IsNil)
	defer conn.Close()
	action := map[string]interface{}{
		"user":   "glenda@tsuru.io",
		"action": "run-command",
		"extra":  []interface{}{"rm", "-rf", "/"},
		"date":   time.Now(),
	}
	err = conn.UserActions().Insert(action)
	c.Assert(err, check.IsNil)
	actionNoDate := map[string]interface{}{
		"user":   "glenda@tsuru.io",
		"action": "list-apps",
		"extra":  nil,
		"date":   nil,
	}
	err = conn.UserActions().Insert(actionNoDate)
	c.Assert(err, check.IsNil)
}

func (CheckerSuite) TearDownTest(c *check.C) {
	conn, err := db.Conn()
	c.Assert(err, check.IsNil)
	defer conn.Close()
	dbtest.ClearAllCollections(conn.Apps().Database)
}

func (CheckerSuite) TestIsRecordedInfo(c *check.C) {
	expected := &check.CheckerInfo{
		Name:   "IsRecorded",
		Params: []string{"action"},
	}
	c.Assert(isRecordedChecker{}.Info(), check.DeepEquals, expected)
}

func (CheckerSuite) TestIsRecordedCheckInvalidParameter(c *check.C) {
	result, error := isRecordedChecker{}.Check([]interface{}{"action"}, []string{"action"})
	c.Assert(result, check.Equals, false)
	c.Assert(error, check.Equals, "First parameter must be of type Action or *Action")
}

func (CheckerSuite) TestIsRecordedCheckWithValue(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "run-command",
		Extra:  []interface{}{"rm", "-rf", "/"},
	}
	result, error := isRecordedChecker{}.Check([]interface{}{action}, []string{})
	c.Assert(result, check.Equals, true)
	c.Assert(error, check.Equals, "")
}

func (CheckerSuite) TestIsRecordedCheckWithReference(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "run-command",
		Extra:  []interface{}{"rm", "-rf", "/"},
	}
	result, error := isRecordedChecker{}.Check([]interface{}{&action}, []string{})
	c.Assert(result, check.Equals, true)
	c.Assert(error, check.Equals, "")
}

func (CheckerSuite) TestIsRecordedNotInDatabase(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "run-command",
		Extra:  []interface{}{"rm", "-rf", "/home"},
	}
	result, error := isRecordedChecker{}.Check([]interface{}{action}, []string{})
	c.Assert(result, check.Equals, false)
	c.Assert(error, check.Equals, "Action not in the database")
}

func (CheckerSuite) TestIsRecordedWithoutDate(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "list-apps",
	}
	result, error := isRecordedChecker{}.Check([]interface{}{action}, []string{})
	c.Assert(result, check.Equals, false)
	c.Assert(error, check.Equals, "Action was not recorded using rec.Log")
}
<commit_msg>rec/rectest: change name of actions stored in the database<commit_after>// Copyright 2015 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rectest

import (
	"testing"
	"time"

	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/db/dbtest"
	"gopkg.in/check.v1"
)

func Test(t *testing.T) {
	check.TestingT(t)
}

type CheckerSuite struct{}

var _ = check.Suite(CheckerSuite{})

func (CheckerSuite) SetUpTest(c *check.C) {
	config.Set("database:url", "localhost:27017")
	config.Set("database:name", "tsuru_rectest_test")
	conn, err := db.Conn()
	c.Assert(err, check.IsNil)
	defer conn.Close()
	action := map[string]interface{}{
		"user":   "glenda@tsuru.io",
		"action": "test-run-command",
		"extra":  []interface{}{"rm", "-rf", "/"},
		"date":   time.Now(),
	}
	err = conn.UserActions().Insert(action)
	c.Assert(err, check.IsNil)
	actionNoDate := map[string]interface{}{
		"user":   "glenda@tsuru.io",
		"action": "test-list-apps",
		"extra":  nil,
		"date":   nil,
	}
	err = conn.UserActions().Insert(actionNoDate)
	c.Assert(err, check.IsNil)
}

func (CheckerSuite) TearDownTest(c *check.C) {
	conn, err := db.Conn()
	c.Assert(err, check.IsNil)
	defer conn.Close()
	dbtest.ClearAllCollections(conn.Apps().Database)
}

func (CheckerSuite) TestIsRecordedInfo(c *check.C) {
	expected := &check.CheckerInfo{
		Name:   "IsRecorded",
		Params: []string{"action"},
	}
	c.Assert(isRecordedChecker{}.Info(), check.DeepEquals, expected)
}

func (CheckerSuite) TestIsRecordedCheckInvalidParameter(c *check.C) {
	result, error := isRecordedChecker{}.Check([]interface{}{"action"}, []string{"action"})
	c.Assert(result, check.Equals, false)
	c.Assert(error, check.Equals, "First parameter must be of type Action or *Action")
}

func (CheckerSuite) TestIsRecordedCheckWithValue(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "test-run-command",
		Extra:  []interface{}{"rm", "-rf", "/"},
	}
	result, error := isRecordedChecker{}.Check([]interface{}{action}, []string{})
	c.Assert(result, check.Equals, true)
	c.Assert(error, check.Equals, "")
}

func (CheckerSuite) TestIsRecordedCheckWithReference(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "test-run-command",
		Extra:  []interface{}{"rm", "-rf", "/"},
	}
	result, error := isRecordedChecker{}.Check([]interface{}{&action}, []string{})
	c.Assert(result, check.Equals, true)
	c.Assert(error, check.Equals, "")
}

func (CheckerSuite) TestIsRecordedNotInDatabase(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "test-run-command",
		Extra:  []interface{}{"rm", "-rf", "/home"},
	}
	result, error := isRecordedChecker{}.Check([]interface{}{action}, []string{})
	c.Assert(result, check.Equals, false)
	c.Assert(error, check.Equals, "Action not in the database")
}

func (CheckerSuite) TestIsRecordedWithoutDate(c *check.C) {
	action := Action{
		User:   "glenda@tsuru.io",
		Action: "test-list-apps",
	}
	result, error := isRecordedChecker{}.Check([]interface{}{action}, []string{})
	c.Assert(result, check.Equals, false)
	c.Assert(error, check.Equals, "Action was not recorded using rec.Log")
}
<|endoftext|>


<commit_before>// Copyright 2016 The fleet 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 rpc

import (
	"errors"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/coreos/go-semver/semver"

	"github.com/coreos/fleet/engine"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/log"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg/lease"
	"github.com/coreos/fleet/registry"
	"github.com/coreos/fleet/unit"
)

type RegistryMux struct {
	etcdRegistry    *registry.EtcdRegistry
	localMachine    machine.Machine
	rpcserver       *rpcserver
	currentRegistry registry.Registry
	rpcRegistry     *RPCRegistry
	currentEngine   machine.MachineState
	leaseManager    lease.Manager

	handlingEngineChange *sync.RWMutex
}

const (
	dialRegistryReconnectTimeout = 200 * time.Millisecond

	engineLeaderKeyPath = "engine-leader"
)

func NewRegistryMux(etcdRegistry *registry.EtcdRegistry, localMachine machine.Machine, leaseManager lease.Manager) *RegistryMux {
	return &RegistryMux{
		etcdRegistry:         etcdRegistry,
		localMachine:         localMachine,
		handlingEngineChange: new(sync.RWMutex),
		leaseManager:         leaseManager,
	}
}

// ConnectToRegistry allows to disable_engine fleet agents to adapt its Registry
// to fleet leader changes regardless of whether is etcd or gRPC based.
func (r *RegistryMux) ConnectToRegistry(e *engine.Engine) {
	for {
		// We have to check if the leader has changed to etcd otherwise keep grpc connection
		isGrpc, err := e.IsGrpcLeader()
		// If there is not error then we are able to get the leader state and continue
		// otherwise we have to wait
		if err == nil {
			if isGrpc {
				if r.rpcRegistry != nil && r.rpcRegistry.IsRegistryReady() {
					log.Infof("Reusing gRPC engine, connection is READY\n")
					r.currentRegistry = r.rpcRegistry
				} else {
					if r.rpcRegistry != nil {
						r.rpcRegistry.Close()
					}
					log.Infof("New engine supports gRPC, connecting\n")
					r.rpcRegistry = NewRPCRegistry(r.rpcDialerNoEngine)
					// connect to rpc registry
					r.rpcRegistry.Connect()
					r.currentRegistry = r.rpcRegistry
				}
			} else {
				if r.rpcRegistry != nil {
					r.rpcRegistry.Close()
				}
				// new leader is etcd-based
				r.currentRegistry = r.etcdRegistry
			}
		}
		time.Sleep(5 * time.Second)
	}
}

func (r *RegistryMux) rpcDialerNoEngine(_ string, timeout time.Duration) (net.Conn, error) {
	ticker := time.Tick(dialRegistryReconnectTimeout)
	// Timeout re-defined to call etcd every 5secs to get the leader
	timeout = 5 * time.Second
	check := time.After(timeout)

	for {
		select {
		case <-check:
			log.Errorf("Unable to connect to engine %s\n", r.currentEngine.PublicIP)
			// Get the new engine leader of the cluster out of etcd
			lease, err := r.leaseManager.GetLease(engineLeaderKeyPath)
			// Key found
			if err == nil && lease != nil {
				var err error
				machines, err := r.etcdRegistry.Machines()
				if err != nil {
					log.Errorf("Unable to get the machines of the cluster %v\n", err)
					return nil, errors.New("Unable to get the machines of the cluster")
				}
				for _, s := range machines {
					// Update the currentEngine with the new one... otherwise wait until
					// there is one
					if s.ID == lease.MachineID() {
						// New leader has not gRPC capabilities enabled.
						if !s.Capabilities.Has(machine.CapGRPC) {
							log.Error("New leader engine has not gRPC enabled!")
							return nil, errors.New("New leader engine has not gRPC enabled!")
						}
						r.currentEngine = s
						log.Infof("Found a new engine to connect to: %s\n", r.currentEngine.PublicIP)
						// Restore initial check configuration
						timeout = 5 * time.Second
						check = time.After(timeout)
					}
				}
			} else {
				timeout = 2 * time.Second
				log.Errorf("Unable to get the leader engine, retrying in %v...", timeout)
				check = time.After(timeout)
			}
		case <-ticker:
			addr := fmt.Sprintf("%s:%d", r.currentEngine.PublicIP, rpcServerPort)
			conn, err := net.Dial("tcp", addr)
			if err == nil {
				log.Infof("Connected to engine on %s\n", r.currentEngine.PublicIP)
				return conn, nil
			}
			log.Errorf("Retry to connect to new engine: %+v", err)
		}
	}
}

func (r *RegistryMux) rpcDialer(_ string, timeout time.Duration) (net.Conn, error) {
	ticker := time.Tick(dialRegistryReconnectTimeout)
	alert := time.After(timeout)

	for {
		select {
		case <-alert:
			log.Errorf("Unable to connect to engine %s\n", r.currentEngine.PublicIP)
			return nil, errors.New("Unable to connect to new engine, the client connection is closing")
		case <-ticker:
			addr := fmt.Sprintf("%s:%d", r.currentEngine.PublicIP, rpcServerPort)
			conn, err := net.Dial("tcp", addr)
			if err == nil {
				log.Infof("Connected to engine on %s\n", r.currentEngine.PublicIP)
				return conn, nil
			}
			log.Errorf("Retry to connect to new engine: %+v", err)
		}
	}
}

func (r *RegistryMux) EngineChanged(newEngine machine.MachineState) {
	r.handlingEngineChange.Lock()
	defer r.handlingEngineChange.Unlock()

	stopServer := false
	if r.currentEngine.ID != newEngine.ID {
		stopServer = true
	}
	r.currentEngine = newEngine
	log.Infof("Engine changed, checking capabilities %+v", newEngine)
	if r.localMachine.State().Capabilities.Has(machine.CapGRPC) {
		if r.rpcserver != nil && ((r.rpcRegistry != nil && !r.rpcRegistry.IsRegistryReady()) || stopServer) {
			// If the engine changed, we need to stop the rpc server
			r.rpcserver.Stop()
			r.rpcserver = nil
		}
		if newEngine.ID == r.localMachine.State().ID {
			if r.rpcserver == nil {
				// start rpc server
				log.Infof("Starting rpc server...\n")
				var err error
				r.rpcserver, err = NewRPCServer(r.etcdRegistry, newEngine.PublicIP)
				if err != nil {
					log.Fatalf("Unable to create rpc server %+v", err)
				}

				go func() {
					errc := make(chan error, 1)
					if errc <- r.rpcserver.Start(); <-errc != nil {
						log.Fatalf("Failed to serve gRPC requests on listener: %v", <-errc)
					}
				}()
			}

		}
		if newEngine.Capabilities.Has(machine.CapGRPC) {
			if r.rpcRegistry != nil && r.rpcRegistry.IsRegistryReady() {
				log.Infof("Reusing gRPC engine, connection is READY\n")
				r.currentRegistry = r.rpcRegistry
			} else {
				log.Infof("New engine supports gRPC, connecting\n")
				r.rpcRegistry = NewRPCRegistry(r.rpcDialer)
				// connect to rpc registry
				r.rpcRegistry.Connect()
				r.currentRegistry = r.rpcRegistry
			}
		} else {
			log.Infof("Falling back to etcd registry\n")
			if r.rpcserver != nil {
				// If the engine changed to a non gRPC leader, we need to stop the server
				r.rpcserver.Stop()
			}
			r.currentRegistry = r.etcdRegistry
		}

	} else {
		log.Infof("Falling back to etcd registry\n")
		r.currentRegistry = r.etcdRegistry
	}
}

func (r *RegistryMux) getRegistry() registry.Registry {
	r.handlingEngineChange.RLock()
	defer r.handlingEngineChange.RUnlock()
	if r.currentRegistry == nil {
		return r.etcdRegistry
	}
	return r.currentRegistry
}

func (r *RegistryMux) IsRegistryReady() bool {
	return r.getRegistry().IsRegistryReady()
}

func (r *RegistryMux) UseEtcdRegistry() bool {
	return r.getRegistry().UseEtcdRegistry()
}

func (r *RegistryMux) ClearUnitHeartbeat(name string) {
	r.getRegistry().ClearUnitHeartbeat(name)
}

func (r *RegistryMux) CreateUnit(unit *job.Unit) error {
	return r.getRegistry().CreateUnit(unit)
}

func (r *RegistryMux) CreateMachineState(ms machine.MachineState, ttl time.Duration) (uint64, error) {
	return r.getRegistry().CreateMachineState(ms, ttl)
}

func (r *RegistryMux) DestroyUnit(unit string) error {
	return r.getRegistry().DestroyUnit(unit)
}

func (r *RegistryMux) UnitHeartbeat(name string, machID string, ttl time.Duration) error {
	return r.getRegistry().UnitHeartbeat(name, machID, ttl)
}

func (r *RegistryMux) Machines() ([]machine.MachineState, error) {
	return r.etcdRegistry.Machines()
}

func (r *RegistryMux) RemoveMachineState(machID string) error {
	return r.etcdRegistry.RemoveMachineState(machID)
}

func (r *RegistryMux) RemoveUnitState(jobName string) error {
	return r.getRegistry().RemoveUnitState(jobName)
}

func (r *RegistryMux) SaveUnitState(jobName string, unitState *unit.UnitState, ttl time.Duration) {
	r.getRegistry().SaveUnitState(jobName, unitState, ttl)
}

func (r *RegistryMux) ScheduleUnit(name string, machID string) error {
	return r.getRegistry().ScheduleUnit(name, machID)
}

func (r *RegistryMux) SetUnitTargetState(name string, state job.JobState) error {
	return r.getRegistry().SetUnitTargetState(name, state)
}

func (r *RegistryMux) MachineState(machID string) (machine.MachineState, error) {
	return r.getRegistry().MachineState(machID)
}

func (r *RegistryMux) SetMachineState(ms machine.MachineState, ttl time.Duration) (uint64, error) {
	return r.etcdRegistry.SetMachineState(ms, ttl)
}

func (r *RegistryMux) UnscheduleUnit(name string, machID string) error {
	return r.getRegistry().UnscheduleUnit(name, machID)
}

func (r *RegistryMux) Schedule() ([]job.ScheduledUnit, error) {
	return r.getRegistry().Schedule()
}

func (r *RegistryMux) ScheduledUnit(name string) (*job.ScheduledUnit, error) {
	return r.getRegistry().ScheduledUnit(name)
}

func (r *RegistryMux) Unit(name string) (*job.Unit, error) {
	return r.getRegistry().Unit(name)
}

func (r *RegistryMux) Units() ([]job.Unit, error) {
	return r.getRegistry().Units()
}

func (r *RegistryMux) UnitState(name string) (*unit.UnitState, error) {
	return r.getRegistry().UnitState(name)
}

func (r *RegistryMux) UnitStates() ([]*unit.UnitState, error) {
	return r.getRegistry().UnitStates()
}

func (r *RegistryMux) LatestDaemonVersion() (*semver.Version, error) {
	return r.etcdRegistry.LatestDaemonVersion()
}

func (r *RegistryMux) EngineVersion() (int, error) {
	return r.etcdRegistry.EngineVersion()
}

func (r *RegistryMux) UpdateEngineVersion(from int, to int) error {
	return r.etcdRegistry.UpdateEngineVersion(from, to)
}

func (r *RegistryMux) SetMachineMetadata(machID string, key string, value string) error {
	return r.etcdRegistry.SetMachineMetadata(machID, key, value)
}

func (r *RegistryMux) DeleteMachineMetadata(machID string, key string) error {
	return r.etcdRegistry.DeleteMachineMetadata(machID, key)
}
<commit_msg>registry/rpc: call etcdRegistry.MachineState() for every case<commit_after>// Copyright 2016 The fleet 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 rpc

import (
	"errors"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/coreos/go-semver/semver"

	"github.com/coreos/fleet/engine"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/log"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg/lease"
	"github.com/coreos/fleet/registry"
	"github.com/coreos/fleet/unit"
)

type RegistryMux struct {
	etcdRegistry    *registry.EtcdRegistry
	localMachine    machine.Machine
	rpcserver       *rpcserver
	currentRegistry registry.Registry
	rpcRegistry     *RPCRegistry
	currentEngine   machine.MachineState
	leaseManager    lease.Manager

	handlingEngineChange *sync.RWMutex
}

const (
	dialRegistryReconnectTimeout = 200 * time.Millisecond

	engineLeaderKeyPath = "engine-leader"
)

func NewRegistryMux(etcdRegistry *registry.EtcdRegistry, localMachine machine.Machine, leaseManager lease.Manager) *RegistryMux {
	return &RegistryMux{
		etcdRegistry:         etcdRegistry,
		localMachine:         localMachine,
		handlingEngineChange: new(sync.RWMutex),
		leaseManager:         leaseManager,
	}
}

// ConnectToRegistry allows to disable_engine fleet agents to adapt its Registry
// to fleet leader changes regardless of whether is etcd or gRPC based.
func (r *RegistryMux) ConnectToRegistry(e *engine.Engine) {
	for {
		// We have to check if the leader has changed to etcd otherwise keep grpc connection
		isGrpc, err := e.IsGrpcLeader()
		// If there is not error then we are able to get the leader state and continue
		// otherwise we have to wait
		if err == nil {
			if isGrpc {
				if r.rpcRegistry != nil && r.rpcRegistry.IsRegistryReady() {
					log.Infof("Reusing gRPC engine, connection is READY\n")
					r.currentRegistry = r.rpcRegistry
				} else {
					if r.rpcRegistry != nil {
						r.rpcRegistry.Close()
					}
					log.Infof("New engine supports gRPC, connecting\n")
					r.rpcRegistry = NewRPCRegistry(r.rpcDialerNoEngine)
					// connect to rpc registry
					r.rpcRegistry.Connect()
					r.currentRegistry = r.rpcRegistry
				}
			} else {
				if r.rpcRegistry != nil {
					r.rpcRegistry.Close()
				}
				// new leader is etcd-based
				r.currentRegistry = r.etcdRegistry
			}
		}
		time.Sleep(5 * time.Second)
	}
}

func (r *RegistryMux) rpcDialerNoEngine(_ string, timeout time.Duration) (net.Conn, error) {
	ticker := time.Tick(dialRegistryReconnectTimeout)
	// Timeout re-defined to call etcd every 5secs to get the leader
	timeout = 5 * time.Second
	check := time.After(timeout)

	for {
		select {
		case <-check:
			log.Errorf("Unable to connect to engine %s\n", r.currentEngine.PublicIP)
			// Get the new engine leader of the cluster out of etcd
			lease, err := r.leaseManager.GetLease(engineLeaderKeyPath)
			// Key found
			if err == nil && lease != nil {
				var err error
				machines, err := r.etcdRegistry.Machines()
				if err != nil {
					log.Errorf("Unable to get the machines of the cluster %v\n", err)
					return nil, errors.New("Unable to get the machines of the cluster")
				}
				for _, s := range machines {
					// Update the currentEngine with the new one... otherwise wait until
					// there is one
					if s.ID == lease.MachineID() {
						// New leader has not gRPC capabilities enabled.
						if !s.Capabilities.Has(machine.CapGRPC) {
							log.Error("New leader engine has not gRPC enabled!")
							return nil, errors.New("New leader engine has not gRPC enabled!")
						}
						r.currentEngine = s
						log.Infof("Found a new engine to connect to: %s\n", r.currentEngine.PublicIP)
						// Restore initial check configuration
						timeout = 5 * time.Second
						check = time.After(timeout)
					}
				}
			} else {
				timeout = 2 * time.Second
				log.Errorf("Unable to get the leader engine, retrying in %v...", timeout)
				check = time.After(timeout)
			}
		case <-ticker:
			addr := fmt.Sprintf("%s:%d", r.currentEngine.PublicIP, rpcServerPort)
			conn, err := net.Dial("tcp", addr)
			if err == nil {
				log.Infof("Connected to engine on %s\n", r.currentEngine.PublicIP)
				return conn, nil
			}
			log.Errorf("Retry to connect to new engine: %+v", err)
		}
	}
}

func (r *RegistryMux) rpcDialer(_ string, timeout time.Duration) (net.Conn, error) {
	ticker := time.Tick(dialRegistryReconnectTimeout)
	alert := time.After(timeout)

	for {
		select {
		case <-alert:
			log.Errorf("Unable to connect to engine %s\n", r.currentEngine.PublicIP)
			return nil, errors.New("Unable to connect to new engine, the client connection is closing")
		case <-ticker:
			addr := fmt.Sprintf("%s:%d", r.currentEngine.PublicIP, rpcServerPort)
			conn, err := net.Dial("tcp", addr)
			if err == nil {
				log.Infof("Connected to engine on %s\n", r.currentEngine.PublicIP)
				return conn, nil
			}
			log.Errorf("Retry to connect to new engine: %+v", err)
		}
	}
}

func (r *RegistryMux) EngineChanged(newEngine machine.MachineState) {
	r.handlingEngineChange.Lock()
	defer r.handlingEngineChange.Unlock()

	stopServer := false
	if r.currentEngine.ID != newEngine.ID {
		stopServer = true
	}
	r.currentEngine = newEngine
	log.Infof("Engine changed, checking capabilities %+v", newEngine)
	if r.localMachine.State().Capabilities.Has(machine.CapGRPC) {
		if r.rpcserver != nil && ((r.rpcRegistry != nil && !r.rpcRegistry.IsRegistryReady()) || stopServer) {
			// If the engine changed, we need to stop the rpc server
			r.rpcserver.Stop()
			r.rpcserver = nil
		}
		if newEngine.ID == r.localMachine.State().ID {
			if r.rpcserver == nil {
				// start rpc server
				log.Infof("Starting rpc server...\n")
				var err error
				r.rpcserver, err = NewRPCServer(r.etcdRegistry, newEngine.PublicIP)
				if err != nil {
					log.Fatalf("Unable to create rpc server %+v", err)
				}

				go func() {
					errc := make(chan error, 1)
					if errc <- r.rpcserver.Start(); <-errc != nil {
						log.Fatalf("Failed to serve gRPC requests on listener: %v", <-errc)
					}
				}()
			}

		}
		if newEngine.Capabilities.Has(machine.CapGRPC) {
			if r.rpcRegistry != nil && r.rpcRegistry.IsRegistryReady() {
				log.Infof("Reusing gRPC engine, connection is READY\n")
				r.currentRegistry = r.rpcRegistry
			} else {
				log.Infof("New engine supports gRPC, connecting\n")
				r.rpcRegistry = NewRPCRegistry(r.rpcDialer)
				// connect to rpc registry
				r.rpcRegistry.Connect()
				r.currentRegistry = r.rpcRegistry
			}
		} else {
			log.Infof("Falling back to etcd registry\n")
			if r.rpcserver != nil {
				// If the engine changed to a non gRPC leader, we need to stop the server
				r.rpcserver.Stop()
			}
			r.currentRegistry = r.etcdRegistry
		}

	} else {
		log.Infof("Falling back to etcd registry\n")
		r.currentRegistry = r.etcdRegistry
	}
}

func (r *RegistryMux) getRegistry() registry.Registry {
	r.handlingEngineChange.RLock()
	defer r.handlingEngineChange.RUnlock()
	if r.currentRegistry == nil {
		return r.etcdRegistry
	}
	return r.currentRegistry
}

func (r *RegistryMux) IsRegistryReady() bool {
	return r.getRegistry().IsRegistryReady()
}

func (r *RegistryMux) UseEtcdRegistry() bool {
	return r.getRegistry().UseEtcdRegistry()
}

func (r *RegistryMux) ClearUnitHeartbeat(name string) {
	r.getRegistry().ClearUnitHeartbeat(name)
}

func (r *RegistryMux) CreateUnit(unit *job.Unit) error {
	return r.getRegistry().CreateUnit(unit)
}

func (r *RegistryMux) CreateMachineState(ms machine.MachineState, ttl time.Duration) (uint64, error) {
	return r.etcdRegistry.CreateMachineState(ms, ttl)
}

func (r *RegistryMux) DestroyUnit(unit string) error {
	return r.getRegistry().DestroyUnit(unit)
}

func (r *RegistryMux) UnitHeartbeat(name string, machID string, ttl time.Duration) error {
	return r.getRegistry().UnitHeartbeat(name, machID, ttl)
}

func (r *RegistryMux) Machines() ([]machine.MachineState, error) {
	return r.etcdRegistry.Machines()
}

func (r *RegistryMux) RemoveMachineState(machID string) error {
	return r.etcdRegistry.RemoveMachineState(machID)
}

func (r *RegistryMux) RemoveUnitState(jobName string) error {
	return r.getRegistry().RemoveUnitState(jobName)
}

func (r *RegistryMux) SaveUnitState(jobName string, unitState *unit.UnitState, ttl time.Duration) {
	r.getRegistry().SaveUnitState(jobName, unitState, ttl)
}

func (r *RegistryMux) ScheduleUnit(name string, machID string) error {
	return r.getRegistry().ScheduleUnit(name, machID)
}

func (r *RegistryMux) SetUnitTargetState(name string, state job.JobState) error {
	return r.getRegistry().SetUnitTargetState(name, state)
}

func (r *RegistryMux) MachineState(machID string) (machine.MachineState, error) {
	return r.etcdRegistry.MachineState(machID)
}

func (r *RegistryMux) SetMachineState(ms machine.MachineState, ttl time.Duration) (uint64, error) {
	return r.etcdRegistry.SetMachineState(ms, ttl)
}

func (r *RegistryMux) UnscheduleUnit(name string, machID string) error {
	return r.getRegistry().UnscheduleUnit(name, machID)
}

func (r *RegistryMux) Schedule() ([]job.ScheduledUnit, error) {
	return r.getRegistry().Schedule()
}

func (r *RegistryMux) ScheduledUnit(name string) (*job.ScheduledUnit, error) {
	return r.getRegistry().ScheduledUnit(name)
}

func (r *RegistryMux) Unit(name string) (*job.Unit, error) {
	return r.getRegistry().Unit(name)
}

func (r *RegistryMux) Units() ([]job.Unit, error) {
	return r.getRegistry().Units()
}

func (r *RegistryMux) UnitState(name string) (*unit.UnitState, error) {
	return r.getRegistry().UnitState(name)
}

func (r *RegistryMux) UnitStates() ([]*unit.UnitState, error) {
	return r.getRegistry().UnitStates()
}

func (r *RegistryMux) LatestDaemonVersion() (*semver.Version, error) {
	return r.etcdRegistry.LatestDaemonVersion()
}

func (r *RegistryMux) EngineVersion() (int, error) {
	return r.etcdRegistry.EngineVersion()
}

func (r *RegistryMux) UpdateEngineVersion(from int, to int) error {
	return r.etcdRegistry.UpdateEngineVersion(from, to)
}

func (r *RegistryMux) SetMachineMetadata(machID string, key string, value string) error {
	return r.etcdRegistry.SetMachineMetadata(machID, key, value)
}

func (r *RegistryMux) DeleteMachineMetadata(machID string, key string) error {
	return r.etcdRegistry.DeleteMachineMetadata(machID, key)
}
<|endoftext|>