<commit_before>// Copyright 2013 Google Inc.  All rights reserved.
// Copyright 2016 the gousb Authors.  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 usb

// #include <libusb.h>
import "C"

import (
	"fmt"
	"reflect"
	"sync"
	"time"
	"unsafe"
)

var DefaultReadTimeout = 1 * time.Second
var DefaultWriteTimeout = 1 * time.Second
var DefaultControlTimeout = 250 * time.Millisecond //5 * time.Second

type Device struct {
	handle *C.libusb_device_handle

	// Embed the device information for easy access
	*Descriptor

	// Timeouts
	ReadTimeout    time.Duration
	WriteTimeout   time.Duration
	ControlTimeout time.Duration

	// Claimed interfaces
	lock    *sync.Mutex
	claimed map[uint8]int
}

func newDevice(handle *C.libusb_device_handle, desc *Descriptor) *Device {
	ifaces := 0
	d := &Device{
		handle:         handle,
		Descriptor:     desc,
		ReadTimeout:    DefaultReadTimeout,
		WriteTimeout:   DefaultWriteTimeout,
		ControlTimeout: DefaultControlTimeout,
		lock:           new(sync.Mutex),
		claimed:        make(map[uint8]int, ifaces),
	}

	return d
}

func (d *Device) Reset() error {
	if errno := C.libusb_reset_device(d.handle); errno != 0 {
		return usbError(errno)
	}
	return nil
}

func (d *Device) Control(rType, request uint8, val, idx uint16, data []byte) (int, error) {
	//log.Printf("control xfer: %d:%d/%d:%d %x", idx, rType, request, val, string(data))
	dataSlice := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	n := C.libusb_control_transfer(
		d.handle,
		C.uint8_t(rType),
		C.uint8_t(request),
		C.uint16_t(val),
		C.uint16_t(idx),
		(*C.uchar)(unsafe.Pointer(dataSlice.Data)),
		C.uint16_t(len(data)),
		C.uint(d.ControlTimeout/time.Millisecond))
	if n < 0 {
		return int(n), usbError(n)
	}
	return int(n), nil
}

// ActiveConfig returns the config id (not the index) of the active configuration.
// This corresponds to the ConfigInfo.Config field.
func (d *Device) ActiveConfig() (uint8, error) {
	var cfg C.int
	if errno := C.libusb_get_configuration(d.handle, &cfg); errno < 0 {
		return 0, usbError(errno)
	}
	return uint8(cfg), nil
}

// SetConfig attempts to change the active configuration.
// The cfg provided is the config id (not the index) of the configuration to set,
// which corresponds to the ConfigInfo.Config field.
func (d *Device) SetConfig(cfg uint8) error {
	if errno := C.libusb_set_configuration(d.handle, C.int(cfg)); errno < 0 {
		return usbError(errno)
	}
	return nil
}

// Close the device.
func (d *Device) Close() error {
	if d.handle == nil {
		return fmt.Errorf("usb: double close on device")
	}
	d.lock.Lock()
	defer d.lock.Unlock()
	for iface := range d.claimed {
		C.libusb_release_interface(d.handle, C.int(iface))
	}
	C.libusb_close(d.handle)
	d.handle = nil
	return nil
}

func (d *Device) OpenEndpoint(conf, iface, setup, epoint uint8) (Endpoint, error) {
	end := &endpoint{
		Device: d,
	}

	var setAlternate bool
	for _, c := range d.Configs {
		if c.Config != conf {
			continue
		}
		debug.Printf("found conf: %#v\n", c)
		for _, i := range c.Interfaces {
			if i.Number != iface {
				continue
			}
			debug.Printf("found iface: %#v\n", i)
			for i, s := range i.Setups {
				if s.Alternate != setup {
					continue
				}
				setAlternate = i != 0

				debug.Printf("found setup: %#v [default: %v]\n", s, !setAlternate)
				for _, e := range s.Endpoints {
					debug.Printf("ep %02x search: %#v\n", epoint, s)
					if e.Address != epoint {
						continue
					}
					end.InterfaceSetup = s
					end.EndpointInfo = e
					goto found
				}
				return nil, fmt.Errorf("usb: unknown endpoint %02x", epoint)
			}
			return nil, fmt.Errorf("usb: unknown setup %02x", setup)
		}
		return nil, fmt.Errorf("usb: unknown interface %02x", iface)
	}
	return nil, fmt.Errorf("usb: unknown configuration %02x", conf)

found:

	// Set the configuration
	var activeConf C.int
	if errno := C.libusb_get_configuration(d.handle, &activeConf); errno < 0 {
		return nil, fmt.Errorf("usb: getcfg: %s", usbError(errno))
	}
	if int(activeConf) != int(conf) {
		if errno := C.libusb_set_configuration(d.handle, C.int(conf)); errno < 0 {
			return nil, fmt.Errorf("usb: setcfg: %s", usbError(errno))
		}
	}

	// Claim the interface
	if errno := C.libusb_claim_interface(d.handle, C.int(iface)); errno < 0 {
		return nil, fmt.Errorf("usb: claim: %s", usbError(errno))
	}

	// Increment the claim count
	d.lock.Lock()
	d.claimed[iface]++
	d.lock.Unlock() // unlock immediately because the next calls may block

	// Choose the alternate
	if setAlternate {
		if errno := C.libusb_set_interface_alt_setting(d.handle, C.int(iface), C.int(setup)); errno < 0 {
			debug.Printf("altsetting error: %s", usbError(errno))
			return nil, fmt.Errorf("usb: setalt: %s", usbError(errno))
		}
	}

	return end, nil
}

func (d *Device) GetStringDescriptor(desc_index int) (string, error) {

	// allocate 200-byte array limited the length of string descriptor
	goBuffer := make([]byte, 200)

	// get string descriptor from libusb. if errno < 0 then there are any errors.
	// if errno >= 0; it is a length of result string descriptor
	errno := C.libusb_get_string_descriptor_ascii(
		d.handle,
		C.uint8_t(desc_index),
		(*C.uchar)(unsafe.Pointer(&goBuffer[0])),
		200)

	// if any errors occur
	if errno < 0 {
		return "", fmt.Errorf("usb: getstr: %s", usbError(errno))
	}
	// convert slice of byte to string with limited length from errno
	stringDescriptor := string(goBuffer[:errno])

	return stringDescriptor, nil
}

// SetAutoDetach Enable/disable libusb's automatic kernel driver detachment.
// When this is enabled libusb will automatically detach the kernel driver
// on an interface when claiming the interface, and attach it when releasing the interface.
// Automatic kernel driver detachment is disabled on newly opened device handles by default.
func (d *Device) SetAutoDetach(autodetach bool) error {
	autodetachInt := 0
	if autodetach {
		autodetachInt = 1
	}

	errno := C.libusb_set_auto_detach_kernel_driver(
		d.handle,
		C.int(autodetach),
	)

	// TODO LIBUSB_ERROR_NOT_SUPPORTED (-12) handling
	// if any errors occur
	if errno < 0 {
		return fmt.Errorf("usb: setautodetach: %s", usbError(errno))
	}
	return nil
}
<commit_msg>improve comment<commit_after>// Copyright 2013 Google Inc.  All rights reserved.
// Copyright 2016 the gousb Authors.  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 usb

// #include <libusb.h>
import "C"

import (
	"fmt"
	"reflect"
	"sync"
	"time"
	"unsafe"
)

var DefaultReadTimeout = 1 * time.Second
var DefaultWriteTimeout = 1 * time.Second
var DefaultControlTimeout = 250 * time.Millisecond //5 * time.Second

type Device struct {
	handle *C.libusb_device_handle

	// Embed the device information for easy access
	*Descriptor

	// Timeouts
	ReadTimeout    time.Duration
	WriteTimeout   time.Duration
	ControlTimeout time.Duration

	// Claimed interfaces
	lock    *sync.Mutex
	claimed map[uint8]int
}

func newDevice(handle *C.libusb_device_handle, desc *Descriptor) *Device {
	ifaces := 0
	d := &Device{
		handle:         handle,
		Descriptor:     desc,
		ReadTimeout:    DefaultReadTimeout,
		WriteTimeout:   DefaultWriteTimeout,
		ControlTimeout: DefaultControlTimeout,
		lock:           new(sync.Mutex),
		claimed:        make(map[uint8]int, ifaces),
	}

	return d
}

func (d *Device) Reset() error {
	if errno := C.libusb_reset_device(d.handle); errno != 0 {
		return usbError(errno)
	}
	return nil
}

func (d *Device) Control(rType, request uint8, val, idx uint16, data []byte) (int, error) {
	//log.Printf("control xfer: %d:%d/%d:%d %x", idx, rType, request, val, string(data))
	dataSlice := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	n := C.libusb_control_transfer(
		d.handle,
		C.uint8_t(rType),
		C.uint8_t(request),
		C.uint16_t(val),
		C.uint16_t(idx),
		(*C.uchar)(unsafe.Pointer(dataSlice.Data)),
		C.uint16_t(len(data)),
		C.uint(d.ControlTimeout/time.Millisecond))
	if n < 0 {
		return int(n), usbError(n)
	}
	return int(n), nil
}

// ActiveConfig returns the config id (not the index) of the active configuration.
// This corresponds to the ConfigInfo.Config field.
func (d *Device) ActiveConfig() (uint8, error) {
	var cfg C.int
	if errno := C.libusb_get_configuration(d.handle, &cfg); errno < 0 {
		return 0, usbError(errno)
	}
	return uint8(cfg), nil
}

// SetConfig attempts to change the active configuration.
// The cfg provided is the config id (not the index) of the configuration to set,
// which corresponds to the ConfigInfo.Config field.
func (d *Device) SetConfig(cfg uint8) error {
	if errno := C.libusb_set_configuration(d.handle, C.int(cfg)); errno < 0 {
		return usbError(errno)
	}
	return nil
}

// Close the device.
func (d *Device) Close() error {
	if d.handle == nil {
		return fmt.Errorf("usb: double close on device")
	}
	d.lock.Lock()
	defer d.lock.Unlock()
	for iface := range d.claimed {
		C.libusb_release_interface(d.handle, C.int(iface))
	}
	C.libusb_close(d.handle)
	d.handle = nil
	return nil
}

func (d *Device) OpenEndpoint(conf, iface, setup, epoint uint8) (Endpoint, error) {
	end := &endpoint{
		Device: d,
	}

	var setAlternate bool
	for _, c := range d.Configs {
		if c.Config != conf {
			continue
		}
		debug.Printf("found conf: %#v\n", c)
		for _, i := range c.Interfaces {
			if i.Number != iface {
				continue
			}
			debug.Printf("found iface: %#v\n", i)
			for i, s := range i.Setups {
				if s.Alternate != setup {
					continue
				}
				setAlternate = i != 0

				debug.Printf("found setup: %#v [default: %v]\n", s, !setAlternate)
				for _, e := range s.Endpoints {
					debug.Printf("ep %02x search: %#v\n", epoint, s)
					if e.Address != epoint {
						continue
					}
					end.InterfaceSetup = s
					end.EndpointInfo = e
					goto found
				}
				return nil, fmt.Errorf("usb: unknown endpoint %02x", epoint)
			}
			return nil, fmt.Errorf("usb: unknown setup %02x", setup)
		}
		return nil, fmt.Errorf("usb: unknown interface %02x", iface)
	}
	return nil, fmt.Errorf("usb: unknown configuration %02x", conf)

found:

	// Set the configuration
	var activeConf C.int
	if errno := C.libusb_get_configuration(d.handle, &activeConf); errno < 0 {
		return nil, fmt.Errorf("usb: getcfg: %s", usbError(errno))
	}
	if int(activeConf) != int(conf) {
		if errno := C.libusb_set_configuration(d.handle, C.int(conf)); errno < 0 {
			return nil, fmt.Errorf("usb: setcfg: %s", usbError(errno))
		}
	}

	// Claim the interface
	if errno := C.libusb_claim_interface(d.handle, C.int(iface)); errno < 0 {
		return nil, fmt.Errorf("usb: claim: %s", usbError(errno))
	}

	// Increment the claim count
	d.lock.Lock()
	d.claimed[iface]++
	d.lock.Unlock() // unlock immediately because the next calls may block

	// Choose the alternate
	if setAlternate {
		if errno := C.libusb_set_interface_alt_setting(d.handle, C.int(iface), C.int(setup)); errno < 0 {
			debug.Printf("altsetting error: %s", usbError(errno))
			return nil, fmt.Errorf("usb: setalt: %s", usbError(errno))
		}
	}

	return end, nil
}

func (d *Device) GetStringDescriptor(desc_index int) (string, error) {

	// allocate 200-byte array limited the length of string descriptor
	goBuffer := make([]byte, 200)

	// get string descriptor from libusb. if errno < 0 then there are any errors.
	// if errno >= 0; it is a length of result string descriptor
	errno := C.libusb_get_string_descriptor_ascii(
		d.handle,
		C.uint8_t(desc_index),
		(*C.uchar)(unsafe.Pointer(&goBuffer[0])),
		200)

	// if any errors occur
	if errno < 0 {
		return "", fmt.Errorf("usb: getstr: %s", usbError(errno))
	}
	// convert slice of byte to string with limited length from errno
	stringDescriptor := string(goBuffer[:errno])

	return stringDescriptor, nil
}

// SetAutoDetach Enable/disable libusb's automatic kernel driver detachment.
// When autodetach is enabled libusb will automatically detach the kernel driver
// on the interface and reattach it when releasing the interface.
// Automatic kernel driver detachment is disabled on newly opened device handles by default.
func (d *Device) SetAutoDetach(autodetach bool) error {
	autodetachInt := 0
	if autodetach {
		autodetachInt = 1
	}

	errno := C.libusb_set_auto_detach_kernel_driver(
		d.handle,
		C.int(autodetach),
	)

	// TODO LIBUSB_ERROR_NOT_SUPPORTED (-12) handling
	// if any errors occur
	if errno < 0 {
		return fmt.Errorf("usb: setautodetach: %s", usbError(errno))
	}
	return nil
}
<|endoftext|>


<commit_before>package redisc

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/PuerkitoBio/juggler/internal/redistest"
	"github.com/garyburd/redigo/redis"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestClusterRefreshNormalServer(t *testing.T) {
	cmd, port := redistest.StartServer(t, nil, "")
	defer cmd.Process.Kill()

	c := &Cluster{
		StartupNodes: []string{":" + port},
	}
	err := c.Refresh()
	if assert.Error(t, err, "Refresh") {
		assert.Contains(t, err.Error(), "redisc: all nodes failed", "expected error message")
	}
}

func TestClusterRefresh(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
	}

	err := c.Refresh()
	if assert.NoError(t, err, "Refresh") {
		var prev string
		pix := -1
		for ix, master := range c.mapping {
			if master != prev || ix == len(c.mapping)-1 {
				prev = master
				t.Logf("%5d: %s\n", ix, master)
				pix++
			}
			if assert.NotEmpty(t, master) {
				split := strings.Index(master, ":")
				assert.Contains(t, ports, master[split:], "expected master")
			}
		}
	}
}

func TestClusterClose(t *testing.T) {
	c := &Cluster{
		StartupNodes: []string{":6379"},
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	assert.NoError(t, c.Close(), "Close")
	if err := c.Close(); assert.Error(t, err, "Close after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if conn := c.Get(); assert.Error(t, conn.Err(), "Get after Close") {
		assert.Contains(t, conn.Err().Error(), "redisc: closed", "expected message")
	}
	if _, err := c.Dial(); assert.Error(t, err, "Dial after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if err := c.Refresh(); assert.Error(t, err, "Refresh after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
}

func createPool(addr string, opts ...redis.DialOption) (*redis.Pool, error) {
	return &redis.Pool{
		MaxIdle:     5,
		MaxActive:   10,
		IdleTimeout: time.Minute,
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", addr, opts...)
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}, nil
}

type redisCmd struct {
	name   string
	args   redis.Args
	resp   interface{} // if resp is of type lenResult, asserts that there is a result at least this long
	errMsg string
}

type lenResult int

func TestCommands(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	cmdsPerGroup := map[string][]redisCmd{
		"cluster": {
			{"CLUSTER", redis.Args{"INFO"}, lenResult(10), ""},
			{"READONLY", nil, "OK", ""},
			{"READWRITE", nil, "OK", ""},
			{"CLUSTER", redis.Args{"COUNTKEYSINSLOT", 12345}, int64(0), ""},
			{"CLUSTER", redis.Args{"KEYSLOT", "a"}, int64(15495), ""},
			{"CLUSTER", redis.Args{"NODES"}, lenResult(100), ""},
		},
		"connection": {
			{"AUTH", redis.Args{"pwd"}, nil, "ERR Client sent AUTH, but no password is set"},
			{"ECHO", redis.Args{"a"}, []byte("a"), ""},
			{"PING", nil, "PONG", ""},
			{"SELECT", redis.Args{1}, nil, "ERR SELECT is not allowed in cluster mode"},
			{"QUIT", nil, "OK", ""},
		},
		"hashes": {
			{"HSET", redis.Args{"ha", "f1", "1"}, int64(1), ""},
			{"HLEN", redis.Args{"ha"}, int64(1), ""},
			{"HEXISTS", redis.Args{"ha", "f1"}, int64(1), ""},
			{"HDEL", redis.Args{"ha", "f1", "f2"}, int64(1), ""},
			{"HINCRBY", redis.Args{"hb", "f1", "1"}, int64(1), ""},
			{"HINCRBYFLOAT", redis.Args{"hb", "f2", "0.5"}, []byte("0.5"), ""},
			{"HKEYS", redis.Args{"hb"}, []interface{}{[]byte("f1"), []byte("f2")}, ""},
			{"HMGET", redis.Args{"hb", "f1", "f2"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HMSET", redis.Args{"hc", "f1", "a", "f2", "b"}, "OK", ""},
			{"HSET", redis.Args{"ha", "f1", "2"}, int64(1), ""},
			{"HGET", redis.Args{"ha", "f1"}, []byte("2"), ""},
			{"HGETALL", redis.Args{"ha"}, []interface{}{[]byte("f1"), []byte("2")}, ""},
			{"HSETNX", redis.Args{"ha", "f2", "3"}, int64(1), ""},
			//{"HSTRLEN", redis.Args{"hb", "f2"}, int64(3), ""}, // redis 3.2 only
			{"HVALS", redis.Args{"hb"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HSCAN", redis.Args{"hb", 0}, lenResult(2), ""},
		},
		"hyperloglog": {
			{"PFADD", redis.Args{"hll", "a", "b", "c"}, int64(1), ""},
			{"PFCOUNT", redis.Args{"hll"}, int64(3), ""},
			{"PFADD", redis.Args{"hll2", "d"}, int64(1), ""},
			{"PFMERGE", redis.Args{"hll", "hll2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
		},
		"keys": {
			{"SET", redis.Args{"k1", "z"}, "OK", ""},
			{"EXISTS", redis.Args{"k1"}, int64(1), ""},
			{"DUMP", redis.Args{"k1"}, lenResult(10), ""},
			{"EXPIRE", redis.Args{"k1", 10}, int64(1), ""},
			{"EXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).Unix()}, int64(1), ""},
			{"KEYS", redis.Args{"z*"}, []interface{}{}, ""}, // KEYS is supported, but uses a random node and returns keys from that node (undeterministic)
			{"MOVE", redis.Args{"k1", 2}, nil, "ERR MOVE is not allowed in cluster mode"},
			{"PERSIST", redis.Args{"k1"}, int64(1), ""},
			{"PEXPIRE", redis.Args{"k1", 10000}, int64(1), ""},
			{"PEXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).UnixNano() / int64(time.Millisecond)}, int64(1), ""},
			{"PTTL", redis.Args{"k1"}, lenResult(3500000), ""},
			// RANDOMKEY is not deterministic
			{"RENAME", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"RENAMENX", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"SCAN", redis.Args{0}, lenResult(2), ""}, // works, but only for the keys on that random node
			{"TTL", redis.Args{"k1"}, lenResult(3000), ""},
			{"TYPE", redis.Args{"k1"}, "string", ""},
			{"DEL", redis.Args{"k1"}, int64(1), ""},
		},
	}

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	require.NoError(t, c.Refresh(), "Refresh")

	var wg sync.WaitGroup
	wg.Add(len(cmdsPerGroup))
	for _, cmds := range cmdsPerGroup {
		go func(cmds []redisCmd) {
			defer wg.Done()
			runCommands(t, c, cmds)
		}(cmds)
	}
	wg.Wait()
}

func runCommands(t *testing.T, c *Cluster, cmds []redisCmd) {
	for _, cmd := range cmds {
		conn := c.Get()
		res, err := conn.Do(cmd.name, cmd.args...)
		if cmd.errMsg != "" {
			if assert.Error(t, err, cmd.name) {
				assert.Contains(t, err.Error(), cmd.errMsg, cmd.name)
			}
		} else {
			assert.NoError(t, err, cmd.name)
			if lr, ok := cmd.resp.(lenResult); ok {
				switch res := res.(type) {
				case []byte:
					assert.True(t, len(res) >= int(lr), "result has at least %d bytes, has %d", lr, len(res))
				case []interface{}:
					assert.Equal(t, int(lr), len(res), "result array has %d items, has %d", lr, len(res))
				case int64:
					assert.True(t, res >= int64(lr), "result is at least %d, is %d", lr, res)
				default:
					t.Errorf("unexpected result type %T", res)
				}
			} else {
				if !assert.Equal(t, cmd.resp, res, cmd.name) {
					t.Logf("%T vs %T", cmd.resp, res)
				}
			}
		}
		require.NoError(t, conn.Close(), "Close")
	}
}
<commit_msg>redisc: add commands tests for lists<commit_after>package redisc

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/PuerkitoBio/juggler/internal/redistest"
	"github.com/garyburd/redigo/redis"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestClusterRefreshNormalServer(t *testing.T) {
	cmd, port := redistest.StartServer(t, nil, "")
	defer cmd.Process.Kill()

	c := &Cluster{
		StartupNodes: []string{":" + port},
	}
	err := c.Refresh()
	if assert.Error(t, err, "Refresh") {
		assert.Contains(t, err.Error(), "redisc: all nodes failed", "expected error message")
	}
}

func TestClusterRefresh(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
	}

	err := c.Refresh()
	if assert.NoError(t, err, "Refresh") {
		var prev string
		pix := -1
		for ix, master := range c.mapping {
			if master != prev || ix == len(c.mapping)-1 {
				prev = master
				t.Logf("%5d: %s\n", ix, master)
				pix++
			}
			if assert.NotEmpty(t, master) {
				split := strings.Index(master, ":")
				assert.Contains(t, ports, master[split:], "expected master")
			}
		}
	}
}

func TestClusterClose(t *testing.T) {
	c := &Cluster{
		StartupNodes: []string{":6379"},
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	assert.NoError(t, c.Close(), "Close")
	if err := c.Close(); assert.Error(t, err, "Close after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if conn := c.Get(); assert.Error(t, conn.Err(), "Get after Close") {
		assert.Contains(t, conn.Err().Error(), "redisc: closed", "expected message")
	}
	if _, err := c.Dial(); assert.Error(t, err, "Dial after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if err := c.Refresh(); assert.Error(t, err, "Refresh after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
}

func createPool(addr string, opts ...redis.DialOption) (*redis.Pool, error) {
	return &redis.Pool{
		MaxIdle:     5,
		MaxActive:   10,
		IdleTimeout: time.Minute,
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", addr, opts...)
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}, nil
}

type redisCmd struct {
	name   string
	args   redis.Args
	resp   interface{} // if resp is of type lenResult, asserts that there is a result at least this long
	errMsg string
}

type lenResult int

func TestCommands(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	cmdsPerGroup := map[string][]redisCmd{
		"cluster": {
			{"CLUSTER", redis.Args{"INFO"}, lenResult(10), ""},
			{"READONLY", nil, "OK", ""},
			{"READWRITE", nil, "OK", ""},
			{"CLUSTER", redis.Args{"COUNTKEYSINSLOT", 12345}, int64(0), ""},
			{"CLUSTER", redis.Args{"KEYSLOT", "a"}, int64(15495), ""},
			{"CLUSTER", redis.Args{"NODES"}, lenResult(100), ""},
		},
		"connection": {
			{"AUTH", redis.Args{"pwd"}, nil, "ERR Client sent AUTH, but no password is set"},
			{"ECHO", redis.Args{"a"}, []byte("a"), ""},
			{"PING", nil, "PONG", ""},
			{"SELECT", redis.Args{1}, nil, "ERR SELECT is not allowed in cluster mode"},
			{"QUIT", nil, "OK", ""},
		},
		"hashes": {
			{"HSET", redis.Args{"ha", "f1", "1"}, int64(1), ""},
			{"HLEN", redis.Args{"ha"}, int64(1), ""},
			{"HEXISTS", redis.Args{"ha", "f1"}, int64(1), ""},
			{"HDEL", redis.Args{"ha", "f1", "f2"}, int64(1), ""},
			{"HINCRBY", redis.Args{"hb", "f1", "1"}, int64(1), ""},
			{"HINCRBYFLOAT", redis.Args{"hb", "f2", "0.5"}, []byte("0.5"), ""},
			{"HKEYS", redis.Args{"hb"}, []interface{}{[]byte("f1"), []byte("f2")}, ""},
			{"HMGET", redis.Args{"hb", "f1", "f2"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HMSET", redis.Args{"hc", "f1", "a", "f2", "b"}, "OK", ""},
			{"HSET", redis.Args{"ha", "f1", "2"}, int64(1), ""},
			{"HGET", redis.Args{"ha", "f1"}, []byte("2"), ""},
			{"HGETALL", redis.Args{"ha"}, []interface{}{[]byte("f1"), []byte("2")}, ""},
			{"HSETNX", redis.Args{"ha", "f2", "3"}, int64(1), ""},
			//{"HSTRLEN", redis.Args{"hb", "f2"}, int64(3), ""}, // redis 3.2 only
			{"HVALS", redis.Args{"hb"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HSCAN", redis.Args{"hb", 0}, lenResult(2), ""},
		},
		"hyperloglog": {
			{"PFADD", redis.Args{"hll", "a", "b", "c"}, int64(1), ""},
			{"PFCOUNT", redis.Args{"hll"}, int64(3), ""},
			{"PFADD", redis.Args{"hll2", "d"}, int64(1), ""},
			{"PFMERGE", redis.Args{"hll", "hll2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
		},
		"keys": {
			{"SET", redis.Args{"k1", "z"}, "OK", ""},
			{"EXISTS", redis.Args{"k1"}, int64(1), ""},
			{"DUMP", redis.Args{"k1"}, lenResult(10), ""},
			{"EXPIRE", redis.Args{"k1", 10}, int64(1), ""},
			{"EXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).Unix()}, int64(1), ""},
			{"KEYS", redis.Args{"z*"}, []interface{}{}, ""}, // KEYS is supported, but uses a random node and returns keys from that node (undeterministic)
			{"MOVE", redis.Args{"k1", 2}, nil, "ERR MOVE is not allowed in cluster mode"},
			{"PERSIST", redis.Args{"k1"}, int64(1), ""},
			{"PEXPIRE", redis.Args{"k1", 10000}, int64(1), ""},
			{"PEXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).UnixNano() / int64(time.Millisecond)}, int64(1), ""},
			{"PTTL", redis.Args{"k1"}, lenResult(3500000), ""},
			// RANDOMKEY is not deterministic
			{"RENAME", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"RENAMENX", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"SCAN", redis.Args{0}, lenResult(2), ""}, // works, but only for the keys on that random node
			{"TTL", redis.Args{"k1"}, lenResult(3000), ""},
			{"TYPE", redis.Args{"k1"}, "string", ""},
			{"DEL", redis.Args{"k1"}, int64(1), ""},
		},
		"lists": {
			{"LPUSH", redis.Args{"l1", "a", "b", "c"}, int64(3), ""},
			{"LINDEX", redis.Args{"l1", 1}, []byte("b"), ""},
			{"LINSERT", redis.Args{"l1", "BEFORE", "b", "d"}, int64(4), ""},
			{"LLEN", redis.Args{"l1"}, int64(4), ""},
			{"LPOP", redis.Args{"l1"}, []byte("c"), ""},
			{"LPUSHX", redis.Args{"l1", "e"}, int64(4), ""},
			{"LRANGE", redis.Args{"l1", 0, 1}, []interface{}{[]byte("e"), []byte("d")}, ""},
			{"LREM", redis.Args{"l1", 0, "d"}, int64(1), ""},
			{"LSET", redis.Args{"l1", 0, "f"}, "OK", ""},
			{"LTRIM", redis.Args{"l1", 0, 3}, "OK", ""},
			{"RPOP", redis.Args{"l1"}, []byte("a"), ""},
			{"RPOPLPUSH", redis.Args{"l1", "l2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"RPUSH", redis.Args{"l1", "g"}, int64(3), ""},
			{"RPUSH", redis.Args{"l1", "h"}, int64(4), ""},
			{"BLPOP", redis.Args{"l1", 1}, lenResult(2), ""},
			{"BRPOP", redis.Args{"l1", 1}, lenResult(2), ""},
			{"BRPOPLPUSH", redis.Args{"l1", "l2", 1}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
		},
	}

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	require.NoError(t, c.Refresh(), "Refresh")

	var wg sync.WaitGroup
	wg.Add(len(cmdsPerGroup))
	for _, cmds := range cmdsPerGroup {
		go func(cmds []redisCmd) {
			defer wg.Done()
			runCommands(t, c, cmds)
		}(cmds)
	}
	wg.Wait()
}

func runCommands(t *testing.T, c *Cluster, cmds []redisCmd) {
	for _, cmd := range cmds {
		conn := c.Get()
		res, err := conn.Do(cmd.name, cmd.args...)
		if cmd.errMsg != "" {
			if assert.Error(t, err, cmd.name) {
				assert.Contains(t, err.Error(), cmd.errMsg, cmd.name)
			}
		} else {
			assert.NoError(t, err, cmd.name)
			if lr, ok := cmd.resp.(lenResult); ok {
				switch res := res.(type) {
				case []byte:
					assert.True(t, len(res) >= int(lr), "result has at least %d bytes, has %d", lr, len(res))
				case []interface{}:
					assert.Equal(t, int(lr), len(res), "result array has %d items, has %d", lr, len(res))
				case int64:
					assert.True(t, res >= int64(lr), "result is at least %d, is %d", lr, res)
				default:
					t.Errorf("unexpected result type %T", res)
				}
			} else {
				if !assert.Equal(t, cmd.resp, res, cmd.name) {
					t.Logf("%T vs %T", cmd.resp, res)
				}
			}
		}
		require.NoError(t, conn.Close(), "Close")
	}
}
<|endoftext|>


<commit_before>package cluster

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"runtime"
	"sort"
	"testing"
	"time"

	"github.com/Shopify/sarama"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("PartitionSlice", func() {

	It("should sort correctly", func() {
		p1 := Partition{Addr: "host1:9093", ID: 1}
		p2 := Partition{Addr: "host1:9092", ID: 2}
		p3 := Partition{Addr: "host2:9092", ID: 3}
		p4 := Partition{Addr: "host3:9091", ID: 4}
		p5 := Partition{Addr: "host2:9093", ID: 5}
		p6 := Partition{Addr: "host1:9092", ID: 6}

		slice := PartitionSlice{p1, p2, p3, p4, p5, p6}
		sort.Sort(slice)
		Expect(slice).To(BeEquivalentTo(PartitionSlice{p2, p6, p1, p3, p5, p4}))
	})

})

/*********************************************************************
 * TEST HOOK
 *********************************************************************/

const (
	t_KAFKA_VERSION = "kafka_2.10-0.8.1.1"
	t_CLIENT        = "sarama-cluster-client"
	t_TOPIC         = "sarama-cluster-topic"
	t_GROUP         = "sarama-cluster-group"
	t_DIR           = "/tmp/sarama-cluster-test"
)

var _ = BeforeSuite(func() {
	runner := testDir(t_KAFKA_VERSION, "bin", "kafka-run-class.sh")
	testState.zookeeper = exec.Command(runner, "-name", "zookeeper", "org.apache.zookeeper.server.ZooKeeperServerMain", testDir("zookeeper.properties"))
	testState.kafka = exec.Command(runner, "-name", "kafkaServer", "kafka.Kafka", testDir("server.properties"))
	testState.kafka.Env = []string{"KAFKA_HEAP_OPTS=-Xmx1G -Xms1G"}

	// Create Dir
	Expect(os.MkdirAll(t_DIR, 0775)).NotTo(HaveOccurred())

	// Start ZK
	Expect(testState.zookeeper.Start()).NotTo(HaveOccurred())
	Eventually(func() *os.Process {
		return testState.zookeeper.Process
	}).ShouldNot(BeNil())

	// Start Kafka
	Expect(testState.kafka.Start()).NotTo(HaveOccurred())
	Eventually(func() *os.Process {
		return testState.kafka.Process
	}).ShouldNot(BeNil())
	time.Sleep(3 * time.Second)

	// Create and wait for client
	client, err := newClient()
	Expect(err).NotTo(HaveOccurred())
	defer client.Close()

	Eventually(func() error {
		_, err := client.Partitions(t_TOPIC)
		return err
	}).ShouldNot(HaveOccurred(), "10s")

	// Seed messages
	Expect(seedMessages(client, 10000)).NotTo(HaveOccurred())
})

var _ = AfterSuite(func() {
	if testState.kafka != nil {
		testState.kafka.Process.Kill()
	}
	if testState.zookeeper != nil {
		testState.zookeeper.Process.Kill()
	}
	Expect(os.RemoveAll(t_DIR)).NotTo(HaveOccurred())
})

func TestSuite(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "sarama/cluster")
}

/*******************************************************************
 * TEST HELPERS
 *******************************************************************/

var testState struct{ kafka, zookeeper *exec.Cmd }

func newClient() (*sarama.Client, error) {
	return sarama.NewClient(t_CLIENT, []string{"127.0.0.1:29092"}, sarama.NewClientConfig())
}

func testDir(tokens ...string) string {
	_, filename, _, _ := runtime.Caller(1)
	tokens = append([]string{path.Dir(filename), "test"}, tokens...)
	return path.Join(tokens...)
}

func seedMessages(client *sarama.Client, count int) error {
	producer, err := sarama.NewSimpleProducer(client, t_TOPIC, nil)
	if err != nil {
		return err
	}
	defer producer.Close()

	for i := 0; i < count; i++ {
		kv := sarama.StringEncoder(fmt.Sprintf("PLAINDATA-%08d", i))
		err := producer.SendMessage(kv, kv)
		if err != nil {
			return err
		}
	}
	return nil
}

type mockNotifier struct{ messages []string }

func (n *mockNotifier) RebalanceStart(c *Consumer) {
	n.messages = append(n.messages, "REBALANCE START")
}
func (n *mockNotifier) RebalanceOK(c *Consumer) {
	n.messages = append(n.messages, "REBALANCE OK")
}
func (n *mockNotifier) RebalanceError(c *Consumer, err error) {
	n.messages = append(n.messages, "REBALANCE ERROR")
}
func (n *mockNotifier) CommitError(c *Consumer, err error) {
	n.messages = append(n.messages, "COMMIT ERROR")
}
<commit_msg>Wait a little longer for Kafka to start<commit_after>package cluster

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"runtime"
	"sort"
	"testing"

	"github.com/Shopify/sarama"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("PartitionSlice", func() {

	It("should sort correctly", func() {
		p1 := Partition{Addr: "host1:9093", ID: 1}
		p2 := Partition{Addr: "host1:9092", ID: 2}
		p3 := Partition{Addr: "host2:9092", ID: 3}
		p4 := Partition{Addr: "host3:9091", ID: 4}
		p5 := Partition{Addr: "host2:9093", ID: 5}
		p6 := Partition{Addr: "host1:9092", ID: 6}

		slice := PartitionSlice{p1, p2, p3, p4, p5, p6}
		sort.Sort(slice)
		Expect(slice).To(BeEquivalentTo(PartitionSlice{p2, p6, p1, p3, p5, p4}))
	})

})

/*********************************************************************
 * TEST HOOK
 *********************************************************************/

const (
	t_KAFKA_VERSION = "kafka_2.10-0.8.1.1"
	t_CLIENT        = "sarama-cluster-client"
	t_TOPIC         = "sarama-cluster-topic"
	t_GROUP         = "sarama-cluster-group"
	t_DIR           = "/tmp/sarama-cluster-test"
)

var _ = BeforeSuite(func() {
	runner := testDir(t_KAFKA_VERSION, "bin", "kafka-run-class.sh")
	testState.zookeeper = exec.Command(runner, "-name", "zookeeper", "org.apache.zookeeper.server.ZooKeeperServerMain", testDir("zookeeper.properties"))
	testState.kafka = exec.Command(runner, "-name", "kafkaServer", "kafka.Kafka", testDir("server.properties"))
	testState.kafka.Env = []string{"KAFKA_HEAP_OPTS=-Xmx1G -Xms1G"}

	// Create Dir
	Expect(os.MkdirAll(t_DIR, 0775)).NotTo(HaveOccurred())

	// Start ZK & Kafka
	Expect(testState.zookeeper.Start()).NotTo(HaveOccurred())
	Expect(testState.kafka.Start()).NotTo(HaveOccurred())

	// Wait for client
	var client *sarama.Client
	Eventually(func() error {
		var err error
		client, err = newClient()
		return err
	}, "10s", "1s").ShouldNot(HaveOccurred())
	defer client.Close()

	Eventually(func() error {
		_, err := client.Partitions(t_TOPIC)
		return err
	}, "10s", "1s").ShouldNot(HaveOccurred())

	// Seed messages
	Expect(seedMessages(client, 10000)).NotTo(HaveOccurred())
})

var _ = AfterSuite(func() {
	if testState.kafka != nil {
		testState.kafka.Process.Kill()
	}
	if testState.zookeeper != nil {
		testState.zookeeper.Process.Kill()
	}
	Expect(os.RemoveAll(t_DIR)).NotTo(HaveOccurred())
})

func TestSuite(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "sarama/cluster")
}

/*******************************************************************
 * TEST HELPERS
 *******************************************************************/

var testState struct{ kafka, zookeeper *exec.Cmd }

func newClient() (*sarama.Client, error) {
	return sarama.NewClient(t_CLIENT, []string{"127.0.0.1:29092"}, nil)
}

func testDir(tokens ...string) string {
	_, filename, _, _ := runtime.Caller(1)
	tokens = append([]string{path.Dir(filename), "test"}, tokens...)
	return path.Join(tokens...)
}

func seedMessages(client *sarama.Client, count int) error {
	producer, err := sarama.NewSimpleProducer(client, t_TOPIC, nil)
	if err != nil {
		return err
	}
	defer producer.Close()

	for i := 0; i < count; i++ {
		kv := sarama.StringEncoder(fmt.Sprintf("PLAINDATA-%08d", i))
		err := producer.SendMessage(kv, kv)
		if err != nil {
			return err
		}
	}
	return nil
}

type mockNotifier struct{ messages []string }

func (n *mockNotifier) RebalanceStart(c *Consumer) {
	n.messages = append(n.messages, "REBALANCE START")
}
func (n *mockNotifier) RebalanceOK(c *Consumer) {
	n.messages = append(n.messages, "REBALANCE OK")
}
func (n *mockNotifier) RebalanceError(c *Consumer, err error) {
	n.messages = append(n.messages, "REBALANCE ERROR")
}
func (n *mockNotifier) CommitError(c *Consumer, err error) {
	n.messages = append(n.messages, "COMMIT ERROR")
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"net/url"
	"strings"

	homedir "github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
	"github.com/zerowidth/gh-shorthand/alfred"
	"github.com/zerowidth/gh-shorthand/config"
	"github.com/zerowidth/gh-shorthand/parser"
)

func init() {
	RootCmd.AddCommand(completeCommand)
}

var completeCommand = &cobra.Command{
	Use:   "complete ['input string']",
	Short: "Completion mode",
	Long: `Parse the given input and generate matching Alfred items.

Parses an input string as directly provided by an Alfred script filter innput.
It expects a leading space for the default mode (that is, "space optional" in
the script filter), and uses the first character of the input as a mode string:

  ' ' is default completion mode, for opening repos and issues.
  'i' is issue listing or search
  'n' is new issue in a repo`,
	Run: func(cmd *cobra.Command, args []string) {
		var input string
		var items = []*alfred.Item{}

		if len(args) == 0 {
			input = ""
		} else {
			input = strings.Join(args, " ")
		}

		path, _ := homedir.Expand("~/.gh-shorthand.yml")
		cfg, err := config.LoadFromFile(path)
		if err != nil {
			items = []*alfred.Item{errorItem("when loading ~/.gh-shorthand.yml", err.Error())}
		} else {
			items = completeItems(cfg, input)
		}

		printItems(items)
	},
}

func completeItems(cfg *config.Config, input string) []*alfred.Item {
	items := []*alfred.Item{}
	fullInput := input

	if len(input) == 0 {
		return items
	}

	// input includes leading space or leading mode char followed by a space
	var mode string
	if len(input) > 1 && input[0:1] != " " {
		mode = input[0:1]
		input = input[2:]
	} else if len(input) > 0 && input[0:1] == " " {
		mode = " "
		input = input[1:]
	}

	result := parser.Parse(cfg.RepoMap, input)
	usedDefault := false

	if len(cfg.DefaultRepo) > 0 && len(result.Repo) == 0 && len(result.Path) == 0 &&
		((mode == "i" || mode == "n") || len(result.Query) == 0) {
		result.Repo = cfg.DefaultRepo
		usedDefault = true
	}

	switch mode {
	case " ": // open repo, issue, and/or path
		// repo required, no query allowed
		if len(result.Repo) > 0 && len(result.Query) == 0 {
			items = append(items, openRepoItems(result, usedDefault)...)
		}

		if len(result.Repo) == 0 && len(result.Path) > 0 {
			items = append(items, openPathItem(result.Path))
		}

		if len(input) > 0 && !strings.Contains(input, " ") {
			items = append(items,
				autocompleteItems(cfg, input, result,
					autocompleteOpenItem, openEndedOpenItem)...)
		}
	case "i":
		// repo required, no issue or path, query allowed
		if len(result.Repo) > 0 && len(result.Issue) == 0 && len(result.Path) == 0 {
			items = append(items, openIssueItems(result, usedDefault, fullInput)...)
		}

		if len(input) > 0 && !strings.Contains(input, " ") {
			items = append(items,
				autocompleteItems(cfg, input, result,
					autocompleteIssueItem, openEndedIssueItem)...)
		}
	case "n":
		// repo required, no issue or path, query allowed
		if len(result.Repo) > 0 && len(result.Issue) == 0 && len(result.Path) == 0 {
			items = append(items, newIssueItems(result, usedDefault)...)
		}

		if len(input) > 0 && !strings.Contains(input, " ") {
			items = append(items,
				autocompleteItems(cfg, input, result,
					autocompleteNewIssueItem, openEndedNewIssueItem)...)
		}
	}

	return items
}

func openRepoItems(result *parser.Result, usedDefault bool) (items []*alfred.Item) {
	uid := "gh:" + result.Repo
	title := "Open " + result.Repo
	arg := "open https://github.com/" + result.Repo
	icon := repoIcon

	if len(result.Issue) > 0 {
		uid += "#" + result.Issue
		title += "#" + result.Issue
		arg += "/issues/" + result.Issue
		icon = issueIcon
	}

	if len(result.Path) > 0 {
		uid += result.Path
		title += result.Path
		arg += result.Path
		icon = pathIcon
	}

	if len(result.Match) > 0 {
		title += " (" + result.Match
		if len(result.Issue) > 0 {
			title += "#" + result.Issue
		}
		title += ")"
	} else if usedDefault {
		title += " (default repo)"
	}

	items = append(items, &alfred.Item{
		UID:   uid,
		Title: title + " on GitHub",
		Arg:   arg,
		Valid: true,
		Icon:  icon,
	})
	return items
}

func openPathItem(path string) *alfred.Item {
	return &alfred.Item{
		UID:   "gh:" + path,
		Title: fmt.Sprintf("Open %s on GitHub", path),
		Arg:   "open https://github.com" + path,
		Valid: true,
		Icon:  pathIcon,
	}
}

func openIssueItems(result *parser.Result, usedDefault bool, fullInput string) (items []*alfred.Item) {
	extra := ""
	if len(result.Match) > 0 {
		extra += " (" + result.Match + ")"
	} else if usedDefault {
		extra += " (default repo)"
	}

	if len(result.Query) == 0 {
		items = append(items, &alfred.Item{
			UID:   "ghi:" + result.Repo,
			Title: "Open issues for " + result.Repo + extra,
			Arg:   "open https://github.com/" + result.Repo + "/issues",
			Valid: true,
			Icon:  issueListIcon,
		})
		items = append(items, &alfred.Item{
			Title:        "Search issues in " + result.Repo + extra + " for...",
			Valid:        false,
			Icon:         issueSearchIcon,
			Autocomplete: fullInput + " ",
		})
	} else {
		escaped := url.PathEscape(result.Query)
		arg := "open https://github.com/" + result.Repo + "/search?utf8=✓&type=Issues&q=" + escaped
		items = append(items, &alfred.Item{
			UID:   "ghis:" + result.Repo,
			Title: "Search issues in " + result.Repo + extra + " for " + result.Query,
			Arg:   arg,
			Valid: true,
			Icon:  issueSearchIcon,
		})
	}
	return
}

func newIssueItems(result *parser.Result, usedDefault bool) (items []*alfred.Item) {
	title := "New issue in " + result.Repo
	if len(result.Match) > 0 {
		title += " (" + result.Match + ")"
	} else if usedDefault {
		title += " (default repo)"
	}

	if len(result.Query) == 0 {
		items = append(items, &alfred.Item{
			UID:   "ghn:" + result.Repo,
			Title: title,
			Arg:   "open https://github.com/" + result.Repo + "/issues/new",
			Valid: true,
			Icon:  newIssueIcon,
		})
	} else {
		escaped := url.PathEscape(result.Query)
		arg := "open https://github.com/" + result.Repo + "/issues/new?title=" + escaped
		items = append(items, &alfred.Item{
			UID:   "ghn:" + result.Repo,
			Title: title + ": " + result.Query,
			Arg:   arg,
			Valid: true,
			Icon:  newIssueIcon,
		})
	}
	return
}

func autocompleteOpenItem(key, repo string) *alfred.Item {
	return &alfred.Item{
		UID:          "gh:" + repo,
		Title:        fmt.Sprintf("Open %s (%s) on GitHub", repo, key),
		Arg:          "open https://github.com/" + repo,
		Valid:        true,
		Autocomplete: " " + key,
		Icon:         repoIcon,
	}
}

func autocompleteIssueItem(key, repo string) *alfred.Item {
	return &alfred.Item{
		UID:          "ghi:" + repo,
		Title:        fmt.Sprintf("Open issues for %s (%s)", repo, key),
		Arg:          "open https://github.com/" + repo + "/issues",
		Valid:        true,
		Autocomplete: "i " + key,
		Icon:         issueListIcon,
	}
}

func autocompleteNewIssueItem(key, repo string) *alfred.Item {
	return &alfred.Item{
		UID:          "ghn:" + repo,
		Title:        fmt.Sprintf("New issue in %s (%s)", repo, key),
		Arg:          "open https://github.com/" + repo + "/issues/new",
		Valid:        true,
		Autocomplete: "n " + key,
		Icon:         newIssueIcon,
	}
}

func openEndedOpenItem(input string) *alfred.Item {
	return &alfred.Item{
		Title:        fmt.Sprintf("Open %s... on GitHub", input),
		Autocomplete: " " + input,
		Valid:        false,
	}
}

func openEndedIssueItem(input string) *alfred.Item {
	return &alfred.Item{
		Title:        fmt.Sprintf("Open issues for %s...", input),
		Autocomplete: "i " + input,
		Valid:        false,
		Icon:         issueListIcon,
	}
}

func openEndedNewIssueItem(input string) *alfred.Item {
	return &alfred.Item{
		Title:        fmt.Sprintf("New issue in %s...", input),
		Autocomplete: "n " + input,
		Valid:        false,
		Icon:         newIssueIcon,
	}
}

func autocompleteItems(cfg *config.Config, input string, result *parser.Result,
	autocompleteItem func(string, string) *alfred.Item,
	openEndedItem func(string) *alfred.Item) (items []*alfred.Item) {
	for key, repo := range cfg.RepoMap {
		if strings.HasPrefix(key, input) && key != result.Match && repo != result.Repo {
			items = append(items, autocompleteItem(key, repo))
		}
	}

	if len(input) > 0 && result.Repo != input {
		items = append(items, openEndedItem(input))
	}
	return
}
<commit_msg>Add basic e/o/t character modes to main completion command<commit_after>package cmd

import (
	"fmt"
	"net/url"
	"strings"

	homedir "github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
	"github.com/zerowidth/gh-shorthand/alfred"
	"github.com/zerowidth/gh-shorthand/config"
	"github.com/zerowidth/gh-shorthand/parser"
)

func init() {
	RootCmd.AddCommand(completeCommand)
}

var completeCommand = &cobra.Command{
	Use:   "complete ['input string']",
	Short: "Completion mode",
	Long: `Parse the given input and generate matching Alfred items.

Parses an input string as directly provided by an Alfred script filter innput.
It expects a leading space for the default mode (that is, "space optional" in
the script filter), and uses the first character of the input as a mode string:

  ' ' is default completion mode, for opening repos and issues.
  'i' is issue listing or search
  'n' is new issue in a repo`,
	Run: func(cmd *cobra.Command, args []string) {
		var input string
		var items = []*alfred.Item{}

		if len(args) == 0 {
			input = ""
		} else {
			input = strings.Join(args, " ")
		}

		path, _ := homedir.Expand("~/.gh-shorthand.yml")
		cfg, err := config.LoadFromFile(path)
		if err != nil {
			items = []*alfred.Item{errorItem("when loading ~/.gh-shorthand.yml", err.Error())}
		} else {
			items = completeItems(cfg, input)
		}

		printItems(items)
	},
}

func completeItems(cfg *config.Config, input string) []*alfred.Item {
	items := []*alfred.Item{}
	fullInput := input

	if len(input) == 0 {
		return items
	}

	// input includes leading space or leading mode char followed by a space
	var mode string
	if len(input) > 1 && input[0:1] != " " {
		mode = input[0:1]
		input = input[2:]
	} else if len(input) > 0 && input[0:1] == " " {
		mode = " "
		input = input[1:]
	}

	result := parser.Parse(cfg.RepoMap, input)
	usedDefault := false

	if len(cfg.DefaultRepo) > 0 && len(result.Repo) == 0 && len(result.Path) == 0 &&
		((mode == "i" || mode == "n") || len(result.Query) == 0) {
		result.Repo = cfg.DefaultRepo
		usedDefault = true
	}

	switch mode {
	case " ": // open repo, issue, and/or path
		// repo required, no query allowed
		if len(result.Repo) > 0 && len(result.Query) == 0 {
			items = append(items, openRepoItems(result, usedDefault)...)
		}

		if len(result.Repo) == 0 && len(result.Path) > 0 {
			items = append(items, openPathItem(result.Path))
		}

		if len(input) > 0 && !strings.Contains(input, " ") {
			items = append(items,
				autocompleteItems(cfg, input, result,
					autocompleteOpenItem, openEndedOpenItem)...)
		}
	case "i":
		// repo required, no issue or path, query allowed
		if len(result.Repo) > 0 && len(result.Issue) == 0 && len(result.Path) == 0 {
			items = append(items, openIssueItems(result, usedDefault, fullInput)...)
		}

		if len(input) > 0 && !strings.Contains(input, " ") {
			items = append(items,
				autocompleteItems(cfg, input, result,
					autocompleteIssueItem, openEndedIssueItem)...)
		}
	case "n":
		// repo required, no issue or path, query allowed
		if len(result.Repo) > 0 && len(result.Issue) == 0 && len(result.Path) == 0 {
			items = append(items, newIssueItems(result, usedDefault)...)
		}

		if len(input) > 0 && !strings.Contains(input, " ") {
			items = append(items,
				autocompleteItems(cfg, input, result,
					autocompleteNewIssueItem, openEndedNewIssueItem)...)
		}
	case "e":
		items = append(items, actionItems(cfg.ProjectDirMap(), "ghe", "edit", "Edit", editorIcon)...)
	case "o":
		items = append(items, actionItems(cfg.ProjectDirMap(), "gho", "finder", "Open Finder in", editorIcon)...)
	case "t":
		items = append(items, actionItems(cfg.ProjectDirMap(), "ght", "term", "Open terminal in", editorIcon)...)
	}

	return items
}

func openRepoItems(result *parser.Result, usedDefault bool) (items []*alfred.Item) {
	uid := "gh:" + result.Repo
	title := "Open " + result.Repo
	arg := "open https://github.com/" + result.Repo
	icon := repoIcon

	if len(result.Issue) > 0 {
		uid += "#" + result.Issue
		title += "#" + result.Issue
		arg += "/issues/" + result.Issue
		icon = issueIcon
	}

	if len(result.Path) > 0 {
		uid += result.Path
		title += result.Path
		arg += result.Path
		icon = pathIcon
	}

	if len(result.Match) > 0 {
		title += " (" + result.Match
		if len(result.Issue) > 0 {
			title += "#" + result.Issue
		}
		title += ")"
	} else if usedDefault {
		title += " (default repo)"
	}

	items = append(items, &alfred.Item{
		UID:   uid,
		Title: title + " on GitHub",
		Arg:   arg,
		Valid: true,
		Icon:  icon,
	})
	return items
}

func openPathItem(path string) *alfred.Item {
	return &alfred.Item{
		UID:   "gh:" + path,
		Title: fmt.Sprintf("Open %s on GitHub", path),
		Arg:   "open https://github.com" + path,
		Valid: true,
		Icon:  pathIcon,
	}
}

func openIssueItems(result *parser.Result, usedDefault bool, fullInput string) (items []*alfred.Item) {
	extra := ""
	if len(result.Match) > 0 {
		extra += " (" + result.Match + ")"
	} else if usedDefault {
		extra += " (default repo)"
	}

	if len(result.Query) == 0 {
		items = append(items, &alfred.Item{
			UID:   "ghi:" + result.Repo,
			Title: "Open issues for " + result.Repo + extra,
			Arg:   "open https://github.com/" + result.Repo + "/issues",
			Valid: true,
			Icon:  issueListIcon,
		})
		items = append(items, &alfred.Item{
			Title:        "Search issues in " + result.Repo + extra + " for...",
			Valid:        false,
			Icon:         issueSearchIcon,
			Autocomplete: fullInput + " ",
		})
	} else {
		escaped := url.PathEscape(result.Query)
		arg := "open https://github.com/" + result.Repo + "/search?utf8=✓&type=Issues&q=" + escaped
		items = append(items, &alfred.Item{
			UID:   "ghis:" + result.Repo,
			Title: "Search issues in " + result.Repo + extra + " for " + result.Query,
			Arg:   arg,
			Valid: true,
			Icon:  issueSearchIcon,
		})
	}
	return
}

func newIssueItems(result *parser.Result, usedDefault bool) (items []*alfred.Item) {
	title := "New issue in " + result.Repo
	if len(result.Match) > 0 {
		title += " (" + result.Match + ")"
	} else if usedDefault {
		title += " (default repo)"
	}

	if len(result.Query) == 0 {
		items = append(items, &alfred.Item{
			UID:   "ghn:" + result.Repo,
			Title: title,
			Arg:   "open https://github.com/" + result.Repo + "/issues/new",
			Valid: true,
			Icon:  newIssueIcon,
		})
	} else {
		escaped := url.PathEscape(result.Query)
		arg := "open https://github.com/" + result.Repo + "/issues/new?title=" + escaped
		items = append(items, &alfred.Item{
			UID:   "ghn:" + result.Repo,
			Title: title + ": " + result.Query,
			Arg:   arg,
			Valid: true,
			Icon:  newIssueIcon,
		})
	}
	return
}

func autocompleteOpenItem(key, repo string) *alfred.Item {
	return &alfred.Item{
		UID:          "gh:" + repo,
		Title:        fmt.Sprintf("Open %s (%s) on GitHub", repo, key),
		Arg:          "open https://github.com/" + repo,
		Valid:        true,
		Autocomplete: " " + key,
		Icon:         repoIcon,
	}
}

func autocompleteIssueItem(key, repo string) *alfred.Item {
	return &alfred.Item{
		UID:          "ghi:" + repo,
		Title:        fmt.Sprintf("Open issues for %s (%s)", repo, key),
		Arg:          "open https://github.com/" + repo + "/issues",
		Valid:        true,
		Autocomplete: "i " + key,
		Icon:         issueListIcon,
	}
}

func autocompleteNewIssueItem(key, repo string) *alfred.Item {
	return &alfred.Item{
		UID:          "ghn:" + repo,
		Title:        fmt.Sprintf("New issue in %s (%s)", repo, key),
		Arg:          "open https://github.com/" + repo + "/issues/new",
		Valid:        true,
		Autocomplete: "n " + key,
		Icon:         newIssueIcon,
	}
}

func openEndedOpenItem(input string) *alfred.Item {
	return &alfred.Item{
		Title:        fmt.Sprintf("Open %s... on GitHub", input),
		Autocomplete: " " + input,
		Valid:        false,
	}
}

func openEndedIssueItem(input string) *alfred.Item {
	return &alfred.Item{
		Title:        fmt.Sprintf("Open issues for %s...", input),
		Autocomplete: "i " + input,
		Valid:        false,
		Icon:         issueListIcon,
	}
}

func openEndedNewIssueItem(input string) *alfred.Item {
	return &alfred.Item{
		Title:        fmt.Sprintf("New issue in %s...", input),
		Autocomplete: "n " + input,
		Valid:        false,
		Icon:         newIssueIcon,
	}
}

func autocompleteItems(cfg *config.Config, input string, result *parser.Result,
	autocompleteItem func(string, string) *alfred.Item,
	openEndedItem func(string) *alfred.Item) (items []*alfred.Item) {
	for key, repo := range cfg.RepoMap {
		if strings.HasPrefix(key, input) && key != result.Match && repo != result.Repo {
			items = append(items, autocompleteItem(key, repo))
		}
	}

	if len(input) > 0 && result.Repo != input {
		items = append(items, openEndedItem(input))
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"

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

	"github.com/derekparker/delve/service"
	"github.com/derekparker/delve/service/api"
	"github.com/derekparker/delve/service/rpc"
	"github.com/derekparker/delve/terminal"
	"github.com/spf13/cobra"
)

const version string = "0.6.0.beta"

var (
	Log      bool
	Headless bool
	Addr     string
)

func main() {
	// Main dlv root command.
	rootCommand := &cobra.Command{
		Use:   "dlv",
		Short: "Delve is a debugger for the Go programming language.",
		Long: `Delve is a source level debugger for Go programs.

Delve enables you to interact with your program by controlling the execution of the process,
evaluating variables, and providing information of thread / goroutine state, CPU register state and more.

The goal of this tool is to provide a simple yet powerful interface for debugging Go programs.
`,
	}
	rootCommand.PersistentFlags().StringVarP(&Addr, "listen", "l", "localhost:0", "Debugging server listen address.")
	rootCommand.PersistentFlags().BoolVarP(&Log, "log", "", false, "Enable debugging server logging.")
	rootCommand.PersistentFlags().BoolVarP(&Headless, "headless", "", false, "Run debug server only, in headless mode.")

	// 'version' subcommand.
	versionCommand := &cobra.Command{
		Use:   "version",
		Short: "Prints version.",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println("Delve version: " + version)
		},
	}
	rootCommand.AddCommand(versionCommand)

	// 'run' subcommand.
	runCommand := &cobra.Command{
		Use:   "run",
		Short: "Compile and begin debugging program.",
		Long: `Compiles your program with optimizations disabled,
starts and attaches to it, and enables you to immediately begin debugging your program.`,
		Run: func(cmd *cobra.Command, args []string) {
			status := func() int {
				const debugname = "debug"
				goBuild := exec.Command("go", "build", "-o", debugname, "-gcflags", "-N -l")
				goBuild.Stderr = os.Stderr
				err := goBuild.Run()
				if err != nil {
					return 1
				}
				fp, err := filepath.Abs("./" + debugname)
				if err != nil {
					fmt.Fprintf(os.Stderr, err.Error())
					return 1
				}
				defer os.Remove(fp)

				processArgs := append([]string{"./" + debugname}, args...)
				return execute(0, processArgs)
			}()
			os.Exit(status)
		},
	}
	rootCommand.AddCommand(runCommand)

	// 'trace' subcommand.
	var traceAttachPid int
	traceCommand := &cobra.Command{
		Use:   "trace [regexp]",
		Short: "Compile and begin tracing program.",
		Long:  "Trace program execution. Will set a tracepoint on every function matching [regexp] and output information when tracepoint is hit.",
		Run: func(cmd *cobra.Command, args []string) {
			status := func() int {
				const debugname = "debug"
				var processArgs []string
				if traceAttachPid == 0 {
					goBuild := exec.Command("go", "build", "-o", debugname, "-gcflags", "-N -l")
					goBuild.Stderr = os.Stderr
					err := goBuild.Run()
					if err != nil {
						return 1
					}
					fp, err := filepath.Abs("./" + debugname)
					if err != nil {
						fmt.Fprintf(os.Stderr, err.Error())
						return 1
					}
					defer os.Remove(fp)

					processArgs = append([]string{"./" + debugname}, args...)
				}
				// Make a TCP listener
				listener, err := net.Listen("tcp", Addr)
				if err != nil {
					fmt.Printf("couldn't start listener: %s\n", err)
					return 1
				}
				defer listener.Close()

				// Create and start a debugger server
				server := rpc.NewServer(&service.Config{
					Listener:    listener,
					ProcessArgs: processArgs,
					AttachPid:   traceAttachPid,
				}, Log)
				if err := server.Run(); err != nil {
					fmt.Fprintln(os.Stderr, err)
					return 1
				}
				sigChan := make(chan os.Signal)
				signal.Notify(sigChan, sys.SIGINT)
				client := rpc.NewClient(listener.Addr().String())
				funcs, err := client.ListFunctions(args[0])
				if err != nil {
					fmt.Fprintln(os.Stderr, err)
					return 1
				}
				for i := range funcs {
					_, err := client.CreateBreakpoint(&api.Breakpoint{FunctionName: funcs[i], Tracepoint: true})
					if err != nil {
						fmt.Fprintln(os.Stderr, err)
						return 1
					}
				}
				stateChan := client.Continue()
				for {
					select {
					case state := <-stateChan:
						if state.Err != nil {
							fmt.Fprintln(os.Stderr, state.Err)
							return 0
						}
						var args []string
						var fname string
						if state.CurrentThread != nil && state.CurrentThread.Function != nil {
							fname = state.CurrentThread.Function.Name
						}
						if state.BreakpointInfo != nil {
							for _, arg := range state.BreakpointInfo.Arguments {
								args = append(args, arg.Value)
							}
						}
						fmt.Printf("%s(%s) %s:%d\n", fname, strings.Join(args, ", "), state.CurrentThread.File, state.CurrentThread.Line)
					case <-sigChan:
						server.Stop(traceAttachPid == 0)
						return 1
					}
				}
				return 0
			}()
			os.Exit(status)
		},
	}
	traceCommand.Flags().IntVarP(&traceAttachPid, "pid", "p", 0, "Pid to attach to.")
	rootCommand.AddCommand(traceCommand)

	// 'test' subcommand.
	testCommand := &cobra.Command{
		Use:   "test",
		Short: "Compile test binary and begin debugging program.",
		Long: `Compiles a test binary with optimizations disabled,
starts and attaches to it, and enable you to immediately begin debugging your program.`,
		Run: func(cmd *cobra.Command, args []string) {
			status := func() int {
				wd, err := os.Getwd()
				if err != nil {
					fmt.Fprintf(os.Stderr, err.Error())
					return 1
				}
				base := filepath.Base(wd)
				goTest := exec.Command("go", "test", "-c", "-gcflags", "-N -l")
				goTest.Stderr = os.Stderr
				err = goTest.Run()
				if err != nil {
					return 1
				}
				debugname := "./" + base + ".test"
				defer os.Remove(debugname)
				processArgs := append([]string{debugname}, args...)

				return execute(0, processArgs)
			}()
			os.Exit(status)
		},
	}
	rootCommand.AddCommand(testCommand)

	// 'attach' subcommand.
	attachCommand := &cobra.Command{
		Use:   "attach [pid]",
		Short: "Attach to running process and begin debugging.",
		Long:  "Attach to running process and begin debugging.",
		Run: func(cmd *cobra.Command, args []string) {
			pid, err := strconv.Atoi(args[0])
			if err != nil {
				fmt.Fprintf(os.Stderr, "Invalid pid: %d", args[0])
				os.Exit(1)
			}
			os.Exit(execute(pid, nil))
		},
	}
	rootCommand.AddCommand(attachCommand)

	rootCommand.Execute()
}

func execute(attachPid int, processArgs []string) int {
	// Make a TCP listener
	listener, err := net.Listen("tcp", Addr)
	if err != nil {
		fmt.Printf("couldn't start listener: %s\n", err)
		return 1
	}
	defer listener.Close()

	// Create and start a debugger server
	server := rpc.NewServer(&service.Config{
		Listener:    listener,
		ProcessArgs: processArgs,
		AttachPid:   attachPid,
	}, Log)
	if err := server.Run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return 1
	}

	var status int
	if !Headless {
		// Create and start a terminal
		var client service.Client
		client = rpc.NewClient(listener.Addr().String())
		term := terminal.New(client)
		err, status = term.Run()
	} else {
		ch := make(chan os.Signal)
		signal.Notify(ch, sys.SIGINT)
		<-ch
		err = server.Stop(true)
	}

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

	return status
}
<commit_msg>cmd/dlv: fix bad format string<commit_after>package main

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"

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

	"github.com/derekparker/delve/service"
	"github.com/derekparker/delve/service/api"
	"github.com/derekparker/delve/service/rpc"
	"github.com/derekparker/delve/terminal"
	"github.com/spf13/cobra"
)

const version string = "0.6.0.beta"

var (
	Log      bool
	Headless bool
	Addr     string
)

func main() {
	// Main dlv root command.
	rootCommand := &cobra.Command{
		Use:   "dlv",
		Short: "Delve is a debugger for the Go programming language.",
		Long: `Delve is a source level debugger for Go programs.

Delve enables you to interact with your program by controlling the execution of the process,
evaluating variables, and providing information of thread / goroutine state, CPU register state and more.

The goal of this tool is to provide a simple yet powerful interface for debugging Go programs.
`,
	}
	rootCommand.PersistentFlags().StringVarP(&Addr, "listen", "l", "localhost:0", "Debugging server listen address.")
	rootCommand.PersistentFlags().BoolVarP(&Log, "log", "", false, "Enable debugging server logging.")
	rootCommand.PersistentFlags().BoolVarP(&Headless, "headless", "", false, "Run debug server only, in headless mode.")

	// 'version' subcommand.
	versionCommand := &cobra.Command{
		Use:   "version",
		Short: "Prints version.",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println("Delve version: " + version)
		},
	}
	rootCommand.AddCommand(versionCommand)

	// 'run' subcommand.
	runCommand := &cobra.Command{
		Use:   "run",
		Short: "Compile and begin debugging program.",
		Long: `Compiles your program with optimizations disabled,
starts and attaches to it, and enables you to immediately begin debugging your program.`,
		Run: func(cmd *cobra.Command, args []string) {
			status := func() int {
				const debugname = "debug"
				goBuild := exec.Command("go", "build", "-o", debugname, "-gcflags", "-N -l")
				goBuild.Stderr = os.Stderr
				err := goBuild.Run()
				if err != nil {
					return 1
				}
				fp, err := filepath.Abs("./" + debugname)
				if err != nil {
					fmt.Fprintf(os.Stderr, err.Error())
					return 1
				}
				defer os.Remove(fp)

				processArgs := append([]string{"./" + debugname}, args...)
				return execute(0, processArgs)
			}()
			os.Exit(status)
		},
	}
	rootCommand.AddCommand(runCommand)

	// 'trace' subcommand.
	var traceAttachPid int
	traceCommand := &cobra.Command{
		Use:   "trace [regexp]",
		Short: "Compile and begin tracing program.",
		Long:  "Trace program execution. Will set a tracepoint on every function matching [regexp] and output information when tracepoint is hit.",
		Run: func(cmd *cobra.Command, args []string) {
			status := func() int {
				const debugname = "debug"
				var processArgs []string
				if traceAttachPid == 0 {
					goBuild := exec.Command("go", "build", "-o", debugname, "-gcflags", "-N -l")
					goBuild.Stderr = os.Stderr
					err := goBuild.Run()
					if err != nil {
						return 1
					}
					fp, err := filepath.Abs("./" + debugname)
					if err != nil {
						fmt.Fprintf(os.Stderr, err.Error())
						return 1
					}
					defer os.Remove(fp)

					processArgs = append([]string{"./" + debugname}, args...)
				}
				// Make a TCP listener
				listener, err := net.Listen("tcp", Addr)
				if err != nil {
					fmt.Printf("couldn't start listener: %s\n", err)
					return 1
				}
				defer listener.Close()

				// Create and start a debugger server
				server := rpc.NewServer(&service.Config{
					Listener:    listener,
					ProcessArgs: processArgs,
					AttachPid:   traceAttachPid,
				}, Log)
				if err := server.Run(); err != nil {
					fmt.Fprintln(os.Stderr, err)
					return 1
				}
				sigChan := make(chan os.Signal)
				signal.Notify(sigChan, sys.SIGINT)
				client := rpc.NewClient(listener.Addr().String())
				funcs, err := client.ListFunctions(args[0])
				if err != nil {
					fmt.Fprintln(os.Stderr, err)
					return 1
				}
				for i := range funcs {
					_, err := client.CreateBreakpoint(&api.Breakpoint{FunctionName: funcs[i], Tracepoint: true})
					if err != nil {
						fmt.Fprintln(os.Stderr, err)
						return 1
					}
				}
				stateChan := client.Continue()
				for {
					select {
					case state := <-stateChan:
						if state.Err != nil {
							fmt.Fprintln(os.Stderr, state.Err)
							return 0
						}
						var args []string
						var fname string
						if state.CurrentThread != nil && state.CurrentThread.Function != nil {
							fname = state.CurrentThread.Function.Name
						}
						if state.BreakpointInfo != nil {
							for _, arg := range state.BreakpointInfo.Arguments {
								args = append(args, arg.Value)
							}
						}
						fmt.Printf("%s(%s) %s:%d\n", fname, strings.Join(args, ", "), state.CurrentThread.File, state.CurrentThread.Line)
					case <-sigChan:
						server.Stop(traceAttachPid == 0)
						return 1
					}
				}
				return 0
			}()
			os.Exit(status)
		},
	}
	traceCommand.Flags().IntVarP(&traceAttachPid, "pid", "p", 0, "Pid to attach to.")
	rootCommand.AddCommand(traceCommand)

	// 'test' subcommand.
	testCommand := &cobra.Command{
		Use:   "test",
		Short: "Compile test binary and begin debugging program.",
		Long: `Compiles a test binary with optimizations disabled,
starts and attaches to it, and enable you to immediately begin debugging your program.`,
		Run: func(cmd *cobra.Command, args []string) {
			status := func() int {
				wd, err := os.Getwd()
				if err != nil {
					fmt.Fprintf(os.Stderr, err.Error())
					return 1
				}
				base := filepath.Base(wd)
				goTest := exec.Command("go", "test", "-c", "-gcflags", "-N -l")
				goTest.Stderr = os.Stderr
				err = goTest.Run()
				if err != nil {
					return 1
				}
				debugname := "./" + base + ".test"
				defer os.Remove(debugname)
				processArgs := append([]string{debugname}, args...)

				return execute(0, processArgs)
			}()
			os.Exit(status)
		},
	}
	rootCommand.AddCommand(testCommand)

	// 'attach' subcommand.
	attachCommand := &cobra.Command{
		Use:   "attach [pid]",
		Short: "Attach to running process and begin debugging.",
		Long:  "Attach to running process and begin debugging.",
		Run: func(cmd *cobra.Command, args []string) {
			pid, err := strconv.Atoi(args[0])
			if err != nil {
				fmt.Fprintf(os.Stderr, "Invalid pid: %s\n", args[0])
				os.Exit(1)
			}
			os.Exit(execute(pid, nil))
		},
	}
	rootCommand.AddCommand(attachCommand)

	rootCommand.Execute()
}

func execute(attachPid int, processArgs []string) int {
	// Make a TCP listener
	listener, err := net.Listen("tcp", Addr)
	if err != nil {
		fmt.Printf("couldn't start listener: %s\n", err)
		return 1
	}
	defer listener.Close()

	// Create and start a debugger server
	server := rpc.NewServer(&service.Config{
		Listener:    listener,
		ProcessArgs: processArgs,
		AttachPid:   attachPid,
	}, Log)
	if err := server.Run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		return 1
	}

	var status int
	if !Headless {
		// Create and start a terminal
		var client service.Client
		client = rpc.NewClient(listener.Addr().String())
		term := terminal.New(client)
		err, status = term.Run()
	} else {
		ch := make(chan os.Signal)
		signal.Notify(ch, sys.SIGINT)
		<-ch
		err = server.Stop(true)
	}

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

	return status
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"os/exec"
)

// SSHCommand is responsible for launching a ssh shell on a given unit or machine.
type SSHCommand struct {
	SSHCommon
}

// SSHCommon provides common methods for SSHCommand and SCPCommand.
type SSHCommon struct {
	EnvCommandBase
	Target string
	Args   []string
	*juju.Conn
}

const sshDoc = `
Launch an ssh shell on the machine identified by the <service> parameter.
<service> can be either a machine id or a service name.  Any extra parameters
are treated as extra parameters for the ssh command.
`

func (c *SSHCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "ssh",
		Args:    "<service> [<ssh args>...]",
		Purpose: "launch an ssh shell on a given unit or machine",
		Doc:     sshDoc,
	}
}

func (c *SSHCommand) Init(args []string) error {
	if len(args) == 0 {
		return errors.New("no service name specified")
	}
	c.Target, c.Args = args[0], args[1:]
	return nil
}

// Run resolves c.Target to a machine, to the address of a i
// machine or unit forks ssh passing any arguments provided.
func (c *SSHCommand) Run(ctx *cmd.Context) error {
	var err error
	c.Conn, err = juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer c.Close()
	host, err := c.hostFromTarget(c.Target)
	if err != nil {
		return err
	}
	args := []string{"-l", "ubuntu", "-t", "-o", "StrictHostKeyChecking no", "-o", "PasswordAuthentication no", host}
	args = append(args, c.Args...)
	cmd := exec.Command("ssh", args...)
	cmd.Stdin = ctx.Stdin
	cmd.Stdout = ctx.Stdout
	cmd.Stderr = ctx.Stderr
	c.Close()
	return cmd.Run()
}

func (c *SSHCommon) hostFromTarget(target string) (string, error) {
	// is the target the id of a machine ?
	if state.IsMachineId(target) {
		log.Infof("cmd/juju: looking up address for machine %s...", target)
		// TODO(dfc) maybe we should have machine.PublicAddress() ?
		return c.machinePublicAddress(target)
	}
	// maybe the target is a unit ?
	if state.IsUnitName(target) {
		log.Infof("cmd/juju: Looking up address for unit %q...", c.Target)
		unit, err := c.State.Unit(target)
		if err != nil {
			return "", err
		}
		addr, ok := unit.PublicAddress()
		if !ok {
			return "", fmt.Errorf("unit %q has no public address", unit)
		}
		return addr, nil
	}
	return "", fmt.Errorf("unknown unit or machine %q", target)
}

func (c *SSHCommon) machinePublicAddress(id string) (string, error) {
	machine, err := c.State.Machine(id)
	if err != nil {
		return "", err
	}
	// wait for instance id
	w := machine.Watch()
	for _ = range w.Changes() {
		if instid, ok := machine.InstanceId(); ok {
			w.Stop()
			inst, err := c.Environ.Instances([]state.InstanceId{instid})
			if err != nil {
				return "", err
			}
			return inst[0].WaitDNSName()
		}
		// BUG(dfc) this does not refresh the machine, so
		// this loop will loop forever if it gets to this point.
		// https://bugs.launchpad.net/juju-core/+bug/1130051
	}
	// oops, watcher closed before we could get an answer
	return "", w.Stop()
}
<commit_msg>Fixed a log msg<commit_after>package main

import (
	"errors"
	"fmt"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"os/exec"
)

// SSHCommand is responsible for launching a ssh shell on a given unit or machine.
type SSHCommand struct {
	SSHCommon
}

// SSHCommon provides common methods for SSHCommand and SCPCommand.
type SSHCommon struct {
	EnvCommandBase
	Target string
	Args   []string
	*juju.Conn
}

const sshDoc = `
Launch an ssh shell on the machine identified by the <service> parameter.
<service> can be either a machine id or a service name.  Any extra parameters
are treated as extra parameters for the ssh command.
`

func (c *SSHCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "ssh",
		Args:    "<service> [<ssh args>...]",
		Purpose: "launch an ssh shell on a given unit or machine",
		Doc:     sshDoc,
	}
}

func (c *SSHCommand) Init(args []string) error {
	if len(args) == 0 {
		return errors.New("no service name specified")
	}
	c.Target, c.Args = args[0], args[1:]
	return nil
}

// Run resolves c.Target to a machine, to the address of a i
// machine or unit forks ssh passing any arguments provided.
func (c *SSHCommand) Run(ctx *cmd.Context) error {
	var err error
	c.Conn, err = juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer c.Close()
	host, err := c.hostFromTarget(c.Target)
	if err != nil {
		return err
	}
	args := []string{"-l", "ubuntu", "-t", "-o", "StrictHostKeyChecking no", "-o", "PasswordAuthentication no", host}
	args = append(args, c.Args...)
	cmd := exec.Command("ssh", args...)
	cmd.Stdin = ctx.Stdin
	cmd.Stdout = ctx.Stdout
	cmd.Stderr = ctx.Stderr
	c.Close()
	return cmd.Run()
}

func (c *SSHCommon) hostFromTarget(target string) (string, error) {
	// is the target the id of a machine ?
	if state.IsMachineId(target) {
		log.Infof("cmd/juju: looking up address for machine %s...", target)
		// TODO(dfc) maybe we should have machine.PublicAddress() ?
		return c.machinePublicAddress(target)
	}
	// maybe the target is a unit ?
	if state.IsUnitName(target) {
		log.Infof("cmd/juju: looking up address for unit %q...", c.Target)
		unit, err := c.State.Unit(target)
		if err != nil {
			return "", err
		}
		addr, ok := unit.PublicAddress()
		if !ok {
			return "", fmt.Errorf("unit %q has no public address", unit)
		}
		return addr, nil
	}
	return "", fmt.Errorf("unknown unit or machine %q", target)
}

func (c *SSHCommon) machinePublicAddress(id string) (string, error) {
	machine, err := c.State.Machine(id)
	if err != nil {
		return "", err
	}
	// wait for instance id
	w := machine.Watch()
	for _ = range w.Changes() {
		if instid, ok := machine.InstanceId(); ok {
			w.Stop()
			inst, err := c.Environ.Instances([]state.InstanceId{instid})
			if err != nil {
				return "", err
			}
			return inst[0].WaitDNSName()
		}
		// BUG(dfc) this does not refresh the machine, so
		// this loop will loop forever if it gets to this point.
		// https://bugs.launchpad.net/juju-core/+bug/1130051
	}
	// oops, watcher closed before we could get an answer
	return "", w.Stop()
}
<|endoftext|>


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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strings"
	"time"

	"github.com/couchbaselabs/blance"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
	"github.com/couchbaselabs/cbgt/cmd"
	"github.com/couchbaselabs/cbgt/rebalance"
)

var VERSION = "v0.0.0"

func main() {
	flag.Parse()

	if flags.Help {
		flag.Usage()
		os.Exit(2)
	}

	if flags.Version {
		fmt.Printf("%s main: %s, data: %s\n",
			path.Base(os.Args[0]), VERSION, cbgt.VERSION)
		os.Exit(0)
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	log.Printf("main: %s started (%s/%s)",
		os.Args[0], VERSION, cbgt.VERSION)

	rand.Seed(time.Now().UTC().UnixNano())

	go dumpOnSignalForPlatform()

	MainWelcome(flagAliases)

	nodesToRemove := []string(nil)
	if len(flags.RemoveNodes) > 0 {
		nodesToRemove = strings.Split(flags.RemoveNodes, ",")
	}

	bindHttp := "NO-BIND-HTTP"
	register := "unchanged"
	dataDir := "NO-DATA-DIR"

	// If cfg is down, we error, leaving it to some user-supplied
	// outside watchdog to backoff and restart/retry.
	cfg, err := cmd.MainCfg("mcp", flags.CfgConnect,
		bindHttp, register, dataDir)
	if err != nil {
		if err == cmd.ErrorBindHttp {
			log.Fatalf("%v", err)
			return
		}
		log.Fatalf("main: could not start cfg, cfgConnect: %s, err: %v\n"+
			"  Please check that your -cfg/-cfgConnect parameter (%q)\n"+
			"  is correct and/or that your configuration provider\n"+
			"  is available.",
			flags.CfgConnect, err, flags.CfgConnect)
		return
	}

	r, err := rebalance.StartRebalance(cbgt.VERSION, cfg, flags.Server,
		nodesToRemove,
		rebalance.RebalanceOptions{
			DryRun:  flags.DryRun,
			Verbose: flags.Verbose,
		})
	if err != nil {
		log.Fatalf("main: StartRebalance, err: %v", err)
		return
	}

	reportProgress(r)

	r.Stop()

	log.Printf("main: done")
}

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

type ProgressEntry struct {
	pindex, sourcePartition, node string // Immutable.

	stateOp     rebalance.StateOp
	initUUIDSeq cbgt.UUIDSeq
	currUUIDSeq cbgt.UUIDSeq
	wantUUIDSeq cbgt.UUIDSeq

	move int
	done bool
}

func reportProgress(r *rebalance.Rebalancer) {
	var lastEmit string

	maxNodeLen := 0
	maxPIndexLen := 0

	seenNodes := map[string]bool{}
	seenNodesSorted := []string(nil)

	// Map of pindex -> (source) partition -> node -> *ProgressEntry
	progressEntries := map[string]map[string]map[string]*ProgressEntry{}

	seenPIndexes := map[string]bool{}
	seenPIndexesSorted := []string(nil)

	updateProgressEntry := func(pindex, sourcePartition, node string,
		cb func(*ProgressEntry)) {
		if !seenNodes[node] {
			seenNodes[node] = true
			seenNodesSorted = append(seenNodesSorted, node)
			sort.Strings(seenNodesSorted)

			if maxNodeLen < len(node) {
				maxNodeLen = len(node)
			}
		}

		if maxPIndexLen < len(pindex) {
			maxPIndexLen = len(pindex)
		}

		sourcePartitions, exists := progressEntries[pindex]
		if !exists || sourcePartitions == nil {
			sourcePartitions = map[string]map[string]*ProgressEntry{}
			progressEntries[pindex] = sourcePartitions
		}

		nodes, exists := sourcePartitions[sourcePartition]
		if !exists || nodes == nil {
			nodes = map[string]*ProgressEntry{}
			sourcePartitions[sourcePartition] = nodes
		}

		progressEntry, exists := nodes[node]
		if !exists || progressEntry == nil {
			progressEntry = &ProgressEntry{
				pindex:          pindex,
				sourcePartition: sourcePartition,
				node:            node,
				move:            -1,
			}
			nodes[node] = progressEntry
		}

		cb(progressEntry)

		// TODO: Check UUID matches, too.

		if !seenPIndexes[pindex] {
			seenPIndexes[pindex] = true
			seenPIndexesSorted =
				append(seenPIndexesSorted, pindex)

			sort.Strings(seenPIndexesSorted)
		}
	}

	for progress := range r.ProgressCh() {
		if progress.Error != nil {
			r.Log("main: error, progress: %+v", progress)
			continue
		}

		updateProgressEntries(r, updateProgressEntry)

		currEmit := progressTable(maxNodeLen, maxPIndexLen,
			seenNodes,
			seenNodesSorted,
			seenPIndexes,
			seenPIndexesSorted,
			progressEntries)
		if currEmit != lastEmit {
			r.Log("%s", currEmit)
		}

		lastEmit = currEmit
	}
}

func updateProgressEntries(
	r *rebalance.Rebalancer,
	updateProgressEntry func(pindex, sourcePartition, node string,
		cb func(*ProgressEntry)),
) {
	r.Visit(func(
		currStates rebalance.CurrStates,
		currSeqs rebalance.CurrSeqs,
		wantSeqs rebalance.WantSeqs,
		mapNextMoves map[string]*blance.NextMoves,
	) {
		for _, pindexes := range currStates {
			for pindex, nodes := range pindexes {
				for node, stateOp := range nodes {
					updateProgressEntry(pindex, "", node,
						func(pe *ProgressEntry) {
							pe.stateOp = stateOp
						})
				}
			}
		}

		for pindex, sourcePartitions := range currSeqs {
			for sourcePartition, nodes := range sourcePartitions {
				for node, currUUIDSeq := range nodes {
					updateProgressEntry(pindex,
						sourcePartition, node,
						func(pe *ProgressEntry) {
							pe.currUUIDSeq = currUUIDSeq

							if pe.initUUIDSeq.UUID == "" {
								pe.initUUIDSeq = currUUIDSeq
							}
						})
				}
			}
		}

		for pindex, sourcePartitions := range wantSeqs {
			for sourcePartition, nodes := range sourcePartitions {
				for node, wantUUIDSeq := range nodes {
					updateProgressEntry(pindex,
						sourcePartition, node,
						func(pe *ProgressEntry) {
							pe.wantUUIDSeq = wantUUIDSeq
						})
				}
			}
		}

		for pindex, nextMoves := range mapNextMoves {
			for i, nodeStateOp := range nextMoves.Moves {
				updateProgressEntry(pindex, "", nodeStateOp.Node,
					func(pe *ProgressEntry) {
						pe.move = i
						pe.done = i < nextMoves.Next
					})
			}
		}
	})
}

func progressTable(maxNodeLen, maxPIndexLen int,
	seenNodes map[string]bool,
	seenNodesSorted []string,
	seenPIndexes map[string]bool,
	seenPIndexesSorted []string,
	progressEntries map[string]map[string]map[string]*ProgressEntry,
) string {
	var b bytes.Buffer

	written, _ := b.Write([]byte("%%%"))
	for i := written; i < maxPIndexLen; i++ {
		b.WriteByte(' ')
	}
	b.WriteByte(' ')

	for i, seenNode := range seenNodesSorted {
		if i > 0 {
			b.WriteByte(' ')
		}

		// TODO: Emit node human readable ADDR:PORT.
		b.Write([]byte(seenNode))
	}
	b.WriteByte('\n')

	for _, seenPIndex := range seenPIndexesSorted {
		b.Write([]byte(" %                  "))
		b.Write([]byte(seenPIndex))

		for _, seenNode := range seenNodesSorted {
			b.WriteByte(' ')

			sourcePartitions, exists :=
				progressEntries[seenPIndex]
			if !exists || sourcePartitions == nil {
				progressCell(&b, nil, nil, maxNodeLen)
				continue
			}

			nodes, exists := sourcePartitions[""]
			if !exists || nodes == nil {
				progressCell(&b, nil, nil, maxNodeLen)
				continue
			}

			pe, exists := nodes[seenNode]
			if !exists || pe == nil {
				progressCell(&b, nil, nil, maxNodeLen)
				continue
			}

			progressCell(&b, pe, sourcePartitions, maxNodeLen)
		}

		b.WriteByte('\n')
	}

	return b.String()
}

var opMap = map[string]string{
	"":        ".",
	"add":     "+",
	"del":     "-",
	"promote": "P",
	"demote":  "D",
}

func progressCell(b *bytes.Buffer,
	pe *ProgressEntry,
	sourcePartitions map[string]map[string]*ProgressEntry,
	maxNodeLen int) {
	written := 0

	totPct := 0.0 // To compute average pct.
	numPct := 0

	if pe != nil {
		written, _ = fmt.Fprintf(b, "%d ", pe.move)

		if sourcePartitions != nil {
			n, _ := b.Write([]byte(opMap[pe.stateOp.Op]))
			written = written + n

			for sourcePartition, nodes := range sourcePartitions {
				if sourcePartition == "" {
					continue
				}

				pex := nodes[pe.node]
				if pex == nil || pex.wantUUIDSeq.UUID == "" {
					continue
				}

				if pex.wantUUIDSeq.Seq <= pex.currUUIDSeq.Seq {
					totPct = totPct + 1.0
					numPct = numPct + 1
					continue
				}

				n := pex.currUUIDSeq.Seq - pex.initUUIDSeq.Seq
				d := pex.wantUUIDSeq.Seq - pex.initUUIDSeq.Seq
				if d > 0 {
					pct := float64(n) / float64(d)
					totPct = totPct + pct
					numPct = numPct + 1
				}
			}
		}
	} else {
		b.Write([]byte("  ."))
		written = 3
	}

	if numPct > 0 {
		avgPct := totPct / float64(numPct)

		n, _ := fmt.Fprintf(b, " %.1f%%", avgPct*100.0)
		written = written + n
	}

	for i := written; i < maxNodeLen; i++ {
		b.WriteByte(' ')
	}
}

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

func MainWelcome(flagAliases map[string][]string) {
	flag.VisitAll(func(f *flag.Flag) {
		if flagAliases[f.Name] != nil {
			log.Printf("  -%s=%q\n", f.Name, f.Value)
		}
	})
	log.Printf("  GOMAXPROCS=%d", runtime.GOMAXPROCS(-1))
}

func dumpOnSignal(signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, signals...)
	for _ = range c {
		log.Printf("dump: goroutine...")
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		log.Printf("dump: heap...")
		pprof.Lookup("heap").WriteTo(os.Stderr, 1)
	}
}
<commit_msg>renamed mcp writeProgressTable/writeProgressCell<commit_after>//  Copyright (c) 2015 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the
//  License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing,
//  software distributed under the License is distributed on an "AS
//  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
//  express or implied. See the License for the specific language
//  governing permissions and limitations under the License.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strings"
	"time"

	"github.com/couchbaselabs/blance"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
	"github.com/couchbaselabs/cbgt/cmd"
	"github.com/couchbaselabs/cbgt/rebalance"
)

var VERSION = "v0.0.0"

func main() {
	flag.Parse()

	if flags.Help {
		flag.Usage()
		os.Exit(2)
	}

	if flags.Version {
		fmt.Printf("%s main: %s, data: %s\n",
			path.Base(os.Args[0]), VERSION, cbgt.VERSION)
		os.Exit(0)
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	log.Printf("main: %s started (%s/%s)",
		os.Args[0], VERSION, cbgt.VERSION)

	rand.Seed(time.Now().UTC().UnixNano())

	go dumpOnSignalForPlatform()

	MainWelcome(flagAliases)

	nodesToRemove := []string(nil)
	if len(flags.RemoveNodes) > 0 {
		nodesToRemove = strings.Split(flags.RemoveNodes, ",")
	}

	bindHttp := "NO-BIND-HTTP"
	register := "unchanged"
	dataDir := "NO-DATA-DIR"

	// If cfg is down, we error, leaving it to some user-supplied
	// outside watchdog to backoff and restart/retry.
	cfg, err := cmd.MainCfg("mcp", flags.CfgConnect,
		bindHttp, register, dataDir)
	if err != nil {
		if err == cmd.ErrorBindHttp {
			log.Fatalf("%v", err)
			return
		}
		log.Fatalf("main: could not start cfg, cfgConnect: %s, err: %v\n"+
			"  Please check that your -cfg/-cfgConnect parameter (%q)\n"+
			"  is correct and/or that your configuration provider\n"+
			"  is available.",
			flags.CfgConnect, err, flags.CfgConnect)
		return
	}

	r, err := rebalance.StartRebalance(cbgt.VERSION, cfg, flags.Server,
		nodesToRemove,
		rebalance.RebalanceOptions{
			DryRun:  flags.DryRun,
			Verbose: flags.Verbose,
		})
	if err != nil {
		log.Fatalf("main: StartRebalance, err: %v", err)
		return
	}

	reportProgress(r)

	r.Stop()

	log.Printf("main: done")
}

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

type ProgressEntry struct {
	pindex, sourcePartition, node string // Immutable.

	stateOp     rebalance.StateOp
	initUUIDSeq cbgt.UUIDSeq
	currUUIDSeq cbgt.UUIDSeq
	wantUUIDSeq cbgt.UUIDSeq

	move int
	done bool
}

func reportProgress(r *rebalance.Rebalancer) {
	var lastEmit string

	maxNodeLen := 0
	maxPIndexLen := 0

	seenNodes := map[string]bool{}
	seenNodesSorted := []string(nil)

	// Map of pindex -> (source) partition -> node -> *ProgressEntry
	progressEntries := map[string]map[string]map[string]*ProgressEntry{}

	seenPIndexes := map[string]bool{}
	seenPIndexesSorted := []string(nil)

	updateProgressEntry := func(pindex, sourcePartition, node string,
		cb func(*ProgressEntry)) {
		if !seenNodes[node] {
			seenNodes[node] = true
			seenNodesSorted = append(seenNodesSorted, node)
			sort.Strings(seenNodesSorted)

			if maxNodeLen < len(node) {
				maxNodeLen = len(node)
			}
		}

		if maxPIndexLen < len(pindex) {
			maxPIndexLen = len(pindex)
		}

		sourcePartitions, exists := progressEntries[pindex]
		if !exists || sourcePartitions == nil {
			sourcePartitions = map[string]map[string]*ProgressEntry{}
			progressEntries[pindex] = sourcePartitions
		}

		nodes, exists := sourcePartitions[sourcePartition]
		if !exists || nodes == nil {
			nodes = map[string]*ProgressEntry{}
			sourcePartitions[sourcePartition] = nodes
		}

		progressEntry, exists := nodes[node]
		if !exists || progressEntry == nil {
			progressEntry = &ProgressEntry{
				pindex:          pindex,
				sourcePartition: sourcePartition,
				node:            node,
				move:            -1,
			}
			nodes[node] = progressEntry
		}

		cb(progressEntry)

		// TODO: Check UUID matches, too.

		if !seenPIndexes[pindex] {
			seenPIndexes[pindex] = true
			seenPIndexesSorted =
				append(seenPIndexesSorted, pindex)

			sort.Strings(seenPIndexesSorted)
		}
	}

	for progress := range r.ProgressCh() {
		if progress.Error != nil {
			r.Log("main: error, progress: %+v", progress)
			continue
		}

		updateProgressEntries(r, updateProgressEntry)

		var b bytes.Buffer

		writeProgressTable(&b, maxNodeLen, maxPIndexLen,
			seenNodes,
			seenNodesSorted,
			seenPIndexes,
			seenPIndexesSorted,
			progressEntries)

		currEmit := b.String()
		if currEmit != lastEmit {
			r.Log("%s", currEmit)
		}

		lastEmit = currEmit
	}
}

func updateProgressEntries(
	r *rebalance.Rebalancer,
	updateProgressEntry func(pindex, sourcePartition, node string,
		cb func(*ProgressEntry)),
) {
	r.Visit(func(
		currStates rebalance.CurrStates,
		currSeqs rebalance.CurrSeqs,
		wantSeqs rebalance.WantSeqs,
		mapNextMoves map[string]*blance.NextMoves,
	) {
		for _, pindexes := range currStates {
			for pindex, nodes := range pindexes {
				for node, stateOp := range nodes {
					updateProgressEntry(pindex, "", node,
						func(pe *ProgressEntry) {
							pe.stateOp = stateOp
						})
				}
			}
		}

		for pindex, sourcePartitions := range currSeqs {
			for sourcePartition, nodes := range sourcePartitions {
				for node, currUUIDSeq := range nodes {
					updateProgressEntry(pindex,
						sourcePartition, node,
						func(pe *ProgressEntry) {
							pe.currUUIDSeq = currUUIDSeq

							if pe.initUUIDSeq.UUID == "" {
								pe.initUUIDSeq = currUUIDSeq
							}
						})
				}
			}
		}

		for pindex, sourcePartitions := range wantSeqs {
			for sourcePartition, nodes := range sourcePartitions {
				for node, wantUUIDSeq := range nodes {
					updateProgressEntry(pindex,
						sourcePartition, node,
						func(pe *ProgressEntry) {
							pe.wantUUIDSeq = wantUUIDSeq
						})
				}
			}
		}

		for pindex, nextMoves := range mapNextMoves {
			for i, nodeStateOp := range nextMoves.Moves {
				updateProgressEntry(pindex, "", nodeStateOp.Node,
					func(pe *ProgressEntry) {
						pe.move = i
						pe.done = i < nextMoves.Next
					})
			}
		}
	})
}

func writeProgressTable(b *bytes.Buffer,
	maxNodeLen, maxPIndexLen int,
	seenNodes map[string]bool,
	seenNodesSorted []string,
	seenPIndexes map[string]bool,
	seenPIndexesSorted []string,
	progressEntries map[string]map[string]map[string]*ProgressEntry,
) {
	written, _ := b.Write([]byte("%%%"))
	for i := written; i < maxPIndexLen; i++ {
		b.WriteByte(' ')
	}
	b.WriteByte(' ')

	for i, seenNode := range seenNodesSorted {
		if i > 0 {
			b.WriteByte(' ')
		}

		// TODO: Emit node human readable ADDR:PORT.
		b.Write([]byte(seenNode))
	}
	b.WriteByte('\n')

	for _, seenPIndex := range seenPIndexesSorted {
		b.Write([]byte(" %                  "))
		b.Write([]byte(seenPIndex))

		for _, seenNode := range seenNodesSorted {
			b.WriteByte(' ')

			sourcePartitions, exists :=
				progressEntries[seenPIndex]
			if !exists || sourcePartitions == nil {
				writeProgressCell(b, nil, nil, maxNodeLen)
				continue
			}

			nodes, exists := sourcePartitions[""]
			if !exists || nodes == nil {
				writeProgressCell(b, nil, nil, maxNodeLen)
				continue
			}

			pe, exists := nodes[seenNode]
			if !exists || pe == nil {
				writeProgressCell(b, nil, nil, maxNodeLen)
				continue
			}

			writeProgressCell(b, pe, sourcePartitions, maxNodeLen)
		}

		b.WriteByte('\n')
	}
}

var opMap = map[string]string{
	"":        ".",
	"add":     "+",
	"del":     "-",
	"promote": "P",
	"demote":  "D",
}

func writeProgressCell(b *bytes.Buffer,
	pe *ProgressEntry,
	sourcePartitions map[string]map[string]*ProgressEntry,
	maxNodeLen int) {
	written := 0

	totPct := 0.0 // To compute average pct.
	numPct := 0

	if pe != nil {
		written, _ = fmt.Fprintf(b, "%d ", pe.move)

		if sourcePartitions != nil {
			n, _ := b.Write([]byte(opMap[pe.stateOp.Op]))
			written = written + n

			for sourcePartition, nodes := range sourcePartitions {
				if sourcePartition == "" {
					continue
				}

				pex := nodes[pe.node]
				if pex == nil || pex.wantUUIDSeq.UUID == "" {
					continue
				}

				if pex.wantUUIDSeq.Seq <= pex.currUUIDSeq.Seq {
					totPct = totPct + 1.0
					numPct = numPct + 1
					continue
				}

				n := pex.currUUIDSeq.Seq - pex.initUUIDSeq.Seq
				d := pex.wantUUIDSeq.Seq - pex.initUUIDSeq.Seq
				if d > 0 {
					pct := float64(n) / float64(d)
					totPct = totPct + pct
					numPct = numPct + 1
				}
			}
		}
	} else {
		b.Write([]byte("  ."))
		written = 3
	}

	if numPct > 0 {
		avgPct := totPct / float64(numPct)

		n, _ := fmt.Fprintf(b, " %.1f%%", avgPct*100.0)
		written = written + n
	}

	for i := written; i < maxNodeLen; i++ {
		b.WriteByte(' ')
	}
}

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

func MainWelcome(flagAliases map[string][]string) {
	flag.VisitAll(func(f *flag.Flag) {
		if flagAliases[f.Name] != nil {
			log.Printf("  -%s=%q\n", f.Name, f.Value)
		}
	})
	log.Printf("  GOMAXPROCS=%d", runtime.GOMAXPROCS(-1))
}

func dumpOnSignal(signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, signals...)
	for _ = range c {
		log.Printf("dump: goroutine...")
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		log.Printf("dump: heap...")
		pprof.Lookup("heap").WriteTo(os.Stderr, 1)
	}
}
<|endoftext|>