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

// Unix domain sockets

package net

import (
	"os"
	"syscall"
)

func unixSocket(net string, laddr, raddr *UnixAddr, mode string) (fd *netFD, err os.Error) {
	var proto int
	switch net {
	default:
		return nil, UnknownNetworkError(net)
	case "unix":
		proto = syscall.SOCK_STREAM
	case "unixgram":
		proto = syscall.SOCK_DGRAM
	}

	var la, ra syscall.Sockaddr
	switch mode {
	default:
		panic("unixSocket mode " + mode)

	case "dial":
		if laddr != nil {
			la = &syscall.SockaddrUnix{Name: laddr.Name}
		}
		if raddr != nil {
			ra = &syscall.SockaddrUnix{Name: raddr.Name}
		} else if proto != syscall.SOCK_DGRAM || laddr == nil {
			return nil, &OpError{Op: mode, Net: net, Error: errMissingAddress}
		}

	case "listen":
		if laddr == nil {
			return nil, &OpError{mode, net, nil, errMissingAddress}
		}
		la = &syscall.SockaddrUnix{Name: laddr.Name}
		if raddr != nil {
			return nil, &OpError{Op: mode, Net: net, Addr: raddr, Error: &AddrError{Error: "unexpected remote address", Addr: raddr.String()}}
		}
	}

	f := sockaddrToUnix
	if proto != syscall.SOCK_STREAM {
		f = sockaddrToUnixgram
	}
	fd, oserr := socket(net, syscall.AF_UNIX, proto, 0, la, ra, f)
	if oserr != nil {
		goto Error
	}
	return fd, nil

Error:
	addr := raddr
	if mode == "listen" {
		addr = laddr
	}
	return nil, &OpError{Op: mode, Net: net, Addr: addr, Error: oserr}
}

// UnixAddr represents the address of a Unix domain socket end point.
type UnixAddr struct {
	Name     string
	Datagram bool
}

func sockaddrToUnix(sa syscall.Sockaddr) Addr {
	if s, ok := sa.(*syscall.SockaddrUnix); ok {
		return &UnixAddr{s.Name, false}
	}
	return nil
}

func sockaddrToUnixgram(sa syscall.Sockaddr) Addr {
	if s, ok := sa.(*syscall.SockaddrUnix); ok {
		return &UnixAddr{s.Name, true}
	}
	return nil
}

// Network returns the address's network name, "unix" or "unixgram".
func (a *UnixAddr) Network() string {
	if a == nil || !a.Datagram {
		return "unix"
	}
	return "unixgram"
}

func (a *UnixAddr) String() string {
	if a == nil {
		return "<nil>"
	}
	return a.Name
}

func (a *UnixAddr) toAddr() Addr {
	if a == nil { // nil *UnixAddr
		return nil // nil interface
	}
	return a
}

// ResolveUnixAddr parses addr as a Unix domain socket address.
// The string net gives the network name, "unix" or "unixgram".
func ResolveUnixAddr(net, addr string) (*UnixAddr, os.Error) {
	var datagram bool
	switch net {
	case "unix":
	case "unixgram":
		datagram = true
	default:
		return nil, UnknownNetworkError(net)
	}
	return &UnixAddr{addr, datagram}, nil
}

// UnixConn is an implementation of the Conn interface
// for connections to Unix domain sockets.
type UnixConn struct {
	fd *netFD
}

func newUnixConn(fd *netFD) *UnixConn { return &UnixConn{fd} }

func (c *UnixConn) ok() bool { return c != nil && c.fd != nil }

// Implementation of the Conn interface - see Conn for documentation.

// Read implements the net.Conn Read method.
func (c *UnixConn) Read(b []byte) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	return c.fd.Read(b)
}

// Write implements the net.Conn Write method.
func (c *UnixConn) Write(b []byte) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	return c.fd.Write(b)
}

// Close closes the Unix domain connection.
func (c *UnixConn) Close() os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	err := c.fd.Close()
	c.fd = nil
	return err
}

// LocalAddr returns the local network address, a *UnixAddr.
// Unlike in other protocols, LocalAddr is usually nil for dialed connections.
func (c *UnixConn) LocalAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.laddr
}

// RemoteAddr returns the remote network address, a *UnixAddr.
// Unlike in other protocols, RemoteAddr is usually nil for connections
// accepted by a listener.
func (c *UnixConn) RemoteAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.raddr
}

// SetTimeout implements the net.Conn SetTimeout method.
func (c *UnixConn) SetTimeout(nsec int64) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setTimeout(c.fd, nsec)
}

// SetReadTimeout implements the net.Conn SetReadTimeout method.
func (c *UnixConn) SetReadTimeout(nsec int64) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setReadTimeout(c.fd, nsec)
}

// SetWriteTimeout implements the net.Conn SetWriteTimeout method.
func (c *UnixConn) SetWriteTimeout(nsec int64) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setWriteTimeout(c.fd, nsec)
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
func (c *UnixConn) SetReadBuffer(bytes int) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setReadBuffer(c.fd, bytes)
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
func (c *UnixConn) SetWriteBuffer(bytes int) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setWriteBuffer(c.fd, bytes)
}

// ReadFromUnix reads a packet from c, copying the payload into b.
// It returns the number of bytes copied into b and the return address
// that was on the packet.
//
// ReadFromUnix can be made to time out and return
// an error with Timeout() == true after a fixed time limit;
// see SetTimeout and SetReadTimeout.
func (c *UnixConn) ReadFromUnix(b []byte) (n int, addr *UnixAddr, err os.Error) {
	if !c.ok() {
		return 0, nil, os.EINVAL
	}
	n, sa, err := c.fd.ReadFrom(b)
	switch sa := sa.(type) {
	case *syscall.SockaddrUnix:
		addr = &UnixAddr{sa.Name, c.fd.proto == syscall.SOCK_DGRAM}
	}
	return
}

// ReadFrom implements the net.PacketConn ReadFrom method.
func (c *UnixConn) ReadFrom(b []byte) (n int, addr Addr, err os.Error) {
	if !c.ok() {
		return 0, nil, os.EINVAL
	}
	n, uaddr, err := c.ReadFromUnix(b)
	return n, uaddr.toAddr(), err
}

// WriteToUnix writes a packet to addr via c, copying the payload from b.
//
// WriteToUnix can be made to time out and return
// an error with Timeout() == true after a fixed time limit;
// see SetTimeout and SetWriteTimeout.
// On packet-oriented connections, write timeouts are rare.
func (c *UnixConn) WriteToUnix(b []byte, addr *UnixAddr) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	if addr.Datagram != (c.fd.proto == syscall.SOCK_DGRAM) {
		return 0, os.EAFNOSUPPORT
	}
	sa := &syscall.SockaddrUnix{Name: addr.Name}
	return c.fd.WriteTo(b, sa)
}

// WriteTo implements the net.PacketConn WriteTo method.
func (c *UnixConn) WriteTo(b []byte, addr Addr) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	a, ok := addr.(*UnixAddr)
	if !ok {
		return 0, &OpError{"writeto", "unix", addr, os.EINVAL}
	}
	return c.WriteToUnix(b, a)
}

// DialUnix connects to the remote address raddr on the network net,
// which must be "unix" or "unixdgram".  If laddr is not nil, it is used
// as the local address for the connection.
func DialUnix(net string, laddr, raddr *UnixAddr) (c *UnixConn, err os.Error) {
	fd, e := unixSocket(net, laddr, raddr, "dial")
	if e != nil {
		return nil, e
	}
	return newUnixConn(fd), nil
}

// UnixListener is a Unix domain socket listener.
// Clients should typically use variables of type Listener
// instead of assuming Unix domain sockets.
type UnixListener struct {
	fd   *netFD
	path string
}

// ListenUnix announces on the Unix domain socket laddr and returns a Unix listener.
// Net must be "unix" (stream sockets).
func ListenUnix(net string, laddr *UnixAddr) (l *UnixListener, err os.Error) {
	if net != "unix" && net != "unixgram" {
		return nil, UnknownNetworkError(net)
	}
	if laddr != nil {
		laddr = &UnixAddr{laddr.Name, net == "unixgram"} // make our own copy
	}
	fd, err := unixSocket(net, laddr, nil, "listen")
	if err != nil {
		return nil, err
	}
	e1 := syscall.Listen(fd.sysfd, 8) // listenBacklog());
	if e1 != 0 {
		syscall.Close(fd.sysfd)
		return nil, &OpError{Op: "listen", Net: "unix", Addr: laddr, Error: os.Errno(e1)}
	}
	return &UnixListener{fd, laddr.Name}, nil
}

// AcceptUnix accepts the next incoming call and returns the new connection
// and the remote address.
func (l *UnixListener) AcceptUnix() (c *UnixConn, err os.Error) {
	if l == nil || l.fd == nil {
		return nil, os.EINVAL
	}
	fd, e := l.fd.accept(sockaddrToUnix)
	if e != nil {
		return nil, e
	}
	c = newUnixConn(fd)
	return c, nil
}

// Accept implements the Accept method in the Listener interface;
// it waits for the next call and returns a generic Conn.
func (l *UnixListener) Accept() (c Conn, err os.Error) {
	c1, err := l.AcceptUnix()
	if err != nil {
		return nil, err
	}
	return c1, nil
}

// Close stops listening on the Unix address.
// Already accepted connections are not closed.
func (l *UnixListener) Close() os.Error {
	if l == nil || l.fd == nil {
		return os.EINVAL
	}

	// The operating system doesn't clean up
	// the file that announcing created, so
	// we have to clean it up ourselves.
	// There's a race here--we can't know for
	// sure whether someone else has come along
	// and replaced our socket name already--
	// but this sequence (remove then close)
	// is at least compatible with the auto-remove
	// sequence in ListenUnix.  It's only non-Go
	// programs that can mess us up.
	if l.path[0] != '@' {
		syscall.Unlink(l.path)
	}
	err := l.fd.Close()
	l.fd = nil
	return err
}

// Addr returns the listener's network address.
func (l *UnixListener) Addr() Addr { return l.fd.laddr }

// ListenUnixgram listens for incoming Unix datagram packets addressed to the
// local address laddr.  The returned connection c's ReadFrom
// and WriteTo methods can be used to receive and send UDP
// packets with per-packet addressing.  The network net must be "unixgram".
func ListenUnixgram(net string, laddr *UnixAddr) (c *UDPConn, err os.Error) {
	switch net {
	case "unixgram":
	default:
		return nil, UnknownNetworkError(net)
	}
	if laddr == nil {
		return nil, &OpError{"listen", "unixgram", nil, errMissingAddress}
	}
	fd, e := unixSocket(net, laddr, nil, "listen")
	if e != nil {
		return nil, e
	}
	return newUDPConn(fd), nil
}
<commit_msg>Fix a typo in net/unixsock<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Unix domain sockets

package net

import (
	"os"
	"syscall"
)

func unixSocket(net string, laddr, raddr *UnixAddr, mode string) (fd *netFD, err os.Error) {
	var proto int
	switch net {
	default:
		return nil, UnknownNetworkError(net)
	case "unix":
		proto = syscall.SOCK_STREAM
	case "unixgram":
		proto = syscall.SOCK_DGRAM
	}

	var la, ra syscall.Sockaddr
	switch mode {
	default:
		panic("unixSocket mode " + mode)

	case "dial":
		if laddr != nil {
			la = &syscall.SockaddrUnix{Name: laddr.Name}
		}
		if raddr != nil {
			ra = &syscall.SockaddrUnix{Name: raddr.Name}
		} else if proto != syscall.SOCK_DGRAM || laddr == nil {
			return nil, &OpError{Op: mode, Net: net, Error: errMissingAddress}
		}

	case "listen":
		if laddr == nil {
			return nil, &OpError{mode, net, nil, errMissingAddress}
		}
		la = &syscall.SockaddrUnix{Name: laddr.Name}
		if raddr != nil {
			return nil, &OpError{Op: mode, Net: net, Addr: raddr, Error: &AddrError{Error: "unexpected remote address", Addr: raddr.String()}}
		}
	}

	f := sockaddrToUnix
	if proto != syscall.SOCK_STREAM {
		f = sockaddrToUnixgram
	}
	fd, oserr := socket(net, syscall.AF_UNIX, proto, 0, la, ra, f)
	if oserr != nil {
		goto Error
	}
	return fd, nil

Error:
	addr := raddr
	if mode == "listen" {
		addr = laddr
	}
	return nil, &OpError{Op: mode, Net: net, Addr: addr, Error: oserr}
}

// UnixAddr represents the address of a Unix domain socket end point.
type UnixAddr struct {
	Name     string
	Datagram bool
}

func sockaddrToUnix(sa syscall.Sockaddr) Addr {
	if s, ok := sa.(*syscall.SockaddrUnix); ok {
		return &UnixAddr{s.Name, false}
	}
	return nil
}

func sockaddrToUnixgram(sa syscall.Sockaddr) Addr {
	if s, ok := sa.(*syscall.SockaddrUnix); ok {
		return &UnixAddr{s.Name, true}
	}
	return nil
}

// Network returns the address's network name, "unix" or "unixgram".
func (a *UnixAddr) Network() string {
	if a == nil || !a.Datagram {
		return "unix"
	}
	return "unixgram"
}

func (a *UnixAddr) String() string {
	if a == nil {
		return "<nil>"
	}
	return a.Name
}

func (a *UnixAddr) toAddr() Addr {
	if a == nil { // nil *UnixAddr
		return nil // nil interface
	}
	return a
}

// ResolveUnixAddr parses addr as a Unix domain socket address.
// The string net gives the network name, "unix" or "unixgram".
func ResolveUnixAddr(net, addr string) (*UnixAddr, os.Error) {
	var datagram bool
	switch net {
	case "unix":
	case "unixgram":
		datagram = true
	default:
		return nil, UnknownNetworkError(net)
	}
	return &UnixAddr{addr, datagram}, nil
}

// UnixConn is an implementation of the Conn interface
// for connections to Unix domain sockets.
type UnixConn struct {
	fd *netFD
}

func newUnixConn(fd *netFD) *UnixConn { return &UnixConn{fd} }

func (c *UnixConn) ok() bool { return c != nil && c.fd != nil }

// Implementation of the Conn interface - see Conn for documentation.

// Read implements the net.Conn Read method.
func (c *UnixConn) Read(b []byte) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	return c.fd.Read(b)
}

// Write implements the net.Conn Write method.
func (c *UnixConn) Write(b []byte) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	return c.fd.Write(b)
}

// Close closes the Unix domain connection.
func (c *UnixConn) Close() os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	err := c.fd.Close()
	c.fd = nil
	return err
}

// LocalAddr returns the local network address, a *UnixAddr.
// Unlike in other protocols, LocalAddr is usually nil for dialed connections.
func (c *UnixConn) LocalAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.laddr
}

// RemoteAddr returns the remote network address, a *UnixAddr.
// Unlike in other protocols, RemoteAddr is usually nil for connections
// accepted by a listener.
func (c *UnixConn) RemoteAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.raddr
}

// SetTimeout implements the net.Conn SetTimeout method.
func (c *UnixConn) SetTimeout(nsec int64) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setTimeout(c.fd, nsec)
}

// SetReadTimeout implements the net.Conn SetReadTimeout method.
func (c *UnixConn) SetReadTimeout(nsec int64) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setReadTimeout(c.fd, nsec)
}

// SetWriteTimeout implements the net.Conn SetWriteTimeout method.
func (c *UnixConn) SetWriteTimeout(nsec int64) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setWriteTimeout(c.fd, nsec)
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
func (c *UnixConn) SetReadBuffer(bytes int) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setReadBuffer(c.fd, bytes)
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
func (c *UnixConn) SetWriteBuffer(bytes int) os.Error {
	if !c.ok() {
		return os.EINVAL
	}
	return setWriteBuffer(c.fd, bytes)
}

// ReadFromUnix reads a packet from c, copying the payload into b.
// It returns the number of bytes copied into b and the return address
// that was on the packet.
//
// ReadFromUnix can be made to time out and return
// an error with Timeout() == true after a fixed time limit;
// see SetTimeout and SetReadTimeout.
func (c *UnixConn) ReadFromUnix(b []byte) (n int, addr *UnixAddr, err os.Error) {
	if !c.ok() {
		return 0, nil, os.EINVAL
	}
	n, sa, err := c.fd.ReadFrom(b)
	switch sa := sa.(type) {
	case *syscall.SockaddrUnix:
		addr = &UnixAddr{sa.Name, c.fd.proto == syscall.SOCK_DGRAM}
	}
	return
}

// ReadFrom implements the net.PacketConn ReadFrom method.
func (c *UnixConn) ReadFrom(b []byte) (n int, addr Addr, err os.Error) {
	if !c.ok() {
		return 0, nil, os.EINVAL
	}
	n, uaddr, err := c.ReadFromUnix(b)
	return n, uaddr.toAddr(), err
}

// WriteToUnix writes a packet to addr via c, copying the payload from b.
//
// WriteToUnix can be made to time out and return
// an error with Timeout() == true after a fixed time limit;
// see SetTimeout and SetWriteTimeout.
// On packet-oriented connections, write timeouts are rare.
func (c *UnixConn) WriteToUnix(b []byte, addr *UnixAddr) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	if addr.Datagram != (c.fd.proto == syscall.SOCK_DGRAM) {
		return 0, os.EAFNOSUPPORT
	}
	sa := &syscall.SockaddrUnix{Name: addr.Name}
	return c.fd.WriteTo(b, sa)
}

// WriteTo implements the net.PacketConn WriteTo method.
func (c *UnixConn) WriteTo(b []byte, addr Addr) (n int, err os.Error) {
	if !c.ok() {
		return 0, os.EINVAL
	}
	a, ok := addr.(*UnixAddr)
	if !ok {
		return 0, &OpError{"writeto", "unix", addr, os.EINVAL}
	}
	return c.WriteToUnix(b, a)
}

// DialUnix connects to the remote address raddr on the network net,
// which must be "unix" or "unixgram".  If laddr is not nil, it is used
// as the local address for the connection.
func DialUnix(net string, laddr, raddr *UnixAddr) (c *UnixConn, err os.Error) {
	fd, e := unixSocket(net, laddr, raddr, "dial")
	if e != nil {
		return nil, e
	}
	return newUnixConn(fd), nil
}

// UnixListener is a Unix domain socket listener.
// Clients should typically use variables of type Listener
// instead of assuming Unix domain sockets.
type UnixListener struct {
	fd   *netFD
	path string
}

// ListenUnix announces on the Unix domain socket laddr and returns a Unix listener.
// Net must be "unix" (stream sockets).
func ListenUnix(net string, laddr *UnixAddr) (l *UnixListener, err os.Error) {
	if net != "unix" && net != "unixgram" {
		return nil, UnknownNetworkError(net)
	}
	if laddr != nil {
		laddr = &UnixAddr{laddr.Name, net == "unixgram"} // make our own copy
	}
	fd, err := unixSocket(net, laddr, nil, "listen")
	if err != nil {
		return nil, err
	}
	e1 := syscall.Listen(fd.sysfd, 8) // listenBacklog());
	if e1 != 0 {
		syscall.Close(fd.sysfd)
		return nil, &OpError{Op: "listen", Net: "unix", Addr: laddr, Error: os.Errno(e1)}
	}
	return &UnixListener{fd, laddr.Name}, nil
}

// AcceptUnix accepts the next incoming call and returns the new connection
// and the remote address.
func (l *UnixListener) AcceptUnix() (c *UnixConn, err os.Error) {
	if l == nil || l.fd == nil {
		return nil, os.EINVAL
	}
	fd, e := l.fd.accept(sockaddrToUnix)
	if e != nil {
		return nil, e
	}
	c = newUnixConn(fd)
	return c, nil
}

// Accept implements the Accept method in the Listener interface;
// it waits for the next call and returns a generic Conn.
func (l *UnixListener) Accept() (c Conn, err os.Error) {
	c1, err := l.AcceptUnix()
	if err != nil {
		return nil, err
	}
	return c1, nil
}

// Close stops listening on the Unix address.
// Already accepted connections are not closed.
func (l *UnixListener) Close() os.Error {
	if l == nil || l.fd == nil {
		return os.EINVAL
	}

	// The operating system doesn't clean up
	// the file that announcing created, so
	// we have to clean it up ourselves.
	// There's a race here--we can't know for
	// sure whether someone else has come along
	// and replaced our socket name already--
	// but this sequence (remove then close)
	// is at least compatible with the auto-remove
	// sequence in ListenUnix.  It's only non-Go
	// programs that can mess us up.
	if l.path[0] != '@' {
		syscall.Unlink(l.path)
	}
	err := l.fd.Close()
	l.fd = nil
	return err
}

// Addr returns the listener's network address.
func (l *UnixListener) Addr() Addr { return l.fd.laddr }

// ListenUnixgram listens for incoming Unix datagram packets addressed to the
// local address laddr.  The returned connection c's ReadFrom
// and WriteTo methods can be used to receive and send UDP
// packets with per-packet addressing.  The network net must be "unixgram".
func ListenUnixgram(net string, laddr *UnixAddr) (c *UDPConn, err os.Error) {
	switch net {
	case "unixgram":
	default:
		return nil, UnknownNetworkError(net)
	}
	if laddr == nil {
		return nil, &OpError{"listen", "unixgram", nil, errMissingAddress}
	}
	fd, e := unixSocket(net, laddr, nil, "listen")
	if e != nil {
		return nil, e
	}
	return newUDPConn(fd), nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package cluster

import (
	"testing"

	"reflect"

	"k8s.io/kubernetes/federation/apis/federation"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/testapi"
	apitesting "k8s.io/kubernetes/pkg/api/testing"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
)

func validNewCluster() *federation.Cluster {
	return &federation.Cluster{
		ObjectMeta: api.ObjectMeta{
			Name:            "foo",
			ResourceVersion: "4",
			Labels: map[string]string{
				"name": "foo",
			},
		},
		Spec: federation.ClusterSpec{
			ServerAddressByClientCIDRs: []federation.ServerAddressByClientCIDR{
				{
					ClientCIDR:    "0.0.0.0/0",
					ServerAddress: "localhost:8888",
				},
			},
		},
		Status: federation.ClusterStatus{
			Conditions: []federation.ClusterCondition{
				{Type: federation.ClusterReady, Status: api.ConditionTrue},
			},
		},
	}
}

func invalidNewCluster() *federation.Cluster {
	// Create a cluster with empty ServerAddressByClientCIDRs (which is a required field).
	return &federation.Cluster{
		ObjectMeta: api.ObjectMeta{
			Name:            "foo2",
			ResourceVersion: "5",
		},
		Status: federation.ClusterStatus{
			Conditions: []federation.ClusterCondition{
				{Type: federation.ClusterReady, Status: api.ConditionFalse},
			},
		},
	}
}

func TestClusterStrategy(t *testing.T) {
	ctx := api.NewDefaultContext()
	if Strategy.NamespaceScoped() {
		t.Errorf("Cluster should not be namespace scoped")
	}
	if Strategy.AllowCreateOnUpdate() {
		t.Errorf("Cluster should not allow create on update")
	}

	cluster := validNewCluster()
	Strategy.PrepareForCreate(ctx, cluster)
	if len(cluster.Status.Conditions) != 0 {
		t.Errorf("Cluster should not allow setting conditions on create")
	}
	errs := Strategy.Validate(ctx, cluster)
	if len(errs) != 0 {
		t.Errorf("Unexpected error validating %v", errs)
	}

	invalidCluster := invalidNewCluster()
	Strategy.PrepareForUpdate(ctx, invalidCluster, cluster)
	if reflect.DeepEqual(invalidCluster.Spec, cluster.Spec) ||
		!reflect.DeepEqual(invalidCluster.Status, cluster.Status) {
		t.Error("Only spec is expected being changed")
	}
	errs = Strategy.ValidateUpdate(ctx, invalidCluster, cluster)
	if len(errs) == 0 {
		t.Errorf("Expected a validation error")
	}
	if cluster.ResourceVersion != "4" {
		t.Errorf("Incoming resource version on update should not be mutated")
	}
}

func TestClusterStatusStrategy(t *testing.T) {
	ctx := api.NewDefaultContext()
	if StatusStrategy.NamespaceScoped() {
		t.Errorf("Cluster should not be namespace scoped")
	}
	if StatusStrategy.AllowCreateOnUpdate() {
		t.Errorf("Cluster should not allow create on update")
	}

	cluster := validNewCluster()
	invalidCluster := invalidNewCluster()
	StatusStrategy.PrepareForUpdate(ctx, cluster, invalidCluster)
	if !reflect.DeepEqual(invalidCluster.Spec, cluster.Spec) ||
		reflect.DeepEqual(invalidCluster.Status, cluster.Status) {
		t.Logf("== cluster.Spec: %v\n", cluster.Spec)
		t.Logf("== cluster.Status: %v\n", cluster.Status)
		t.Logf("== invalidCluster.Spec: %v\n", cluster.Spec)
		t.Logf("== invalidCluster.Spec: %v\n", cluster.Status)
		t.Error("Only spec is expected being changed")
	}
	errs := Strategy.ValidateUpdate(ctx, invalidCluster, cluster)
	if len(errs) == 0 {
		t.Errorf("Expected a validation error")
	}
	if cluster.ResourceVersion != "4" {
		t.Errorf("Incoming resource version on update should not be mutated")
	}
}

func TestMatchCluster(t *testing.T) {
	testFieldMap := map[bool][]fields.Set{
		true: {
			{"metadata.name": "foo"},
		},
		false: {
			{"foo": "bar"},
		},
	}

	for expectedResult, fieldSet := range testFieldMap {
		for _, field := range fieldSet {
			m := MatchCluster(labels.Everything(), field.AsSelector())
			_, matchesSingle := m.MatchesSingle()
			if e, a := expectedResult, matchesSingle; e != a {
				t.Errorf("%+v: expected %v, got %v", fieldSet, e, a)
			}
		}
	}
}

func TestSelectableFieldLabelConversions(t *testing.T) {
	apitesting.TestSelectableFieldLabelConversionsOfKind(t,
		testapi.Federation.GroupVersion().String(),
		"Cluster",
		labels.Set(ClusterToSelectableFields(&federation.Cluster{})),
		nil,
	)
}
<commit_msg>change all PredicateFunc to use SelectionPredicate<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package cluster

import (
	"testing"

	"reflect"

	"k8s.io/kubernetes/federation/apis/federation"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/testapi"
	apitesting "k8s.io/kubernetes/pkg/api/testing"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
)

func validNewCluster() *federation.Cluster {
	return &federation.Cluster{
		ObjectMeta: api.ObjectMeta{
			Name:            "foo",
			ResourceVersion: "4",
			Labels: map[string]string{
				"name": "foo",
			},
		},
		Spec: federation.ClusterSpec{
			ServerAddressByClientCIDRs: []federation.ServerAddressByClientCIDR{
				{
					ClientCIDR:    "0.0.0.0/0",
					ServerAddress: "localhost:8888",
				},
			},
		},
		Status: federation.ClusterStatus{
			Conditions: []federation.ClusterCondition{
				{Type: federation.ClusterReady, Status: api.ConditionTrue},
			},
		},
	}
}

func invalidNewCluster() *federation.Cluster {
	// Create a cluster with empty ServerAddressByClientCIDRs (which is a required field).
	return &federation.Cluster{
		ObjectMeta: api.ObjectMeta{
			Name:            "foo2",
			ResourceVersion: "5",
		},
		Status: federation.ClusterStatus{
			Conditions: []federation.ClusterCondition{
				{Type: federation.ClusterReady, Status: api.ConditionFalse},
			},
		},
	}
}

func TestClusterStrategy(t *testing.T) {
	ctx := api.NewDefaultContext()
	if Strategy.NamespaceScoped() {
		t.Errorf("Cluster should not be namespace scoped")
	}
	if Strategy.AllowCreateOnUpdate() {
		t.Errorf("Cluster should not allow create on update")
	}

	cluster := validNewCluster()
	Strategy.PrepareForCreate(ctx, cluster)
	if len(cluster.Status.Conditions) != 0 {
		t.Errorf("Cluster should not allow setting conditions on create")
	}
	errs := Strategy.Validate(ctx, cluster)
	if len(errs) != 0 {
		t.Errorf("Unexpected error validating %v", errs)
	}

	invalidCluster := invalidNewCluster()
	Strategy.PrepareForUpdate(ctx, invalidCluster, cluster)
	if reflect.DeepEqual(invalidCluster.Spec, cluster.Spec) ||
		!reflect.DeepEqual(invalidCluster.Status, cluster.Status) {
		t.Error("Only spec is expected being changed")
	}
	errs = Strategy.ValidateUpdate(ctx, invalidCluster, cluster)
	if len(errs) == 0 {
		t.Errorf("Expected a validation error")
	}
	if cluster.ResourceVersion != "4" {
		t.Errorf("Incoming resource version on update should not be mutated")
	}
}

func TestClusterStatusStrategy(t *testing.T) {
	ctx := api.NewDefaultContext()
	if StatusStrategy.NamespaceScoped() {
		t.Errorf("Cluster should not be namespace scoped")
	}
	if StatusStrategy.AllowCreateOnUpdate() {
		t.Errorf("Cluster should not allow create on update")
	}

	cluster := validNewCluster()
	invalidCluster := invalidNewCluster()
	StatusStrategy.PrepareForUpdate(ctx, cluster, invalidCluster)
	if !reflect.DeepEqual(invalidCluster.Spec, cluster.Spec) ||
		reflect.DeepEqual(invalidCluster.Status, cluster.Status) {
		t.Logf("== cluster.Spec: %v\n", cluster.Spec)
		t.Logf("== cluster.Status: %v\n", cluster.Status)
		t.Logf("== invalidCluster.Spec: %v\n", cluster.Spec)
		t.Logf("== invalidCluster.Spec: %v\n", cluster.Status)
		t.Error("Only spec is expected being changed")
	}
	errs := Strategy.ValidateUpdate(ctx, invalidCluster, cluster)
	if len(errs) == 0 {
		t.Errorf("Expected a validation error")
	}
	if cluster.ResourceVersion != "4" {
		t.Errorf("Incoming resource version on update should not be mutated")
	}
}

func TestMatchCluster(t *testing.T) {
	testFieldMap := map[bool][]fields.Set{
		true: {
			{"metadata.name": "foo"},
		},
		false: {
			{"foo": "bar"},
		},
	}

	for expectedResult, fieldSet := range testFieldMap {
		for _, field := range fieldSet {
			m := MatchCluster(labels.Everything(), field.AsSelector())
			_, matchesSingle := m.MatchesSingle()
			if e, a := expectedResult, matchesSingle; e != a {
				t.Errorf("%+v: expected %v, got %v", fieldSet, e, a)
			}
		}
	}
}

func TestSelectableFieldLabelConversions(t *testing.T) {
	apitesting.TestSelectableFieldLabelConversionsOfKind(t,
		testapi.Federation.GroupVersion().String(),
		"Cluster",
		ClusterToSelectableFields(&federation.Cluster{}),
		nil,
	)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"image"
	"image/png"
	"os"

	"path/filepath"

	"github.com/codegangsta/cli"
	"github.com/landaire/osutil"
)

const (
	version = "0.1.0"
)

func main() {
	app := cli.NewApp()
	app.Name = "paatoimg"
	app.Usage = "Extract satellite PAA files from PBO archives, and convert them to a giant stitched maps"
	app.Author = "Lander Brandt"
	app.Email = "@landaire"
	app.Version = version

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "pbo",
			Usage: "Base PBO file to read (others will be detected automatically)",
		},
		cli.StringFlag{
			Name:  "outdir, od",
			Usage: "Output directory to dump PNG files",
		},
		cli.StringFlag{
			Name:  "outfile, of",
			Usage: "Output output file to write stitched PNG file",
		},
	}

	app.Action = Stitch

	err := app.Run(os.Args)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
	}
}

func Stitch(c *cli.Context) {
	required := []string{"pbo", "outdir", "outfile"}
	for _, flag := range required {
		if !c.GlobalIsSet(flag) {
			c.App.Command("help").Run(c)
			return
		}
	}

	paaFiles, err := DumpPaaFiles(c.String("pbo"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred dumping PAA files: %s", err)
		os.Exit(1)
	}

	// Create the output dir if it doesn't exist
	if exist, _ := osutil.Exists(c.String("outdir")); !exist {
		err := osutil.MkdirIntermediate(c.String("outdir"))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error occurred creating output dir: %s", err)
			os.Exit(1)
		}
	}

	pngs := []string{}
	for _, file := range paaFiles {
		paaFileName := filepath.Base(file)
		pngFileName := paaFileName[:len(paaFileName)-len(filepath.Ext(paaFileName))] + ".png"

		err := ConvertPaaToPng(file, filepath.Join(c.String("outdir"), pngFileName))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error occurred converting PAA to PNG: %s", err)
			os.Exit(1)
		}

		pngs = append(pngs, pngFileName)
	}

	stitchedImage, err := StitchImages(pngs, image.Point{512, 512})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred creating stitched image: %s", err)
		os.Exit(1)
	}

	outFile, err := os.Create(c.String("outfile"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred creating output file: %s", err)
		os.Exit(1)
	}
	defer outFile.Close()

	err = png.Encode(outFile, stitchedImage)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred while writing PNG: %s", err)
		os.Exit(1)
	}
}
<commit_msg>Fix creating intermediate dirs for outdir<commit_after>package main

import (
	"fmt"
	"image"
	"image/png"
	"os"

	"path/filepath"

	"github.com/codegangsta/cli"
	"github.com/landaire/osutil"
)

const (
	version = "0.1.0"
)

func main() {
	app := cli.NewApp()
	app.Name = "paatoimg"
	app.Usage = "Extract satellite PAA files from PBO archives, and convert them to a giant stitched maps"
	app.Author = "Lander Brandt"
	app.Email = "@landaire"
	app.Version = version

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "pbo",
			Usage: "Base PBO file to read (others will be detected automatically)",
		},
		cli.StringFlag{
			Name:  "outdir, od",
			Usage: "Output directory to dump PNG files",
		},
		cli.StringFlag{
			Name:  "outfile, of",
			Usage: "Output output file to write stitched PNG file",
		},
	}

	app.Action = Stitch

	err := app.Run(os.Args)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
	}
}

func Stitch(c *cli.Context) {
	required := []string{"pbo", "outdir", "outfile"}
	for _, flag := range required {
		if !c.GlobalIsSet(flag) {
			c.App.Command("help").Run(c)
			return
		}
	}

	paaFiles, err := DumpPaaFiles(c.String("pbo"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred dumping PAA files: %s", err)
		os.Exit(1)
	}

	// Create the output dir if it doesn't exist
	if exist, _ := osutil.Exists(c.String("outdir")); !exist {
		err := os.MkdirAll(c.String("outdir"), os.ModePerm)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error occurred creating output dir: %s", err)
			os.Exit(1)
		}
	}

	pngs := []string{}
	for _, file := range paaFiles {
		paaFileName := filepath.Base(file)
		pngFileName := paaFileName[:len(paaFileName)-len(filepath.Ext(paaFileName))] + ".png"

		err := ConvertPaaToPng(file, filepath.Join(c.String("outdir"), pngFileName))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error occurred converting PAA to PNG: %s", err)
			os.Exit(1)
		}

		pngs = append(pngs, pngFileName)
	}

	stitchedImage, err := StitchImages(pngs, image.Point{512, 512})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred creating stitched image: %s", err)
		os.Exit(1)
	}

	outFile, err := os.Create(c.String("outfile"))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred creating output file: %s", err)
		os.Exit(1)
	}
	defer outFile.Close()

	err = png.Encode(outFile, stitchedImage)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error occurred while writing PNG: %s", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	geoipc "github.com/rubiojr/freegeoip-client"
	mqttc "./utils/mqtt"
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"gopkg.in/alecthomas/kingpin.v1"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

type Host struct {
	IP          string  `json:"ip"`
	Name        string  `json:"hostname"`
	Hop         int     `json:"hop-number"`
	Sent        int     `json:"sent"`
	LostPercent float64 `json:"lost-percent"`
	Last        float64 `json:"mean"`
	Avg         float64 `json:"mean"`
	Best        float64 `json:"best"`
	Worst       float64 `json:"worst"`
	StDev       float64 `json:"standard-dev"`
}

type Report struct {
	Time        time.Time     `json:"time"`
	Hosts       []*Host       `json:"hosts"`
	Hops        int           `json:"hops"`
	ElapsedTime time.Duration `json:"elapsed_time"`
	Location    geoipc.Location `json:"location"`
}

func NewReport(reportCycles int, host string, args ...string) *Report {
	report := &Report{}
	report.Time = time.Now()
	args = append([]string{"--report", "-c", strconv.Itoa(reportCycles), host}, args...)

	tstart := time.Now()
	rawOutput, err := exec.Command("mtr", args...).Output()

	if err != nil {
		panic("Error running the mtr command")
	}

	buf := bytes.NewBuffer(rawOutput)
	scanner := bufio.NewScanner(buf)
	scanner.Split(bufio.ScanLines)

	skipHeader := 2
	for scanner.Scan() {
		if skipHeader != 0 {
			skipHeader -= 1
			continue
		}

		tokens := strings.Fields(scanner.Text())
		sent, err := strconv.Atoi(tokens[3])
		if err != nil {
			panic("Error parsing sent field")
		}

		host := Host{
			IP:   tokens[1],
			Sent: sent,
		}

		f2F(strings.Replace(tokens[2], "%", "", -1), &host.LostPercent)
		f2F(tokens[4], &host.Last)
		f2F(tokens[5], &host.Avg)
		f2F(tokens[6], &host.Best)
		f2F(tokens[7], &host.Worst)
		f2F(tokens[8], &host.StDev)

		report.Hosts = append(report.Hosts, &host)
	}

	report.Hops = len(report.Hosts)
	report.ElapsedTime = time.Since(tstart)
	loc, err := geoipc.GetLocation()
	if err != nil {
		report.Location = geoipc.Location{}
	} else {
		report.Location = loc
	}

	return report
}

func f2F(val string, field *float64) {
	f, err := strconv.ParseFloat(val, 64)
	*field = f
	if err != nil {
		panic("Error parsing field")
	}
}

func run(count int, host, brokerUrl, topic string) error {
	r := NewReport(count, host, "-n")
	msg, _ := json.Marshal(r)
	err := mqttc.PushMsg("push-mtr", brokerUrl, topic, string(msg))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error sending report: %s", err)
	}
	return err
}

func main() {
	count := kingpin.Flag("count", "Report cycles (mtr -c)").
		Default("10").Int()

	topic := kingpin.Flag("topic", "MTTQ topic").Default("/metrics/mtr").
		String()

	host := kingpin.Arg("host", "Target host").Required().String()

	repeat := kingpin.Flag("repeat", "Send the report every X seconds").
		Default("0").Int()

	brokerUrl := kingpin.Flag("broker-url", "MQTT broker URL").
		Default("").String()

	kingpin.Version("0.1")
	kingpin.Parse()

	if *brokerUrl == "" {
		*brokerUrl = os.Getenv("MQTT_URL")
		if *brokerUrl == "" {
			fmt.Fprintf(os.Stderr, "Invalid MQTT broker URL")
		}
	}

	if *repeat != 0 {
		timer := time.NewTicker(1 * time.Second)
		for _ = range timer.C {
			run(*count, *host, *brokerUrl, *topic)
		}
	} else {
		err := run(*count, *host, *brokerUrl, *topic)
		if err != nil {
			os.Exit(1)
		}
	}
}
<commit_msg>exit 1 if mtr not in path<commit_after>package main

import (
	geoipc "github.com/rubiojr/freegeoip-client"
	mqttc "./utils/mqtt"
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"gopkg.in/alecthomas/kingpin.v1"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

type Host struct {
	IP          string  `json:"ip"`
	Name        string  `json:"hostname"`
	Hop         int     `json:"hop-number"`
	Sent        int     `json:"sent"`
	LostPercent float64 `json:"lost-percent"`
	Last        float64 `json:"mean"`
	Avg         float64 `json:"mean"`
	Best        float64 `json:"best"`
	Worst       float64 `json:"worst"`
	StDev       float64 `json:"standard-dev"`
}

type Report struct {
	Time        time.Time     `json:"time"`
	Hosts       []*Host       `json:"hosts"`
	Hops        int           `json:"hops"`
	ElapsedTime time.Duration `json:"elapsed_time"`
	Location    geoipc.Location `json:"location"`
}

func NewReport(reportCycles int, host string, args ...string) *Report {
	report := &Report{}
	report.Time = time.Now()
	args = append([]string{"--report", "-n", "-c", strconv.Itoa(reportCycles), host}, args...)

	tstart := time.Now()
	mtr := findMtrBin()
	rawOutput, err := exec.Command(mtr, args...).Output()

	if err != nil {
		panic("Error running the mtr command")
	}

	buf := bytes.NewBuffer(rawOutput)
	scanner := bufio.NewScanner(buf)
	scanner.Split(bufio.ScanLines)

	skipHeader := 2
	for scanner.Scan() {
		if skipHeader != 0 {
			skipHeader -= 1
			continue
		}

		tokens := strings.Fields(scanner.Text())
		sent, err := strconv.Atoi(tokens[3])
		if err != nil {
			panic("Error parsing sent field")
		}

		host := Host{
			IP:   tokens[1],
			Sent: sent,
		}

		f2F(strings.Replace(tokens[2], "%", "", -1), &host.LostPercent)
		f2F(tokens[4], &host.Last)
		f2F(tokens[5], &host.Avg)
		f2F(tokens[6], &host.Best)
		f2F(tokens[7], &host.Worst)
		f2F(tokens[8], &host.StDev)

		report.Hosts = append(report.Hosts, &host)
	}

	report.Hops = len(report.Hosts)
	report.ElapsedTime = time.Since(tstart)
	loc, err := geoipc.GetLocation()
	if err != nil {
		report.Location = geoipc.Location{}
	} else {
		report.Location = loc
	}

	return report
}

func f2F(val string, field *float64) {
	f, err := strconv.ParseFloat(val, 64)
	*field = f
	if err != nil {
		panic("Error parsing field")
	}
}

func findMtrBin() string {
	paths := os.Getenv("PATH")
	if paths == "" {
		return ""
	}

	for _, path := range strings.Split(paths, ":") {
		if _, err := os.Stat(path + "/mtr"); err == nil {
			return path + "/mtr"
		}
	}

	return ""
}

func run(count int, host, brokerUrl, topic string) error {
	r := NewReport(count, host)
	msg, _ := json.Marshal(r)
	err := mqttc.PushMsg("push-mtr", brokerUrl, topic, string(msg))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error sending report: %s", err)
	}
	return err
}

func main() {
	count := kingpin.Flag("count", "Report cycles (mtr -c)").
		Default("10").Int()

	topic := kingpin.Flag("topic", "MTTQ topic").Default("/metrics/mtr").
		String()

	host := kingpin.Arg("host", "Target host").Required().String()

	repeat := kingpin.Flag("repeat", "Send the report every X seconds").
		Default("0").Int()

	brokerUrl := kingpin.Flag("broker-url", "MQTT broker URL").
		Default("").String()

	kingpin.Version("0.1")
	kingpin.Parse()

	if findMtrBin() == "" {
		fmt.Fprintf(os.Stderr, "mtr binary not found in path")
		os.Exit(1)
	}

	if *brokerUrl == "" {
		*brokerUrl = os.Getenv("MQTT_URL")
		if *brokerUrl == "" {
			fmt.Fprintf(os.Stderr, "Invalid MQTT broker URL")
		}
	}

	if *repeat != 0 {
		timer := time.NewTicker(1 * time.Second)
		for _ = range timer.C {
			run(*count, *host, *brokerUrl, *topic)
		}
	} else {
		err := run(*count, *host, *brokerUrl, *topic)
		if err != nil {
			os.Exit(1)
		}
	}
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"net/http"

	"github.com/rs/zerolog"

	"github.com/gilcrest/go-api-basic/domain/app"
	"github.com/gilcrest/go-api-basic/domain/audit"
	"github.com/gilcrest/go-api-basic/domain/user"
	"github.com/gilcrest/go-api-basic/service"
)

// CreateMovieService creates a Movie
type CreateMovieService interface {
	Create(ctx context.Context, r *service.CreateMovieRequest, adt audit.Audit) (service.MovieResponse, error)
}

// UpdateMovieService is a service for updating a Movie
type UpdateMovieService interface {
	Update(ctx context.Context, r *service.UpdateMovieRequest, adt audit.Audit) (service.MovieResponse, error)
}

// DeleteMovieService is a service for deleting a Movie
type DeleteMovieService interface {
	Delete(ctx context.Context, extlID string) (service.DeleteMovieResponse, error)
}

// FindMovieService interface reads a Movie form the database
type FindMovieService interface {
	FindMovieByID(ctx context.Context, extlID string) (service.MovieResponse, error)
	FindAllMovies(ctx context.Context) ([]service.MovieResponse, error)
}

// CreateOrgService creates an Org
type CreateOrgService interface {
	Create(ctx context.Context, r *service.CreateOrgRequest, adt audit.Audit) (service.OrgResponse, error)
}

// UpdateOrgService updates an Org
type UpdateOrgService interface {
	Update(ctx context.Context, r *service.UpdateOrgRequest, adt audit.Audit) (service.OrgResponse, error)
}

// FindOrgService retrieves Org information from the datastore
type FindOrgService interface {
	FindAll(ctx context.Context) ([]service.OrgResponse, error)
	FindByExternalID(ctx context.Context, extlID string) (service.OrgResponse, error)
}

// CreateAppService creates an App
type CreateAppService interface {
	Create(ctx context.Context, r *service.CreateAppRequest, adt audit.Audit) (service.AppResponse, error)
}

// FindAppService retrieves an App
type FindAppService interface {
	// FindAppByAPIKey finds an app given its External ID and determines
	// if the given API key is a valid key for it
	FindAppByAPIKey(ctx context.Context, realm, appExtlID, apiKey string) (app.App, error)
}

// RegisterUserService registers a new user
type RegisterUserService interface {
	SelfRegister(ctx context.Context, adt audit.Audit) error
}

// FindUserService retrieves a User
type FindUserService interface {
	FindUserByOauth2Token(ctx context.Context, params service.FindUserParams) (user.User, error)
}

// AuthorizeService determines whether an app and a user can perform
// an action against a resource
type AuthorizeService interface {
	Authorize(lgr zerolog.Logger, r *http.Request, sub audit.Audit) error
}

// LoggerService reads and updates the logger state
type LoggerService interface {
	Read() service.LoggerResponse
	Update(r *service.LoggerRequest) (service.LoggerResponse, error)
}

// PingService pings the database and responds whether it is up or down
type PingService interface {
	Ping(ctx context.Context, logger zerolog.Logger) service.PingResponse
}

// GenesisService initializes the database with dependent data
type GenesisService interface {
	Seed(ctx context.Context, r *service.GenesisRequest) (service.FullGenesisResponse, error)
}

// Services are used by the application service handlers
type Services struct {
	CreateMovieService  CreateMovieService
	UpdateMovieService  UpdateMovieService
	DeleteMovieService  DeleteMovieService
	FindMovieService    FindMovieService
	PingService         PingService
	LoggerService       LoggerService
	CreateOrgService    CreateOrgService
	UpdateOrgService    UpdateOrgService
	FindOrgService      FindOrgService
	CreateAppService    CreateAppService
	FindAppService      FindAppService
	RegisterUserService RegisterUserService
	FindUserService     FindUserService
	AuthorizeService    AuthorizeService
	GenesisService      GenesisService
}
<commit_msg>changes for OrgService<commit_after>package server

import (
	"context"
	"net/http"

	"github.com/rs/zerolog"

	"github.com/gilcrest/go-api-basic/domain/app"
	"github.com/gilcrest/go-api-basic/domain/audit"
	"github.com/gilcrest/go-api-basic/domain/user"
	"github.com/gilcrest/go-api-basic/service"
)

// CreateMovieService creates a Movie
type CreateMovieService interface {
	Create(ctx context.Context, r *service.CreateMovieRequest, adt audit.Audit) (service.MovieResponse, error)
}

// UpdateMovieService is a service for updating a Movie
type UpdateMovieService interface {
	Update(ctx context.Context, r *service.UpdateMovieRequest, adt audit.Audit) (service.MovieResponse, error)
}

// DeleteMovieService is a service for deleting a Movie
type DeleteMovieService interface {
	Delete(ctx context.Context, extlID string) (service.DeleteMovieResponse, error)
}

// FindMovieService interface reads a Movie form the database
type FindMovieService interface {
	FindMovieByID(ctx context.Context, extlID string) (service.MovieResponse, error)
	FindAllMovies(ctx context.Context) ([]service.MovieResponse, error)
}

// OrgService manages the retrieval and manipulation of an Org
type OrgService interface {
	Create(ctx context.Context, r *service.CreateOrgRequest, adt audit.Audit) (service.OrgResponse, error)
	Update(ctx context.Context, r *service.UpdateOrgRequest, adt audit.Audit) (service.OrgResponse, error)
	FindAll(ctx context.Context) ([]service.OrgResponse, error)
	FindByExternalID(ctx context.Context, extlID string) (service.OrgResponse, error)
}

// CreateAppService creates an App
type CreateAppService interface {
	Create(ctx context.Context, r *service.CreateAppRequest, adt audit.Audit) (service.AppResponse, error)
}

// FindAppService retrieves an App
type FindAppService interface {
	// FindAppByAPIKey finds an app given its External ID and determines
	// if the given API key is a valid key for it
	FindAppByAPIKey(ctx context.Context, realm, appExtlID, apiKey string) (app.App, error)
}

// RegisterUserService registers a new user
type RegisterUserService interface {
	SelfRegister(ctx context.Context, adt audit.Audit) error
}

// FindUserService retrieves a User
type FindUserService interface {
	FindUserByOauth2Token(ctx context.Context, params service.FindUserParams) (user.User, error)
}

// AuthorizeService determines whether an app and a user can perform
// an action against a resource
type AuthorizeService interface {
	Authorize(lgr zerolog.Logger, r *http.Request, sub audit.Audit) error
}

// LoggerService reads and updates the logger state
type LoggerService interface {
	Read() service.LoggerResponse
	Update(r *service.LoggerRequest) (service.LoggerResponse, error)
}

// PingService pings the database and responds whether it is up or down
type PingService interface {
	Ping(ctx context.Context, logger zerolog.Logger) service.PingResponse
}

// GenesisService initializes the database with dependent data
type GenesisService interface {
	Seed(ctx context.Context) (service.FullGenesisResponse, error)
}

// Services are used by the application service handlers
type Services struct {
	CreateMovieService  CreateMovieService
	UpdateMovieService  UpdateMovieService
	DeleteMovieService  DeleteMovieService
	FindMovieService    FindMovieService
	OrgService          OrgService
	CreateAppService    CreateAppService
	FindAppService      FindAppService
	RegisterUserService RegisterUserService
	FindUserService     FindUserService
	AuthorizeService    AuthorizeService
	PingService         PingService
	LoggerService       LoggerService
	GenesisService      GenesisService
}
<|endoftext|>


<commit_before>package pixur

import (
	"database/sql"
	"math"
)

const (
	DefaultStartID = math.MaxInt64
	DefaultMaxPics = 60
)

type ReadIndexPicsTask struct {
	// Deps
	db *sql.DB

	// Inputs
	// Only get pics with Pic Id <= than this.  If unset, the latest pics will be returned.
	StartID int64
	// MaxPics is the maximum number of pics to return.  Note that the number of pictures returned
	// may be less than the number requested.  If unset, the de
	MaxPics int64

	// State

	// Results
	Pics []*Pic
}

func (t *ReadIndexPicsTask) Reset() {}

func (t *ReadIndexPicsTask) Run() error {

	var startID int64
	if t.StartID != 0 {
		startID = t.StartID
	} else {
		startID = DefaultStartID
	}

	var maxPics int64
	if t.MaxPics != 0 {
		maxPics = t.MaxPics
	} else {
		maxPics = DefaultMaxPics
	}

	// Technically an initial lookup of the created time of the provided Pic ID id needed.
	// TODO: decide if this is worth the extra DB call.
	rows, err := t.db.Query(
		"SELECT * FROM pics WHERE id <= ? ORDER BY created_time DESC LIMIT ?;",
		startID, maxPics)

	if err != nil {
		return err
	}

	defer rows.Close()

	columnNames, err := rows.Columns()
	if err != nil {
		return err
	}

	var pics []*Pic
	for rows.Next() {
		var p = new(Pic)
		if err := rows.Scan(p.ColumnPointers(columnNames)...); err != nil {
			return err
		}
		pics = append(pics, p)
	}

	if err := rows.Err(); err != nil {
		return err
	}

	t.Pics = pics

	return nil
}
<commit_msg>make read index pics task reset state properly<commit_after>package pixur

import (
	"database/sql"
	"math"
)

const (
	DefaultStartID = math.MaxInt64
	DefaultMaxPics = 60
)

type ReadIndexPicsTask struct {
	// Deps
	db *sql.DB

	// Inputs
	// Only get pics with Pic Id <= than this.  If unset, the latest pics will be returned.
	StartID int64
	// MaxPics is the maximum number of pics to return.  Note that the number of pictures returned
	// may be less than the number requested.  If unset, the de
	MaxPics int64

	// State

	// Results
	Pics []*Pic
}

func (t *ReadIndexPicsTask) Reset() {
	t.Pics = nil
}

func (t *ReadIndexPicsTask) Run() error {

	var startID int64
	if t.StartID != 0 {
		startID = t.StartID
	} else {
		startID = DefaultStartID
	}

	var maxPics int64
	if t.MaxPics != 0 {
		maxPics = t.MaxPics
	} else {
		maxPics = DefaultMaxPics
	}

	// Technically an initial lookup of the created time of the provided Pic ID id needed.
	// TODO: decide if this is worth the extra DB call.
	rows, err := t.db.Query(
		"SELECT * FROM pics WHERE id <= ? ORDER BY created_time DESC LIMIT ?;",
		startID, maxPics)

	if err != nil {
		return err
	}

	defer rows.Close()

	columnNames, err := rows.Columns()
	if err != nil {
		return err
	}

	var pics []*Pic
	for rows.Next() {
		var p = new(Pic)
		if err := rows.Scan(p.ColumnPointers(columnNames)...); err != nil {
			return err
		}
		pics = append(pics, p)
	}

	if err := rows.Err(); err != nil {
		return err
	}

	t.Pics = pics

	return nil
}
<|endoftext|>


<commit_before>package readline

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/mattn/go-tty"
)

var FlushBeforeReadline = false

type Result int

const (
	CONTINUE Result = iota
	ENTER    Result = iota
	ABORT    Result = iota
	INTR     Result = iota
)

func (this Result) String() string {
	switch this {
	case CONTINUE:
		return "CONTINUE"
	case ENTER:
		return "ENTER"
	case ABORT:
		return "ABORT"
	case INTR:
		return "INTR"
	default:
		return "ERROR"
	}
}

type KeyFuncT interface {
	Call(ctx context.Context, buffer *Buffer) Result
}

type KeyGoFuncT struct {
	Func func(ctx context.Context, buffer *Buffer) Result
	Name string
}

func (this *KeyGoFuncT) Call(ctx context.Context, buffer *Buffer) Result {
	if this.Func == nil {
		return CONTINUE
	}
	return this.Func(ctx, buffer)
}

func (this KeyGoFuncT) String() string {
	return this.Name
}

var keyMap = map[string]KeyFuncT{
	name2char[K_CTRL_A]:    name2func(F_BEGINNING_OF_LINE),
	name2char[K_CTRL_B]:    name2func(F_BACKWARD_CHAR),
	name2char[K_BACKSPACE]: name2func(F_BACKWARD_DELETE_CHAR),
	name2char[K_CTRL_C]:    name2func(F_INTR),
	name2char[K_CTRL_D]:    name2func(F_DELETE_OR_ABORT),
	name2char[K_CTRL_E]:    name2func(F_END_OF_LINE),
	name2char[K_CTRL_F]:    name2func(F_FORWARD_CHAR),
	name2char[K_CTRL_H]:    name2func(F_BACKWARD_DELETE_CHAR),
	name2char[K_CTRL_K]:    name2func(F_KILL_LINE),
	name2char[K_CTRL_L]:    name2func(F_CLEAR_SCREEN),
	name2char[K_CTRL_M]:    name2func(F_ACCEPT_LINE),
	name2char[K_CTRL_R]:    name2func(F_ISEARCH_BACKWARD),
	name2char[K_CTRL_U]:    name2func(F_UNIX_LINE_DISCARD),
	name2char[K_CTRL_Y]:    name2func(F_YANK),
	name2char[K_DELETE]:    name2func(F_DELETE_CHAR),
	name2char[K_ENTER]:     name2func(F_ACCEPT_LINE),
	name2char[K_ESCAPE]:    name2func(F_KILL_WHOLE_LINE),
	name2char[K_CTRL_N]:    name2func(F_HISTORY_DOWN),
	name2char[K_CTRL_P]:    name2func(F_HISTORY_UP),
	name2char[K_CTRL_Q]:    name2func(F_QUOTED_INSERT),
	name2char[K_CTRL_T]:    name2func(F_SWAPCHAR),
	name2char[K_CTRL_V]:    name2func(F_QUOTED_INSERT),
	name2char[K_CTRL_W]:    name2func(F_UNIX_WORD_RUBOUT),
	name2char[K_CTRL]:      name2func(F_PASS),
	name2char[K_DELETE]:    name2func(F_DELETE_CHAR),
	name2char[K_END]:       name2func(F_END_OF_LINE),
	name2char[K_HOME]:      name2func(F_BEGINNING_OF_LINE),
	name2char[K_LEFT]:      name2func(F_BACKWARD_CHAR),
	name2char[K_RIGHT]:     name2func(F_FORWARD_CHAR),
	name2char[K_SHIFT]:     name2func(F_PASS),
	name2char[K_DOWN]:      name2func(F_HISTORY_DOWN),
	name2char[K_UP]:        name2func(F_HISTORY_UP),
	name2char[K_ALT_V]:     name2func(F_YANK),
	name2char[K_ALT_Y]:     name2func(F_YANK_WITH_QUOTE),
}

func normWord(src string) string {
	return strings.Replace(strings.ToUpper(src), "-", "_", -1)
}

func BindKeyFunc(keyName string, funcValue KeyFuncT) error {
	keyName_ := normWord(keyName)
	if charValue, charOk := name2char[keyName_]; charOk {
		keyMap[charValue] = funcValue
		return nil
	}
	return fmt.Errorf("%s: no such keyname", keyName)
}

func BindKeyClosure(name string, f func(context.Context, *Buffer) Result) error {
	return BindKeyFunc(name, &KeyGoFuncT{Func: f, Name: "annonymous"})
}

func GetBindKey(keyName string) KeyFuncT {
	keyName_ := normWord(keyName)
	if charValue, charOk := name2char[keyName_]; charOk {
		return keyMap[charValue]
	} else {
		return nil
	}
}

func GetFunc(funcName string) (KeyFuncT, error) {
	rc := name2func(normWord(funcName))
	if rc != nil {
		return rc, nil
	} else {
		return nil, fmt.Errorf("%s: not found in the function-list", funcName)
	}
}

func BindKeySymbol(keyName, funcName string) error {
	funcValue := name2func(normWord(funcName))
	if funcValue == nil {
		return fmt.Errorf("%s: no such function.", funcName)
	}
	return BindKeyFunc(keyName, funcValue)
}

type EmptyHistory struct{}

func (this *EmptyHistory) Len() int      { return 0 }
func (this *EmptyHistory) At(int) string { return "" }

const (
	ansiCursorOff = "\x1B[?25l"
	ansiCursorOn  = "\x1B[?25h\x1B[s\x1B[u"
)

var CtrlC = errors.New("^C")

func getKey(tty1 *tty.TTY) (string, error) {
	var buffer strings.Builder
	for {
		r, err := tty1.ReadRune()
		if err != nil {
			return "", err
		}
		if r == 0 {
			continue
		}
		buffer.WriteRune(r)
		if !tty1.Buffered() {
			return buffer.String(), nil
		}
	}
}

var mu sync.Mutex

// Call LineEditor
// - ENTER typed -> returns TEXT and nil
// - CTRL-C typed -> returns "" and readline.CtrlC
// - CTRL-D typed -> returns "" and io.EOF
func (session *Editor) ReadLine(ctx context.Context) (string, error) {
	if session.Writer == nil {
		panic("readline.Editor.Writer is not set. Set an instance such as go-colorable.NewColorableStdout()")
	}
	if session.Out == nil {
		session.Out = bufio.NewWriter(session.Writer)
	}
	defer func() {
		session.Out.WriteString(ansiCursorOn)
		session.Out.Flush()
	}()

	if session.Prompt == nil {
		session.Prompt = func() (int, error) {
			session.Out.WriteString("\n> ")
			return 2, nil
		}
	}
	if session.History == nil {
		session.History = new(EmptyHistory)
	}
	this := Buffer{
		Editor:         session,
		Buffer:         make([]rune, 20),
		HistoryPointer: session.History.Len(),
	}

	tty1, err := tty.Open()
	if err != nil {
		return "", fmt.Errorf("go-tty.Open: %s", err.Error())
	}
	this.TTY = tty1
	defer tty1.Close()

	this.TermWidth, _, err = tty1.Size()
	if err != nil {
		return "", fmt.Errorf("go-tty.Size: %s", err.Error())
	}

	var err1 error
	this.TopColumn, err1 = session.Prompt()
	if err1 != nil {
		// unable to get prompt-string.
		fmt.Fprintf(this.Out, "%s\n$ ", err1.Error())
		this.TopColumn = 2
	} else if this.TopColumn >= this.TermWidth-3 {
		// ViewWidth is too narrow to edit.
		io.WriteString(this.Out, "\n")
		this.TopColumn = 0
	}
	this.InsertString(0, session.Default)
	if this.Cursor > this.Length {
		this.Cursor = this.Length
	}
	this.RepaintAfterPrompt()

	cursorOnSwitch := false

	ws := tty1.SIGWINCH()
	go func(lastw int) {
		for ws1 := range ws {
			w := ws1.W
			if lastw != w {
				mu.Lock()
				this.TermWidth = w
				fmt.Fprintf(this.Out, "\x1B[%dG", this.TopColumn+1)
				this.RepaintAfterPrompt()
				mu.Unlock()
				lastw = w
			}
		}
	}(this.TermWidth)

	for {
		mu.Lock()
		if !cursorOnSwitch {
			io.WriteString(this.Out, ansiCursorOn)
			cursorOnSwitch = true
		}
		this.Out.Flush()

		mu.Unlock()
		key1, err := getKey(tty1)
		if err != nil {
			return "", err
		}
		mu.Lock()
		f, ok := keyMap[key1]
		if !ok {
			f = &KeyGoFuncT{
				Func: func(ctx context.Context, this *Buffer) Result {
					return keyFuncInsertSelf(ctx, this, key1)
				},
				Name: key1,
			}
		}

		if fg, ok := f.(*KeyGoFuncT); !ok || fg.Func != nil {
			io.WriteString(this.Out, ansiCursorOff)
			cursorOnSwitch = false
		}
		rc := f.Call(ctx, &this)
		if rc != CONTINUE {
			this.Out.WriteByte('\n')
			if !cursorOnSwitch {
				io.WriteString(this.Out, ansiCursorOn)
			}
			this.Out.Flush()
			result := this.String()
			mu.Unlock()
			if rc == ENTER {
				return result, nil
			} else if rc == INTR {
				return result, CtrlC
			} else {
				return result, io.EOF
			}
		}
		mu.Unlock()
	}
}
<commit_msg>On linux, fix the bug that cursor position got wrong via input method editor.<commit_after>package readline

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/mattn/go-tty"
)

var FlushBeforeReadline = false

type Result int

const (
	CONTINUE Result = iota
	ENTER    Result = iota
	ABORT    Result = iota
	INTR     Result = iota
)

func (this Result) String() string {
	switch this {
	case CONTINUE:
		return "CONTINUE"
	case ENTER:
		return "ENTER"
	case ABORT:
		return "ABORT"
	case INTR:
		return "INTR"
	default:
		return "ERROR"
	}
}

type KeyFuncT interface {
	Call(ctx context.Context, buffer *Buffer) Result
}

type KeyGoFuncT struct {
	Func func(ctx context.Context, buffer *Buffer) Result
	Name string
}

func (this *KeyGoFuncT) Call(ctx context.Context, buffer *Buffer) Result {
	if this.Func == nil {
		return CONTINUE
	}
	return this.Func(ctx, buffer)
}

func (this KeyGoFuncT) String() string {
	return this.Name
}

var keyMap = map[string]KeyFuncT{
	name2char[K_CTRL_A]:    name2func(F_BEGINNING_OF_LINE),
	name2char[K_CTRL_B]:    name2func(F_BACKWARD_CHAR),
	name2char[K_BACKSPACE]: name2func(F_BACKWARD_DELETE_CHAR),
	name2char[K_CTRL_C]:    name2func(F_INTR),
	name2char[K_CTRL_D]:    name2func(F_DELETE_OR_ABORT),
	name2char[K_CTRL_E]:    name2func(F_END_OF_LINE),
	name2char[K_CTRL_F]:    name2func(F_FORWARD_CHAR),
	name2char[K_CTRL_H]:    name2func(F_BACKWARD_DELETE_CHAR),
	name2char[K_CTRL_K]:    name2func(F_KILL_LINE),
	name2char[K_CTRL_L]:    name2func(F_CLEAR_SCREEN),
	name2char[K_CTRL_M]:    name2func(F_ACCEPT_LINE),
	name2char[K_CTRL_R]:    name2func(F_ISEARCH_BACKWARD),
	name2char[K_CTRL_U]:    name2func(F_UNIX_LINE_DISCARD),
	name2char[K_CTRL_Y]:    name2func(F_YANK),
	name2char[K_DELETE]:    name2func(F_DELETE_CHAR),
	name2char[K_ENTER]:     name2func(F_ACCEPT_LINE),
	name2char[K_ESCAPE]:    name2func(F_KILL_WHOLE_LINE),
	name2char[K_CTRL_N]:    name2func(F_HISTORY_DOWN),
	name2char[K_CTRL_P]:    name2func(F_HISTORY_UP),
	name2char[K_CTRL_Q]:    name2func(F_QUOTED_INSERT),
	name2char[K_CTRL_T]:    name2func(F_SWAPCHAR),
	name2char[K_CTRL_V]:    name2func(F_QUOTED_INSERT),
	name2char[K_CTRL_W]:    name2func(F_UNIX_WORD_RUBOUT),
	name2char[K_CTRL]:      name2func(F_PASS),
	name2char[K_DELETE]:    name2func(F_DELETE_CHAR),
	name2char[K_END]:       name2func(F_END_OF_LINE),
	name2char[K_HOME]:      name2func(F_BEGINNING_OF_LINE),
	name2char[K_LEFT]:      name2func(F_BACKWARD_CHAR),
	name2char[K_RIGHT]:     name2func(F_FORWARD_CHAR),
	name2char[K_SHIFT]:     name2func(F_PASS),
	name2char[K_DOWN]:      name2func(F_HISTORY_DOWN),
	name2char[K_UP]:        name2func(F_HISTORY_UP),
	name2char[K_ALT_V]:     name2func(F_YANK),
	name2char[K_ALT_Y]:     name2func(F_YANK_WITH_QUOTE),
}

func normWord(src string) string {
	return strings.Replace(strings.ToUpper(src), "-", "_", -1)
}

func BindKeyFunc(keyName string, funcValue KeyFuncT) error {
	keyName_ := normWord(keyName)
	if charValue, charOk := name2char[keyName_]; charOk {
		keyMap[charValue] = funcValue
		return nil
	}
	return fmt.Errorf("%s: no such keyname", keyName)
}

func BindKeyClosure(name string, f func(context.Context, *Buffer) Result) error {
	return BindKeyFunc(name, &KeyGoFuncT{Func: f, Name: "annonymous"})
}

func GetBindKey(keyName string) KeyFuncT {
	keyName_ := normWord(keyName)
	if charValue, charOk := name2char[keyName_]; charOk {
		return keyMap[charValue]
	} else {
		return nil
	}
}

func GetFunc(funcName string) (KeyFuncT, error) {
	rc := name2func(normWord(funcName))
	if rc != nil {
		return rc, nil
	} else {
		return nil, fmt.Errorf("%s: not found in the function-list", funcName)
	}
}

func BindKeySymbol(keyName, funcName string) error {
	funcValue := name2func(normWord(funcName))
	if funcValue == nil {
		return fmt.Errorf("%s: no such function.", funcName)
	}
	return BindKeyFunc(keyName, funcValue)
}

type EmptyHistory struct{}

func (this *EmptyHistory) Len() int      { return 0 }
func (this *EmptyHistory) At(int) string { return "" }

const (
	ansiCursorOff = "\x1B[?25l"
	ansiCursorOn  = "\x1B[?25h\x1B[s\x1B[u"
)

var CtrlC = errors.New("^C")

func getKey(tty1 *tty.TTY) (string, error) {
	var buffer strings.Builder
	escape := false
	for {
		r, err := tty1.ReadRune()
		if err != nil {
			return "", err
		}
		if r == 0 {
			continue
		}
		buffer.WriteRune(r)
		if r == '\x1B' {
			escape = true
		}
		if !(escape && tty1.Buffered()) {
			return buffer.String(), nil
		}
	}
}

var mu sync.Mutex

// Call LineEditor
// - ENTER typed -> returns TEXT and nil
// - CTRL-C typed -> returns "" and readline.CtrlC
// - CTRL-D typed -> returns "" and io.EOF
func (session *Editor) ReadLine(ctx context.Context) (string, error) {
	if session.Writer == nil {
		panic("readline.Editor.Writer is not set. Set an instance such as go-colorable.NewColorableStdout()")
	}
	if session.Out == nil {
		session.Out = bufio.NewWriter(session.Writer)
	}
	defer func() {
		session.Out.WriteString(ansiCursorOn)
		session.Out.Flush()
	}()

	if session.Prompt == nil {
		session.Prompt = func() (int, error) {
			session.Out.WriteString("\n> ")
			return 2, nil
		}
	}
	if session.History == nil {
		session.History = new(EmptyHistory)
	}
	this := Buffer{
		Editor:         session,
		Buffer:         make([]rune, 20),
		HistoryPointer: session.History.Len(),
	}

	tty1, err := tty.Open()
	if err != nil {
		return "", fmt.Errorf("go-tty.Open: %s", err.Error())
	}
	this.TTY = tty1
	defer tty1.Close()

	this.TermWidth, _, err = tty1.Size()
	if err != nil {
		return "", fmt.Errorf("go-tty.Size: %s", err.Error())
	}

	var err1 error
	this.TopColumn, err1 = session.Prompt()
	if err1 != nil {
		// unable to get prompt-string.
		fmt.Fprintf(this.Out, "%s\n$ ", err1.Error())
		this.TopColumn = 2
	} else if this.TopColumn >= this.TermWidth-3 {
		// ViewWidth is too narrow to edit.
		io.WriteString(this.Out, "\n")
		this.TopColumn = 0
	}
	this.InsertString(0, session.Default)
	if this.Cursor > this.Length {
		this.Cursor = this.Length
	}
	this.RepaintAfterPrompt()

	cursorOnSwitch := false

	ws := tty1.SIGWINCH()
	go func(lastw int) {
		for ws1 := range ws {
			w := ws1.W
			if lastw != w {
				mu.Lock()
				this.TermWidth = w
				fmt.Fprintf(this.Out, "\x1B[%dG", this.TopColumn+1)
				this.RepaintAfterPrompt()
				mu.Unlock()
				lastw = w
			}
		}
	}(this.TermWidth)

	for {
		mu.Lock()
		if !cursorOnSwitch {
			io.WriteString(this.Out, ansiCursorOn)
			cursorOnSwitch = true
		}
		this.Out.Flush()

		mu.Unlock()
		key1, err := getKey(tty1)
		if err != nil {
			return "", err
		}
		mu.Lock()
		f, ok := keyMap[key1]
		if !ok {
			f = &KeyGoFuncT{
				Func: func(ctx context.Context, this *Buffer) Result {
					return keyFuncInsertSelf(ctx, this, key1)
				},
				Name: key1,
			}
		}

		if fg, ok := f.(*KeyGoFuncT); !ok || fg.Func != nil {
			io.WriteString(this.Out, ansiCursorOff)
			cursorOnSwitch = false
		}
		rc := f.Call(ctx, &this)
		if rc != CONTINUE {
			this.Out.WriteByte('\n')
			if !cursorOnSwitch {
				io.WriteString(this.Out, ansiCursorOn)
			}
			this.Out.Flush()
			result := this.String()
			mu.Unlock()
			if rc == ENTER {
				return result, nil
			} else if rc == INTR {
				return result, CtrlC
			} else {
				return result, io.EOF
			}
		}
		mu.Unlock()
	}
}
<|endoftext|>


<commit_before>package ring

import (
	"sort"
)

type rebalanceContextImpl struct {
	builder                         *ringBuilderImpl
	first                           bool
	nodeIndex2DesiredPartitionCount []int32
	nodeIndexesByDesire             []int32
	nodeIndex2Used                  []bool
	tierCount                       int
	tier2TierIDs                    [][]*tierIDImpl
	tier2NodeIndex2TierID           [][]*tierIDImpl
}

func newRebalanceContext(builder *ringBuilderImpl) *rebalanceContextImpl {
	rebalanceContext := &rebalanceContextImpl{builder: builder}
	rebalanceContext.initTierCount()
	rebalanceContext.initNodeIndex2DesiredPartitionCount()
	rebalanceContext.initTier2NodeIndex2TierID()
	return rebalanceContext
}

func (rebalanceContext *rebalanceContextImpl) initTierCount() {
	rebalanceContext.tierCount = 0
	for _, node := range rebalanceContext.builder.nodes {
		if !node.Active() {
			continue
		}
		nodeTierCount := len(node.TierValues())
		if nodeTierCount > rebalanceContext.tierCount {
			rebalanceContext.tierCount = nodeTierCount
		}
	}
}

func (rebalanceContext *rebalanceContextImpl) initNodeIndex2DesiredPartitionCount() {
	totalCapacity := uint64(0)
	for _, node := range rebalanceContext.builder.nodes {
		if node.Active() {
			totalCapacity += node.Capacity()
		}
	}
	nodeIndex2PartitionCount := make([]int32, len(rebalanceContext.builder.nodes))
	rebalanceContext.first = true
	for _, partition2NodeIndex := range rebalanceContext.builder.replica2Partition2NodeIndex {
		for _, nodeIndex := range partition2NodeIndex {
			if nodeIndex >= 0 {
				nodeIndex2PartitionCount[nodeIndex]++
				rebalanceContext.first = false
			}
		}
	}
	rebalanceContext.nodeIndex2DesiredPartitionCount = make([]int32, len(rebalanceContext.builder.nodes))
	allPartitionsCount := len(rebalanceContext.builder.replica2Partition2NodeIndex) * len(rebalanceContext.builder.replica2Partition2NodeIndex[0])
	for nodeIndex, node := range rebalanceContext.builder.nodes {
		if node.Active() {
			rebalanceContext.nodeIndex2DesiredPartitionCount[nodeIndex] = int32(float64(node.Capacity())/float64(totalCapacity)*float64(allPartitionsCount)+0.5) - nodeIndex2PartitionCount[nodeIndex]
		} else {
			rebalanceContext.nodeIndex2DesiredPartitionCount[nodeIndex] = -2147483648
		}
	}
	rebalanceContext.nodeIndexesByDesire = make([]int32, 0, len(rebalanceContext.builder.nodes))
	for nodeIndex, node := range rebalanceContext.builder.nodes {
		if node.Active() {
			rebalanceContext.nodeIndexesByDesire = append(rebalanceContext.nodeIndexesByDesire, int32(nodeIndex))
		}
	}
	sort.Sort(&nodeIndexByDesireSorterImpl{
		nodeIndexesByDesire:             rebalanceContext.nodeIndexesByDesire,
		nodeIndex2DesiredPartitionCount: rebalanceContext.nodeIndex2DesiredPartitionCount,
	})
}

func (rebalanceContext *rebalanceContextImpl) initTier2NodeIndex2TierID() {
	rebalanceContext.tier2NodeIndex2TierID = make([][]*tierIDImpl, rebalanceContext.tierCount)
	rebalanceContext.tier2TierIDs = make([][]*tierIDImpl, rebalanceContext.tierCount)
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		rebalanceContext.tier2NodeIndex2TierID[tier] = make([]*tierIDImpl, len(rebalanceContext.builder.nodes))
		rebalanceContext.tier2TierIDs[tier] = make([]*tierIDImpl, 0)
	}
	for nodeIndex, node := range rebalanceContext.builder.nodes {
		nodeTierValues := node.TierValues()
		for tier := 1; tier < rebalanceContext.tierCount; tier++ {
			var tierID *tierIDImpl
			for _, candidateTierID := range rebalanceContext.tier2TierIDs[tier] {
				tierID = candidateTierID
				for valueIndex := 0; valueIndex < rebalanceContext.tierCount-tier; valueIndex++ {
					value := 0
					if valueIndex+tier < len(nodeTierValues) {
						value = nodeTierValues[valueIndex+tier]
					}
					if tierID.values[valueIndex] != value {
						tierID = nil
						break
					}
				}
				if tierID != nil {
					break
				}
			}
			if tierID == nil {
				tierID = &tierIDImpl{values: make([]int, rebalanceContext.tierCount-tier), nodeIndexesByDesire: []int32{int32(nodeIndex)}}
				for valueIndex := 0; valueIndex < rebalanceContext.tierCount-tier; valueIndex++ {
					value := 0
					if valueIndex+tier < len(nodeTierValues) {
						value = nodeTierValues[valueIndex+tier]
					}
					tierID.values[valueIndex] = value
				}
				rebalanceContext.tier2TierIDs[tier] = append(rebalanceContext.tier2TierIDs[tier], tierID)
			} else {
				tierID.nodeIndexesByDesire = append(tierID.nodeIndexesByDesire, int32(nodeIndex))
			}
			rebalanceContext.tier2NodeIndex2TierID[tier][int32(nodeIndex)] = tierID
		}
	}
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		for _, tierID := range rebalanceContext.tier2TierIDs[tier] {
			sort.Sort(&nodeIndexByDesireSorterImpl{
				nodeIndexesByDesire:             tierID.nodeIndexesByDesire,
				nodeIndex2DesiredPartitionCount: rebalanceContext.nodeIndex2DesiredPartitionCount,
			})
		}
	}
}

func (rebalanceContext *rebalanceContextImpl) rebalance() bool {
	if rebalanceContext.first {
		rebalanceContext.firstRebalance()
		return true
	}
	return rebalanceContext.subsequentRebalance()
}

// firstRebalance is much simpler than what we have to do to rebalance existing
// assignments. Here, we just assign each partition in order, giving each
// replica of that partition to the next most-desired node, keeping in mind
// tier separation preferences.
func (rebalanceContext *rebalanceContextImpl) firstRebalance() {
	replicaCount := len(rebalanceContext.builder.replica2Partition2NodeIndex)
	partitionCount := len(rebalanceContext.builder.replica2Partition2NodeIndex[0])
	// We track the other nodes and tiers we've assigned partition replicas to
	// so that we can try to avoid assigning further replicas to similar nodes.
	otherNodeIndexes := make([]int32, replicaCount)
	rebalanceContext.nodeIndex2Used = make([]bool, len(rebalanceContext.builder.nodes))
	tier2OtherTierIDs := make([][]*tierIDImpl, rebalanceContext.tierCount)
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		tier2OtherTierIDs[tier] = make([]*tierIDImpl, replicaCount)
	}
	for partition := 0; partition < partitionCount; partition++ {
		for replica := 0; replica < replicaCount; replica++ {
			if otherNodeIndexes[replica] != -1 {
				rebalanceContext.nodeIndex2Used[otherNodeIndexes[replica]] = false
			}
			otherNodeIndexes[replica] = -1
		}
		for tier := 1; tier < rebalanceContext.tierCount; tier++ {
			for replica := 0; replica < replicaCount; replica++ {
				if tier2OtherTierIDs[tier][replica] != nil {
					tier2OtherTierIDs[tier][replica].used = false
				}
				tier2OtherTierIDs[tier][replica] = nil
			}
		}
		for replica := 0; replica < replicaCount; replica++ {
			nodeIndex := rebalanceContext.bestNodeIndex()
			rebalanceContext.builder.replica2Partition2NodeIndex[replica][partition] = nodeIndex
			rebalanceContext.decrementDesire(nodeIndex)
			rebalanceContext.nodeIndex2Used[nodeIndex] = true
			otherNodeIndexes[replica] = nodeIndex
			for tier := 1; tier < rebalanceContext.tierCount; tier++ {
				tierID := rebalanceContext.tier2NodeIndex2TierID[tier][nodeIndex]
				tierID.used = true
				tier2OtherTierIDs[tier][replica] = tierID
			}
		}
	}
}

// subsequentRebalance is much more complicated than firstRebalance.
// First we'll reassign any partition replicas assigned to nodes with a
// weight less than 0, as this indicates a deleted node.
// Then we'll attempt to reassign partition replicas that are at extremely
// high risk because they're on the exact same node.
// Next we'll attempt to reassign partition replicas that are at some risk
// because they are currently assigned within the same tier separation.
// Then, we'll attempt to reassign replicas within tiers to achieve better
// distribution, as usually such intra-tier movements are more efficient
// for users of the ring.
// Finally, one last pass will be done to reassign replicas to still
// underweight nodes.
func (rebalanceContext *rebalanceContextImpl) subsequentRebalance() bool {
	replicaCount := len(rebalanceContext.builder.replica2Partition2NodeIndex)
	partitionCount := len(rebalanceContext.builder.replica2Partition2NodeIndex[0])
	// We'll track how many times we can move replicas for a given partition;
	// we want to leave at least half a partition's replicas in place.
	movementsPerPartition := byte(replicaCount / 2)
	if movementsPerPartition < 1 {
		movementsPerPartition = 1
	}
	partition2MovementsLeft := make([]byte, partitionCount)
	for partition := 0; partition < partitionCount; partition++ {
		partition2MovementsLeft[partition] = movementsPerPartition
	}
	altered := false
	// First we'll reassign any partition replicas assigned to nodes with a
	// weight less than 0, as this indicates a deleted node.
	for deletedNodeIndex, deletedNode := range rebalanceContext.builder.nodes {
		if deletedNode.Active() {
			continue
		}
		for replica := 0; replica < replicaCount; replica++ {
			partition2NodeIndex := rebalanceContext.builder.replica2Partition2NodeIndex[replica]
			for partition := 0; partition < partitionCount; partition++ {
				if partition2NodeIndex[partition] != int32(deletedNodeIndex) {
					continue
				}
				// We track the other nodes and tiers we've assigned partition
				// replicas to so that we can try to avoid assigning further
				// replicas to similar nodes.
				otherNodeIndexes := make([]int32, replicaCount)
				rebalanceContext.nodeIndex2Used = make([]bool, len(rebalanceContext.builder.nodes))
				tier2OtherTierIDs := make([][]*tierIDImpl, rebalanceContext.tierCount)
				for tier := 1; tier < rebalanceContext.tierCount; tier++ {
					tier2OtherTierIDs[tier] = make([]*tierIDImpl, replicaCount)
				}
				for replicaB := 0; replicaB < replicaCount; replicaB++ {
					otherNodeIndexes[replicaB] = rebalanceContext.builder.replica2Partition2NodeIndex[replicaB][partition]
					for tier := 1; tier < rebalanceContext.tierCount; tier++ {
						tierID := rebalanceContext.tier2NodeIndex2TierID[tier][otherNodeIndexes[replicaB]]
						tierID.used = true
						tier2OtherTierIDs[tier][replicaB] = tierID
					}
				}
				nodeIndex := rebalanceContext.bestNodeIndex()
				partition2NodeIndex[partition] = nodeIndex
				altered = true
				rebalanceContext.decrementDesire(nodeIndex)
			}
		}
	}
	return altered
}

func (rebalanceContext *rebalanceContextImpl) bestNodeIndex() int32 {
	bestNodeIndex := int32(-1)
	bestNodeDesiredPartitionCount := ^int32(0)
	nodeIndex2DesiredPartitionCount := rebalanceContext.nodeIndex2DesiredPartitionCount
	var tierID *tierIDImpl
	var nodeIndex int32
	tier2TierIDs := rebalanceContext.tier2TierIDs
	for tier := rebalanceContext.tierCount - 1; tier > 0; tier-- {
		// We will go through all tierIDs for a tier to get the
		// best node at that tier.
		for _, tierID = range tier2TierIDs[tier] {
			if !tierID.used {
				nodeIndex = tierID.nodeIndexesByDesire[0]
				if bestNodeDesiredPartitionCount < nodeIndex2DesiredPartitionCount[nodeIndex] {
					bestNodeIndex = nodeIndex
					bestNodeDesiredPartitionCount = nodeIndex2DesiredPartitionCount[nodeIndex]
				}
			}
		}
		// If we found a node at this tier, we don't need to check the lower
		// tiers.
		if bestNodeIndex >= 0 {
			return bestNodeIndex
		}
	}
	// If we found no good higher tiered candidates, we'll have to just
	// take the node with the highest desire that hasn't already been
	// selected.
	for _, nodeIndex := range rebalanceContext.nodeIndexesByDesire {
		if !rebalanceContext.nodeIndex2Used[nodeIndex] {
			return nodeIndex
		}
	}
	// If we still found no good candidates, we'll have to just take the
	// node with the highest desire.
	return rebalanceContext.nodeIndexesByDesire[0]
}

func (rebalanceContext *rebalanceContextImpl) decrementDesire(nodeIndex int32) {
	nodeIndex2DesiredPartitionCount := rebalanceContext.nodeIndex2DesiredPartitionCount
	nodeIndexesByDesire := rebalanceContext.nodeIndexesByDesire
	scanDesiredPartitionCount := nodeIndex2DesiredPartitionCount[nodeIndex] - 1
	swapWith := 0
	hi := len(nodeIndexesByDesire)
	mid := 0
	for swapWith < hi {
		mid = (swapWith + hi) / 2
		if nodeIndex2DesiredPartitionCount[nodeIndexesByDesire[mid]] > scanDesiredPartitionCount {
			swapWith = mid + 1
		} else {
			hi = mid
		}
	}
	prev := swapWith
	if prev >= len(nodeIndexesByDesire) {
		prev--
	}
	swapWith--
	for nodeIndexesByDesire[prev] != nodeIndex {
		prev--
	}
	if prev != swapWith {
		nodeIndexesByDesire[prev], nodeIndexesByDesire[swapWith] = nodeIndexesByDesire[swapWith], nodeIndexesByDesire[prev]
	}
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		nodeIndexesByDesire := rebalanceContext.tier2NodeIndex2TierID[tier][nodeIndex].nodeIndexesByDesire
		swapWith = 0
		hi = len(nodeIndexesByDesire)
		mid = 0
		for swapWith < hi {
			mid = (swapWith + hi) / 2
			if nodeIndex2DesiredPartitionCount[nodeIndexesByDesire[mid]] > scanDesiredPartitionCount {
				swapWith = mid + 1
			} else {
				hi = mid
			}
		}
		prev = swapWith
		if prev >= len(nodeIndexesByDesire) {
			prev--
		}
		swapWith--
		for nodeIndexesByDesire[prev] != nodeIndex {
			prev--
		}
		if prev != swapWith {
			nodeIndexesByDesire[prev], nodeIndexesByDesire[swapWith] = nodeIndexesByDesire[swapWith], nodeIndexesByDesire[prev]
		}
	}
	nodeIndex2DesiredPartitionCount[nodeIndex]--
}

type tierIDImpl struct {
	values              []int
	nodeIndexesByDesire []int32
	used                bool
}
<commit_msg>Added some TODOs<commit_after>package ring

import (
	"sort"
)

type rebalanceContextImpl struct {
	builder                         *ringBuilderImpl
	first                           bool
	nodeIndex2DesiredPartitionCount []int32
	nodeIndexesByDesire             []int32
	nodeIndex2Used                  []bool
	tierCount                       int
	tier2TierIDs                    [][]*tierIDImpl
	tier2NodeIndex2TierID           [][]*tierIDImpl
}

func newRebalanceContext(builder *ringBuilderImpl) *rebalanceContextImpl {
	rebalanceContext := &rebalanceContextImpl{builder: builder}
	rebalanceContext.initTierCount()
	rebalanceContext.initNodeIndex2DesiredPartitionCount()
	rebalanceContext.initTier2NodeIndex2TierID()
	return rebalanceContext
}

func (rebalanceContext *rebalanceContextImpl) initTierCount() {
	rebalanceContext.tierCount = 0
	for _, node := range rebalanceContext.builder.nodes {
		if !node.Active() {
			continue
		}
		nodeTierCount := len(node.TierValues())
		if nodeTierCount > rebalanceContext.tierCount {
			rebalanceContext.tierCount = nodeTierCount
		}
	}
}

func (rebalanceContext *rebalanceContextImpl) initNodeIndex2DesiredPartitionCount() {
	totalCapacity := uint64(0)
	for _, node := range rebalanceContext.builder.nodes {
		if node.Active() {
			totalCapacity += node.Capacity()
		}
	}
	nodeIndex2PartitionCount := make([]int32, len(rebalanceContext.builder.nodes))
	rebalanceContext.first = true
	for _, partition2NodeIndex := range rebalanceContext.builder.replica2Partition2NodeIndex {
		for _, nodeIndex := range partition2NodeIndex {
			if nodeIndex >= 0 {
				nodeIndex2PartitionCount[nodeIndex]++
				rebalanceContext.first = false
			}
		}
	}
	rebalanceContext.nodeIndex2DesiredPartitionCount = make([]int32, len(rebalanceContext.builder.nodes))
	allPartitionsCount := len(rebalanceContext.builder.replica2Partition2NodeIndex) * len(rebalanceContext.builder.replica2Partition2NodeIndex[0])
	for nodeIndex, node := range rebalanceContext.builder.nodes {
		if node.Active() {
			rebalanceContext.nodeIndex2DesiredPartitionCount[nodeIndex] = int32(float64(node.Capacity())/float64(totalCapacity)*float64(allPartitionsCount)+0.5) - nodeIndex2PartitionCount[nodeIndex]
		} else {
			rebalanceContext.nodeIndex2DesiredPartitionCount[nodeIndex] = -2147483648
		}
	}
	rebalanceContext.nodeIndexesByDesire = make([]int32, 0, len(rebalanceContext.builder.nodes))
	for nodeIndex, node := range rebalanceContext.builder.nodes {
		if node.Active() {
			rebalanceContext.nodeIndexesByDesire = append(rebalanceContext.nodeIndexesByDesire, int32(nodeIndex))
		}
	}
	sort.Sort(&nodeIndexByDesireSorterImpl{
		nodeIndexesByDesire:             rebalanceContext.nodeIndexesByDesire,
		nodeIndex2DesiredPartitionCount: rebalanceContext.nodeIndex2DesiredPartitionCount,
	})
}

func (rebalanceContext *rebalanceContextImpl) initTier2NodeIndex2TierID() {
	rebalanceContext.tier2NodeIndex2TierID = make([][]*tierIDImpl, rebalanceContext.tierCount)
	rebalanceContext.tier2TierIDs = make([][]*tierIDImpl, rebalanceContext.tierCount)
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		rebalanceContext.tier2NodeIndex2TierID[tier] = make([]*tierIDImpl, len(rebalanceContext.builder.nodes))
		rebalanceContext.tier2TierIDs[tier] = make([]*tierIDImpl, 0)
	}
	for nodeIndex, node := range rebalanceContext.builder.nodes {
		nodeTierValues := node.TierValues()
		for tier := 1; tier < rebalanceContext.tierCount; tier++ {
			var tierID *tierIDImpl
			for _, candidateTierID := range rebalanceContext.tier2TierIDs[tier] {
				tierID = candidateTierID
				for valueIndex := 0; valueIndex < rebalanceContext.tierCount-tier; valueIndex++ {
					value := 0
					if valueIndex+tier < len(nodeTierValues) {
						value = nodeTierValues[valueIndex+tier]
					}
					if tierID.values[valueIndex] != value {
						tierID = nil
						break
					}
				}
				if tierID != nil {
					break
				}
			}
			if tierID == nil {
				tierID = &tierIDImpl{values: make([]int, rebalanceContext.tierCount-tier), nodeIndexesByDesire: []int32{int32(nodeIndex)}}
				for valueIndex := 0; valueIndex < rebalanceContext.tierCount-tier; valueIndex++ {
					value := 0
					if valueIndex+tier < len(nodeTierValues) {
						value = nodeTierValues[valueIndex+tier]
					}
					tierID.values[valueIndex] = value
				}
				rebalanceContext.tier2TierIDs[tier] = append(rebalanceContext.tier2TierIDs[tier], tierID)
			} else {
				tierID.nodeIndexesByDesire = append(tierID.nodeIndexesByDesire, int32(nodeIndex))
			}
			rebalanceContext.tier2NodeIndex2TierID[tier][int32(nodeIndex)] = tierID
		}
	}
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		for _, tierID := range rebalanceContext.tier2TierIDs[tier] {
			sort.Sort(&nodeIndexByDesireSorterImpl{
				nodeIndexesByDesire:             tierID.nodeIndexesByDesire,
				nodeIndex2DesiredPartitionCount: rebalanceContext.nodeIndex2DesiredPartitionCount,
			})
		}
	}
}

func (rebalanceContext *rebalanceContextImpl) rebalance() bool {
	if rebalanceContext.first {
		rebalanceContext.firstRebalance()
		return true
	}
	return rebalanceContext.subsequentRebalance()
}

// firstRebalance is much simpler than what we have to do to rebalance existing
// assignments. Here, we just assign each partition in order, giving each
// replica of that partition to the next most-desired node, keeping in mind
// tier separation preferences.
func (rebalanceContext *rebalanceContextImpl) firstRebalance() {
	replicaCount := len(rebalanceContext.builder.replica2Partition2NodeIndex)
	partitionCount := len(rebalanceContext.builder.replica2Partition2NodeIndex[0])
	// We track the other nodes and tiers we've assigned partition replicas to
	// so that we can try to avoid assigning further replicas to similar nodes.
	otherNodeIndexes := make([]int32, replicaCount)
	rebalanceContext.nodeIndex2Used = make([]bool, len(rebalanceContext.builder.nodes))
	tier2OtherTierIDs := make([][]*tierIDImpl, rebalanceContext.tierCount)
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		tier2OtherTierIDs[tier] = make([]*tierIDImpl, replicaCount)
	}
	for partition := 0; partition < partitionCount; partition++ {
		for replica := 0; replica < replicaCount; replica++ {
			if otherNodeIndexes[replica] != -1 {
				rebalanceContext.nodeIndex2Used[otherNodeIndexes[replica]] = false
			}
			otherNodeIndexes[replica] = -1
		}
		for tier := 1; tier < rebalanceContext.tierCount; tier++ {
			for replica := 0; replica < replicaCount; replica++ {
				if tier2OtherTierIDs[tier][replica] != nil {
					tier2OtherTierIDs[tier][replica].used = false
				}
				tier2OtherTierIDs[tier][replica] = nil
			}
		}
		for replica := 0; replica < replicaCount; replica++ {
			nodeIndex := rebalanceContext.bestNodeIndex()
			rebalanceContext.builder.replica2Partition2NodeIndex[replica][partition] = nodeIndex
			rebalanceContext.decrementDesire(nodeIndex)
			rebalanceContext.nodeIndex2Used[nodeIndex] = true
			otherNodeIndexes[replica] = nodeIndex
			for tier := 1; tier < rebalanceContext.tierCount; tier++ {
				tierID := rebalanceContext.tier2NodeIndex2TierID[tier][nodeIndex]
				tierID.used = true
				tier2OtherTierIDs[tier][replica] = tierID
			}
		}
	}
}

// subsequentRebalance is much more complicated than firstRebalance.
// First we'll reassign any partition replicas assigned to nodes with a
// weight less than 0, as this indicates a deleted node.
// Then we'll attempt to reassign partition replicas that are at extremely
// high risk because they're on the exact same node.
// Next we'll attempt to reassign partition replicas that are at some risk
// because they are currently assigned within the same tier separation.
// Then, we'll attempt to reassign replicas within tiers to achieve better
// distribution, as usually such intra-tier movements are more efficient
// for users of the ring.
// Finally, one last pass will be done to reassign replicas to still
// underweight nodes.
func (rebalanceContext *rebalanceContextImpl) subsequentRebalance() bool {
	replicaCount := len(rebalanceContext.builder.replica2Partition2NodeIndex)
	partitionCount := len(rebalanceContext.builder.replica2Partition2NodeIndex[0])
	// We'll track how many times we can move replicas for a given partition;
	// we want to leave at least half a partition's replicas in place.
	movementsPerPartition := byte(replicaCount / 2)
	if movementsPerPartition < 1 {
		movementsPerPartition = 1
	}
	partition2MovementsLeft := make([]byte, partitionCount)
	for partition := 0; partition < partitionCount; partition++ {
		partition2MovementsLeft[partition] = movementsPerPartition
	}
	altered := false
	// First we'll reassign any partition replicas assigned to nodes with a
	// weight less than 0, as this indicates a deleted node.
	for deletedNodeIndex, deletedNode := range rebalanceContext.builder.nodes {
		if deletedNode.Active() {
			continue
		}
		for replica := 0; replica < replicaCount; replica++ {
			partition2NodeIndex := rebalanceContext.builder.replica2Partition2NodeIndex[replica]
			for partition := 0; partition < partitionCount; partition++ {
				if partition2NodeIndex[partition] != int32(deletedNodeIndex) {
					continue
				}
				// We track the other nodes and tiers we've assigned partition
				// replicas to so that we can try to avoid assigning further
				// replicas to similar nodes.
				otherNodeIndexes := make([]int32, replicaCount)
				rebalanceContext.nodeIndex2Used = make([]bool, len(rebalanceContext.builder.nodes))
				tier2OtherTierIDs := make([][]*tierIDImpl, rebalanceContext.tierCount)
				for tier := 1; tier < rebalanceContext.tierCount; tier++ {
					tier2OtherTierIDs[tier] = make([]*tierIDImpl, replicaCount)
				}
				for replicaB := 0; replicaB < replicaCount; replicaB++ {
					otherNodeIndexes[replicaB] = rebalanceContext.builder.replica2Partition2NodeIndex[replicaB][partition]
					for tier := 1; tier < rebalanceContext.tierCount; tier++ {
						tierID := rebalanceContext.tier2NodeIndex2TierID[tier][otherNodeIndexes[replicaB]]
						tierID.used = true
						tier2OtherTierIDs[tier][replicaB] = tierID
					}
				}
				nodeIndex := rebalanceContext.bestNodeIndex()
				partition2NodeIndex[partition] = nodeIndex
				altered = true
				rebalanceContext.decrementDesire(nodeIndex)
			}
		}
	}
    // TODO: Then we attempt to reassign at risk partitions. Partitions are
    // considered at risk if they have multiple replicas on the same node or
    // within the same tier separation.

    // TODO: Attempt to reassign replicas within tiers, from innermost tier to
    // outermost, as usually such movements are more efficient for users of the
    // ring. We do this by selecting the most needy node, and then look for
    // overweight nodes in the same tier to steal replicas from.

    // TODO: Lastly, we try to reassign replicas from overweight nodes to
    // underweight ones.
	return altered
}

func (rebalanceContext *rebalanceContextImpl) bestNodeIndex() int32 {
	bestNodeIndex := int32(-1)
	bestNodeDesiredPartitionCount := ^int32(0)
	nodeIndex2DesiredPartitionCount := rebalanceContext.nodeIndex2DesiredPartitionCount
	var tierID *tierIDImpl
	var nodeIndex int32
	tier2TierIDs := rebalanceContext.tier2TierIDs
	for tier := rebalanceContext.tierCount - 1; tier > 0; tier-- {
		// We will go through all tierIDs for a tier to get the
		// best node at that tier.
		for _, tierID = range tier2TierIDs[tier] {
			if !tierID.used {
				nodeIndex = tierID.nodeIndexesByDesire[0]
				if bestNodeDesiredPartitionCount < nodeIndex2DesiredPartitionCount[nodeIndex] {
					bestNodeIndex = nodeIndex
					bestNodeDesiredPartitionCount = nodeIndex2DesiredPartitionCount[nodeIndex]
				}
			}
		}
		// If we found a node at this tier, we don't need to check the lower
		// tiers.
		if bestNodeIndex >= 0 {
			return bestNodeIndex
		}
	}
	// If we found no good higher tiered candidates, we'll have to just
	// take the node with the highest desire that hasn't already been
	// selected.
	for _, nodeIndex := range rebalanceContext.nodeIndexesByDesire {
		if !rebalanceContext.nodeIndex2Used[nodeIndex] {
			return nodeIndex
		}
	}
	// If we still found no good candidates, we'll have to just take the
	// node with the highest desire.
	return rebalanceContext.nodeIndexesByDesire[0]
}

func (rebalanceContext *rebalanceContextImpl) decrementDesire(nodeIndex int32) {
	nodeIndex2DesiredPartitionCount := rebalanceContext.nodeIndex2DesiredPartitionCount
	nodeIndexesByDesire := rebalanceContext.nodeIndexesByDesire
	scanDesiredPartitionCount := nodeIndex2DesiredPartitionCount[nodeIndex] - 1
	swapWith := 0
	hi := len(nodeIndexesByDesire)
	mid := 0
	for swapWith < hi {
		mid = (swapWith + hi) / 2
		if nodeIndex2DesiredPartitionCount[nodeIndexesByDesire[mid]] > scanDesiredPartitionCount {
			swapWith = mid + 1
		} else {
			hi = mid
		}
	}
	prev := swapWith
	if prev >= len(nodeIndexesByDesire) {
		prev--
	}
	swapWith--
	for nodeIndexesByDesire[prev] != nodeIndex {
		prev--
	}
	if prev != swapWith {
		nodeIndexesByDesire[prev], nodeIndexesByDesire[swapWith] = nodeIndexesByDesire[swapWith], nodeIndexesByDesire[prev]
	}
	for tier := 1; tier < rebalanceContext.tierCount; tier++ {
		nodeIndexesByDesire := rebalanceContext.tier2NodeIndex2TierID[tier][nodeIndex].nodeIndexesByDesire
		swapWith = 0
		hi = len(nodeIndexesByDesire)
		mid = 0
		for swapWith < hi {
			mid = (swapWith + hi) / 2
			if nodeIndex2DesiredPartitionCount[nodeIndexesByDesire[mid]] > scanDesiredPartitionCount {
				swapWith = mid + 1
			} else {
				hi = mid
			}
		}
		prev = swapWith
		if prev >= len(nodeIndexesByDesire) {
			prev--
		}
		swapWith--
		for nodeIndexesByDesire[prev] != nodeIndex {
			prev--
		}
		if prev != swapWith {
			nodeIndexesByDesire[prev], nodeIndexesByDesire[swapWith] = nodeIndexesByDesire[swapWith], nodeIndexesByDesire[prev]
		}
	}
	nodeIndex2DesiredPartitionCount[nodeIndex]--
}

type tierIDImpl struct {
	values              []int
	nodeIndexesByDesire []int32
	used                bool
}
<|endoftext|>