<commit_before>/*
Copyright 2020 The Knative Authors

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

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

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

package trigger

import (
	"context"
	"fmt"

	"go.uber.org/zap"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/equality"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"knative.dev/pkg/logging"

	"knative.dev/eventing-autoscaler-keda/pkg/reconciler/keda"
	"knative.dev/eventing-autoscaler-keda/pkg/reconciler/trigger/resources"
	"knative.dev/eventing/pkg/apis/eventing"
	v1 "knative.dev/eventing/pkg/apis/eventing/v1"
	triggerreconciler "knative.dev/eventing/pkg/client/injection/reconciler/eventing/v1/trigger"
	eventinglisters "knative.dev/eventing/pkg/client/listers/eventing/v1"

	pkgreconciler "knative.dev/pkg/reconciler"

	kedaclientset "github.com/kedacore/keda/pkg/generated/clientset/versioned"
	kedalisters "github.com/kedacore/keda/pkg/generated/listers/keda/v1alpha1"
)

const (
	// Name of the corev1.Events emitted from the Trigger reconciliation process.
	triggerReconciled = "TriggerReconciled"
)

// This has to stay in sync with:
// https://github.com/knative-sandbox/eventing-rabbitmq/blob/master/pkg/reconciler/broker/resources/secret.go#L49
func secretName(brokerName string) string {
	return fmt.Sprintf("%s-broker-rabbit", brokerName)
}

type Reconciler struct {
	kedaClientset kedaclientset.Interface

	// listers index properties about resources
	brokerLister       eventinglisters.BrokerLister
	scaledObjectLister kedalisters.ScaledObjectLister
	brokerClass        string
}

// Check that our Reconciler implements Interface
var _ triggerreconciler.Interface = (*Reconciler)(nil)

// ReconcilerArgs are the arguments needed to create a broker.Reconciler.
type ReconcilerArgs struct {
	DispatcherImage              string
	DispatcherServiceAccountName string
}

func newReconciledNormal(namespace, name string) pkgreconciler.Event {
	return pkgreconciler.NewEvent(corev1.EventTypeNormal, triggerReconciled, "Trigger reconciled: \"%s/%s\"", namespace, name)
}

func (r *Reconciler) ReconcileKind(ctx context.Context, t *v1.Trigger) pkgreconciler.Event {
	logging.FromContext(ctx).Debug("Reconciling", zap.Any("Trigger", t))

	broker, err := r.brokerLister.Brokers(t.Namespace).Get(t.Spec.Broker)
	if err != nil {
		if apierrs.IsNotFound(err) {
			// Ok to return nil here. Once the Broker comes available, or Trigger changes, we get requeued.
			return nil
		}
		logging.FromContext(ctx).Errorf("Failed to get Broker: \"%s/%s\" : %s", t.Spec.Broker, t.Namespace, err)
		return nil
	}

	// If it's not my brokerclass, ignore
	if broker.Annotations[eventing.BrokerClassKey] != r.brokerClass {
		logging.FromContext(ctx).Infof("Ignoring trigger %s/%s", t.Namespace, t.Name)
		return nil
	}

	return r.reconcileScaledObject(ctx, broker, t)
}

func (r *Reconciler) reconcileScaledObject(ctx context.Context, broker *v1.Broker, trigger *v1.Trigger) error {
	// Check the annotation to see if the Brokers Triggers should even be scaled.
	doAutoscale := broker.GetAnnotations()[keda.AutoscalingClassAnnotation] == keda.KEDA

	so, err := resources.MakeDispatcherScaledObject(ctx, broker, trigger)
	if err != nil {
		logging.FromContext(ctx).Errorw("Failed to create scaled object resource", zap.Error(err))
		return err
	}

	current, err := r.scaledObjectLister.ScaledObjects(so.Namespace).Get(so.Name)
	if apierrs.IsNotFound(err) {
		if !doAutoscale {
			// Ok, not there, not wanted, we're good...
			return nil
		}
		logging.FromContext(ctx).Infof("Creating ScaledObject %s/%s", so.Namespace, so.Name)
		_, err = r.kedaClientset.KedaV1alpha1().ScaledObjects(so.Namespace).Create(ctx, so, metav1.CreateOptions{})
		return err
	}
	if err != nil {
		return err
	}

	// It's there, should it be?
	if !doAutoscale {
		logging.FromContext(ctx).Infof("Deleting ScaledObject %s/%s", so.Namespace, so.Name)
		err = r.kedaClientset.KedaV1alpha1().ScaledObjects(so.Namespace).Delete(ctx, so.Name, metav1.DeleteOptions{})
		if err != nil {
			logging.FromContext(ctx).Errorw("Failed to delete ScaledObject", zap.Error(err))
		}
		return err
	}
	if !equality.Semantic.DeepDerivative(so.Spec, current.Spec) {
		// Don't modify the informers copy.
		desired := current.DeepCopy()
		desired.Spec = so.Spec
		_, err = r.kedaClientset.KedaV1alpha1().ScaledObjects(desired.Namespace).Update(ctx, desired, metav1.UpdateOptions{})
		return err
	}
	return nil
}
<commit_msg>Remove dead code (#33)<commit_after>/*
Copyright 2020 The Knative Authors

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

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

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

package trigger

import (
	"context"

	"go.uber.org/zap"

	"k8s.io/apimachinery/pkg/api/equality"
	apierrs "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"knative.dev/pkg/logging"

	"knative.dev/eventing-autoscaler-keda/pkg/reconciler/keda"
	"knative.dev/eventing-autoscaler-keda/pkg/reconciler/trigger/resources"
	"knative.dev/eventing/pkg/apis/eventing"
	v1 "knative.dev/eventing/pkg/apis/eventing/v1"
	triggerreconciler "knative.dev/eventing/pkg/client/injection/reconciler/eventing/v1/trigger"
	eventinglisters "knative.dev/eventing/pkg/client/listers/eventing/v1"

	pkgreconciler "knative.dev/pkg/reconciler"

	kedaclientset "github.com/kedacore/keda/pkg/generated/clientset/versioned"
	kedalisters "github.com/kedacore/keda/pkg/generated/listers/keda/v1alpha1"
)

const (
	// Name of the corev1.Events emitted from the Trigger reconciliation process.
	triggerReconciled = "TriggerReconciled"
)

type Reconciler struct {
	kedaClientset kedaclientset.Interface

	// listers index properties about resources
	brokerLister       eventinglisters.BrokerLister
	scaledObjectLister kedalisters.ScaledObjectLister
	brokerClass        string
}

// Check that our Reconciler implements Interface
var _ triggerreconciler.Interface = (*Reconciler)(nil)

// ReconcilerArgs are the arguments needed to create a broker.Reconciler.
type ReconcilerArgs struct {
	DispatcherImage              string
	DispatcherServiceAccountName string
}

func (r *Reconciler) ReconcileKind(ctx context.Context, t *v1.Trigger) pkgreconciler.Event {
	logging.FromContext(ctx).Debug("Reconciling", zap.Any("Trigger", t))

	broker, err := r.brokerLister.Brokers(t.Namespace).Get(t.Spec.Broker)
	if err != nil {
		if apierrs.IsNotFound(err) {
			// Ok to return nil here. Once the Broker comes available, or Trigger changes, we get requeued.
			return nil
		}
		logging.FromContext(ctx).Errorf("Failed to get Broker: \"%s/%s\" : %s", t.Spec.Broker, t.Namespace, err)
		return nil
	}

	// If it's not my brokerclass, ignore
	if broker.Annotations[eventing.BrokerClassKey] != r.brokerClass {
		logging.FromContext(ctx).Infof("Ignoring trigger %s/%s", t.Namespace, t.Name)
		return nil
	}

	return r.reconcileScaledObject(ctx, broker, t)
}

func (r *Reconciler) reconcileScaledObject(ctx context.Context, broker *v1.Broker, trigger *v1.Trigger) error {
	// Check the annotation to see if the Brokers Triggers should even be scaled.
	doAutoscale := broker.GetAnnotations()[keda.AutoscalingClassAnnotation] == keda.KEDA

	so, err := resources.MakeDispatcherScaledObject(ctx, broker, trigger)
	if err != nil {
		logging.FromContext(ctx).Errorw("Failed to create scaled object resource", zap.Error(err))
		return err
	}

	current, err := r.scaledObjectLister.ScaledObjects(so.Namespace).Get(so.Name)
	if apierrs.IsNotFound(err) {
		if !doAutoscale {
			// Ok, not there, not wanted, we're good...
			return nil
		}
		logging.FromContext(ctx).Infof("Creating ScaledObject %s/%s", so.Namespace, so.Name)
		_, err = r.kedaClientset.KedaV1alpha1().ScaledObjects(so.Namespace).Create(ctx, so, metav1.CreateOptions{})
		return err
	}
	if err != nil {
		return err
	}

	// It's there, should it be?
	if !doAutoscale {
		logging.FromContext(ctx).Infof("Deleting ScaledObject %s/%s", so.Namespace, so.Name)
		err = r.kedaClientset.KedaV1alpha1().ScaledObjects(so.Namespace).Delete(ctx, so.Name, metav1.DeleteOptions{})
		if err != nil {
			logging.FromContext(ctx).Errorw("Failed to delete ScaledObject", zap.Error(err))
		}
		return err
	}
	if !equality.Semantic.DeepDerivative(so.Spec, current.Spec) {
		// Don't modify the informers copy.
		desired := current.DeepCopy()
		desired.Spec = so.Spec
		_, err = r.kedaClientset.KedaV1alpha1().ScaledObjects(desired.Namespace).Update(ctx, desired, metav1.UpdateOptions{})
		return err
	}
	return nil
}
<|endoftext|>


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

package path

import (
	"os";
	"testing";
)

type CleanTest struct {
	path, clean string;
}

var cleantests = []CleanTest{
	// Already clean
	CleanTest{"", "."},
	CleanTest{"abc", "abc"},
	CleanTest{"abc/def", "abc/def"},
	CleanTest{"a/b/c", "a/b/c"},
	CleanTest{".", "."},
	CleanTest{"..", ".."},
	CleanTest{"../..", "../.."},
	CleanTest{"../../abc", "../../abc"},
	CleanTest{"/abc", "/abc"},
	CleanTest{"/", "/"},

	// Remove trailing slash
	CleanTest{"abc/", "abc"},
	CleanTest{"abc/def/", "abc/def"},
	CleanTest{"a/b/c/", "a/b/c"},
	CleanTest{"./", "."},
	CleanTest{"../", ".."},
	CleanTest{"../../", "../.."},
	CleanTest{"/abc/", "/abc"},

	// Remove doubled slash
	CleanTest{"abc//def//ghi", "abc/def/ghi"},
	CleanTest{"//abc", "/abc"},
	CleanTest{"///abc", "/abc"},
	CleanTest{"//abc//", "/abc"},
	CleanTest{"abc//", "abc"},

	// Remove . elements
	CleanTest{"abc/./def", "abc/def"},
	CleanTest{"/./abc/def", "/abc/def"},
	CleanTest{"abc/.", "abc"},

	// Remove .. elements
	CleanTest{"abc/def/ghi/../jkl", "abc/def/jkl"},
	CleanTest{"abc/def/../ghi/../jkl", "abc/jkl"},
	CleanTest{"abc/def/..", "abc"},
	CleanTest{"abc/def/../..", "."},
	CleanTest{"/abc/def/../..", "/"},
	CleanTest{"abc/def/../../..", ".."},
	CleanTest{"/abc/def/../../..", "/"},
	CleanTest{"abc/def/../../../ghi/jkl/../../../mno", "../../mno"},

	// Combinations
	CleanTest{"abc/./../def", "def"},
	CleanTest{"abc//./../def", "def"},
	CleanTest{"abc/../../././../def", "../../def"},
}

func TestClean(t *testing.T) {
	for _, test := range cleantests {
		if s := Clean(test.path); s != test.clean {
			t.Errorf("Clean(%q) = %q, want %q", test.path, s, test.clean)
		}
	}
}

type SplitTest struct {
	path, dir, file string;
}

var splittests = []SplitTest{
	SplitTest{"a/b", "a/", "b"},
	SplitTest{"a/b/", "a/b/", ""},
	SplitTest{"a/", "a/", ""},
	SplitTest{"a", "", "a"},
	SplitTest{"/", "/", ""},
}

func TestSplit(t *testing.T) {
	for _, test := range splittests {
		if d, f := Split(test.path); d != test.dir || f != test.file {
			t.Errorf("Split(%q) = %q, %q, want %q, %q", test.path, d, f, test.dir, test.file)
		}
	}
}

type JoinTest struct {
	dir, file, path string;
}

var jointests = []JoinTest{
	JoinTest{"a", "b", "a/b"},
	JoinTest{"a", "", "a"},
	JoinTest{"", "b", "b"},
	JoinTest{"/", "a", "/a"},
	JoinTest{"/", "", "/"},
	JoinTest{"a/", "b", "a/b"},
	JoinTest{"a/", "", "a"},
}

func TestJoin(t *testing.T) {
	for _, test := range jointests {
		if p := Join(test.dir, test.file); p != test.path {
			t.Errorf("Join(%q, %q) = %q, want %q", test.dir, test.file, p, test.path)
		}
	}
}

type ExtTest struct {
	path, ext string;
}

var exttests = []ExtTest{
	ExtTest{"path.go", ".go"},
	ExtTest{"path.pb.go", ".go"},
	ExtTest{"a.dir/b", ""},
	ExtTest{"a.dir/b.go", ".go"},
	ExtTest{"a.dir/", ""},
}

func TestExt(t *testing.T) {
	for _, test := range exttests {
		if x := Ext(test.path); x != test.ext {
			t.Errorf("Ext(%q) = %q, want %q", test.path, x, test.ext)
		}
	}
}

type Node struct {
	name	string;
	entries	[]*Node;	// nil if the entry is a file
	mark	int;
}

var tree = &Node{
	"testdata",
	[]*Node{
		&Node{"a", nil, 0},
		&Node{"b", []*Node{}, 0},
		&Node{"c", nil, 0},
		&Node{
			"d",
			[]*Node{
				&Node{"x", nil, 0},
				&Node{"y", []*Node{}, 0},
				&Node{
					"z",
					[]*Node{
						&Node{"u", nil, 0},
						&Node{"v", nil, 0},
					},
					0,
				},
			},
			0,
		},
	},
	0,
}

func walkTree(n *Node, path string, f func(path string, n *Node)) {
	f(path, n);
	for _, e := range n.entries {
		walkTree(e, Join(path, e.name), f)
	}
}

func makeTree(t *testing.T) {
	walkTree(tree, tree.name, func(path string, n *Node) {
		if n.entries == nil {
			fd, err := os.Open(path, os.O_CREAT, 0660);
			if err != nil {
				t.Errorf("makeTree: %v", err)
			}
			fd.Close();
		} else {
			os.Mkdir(path, 0770)
		}
	})
}

func markTree(n *Node)	{ walkTree(n, "", func(path string, n *Node) { n.mark++ }) }

func checkMarks(t *testing.T) {
	walkTree(tree, tree.name, func(path string, n *Node) {
		if n.mark != 1 {
			t.Errorf("node %s mark = %d; expected 1", path, n.mark)
		}
		n.mark = 0;
	})
}

// Assumes that each node name is unique. Good enough for a test.
func mark(name string) {
	walkTree(tree, tree.name, func(path string, n *Node) {
		if n.name == name {
			n.mark++
		}
	})
}

type TestVisitor struct{}

func (v *TestVisitor) VisitDir(path string, d *os.Dir) bool {
	mark(d.Name);
	return true;
}

func (v *TestVisitor) VisitFile(path string, d *os.Dir) {
	mark(d.Name)
}

func TestWalk(t *testing.T) {
	makeTree(t);

	// 1) ignore error handling, expect none
	v := &TestVisitor{};
	Walk(tree.name, v, nil);
	checkMarks(t);

	// 2) handle errors, expect none
	errors := make(chan os.Error, 64);
	Walk(tree.name, v, errors);
	if err, ok := <-errors; ok {
		t.Errorf("no error expected, found: s", err)
	}
	checkMarks(t);

	// introduce 2 errors: chmod top-level directories to 0
	os.Chmod(Join(tree.name, tree.entries[1].name), 0);
	os.Chmod(Join(tree.name, tree.entries[3].name), 0);
	// mark respective subtrees manually
	markTree(tree.entries[1]);
	markTree(tree.entries[3]);
	// correct double-marking of directory itself
	tree.entries[1].mark--;
	tree.entries[3].mark--;

	// 3) handle errors, expect two
	errors = make(chan os.Error, 64);
	os.Chmod(Join(tree.name, tree.entries[1].name), 0);
	Walk(tree.name, v, errors);
	for i := 1; i <= 2; i++ {
		if _, ok := <-errors; !ok {
			t.Errorf("%d. error expected, none found", i);
			break;
		}
	}
	if err, ok := <-errors; ok {
		t.Errorf("only two errors expected, found 3rd: %v", err)
	}
	// the inaccessible subtrees were marked manually
	checkMarks(t);

	// cleanup
	os.Chmod(Join(tree.name, tree.entries[1].name), 0770);
	os.Chmod(Join(tree.name, tree.entries[3].name), 0770);
	if err := os.RemoveAll(tree.name); err != nil {
		t.Errorf("removeTree: %v", err)
	}
}
<commit_msg>path.TestWalk: disable error case if root (chmod 0 doesn't cause errors for root)<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package path

import (
	"os";
	"testing";
)

type CleanTest struct {
	path, clean string;
}

var cleantests = []CleanTest{
	// Already clean
	CleanTest{"", "."},
	CleanTest{"abc", "abc"},
	CleanTest{"abc/def", "abc/def"},
	CleanTest{"a/b/c", "a/b/c"},
	CleanTest{".", "."},
	CleanTest{"..", ".."},
	CleanTest{"../..", "../.."},
	CleanTest{"../../abc", "../../abc"},
	CleanTest{"/abc", "/abc"},
	CleanTest{"/", "/"},

	// Remove trailing slash
	CleanTest{"abc/", "abc"},
	CleanTest{"abc/def/", "abc/def"},
	CleanTest{"a/b/c/", "a/b/c"},
	CleanTest{"./", "."},
	CleanTest{"../", ".."},
	CleanTest{"../../", "../.."},
	CleanTest{"/abc/", "/abc"},

	// Remove doubled slash
	CleanTest{"abc//def//ghi", "abc/def/ghi"},
	CleanTest{"//abc", "/abc"},
	CleanTest{"///abc", "/abc"},
	CleanTest{"//abc//", "/abc"},
	CleanTest{"abc//", "abc"},

	// Remove . elements
	CleanTest{"abc/./def", "abc/def"},
	CleanTest{"/./abc/def", "/abc/def"},
	CleanTest{"abc/.", "abc"},

	// Remove .. elements
	CleanTest{"abc/def/ghi/../jkl", "abc/def/jkl"},
	CleanTest{"abc/def/../ghi/../jkl", "abc/jkl"},
	CleanTest{"abc/def/..", "abc"},
	CleanTest{"abc/def/../..", "."},
	CleanTest{"/abc/def/../..", "/"},
	CleanTest{"abc/def/../../..", ".."},
	CleanTest{"/abc/def/../../..", "/"},
	CleanTest{"abc/def/../../../ghi/jkl/../../../mno", "../../mno"},

	// Combinations
	CleanTest{"abc/./../def", "def"},
	CleanTest{"abc//./../def", "def"},
	CleanTest{"abc/../../././../def", "../../def"},
}

func TestClean(t *testing.T) {
	for _, test := range cleantests {
		if s := Clean(test.path); s != test.clean {
			t.Errorf("Clean(%q) = %q, want %q", test.path, s, test.clean)
		}
	}
}

type SplitTest struct {
	path, dir, file string;
}

var splittests = []SplitTest{
	SplitTest{"a/b", "a/", "b"},
	SplitTest{"a/b/", "a/b/", ""},
	SplitTest{"a/", "a/", ""},
	SplitTest{"a", "", "a"},
	SplitTest{"/", "/", ""},
}

func TestSplit(t *testing.T) {
	for _, test := range splittests {
		if d, f := Split(test.path); d != test.dir || f != test.file {
			t.Errorf("Split(%q) = %q, %q, want %q, %q", test.path, d, f, test.dir, test.file)
		}
	}
}

type JoinTest struct {
	dir, file, path string;
}

var jointests = []JoinTest{
	JoinTest{"a", "b", "a/b"},
	JoinTest{"a", "", "a"},
	JoinTest{"", "b", "b"},
	JoinTest{"/", "a", "/a"},
	JoinTest{"/", "", "/"},
	JoinTest{"a/", "b", "a/b"},
	JoinTest{"a/", "", "a"},
}

func TestJoin(t *testing.T) {
	for _, test := range jointests {
		if p := Join(test.dir, test.file); p != test.path {
			t.Errorf("Join(%q, %q) = %q, want %q", test.dir, test.file, p, test.path)
		}
	}
}

type ExtTest struct {
	path, ext string;
}

var exttests = []ExtTest{
	ExtTest{"path.go", ".go"},
	ExtTest{"path.pb.go", ".go"},
	ExtTest{"a.dir/b", ""},
	ExtTest{"a.dir/b.go", ".go"},
	ExtTest{"a.dir/", ""},
}

func TestExt(t *testing.T) {
	for _, test := range exttests {
		if x := Ext(test.path); x != test.ext {
			t.Errorf("Ext(%q) = %q, want %q", test.path, x, test.ext)
		}
	}
}

type Node struct {
	name	string;
	entries	[]*Node;	// nil if the entry is a file
	mark	int;
}

var tree = &Node{
	"testdata",
	[]*Node{
		&Node{"a", nil, 0},
		&Node{"b", []*Node{}, 0},
		&Node{"c", nil, 0},
		&Node{
			"d",
			[]*Node{
				&Node{"x", nil, 0},
				&Node{"y", []*Node{}, 0},
				&Node{
					"z",
					[]*Node{
						&Node{"u", nil, 0},
						&Node{"v", nil, 0},
					},
					0,
				},
			},
			0,
		},
	},
	0,
}

func walkTree(n *Node, path string, f func(path string, n *Node)) {
	f(path, n);
	for _, e := range n.entries {
		walkTree(e, Join(path, e.name), f)
	}
}

func makeTree(t *testing.T) {
	walkTree(tree, tree.name, func(path string, n *Node) {
		if n.entries == nil {
			fd, err := os.Open(path, os.O_CREAT, 0660);
			if err != nil {
				t.Errorf("makeTree: %v", err)
			}
			fd.Close();
		} else {
			os.Mkdir(path, 0770)
		}
	})
}

func markTree(n *Node)	{ walkTree(n, "", func(path string, n *Node) { n.mark++ }) }

func checkMarks(t *testing.T) {
	walkTree(tree, tree.name, func(path string, n *Node) {
		if n.mark != 1 {
			t.Errorf("node %s mark = %d; expected 1", path, n.mark)
		}
		n.mark = 0;
	})
}

// Assumes that each node name is unique. Good enough for a test.
func mark(name string) {
	walkTree(tree, tree.name, func(path string, n *Node) {
		if n.name == name {
			n.mark++
		}
	})
}

type TestVisitor struct{}

func (v *TestVisitor) VisitDir(path string, d *os.Dir) bool {
	mark(d.Name);
	return true;
}

func (v *TestVisitor) VisitFile(path string, d *os.Dir) {
	mark(d.Name)
}

func TestWalk(t *testing.T) {
	makeTree(t);

	// 1) ignore error handling, expect none
	v := &TestVisitor{};
	Walk(tree.name, v, nil);
	checkMarks(t);

	// 2) handle errors, expect none
	errors := make(chan os.Error, 64);
	Walk(tree.name, v, errors);
	if err, ok := <-errors; ok {
		t.Errorf("no error expected, found: s", err)
	}
	checkMarks(t);

	if os.Getuid() != 0 {
		// introduce 2 errors: chmod top-level directories to 0
		os.Chmod(Join(tree.name, tree.entries[1].name), 0);
		os.Chmod(Join(tree.name, tree.entries[3].name), 0);
		// mark respective subtrees manually
		markTree(tree.entries[1]);
		markTree(tree.entries[3]);
		// correct double-marking of directory itself
		tree.entries[1].mark--;
		tree.entries[3].mark--;

		// 3) handle errors, expect two
		errors = make(chan os.Error, 64);
		os.Chmod(Join(tree.name, tree.entries[1].name), 0);
		Walk(tree.name, v, errors);
		for i := 1; i <= 2; i++ {
			if _, ok := <-errors; !ok {
				t.Errorf("%d. error expected, none found", i);
				break;
			}
		}
		if err, ok := <-errors; ok {
			t.Errorf("only two errors expected, found 3rd: %v", err)
		}
		// the inaccessible subtrees were marked manually
		checkMarks(t);
	}

	// cleanup
	os.Chmod(Join(tree.name, tree.entries[1].name), 0770);
	os.Chmod(Join(tree.name, tree.entries[3].name), 0770);
	if err := os.RemoveAll(tree.name); err != nil {
		t.Errorf("removeTree: %v", err)
	}
}
<|endoftext|>


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

// +build !privileged_tests

package ipmasq

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"testing"
	"time"

	"gopkg.in/check.v1"

	"github.com/cilium/cilium/pkg/lock"
)

func Test(t *testing.T) {
	check.TestingT(t)
}

type ipMasqMapMock struct {
	lock.RWMutex
	cidrs map[string]net.IPNet
}

func (m *ipMasqMapMock) Update(cidr net.IPNet) error {
	m.Lock()
	defer m.Unlock()

	cidrStr := cidr.String()
	if _, ok := m.cidrs[cidrStr]; ok {
		return fmt.Errorf("CIDR already exists: %s", cidrStr)
	}
	m.cidrs[cidrStr] = cidr

	return nil
}

func (m *ipMasqMapMock) Delete(cidr net.IPNet) error {
	m.Lock()
	defer m.Unlock()

	cidrStr := cidr.String()
	if _, ok := m.cidrs[cidrStr]; !ok {
		return fmt.Errorf("CIDR not found: %s", cidrStr)
	}
	delete(m.cidrs, cidrStr)

	return nil
}

func (m *ipMasqMapMock) Dump() ([]net.IPNet, error) {
	m.RLock()
	defer m.RUnlock()

	cidrs := make([]net.IPNet, 0, len(m.cidrs))
	for _, cidr := range m.cidrs {
		cidrs = append(cidrs, cidr)
	}

	return cidrs, nil
}

func (m *ipMasqMapMock) dumpToSet() map[string]struct{} {
	m.RLock()
	defer m.RUnlock()

	cidrs := make(map[string]struct{}, len(m.cidrs))
	for cidrStr := range m.cidrs {
		cidrs[cidrStr] = struct{}{}
	}

	return cidrs
}

type IPMasqTestSuite struct {
	ipMasqMap   *ipMasqMapMock
	ipMasqAgent *IPMasqAgent
	configFile  *os.File
}

var _ = check.Suite(&IPMasqTestSuite{})

func (i *IPMasqTestSuite) SetUpTest(c *check.C) {
	i.ipMasqMap = &ipMasqMapMock{cidrs: map[string]net.IPNet{}}

	configFile, err := ioutil.TempFile("", "ipmasq-test")
	c.Assert(err, check.IsNil)
	i.configFile = configFile

	agent, err := newIPMasqAgent(configFile.Name(), i.ipMasqMap)
	c.Assert(err, check.IsNil)
	i.ipMasqAgent = agent
	i.ipMasqAgent.Start()
}

func (i *IPMasqTestSuite) TearDownTest(c *check.C) {
	i.ipMasqAgent.Stop()
	os.Remove(i.configFile.Name())
}

func (i *IPMasqTestSuite) TestUpdate(c *check.C) {
	_, err := i.configFile.WriteString("nonMasqueradeCIDRs:\n- 1.1.1.1/32\n- 2.2.2.2/16")
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)

	ipnets := i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 2)
	_, ok := ipnets["1.1.1.1/32"]
	c.Assert(ok, check.Equals, true)
	_, ok = ipnets["2.2.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Write new config
	_, err = i.configFile.Seek(0, 0)
	c.Assert(err, check.IsNil)
	_, err = i.configFile.WriteString("nonMasqueradeCIDRs:\n- 8.8.0.0/16\n- 2.2.2.2/16")
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)

	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 2)
	_, ok = ipnets["8.8.0.0/16"]
	c.Assert(ok, check.Equals, true)
	_, ok = ipnets["2.2.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Write new config in JSON
	_, err = i.configFile.Seek(0, 0)
	c.Assert(err, check.IsNil)
	_, err = i.configFile.WriteString(`{"nonMasqueradeCIDRs": ["8.8.0.0/16", "1.1.2.3/16"]}`)
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)

	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 2)
	_, ok = ipnets["8.8.0.0/16"]
	c.Assert(ok, check.Equals, true)
	_, ok = ipnets["1.1.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Delete file, should remove the CIDRs
	err = os.Remove(i.configFile.Name())
	c.Assert(err, check.IsNil)
	err = i.configFile.Close()
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)
	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 0)
}

func (i *IPMasqTestSuite) TestRestore(c *check.C) {
	// Stop ip-masq-agent goroutine (we can't use i.ipMasqMap.Stop(), as it stops
	// the watcher)
	close(i.ipMasqAgent.stop)

	_, cidr, _ := net.ParseCIDR("3.3.3.0/24")
	i.ipMasqMap.cidrs[cidr.String()] = *cidr
	_, cidr, _ = net.ParseCIDR("4.4.0.0/16")
	i.ipMasqMap.cidrs[cidr.String()] = *cidr

	_, err := i.configFile.WriteString("nonMasqueradeCIDRs:\n- 4.4.0.0/16")
	c.Assert(err, check.IsNil)

	i.ipMasqAgent.Start()
	time.Sleep(300 * time.Millisecond)

	ipnets := i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 1)
	_, ok := ipnets["4.4.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Now stop the goroutine, and also remove the maps. It should bootstrap from
	// the config
	close(i.ipMasqAgent.stop)
	i.ipMasqMap = &ipMasqMapMock{cidrs: map[string]net.IPNet{}}
	i.ipMasqAgent.ipMasqMap = i.ipMasqMap
	_, err = i.configFile.Seek(0, 0)
	c.Assert(err, check.IsNil)
	_, err = i.configFile.WriteString("nonMasqueradeCIDRs:\n- 3.3.0.0/16")
	c.Assert(err, check.IsNil)
	i.ipMasqAgent.Start()

	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 1)
	_, ok = ipnets["3.3.0.0/16"]
	c.Assert(ok, check.Equals, true)
}
<commit_msg>ipmasq: Stop and wait until goroutine is finished in unit tests<commit_after>// Copyright 2020 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build !privileged_tests

package ipmasq

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"testing"
	"time"

	"gopkg.in/check.v1"

	"github.com/cilium/cilium/pkg/lock"
)

func Test(t *testing.T) {
	check.TestingT(t)
}

type ipMasqMapMock struct {
	lock.RWMutex
	cidrs map[string]net.IPNet
}

func (m *ipMasqMapMock) Update(cidr net.IPNet) error {
	m.Lock()
	defer m.Unlock()

	cidrStr := cidr.String()
	if _, ok := m.cidrs[cidrStr]; ok {
		return fmt.Errorf("CIDR already exists: %s", cidrStr)
	}
	m.cidrs[cidrStr] = cidr

	return nil
}

func (m *ipMasqMapMock) Delete(cidr net.IPNet) error {
	m.Lock()
	defer m.Unlock()

	cidrStr := cidr.String()
	if _, ok := m.cidrs[cidrStr]; !ok {
		return fmt.Errorf("CIDR not found: %s", cidrStr)
	}
	delete(m.cidrs, cidrStr)

	return nil
}

func (m *ipMasqMapMock) Dump() ([]net.IPNet, error) {
	m.RLock()
	defer m.RUnlock()

	cidrs := make([]net.IPNet, 0, len(m.cidrs))
	for _, cidr := range m.cidrs {
		cidrs = append(cidrs, cidr)
	}

	return cidrs, nil
}

func (m *ipMasqMapMock) dumpToSet() map[string]struct{} {
	m.RLock()
	defer m.RUnlock()

	cidrs := make(map[string]struct{}, len(m.cidrs))
	for cidrStr := range m.cidrs {
		cidrs[cidrStr] = struct{}{}
	}

	return cidrs
}

type IPMasqTestSuite struct {
	ipMasqMap   *ipMasqMapMock
	ipMasqAgent *IPMasqAgent
	configFile  *os.File
}

var _ = check.Suite(&IPMasqTestSuite{})

func (i *IPMasqTestSuite) SetUpTest(c *check.C) {
	i.ipMasqMap = &ipMasqMapMock{cidrs: map[string]net.IPNet{}}

	configFile, err := ioutil.TempFile("", "ipmasq-test")
	c.Assert(err, check.IsNil)
	i.configFile = configFile

	agent, err := newIPMasqAgent(configFile.Name(), i.ipMasqMap)
	c.Assert(err, check.IsNil)
	i.ipMasqAgent = agent
	i.ipMasqAgent.Start()
}

func (i *IPMasqTestSuite) TearDownTest(c *check.C) {
	i.ipMasqAgent.Stop()
	os.Remove(i.configFile.Name())
}

func (i *IPMasqTestSuite) TestUpdate(c *check.C) {
	_, err := i.configFile.WriteString("nonMasqueradeCIDRs:\n- 1.1.1.1/32\n- 2.2.2.2/16")
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)

	ipnets := i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 2)
	_, ok := ipnets["1.1.1.1/32"]
	c.Assert(ok, check.Equals, true)
	_, ok = ipnets["2.2.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Write new config
	_, err = i.configFile.Seek(0, 0)
	c.Assert(err, check.IsNil)
	_, err = i.configFile.WriteString("nonMasqueradeCIDRs:\n- 8.8.0.0/16\n- 2.2.2.2/16")
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)

	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 2)
	_, ok = ipnets["8.8.0.0/16"]
	c.Assert(ok, check.Equals, true)
	_, ok = ipnets["2.2.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Write new config in JSON
	_, err = i.configFile.Seek(0, 0)
	c.Assert(err, check.IsNil)
	_, err = i.configFile.WriteString(`{"nonMasqueradeCIDRs": ["8.8.0.0/16", "1.1.2.3/16"]}`)
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)

	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 2)
	_, ok = ipnets["8.8.0.0/16"]
	c.Assert(ok, check.Equals, true)
	_, ok = ipnets["1.1.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Delete file, should remove the CIDRs
	err = os.Remove(i.configFile.Name())
	c.Assert(err, check.IsNil)
	err = i.configFile.Close()
	c.Assert(err, check.IsNil)
	time.Sleep(300 * time.Millisecond)
	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 0)
}

func (i *IPMasqTestSuite) TestRestore(c *check.C) {
	// Check that stale entry is removed from the map after restore
	i.ipMasqAgent.Stop()

	_, cidr, _ := net.ParseCIDR("3.3.3.0/24")
	i.ipMasqMap.cidrs[cidr.String()] = *cidr
	_, cidr, _ = net.ParseCIDR("4.4.0.0/16")
	i.ipMasqMap.cidrs[cidr.String()] = *cidr

	_, err := i.configFile.WriteString("nonMasqueradeCIDRs:\n- 4.4.0.0/16")
	c.Assert(err, check.IsNil)

	i.ipMasqAgent, err = newIPMasqAgent(i.configFile.Name(), i.ipMasqMap)
	c.Assert(err, check.IsNil)
	i.ipMasqAgent.Start()
	time.Sleep(300 * time.Millisecond)

	ipnets := i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 1)
	_, ok := ipnets["4.4.0.0/16"]
	c.Assert(ok, check.Equals, true)

	// Now stop the goroutine, and also remove the maps. It should bootstrap from
	// the config
	i.ipMasqAgent.Stop()
	i.ipMasqMap = &ipMasqMapMock{cidrs: map[string]net.IPNet{}}
	i.ipMasqAgent.ipMasqMap = i.ipMasqMap
	_, err = i.configFile.Seek(0, 0)
	c.Assert(err, check.IsNil)
	_, err = i.configFile.WriteString("nonMasqueradeCIDRs:\n- 3.3.0.0/16")
	c.Assert(err, check.IsNil)
	i.ipMasqAgent, err = newIPMasqAgent(i.configFile.Name(), i.ipMasqMap)
	c.Assert(err, check.IsNil)
	i.ipMasqAgent.Start()

	ipnets = i.ipMasqMap.dumpToSet()
	c.Assert(len(ipnets), check.Equals, 1)
	_, ok = ipnets["3.3.0.0/16"]
	c.Assert(ok, check.Equals, true)
}
<|endoftext|>


<commit_before>package watcher

import (
	"fmt"
	"launchpad.net/gozk/zookeeper"
	"launchpad.net/tomb"
)

// ContentChange holds information on the existence
// and contents of a node. Content will be empty when the
// node does not exist.
type ContentChange struct {
	Exists  bool
	Content string
}

// ContentWatcher observes a ZooKeeper node and delivers a
// notification when a content change is detected.
type ContentWatcher struct {
	zk         *zookeeper.Conn
	path       string
	tomb       tomb.Tomb
	changeChan chan ContentChange
	content    ContentChange
}

// NewContentWatcher creates a ContentWatcher observing
// the ZooKeeper node at watchedPath.
func NewContentWatcher(zk *zookeeper.Conn, watchedPath string) *ContentWatcher {
	w := &ContentWatcher{
		zk:         zk,
		path:       watchedPath,
		changeChan: make(chan ContentChange),
	}
	go w.loop()
	return w
}

// Changes returns a channel that will receive the new node
// content when a change is detected. Note that multiple
// changes may be observed as a single event in the channel.
func (w *ContentWatcher) Changes() <-chan ContentChange {
	return w.changeChan
}

// Dying returns a channel that is closed when the
// watcher has stopped or is about to stop.
func (w *ContentWatcher) Dying() <-chan struct{} {
	return w.tomb.Dying()
}

// Stop stops the watch and returns any error encountered
// while watching. This method should always be called before
// discarding the watcher.
func (w *ContentWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

// loop is the backend for watching.
func (w *ContentWatcher) loop() {
	defer w.tomb.Done()
	defer close(w.changeChan)

	watch, err := w.update(zookeeper.EVENT_CHANGED)
	if err != nil {
		w.tomb.Kill(err)
		return
	}

	for {
		select {
		case <-w.tomb.Dying():
			return
		case evt := <-watch:
			if !evt.Ok() {
				w.tomb.Killf("watcher: critical session event: %v", evt)
				return
			}
			watch, err = w.update(evt.Type)
			if err != nil {
				w.tomb.Kill(err)
				return
			}
		}
	}
}

// update retrieves the node content and emits it to the change
// channel if it has changed. It returns the next watch.
func (w *ContentWatcher) update(lastEventType int) (nextWatch <-chan zookeeper.Event, err error) {
	var content string
	var stat *zookeeper.Stat
	var watch <-chan zookeeper.Event
	var exists bool
	// Repeat until we have a valid watch or an error.
	eventType := lastEventType
	for {
		if eventType != zookeeper.EVENT_DELETED {
			content, stat, watch, err = w.zk.GetW(w.path)
			if err == nil {
				// Node exists, so leave the loop.
				exists = true
				break
			}
		}
		if zookeeper.IsError(err, zookeeper.ZNONODE) || eventType == zookeeper.EVENT_DELETED {
			// Need a new watch to receive a signal when the node is created.
			stat, watch, err = w.zk.ExistsW(w.path)
			if stat != nil {
				// Node has been created just before ExistsW(),
				// so call GetW() with new loop run again.
				eventType = zookeeper.EVENT_CREATED
				continue
			}
			if err == nil {
				// Got a valid watch, so leave loop.
				exists = false
				break
			}
		}
		// Any other error during GetW() or ExistsW().
		return nil, fmt.Errorf("watcher: can't get content of node %q: %v", w.path, err)
	}
	if exists {
		// Check if already exists and content has changed.
		if w.content.Exists && content == w.content.Content {
			return watch, nil
		}
		w.content.Exists = true
		w.content.Content = content
	} else {
		// Check if not yet exists..
		if !w.content.Exists {
			return watch, nil
		}
		w.content.Exists = false
		w.content.Content = ""
	}
	select {
	case <-w.tomb.Dying():
		return nil, tomb.ErrDying
	case w.changeChan <- w.content:
	}
	return watch, nil
}

// ChildrenChange contains information about
// children that have been created or deleted.
type ChildrenChange struct {
	Added   []string
	Deleted []string
}

// ChildrenWatcher observes a ZooKeeper node and delivers a
// notification when child nodes are added or removed.
type ChildrenWatcher struct {
	zk         *zookeeper.Conn
	path       string
	tomb       tomb.Tomb
	changeChan chan ChildrenChange
	children   map[string]bool
}

// NewChildrenWatcher creates a ChildrenWatcher observing
// the ZooKeeper node at watchedPath.
func NewChildrenWatcher(zk *zookeeper.Conn, watchedPath string) *ChildrenWatcher {
	w := &ChildrenWatcher{
		zk:         zk,
		path:       watchedPath,
		changeChan: make(chan ChildrenChange),
		children:   make(map[string]bool),
	}
	go w.loop()
	return w
}

// Changes returns a channel that will receive the changes
// performed to the set of children of the watched node.
// Note that multiple changes may be observed as a single
// event in the channel.
func (w *ChildrenWatcher) Changes() <-chan ChildrenChange {
	return w.changeChan
}

// Dying returns a channel that is closed when the
// watcher has stopped or is about to stop.
func (w *ChildrenWatcher) Dying() <-chan struct{} {
	return w.tomb.Dying()
}

// Stop stops the watch and returns any error encountered
// while watching. This method should always be called before
// discarding the watcher.
func (w *ChildrenWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

// loop is the backend for watching.
func (w *ChildrenWatcher) loop() {
	defer w.tomb.Done()
	defer close(w.changeChan)

	watch, err := w.update(zookeeper.EVENT_CHILD)
	if err != nil {
		w.tomb.Kill(err)
		return
	}

	for {
		select {
		case <-w.tomb.Dying():
			return
		case evt := <-watch:
			if !evt.Ok() {
				w.tomb.Killf("watcher: critical session event: %v", evt)
				return
			}
			watch, err = w.update(evt.Type)
			if err != nil {
				w.tomb.Kill(err)
				return
			}
		}
	}
}

// update retrieves the node children and emits the added or deleted children to
// the change channel if it has changed. It returns the next watch.
func (w *ChildrenWatcher) update(eventType int) (nextWatch <-chan zookeeper.Event, err error) {
	if eventType == zookeeper.EVENT_DELETED {
		return nil, fmt.Errorf("watcher: node %q has been deleted", w.path)
	}
	retrievedChildren, _, watch, err := w.zk.ChildrenW(w.path)
	if err != nil {
		return nil, fmt.Errorf("watcher: can't get children of node %q: %v", w.path, err)
	}
	children := make(map[string]bool)
	for _, child := range retrievedChildren {
		children[child] = true
	}
	var change ChildrenChange
	for child, _ := range w.children {
		if !children[child] {
			change.Deleted = append(change.Deleted, child)
			delete(w.children, child)
		}
	}
	for child, _ := range children {
		if !w.children[child] {
			change.Added = append(change.Added, child)
			w.children[child] = true
		}
	}
	if len(change.Deleted) == 0 && len(change.Added) == 0 {
		return watch, nil
	}
	select {
	case <-w.tomb.Dying():
		return nil, tomb.ErrDying
	case w.changeChan <- change:
	}
	return watch, nil
}
<commit_msg>Optimized ContentWatcher update method after hints of Gustavo.<commit_after>package watcher

import (
	"fmt"
	"launchpad.net/gozk/zookeeper"
	"launchpad.net/tomb"
)

// ContentChange holds information on the existence
// and contents of a node. Content will be empty when the
// node does not exist.
type ContentChange struct {
	Exists  bool
	Content string
}

// ContentWatcher observes a ZooKeeper node and delivers a
// notification when a content change is detected.
type ContentWatcher struct {
	zk         *zookeeper.Conn
	path       string
	tomb       tomb.Tomb
	changeChan chan ContentChange
	content    ContentChange
}

// NewContentWatcher creates a ContentWatcher observing
// the ZooKeeper node at watchedPath.
func NewContentWatcher(zk *zookeeper.Conn, watchedPath string) *ContentWatcher {
	w := &ContentWatcher{
		zk:         zk,
		path:       watchedPath,
		changeChan: make(chan ContentChange),
	}
	go w.loop()
	return w
}

// Changes returns a channel that will receive the new node
// content when a change is detected. Note that multiple
// changes may be observed as a single event in the channel.
func (w *ContentWatcher) Changes() <-chan ContentChange {
	return w.changeChan
}

// Dying returns a channel that is closed when the
// watcher has stopped or is about to stop.
func (w *ContentWatcher) Dying() <-chan struct{} {
	return w.tomb.Dying()
}

// Stop stops the watch and returns any error encountered
// while watching. This method should always be called before
// discarding the watcher.
func (w *ContentWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

// loop is the backend for watching.
func (w *ContentWatcher) loop() {
	defer w.tomb.Done()
	defer close(w.changeChan)

	watch, err := w.update()
	if err != nil {
		w.tomb.Kill(err)
		return
	}

	for {
		select {
		case <-w.tomb.Dying():
			return
		case evt := <-watch:
			if !evt.Ok() {
				w.tomb.Killf("watcher: critical session event: %v", evt)
				return
			}
			watch, err = w.update()
			if err != nil {
				w.tomb.Kill(err)
				return
			}
		}
	}
}

// update retrieves the node content and emits it as well as an existence
// flag to the change channel if it has changed. It returns the next watch.
func (w *ContentWatcher) update() (nextWatch <-chan zookeeper.Event, err error) {
	var content string
	var stat *zookeeper.Stat
	var exists bool
	// Repeat until we have a valid watch or an error.
	for {
		content, stat, nextWatch, err = w.zk.GetW(w.path)
		if err == nil {
			// Node exists, so leave the loop.
			exists = true
			break
		}
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			// Need a new watch to receive a signal when the node is created.
			stat, nextWatch, err = w.zk.ExistsW(w.path)
			if stat != nil {
				// Node has been created just before ExistsW(),
				// so call GetW() with new loop run again.
				continue
			}
			if err == nil {
				// Got a valid watch, so leave loop.
				exists = false
				break
			}
		}
		// Any other error during GetW() or ExistsW().
		return nil, fmt.Errorf("watcher: can't get content of node %q: %v", w.path, err)
	}
	if exists {
		if w.content.Exists && content == w.content.Content {
			return nextWatch, nil
		}
		w.content.Exists = true
		w.content.Content = content
	} else {
		if !w.content.Exists {
			return nextWatch, nil
		}
		w.content.Exists = false
		w.content.Content = ""
	}
	select {
	case <-w.tomb.Dying():
		return nil, tomb.ErrDying
	case w.changeChan <- w.content:
	}
	return nextWatch, nil
}

// ChildrenChange contains information about
// children that have been created or deleted.
type ChildrenChange struct {
	Added   []string
	Deleted []string
}

// ChildrenWatcher observes a ZooKeeper node and delivers a
// notification when child nodes are added or removed.
type ChildrenWatcher struct {
	zk         *zookeeper.Conn
	path       string
	tomb       tomb.Tomb
	changeChan chan ChildrenChange
	children   map[string]bool
}

// NewChildrenWatcher creates a ChildrenWatcher observing
// the ZooKeeper node at watchedPath.
func NewChildrenWatcher(zk *zookeeper.Conn, watchedPath string) *ChildrenWatcher {
	w := &ChildrenWatcher{
		zk:         zk,
		path:       watchedPath,
		changeChan: make(chan ChildrenChange),
		children:   make(map[string]bool),
	}
	go w.loop()
	return w
}

// Changes returns a channel that will receive the changes
// performed to the set of children of the watched node.
// Note that multiple changes may be observed as a single
// event in the channel.
func (w *ChildrenWatcher) Changes() <-chan ChildrenChange {
	return w.changeChan
}

// Dying returns a channel that is closed when the
// watcher has stopped or is about to stop.
func (w *ChildrenWatcher) Dying() <-chan struct{} {
	return w.tomb.Dying()
}

// Stop stops the watch and returns any error encountered
// while watching. This method should always be called before
// discarding the watcher.
func (w *ChildrenWatcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

// loop is the backend for watching.
func (w *ChildrenWatcher) loop() {
	defer w.tomb.Done()
	defer close(w.changeChan)

	watch, err := w.update(zookeeper.EVENT_CHILD)
	if err != nil {
		w.tomb.Kill(err)
		return
	}

	for {
		select {
		case <-w.tomb.Dying():
			return
		case evt := <-watch:
			if !evt.Ok() {
				w.tomb.Killf("watcher: critical session event: %v", evt)
				return
			}
			watch, err = w.update(evt.Type)
			if err != nil {
				w.tomb.Kill(err)
				return
			}
		}
	}
}

// update retrieves the node children and emits the added or deleted children to
// the change channel if it has changed. It returns the next watch.
func (w *ChildrenWatcher) update(eventType int) (nextWatch <-chan zookeeper.Event, err error) {
	if eventType == zookeeper.EVENT_DELETED {
		return nil, fmt.Errorf("watcher: node %q has been deleted", w.path)
	}
	retrievedChildren, _, watch, err := w.zk.ChildrenW(w.path)
	if err != nil {
		return nil, fmt.Errorf("watcher: can't get children of node %q: %v", w.path, err)
	}
	children := make(map[string]bool)
	for _, child := range retrievedChildren {
		children[child] = true
	}
	var change ChildrenChange
	for child, _ := range w.children {
		if !children[child] {
			change.Deleted = append(change.Deleted, child)
			delete(w.children, child)
		}
	}
	for child, _ := range children {
		if !w.children[child] {
			change.Added = append(change.Added, child)
			w.children[child] = true
		}
	}
	if len(change.Deleted) == 0 && len(change.Added) == 0 {
		return watch, nil
	}
	select {
	case <-w.tomb.Dying():
		return nil, tomb.ErrDying
	case w.changeChan <- change:
	}
	return watch, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package etcd3

import (
	"errors"
	"fmt"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/integration"
	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/testapi"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/storage"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/watch"
)

func TestWatch(t *testing.T) {
	testWatch(t, false)
}

func TestWatchList(t *testing.T) {
	testWatch(t, true)
}

// It tests that
// - first occurrence of objects should notify Add event
// - update should trigger Modified event
// - update that gets filtered should trigger Deleted event
func testWatch(t *testing.T, recursive bool) {
	podFoo := &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}}
	podBar := &api.Pod{ObjectMeta: api.ObjectMeta{Name: "bar"}}

	tests := []struct {
		key        string
		pred       storage.SelectionPredicate
		watchTests []*testWatchStruct
	}{{ // create a key
		key:        "/somekey-1",
		watchTests: []*testWatchStruct{{podFoo, true, watch.Added}},
		pred:       storage.Everything,
	}, { // create a key but obj gets filtered. Then update it with unfiltered obj
		key:        "/somekey-3",
		watchTests: []*testWatchStruct{{podFoo, false, ""}, {podBar, true, watch.Added}},
		pred: storage.SelectionPredicate{
			Label: labels.Everything(),
			Field: fields.ParseSelectorOrDie("metadata.name=bar"),
			GetAttrs: func(obj runtime.Object) (labels.Set, fields.Set, error) {
				pod := obj.(*api.Pod)
				return nil, fields.Set{"metadata.name": pod.Name}, nil
			},
		},
	}, { // update
		key:        "/somekey-4",
		watchTests: []*testWatchStruct{{podFoo, true, watch.Added}, {podBar, true, watch.Modified}},
		pred:       storage.Everything,
	}, { // delete because of being filtered
		key:        "/somekey-5",
		watchTests: []*testWatchStruct{{podFoo, true, watch.Added}, {podBar, true, watch.Deleted}},
		pred: storage.SelectionPredicate{
			Label: labels.Everything(),
			Field: fields.ParseSelectorOrDie("metadata.name!=bar"),
			GetAttrs: func(obj runtime.Object) (labels.Set, fields.Set, error) {
				pod := obj.(*api.Pod)
				return nil, fields.Set{"metadata.name": pod.Name}, nil
			},
		},
	}}
	for i, tt := range tests {
		ctx, store, cluster := testSetup(t)
		w, err := store.watch(ctx, tt.key, "0", storage.SimpleFilter(tt.pred), recursive)
		if err != nil {
			t.Fatalf("Watch failed: %v", err)
		}
		var prevObj *api.Pod
		for _, watchTest := range tt.watchTests {
			out := &api.Pod{}
			key := tt.key
			if recursive {
				key = key + "/item"
			}
			err := store.GuaranteedUpdate(ctx, key, out, true, nil, storage.SimpleUpdate(
				func(runtime.Object) (runtime.Object, error) {
					return watchTest.obj, nil
				}))
			if err != nil {
				t.Fatalf("GuaranteedUpdate failed: %v", err)
			}
			if watchTest.expectEvent {
				expectObj := out
				if watchTest.watchType == watch.Deleted {
					expectObj = prevObj
					expectObj.ResourceVersion = out.ResourceVersion
				}
				testCheckResult(t, i, watchTest.watchType, w, expectObj)
			}
			prevObj = out
		}
		w.Stop()
		testCheckStop(t, i, w)
		cluster.Terminate(t)
	}
}

func TestDeleteTriggerWatch(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})
	w, err := store.Watch(ctx, key, storedObj.ResourceVersion, storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	if err := store.Delete(ctx, key, &api.Pod{}, nil); err != nil {
		t.Fatalf("Delete failed: %v", err)
	}
	testCheckEventType(t, watch.Deleted, w)
}

// TestWatchFromZero tests that
// - watch from 0 should sync up and grab the object added before
// - watch from non-0 should just watch changes after given version
func TestWatchFromZero(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})

	w, err := store.Watch(ctx, key, "0", storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	testCheckResult(t, 0, watch.Added, w, storedObj)
}

// TestWatchFromNoneZero tests that
// - watch from non-0 should just watch changes after given version
func TestWatchFromNoneZero(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})

	w, err := store.Watch(ctx, key, storedObj.ResourceVersion, storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	out := &api.Pod{}
	store.GuaranteedUpdate(ctx, key, out, true, nil, storage.SimpleUpdate(
		func(runtime.Object) (runtime.Object, error) {
			return &api.Pod{ObjectMeta: api.ObjectMeta{Name: "bar"}}, err
		}))
	testCheckResult(t, 0, watch.Modified, w, out)
}

func TestWatchError(t *testing.T) {
	cluster := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer cluster.Terminate(t)
	invalidStore := newStore(cluster.RandClient(), &testCodec{testapi.Default.Codec()}, "")
	ctx := context.Background()
	w, err := invalidStore.Watch(ctx, "/abc", "0", storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	validStore := newStore(cluster.RandClient(), testapi.Default.Codec(), "")
	validStore.GuaranteedUpdate(ctx, "/abc", &api.Pod{}, true, nil, storage.SimpleUpdate(
		func(runtime.Object) (runtime.Object, error) {
			return &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}}, nil
		}))
	testCheckEventType(t, watch.Error, w)
}

func TestWatchContextCancel(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	canceledCtx, cancel := context.WithCancel(ctx)
	cancel()
	// When we watch with a canceled context, we should detect that it's context canceled.
	// We won't take it as error and also close the watcher.
	w, err := store.watcher.Watch(canceledCtx, "/abc", 0, false, storage.SimpleFilter(storage.Everything))
	if err != nil {
		t.Fatal(err)
	}

	select {
	case _, ok := <-w.ResultChan():
		if ok {
			t.Error("ResultChan() should be closed")
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("timeout after %v", wait.ForeverTestTimeout)
	}
}

func TestWatchErrResultNotBlockAfterCancel(t *testing.T) {
	origCtx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	ctx, cancel := context.WithCancel(origCtx)
	w := store.watcher.createWatchChan(ctx, "/abc", 0, false, storage.SimpleFilter(storage.Everything))
	// make resutlChan and errChan blocking to ensure ordering.
	w.resultChan = make(chan watch.Event)
	w.errChan = make(chan error)
	// The event flow goes like:
	// - first we send an error, it should block on resultChan.
	// - Then we cancel ctx. The blocking on resultChan should be freed up
	//   and run() goroutine should return.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		w.run()
		wg.Done()
	}()
	w.errChan <- fmt.Errorf("some error")
	cancel()
	wg.Wait()
}

func TestWatchDeleteEventObjectHaveLatestRV(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})

	w, err := store.Watch(ctx, key, storedObj.ResourceVersion, storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	etcdW := cluster.RandClient().Watch(ctx, "/", clientv3.WithPrefix())

	if err := store.Delete(ctx, key, &api.Pod{}, &storage.Preconditions{}); err != nil {
		t.Fatalf("Delete failed: %v", err)
	}

	e := <-w.ResultChan()
	watchedDeleteObj := e.Object.(*api.Pod)
	var wres clientv3.WatchResponse
	wres = <-etcdW

	watchedDeleteRev, err := storage.ParseWatchResourceVersion(watchedDeleteObj.ResourceVersion)
	if err != nil {
		t.Fatalf("ParseWatchResourceVersion failed: %v", err)
	}
	if int64(watchedDeleteRev) != wres.Events[0].Kv.ModRevision {
		t.Errorf("Object from delete event have version: %v, should be the same as etcd delete's mod rev: %d",
			watchedDeleteRev, wres.Events[0].Kv.ModRevision)
	}
}

type testWatchStruct struct {
	obj         *api.Pod
	expectEvent bool
	watchType   watch.EventType
}

type testCodec struct {
	runtime.Codec
}

func (c *testCodec) Decode(data []byte, defaults *unversioned.GroupVersionKind, into runtime.Object) (runtime.Object, *unversioned.GroupVersionKind, error) {
	return nil, nil, errors.New("Expected decoding failure")
}

func testCheckEventType(t *testing.T, expectEventType watch.EventType, w watch.Interface) {
	select {
	case res := <-w.ResultChan():
		if res.Type != expectEventType {
			t.Errorf("event type want=%v, get=%v", expectEventType, res.Type)
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("time out after waiting %v on ResultChan", wait.ForeverTestTimeout)
	}
}

func testCheckResult(t *testing.T, i int, expectEventType watch.EventType, w watch.Interface, expectObj *api.Pod) {
	select {
	case res := <-w.ResultChan():
		if res.Type != expectEventType {
			t.Errorf("#%d: event type want=%v, get=%v", i, expectEventType, res.Type)
			return
		}
		if !reflect.DeepEqual(expectObj, res.Object) {
			t.Errorf("#%d: obj want=\n%#v\nget=\n%#v", i, expectObj, res.Object)
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("#%d: time out after waiting %v on ResultChan", i, wait.ForeverTestTimeout)
	}
}

func testCheckStop(t *testing.T, i int, w watch.Interface) {
	select {
	case e, ok := <-w.ResultChan():
		if ok {
			var obj string
			switch e.Object.(type) {
			case *api.Pod:
				obj = e.Object.(*api.Pod).Name
			case *unversioned.Status:
				obj = e.Object.(*unversioned.Status).Message
			}
			t.Errorf("#%d: ResultChan should have been closed. Event: %s. Object: %s", i, e.Type, obj)
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("#%d: time out after waiting 1s on ResultChan", i)
	}
}
<commit_msg>revert #32012 (Unshare cluster in each test suite)<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package etcd3

import (
	"errors"
	"fmt"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/integration"
	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/testapi"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/storage"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/watch"
)

func TestWatch(t *testing.T) {
	testWatch(t, false)
}

func TestWatchList(t *testing.T) {
	testWatch(t, true)
}

// It tests that
// - first occurrence of objects should notify Add event
// - update should trigger Modified event
// - update that gets filtered should trigger Deleted event
func testWatch(t *testing.T, recursive bool) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	podFoo := &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}}
	podBar := &api.Pod{ObjectMeta: api.ObjectMeta{Name: "bar"}}

	tests := []struct {
		key        string
		pred       storage.SelectionPredicate
		watchTests []*testWatchStruct
	}{{ // create a key
		key:        "/somekey-1",
		watchTests: []*testWatchStruct{{podFoo, true, watch.Added}},
		pred:       storage.Everything,
	}, { // create a key but obj gets filtered. Then update it with unfiltered obj
		key:        "/somekey-3",
		watchTests: []*testWatchStruct{{podFoo, false, ""}, {podBar, true, watch.Added}},
		pred: storage.SelectionPredicate{
			Label: labels.Everything(),
			Field: fields.ParseSelectorOrDie("metadata.name=bar"),
			GetAttrs: func(obj runtime.Object) (labels.Set, fields.Set, error) {
				pod := obj.(*api.Pod)
				return nil, fields.Set{"metadata.name": pod.Name}, nil
			},
		},
	}, { // update
		key:        "/somekey-4",
		watchTests: []*testWatchStruct{{podFoo, true, watch.Added}, {podBar, true, watch.Modified}},
		pred:       storage.Everything,
	}, { // delete because of being filtered
		key:        "/somekey-5",
		watchTests: []*testWatchStruct{{podFoo, true, watch.Added}, {podBar, true, watch.Deleted}},
		pred: storage.SelectionPredicate{
			Label: labels.Everything(),
			Field: fields.ParseSelectorOrDie("metadata.name!=bar"),
			GetAttrs: func(obj runtime.Object) (labels.Set, fields.Set, error) {
				pod := obj.(*api.Pod)
				return nil, fields.Set{"metadata.name": pod.Name}, nil
			},
		},
	}}
	for i, tt := range tests {
		w, err := store.watch(ctx, tt.key, "0", storage.SimpleFilter(tt.pred), recursive)
		if err != nil {
			t.Fatalf("Watch failed: %v", err)
		}
		var prevObj *api.Pod
		for _, watchTest := range tt.watchTests {
			out := &api.Pod{}
			key := tt.key
			if recursive {
				key = key + "/item"
			}
			err := store.GuaranteedUpdate(ctx, key, out, true, nil, storage.SimpleUpdate(
				func(runtime.Object) (runtime.Object, error) {
					return watchTest.obj, nil
				}))
			if err != nil {
				t.Fatalf("GuaranteedUpdate failed: %v", err)
			}
			if watchTest.expectEvent {
				expectObj := out
				if watchTest.watchType == watch.Deleted {
					expectObj = prevObj
					expectObj.ResourceVersion = out.ResourceVersion
				}
				testCheckResult(t, i, watchTest.watchType, w, expectObj)
			}
			prevObj = out
		}
		w.Stop()
		testCheckStop(t, i, w)
	}
}

func TestDeleteTriggerWatch(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})
	w, err := store.Watch(ctx, key, storedObj.ResourceVersion, storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	if err := store.Delete(ctx, key, &api.Pod{}, nil); err != nil {
		t.Fatalf("Delete failed: %v", err)
	}
	testCheckEventType(t, watch.Deleted, w)
}

// TestWatchFromZero tests that
// - watch from 0 should sync up and grab the object added before
// - watch from non-0 should just watch changes after given version
func TestWatchFromZero(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})

	w, err := store.Watch(ctx, key, "0", storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	testCheckResult(t, 0, watch.Added, w, storedObj)
}

// TestWatchFromNoneZero tests that
// - watch from non-0 should just watch changes after given version
func TestWatchFromNoneZero(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})

	w, err := store.Watch(ctx, key, storedObj.ResourceVersion, storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	out := &api.Pod{}
	store.GuaranteedUpdate(ctx, key, out, true, nil, storage.SimpleUpdate(
		func(runtime.Object) (runtime.Object, error) {
			return &api.Pod{ObjectMeta: api.ObjectMeta{Name: "bar"}}, err
		}))
	testCheckResult(t, 0, watch.Modified, w, out)
}

func TestWatchError(t *testing.T) {
	cluster := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 1})
	defer cluster.Terminate(t)
	invalidStore := newStore(cluster.RandClient(), &testCodec{testapi.Default.Codec()}, "")
	ctx := context.Background()
	w, err := invalidStore.Watch(ctx, "/abc", "0", storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	validStore := newStore(cluster.RandClient(), testapi.Default.Codec(), "")
	validStore.GuaranteedUpdate(ctx, "/abc", &api.Pod{}, true, nil, storage.SimpleUpdate(
		func(runtime.Object) (runtime.Object, error) {
			return &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}}, nil
		}))
	testCheckEventType(t, watch.Error, w)
}

func TestWatchContextCancel(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	canceledCtx, cancel := context.WithCancel(ctx)
	cancel()
	// When we watch with a canceled context, we should detect that it's context canceled.
	// We won't take it as error and also close the watcher.
	w, err := store.watcher.Watch(canceledCtx, "/abc", 0, false, storage.SimpleFilter(storage.Everything))
	if err != nil {
		t.Fatal(err)
	}

	select {
	case _, ok := <-w.ResultChan():
		if ok {
			t.Error("ResultChan() should be closed")
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("timeout after %v", wait.ForeverTestTimeout)
	}
}

func TestWatchErrResultNotBlockAfterCancel(t *testing.T) {
	origCtx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	ctx, cancel := context.WithCancel(origCtx)
	w := store.watcher.createWatchChan(ctx, "/abc", 0, false, storage.SimpleFilter(storage.Everything))
	// make resutlChan and errChan blocking to ensure ordering.
	w.resultChan = make(chan watch.Event)
	w.errChan = make(chan error)
	// The event flow goes like:
	// - first we send an error, it should block on resultChan.
	// - Then we cancel ctx. The blocking on resultChan should be freed up
	//   and run() goroutine should return.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		w.run()
		wg.Done()
	}()
	w.errChan <- fmt.Errorf("some error")
	cancel()
	wg.Wait()
}

func TestWatchDeleteEventObjectHaveLatestRV(t *testing.T) {
	ctx, store, cluster := testSetup(t)
	defer cluster.Terminate(t)
	key, storedObj := testPropogateStore(t, store, ctx, &api.Pod{ObjectMeta: api.ObjectMeta{Name: "foo"}})

	w, err := store.Watch(ctx, key, storedObj.ResourceVersion, storage.Everything)
	if err != nil {
		t.Fatalf("Watch failed: %v", err)
	}
	etcdW := cluster.RandClient().Watch(ctx, "/", clientv3.WithPrefix())

	if err := store.Delete(ctx, key, &api.Pod{}, &storage.Preconditions{}); err != nil {
		t.Fatalf("Delete failed: %v", err)
	}

	e := <-w.ResultChan()
	watchedDeleteObj := e.Object.(*api.Pod)
	var wres clientv3.WatchResponse
	wres = <-etcdW

	watchedDeleteRev, err := storage.ParseWatchResourceVersion(watchedDeleteObj.ResourceVersion)
	if err != nil {
		t.Fatalf("ParseWatchResourceVersion failed: %v", err)
	}
	if int64(watchedDeleteRev) != wres.Events[0].Kv.ModRevision {
		t.Errorf("Object from delete event have version: %v, should be the same as etcd delete's mod rev: %d",
			watchedDeleteRev, wres.Events[0].Kv.ModRevision)
	}
}

type testWatchStruct struct {
	obj         *api.Pod
	expectEvent bool
	watchType   watch.EventType
}

type testCodec struct {
	runtime.Codec
}

func (c *testCodec) Decode(data []byte, defaults *unversioned.GroupVersionKind, into runtime.Object) (runtime.Object, *unversioned.GroupVersionKind, error) {
	return nil, nil, errors.New("Expected decoding failure")
}

func testCheckEventType(t *testing.T, expectEventType watch.EventType, w watch.Interface) {
	select {
	case res := <-w.ResultChan():
		if res.Type != expectEventType {
			t.Errorf("event type want=%v, get=%v", expectEventType, res.Type)
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("time out after waiting %v on ResultChan", wait.ForeverTestTimeout)
	}
}

func testCheckResult(t *testing.T, i int, expectEventType watch.EventType, w watch.Interface, expectObj *api.Pod) {
	select {
	case res := <-w.ResultChan():
		if res.Type != expectEventType {
			t.Errorf("#%d: event type want=%v, get=%v", i, expectEventType, res.Type)
			return
		}
		if !reflect.DeepEqual(expectObj, res.Object) {
			t.Errorf("#%d: obj want=\n%#v\nget=\n%#v", i, expectObj, res.Object)
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("#%d: time out after waiting %v on ResultChan", i, wait.ForeverTestTimeout)
	}
}

func testCheckStop(t *testing.T, i int, w watch.Interface) {
	select {
	case e, ok := <-w.ResultChan():
		if ok {
			var obj string
			switch e.Object.(type) {
			case *api.Pod:
				obj = e.Object.(*api.Pod).Name
			case *unversioned.Status:
				obj = e.Object.(*unversioned.Status).Message
			}
			t.Errorf("#%d: ResultChan should have been closed. Event: %s. Object: %s", i, e.Type, obj)
		}
	case <-time.After(wait.ForeverTestTimeout):
		t.Errorf("#%d: time out after waiting 1s on ResultChan", i)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Mirantis

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

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

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

package libvirttools

import (
	"crypto/sha1"
	"fmt"
	"io"
	"net/url"
	"os"
	"strings"

	"github.com/golang/glog"
	libvirtxml "github.com/libvirt/libvirt-go-xml"

	"github.com/Mirantis/virtlet/pkg/imagetranslation"
	"github.com/Mirantis/virtlet/pkg/utils"
	"github.com/Mirantis/virtlet/pkg/virt"
)

type ImageTool struct {
	pool       virt.VirtStoragePool
	downloader utils.Downloader
}

type ImagePullError struct {
	message    string
	InnerError error
}

func (e ImagePullError) Error() string {
	if e.InnerError == nil {
		return e.message
	}
	return fmt.Sprintf("%s: %v", e.message, e.InnerError)
}

var _ ImageManager = &ImageTool{}

func NewImageTool(conn virt.VirtStorageConnection, downloader utils.Downloader, poolName string) (*ImageTool, error) {
	pool, err := ensureStoragePool(conn, poolName)
	if err != nil {
		return nil, err
	}
	return &ImageTool{pool: pool, downloader: downloader}, nil
}

func (i *ImageTool) ListVolumes() ([]virt.VirtStorageVolume, error) {
	return i.pool.ListAllVolumes()
}

func (i *ImageTool) ImageAsVolume(volumeName string) (virt.VirtStorageVolume, error) {
	return i.pool.LookupVolumeByName(volumeName)
}

func (i *ImageTool) fileToVolume(path, volumeName string) (virt.VirtStorageVolume, error) {
	imageSize, err := getFileSize(path)
	if err != nil {
		return nil, err
	}
	libvirtFilePath := fmt.Sprintf("/var/lib/libvirt/images/%s", volumeName)
	return i.pool.ImageToVolume(&libvirtxml.StorageVolume{
		Name:       volumeName,
		Allocation: &libvirtxml.StorageVolumeSize{Value: 0},
		Capacity:   &libvirtxml.StorageVolumeSize{Unit: "b", Value: imageSize},
		Target:     &libvirtxml.StorageVolumeTarget{Path: libvirtFilePath},
	}, path)
}

func (i *ImageTool) PullRemoteImageToVolume(imageName, volumeName string, nameTranslator imagetranslation.ImageNameTranslator) (virt.VirtStorageVolume, error) {
	imageName = stripTagFromImageName(imageName)
	endpoint := nameTranslator.Translate(imageName)
	if endpoint.Url == "" {
		endpoint = utils.Endpoint{Url: imageName}
		glog.V(1).Infof("Using URL %q without translation", imageName)
	} else {
		glog.V(1).Infof("URL %q was translated to %q", imageName, endpoint.Url)
	}

	// TODO(nhlfr): Handle AuthConfig from PullImageRequest.
	path, err := i.downloader.DownloadFile(endpoint)
	if err == nil {
		defer os.Remove(path)
		var vsv virt.VirtStorageVolume
		vsv, err = i.fileToVolume(path, volumeName)
		if err == nil {
			return vsv, nil
		}
	}
	return nil, ImagePullError{
		message:    fmt.Sprintf("error pulling image %q from %q", imageName, endpoint.Url),
		InnerError: err,
	}
}

func (i *ImageTool) RemoveImage(volumeName string) error {
	return i.pool.RemoveVolumeByName(volumeName)
}

func (i *ImageTool) GetImageVolume(imageName string) (virt.VirtStorageVolume, error) {
	imageVolumeName, err := ImageNameToVolumeName(imageName)
	if err != nil {
		return nil, err
	}

	return i.pool.LookupVolumeByName(imageVolumeName)
}

func stripTagFromImageName(imageName string) string {
	return strings.Split(imageName, ":")[0]
}

func ImageNameToVolumeName(imageName string) (string, error) {
	u, err := url.Parse(stripTagFromImageName(imageName))
	if err != nil {
		return "", err
	}

	h := sha1.New()
	io.WriteString(h, u.String())

	segments := strings.Split(u.Path, "/")

	volumeName := fmt.Sprintf("%x_%s", h.Sum(nil), segments[len(segments)-1])

	return volumeName, nil
}

func getFileSize(path string) (uint64, error) {
	fileInfo, err := os.Stat(path)
	if err != nil {
		return 0, err
	}
	return uint64(fileInfo.Size()), nil
}
<commit_msg>Fixes a bug that prevented virtlet to download images from github<commit_after>/*
Copyright 2017 Mirantis

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

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

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

package libvirttools

import (
	"crypto/sha1"
	"fmt"
	"io"
	"net/url"
	"os"
	"strings"

	"github.com/golang/glog"
	libvirtxml "github.com/libvirt/libvirt-go-xml"

	"github.com/Mirantis/virtlet/pkg/imagetranslation"
	"github.com/Mirantis/virtlet/pkg/utils"
	"github.com/Mirantis/virtlet/pkg/virt"
)

type ImageTool struct {
	pool       virt.VirtStoragePool
	downloader utils.Downloader
}

type ImagePullError struct {
	message    string
	InnerError error
}

func (e ImagePullError) Error() string {
	if e.InnerError == nil {
		return e.message
	}
	return fmt.Sprintf("%s: %v", e.message, e.InnerError)
}

var _ ImageManager = &ImageTool{}

func NewImageTool(conn virt.VirtStorageConnection, downloader utils.Downloader, poolName string) (*ImageTool, error) {
	pool, err := ensureStoragePool(conn, poolName)
	if err != nil {
		return nil, err
	}
	return &ImageTool{pool: pool, downloader: downloader}, nil
}

func (i *ImageTool) ListVolumes() ([]virt.VirtStorageVolume, error) {
	return i.pool.ListAllVolumes()
}

func (i *ImageTool) ImageAsVolume(volumeName string) (virt.VirtStorageVolume, error) {
	return i.pool.LookupVolumeByName(volumeName)
}

func (i *ImageTool) fileToVolume(path, volumeName string) (virt.VirtStorageVolume, error) {
	imageSize, err := getFileSize(path)
	if err != nil {
		return nil, err
	}
	libvirtFilePath := fmt.Sprintf("/var/lib/libvirt/images/%s", volumeName)
	return i.pool.ImageToVolume(&libvirtxml.StorageVolume{
		Name:       volumeName,
		Allocation: &libvirtxml.StorageVolumeSize{Value: 0},
		Capacity:   &libvirtxml.StorageVolumeSize{Unit: "b", Value: imageSize},
		Target:     &libvirtxml.StorageVolumeTarget{Path: libvirtFilePath},
	}, path)
}

func (i *ImageTool) PullRemoteImageToVolume(imageName, volumeName string, nameTranslator imagetranslation.ImageNameTranslator) (virt.VirtStorageVolume, error) {
	imageName = stripTagFromImageName(imageName)
	endpoint := nameTranslator.Translate(imageName)
	if endpoint.Url == "" {
		endpoint = utils.Endpoint{Url: imageName, MaxRedirects: -1}
		glog.V(1).Infof("Using URL %q without translation", imageName)
	} else {
		glog.V(1).Infof("URL %q was translated to %q", imageName, endpoint.Url)
	}

	// TODO(nhlfr): Handle AuthConfig from PullImageRequest.
	path, err := i.downloader.DownloadFile(endpoint)
	if err == nil {
		defer os.Remove(path)
		var vsv virt.VirtStorageVolume
		vsv, err = i.fileToVolume(path, volumeName)
		if err == nil {
			return vsv, nil
		}
	}
	return nil, ImagePullError{
		message:    fmt.Sprintf("error pulling image %q from %q", imageName, endpoint.Url),
		InnerError: err,
	}
}

func (i *ImageTool) RemoveImage(volumeName string) error {
	return i.pool.RemoveVolumeByName(volumeName)
}

func (i *ImageTool) GetImageVolume(imageName string) (virt.VirtStorageVolume, error) {
	imageVolumeName, err := ImageNameToVolumeName(imageName)
	if err != nil {
		return nil, err
	}

	return i.pool.LookupVolumeByName(imageVolumeName)
}

func stripTagFromImageName(imageName string) string {
	return strings.Split(imageName, ":")[0]
}

func ImageNameToVolumeName(imageName string) (string, error) {
	u, err := url.Parse(stripTagFromImageName(imageName))
	if err != nil {
		return "", err
	}

	h := sha1.New()
	io.WriteString(h, u.String())

	segments := strings.Split(u.Path, "/")

	volumeName := fmt.Sprintf("%x_%s", h.Sum(nil), segments[len(segments)-1])

	return volumeName, nil
}

func getFileSize(path string) (uint64, error) {
	fileInfo, err := os.Stat(path)
	if err != nil {
		return 0, err
	}
	return uint64(fileInfo.Size()), nil
}
<|endoftext|>


<commit_before>package truncindex

import "testing"

// Test the behavior of TruncIndex, an index for querying IDs from a non-conflicting prefix.
func TestTruncIndex(t *testing.T) {
	ids := []string{}
	index := NewTruncIndex(ids)
	// Get on an empty index
	if _, err := index.Get("foobar"); err == nil {
		t.Fatal("Get on an empty index should return an error")
	}

	// Spaces should be illegal in an id
	if err := index.Add("I have a space"); err == nil {
		t.Fatalf("Adding an id with ' ' should return an error")
	}

	id := "99b36c2c326ccc11e726eee6ee78a0baf166ef96"
	// Add an id
	if err := index.Add(id); err != nil {
		t.Fatal(err)
	}
	// Get a non-existing id
	assertIndexGet(t, index, "abracadabra", "", true)
	// Get the exact id
	assertIndexGet(t, index, id, id, false)
	// The first letter should match
	assertIndexGet(t, index, id[:1], id, false)
	// The first half should match
	assertIndexGet(t, index, id[:len(id)/2], id, false)
	// The second half should NOT match
	assertIndexGet(t, index, id[len(id)/2:], "", true)

	id2 := id[:6] + "blabla"
	// Add an id
	if err := index.Add(id2); err != nil {
		t.Fatal(err)
	}
	// Both exact IDs should work
	assertIndexGet(t, index, id, id, false)
	assertIndexGet(t, index, id2, id2, false)

	// 6 characters or less should conflict
	assertIndexGet(t, index, id[:6], "", true)
	assertIndexGet(t, index, id[:4], "", true)
	assertIndexGet(t, index, id[:1], "", true)

	// 7 characters should NOT conflict
	assertIndexGet(t, index, id[:7], id, false)
	assertIndexGet(t, index, id2[:7], id2, false)

	// Deleting a non-existing id should return an error
	if err := index.Delete("non-existing"); err == nil {
		t.Fatalf("Deleting a non-existing id should return an error")
	}

	// Deleting id2 should remove conflicts
	if err := index.Delete(id2); err != nil {
		t.Fatal(err)
	}
	// id2 should no longer work
	assertIndexGet(t, index, id2, "", true)
	assertIndexGet(t, index, id2[:7], "", true)
	assertIndexGet(t, index, id2[:11], "", true)

	// conflicts between id and id2 should be gone
	assertIndexGet(t, index, id[:6], id, false)
	assertIndexGet(t, index, id[:4], id, false)
	assertIndexGet(t, index, id[:1], id, false)

	// non-conflicting substrings should still not conflict
	assertIndexGet(t, index, id[:7], id, false)
	assertIndexGet(t, index, id[:15], id, false)
	assertIndexGet(t, index, id, id, false)
}

func assertIndexGet(t *testing.T, index *TruncIndex, input, expectedResult string, expectError bool) {
	if result, err := index.Get(input); err != nil && !expectError {
		t.Fatalf("Unexpected error getting '%s': %s", input, err)
	} else if err == nil && expectError {
		t.Fatalf("Getting '%s' should return an error", input)
	} else if result != expectedResult {
		t.Fatalf("Getting '%s' returned '%s' instead of '%s'", input, result, expectedResult)
	}
}

func BenchmarkTruncIndexAdd(b *testing.B) {
	ids := []string{"banana", "bananaa", "bananab"}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		index := NewTruncIndex([]string{})
		for _, id := range ids {
			index.Add(id)
		}
	}
}

func BenchmarkTruncIndexNew(b *testing.B) {
	ids := []string{"banana", "bananaa", "bananab"}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		NewTruncIndex(ids)
	}
}
<commit_msg>Improve truncindex benchmarks<commit_after>package truncindex

import (
	"math/rand"
	"testing"

	"github.com/dotcloud/docker/utils"
)

// Test the behavior of TruncIndex, an index for querying IDs from a non-conflicting prefix.
func TestTruncIndex(t *testing.T) {
	ids := []string{}
	index := NewTruncIndex(ids)
	// Get on an empty index
	if _, err := index.Get("foobar"); err == nil {
		t.Fatal("Get on an empty index should return an error")
	}

	// Spaces should be illegal in an id
	if err := index.Add("I have a space"); err == nil {
		t.Fatalf("Adding an id with ' ' should return an error")
	}

	id := "99b36c2c326ccc11e726eee6ee78a0baf166ef96"
	// Add an id
	if err := index.Add(id); err != nil {
		t.Fatal(err)
	}
	// Get a non-existing id
	assertIndexGet(t, index, "abracadabra", "", true)
	// Get the exact id
	assertIndexGet(t, index, id, id, false)
	// The first letter should match
	assertIndexGet(t, index, id[:1], id, false)
	// The first half should match
	assertIndexGet(t, index, id[:len(id)/2], id, false)
	// The second half should NOT match
	assertIndexGet(t, index, id[len(id)/2:], "", true)

	id2 := id[:6] + "blabla"
	// Add an id
	if err := index.Add(id2); err != nil {
		t.Fatal(err)
	}
	// Both exact IDs should work
	assertIndexGet(t, index, id, id, false)
	assertIndexGet(t, index, id2, id2, false)

	// 6 characters or less should conflict
	assertIndexGet(t, index, id[:6], "", true)
	assertIndexGet(t, index, id[:4], "", true)
	assertIndexGet(t, index, id[:1], "", true)

	// 7 characters should NOT conflict
	assertIndexGet(t, index, id[:7], id, false)
	assertIndexGet(t, index, id2[:7], id2, false)

	// Deleting a non-existing id should return an error
	if err := index.Delete("non-existing"); err == nil {
		t.Fatalf("Deleting a non-existing id should return an error")
	}

	// Deleting id2 should remove conflicts
	if err := index.Delete(id2); err != nil {
		t.Fatal(err)
	}
	// id2 should no longer work
	assertIndexGet(t, index, id2, "", true)
	assertIndexGet(t, index, id2[:7], "", true)
	assertIndexGet(t, index, id2[:11], "", true)

	// conflicts between id and id2 should be gone
	assertIndexGet(t, index, id[:6], id, false)
	assertIndexGet(t, index, id[:4], id, false)
	assertIndexGet(t, index, id[:1], id, false)

	// non-conflicting substrings should still not conflict
	assertIndexGet(t, index, id[:7], id, false)
	assertIndexGet(t, index, id[:15], id, false)
	assertIndexGet(t, index, id, id, false)
}

func assertIndexGet(t *testing.T, index *TruncIndex, input, expectedResult string, expectError bool) {
	if result, err := index.Get(input); err != nil && !expectError {
		t.Fatalf("Unexpected error getting '%s': %s", input, err)
	} else if err == nil && expectError {
		t.Fatalf("Getting '%s' should return an error", input)
	} else if result != expectedResult {
		t.Fatalf("Getting '%s' returned '%s' instead of '%s'", input, result, expectedResult)
	}
}

func BenchmarkTruncIndexAdd100(b *testing.B) {
	var testSet []string
	for i := 0; i < 100; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		index := NewTruncIndex([]string{})
		for _, id := range testSet {
			if err := index.Add(id); err != nil {
				b.Fatal(err)
			}
		}
	}
}

func BenchmarkTruncIndexAdd250(b *testing.B) {
	var testSet []string
	for i := 0; i < 250; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		index := NewTruncIndex([]string{})
		for _, id := range testSet {
			if err := index.Add(id); err != nil {
				b.Fatal(err)
			}
		}
	}
}

func BenchmarkTruncIndexAdd500(b *testing.B) {
	var testSet []string
	for i := 0; i < 500; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		index := NewTruncIndex([]string{})
		for _, id := range testSet {
			if err := index.Add(id); err != nil {
				b.Fatal(err)
			}
		}
	}
}

func BenchmarkTruncIndexGet100(b *testing.B) {
	var testSet []string
	var testKeys []string
	for i := 0; i < 100; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	index := NewTruncIndex([]string{})
	for _, id := range testSet {
		if err := index.Add(id); err != nil {
			b.Fatal(err)
		}
		l := rand.Intn(12) + 12
		testKeys = append(testKeys, id[:l])
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for _, id := range testKeys {
			if res, err := index.Get(id); err != nil {
				b.Fatal(res, err)
			}
		}
	}
}

func BenchmarkTruncIndexGet250(b *testing.B) {
	var testSet []string
	var testKeys []string
	for i := 0; i < 250; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	index := NewTruncIndex([]string{})
	for _, id := range testSet {
		if err := index.Add(id); err != nil {
			b.Fatal(err)
		}
		l := rand.Intn(12) + 12
		testKeys = append(testKeys, id[:l])
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for _, id := range testKeys {
			if res, err := index.Get(id); err != nil {
				b.Fatal(res, err)
			}
		}
	}
}

func BenchmarkTruncIndexGet500(b *testing.B) {
	var testSet []string
	var testKeys []string
	for i := 0; i < 500; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	index := NewTruncIndex([]string{})
	for _, id := range testSet {
		if err := index.Add(id); err != nil {
			b.Fatal(err)
		}
		l := rand.Intn(12) + 12
		testKeys = append(testKeys, id[:l])
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for _, id := range testKeys {
			if res, err := index.Get(id); err != nil {
				b.Fatal(res, err)
			}
		}
	}
}

func BenchmarkTruncIndexNew100(b *testing.B) {
	var testSet []string
	for i := 0; i < 100; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		NewTruncIndex(testSet)
	}
}

func BenchmarkTruncIndexNew250(b *testing.B) {
	var testSet []string
	for i := 0; i < 250; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		NewTruncIndex(testSet)
	}
}

func BenchmarkTruncIndexNew500(b *testing.B) {
	var testSet []string
	for i := 0; i < 500; i++ {
		testSet = append(testSet, utils.GenerateRandomID())
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		NewTruncIndex(testSet)
	}
}
<|endoftext|>


<commit_before>package scaler

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os"

	"fmt"

	"bytes"

	"github.com/Gamebuildr/Hal/pkg/config"
	jwt "github.com/dgrijalva/jwt-go"
)

// HTTPScaler is a scaling system to increase remote api system scale
type HTTPScaler struct {
	LoadAPIUrl    string
	AddLoadAPIUrl string
	Client        *http.Client
}

// Response is the data that is returned from the API
type Response struct {
	LoadCount int
}

// GetSystemLoad returns to load count of the running system
func (system HTTPScaler) GetSystemLoad() (int, error) {
	r, err := http.NewRequest(http.MethodPost, system.LoadAPIUrl, nil)
	if err != nil {
		return 0, err
	}
	r.Header.Add("Content-Type", "application/json")
	if err = authenticateRoute(r); err != nil {
		return 0, err
	}

	w, err := system.Client.Do(r)
	if err != nil {
		return 0, err
	}
	resp, err := ioutil.ReadAll(w.Body)
	if err != nil {
		return 0, err
	}
	jsonResp := Response{}
	if err := json.Unmarshal(resp, &jsonResp); err != nil {
		return 0, fmt.Errorf("Container Load Response: %v, Error: %v", string(resp), err.Error())
	}
	return jsonResp.LoadCount, nil
}

// AddSystemLoad will increase the systems load by one
func (system HTTPScaler) AddSystemLoad(message string) (*http.Response, error) {
	r, err := http.NewRequest(http.MethodPost, system.AddLoadAPIUrl, bytes.NewBuffer([]byte(message)))
	r.Header.Set("Content-Type", "application/json")
	if err != nil {
		return nil, err
	}
	if err = authenticateRoute(r); err != nil {
		return nil, err
	}
	w, err := system.Client.Do(r)
	if err != nil {
		return w, err
	}
	return w, nil
}

func authenticateRoute(r *http.Request) error {
	token, err := getStringToken()
	if err != nil {
		return err
	}
	bearer := "Bearer " + token
	r.Header.Add("Authorization", bearer)
	return nil
}

func getStringToken() (string, error) {
	tokenValue := os.Getenv(config.Auth0ClientSecret)
	secretKey := []byte(tokenValue)
	token := jwt.New(jwt.SigningMethodHS256)
	tokenString, err := token.SignedString(secretKey)
	if err != nil {
		return "", err
	}
	return tokenString, nil
}
<commit_msg>Change getting counts from hal to a get request<commit_after>package scaler

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"os"

	"fmt"

	"bytes"

	"github.com/Gamebuildr/Hal/pkg/config"
	jwt "github.com/dgrijalva/jwt-go"
)

// HTTPScaler is a scaling system to increase remote api system scale
type HTTPScaler struct {
	LoadAPIUrl    string
	AddLoadAPIUrl string
	Client        *http.Client
}

// Response is the data that is returned from the API
type Response struct {
	LoadCount int
}

// GetSystemLoad returns to load count of the running system
func (system HTTPScaler) GetSystemLoad() (int, error) {
	r, err := http.NewRequest(http.MethodGet, system.LoadAPIUrl, nil)
	if err != nil {
		return 0, err
	}
	r.Header.Add("Content-Type", "application/json")
	if err = authenticateRoute(r); err != nil {
		return 0, err
	}

	w, err := system.Client.Do(r)
	if err != nil {
		return 0, err
	}
	resp, err := ioutil.ReadAll(w.Body)
	if err != nil {
		return 0, err
	}
	jsonResp := Response{}
	if err := json.Unmarshal(resp, &jsonResp); err != nil {
		return 0, fmt.Errorf("Container Load Response: %v, Error: %v", string(resp), err.Error())
	}
	return jsonResp.LoadCount, nil
}

// AddSystemLoad will increase the systems load by one
func (system HTTPScaler) AddSystemLoad(message string) (*http.Response, error) {
	r, err := http.NewRequest(http.MethodPost, system.AddLoadAPIUrl, bytes.NewBuffer([]byte(message)))
	r.Header.Set("Content-Type", "application/json")
	if err != nil {
		return nil, err
	}
	if err = authenticateRoute(r); err != nil {
		return nil, err
	}
	w, err := system.Client.Do(r)
	if err != nil {
		return w, err
	}
	return w, nil
}

func authenticateRoute(r *http.Request) error {
	token, err := getStringToken()
	if err != nil {
		return err
	}
	bearer := "Bearer " + token
	r.Header.Add("Authorization", bearer)
	return nil
}

func getStringToken() (string, error) {
	tokenValue := os.Getenv(config.Auth0ClientSecret)
	secretKey := []byte(tokenValue)
	token := jwt.New(jwt.SigningMethodHS256)
	tokenString, err := token.SignedString(secretKey)
	if err != nil {
		return "", err
	}
	return tokenString, nil
}
<|endoftext|>