<commit_before>/*
 Copyright 2016 Padduck, LLC

 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 docker

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
	v1 "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pufferpanel/pufferpanel/v2"
	"github.com/pufferpanel/pufferpanel/v2/logging"
	"github.com/pufferpanel/pufferpanel/v2/messages"
	"io"
	"os"
	"path"
	"runtime"
	"syscall"
	"time"
)

type docker struct {
	*pufferpanel.BaseEnvironment
	ContainerId  string              `json:"-"`
	ImageName    string              `json:"image"`
	Binds        map[string]string   `json:"bindings,omitempty"`
	NetworkMode  string              `json:"networkMode,omitempty"`
	Network      string              `json:"networkName,omitempty"`
	Ports        []string            `json:"portBindings,omitempty"`
	Resources    container.Resources `json:"resources,omitempty"`
	ExposedPorts nat.PortSet         `json:"exposedPorts,omitempty"`

	connection       types.HijackedResponse
	cli              *client.Client
	downloadingImage bool
}

func (d *docker) dockerExecuteAsync(steps pufferpanel.ExecutionData) error {
	running, err := d.IsRunning()
	if err != nil {
		return err
	}
	if running {
		return pufferpanel.ErrContainerRunning
	}

	d.Wait.Wait()

	if d.downloadingImage {
		return pufferpanel.ErrImageDownloading
	}

	dockerClient, err := d.getClient()
	ctx := context.Background()

	//TODO: This logic may not work anymore, it's complicated to use an existing container with install/uninstall
	exists, err := d.doesContainerExist(dockerClient, ctx)

	if err != nil {
		return err
	}

	if exists {
		return errors.New("docker container already exists")
	}

	err = d.createContainer(dockerClient, ctx, steps.Command, steps.Arguments, steps.Environment, steps.WorkingDirectory)
	if err != nil {
		return err
	}

	config := types.ContainerAttachOptions{
		Stdin:  true,
		Stdout: true,
		Stderr: true,
		Stream: true,
	}

	d.connection, err = dockerClient.ContainerAttach(ctx, d.ContainerId, config)
	if err != nil {
		return err
	}

	d.Wait.Add(1)

	go func() {
		defer d.connection.Close()
		wrapper := d.CreateWrapper()
		_, _ = io.Copy(wrapper, d.connection.Reader)
		//because we use the auto-delete, we don't manually stop the container
		//c, _ := d.getClient()
		//err = c.ContainerStop(context.Background(), d.ContainerId, nil)

		// The container hasn't actually been removed yet, we're sending an explicit remove to wait for it finish
		_ = dockerClient.ContainerRemove(ctx, d.ContainerId, types.ContainerRemoveOptions{})
		d.Wait.Done()
		if err != nil {
			logging.Error.Printf("Error stopping container "+d.ContainerId, err)
		}

		msg := messages.Status{Running: false}
		_ = d.WSManager.WriteMessage(msg)

		if steps.Callback != nil {
			steps.Callback(err == nil)
		}
	}()

	startOpts := types.ContainerStartOptions{}

	msg := messages.Status{Running: true}
	_ = d.WSManager.WriteMessage(msg)

	d.DisplayToConsole(true, "Starting container\n")
	err = dockerClient.ContainerStart(ctx, d.ContainerId, startOpts)
	if err != nil {
		return err
	}
	return err
}

func (d *docker) ExecuteInMainProcess(cmd string) (err error) {
	running, err := d.IsRunning()
	if err != nil {
		return
	}
	if !running {
		err = pufferpanel.ErrServerOffline
		return
	}

	_, _ = d.connection.Conn.Write([]byte(cmd + "\n"))
	return
}

func (d *docker) Kill() (err error) {
	running, err := d.IsRunning()
	if err != nil {
		return err
	}

	if !running {
		return
	}

	dockerClient, err := d.getClient()
	if err != nil {
		return err
	}
	err = dockerClient.ContainerKill(context.Background(), d.ContainerId, "SIGKILL")
	return
}

func (d *docker) Create() error {
	go func() {
		c, err := d.getClient()
		if err != nil {
			logging.Error.Printf("Error getting docker client: %s\n", err.Error())
			d.DisplayToConsole(true, "Error downloading image")
			return
		}
		err = d.pullImage(c, context.Background(), false)
		if err != nil {
			logging.Error.Printf("Error downloading image: %s\n", err.Error())
			d.DisplayToConsole(true, "Error downloading image")
			return
		}
	}()
	return os.Mkdir(d.RootDirectory, 0755)
}

func (d *docker) IsRunning() (bool, error) {
	dockerClient, err := d.getClient()
	if err != nil {
		return false, err
	}

	ctx := context.Background()

	exists, err := d.doesContainerExist(dockerClient, ctx)
	if !exists {
		return false, err
	}

	stats, err := dockerClient.ContainerInspect(ctx, d.ContainerId)
	if err != nil {
		return false, err
	}
	return stats.State.Running, nil
}

func (d *docker) GetStats() (*pufferpanel.ServerStats, error) {
	running, err := d.IsRunning()
	if err != nil {
		return nil, err
	}

	if !running {
		return nil, pufferpanel.ErrServerOffline
	}

	dockerClient, err := d.getClient()

	if err != nil {
		return nil, err
	}

	ctx := context.Background()
	res, err := dockerClient.ContainerStats(ctx, d.ContainerId, false)
	defer func() {
		if res.Body != nil {
			pufferpanel.Close(res.Body)
		}
	}()
	if err != nil {
		return nil, err
	}

	data := &types.StatsJSON{}
	err = json.NewDecoder(res.Body).Decode(&data)
	if err != nil {
		return nil, err
	}

	return &pufferpanel.ServerStats{
		Memory: calculateMemoryPercent(data),
		Cpu:    calculateCPUPercent(data),
	}, nil
}

func (d *docker) WaitForMainProcess() error {
	return d.WaitForMainProcessFor(0)
}

func (d *docker) WaitForMainProcessFor(timeout int) (err error) {
	running, err := d.IsRunning()
	if err != nil {
		return
	}
	if running {
		if timeout > 0 {
			var timer = time.AfterFunc(time.Duration(timeout)*time.Millisecond, func() {
				err = d.Kill()
			})
			d.Wait.Wait()
			timer.Stop()
		} else {
			d.Wait.Wait()
		}
	}
	return
}

func (d *docker) getClient() (*client.Client, error) {
	var err error = nil
	if d.cli == nil {
		d.cli, err = client.NewClientWithOpts(client.FromEnv)
		ctx := context.Background()
		d.cli.NegotiateAPIVersion(ctx)
	}
	return d.cli, err
}

func (d *docker) doesContainerExist(client *client.Client, ctx context.Context) (bool, error) {
	opts := types.ContainerListOptions{
		Filters: filters.NewArgs(),
	}

	opts.All = true
	opts.Filters.Add("name", d.ContainerId)

	existingContainers, err := client.ContainerList(ctx, opts)

	if len(existingContainers) == 0 {
		return false, err
	} else {
		return true, err
	}
}

func (d *docker) pullImage(client *client.Client, ctx context.Context, force bool) error {
	if d.downloadingImage {
		return pufferpanel.ErrImageDownloading
	}

	exists := false

	opts := types.ImageListOptions{
		All:     true,
		Filters: filters.NewArgs(),
	}
	opts.Filters.Add("reference", d.ImageName)
	images, err := client.ImageList(ctx, opts)

	if err != nil {
		return err
	}

	if len(images) >= 1 {
		exists = true
	}

	logging.Debug.Printf("Does image %v exist? %v", d.ImageName, exists)

	if exists && !force {
		return nil
	}

	op := types.ImagePullOptions{}

	logging.Debug.Printf("Downloading image %v", d.ImageName)
	d.DisplayToConsole(true, "Downloading image for container, please wait\n")

	d.downloadingImage = true
	defer func() {
		d.downloadingImage = false
	}()

	r, err := client.ImagePull(ctx, d.ImageName, op)
	defer pufferpanel.Close(r)
	if err != nil {
		return err
	}

	w := &ImageWriter{Parent: d.WSManager}
	_, err = io.Copy(w, r)

	if err != nil {
		return err
	}

	logging.Debug.Printf("Downloaded image %v", d.ImageName)
	d.DisplayToConsole(true, "Downloaded image for container\n")
	return err
}

func (d *docker) createContainer(client *client.Client, ctx context.Context, cmd string, args []string, env map[string]string, workDir string) error {
	logging.Debug.Printf("Creating container")
	containerRoot := "/pufferpanel"
	err := d.pullImage(client, ctx, false)

	if err != nil {
		return err
	}

	cmdSlice := strslice.StrSlice{}

	cmdSlice = append(cmdSlice, cmd)

	for _, v := range args {
		cmdSlice = append(cmdSlice, v)
	}

	//newEnv := os.Environ()
	newEnv := []string{"HOME=" + containerRoot}

	for k, v := range env {
		newEnv = append(newEnv, fmt.Sprintf("%s=%s", k, v))
	}

	if workDir == "" {
		workDir = containerRoot
	}

	logging.Debug.Printf("Container command: %s\n", cmdSlice)

	containerConfig := &container.Config{
		AttachStderr:    true,
		AttachStdin:     true,
		AttachStdout:    true,
		Tty:             true,
		OpenStdin:       true,
		NetworkDisabled: false,
		Cmd:             cmdSlice,
		Image:           d.ImageName,
		WorkingDir:      workDir,
		Env:             newEnv,
		ExposedPorts:    d.ExposedPorts,
	}

	if runtime.GOOS == "linux" {
		containerConfig.User = fmt.Sprintf("%d:%d", os.Getuid(), os.Getgid())
	}

	dir := d.RootDirectory

	//convert root dir to a full path, so we can bind it
	if !path.IsAbs(dir) {
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}
		dir = path.Join(pwd, dir)
	}

	hostConfig := &container.HostConfig{
		AutoRemove:   true,
		NetworkMode:  container.NetworkMode(d.NetworkMode),
		Resources:    d.Resources,
		Binds:        []string{dir + ":" + containerRoot},
		PortBindings: nat.PortMap{},
	}

	for k, v := range d.Binds {
		hostConfig.Binds = append(hostConfig.Binds, k+":"+v)
	}

	networkConfig := &network.NetworkingConfig{}

	_, bindings, err := nat.ParsePortSpecs(d.Ports)
	if err != nil {
		return err
	}
	hostConfig.PortBindings = bindings

	//for now, default to linux across the board. This resolves problems that Windows has when you use it and docker
	_, err = client.ContainerCreate(ctx, containerConfig, hostConfig, networkConfig, &v1.Platform{OS: "linux"}, d.ContainerId)
	return err
}

func (d *docker) SendCode(code int) error {
	running, err := d.IsRunning()

	if err != nil || !running {
		return err
	}

	dockerClient, err := d.getClient()

	if err != nil {
		return err
	}

	ctx := context.Background()
	return dockerClient.ContainerKill(ctx, d.ContainerId, syscall.Signal(code).String())
}

func calculateCPUPercent(v *types.StatsJSON) float64 {
	// Max number of 100ns intervals between the previous time read and now
	possIntervals := uint64(v.Read.Sub(v.PreRead).Nanoseconds()) // Start with number of ns intervals
	possIntervals /= 100                                         // Convert to number of 100ns intervals
	//possIntervals *= uint64(v.NumProcs)                          // Multiple by the number of processors

	// Intervals used
	intervalsUsed := v.CPUStats.CPUUsage.TotalUsage - v.PreCPUStats.CPUUsage.TotalUsage

	// Percentage avoiding divide-by-zero
	if possIntervals > 0 {
		return float64(intervalsUsed) / float64(possIntervals)
	}
	return 0.00
}

func calculateMemoryPercent(v *types.StatsJSON) float64 {
	return float64(v.MemoryStats.Usage)
}
<commit_msg>Fix exposedPorts being required, we will now use portBindings to resolve this<commit_after>/*
 Copyright 2016 Padduck, LLC

 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 docker

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
	v1 "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pufferpanel/pufferpanel/v2"
	"github.com/pufferpanel/pufferpanel/v2/logging"
	"github.com/pufferpanel/pufferpanel/v2/messages"
	"io"
	"os"
	"path"
	"runtime"
	"syscall"
	"time"
)

type docker struct {
	*pufferpanel.BaseEnvironment
	ContainerId string              `json:"-"`
	ImageName   string              `json:"image"`
	Binds       map[string]string   `json:"bindings,omitempty"`
	NetworkMode string              `json:"networkMode,omitempty"`
	Network     string              `json:"networkName,omitempty"`
	Ports       []string            `json:"portBindings,omitempty"`
	Resources   container.Resources `json:"resources,omitempty"`

	connection       types.HijackedResponse
	cli              *client.Client
	downloadingImage bool
}

func (d *docker) dockerExecuteAsync(steps pufferpanel.ExecutionData) error {
	running, err := d.IsRunning()
	if err != nil {
		return err
	}
	if running {
		return pufferpanel.ErrContainerRunning
	}

	d.Wait.Wait()

	if d.downloadingImage {
		return pufferpanel.ErrImageDownloading
	}

	dockerClient, err := d.getClient()
	ctx := context.Background()

	//TODO: This logic may not work anymore, it's complicated to use an existing container with install/uninstall
	exists, err := d.doesContainerExist(dockerClient, ctx)

	if err != nil {
		return err
	}

	if exists {
		return errors.New("docker container already exists")
	}

	err = d.createContainer(dockerClient, ctx, steps.Command, steps.Arguments, steps.Environment, steps.WorkingDirectory)
	if err != nil {
		return err
	}

	config := types.ContainerAttachOptions{
		Stdin:  true,
		Stdout: true,
		Stderr: true,
		Stream: true,
	}

	d.connection, err = dockerClient.ContainerAttach(ctx, d.ContainerId, config)
	if err != nil {
		return err
	}

	d.Wait.Add(1)

	go func() {
		defer d.connection.Close()
		wrapper := d.CreateWrapper()
		_, _ = io.Copy(wrapper, d.connection.Reader)
		//because we use the auto-delete, we don't manually stop the container
		//c, _ := d.getClient()
		//err = c.ContainerStop(context.Background(), d.ContainerId, nil)

		// The container hasn't actually been removed yet, we're sending an explicit remove to wait for it finish
		_ = dockerClient.ContainerRemove(ctx, d.ContainerId, types.ContainerRemoveOptions{})
		d.Wait.Done()
		if err != nil {
			logging.Error.Printf("Error stopping container "+d.ContainerId, err)
		}

		msg := messages.Status{Running: false}
		_ = d.WSManager.WriteMessage(msg)

		if steps.Callback != nil {
			steps.Callback(err == nil)
		}
	}()

	startOpts := types.ContainerStartOptions{}

	msg := messages.Status{Running: true}
	_ = d.WSManager.WriteMessage(msg)

	d.DisplayToConsole(true, "Starting container\n")
	err = dockerClient.ContainerStart(ctx, d.ContainerId, startOpts)
	if err != nil {
		return err
	}
	return err
}

func (d *docker) ExecuteInMainProcess(cmd string) (err error) {
	running, err := d.IsRunning()
	if err != nil {
		return
	}
	if !running {
		err = pufferpanel.ErrServerOffline
		return
	}

	_, _ = d.connection.Conn.Write([]byte(cmd + "\n"))
	return
}

func (d *docker) Kill() (err error) {
	running, err := d.IsRunning()
	if err != nil {
		return err
	}

	if !running {
		return
	}

	dockerClient, err := d.getClient()
	if err != nil {
		return err
	}
	err = dockerClient.ContainerKill(context.Background(), d.ContainerId, "SIGKILL")
	return
}

func (d *docker) Create() error {
	go func() {
		c, err := d.getClient()
		if err != nil {
			logging.Error.Printf("Error getting docker client: %s\n", err.Error())
			d.DisplayToConsole(true, "Error downloading image")
			return
		}
		err = d.pullImage(c, context.Background(), false)
		if err != nil {
			logging.Error.Printf("Error downloading image: %s\n", err.Error())
			d.DisplayToConsole(true, "Error downloading image")
			return
		}
	}()
	return os.Mkdir(d.RootDirectory, 0755)
}

func (d *docker) IsRunning() (bool, error) {
	dockerClient, err := d.getClient()
	if err != nil {
		return false, err
	}

	ctx := context.Background()

	exists, err := d.doesContainerExist(dockerClient, ctx)
	if !exists {
		return false, err
	}

	stats, err := dockerClient.ContainerInspect(ctx, d.ContainerId)
	if err != nil {
		return false, err
	}
	return stats.State.Running, nil
}

func (d *docker) GetStats() (*pufferpanel.ServerStats, error) {
	running, err := d.IsRunning()
	if err != nil {
		return nil, err
	}

	if !running {
		return nil, pufferpanel.ErrServerOffline
	}

	dockerClient, err := d.getClient()

	if err != nil {
		return nil, err
	}

	ctx := context.Background()
	res, err := dockerClient.ContainerStats(ctx, d.ContainerId, false)
	defer func() {
		if res.Body != nil {
			pufferpanel.Close(res.Body)
		}
	}()
	if err != nil {
		return nil, err
	}

	data := &types.StatsJSON{}
	err = json.NewDecoder(res.Body).Decode(&data)
	if err != nil {
		return nil, err
	}

	return &pufferpanel.ServerStats{
		Memory: calculateMemoryPercent(data),
		Cpu:    calculateCPUPercent(data),
	}, nil
}

func (d *docker) WaitForMainProcess() error {
	return d.WaitForMainProcessFor(0)
}

func (d *docker) WaitForMainProcessFor(timeout int) (err error) {
	running, err := d.IsRunning()
	if err != nil {
		return
	}
	if running {
		if timeout > 0 {
			var timer = time.AfterFunc(time.Duration(timeout)*time.Millisecond, func() {
				err = d.Kill()
			})
			d.Wait.Wait()
			timer.Stop()
		} else {
			d.Wait.Wait()
		}
	}
	return
}

func (d *docker) getClient() (*client.Client, error) {
	var err error = nil
	if d.cli == nil {
		d.cli, err = client.NewClientWithOpts(client.FromEnv)
		ctx := context.Background()
		d.cli.NegotiateAPIVersion(ctx)
	}
	return d.cli, err
}

func (d *docker) doesContainerExist(client *client.Client, ctx context.Context) (bool, error) {
	opts := types.ContainerListOptions{
		Filters: filters.NewArgs(),
	}

	opts.All = true
	opts.Filters.Add("name", d.ContainerId)

	existingContainers, err := client.ContainerList(ctx, opts)

	if len(existingContainers) == 0 {
		return false, err
	} else {
		return true, err
	}
}

func (d *docker) pullImage(client *client.Client, ctx context.Context, force bool) error {
	if d.downloadingImage {
		return pufferpanel.ErrImageDownloading
	}

	exists := false

	opts := types.ImageListOptions{
		All:     true,
		Filters: filters.NewArgs(),
	}
	opts.Filters.Add("reference", d.ImageName)
	images, err := client.ImageList(ctx, opts)

	if err != nil {
		return err
	}

	if len(images) >= 1 {
		exists = true
	}

	logging.Debug.Printf("Does image %v exist? %v", d.ImageName, exists)

	if exists && !force {
		return nil
	}

	op := types.ImagePullOptions{}

	logging.Debug.Printf("Downloading image %v", d.ImageName)
	d.DisplayToConsole(true, "Downloading image for container, please wait\n")

	d.downloadingImage = true
	defer func() {
		d.downloadingImage = false
	}()

	r, err := client.ImagePull(ctx, d.ImageName, op)
	defer pufferpanel.Close(r)
	if err != nil {
		return err
	}

	w := &ImageWriter{Parent: d.WSManager}
	_, err = io.Copy(w, r)

	if err != nil {
		return err
	}

	logging.Debug.Printf("Downloaded image %v", d.ImageName)
	d.DisplayToConsole(true, "Downloaded image for container\n")
	return err
}

func (d *docker) createContainer(client *client.Client, ctx context.Context, cmd string, args []string, env map[string]string, workDir string) error {
	logging.Debug.Printf("Creating container")
	containerRoot := "/pufferpanel"
	err := d.pullImage(client, ctx, false)

	if err != nil {
		return err
	}

	cmdSlice := strslice.StrSlice{}

	cmdSlice = append(cmdSlice, cmd)

	for _, v := range args {
		cmdSlice = append(cmdSlice, v)
	}

	//newEnv := os.Environ()
	newEnv := []string{"HOME=" + containerRoot}

	for k, v := range env {
		newEnv = append(newEnv, fmt.Sprintf("%s=%s", k, v))
	}

	if workDir == "" {
		workDir = containerRoot
	}

	logging.Debug.Printf("Container command: %s\n", cmdSlice)

	containerConfig := &container.Config{
		AttachStderr:    true,
		AttachStdin:     true,
		AttachStdout:    true,
		Tty:             true,
		OpenStdin:       true,
		NetworkDisabled: false,
		Cmd:             cmdSlice,
		Image:           d.ImageName,
		WorkingDir:      workDir,
		Env:             newEnv,
	}

	if runtime.GOOS == "linux" {
		containerConfig.User = fmt.Sprintf("%d:%d", os.Getuid(), os.Getgid())
	}

	dir := d.RootDirectory

	//convert root dir to a full path, so we can bind it
	if !path.IsAbs(dir) {
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}
		dir = path.Join(pwd, dir)
	}

	hostConfig := &container.HostConfig{
		AutoRemove:   true,
		NetworkMode:  container.NetworkMode(d.NetworkMode),
		Resources:    d.Resources,
		Binds:        []string{dir + ":" + containerRoot},
		PortBindings: nat.PortMap{},
	}

	for k, v := range d.Binds {
		hostConfig.Binds = append(hostConfig.Binds, k+":"+v)
	}

	networkConfig := &network.NetworkingConfig{}

	_, bindings, err := nat.ParsePortSpecs(d.Ports)
	if err != nil {
		return err
	}
	hostConfig.PortBindings = bindings

	exposedPorts := make(nat.PortSet)
	for k, _ := range bindings {
		exposedPorts[k] = struct{}{}
	}
	containerConfig.ExposedPorts = exposedPorts

	//for now, default to linux across the board. This resolves problems that Windows has when you use it and docker
	_, err = client.ContainerCreate(ctx, containerConfig, hostConfig, networkConfig, &v1.Platform{OS: "linux"}, d.ContainerId)
	return err
}

func (d *docker) SendCode(code int) error {
	running, err := d.IsRunning()

	if err != nil || !running {
		return err
	}

	dockerClient, err := d.getClient()

	if err != nil {
		return err
	}

	ctx := context.Background()
	return dockerClient.ContainerKill(ctx, d.ContainerId, syscall.Signal(code).String())
}

func calculateCPUPercent(v *types.StatsJSON) float64 {
	// Max number of 100ns intervals between the previous time read and now
	possIntervals := uint64(v.Read.Sub(v.PreRead).Nanoseconds()) // Start with number of ns intervals
	possIntervals /= 100                                         // Convert to number of 100ns intervals
	//possIntervals *= uint64(v.NumProcs)                          // Multiple by the number of processors

	// Intervals used
	intervalsUsed := v.CPUStats.CPUUsage.TotalUsage - v.PreCPUStats.CPUUsage.TotalUsage

	// Percentage avoiding divide-by-zero
	if possIntervals > 0 {
		return float64(intervalsUsed) / float64(possIntervals)
	}
	return 0.00
}

func calculateMemoryPercent(v *types.StatsJSON) float64 {
	return float64(v.MemoryStats.Usage)
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"fmt"
	"time"
	"bufio"
	"strings"
	"path/filepath"
	"text/template"
)

func main() {
	fmt.Println("Starting...")
	generate(findProjects())
}

type User struct {
	Name string
	Projects []Project
}

type Project struct {
	Name string
	Path string
}

func findProjects() map[string]User {
	files, _ := filepath.Glob("/home/*/Code/*")
	users := make(map[string]User)

	for _, path := range files {
		pparts := strings.Split(path, "/")
		uname := pparts[2]
		proj := &Project{Name: filepath.Base(path), Path: strings.Replace(path, "/home/", "~", -1)}
		u, exists := users[uname]
		if !exists {
			fmt.Printf("Found Code for ~%s.\n", uname)
			projs := []Project{*proj}
			u = User{Name: uname, Projects: projs}
		} else {
			u.Projects = append(u.Projects, *proj)
		}
		users[uname] = u
	}
	return users
}

type Page struct {
	Host string
	Users map[string]User
	Updated string
	UpdatedForHumans string
}

func generate(users map[string]User) {
	fmt.Println("Generating page.")

	f, err := os.Create(os.Getenv("HOME") + "/public_html/code.html")
	if err != nil {
		panic(err)
	}

	defer f.Close()

	w := bufio.NewWriter(f)
	template, err := template.ParseFiles("templates/code.html")
	if err != nil {
		panic(err)
	}

	// Extra page data
	host, _ := os.Hostname()
	curTime := time.Now()
	updatedReadable := curTime.UTC().Format(time.RFC1123)
	updated := curTime.UTC().Format(time.RFC3339)

	// Generate the page
	page := &Page{Host: host, UpdatedForHumans: updatedReadable, Updated: updated, Users: users}
	template.ExecuteTemplate(w, "code", page)
	w.Flush()

	fmt.Println("DONE!")
}
<commit_msg>Exclude hidden and .log files<commit_after>package main

import (
	"os"
	"fmt"
	"time"
	"bufio"
	"strings"
	"path/filepath"
	"text/template"
)

func main() {
	fmt.Println("Starting...")
	generate(findProjects())
}

type User struct {
	Name string
	Projects []Project
}

type Project struct {
	Name string
	Path string
}

func findProjects() map[string]User {
	files, _ := filepath.Glob("/home/*/Code/*")
	users := make(map[string]User)

	for _, path := range files {
		pparts := strings.Split(path, "/")
		uname := pparts[2]
		fname := filepath.Base(path)

		// Exclude certain files
		if strings.HasPrefix(fname, ".") || strings.HasSuffix(fname, ".log") {
			continue
		}

		proj := &Project{Name: fname, Path: strings.Replace(path, "/home/", "~", -1)}
		u, exists := users[uname]
		if !exists {
			fmt.Printf("Found Code for ~%s.\n", uname)
			projs := []Project{*proj}
			u = User{Name: uname, Projects: projs}
		} else {
			u.Projects = append(u.Projects, *proj)
		}
		users[uname] = u
	}
	return users
}

type Page struct {
	Host string
	Users map[string]User
	Updated string
	UpdatedForHumans string
}

func generate(users map[string]User) {
	fmt.Println("Generating page.")

	f, err := os.Create(os.Getenv("HOME") + "/public_html/code.html")
	if err != nil {
		panic(err)
	}

	defer f.Close()

	w := bufio.NewWriter(f)
	template, err := template.ParseFiles("templates/code.html")
	if err != nil {
		panic(err)
	}

	// Extra page data
	host, _ := os.Hostname()
	curTime := time.Now()
	updatedReadable := curTime.UTC().Format(time.RFC1123)
	updated := curTime.UTC().Format(time.RFC3339)

	// Generate the page
	page := &Page{Host: host, UpdatedForHumans: updatedReadable, Updated: updated, Users: users}
	template.ExecuteTemplate(w, "code", page)
	w.Flush()

	fmt.Println("DONE!")
}
<|endoftext|>


<commit_before>package builder

import (
	"errors"
	"strings"
)

type Cond interface {
	ToSQL() (string, []interface{}, error)
	And(...Cond) Cond
	Or(...Cond) Cond
}

type condAnd []Cond

func And(conds ...Cond) Cond {
	return condAnd(conds)
}

func (and condAnd) ToSQL() (string, []interface{}, error) {
	var sqls = make([]string, 0, len(and))
	var args = make([]interface{}, 0, 2*len(and))

	for _, cond := range and {
		sql, args1, err := cond.ToSQL()
		if err != nil {
			return "", nil, err
		}

		switch cond.(type) {
		case condOr:
			sqls = append(sqls, "("+sql+")")
		default:
			sqls = append(sqls, sql)
		}

		args = append(args, args1...)
	}

	return strings.Join(sqls, " AND "), args, nil
}

func (and condAnd) And(conds ...Cond) Cond {
	return condAnd(append(and, conds...))
}

func (and condAnd) Or(conds ...Cond) Cond {
	return append(condOr{and}, conds...)
}

type condOr []Cond

func Or(conds ...Cond) Cond {
	return condOr(conds)
}

func (or condOr) ToSQL() (string, []interface{}, error) {
	var sqls = make([]string, 0, len(or))
	var args = make([]interface{}, 0, 2*len(or))

	for _, cond := range or {
		sql, args1, err := cond.ToSQL()
		if err != nil {
			return "", nil, err
		}

		switch cond.(type) {
		case condAnd:
			sqls = append(sqls, "("+sql+")")
		default:
			sqls = append(sqls, sql)
		}

		args = append(args, args1...)
	}

	return strings.Join(sqls, " OR "), args, nil
}

func (o condOr) And(conds ...Cond) Cond {
	return append(condAnd{o}, conds...)
}

func (o condOr) Or(conds ...Cond) Cond {
	return append(o, conds...)
}

type Expr string

func (expr Expr) ToSQL() (string, []interface{}, error) {
	return string(expr), nil, nil
}

func (expr Expr) And(conds ...Cond) Cond {
	return append(condAnd{expr}, conds...)
}

func (expr Expr) Or(conds ...Cond) Cond {
	return append(condOr{expr}, conds...)
}

type Eq map[string]interface{}

func (eq Eq) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(eq))
	for k, v := range eq {
		conds = append(conds, k+"=?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

func (eq Eq) And(conds ...Cond) Cond {
	return append(condAnd{eq}, conds...)
}

func (eq Eq) Or(conds ...Cond) Cond {
	return append(condOr{eq}, conds...)
}

type Neq map[string]interface{}

func (neq Neq) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(neq))
	for k, v := range neq {
		conds = append(conds, k+"<>?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

type Lt map[string]interface{}

func (lt Lt) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(lt))
	for k, v := range lt {
		conds = append(conds, k+"<?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

type Gt map[string]interface{}

func (gt Gt) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(gt))
	for k, v := range gt {
		conds = append(conds, k+">?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

type condIn struct {
	col  string
	vals []interface{}
}

func In(col string, values ...interface{}) Cond {
	return condIn{col, values}
}

func (condIn condIn) ToSQL() (string, []interface{}, error) {
	if len(condIn.vals) <= 0 {
		return "", nil, errors.New("No in conditions")
	}

	questionMark := strings.Repeat("?,", len(condIn.vals))
	return condIn.col + "IN (" + questionMark[:len(questionMark)-1] + ")", condIn.vals, nil
}

func (condIn condIn) And(conds ...Cond) Cond {
	return append(condAnd{condIn}, conds...)
}

func (condIn condIn) Or(conds ...Cond) Cond {
	return append(condOr{condIn}, conds...)
}

type condNotIn condIn

func NotIn(col string, values ...interface{}) Cond {
	return condNotIn{col, values}
}

func (condNotIn condNotIn) ToSQL() (string, []interface{}, error) {
	if len(condNotIn.vals) <= 0 {
		return "", nil, errors.New("No in conditions")
	}

	questionMark := strings.Repeat("?,", len(condNotIn.vals))
	return condNotIn.col + "NOT IN (" + questionMark[:len(questionMark)-1] + ")", condNotIn.vals, nil
}

func (condNotIn condNotIn) And(conds ...Cond) Cond {
	return append(condAnd{condNotIn}, conds...)
}

func (condNotIn condNotIn) Or(conds ...Cond) Cond {
	return append(condOr{condNotIn}, conds...)
}

type Like [2]string

func (like Like) ToSQL() (string, []interface{}, error) {
	return like[0] + " LIKE ?", []interface{}{"%" + like[1] + "%"}, nil
}

func (like Like) And(conds ...Cond) Cond {
	return append(condAnd{like}, conds...)
}

func (like Like) Or(conds ...Cond) Cond {
	return append(condOr{like}, conds...)
}

type IsNull [1]string

func (isNull IsNull) ToSQL() (string, []interface{}, error) {
	return string(isNull[0]) + " IS NULL", []interface{}{}, nil
}

func (isNull IsNull) And(conds ...Cond) Cond {
	return append(condAnd{isNull}, conds...)
}

func (isNull IsNull) Or(conds ...Cond) Cond {
	return append(condOr{isNull}, conds...)
}

type NotNull [1]string

func (notNull NotNull) ToSQL() (string, []interface{}, error) {
	return notNull[0] + " IS NOT NULL", []interface{}{}, nil
}

func (notNull NotNull) And(conds ...Cond) Cond {
	return append(condAnd{notNull}, conds...)
}

func (notNull NotNull) Or(conds ...Cond) Cond {
	return append(condOr{notNull}, conds...)
}

type Not [1]Cond

func (not Not) ToSQL() (string, []interface{}, error) {
	sql, args, err := not[0].ToSQL()
	if err != nil {
		return "", nil, err
	}
	switch not[0].(type) {
	case condAnd, condOr:
		return "NOT (" + sql + ")", args, nil
	}
	return " NOT " + sql, args, nil
}

func (not Not) And(conds ...Cond) Cond {
	return append(condAnd{not}, conds...)
}

func (not Not) Or(conds ...Cond) Cond {
	return append(condOr{not}, conds...)
}
<commit_msg>added more conds<commit_after>package builder

import (
	"errors"
	"strings"
)

type Cond interface {
	ToSQL() (string, []interface{}, error)
	And(...Cond) Cond
	Or(...Cond) Cond
}

type condAnd []Cond

func And(conds ...Cond) Cond {
	return condAnd(conds)
}

func (and condAnd) ToSQL() (string, []interface{}, error) {
	var sqls = make([]string, 0, len(and))
	var args = make([]interface{}, 0, 2*len(and))

	for _, cond := range and {
		sql, args1, err := cond.ToSQL()
		if err != nil {
			return "", nil, err
		}

		switch cond.(type) {
		case condOr:
			sqls = append(sqls, "("+sql+")")
		default:
			sqls = append(sqls, sql)
		}

		args = append(args, args1...)
	}

	return strings.Join(sqls, " AND "), args, nil
}

func (and condAnd) And(conds ...Cond) Cond {
	return condAnd(append(and, conds...))
}

func (and condAnd) Or(conds ...Cond) Cond {
	return append(condOr{and}, conds...)
}

type condOr []Cond

func Or(conds ...Cond) Cond {
	return condOr(conds)
}

func (or condOr) ToSQL() (string, []interface{}, error) {
	var sqls = make([]string, 0, len(or))
	var args = make([]interface{}, 0, 2*len(or))

	for _, cond := range or {
		sql, args1, err := cond.ToSQL()
		if err != nil {
			return "", nil, err
		}

		switch cond.(type) {
		case condAnd:
			sqls = append(sqls, "("+sql+")")
		default:
			sqls = append(sqls, sql)
		}

		args = append(args, args1...)
	}

	return strings.Join(sqls, " OR "), args, nil
}

func (o condOr) And(conds ...Cond) Cond {
	return append(condAnd{o}, conds...)
}

func (o condOr) Or(conds ...Cond) Cond {
	return append(o, conds...)
}

type expr struct {
	sql  string
	args []interface{}
}

func Expr(sql string, args ...interface{}) Cond {
	return expr{sql, args}
}

func (expr expr) ToSQL() (string, []interface{}, error) {
	return expr.sql, expr.args, nil
}

func (expr expr) And(conds ...Cond) Cond {
	return append(condAnd{expr}, conds...)
}

func (expr expr) Or(conds ...Cond) Cond {
	return append(condOr{expr}, conds...)
}

type Eq map[string]interface{}

func (eq Eq) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(eq))
	for k, v := range eq {
		conds = append(conds, k+"=?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

func (eq Eq) And(conds ...Cond) Cond {
	return append(condAnd{eq}, conds...)
}

func (eq Eq) Or(conds ...Cond) Cond {
	return append(condOr{eq}, conds...)
}

type Neq map[string]interface{}

func (neq Neq) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(neq))
	for k, v := range neq {
		conds = append(conds, k+"<>?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

func (neq Neq) And(conds ...Cond) Cond {
	return append(condAnd{neq}, conds...)
}

func (neq Neq) Or(conds ...Cond) Cond {
	return append(condOr{neq}, conds...)
}

type Lt map[string]interface{}

func (lt Lt) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(lt))
	for k, v := range lt {
		conds = append(conds, k+"<?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

func (lt Lt) And(conds ...Cond) Cond {
	return append(condAnd{lt}, conds...)
}

func (lt Lt) Or(conds ...Cond) Cond {
	return append(condOr{lt}, conds...)
}

type Lte map[string]interface{}

func (lte Lte) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(lte))
	for k, v := range lte {
		conds = append(conds, k+"<=?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

func (lte Lte) And(conds ...Cond) Cond {
	return append(condAnd{lte}, conds...)
}

func (lte Lte) Or(conds ...Cond) Cond {
	return append(condOr{lte}, conds...)
}

type Gt map[string]interface{}

func (gt Gt) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(gt))
	for k, v := range gt {
		conds = append(conds, k+">?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

func (gt Gt) And(conds ...Cond) Cond {
	return append(condAnd{gt}, conds...)
}

func (gt Gt) Or(conds ...Cond) Cond {
	return append(condOr{gt}, conds...)
}

type Gte map[string]interface{}

func (gte Gte) ToSQL() (string, []interface{}, error) {
	var conds []string
	var args = make([]interface{}, 0, len(gte))
	for k, v := range gte {
		conds = append(conds, k+">=?")
		args = append(args, v)
	}
	return strings.Join(conds, " AND "), args, nil
}

func (gte Gte) And(conds ...Cond) Cond {
	return append(condAnd{gte}, conds...)
}

func (gte Gte) Or(conds ...Cond) Cond {
	return append(condOr{gte}, conds...)
}

type Between struct {
	col     string
	lessVal interface{}
	moreVal interface{}
}

func (between Between) ToSQL() (string, []interface{}, error) {
	return between.col + " BETWEEN ? AND ?", []interface{}{between.lessVal, between.moreVal}, nil
}

func (between Between) And(conds ...Cond) Cond {
	return append(condAnd{between}, conds...)
}

func (between Between) Or(conds ...Cond) Cond {
	return append(condOr{between}, conds...)
}

type condIn struct {
	col  string
	vals []interface{}
}

func In(col string, values ...interface{}) Cond {
	return condIn{col, values}
}

func (condIn condIn) ToSQL() (string, []interface{}, error) {
	if len(condIn.vals) <= 0 {
		return "", nil, errors.New("No in conditions")
	}

	questionMark := strings.Repeat("?,", len(condIn.vals))
	return condIn.col + "IN (" + questionMark[:len(questionMark)-1] + ")", condIn.vals, nil
}

func (condIn condIn) And(conds ...Cond) Cond {
	return append(condAnd{condIn}, conds...)
}

func (condIn condIn) Or(conds ...Cond) Cond {
	return append(condOr{condIn}, conds...)
}

type condNotIn condIn

func NotIn(col string, values ...interface{}) Cond {
	return condNotIn{col, values}
}

func (condNotIn condNotIn) ToSQL() (string, []interface{}, error) {
	if len(condNotIn.vals) <= 0 {
		return "", nil, errors.New("No in conditions")
	}

	questionMark := strings.Repeat("?,", len(condNotIn.vals))
	return condNotIn.col + "NOT IN (" + questionMark[:len(questionMark)-1] + ")", condNotIn.vals, nil
}

func (condNotIn condNotIn) And(conds ...Cond) Cond {
	return append(condAnd{condNotIn}, conds...)
}

func (condNotIn condNotIn) Or(conds ...Cond) Cond {
	return append(condOr{condNotIn}, conds...)
}

type Like [2]string

func (like Like) ToSQL() (string, []interface{}, error) {
	return like[0] + " LIKE ?", []interface{}{"%" + like[1] + "%"}, nil
}

func (like Like) And(conds ...Cond) Cond {
	return append(condAnd{like}, conds...)
}

func (like Like) Or(conds ...Cond) Cond {
	return append(condOr{like}, conds...)
}

type IsNull [1]string

func (isNull IsNull) ToSQL() (string, []interface{}, error) {
	return string(isNull[0]) + " IS NULL", []interface{}{}, nil
}

func (isNull IsNull) And(conds ...Cond) Cond {
	return append(condAnd{isNull}, conds...)
}

func (isNull IsNull) Or(conds ...Cond) Cond {
	return append(condOr{isNull}, conds...)
}

type NotNull [1]string

func (notNull NotNull) ToSQL() (string, []interface{}, error) {
	return notNull[0] + " IS NOT NULL", []interface{}{}, nil
}

func (notNull NotNull) And(conds ...Cond) Cond {
	return append(condAnd{notNull}, conds...)
}

func (notNull NotNull) Or(conds ...Cond) Cond {
	return append(condOr{notNull}, conds...)
}

type Not [1]Cond

func (not Not) ToSQL() (string, []interface{}, error) {
	sql, args, err := not[0].ToSQL()
	if err != nil {
		return "", nil, err
	}
	switch not[0].(type) {
	case condAnd, condOr:
		return "NOT (" + sql + ")", args, nil
	}
	return " NOT " + sql, args, nil
}

func (not Not) And(conds ...Cond) Cond {
	return append(condAnd{not}, conds...)
}

func (not Not) Or(conds ...Cond) Cond {
	return append(condOr{not}, conds...)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
)

type LogRedirects struct {
	Transport http.RoundTripper
}

func main() {
	// get URL from cli arguments
	url := os.Args[1]

	// new http client with custom round tripper
	client := &http.Client{Transport: LogRedirects{}}

	// prepare request
	req, err := http.NewRequest(
		"GET",
		url,
		nil,
	)

	// … with custom user agent
	req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_3) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/36.0.1985.125 Safari/537.36")
	resp, err := client.Do(req)

	if err != nil {
		log.Fatalln(err)
	}

	fmt.Printf("[%d] %s\n", resp.StatusCode, resp.Request.URL)
}

// RoundTrip only receives a request and returns the response without modifying it
// Basically we’re doing a regular request and if the code status code is 3xx then we print it out
func (l LogRedirects) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	t := l.Transport

	if t == nil {
		t = http.DefaultTransport
	}

	resp, err = t.RoundTrip(req)
	if err != nil {
		return
	}

	switch resp.StatusCode {
	case http.StatusMovedPermanently,
		http.StatusFound,
		http.StatusSeeOther,
		http.StatusTemporaryRedirect:

		// other way of displaying it
		// newLoc, _ := resp.Location()
		// fmt.Printf("[%d] %s\n", resp.StatusCode, newLoc.String())

		fmt.Printf("[%d] %s\n", resp.StatusCode, req.URL)
	}

	return
}
<commit_msg>refactoring<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
)

type redirectLogger struct{}

func main() {
	// get URL from cli arguments
	url := os.Args[1]

	// new http client with custom round tripper which
	// itself uses the default transport policy
	client := &http.Client{
		Transport:     redirectLogger{},
		CheckRedirect: redirectCounter}

	// create request
	req, _ := http.NewRequest("GET", url, nil)

	// … with custom user agent
	req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_3) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/36.0.1985.125 Safari/537.36")

	resp, err := client.Do(req)
	if err != nil {
		log.Fatalln(err)
	}

	printResponse(resp)
}

// RoundTrip only receives a request, executes it and returns the response
// without modifications. If the response code is a redirect of some form
// it gets send down the responses channel to be logged
func (l redirectLogger) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	t := http.DefaultTransport

	resp, err = t.RoundTrip(req)
	if err != nil {
		return
	}

	// If the response code is 3xx we print it out
	switch resp.StatusCode {
	case http.StatusMovedPermanently,
		http.StatusFound,
		http.StatusSeeOther,
		http.StatusTemporaryRedirect:

		printResponse(resp)
	}
	return
}

// Increases the max redirect count to 20
var redirectCounter = func(req *http.Request, via []*http.Request) error {
	if len(via) >= 20 {
		return errors.New("Stopped after 20 redirects")
	}
	return nil
}

func printResponse(resp *http.Response) {
	fmt.Printf("[%d] %s\n", resp.StatusCode, resp.Request.URL)
}
<|endoftext|>


<commit_before>package tasks

import (
	"fmt"

	"github.com/DVI-GI-2017/Jira__backend/db"
	"github.com/DVI-GI-2017/Jira__backend/models"
	"gopkg.in/mgo.v2/bson"
)

// Returns all labels from given task.
func AllLabels(source db.DataSource, task_id bson.ObjectId) (models.LabelsList, error) {
	var container = struct {
		LabelsList models.LabelsList `bson:"labels"`
	}{}

	err := queryLabels(source.C(cTasks), task_id).One(&container)

	if err != nil {
		return models.LabelsList{},
			fmt.Errorf("can not retrieve all labels on task %s: %v", task_id.Hex(), err)
	}

	return container.LabelsList, nil
}

// Selects labels from tasks query
func queryLabels(collection db.Collection, task_id bson.ObjectId) db.Query {
	return collection.Find(
		bson.M{
			"_id": task_id,
		}).Select(bson.M{"labels": 1})
}

// Checks if label already set on this task.
func CheckLabelAlreadySet(source db.DataSource, id bson.ObjectId, label models.Label) (bool, error) {
	notset, err := queryLabel(source.C(cTasks), id, label).IsEmpty()
	if err != nil {
		return false, err
	}

	return !notset, nil
}

// Selects label from collection.
func queryLabel(collection db.Collection, task_id bson.ObjectId, label models.Label) db.Query {
	return collection.Find(
		bson.M{
			"_id":    task_id,
			"labels": label,
		}).Select(bson.M{"labels": 1})
}

// Adds label to task and returns new list of labels on this task.
func AddLabelToTask(source db.DataSource, task_id bson.ObjectId, label models.Label) (models.LabelsList, error) {
	task, err := UpdateTask(source, task_id, pushLabel(label))
	if err != nil {
		return models.LabelsList{},
			fmt.Errorf("can not add label '%v' to task '%s': %v", label, task.Id.Hex(), err)
	}

	return task.Labels, nil
}

// Returns "updater" that performs push in task labels array.
func pushLabel(label models.Label) interface{} {
	return bson.M{"$push": bson.M{"labels": label}}
}

// Deletes label from task and returns new list of labels on this task
func DeleteLabelFromTask(source db.DataSource, task_id bson.ObjectId, label models.Label) (models.LabelsList, error) {
	task, err := UpdateTask(source, task_id, bson.M{"$pull": bson.M{"labels": label}})
	if err != nil {
		return models.LabelsList{},
			fmt.Errorf("can not delete label '%v' from task '%s': %v", label, task.Id.Hex(), err)
	}

	return task.Labels, nil
}
<commit_msg>Extract helper functions and container.<commit_after>package tasks

import (
	"fmt"

	"github.com/DVI-GI-2017/Jira__backend/db"
	"github.com/DVI-GI-2017/Jira__backend/models"
	"gopkg.in/mgo.v2/bson"
)

type LabelsContainer struct {
	LabelsList models.LabelsList `bson:"labels"`
}

// Returns all labels from given task.
func AllLabels(source db.DataSource, taskId bson.ObjectId) (models.LabelsList, error) {
	var container LabelsContainer

	err := queryLabels(source.C(cTasks), taskId).One(&container)
	if err != nil {
		return models.LabelsList{},
			fmt.Errorf("can not retrieve all labels on task %s: %v", taskId.Hex(), err)
	}

	return container.LabelsList, nil
}

// Selects labels from tasks query
func queryLabels(collection db.Collection, taskId bson.ObjectId) db.Query {
	return collection.Find(
		bson.M{
			"_id": taskId,
		}).Select(bson.M{"labels": 1})
}

// Checks if label already set on this task.
func CheckLabelAlreadySet(source db.DataSource, id bson.ObjectId, label models.Label) (bool, error) {
	notset, err := queryLabel(source.C(cTasks), id, label).IsEmpty()
	if err != nil {
		return false, err
	}

	return !notset, nil
}

// Selects label from collection.
func queryLabel(collection db.Collection, task_id bson.ObjectId, label models.Label) db.Query {
	return collection.Find(
		bson.M{
			"_id":    task_id,
			"labels": label,
		}).Select(bson.M{"labels": 1})
}

// Adds label to task and returns new list of labels on this task.
func AddLabelToTask(source db.DataSource, taskId bson.ObjectId, label models.Label) (models.LabelsList, error) {
	err := pushLabel(source.C(cTasks), taskId, label)
	if err != nil {
		return models.LabelsList{},
			fmt.Errorf("can not add label '%v' to task '%s': %v", label, taskId.Hex(), err)
	}

	return AllLabels(source, taskId)
}

// Pushes label in task's labels array.
func pushLabel(collection db.Collection, task_id bson.ObjectId, label models.Label) error {
	return collection.Update(
		bson.M{"_id": task_id},
		bson.M{"$push": bson.M{"labels": label}},
	)
}

// Deletes label from task and returns new list of labels on this task
func DeleteLabelFromTask(source db.DataSource, task_id bson.ObjectId, label models.Label) (models.LabelsList, error) {
	task, err := UpdateTask(source, task_id, bson.M{"$pull": bson.M{"labels": label}})
	if err != nil {
		return models.LabelsList{},
			fmt.Errorf("can not delete label '%v' from task '%s': %v", label, task.Id.Hex(), err)
	}

	return task.Labels, nil
}
<|endoftext|>


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

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	/*Log.Debugc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})*/
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} /*else {
				Log.Debugc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}*/
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	/*Log.Debugc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})*/
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err == ErrBufferNotNewData {
					Log.Debugc(func() string {
						return fmt.Sprintf("sender debug(%s)", err)
					})
					time.Sleep(10 * time.Millisecond)
					continue
				}
				Log.Errorc(func() string {
					return fmt.Sprintf("(%s)Error Sender  writing data: %v", this.cid(), err)
				})
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	//fmt.Println("(", this.cid(), ")", "this.in.buf=", this.in.buf)
	//fmt.Println("(", this.cid(), ")", "this.out.buf=", this.out.buf)
	//Log.Infoc(func() string {
	//	return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	//})
	for {

		if this.isDone() {
			return nil, 0, io.EOF
		}
		var (
			b   []byte
			err error
		)

		if this.in == nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
			})
			err = ErrBufferNotReady
			return nil, 0, err
		}

		// Peek cnt bytes from the input buffer.
		b, err = this.in.ReadWait()
		if err != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure:%v", this.cid(), err)
			})
			return nil, 0, err
		}
		if len(b) > 0 {
			//total := int(remlen) + 1 + m
			mtype := message.MessageType((b)[0] >> 4)
			//return mtype, total, err
			var msg message.Message
			msg, err = mtype.New()
			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure:%v", this.cid(), err)
				})
				return nil, 0, err
			}
			/*Log.Debugc(func() string {
				return fmt.Sprintf("(%s) 开始创建对象(%s)", this.cid(), msg.Name())
			})*/
			_, err = msg.Decode(b)
			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure:%v", this.cid(), err)
				})
				return nil, 0, err
			}
			/*Log.Debugc(func() string {
				return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
			})*/
			//fmt.Println("this.in.readwait=", msg.Len())
			return msg, len(b), err
		}
	}
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
/*func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}*/

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
/*
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	//l := 0
	//for l < total {
	n, err = this.in.Read(this.intmp[0:])
	//l += n
	Log.Debugc(func() string {
		return fmt.Sprintf("read %d bytes", n)
	})
	if err != nil {
		return nil, 0, err
	}
	//}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}
*/

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l     int = msg.Len()
		m, n  int
		err   error
		buf   []byte
		start int64
	//wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	//buf, wrap, err = this.out.WriteWait(l)//zheliyouwenti
	start, _, err = this.out.waitForWriteSpace(l)
	if err != nil {
		return 0, err
	}
	buf = make([]byte, l)
	n, err = msg.Encode(buf[0:])
	if err != nil {
		return 0, err
	}
	index := start & this.out.mask
	this.out.buf[index] = ByteArray{index: index, bArray: buf}
	m, err = this.out.WriteCommit(n)
	if err != nil {
		return 0, err
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<commit_msg>修改buffer读写限制<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	/*Log.Debugc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})*/
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} /*else {
				Log.Debugc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}*/
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	/*Log.Debugc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})*/
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err == ErrBufferNotNewData {
					Log.Debugc(func() string {
						return fmt.Sprintf("sender debug(%s)", err)
					})
					time.Sleep(10 * time.Millisecond)
					continue
				}
				Log.Errorc(func() string {
					return fmt.Sprintf("(%s)Error Sender  writing data: %v", this.cid(), err)
				})
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	//fmt.Println("(", this.cid(), ")", "this.in.buf=", this.in.buf)
	//fmt.Println("(", this.cid(), ")", "this.out.buf=", this.out.buf)
	//Log.Infoc(func() string {
	//	return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	//})
	for {

		if this.isDone() {
			return nil, 0, io.EOF
		}
		var (
			b   []byte
			err error
		)

		if this.in == nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
			})
			err = ErrBufferNotReady
			return nil, 0, err
		}

		// Peek cnt bytes from the input buffer.
		b, err = this.in.ReadWait()
		if err != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure:%v", this.cid(), err)
			})
			return nil, 0, err
		}
		if len(b) > 0 {
			//total := int(remlen) + 1 + m
			mtype := message.MessageType((b)[0] >> 4)
			//return mtype, total, err
			var msg message.Message
			msg, err = mtype.New()
			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure:%v", this.cid(), err)
				})
				return nil, 0, err
			}
			/*Log.Debugc(func() string {
				return fmt.Sprintf("(%s) 开始创建对象(%s)", this.cid(), msg.Name())
			})*/
			_, err = msg.Decode(b)
			if err != nil {
				Log.Errorc(func() string {
					return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure:%v", this.cid(), err)
				})
				return nil, 0, err
			}
			/*Log.Debugc(func() string {
				return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
			})*/
			//fmt.Println("this.in.readwait=", msg.Len())
			return msg, len(b), err
		} else {
			time.Sleep(2 * time.Microsecond)
		}
	}
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
/*func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}*/

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
/*
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	//l := 0
	//for l < total {
	n, err = this.in.Read(this.intmp[0:])
	//l += n
	Log.Debugc(func() string {
		return fmt.Sprintf("read %d bytes", n)
	})
	if err != nil {
		return nil, 0, err
	}
	//}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}
*/

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l     int = msg.Len()
		m, n  int
		err   error
		buf   []byte
		start int64
	//wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	//buf, wrap, err = this.out.WriteWait(l)//zheliyouwenti
	start, _, err = this.out.waitForWriteSpace(l)
	if err != nil {
		return 0, err
	}
	buf = make([]byte, l)
	n, err = msg.Encode(buf[0:])
	if err != nil {
		return 0, err
	}
	index := start & this.out.mask
	this.out.buf[index] = ByteArray{index: index, bArray: buf}
	m, err = this.out.WriteCommit(n)
	if err != nil {
		return 0, err
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<|endoftext|>


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

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   []byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}

	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType(b[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}

	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  []byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode(buf[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<commit_msg>修改了buffer.go;sendrecv.go;process.go<commit_after>// Copyright (c) 2014 The SurgeMQ Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"fmt"
	"io"
	"net"
	"time"

	"github.com/surgemq/message"
)

type netReader interface {
	io.Reader
	SetReadDeadline(t time.Time) error
}

type timeoutReader struct {
	d    time.Duration
	conn netReader
}

func (r timeoutReader) Read(b []byte) (int, error) {
	if err := r.conn.SetReadDeadline(time.Now().Add(r.d)); err != nil {
		return 0, err
	}
	return r.conn.Read(b)
}

// receiver() reads data from the network, and writes the data into the incoming buffer
func (this *service) receiver() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) receiver开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping receiver", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting receiver", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		//Log.Debugc(func() string{ return fmt.Sprintf("server/handleConnection: Setting read deadline to %d", time.Second*time.Duration(this.keepAlive))})
		keepAlive := time.Second * time.Duration(this.keepAlive)
		r := timeoutReader{
			d:    keepAlive + (keepAlive / 2),
			conn: conn,
		}

		for {
			_, err := this.in.ReadFrom(r)
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sess is: %v", this.sess)})
			//       Log.Errorc(func() string{ return fmt.Sprintf("this.sessMgr is: %v", this.sessMgr)})

			/*if err != nil {
				Log.Infoc(func() string { return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err) })
				//         if err != io.EOF {
				//         }
				return
			}*/
			if err != nil {

				if err == io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error reading from connection: %v", this.cid(), err)
					})
					return
				} else {
					Log.Infoc(func() string {
						return fmt.Sprintf("(%s) info reading from connection: %v", this.cid(), err)
					})
				}

				continue
			} else {
				Log.Infoc(func() string {
					return fmt.Sprintf("(%s)向ringbuffer些数据成功！", this.cid())
				})
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket: %v", this.cid(), ErrInvalidConnectionType)})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) %v", this.cid(), ErrInvalidConnectionType)
		})
	}
}

// sender() writes data from the outgoing buffer to the network
func (this *service) sender() {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) sender开始", this.cid())
	})
	defer func() {
		// Let's recover from panic
		if r := recover(); r != nil {
			Log.Errorc(func() string {
				return fmt.Sprintf("(%s) Recovering from panic: %v", this.cid(), r)
			})
		}

		this.wgStopped.Done()

		Log.Debugc(func() string {
			return fmt.Sprintf("(%s) Stopping sender", this.cid())
		})
	}()

	//   Log.Debugc(func() string{ return fmt.Sprintf("(%s) Starting sender", this.cid())})

	this.wgStarted.Done()

	switch conn := this.conn.(type) {
	case net.Conn:
		for {
			_, err := this.out.WriteTo(conn)

			if err != nil {
				if err != io.EOF {
					Log.Errorc(func() string {
						return fmt.Sprintf("(%s) error writing data: %v", this.cid(), err)
					})
				}
				return
			}
		}

	//case *websocket.Conn:
	//	Log.Errorc(func() string{ return fmt.Sprintf("(%s) Websocket not supported", this.cid())})

	default:
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) Invalid connection type", this.cid())
		})
	}
}

// peekMessageSize() reads, but not commits, enough bytes to determine the size of
// the next message and returns the type and size.
func (this *service) peekMessageSize() (message.Message, int, error) {
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize开始", this.cid())
	})
	var (
		b   []byte
		err error
		cnt int = 2
	)

	if this.in == nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in is nil", this.cid())
		})
		err = ErrBufferNotReady
		return nil, 0, err
	}

	// Let's read enough bytes to get the message header (msg type, remaining length)
	//for {
	// If we have read 5 bytes and still not done, then there's a problem.
	//if cnt > 5 {
	//	return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//}

	// Peek cnt bytes from the input buffer.
	b, err = this.in.ReadWait(cnt)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize this.in.ReadWait falure", this.cid())
		})
		return nil, 0, err
	}

	//// If not enough bytes are returned, then continue until there's enough.
	//if len(b) < cnt {
	//	continue
	//}

	// If we got enough bytes, then check the last byte to see if the continuation
	// bit is set. If so, increment cnt and continue peeking
	/*if b[cnt-1] >= 0x80 {
		cnt++
	} else {
		break
	}*/
	//for {
	//	if cnt > 5 {
	//		return 0, 0, fmt.Errorf("sendrecv/peekMessageSize: 4th byte of remaining length has continuation bit set")
	//	}
	//
	//	if b[cnt-1] >= 0x80 {
	//		cnt++
	//	} else {
	//		break
	//	}
	//}
	//}

	// Get the remaining length of the message
	//remlen, m := binary.Uvarint(b[1:cnt])

	// Total message length is remlen + 1 (msg type) + m (remlen bytes)
	//total := int(remlen) + 1 + m

	mtype := message.MessageType(b[0] >> 4)

	//return mtype, total, err
	var msg message.Message
	msg, err = mtype.New()
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize mtype.New() falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) 开始创建对象(%s)", this.cid(), msg.Name())
	})
	_, err = msg.Decode(b)
	if err != nil {
		Log.Errorc(func() string {
			return fmt.Sprintf("(%s) peekMessageSize msg.Decode falure", this.cid())
		})
		return nil, 0, err
	}
	Log.Infoc(func() string {
		return fmt.Sprintf("(%s) peekMessageSize结束(%s)", this.cid(), msg.Name())
	})
	return msg, len(b), err
}

// peekMessage() reads a message from the buffer, but the bytes are NOT committed.
// This means the buffer still thinks the bytes are not read yet.
func (this *service) peekMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b    []byte
		err  error
		i, n int
		msg  message.Message
	)

	if this.in == nil {
		return nil, 0, ErrBufferNotReady
	}

	// Peek until we get total bytes
	for i = 0; ; i++ {
		// Peek remlen bytes from the input buffer.
		b, err = this.in.ReadWait(total)
		if err != nil && err != ErrBufferInsufficientData {
			return nil, 0, err
		}

		// If not enough bytes are returned, then continue until there's enough.
		if len(b) >= total {
			break
		}
	}

	msg, err = mtype.New()
	if err != nil {
		return nil, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// readMessage() reads and copies a message from the buffer. The buffer bytes are
// committed as a result of the read.
func (this *service) readMessage(mtype message.MessageType, total int) (message.Message, int, error) {
	var (
		b   []byte
		err error
		n   int
		msg message.Message
	)

	if this.in == nil {
		err = ErrBufferNotReady
		return nil, 0, err
	}

	if len(this.intmp) < total {
		this.intmp = make([]byte, total)
	}

	// Read until we get total bytes
	l := 0
	for l < total {
		n, err = this.in.Read(this.intmp[l:])
		l += n
		Log.Debugc(func() string {
			return fmt.Sprintf("read %d bytes, total %d", n, l)
		})
		if err != nil {
			return nil, 0, err
		}
	}

	b = this.intmp[:total]

	msg, err = mtype.New()
	if err != nil {
		return msg, 0, err
	}

	n, err = msg.Decode(b)
	return msg, n, err
}

// writeMessage() writes a message to the outgoing buffer
func (this *service) writeMessage(msg message.Message) (int, error) {
	var (
		l    int = msg.Len()
		m, n int
		err  error
		buf  []byte
		wrap bool
	)

	if this.out == nil {
		return 0, ErrBufferNotReady
	}

	// This is to serialize writes to the underlying buffer. Multiple goroutines could
	// potentially get here because of calling Publish() or Subscribe() or other
	// functions that will send messages. For example, if a message is received in
	// another connetion, and the message needs to be published to this client, then
	// the Publish() function is called, and at the same time, another client could
	// do exactly the same thing.
	//
	// Not an ideal fix though. If possible we should remove mutex and be lockfree.
	// Mainly because when there's a large number of goroutines that want to publish
	// to this client, then they will all block. However, this will do for now.
	//
	// FIXME: Try to find a better way than a mutex...if possible.
	this.wmu.Lock()
	defer this.wmu.Unlock()

	buf, wrap, err = this.out.WriteWait(l)
	if err != nil {
		return 0, err
	}

	if wrap {
		if len(this.outtmp) < l {
			this.outtmp = make([]byte, l)
		}

		n, err = msg.Encode(this.outtmp[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.Write(this.outtmp[0:n])
		if err != nil {
			return m, err
		}
	} else {
		n, err = msg.Encode(buf[0:])
		if err != nil {
			return 0, err
		}

		m, err = this.out.WriteCommit(n)
		if err != nil {
			return 0, err
		}
	}

	this.outStat.increment(int64(m))

	return m, nil
}
<|endoftext|>


<commit_before>package newman

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

	"golang.org/x/net/context"
)

type ConnOption uint64

const (
	DefaultBufferSize            = 1024 * 100 // 100k
	Compression       ConnOption = 1 << iota
	Encrypted
)

// NoopCloser io.ReadWriteCloser has a noop for "Close"
type NoopCloser struct {
	io.ReadWriter
}

// Close noop
func (n *NoopCloser) Close() error {
	return nil
}

// WrapNoopCloser wraps a io.ReadWriter in a io.ReadWriteCloser that does nothing on "Close"
func WrapNoopCloser(rw io.ReadWriter) io.ReadWriteCloser {
	return &NoopCloser{rw}
}

// Conn provides a wrapper around an io.ReadWriteCloser for sending size delimited data over any reader/writer
type Conn struct {
	buff         []byte
	nextSizeBuff []byte
	rwc          io.ReadWriteCloser
	options      ConnOption
	ctx          context.Context
}

// NewConn returns a new Conn with the appropriate configuration
func NewConn(rwc io.ReadWriteCloser, options ...ConnOption) *Conn {
	c := &Conn{
		buff:         make([]byte, DefaultBufferSize),
		nextSizeBuff: make([]byte, 8),
		rwc:          rwc,
		ctx:          context.Background(),
	}

	// assign all options
	for _, o := range options {
		c.options |= o
	}

	return c
}

// Write a message to the connection
func (c *Conn) Write(m Message) error {
	// encode the message as binary
	buff, err := m.MarshalBinary()
	if err != nil {
		return err
	}

	// write out the message on the connection
	return c.writeNextBuffer(buff)
}

// Next reads the next message off the line and unmarshals it into the given message
func (c *Conn) Next(m Message) error {

	// fill the message buffer
	size, err := c.readNextBuffer()
	if err != nil {
		return err
	}

	// unmarshal the message
	return m.UnmarshalBinary(c.buff[:size])
}

// Generate returns a channel that will recieve messages as they come in off the line
// Only safe to run once per Conn
func (c *Conn) Generate(f func() Message) (<-chan Message, context.CancelFunc) {

	// output channel
	out := make(chan Message, 10)

	// set up the context to cancel the generation
	ctx, stop := context.WithCancel(c.ctx)
	c.ctx = ctx

	// read messages as they come off the line and send them on the output channel
	go func() {
		var err error
		var m Message

		// loop until something bad happens
		for {
			select {

			// check to see if the generation should be stopped
			case <-c.ctx.Done():
				close(out)
				return

			default:
				// make a new message
				m = f()

				// read the next message off the line
				err = c.Next(m)
				if err != nil {
					// close the out put channel and return
					close(out)
					return
				}

				out <- m
			}
		}
	}()

	return out, stop
}

// writeNextBuffer writes the next message legnth and message to the enclosed writer
func (c *Conn) writeNextBuffer(buff []byte) error {
	// write out the size of the next message
	binary.LittleEndian.PutUint64(c.nextSizeBuff, uint64(len(buff)))
	err := writeIntoBuffer(c.nextSizeBuff, c.rwc)
	if err != nil {
		return err
	}

	// write the next message
	return writeIntoBuffer(buff, c.rwc)
}

// readNextBuffer reads the next message into the message buffer
func (c *Conn) readNextBuffer() (int, error) {
	// read the size of the next message
	err := readIntoBuffer(c.nextSizeBuff, c.rwc)
	if err != nil {
		return 0, err
	}

	// find the size of the next message
	size := int(binary.LittleEndian.Uint64(c.nextSizeBuff))

	// make sure the buffer is large enough to hold the next message
	if size > len(c.buff) {

		// make a new buffer twice as large as we need to avoid reallocating
		c.buff = make([]byte, size*2)
	}

	// read the next message into the buffer
	return size, readIntoBuffer(c.buff[:size], c.rwc)
}

// writeIntoBuffer writes an entire message to the io.Writer
func writeIntoBuffer(buff []byte, w io.Writer) error {

	// this will be used to back off the connection if the connection's buffers are full
	backOffTime := time.Millisecond

	x := len(buff)
	var n, l int
	var err error
	// read until the buffer size is met
	for l < x {
		n, err = w.Write(buff[l:])
		if err != nil {
			return err
		}

		// back off if nothing could be written
		if n == 0 {
			time.Sleep(backOffTime)
			backOffTime = backOffTime * 2

		} else {
			backOffTime = time.Millisecond

			// add the written amount to the write counter
			l += n
		}

	}

	return nil

}

// readIntoBuffer will read off of a rwc until it the buffer has been filled
func readIntoBuffer(buff []byte, r io.Reader) error {

	// this will be used to back off the connection when there is nothing to read
	backOffTime := time.Millisecond

	x := len(buff)
	var n, l int
	var err error
	// read until the buffer size is met
	for l < x {
		n, err = r.Read(buff[l:])
		if err != nil {
			return err
		}

		// back off if we got nothing
		if n == 0 {
			time.Sleep(backOffTime)
			backOffTime = backOffTime * 2
		} else {
			// reset the backOffTime
			backOffTime = time.Millisecond

			// add the read amount to the length counter
			l += n
		}

	}

	return nil
}
<commit_msg>dont' need an option type<commit_after>package newman

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

	"golang.org/x/net/context"
)

const (
	DefaultBufferSize        = 1024 * 100 // 100k
	Compression       uint64 = 1 << iota
	Encrypted
)

// NoopCloser io.ReadWriteCloser has a noop for "Close"
type NoopCloser struct {
	io.ReadWriter
}

// Close noop
func (n *NoopCloser) Close() error {
	return nil
}

// WrapNoopCloser wraps a io.ReadWriter in a io.ReadWriteCloser that does nothing on "Close"
func WrapNoopCloser(rw io.ReadWriter) io.ReadWriteCloser {
	return &NoopCloser{rw}
}

// Conn provides a wrapper around an io.ReadWriteCloser for sending size delimited data over any reader/writer
type Conn struct {
	buff         []byte
	nextSizeBuff []byte
	rwc          io.ReadWriteCloser
	options      uint64
	ctx          context.Context
}

// NewConn returns a new Conn with the appropriate configuration
func NewConn(rwc io.ReadWriteCloser, options ...uint64) *Conn {
	c := &Conn{
		buff:         make([]byte, DefaultBufferSize),
		nextSizeBuff: make([]byte, 8),
		rwc:          rwc,
		ctx:          context.Background(),
	}

	// assign all options
	for _, o := range options {
		c.options |= o
	}

	return c
}

// Write a message to the connection
func (c *Conn) Write(m Message) error {
	// encode the message as binary
	buff, err := m.MarshalBinary()
	if err != nil {
		return err
	}

	// write out the message on the connection
	return c.writeNextBuffer(buff)
}

// Next reads the next message off the line and unmarshals it into the given message
func (c *Conn) Next(m Message) error {

	// fill the message buffer
	size, err := c.readNextBuffer()
	if err != nil {
		return err
	}

	// unmarshal the message
	return m.UnmarshalBinary(c.buff[:size])
}

// Generate returns a channel that will recieve messages as they come in off the line
// Only safe to run once per Conn
func (c *Conn) Generate(f func() Message) (<-chan Message, context.CancelFunc) {

	// output channel
	out := make(chan Message, 10)

	// set up the context to cancel the generation
	ctx, stop := context.WithCancel(c.ctx)
	c.ctx = ctx

	// read messages as they come off the line and send them on the output channel
	go func() {
		var err error
		var m Message

		// loop until something bad happens
		for {
			select {

			// check to see if the generation should be stopped
			case <-c.ctx.Done():
				close(out)
				return

			default:
				// make a new message
				m = f()

				// read the next message off the line
				err = c.Next(m)
				if err != nil {
					// close the out put channel and return
					close(out)
					return
				}

				out <- m
			}
		}
	}()

	return out, stop
}

// writeNextBuffer writes the next message legnth and message to the enclosed writer
func (c *Conn) writeNextBuffer(buff []byte) error {
	// write out the size of the next message
	binary.LittleEndian.PutUint64(c.nextSizeBuff, uint64(len(buff)))
	err := writeIntoBuffer(c.nextSizeBuff, c.rwc)
	if err != nil {
		return err
	}

	// write the next message
	return writeIntoBuffer(buff, c.rwc)
}

// readNextBuffer reads the next message into the message buffer
func (c *Conn) readNextBuffer() (int, error) {
	// read the size of the next message
	err := readIntoBuffer(c.nextSizeBuff, c.rwc)
	if err != nil {
		return 0, err
	}

	// find the size of the next message
	size := int(binary.LittleEndian.Uint64(c.nextSizeBuff))

	// make sure the buffer is large enough to hold the next message
	if size > len(c.buff) {

		// make a new buffer twice as large as we need to avoid reallocating
		c.buff = make([]byte, size*2)
	}

	// read the next message into the buffer
	return size, readIntoBuffer(c.buff[:size], c.rwc)
}

// writeIntoBuffer writes an entire message to the io.Writer
func writeIntoBuffer(buff []byte, w io.Writer) error {

	// this will be used to back off the connection if the connection's buffers are full
	backOffTime := time.Millisecond

	x := len(buff)
	var n, l int
	var err error
	// read until the buffer size is met
	for l < x {
		n, err = w.Write(buff[l:])
		if err != nil {
			return err
		}

		// back off if nothing could be written
		if n == 0 {
			time.Sleep(backOffTime)
			backOffTime = backOffTime * 2

		} else {
			backOffTime = time.Millisecond

			// add the written amount to the write counter
			l += n
		}

	}

	return nil

}

// readIntoBuffer will read off of a rwc until it the buffer has been filled
func readIntoBuffer(buff []byte, r io.Reader) error {

	// this will be used to back off the connection when there is nothing to read
	backOffTime := time.Millisecond

	x := len(buff)
	var n, l int
	var err error
	// read until the buffer size is met
	for l < x {
		n, err = r.Read(buff[l:])
		if err != nil {
			return err
		}

		// back off if we got nothing
		if n == 0 {
			time.Sleep(backOffTime)
			backOffTime = backOffTime * 2
		} else {
			// reset the backOffTime
			backOffTime = time.Millisecond

			// add the read amount to the length counter
			l += n
		}

	}

	return nil
}
<|endoftext|>