<commit_before>// gRPC client
package gcli

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"time"

	"github.com/glycerine/blake2b" // vendor https://github.com/dchest/blake2b"
	pb "github.com/glycerine/hnatsd/peer/streambigfile"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
	"hash"
)

type client struct {
	hasher     hash.Hash
	nextChunk  int64
	peerClient pb.PeerClient
}

func newClient(conn *grpc.ClientConn) *client {
	h, err := blake2b.New(nil)
	panicOn(err)
	return &client{
		hasher:     h,
		peerClient: pb.NewPeerClient(conn),
	}
}

func (c *client) startNewFile() {
	c.hasher.Reset()
	c.nextChunk = 0
}

func (c *client) runSendFile(path string, data []byte, maxChunkSize int, isBcastSet bool) error {
	//p("client runSendFile(path='%s') starting", path)

	c.startNewFile()
	stream, err := c.peerClient.SendFile(context.Background())
	if err != nil {
		grpclog.Fatalf("%v.SendFile(_) = _, %v", c.peerClient, err)
	}
	n := len(data)
	numChunk := n / maxChunkSize
	if n%maxChunkSize > 0 {
		numChunk++
	}
	nextByte := 0
	lastChunk := numChunk - 1
	//p("'%s' client sees %v chunks of size ~ %v bytes", path, numChunk, intMin(n, maxChunkSize))
	for i := 0; i < numChunk; i++ {
		sendLen := intMin(maxChunkSize, n-(i*maxChunkSize))
		chunk := data[nextByte:(nextByte + sendLen)]
		nextByte += sendLen

		var nk pb.BigFileChunk
		nk.IsBcastSet = isBcastSet
		nk.Filepath = path
		nk.SizeInBytes = int64(sendLen)
		nk.SendTime = uint64(time.Now().UnixNano())

		// checksums
		c.hasher.Write(chunk)
		nk.Blake2B = blake2bOfBytes(chunk)
		nk.Blake2BCumulative = []byte(c.hasher.Sum(nil))

		nk.Data = chunk
		nk.ChunkNumber = c.nextChunk
		c.nextChunk++
		nk.IsLastChunk = (i == lastChunk)

		//		if nk.ChunkNumber%100 == 0 {
		//p("client, on chunk %v of '%s', checksum='%x', and cumul='%x'", nk.ChunkNumber, nk.Filepath, nk.Blake2B, nk.Blake2BCumulative)
		//		}

		if err := stream.Send(&nk); err != nil {
			// EOF?
			if err == io.EOF {
				if !nk.IsLastChunk {
					panic(fmt.Sprintf("'%s' we got io.EOF before "+
						"the last chunk! At: %v of %v", path, nk.ChunkNumber, numChunk))
				} else {
					break
				}
			}
			panic(err)
			//grpclog.Fatalf("%v.Send() = %v", stream, err)
		}
	}
	reply, err := stream.CloseAndRecv()
	if err != nil {
		// EOF ??
		grpclog.Printf("%v.CloseAndRecv() got error %v, want %v. reply=%v", stream, err, nil, reply)
		return err
	}

	compared := bytes.Compare(reply.WholeFileBlake2B, []byte(c.hasher.Sum(nil)))
	grpclog.Printf("Reply saw checksum: '%x' match: %v; size sent = %v, size received = %v", reply.WholeFileBlake2B, compared == 0, len(data), reply.SizeInBytes)

	if int64(len(data)) != reply.SizeInBytes {
		panic("size mismatch")
	}

	return nil
}

func blake2bOfBytes(by []byte) []byte {
	h, err := blake2b.New(nil)
	panicOn(err)
	h.Write(by)
	return []byte(h.Sum(nil))
}

func intMin(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func SequentialPayload(n int64) []byte {
	if n%8 != 0 {
		panic(fmt.Sprintf("n == %v must be a multiple of 8; has remainder %v", n, n%8))
	}

	k := uint64(n / 8)
	by := make([]byte, n)
	j := uint64(0)
	for i := uint64(0); i < k; i++ {
		j = i * 8
		binary.LittleEndian.PutUint64(by[j:j+8], j)
	}
	return by
}

func (cfg *ClientConfig) ClientSendFile(path string, data []byte, isBcastSet bool) error {

	var opts []grpc.DialOption
	if cfg.UseTLS {
		cfg.setupTLS(&opts)
	} else {
		cfg.setupSSH(&opts)
	}

	serverAddr := fmt.Sprintf("%v:%v", cfg.ServerHost, cfg.ServerPort)

	conn, err := grpc.Dial(serverAddr, opts...)
	if err != nil {
		return err
	}
	defer conn.Close()

	// SendFile
	c := newClient(conn)

	chunkSz := 1 << 20

	t0 := time.Now()
	err = c.runSendFile(path, data, chunkSz, isBcastSet)
	t1 := time.Now()
	if err != nil {
		return err
	}
	mb := float64(len(data)) / float64(1<<20)
	elap := t1.Sub(t0)
	//p("c: elap time to send %v MB was %v => %.03f MB/sec", mb, elap, mb/(float64(elap)/1e9))
	return nil
}

func (cfg *ClientConfig) setupTLS(opts *[]grpc.DialOption) {
	var sn string
	if cfg.ServerHostOverride != "" {
		sn = cfg.ServerHostOverride
	}
	var creds credentials.TransportCredentials
	if cfg.CertPath != "" {
		var err error
		creds, err = credentials.NewClientTLSFromFile(cfg.CertPath, sn)
		if err != nil {
			grpclog.Fatalf("Failed to create TLS credentials %v", err)
		}
	} else {
		creds = credentials.NewClientTLSFromCert(nil, sn)
	}
	*opts = append(*opts, grpc.WithTransportCredentials(creds))
}

func (cfg *ClientConfig) setupSSH(opts *[]grpc.DialOption) {

	destAddr := fmt.Sprintf("%v:%v", cfg.ServerInternalHost, cfg.ServerInternalPort)

	dialer, err := clientSshMain(cfg.AllowNewServer, cfg.TestAllowOneshotConnect, cfg.PrivateKeyPath, cfg.ClientKnownHostsPath, cfg.Username, cfg.ServerHost, destAddr, int64(cfg.ServerPort))
	panicOn(err)

	*opts = append(*opts, grpc.WithDialer(dialer))

	// have to do this too, since we are using an SSH tunnel
	// that grpc doesn't know about:
	*opts = append(*opts, grpc.WithInsecure())
}
<commit_msg>atg. quiet down. actually builds.<commit_after>// gRPC client
package gcli

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"time"

	"github.com/glycerine/blake2b" // vendor https://github.com/dchest/blake2b"
	pb "github.com/glycerine/hnatsd/peer/streambigfile"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/grpclog"
	"hash"
)

type client struct {
	hasher     hash.Hash
	nextChunk  int64
	peerClient pb.PeerClient
}

func newClient(conn *grpc.ClientConn) *client {
	h, err := blake2b.New(nil)
	panicOn(err)
	return &client{
		hasher:     h,
		peerClient: pb.NewPeerClient(conn),
	}
}

func (c *client) startNewFile() {
	c.hasher.Reset()
	c.nextChunk = 0
}

func (c *client) runSendFile(path string, data []byte, maxChunkSize int, isBcastSet bool) error {
	//p("client runSendFile(path='%s') starting", path)

	c.startNewFile()
	stream, err := c.peerClient.SendFile(context.Background())
	if err != nil {
		grpclog.Fatalf("%v.SendFile(_) = _, %v", c.peerClient, err)
	}
	n := len(data)
	numChunk := n / maxChunkSize
	if n%maxChunkSize > 0 {
		numChunk++
	}
	nextByte := 0
	lastChunk := numChunk - 1
	//p("'%s' client sees %v chunks of size ~ %v bytes", path, numChunk, intMin(n, maxChunkSize))
	for i := 0; i < numChunk; i++ {
		sendLen := intMin(maxChunkSize, n-(i*maxChunkSize))
		chunk := data[nextByte:(nextByte + sendLen)]
		nextByte += sendLen

		var nk pb.BigFileChunk
		nk.IsBcastSet = isBcastSet
		nk.Filepath = path
		nk.SizeInBytes = int64(sendLen)
		nk.SendTime = uint64(time.Now().UnixNano())

		// checksums
		c.hasher.Write(chunk)
		nk.Blake2B = blake2bOfBytes(chunk)
		nk.Blake2BCumulative = []byte(c.hasher.Sum(nil))

		nk.Data = chunk
		nk.ChunkNumber = c.nextChunk
		c.nextChunk++
		nk.IsLastChunk = (i == lastChunk)

		//		if nk.ChunkNumber%100 == 0 {
		//p("client, on chunk %v of '%s', checksum='%x', and cumul='%x'", nk.ChunkNumber, nk.Filepath, nk.Blake2B, nk.Blake2BCumulative)
		//		}

		if err := stream.Send(&nk); err != nil {
			// EOF?
			if err == io.EOF {
				if !nk.IsLastChunk {
					panic(fmt.Sprintf("'%s' we got io.EOF before "+
						"the last chunk! At: %v of %v", path, nk.ChunkNumber, numChunk))
				} else {
					break
				}
			}
			panic(err)
			//grpclog.Fatalf("%v.Send() = %v", stream, err)
		}
	}
	reply, err := stream.CloseAndRecv()
	if err != nil {
		// EOF ??
		grpclog.Printf("%v.CloseAndRecv() got error %v, want %v. reply=%v", stream, err, nil, reply)
		return err
	}

	compared := bytes.Compare(reply.WholeFileBlake2B, []byte(c.hasher.Sum(nil)))
	grpclog.Printf("Reply saw checksum: '%x' match: %v; size sent = %v, size received = %v", reply.WholeFileBlake2B, compared == 0, len(data), reply.SizeInBytes)

	if int64(len(data)) != reply.SizeInBytes {
		panic("size mismatch")
	}

	return nil
}

func blake2bOfBytes(by []byte) []byte {
	h, err := blake2b.New(nil)
	panicOn(err)
	h.Write(by)
	return []byte(h.Sum(nil))
}

func intMin(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func SequentialPayload(n int64) []byte {
	if n%8 != 0 {
		panic(fmt.Sprintf("n == %v must be a multiple of 8; has remainder %v", n, n%8))
	}

	k := uint64(n / 8)
	by := make([]byte, n)
	j := uint64(0)
	for i := uint64(0); i < k; i++ {
		j = i * 8
		binary.LittleEndian.PutUint64(by[j:j+8], j)
	}
	return by
}

func (cfg *ClientConfig) ClientSendFile(path string, data []byte, isBcastSet bool) error {

	var opts []grpc.DialOption
	if cfg.UseTLS {
		cfg.setupTLS(&opts)
	} else {
		cfg.setupSSH(&opts)
	}

	serverAddr := fmt.Sprintf("%v:%v", cfg.ServerHost, cfg.ServerPort)

	conn, err := grpc.Dial(serverAddr, opts...)
	if err != nil {
		return err
	}
	defer conn.Close()

	// SendFile
	c := newClient(conn)

	chunkSz := 1 << 20

	t0 := time.Now()
	err = c.runSendFile(path, data, chunkSz, isBcastSet)
	t1 := time.Now()
	if err != nil {
		return err
	}
	mb := float64(len(data)) / float64(1<<20)
	elap := t1.Sub(t0)
	_ = mb
	_ = elap
	//p("c: elap time to send %v MB was %v => %.03f MB/sec", mb, elap, mb/(float64(elap)/1e9))
	return nil
}

func (cfg *ClientConfig) setupTLS(opts *[]grpc.DialOption) {
	var sn string
	if cfg.ServerHostOverride != "" {
		sn = cfg.ServerHostOverride
	}
	var creds credentials.TransportCredentials
	if cfg.CertPath != "" {
		var err error
		creds, err = credentials.NewClientTLSFromFile(cfg.CertPath, sn)
		if err != nil {
			grpclog.Fatalf("Failed to create TLS credentials %v", err)
		}
	} else {
		creds = credentials.NewClientTLSFromCert(nil, sn)
	}
	*opts = append(*opts, grpc.WithTransportCredentials(creds))
}

func (cfg *ClientConfig) setupSSH(opts *[]grpc.DialOption) {

	destAddr := fmt.Sprintf("%v:%v", cfg.ServerInternalHost, cfg.ServerInternalPort)

	dialer, err := clientSshMain(cfg.AllowNewServer, cfg.TestAllowOneshotConnect, cfg.PrivateKeyPath, cfg.ClientKnownHostsPath, cfg.Username, cfg.ServerHost, destAddr, int64(cfg.ServerPort))
	panicOn(err)

	*opts = append(*opts, grpc.WithDialer(dialer))

	// have to do this too, since we are using an SSH tunnel
	// that grpc doesn't know about:
	*opts = append(*opts, grpc.WithInsecure())
}
<|endoftext|>


<commit_before>// +build linux

package sandbox

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sync"

	nspkg "github.com/containernetworking/plugins/pkg/ns"
	"github.com/containers/storage/pkg/idtools"
	"github.com/cri-o/cri-o/pkg/config"
	"github.com/google/uuid"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/sys/unix"
)

// Namespace handles data pertaining to a namespace
type Namespace struct {
	sync.Mutex
	ns          NS
	closed      bool
	initialized bool
	nsType      NSType
	nsPath      string
}

// NS is a wrapper for the containernetworking plugin's NetNS interface
// It exists because while NetNS is specifically called such, it is really a generic
// namespace, and can be used for other namespaces
type NS interface {
	nspkg.NetNS
}

// Get returns the Namespace for a given NsIface
func (n *Namespace) Get() *Namespace {
	return n
}

// Initialized returns true if the Namespace is already initialized
func (n *Namespace) Initialized() bool {
	return n.initialized
}

// Initialize does the necessary setup for a Namespace
// It does not do the bind mounting and nspinning
func (n *Namespace) Initialize() NamespaceIface {
	n.closed = false
	n.initialized = true
	return n
}

func getMappingsForPinns(mappings []idtools.IDMap) string {
	g := new(bytes.Buffer)
	for _, m := range mappings {
		fmt.Fprintf(g, "%d-%d-%d@", m.ContainerID, m.HostID, m.Size)
	}
	return g.String()
}

// Creates a new persistent namespace and returns an object
// representing that namespace, without switching to it
func pinNamespaces(nsTypes []NSType, cfg *config.Config, idMappings *idtools.IDMappings) ([]NamespaceIface, error) {
	typeToArg := map[NSType]string{
		IPCNS:  "-i",
		UTSNS:  "-u",
		USERNS: "-U",
		NETNS:  "-n",
	}

	pinnedNamespace := uuid.New().String()
	pinnsArgs := []string{
		"-d", cfg.NamespacesDir,
		"-f", pinnedNamespace,
	}
	type namespaceInfo struct {
		path   string
		nsType NSType
	}

	mountedNamespaces := make([]namespaceInfo, 0, len(nsTypes))

	var rootPair idtools.IDPair
	if idMappings != nil {
		rootPair = idMappings.RootPair()
	}

	for _, nsType := range nsTypes {
		arg, ok := typeToArg[nsType]
		if !ok {
			return nil, errors.Errorf("Invalid namespace type: %s", nsType)
		}
		pinnsArgs = append(pinnsArgs, arg)
		pinPath := filepath.Join(cfg.NamespacesDir, string(nsType)+"ns", pinnedNamespace)
		mountedNamespaces = append(mountedNamespaces, namespaceInfo{
			path:   pinPath,
			nsType: nsType,
		})
		if idMappings != nil {
			err := os.MkdirAll(filepath.Dir(pinPath), 0o755)
			if err != nil {
				return nil, err
			}
			f, err := os.Create(pinPath)
			if err != nil {
				return nil, err
			}
			f.Close()
			if err := os.Chown(pinPath, rootPair.UID, rootPair.GID); err != nil {
				return nil, err
			}
		}
	}

	if idMappings != nil {
		pinnsArgs = append(pinnsArgs,
			fmt.Sprintf("--uid-mapping=%s", getMappingsForPinns(idMappings.UIDs())),
			fmt.Sprintf("--gid-mapping=%s", getMappingsForPinns(idMappings.GIDs())))
	}

	pinns := cfg.PinnsPath

	logrus.Debugf("calling pinns with %v", pinnsArgs)
	output, err := exec.Command(pinns, pinnsArgs...).Output()
	if len(output) != 0 {
		logrus.Debugf("pinns output: %s", string(output))
	}
	if err != nil {
		// cleanup after ourselves
		failedUmounts := make([]string, 0)
		for _, info := range mountedNamespaces {
			if unmountErr := unix.Unmount(info.path, unix.MNT_DETACH); unmountErr != nil {
				failedUmounts = append(failedUmounts, info.path)
			}
		}
		if len(failedUmounts) != 0 {
			return nil, fmt.Errorf("failed to cleanup %v after pinns failure %s %v", failedUmounts, output, err)
		}
		return nil, fmt.Errorf("failed to pin namespaces %v: %s %v", nsTypes, output, err)
	}

	returnedNamespaces := make([]NamespaceIface, 0)
	for _, info := range mountedNamespaces {
		ret, err := nspkg.GetNS(info.path)
		if err != nil {
			return nil, err
		}

		returnedNamespaces = append(returnedNamespaces, &Namespace{
			ns:     ret.(NS),
			nsType: info.nsType,
			nsPath: info.path,
		})
	}
	return returnedNamespaces, nil
}

// getNamespace takes a path, checks if it is a namespace, and if so
// returns a Namespace
func getNamespace(nsPath string) (*Namespace, error) {
	if err := nspkg.IsNSorErr(nsPath); err != nil {
		return nil, err
	}

	ns, err := nspkg.GetNS(nsPath)
	if err != nil {
		return nil, err
	}

	return &Namespace{ns: ns, closed: false, nsPath: nsPath}, nil
}

// Path returns the path of the namespace handle
func (n *Namespace) Path() string {
	if n == nil || n.ns == nil {
		return ""
	}
	return n.nsPath
}

// Type returns which namespace this structure represents
func (n *Namespace) Type() NSType {
	return n.nsType
}

// Close closes this namespace
func (n *Namespace) Close() error {
	if n == nil || n.ns == nil {
		return nil
	}
	return n.ns.Close()
}

// Remove ensures this namespace handle is closed and removed
func (n *Namespace) Remove() error {
	n.Lock()
	defer n.Unlock()

	if n.closed {
		// nsRemove() can be called multiple
		// times without returning an error.
		return nil
	}

	if err := n.Close(); err != nil {
		return err
	}

	n.closed = true

	fp := n.Path()
	if fp == "" {
		return nil
	}

	// try to unmount, ignoring "not mounted" (EINVAL) error
	if err := unix.Unmount(fp, unix.MNT_DETACH); err != nil && err != unix.EINVAL {
		return errors.Wrapf(err, "unable to unmount %s", fp)
	}
	return os.RemoveAll(fp)
}
<commit_msg>pinNamespaces: set capacity for returnedNamespaces<commit_after>// +build linux

package sandbox

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sync"

	nspkg "github.com/containernetworking/plugins/pkg/ns"
	"github.com/containers/storage/pkg/idtools"
	"github.com/cri-o/cri-o/pkg/config"
	"github.com/google/uuid"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/sys/unix"
)

// Namespace handles data pertaining to a namespace
type Namespace struct {
	sync.Mutex
	ns          NS
	closed      bool
	initialized bool
	nsType      NSType
	nsPath      string
}

// NS is a wrapper for the containernetworking plugin's NetNS interface
// It exists because while NetNS is specifically called such, it is really a generic
// namespace, and can be used for other namespaces
type NS interface {
	nspkg.NetNS
}

// Get returns the Namespace for a given NsIface
func (n *Namespace) Get() *Namespace {
	return n
}

// Initialized returns true if the Namespace is already initialized
func (n *Namespace) Initialized() bool {
	return n.initialized
}

// Initialize does the necessary setup for a Namespace
// It does not do the bind mounting and nspinning
func (n *Namespace) Initialize() NamespaceIface {
	n.closed = false
	n.initialized = true
	return n
}

func getMappingsForPinns(mappings []idtools.IDMap) string {
	g := new(bytes.Buffer)
	for _, m := range mappings {
		fmt.Fprintf(g, "%d-%d-%d@", m.ContainerID, m.HostID, m.Size)
	}
	return g.String()
}

// Creates a new persistent namespace and returns an object
// representing that namespace, without switching to it
func pinNamespaces(nsTypes []NSType, cfg *config.Config, idMappings *idtools.IDMappings) ([]NamespaceIface, error) {
	typeToArg := map[NSType]string{
		IPCNS:  "-i",
		UTSNS:  "-u",
		USERNS: "-U",
		NETNS:  "-n",
	}

	pinnedNamespace := uuid.New().String()
	pinnsArgs := []string{
		"-d", cfg.NamespacesDir,
		"-f", pinnedNamespace,
	}
	type namespaceInfo struct {
		path   string
		nsType NSType
	}

	mountedNamespaces := make([]namespaceInfo, 0, len(nsTypes))

	var rootPair idtools.IDPair
	if idMappings != nil {
		rootPair = idMappings.RootPair()
	}

	for _, nsType := range nsTypes {
		arg, ok := typeToArg[nsType]
		if !ok {
			return nil, errors.Errorf("Invalid namespace type: %s", nsType)
		}
		pinnsArgs = append(pinnsArgs, arg)
		pinPath := filepath.Join(cfg.NamespacesDir, string(nsType)+"ns", pinnedNamespace)
		mountedNamespaces = append(mountedNamespaces, namespaceInfo{
			path:   pinPath,
			nsType: nsType,
		})
		if idMappings != nil {
			err := os.MkdirAll(filepath.Dir(pinPath), 0o755)
			if err != nil {
				return nil, err
			}
			f, err := os.Create(pinPath)
			if err != nil {
				return nil, err
			}
			f.Close()
			if err := os.Chown(pinPath, rootPair.UID, rootPair.GID); err != nil {
				return nil, err
			}
		}
	}

	if idMappings != nil {
		pinnsArgs = append(pinnsArgs,
			fmt.Sprintf("--uid-mapping=%s", getMappingsForPinns(idMappings.UIDs())),
			fmt.Sprintf("--gid-mapping=%s", getMappingsForPinns(idMappings.GIDs())))
	}

	pinns := cfg.PinnsPath

	logrus.Debugf("calling pinns with %v", pinnsArgs)
	output, err := exec.Command(pinns, pinnsArgs...).Output()
	if len(output) != 0 {
		logrus.Debugf("pinns output: %s", string(output))
	}
	if err != nil {
		// cleanup after ourselves
		failedUmounts := make([]string, 0)
		for _, info := range mountedNamespaces {
			if unmountErr := unix.Unmount(info.path, unix.MNT_DETACH); unmountErr != nil {
				failedUmounts = append(failedUmounts, info.path)
			}
		}
		if len(failedUmounts) != 0 {
			return nil, fmt.Errorf("failed to cleanup %v after pinns failure %s %v", failedUmounts, output, err)
		}
		return nil, fmt.Errorf("failed to pin namespaces %v: %s %v", nsTypes, output, err)
	}

	returnedNamespaces := make([]NamespaceIface, 0, len(nsTypes))
	for _, info := range mountedNamespaces {
		ret, err := nspkg.GetNS(info.path)
		if err != nil {
			return nil, err
		}

		returnedNamespaces = append(returnedNamespaces, &Namespace{
			ns:     ret.(NS),
			nsType: info.nsType,
			nsPath: info.path,
		})
	}
	return returnedNamespaces, nil
}

// getNamespace takes a path, checks if it is a namespace, and if so
// returns a Namespace
func getNamespace(nsPath string) (*Namespace, error) {
	if err := nspkg.IsNSorErr(nsPath); err != nil {
		return nil, err
	}

	ns, err := nspkg.GetNS(nsPath)
	if err != nil {
		return nil, err
	}

	return &Namespace{ns: ns, closed: false, nsPath: nsPath}, nil
}

// Path returns the path of the namespace handle
func (n *Namespace) Path() string {
	if n == nil || n.ns == nil {
		return ""
	}
	return n.nsPath
}

// Type returns which namespace this structure represents
func (n *Namespace) Type() NSType {
	return n.nsType
}

// Close closes this namespace
func (n *Namespace) Close() error {
	if n == nil || n.ns == nil {
		return nil
	}
	return n.ns.Close()
}

// Remove ensures this namespace handle is closed and removed
func (n *Namespace) Remove() error {
	n.Lock()
	defer n.Unlock()

	if n.closed {
		// nsRemove() can be called multiple
		// times without returning an error.
		return nil
	}

	if err := n.Close(); err != nil {
		return err
	}

	n.closed = true

	fp := n.Path()
	if fp == "" {
		return nil
	}

	// try to unmount, ignoring "not mounted" (EINVAL) error
	if err := unix.Unmount(fp, unix.MNT_DETACH); err != nil && err != unix.EINVAL {
		return errors.Wrapf(err, "unable to unmount %s", fp)
	}
	return os.RemoveAll(fp)
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	_ "github.com/lib/pq"
)

type PostgresDB struct {
	c          *sql.DB
	SchemaName string
	TableName  string
}

func NewPostgresDB() ShortenBackend {
	pgHost := GetOrDefault("PG_HOST", "localhost")
	pgPort := GetOrDefault("PG_HOST", "5432")
	pgUser := GetOrDefault("PG_USER", "shortener")
	pgPass := GetOrDefault("PG_PASS", "NOPE")
	pgDatabase := GetOrDefault("PG_DB", "shortener")
	pgSchema := GetOrDefault("PG_SCHEMA", "shortener")
	pgTable := GetOrDefault("PG_TABLE", "shortener")
	pgSSLMode := GetOrDefault("PG_SSL", "disable")

	connString := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s", pgHost, pgPort, pgUser, pgPass, pgDatabase, pgSSLMode)
	db, err := sql.Open("postgres", connString)
	if err != nil {
		log.Fatal(err)
	}
	return &PostgresDB{
		c:          db,
		SchemaName: pgSchema,
		TableName:  pgTable,
	}
}

func (pgDB *PostgresDB) DeleteURL(slug string) error {
	_, err := pgDB.c.Query(fmt.Sprintf("DELETE FROM %s.%s WHERE slug=$1", pgDB.SchemaName, pgDB.TableName), slug)
	if err != nil {
		return err
	}
	log.Printf("Successfully deleted slug: %s", slug)
	return nil

}
func (pgDB *PostgresDB) ShortenURL(slug, longURL, owner string, expires time.Time) error {
	// postgres & redshift don't have an upsert method yet
	existingLong, err := pgDB.GetLongURL(slug)
	if existingLong == "" || err != nil { // TODO figure out what happens on nothing, err?
		//q := fmt.Sprintf("INSERT INTO %s.%s(slug, long_url, expires, modified) VALUES($1, $2, $3, $4)")
		q := fmt.Sprintf("INSERT INTO %s.%s(slug, long_url, owner) VALUES($1, $2, $3)", pgDB.SchemaName, pgDB.TableName)
		_, err := pgDB.c.Query(q, slug, longURL, owner)
		if err != nil {
			return fmt.Errorf("Issue inserting new row for slug: %s, err is: %s", slug, err)
		}
		return nil
	}
	// Otherwise, upsert
	q := fmt.Sprintf("UPDATE %s.%s SET long_url=$2, owner=$3 WHERE slug=$1", pgDB.SchemaName, pgDB.TableName)
	_, err = pgDB.c.Query(q, slug, longURL, owner)
	if err != nil {
		return err
	}
	log.Printf("Successfully updated slug: %s", slug)
	return nil
}

func (pgDB *PostgresDB) GetLongURL(slug string) (string, error) {
	//var retObj ShortenObject
	q := fmt.Sprintf("SELECT long_url FROM %s.%s WHERE slug = $1", pgDB.SchemaName, pgDB.TableName)
	var long_url string
	err := pgDB.c.QueryRow(q, slug).Scan(&long_url)
	if err != nil {
		return "", err
	}
	log.Println("long: ", long_url)
	return long_url, nil

}

func (pgDB *PostgresDB) GetList() ([]ShortenObject, error) {
	rows, err := pgDB.c.Query(fmt.Sprintf("SELECT slug, long_url, owner FROM %s.%s", pgDB.SchemaName, pgDB.TableName))
	if err != nil {
		return nil, err
	}
	var retObjs []ShortenObject
	var slug string
	var long_url string
	var owner string
	//tags :=  []string{}
	//var expires time.Time
	//var modified time.Time
	for rows.Next() {
		err = rows.Scan(&slug, &long_url, &owner) //, &modified, &expires)
		if err != nil {
			return nil, fmt.Errorf("issue scanning row for list: %s", err)
		}
		retObjs = append(retObjs, ShortenObject{
			Slug:    slug,
			LongURL: long_url,
			Owner:   owner,
			//Expires:  expires,
			//Modified: modified,
		})
	}
	return retObjs, nil
}
<commit_msg>fix postgres environment variables<commit_after>package db

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	_ "github.com/lib/pq"
)

type PostgresDB struct {
	c          *sql.DB
	SchemaName string
	TableName  string
}

func NewPostgresDB() ShortenBackend {
	pgHost := GetOrDefault("PG_HOST", "localhost")
	pgPort := GetOrDefault("PG_PORT", "5432")
	pgUser := GetOrDefault("PG_USER", "shortener")
	pgPass := GetOrDefault("PG_PASSWORD", "NOPE")
	pgDatabase := GetOrDefault("PG_DATABASE", "shortener")
	pgSchema := GetOrDefault("PG_SCHEMA", "shortener")
	pgTable := GetOrDefault("PG_TABLE", "shortener")
	pgSSLMode := GetOrDefault("PG_SSL", "disable")

	connString := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s", pgHost, pgPort, pgUser, pgPass, pgDatabase, pgSSLMode)
	db, err := sql.Open("postgres", connString)
	if err != nil {
		log.Fatal(err)
	}
	return &PostgresDB{
		c:          db,
		SchemaName: pgSchema,
		TableName:  pgTable,
	}
}

func (pgDB *PostgresDB) DeleteURL(slug string) error {
	_, err := pgDB.c.Query(fmt.Sprintf("DELETE FROM %s.%s WHERE slug=$1", pgDB.SchemaName, pgDB.TableName), slug)
	if err != nil {
		return err
	}
	log.Printf("Successfully deleted slug: %s", slug)
	return nil

}
func (pgDB *PostgresDB) ShortenURL(slug, longURL, owner string, expires time.Time) error {
	// postgres & redshift don't have an upsert method yet
	existingLong, err := pgDB.GetLongURL(slug)
	if existingLong == "" || err != nil { // TODO figure out what happens on nothing, err?
		//q := fmt.Sprintf("INSERT INTO %s.%s(slug, long_url, expires, modified) VALUES($1, $2, $3, $4)")
		q := fmt.Sprintf("INSERT INTO %s.%s(slug, long_url, owner) VALUES($1, $2, $3)", pgDB.SchemaName, pgDB.TableName)
		_, err := pgDB.c.Query(q, slug, longURL, owner)
		if err != nil {
			return fmt.Errorf("Issue inserting new row for slug: %s, err is: %s", slug, err)
		}
		return nil
	}
	// Otherwise, upsert
	q := fmt.Sprintf("UPDATE %s.%s SET long_url=$2, owner=$3 WHERE slug=$1", pgDB.SchemaName, pgDB.TableName)
	_, err = pgDB.c.Query(q, slug, longURL, owner)
	if err != nil {
		return err
	}
	log.Printf("Successfully updated slug: %s", slug)
	return nil
}

func (pgDB *PostgresDB) GetLongURL(slug string) (string, error) {
	//var retObj ShortenObject
	q := fmt.Sprintf("SELECT long_url FROM %s.%s WHERE slug = $1", pgDB.SchemaName, pgDB.TableName)
	var long_url string
	err := pgDB.c.QueryRow(q, slug).Scan(&long_url)
	if err != nil {
		return "", err
	}
	log.Println("long: ", long_url)
	return long_url, nil

}

func (pgDB *PostgresDB) GetList() ([]ShortenObject, error) {
	rows, err := pgDB.c.Query(fmt.Sprintf("SELECT slug, long_url, owner FROM %s.%s", pgDB.SchemaName, pgDB.TableName))
	if err != nil {
		return nil, err
	}
	var retObjs []ShortenObject
	var slug string
	var long_url string
	var owner string
	//tags :=  []string{}
	//var expires time.Time
	//var modified time.Time
	for rows.Next() {
		err = rows.Scan(&slug, &long_url, &owner) //, &modified, &expires)
		if err != nil {
			return nil, fmt.Errorf("issue scanning row for list: %s", err)
		}
		retObjs = append(retObjs, ShortenObject{
			Slug:    slug,
			LongURL: long_url,
			Owner:   owner,
			//Expires:  expires,
			//Modified: modified,
		})
	}
	return retObjs, nil
}
<|endoftext|>


<commit_before>package defaults

import (
	"context"
	"encoding/json"

	"github.com/volatiletech/authboss"
)

var (
	jsonDefaultFailures = []string{authboss.DataErr, authboss.DataValidation}
)

// There is a separate package that does HTML Rendering authboss-renderer

// JSONRenderer simply renders the data provided in JSON.
// Known failure keys in the HTMLData can be passed in to force a
// status: failure in the JSON when they appear.
type JSONRenderer struct {
	Failures []string
}

// Load is a no-op since json doesn't require any templates
func (JSONRenderer) Load(names ...string) error {
	return nil
}

// Render the data
func (j JSONRenderer) Render(ctx context.Context, page string, data authboss.HTMLData) (output []byte, contentType string, err error) {
	if _, hasStatus := data["status"]; !hasStatus {
		failures := j.Failures
		if len(failures) == 0 {
			failures = jsonDefaultFailures
		}

		status := "success"
		for _, failure := range failures {
			val, has := data[failure]
			if has && val != nil {
				status = "failure"
				break
			}
		}

		data["status"] = status
	}

	b, err := json.Marshal(data)
	if err != nil {
		return nil, "", err
	}

	return b, "application/json", nil
}
<commit_msg>Fix panic in JSON renderer on empty body<commit_after>package defaults

import (
	"context"
	"encoding/json"

	"github.com/volatiletech/authboss"
)

var (
	jsonDefaultFailures = []string{authboss.DataErr, authboss.DataValidation}
)

// There is a separate package that does HTML Rendering authboss-renderer

// JSONRenderer simply renders the data provided in JSON.
// Known failure keys in the HTMLData can be passed in to force a
// status: failure in the JSON when they appear.
type JSONRenderer struct {
	Failures []string
}

// Load is a no-op since json doesn't require any templates
func (JSONRenderer) Load(names ...string) error {
	return nil
}

// Render the data
func (j JSONRenderer) Render(ctx context.Context, page string, data authboss.HTMLData) (output []byte, contentType string, err error) {
	if data == nil {
		return []byte(`{"status":"success"}`), "application/json", nil
	}

	if _, hasStatus := data["status"]; !hasStatus {
		failures := j.Failures
		if len(failures) == 0 {
			failures = jsonDefaultFailures
		}

		status := "success"
		for _, failure := range failures {
			val, has := data[failure]
			if has && val != nil {
				status = "failure"
				break
			}
		}

		data["status"] = status
	}

	b, err := json.Marshal(data)
	if err != nil {
		return nil, "", err
	}

	return b, "application/json", nil
}
<|endoftext|>


<commit_before>package redis

import (
	"fmt"
	"reflect"
	"strconv"
	"time"
)

type (
	HashValue   map[string][]byte
	HashHash    map[string]HashValue
	HashSub     map[string][]*ChannelWriter
	HashBrStack map[string]*Stack
)

type Database struct {
	children map[int]*Database
	parent   *Database

	values  HashValue
	hvalues HashHash
	brstack HashBrStack

	sub HashSub
}

func NewDatabase(parent *Database) *Database {
	db := &Database{
		values:   make(HashValue),
		sub:      make(HashSub),
		brstack:  make(HashBrStack),
		children: map[int]*Database{},
		parent:   parent,
	}
	db.children[0] = db
	return db
}

type DefaultHandler struct {
	*Database
	currentDb int
	dbs       map[int]*Database
}

func (h *DefaultHandler) Rpush(key string, value []byte, values ...[]byte) (int, error) {
	values = append([][]byte{value}, values...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}
	for _, value := range values {
		h.brstack[key].PushBack(value)
	}
	return h.brstack[key].Len(), nil
}

func (h *DefaultHandler) Brpop(key string, keys ...string) (data [][]byte, err error) {
	keys = append([]string{key}, keys...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	if len(keys) == 0 {
		return nil, ErrParseTimeout
	}

	timeout, err := strconv.Atoi(keys[len(keys)-1])
	if err != nil {
		return nil, ErrParseTimeout
	}
	keys = keys[:len(keys)-1]

	var timeoutChan <-chan time.Time
	if timeout > 0 {
		timeoutChan = time.After(time.Duration(timeout) * time.Second)
	} else {
		timeoutChan = make(chan time.Time)
	}

	finishedChan := make(chan struct{})
	go func() {
		defer close(finishedChan)
		selectCases := []reflect.SelectCase{}
		for _, k := range keys {
			key := string(k)
			if _, exists := h.brstack[key]; !exists {
				h.brstack[key] = NewStack(k)
			}
			selectCases = append(selectCases, reflect.SelectCase{
				Dir:  reflect.SelectRecv,
				Chan: reflect.ValueOf(h.brstack[key].Chan),
			})
		}
		_, recv, _ := reflect.Select(selectCases)
		s, ok := recv.Interface().(*Stack)
		if !ok {
			err = fmt.Errorf("Impossible to retrieve data. Wrong type.")
			return
		}
		data = [][]byte{[]byte(s.Key), s.PopBack()}
	}()

	select {
	case <-finishedChan:
		return data, err
	case <-timeoutChan:
		return nil, nil
	}
	return nil, nil
}

func (h *DefaultHandler) Lrange(key string, start, stop int) ([][]byte, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}

	if start < 0 {
		if start = h.brstack[key].Len() + start; start < 0 {
			start = 0
		}
	}

	var ret [][]byte
	for i := start; i <= stop; i++ {
		if val := h.brstack[key].GetIndex(i); val != nil {
			ret = append(ret, val)
		}
	}
	return ret, nil
}

func (h *DefaultHandler) Lindex(key string, index int) ([]byte, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}
	return h.brstack[key].GetIndex(index), nil
}

func (h *DefaultHandler) Lpush(key string, value []byte, values ...[]byte) (int, error) {
	values = append([][]byte{value}, values...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}
	for _, value := range values {
		h.brstack[key].PushFront(value)
	}
	return h.brstack[key].Len(), nil
}

func (h *DefaultHandler) Blpop(key string, keys ...string) (data [][]byte, err error) {
	keys = append([]string{key}, keys...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	if len(keys) == 0 {
		return nil, ErrParseTimeout
	}

	timeout, err := strconv.Atoi(keys[len(keys)-1])
	if err != nil {
		return nil, ErrParseTimeout
	}
	keys = keys[:len(keys)-1]

	var timeoutChan <-chan time.Time
	if timeout > 0 {
		timeoutChan = time.After(time.Duration(timeout) * time.Second)
	} else {
		timeoutChan = make(chan time.Time)
	}

	finishedChan := make(chan struct{})

	go func() {
		defer close(finishedChan)
		selectCases := []reflect.SelectCase{}
		for _, k := range keys {
			key := string(k)
			if _, exists := h.brstack[key]; !exists {
				h.brstack[key] = NewStack(k)
			}
			selectCases = append(selectCases, reflect.SelectCase{
				Dir:  reflect.SelectRecv,
				Chan: reflect.ValueOf(h.brstack[key].Chan),
			})
		}
		_, recv, _ := reflect.Select(selectCases)
		s, ok := recv.Interface().(*Stack)
		if !ok {
			err = fmt.Errorf("Impossible to retrieve data. Wrong type.")
			return
		}
		data = [][]byte{[]byte(s.Key), s.PopFront()}
	}()

	select {
	case <-finishedChan:
		return data, err
	case <-timeoutChan:
		return nil, nil
	}
	return nil, nil
}

func (h *DefaultHandler) Hget(key, subkey string) ([]byte, error) {
	if h.Database == nil || h.hvalues == nil {
		return nil, nil
	}

	if v, exists := h.hvalues[key]; exists {
		if v, exists := v[subkey]; exists {
			return v, nil
		}
	}
	return nil, nil
}

func (h *DefaultHandler) Hset(key, subkey string, value []byte) (int, error) {
	ret := 0

	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.hvalues[key]; !exists {
		h.hvalues[key] = make(HashValue)
		ret = 1
	}

	if _, exists := h.hvalues[key][subkey]; !exists {
		ret = 1
	}

	h.hvalues[key][subkey] = value

	return ret, nil
}

func (h *DefaultHandler) Hgetall(key string) (HashValue, error) {
	if h.Database == nil || h.hvalues == nil {
		return nil, nil
	}
	return h.hvalues[key], nil
}

func (h *DefaultHandler) Get(key string) ([]byte, error) {
	if h.Database == nil || h.values == nil {
		return nil, nil
	}
	return h.values[key], nil
}

func (h *DefaultHandler) Set(key string, value []byte) error {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	h.values[key] = value
	return nil
}

func (h *DefaultHandler) Del(key string, keys ...string) (int, error) {
	keys = append([]string{key}, keys...)
	if h.Database == nil {
		return 0, nil
	}
	count := 0
	for _, k := range keys {
		if _, exists := h.values[k]; exists {
			delete(h.values, k)
			count++
		}
		if _, exists := h.hvalues[key]; exists {
			delete(h.hvalues, k)
			count++
		}
	}
	return count, nil
}

func (h *DefaultHandler) Ping() (*StatusReply, error) {
	return &StatusReply{code: "PONG"}, nil
}

func (h *DefaultHandler) Subscribe(channels ...[]byte) (*MultiChannelWriter, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	ret := &MultiChannelWriter{Chans: make([]*ChannelWriter, 0, len(channels))}
	for _, key := range channels {
		Debugf("SUBSCRIBE on %s\n", key)
		cw := &ChannelWriter{
			FirstReply: []interface{}{
				"subscribe",
				key,
				1,
			},
			Channel: make(chan []interface{}),
		}
		if h.sub[string(key)] == nil {
			h.sub[string(key)] = []*ChannelWriter{cw}
		} else {
			h.sub[string(key)] = append(h.sub[string(key)], cw)
		}
		ret.Chans = append(ret.Chans, cw)
	}
	return ret, nil
}

func (h *DefaultHandler) Publish(key string, value []byte) (int, error) {
	if h.Database == nil || h.sub == nil {
		return 0, nil
	}
	//	Debugf("Publishing %s on %s\n", value, key)
	v, exists := h.sub[key]
	if !exists {
		return 0, nil
	}
	i := 0
	for _, c := range v {
		select {
		case c.Channel <- []interface{}{
			"message",
			key,
			value,
		}:
			i++
		default:
		}
	}
	return i, nil
}

func (h *DefaultHandler) Select(key string) error {
	if h.dbs == nil {
		h.dbs = map[int]*Database{0: h.Database}
	}
	index, err := strconv.Atoi(key)
	if err != nil {
		return err
	}
	h.dbs[h.currentDb] = h.Database
	h.currentDb = index
	if _, exists := h.dbs[index]; !exists {
		println("DB not exits, create ", index)
		h.dbs[index] = NewDatabase(nil)
	}
	h.Database = h.dbs[index]
	return nil
}

func (h *DefaultHandler) Monitor() (*MonitorReply, error) {
	return &MonitorReply{}, nil
}

var lock = make(chan bool, 1)

func (h *DefaultHandler) Incr(key string) (int, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	lock <- true

	temp, _ := strconv.Atoi(string(h.values[key]))
	temp = temp + 1
	h.values[key] = []byte(strconv.Itoa(temp))

	<-lock

	return temp, nil
}

func (h *DefaultHandler) Decr(key string) (int, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	lock <- true

	temp, _ := strconv.Atoi(string(h.values[key]))
	temp = temp - 1
	h.values[key] = []byte(strconv.Itoa(temp))

	<-lock

	return temp, nil
}

func (h *DefaultHandler) Expire(key, after string) (error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	d, _ := strconv.Atoi(after)

	time.AfterFunc(time.Duration(d) * time.Second, func() {
		h.Del(key)
	})

	return nil
}

func NewDefaultHandler() *DefaultHandler {
	db := NewDatabase(nil)
	ret := &DefaultHandler{
		Database:  db,
		currentDb: 0,
		dbs:       map[int]*Database{0: db},
	}
	return ret
}
<commit_msg>Add 'exists' method.<commit_after>package redis

import (
	"fmt"
	"reflect"
	"strconv"
	"time"
)

type (
	HashValue   map[string][]byte
	HashHash    map[string]HashValue
	HashSub     map[string][]*ChannelWriter
	HashBrStack map[string]*Stack
)

type Database struct {
	children map[int]*Database
	parent   *Database

	values  HashValue
	hvalues HashHash
	brstack HashBrStack

	sub HashSub
}

func NewDatabase(parent *Database) *Database {
	db := &Database{
		values:   make(HashValue),
		sub:      make(HashSub),
		brstack:  make(HashBrStack),
		children: map[int]*Database{},
		parent:   parent,
	}
	db.children[0] = db
	return db
}

type DefaultHandler struct {
	*Database
	currentDb int
	dbs       map[int]*Database
}

func (h *DefaultHandler) Rpush(key string, value []byte, values ...[]byte) (int, error) {
	values = append([][]byte{value}, values...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}
	for _, value := range values {
		h.brstack[key].PushBack(value)
	}
	return h.brstack[key].Len(), nil
}

func (h *DefaultHandler) Brpop(key string, keys ...string) (data [][]byte, err error) {
	keys = append([]string{key}, keys...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	if len(keys) == 0 {
		return nil, ErrParseTimeout
	}

	timeout, err := strconv.Atoi(keys[len(keys)-1])
	if err != nil {
		return nil, ErrParseTimeout
	}
	keys = keys[:len(keys)-1]

	var timeoutChan <-chan time.Time
	if timeout > 0 {
		timeoutChan = time.After(time.Duration(timeout) * time.Second)
	} else {
		timeoutChan = make(chan time.Time)
	}

	finishedChan := make(chan struct{})
	go func() {
		defer close(finishedChan)
		selectCases := []reflect.SelectCase{}
		for _, k := range keys {
			key := string(k)
			if _, exists := h.brstack[key]; !exists {
				h.brstack[key] = NewStack(k)
			}
			selectCases = append(selectCases, reflect.SelectCase{
				Dir:  reflect.SelectRecv,
				Chan: reflect.ValueOf(h.brstack[key].Chan),
			})
		}
		_, recv, _ := reflect.Select(selectCases)
		s, ok := recv.Interface().(*Stack)
		if !ok {
			err = fmt.Errorf("Impossible to retrieve data. Wrong type.")
			return
		}
		data = [][]byte{[]byte(s.Key), s.PopBack()}
	}()

	select {
	case <-finishedChan:
		return data, err
	case <-timeoutChan:
		return nil, nil
	}
	return nil, nil
}

func (h *DefaultHandler) Lrange(key string, start, stop int) ([][]byte, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}

	if start < 0 {
		if start = h.brstack[key].Len() + start; start < 0 {
			start = 0
		}
	}

	var ret [][]byte
	for i := start; i <= stop; i++ {
		if val := h.brstack[key].GetIndex(i); val != nil {
			ret = append(ret, val)
		}
	}
	return ret, nil
}

func (h *DefaultHandler) Lindex(key string, index int) ([]byte, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}
	return h.brstack[key].GetIndex(index), nil
}

func (h *DefaultHandler) Lpush(key string, value []byte, values ...[]byte) (int, error) {
	values = append([][]byte{value}, values...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.brstack[key]; !exists {
		h.brstack[key] = NewStack(key)
	}
	for _, value := range values {
		h.brstack[key].PushFront(value)
	}
	return h.brstack[key].Len(), nil
}

func (h *DefaultHandler) Blpop(key string, keys ...string) (data [][]byte, err error) {
	keys = append([]string{key}, keys...)
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	if len(keys) == 0 {
		return nil, ErrParseTimeout
	}

	timeout, err := strconv.Atoi(keys[len(keys)-1])
	if err != nil {
		return nil, ErrParseTimeout
	}
	keys = keys[:len(keys)-1]

	var timeoutChan <-chan time.Time
	if timeout > 0 {
		timeoutChan = time.After(time.Duration(timeout) * time.Second)
	} else {
		timeoutChan = make(chan time.Time)
	}

	finishedChan := make(chan struct{})

	go func() {
		defer close(finishedChan)
		selectCases := []reflect.SelectCase{}
		for _, k := range keys {
			key := string(k)
			if _, exists := h.brstack[key]; !exists {
				h.brstack[key] = NewStack(k)
			}
			selectCases = append(selectCases, reflect.SelectCase{
				Dir:  reflect.SelectRecv,
				Chan: reflect.ValueOf(h.brstack[key].Chan),
			})
		}
		_, recv, _ := reflect.Select(selectCases)
		s, ok := recv.Interface().(*Stack)
		if !ok {
			err = fmt.Errorf("Impossible to retrieve data. Wrong type.")
			return
		}
		data = [][]byte{[]byte(s.Key), s.PopFront()}
	}()

	select {
	case <-finishedChan:
		return data, err
	case <-timeoutChan:
		return nil, nil
	}
	return nil, nil
}

func (h *DefaultHandler) Hget(key, subkey string) ([]byte, error) {
	if h.Database == nil || h.hvalues == nil {
		return nil, nil
	}

	if v, exists := h.hvalues[key]; exists {
		if v, exists := v[subkey]; exists {
			return v, nil
		}
	}
	return nil, nil
}

func (h *DefaultHandler) Hset(key, subkey string, value []byte) (int, error) {
	ret := 0

	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	if _, exists := h.hvalues[key]; !exists {
		h.hvalues[key] = make(HashValue)
		ret = 1
	}

	if _, exists := h.hvalues[key][subkey]; !exists {
		ret = 1
	}

	h.hvalues[key][subkey] = value

	return ret, nil
}

func (h *DefaultHandler) Hgetall(key string) (HashValue, error) {
	if h.Database == nil || h.hvalues == nil {
		return nil, nil
	}
	return h.hvalues[key], nil
}

func (h *DefaultHandler) Get(key string) ([]byte, error) {
	if h.Database == nil || h.values == nil {
		return nil, nil
	}
	return h.values[key], nil
}

func (h *DefaultHandler) Set(key string, value []byte) error {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	h.values[key] = value
	return nil
}

func (h *DefaultHandler) Del(key string, keys ...string) (int, error) {
	keys = append([]string{key}, keys...)
	if h.Database == nil {
		return 0, nil
	}
	count := 0
	for _, k := range keys {
		if _, exists := h.values[k]; exists {
			delete(h.values, k)
			count++
		}
		if _, exists := h.hvalues[key]; exists {
			delete(h.hvalues, k)
			count++
		}
	}
	return count, nil
}

func (h *DefaultHandler) Ping() (*StatusReply, error) {
	return &StatusReply{code: "PONG"}, nil
}

func (h *DefaultHandler) Subscribe(channels ...[]byte) (*MultiChannelWriter, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}
	ret := &MultiChannelWriter{Chans: make([]*ChannelWriter, 0, len(channels))}
	for _, key := range channels {
		Debugf("SUBSCRIBE on %s\n", key)
		cw := &ChannelWriter{
			FirstReply: []interface{}{
				"subscribe",
				key,
				1,
			},
			Channel: make(chan []interface{}),
		}
		if h.sub[string(key)] == nil {
			h.sub[string(key)] = []*ChannelWriter{cw}
		} else {
			h.sub[string(key)] = append(h.sub[string(key)], cw)
		}
		ret.Chans = append(ret.Chans, cw)
	}
	return ret, nil
}

func (h *DefaultHandler) Publish(key string, value []byte) (int, error) {
	if h.Database == nil || h.sub == nil {
		return 0, nil
	}
	//	Debugf("Publishing %s on %s\n", value, key)
	v, exists := h.sub[key]
	if !exists {
		return 0, nil
	}
	i := 0
	for _, c := range v {
		select {
		case c.Channel <- []interface{}{
			"message",
			key,
			value,
		}:
			i++
		default:
		}
	}
	return i, nil
}

func (h *DefaultHandler) Select(key string) error {
	if h.dbs == nil {
		h.dbs = map[int]*Database{0: h.Database}
	}
	index, err := strconv.Atoi(key)
	if err != nil {
		return err
	}
	h.dbs[h.currentDb] = h.Database
	h.currentDb = index
	if _, exists := h.dbs[index]; !exists {
		println("DB not exits, create ", index)
		h.dbs[index] = NewDatabase(nil)
	}
	h.Database = h.dbs[index]
	return nil
}

func (h *DefaultHandler) Monitor() (*MonitorReply, error) {
	return &MonitorReply{}, nil
}

var lock = make(chan bool, 1)

func (h *DefaultHandler) Incr(key string) (int, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	lock <- true

	temp, _ := strconv.Atoi(string(h.values[key]))
	temp = temp + 1
	h.values[key] = []byte(strconv.Itoa(temp))

	<-lock

	return temp, nil
}

func (h *DefaultHandler) Decr(key string) (int, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	lock <- true

	temp, _ := strconv.Atoi(string(h.values[key]))
	temp = temp - 1
	h.values[key] = []byte(strconv.Itoa(temp))

	<-lock

	return temp, nil
}

func (h *DefaultHandler) Expire(key, after string) (error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	d, _ := strconv.Atoi(after)

	time.AfterFunc(time.Duration(d) * time.Second, func() {
		h.Del(key)
	})

	return nil
}

func (h *DefaultHandler) Exists(key string) (int, error) {
	if h.Database == nil {
		h.Database = NewDatabase(nil)
	}

	_, exists := h.values[key]
	if exists {
		return 1, nil
	} else {
		return 0, nil
	}
}

func NewDefaultHandler() *DefaultHandler {
	db := NewDatabase(nil)
	ret := &DefaultHandler{
		Database:  db,
		currentDb: 0,
		dbs:       map[int]*Database{0: db},
	}
	return ret
}
<|endoftext|>


<commit_before>package device

import (
	"encoding/json"
	"github.com/Comcast/webpa-common/convey/conveymetric"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/Comcast/webpa-common/convey/conveyhttp"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/wrp"
	"github.com/Comcast/webpa-common/xhttp"
	"github.com/go-kit/kit/log"
	"github.com/gorilla/websocket"
)

const MaxDevicesHeader = "X-Xmidt-Max-Devices"

// Connector is a strategy interface for managing device connections to a server.
// Implementations are responsible for upgrading websocket connections and providing
// for explicit disconnection.
type Connector interface {
	// Connect upgrade an HTTP connection to a websocket and begins concurrent
	// management of the device.
	Connect(http.ResponseWriter, *http.Request, http.Header) (Interface, error)

	// Disconnect disconnects the device associated with the given id.
	// If the id was found, this method returns true.
	Disconnect(ID) bool

	// DisconnectIf iterates over all devices known to this manager, applying the
	// given predicate.  For any devices that result in true, this method disconnects them.
	// Note that this method may pause connections and disconnections while it is executing.
	// This method returns the number of devices that were disconnected.
	//
	// Only disconnection by ID is supported, which means that any identifier matching
	// the predicate will result in *all* duplicate devices under that ID being removed.
	//
	// No methods on this Manager should be called from within the predicate function, or
	// a deadlock will likely occur.
	DisconnectIf(func(ID) bool) int

	// DisconnectAll disconnects all devices from this instance, and returns the count of
	// devices disconnected.
	DisconnectAll() int
}

// Router handles dispatching messages to devices.
type Router interface {
	// Route dispatches a WRP request to exactly one device, identified by the ID
	// field of the request.  Route is synchronous, and honors the cancellation semantics
	// of the Request's context.
	Route(*Request) (*Response, error)
}

// Registry is the strategy interface for querying the set of connected devices.  Methods
// in this interface follow the Visitor pattern and are typically executed under a read lock.
type Registry interface {
	// Len returns the count of devices currently in this registry
	Len() int

	// Get returns the device associated with the given ID, if any
	Get(ID) (Interface, bool)

	// VisitAll applies the given visitor function to each device known to this manager.
	//
	// No methods on this Manager should be called from within the visitor function, or
	// a deadlock will likely occur.
	VisitAll(func(Interface) bool) int
}

// Manager supplies a hub for connecting and disconnecting devices as well as
// an access point for obtaining device metadata.
type Manager interface {
	Connector
	Router
	Registry
}

// NewManager constructs a Manager from a set of options.  A ConnectionFactory will be
// created from the options if one is not supplied.
func NewManager(o *Options) Manager {
	var (
		logger   = o.logger()
		measures = NewMeasures(o.metricsProvider())
	)

	return &manager{
		logger:   logger,
		errorLog: logging.Error(logger),
		debugLog: logging.Debug(logger),

		readDeadline:     NewDeadline(o.idlePeriod(), o.now()),
		writeDeadline:    NewDeadline(o.writeTimeout(), o.now()),
		upgrader:         o.upgrader(),
		conveyTranslator: conveyhttp.NewHeaderTranslator("", nil),
		devices: newRegistry(registryOptions{
			Logger:   logger,
			Limit:    o.maxDevices(),
			Measures: measures,
		}),
		conveyHWMetric: conveymetric.NewConveyMetric(o.metricsProvider(), "hw-model", "hardware_model"),

		deviceMessageQueueSize: o.deviceMessageQueueSize(),
		pingPeriod:             o.pingPeriod(),

		listeners: o.listeners(),
		measures:  measures,
	}
}

// manager is the internal Manager implementation.
type manager struct {
	logger   log.Logger
	errorLog log.Logger
	debugLog log.Logger

	readDeadline     func() time.Time
	writeDeadline    func() time.Time
	upgrader         *websocket.Upgrader
	conveyTranslator conveyhttp.HeaderTranslator

	devices        *registry
	conveyHWMetric conveymetric.CMetric

	deviceMessageQueueSize int
	pingPeriod             time.Duration

	listeners []Listener
	measures  Measures
}

func (m *manager) Connect(response http.ResponseWriter, request *http.Request, responseHeader http.Header) (Interface, error) {
	m.debugLog.Log(logging.MessageKey(), "device connect", "url", request.URL)
	id, ok := GetID(request.Context())
	if !ok {
		xhttp.WriteError(
			response,
			http.StatusInternalServerError,
			ErrorMissingDeviceNameContext,
		)

		return nil, ErrorMissingDeviceNameContext
	}

	d := newDevice(deviceOptions{ID: id, QueueSize: m.deviceMessageQueueSize, Logger: m.logger})
	convey, conveyErr := m.conveyTranslator.FromHeader(request.Header)
	if conveyErr == nil {
		d.infoLog.Log("convey", convey)
	} else if conveyErr != conveyhttp.ErrMissingHeader {
		d.errorLog.Log(logging.MessageKey(), "badly formatted convey data", logging.ErrorKey(), conveyErr)
	}

	c, err := m.upgrader.Upgrade(response, request, responseHeader)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "failed websocket upgrade", logging.ErrorKey(), err)
		return nil, err
	}

	d.debugLog.Log(logging.MessageKey(), "websocket upgrade complete", "localAddress", c.LocalAddr().String())

	pinger, err := NewPinger(c, m.measures.Ping, []byte(d.ID()), m.writeDeadline)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to create pinger", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	if err := m.devices.add(d); err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to register device", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	event := &Event{
		Type:   Connect,
		Device: d,
	}

	if conveyErr == nil {
		bytes, err := json.Marshal(convey)
		if err == nil {
			event.Format = wrp.JSON
			event.Contents = bytes
		} else {
			d.errorLog.Log(logging.MessageKey(), "unable to marshal the convey header", logging.ErrorKey(), err)
		}
	}

	metricClosure, err := m.conveyHWMetric.Update(convey)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "failed to update convey metrics", logging.ErrorKey(), err)
	}

	d.conveyClosure = metricClosure
	m.dispatch(event)

	SetPongHandler(c, m.measures.Pong, m.readDeadline)
	closeOnce := new(sync.Once)
	go m.readPump(d, InstrumentReader(c, d.statistics), closeOnce)
	go m.writePump(d, InstrumentWriter(c, d.statistics), pinger, closeOnce)

	return d, nil
}

func (m *manager) dispatch(e *Event) {
	for _, listener := range m.listeners {
		listener(e)
	}
}

// pumpClose handles the proper shutdown and logging of a device's pumps.
// This method should be executed within a sync.Once, so that it only executes
// once for a given device.
//
// Note that the write pump does additional cleanup.  In particular, the write pump
// dispatches message failed events for any messages that were waiting to be delivered
// at the time of pump closure.
func (m *manager) pumpClose(d *device, c io.Closer, pumpError error) {
	// remove will invoke requestClose()
	m.devices.remove(d.id)

	closeError := c.Close()

	d.errorLog.Log(logging.MessageKey(), "Closed device connection",
		"closeError", closeError, "pumpError", pumpError,
		"finalStatistics", d.Statistics().String())

	m.dispatch(
		&Event{
			Type:   Disconnect,
			Device: d,
		},
	)
	d.conveyClosure()
}

// readPump is the goroutine which handles the stream of WRP messages from a device.
// This goroutine exits when any error occurs on the connection.
func (m *manager) readPump(d *device, r ReadCloser, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "readPump exiting")
	d.debugLog.Log(logging.MessageKey(), "readPump starting")

	var (
		readError error
		decoder   = wrp.NewDecoder(nil, wrp.Msgpack)
	)

	// all the read pump has to do is ensure the device and the connection are closed
	// it is the write pump's responsibility to do further cleanup
	defer closeOnce.Do(func() { m.pumpClose(d, r, readError) })

	for {
		messageType, data, readError := r.ReadMessage()
		if readError != nil {
			d.errorLog.Log(logging.MessageKey(), "read error", logging.ErrorKey(), readError)
			return
		}

		if messageType != websocket.BinaryMessage {
			d.errorLog.Log(logging.MessageKey(), "skipping non-binary frame", "messageType", messageType)
			continue
		}

		var (
			message = new(wrp.Message)
			event   = Event{
				Type:     MessageReceived,
				Device:   d,
				Message:  message,
				Format:   wrp.Msgpack,
				Contents: data,
			}
		)

		decoder.ResetBytes(data)
		err := decoder.Decode(message)
		decoder.ResetBytes(nil)
		if err != nil {
			d.errorLog.Log(logging.MessageKey(), "skipping malformed WRP message", logging.ErrorKey(), err)
			continue
		}

		if message.Type == wrp.SimpleRequestResponseMessageType {
			m.measures.RequestResponse.Add(1.0)
		}

		// update any waiting transaction
		if message.IsTransactionPart() {
			err := d.transactions.Complete(
				message.TransactionKey(),
				&Response{
					Device:   d,
					Message:  message,
					Format:   wrp.Msgpack,
					Contents: data,
				},
			)

			if err != nil {
				d.errorLog.Log(logging.MessageKey(), "Error while completing transaction", "transactionKey", message.TransactionKey(), logging.ErrorKey(), err)
				event.Type = TransactionBroken
				event.Error = err
			} else {
				event.Type = TransactionComplete
			}
		}

		m.dispatch(&event)
	}
}

// writePump is the goroutine which services messages addressed to the device.
// this goroutine exits when either an explicit shutdown is requested or any
// error occurs on the connection.
func (m *manager) writePump(d *device, w WriteCloser, pinger func() error, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "writePump exiting")
	d.debugLog.Log(logging.MessageKey(), "writePump starting")

	var (
		envelope   *envelope
		encoder    = wrp.NewEncoder(nil, wrp.Msgpack)
		writeError error

		pingTicker = time.NewTicker(m.pingPeriod)
	)

	// cleanup: we not only ensure that the device and connection are closed but also
	// ensure that any messages that were waiting and/or failed are dispatched to
	// the configured listener
	defer func() {
		pingTicker.Stop()
		closeOnce.Do(func() { m.pumpClose(d, w, writeError) })

		// notify listener of any message that just now failed
		// any writeError is passed via this event
		if envelope != nil {
			m.dispatch(&Event{
				Type:     MessageFailed,
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			})
		}

		// drain the messages, dispatching them as message failed events.  we never close
		// the message channel, so just drain until a receive would block.
		//
		// Nil is passed explicitly as the error to indicate that these messages failed due
		// to the device disconnecting, not due to an actual I/O error.
		for {
			select {
			case undeliverable := <-d.messages:
				d.errorLog.Log(logging.MessageKey(), "undeliverable message", "deviceMessage", undeliverable)
				m.dispatch(&Event{
					Type:     MessageFailed,
					Device:   d,
					Message:  undeliverable.request.Message,
					Format:   undeliverable.request.Format,
					Contents: undeliverable.request.Contents,
					Error:    writeError,
				})
			default:
				return
			}
		}
	}()

	for writeError == nil {
		envelope = nil

		select {
		case <-d.shutdown:
			d.debugLog.Log(logging.MessageKey(), "explicit shutdown")
			writeError = w.Close()
			return

		case envelope = <-d.messages:
			var frameContents []byte
			if envelope.request.Format == wrp.Msgpack && len(envelope.request.Contents) > 0 {
				frameContents = envelope.request.Contents
			} else {
				// if the request was in a format other than Msgpack, or if the caller did not pass
				// Contents, then do the encoding here.
				encoder.ResetBytes(&frameContents)
				writeError = encoder.Encode(envelope.request.Message)
				encoder.ResetBytes(nil)
			}

			if writeError == nil {
				writeError = w.WriteMessage(websocket.BinaryMessage, frameContents)
			}

			event := Event{
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			}

			if writeError != nil {
				envelope.complete <- writeError
				event.Type = MessageFailed
			} else {
				event.Type = MessageSent
			}

			close(envelope.complete)
			m.dispatch(&event)

		case <-pingTicker.C:
			writeError = pinger()
		}
	}
}

func (m *manager) Disconnect(id ID) bool {
	_, ok := m.devices.remove(id)
	return ok
}

func (m *manager) DisconnectIf(filter func(ID) bool) int {
	return m.devices.removeIf(func(d *device) bool {
		return filter(d.id)
	})
}

func (m *manager) DisconnectAll() int {
	return m.devices.removeAll()
}

func (m *manager) Len() int {
	return m.devices.len()
}

func (m *manager) Get(id ID) (Interface, bool) {
	return m.devices.get(id)
}

func (m *manager) VisitAll(visitor func(Interface) bool) int {
	return m.devices.visit(func(d *device) bool {
		return visitor(d)
	})
}

func (m *manager) Route(request *Request) (*Response, error) {
	if destination, err := request.ID(); err != nil {
		return nil, err
	} else if d, ok := m.devices.get(destination); ok {
		return d.Send(request)
	} else {
		return nil, ErrorDeviceNotFound
	}
}
<commit_msg>update manager to provide a gauge<commit_after>package device

import (
	"encoding/json"
	"github.com/Comcast/webpa-common/convey/conveymetric"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/Comcast/webpa-common/convey/conveyhttp"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/wrp"
	"github.com/Comcast/webpa-common/xhttp"
	"github.com/go-kit/kit/log"
	"github.com/gorilla/websocket"
)

const MaxDevicesHeader = "X-Xmidt-Max-Devices"

// Connector is a strategy interface for managing device connections to a server.
// Implementations are responsible for upgrading websocket connections and providing
// for explicit disconnection.
type Connector interface {
	// Connect upgrade an HTTP connection to a websocket and begins concurrent
	// management of the device.
	Connect(http.ResponseWriter, *http.Request, http.Header) (Interface, error)

	// Disconnect disconnects the device associated with the given id.
	// If the id was found, this method returns true.
	Disconnect(ID) bool

	// DisconnectIf iterates over all devices known to this manager, applying the
	// given predicate.  For any devices that result in true, this method disconnects them.
	// Note that this method may pause connections and disconnections while it is executing.
	// This method returns the number of devices that were disconnected.
	//
	// Only disconnection by ID is supported, which means that any identifier matching
	// the predicate will result in *all* duplicate devices under that ID being removed.
	//
	// No methods on this Manager should be called from within the predicate function, or
	// a deadlock will likely occur.
	DisconnectIf(func(ID) bool) int

	// DisconnectAll disconnects all devices from this instance, and returns the count of
	// devices disconnected.
	DisconnectAll() int
}

// Router handles dispatching messages to devices.
type Router interface {
	// Route dispatches a WRP request to exactly one device, identified by the ID
	// field of the request.  Route is synchronous, and honors the cancellation semantics
	// of the Request's context.
	Route(*Request) (*Response, error)
}

// Registry is the strategy interface for querying the set of connected devices.  Methods
// in this interface follow the Visitor pattern and are typically executed under a read lock.
type Registry interface {
	// Len returns the count of devices currently in this registry
	Len() int

	// Get returns the device associated with the given ID, if any
	Get(ID) (Interface, bool)

	// VisitAll applies the given visitor function to each device known to this manager.
	//
	// No methods on this Manager should be called from within the visitor function, or
	// a deadlock will likely occur.
	VisitAll(func(Interface) bool) int
}

// Manager supplies a hub for connecting and disconnecting devices as well as
// an access point for obtaining device metadata.
type Manager interface {
	Connector
	Router
	Registry
}

// NewManager constructs a Manager from a set of options.  A ConnectionFactory will be
// created from the options if one is not supplied.
func NewManager(o *Options) Manager {
	var (
		logger   = o.logger()
		measures = NewMeasures(o.metricsProvider())
	)

	return &manager{
		logger:   logger,
		errorLog: logging.Error(logger),
		debugLog: logging.Debug(logger),

		readDeadline:     NewDeadline(o.idlePeriod(), o.now()),
		writeDeadline:    NewDeadline(o.writeTimeout(), o.now()),
		upgrader:         o.upgrader(),
		conveyTranslator: conveyhttp.NewHeaderTranslator("", nil),
		devices: newRegistry(registryOptions{
			Logger:   logger,
			Limit:    o.maxDevices(),
			Measures: measures,
		}),
		conveyHWMetric: conveymetric.NewConveyMetric(o.metricsProvider().NewGauge("hardware"), "hw-model", "model"),

		deviceMessageQueueSize: o.deviceMessageQueueSize(),
		pingPeriod:             o.pingPeriod(),

		listeners: o.listeners(),
		measures:  measures,
	}
}

// manager is the internal Manager implementation.
type manager struct {
	logger   log.Logger
	errorLog log.Logger
	debugLog log.Logger

	readDeadline     func() time.Time
	writeDeadline    func() time.Time
	upgrader         *websocket.Upgrader
	conveyTranslator conveyhttp.HeaderTranslator

	devices        *registry
	conveyHWMetric conveymetric.CMetric

	deviceMessageQueueSize int
	pingPeriod             time.Duration

	listeners []Listener
	measures  Measures
}

func (m *manager) Connect(response http.ResponseWriter, request *http.Request, responseHeader http.Header) (Interface, error) {
	m.debugLog.Log(logging.MessageKey(), "device connect", "url", request.URL)
	id, ok := GetID(request.Context())
	if !ok {
		xhttp.WriteError(
			response,
			http.StatusInternalServerError,
			ErrorMissingDeviceNameContext,
		)

		return nil, ErrorMissingDeviceNameContext
	}

	d := newDevice(deviceOptions{ID: id, QueueSize: m.deviceMessageQueueSize, Logger: m.logger})
	convey, conveyErr := m.conveyTranslator.FromHeader(request.Header)
	if conveyErr == nil {
		d.infoLog.Log("convey", convey)
	} else if conveyErr != conveyhttp.ErrMissingHeader {
		d.errorLog.Log(logging.MessageKey(), "badly formatted convey data", logging.ErrorKey(), conveyErr)
	}

	c, err := m.upgrader.Upgrade(response, request, responseHeader)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "failed websocket upgrade", logging.ErrorKey(), err)
		return nil, err
	}

	d.debugLog.Log(logging.MessageKey(), "websocket upgrade complete", "localAddress", c.LocalAddr().String())

	pinger, err := NewPinger(c, m.measures.Ping, []byte(d.ID()), m.writeDeadline)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to create pinger", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	if err := m.devices.add(d); err != nil {
		d.errorLog.Log(logging.MessageKey(), "unable to register device", logging.ErrorKey(), err)
		c.Close()
		return nil, err
	}

	event := &Event{
		Type:   Connect,
		Device: d,
	}

	if conveyErr == nil {
		bytes, err := json.Marshal(convey)
		if err == nil {
			event.Format = wrp.JSON
			event.Contents = bytes
		} else {
			d.errorLog.Log(logging.MessageKey(), "unable to marshal the convey header", logging.ErrorKey(), err)
		}
	}

	metricClosure, err := m.conveyHWMetric.Update(convey)
	if err != nil {
		d.errorLog.Log(logging.MessageKey(), "failed to update convey metrics", logging.ErrorKey(), err)
	}

	d.conveyClosure = metricClosure
	m.dispatch(event)

	SetPongHandler(c, m.measures.Pong, m.readDeadline)
	closeOnce := new(sync.Once)
	go m.readPump(d, InstrumentReader(c, d.statistics), closeOnce)
	go m.writePump(d, InstrumentWriter(c, d.statistics), pinger, closeOnce)

	return d, nil
}

func (m *manager) dispatch(e *Event) {
	for _, listener := range m.listeners {
		listener(e)
	}
}

// pumpClose handles the proper shutdown and logging of a device's pumps.
// This method should be executed within a sync.Once, so that it only executes
// once for a given device.
//
// Note that the write pump does additional cleanup.  In particular, the write pump
// dispatches message failed events for any messages that were waiting to be delivered
// at the time of pump closure.
func (m *manager) pumpClose(d *device, c io.Closer, pumpError error) {
	// remove will invoke requestClose()
	m.devices.remove(d.id)

	closeError := c.Close()

	d.errorLog.Log(logging.MessageKey(), "Closed device connection",
		"closeError", closeError, "pumpError", pumpError,
		"finalStatistics", d.Statistics().String())

	m.dispatch(
		&Event{
			Type:   Disconnect,
			Device: d,
		},
	)
	d.conveyClosure()
}

// readPump is the goroutine which handles the stream of WRP messages from a device.
// This goroutine exits when any error occurs on the connection.
func (m *manager) readPump(d *device, r ReadCloser, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "readPump exiting")
	d.debugLog.Log(logging.MessageKey(), "readPump starting")

	var (
		readError error
		decoder   = wrp.NewDecoder(nil, wrp.Msgpack)
	)

	// all the read pump has to do is ensure the device and the connection are closed
	// it is the write pump's responsibility to do further cleanup
	defer closeOnce.Do(func() { m.pumpClose(d, r, readError) })

	for {
		messageType, data, readError := r.ReadMessage()
		if readError != nil {
			d.errorLog.Log(logging.MessageKey(), "read error", logging.ErrorKey(), readError)
			return
		}

		if messageType != websocket.BinaryMessage {
			d.errorLog.Log(logging.MessageKey(), "skipping non-binary frame", "messageType", messageType)
			continue
		}

		var (
			message = new(wrp.Message)
			event   = Event{
				Type:     MessageReceived,
				Device:   d,
				Message:  message,
				Format:   wrp.Msgpack,
				Contents: data,
			}
		)

		decoder.ResetBytes(data)
		err := decoder.Decode(message)
		decoder.ResetBytes(nil)
		if err != nil {
			d.errorLog.Log(logging.MessageKey(), "skipping malformed WRP message", logging.ErrorKey(), err)
			continue
		}

		if message.Type == wrp.SimpleRequestResponseMessageType {
			m.measures.RequestResponse.Add(1.0)
		}

		// update any waiting transaction
		if message.IsTransactionPart() {
			err := d.transactions.Complete(
				message.TransactionKey(),
				&Response{
					Device:   d,
					Message:  message,
					Format:   wrp.Msgpack,
					Contents: data,
				},
			)

			if err != nil {
				d.errorLog.Log(logging.MessageKey(), "Error while completing transaction", "transactionKey", message.TransactionKey(), logging.ErrorKey(), err)
				event.Type = TransactionBroken
				event.Error = err
			} else {
				event.Type = TransactionComplete
			}
		}

		m.dispatch(&event)
	}
}

// writePump is the goroutine which services messages addressed to the device.
// this goroutine exits when either an explicit shutdown is requested or any
// error occurs on the connection.
func (m *manager) writePump(d *device, w WriteCloser, pinger func() error, closeOnce *sync.Once) {
	defer d.debugLog.Log(logging.MessageKey(), "writePump exiting")
	d.debugLog.Log(logging.MessageKey(), "writePump starting")

	var (
		envelope   *envelope
		encoder    = wrp.NewEncoder(nil, wrp.Msgpack)
		writeError error

		pingTicker = time.NewTicker(m.pingPeriod)
	)

	// cleanup: we not only ensure that the device and connection are closed but also
	// ensure that any messages that were waiting and/or failed are dispatched to
	// the configured listener
	defer func() {
		pingTicker.Stop()
		closeOnce.Do(func() { m.pumpClose(d, w, writeError) })

		// notify listener of any message that just now failed
		// any writeError is passed via this event
		if envelope != nil {
			m.dispatch(&Event{
				Type:     MessageFailed,
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			})
		}

		// drain the messages, dispatching them as message failed events.  we never close
		// the message channel, so just drain until a receive would block.
		//
		// Nil is passed explicitly as the error to indicate that these messages failed due
		// to the device disconnecting, not due to an actual I/O error.
		for {
			select {
			case undeliverable := <-d.messages:
				d.errorLog.Log(logging.MessageKey(), "undeliverable message", "deviceMessage", undeliverable)
				m.dispatch(&Event{
					Type:     MessageFailed,
					Device:   d,
					Message:  undeliverable.request.Message,
					Format:   undeliverable.request.Format,
					Contents: undeliverable.request.Contents,
					Error:    writeError,
				})
			default:
				return
			}
		}
	}()

	for writeError == nil {
		envelope = nil

		select {
		case <-d.shutdown:
			d.debugLog.Log(logging.MessageKey(), "explicit shutdown")
			writeError = w.Close()
			return

		case envelope = <-d.messages:
			var frameContents []byte
			if envelope.request.Format == wrp.Msgpack && len(envelope.request.Contents) > 0 {
				frameContents = envelope.request.Contents
			} else {
				// if the request was in a format other than Msgpack, or if the caller did not pass
				// Contents, then do the encoding here.
				encoder.ResetBytes(&frameContents)
				writeError = encoder.Encode(envelope.request.Message)
				encoder.ResetBytes(nil)
			}

			if writeError == nil {
				writeError = w.WriteMessage(websocket.BinaryMessage, frameContents)
			}

			event := Event{
				Device:   d,
				Message:  envelope.request.Message,
				Format:   envelope.request.Format,
				Contents: envelope.request.Contents,
				Error:    writeError,
			}

			if writeError != nil {
				envelope.complete <- writeError
				event.Type = MessageFailed
			} else {
				event.Type = MessageSent
			}

			close(envelope.complete)
			m.dispatch(&event)

		case <-pingTicker.C:
			writeError = pinger()
		}
	}
}

func (m *manager) Disconnect(id ID) bool {
	_, ok := m.devices.remove(id)
	return ok
}

func (m *manager) DisconnectIf(filter func(ID) bool) int {
	return m.devices.removeIf(func(d *device) bool {
		return filter(d.id)
	})
}

func (m *manager) DisconnectAll() int {
	return m.devices.removeAll()
}

func (m *manager) Len() int {
	return m.devices.len()
}

func (m *manager) Get(id ID) (Interface, bool) {
	return m.devices.get(id)
}

func (m *manager) VisitAll(visitor func(Interface) bool) int {
	return m.devices.visit(func(d *device) bool {
		return visitor(d)
	})
}

func (m *manager) Route(request *Request) (*Response, error) {
	if destination, err := request.ID(); err != nil {
		return nil, err
	} else if d, ok := m.devices.get(destination); ok {
		return d.Send(request)
	} else {
		return nil, ErrorDeviceNotFound
	}
}
<|endoftext|>


<commit_before>// Package beta provides algorithms for working with beta distributions.
//
// https://en.wikipedia.org/wiki/Beta_distribution
package beta

import (
	"math"
)

// Self represents a particular distribution from the family.
type Self struct {
	α float64
	β float64
	a float64
	b float64
}

// New returns a beta distribution with α and β on [a, b].
func New(α, β, a, b float64) *Self {
	return &Self{α, β, a, b}
}

// CDF evaluates the CDF of the distribution.
func (s *Self) CDF(points []float64) []float64 {
	values := make([]float64, len(points))

	α, β, k, b := s.α, s.β, s.b-s.a, s.a
	logB := logBeta(α, β)

	for i, x := range points {
		values[i] = incBeta((x-b)/k, α, β, logB)
	}

	return values
}

// InvCDF evaluates the inverse CDF of the distribution.
func (s *Self) InvCDF(points []float64) []float64 {
	values := make([]float64, len(points))

	α, β, k, b := s.α, s.β, s.b-s.a, s.a
	logB := logBeta(α, β)

	for i, p := range points {
		values[i] = k*invIncBeta(p, α, β, logB) + b
	}

	return values
}

func incBeta(x, p, q, logB float64) float64 {
	// The code is based on a C implementation by John Burkardt.
	// http://people.sc.fsu.edu/~jburkardt/c_src/asa109/asa109.html

	const (
		acu = 0.1e-14
	)

	if x <= 0 {
		return 0
	}
	if 1 <= x {
		return 1
	}

	sum := p + q
	px, qx := x, 1-x

	var flip bool
	if p < sum*x {
		p, px, q, qx = q, qx, p, px
		flip = true
	}

	// Use Soper’s reduction formula.
	rx := px / qx

	ns := int(q + qx*sum)
	if ns == 0 {
		rx = px
	}

	ai := 1
	temp := q - float64(ai)
	term := 1.0

	α := 1.0

	for {
		term = term * temp * rx / (p + float64(ai))

		α += term

		temp = math.Abs(term)
		if temp <= acu && temp <= acu*α {
			break
		}

		ai++
		ns--

		if 0 < ns {
			temp = q - float64(ai)
		} else if ns == 0 {
			temp = q - float64(ai)
			rx = px
		} else {
			temp = sum
			sum += 1
		}
	}

	// Applied Statistics. Algorithm AS 109
	// http://www.jstor.org/discover/10.2307/2346887
	α = α * math.Exp(p*math.Log(px)+(q-1)*math.Log(qx)-logB) / p

	if flip {
		return 1 - α
	} else {
		return α
	}
}

func invIncBeta(α, p, q, logB float64) float64 {
	// The code is based on a C implementation by John Burkardt.
	// http://people.sc.fsu.edu/~jburkardt/c_src/asa109/asa109.html

	const (
		// Applied Statistics. Algorithm AS R83
		// http://www.jstor.org/discover/10.2307/2347779
		//
		// The machine-dependent smallest allowable exponent of 10 to avoid
		// floating-point underflow error.
		sae = -30
	)

	if α <= 0 {
		return 0
	}
	if 1 <= α {
		return 1
	}

	var flip bool
	if 0.5 < α {
		α = 1 - α
		p, q = q, p
		flip = true
	}

	// An approximation x₀ to x if found from (cf. Scheffé and Tukey, 1944)
	//
	//     (1 + x₀)/(1 - x₀) = (4*p + 2*q - 2)/χ²(α)
	//
	// where χ²(α) is the upper α point of the χ² distribution with 2*q degrees
	// of freedom and is obtained from Wilson and Hilferty’s approximation (cf.
	// Wilson and Hilferty, 1931)
	//
	//     χ²(α) = 2*q*(1 - 1/(9*q) + y(α) * sqrt(1/(9*q)))**3,
	//
	// y(α) being Hastings’ approximation (cf. Hastings, 1955) for the upper α
	// point of the standard normal distribution. If χ²(α) < 0, then
	//
	//     x₀ = 1 - ((1 - α)*q*B(p, q))**(1/q).
	//
	// Again if (4*p + 2*q - 2)/χ²(α) does not exceed 1, x₀ is obtained from
	//
	//     x₀ = (α*p*B(p, q))**(1/p).
	//
	// Applied Statistics. Algorithm AS 46
	// http://www.jstor.org/discover/10.2307/2346798
	var x float64

	var y, r, t float64

	r = math.Sqrt(-math.Log(α * α))
	y = r - (2.30753+0.27061*r)/(1+(0.99229+0.04481*r)*r)

	if 1 < p && 1 < q {
		// For p and q > 1, the approximation given by Carter (1947), which
		// improves the Fisher–Cochran formula, is generally better. For other
		// values of p and q en empirical investigation has shown that the
		// approximation given in AS 64 is adequate.
		//
		// Applied Statistics. Algorithm AS 109
		// http://www.jstor.org/discover/10.2307/2346887
		r = (y*y - 3) / 6
		s := 1 / (2*p - 1)
		t = 1 / (2*q - 1)
		h := 2 / (s + t)
		w := y*math.Sqrt(h+r)/h - (t-s)*(r+5/6-2/(3*h))
		x = p / (p + q*math.Exp(2*w))
	} else {
		t = 1 / (9 * q)
		t = 2 * q * math.Pow(1-t+y*math.Sqrt(t), 3)
		if t <= 0 {
			x = 1 - math.Exp((math.Log((1-α)*q)+logB)/q)
		} else {
			t = 2 * (2*p + q - 1) / t
			if t <= 1 {
				x = math.Exp((math.Log(α*p) + logB) / p)
			} else {
				x = 1 - 2/(t+1)
			}
		}
	}

	if x < 0.0001 {
		x = 0.0001
	} else if 0.9999 < x {
		x = 0.9999
	}

	// The final solution is obtained by the Newton–Raphson method from the
	// relation
	//
	//     x[i] = x[i-1] - f(x[i-1])/f'(x[i-1])
	//
	// where
	//
	//     f(x) = I(x, p, q) - α.
	//
	// Applied Statistics. Algorithm AS 46
	// http://www.jstor.org/discover/10.2307/2346798
	r = 1 - p
	t = 1 - q
	yprev := 0.0
	sq := 1.0
	prev := 1.0

	// Applied Statistics. Algorithm AS R83
	// http://www.jstor.org/discover/10.2307/2347779
	fpu := math.Pow10(sae)
	acu := fpu
	if e := int(-5/p/p - 1/math.Pow(α, 0.2) - 13); e > sae {
		acu = math.Pow10(e)
	}

	var tx, g, adj float64

outer:
	for {
		// Applied Statistics. Algorithm AS 109
		// http://www.jstor.org/discover/10.2307/2346887
		y = incBeta(x, p, q, logB)
		y = (y - α) * math.Exp(logB+r*math.Log(x)+t*math.Log(1-x))

		// Applied Statistics. Algorithm AS R83
		// http://www.jstor.org/discover/10.2307/2347779
		if y*yprev <= 0 {
			prev = math.Max(sq, fpu)
		}

		// Applied Statistics. Algorithm AS 109
		// http://www.jstor.org/discover/10.2307/2346887
		g = 1
		for {
			for {
				adj = g * y
				sq = adj * adj

				if sq < prev {
					tx = x - adj

					if 0 <= tx && tx <= 1 {
						break
					}
				}
				g /= 3
			}

			if prev <= acu || y*y <= acu {
				x = tx
				break outer
			}

			if tx != 0 && tx != 1 {
				break
			}

			g /= 3
		}

		if tx == x {
			break
		}

		x = tx
		yprev = y
	}

	if flip {
		return 1 - x
	} else {
		return x
	}
}

func logBeta(x, y float64) float64 {
	z, _ := math.Lgamma(x + y)
	x, _ = math.Lgamma(x)
	y, _ = math.Lgamma(y)

	return x + y - z
}
<commit_msg>Polishing the comments in beta<commit_after>// Package beta provides algorithms for working with beta distributions.
//
// https://en.wikipedia.org/wiki/Beta_distribution
package beta

import (
	"math"
)

// Self represents a particular distribution from the family.
type Self struct {
	α float64
	β float64
	a float64
	b float64
}

// New returns a beta distribution with α and β on [a, b].
func New(α, β, a, b float64) *Self {
	return &Self{α, β, a, b}
}

// CDF evaluates the CDF of the distribution.
func (s *Self) CDF(points []float64) []float64 {
	values := make([]float64, len(points))

	α, β, k, b := s.α, s.β, s.b-s.a, s.a
	logB := logBeta(α, β)

	for i, x := range points {
		values[i] = incBeta((x-b)/k, α, β, logB)
	}

	return values
}

// InvCDF evaluates the inverse CDF of the distribution.
func (s *Self) InvCDF(points []float64) []float64 {
	values := make([]float64, len(points))

	α, β, k, b := s.α, s.β, s.b-s.a, s.a
	logB := logBeta(α, β)

	for i, p := range points {
		values[i] = k*invIncBeta(p, α, β, logB) + b
	}

	return values
}

// incBeta computes the incomplete beta function.
//
// The code is based on a C implementation by John Burkardt.
// http://people.sc.fsu.edu/~jburkardt/c_src/asa109/asa109.html
//
// The original algorithm was published in Applied Statistics and described
// below.
//
// Algorithm AS 63
// http://www.jstor.org/stable/2346797
//
// The function uses the method discussed by Soper (1921). If p is not less
// than (p + q)x and the integral part of q + (1 - x)(p + q) is a positive
// integer, say s, reductions are made up to s times “by parts” using the
// recurrence relation
//
//                     Γ(p+q)
//     I(x, p, q) = ----------- x^p (1-x)^(q-1) + I(x, p+1, q-1)
//                  Γ(p+1) Γ(q)
//
// and then reductions are continued by “raising p” with the recurrence
// relation
//
//                          Γ(p+q)
//     I(x, p+s, q-s) = --------------- x^(p+s) (1-x)^(q-s) + I(x, p+s+1, q-s)
//                      Γ(p+s+1) Γ(q-s)
//
// If s is not a positive integer, reductions are made only by “raising p.”
// The process of reduction is terminated when the relative contribution to the
// integral is not greater than the value of ACU. If p is less than (p + q)x,
// I(1-x, q, p) is first calculated by the above procedure and then I(x, p, q)
// is obtained from the relation
//
//     I(x, p, q) = 1 - I(1-x, p, q).
//
// Soper (1921) demonstrated that the expansion of I(x, p, q) by “parts” and
// “raising p” method as described above converges more rapidly than any other
// series expansions.
func incBeta(x, p, q, logB float64) float64 {
	const (
		acu = 0.1e-14
	)

	if x <= 0 {
		return 0
	}
	if 1 <= x {
		return 1
	}

	psq := p + q
	pbase, qbase := x, 1-x

	flip := false
	if p < psq*x {
		p, q, pbase, qbase = q, p, qbase, pbase
		flip = true
	}

	term, ai := 1.0, 1.0
	temp := q - ai

	rx := pbase / qbase
	ns := int(q + qbase*psq)
	if ns == 0 {
		rx = pbase
	}

	α := 1.0

	for {
		term = term * temp * rx / (p + ai)

		α += term

		temp = math.Abs(term)
		if temp <= acu && temp <= acu*α {
			break
		}

		ai++
		ns--

		if 0 < ns {
			temp = q - ai
		} else if ns == 0 {
			temp = q - ai
			rx = pbase
		} else {
			temp = psq
			psq++
		}
	}

	// Remark AS R19 and Algorithm AS 109
	// http://www.jstor.org/stable/2346887
	α = α * math.Exp(p*math.Log(pbase)+(q-1)*math.Log(qbase)-logB) / p

	if flip {
		return 1 - α
	} else {
		return α
	}
}

// invIncBeta computes the inverse of the incomplete beta function.
//
// The code is based on a C implementation by John Burkardt.
// http://people.sc.fsu.edu/~jburkardt/c_src/asa109/asa109.html
//
// The original algorithm was published in Applied Statistics and described
// below.
//
// Algorithm AS 64
// http://www.jstor.org/stable/2346798
//
// An approximation x₀ to x if found from (cf. Scheffé and Tukey, 1944)
//
//     (1 + x₀)/(1 - x₀) = (4*p + 2*q - 2)/χ²(α)
//
// where χ²(α) is the upper α point of the χ² distribution with 2*q degrees
// of freedom and is obtained from Wilson and Hilferty’s approximation (cf.
// Wilson and Hilferty, 1931)
//
//     χ²(α) = 2*q*(1 - 1/(9*q) + y(α) * sqrt(1/(9*q)))**3,
//
// y(α) being Hastings’ approximation (cf. Hastings, 1955) for the upper α
// point of the standard normal distribution. If χ²(α) < 0, then
//
//     x₀ = 1 - ((1 - α)*q*B(p, q))**(1/q).
//
// Again if (4*p + 2*q - 2)/χ²(α) does not exceed 1, x₀ is obtained from
//
//     x₀ = (α*p*B(p, q))**(1/p).
//
// The final solution is obtained by the Newton–Raphson method from the
// relation
//
//     x[i] = x[i-1] - f(x[i-1])/f'(x[i-1])
//
// where
//
//     f(x) = I(x, p, q) - α.
func invIncBeta(α, p, q, logB float64) float64 {
	const (
		// Remark AS R83
		// http://www.jstor.org/stable/2347779
		sae = -30
	)

	if α <= 0 {
		return 0
	}
	if 1 <= α {
		return 1
	}

	flip := false
	if 0.5 < α {
		α = 1 - α
		p, q = q, p
		flip = true
	}

	x := math.Sqrt(-math.Log(α * α))
	y := x - (2.30753+0.27061*x)/(1+(0.99229+0.04481*x)*x)

	if 1 < p && 1 < q {
		// Remark AS R19 and Algorithm AS 109
		// http://www.jstor.org/stable/2346887
		//
		// For p and q > 1, the approximation given by Carter (1947), which
		// improves the Fisher–Cochran formula, is generally better. For other
		// values of p and q en empirical investigation has shown that the
		// approximation given in AS 64 is adequate.
		r := (y*y - 3) / 6
		s := 1 / (2*p - 1)
		t := 1 / (2*q - 1)
		h := 2 / (s + t)
		w := y*math.Sqrt(h+r)/h - (t-s)*(r+5/6-2/(3*h))
		x = p / (p + q*math.Exp(2*w))
	} else {
		t := 1 / (9 * q)
		t = 2 * q * math.Pow(1-t+y*math.Sqrt(t), 3)
		if t <= 0 {
			x = 1 - math.Exp((math.Log((1-α)*q)+logB)/q)
		} else {
			t = 2 * (2*p + q - 1) / t
			if t <= 1 {
				x = math.Exp((math.Log(α*p) + logB) / p)
			} else {
				x = 1 - 2/(t+1)
			}
		}
	}

	if x < 0.0001 {
		x = 0.0001
	} else if 0.9999 < x {
		x = 0.9999
	}

	// Remark AS R83
	// http://www.jstor.org/stable/2347779
	fpu := math.Pow10(sae)
	acu := fpu
	if exp := int(-5/p/p - 1/math.Pow(α, 0.2) - 13); exp > sae {
		acu = math.Pow10(exp)
	}

	tx, yprev, sq, prev := 0.0, 0.0, 1.0, 1.0

outer:
	for {
		// Remark AS R19 and Algorithm AS 109
		// http://www.jstor.org/stable/2346887
		y = incBeta(x, p, q, logB)
		y = (y - α) * math.Exp(logB+(1-p)*math.Log(x)+(1-q)*math.Log(1-x))

		// Remark AS R83
		// http://www.jstor.org/stable/2347779
		if y*yprev <= 0 {
			prev = math.Max(sq, fpu)
		}

		// Remark AS R19 and Algorithm AS 109
		// http://www.jstor.org/stable/2346887
		for g := 1.0; ; {
			for {
				adj := g * y
				sq = adj * adj

				if sq < prev {
					tx = x - adj

					if 0 <= tx && tx <= 1 {
						break
					}
				}
				g /= 3
			}

			if prev <= acu || y*y <= acu {
				x = tx
				break outer
			}

			if tx != 0 && tx != 1 {
				break
			}

			g /= 3
		}

		if tx == x {
			break
		}

		x = tx
		yprev = y
	}

	if flip {
		return 1 - x
	} else {
		return x
	}
}

func logBeta(x, y float64) float64 {
	z, _ := math.Lgamma(x + y)
	x, _ = math.Lgamma(x)
	y, _ = math.Lgamma(y)

	return x + y - z
}
<|endoftext|>


<commit_before>// errorcheck

// Copyright 2014 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.

// Issue 7525: self-referential array types.

package main

import "unsafe"

var x struct {
	a [unsafe.Sizeof(x.a)]int // ERROR "array bound|invalid array"
	b [unsafe.Offsetof(x.b)]int // ERROR "array bound|invalid array"
	c [unsafe.Alignof(x.c)]int // ERROR "array bound|invalid array"
	d [len(x.d)]int // ERROR "array bound|invalid array"
	e [cap(x.e)]int // ERROR "array bound|invalid array"
}
<commit_msg>undo CL 77050045 / 073d79675aae<commit_after><|endoftext|>


<commit_before>package route53_backup_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/vito/cmdtest"
	. "github.com/vito/cmdtest/matchers"

	"fmt"
	. "github.com/cloudfoundry/bosh-aws-tools/test/bosh"
)

var _ = Describe("route53_backup job", func() {
	type deploymentManifestVals struct {
		DirectorUUID string
		DirectorHost string

		DeploymentName string
		JobName        string

		AwsAccessId         string
		AwsSecretAcccessKey string

		Route53ZoneNames []string
	}

	deploymentConfig := deploymentManifestVals{
		DirectorUUID: Config.DirectorUUID,
		DirectorHost: Config.DirectorHost,

		DeploymentName: "route53_backup",
		JobName:        "route53_backup",

		AwsAccessId:         Config.AwsAccessId,
		AwsSecretAcccessKey: Config.AwsSecretAcccessKey,

		Route53ZoneNames: Config.Route53ZoneNames,
	}

	boshSshToReturnBackupCounts := func() *cmdtest.Session {
		return Bosh(
			"-n", "ssh",
			deploymentConfig.JobName,
			"--gateway_host", deploymentConfig.DirectorHost,
			"--gateway_user", "vcap",
			"ls /var/vcap/store/route53_backup/* | xargs head -1 -q | sort -n | uniq -c",
		)
	}

	cleanUp := func() {
		BoshDeleteDeployment(deploymentConfig.DeploymentName)
		BoshDeleteRelease()
	}

	BeforeEach(cleanUp)
	AfterEach(cleanUp)

	BeforeEach(func() {
		BoshCreateRelease()
		BoshUploadRelease()
	})

	It("starts route53_backup job which periodically backs up route53 zones", func() {
		manifestFile := BoshDeployDeployment(deploymentManifestTplStr, deploymentConfig)

		defer manifestFile.Close()

		for _, zone := range deploymentConfig.Route53ZoneNames {
			// 1+ count for this zone
			countMatch := fmt.Sprintf("\\s+([2-9]|\\d{2,})\\s+\\$ORIGIN %s", zone)

			Eventually(
				boshSshToReturnBackupCounts,
				float64(3*60), // timeout in secs
				float64(10),   // interval in secs
			).Should(Say(countMatch))
		}
	})
})

// Keep template localized to this file
// to stop sharing templates between unrelated tests!
const deploymentManifestTplStr = `
---
name: {{ .DeploymentName }}
director_uuid: {{ .DirectorUUID }}

releases:
- name: bosh-aws-tools
  version: latest

networks:
- name: default
  type: manual
  subnets:
  - range: 10.10.16.0/24
    gateway: 10.10.16.1
    reserved:
    - 10.10.16.2 - 10.10.16.10 # full bosh is .7
    dns:
    - 10.10.16.6
    cloud_properties:
      subnet: subnet-f8744a8c

resource_pools:
- name: default
  stemcell:
    name: bosh-aws-xen-ubuntu
    version: latest
  network: default
  size: 1
  cloud_properties:
    instance_type: m1.small
    availability_zone: us-east-1b

compilation:
  reuse_compilation_vms: true
  workers: 1
  network: default
  cloud_properties:
    instance_type: c1.medium
    availability_zone: us-east-1b

update:
  canaries: 1
  canary_watch_time: 1000 - 90000
  update_watch_time: 1000 - 90000
  max_in_flight: 1
  max_errors: 1

jobs:
- name: {{ .JobName }}
  template: route53_backup
  resource_pool: default
  instances: 1
  networks:
  - name: default

properties:
  route53_backup:
    aws_access_key_id: {{ .AwsAccessId }}
    aws_secret_access_key: {{ .AwsSecretAcccessKey }}
    schedule: "*/1 * * * *"
`
<commit_msg>use test_ in names for deployment/job<commit_after>package route53_backup_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"github.com/vito/cmdtest"
	. "github.com/vito/cmdtest/matchers"

	"fmt"
	. "github.com/cloudfoundry/bosh-aws-tools/test/bosh"
)

var _ = Describe("route53_backup job", func() {
	type deploymentManifestVals struct {
		DirectorUUID string
		DirectorHost string

		DeploymentName string
		JobName        string

		AwsAccessId         string
		AwsSecretAcccessKey string

		Route53ZoneNames []string
	}

	deploymentConfig := deploymentManifestVals{
		DirectorUUID: Config.DirectorUUID,
		DirectorHost: Config.DirectorHost,

		// Use test in names to make sure we do not delete real deployments
		DeploymentName: "test_route53_backup",
		JobName:        "test_route53_backup",

		AwsAccessId:         Config.AwsAccessId,
		AwsSecretAcccessKey: Config.AwsSecretAcccessKey,

		Route53ZoneNames: Config.Route53ZoneNames,
	}

	boshSshToReturnBackupCounts := func() *cmdtest.Session {
		return Bosh(
			"-n", "ssh",
			deploymentConfig.JobName,
			"--gateway_host", deploymentConfig.DirectorHost,
			"--gateway_user", "vcap",
			"ls /var/vcap/store/route53_backup/* | xargs head -1 -q | sort -n | uniq -c",
		)
	}

	cleanUp := func() {
		BoshDeleteDeployment(deploymentConfig.DeploymentName)
		BoshDeleteRelease()
	}

	BeforeEach(cleanUp)
	AfterEach(cleanUp)

	BeforeEach(func() {
		BoshCreateRelease()
		BoshUploadRelease()
	})

	It("starts route53_backup job which periodically backs up route53 zones", func() {
		manifestFile := BoshDeployDeployment(deploymentManifestTplStr, deploymentConfig)

		defer manifestFile.Close()

		for _, zone := range deploymentConfig.Route53ZoneNames {
			// 1+ count for this zone
			countMatch := fmt.Sprintf("\\s+([2-9]|\\d{2,})\\s+\\$ORIGIN %s", zone)

			Eventually(
				boshSshToReturnBackupCounts,
				float64(3*60), // timeout in secs
				float64(10),   // interval in secs
			).Should(Say(countMatch))
		}
	})
})

// Keep template localized to this file
// to stop sharing templates between unrelated tests!
const deploymentManifestTplStr = `
---
name: {{ .DeploymentName }}
director_uuid: {{ .DirectorUUID }}

releases:
- name: bosh-aws-tools
  version: latest

networks:
- name: default
  type: manual
  subnets:
  - range: 10.10.16.0/24
    gateway: 10.10.16.1
    reserved:
    - 10.10.16.2 - 10.10.16.10 # full bosh is .7
    dns:
    - 10.10.16.6
    cloud_properties:
      subnet: subnet-f8744a8c

resource_pools:
- name: default
  stemcell:
    name: bosh-aws-xen-ubuntu
    version: latest
  network: default
  size: 1
  cloud_properties:
    instance_type: m1.small
    availability_zone: us-east-1b

compilation:
  reuse_compilation_vms: true
  workers: 1
  network: default
  cloud_properties:
    instance_type: c1.medium
    availability_zone: us-east-1b

update:
  canaries: 1
  canary_watch_time: 1000 - 90000
  update_watch_time: 1000 - 90000
  max_in_flight: 1
  max_errors: 1

jobs:
- name: {{ .JobName }}
  template: route53_backup
  resource_pool: default
  instances: 1
  networks:
  - name: default

properties:
  route53_backup:
    aws_access_key_id: {{ .AwsAccessId }}
    aws_secret_access_key: {{ .AwsSecretAcccessKey }}
    schedule: "*/1 * * * *"
`
<|endoftext|>