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

package push

import (
	"context"
	"fmt"

	"github.com/bufbuild/buf/private/buf/bufcli"
	"github.com/bufbuild/buf/private/bufpkg/bufanalysis"
	"github.com/bufbuild/buf/private/bufpkg/bufmodule"
	"github.com/bufbuild/buf/private/bufpkg/bufmodule/bufmoduleref"
	"github.com/bufbuild/buf/private/pkg/app/appcmd"
	"github.com/bufbuild/buf/private/pkg/app/appflag"
	"github.com/bufbuild/buf/private/pkg/command"
	"github.com/bufbuild/buf/private/pkg/rpc"
	"github.com/bufbuild/buf/private/pkg/stringutil"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

const (
	trackFlagName           = "track"
	tagFlagName             = "tag"
	tagFlagShortName        = "t"
	errorFormatFlagName     = "error-format"
	disableSymlinksFlagName = "disable-symlinks"
)

// NewCommand returns a new Command.
func NewCommand(
	name string,
	builder appflag.Builder,
) *appcmd.Command {
	flags := newFlags()
	return &appcmd.Command{
		Use:   name + " <source>",
		Short: "Push a module to a registry.",
		Long:  bufcli.GetSourceLong(`the source to push`),
		Args:  cobra.MaximumNArgs(1),
		Run: builder.NewRunFunc(
			func(ctx context.Context, container appflag.Container) error {
				return run(ctx, container, flags)
			},
			bufcli.NewErrorInterceptor(),
		),
		BindFlags: flags.Bind,
	}
}

type flags struct {
	Tracks          []string
	Tags            []string
	ErrorFormat     string
	DisableSymlinks bool
	// special
	InputHashtag string
}

func newFlags() *flags {
	return &flags{}
}

func (f *flags) Bind(flagSet *pflag.FlagSet) {
	bufcli.BindInputHashtag(flagSet, &f.InputHashtag)
	bufcli.BindDisableSymlinks(flagSet, &f.DisableSymlinks, disableSymlinksFlagName)
	flagSet.StringSliceVar(
		&f.Tracks,
		trackFlagName,
		nil,
		"Append the pushed module to this track. Multiple tracks are appended if specified multiple times.",
	)
	flagSet.StringSliceVarP(
		&f.Tags,
		tagFlagName,
		tagFlagShortName,
		nil,
		"Create a tag for the pushed commit. Multiple tracks are appended if specified multiple times.",
	)
	flagSet.StringVar(
		&f.ErrorFormat,
		errorFormatFlagName,
		"text",
		fmt.Sprintf(
			"The format for build errors printed to stderr. Must be one of %s.",
			stringutil.SliceToString(bufanalysis.AllFormatStrings),
		),
	)
}

func run(
	ctx context.Context,
	container appflag.Container,
	flags *flags,
) (retErr error) {
	if err := bufcli.ValidateErrorFormatFlag(flags.ErrorFormat, errorFormatFlagName); err != nil {
		return err
	}
	source, err := bufcli.GetInputValue(container, flags.InputHashtag, ".")
	if err != nil {
		return err
	}
	storageosProvider := bufcli.NewStorageosProvider(flags.DisableSymlinks)
	runner := command.NewRunner()
	// We are pushing to the BSR, this module has to be independently buildable
	// given the configuration it has without any enclosing workspace.
	module, moduleIdentity, err := bufcli.ReadModuleWithWorkspacesDisabled(
		ctx,
		container,
		storageosProvider,
		runner,
		source,
	)
	if err != nil {
		return err
	}
	protoModule, err := bufmodule.ModuleToProtoModule(ctx, module)
	if err != nil {
		return err
	}
	apiProvider, err := bufcli.NewRegistryProvider(ctx, container)
	if err != nil {
		return err
	}
	service, err := apiProvider.NewPushService(ctx, moduleIdentity.Remote())
	if err != nil {
		return err
	}
	tracks := flags.Tracks
	if tracks == nil {
		tracks = []string{bufmoduleref.MainTrack}
	}
	localModulePin, err := service.Push(
		ctx,
		moduleIdentity.Owner(),
		moduleIdentity.Repository(),
		"",
		protoModule,
		flags.Tags,
		tracks,
	)
	if err != nil {
		if rpc.GetErrorCode(err) == rpc.ErrorCodeAlreadyExists {
			if _, err := container.Stderr().Write(
				[]byte("The latest commit has the same content; not creating a new commit.\n"),
			); err != nil {
				return err
			}
			return nil
		}
		return err
	}
	if _, err := container.Stdout().Write([]byte(localModulePin.Commit + "\n")); err != nil {
		return err
	}
	return nil
}
<commit_msg>buf push: fix --tag usage typo (#969)<commit_after>// Copyright 2020-2022 Buf Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package push

import (
	"context"
	"fmt"

	"github.com/bufbuild/buf/private/buf/bufcli"
	"github.com/bufbuild/buf/private/bufpkg/bufanalysis"
	"github.com/bufbuild/buf/private/bufpkg/bufmodule"
	"github.com/bufbuild/buf/private/bufpkg/bufmodule/bufmoduleref"
	"github.com/bufbuild/buf/private/pkg/app/appcmd"
	"github.com/bufbuild/buf/private/pkg/app/appflag"
	"github.com/bufbuild/buf/private/pkg/command"
	"github.com/bufbuild/buf/private/pkg/rpc"
	"github.com/bufbuild/buf/private/pkg/stringutil"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

const (
	trackFlagName           = "track"
	tagFlagName             = "tag"
	tagFlagShortName        = "t"
	errorFormatFlagName     = "error-format"
	disableSymlinksFlagName = "disable-symlinks"
)

// NewCommand returns a new Command.
func NewCommand(
	name string,
	builder appflag.Builder,
) *appcmd.Command {
	flags := newFlags()
	return &appcmd.Command{
		Use:   name + " <source>",
		Short: "Push a module to a registry.",
		Long:  bufcli.GetSourceLong(`the source to push`),
		Args:  cobra.MaximumNArgs(1),
		Run: builder.NewRunFunc(
			func(ctx context.Context, container appflag.Container) error {
				return run(ctx, container, flags)
			},
			bufcli.NewErrorInterceptor(),
		),
		BindFlags: flags.Bind,
	}
}

type flags struct {
	Tracks          []string
	Tags            []string
	ErrorFormat     string
	DisableSymlinks bool
	// special
	InputHashtag string
}

func newFlags() *flags {
	return &flags{}
}

func (f *flags) Bind(flagSet *pflag.FlagSet) {
	bufcli.BindInputHashtag(flagSet, &f.InputHashtag)
	bufcli.BindDisableSymlinks(flagSet, &f.DisableSymlinks, disableSymlinksFlagName)
	flagSet.StringSliceVar(
		&f.Tracks,
		trackFlagName,
		nil,
		"Append the pushed module to this track. Multiple tracks are appended if specified multiple times.",
	)
	flagSet.StringSliceVarP(
		&f.Tags,
		tagFlagName,
		tagFlagShortName,
		nil,
		"Create a tag for the pushed commit. Multiple tags are created if specified multiple times.",
	)
	flagSet.StringVar(
		&f.ErrorFormat,
		errorFormatFlagName,
		"text",
		fmt.Sprintf(
			"The format for build errors printed to stderr. Must be one of %s.",
			stringutil.SliceToString(bufanalysis.AllFormatStrings),
		),
	)
}

func run(
	ctx context.Context,
	container appflag.Container,
	flags *flags,
) (retErr error) {
	if err := bufcli.ValidateErrorFormatFlag(flags.ErrorFormat, errorFormatFlagName); err != nil {
		return err
	}
	source, err := bufcli.GetInputValue(container, flags.InputHashtag, ".")
	if err != nil {
		return err
	}
	storageosProvider := bufcli.NewStorageosProvider(flags.DisableSymlinks)
	runner := command.NewRunner()
	// We are pushing to the BSR, this module has to be independently buildable
	// given the configuration it has without any enclosing workspace.
	module, moduleIdentity, err := bufcli.ReadModuleWithWorkspacesDisabled(
		ctx,
		container,
		storageosProvider,
		runner,
		source,
	)
	if err != nil {
		return err
	}
	protoModule, err := bufmodule.ModuleToProtoModule(ctx, module)
	if err != nil {
		return err
	}
	apiProvider, err := bufcli.NewRegistryProvider(ctx, container)
	if err != nil {
		return err
	}
	service, err := apiProvider.NewPushService(ctx, moduleIdentity.Remote())
	if err != nil {
		return err
	}
	tracks := flags.Tracks
	if tracks == nil {
		tracks = []string{bufmoduleref.MainTrack}
	}
	localModulePin, err := service.Push(
		ctx,
		moduleIdentity.Owner(),
		moduleIdentity.Repository(),
		"",
		protoModule,
		flags.Tags,
		tracks,
	)
	if err != nil {
		if rpc.GetErrorCode(err) == rpc.ErrorCodeAlreadyExists {
			if _, err := container.Stderr().Write(
				[]byte("The latest commit has the same content; not creating a new commit.\n"),
			); err != nil {
				return err
			}
			return nil
		}
		return err
	}
	if _, err := container.Stdout().Write([]byte(localModulePin.Commit + "\n")); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


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

package sdk

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"text/template"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/pkg/capnslog"
	"github.com/coreos/mantle/system/exec"
	"github.com/coreos/mantle/system/user"
	"github.com/coreos/mantle/util"
)

// Must run inside the SDK chroot, easiest to just assemble a script to do it
const (
	safePath   = "PATH=/usr/sbin:/usr/bin:/sbin:/bin"
	sudoPrompt = "--prompt=sudo password for %p: "
	script     = `#!/bin/bash
set -e

# make sure user/group database files exist
touch /etc/{group,gshadow,passwd,shadow}
chmod 0640 /etc/{gshadow,shadow}

# add group if it doesn't exist already
if ! getent group {{printf "%q" .Groupname}} >/dev/null; then
	echo Adding group {{printf "%q" .Groupname}}
	groupadd -o -g {{.Gid}} {{printf "%q" .Groupname}}
fi

echo Adding user {{printf "%q" .Username}}
useradd -o -g {{.Gid}} -u {{.Uid}} -s /bin/bash -m \
	-c {{printf "%q" .Name}} {{printf "%q" .Username}}

for g in kvm portage sudo; do
	# copy system group from /usr to /etc if needed
	if getent -s usrfiles group "$g" >/dev/null && \
	   ! getent -s files group "$g" >/dev/null; then
		getent -s usrfiles group "$g" >> /etc/group
	fi
	gpasswd -a {{printf "%q" .Username}} "$g"
done

echo Setting up sudoers
cat >/etc/sudoers.d/90_env_keep <<EOF
Defaults env_keep += "\
COREOS_BUILD_ID COREOS_OFFICIAL \
EMAIL GIT_AUTHOR_EMAIL GIT_AUTHOR_NAME \
GIT_COMMITTER_EMAIL GIT_COMMITTER_NAME \
GIT_PROXY_COMMAND GIT_SSH RSYNC_PROXY \
GPG_AGENT_INFO SSH_AGENT_PID SSH_AUTH_SOCK \
USE FEATURES PORTAGE_USERNAME \
all_proxy ftp_proxy http_proxy https_proxy no_proxy"
EOF
chmod 0440 /etc/sudoers.d/90_env_keep

echo Setting default enviornment variables
cat >/etc/env.d/90portage_username <<EOF
PORTAGE_USERNAME={{printf "%q" .Username}}
EOF
# needlessly noisy since portage isn't set up yet
env-update &>/dev/null

echo Setting up home directory
HOME=/home/{{printf "%q" .Username}}

# Create ~/trunk symlink
ln -sfT /mnt/host/source "$HOME"/trunk

rm -f "$HOME"/.bash{_logout,_profile,rc}
cat >"$HOME"/.bash_logout <<EOF
# .bash_logout

# This file is sourced when a login shell terminates.
EOF

cat >"$HOME"/.bash_profile <<EOF
# .bash_profile

# This file is sourced by bash for login shells.  The following line
# runs your .bashrc and is recommended by the bash info pages.
[[ -f ~/.bashrc ]] && . ~/.bashrc

# Automatically change to scripts directory.
cd ${CHROOT_CWD:-~/trunk/src/scripts}
EOF

cat >"$HOME"/.bashrc <<EOF
# .bashrc

# This file is sourced by all *interactive* bash shells on startup,
# including some apparently interactive shells such as scp and rcp
# that can't tolerate any output.  So make sure this doesn't display
# anything or bad things will happen !

# Test for an interactive shell.  There is no need to set anything
# past this point for scp and rcp, and it's important to refrain from
# outputting anything in those cases.
if [[ $- != *i* ]] ; then
	# Shell is non-interactive.  Be done now!
	return
fi

# Enable bash completion for build scripts.
source ~/trunk/src/scripts/bash_completion

# Put your fun stuff here.
EOF

chown -R {{.Uid}}:{{.Gid}} "$HOME"

# Checked in src/scripts/common.sh
touch /etc/debian_chroot
`
)

var scriptTemplate = template.Must(template.New("script").Parse(script))

func Setup(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	u, err := user.Current()
	if err != nil {
		return err
	}

	var sc bytes.Buffer
	if err := scriptTemplate.Execute(&sc, u); err != nil {
		return err
	}

	plog.Info("Configuring SDK chroot")
	sh := exec.Command(
		"sudo", sudoPrompt,
		"chroot", chroot,
		"/usr/bin/env", "-i",
		"/bin/bash", "--login")
	sh.Stdin = &sc
	sh.Stderr = os.Stderr
	if plog.LevelAt(capnslog.INFO) {
		out, err := sh.StdoutPipe()
		if err != nil {
			return err
		}
		go util.LogFrom(capnslog.INFO, out)
	}
	if plog.LevelAt(capnslog.DEBUG) {
		sh.Args = append(sh.Args, "-x")
	}
	return sh.Run()
}

func extract(tar, dir string) error {
	in, err := os.Open(tar)
	if err != nil {
		return err
	}
	defer in.Close()

	unzipper, err := exec.LookPath("lbzcat")
	if err != nil {
		unzipper = "bzcat"
	}

	unzip := exec.Command(unzipper)
	unzip.Stdin = in
	unzip.Stderr = os.Stderr
	unzipped, err := unzip.StdoutPipe()
	if err != nil {
		return err
	}

	untar := exec.Command("sudo", sudoPrompt,
		"tar", "--numeric-owner", "-x")
	untar.Dir = dir
	untar.Stdin = unzipped
	untar.Stderr = os.Stderr

	if err := unzip.Start(); err != nil {
		return err
	}

	if err := untar.Start(); err != nil {
		unzip.Kill()
		return err
	}

	if err := untar.Wait(); err != nil {
		unzip.Kill()
		return err
	}

	if err := unzip.Wait(); err != nil {
		return err
	}

	return nil
}

func Unpack(version, name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); !os.IsNotExist(err) {
		if err == nil {
			err = fmt.Errorf("Path already exists: %s", chroot)
		}
		return err
	}

	plog.Noticef("Unpacking SDK into %s", chroot)
	if err := os.MkdirAll(chroot, 0777); err != nil {
		return err
	}

	tar := filepath.Join(RepoCache(), "sdk", TarballName(version))
	plog.Infof("Using %s", tar)
	if err := extract(tar, chroot); err != nil {
		plog.Errorf("Extracting %s to %s failed: %v", tar, chroot, err)
		return err
	}
	plog.Notice("Unpacked")

	return nil
}

func Delete(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); err != nil {
		if os.IsNotExist(err) {
			plog.Infof("Path does not exist: %s", chroot)
			return nil
		}
		return err
	}

	plog.Noticef("Removing SDK at %s", chroot)
	rm := exec.Command("sudo", sudoPrompt, "rm", "-rf", chroot)
	rm.Stderr = os.Stderr
	if err := rm.Run(); err != nil {
		return err
	}
	plog.Notice("Removed")

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

package sdk

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"text/template"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/pkg/capnslog"
	"github.com/coreos/mantle/system/exec"
	"github.com/coreos/mantle/system/user"
	"github.com/coreos/mantle/util"
)

// Must run inside the SDK chroot, easiest to just assemble a script to do it
const (
	safePath   = "PATH=/usr/sbin:/usr/bin:/sbin:/bin"
	sudoPrompt = "--prompt=sudo password for %p: "
	script     = `#!/bin/bash
set -e

# make sure user/group database files exist
touch /etc/{group,gshadow,passwd,shadow}
chmod 0640 /etc/{gshadow,shadow}

# add group if it doesn't exist already
if ! getent group {{printf "%q" .Groupname}} >/dev/null; then
	echo Adding group {{printf "%q" .Groupname}}
	groupadd -o -g {{.Gid}} {{printf "%q" .Groupname}}
fi

echo Adding user {{printf "%q" .Username}}
useradd -o -g {{.Gid}} -u {{.Uid}} -s /bin/bash -m \
	-c {{printf "%q" .Name}} {{printf "%q" .Username}}

for g in kvm portage sudo; do
	# copy system group from /usr to /etc if needed
	if getent -s usrfiles group "$g" >/dev/null && \
	   ! getent -s files group "$g" >/dev/null; then
		getent -s usrfiles group "$g" >> /etc/group
	fi
	gpasswd -a {{printf "%q" .Username}} "$g"
done

echo Setting up sudoers
cat >/etc/sudoers.d/90_env_keep <<EOF
Defaults env_keep += "\
COREOS_BUILD_ID COREOS_OFFICIAL \
EMAIL GIT_AUTHOR_EMAIL GIT_AUTHOR_NAME \
GIT_COMMITTER_EMAIL GIT_COMMITTER_NAME \
GIT_PROXY_COMMAND GIT_SSH RSYNC_PROXY \
GPG_AGENT_INFO SSH_AUTH_SOCK \
USE FEATURES PORTAGE_USERNAME \
all_proxy ftp_proxy http_proxy https_proxy no_proxy"
EOF
chmod 0440 /etc/sudoers.d/90_env_keep

echo Setting default enviornment variables
cat >/etc/env.d/90portage_username <<EOF
PORTAGE_USERNAME={{printf "%q" .Username}}
EOF
# needlessly noisy since portage isn't set up yet
env-update &>/dev/null

echo Setting up home directory
HOME=/home/{{printf "%q" .Username}}

# Create ~/trunk symlink
ln -sfT /mnt/host/source "$HOME"/trunk

rm -f "$HOME"/.bash{_logout,_profile,rc}
cat >"$HOME"/.bash_logout <<EOF
# .bash_logout

# This file is sourced when a login shell terminates.
EOF

cat >"$HOME"/.bash_profile <<EOF
# .bash_profile

# This file is sourced by bash for login shells.  The following line
# runs your .bashrc and is recommended by the bash info pages.
[[ -f ~/.bashrc ]] && . ~/.bashrc

# Automatically change to scripts directory.
cd ${CHROOT_CWD:-~/trunk/src/scripts}
EOF

cat >"$HOME"/.bashrc <<EOF
# .bashrc

# This file is sourced by all *interactive* bash shells on startup,
# including some apparently interactive shells such as scp and rcp
# that can't tolerate any output.  So make sure this doesn't display
# anything or bad things will happen !

# Test for an interactive shell.  There is no need to set anything
# past this point for scp and rcp, and it's important to refrain from
# outputting anything in those cases.
if [[ $- != *i* ]] ; then
	# Shell is non-interactive.  Be done now!
	return
fi

# Enable bash completion for build scripts.
source ~/trunk/src/scripts/bash_completion

# Put your fun stuff here.
EOF

chown -R {{.Uid}}:{{.Gid}} "$HOME"

# Checked in src/scripts/common.sh
touch /etc/debian_chroot
`
)

var scriptTemplate = template.Must(template.New("script").Parse(script))

func Setup(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	u, err := user.Current()
	if err != nil {
		return err
	}

	var sc bytes.Buffer
	if err := scriptTemplate.Execute(&sc, u); err != nil {
		return err
	}

	plog.Info("Configuring SDK chroot")
	sh := exec.Command(
		"sudo", sudoPrompt,
		"chroot", chroot,
		"/usr/bin/env", "-i",
		"/bin/bash", "--login")
	sh.Stdin = &sc
	sh.Stderr = os.Stderr
	if plog.LevelAt(capnslog.INFO) {
		out, err := sh.StdoutPipe()
		if err != nil {
			return err
		}
		go util.LogFrom(capnslog.INFO, out)
	}
	if plog.LevelAt(capnslog.DEBUG) {
		sh.Args = append(sh.Args, "-x")
	}
	return sh.Run()
}

func extract(tar, dir string) error {
	in, err := os.Open(tar)
	if err != nil {
		return err
	}
	defer in.Close()

	unzipper, err := exec.LookPath("lbzcat")
	if err != nil {
		unzipper = "bzcat"
	}

	unzip := exec.Command(unzipper)
	unzip.Stdin = in
	unzip.Stderr = os.Stderr
	unzipped, err := unzip.StdoutPipe()
	if err != nil {
		return err
	}

	untar := exec.Command("sudo", sudoPrompt,
		"tar", "--numeric-owner", "-x")
	untar.Dir = dir
	untar.Stdin = unzipped
	untar.Stderr = os.Stderr

	if err := unzip.Start(); err != nil {
		return err
	}

	if err := untar.Start(); err != nil {
		unzip.Kill()
		return err
	}

	if err := untar.Wait(); err != nil {
		unzip.Kill()
		return err
	}

	if err := unzip.Wait(); err != nil {
		return err
	}

	return nil
}

func Unpack(version, name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); !os.IsNotExist(err) {
		if err == nil {
			err = fmt.Errorf("Path already exists: %s", chroot)
		}
		return err
	}

	plog.Noticef("Unpacking SDK into %s", chroot)
	if err := os.MkdirAll(chroot, 0777); err != nil {
		return err
	}

	tar := filepath.Join(RepoCache(), "sdk", TarballName(version))
	plog.Infof("Using %s", tar)
	if err := extract(tar, chroot); err != nil {
		plog.Errorf("Extracting %s to %s failed: %v", tar, chroot, err)
		return err
	}
	plog.Notice("Unpacked")

	return nil
}

func Delete(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); err != nil {
		if os.IsNotExist(err) {
			plog.Infof("Path does not exist: %s", chroot)
			return nil
		}
		return err
	}

	plog.Noticef("Removing SDK at %s", chroot)
	rm := exec.Command("sudo", sudoPrompt, "rm", "-rf", chroot)
	rm.Stderr = os.Stderr
	if err := rm.Run(); err != nil {
		return err
	}
	plog.Notice("Removed")

	return nil
}
<|endoftext|>


<commit_before>package tool

import (
	"os/exec"
	"reflect"
	"testing"

	"github.com/bouk/monkey"
	. "github.com/smartystreets/goconvey/convey"
	"github.com/voidint/gbb/config"
)

func TestBuild(t *testing.T) {
	var cmd *exec.Cmd
	monkey.PatchInstanceMethod(reflect.TypeOf(cmd), "Run", func(c *exec.Cmd) error {
		return nil
	})
	defer monkey.UnpatchInstanceMethod(reflect.TypeOf(cmd), "Run")

	Convey("调用gb build编译", t, func() {
		Convey("包含变量表达式", func() {
			c := &config.Config{
				Version: "0.3.0",
				Tool:    "gb build",
				Package: "github.com/voidint/gbb/build",
			}

			Convey("包含非法变量表达式", func() {
				c.Variables = []config.Variable{
					{Variable: "Date", Value: "xxxx"},
				}
				err := Build(c, true, "./")
				So(err, ShouldNotBeNil)
			})

			Convey("包含合法变量表达式", func() {
				c.Variables = []config.Variable{
					{Variable: "Date", Value: "{{.Date}}"},
				}
				err := Build(c, true, "./")
				So(err, ShouldBeNil)
			})

		})

		Convey("不包含变量表达式", func() {
			c := &config.Config{
				Version: "0.3.0",
				Tool:    "gb build",
			}

			err := Build(c, true, "./")
			So(err, ShouldBeNil)
		})
	})

	Convey("调用go build编译", t, func() {
		c := &config.Config{
			Version: "0.3.0",
			Tool:    "go build",
			Package: "github.com/voidint/gbb/build",
			Variables: []config.Variable{
				{Variable: "Date", Value: "{{.Date}}"},
			},
		}

		err := Build(c, true, "./")
		So(err, ShouldBeNil)

	})

	Convey("调用非法的编译工具编译", t, func() {
		c := &config.Config{
			Version: "0.3.0",
			Tool:    "unsupported tool",
		}

		err := Build(c, true, "./")
		So(err, ShouldEqual, ErrBuildTool)
	})
}
<commit_msg>Fix unit test cases<commit_after>package tool

import (
	"os/exec"
	"reflect"
	"testing"

	"github.com/bouk/monkey"
	. "github.com/smartystreets/goconvey/convey"
	"github.com/voidint/gbb/config"
)

func TestBuild(t *testing.T) {
	var cmd *exec.Cmd
	monkey.PatchInstanceMethod(reflect.TypeOf(cmd), "Run", func(c *exec.Cmd) error {
		return nil
	})
	defer monkey.UnpatchInstanceMethod(reflect.TypeOf(cmd), "Run")

	Convey("调用gb build编译", t, func() {
		Convey("包含变量表达式", func() {
			c := &config.Config{
				Version: "0.3.0",
				Tool:    "gb build",
				Package: "github.com/voidint/gbb/build",
				Debug:   true,
			}

			Convey("包含非法变量表达式", func() {
				c.Variables = []config.Variable{
					{Variable: "Date", Value: "xxxx"},
				}
				err := Build(c, "./")
				So(err, ShouldNotBeNil)
			})

			Convey("包含合法变量表达式", func() {
				c.Variables = []config.Variable{
					{Variable: "Date", Value: "{{.Date}}"},
				}
				err := Build(c, "./")
				So(err, ShouldBeNil)
			})

		})

		Convey("不包含变量表达式", func() {
			c := &config.Config{
				Version: "0.3.0",
				Tool:    "gb build",
				Debug:   true,
			}

			err := Build(c, "./")
			So(err, ShouldBeNil)
		})
	})

	Convey("调用go build编译", t, func() {
		c := &config.Config{
			Version: "0.3.0",
			Tool:    "go build",
			Package: "github.com/voidint/gbb/build",
			Variables: []config.Variable{
				{Variable: "Date", Value: "{{.Date}}"},
			},
			Debug: true,
		}

		err := Build(c, "./")
		So(err, ShouldBeNil)

	})

	Convey("调用非法的编译工具编译", t, func() {
		c := &config.Config{
			Version: "0.3.0",
			Tool:    "unsupported tool",
			Debug:   true,
		}

		err := Build(c, "./")
		So(err, ShouldEqual, ErrBuildTool)
	})
}
<|endoftext|>


<commit_before>package manta

import (
	"strconv"
	"strings"

	"github.com/dotabuff/manta/dota"
	"github.com/golang/protobuf/proto"
)

// Internal callback for OnCDemoSendTables.
func (p *Parser) onCDemoSendTables(m *dota.CDemoSendTables) error {
	// Parse the send tables
	st, err := parseSendTables(m)
	if err != nil {
		return err
	}

	// Update the parser state
	p.sendTables = st

	return nil
}

// Internal callback for OnCSVCMsg_SendTable.
// So far we haven't seen any of these, they very well may not exist.
func (p *Parser) onCSVCMsg_SendTable(m *dota.CSVCMsg_SendTable) error {
	_panicf("got a CSVCMsg_SendTable")
	return nil
}

// Holds and maintains send tables for an instance of Parser.
type sendTables struct {
	tables map[string]*sendTable
	props  []*sendProp
}

// Get a send table by name.
func (ts *sendTables) getTableByName(name string) (*sendTable, bool) {
	t, ok := ts.tables[name]
	return t, ok
}

// Represents a single send table.
type sendTable struct {
	index   int32
	name    string
	version int32
	props   []*sendProp
}

// Represents a property in a send table.
type sendProp struct {
	dtIndex                int32
	dtName                 string
	varIndex               int32
	varName                string
	bitCount               *int32
	lowValue               *float32
	highValue              *float32
	encodeFlags            *uint32
	fieldSerializerIndex   *int32
	fieldSerializerName    *string
	fieldSerializerVersion *int32
	serializedFromIndex    *int32
	serializedFromName     *string
	sendNodeIndex          int32
	sendNodeName           string
}

// Copies a sendProp.
func (p *sendProp) copy() *sendProp {
	return &sendProp{
		dtIndex:                p.dtIndex,
		dtName:                 p.dtName,
		varIndex:               p.varIndex,
		varName:                p.varName,
		bitCount:               p.bitCount,
		lowValue:               p.lowValue,
		highValue:              p.highValue,
		encodeFlags:            p.encodeFlags,
		fieldSerializerIndex:   p.fieldSerializerIndex,
		fieldSerializerName:    p.fieldSerializerName,
		fieldSerializerVersion: p.fieldSerializerVersion,
		serializedFromIndex:    p.serializedFromIndex,
		serializedFromName:     p.serializedFromName,
		sendNodeIndex:          p.sendNodeIndex,
		sendNodeName:           p.sendNodeName,
	}
}

// Debugging method to describe the table.
func (t *sendTable) Describe() {
	_debugf("table %d (%s) version %d with %d raw props", t.index, t.name, t.version, len(t.props))
	for i, p := range t.props {
		_debugf("-> prop %d: %s", i, p.Describe())
	}
}

// Returns the type name and count for a sendprop.
func (p *sendProp) typeInfo() (s string, n int, err error) {
	ss := strings.Split(strings.Replace(p.dtName, "]", "[", 1), "[")
	s = ss[0]
	switch s {
	case "char":
		n = 1
	default:
		if len(ss) >= 2 {
			n, err = strconv.Atoi(ss[1])
		} else {
			n = 1
		}
	}
	return
}

func (p *sendProp) Describe() string {
	out := _sprintf("type:%s(%d) name:%s(%d)",
		p.dtName, p.dtIndex, p.varName, p.varIndex)

	// This doesn't seem to be very helpful yet.
	// out += _sprintf(" sendNode: %s(%d)", p.sendNodeName, p.sendNodeIndex)

	if p.fieldSerializerIndex != nil {
		out += _sprintf(" serializer:%s(%d)", *p.fieldSerializerName, *p.fieldSerializerIndex)
	}

	if p.serializedFromIndex != nil {
		out += _sprintf(" serializedFrom:%s(%d)", *p.serializedFromName, *p.serializedFromIndex)
	}

	if p.bitCount != nil {
		out += _sprintf(" bitCount:%d", *p.bitCount)
	}
	if p.lowValue != nil {
		out += _sprintf(" lowValue:%f", *p.lowValue)
	}
	if p.highValue != nil {
		out += _sprintf(" highValue:%f", *p.highValue)
	}
	if p.encodeFlags != nil {
		out += _sprintf(" encodeFlags:%d", *p.encodeFlags)
	}
	if p.fieldSerializerVersion != nil {
		out += _sprintf(" fieldSerVer:%d", *p.fieldSerializerVersion)
	}
	return out
}

// Parses a CDemoSendTables buffer, producing a sendTables object.
func parseSendTables(m *dota.CDemoSendTables) (*sendTables, error) {
	// This packet just contains a single large buffer
	r := newReader(m.GetData())

	// The buffer starts with a varint encoded length
	size := int(r.readVarUint32())
	if size != r.remBytes() {
		_panicf("expected %d additional bytes, got %d", size, r.remBytes())
	}

	// Read the rest of the buffer as a CSVCMsg_FlattenedSerializer.
	buf := r.readBytes(size)
	msg := &dota.CSVCMsg_FlattenedSerializer{}
	if err := proto.Unmarshal(buf, msg); err != nil {
		_panicf("cannot decode proto: %s", err)
	}

	// Create a list of sendProps
	props := make([]*sendProp, 0)
	for _, o := range msg.GetFields() {
		p := &sendProp{
			dtIndex:                o.GetVarTypeSym(),
			dtName:                 msg.GetSymbols()[o.GetVarTypeSym()],
			varIndex:               o.GetVarNameSym(),
			varName:                msg.GetSymbols()[o.GetVarNameSym()],
			bitCount:               o.BitCount,
			lowValue:               o.LowValue,
			highValue:              o.HighValue,
			fieldSerializerVersion: o.FieldSerializerVersion,
			sendNodeIndex:          o.GetSendNodeSym(),
			sendNodeName:           msg.GetSymbols()[o.GetSendNodeSym()],
		}

		if o.EncodeFlags != nil {
			p.encodeFlags = proto.Uint32(uint32(*o.EncodeFlags))
		}

		if o.FieldSerializerNameSym != nil {
			p.fieldSerializerIndex = o.FieldSerializerNameSym
			p.fieldSerializerName = proto.String(msg.GetSymbols()[o.GetFieldSerializerNameSym()])
		}

		props = append(props, p)
	}

	// Create a map of sendTables
	tables := make(map[string]*sendTable)
	for _, o := range msg.GetSerializers() {
		// Create the basic table.
		t := &sendTable{
			index:   o.GetSerializerNameSym(),
			name:    msg.GetSymbols()[o.GetSerializerNameSym()],
			version: o.GetSerializerVersion(),
			props:   make([]*sendProp, 0),
		}

		// Iterate through prop field indexes.
		for _, pid := range o.GetFieldsIndex() {
			// Get the property at the given index.
			prop := props[int(pid)]

			// If the prop has a serializer, inherit its properties
			if prop.fieldSerializerIndex != nil {
				// Find the serializer.
				ser, ok := tables[*prop.fieldSerializerName]
				if !ok {
					_panicf("unable to find serializer %d (%s)", *prop.fieldSerializerIndex, *prop.fieldSerializerName)
				}

				// Iterate through serializer props, adding them to the table.
				// Property names are subclassed as "%propVarName.%serializerVarName"
				for _, p := range ser.props {
					p2 := p.copy()
					p2.serializedFromIndex = prop.fieldSerializerIndex
					p2.serializedFromName = prop.fieldSerializerName
					p2.varName = _sprintf("%s.%s", prop.varName, p.varName)
					t.props = append(t.props, p2)
				}
				continue
			}

			// For normal props (without serializers), extract base type name
			// and element count, then store those as props in the table.
			typName, typCount, err := prop.typeInfo()
			if err != nil {
				_panicf("unable to flatten property %s: %s", prop.Describe(), err)
			}

			// Iterate through elements, adding them as new properties to the table.
			// Single element properties are named using the property var name.
			// Multiple element properties are named as "%propVarName.N".
			for i := 0; i < typCount; i++ {
				p2 := prop.copy()
				p2.dtName = typName
				if typCount > 1 {
					p2.varName = _sprintf("%s.%04d", prop.varName, i)
				}
				t.props = append(t.props, p2)
			}
		}

		tables[t.name] = t
	}

	// Return a sendTables object
	return &sendTables{tables: tables, props: props}, nil
}
<commit_msg>Added dumpSendTables function return a json string<commit_after>package manta

import (
	"encoding/json"
	"strconv"
	"strings"

	"github.com/dotabuff/manta/dota"
	"github.com/golang/protobuf/proto"
)

// Internal callback for OnCDemoSendTables.
func (p *Parser) onCDemoSendTables(m *dota.CDemoSendTables) error {
	// Parse the send tables
	st, err := parseSendTables(m)
	if err != nil {
		return err
	}

	// Update the parser state
	p.sendTables = st

	return nil
}

// Internal callback for OnCSVCMsg_SendTable.
// So far we haven't seen any of these, they very well may not exist.
func (p *Parser) onCSVCMsg_SendTable(m *dota.CSVCMsg_SendTable) error {
	_panicf("got a CSVCMsg_SendTable")
	return nil
}

// Holds and maintains send tables for an instance of Parser.
type sendTables struct {
	tables map[string]*sendTable
	props  []*sendProp
}

// Get a send table by name.
func (ts *sendTables) getTableByName(name string) (*sendTable, bool) {
	t, ok := ts.tables[name]
	return t, ok
}

// Represents a single send table.
type sendTable struct {
	index   int32
	name    string
	version int32
	props   []*sendProp
}

// Represents a property in a send table.
type sendProp struct {
	dtIndex                int32
	dtName                 string
	varIndex               int32
	varName                string
	bitCount               *int32
	lowValue               *float32
	highValue              *float32
	encodeFlags            *uint32
	fieldSerializerIndex   *int32
	fieldSerializerName    *string
	fieldSerializerVersion *int32
	serializedFromIndex    *int32
	serializedFromName     *string
	sendNodeIndex          int32
	sendNodeName           string
}

// Copies a sendProp.
func (p *sendProp) copy() *sendProp {
	return &sendProp{
		dtIndex:                p.dtIndex,
		dtName:                 p.dtName,
		varIndex:               p.varIndex,
		varName:                p.varName,
		bitCount:               p.bitCount,
		lowValue:               p.lowValue,
		highValue:              p.highValue,
		encodeFlags:            p.encodeFlags,
		fieldSerializerIndex:   p.fieldSerializerIndex,
		fieldSerializerName:    p.fieldSerializerName,
		fieldSerializerVersion: p.fieldSerializerVersion,
		serializedFromIndex:    p.serializedFromIndex,
		serializedFromName:     p.serializedFromName,
		sendNodeIndex:          p.sendNodeIndex,
		sendNodeName:           p.sendNodeName,
	}
}

// Debugging method to describe the table.
func (t *sendTable) Describe() {
	_debugf("table %d (%s) version %d with %d raw props", t.index, t.name, t.version, len(t.props))
	for i, p := range t.props {
		_debugf("-> prop %d: %s", i, p.Describe())
	}
}

// Returns the type name and count for a sendprop.
func (p *sendProp) typeInfo() (s string, n int, err error) {
	ss := strings.Split(strings.Replace(p.dtName, "]", "[", 1), "[")
	s = ss[0]
	switch s {
	case "char":
		n = 1
	default:
		if len(ss) >= 2 {
			n, err = strconv.Atoi(ss[1])
		} else {
			n = 1
		}
	}
	return
}

func (p *sendProp) Describe() string {
	out := _sprintf("type:%s(%d) name:%s(%d)",
		p.dtName, p.dtIndex, p.varName, p.varIndex)

	// This doesn't seem to be very helpful yet.
	// out += _sprintf(" sendNode: %s(%d)", p.sendNodeName, p.sendNodeIndex)

	if p.fieldSerializerIndex != nil {
		out += _sprintf(" serializer:%s(%d)", *p.fieldSerializerName, *p.fieldSerializerIndex)
	}

	if p.serializedFromIndex != nil {
		out += _sprintf(" serializedFrom:%s(%d)", *p.serializedFromName, *p.serializedFromIndex)
	}

	if p.bitCount != nil {
		out += _sprintf(" bitCount:%d", *p.bitCount)
	}
	if p.lowValue != nil {
		out += _sprintf(" lowValue:%f", *p.lowValue)
	}
	if p.highValue != nil {
		out += _sprintf(" highValue:%f", *p.highValue)
	}
	if p.encodeFlags != nil {
		out += _sprintf(" encodeFlags:%d", *p.encodeFlags)
	}
	if p.fieldSerializerVersion != nil {
		out += _sprintf(" fieldSerVer:%d", *p.fieldSerializerVersion)
	}
	return out
}

// Dumps the json representation of the new FlattenedSerializer Packets
func dumpSendTables(m *dota.CDemoSendTables) string {
	// This packet just contains a single large buffer
	r := newReader(m.GetData())

	// The buffer starts with a varint encoded length
	size := int(r.readVarUint32())
	if size != r.remBytes() {
		_panicf("expected %d additional bytes, got %d", size, r.remBytes())
	}

	// Read the rest of the buffer as a CSVCMsg_FlattenedSerializer.
	buf := r.readBytes(size)
	msg := &dota.CSVCMsg_FlattenedSerializer{}
	if err := proto.Unmarshal(buf, msg); err != nil {
		_panicf("cannot decode proto: %s", err)
	}

	str, _ := json.MarshalIndent(msg, "", "  ") // two space ident
	return string(str)
}

// Parses a CDemoSendTables buffer, producing a sendTables object.
func parseSendTables(m *dota.CDemoSendTables) (*sendTables, error) {
	// This packet just contains a single large buffer
	r := newReader(m.GetData())

	// The buffer starts with a varint encoded length
	size := int(r.readVarUint32())
	if size != r.remBytes() {
		_panicf("expected %d additional bytes, got %d", size, r.remBytes())
	}

	// Read the rest of the buffer as a CSVCMsg_FlattenedSerializer.
	buf := r.readBytes(size)
	msg := &dota.CSVCMsg_FlattenedSerializer{}
	if err := proto.Unmarshal(buf, msg); err != nil {
		_panicf("cannot decode proto: %s", err)
	}

	// Create a list of sendProps
	props := make([]*sendProp, 0)
	for _, o := range msg.GetFields() {
		p := &sendProp{
			dtIndex:                o.GetVarTypeSym(),
			dtName:                 msg.GetSymbols()[o.GetVarTypeSym()],
			varIndex:               o.GetVarNameSym(),
			varName:                msg.GetSymbols()[o.GetVarNameSym()],
			bitCount:               o.BitCount,
			lowValue:               o.LowValue,
			highValue:              o.HighValue,
			fieldSerializerVersion: o.FieldSerializerVersion,
			sendNodeIndex:          o.GetSendNodeSym(),
			sendNodeName:           msg.GetSymbols()[o.GetSendNodeSym()],
		}

		if o.EncodeFlags != nil {
			p.encodeFlags = proto.Uint32(uint32(*o.EncodeFlags))
		}

		if o.FieldSerializerNameSym != nil {
			p.fieldSerializerIndex = o.FieldSerializerNameSym
			p.fieldSerializerName = proto.String(msg.GetSymbols()[o.GetFieldSerializerNameSym()])
		}

		props = append(props, p)
	}

	// Create a map of sendTables
	tables := make(map[string]*sendTable)
	for _, o := range msg.GetSerializers() {
		// Create the basic table.
		t := &sendTable{
			index:   o.GetSerializerNameSym(),
			name:    msg.GetSymbols()[o.GetSerializerNameSym()],
			version: o.GetSerializerVersion(),
			props:   make([]*sendProp, 0),
		}

		// Iterate through prop field indexes.
		for _, pid := range o.GetFieldsIndex() {
			// Get the property at the given index.
			prop := props[int(pid)]

			// If the prop has a serializer, inherit its properties
			if prop.fieldSerializerIndex != nil {
				// Find the serializer.
				ser, ok := tables[*prop.fieldSerializerName]
				if !ok {
					_panicf("unable to find serializer %d (%s)", *prop.fieldSerializerIndex, *prop.fieldSerializerName)
				}

				// Iterate through serializer props, adding them to the table.
				// Property names are subclassed as "%propVarName.%serializerVarName"
				for _, p := range ser.props {
					p2 := p.copy()
					p2.serializedFromIndex = prop.fieldSerializerIndex
					p2.serializedFromName = prop.fieldSerializerName
					p2.varName = _sprintf("%s.%s", prop.varName, p.varName)
					t.props = append(t.props, p2)
				}
				continue
			}

			// For normal props (without serializers), extract base type name
			// and element count, then store those as props in the table.
			typName, typCount, err := prop.typeInfo()
			if err != nil {
				_panicf("unable to flatten property %s: %s", prop.Describe(), err)
			}

			// Iterate through elements, adding them as new properties to the table.
			// Single element properties are named using the property var name.
			// Multiple element properties are named as "%propVarName.N".
			for i := 0; i < typCount; i++ {
				p2 := prop.copy()
				p2.dtName = typName
				if typCount > 1 {
					p2.varName = _sprintf("%s.%04d", prop.varName, i)
				}
				t.props = append(t.props, p2)
			}
		}

		tables[t.name] = t
	}

	// Return a sendTables object
	return &sendTables{tables: tables, props: props}, nil
}
<|endoftext|>


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

package server

import (
	"bytes"
	"fmt"
	"os"
	"time"

	"github.com/osrg/gobgp/config"
	"github.com/osrg/gobgp/packet/bgp"
	"github.com/osrg/gobgp/packet/mrt"
	"github.com/osrg/gobgp/table"
	log "github.com/sirupsen/logrus"
)

const (
	MIN_ROTATION_INTERVAL = 60
	MIN_DUMP_INTERVAL     = 60
)

type mrtWriter struct {
	dead             chan struct{}
	s                *BgpServer
	c                *config.MrtConfig
	file             *os.File
	rotationInterval uint64
	dumpInterval     uint64
}

func (m *mrtWriter) Stop() {
	close(m.dead)
}

func (m *mrtWriter) loop() error {
	ops := []WatchOption{}
	switch m.c.DumpType {
	case config.MRT_TYPE_UPDATES:
		ops = append(ops, WatchUpdate(false))
	case config.MRT_TYPE_TABLE:
		if len(m.c.TableName) > 0 {
			ops = append(ops, WatchTableName(m.c.TableName))
		}
	}
	w := m.s.Watch(ops...)
	rotator := func() *time.Ticker {
		if m.rotationInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.rotationInterval))
	}()
	dump := func() *time.Ticker {
		if m.dumpInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.dumpInterval))
	}()

	defer func() {
		if m.file != nil {
			m.file.Close()
		}
		if m.rotationInterval != 0 {
			rotator.Stop()
		}
		if m.dumpInterval != 0 {
			dump.Stop()
		}
		w.Stop()
	}()

	for {
		serialize := func(ev WatchEvent) []*mrt.MRTMessage {
			msg := make([]*mrt.MRTMessage, 0, 1)
			switch m := ev.(type) {
			case *WatchEventUpdate:
				subtype := mrt.MESSAGE_AS4
				mp := mrt.NewBGP4MPMessage(m.PeerAS, m.LocalAS, 0, m.PeerAddress.String(), m.LocalAddress.String(), m.FourBytesAs, nil)
				mp.BGPMessagePayload = m.Payload
				if m.FourBytesAs == false {
					subtype = mrt.MESSAGE
				}
				if bm, err := mrt.NewMRTMessage(uint32(m.Timestamp.Unix()), mrt.BGP4MP, subtype, mp); err != nil {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Data":  m,
						"Error": err,
					}).Warn("Failed to create MRT message in serialize()")
				} else {
					msg = append(msg, bm)
				}
			case *WatchEventTable:
				t := uint32(time.Now().Unix())
				peers := make([]*mrt.Peer, 0, len(m.Neighbor))
				for _, pconf := range m.Neighbor {
					peers = append(peers, mrt.NewPeer(pconf.State.RemoteRouterId, pconf.State.NeighborAddress, pconf.Config.PeerAs, true))
				}
				if bm, err := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, mrt.PEER_INDEX_TABLE, mrt.NewPeerIndexTable(m.RouterId, "", peers)); err != nil {
					break
				} else {
					msg = append(msg, bm)
				}

				idx := func(p *table.Path) uint16 {
					for i, pconf := range m.Neighbor {
						if p.GetSource().Address.String() == pconf.State.NeighborAddress {
							return uint16(i)
						}
					}
					return uint16(len(m.Neighbor))
				}

				subtype := func(p *table.Path) mrt.MRTSubTypeTableDumpv2 {
					switch p.GetRouteFamily() {
					case bgp.RF_IPv4_UC:
						return mrt.RIB_IPV4_UNICAST
					case bgp.RF_IPv4_MC:
						return mrt.RIB_IPV4_MULTICAST
					case bgp.RF_IPv6_UC:
						return mrt.RIB_IPV6_UNICAST
					case bgp.RF_IPv6_MC:
						return mrt.RIB_IPV6_MULTICAST
					}
					return mrt.RIB_GENERIC
				}

				seq := uint32(0)
				for _, pathList := range m.PathList {
					entries := make([]*mrt.RibEntry, 0, len(pathList))
					for _, path := range pathList {
						if path.IsLocal() {
							continue
						}
						entries = append(entries, mrt.NewRibEntry(idx(path), uint32(path.GetTimestamp().Unix()), 0, path.GetPathAttrs()))
					}
					if len(entries) > 0 {
						bm, _ := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, subtype(pathList[0]), mrt.NewRib(seq, pathList[0].GetNlri(), entries))
						msg = append(msg, bm)
						seq++
					}
				}
			}
			return msg
		}

		drain := func(ev WatchEvent) {
			events := make([]WatchEvent, 0, 1+len(w.Event()))
			if ev != nil {
				events = append(events, ev)
			}

			for len(w.Event()) > 0 {
				events = append(events, <-w.Event())
			}

			w := func(buf []byte) {
				if _, err := m.file.Write(buf); err == nil {
					m.file.Sync()
				} else {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Error": err,
					}).Warn("Can't write to destination MRT file")
				}
			}

			var b bytes.Buffer
			for _, e := range events {
				for _, m := range serialize(e) {
					if buf, err := m.Serialize(); err != nil {
						log.WithFields(log.Fields{
							"Topic": "mrt",
							"Data":  e,
							"Error": err,
						}).Warn("Failed to serialize event")
					} else {
						b.Write(buf)
						if b.Len() > 1*1000*1000 {
							w(b.Bytes())
							b.Reset()
						}
					}
				}
			}
			if b.Len() > 0 {
				w(b.Bytes())
			}
		}
		rotate := func() {
			m.file.Close()
			file, err := mrtFileOpen(m.c.FileName, m.rotationInterval)
			if err == nil {
				m.file = file
			} else {
				log.WithFields(log.Fields{
					"Topic": "mrt",
					"Error": err,
				}).Warn("can't rotate MRT file")
			}
		}

		select {
		case <-m.dead:
			drain(nil)
			return nil
		case e := <-w.Event():
			drain(e)
			if m.c.DumpType == config.MRT_TYPE_TABLE && m.rotationInterval != 0 {
				rotate()
			}
		case <-rotator.C:
			if m.c.DumpType == config.MRT_TYPE_UPDATES {
				rotate()
			} else {
				w.Generate(WATCH_EVENT_TYPE_TABLE)
			}
		case <-dump.C:
			w.Generate(WATCH_EVENT_TYPE_TABLE)
		}
	}
}

func mrtFileOpen(filename string, interval uint64) (*os.File, error) {
	realname := filename
	if interval != 0 {
		realname = time.Now().Format(filename)
	}
	log.WithFields(log.Fields{
		"Topic":         "mrt",
		"Filename":      realname,
		"Dump Interval": interval,
	}).Debug("Setting new MRT destination file")

	i := len(realname)
	for i > 0 && os.IsPathSeparator(realname[i-1]) {
		// skip trailing path separators
		i--
	}
	j := i

	for j > 0 && !os.IsPathSeparator(realname[j-1]) {
		j--
	}

	if j > 0 {
		if err := os.MkdirAll(realname[0:j-1], 0755); err != nil {
			log.WithFields(log.Fields{
				"Topic": "mrt",
				"Error": err,
			}).Warn("can't create MRT destination directory")
			return nil, err
		}
	}

	file, err := os.OpenFile(realname, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
	if err != nil {
		log.WithFields(log.Fields{
			"Topic": "mrt",
			"Error": err,
		}).Warn("can't create MRT destination file")
	}
	return file, err
}

func newMrtWriter(s *BgpServer, c *config.MrtConfig, rInterval, dInterval uint64) (*mrtWriter, error) {
	file, err := mrtFileOpen(c.FileName, rInterval)
	if err != nil {
		return nil, err
	}
	m := mrtWriter{
		s:                s,
		c:                c,
		file:             file,
		rotationInterval: rInterval,
		dumpInterval:     dInterval,
	}
	go m.loop()
	return &m, nil
}

type mrtManager struct {
	bgpServer *BgpServer
	writer    map[string]*mrtWriter
}

func (m *mrtManager) enable(c *config.MrtConfig) error {
	if _, ok := m.writer[c.FileName]; ok {
		return fmt.Errorf("%s already exists", c.FileName)
	}

	rInterval := c.RotationInterval
	dInterval := c.DumpInterval

	setRotationMin := func() {
		if rInterval < MIN_ROTATION_INTERVAL {
			log.WithFields(log.Fields{
				"Topic": "MRT",
			}).Info("minimum mrt rotation interval is %d seconds", MIN_ROTATION_INTERVAL)
			rInterval = MIN_ROTATION_INTERVAL
		}
	}

	if c.DumpType == config.MRT_TYPE_TABLE {
		if rInterval == 0 {
			if dInterval < MIN_DUMP_INTERVAL {
				log.WithFields(log.Fields{
					"Topic": "MRT",
				}).Info("minimum mrt dump interval is %d seconds", MIN_DUMP_INTERVAL)
				dInterval = MIN_DUMP_INTERVAL
			}
		} else if dInterval == 0 {
			setRotationMin()
		} else {
			return fmt.Errorf("can't specify both intervals in the table dump type")
		}
	} else if c.DumpType == config.MRT_TYPE_UPDATES {
		// ignore the dump interval
		dInterval = 0
		if len(c.TableName) > 0 {
			return fmt.Errorf("can't specify the table name with the update dump type")
		}
		setRotationMin()
	}

	w, err := newMrtWriter(m.bgpServer, c, rInterval, dInterval)
	if err == nil {
		m.writer[c.FileName] = w
	}
	return err
}

func (m *mrtManager) disable(c *config.MrtConfig) error {
	if w, ok := m.writer[c.FileName]; !ok {
		return fmt.Errorf("%s doesn't exists", c.FileName)
	} else {
		w.Stop()
		delete(m.writer, c.FileName)
	}
	return nil
}

func newMrtManager(s *BgpServer) *mrtManager {
	return &mrtManager{
		bgpServer: s,
		writer:    make(map[string]*mrtWriter),
	}
}
<commit_msg>server/mrt: Rename a variable to avoid name collision<commit_after>// Copyright (C) 2016 Nippon Telegraph and Telephone Corporation.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"bytes"
	"fmt"
	"os"
	"time"

	"github.com/osrg/gobgp/config"
	"github.com/osrg/gobgp/packet/bgp"
	"github.com/osrg/gobgp/packet/mrt"
	"github.com/osrg/gobgp/table"
	log "github.com/sirupsen/logrus"
)

const (
	MIN_ROTATION_INTERVAL = 60
	MIN_DUMP_INTERVAL     = 60
)

type mrtWriter struct {
	dead             chan struct{}
	s                *BgpServer
	c                *config.MrtConfig
	file             *os.File
	rotationInterval uint64
	dumpInterval     uint64
}

func (m *mrtWriter) Stop() {
	close(m.dead)
}

func (m *mrtWriter) loop() error {
	ops := []WatchOption{}
	switch m.c.DumpType {
	case config.MRT_TYPE_UPDATES:
		ops = append(ops, WatchUpdate(false))
	case config.MRT_TYPE_TABLE:
		if len(m.c.TableName) > 0 {
			ops = append(ops, WatchTableName(m.c.TableName))
		}
	}
	w := m.s.Watch(ops...)
	rotator := func() *time.Ticker {
		if m.rotationInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.rotationInterval))
	}()
	dump := func() *time.Ticker {
		if m.dumpInterval == 0 {
			return &time.Ticker{}
		}
		return time.NewTicker(time.Second * time.Duration(m.dumpInterval))
	}()

	defer func() {
		if m.file != nil {
			m.file.Close()
		}
		if m.rotationInterval != 0 {
			rotator.Stop()
		}
		if m.dumpInterval != 0 {
			dump.Stop()
		}
		w.Stop()
	}()

	for {
		serialize := func(ev WatchEvent) []*mrt.MRTMessage {
			msg := make([]*mrt.MRTMessage, 0, 1)
			switch e := ev.(type) {
			case *WatchEventUpdate:
				subtype := mrt.MESSAGE_AS4
				mp := mrt.NewBGP4MPMessage(e.PeerAS, e.LocalAS, 0, e.PeerAddress.String(), e.LocalAddress.String(), e.FourBytesAs, nil)
				mp.BGPMessagePayload = e.Payload
				if e.FourBytesAs == false {
					subtype = mrt.MESSAGE
				}
				if bm, err := mrt.NewMRTMessage(uint32(e.Timestamp.Unix()), mrt.BGP4MP, subtype, mp); err != nil {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Data":  e,
						"Error": err,
					}).Warn("Failed to create MRT message in serialize()")
				} else {
					msg = append(msg, bm)
				}
			case *WatchEventTable:
				t := uint32(time.Now().Unix())
				peers := make([]*mrt.Peer, 0, len(e.Neighbor))
				for _, pconf := range e.Neighbor {
					peers = append(peers, mrt.NewPeer(pconf.State.RemoteRouterId, pconf.State.NeighborAddress, pconf.Config.PeerAs, true))
				}
				if bm, err := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, mrt.PEER_INDEX_TABLE, mrt.NewPeerIndexTable(e.RouterId, "", peers)); err != nil {
					break
				} else {
					msg = append(msg, bm)
				}

				idx := func(p *table.Path) uint16 {
					for i, pconf := range e.Neighbor {
						if p.GetSource().Address.String() == pconf.State.NeighborAddress {
							return uint16(i)
						}
					}
					return uint16(len(e.Neighbor))
				}

				subtype := func(p *table.Path) mrt.MRTSubTypeTableDumpv2 {
					switch p.GetRouteFamily() {
					case bgp.RF_IPv4_UC:
						return mrt.RIB_IPV4_UNICAST
					case bgp.RF_IPv4_MC:
						return mrt.RIB_IPV4_MULTICAST
					case bgp.RF_IPv6_UC:
						return mrt.RIB_IPV6_UNICAST
					case bgp.RF_IPv6_MC:
						return mrt.RIB_IPV6_MULTICAST
					}
					return mrt.RIB_GENERIC
				}

				seq := uint32(0)
				for _, pathList := range e.PathList {
					entries := make([]*mrt.RibEntry, 0, len(pathList))
					for _, path := range pathList {
						if path.IsLocal() {
							continue
						}
						entries = append(entries, mrt.NewRibEntry(idx(path), uint32(path.GetTimestamp().Unix()), 0, path.GetPathAttrs()))
					}
					if len(entries) > 0 {
						bm, _ := mrt.NewMRTMessage(t, mrt.TABLE_DUMPv2, subtype(pathList[0]), mrt.NewRib(seq, pathList[0].GetNlri(), entries))
						msg = append(msg, bm)
						seq++
					}
				}
			}
			return msg
		}

		drain := func(ev WatchEvent) {
			events := make([]WatchEvent, 0, 1+len(w.Event()))
			if ev != nil {
				events = append(events, ev)
			}

			for len(w.Event()) > 0 {
				events = append(events, <-w.Event())
			}

			w := func(buf []byte) {
				if _, err := m.file.Write(buf); err == nil {
					m.file.Sync()
				} else {
					log.WithFields(log.Fields{
						"Topic": "mrt",
						"Error": err,
					}).Warn("Can't write to destination MRT file")
				}
			}

			var b bytes.Buffer
			for _, e := range events {
				for _, m := range serialize(e) {
					if buf, err := m.Serialize(); err != nil {
						log.WithFields(log.Fields{
							"Topic": "mrt",
							"Data":  e,
							"Error": err,
						}).Warn("Failed to serialize event")
					} else {
						b.Write(buf)
						if b.Len() > 1*1000*1000 {
							w(b.Bytes())
							b.Reset()
						}
					}
				}
			}
			if b.Len() > 0 {
				w(b.Bytes())
			}
		}
		rotate := func() {
			m.file.Close()
			file, err := mrtFileOpen(m.c.FileName, m.rotationInterval)
			if err == nil {
				m.file = file
			} else {
				log.WithFields(log.Fields{
					"Topic": "mrt",
					"Error": err,
				}).Warn("can't rotate MRT file")
			}
		}

		select {
		case <-m.dead:
			drain(nil)
			return nil
		case e := <-w.Event():
			drain(e)
			if m.c.DumpType == config.MRT_TYPE_TABLE && m.rotationInterval != 0 {
				rotate()
			}
		case <-rotator.C:
			if m.c.DumpType == config.MRT_TYPE_UPDATES {
				rotate()
			} else {
				w.Generate(WATCH_EVENT_TYPE_TABLE)
			}
		case <-dump.C:
			w.Generate(WATCH_EVENT_TYPE_TABLE)
		}
	}
}

func mrtFileOpen(filename string, interval uint64) (*os.File, error) {
	realname := filename
	if interval != 0 {
		realname = time.Now().Format(filename)
	}
	log.WithFields(log.Fields{
		"Topic":         "mrt",
		"Filename":      realname,
		"Dump Interval": interval,
	}).Debug("Setting new MRT destination file")

	i := len(realname)
	for i > 0 && os.IsPathSeparator(realname[i-1]) {
		// skip trailing path separators
		i--
	}
	j := i

	for j > 0 && !os.IsPathSeparator(realname[j-1]) {
		j--
	}

	if j > 0 {
		if err := os.MkdirAll(realname[0:j-1], 0755); err != nil {
			log.WithFields(log.Fields{
				"Topic": "mrt",
				"Error": err,
			}).Warn("can't create MRT destination directory")
			return nil, err
		}
	}

	file, err := os.OpenFile(realname, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
	if err != nil {
		log.WithFields(log.Fields{
			"Topic": "mrt",
			"Error": err,
		}).Warn("can't create MRT destination file")
	}
	return file, err
}

func newMrtWriter(s *BgpServer, c *config.MrtConfig, rInterval, dInterval uint64) (*mrtWriter, error) {
	file, err := mrtFileOpen(c.FileName, rInterval)
	if err != nil {
		return nil, err
	}
	m := mrtWriter{
		s:                s,
		c:                c,
		file:             file,
		rotationInterval: rInterval,
		dumpInterval:     dInterval,
	}
	go m.loop()
	return &m, nil
}

type mrtManager struct {
	bgpServer *BgpServer
	writer    map[string]*mrtWriter
}

func (m *mrtManager) enable(c *config.MrtConfig) error {
	if _, ok := m.writer[c.FileName]; ok {
		return fmt.Errorf("%s already exists", c.FileName)
	}

	rInterval := c.RotationInterval
	dInterval := c.DumpInterval

	setRotationMin := func() {
		if rInterval < MIN_ROTATION_INTERVAL {
			log.WithFields(log.Fields{
				"Topic": "MRT",
			}).Info("minimum mrt rotation interval is %d seconds", MIN_ROTATION_INTERVAL)
			rInterval = MIN_ROTATION_INTERVAL
		}
	}

	if c.DumpType == config.MRT_TYPE_TABLE {
		if rInterval == 0 {
			if dInterval < MIN_DUMP_INTERVAL {
				log.WithFields(log.Fields{
					"Topic": "MRT",
				}).Info("minimum mrt dump interval is %d seconds", MIN_DUMP_INTERVAL)
				dInterval = MIN_DUMP_INTERVAL
			}
		} else if dInterval == 0 {
			setRotationMin()
		} else {
			return fmt.Errorf("can't specify both intervals in the table dump type")
		}
	} else if c.DumpType == config.MRT_TYPE_UPDATES {
		// ignore the dump interval
		dInterval = 0
		if len(c.TableName) > 0 {
			return fmt.Errorf("can't specify the table name with the update dump type")
		}
		setRotationMin()
	}

	w, err := newMrtWriter(m.bgpServer, c, rInterval, dInterval)
	if err == nil {
		m.writer[c.FileName] = w
	}
	return err
}

func (m *mrtManager) disable(c *config.MrtConfig) error {
	if w, ok := m.writer[c.FileName]; !ok {
		return fmt.Errorf("%s doesn't exists", c.FileName)
	} else {
		w.Stop()
		delete(m.writer, c.FileName)
	}
	return nil
}

func newMrtManager(s *BgpServer) *mrtManager {
	return &mrtManager{
		bgpServer: s,
		writer:    make(map[string]*mrtWriter),
	}
}
<|endoftext|>


<commit_before>/*
	This file handles RPC connections, usually from DVID clients.
*/

package server

import (
	"fmt"
	"os"
	"time"

	"github.com/janelia-flyem/dvid/datastore"
	"github.com/janelia-flyem/dvid/dvid"
)

const RPCHelpMessage = `Commands executed on the server (rpc address = %s):

	help
	about
	shutdown

	types
	types <datatype name> help

	datasets info
	datasets new         (returns UUID of dataset's root node)

	dataset <UUID> new <datatype name> <data name> <datatype-specific config>...
	dataset <UUID> new <datatype name> <data name> <datatype-specific config>...

	dataset <UUID> <data name> help

	node <UUID> lock
	node <UUID> branch   (returns UUID of new child node)
	node <UUID> <data name> <type-specific commands>

%s

For further information, use a web browser to visit the server for this
datastore:

	http://%s
`

// RPCConnection will export all of its functions for rpc access.
type RPCConnection struct{}

// Do acts as a switchboard for remote command execution
func (c *RPCConnection) Do(cmd datastore.Request, reply *datastore.Response) error {
	if reply == nil {
		dvid.Log(dvid.Debug, "reply is nil coming in!\n")
		return nil
	}
	if cmd.Name() == "" {
		return fmt.Errorf("Server error: got empty command!")
	}
	if runningService.Service == nil {
		return fmt.Errorf("Datastore not open!  Cannot execute command.")
	}

	switch cmd.Name() {

	case "help":
		reply.Text = fmt.Sprintf(RPCHelpMessage,
			runningService.RPCAddress, runningService.SupportedDataChart(),
			runningService.WebAddress)

	case "about":
		reply.Text = fmt.Sprintf("%s\n", runningService.About())

	case "shutdown":
		Shutdown()
		// Make this process shutdown in a second to allow time for RPC to finish.
		// TODO -- Better way to do this?
		go func() {
			time.Sleep(1 * time.Second)
			os.Exit(0)
		}()
		reply.Text = fmt.Sprintf("DVID server at %s has been halted.",
			runningService.RPCAddress)

	case "types":
		if len(cmd.Command) == 1 {
			reply.Text = runningService.SupportedDataChart()
		} else {
			if len(cmd.Command) != 3 || cmd.Command[2] != "help" {
				return fmt.Errorf("Unknown types command: %q", cmd.Command)
			}
			var typename string
			cmd.CommandArgs(1, &typename)
			typeservice, err := datastore.TypeServiceByName(typename)
			if err != nil {
				return err
			}
			reply.Text = typeservice.Help()
		}

	case "datasets":
		var subcommand string
		cmd.CommandArgs(1, &subcommand)
		switch subcommand {
		case "info":
			jsonStr, err := runningService.DatasetsListJSON()
			if err != nil {
				return err
			}
			reply.Text = jsonStr
		case "new":
			uuid, _, err := runningService.NewDataset()
			if err != nil {
				return err
			}
			reply.Text = string(uuid)
		default:
			return fmt.Errorf("Unknown datasets command: %q", subcommand)
		}

	case "dataset":
		var uuidStr, subcommand, typename, dataname string
		cmd.CommandArgs(1, &uuidStr, &subcommand)
		uuid, _, _, err := runningService.NodeIDFromString(uuidStr)
		if err != nil {
			return err
		}
		switch subcommand {
		case "new":
			cmd.CommandArgs(3, &typename, &dataname)
			config := cmd.Settings()
			err = runningService.NewData(uuid, typename, dataname, config)
			if err != nil {
				return err
			}
			reply.Text = fmt.Sprintf("Data %q [%s] added to node %s", dataname, typename, uuidStr)
		default:
			dataname := datastore.DataString(subcommand)
			dataservice, err := runningService.DataService(uuid, dataname)
			if err != nil {
				return err
			}
			var subcommand2 string
			cmd.CommandArgs(3, &subcommand2)
			if subcommand2 == "help" {
				reply.Text = dataservice.Help()
			} else {
				return fmt.Errorf("Unknown command: %q", cmd)
			}
		}

	case "node":
		var uuidStr, descriptor string
		cmd.CommandArgs(1, &uuidStr, &descriptor)
		uuid, _, _, err := runningService.NodeIDFromString(uuidStr)
		if err != nil {
			return err
		}
		switch descriptor {
		case "lock":
			err := runningService.Lock(uuid)
			if err != nil {
				return err
			}
		case "branch":
			newuuid, err := runningService.NewVersion(uuid)
			if err != nil {
				return err
			}
			reply.Text = string(newuuid)

		default:
			dataname := datastore.DataString(descriptor)
			var subcommand string
			cmd.CommandArgs(3, &subcommand)
			dataservice, err := runningService.DataService(uuid, dataname)
			if err != nil {
				return err
			}
			if subcommand == "help" {
				reply.Text = dataservice.Help()
				return nil
			}
			return dataservice.DoRPC(cmd, reply)
		}

	default:
		return fmt.Errorf("Unknown command: '%s'", cmd)
	}
	return nil
}
<commit_msg>Add shutdown log message<commit_after>/*
	This file handles RPC connections, usually from DVID clients.
*/

package server

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/janelia-flyem/dvid/datastore"
	"github.com/janelia-flyem/dvid/dvid"
)

const RPCHelpMessage = `Commands executed on the server (rpc address = %s):

	help
	about
	shutdown

	types
	types <datatype name> help

	datasets info
	datasets new         (returns UUID of dataset's root node)

	dataset <UUID> new <datatype name> <data name> <datatype-specific config>...
	dataset <UUID> new <datatype name> <data name> <datatype-specific config>...

	dataset <UUID> <data name> help

	node <UUID> lock
	node <UUID> branch   (returns UUID of new child node)
	node <UUID> <data name> <type-specific commands>

%s

For further information, use a web browser to visit the server for this
datastore:

	http://%s
`

// RPCConnection will export all of its functions for rpc access.
type RPCConnection struct{}

// Do acts as a switchboard for remote command execution
func (c *RPCConnection) Do(cmd datastore.Request, reply *datastore.Response) error {
	if reply == nil {
		dvid.Log(dvid.Debug, "reply is nil coming in!\n")
		return nil
	}
	if cmd.Name() == "" {
		return fmt.Errorf("Server error: got empty command!")
	}
	if runningService.Service == nil {
		return fmt.Errorf("Datastore not open!  Cannot execute command.")
	}

	switch cmd.Name() {

	case "help":
		reply.Text = fmt.Sprintf(RPCHelpMessage,
			runningService.RPCAddress, runningService.SupportedDataChart(),
			runningService.WebAddress)

	case "about":
		reply.Text = fmt.Sprintf("%s\n", runningService.About())

	case "shutdown":
		Shutdown()
		// Make this process shutdown in a second to allow time for RPC to finish.
		// TODO -- Better way to do this?
		log.Printf("DVID server halted due to 'shutdown' command.")
		go func() {
			time.Sleep(1 * time.Second)
			os.Exit(0)
		}()
		reply.Text = fmt.Sprintf("DVID server at %s has been halted.",
			runningService.RPCAddress)

	case "types":
		if len(cmd.Command) == 1 {
			reply.Text = runningService.SupportedDataChart()
		} else {
			if len(cmd.Command) != 3 || cmd.Command[2] != "help" {
				return fmt.Errorf("Unknown types command: %q", cmd.Command)
			}
			var typename string
			cmd.CommandArgs(1, &typename)
			typeservice, err := datastore.TypeServiceByName(typename)
			if err != nil {
				return err
			}
			reply.Text = typeservice.Help()
		}

	case "datasets":
		var subcommand string
		cmd.CommandArgs(1, &subcommand)
		switch subcommand {
		case "info":
			jsonStr, err := runningService.DatasetsListJSON()
			if err != nil {
				return err
			}
			reply.Text = jsonStr
		case "new":
			uuid, _, err := runningService.NewDataset()
			if err != nil {
				return err
			}
			reply.Text = string(uuid)
		default:
			return fmt.Errorf("Unknown datasets command: %q", subcommand)
		}

	case "dataset":
		var uuidStr, subcommand, typename, dataname string
		cmd.CommandArgs(1, &uuidStr, &subcommand)
		uuid, _, _, err := runningService.NodeIDFromString(uuidStr)
		if err != nil {
			return err
		}
		switch subcommand {
		case "new":
			cmd.CommandArgs(3, &typename, &dataname)
			config := cmd.Settings()
			err = runningService.NewData(uuid, typename, dataname, config)
			if err != nil {
				return err
			}
			reply.Text = fmt.Sprintf("Data %q [%s] added to node %s", dataname, typename, uuidStr)
		default:
			dataname := datastore.DataString(subcommand)
			dataservice, err := runningService.DataService(uuid, dataname)
			if err != nil {
				return err
			}
			var subcommand2 string
			cmd.CommandArgs(3, &subcommand2)
			if subcommand2 == "help" {
				reply.Text = dataservice.Help()
			} else {
				return fmt.Errorf("Unknown command: %q", cmd)
			}
		}

	case "node":
		var uuidStr, descriptor string
		cmd.CommandArgs(1, &uuidStr, &descriptor)
		uuid, _, _, err := runningService.NodeIDFromString(uuidStr)
		if err != nil {
			return err
		}
		switch descriptor {
		case "lock":
			err := runningService.Lock(uuid)
			if err != nil {
				return err
			}
		case "branch":
			newuuid, err := runningService.NewVersion(uuid)
			if err != nil {
				return err
			}
			reply.Text = string(newuuid)

		default:
			dataname := datastore.DataString(descriptor)
			var subcommand string
			cmd.CommandArgs(3, &subcommand)
			dataservice, err := runningService.DataService(uuid, dataname)
			if err != nil {
				return err
			}
			if subcommand == "help" {
				reply.Text = dataservice.Help()
				return nil
			}
			return dataservice.DoRPC(cmd, reply)
		}

	default:
		return fmt.Errorf("Unknown command: '%s'", cmd)
	}
	return nil
}
<|endoftext|>


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

package bfs

import (
	"math/rand"
	"sync"
	"time"

	. "../../equtils"
	. "../../historystorage"
)

type BFSParam struct {
	interval time.Duration // in millisecond
}

type BFS struct {
	nextEventChan chan *Event
	randGen       *rand.Rand
	queueMutex    *sync.Mutex

	eventQueue []*Event // high priority
	dumb       bool

	param *BFSParam
}

func constrBFSParam(rawParam map[string]interface{}) *BFSParam {
	var param BFSParam

	if _, ok := rawParam["interval"]; ok {
		param.interval = time.Duration(int(rawParam["interval"].(float64)))
	} else {
		param.interval = time.Duration(100) // default: 100ms
	}

	return &param
}

func (this *BFS) Init(storage HistoryStorage, param map[string]interface{}) {
	this.param = constrBFSParam(param)

	prefix := make([]Event, 0)

	go func() {
		for {
			time.Sleep(this.param.interval * time.Millisecond)

			this.queueMutex.Lock()
			if len(this.eventQueue) == 0 {
				Log("no event is queued")
				this.queueMutex.Unlock()
				continue
			}

			nextIdx := -1
			for i := 0; i < len(this.eventQueue); i++ {
				// TODO: match every event in the queue, wait for more intervals, etc
				tmpprefix := append(prefix, *this.eventQueue[i])

				matched := storage.Search(tmpprefix)
				if len(matched) == 0 {
					nextIdx = i
					break
				}
			}

			if nextIdx == -1 {
				// no seen events in the past

				Log("all histories contain current prefix, go dumb")

				evQ := this.eventQueue
				this.eventQueue = []*Event{}
				this.dumb = true
				this.queueMutex.Unlock()

				go func() {
					for _, e := range evQ {
						this.nextEventChan <- e
					}
				}()

				// end this goroutine, below events are processed in a manner of dumb policy
				break
			}

			next := this.eventQueue[nextIdx]
			this.eventQueue = append(this.eventQueue[:nextIdx], this.eventQueue[nextIdx+1:]...)

			this.queueMutex.Unlock()

			prefix = append(prefix, *next)

			this.nextEventChan <- next
		}
	}()
}

func (this *BFS) Name() string {
	return "BFS"
}

func (this *BFS) GetNextEventChan() chan *Event {
	return this.nextEventChan
}

func (this *BFS) QueueNextEvent(procId string, ev *Event) {
	this.queueMutex.Lock()

	if !this.dumb {
		this.eventQueue = append(this.eventQueue, ev)
	} else {
		go func() {
			this.nextEventChan <- ev
		}()
	}

	this.queueMutex.Unlock()
}

func BFSNew() *BFS {
	nextEventChan := make(chan *Event)
	eventQueue := make([]*Event, 0)
	mutex := new(sync.Mutex)
	r := rand.New(rand.NewSource(time.Now().Unix()))

	return &BFS{
		nextEventChan,
		r,
		mutex,
		eventQueue,
		false,
		nil,
	}
}
<commit_msg>explorerpolicy, bfs: be more greedy<commit_after>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package bfs

import (
	"math/rand"
	"sync"
	"time"

	. "../../equtils"
	. "../../historystorage"
)

type BFSParam struct {
	interval time.Duration // in millisecond
}

type BFS struct {
	nextEventChan chan *Event
	randGen       *rand.Rand
	queueMutex    *sync.Mutex

	eventQueue []*Event // high priority

	param *BFSParam
}

func constrBFSParam(rawParam map[string]interface{}) *BFSParam {
	var param BFSParam

	if _, ok := rawParam["interval"]; ok {
		param.interval = time.Duration(int(rawParam["interval"].(float64)))
	} else {
		param.interval = time.Duration(100) // default: 100ms
	}

	return &param
}

func (this *BFS) Init(storage HistoryStorage, param map[string]interface{}) {
	this.param = constrBFSParam(param)

	prefix := make([]Event, 0)

	go func() {
		for {
			time.Sleep(this.param.interval * time.Millisecond)

			this.queueMutex.Lock()
			if len(this.eventQueue) == 0 {
				Log("no event is queued")
				this.queueMutex.Unlock()
				continue
			}

			nextIdx := -1
			for i := 0; i < len(this.eventQueue); i++ {
				// TODO: match every event in the queue, wait for more intervals, etc
				tmpprefix := append(prefix, *this.eventQueue[i])

				matched := storage.Search(tmpprefix)
				if len(matched) == 0 {
					nextIdx = i
					break
				}
			}

			if nextIdx == -1 {
				nextIdx = this.randGen.Int() % len(this.eventQueue)
			}

			next := this.eventQueue[nextIdx]
			this.eventQueue = append(this.eventQueue[:nextIdx], this.eventQueue[nextIdx+1:]...)

			this.queueMutex.Unlock()

			prefix = append(prefix, *next)

			this.nextEventChan <- next
		}
	}()
}

func (this *BFS) Name() string {
	return "BFS"
}

func (this *BFS) GetNextEventChan() chan *Event {
	return this.nextEventChan
}

func (this *BFS) QueueNextEvent(procId string, ev *Event) {
	this.queueMutex.Lock()
	this.eventQueue = append(this.eventQueue, ev)
	this.queueMutex.Unlock()
}

func BFSNew() *BFS {
	nextEventChan := make(chan *Event)
	eventQueue := make([]*Event, 0)
	mutex := new(sync.Mutex)
	r := rand.New(rand.NewSource(time.Now().Unix()))

	return &BFS{
		nextEventChan,
		r,
		mutex,
		eventQueue,
		nil,
	}
}
<|endoftext|>


<commit_before>package logtap

import (
	"github.com/jeromer/syslogparser/rfc3164"
	"github.com/jeromer/syslogparser/rfc5424"
	"github.com/pagodabox/golang-hatchet"
	"net"
	"strconv"
	"time"
)

type SyslogCollector struct {
	wChan chan Message
	log   hatchet.Logger

	Port string
}

// NewSyslogCollector creates a new syslog collector
func NewSyslogCollector(port string) *SyslogCollector {
	return &SyslogCollector{
		Port:  port,
		wChan: make(chan Message),
	}
}

// SetLogger really allows the logtap main struct
// to assign its own logger to the syslog collector
func (s *SyslogCollector) SetLogger(l hatchet.Logger) {
	s.log = l
}

// CollectChan grats access to the collection channel
// any logs collected from syslog will be translated into a message
// and dropped on this channel for processing
func (s *SyslogCollector) CollectChan() chan Message {
	return s.wChan
}

// Start begins listening to the syslog port transfers all
// syslog messages on the wChan
func (s *SyslogCollector) Start() {
	go func() {

		udpAddress, err := net.ResolveUDPAddr("udp4", ("0.0.0.0:" + s.Port))
		if err != nil {
			s.log.Error("[LOGTAP][syslog]resolving UDP address on ", s.Port)
			s.log.Error("[LOGTAP][syslog]" + err.Error())
			return
		}

		conn, err := net.ListenUDP("udp", udpAddress)
		if err != nil {
			s.log.Error("[LOGTAP][syslog]listening on UDP port ", s.Port)
			s.log.Error("[LOGTAP][syslog]" + err.Error())
			return
		}
		defer conn.Close()

		var buf []byte = make([]byte, 1024)
		for {
			s.log.Debug("[LOGTAP][syslog][start] listen")
			n, address, err := conn.ReadFromUDP(buf)
			if err != nil {
				s.log.Error("[LOGTAP][syslog]reading data from connection")
				s.log.Error("[LOGTAP][syslog]" + err.Error())
			}
			if address != nil {
				s.log.Debug("[LOGTAP][syslog][start] got message from " + address.String() + " with n = " + strconv.Itoa(n))
				if n > 0 {
					msg := s.parseMessage(buf[0:n])
					s.log.Debug("[LOGTAP][syslog][start] msg content: " + msg.Content)
					s.wChan <- msg
				}
			}
		}
	}()
}

// parseMessage parses the syslog message and returns a msg
// if the msg is not parsable or a standard formatted syslog message
// it will drop the whole message into the content and make up a timestamp
// and a priority
func (s *SyslogCollector) parseMessage(b []byte) (msg Message) {
	msg.Type = "syslog"
	p := rfc3164.NewParser(b)
	err := p.Parse()
	if err == nil {
		parsedData := p.Dump()
		msg.Time = parsedData["timestamp"].(time.Time)
		msg.Priority = adjustInt(parsedData["priority"].(int) % 8)
		msg.Content = parsedData["content"].(string)
	} else {
		p := rfc5424.NewParser(b)
		err := p.Parse()
		if err == nil {
			parsedData := p.Dump()
			msg.Time = parsedData["timestamp"].(time.Time)
			msg.Priority = adjustInt(parsedData["priority"].(int) % 8)
			msg.Content = parsedData["content"].(string)
		} else {
			s.log.Error("[LOGTAP]Unable to parse data: " + string(b))
			msg.Time = time.Now()
			msg.Priority = 1
			msg.Content = string(b)
		}
	}
	return
}

// I need to adjust the possible prioritys from rfc3164 and rfc5424
// to the 5 priority options.
func adjustInt(in int) int {
	if in < 3 {
		return 0
	}
	if in < 5 {
		return in - 2
	}
	return in - 3
}
<commit_msg>put the tag into the context of the message<commit_after>package logtap

import (
	"github.com/jeromer/syslogparser/rfc3164"
	"github.com/jeromer/syslogparser/rfc5424"
	"github.com/pagodabox/golang-hatchet"
	"net"
	"strconv"
	"time"
)

type SyslogCollector struct {
	wChan chan Message
	log   hatchet.Logger

	Port string
}

// NewSyslogCollector creates a new syslog collector
func NewSyslogCollector(port string) *SyslogCollector {
	return &SyslogCollector{
		Port:  port,
		wChan: make(chan Message),
	}
}

// SetLogger really allows the logtap main struct
// to assign its own logger to the syslog collector
func (s *SyslogCollector) SetLogger(l hatchet.Logger) {
	s.log = l
}

// CollectChan grats access to the collection channel
// any logs collected from syslog will be translated into a message
// and dropped on this channel for processing
func (s *SyslogCollector) CollectChan() chan Message {
	return s.wChan
}

// Start begins listening to the syslog port transfers all
// syslog messages on the wChan
func (s *SyslogCollector) Start() {
	go func() {

		udpAddress, err := net.ResolveUDPAddr("udp4", ("0.0.0.0:" + s.Port))
		if err != nil {
			s.log.Error("[LOGTAP][syslog]resolving UDP address on ", s.Port)
			s.log.Error("[LOGTAP][syslog]" + err.Error())
			return
		}

		conn, err := net.ListenUDP("udp", udpAddress)
		if err != nil {
			s.log.Error("[LOGTAP][syslog]listening on UDP port ", s.Port)
			s.log.Error("[LOGTAP][syslog]" + err.Error())
			return
		}
		defer conn.Close()

		var buf []byte = make([]byte, 1024)
		for {
			s.log.Debug("[LOGTAP][syslog][start] listen")
			n, address, err := conn.ReadFromUDP(buf)
			if err != nil {
				s.log.Error("[LOGTAP][syslog]reading data from connection")
				s.log.Error("[LOGTAP][syslog]" + err.Error())
			}
			if address != nil {
				s.log.Debug("[LOGTAP][syslog][start] got message from " + address.String() + " with n = " + strconv.Itoa(n))
				if n > 0 {
					msg := s.parseMessage(buf[0:n])
					s.log.Debug("[LOGTAP][syslog][start] msg content: " + msg.Content)
					s.wChan <- msg
				}
			}
		}
	}()
}

// parseMessage parses the syslog message and returns a msg
// if the msg is not parsable or a standard formatted syslog message
// it will drop the whole message into the content and make up a timestamp
// and a priority
func (s *SyslogCollector) parseMessage(b []byte) (msg Message) {
	msg.Type = "syslog"
	p := rfc3164.NewParser(b)
	err := p.Parse()
	if err == nil {
		parsedData := p.Dump()
		// fmt.Printf("%#v\n",parsedData)
		msg.Time = parsedData["timestamp"].(time.Time)
		msg.Priority = adjustInt(parsedData["priority"].(int) % 8)
		msg.Content = parsedData["tag"].(string) + " " + parsedData["content"].(string)
	} else {
		p := rfc5424.NewParser(b)
		err := p.Parse()
		if err == nil {
			parsedData := p.Dump()
			// fmt.Printf("%#v\n",parsedData)
			msg.Time = parsedData["timestamp"].(time.Time)
			msg.Priority = adjustInt(parsedData["priority"].(int) % 8)
			msg.Content = parsedData["tag"].(string) + " " + parsedData["content"].(string)
		} else {
			s.log.Error("[LOGTAP]Unable to parse data: " + string(b))
			msg.Time = time.Now()
			msg.Priority = 1
			msg.Content = string(b)
		}
	}
	return
}

// I need to adjust the possible prioritys from rfc3164 and rfc5424
// to the 5 priority options.
func adjustInt(in int) int {
	if in < 3 {
		return 0
	}
	if in < 5 {
		return in - 2
	}
	return in - 3
}
<|endoftext|>


<commit_before>package main

import "sync"
import "sync/atomic"
import "runtime"
import "flag"
import "reflect"
import "unsafe"
import "os"
import "log"
import "sort"
import "time"
import "strings"
import "strconv"
import "net"
import "fmt"
import "compress/flate"
import "runtime/pprof"

import "github.com/prataprc/gofast"

var options struct {
	do       string
	count    int
	routines int
	conns    int
	addr     string
	payload  int
	log      string
}

func argParse() {
	flag.StringVar(&options.do, "do", "post",
		"post / request / stream benchmark to do")
	flag.IntVar(&options.conns, "conns", 1,
		"number of connections to use")
	flag.IntVar(&options.routines, "routines", 1,
		"number of concurrent routines per connection")
	flag.IntVar(&options.count, "count", 1,
		"number of requests per routine")
	flag.StringVar(&options.addr, "addr", "127.0.0.1:9998",
		"number of concurrent routines")
	flag.StringVar(&options.log, "log", "error",
		"number of concurrent routines")
	flag.IntVar(&options.payload, "payload", 10,
		"payload size to ping pong.")
	flag.Parse()
}

var av = &Average{}

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

	// start cpu profile.
	fname := "client.pprof"
	fd, err := os.Create(fname)
	if err != nil {
		log.Fatalf("unable to create %q: %v\n", fname, err)
	}
	defer fd.Close()
	pprof.StartCPUProfile(fd)
	defer pprof.StopCPUProfile()

	switch options.do {
	case "post":
		doTransport(doPost)
	case "request":
		doTransport(doRequest)
	case "stream":
		doTransport(doStream)
	}

	// take memory profile.
	fname = "client.mprof"
	fd, err = os.Create(fname)
	if err != nil {
		log.Fatal(err)
	}
	defer fd.Close()
	pprof.WriteHeapProfile(fd)
}

func doTransport(callb func(trans *gofast.Transport)) {
	var wg sync.WaitGroup

	ver := testVersion(1)
	n_trans := make([]*gofast.Transport, 0, 16)
	for i := 0; i < options.conns; i++ {
		wg.Add(1)
		config := newconfig("client", 3000, 4000)
		config["tags"] = ""
		conn, err := net.Dial("tcp", options.addr)
		if err != nil {
			panic(err)
		}
		trans, err := gofast.NewTransport(conn, &ver, nil, config)
		if err != nil {
			panic(err)
		}
		n_trans = append(n_trans, trans)
		go func(trans *gofast.Transport) {
			trans.FlushPeriod(100 * time.Millisecond)
			callb(trans)
			wg.Done()
			trans.Close()
		}(trans)
	}
	wg.Wait()
	printCounts(addCounts(n_trans...))
	fmsg := "request stats: n:%v mean:%v var:%v sd:%v\n"
	n, m := av.Count(), time.Duration(av.Mean())
	v, s := time.Duration(av.Variance()), time.Duration(av.Sd())
	fmt.Printf(fmsg, n, m, v, s)
}

func doPost(trans *gofast.Transport) {
	var wg sync.WaitGroup

	trans.SubscribeMessage(&msgPost{}, nil).Handshake()
	msg := &msgPost{data: make([]byte, options.payload)}
	for i := 0; i < options.payload; i++ {
		msg.data[i] = 'a'
	}

	for i := 0; i < options.routines; i++ {
		wg.Add(1)
		go func() {
			for j := 0; j < options.count; j++ {
				since := time.Now()
				if err := trans.Post(msg, true); err != nil {
					fmt.Printf("%v\n", err)
					panic("exit")
				}
				av.Add(uint64(time.Since(since)))
			}
			wg.Done()
		}()
	}
	wg.Wait()
	if _, err := trans.Whoami(); err != nil {
		log.Fatal(err)
	}
}

func doRequest(trans *gofast.Transport) {
	var wg sync.WaitGroup

	trans.SubscribeMessage(&msgReqsp{}, nil).Handshake()
	msg := &msgReqsp{data: make([]byte, options.payload)}
	for i := 0; i < options.payload; i++ {
		msg.data[i] = 'a'
	}

	for i := 0; i < options.routines; i++ {
		wg.Add(1)
		go func() {
			for j := 0; j < options.count; j++ {
				since := time.Now()
				if rmsg, err := trans.Request(msg, true); err != nil {
					fmt.Printf("%v\n", err)
					panic("exit")
				} else {
					trans.Free(rmsg)
				}
				av.Add(uint64(time.Since(since)))
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func doStream(trans *gofast.Transport) {
	var wg sync.WaitGroup

	trans.SubscribeMessage(&msgStream{}, nil).Handshake()

	for i := 0; i < options.routines; i++ {
		wg.Add(1)
		go func() {
			msg := &msgStream{data: make([]byte, 0, options.payload*2)}
			ch := make(chan gofast.Message, 100)
			since := time.Now()
			stream, err := trans.Stream(msg, true, ch)
			if err != nil {
				log.Fatal(err)
			}
			//fmt.Println("started", msg)
			var received int64
			go func() {
				for {
					msg, ok := <-ch
					if ok {
						trans.Free(msg)
						atomic.AddInt64(&received, 1)
						if atomic.LoadInt64(&received) == int64(options.count+1) {
							stream.Close()
							break
						}
						//fmt.Println("received", msg)
					} else {
						break
					}
				}
				wg.Done()
			}()
			av.Add(uint64(time.Since(since)))
			for j := 0; j < options.count; j++ {
				msg.data = msg.data[:cap(msg.data)]
				for i := 0; i < options.payload; i++ {
					msg.data[i] = 'a'
				}
				n := options.payload
				tmp := strconv.AppendInt(msg.data[n:n], int64(j), 10)
				msg.data = msg.data[:n+len(tmp)]

				since := time.Now()
				if err := stream.Stream(msg, true); err != nil {
					fmt.Printf("%v\n", err)
					panic("exit")
				}
				//fmt.Println("count", msg)
				av.Add(uint64(time.Since(since)))
			}
		}()
	}
	wg.Wait()
}

func newconfig(name string, start, end int) map[string]interface{} {
	return map[string]interface{}{
		"name":         name,
		"buffersize":   1024,
		"chansize":     100000,
		"batchsize":    100,
		"tags":         "",
		"opaque.start": start,
		"opaque.end":   end,
		"log.level":    options.log,
		"gzip.file":    flate.BestSpeed,
	}
}

type testVersion int

func (v *testVersion) Less(ver gofast.Version) bool {
	return (*v) < (*ver.(*testVersion))
}

func (v *testVersion) Equal(ver gofast.Version) bool {
	return (*v) == (*ver.(*testVersion))
}

func (v *testVersion) String() string {
	return fmt.Sprintf("%v", int(*v))
}

func (v *testVersion) Marshal(out []byte) int {
	return valuint642cbor(uint64(*v), out)
}

func (v *testVersion) Unmarshal(in []byte) int {
	ln, n := cborItemLength(in)
	*v = testVersion(ln)
	return n
}

func printCounts(counts map[string]uint64) {
	keys := []string{}
	for key := range counts {
		keys = append(keys, key)
	}
	sort.Sort(sort.StringSlice(keys))
	s := []string{}
	for _, key := range keys {
		s = append(s, fmt.Sprintf("%v:%v", key, counts[key]))
	}
	fmt.Println(strings.Join(s, ", "))
}

func addCounts(n_trans ...*gofast.Transport) map[string]uint64 {
	counts := n_trans[0].Counts()
	for _, trans := range n_trans[1:] {
		for k, v := range trans.Counts() {
			counts[k] += v
		}
	}
	return counts
}

func bytes2str(bytes []byte) string {
	if bytes == nil {
		return ""
	}
	sl := (*reflect.SliceHeader)(unsafe.Pointer(&bytes))
	st := &reflect.StringHeader{Data: sl.Data, Len: sl.Len}
	return *(*string)(unsafe.Pointer(st))
}

//-- post message for benchmarking

type msgPost struct {
	data []byte
}

func (msg *msgPost) Id() uint64 {
	return 111
}

func (msg *msgPost) Encode(out []byte) int {
	return valbytes2cbor(msg.data, out)
}

func (msg *msgPost) Decode(in []byte) {
	ln, m := cborItemLength(in)
	if msg.data == nil {
		msg.data = make([]byte, 0, options.payload)
	}
	msg.data = append(msg.data[:0], in[m:m+ln]...)
}

func (msg *msgPost) String() string {
	return "msgPost"
}

//-- reqsp message for benchmarking

type msgReqsp struct {
	data []byte
}

func (msg *msgReqsp) Id() uint64 {
	return 112
}

func (msg *msgReqsp) Encode(out []byte) int {
	return valbytes2cbor(msg.data, out)
}

func (msg *msgReqsp) Decode(in []byte) {
	ln, m := cborItemLength(in)
	if msg.data == nil {
		msg.data = make([]byte, 0, options.payload)
	}
	msg.data = append(msg.data[:0], in[m:m+ln]...)
}

func (msg *msgReqsp) String() string {
	return "msgReqsp"
}

//-- stream message for benchmarking

type msgStream struct {
	data []byte
}

func (msg *msgStream) Id() uint64 {
	return 113
}

func (msg *msgStream) Encode(out []byte) int {
	return valbytes2cbor(msg.data, out)
}

func (msg *msgStream) Decode(in []byte) {
	ln, m := cborItemLength(in)
	if msg.data == nil {
		msg.data = make([]byte, 0, options.payload)
	}
	copy(msg.data[:0], in[m:m+ln])
}

func (msg *msgStream) String() string {
	return "msgStream"
}
<commit_msg>tools/client.go: to batch doPost() and doRequest().<commit_after>package main

import "sync"
import "sync/atomic"
import "runtime"
import "flag"
import "reflect"
import "unsafe"
import "os"
import "log"
import "sort"
import "time"
import "strings"
import "strconv"
import "net"
import "fmt"
import "compress/flate"
import "runtime/pprof"

import "github.com/prataprc/gofast"

var options struct {
	do       string
	count    int
	routines int
	conns    int
	addr     string
	payload  int
	log      string
}

func argParse() {
	flag.StringVar(&options.do, "do", "post",
		"post / request / stream benchmark to do")
	flag.IntVar(&options.conns, "conns", 1,
		"number of connections to use")
	flag.IntVar(&options.routines, "routines", 1,
		"number of concurrent routines per connection")
	flag.IntVar(&options.count, "count", 1,
		"number of requests per routine")
	flag.StringVar(&options.addr, "addr", "127.0.0.1:9998",
		"number of concurrent routines")
	flag.StringVar(&options.log, "log", "error",
		"number of concurrent routines")
	flag.IntVar(&options.payload, "payload", 10,
		"payload size to ping pong.")
	flag.Parse()
}

var av = &Average{}

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

	// start cpu profile.
	fname := "client.pprof"
	fd, err := os.Create(fname)
	if err != nil {
		log.Fatalf("unable to create %q: %v\n", fname, err)
	}
	defer fd.Close()
	pprof.StartCPUProfile(fd)
	defer pprof.StopCPUProfile()

	switch options.do {
	case "post":
		doTransport(doPost)
	case "request":
		doTransport(doRequest)
	case "stream":
		doTransport(doStream)
	}

	// take memory profile.
	fname = "client.mprof"
	fd, err = os.Create(fname)
	if err != nil {
		log.Fatal(err)
	}
	defer fd.Close()
	pprof.WriteHeapProfile(fd)
}

func doTransport(callb func(trans *gofast.Transport)) {
	var wg sync.WaitGroup

	ver := testVersion(1)
	n_trans := make([]*gofast.Transport, 0, 16)
	for i := 0; i < options.conns; i++ {
		wg.Add(1)
		config := newconfig("client", 3000, 4000)
		config["tags"] = ""
		conn, err := net.Dial("tcp", options.addr)
		if err != nil {
			panic(err)
		}
		trans, err := gofast.NewTransport(conn, &ver, nil, config)
		if err != nil {
			panic(err)
		}
		n_trans = append(n_trans, trans)
		go func(trans *gofast.Transport) {
			trans.FlushPeriod(100 * time.Millisecond)
			callb(trans)
			wg.Done()
			trans.Close()
		}(trans)
	}
	wg.Wait()
	printCounts(addCounts(n_trans...))
	fmsg := "request stats: n:%v mean:%v var:%v sd:%v\n"
	n, m := av.Count(), time.Duration(av.Mean())
	v, s := time.Duration(av.Variance()), time.Duration(av.Sd())
	fmt.Printf(fmsg, n, m, v, s)
}

func doPost(trans *gofast.Transport) {
	var wg sync.WaitGroup

	trans.SubscribeMessage(&msgPost{}, nil).Handshake()
	msg := &msgPost{data: make([]byte, options.payload)}
	for i := 0; i < options.payload; i++ {
		msg.data[i] = 'a'
	}

	for i := 0; i < options.routines; i++ {
		wg.Add(1)
		go func() {
			for j := 0; j < options.count; j++ {
				since := time.Now()
				if err := trans.Post(msg, false); err != nil {
					fmt.Printf("%v\n", err)
					panic("exit")
				}
				av.Add(uint64(time.Since(since)))
			}
			wg.Done()
		}()
	}
	wg.Wait()
	if _, err := trans.Whoami(); err != nil {
		log.Fatal(err)
	}
}

func doRequest(trans *gofast.Transport) {
	var wg sync.WaitGroup

	trans.SubscribeMessage(&msgReqsp{}, nil).Handshake()
	msg := &msgReqsp{data: make([]byte, options.payload)}
	for i := 0; i < options.payload; i++ {
		msg.data[i] = 'a'
	}

	for i := 0; i < options.routines; i++ {
		wg.Add(1)
		go func() {
			for j := 0; j < options.count; j++ {
				since := time.Now()
				if rmsg, err := trans.Request(msg, false); err != nil {
					fmt.Printf("%v\n", err)
					panic("exit")
				} else {
					trans.Free(rmsg)
				}
				av.Add(uint64(time.Since(since)))
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func doStream(trans *gofast.Transport) {
	var wg sync.WaitGroup

	trans.SubscribeMessage(&msgStream{}, nil).Handshake()

	for i := 0; i < options.routines; i++ {
		wg.Add(1)
		go func() {
			msg := &msgStream{data: make([]byte, 0, options.payload*2)}
			ch := make(chan gofast.Message, 100)
			since := time.Now()
			stream, err := trans.Stream(msg, true, ch)
			if err != nil {
				log.Fatal(err)
			}
			//fmt.Println("started", msg)
			var received int64
			go func() {
				for {
					msg, ok := <-ch
					if ok {
						trans.Free(msg)
						atomic.AddInt64(&received, 1)
						if atomic.LoadInt64(&received) == int64(options.count+1) {
							stream.Close()
							break
						}
						//fmt.Println("received", msg)
					} else {
						break
					}
				}
				wg.Done()
			}()
			av.Add(uint64(time.Since(since)))
			for j := 0; j < options.count; j++ {
				msg.data = msg.data[:cap(msg.data)]
				for i := 0; i < options.payload; i++ {
					msg.data[i] = 'a'
				}
				n := options.payload
				tmp := strconv.AppendInt(msg.data[n:n], int64(j), 10)
				msg.data = msg.data[:n+len(tmp)]

				since := time.Now()
				if err := stream.Stream(msg, true); err != nil {
					fmt.Printf("%v\n", err)
					panic("exit")
				}
				//fmt.Println("count", msg)
				av.Add(uint64(time.Since(since)))
			}
		}()
	}
	wg.Wait()
}

func newconfig(name string, start, end int) map[string]interface{} {
	return map[string]interface{}{
		"name":         name,
		"buffersize":   1024,
		"chansize":     100000,
		"batchsize":    100,
		"tags":         "",
		"opaque.start": start,
		"opaque.end":   end,
		"log.level":    options.log,
		"gzip.file":    flate.BestSpeed,
	}
}

type testVersion int

func (v *testVersion) Less(ver gofast.Version) bool {
	return (*v) < (*ver.(*testVersion))
}

func (v *testVersion) Equal(ver gofast.Version) bool {
	return (*v) == (*ver.(*testVersion))
}

func (v *testVersion) String() string {
	return fmt.Sprintf("%v", int(*v))
}

func (v *testVersion) Marshal(out []byte) int {
	return valuint642cbor(uint64(*v), out)
}

func (v *testVersion) Unmarshal(in []byte) int {
	ln, n := cborItemLength(in)
	*v = testVersion(ln)
	return n
}

func printCounts(counts map[string]uint64) {
	keys := []string{}
	for key := range counts {
		keys = append(keys, key)
	}
	sort.Sort(sort.StringSlice(keys))
	s := []string{}
	for _, key := range keys {
		s = append(s, fmt.Sprintf("%v:%v", key, counts[key]))
	}
	fmt.Println(strings.Join(s, ", "))
}

func addCounts(n_trans ...*gofast.Transport) map[string]uint64 {
	counts := n_trans[0].Counts()
	for _, trans := range n_trans[1:] {
		for k, v := range trans.Counts() {
			counts[k] += v
		}
	}
	return counts
}

func bytes2str(bytes []byte) string {
	if bytes == nil {
		return ""
	}
	sl := (*reflect.SliceHeader)(unsafe.Pointer(&bytes))
	st := &reflect.StringHeader{Data: sl.Data, Len: sl.Len}
	return *(*string)(unsafe.Pointer(st))
}

//-- post message for benchmarking

type msgPost struct {
	data []byte
}

func (msg *msgPost) Id() uint64 {
	return 111
}

func (msg *msgPost) Encode(out []byte) int {
	return valbytes2cbor(msg.data, out)
}

func (msg *msgPost) Decode(in []byte) {
	ln, m := cborItemLength(in)
	if msg.data == nil {
		msg.data = make([]byte, 0, options.payload)
	}
	msg.data = append(msg.data[:0], in[m:m+ln]...)
}

func (msg *msgPost) String() string {
	return "msgPost"
}

//-- reqsp message for benchmarking

type msgReqsp struct {
	data []byte
}

func (msg *msgReqsp) Id() uint64 {
	return 112
}

func (msg *msgReqsp) Encode(out []byte) int {
	return valbytes2cbor(msg.data, out)
}

func (msg *msgReqsp) Decode(in []byte) {
	ln, m := cborItemLength(in)
	if msg.data == nil {
		msg.data = make([]byte, 0, options.payload)
	}
	msg.data = append(msg.data[:0], in[m:m+ln]...)
}

func (msg *msgReqsp) String() string {
	return "msgReqsp"
}

//-- stream message for benchmarking

type msgStream struct {
	data []byte
}

func (msg *msgStream) Id() uint64 {
	return 113
}

func (msg *msgStream) Encode(out []byte) int {
	return valbytes2cbor(msg.data, out)
}

func (msg *msgStream) Decode(in []byte) {
	ln, m := cborItemLength(in)
	if msg.data == nil {
		msg.data = make([]byte, 0, options.payload)
	}
	copy(msg.data[:0], in[m:m+ln])
}

func (msg *msgStream) String() string {
	return "msgStream"
}
<|endoftext|>