<commit_before>package nodepool

import (
	"fmt"

	v1alpha1 "github.com/jetstack/navigator/pkg/apis/navigator/v1alpha1"
	"github.com/jetstack/navigator/pkg/controllers/cassandra/util"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	appslisters "k8s.io/client-go/listers/apps/v1beta2"
	"k8s.io/client-go/tools/record"
)

type Interface interface {
	Sync(*v1alpha1.CassandraCluster) error
}

type defaultCassandraClusterNodepoolControl struct {
	kubeClient        kubernetes.Interface
	statefulsetLister appslisters.StatefulSetLister
	recorder          record.EventRecorder
}

var _ Interface = &defaultCassandraClusterNodepoolControl{}

func NewControl(
	kubeClient kubernetes.Interface,
	statefulsetLister appslisters.StatefulSetLister,
	recorder record.EventRecorder,
) Interface {
	return &defaultCassandraClusterNodepoolControl{
		kubeClient:        kubeClient,
		statefulsetLister: statefulsetLister,
		recorder:          recorder,
	}
}

func (e *defaultCassandraClusterNodepoolControl) removeUnusedStatefulSets(
	cluster *v1alpha1.CassandraCluster,
) error {
	expectedStatefulSetNames := map[string]bool{}
	for _, pool := range cluster.Spec.NodePools {
		name := util.NodePoolResourceName(cluster, &pool)
		expectedStatefulSetNames[name] = true
	}
	client := e.kubeClient.AppsV1beta2().StatefulSets(cluster.Namespace)
	selector, err := util.SelectorForCluster(cluster)
	if err != nil {
		return err
	}
	existingSets, err := e.statefulsetLister.
		StatefulSets(cluster.Namespace).
		List(selector)
	if err != nil {
		return err
	}
	for _, set := range existingSets {
		if !metav1.IsControlledBy(set, cluster) {
			ownerRef := metav1.GetControllerOf(set)
			return fmt.Errorf(
				"Foreign owned StatefulSet: "+
					"A StatefulSet with name '%s/%s' already exists, "+
					"but it is controlled by '%v', not '%s/%s'.",
				set.Namespace, set.Name, ownerRef,
				cluster.Namespace, cluster.Name,
			)
		}
		_, found := expectedStatefulSetNames[set.Name]
		if !found {
			err := client.Delete(set.Name, nil)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (e *defaultCassandraClusterNodepoolControl) createOrUpdateStatefulSet(
	cluster *v1alpha1.CassandraCluster,
	nodePool *v1alpha1.CassandraClusterNodePool,
) error {
	client := e.kubeClient.AppsV1beta2().StatefulSets(cluster.Namespace)
	desiredSet := StatefulSetForCluster(cluster, nodePool)
	existingSet, err := e.statefulsetLister.
		StatefulSets(desiredSet.Namespace).
		Get(desiredSet.Name)
	if k8sErrors.IsNotFound(err) {
		_, err = client.Create(desiredSet)
		return err
	}
	if err != nil {
		return err
	}
	if !metav1.IsControlledBy(existingSet, cluster) {
		ownerRef := metav1.GetControllerOf(existingSet)
		return fmt.Errorf(
			"Foreign owned StatefulSet: "+
				"A StatefulSet with name '%s/%s' already exists, "+
				"but it is controlled by '%v', not '%s/%s'.",
			existingSet.Namespace, existingSet.Name, ownerRef,
			cluster.Namespace, cluster.Name,
		)
	}
	_, err = client.Update(desiredSet)
	return err
}

func (e *defaultCassandraClusterNodepoolControl) syncStatefulSets(
	cluster *v1alpha1.CassandraCluster,
) error {
	for _, pool := range cluster.Spec.NodePools {
		err := e.createOrUpdateStatefulSet(cluster, &pool)
		if err != nil {
			return err
		}
	}
	err := e.removeUnusedStatefulSets(cluster)
	return err
}

func (e *defaultCassandraClusterNodepoolControl) Sync(cluster *v1alpha1.CassandraCluster) error {
	return e.syncStatefulSets(cluster)
}
<commit_msg>refactor ownerCheck into a function<commit_after>package nodepool

import (
	"fmt"

	v1alpha1 "github.com/jetstack/navigator/pkg/apis/navigator/v1alpha1"
	"github.com/jetstack/navigator/pkg/controllers/cassandra/util"
	"k8s.io/api/apps/v1beta2"
	k8sErrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	appslisters "k8s.io/client-go/listers/apps/v1beta2"
	"k8s.io/client-go/tools/record"
)

type Interface interface {
	Sync(*v1alpha1.CassandraCluster) error
}

type defaultCassandraClusterNodepoolControl struct {
	kubeClient        kubernetes.Interface
	statefulsetLister appslisters.StatefulSetLister
	recorder          record.EventRecorder
}

var _ Interface = &defaultCassandraClusterNodepoolControl{}

func NewControl(
	kubeClient kubernetes.Interface,
	statefulsetLister appslisters.StatefulSetLister,
	recorder record.EventRecorder,
) Interface {
	return &defaultCassandraClusterNodepoolControl{
		kubeClient:        kubeClient,
		statefulsetLister: statefulsetLister,
		recorder:          recorder,
	}
}

func ownerCheck(
	set *v1beta2.StatefulSet,
	cluster *v1alpha1.CassandraCluster,
) error {
	if !metav1.IsControlledBy(set, cluster) {
		ownerRef := metav1.GetControllerOf(set)
		return fmt.Errorf(
			"Foreign owned StatefulSet: "+
				"A StatefulSet with name '%s/%s' already exists, "+
				"but it is controlled by '%v', not '%s/%s'.",
			set.Namespace, set.Name, ownerRef,
			cluster.Namespace, cluster.Name,
		)
	}
	return nil
}

func (e *defaultCassandraClusterNodepoolControl) removeUnusedStatefulSets(
	cluster *v1alpha1.CassandraCluster,
) error {
	expectedStatefulSetNames := map[string]bool{}
	for _, pool := range cluster.Spec.NodePools {
		name := util.NodePoolResourceName(cluster, &pool)
		expectedStatefulSetNames[name] = true
	}
	client := e.kubeClient.AppsV1beta2().StatefulSets(cluster.Namespace)
	selector, err := util.SelectorForCluster(cluster)
	if err != nil {
		return err
	}
	existingSets, err := e.statefulsetLister.
		StatefulSets(cluster.Namespace).
		List(selector)
	if err != nil {
		return err
	}
	for _, set := range existingSets {
		err := ownerCheck(set, cluster)
		if err != nil {
			return err
		}
		_, found := expectedStatefulSetNames[set.Name]
		if !found {
			err := client.Delete(set.Name, nil)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (e *defaultCassandraClusterNodepoolControl) createOrUpdateStatefulSet(
	cluster *v1alpha1.CassandraCluster,
	nodePool *v1alpha1.CassandraClusterNodePool,
) error {
	client := e.kubeClient.AppsV1beta2().StatefulSets(cluster.Namespace)
	desiredSet := StatefulSetForCluster(cluster, nodePool)
	existingSet, err := e.statefulsetLister.
		StatefulSets(desiredSet.Namespace).
		Get(desiredSet.Name)
	if k8sErrors.IsNotFound(err) {
		_, err = client.Create(desiredSet)
		return err
	}
	if err != nil {
		return err
	}
	err = ownerCheck(existingSet, cluster)
	if err != nil {
		return err
	}
	_, err = client.Update(desiredSet)
	return err
}

func (e *defaultCassandraClusterNodepoolControl) syncStatefulSets(
	cluster *v1alpha1.CassandraCluster,
) error {
	for _, pool := range cluster.Spec.NodePools {
		err := e.createOrUpdateStatefulSet(cluster, &pool)
		if err != nil {
			return err
		}
	}
	err := e.removeUnusedStatefulSets(cluster)
	return err
}

func (e *defaultCassandraClusterNodepoolControl) Sync(cluster *v1alpha1.CassandraCluster) error {
	return e.syncStatefulSets(cluster)
}
<|endoftext|>


<commit_before>package vendor

import (
	"fmt"
	"reflect"
	"testing"
)

func TestDeduceRemoteRepo(t *testing.T) {
	if testing.Short() {
		t.Skipf("skipping network tests in -short mode")
	}
	tests := []struct {
		path     string
		want     RemoteRepo
		extra    string
		err      error
		insecure bool
	}{{
		path: "",
		err:  fmt.Errorf(`"" is not a valid import path`),
	}, {
		path: "corporate",
		err:  fmt.Errorf(`"corporate" is not a valid import path`),
	}, {
		path: "github.com/cznic/b",
		want: &gitrepo{
			url: "https://github.com/cznic/b",
		},
	}, {
		path: "github.com/pkg/sftp",
		want: &gitrepo{
			url: "https://github.com/pkg/sftp",
		},
	}, {
		path: "github.com/pkg/sftp/examples/gsftp",
		want: &gitrepo{
			url: "https://github.com/pkg/sftp",
		},
		extra: "/examples/gsftp",
	}, {
		path: "github.com/coreos/go-etcd",
		want: &gitrepo{
			url: "https://github.com/coreos/go-etcd",
		},
	}, {
		path: "bitbucket.org/davecheney/gitrepo/cmd/main",
		want: &gitrepo{
			url: "https://bitbucket.org/davecheney/gitrepo",
		},
		extra: "/cmd/main",
	}, {
		path: "bitbucket.org/davecheney/hgrepo/cmd/main",
		want: &hgrepo{
			url: "https://bitbucket.org/davecheney/hgrepo",
		},
		extra: "/cmd/main",
	}, {
		path: "code.google.com/p/goauth2/oauth",
		want: &hgrepo{
			url: "https://code.google.com/p/goauth2",
		},
		extra: "/oauth",
	}, {
		path: "code.google.com/p/gami",
		want: &gitrepo{
			url: "https://code.google.com/p/gami",
		},
	}, {
		path: "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git",
		want: &gitrepo{
			url: "https://git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git",
		},
	}, {
		path: "git.apache.org/thrift.git/lib/go/thrift",
		want: &gitrepo{
			url: "https://git.apache.org/thrift.git",
		},
		extra: "/lib/go/thrift",
	}, {
		path: "gopkg.in/check.v1",
		want: &gitrepo{
			url: "https://gopkg.in/check.v1",
		},
		extra: "",
	}, {
		path: "golang.org/x/tools/go/vcs",
		want: &gitrepo{
			url: "https://go.googlesource.com/tools",
		},
		extra: "/go/vcs",
	}, {
		path: "labix.org/v2/mgo",
		want: &bzrrepo{
			url: "https://launchpad.net/mgo/v2",
		},
		insecure: true,
	}, {
		path: "launchpad.net/gnuflag",
		want: &bzrrepo{
			url: "https://launchpad.net/gnuflag",
		},
	}, {
		path: "https://github.com/pkg/sftp",
		want: &gitrepo{
			url: "https://github.com/pkg/sftp",
		},
	}, {
		path: "git://github.com/pkg/sftp",
		want: &gitrepo{
			url: "git://github.com/pkg/sftp",
		},
		insecure: true,
	}, {
		path: "code.google.com/p/google-api-go-client/bigquery/v2",
		want: &hgrepo{
			url: "https://code.google.com/p/google-api-go-client",
		},
		extra: "/bigquery/v2",
	}}

	for _, tt := range tests {
		t.Logf("DeduceRemoteRepo(%q, %v)", tt.path, tt.insecure)
		got, extra, err := DeduceRemoteRepo(tt.path, tt.insecure)
		if !reflect.DeepEqual(err, tt.err) {
			t.Errorf("DeduceRemoteRepo(%q): want err: %v, got err: %v", tt.path, tt.err, err)
			continue
		}
		if !reflect.DeepEqual(got, tt.want) || extra != tt.extra {
			t.Errorf("DeduceRemoteRepo(%q): want %#v, %v, got %#v, %v", tt.path, tt.want, tt.extra, got, extra)
		}
	}
}
<commit_msg>Extra tests<commit_after>package vendor

import (
	"fmt"
	"reflect"
	"testing"
)

func TestDeduceRemoteRepo(t *testing.T) {
	if testing.Short() {
		t.Skipf("skipping network tests in -short mode")
	}
	tests := []struct {
		path     string
		want     RemoteRepo
		extra    string
		err      error
		insecure bool
	}{{
		path: "",
		err:  fmt.Errorf(`"" is not a valid import path`),
	}, {
		path: "corporate",
		err:  fmt.Errorf(`"corporate" is not a valid import path`),
	}, {
		path: "github.com/cznic/b",
		want: &gitrepo{
			url: "https://github.com/cznic/b",
		},
	}, {
		path: "github.com/pkg/sftp",
		want: &gitrepo{
			url: "https://github.com/pkg/sftp",
		},
	}, {
		path: "github.com/pkg/sftp/examples/gsftp",
		want: &gitrepo{
			url: "https://github.com/pkg/sftp",
		},
		extra: "/examples/gsftp",
	}, {
		path: "github.com/coreos/go-etcd",
		want: &gitrepo{
			url: "https://github.com/coreos/go-etcd",
		},
	}, {
		path: "bitbucket.org/davecheney/gitrepo/cmd/main",
		want: &gitrepo{
			url: "https://bitbucket.org/davecheney/gitrepo",
		},
		extra: "/cmd/main",
	}, {
		path: "bitbucket.org/davecheney/hgrepo/cmd/main",
		want: &hgrepo{
			url: "https://bitbucket.org/davecheney/hgrepo",
		},
		extra: "/cmd/main",
	}, {
		path: "code.google.com/p/goauth2/oauth",
		want: &hgrepo{
			url: "https://code.google.com/p/goauth2",
		},
		extra: "/oauth",
	}, {
		path: "code.google.com/p/gami",
		want: &gitrepo{
			url: "https://code.google.com/p/gami",
		},
	}, {
		path: "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git",
		want: &gitrepo{
			url: "https://git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git",
		},
	}, {
		path: "git.apache.org/thrift.git/lib/go/thrift",
		want: &gitrepo{
			url: "https://git.apache.org/thrift.git",
		},
		extra: "/lib/go/thrift",
	}, {
		path: "gopkg.in/check.v1",
		want: &gitrepo{
			url: "https://gopkg.in/check.v1",
		},
		extra: "",
	}, {
		path: "golang.org/x/tools/go/vcs",
		want: &gitrepo{
			url: "https://go.googlesource.com/tools",
		},
		extra: "/go/vcs",
	}, {
		path: "labix.org/v2/mgo",
		want: &bzrrepo{
			url: "https://launchpad.net/mgo/v2",
		},
		insecure: true,
	}, {
		path: "launchpad.net/gnuflag",
		want: &bzrrepo{
			url: "https://launchpad.net/gnuflag",
		},
	}, {
		path: "https://github.com/pkg/sftp",
		want: &gitrepo{
			url: "https://github.com/pkg/sftp",
		},
	}, {
		path: "git://github.com/pkg/sftp",
		want: &gitrepo{
			url: "git://github.com/pkg/sftp",
		},
		insecure: true,
	}, {
		path: "code.google.com/p/google-api-go-client/bigquery/v2",
		want: &hgrepo{
			url: "https://code.google.com/p/google-api-go-client",
		},
		extra: "/bigquery/v2",
	}, {
		path: "code.google.com/p/go-sqlite/go1/sqlite3",
		want: &hgrepo{
			url: "https://code.google.com/p/go-sqlite",
		},
		extra: "/go1/sqlite3",
	}}

	for _, tt := range tests {
		t.Logf("DeduceRemoteRepo(%q, %v)", tt.path, tt.insecure)
		got, extra, err := DeduceRemoteRepo(tt.path, tt.insecure)
		if !reflect.DeepEqual(err, tt.err) {
			t.Errorf("DeduceRemoteRepo(%q): want err: %v, got err: %v", tt.path, tt.err, err)
			continue
		}
		if !reflect.DeepEqual(got, tt.want) || extra != tt.extra {
			t.Errorf("DeduceRemoteRepo(%q): want %#v, %v, got %#v, %v", tt.path, tt.want, tt.extra, got, extra)
		}
	}
}
<|endoftext|>


<commit_before>package cryptica

import "bytes"

type Solution []Direction

func (solution *Solution) String() string {
	var buffer bytes.Buffer
	for i, dir := range *solution {
		switch {
		case i > 0:
			buffer.WriteRune(' ')
		default:
			buffer.WriteString(dir.String())
		}
	}
	return buffer.String()
}

type Game struct {
	State State
	Goal  Goal
	Steps int
}

type Solver interface {
	solve(game Game, solutions chan Solution)
}

type DepthFirstSolver struct {
}

func (solver DepthFirstSolver) solve(game Game, solutions chan Solution) {
	cache := make(map[uint64]int)
	f := func(state State, solution Solution, k interface{}) {
		n := len(solution)
		if n > game.Steps {
			return
		}
		code := state.Board.Encode(state)
		if m, exists := cache[code]; exists && m <= n {
			return
		}
		cache[code] = n
		if state.Match(game.Goal) {
			newsolution := make(Solution, len(solution))
			copy(newsolution, solution)
			solutions <- newsolution
			return
		}
		f := k.(func(State, Solution, interface{}))
		for i := Up; i <= Right; i++ {
			dir := Direction(i)
			f(state.Move(dir), append(solution, dir), k)
		}
	}
	f(game.State, make(Solution, 0), f)
}

type BreadthFirstSolver struct {
}

func (solver BreadthFirstSolver) solve(game Game, solutions chan Solution) {
	cache := make(map[uint64]Solution)
	var current, next []uint64
	{
		code := game.State.Board.Encode(game.State)
		cache[code] = make(Solution, 0)
		current = []uint64{code}
	}
	for len(current) > 0 {
		next = make([]uint64, 0)
		for _, code := range current {
			solution := cache[code]
			state := game.State.Board.Decode(code)
			if state.Match(game.Goal) {
				solutions <- solution
				continue
			}
			for i := Up; i <= Right; i++ {
				dir := Direction(i)
				newstate := state.Move(dir)
				newcode := game.State.Board.Encode(newstate)
				if _, exists := cache[newcode]; !exists {
					newsolution := make(Solution, len(solution))
					copy(newsolution, solution)
					newsolution = append(newsolution, dir)
					cache[newcode] = newsolution
					next = append(next, newcode)
				}
			}
		}
		current = next
	}
}
<commit_msg>breadth-first search now also stops at Steps<commit_after>package cryptica

import "bytes"

type Solution []Direction

func (solution *Solution) String() string {
	var buffer bytes.Buffer
	for i, dir := range *solution {
		switch {
		case i > 0:
			buffer.WriteRune(' ')
		default:
			buffer.WriteString(dir.String())
		}
	}
	return buffer.String()
}

type Game struct {
	State State
	Goal  Goal
	Steps int
}

type Solver interface {
	solve(game Game, solutions chan Solution)
}

type DepthFirstSolver struct {
}

func (solver DepthFirstSolver) solve(game Game, solutions chan Solution) {
	cache := make(map[uint64]int)
	f := func(state State, solution Solution, k interface{}) {
		n := len(solution)
		if n > game.Steps {
			return
		}
		code := state.Board.Encode(state)
		if m, exists := cache[code]; exists && m <= n {
			return
		}
		cache[code] = n
		if state.Match(game.Goal) {
			newsolution := make(Solution, len(solution))
			copy(newsolution, solution)
			solutions <- newsolution
			return
		}
		f := k.(func(State, Solution, interface{}))
		for i := Up; i <= Right; i++ {
			dir := Direction(i)
			f(state.Move(dir), append(solution, dir), k)
		}
	}
	f(game.State, make(Solution, 0), f)
}

type BreadthFirstSolver struct {
}

func (solver BreadthFirstSolver) solve(game Game, solutions chan Solution) {
	cache := make(map[uint64]Solution)
	var current, next []uint64
	{
		code := game.State.Board.Encode(game.State)
		cache[code] = make(Solution, 0)
		current = []uint64{code}
	}
	for depth := 0; depth < game.Steps && len(current) > 0; depth++ {
		next = make([]uint64, 0)
		for _, code := range current {
			solution := cache[code]
			state := game.State.Board.Decode(code)
			if state.Match(game.Goal) {
				solutions <- solution
				continue
			}
			for i := Up; i <= Right; i++ {
				dir := Direction(i)
				newstate := state.Move(dir)
				newcode := game.State.Board.Encode(newstate)
				if _, exists := cache[newcode]; !exists {
					newsolution := make(Solution, len(solution))
					copy(newsolution, solution)
					newsolution = append(newsolution, dir)
					cache[newcode] = newsolution
					next = append(next, newcode)
				}
			}
		}
		current = next
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"testing"
)

func TestInitialiseKubeconfig(t *testing.T) {

	cmd := make([]string, 2)
	cmd[0] = "install"
	cmd[1] = "--debug"

	plugin := Plugin{
		Config: Config{
			APIServer:     "http://myapiserver",
			Token:         "secret-token",
			HelmCommand:   cmd,
			Namespace:     "default",
			SkipTLSVerify: true,
		},
	}

	initialiseKubeconfig(&plugin.Config, "kubeconfig", "config3.test")

}

func TestGetHelmCommand(t *testing.T) {
	plugin := &Plugin{
		Config: Config{
			APIServer:     "http://myapiserver",
			Token:         "secret-token",
			HelmCommand:   nil,
			Namespace:     "default",
			SkipTLSVerify: true,
			Debug:         true,
			DryRun:        true,
			Chart:         "./chart/test",
			Release:       "test-release",
		},
	}
	setHelmCommand(plugin)
	fmt.Println(plugin.Config.HelmCommand)

}
<commit_msg>remove debug from test<commit_after>package main

import (
	"io/ioutil"
	"strings"
	"testing"
)

func TestInitialiseKubeconfig(t *testing.T) {

	cmd := make([]string, 2)
	cmd[0] = "install"
	cmd[1] = "--debug"

	plugin := Plugin{
		Config: Config{
			APIServer:     "http://myapiserver",
			Token:         "secret-token",
			HelmCommand:   cmd,
			Namespace:     "default",
			SkipTLSVerify: true,
		},
	}

	configfile := "config3.test"
	initialiseKubeconfig(&plugin.Config, "kubeconfig", configfile)
	data, err := ioutil.ReadFile(configfile)
	if err != nil {
		t.Errorf("Error reading file %v", err)
	}
	kubeConfigStr := string(data)

	if !strings.Contains(kubeConfigStr, "secret-token") {
		t.Errorf("Kubeconfig doesn't render token")
	}
	if !strings.Contains(kubeConfigStr, "http://myapiserver") {
		t.Errorf("Kubeconfig doesn't render APIServer")
	}

}

func TestGetHelmCommand(t *testing.T) {
	plugin := &Plugin{
		Config: Config{
			APIServer:     "http://myapiserver",
			Token:         "secret-token",
			HelmCommand:   nil,
			Namespace:     "default",
			SkipTLSVerify: true,
			Debug:         true,
			DryRun:        true,
			Chart:         "./chart/test",
			Release:       "test-release",
		},
	}
	setHelmCommand(plugin)
	res := strings.Join(plugin.Config.HelmCommand[:], " ")
	expected := "upgrade --install test-release ./chart/test --debug --dry-run"
	if res != expected {
		t.Errorf("Result is %s and we expected %s", res, expected)
	}
}
<|endoftext|>


<commit_before>package ws

import (
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

var wsUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

// InitWS add ws handler for api
func InitWS(router *gin.Engine) {
	router.GET("/ws", wsHandler)
}

func wsHandler(c *gin.Context) {
	conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logrus.Warnf("Failed to set websocket upgrade %+v", err)
		c.AbortWithStatus(http.StatusBadRequest)
	}
	defer conn.Close()

	for {
		t, msg, err := conn.ReadMessage()
		if err != nil {
			break
		}
		logrus.Debugf("Got ws message type=%d %s", t, msg)
		conn.WriteMessage(t, msg)
	}
}
<commit_msg>Added ws authentication<commit_after>package ws

import (
	"encoding/json"
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/da4nik/swanager/core/auth"
	"github.com/da4nik/swanager/core/entities"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

type authMessage struct {
	Token string `json:"token"`
}

type answer struct {
	AnswerType string `json:"type"`
	Data       string
}

type connContext struct {
	State     string
	User      *entities.User
	Conn      *websocket.Conn
	AuthError error
}

const (
	stateWorking         = "working"
	stateUnauthenticated = "unauthenticated"
)

var wsUpgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

// InitWS add ws handler for api
func InitWS(router *gin.Engine) {
	router.GET("/ws", wsHandler)
}

func wsHandler(c *gin.Context) {
	conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logrus.Warnf("Failed to set websocket upgrade %+v", err)
		c.AbortWithStatus(http.StatusBadRequest)
	}
	defer conn.Close()

	context := connContext{
		State: stateUnauthenticated,
		Conn:  conn,
	}

	for {
		t, msg, err := conn.ReadMessage()
		if err != nil {
			break
		}

		logrus.Debugf("[%s] Got ws message type=%d %s", context.State, t, msg)

		switch context.State {
		case stateUnauthenticated:
			context.authenticate(msg)
			break
		case stateWorking:
			context.processMessage(msg)
		}
	}
}

func (c *connContext) processMessage(msg []byte) {
	c.Conn.WriteMessage(1, msg)
}

func (c *connContext) authenticate(msg []byte) {
	var message authMessage
	c.AuthError = json.Unmarshal(msg, &message)
	if c.AuthError != nil {
		c.authError()
		return
	}

	c.User, c.AuthError = auth.WithToken(message.Token)
	if c.AuthError != nil {
		c.authError()
		return
	}

	logrus.Debugf("[WS] Authenticated, proceeding with normal mode")

	c.State = stateWorking

	c.sendAnswer(answer{
		AnswerType: "authenticated",
		Data:       "Ok",
	})
}

func (c *connContext) authError() {
	logrus.Debugf("[WS] Auth error: %s", c.AuthError.Error())

	c.sendAnswer(answer{
		AnswerType: "error",
		Data:       c.AuthError.Error(),
	})

	c.Conn.Close()
}

func (c *connContext) sendAnswer(ans answer) {
	result, _ := json.Marshal(ans)
	c.Conn.WriteMessage(1, result)
}
<|endoftext|>


<commit_before>package main

import (
	"sync"
	"sync/atomic"
)

var wg sync.WaitGroup

type CallFrame struct {
	parent    *CallFrame
	var_table map[string]Object
	stack     []Object
	me        Object
}

type Transaction struct {
	startInst        Instruction
	instList         []Instruction
	objectSet        map[*RObject]*RObject
	transactionStack []*CallFrame
	env              *ThreadEnv
	inevitable       bool
	rev              int64
}

type ThreadEnv struct {
	instList      []Instruction
	threadStack   []*CallFrame
	transactionPC *Transaction
	id            int
}

type GobiesVM struct {
	instList       []Instruction
	callFrameStack []*CallFrame
	consts         map[string]*RObject
	symbols        map[string]int
	rev            int64
	globalLock     sync.RWMutex
	has_inevitable int32
}

func initVM() *GobiesVM {
	VM := &GobiesVM{}
	top := initCallFrame()
	VM.callFrameStack = append(VM.callFrameStack, top)
	VM.consts = make(map[string]*RObject)
	VM.initConsts()
	VM.symbols = make(map[string]int)
	top.me = initRKernel()
	return VM
}

func initCallFrame() *CallFrame {
	frame := &CallFrame{}
	frame.var_table = make(map[string]Object)
	return frame
}

func (VM *GobiesVM) initConsts() {
	VM.consts["RString"] = initRString()
	VM.consts["RFixnum"] = initRFixnum()
	VM.consts["RArray"] = initRArray()
	VM.consts["RHash"] = initRHash()
	VM.consts["IO"] = initRIO()
	VM.consts["Thread"] = initRThread()
}

func (obj *RObject) methodLookup(method_name string) *RMethod {
	if val, ok := obj.methods[method_name]; ok {
		return val
	}
	if obj.class != nil {
		return obj.class.methodLookup(method_name)
	}
	return nil
}

func (currentCallFrame *CallFrame) variableLookup(var_name string) Object {
	if obj, ok := currentCallFrame.var_table[var_name]; ok {
		return obj
	}
	if currentCallFrame.parent != nil {
		return currentCallFrame.parent.variableLookup(var_name)
	}
	return nil
}

func (t *Transaction) initTransaction(env *ThreadEnv, instList []Instruction) *Transaction {
	t.instList = []Instruction{}
	for _, inst := range instList {
		t.instList = append(t.instList, inst)
	}
	t.transactionStack = copyFrames(env.threadStack)
	t.objectSet = make(map[*RObject]*RObject)
	return t
}

func copyFrames(src []*CallFrame) []*CallFrame {
	newStack := []*CallFrame{}
	for _, frame := range src {
		newFrame := initCallFrame()
		newFrame.me = frame.me
		newFrame.parent = frame.parent

		// Copy every object "pointer" in old frame stack
		for _, obj := range frame.stack {
			newFrame.stack = append(newFrame.stack, obj)
		}

		// Copy every object "pointer" in instance variable table
		for key, obj := range frame.var_table {
			newFrame.var_table[key] = obj
		}

		newStack = append(newStack, newFrame)
	}
	return newStack
}

func (obj *RObject) copyObject(src *RObject) {
	// Copy instance variables
	ivars := src.ivars
	for k, v := range ivars {
		obj.ivars[k] = v
	}
}

func addRObjectToSet(obj *RObject, env *ThreadEnv) *RObject {
	if env == nil { // created during complication
		return obj
	}
	if _, ok := env.transactionPC.objectSet[obj]; !ok {
		env.transactionPC.objectSet[obj] = obj
	}
	return obj
}

func (VM *GobiesVM) execute() {
	// 	root := initTransaction(VM.instList)
	// 	root.inevitable = true

	VM.rev = 0

	// Execute root transaction which is inevitable
	wg.Add(1)
	go VM.executeThread(VM.instList, nil)
	wg.Wait()
}

func (VM *GobiesVM) executeThread(instList []Instruction, parentScope *ThreadEnv) {
	// Create clean call frame without pushing back to VM stack
	currentCallFrame := initCallFrame()

	env := &ThreadEnv{instList: instList}
	if parentScope == nil {
		currentCallFrame.parent = VM.callFrameStack[len(VM.callFrameStack)-1]
	} else {
		env.threadStack = copyFrames(parentScope.threadStack)
		currentCallFrame.parent = env.threadStack[len(env.threadStack)-1]
		// currentCallFrame.parent = parentScope.threadStack[len(parentScope.threadStack)-1]
	}
	currentCallFrame.me = currentCallFrame.parent.me
	env.threadStack = append(env.threadStack, currentCallFrame)

	// t.inevitable = true
	VM.executeBytecodes(nil, env)
	wg.Done()
}

func (VM *GobiesVM) transactionBegin(env *ThreadEnv, inst []Instruction) *Transaction {
	t := &Transaction{}
	t.initTransaction(env, inst)
	t.rev = atomic.LoadInt64(&VM.rev)
	env.transactionPC = t
	t.env = env

	// Initialize environment
	if t.inevitable {
		VM.globalLock.Lock()
		atomic.StoreInt32(&VM.has_inevitable, 1)
	} else {
		VM.globalLock.RLock()
	}

	return t
}

func (VM *GobiesVM) transactionEnd(env *ThreadEnv) bool {
	t := env.transactionPC

	// Lock the write-set
	locked := []*RObject{}
	for orig_obj, new_obj := range t.objectSet {
		// fmt.Println(orig_obj, new_obj)
		if orig_obj != new_obj {
			if orig_obj.writeLock.TryLock() {
				locked = append(locked, orig_obj)
			} else { // Attempt to acquire lock failed
				// Release all locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY

			}
		}
	}

	// Increment global revision
	for !atomic.CompareAndSwapInt64(&VM.rev, VM.rev, VM.rev+1) {
	}

	// Validate the read-set
	for orig_obj, new_obj := range t.objectSet {
		if orig_obj == new_obj {
			if orig_obj.rev > t.rev {
				// Release write locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY
			}
		}
	}

	// Commit then release the locks
	for _, locked_obj := range locked {
		src := t.objectSet[locked_obj]
		locked_obj.copyObject(src)
	}
	for _, locked_obj := range locked {
		locked_obj.writeLock.Unlock()
	}

	env.threadStack = copyFrames(t.transactionStack)

	if t.inevitable {
		VM.globalLock.Unlock()
		atomic.StoreInt32(&VM.has_inevitable, 0)
	} else {
		VM.globalLock.RUnlock()
	}

	env.transactionPC = nil
	return true

TRANSACTION_RETRY:
	t.initTransaction(env, t.instList)
	t.rev = atomic.LoadInt64(&VM.rev)
	return false
}

func (VM *GobiesVM) executeBytecodes(instList []Instruction, env *ThreadEnv) {
	t := env.transactionPC

	if instList == nil {
		instList = env.instList
	}

	// SPECULATIVE_EXEC:
	// Speculative execution
	for i, v := range instList {
		t = env.transactionPC
		if t == nil {
			t = VM.transactionBegin(env, instList[i:])
		}
		currentCallFrame := t.transactionStack[len(t.transactionStack)-1]

		switch v.inst_type {
		case BC_PUTSELF:
			currentCallFrame.stack = append(currentCallFrame.stack, currentCallFrame.me)
		case BC_PUTNIL:
			currentCallFrame.stack = append(currentCallFrame.stack, nil)
		case BC_PUTOBJ:
			currentCallFrame.stack = append(currentCallFrame.stack, v.obj)
		case BC_PUTFIXNUM:
			currentCallFrame.stack = append(currentCallFrame.stack, RFixnum_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTSTRING:
			currentCallFrame.stack = append(currentCallFrame.stack, RString_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTTRUE:
		case BC_PUTFALSE:
		case BC_SETLOCAL:
			top := currentCallFrame.stack[len(currentCallFrame.stack)-1]
			currentCallFrame.var_table[v.obj.(string)] = top
			currentCallFrame.stack = currentCallFrame.stack[0 : len(currentCallFrame.stack)-1] // Pop object from stack
		case BC_GETLOCAL:
			obj := currentCallFrame.variableLookup(v.obj.(string)).(*RObject)
			if _, ok := t.objectSet[obj]; ok {
				obj = t.objectSet[obj]
			}
			currentCallFrame.stack = append(currentCallFrame.stack, obj)
			if VM.transactionEnd(env) == false {

			}
		case BC_SETGLOBAL:
		case BC_GETGLOBAL:
		case BC_SETSYMBOL:
		case BC_GETSYMBOL:
		case BC_SETCONST:
		case BC_GETCONST:
			currentCallFrame.stack = append(currentCallFrame.stack, VM.consts[v.obj.(string)])
		case BC_SETIVAR:
		case BC_GETIVAR:
		case BC_SETCVAR:
		case BC_GETCVAR:
		case BC_SEND:
			argLists := currentCallFrame.stack[len(currentCallFrame.stack)-(v.argc+1):] // argc + 1 ensures inclusion of receiver
			currentCallFrame.stack = currentCallFrame.stack[:len(currentCallFrame.stack)-(v.argc+1)]
			recv := argLists[0].(*RObject)
			argLists = argLists[1:]
			return_val := recv.methodLookup(v.obj.(string)).gofunc(VM, env, recv, argLists)
			// fmt.Println(env.transactionPC)
			// Update address since some functions might init new transaction
			currentCallFrame = env.transactionPC.transactionStack[len(env.transactionPC.transactionStack)-1]
			if return_val != nil {
				currentCallFrame.stack = append(currentCallFrame.stack, return_val)
			}
			if VM.transactionEnd(env) == false {

			}
		case BC_JUMP:

		}
	}
	// End transaction if any
	if env.transactionPC != nil {
		VM.transactionEnd(env)
	}
}

func (VM *GobiesVM) executeBlock(env *ThreadEnv, block *RObject, args []*RObject) {
	if env.transactionPC != nil { // Before
		VM.transactionEnd(env)
	}
	t := VM.transactionBegin(env, block.methods["def"].def)

	// Create clean call frame
	currentCallFrame := initCallFrame()
	currentCallFrame.parent = t.transactionStack[len(t.transactionStack)-1]
	currentCallFrame.me = currentCallFrame.parent.me
	t.transactionStack = append(t.transactionStack, currentCallFrame)

	// Fill in arguments to current call frame
	if block.ivars["params"] != nil {
		params := block.ivars["params"].(*RObject).ivars["array"].([]*RObject)
		for i, v := range params {
			var_name := v.val.str
			currentCallFrame.var_table[var_name] = args[i]
		}
	}

	// Execute block definition
	VM.executeBytecodes(block.methods["def"].def, env)

	VM.transactionBegin(env, []Instruction{})
}
<commit_msg>Eliminate extra "main" thread overhead<commit_after>package main

import (
	"sync"
	"sync/atomic"
)

var wg sync.WaitGroup

type CallFrame struct {
	parent    *CallFrame
	var_table map[string]Object
	stack     []Object
	me        Object
}

type Transaction struct {
	startInst        Instruction
	instList         []Instruction
	objectSet        map[*RObject]*RObject
	transactionStack []*CallFrame
	env              *ThreadEnv
	inevitable       bool
	rev              int64
}

type ThreadEnv struct {
	instList      []Instruction
	threadStack   []*CallFrame
	transactionPC *Transaction
	id            int
}

type GobiesVM struct {
	instList       []Instruction
	callFrameStack []*CallFrame
	consts         map[string]*RObject
	symbols        map[string]int
	rev            int64
	globalLock     sync.RWMutex
	has_inevitable int32
}

func initVM() *GobiesVM {
	VM := &GobiesVM{}
	top := initCallFrame()
	VM.callFrameStack = append(VM.callFrameStack, top)
	VM.consts = make(map[string]*RObject)
	VM.initConsts()
	VM.symbols = make(map[string]int)
	top.me = initRKernel()
	return VM
}

func initCallFrame() *CallFrame {
	frame := &CallFrame{}
	frame.var_table = make(map[string]Object)
	return frame
}

func (VM *GobiesVM) initConsts() {
	VM.consts["RString"] = initRString()
	VM.consts["RFixnum"] = initRFixnum()
	VM.consts["RArray"] = initRArray()
	VM.consts["RHash"] = initRHash()
	VM.consts["IO"] = initRIO()
	VM.consts["Thread"] = initRThread()
}

func (obj *RObject) methodLookup(method_name string) *RMethod {
	if val, ok := obj.methods[method_name]; ok {
		return val
	}
	if obj.class != nil {
		return obj.class.methodLookup(method_name)
	}
	return nil
}

func (currentCallFrame *CallFrame) variableLookup(var_name string) Object {
	if obj, ok := currentCallFrame.var_table[var_name]; ok {
		return obj
	}
	if currentCallFrame.parent != nil {
		return currentCallFrame.parent.variableLookup(var_name)
	}
	return nil
}

func (t *Transaction) initTransaction(env *ThreadEnv, instList []Instruction) *Transaction {
	t.instList = []Instruction{}
	for _, inst := range instList {
		t.instList = append(t.instList, inst)
	}
	t.transactionStack = copyFrames(env.threadStack)
	t.objectSet = make(map[*RObject]*RObject)
	return t
}

func copyFrames(src []*CallFrame) []*CallFrame {
	newStack := []*CallFrame{}
	for _, frame := range src {
		newFrame := initCallFrame()
		newFrame.me = frame.me
		newFrame.parent = frame.parent

		// Copy every object "pointer" in old frame stack
		for _, obj := range frame.stack {
			newFrame.stack = append(newFrame.stack, obj)
		}

		// Copy every object "pointer" in instance variable table
		for key, obj := range frame.var_table {
			newFrame.var_table[key] = obj
		}

		newStack = append(newStack, newFrame)
	}
	return newStack
}

func (obj *RObject) copyObject(src *RObject) {
	// Copy instance variables
	ivars := src.ivars
	for k, v := range ivars {
		obj.ivars[k] = v
	}
}

func addRObjectToSet(obj *RObject, env *ThreadEnv) *RObject {
	if env == nil { // created during complication
		return obj
	}
	if _, ok := env.transactionPC.objectSet[obj]; !ok {
		env.transactionPC.objectSet[obj] = obj
	}
	return obj
}

func (VM *GobiesVM) execute() {
	// 	root := initTransaction(VM.instList)
	// 	root.inevitable = true

	VM.rev = 0

	// Execute root transaction which is inevitable
	VM.executeThread(VM.instList, nil)
	wg.Wait()
}

func (VM *GobiesVM) executeThread(instList []Instruction, parentScope *ThreadEnv) {
	// Create clean call frame without pushing back to VM stack
	currentCallFrame := initCallFrame()

	env := &ThreadEnv{instList: instList}
	if parentScope == nil {
		currentCallFrame.parent = VM.callFrameStack[len(VM.callFrameStack)-1]
	} else {
		env.threadStack = copyFrames(parentScope.threadStack)
		currentCallFrame.parent = env.threadStack[len(env.threadStack)-1]
		// currentCallFrame.parent = parentScope.threadStack[len(parentScope.threadStack)-1]
	}
	currentCallFrame.me = currentCallFrame.parent.me
	env.threadStack = append(env.threadStack, currentCallFrame)

	// t.inevitable = true
	VM.executeBytecodes(nil, env)
	if parentScope != nil {
		wg.Done()
	}
}

func (VM *GobiesVM) transactionBegin(env *ThreadEnv, inst []Instruction) *Transaction {
	t := &Transaction{}
	t.initTransaction(env, inst)
	t.rev = atomic.LoadInt64(&VM.rev)
	env.transactionPC = t
	t.env = env

	// Initialize environment
	if t.inevitable {
		VM.globalLock.Lock()
		atomic.StoreInt32(&VM.has_inevitable, 1)
	} else {
		VM.globalLock.RLock()
	}

	return t
}

func (VM *GobiesVM) transactionEnd(env *ThreadEnv) bool {
	t := env.transactionPC

	// Lock the write-set
	locked := []*RObject{}
	for orig_obj, new_obj := range t.objectSet {
		// fmt.Println(orig_obj, new_obj)
		if orig_obj != new_obj {
			if orig_obj.writeLock.TryLock() {
				locked = append(locked, orig_obj)
			} else { // Attempt to acquire lock failed
				// Release all locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY

			}
		}
	}

	// Increment global revision
	for !atomic.CompareAndSwapInt64(&VM.rev, VM.rev, VM.rev+1) {
	}

	// Validate the read-set
	for orig_obj, new_obj := range t.objectSet {
		if orig_obj == new_obj {
			if orig_obj.rev > t.rev {
				// Release write locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY
			}
		}
	}

	// Commit then release the locks
	for _, locked_obj := range locked {
		src := t.objectSet[locked_obj]
		locked_obj.copyObject(src)
	}
	for _, locked_obj := range locked {
		locked_obj.writeLock.Unlock()
	}

	env.threadStack = copyFrames(t.transactionStack)

	if t.inevitable {
		VM.globalLock.Unlock()
		atomic.StoreInt32(&VM.has_inevitable, 0)
	} else {
		VM.globalLock.RUnlock()
	}

	env.transactionPC = nil
	return true

TRANSACTION_RETRY:
	t.initTransaction(env, t.instList)
	t.rev = atomic.LoadInt64(&VM.rev)
	return false
}

func (VM *GobiesVM) executeBytecodes(instList []Instruction, env *ThreadEnv) {
	t := env.transactionPC

	if instList == nil {
		instList = env.instList
	}

	// SPECULATIVE_EXEC:
	// Speculative execution
	for i, v := range instList {
		t = env.transactionPC
		if t == nil {
			t = VM.transactionBegin(env, instList[i:])
		}
		currentCallFrame := t.transactionStack[len(t.transactionStack)-1]

		switch v.inst_type {
		case BC_PUTSELF:
			currentCallFrame.stack = append(currentCallFrame.stack, currentCallFrame.me)
		case BC_PUTNIL:
			currentCallFrame.stack = append(currentCallFrame.stack, nil)
		case BC_PUTOBJ:
			currentCallFrame.stack = append(currentCallFrame.stack, v.obj)
		case BC_PUTFIXNUM:
			currentCallFrame.stack = append(currentCallFrame.stack, RFixnum_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTSTRING:
			currentCallFrame.stack = append(currentCallFrame.stack, RString_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTTRUE:
		case BC_PUTFALSE:
		case BC_SETLOCAL:
			top := currentCallFrame.stack[len(currentCallFrame.stack)-1]
			currentCallFrame.var_table[v.obj.(string)] = top
			currentCallFrame.stack = currentCallFrame.stack[0 : len(currentCallFrame.stack)-1] // Pop object from stack
		case BC_GETLOCAL:
			obj := currentCallFrame.variableLookup(v.obj.(string)).(*RObject)
			if _, ok := t.objectSet[obj]; ok {
				obj = t.objectSet[obj]
			}
			currentCallFrame.stack = append(currentCallFrame.stack, obj)
			if VM.transactionEnd(env) == false {

			}
		case BC_SETGLOBAL:
		case BC_GETGLOBAL:
		case BC_SETSYMBOL:
		case BC_GETSYMBOL:
		case BC_SETCONST:
		case BC_GETCONST:
			currentCallFrame.stack = append(currentCallFrame.stack, VM.consts[v.obj.(string)])
		case BC_SETIVAR:
		case BC_GETIVAR:
		case BC_SETCVAR:
		case BC_GETCVAR:
		case BC_SEND:
			argLists := currentCallFrame.stack[len(currentCallFrame.stack)-(v.argc+1):] // argc + 1 ensures inclusion of receiver
			currentCallFrame.stack = currentCallFrame.stack[:len(currentCallFrame.stack)-(v.argc+1)]
			recv := argLists[0].(*RObject)
			argLists = argLists[1:]
			return_val := recv.methodLookup(v.obj.(string)).gofunc(VM, env, recv, argLists)
			// fmt.Println(env.transactionPC)
			// Update address since some functions might init new transaction
			currentCallFrame = env.transactionPC.transactionStack[len(env.transactionPC.transactionStack)-1]
			if return_val != nil {
				currentCallFrame.stack = append(currentCallFrame.stack, return_val)
			}
			if VM.transactionEnd(env) == false {

			}
		case BC_JUMP:

		}
	}
	// End transaction if any
	if env.transactionPC != nil {
		VM.transactionEnd(env)
	}
}

func (VM *GobiesVM) executeBlock(env *ThreadEnv, block *RObject, args []*RObject) {
	if env.transactionPC != nil { // Before
		VM.transactionEnd(env)
	}
	t := VM.transactionBegin(env, block.methods["def"].def)

	// Create clean call frame
	currentCallFrame := initCallFrame()
	currentCallFrame.parent = t.transactionStack[len(t.transactionStack)-1]
	currentCallFrame.me = currentCallFrame.parent.me
	t.transactionStack = append(t.transactionStack, currentCallFrame)

	// Fill in arguments to current call frame
	if block.ivars["params"] != nil {
		params := block.ivars["params"].(*RObject).ivars["array"].([]*RObject)
		for i, v := range params {
			var_name := v.val.str
			currentCallFrame.var_table[var_name] = args[i]
		}
	}

	// Execute block definition
	VM.executeBytecodes(block.methods["def"].def, env)

	VM.transactionBegin(env, []Instruction{})
}
<|endoftext|>


<commit_before>package generator

var ProjectTemplate = `
{{ range $i, $store := .projectTargets }}
{{- if $store.CurrentTarget }}

  {{- if eq $store.IsSystemProject true }}
  <source>
    @type  tail
    path  /var/lib/rancher/rke/log/*.log
    pos_file  /fluentd/log/fluentd-rke-logging-system-project.pos
    time_format  %Y-%m-%dT%H:%M:%S
    tag  rke.*
    format  json
    read_from_head  true
  </source>

  <filter rke.**>
    @type record_transformer
    enable_ruby true
    <record>
      tag ${tag}
      log_type k8s_infrastructure_container
      driver rke
      component ${tag_suffix[6].split("_")[0]}
      container_id ${tag_suffix[6].split(".")[0]}
    </record>
  </filter>
  {{end }}

  <source>
     @type  tail
     path  /var/log/containers/*.log
     pos_file  /fluentd/log/fluentd-project-{{$store.ProjectName}}-logging.pos
     time_format  %Y-%m-%dT%H:%M:%S
     tag  {{$store.ProjectName}}.*
     format  json
     read_from_head  true
  </source>

  <filter {{$store.ProjectName}}.**>
     @type  kubernetes_metadata
     merge_json_log  true
     preserve_json_log  true
  </filter>

  <filter {{$store.ProjectName}}.**>
    @type record_transformer
    enable_ruby  true
    <record>
      tag ${tag}
      namespace ${record["kubernetes"]["namespace_name"]}
      {{range $k, $val := $store.OutputTags -}}
      {{$k}} {{$val}}
      {{end -}}
      projectID {{$store.ProjectName}}
    </record>
  </filter>

  <filter {{$store.ProjectName}}.**>
    @type grep
    <regexp>
      key namespace
      pattern {{$store.GrepNamespace}}
    </regexp>
  </filter>

  <filter {{$store.ProjectName}}.**>
    @type record_transformer
    remove_keys namespace
  </filter>

  {{- if eq $store.CurrentTarget "syslog"}}
  {{- if $store.SyslogConfig.Token}}
  <filter {{$store.ProjectName}}.** project-custom.{{$store.ProjectName}}.** {{ if eq $store.IsSystemProject true }}rke.**{{end }} >
    @type record_transformer
    <record>
      tag ${tag} {{$store.SyslogConfig.Token}}
    </record>
  </filter>
  {{end }}
  {{end }}

  <filter {{$store.ProjectName}}.**>
    @type prometheus
    <metric>
      name fluentd_input_status_num_records_total
      type counter
      desc The total number of incoming records
      <labels>
        tag ${tag}
        hostname ${hostname}
      </labels>
    </metric>
  </filter>

  <match  {{$store.ProjectName}}.** project-custom.{{$store.ProjectName}}.** {{ if eq $store.IsSystemProject true }}rke.**{{end }}>
    @type copy
    <store>
      {{- if or (eq $store.CurrentTarget "elasticsearch") (eq $store.CurrentTarget "splunk") (eq $store.CurrentTarget "syslog") (eq $store.CurrentTarget "kafka") (eq $store.CurrentTarget "fluentforwarder")}}

      {{- if eq $store.CurrentTarget "elasticsearch"}}
      @type elasticsearch
      include_tag_key  true
      {{- if and $store.ElasticsearchConfig.AuthUserName $store.ElasticsearchConfig.AuthPassword}}
      hosts {{$store.ElasticsearchTemplateWrap.Scheme}}://{{$store.ElasticsearchConfig.AuthUserName}}:{{$store.ElasticsearchConfig.AuthPassword}}@{{$store.ElasticsearchTemplateWrap.Host}}
      {{else }}
      hosts {{$store.ElasticsearchConfig.Endpoint}}
      {{end }}

      logstash_prefix "{{$store.ElasticsearchConfig.IndexPrefix}}"
      logstash_format true
      logstash_dateformat  {{$store.ElasticsearchTemplateWrap.DateFormat}}
      type_name  "container_log"

      {{- if eq $store.ElasticsearchTemplateWrap.Scheme "https"}}
      ssl_verify {{$store.ElasticsearchConfig.SSLVerify}}
      ssl_version {{ $store.ElasticsearchConfig.SSLVersion }}

      {{- if $store.ElasticsearchConfig.Certificate }}
      ca_file /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.ElasticsearchConfig.ClientCert $store.ElasticsearchConfig.ClientKey}}
      client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      client_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}

      {{- if $store.ElasticsearchConfig.ClientKeyPass}}
      client_key_pass {{$store.ElasticsearchConfig.ClientKeyPass}}
      {{end }}
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "splunk"}}
      @type splunk_hec
      host {{$store.SplunkTemplateWrap.Host}}
      port {{$store.SplunkTemplateWrap.Port}}
      token {{$store.SplunkConfig.Token}}
      {{- if $store.SplunkConfig.Source}}
      sourcetype {{$store.SplunkConfig.Source}}
      {{end }}
      {{- if $store.SplunkConfig.Index}}
      default_index {{ $store.SplunkConfig.Index }}
      {{end }}

      {{- if eq $store.SplunkTemplateWrap.Scheme "https"}}
      use_ssl true
      ssl_verify {{$store.SplunkConfig.SSLVerify}}

      {{- if $store.SplunkConfig.Certificate }}
      ca_file /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.SplunkConfig.ClientCert $store.SplunkConfig.ClientKey}}
      client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      client_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}

      {{- if $store.SplunkConfig.ClientKeyPass}}
      client_key_pass {{ $store.SplunkConfig.ClientKeyPass }}
      {{end }}
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "kafka"}}
      @type kafka_buffered
      {{- if $store.KafkaConfig.ZookeeperEndpoint }}
      zookeeper {{$store.WrapKafka.Zookeeper}}
      {{else}}
      brokers {{$store.WrapKafka.Brokers}}
      {{end }}
      default_topic {{$store.KafkaConfig.Topic}}
      output_data_type  "json"
      output_include_tag  true
      output_include_time  true
      max_send_retries 3

      {{- if $store.KafkaConfig.Certificate }}
      ssl_ca_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.KafkaConfig.ClientCert $store.KafkaConfig.ClientKey}}
      ssl_client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      ssl_client_cert_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}

      {{- if and $store.KafkaConfig.SaslUsername $store.KafkaConfig.SaslPassword}}
      username {{$store.KafkaConfig.SaslUsername}}
      password {{$store.KafkaConfig.SaslPassword}}
      {{end }}

      {{- if and (eq $store.KafkaConfig.SaslType "scram") $store.KafkaConfig.SaslScramMechanism}}
      scram_mechanism {{$store.KafkaConfig.SaslScramMechanism}}
      {{- if eq $store.KafkaTemplateWrap.IsSSL false}}
      sasl_over_ssl false
      {{end}}
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "syslog"}}
      @type remote_syslog
      host {{$store.SyslogTemplateWrap.Host}}
      port {{$store.SyslogTemplateWrap.Port}}
      severity {{$store.SyslogConfig.Severity}}
      protocol {{$store.SyslogConfig.Protocol}}
      {{- if $store.SyslogConfig.Program }}
      program {{$store.SyslogConfig.Program}}
      {{end }}
      packet_size 65535

      {{- if eq $store.SyslogConfig.SSLVerify true}}
      verify_mode 1
      {{else }}
      verify_mode 0
      {{end }}

      {{- if $store.SyslogConfig.Certificate }}
      tls true
      ca_file /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.SyslogConfig.ClientCert $store.SyslogConfig.ClientKey}}
      client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      client_cert_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "fluentforwarder"}}
      @type forward
      {{- if $store.FluentForwarderConfig.EnableTLS }}
      transport tls
      tls_allow_self_signed_cert true
      tls_verify_hostname true
      {{end }}
      {{- if $store.FluentForwarderConfig.Certificate }}
      tls_cert_path /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if $store.FluentForwarderConfig.Compress }}
      compress gzip
      {{end }}

      {{- if $store.FluentForwarderTemplateWrap.EnableShareKey }}
      <security>
        self_hostname "#{Socket.gethostname}"
        shared_key true
      </security>
      {{end }}

      {{- range $k, $val := $store.FluentForwarderTemplateWrap.FluentServers }}
      <server>
        {{if $val.Hostname}}
        name {{$val.Hostname}}
        {{end }}
        host {{$val.Host}}
        port {{$val.Port}}
        {{ if $val.SharedKey}}
        shared_key {{$val.SharedKey}}
        {{end }}
        {{ if $val.Username}}
        username  {{$val.Username}}
        {{end }}
        {{ if $val.Password}}
        password  {{$val.Password}}
        {{end }}
        weight  {{$val.Weight}}
        {{if $val.Standby}}
        standby
        {{end }}
      </server>
      {{end }}
      {{end }}

      <buffer>
        @type file
        path /fluentd/log/buffer/project.{{$store.WrapProjectName}}.buffer
        flush_mode interval
        flush_interval {{$store.OutputFlushInterval}}s
        flush_thread_count 8
        {{- if eq $store.CurrentTarget "splunk"}}
        chunk_limit_size 8m
        {{end }}
      </buffer>

      slow_flush_log_threshold 40.0
      queued_chunks_limit_size 200
      {{end }}

      {{- if eq $store.CurrentTarget "customtarget"}}
      {{$store.CustomTargetWrap.Content}}
      {{end }}
    </store>

    <store>
    @type prometheus
    <metric>
      name fluentd_output_status_num_records_total
      type counter
      desc The total number of outgoing records
      <labels>
        tag ${tag}
        hostname ${hostname}
      </labels>
    </metric>
    </store>
  </match>

{{end }}
{{end }}
`
<commit_msg>fixed the project config template bug<commit_after>package generator

var ProjectTemplate = `
{{ range $i, $store := .projectTargets }}
{{- if $store.CurrentTarget }}

  {{- if eq $store.IsSystemProject true }}
  <source>
    @type  tail
    path  /var/lib/rancher/rke/log/*.log
    pos_file  /fluentd/log/fluentd-rke-logging-system-project.pos
    time_format  %Y-%m-%dT%H:%M:%S
    tag  rke.*
    format  json
    read_from_head  true
  </source>

  <filter rke.**>
    @type record_transformer
    enable_ruby true
    <record>
      tag ${tag}
      log_type k8s_infrastructure_container
      driver rke
      component ${tag_suffix[6].split("_")[0]}
      container_id ${tag_suffix[6].split(".")[0]}
    </record>
  </filter>
  {{end }}

  <source>
     @type  tail
     path  /var/log/containers/*.log
     pos_file  /fluentd/log/fluentd-project-{{$store.ProjectName}}-logging.pos
     time_format  %Y-%m-%dT%H:%M:%S
     tag  {{$store.ProjectName}}.*
     format  json
     read_from_head  true
  </source>

  <filter {{$store.ProjectName}}.**>
     @type  kubernetes_metadata
     merge_json_log  true
     preserve_json_log  true
  </filter>

  <filter {{$store.ProjectName}}.**>
    @type record_transformer
    enable_ruby  true
    <record>
      tag ${tag}
      namespace ${record["kubernetes"]["namespace_name"]}
      {{range $k, $val := $store.OutputTags -}}
      {{$k}} {{$val}}
      {{end -}}
      projectID {{$store.ProjectName}}
    </record>
  </filter>

  <filter {{$store.ProjectName}}.**>
    @type grep
    <regexp>
      key namespace
      pattern {{$store.GrepNamespace}}
    </regexp>
  </filter>

  <filter {{$store.ProjectName}}.**>
    @type record_transformer
    remove_keys namespace
  </filter>

  {{- if eq $store.CurrentTarget "syslog"}}
  {{- if $store.SyslogConfig.Token}}
  <filter {{$store.ProjectName}}.** project-custom.{{$store.ProjectName}}.** {{ if eq $store.IsSystemProject true }}rke.**{{end }} >
    @type record_transformer
    <record>
      tag ${tag} {{$store.SyslogConfig.Token}}
    </record>
  </filter>
  {{end }}
  {{end }}

  <filter {{$store.ProjectName}}.**>
    @type prometheus
    <metric>
      name fluentd_input_status_num_records_total
      type counter
      desc The total number of incoming records
      <labels>
        tag ${tag}
        hostname ${hostname}
      </labels>
    </metric>
  </filter>

  <match  {{$store.ProjectName}}.** project-custom.{{$store.ProjectName}}.** {{ if eq $store.IsSystemProject true }}rke.**{{end }}>
    @type copy
    <store>
      {{- if or (eq $store.CurrentTarget "elasticsearch") (eq $store.CurrentTarget "splunk") (eq $store.CurrentTarget "syslog") (eq $store.CurrentTarget "kafka") (eq $store.CurrentTarget "fluentforwarder")}}

      {{- if eq $store.CurrentTarget "elasticsearch"}}
      @type elasticsearch
      include_tag_key  true
      {{- if and $store.ElasticsearchConfig.AuthUserName $store.ElasticsearchConfig.AuthPassword}}
      hosts {{$store.ElasticsearchTemplateWrap.Scheme}}://{{$store.ElasticsearchConfig.AuthUserName}}:{{$store.ElasticsearchConfig.AuthPassword}}@{{$store.ElasticsearchTemplateWrap.Host}}
      {{else }}
      hosts {{$store.ElasticsearchConfig.Endpoint}}
      {{end }}

      logstash_prefix "{{$store.ElasticsearchConfig.IndexPrefix}}"
      logstash_format true
      logstash_dateformat  {{$store.ElasticsearchTemplateWrap.DateFormat}}
      type_name  "container_log"

      {{- if eq $store.ElasticsearchTemplateWrap.Scheme "https"}}
      ssl_verify {{$store.ElasticsearchConfig.SSLVerify}}
      ssl_version {{ $store.ElasticsearchConfig.SSLVersion }}

      {{- if $store.ElasticsearchConfig.Certificate }}
      ca_file /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.ElasticsearchConfig.ClientCert $store.ElasticsearchConfig.ClientKey}}
      client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      client_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}

      {{- if $store.ElasticsearchConfig.ClientKeyPass}}
      client_key_pass {{$store.ElasticsearchConfig.ClientKeyPass}}
      {{end }}
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "splunk"}}
      @type splunk_hec
      host {{$store.SplunkTemplateWrap.Host}}
      port {{$store.SplunkTemplateWrap.Port}}
      token {{$store.SplunkConfig.Token}}
      {{- if $store.SplunkConfig.Source}}
      sourcetype {{$store.SplunkConfig.Source}}
      {{end }}
      {{- if $store.SplunkConfig.Index}}
      default_index {{ $store.SplunkConfig.Index }}
      {{end }}

      {{- if eq $store.SplunkTemplateWrap.Scheme "https"}}
      use_ssl true
      ssl_verify {{$store.SplunkConfig.SSLVerify}}

      {{- if $store.SplunkConfig.Certificate }}
      ca_file /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.SplunkConfig.ClientCert $store.SplunkConfig.ClientKey}}
      client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      client_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}

      {{- if $store.SplunkConfig.ClientKeyPass}}
      client_key_pass {{ $store.SplunkConfig.ClientKeyPass }}
      {{end }}
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "kafka"}}
      @type kafka_buffered
      {{- if $store.KafkaConfig.ZookeeperEndpoint }}
      zookeeper {{$store.KafkaTemplateWrap.Zookeeper}}
      {{else}}
      brokers {{$store.KafkaTemplateWrap.Brokers}}
      {{end }}
      default_topic {{$store.KafkaConfig.Topic}}
      output_data_type  "json"
      output_include_tag  true
      output_include_time  true
      max_send_retries 3

      {{- if $store.KafkaConfig.Certificate }}
      ssl_ca_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.KafkaConfig.ClientCert $store.KafkaConfig.ClientKey}}
      ssl_client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      ssl_client_cert_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}

      {{- if and $store.KafkaConfig.SaslUsername $store.KafkaConfig.SaslPassword}}
      username {{$store.KafkaConfig.SaslUsername}}
      password {{$store.KafkaConfig.SaslPassword}}
      {{end }}

      {{- if and (eq $store.KafkaConfig.SaslType "scram") $store.KafkaConfig.SaslScramMechanism}}
      scram_mechanism {{$store.KafkaConfig.SaslScramMechanism}}
      {{- if eq $store.KafkaTemplateWrap.IsSSL false}}
      sasl_over_ssl false
      {{end}}
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "syslog"}}
      @type remote_syslog
      host {{$store.SyslogTemplateWrap.Host}}
      port {{$store.SyslogTemplateWrap.Port}}
      severity {{$store.SyslogConfig.Severity}}
      protocol {{$store.SyslogConfig.Protocol}}
      {{- if $store.SyslogConfig.Program }}
      program {{$store.SyslogConfig.Program}}
      {{end }}
      packet_size 65535

      {{- if eq $store.SyslogConfig.SSLVerify true}}
      verify_mode 1
      {{else }}
      verify_mode 0
      {{end }}

      {{- if $store.SyslogConfig.Certificate }}
      tls true
      ca_file /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if and $store.SyslogConfig.ClientCert $store.SyslogConfig.ClientKey}}
      client_cert /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-cert.pem
      client_cert_key /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_client-key.pem
      {{end }}
      {{end }}

      {{- if eq $store.CurrentTarget "fluentforwarder"}}
      @type forward
      {{- if $store.FluentForwarderConfig.EnableTLS }}
      transport tls
      tls_allow_self_signed_cert true
      tls_verify_hostname true
      {{end }}
      {{- if $store.FluentForwarderConfig.Certificate }}
      tls_cert_path /fluentd/etc/config/ssl/project_{{$store.WrapProjectName}}_ca.pem
      {{end }}

      {{- if $store.FluentForwarderConfig.Compress }}
      compress gzip
      {{end }}

      {{- if $store.FluentForwarderTemplateWrap.EnableShareKey }}
      <security>
        self_hostname "#{Socket.gethostname}"
        shared_key true
      </security>
      {{end }}

      {{- range $k, $val := $store.FluentForwarderTemplateWrap.FluentServers }}
      <server>
        {{if $val.Hostname}}
        name {{$val.Hostname}}
        {{end }}
        host {{$val.Host}}
        port {{$val.Port}}
        {{ if $val.SharedKey}}
        shared_key {{$val.SharedKey}}
        {{end }}
        {{ if $val.Username}}
        username  {{$val.Username}}
        {{end }}
        {{ if $val.Password}}
        password  {{$val.Password}}
        {{end }}
        weight  {{$val.Weight}}
        {{if $val.Standby}}
        standby
        {{end }}
      </server>
      {{end }}
      {{end }}

      <buffer>
        @type file
        path /fluentd/log/buffer/project.{{$store.WrapProjectName}}.buffer
        flush_mode interval
        flush_interval {{$store.OutputFlushInterval}}s
        flush_thread_count 8
        {{- if eq $store.CurrentTarget "splunk"}}
        chunk_limit_size 8m
        {{end }}
      </buffer>

      slow_flush_log_threshold 40.0
      queued_chunks_limit_size 200
      {{end }}

      {{- if eq $store.CurrentTarget "customtarget"}}
      {{$store.CustomTargetWrap.Content}}
      {{end }}
    </store>

    <store>
    @type prometheus
    <metric>
      name fluentd_output_status_num_records_total
      type counter
      desc The total number of outgoing records
      <labels>
        tag ${tag}
        hostname ${hostname}
      </labels>
    </metric>
    </store>
  </match>

{{end }}
{{end }}
`
<|endoftext|>


<commit_before>package main

import (
	"context"
	"net/http"
	"net/http/httptest"
	"runtime"
	"strconv"
	"strings"
	"testing"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/docker/docker/integration-cli/cli"
	"github.com/docker/docker/integration-cli/cli/build"
	"github.com/docker/docker/pkg/parsers/kernel"
	"github.com/docker/docker/testutil/request"
	"gotest.tools/assert"
)

func (s *DockerSuite) TestAPIImagesFilter(c *testing.T) {
	cli, err := client.NewClientWithOpts(client.FromEnv)
	assert.NilError(c, err)
	defer cli.Close()

	name := "utest:tag1"
	name2 := "utest/docker:tag2"
	name3 := "utest:5000/docker:tag3"
	for _, n := range []string{name, name2, name3} {
		dockerCmd(c, "tag", "busybox", n)
	}
	getImages := func(filter string) []types.ImageSummary {
		filters := filters.NewArgs()
		filters.Add("reference", filter)
		options := types.ImageListOptions{
			All:     false,
			Filters: filters,
		}
		images, err := cli.ImageList(context.Background(), options)
		assert.NilError(c, err)

		return images
	}

	//incorrect number of matches returned
	images := getImages("utest*/*")
	assert.Equal(c, len(images[0].RepoTags), 2)

	images = getImages("utest")
	assert.Equal(c, len(images[0].RepoTags), 1)

	images = getImages("utest*")
	assert.Equal(c, len(images[0].RepoTags), 1)

	images = getImages("*5000*/*")
	assert.Equal(c, len(images[0].RepoTags), 1)
}

func (s *DockerSuite) TestAPIImagesSaveAndLoad(c *testing.T) {
	if runtime.GOOS == "windows" {
		v, err := kernel.GetKernelVersion()
		assert.NilError(c, err)
		build, _ := strconv.Atoi(strings.Split(strings.SplitN(v.String(), " ", 3)[2][1:], ".")[0])
		if build == 16299 {
			c.Skip("Temporarily disabled on RS3 builds")
		}
	}

	testRequires(c, Network)
	buildImageSuccessfully(c, "saveandload", build.WithDockerfile("FROM busybox\nENV FOO bar"))
	id := getIDByName(c, "saveandload")

	res, body, err := request.Get("/images/" + id + "/get")
	assert.NilError(c, err)
	defer body.Close()
	assert.Equal(c, res.StatusCode, http.StatusOK)

	dockerCmd(c, "rmi", id)

	res, loadBody, err := request.Post("/images/load", request.RawContent(body), request.ContentType("application/x-tar"))
	assert.NilError(c, err)
	defer loadBody.Close()
	assert.Equal(c, res.StatusCode, http.StatusOK)

	inspectOut := cli.InspectCmd(c, id, cli.Format(".Id")).Combined()
	assert.Equal(c, strings.TrimSpace(string(inspectOut)), id, "load did not work properly")
}

func (s *DockerSuite) TestAPIImagesDelete(c *testing.T) {
	cli, err := client.NewClientWithOpts(client.FromEnv)
	assert.NilError(c, err)
	defer cli.Close()

	if testEnv.OSType != "windows" {
		testRequires(c, Network)
	}
	name := "test-api-images-delete"
	buildImageSuccessfully(c, name, build.WithDockerfile("FROM busybox\nENV FOO bar"))
	id := getIDByName(c, name)

	dockerCmd(c, "tag", name, "test:tag1")

	_, err = cli.ImageRemove(context.Background(), id, types.ImageRemoveOptions{})
	assert.ErrorContains(c, err, "unable to delete")

	_, err = cli.ImageRemove(context.Background(), "test:noexist", types.ImageRemoveOptions{})
	assert.ErrorContains(c, err, "No such image")

	_, err = cli.ImageRemove(context.Background(), "test:tag1", types.ImageRemoveOptions{})
	assert.NilError(c, err)
}

func (s *DockerSuite) TestAPIImagesHistory(c *testing.T) {
	cli, err := client.NewClientWithOpts(client.FromEnv)
	assert.NilError(c, err)
	defer cli.Close()

	if testEnv.OSType != "windows" {
		testRequires(c, Network)
	}
	name := "test-api-images-history"
	buildImageSuccessfully(c, name, build.WithDockerfile("FROM busybox\nENV FOO bar"))
	id := getIDByName(c, name)

	historydata, err := cli.ImageHistory(context.Background(), id)
	assert.NilError(c, err)

	assert.Assert(c, len(historydata) != 0)
	var found bool
	for _, tag := range historydata[0].Tags {
		if tag == "test-api-images-history:latest" {
			found = true
			break
		}
	}
	assert.Assert(c, found)
}

func (s *DockerSuite) TestAPIImagesImportBadSrc(c *testing.T) {
	if runtime.GOOS == "windows" {
		v, err := kernel.GetKernelVersion()
		assert.NilError(c, err)
		build, _ := strconv.Atoi(strings.Split(strings.SplitN(v.String(), " ", 3)[2][1:], ".")[0])
		if build == 16299 {
			c.Skip("Temporarily disabled on RS3 builds")
		}
	}

	testRequires(c, Network, testEnv.IsLocalDaemon)

	server := httptest.NewServer(http.NewServeMux())
	defer server.Close()

	tt := []struct {
		statusExp int
		fromSrc   string
	}{
		{http.StatusNotFound, server.URL + "/nofile.tar"},
		{http.StatusNotFound, strings.TrimPrefix(server.URL, "http://") + "/nofile.tar"},
		{http.StatusNotFound, strings.TrimPrefix(server.URL, "http://") + "%2Fdata%2Ffile.tar"},
		{http.StatusInternalServerError, "%2Fdata%2Ffile.tar"},
	}

	for _, te := range tt {
		res, _, err := request.Post(strings.Join([]string{"/images/create?fromSrc=", te.fromSrc}, ""), request.JSON)
		assert.NilError(c, err)
		assert.Equal(c, res.StatusCode, te.statusExp)
		assert.Equal(c, res.Header.Get("Content-Type"), "application/json")
	}

}

// #14846
func (s *DockerSuite) TestAPIImagesSearchJSONContentType(c *testing.T) {
	testRequires(c, Network)

	res, b, err := request.Get("/images/search?term=test", request.JSON)
	assert.NilError(c, err)
	b.Close()
	assert.Equal(c, res.StatusCode, http.StatusOK)
	assert.Equal(c, res.Header.Get("Content-Type"), "application/json")
}

// Test case for 30027: image size reported as -1 in v1.12 client against v1.13 daemon.
// This test checks to make sure both v1.12 and v1.13 client against v1.13 daemon get correct `Size` after the fix.
func (s *DockerSuite) TestAPIImagesSizeCompatibility(c *testing.T) {
	apiclient := testEnv.APIClient()
	defer apiclient.Close()

	images, err := apiclient.ImageList(context.Background(), types.ImageListOptions{})
	assert.NilError(c, err)
	assert.Assert(c, len(images) != 0)
	for _, image := range images {
		assert.Assert(c, image.Size != int64(-1))
	}

	apiclient, err = client.NewClientWithOpts(client.FromEnv, client.WithVersion("v1.24"))
	assert.NilError(c, err)
	defer apiclient.Close()

	v124Images, err := apiclient.ImageList(context.Background(), types.ImageListOptions{})
	assert.NilError(c, err)
	assert.Assert(c, len(v124Images) != 0)
	for _, image := range v124Images {
		assert.Assert(c, image.Size != int64(-1))
	}
}
<commit_msg>integration-cli: Skip TestAPIImagesSaveAndLoad on RS3 and older<commit_after>package main

import (
	"context"
	"net/http"
	"net/http/httptest"
	"runtime"
	"strconv"
	"strings"
	"testing"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/docker/docker/integration-cli/cli"
	"github.com/docker/docker/integration-cli/cli/build"
	"github.com/docker/docker/pkg/parsers/kernel"
	"github.com/docker/docker/testutil/request"
	"gotest.tools/assert"
)

func (s *DockerSuite) TestAPIImagesFilter(c *testing.T) {
	cli, err := client.NewClientWithOpts(client.FromEnv)
	assert.NilError(c, err)
	defer cli.Close()

	name := "utest:tag1"
	name2 := "utest/docker:tag2"
	name3 := "utest:5000/docker:tag3"
	for _, n := range []string{name, name2, name3} {
		dockerCmd(c, "tag", "busybox", n)
	}
	getImages := func(filter string) []types.ImageSummary {
		filters := filters.NewArgs()
		filters.Add("reference", filter)
		options := types.ImageListOptions{
			All:     false,
			Filters: filters,
		}
		images, err := cli.ImageList(context.Background(), options)
		assert.NilError(c, err)

		return images
	}

	//incorrect number of matches returned
	images := getImages("utest*/*")
	assert.Equal(c, len(images[0].RepoTags), 2)

	images = getImages("utest")
	assert.Equal(c, len(images[0].RepoTags), 1)

	images = getImages("utest*")
	assert.Equal(c, len(images[0].RepoTags), 1)

	images = getImages("*5000*/*")
	assert.Equal(c, len(images[0].RepoTags), 1)
}

func (s *DockerSuite) TestAPIImagesSaveAndLoad(c *testing.T) {
	if runtime.GOOS == "windows" {
		v, err := kernel.GetKernelVersion()
		assert.NilError(c, err)
		build, _ := strconv.Atoi(strings.Split(strings.SplitN(v.String(), " ", 3)[2][1:], ".")[0])
		if build <= 16299 {
			c.Skip("Temporarily disabled on RS3 and older because they are too slow. See #39909")
		}
	}

	testRequires(c, Network)
	buildImageSuccessfully(c, "saveandload", build.WithDockerfile("FROM busybox\nENV FOO bar"))
	id := getIDByName(c, "saveandload")

	res, body, err := request.Get("/images/" + id + "/get")
	assert.NilError(c, err)
	defer body.Close()
	assert.Equal(c, res.StatusCode, http.StatusOK)

	dockerCmd(c, "rmi", id)

	res, loadBody, err := request.Post("/images/load", request.RawContent(body), request.ContentType("application/x-tar"))
	assert.NilError(c, err)
	defer loadBody.Close()
	assert.Equal(c, res.StatusCode, http.StatusOK)

	inspectOut := cli.InspectCmd(c, id, cli.Format(".Id")).Combined()
	assert.Equal(c, strings.TrimSpace(string(inspectOut)), id, "load did not work properly")
}

func (s *DockerSuite) TestAPIImagesDelete(c *testing.T) {
	cli, err := client.NewClientWithOpts(client.FromEnv)
	assert.NilError(c, err)
	defer cli.Close()

	if testEnv.OSType != "windows" {
		testRequires(c, Network)
	}
	name := "test-api-images-delete"
	buildImageSuccessfully(c, name, build.WithDockerfile("FROM busybox\nENV FOO bar"))
	id := getIDByName(c, name)

	dockerCmd(c, "tag", name, "test:tag1")

	_, err = cli.ImageRemove(context.Background(), id, types.ImageRemoveOptions{})
	assert.ErrorContains(c, err, "unable to delete")

	_, err = cli.ImageRemove(context.Background(), "test:noexist", types.ImageRemoveOptions{})
	assert.ErrorContains(c, err, "No such image")

	_, err = cli.ImageRemove(context.Background(), "test:tag1", types.ImageRemoveOptions{})
	assert.NilError(c, err)
}

func (s *DockerSuite) TestAPIImagesHistory(c *testing.T) {
	cli, err := client.NewClientWithOpts(client.FromEnv)
	assert.NilError(c, err)
	defer cli.Close()

	if testEnv.OSType != "windows" {
		testRequires(c, Network)
	}
	name := "test-api-images-history"
	buildImageSuccessfully(c, name, build.WithDockerfile("FROM busybox\nENV FOO bar"))
	id := getIDByName(c, name)

	historydata, err := cli.ImageHistory(context.Background(), id)
	assert.NilError(c, err)

	assert.Assert(c, len(historydata) != 0)
	var found bool
	for _, tag := range historydata[0].Tags {
		if tag == "test-api-images-history:latest" {
			found = true
			break
		}
	}
	assert.Assert(c, found)
}

func (s *DockerSuite) TestAPIImagesImportBadSrc(c *testing.T) {
	if runtime.GOOS == "windows" {
		v, err := kernel.GetKernelVersion()
		assert.NilError(c, err)
		build, _ := strconv.Atoi(strings.Split(strings.SplitN(v.String(), " ", 3)[2][1:], ".")[0])
		if build == 16299 {
			c.Skip("Temporarily disabled on RS3 builds")
		}
	}

	testRequires(c, Network, testEnv.IsLocalDaemon)

	server := httptest.NewServer(http.NewServeMux())
	defer server.Close()

	tt := []struct {
		statusExp int
		fromSrc   string
	}{
		{http.StatusNotFound, server.URL + "/nofile.tar"},
		{http.StatusNotFound, strings.TrimPrefix(server.URL, "http://") + "/nofile.tar"},
		{http.StatusNotFound, strings.TrimPrefix(server.URL, "http://") + "%2Fdata%2Ffile.tar"},
		{http.StatusInternalServerError, "%2Fdata%2Ffile.tar"},
	}

	for _, te := range tt {
		res, _, err := request.Post(strings.Join([]string{"/images/create?fromSrc=", te.fromSrc}, ""), request.JSON)
		assert.NilError(c, err)
		assert.Equal(c, res.StatusCode, te.statusExp)
		assert.Equal(c, res.Header.Get("Content-Type"), "application/json")
	}

}

// #14846
func (s *DockerSuite) TestAPIImagesSearchJSONContentType(c *testing.T) {
	testRequires(c, Network)

	res, b, err := request.Get("/images/search?term=test", request.JSON)
	assert.NilError(c, err)
	b.Close()
	assert.Equal(c, res.StatusCode, http.StatusOK)
	assert.Equal(c, res.Header.Get("Content-Type"), "application/json")
}

// Test case for 30027: image size reported as -1 in v1.12 client against v1.13 daemon.
// This test checks to make sure both v1.12 and v1.13 client against v1.13 daemon get correct `Size` after the fix.
func (s *DockerSuite) TestAPIImagesSizeCompatibility(c *testing.T) {
	apiclient := testEnv.APIClient()
	defer apiclient.Close()

	images, err := apiclient.ImageList(context.Background(), types.ImageListOptions{})
	assert.NilError(c, err)
	assert.Assert(c, len(images) != 0)
	for _, image := range images {
		assert.Assert(c, image.Size != int64(-1))
	}

	apiclient, err = client.NewClientWithOpts(client.FromEnv, client.WithVersion("v1.24"))
	assert.NilError(c, err)
	defer apiclient.Close()

	v124Images, err := apiclient.ImageList(context.Background(), types.ImageListOptions{})
	assert.NilError(c, err)
	assert.Assert(c, len(v124Images) != 0)
	for _, image := range v124Images {
		assert.Assert(c, image.Size != int64(-1))
	}
}
<|endoftext|>


<commit_before>package daemon

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"

	"github.com/gdamore/mangos"
	"github.com/gdamore/mangos/protocol/rep"
	"github.com/gdamore/mangos/transport/tcp"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/activemq"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/amqp"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/amqp/rabbitmq"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/beanstalkd"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/kafka"
//	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/kestrel"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/nats"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/nsq"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/pubsub"
)

type daemon string
type operation string

const (
	start    operation = "start"
	stop     operation = "stop"
	run      operation = "run"
	sub      operation = "subscribers"
	pub      operation = "publishers"
	results  operation = "results"
	teardown operation = "teardown"
)

// These are supported message brokers.
const (
	NATS        = "nats"
	Beanstalkd  = "beanstalkd"
	Kafka       = "kafka"
	Kestrel     = "kestrel"
	ActiveMQ    = "activemq"
	RabbitMQ    = "rabbitmq"
	NSQ         = "nsq"
	CloudPubSub = "pubsub"
)

type request struct {
	Operation   operation `json:"operation"`
	Broker      string    `json:"broker"`
	Port        string    `json:"port"`
	NumMessages int       `json:"num_messages"`
	MessageSize int64     `json:"message_size"`
	Count       int       `json:"count"`
	Host        string    `json:"host"`
}

type response struct {
	Success    bool        `json:"success"`
	Message    string      `json:"message"`
	Result     interface{} `json:"result"`
	PubResults []*result   `json:"pub_results,omitempty"`
	SubResults []*result   `json:"sub_results,omitempty"`
}

type result struct {
	Duration   float32         `json:"duration,omitempty"`
	Throughput float32         `json:"throughput,omitempty"`
	Latency    *latencyResults `json:"latency,omitempty"`
	Err        string          `json:"error,omitempty"`
}

// broker handles configuring the message broker for testing.
type broker interface {
	// Start will start the message broker and prepare it for testing.
	Start(string, string) (interface{}, error)

	// Stop will stop the message broker.
	Stop() (interface{}, error)
}

// peer is a single producer or consumer in the test.
type peer interface {
	// Subscribe prepares the peer to consume messages.
	Subscribe() error

	// Recv returns a single message consumed by the peer. Subscribe must be
	// called before this. It returns an error if the receive failed.
	Recv() ([]byte, error)

	// Send returns a channel on which messages can be sent for publishing.
	Send() chan<- []byte

	// Errors returns the channel on which the peer sends publish errors.
	Errors() <-chan error

	// Done signals to the peer that message publishing has completed.
	Done()

	// Setup prepares the peer for testing.
	Setup()

	// Teardown performs any cleanup logic that needs to be performed after the
	// test is complete.
	Teardown()
}

// Config contains configuration settings for the Flotilla daemon.
type Config struct {
	GoogleCloudProjectID string
	GoogleCloudJSONKey   string
}

// Daemon is the server portion of Flotilla which runs on machines we want to
// communicate with and include in our benchmarks.
type Daemon struct {
	mangos.Socket
	broker      broker
	publishers  []*publisher
	subscribers []*subscriber
	config      *Config
}

// NewDaemon creates and returns a new Daemon from the provided Config. An
// error is returned if the Daemon cannot be created.
func NewDaemon(config *Config) (*Daemon, error) {
	rep, err := rep.NewSocket()
	if err != nil {
		return nil, err
	}
	rep.AddTransport(tcp.NewTransport())
	return &Daemon{rep, nil, []*publisher{}, []*subscriber{}, config}, nil
}

// Start will allow the Daemon to begin processing requests. This is a blocking
// call.
func (d *Daemon) Start(port int) error {
	if err := d.Listen(fmt.Sprintf("tcp://:%d", port)); err != nil {
		return err
	}
	return d.loop()
}

func (d *Daemon) loop() error {
	for {
		msg, err := d.Recv()
		if err != nil {
			log.Println(err)
			continue
		}

		var req request
		if err := json.Unmarshal(msg, &req); err != nil {
			log.Println("Invalid peer request:", err)
			d.sendResponse(response{
				Success: false,
				Message: fmt.Sprintf("Invalid request: %s", err.Error()),
			})
			continue
		}

		resp := d.processRequest(req)
		d.sendResponse(resp)
	}
}

func (d *Daemon) sendResponse(rep response) {
	repJSON, err := json.Marshal(rep)
	if err != nil {
		// This is not recoverable.
		panic(err)
	}

	if err := d.Send(repJSON); err != nil {
		log.Println(err)
	}
}

func (d *Daemon) processRequest(req request) response {
	var (
		response response
		err      error
	)
	switch req.Operation {
	case start:
		response.Result, err = d.processBrokerStart(req.Broker, req.Host, req.Port)
	case stop:
		response.Result, err = d.processBrokerStop()
	case pub:
		err = d.processPub(req)
	case sub:
		err = d.processSub(req)
	case run:
		err = d.processPublisherStart()
	case results:
		response.PubResults, response.SubResults, err = d.processResults()
		if err != nil {
			response.Message = err.Error()
			err = nil
		}
	case teardown:
		d.processTeardown()
	default:
		err = fmt.Errorf("Invalid operation %s", req.Operation)
	}

	if err != nil {
		response.Message = err.Error()
	} else {
		response.Success = true
	}

	return response
}
func (d *Daemon) processBrokerStart(broker, host, port string) (interface{}, error) {
	if d.broker != nil {
		return "", errors.New("Broker already running")
	}

	switch broker {
	case NATS:
		d.broker = &nats.Broker{}
	case Beanstalkd:
		d.broker = &beanstalkd.Broker{}
	case Kafka:
		d.broker = &kafka.Broker{}
	case Kestrel:
		d.broker = &kestrel.Broker{}
	case ActiveMQ:
		d.broker = &activemq.Broker{}
	case RabbitMQ:
		d.broker = &rabbitmq.Broker{}
	case NSQ:
		d.broker = &nsq.Broker{}
	case CloudPubSub:
		d.broker = &pubsub.Broker{
			ProjectID: d.config.GoogleCloudProjectID,
			JSONKey:   d.config.GoogleCloudJSONKey,
		}
	default:
		return "", fmt.Errorf("Invalid broker %s", broker)
	}

	result, err := d.broker.Start(host, port)
	if err != nil {
		d.broker = nil
	}
	return result, err
}

func (d *Daemon) processBrokerStop() (interface{}, error) {
	if d.broker == nil {
		return "", errors.New("No broker running")
	}

	result, err := d.broker.Stop()
	if err == nil {
		d.broker = nil
	}
	return result, err
}

func (d *Daemon) processPub(req request) error {
	for i := 0; i < req.Count; i++ {
		sender, err := d.newPeer(req.Broker, req.Host)
		if err != nil {
			return err
		}

		d.publishers = append(d.publishers, &publisher{
			peer:        sender,
			id:          i,
			numMessages: req.NumMessages,
			messageSize: req.MessageSize,
		})
	}

	return nil
}

func (d *Daemon) processSub(req request) error {
	for i := 0; i < req.Count; i++ {
		receiver, err := d.newPeer(req.Broker, req.Host)
		if err != nil {
			return err
		}

		if err := receiver.Subscribe(); err != nil {
			return err
		}

		subscriber := &subscriber{
			peer:        receiver,
			id:          i,
			numMessages: req.NumMessages,
			messageSize: req.MessageSize,
		}
		d.subscribers = append(d.subscribers, subscriber)
		go subscriber.start()
	}

	return nil
}

func (d *Daemon) processPublisherStart() error {
	for _, publisher := range d.publishers {
		go publisher.start()
	}

	return nil
}

func (d *Daemon) processResults() ([]*result, []*result, error) {
	subResults := make([]*result, 0, len(d.subscribers))
	for _, subscriber := range d.subscribers {
		result, err := subscriber.getResults()
		if err != nil {
			return nil, nil, err
		}
		subResults = append(subResults, result)
	}

	pubResults := make([]*result, 0, len(d.publishers))
	for _, publisher := range d.publishers {
		result, err := publisher.getResults()
		if err != nil {
			return nil, nil, err
		}
		pubResults = append(pubResults, result)
	}

	log.Println("Benchmark completed")
	return pubResults, subResults, nil
}

func (d *Daemon) processTeardown() {
	for _, subscriber := range d.subscribers {
		subscriber.Teardown()
	}
	d.subscribers = d.subscribers[:0]

	for _, publisher := range d.publishers {
		publisher.Teardown()
	}
	d.publishers = d.publishers[:0]
}

func (d *Daemon) newPeer(broker, host string) (peer, error) {
	switch broker {
	case NATS:
		return nats.NewPeer(host)
	case Beanstalkd:
		return beanstalkd.NewPeer(host)
	case Kafka:
		return kafka.NewPeer(host)
	case Kestrel:
		return kestrel.NewPeer(host)
	case ActiveMQ:
		return activemq.NewPeer(host)
	case RabbitMQ:
		return amqp.NewPeer(host)
	case NSQ:
		return nsq.NewPeer(host)
	case CloudPubSub:
		return pubsub.NewPeer(
			d.config.GoogleCloudProjectID,
			d.config.GoogleCloudJSONKey,
		)
	default:
		return nil, fmt.Errorf("Invalid broker: %s", broker)
	}
}
<commit_msg>Removed kestrel, since it is not able to download some dependencies.<commit_after>package daemon

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"

	"github.com/gdamore/mangos"
	"github.com/gdamore/mangos/protocol/rep"
	"github.com/gdamore/mangos/transport/tcp"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/activemq"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/amqp"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/amqp/rabbitmq"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/beanstalkd"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/kafka"
//	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/kestrel"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/nats"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/nsq"
	"github.com/davidgev/Flotilla/flotilla-server/daemon/broker/pubsub"
)

type daemon string
type operation string

const (
	start    operation = "start"
	stop     operation = "stop"
	run      operation = "run"
	sub      operation = "subscribers"
	pub      operation = "publishers"
	results  operation = "results"
	teardown operation = "teardown"
)

// These are supported message brokers.
const (
	NATS        = "nats"
	Beanstalkd  = "beanstalkd"
	Kafka       = "kafka"
	Kestrel     = "kestrel"
	ActiveMQ    = "activemq"
	RabbitMQ    = "rabbitmq"
	NSQ         = "nsq"
	CloudPubSub = "pubsub"
)

type request struct {
	Operation   operation `json:"operation"`
	Broker      string    `json:"broker"`
	Port        string    `json:"port"`
	NumMessages int       `json:"num_messages"`
	MessageSize int64     `json:"message_size"`
	Count       int       `json:"count"`
	Host        string    `json:"host"`
}

type response struct {
	Success    bool        `json:"success"`
	Message    string      `json:"message"`
	Result     interface{} `json:"result"`
	PubResults []*result   `json:"pub_results,omitempty"`
	SubResults []*result   `json:"sub_results,omitempty"`
}

type result struct {
	Duration   float32         `json:"duration,omitempty"`
	Throughput float32         `json:"throughput,omitempty"`
	Latency    *latencyResults `json:"latency,omitempty"`
	Err        string          `json:"error,omitempty"`
}

// broker handles configuring the message broker for testing.
type broker interface {
	// Start will start the message broker and prepare it for testing.
	Start(string, string) (interface{}, error)

	// Stop will stop the message broker.
	Stop() (interface{}, error)
}

// peer is a single producer or consumer in the test.
type peer interface {
	// Subscribe prepares the peer to consume messages.
	Subscribe() error

	// Recv returns a single message consumed by the peer. Subscribe must be
	// called before this. It returns an error if the receive failed.
	Recv() ([]byte, error)

	// Send returns a channel on which messages can be sent for publishing.
	Send() chan<- []byte

	// Errors returns the channel on which the peer sends publish errors.
	Errors() <-chan error

	// Done signals to the peer that message publishing has completed.
	Done()

	// Setup prepares the peer for testing.
	Setup()

	// Teardown performs any cleanup logic that needs to be performed after the
	// test is complete.
	Teardown()
}

// Config contains configuration settings for the Flotilla daemon.
type Config struct {
	GoogleCloudProjectID string
	GoogleCloudJSONKey   string
}

// Daemon is the server portion of Flotilla which runs on machines we want to
// communicate with and include in our benchmarks.
type Daemon struct {
	mangos.Socket
	broker      broker
	publishers  []*publisher
	subscribers []*subscriber
	config      *Config
}

// NewDaemon creates and returns a new Daemon from the provided Config. An
// error is returned if the Daemon cannot be created.
func NewDaemon(config *Config) (*Daemon, error) {
	rep, err := rep.NewSocket()
	if err != nil {
		return nil, err
	}
	rep.AddTransport(tcp.NewTransport())
	return &Daemon{rep, nil, []*publisher{}, []*subscriber{}, config}, nil
}

// Start will allow the Daemon to begin processing requests. This is a blocking
// call.
func (d *Daemon) Start(port int) error {
	if err := d.Listen(fmt.Sprintf("tcp://:%d", port)); err != nil {
		return err
	}
	return d.loop()
}

func (d *Daemon) loop() error {
	for {
		msg, err := d.Recv()
		if err != nil {
			log.Println(err)
			continue
		}

		var req request
		if err := json.Unmarshal(msg, &req); err != nil {
			log.Println("Invalid peer request:", err)
			d.sendResponse(response{
				Success: false,
				Message: fmt.Sprintf("Invalid request: %s", err.Error()),
			})
			continue
		}

		resp := d.processRequest(req)
		d.sendResponse(resp)
	}
}

func (d *Daemon) sendResponse(rep response) {
	repJSON, err := json.Marshal(rep)
	if err != nil {
		// This is not recoverable.
		panic(err)
	}

	if err := d.Send(repJSON); err != nil {
		log.Println(err)
	}
}

func (d *Daemon) processRequest(req request) response {
	var (
		response response
		err      error
	)
	switch req.Operation {
	case start:
		response.Result, err = d.processBrokerStart(req.Broker, req.Host, req.Port)
	case stop:
		response.Result, err = d.processBrokerStop()
	case pub:
		err = d.processPub(req)
	case sub:
		err = d.processSub(req)
	case run:
		err = d.processPublisherStart()
	case results:
		response.PubResults, response.SubResults, err = d.processResults()
		if err != nil {
			response.Message = err.Error()
			err = nil
		}
	case teardown:
		d.processTeardown()
	default:
		err = fmt.Errorf("Invalid operation %s", req.Operation)
	}

	if err != nil {
		response.Message = err.Error()
	} else {
		response.Success = true
	}

	return response
}
func (d *Daemon) processBrokerStart(broker, host, port string) (interface{}, error) {
	if d.broker != nil {
		return "", errors.New("Broker already running")
	}

	switch broker {
	case NATS:
		d.broker = &nats.Broker{}
	case Beanstalkd:
		d.broker = &beanstalkd.Broker{}
	case Kafka:
		d.broker = &kafka.Broker{}
//	case Kestrel:
//		d.broker = &kestrel.Broker{}
	case ActiveMQ:
		d.broker = &activemq.Broker{}
	case RabbitMQ:
		d.broker = &rabbitmq.Broker{}
	case NSQ:
		d.broker = &nsq.Broker{}
	case CloudPubSub:
		d.broker = &pubsub.Broker{
			ProjectID: d.config.GoogleCloudProjectID,
			JSONKey:   d.config.GoogleCloudJSONKey,
		}
	default:
		return "", fmt.Errorf("Invalid broker %s", broker)
	}

	result, err := d.broker.Start(host, port)
	if err != nil {
		d.broker = nil
	}
	return result, err
}

func (d *Daemon) processBrokerStop() (interface{}, error) {
	if d.broker == nil {
		return "", errors.New("No broker running")
	}

	result, err := d.broker.Stop()
	if err == nil {
		d.broker = nil
	}
	return result, err
}

func (d *Daemon) processPub(req request) error {
	for i := 0; i < req.Count; i++ {
		sender, err := d.newPeer(req.Broker, req.Host)
		if err != nil {
			return err
		}

		d.publishers = append(d.publishers, &publisher{
			peer:        sender,
			id:          i,
			numMessages: req.NumMessages,
			messageSize: req.MessageSize,
		})
	}

	return nil
}

func (d *Daemon) processSub(req request) error {
	for i := 0; i < req.Count; i++ {
		receiver, err := d.newPeer(req.Broker, req.Host)
		if err != nil {
			return err
		}

		if err := receiver.Subscribe(); err != nil {
			return err
		}

		subscriber := &subscriber{
			peer:        receiver,
			id:          i,
			numMessages: req.NumMessages,
			messageSize: req.MessageSize,
		}
		d.subscribers = append(d.subscribers, subscriber)
		go subscriber.start()
	}

	return nil
}

func (d *Daemon) processPublisherStart() error {
	for _, publisher := range d.publishers {
		go publisher.start()
	}

	return nil
}

func (d *Daemon) processResults() ([]*result, []*result, error) {
	subResults := make([]*result, 0, len(d.subscribers))
	for _, subscriber := range d.subscribers {
		result, err := subscriber.getResults()
		if err != nil {
			return nil, nil, err
		}
		subResults = append(subResults, result)
	}

	pubResults := make([]*result, 0, len(d.publishers))
	for _, publisher := range d.publishers {
		result, err := publisher.getResults()
		if err != nil {
			return nil, nil, err
		}
		pubResults = append(pubResults, result)
	}

	log.Println("Benchmark completed")
	return pubResults, subResults, nil
}

func (d *Daemon) processTeardown() {
	for _, subscriber := range d.subscribers {
		subscriber.Teardown()
	}
	d.subscribers = d.subscribers[:0]

	for _, publisher := range d.publishers {
		publisher.Teardown()
	}
	d.publishers = d.publishers[:0]
}

func (d *Daemon) newPeer(broker, host string) (peer, error) {
	switch broker {
	case NATS:
		return nats.NewPeer(host)
	case Beanstalkd:
		return beanstalkd.NewPeer(host)
	case Kafka:
		return kafka.NewPeer(host)
//	case Kestrel:
//		return kestrel.NewPeer(host)
	case ActiveMQ:
		return activemq.NewPeer(host)
	case RabbitMQ:
		return amqp.NewPeer(host)
	case NSQ:
		return nsq.NewPeer(host)
	case CloudPubSub:
		return pubsub.NewPeer(
			d.config.GoogleCloudProjectID,
			d.config.GoogleCloudJSONKey,
		)
	default:
		return nil, fmt.Errorf("Invalid broker: %s", broker)
	}
}
<|endoftext|>


<commit_before>package rpc

import (
	"cgl.tideland.biz/asserts"
	"net/rpc"
	"testing"
)

type testUi struct {
	sayCalled bool
	sayFormat string
	sayVars []interface{}
}

func (u *testUi) Say(format string, a ...interface{}) {
	u.sayCalled = true
	u.sayFormat = format
	u.sayVars = a
}

func TestUiRPC(t *testing.T) {
	assert := asserts.NewTestingAsserts(t, true)

	// Create the UI to test
	ui := new(testUi)
	uiServer := &UiServer{ui}

	// Start the RPC server
	readyChan := make(chan int)
	stopChan := make(chan int)
	defer func() { stopChan <- 1 }()
	go testRPCServer(":1234", "Ui", uiServer, readyChan, stopChan)
	<-readyChan

	// Create the client over RPC and run some methods to verify it works
	client, err := rpc.Dial("tcp", ":1234")
	if err != nil {
		panic(err)
	}

	uiClient := &Ui{client}
	uiClient.Say("format", "arg0", 42)

	assert.Equal(ui.sayFormat, "format", "format should be correct")
}
<commit_msg>packer/rpc: Use the proper Server for tests<commit_after>package rpc

import (
	"cgl.tideland.biz/asserts"
	"net/rpc"
	"testing"
)

type testUi struct {
	sayCalled bool
	sayFormat string
	sayVars []interface{}
}

func (u *testUi) Say(format string, a ...interface{}) {
	u.sayCalled = true
	u.sayFormat = format
	u.sayVars = a
}

func TestUiRPC(t *testing.T) {
	assert := asserts.NewTestingAsserts(t, true)

	// Create the UI to test
	ui := new(testUi)

	// Start the RPC server
	server := NewServer()
	server.RegisterUi(ui)
	server.Start()
	defer server.Stop()

	// Create the client over RPC and run some methods to verify it works
	client, err := rpc.Dial("tcp", server.Address())
	if err != nil {
		panic(err)
	}

	uiClient := &Ui{client}
	uiClient.Say("format", "arg0", 42)

	assert.Equal(ui.sayFormat, "format", "format should be correct")
}
<|endoftext|>