<commit_before>package plugin_test

import (
	"fmt"
	"strings"
	"testing"

	"github.com/neovim/go-client/nvim"
	"github.com/neovim/go-client/nvim/nvimtest"
	"github.com/neovim/go-client/nvim/plugin"
)

func TestRegister(t *testing.T) {
	t.Parallel()

	type testEval struct {
		BaseDir string `eval:"fnamemodify(getcwd(), ':t')"`
	}

	p := plugin.New(nvimtest.NewChildProcess(t))

	// SimpleHandler
	p.Handle("hello",
		func(s string) (string, error) {
			return "Hello, " + s, nil
		},
	)

	// FunctionHandler
	p.HandleFunction(
		&plugin.FunctionOptions{Name: "Hello"},
		func(args []string) (string, error) {
			return "Hello, " + strings.Join(args, " "), nil
		},
	)

	// FunctionEvalHandler
	p.HandleFunction(
		&plugin.FunctionOptions{Name: "TestEval", Eval: "*"},
		func(_ []string, eval *testEval) (string, error) {
			return eval.BaseDir, nil
		},
	)

	// CommandHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name:     "Hello",
			NArgs:    "*",
			Range:    "%",
			Addr:     "buffers",
			Complete: "buffer",
			Bang:     true,
			Register: true,
			Bar:      true,
		},
		func(n *nvim.Nvim, args []string) error {
			chunks := []nvim.TextChunk{
				{Text: `Hello`},
			}
			for _, arg := range args {
				chunks = append(chunks, nvim.TextChunk{Text: arg})
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// CommandRangeDotHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name:  "HelloRangeDot",
			Range: ".",
		},
		func(n *nvim.Nvim) error {
			chunks := []nvim.TextChunk{
				{Text: `Hello`},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// CommandCountHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name:  "HelloCount",
			Count: "0",
		},
		func(n *nvim.Nvim) error {
			chunks := []nvim.TextChunk{
				{Text: `Hello`},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// CommandEvalHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name: "HelloEval",
			Eval: "*",
		},
		func(n *nvim.Nvim, eval *testEval) error {
			chunks := []nvim.TextChunk{
				{
					Text: eval.BaseDir,
				},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	if err := p.RegisterForTests(); err != nil {
		t.Fatalf("register handlers for test: %v", err)
	}

	t.Run("SimpleHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`:echo Hello('John', 'Doe')`, true)
		if err != nil {
			t.Fatalf("exec 'echo Hello' function: %v", err)
		}

		expected := `Hello, John Doe`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("FunctionHandler", func(t *testing.T) {
		cid := p.Nvim.ChannelID()
		var result string
		if err := p.Nvim.Call(`rpcrequest`, &result, cid, `hello`, `world`); err != nil {
			t.Fatalf("call rpcrequest(%v, %v, %v, %v): %v", &result, cid, "hello", "world", err)
		}

		expected2 := `Hello, world`
		if result != expected2 {
			t.Fatalf("hello returned %q, want %q", result, expected2)
		}
	})

	t.Run("FunctionEvalHandler", func(t *testing.T) {
		var result string
		if err := p.Nvim.Eval(`TestEval()`, &result); err != nil {
			t.Fatalf("eval 'TestEval()' function: %v", err)
		}

		expected3 := `plugin`
		if result != expected3 {
			t.Fatalf("EvalTest returned %q, want %q", result, expected3)
		}
	})

	t.Run("CommandHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`Hello World`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `Helloorld`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("CommandRangeDotHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`HelloRangeDot`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `Hello`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("CommandCountHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`HelloCount`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `Hello`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("CommandEvalHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`HelloEval`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `plugin`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})
}

func TestSubscribe(t *testing.T) {
	t.Parallel()

	p := plugin.New(nvimtest.NewChildProcess(t))

	const event1 = "event1"
	eventFn1 := func(t *testing.T, v *nvim.Nvim) error {
		return v.RegisterHandler(event1, func(event ...interface{}) {
			if event[0] != int64(1) {
				t.Fatalf("expected event[0] is 1 but got %d", event[0])
			}
			if event[1] != int64(2) {
				t.Fatalf("expected event[1] is 2 but got %d", event[1])
			}
			if event[2] != int64(3) {
				t.Fatalf("expected event[2] is 3 but got %d", event[2])
			}
		})
	}
	p.Handle(event1, func() error { return eventFn1(t, p.Nvim) })

	const event2 = "event2"
	eventFn2 := func(t *testing.T, v *nvim.Nvim) error {
		return v.RegisterHandler(event1, func(event ...interface{}) {
			if event[0] != int64(4) {
				t.Fatalf("expected event[0] is 4 but got %d", event[0])
			}
			if event[1] != int64(5) {
				t.Fatalf("expected event[1] is 5 but got %d", event[1])
			}
			if event[2] != int64(6) {
				t.Fatalf("expected event[2] is 6 but got %d", event[2])
			}
		})
	}
	p.Handle(event2, func() error { return eventFn2(t, p.Nvim) })

	if err := p.RegisterForTests(); err != nil {
		t.Fatalf("register for test: %v", err)
	}

	if err := p.Nvim.Subscribe(event1); err != nil {
		t.Fatalf("subscribe(%v): %v", event1, err)
	}

	b := p.Nvim.NewBatch()
	b.Subscribe(event2)
	if err := b.Execute(); err != nil {
		t.Fatalf("batch execute: %v", err)
	}

	// warm-up
	var result int
	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q)`, event1), &result); err != nil {
		t.Fatalf("eval rpcnotify for warm-up of event1: %v", err)
	}
	if result != 1 {
		t.Fatalf("expect 1 but got %d", result)
	}

	var result2 int
	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 1, 2, 3)`, event1), &result2); err != nil {
		t.Fatalf("eval rpcnotify for event1: %v", err)
	}
	if result2 != 1 {
		t.Fatalf("expect 1 but got %d", result2)
	}

	var result3 int
	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 4, 5, 6)`, event2), &result3); err != nil {
		t.Fatalf("eval rpcnotify for event2: %v", err)
	}
	if result3 != 1 {
		t.Fatalf("expect 1 but got %d", result3)
	}

	if err := p.Nvim.Unsubscribe(event1); err != nil {
		t.Fatalf("unsubscribe event1: %v", err)
	}

	b.Unsubscribe(event2)
	if err := b.Execute(); err != nil {
		t.Fatalf("unsubscribe event2: %v", err)
	}

	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 7, 8, 9)`, event1), nil); err != nil {
		t.Fatalf("ensure rpcnotify to event1 is no-op: %v", err)
	}

	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 10, 11, 12)`, event2), nil); err != nil {
		t.Fatalf("ensure rpcnotify to event2 is no-op: %v", err)
	}
}
<commit_msg>nvim/plugin: add Autocmd handlers testcases<commit_after>package plugin_test

import (
	"fmt"
	"strings"
	"testing"

	"github.com/neovim/go-client/nvim"
	"github.com/neovim/go-client/nvim/nvimtest"
	"github.com/neovim/go-client/nvim/plugin"
)

func TestRegister(t *testing.T) {
	t.Parallel()

	type testEval struct {
		BaseDir string `eval:"fnamemodify(getcwd(), ':t')"`
	}

	p := plugin.New(nvimtest.NewChildProcess(t))

	// SimpleHandler
	p.Handle("hello",
		func(s string) (string, error) {
			return "Hello, " + s, nil
		},
	)

	// FunctionHandler
	p.HandleFunction(
		&plugin.FunctionOptions{Name: "Hello"},
		func(args []string) (string, error) {
			return "Hello, " + strings.Join(args, " "), nil
		},
	)

	// FunctionEvalHandler
	p.HandleFunction(
		&plugin.FunctionOptions{Name: "TestEval", Eval: "*"},
		func(_ []string, eval *testEval) (string, error) {
			return eval.BaseDir, nil
		},
	)

	// CommandHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name:     "Hello",
			NArgs:    "*",
			Range:    "%",
			Addr:     "buffers",
			Complete: "buffer",
			Bang:     true,
			Register: true,
			Bar:      true,
		},
		func(n *nvim.Nvim, args []string) error {
			chunks := []nvim.TextChunk{
				{Text: `Hello`},
			}
			for _, arg := range args {
				chunks = append(chunks, nvim.TextChunk{Text: arg})
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// CommandRangeDotHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name:  "HelloRangeDot",
			Range: ".",
		},
		func(n *nvim.Nvim) error {
			chunks := []nvim.TextChunk{
				{Text: `Hello`},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// CommandCountHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name:  "HelloCount",
			Count: "0",
		},
		func(n *nvim.Nvim) error {
			chunks := []nvim.TextChunk{
				{Text: `Hello`},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// CommandEvalHandler
	p.HandleCommand(
		&plugin.CommandOptions{
			Name: "HelloEval",
			Eval: "*",
		},
		func(n *nvim.Nvim, eval *testEval) error {
			chunks := []nvim.TextChunk{
				{
					Text: eval.BaseDir,
				},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// AutocmdHandler
	p.HandleAutocmd(
		&plugin.AutocmdOptions{
			Event:   "User",
			Group:   "Test",
			Pattern: "Test",
			Nested:  true,
			Once:    false,
		},
		func(n *nvim.Nvim, args []string) error {
			chunks := []nvim.TextChunk{
				{
					Text: "Hello",
				},
				{
					Text: "Autocmd",
				},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// AutocmdEvalHandler
	p.HandleAutocmd(
		&plugin.AutocmdOptions{
			Event:   "User",
			Group:   "Eval",
			Pattern: "Eval",
			Eval:    "*",
		},
		func(n *nvim.Nvim, eval *testEval) error {
			chunks := []nvim.TextChunk{
				{
					Text: eval.BaseDir,
				},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	// AutocmdOnceHandler
	p.HandleAutocmd(
		&plugin.AutocmdOptions{
			Event:   "User",
			Group:   "TestOnce",
			Pattern: "Once",
			Nested:  true,
			Once:    true,
		},
		func(n *nvim.Nvim, args []string) error {
			chunks := []nvim.TextChunk{
				{
					Text: "Hello",
				},
				{
					Text: "Autocmd",
				},
				{
					Text: "Once",
				},
			}

			return n.Echo(chunks, true, make(map[string]interface{}))
		},
	)

	if err := p.RegisterForTests(); err != nil {
		t.Fatalf("register handlers for test: %v", err)
	}

	t.Run("SimpleHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`:echo Hello('John', 'Doe')`, true)
		if err != nil {
			t.Fatalf("exec 'echo Hello' function: %v", err)
		}

		expected := `Hello, John Doe`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("FunctionHandler", func(t *testing.T) {
		cid := p.Nvim.ChannelID()
		var result string
		if err := p.Nvim.Call(`rpcrequest`, &result, cid, `hello`, `world`); err != nil {
			t.Fatalf("call rpcrequest(%v, %v, %v, %v): %v", &result, cid, "hello", "world", err)
		}

		expected2 := `Hello, world`
		if result != expected2 {
			t.Fatalf("hello returned %q, want %q", result, expected2)
		}
	})

	t.Run("FunctionEvalHandler", func(t *testing.T) {
		var result string
		if err := p.Nvim.Eval(`TestEval()`, &result); err != nil {
			t.Fatalf("eval 'TestEval()' function: %v", err)
		}

		expected3 := `plugin`
		if result != expected3 {
			t.Fatalf("EvalTest returned %q, want %q", result, expected3)
		}
	})

	t.Run("CommandHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`Hello World`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `Helloorld`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("CommandRangeDotHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`HelloRangeDot`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `Hello`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("CommandCountHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`HelloCount`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `Hello`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("CommandEvalHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`HelloEval`, true)
		if err != nil {
			t.Fatalf("exec 'Hello' command: %v", err)
		}

		expected := `plugin`
		if result != expected {
			t.Fatalf("Hello returned %q, want %q", result, expected)
		}
	})

	t.Run("AutocmdHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`doautocmd User Test`, true)
		if err != nil {
			t.Fatalf("exec 'doautocmd User Test' command: %v", err)
		}

		expected := `HelloAutocmd`
		if result != expected {
			t.Fatalf("'doautocmd User Test' returned %q, want %q", result, expected)
		}

		result2, err := p.Nvim.Exec(`doautocmd User Test`, true)
		if err != nil {
			t.Fatalf("exec 'doautocmd User Test' command: %v", err)
		}

		if result2 != expected {
			t.Fatalf("'doautocmd User Test' returned %q, want %q", result, expected)
		}
	})

	t.Run("AutocmdEvalHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`doautocmd User Eval`, true)
		if err != nil {
			t.Fatalf("exec 'doautocmd User Eval' command: %v", err)
		}

		expected := `plugin`
		if result != expected {
			t.Fatalf("'doautocmd User Eval' returned %q, want %q", result, expected)
		}
	})

	t.Run("AutocmdOnceHandler", func(t *testing.T) {
		result, err := p.Nvim.Exec(`doautocmd User Once`, true)
		if err != nil {
			t.Fatalf("exec 'doautocmd User Once' command: %v", err)
		}

		expected := `HelloAutocmdOnce`
		if result != expected {
			t.Fatalf("'doautocmd User Once' returned %q, want %q", result, expected)
		}

		result2, err := p.Nvim.Exec(`doautocmd User Once`, true)
		if err != nil {
			t.Fatalf("exec 'doautocmd User Once' command: %v", err)
		}

		if result2 == expected {
			t.Fatalf("'doautocmd User Once' returned %q, want %q", result, "")
		}
	})
}

func TestSubscribe(t *testing.T) {
	t.Parallel()

	p := plugin.New(nvimtest.NewChildProcess(t))

	const event1 = "event1"
	eventFn1 := func(t *testing.T, v *nvim.Nvim) error {
		return v.RegisterHandler(event1, func(event ...interface{}) {
			if event[0] != int64(1) {
				t.Fatalf("expected event[0] is 1 but got %d", event[0])
			}
			if event[1] != int64(2) {
				t.Fatalf("expected event[1] is 2 but got %d", event[1])
			}
			if event[2] != int64(3) {
				t.Fatalf("expected event[2] is 3 but got %d", event[2])
			}
		})
	}
	p.Handle(event1, func() error { return eventFn1(t, p.Nvim) })

	const event2 = "event2"
	eventFn2 := func(t *testing.T, v *nvim.Nvim) error {
		return v.RegisterHandler(event1, func(event ...interface{}) {
			if event[0] != int64(4) {
				t.Fatalf("expected event[0] is 4 but got %d", event[0])
			}
			if event[1] != int64(5) {
				t.Fatalf("expected event[1] is 5 but got %d", event[1])
			}
			if event[2] != int64(6) {
				t.Fatalf("expected event[2] is 6 but got %d", event[2])
			}
		})
	}
	p.Handle(event2, func() error { return eventFn2(t, p.Nvim) })

	if err := p.RegisterForTests(); err != nil {
		t.Fatalf("register for test: %v", err)
	}

	if err := p.Nvim.Subscribe(event1); err != nil {
		t.Fatalf("subscribe(%v): %v", event1, err)
	}

	b := p.Nvim.NewBatch()
	b.Subscribe(event2)
	if err := b.Execute(); err != nil {
		t.Fatalf("batch execute: %v", err)
	}

	// warm-up
	var result int
	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q)`, event1), &result); err != nil {
		t.Fatalf("eval rpcnotify for warm-up of event1: %v", err)
	}
	if result != 1 {
		t.Fatalf("expect 1 but got %d", result)
	}

	var result2 int
	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 1, 2, 3)`, event1), &result2); err != nil {
		t.Fatalf("eval rpcnotify for event1: %v", err)
	}
	if result2 != 1 {
		t.Fatalf("expect 1 but got %d", result2)
	}

	var result3 int
	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 4, 5, 6)`, event2), &result3); err != nil {
		t.Fatalf("eval rpcnotify for event2: %v", err)
	}
	if result3 != 1 {
		t.Fatalf("expect 1 but got %d", result3)
	}

	if err := p.Nvim.Unsubscribe(event1); err != nil {
		t.Fatalf("unsubscribe event1: %v", err)
	}

	b.Unsubscribe(event2)
	if err := b.Execute(); err != nil {
		t.Fatalf("unsubscribe event2: %v", err)
	}

	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 7, 8, 9)`, event1), nil); err != nil {
		t.Fatalf("ensure rpcnotify to event1 is no-op: %v", err)
	}

	if err := p.Nvim.Eval(fmt.Sprintf(`rpcnotify(0, %q, 10, 11, 12)`, event2), nil); err != nil {
		t.Fatalf("ensure rpcnotify to event2 is no-op: %v", err)
	}
}
<|endoftext|>


<commit_before>package acceptance_test

import (
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"runtime"
	"time"

	acceptance "github.com/cloudfoundry/bosh-bootloader/acceptance-tests"
	"github.com/cloudfoundry/bosh-bootloader/acceptance-tests/actors"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("Upgrade", func() {
	var (
		oldBBL  actors.BBL
		newBBL  actors.BBL
		boshcli actors.BOSHCLI
		state   acceptance.State
		f       *os.File
	)

	BeforeEach(func() {
		acceptance.SkipUnless("upgrade")

		configuration, err := acceptance.LoadConfig()
		Expect(err).NotTo(HaveOccurred())

		var bblBinaryLocation string
		if runtime.GOOS == "darwin" {
			bblBinaryLocation = "https://github.com/cloudfoundry/bosh-bootloader/releases/download/v5.11.5/bbl-v5.11.5_osx"
		} else {
			bblBinaryLocation = "https://github.com/cloudfoundry/bosh-bootloader/releases/download/v5.11.5/bbl-v5.11.5_linux_x86-64"
		}

		resp, err := http.Get(bblBinaryLocation)
		Expect(err).NotTo(HaveOccurred())

		f, err = ioutil.TempFile("", "")
		Expect(err).NotTo(HaveOccurred())

		_, err = io.Copy(f, resp.Body)
		Expect(err).NotTo(HaveOccurred())

		err = os.Chmod(f.Name(), 0700)
		Expect(err).NotTo(HaveOccurred())

		err = f.Close()
		Expect(err).NotTo(HaveOccurred())

		envName := "upgrade-env"
		testName := os.Getenv("RUN_TEST")
		if testName != "" {
			envName = testName
		}
		oldBBL = actors.NewBBL(configuration.StateFileDir, f.Name(), configuration, envName)
		newBBL = actors.NewBBL(configuration.StateFileDir, pathToBBL, configuration, envName)
		boshcli = actors.NewBOSHCLI()
		state = acceptance.NewState(configuration.StateFileDir)
	})

	AfterEach(func() {
		acceptance.SkipUnless("upgrade")

		By("destroying with the old bbl", func() {
			session := oldBBL.Destroy()
			Eventually(session, 10*time.Minute).Should(gexec.Exit())
		})

		By("destroying with the latest bbl", func() {
			session := newBBL.Destroy()
			Eventually(session, 10*time.Minute).Should(gexec.Exit())
		})

		err := os.Remove(f.Name())
		Expect(err).NotTo(HaveOccurred())
	})

	It("is able to upgrade from an environment bbl'd up with an older version of bbl", func() {
		By("cleaning up any leftovers", func() {
			session := newBBL.CleanupLeftovers(newBBL.PredefinedEnvID())
			Eventually(session, 10*time.Minute).Should(gexec.Exit())
		})

		By("bbl'ing up with old bbl", func() {
			session := oldBBL.Up("--name", oldBBL.PredefinedEnvID())
			Eventually(session, 40*time.Minute).Should(gexec.Exit(0))
		})

		By("verifying the director has a private ip", func() {
			Expect(oldBBL.DirectorAddress()).To(Equal("https://10.0.0.6:25555"))
		})

		By("verifying the director exists", func() {
			exists, err := boshcli.DirectorExists(oldBBL.DirectorAddress(), oldBBL.DirectorUsername(), oldBBL.DirectorPassword(), oldBBL.SaveDirectorCA())
			Expect(err).NotTo(HaveOccurred())
			Expect(exists).To(BeTrue())
		})

		By("upgrading to the latest bbl", func() {
			session := newBBL.Up()
			Eventually(session, 60*time.Minute).Should(gexec.Exit(0))
		})

		By("exporting environment variables to talk to the director", func() {
			newBBL.ExportBoshAllProxy()
		})

		directorAddress := newBBL.DirectorAddress()

		By("verifying the director has the same private ip", func() {
			Expect(directorAddress).To(Equal("https://10.0.0.6:25555"))
		})

		By("verifying the director still exists", func() {
			exists, err := boshcli.DirectorExists(directorAddress, newBBL.DirectorUsername(), newBBL.DirectorPassword(), newBBL.SaveDirectorCA())
			Expect(err).NotTo(HaveOccurred())
			Expect(exists).To(BeTrue())
		})
	})
})
<commit_msg>Update bbl upgrade test to set BOSH_ALL_PROXY<commit_after>package acceptance_test

import (
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"runtime"
	"time"

	acceptance "github.com/cloudfoundry/bosh-bootloader/acceptance-tests"
	"github.com/cloudfoundry/bosh-bootloader/acceptance-tests/actors"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("Upgrade", func() {
	var (
		oldBBL  actors.BBL
		newBBL  actors.BBL
		boshcli actors.BOSHCLI
		state   acceptance.State
		f       *os.File
	)

	BeforeEach(func() {
		acceptance.SkipUnless("upgrade")

		configuration, err := acceptance.LoadConfig()
		Expect(err).NotTo(HaveOccurred())

		var bblBinaryLocation string
		if runtime.GOOS == "darwin" {
			bblBinaryLocation = "https://github.com/cloudfoundry/bosh-bootloader/releases/download/v5.11.5/bbl-v5.11.5_osx"
		} else {
			bblBinaryLocation = "https://github.com/cloudfoundry/bosh-bootloader/releases/download/v5.11.5/bbl-v5.11.5_linux_x86-64"
		}

		resp, err := http.Get(bblBinaryLocation)
		Expect(err).NotTo(HaveOccurred())

		f, err = ioutil.TempFile("", "")
		Expect(err).NotTo(HaveOccurred())

		_, err = io.Copy(f, resp.Body)
		Expect(err).NotTo(HaveOccurred())

		err = os.Chmod(f.Name(), 0700)
		Expect(err).NotTo(HaveOccurred())

		err = f.Close()
		Expect(err).NotTo(HaveOccurred())

		envName := "upgrade-env"
		testName := os.Getenv("RUN_TEST")
		if testName != "" {
			envName = testName
		}
		oldBBL = actors.NewBBL(configuration.StateFileDir, f.Name(), configuration, envName)
		newBBL = actors.NewBBL(configuration.StateFileDir, pathToBBL, configuration, envName)
		boshcli = actors.NewBOSHCLI()
		state = acceptance.NewState(configuration.StateFileDir)
	})

	AfterEach(func() {
		acceptance.SkipUnless("upgrade")

		By("destroying with the old bbl", func() {
			session := oldBBL.Destroy()
			Eventually(session, 10*time.Minute).Should(gexec.Exit())
		})

		By("destroying with the latest bbl", func() {
			session := newBBL.Destroy()
			Eventually(session, 10*time.Minute).Should(gexec.Exit())
		})

		err := os.Remove(f.Name())
		Expect(err).NotTo(HaveOccurred())
	})

	It("is able to upgrade from an environment bbl'd up with an older version of bbl", func() {
		By("cleaning up any leftovers", func() {
			session := newBBL.CleanupLeftovers(newBBL.PredefinedEnvID())
			Eventually(session, 10*time.Minute).Should(gexec.Exit())
		})

		By("bbl'ing up with old bbl", func() {
			session := oldBBL.Up("--name", oldBBL.PredefinedEnvID())
			Eventually(session, 40*time.Minute).Should(gexec.Exit(0))
		})

		By("verifying the director has a private ip", func() {
			Expect(oldBBL.DirectorAddress()).To(Equal("https://10.0.0.6:25555"))
		})

		By("exporting environment variables to talk to the director", func() {
			newBBL.ExportBoshAllProxy()
		})

		By("verifying the director exists", func() {
			exists, err := boshcli.DirectorExists(oldBBL.DirectorAddress(), oldBBL.DirectorUsername(), oldBBL.DirectorPassword(), oldBBL.SaveDirectorCA())
			Expect(err).NotTo(HaveOccurred())
			Expect(exists).To(BeTrue())
		})

		By("upgrading to the latest bbl", func() {
			session := newBBL.Up()
			Eventually(session, 20*time.Minute).Should(gexec.Exit(0))
		})

		By("verifying the director still exists", func() {
			exists, err := boshcli.DirectorExists(newBBL.DirectorAddress(), newBBL.DirectorUsername(), newBBL.DirectorPassword(), newBBL.SaveDirectorCA())
			Expect(err).NotTo(HaveOccurred())
			Expect(exists).To(BeTrue())
		})
	})
})
<|endoftext|>


<commit_before>package osdn

import (
	"fmt"
	"time"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client"
	"k8s.io/kubernetes/pkg/client/cache"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/types"
	"k8s.io/kubernetes/pkg/watch"

	osdnapi "github.com/openshift/openshift-sdn/ovssubnet/api"

	osclient "github.com/openshift/origin/pkg/client"
	oscache "github.com/openshift/origin/pkg/client/cache"
	"github.com/openshift/origin/pkg/sdn/api"
)

type OsdnRegistryInterface struct {
	oClient osclient.Interface
	kClient kclient.Interface
}

func NewOsdnRegistryInterface(osClient *osclient.Client, kClient *kclient.Client) OsdnRegistryInterface {
	return OsdnRegistryInterface{osClient, kClient}
}

func (oi *OsdnRegistryInterface) InitSubnets() error {
	return nil
}

func (oi *OsdnRegistryInterface) GetSubnets() (*[]osdnapi.Subnet, error) {
	hostSubnetList, err := oi.oClient.HostSubnets().List()
	if err != nil {
		return nil, err
	}
	// convert HostSubnet to osdnapi.Subnet
	subList := make([]osdnapi.Subnet, 0)
	for _, subnet := range hostSubnetList.Items {
		subList = append(subList, osdnapi.Subnet{NodeIP: subnet.HostIP, Sub: subnet.Subnet})
	}
	return &subList, nil
}

func (oi *OsdnRegistryInterface) GetSubnet(node string) (*osdnapi.Subnet, error) {
	hs, err := oi.oClient.HostSubnets().Get(node)
	if err != nil {
		return nil, err
	}
	return &osdnapi.Subnet{NodeIP: hs.HostIP, Sub: hs.Subnet}, nil
}

func (oi *OsdnRegistryInterface) DeleteSubnet(node string) error {
	return oi.oClient.HostSubnets().Delete(node)
}

func (oi *OsdnRegistryInterface) CreateSubnet(node string, sub *osdnapi.Subnet) error {
	hs := &api.HostSubnet{
		TypeMeta:   kapi.TypeMeta{Kind: "HostSubnet"},
		ObjectMeta: kapi.ObjectMeta{Name: node},
		Host:       node,
		HostIP:     sub.NodeIP,
		Subnet:     sub.Sub,
	}
	_, err := oi.oClient.HostSubnets().Create(hs)
	return err
}

func (oi *OsdnRegistryInterface) WatchSubnets(receiver chan *osdnapi.SubnetEvent, stop chan bool) error {
	subnetEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.oClient.HostSubnets().List()
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.oClient.HostSubnets().Watch(resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &api.HostSubnet{}, subnetEventQueue, 4*time.Minute).Run()

	for {
		eventType, obj, err := subnetEventQueue.Pop()
		if err != nil {
			return err
		}
		switch eventType {
		case watch.Added, watch.Modified:
			// create SubnetEvent
			hs := obj.(*api.HostSubnet)
			receiver <- &osdnapi.SubnetEvent{Type: osdnapi.Added, Node: hs.Host, Sub: osdnapi.Subnet{NodeIP: hs.HostIP, Sub: hs.Subnet}}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			hs := obj.(*api.HostSubnet)
			receiver <- &osdnapi.SubnetEvent{Type: osdnapi.Deleted, Node: hs.Host, Sub: osdnapi.Subnet{NodeIP: hs.HostIP, Sub: hs.Subnet}}
		}
	}
}

func (oi *OsdnRegistryInterface) InitNodes() error {
	// return no error, as this gets initialized by apiserver
	return nil
}

func (oi *OsdnRegistryInterface) GetNodes() (*[]string, error) {
	nodes, err := oi.kClient.Nodes().List(labels.Everything(), fields.Everything())
	if err != nil {
		return nil, err
	}
	// convert kapi.NodeList to []string
	nodeList := make([]string, 0)
	for _, node := range nodes.Items {
		nodeList = append(nodeList, node.Name)
	}
	return &nodeList, nil
}

func (oi *OsdnRegistryInterface) CreateNode(node string, data string) error {
	return fmt.Errorf("Feature not supported in native mode. SDN cannot create/register nodes.")
}

func (oi *OsdnRegistryInterface) getNodeAddressMap() (map[types.UID]string, error) {
	nodeAddressMap := map[types.UID]string{}

	nodes, err := oi.kClient.Nodes().List(labels.Everything(), fields.Everything())
	if err != nil {
		return nodeAddressMap, err
	}
	for _, node := range nodes.Items {
		if len(node.Status.Addresses) > 0 {
			nodeAddressMap[node.ObjectMeta.UID] = node.Status.Addresses[0].Address
		}
	}
	return nodeAddressMap, nil
}

func (oi *OsdnRegistryInterface) WatchNodes(receiver chan *osdnapi.NodeEvent, stop chan bool) error {
	nodeEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.kClient.Nodes().List(labels.Everything(), fields.Everything())
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.kClient.Nodes().Watch(labels.Everything(), fields.Everything(), resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &kapi.Node{}, nodeEventQueue, 4*time.Minute).Run()

	nodeAddressMap, err := oi.getNodeAddressMap()
	if err != nil {
		return err
	}

	for {
		eventType, obj, err := nodeEventQueue.Pop()
		if err != nil {
			return err
		}
		node := obj.(*kapi.Node)
		nodeIP := ""
		if len(node.Status.Addresses) > 0 {
			nodeIP = node.Status.Addresses[0].Address
		}

		switch eventType {
		case watch.Added:
			receiver <- &osdnapi.NodeEvent{Type: osdnapi.Added, Node: node.ObjectMeta.Name, NodeIP: nodeIP}
			nodeAddressMap[node.ObjectMeta.UID] = nodeIP
		case watch.Modified:
			oldNodeIP, ok := nodeAddressMap[node.ObjectMeta.UID]
			if ok && oldNodeIP != nodeIP {
				// Node Added event will handle update subnet if there is ip mismatch
				receiver <- &osdnapi.NodeEvent{Type: osdnapi.Added, Node: node.ObjectMeta.Name, NodeIP: nodeIP}
				nodeAddressMap[node.ObjectMeta.UID] = nodeIP
			}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			receiver <- &osdnapi.NodeEvent{Type: osdnapi.Deleted, Node: node.ObjectMeta.Name}
			delete(nodeAddressMap, node.ObjectMeta.UID)
		}
	}
}

func (oi *OsdnRegistryInterface) WriteNetworkConfig(network string, subnetLength uint) error {
	cn, err := oi.oClient.ClusterNetwork().Get("default")
	if err == nil {
		if cn.Network == network && cn.HostSubnetLength == int(subnetLength) {
			return nil
		} else {
			return fmt.Errorf("A network already exists and does not match the new network's parameters - Existing: (%s, %d); New: (%s, %d) ", cn.Network, cn.HostSubnetLength, network, subnetLength)
		}
	}
	cn = &api.ClusterNetwork{
		TypeMeta:         kapi.TypeMeta{Kind: "ClusterNetwork"},
		ObjectMeta:       kapi.ObjectMeta{Name: "default"},
		Network:          network,
		HostSubnetLength: int(subnetLength),
	}
	_, err = oi.oClient.ClusterNetwork().Create(cn)
	return err
}

func (oi *OsdnRegistryInterface) GetContainerNetwork() (string, error) {
	cn, err := oi.oClient.ClusterNetwork().Get("default")
	return cn.Network, err
}

func (oi *OsdnRegistryInterface) GetSubnetLength() (uint64, error) {
	cn, err := oi.oClient.ClusterNetwork().Get("default")
	return uint64(cn.HostSubnetLength), err
}

func (oi *OsdnRegistryInterface) CheckEtcdIsAlive(seconds uint64) bool {
	// always assumed to be true as we run through the apiserver
	return true
}

func (oi *OsdnRegistryInterface) WatchNamespaces(receiver chan *osdnapi.NamespaceEvent, stop chan bool) error {
	nsEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.kClient.Namespaces().List(labels.Everything(), fields.Everything())
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.kClient.Namespaces().Watch(labels.Everything(), fields.Everything(), resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &kapi.Namespace{}, nsEventQueue, 4*time.Minute).Run()

	for {
		eventType, obj, err := nsEventQueue.Pop()
		if err != nil {
			return err
		}
		switch eventType {
		case watch.Added:
			// we should ignore the modified event because status updates cause unnecessary noise
			// the only time we would care about modified would be if the node changes its IP address
			// and hence all nodes need to update their vtep entries for the respective subnet
			// create nodeEvent
			ns := obj.(*kapi.Namespace)
			receiver <- &osdnapi.NamespaceEvent{Type: osdnapi.Added, Name: ns.ObjectMeta.Name}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			ns := obj.(*kapi.Namespace)
			receiver <- &osdnapi.NamespaceEvent{Type: osdnapi.Deleted, Name: ns.ObjectMeta.Name}
		}
	}
}

func (oi *OsdnRegistryInterface) WatchNetNamespaces(receiver chan *osdnapi.NetNamespaceEvent, stop chan bool) error {
	netNsEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.oClient.NetNamespaces().List()
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.oClient.NetNamespaces().Watch(resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &api.NetNamespace{}, netNsEventQueue, 4*time.Minute).Run()

	for {
		eventType, obj, err := netNsEventQueue.Pop()
		if err != nil {
			return err
		}
		switch eventType {
		case watch.Added:
			// we should ignore the modified event because status updates cause unnecessary noise
			// the only time we would care about modified would be if the node changes its IP address
			// and hence all nodes need to update their vtep entries for the respective subnet
			// create nodeEvent
			netns := obj.(*api.NetNamespace)
			receiver <- &osdnapi.NetNamespaceEvent{Type: osdnapi.Added, Name: netns.NetName, NetID: netns.NetID}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			netns := obj.(*api.NetNamespace)
			receiver <- &osdnapi.NetNamespaceEvent{Type: osdnapi.Deleted, Name: netns.NetName}
		}
	}
}

func (oi *OsdnRegistryInterface) GetNetNamespaces() ([]osdnapi.NetNamespace, error) {
	netNamespaceList, err := oi.oClient.NetNamespaces().List()
	if err != nil {
		return nil, err
	}
	// convert api.NetNamespace to osdnapi.NetNamespace
	nsList := make([]osdnapi.NetNamespace, 0)
	for _, netns := range netNamespaceList.Items {
		nsList = append(nsList, osdnapi.NetNamespace{Name: netns.Name, NetID: netns.NetID})
	}
	return nsList, nil
}

func (oi *OsdnRegistryInterface) GetNetNamespace(name string) (osdnapi.NetNamespace, error) {
	netns, err := oi.oClient.NetNamespaces().Get(name)
	if err != nil {
		return osdnapi.NetNamespace{}, err
	}
	return osdnapi.NetNamespace{Name: netns.Name, NetID: netns.NetID}, nil
}

func (oi *OsdnRegistryInterface) WriteNetNamespace(name string, id uint) error {
	netns := &api.NetNamespace{
		TypeMeta:   kapi.TypeMeta{Kind: "NetNamespace"},
		ObjectMeta: kapi.ObjectMeta{Name: name},
		NetName:    name,
		NetID:      id,
	}
	_, err := oi.oClient.NetNamespaces().Create(netns)
	return err
}

func (oi *OsdnRegistryInterface) DeleteNetNamespace(name string) error {
	return oi.oClient.NetNamespaces().Delete(name)
}
<commit_msg>Fix sdn api field names to match openshift-sdn repo<commit_after>package osdn

import (
	"fmt"
	"time"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client"
	"k8s.io/kubernetes/pkg/client/cache"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/types"
	"k8s.io/kubernetes/pkg/watch"

	osdnapi "github.com/openshift/openshift-sdn/ovssubnet/api"

	osclient "github.com/openshift/origin/pkg/client"
	oscache "github.com/openshift/origin/pkg/client/cache"
	"github.com/openshift/origin/pkg/sdn/api"
)

type OsdnRegistryInterface struct {
	oClient osclient.Interface
	kClient kclient.Interface
}

func NewOsdnRegistryInterface(osClient *osclient.Client, kClient *kclient.Client) OsdnRegistryInterface {
	return OsdnRegistryInterface{osClient, kClient}
}

func (oi *OsdnRegistryInterface) InitSubnets() error {
	return nil
}

func (oi *OsdnRegistryInterface) GetSubnets() (*[]osdnapi.Subnet, error) {
	hostSubnetList, err := oi.oClient.HostSubnets().List()
	if err != nil {
		return nil, err
	}
	// convert HostSubnet to osdnapi.Subnet
	subList := make([]osdnapi.Subnet, 0)
	for _, subnet := range hostSubnetList.Items {
		subList = append(subList, osdnapi.Subnet{NodeIP: subnet.HostIP, SubnetIP: subnet.Subnet})
	}
	return &subList, nil
}

func (oi *OsdnRegistryInterface) GetSubnet(nodeName string) (*osdnapi.Subnet, error) {
	hs, err := oi.oClient.HostSubnets().Get(nodeName)
	if err != nil {
		return nil, err
	}
	return &osdnapi.Subnet{NodeIP: hs.HostIP, SubnetIP: hs.Subnet}, nil
}

func (oi *OsdnRegistryInterface) DeleteSubnet(nodeName string) error {
	return oi.oClient.HostSubnets().Delete(nodeName)
}

func (oi *OsdnRegistryInterface) CreateSubnet(nodeName string, sub *osdnapi.Subnet) error {
	hs := &api.HostSubnet{
		TypeMeta:   kapi.TypeMeta{Kind: "HostSubnet"},
		ObjectMeta: kapi.ObjectMeta{Name: nodeName},
		Host:       nodeName,
		HostIP:     sub.NodeIP,
		Subnet:     sub.SubnetIP,
	}
	_, err := oi.oClient.HostSubnets().Create(hs)
	return err
}

func (oi *OsdnRegistryInterface) WatchSubnets(receiver chan *osdnapi.SubnetEvent, stop chan bool) error {
	subnetEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.oClient.HostSubnets().List()
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.oClient.HostSubnets().Watch(resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &api.HostSubnet{}, subnetEventQueue, 4*time.Minute).Run()

	for {
		eventType, obj, err := subnetEventQueue.Pop()
		if err != nil {
			return err
		}
		switch eventType {
		case watch.Added, watch.Modified:
			// create SubnetEvent
			hs := obj.(*api.HostSubnet)
			receiver <- &osdnapi.SubnetEvent{Type: osdnapi.Added, NodeName: hs.Host, Subnet: osdnapi.Subnet{NodeIP: hs.HostIP, SubnetIP: hs.Subnet}}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			hs := obj.(*api.HostSubnet)
			receiver <- &osdnapi.SubnetEvent{Type: osdnapi.Deleted, NodeName: hs.Host, Subnet: osdnapi.Subnet{NodeIP: hs.HostIP, SubnetIP: hs.Subnet}}
		}
	}
}

func (oi *OsdnRegistryInterface) InitNodes() error {
	// return no error, as this gets initialized by apiserver
	return nil
}

func (oi *OsdnRegistryInterface) GetNodes() (*[]string, error) {
	nodes, err := oi.kClient.Nodes().List(labels.Everything(), fields.Everything())
	if err != nil {
		return nil, err
	}
	// convert kapi.NodeList to []string
	nodeList := make([]string, 0)
	for _, node := range nodes.Items {
		nodeList = append(nodeList, node.Name)
	}
	return &nodeList, nil
}

func (oi *OsdnRegistryInterface) CreateNode(nodeName string, data string) error {
	return fmt.Errorf("Feature not supported in native mode. SDN cannot create/register nodes.")
}

func (oi *OsdnRegistryInterface) getNodeAddressMap() (map[types.UID]string, error) {
	nodeAddressMap := map[types.UID]string{}

	nodes, err := oi.kClient.Nodes().List(labels.Everything(), fields.Everything())
	if err != nil {
		return nodeAddressMap, err
	}
	for _, node := range nodes.Items {
		if len(node.Status.Addresses) > 0 {
			nodeAddressMap[node.ObjectMeta.UID] = node.Status.Addresses[0].Address
		}
	}
	return nodeAddressMap, nil
}

func (oi *OsdnRegistryInterface) WatchNodes(receiver chan *osdnapi.NodeEvent, stop chan bool) error {
	nodeEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.kClient.Nodes().List(labels.Everything(), fields.Everything())
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.kClient.Nodes().Watch(labels.Everything(), fields.Everything(), resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &kapi.Node{}, nodeEventQueue, 4*time.Minute).Run()

	nodeAddressMap, err := oi.getNodeAddressMap()
	if err != nil {
		return err
	}

	for {
		eventType, obj, err := nodeEventQueue.Pop()
		if err != nil {
			return err
		}
		node := obj.(*kapi.Node)
		nodeIP := ""
		if len(node.Status.Addresses) > 0 {
			nodeIP = node.Status.Addresses[0].Address
		}

		switch eventType {
		case watch.Added:
			receiver <- &osdnapi.NodeEvent{Type: osdnapi.Added, NodeName: node.ObjectMeta.Name, NodeIP: nodeIP}
			nodeAddressMap[node.ObjectMeta.UID] = nodeIP
		case watch.Modified:
			oldNodeIP, ok := nodeAddressMap[node.ObjectMeta.UID]
			if ok && oldNodeIP != nodeIP {
				// Node Added event will handle update subnet if there is ip mismatch
				receiver <- &osdnapi.NodeEvent{Type: osdnapi.Added, NodeName: node.ObjectMeta.Name, NodeIP: nodeIP}
				nodeAddressMap[node.ObjectMeta.UID] = nodeIP
			}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			receiver <- &osdnapi.NodeEvent{Type: osdnapi.Deleted, NodeName: node.ObjectMeta.Name}
			delete(nodeAddressMap, node.ObjectMeta.UID)
		}
	}
}

func (oi *OsdnRegistryInterface) WriteNetworkConfig(network string, subnetLength uint) error {
	cn, err := oi.oClient.ClusterNetwork().Get("default")
	if err == nil {
		if cn.Network == network && cn.HostSubnetLength == int(subnetLength) {
			return nil
		} else {
			return fmt.Errorf("A network already exists and does not match the new network's parameters - Existing: (%s, %d); New: (%s, %d) ", cn.Network, cn.HostSubnetLength, network, subnetLength)
		}
	}
	cn = &api.ClusterNetwork{
		TypeMeta:         kapi.TypeMeta{Kind: "ClusterNetwork"},
		ObjectMeta:       kapi.ObjectMeta{Name: "default"},
		Network:          network,
		HostSubnetLength: int(subnetLength),
	}
	_, err = oi.oClient.ClusterNetwork().Create(cn)
	return err
}

func (oi *OsdnRegistryInterface) GetContainerNetwork() (string, error) {
	cn, err := oi.oClient.ClusterNetwork().Get("default")
	return cn.Network, err
}

func (oi *OsdnRegistryInterface) GetSubnetLength() (uint64, error) {
	cn, err := oi.oClient.ClusterNetwork().Get("default")
	return uint64(cn.HostSubnetLength), err
}

func (oi *OsdnRegistryInterface) CheckEtcdIsAlive(seconds uint64) bool {
	// always assumed to be true as we run through the apiserver
	return true
}

func (oi *OsdnRegistryInterface) WatchNamespaces(receiver chan *osdnapi.NamespaceEvent, stop chan bool) error {
	nsEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.kClient.Namespaces().List(labels.Everything(), fields.Everything())
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.kClient.Namespaces().Watch(labels.Everything(), fields.Everything(), resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &kapi.Namespace{}, nsEventQueue, 4*time.Minute).Run()

	for {
		eventType, obj, err := nsEventQueue.Pop()
		if err != nil {
			return err
		}
		switch eventType {
		case watch.Added:
			// we should ignore the modified event because status updates cause unnecessary noise
			// the only time we would care about modified would be if the node changes its IP address
			// and hence all nodes need to update their vtep entries for the respective subnet
			// create nodeEvent
			ns := obj.(*kapi.Namespace)
			receiver <- &osdnapi.NamespaceEvent{Type: osdnapi.Added, Name: ns.ObjectMeta.Name}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			ns := obj.(*kapi.Namespace)
			receiver <- &osdnapi.NamespaceEvent{Type: osdnapi.Deleted, Name: ns.ObjectMeta.Name}
		}
	}
}

func (oi *OsdnRegistryInterface) WatchNetNamespaces(receiver chan *osdnapi.NetNamespaceEvent, stop chan bool) error {
	netNsEventQueue := oscache.NewEventQueue(cache.MetaNamespaceKeyFunc)
	listWatch := &cache.ListWatch{
		ListFunc: func() (runtime.Object, error) {
			return oi.oClient.NetNamespaces().List()
		},
		WatchFunc: func(resourceVersion string) (watch.Interface, error) {
			return oi.oClient.NetNamespaces().Watch(resourceVersion)
		},
	}
	cache.NewReflector(listWatch, &api.NetNamespace{}, netNsEventQueue, 4*time.Minute).Run()

	for {
		eventType, obj, err := netNsEventQueue.Pop()
		if err != nil {
			return err
		}
		switch eventType {
		case watch.Added:
			// we should ignore the modified event because status updates cause unnecessary noise
			// the only time we would care about modified would be if the node changes its IP address
			// and hence all nodes need to update their vtep entries for the respective subnet
			// create nodeEvent
			netns := obj.(*api.NetNamespace)
			receiver <- &osdnapi.NetNamespaceEvent{Type: osdnapi.Added, Name: netns.NetName, NetID: netns.NetID}
		case watch.Deleted:
			// TODO: There is a chance that a Delete event will not get triggered.
			// Need to use a periodic sync loop that lists and compares.
			netns := obj.(*api.NetNamespace)
			receiver <- &osdnapi.NetNamespaceEvent{Type: osdnapi.Deleted, Name: netns.NetName}
		}
	}
}

func (oi *OsdnRegistryInterface) GetNetNamespaces() ([]osdnapi.NetNamespace, error) {
	netNamespaceList, err := oi.oClient.NetNamespaces().List()
	if err != nil {
		return nil, err
	}
	// convert api.NetNamespace to osdnapi.NetNamespace
	nsList := make([]osdnapi.NetNamespace, 0)
	for _, netns := range netNamespaceList.Items {
		nsList = append(nsList, osdnapi.NetNamespace{Name: netns.Name, NetID: netns.NetID})
	}
	return nsList, nil
}

func (oi *OsdnRegistryInterface) GetNetNamespace(name string) (osdnapi.NetNamespace, error) {
	netns, err := oi.oClient.NetNamespaces().Get(name)
	if err != nil {
		return osdnapi.NetNamespace{}, err
	}
	return osdnapi.NetNamespace{Name: netns.Name, NetID: netns.NetID}, nil
}

func (oi *OsdnRegistryInterface) WriteNetNamespace(name string, id uint) error {
	netns := &api.NetNamespace{
		TypeMeta:   kapi.TypeMeta{Kind: "NetNamespace"},
		ObjectMeta: kapi.ObjectMeta{Name: name},
		NetName:    name,
		NetID:      id,
	}
	_, err := oi.oClient.NetNamespaces().Create(netns)
	return err
}

func (oi *OsdnRegistryInterface) DeleteNetNamespace(name string) error {
	return oi.oClient.NetNamespaces().Delete(name)
}
<|endoftext|>


<commit_before>// +build windows

package enforcer

import (
	"fmt"
	"strings"
	"sync"

	"github.com/scjalliance/resourceful/guardian"
	"github.com/scjalliance/resourceful/lease"
	"github.com/scjalliance/resourceful/policy"
)

// ProcessManager enforces a set of policies on local processes.
type ProcessManager struct {
	client   *guardian.Client
	hostname string
	passive  bool // Don't kill processes if true
	sessions *SessionManager
	logger   Logger

	mutex       sync.RWMutex
	managed     map[UniqueID]lease.Instance
	skipped     map[UniqueID]struct{}
	invocations map[lease.Instance]*Invocation // Keyed by resource consumed
}

// NewProcessManager returns a new process manager that is ready for use.
func NewProcessManager(client *guardian.Client, hostname string, passive bool, sessions *SessionManager, logger Logger) *ProcessManager {
	return &ProcessManager{
		client:      client,
		hostname:    hostname,
		passive:     passive,
		sessions:    sessions,
		logger:      logger,
		managed:     make(map[UniqueID]lease.Instance, 8),
		skipped:     make(map[UniqueID]struct{}),
		invocations: make(map[lease.Instance]*Invocation, 8),
	}
}

// Enforce causes the process manager to enforce the given policy set.
func (m *ProcessManager) Enforce(policies policy.Set) error {
	procs, err := Scan(policies)
	if err != nil {
		return err
	}

	scanned := make(map[UniqueID]struct{}, len(procs))

	m.mutex.Lock()
	defer m.mutex.Unlock()

	var pending []ProcessData
	for _, proc := range procs {
		id := proc.UniqueID()
		scanned[id] = struct{}{} // Record the ID in the map of scanned procs

		// Don't manage blacklisted processes
		if Blacklisted(proc) {
			if instance, exists := m.managed[id]; exists {
				// Stop the invocation
				if inv := m.invocations[instance]; inv != nil {
					inv.Stop()
					delete(m.invocations, instance)
					m.log("Stopped management of blacklisted invocation %s", instance)
				}
				// Remove from managed and add to skipped
				delete(m.managed, id)
				m.skipped[id] = struct{}{}
				m.log("Stopped management of blacklisted process %s", id)
			} else if _, exists := m.skipped[id]; !exists {
				// Add to skipped
				m.log("Skipped management of blacklisted process %s", id)
				m.skipped[id] = struct{}{}
			}
			continue
		} else {
			if _, exists := m.skipped[id]; exists {
				// Remove from skipped
				delete(m.skipped, id)
			}
		}

		// Don't re-manage processes that are already managed
		if _, exists := m.managed[id]; exists {
			continue
		}

		// If it matches a policy add it to the pending slice
		matches := policies.Match(Properties(proc, m.hostname))
		if len(matches) > 0 {
			pending = append(pending, proc)
		}
	}

	// Bookkeeping for dead processes
	for id, instance := range m.managed {
		if _, exists := scanned[id]; !exists {
			if inv := m.invocations[instance]; inv != nil {
				// TODO: Ask the invocation whether it's still alive. Check
				// if it has stale policies. Remove if necessary.
				if inv.Done() {
					inv.Stop()
					delete(m.invocations, instance)
					m.debug("Stopped management of invocation %s", instance.ID)
				}
			}
			delete(m.managed, id)
			m.debug("Stopped management of process %s", id)
		}
	}

	for id := range m.skipped {
		if _, exists := scanned[id]; !exists {
			delete(m.skipped, id)
		}
	}

	// Bookkeeping for dead invocations
	for instance, inv := range m.invocations {
		if !inv.Done() {
			continue
		}
		inv.Stop()
		delete(m.invocations, instance)
		m.debug("Stopped management of invocation %s", instance.ID)
	}

	// Exit early if nothing is pending
	if len(pending) == 0 {
		return nil
	}

	// Begin management of newly discovered processes
	if len(pending) == 1 {
		m.debug("Enforcement found 1 new process: %s", pending[0].ID)
	} else {
		ids := make([]string, len(pending))
		for i := range pending {
			ids[i] = pending[i].ID.String()
		}
		m.debug("Enforcement found %d new processes: %s", len(pending), strings.Join(ids, ", "))
	}

	for _, proc := range pending {
		id := proc.UniqueID()

		// Open a reference to the process.
		process, err := NewProcess(proc, m.passive, m.logger)
		if err != nil {
			// TODO: Retry on some interval with backoff so we don't spam the logs
			m.log("Unable to manage process %s: %v", id, err)
			continue
		}

		// Look for an existing invocation that can absorb the process
		absorbed := false
		for instance, inv := range m.invocations {
			absorbed = inv.Absorb(process)
			if absorbed {
				m.managed[id] = instance
				m.debug("Process %s absorbed into instance %s", id, instance.ID)
				break
			}
		}
		if absorbed {
			continue
		}

		// Create a new invocation
		instance := Instance(m.hostname, proc, NewInstanceID(proc))
		m.debug("Started management of invocation %s", instance.ID)
		m.debug("Started management of process %s", id)
		invocation := NewInvocation(m.client, instance, process, m.sessions.Session(SessionID(proc.SessionID)), m.logger)
		m.invocations[instance] = invocation
		m.managed[id] = instance

	}

	return nil
}

// Stop causes the process manager to stop all process management.
func (m *ProcessManager) Stop() {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	for instance, inv := range m.invocations {
		m.log("Stopping management of invocation %s", instance.ID)
		inv.Stop()
		delete(m.invocations, instance)
		m.log("Stopped management of invocation %s", instance.ID)
	}
	for id := range m.managed {
		delete(m.managed, id)
	}
}

func (m *ProcessManager) log(format string, v ...interface{}) {
	if m.logger == nil {
		return
	}
	m.logger.Log(ServiceEvent{
		Msg: fmt.Sprintf(format, v...),
	})
}

func (m *ProcessManager) debug(format string, v ...interface{}) {
	if m.logger == nil {
		return
	}
	m.logger.Log(ServiceEvent{
		Msg:   fmt.Sprintf(format, v...),
		Debug: true,
	})
}
<commit_msg>enforcer: Improved logging of blacklisted processes<commit_after>// +build windows

package enforcer

import (
	"fmt"
	"strings"
	"sync"

	"github.com/scjalliance/resourceful/guardian"
	"github.com/scjalliance/resourceful/lease"
	"github.com/scjalliance/resourceful/policy"
)

// ProcessManager enforces a set of policies on local processes.
type ProcessManager struct {
	client   *guardian.Client
	hostname string
	passive  bool // Don't kill processes if true
	sessions *SessionManager
	logger   Logger

	mutex       sync.RWMutex
	managed     map[UniqueID]lease.Instance
	skipped     map[UniqueID]struct{}
	invocations map[lease.Instance]*Invocation // Keyed by resource consumed
}

// NewProcessManager returns a new process manager that is ready for use.
func NewProcessManager(client *guardian.Client, hostname string, passive bool, sessions *SessionManager, logger Logger) *ProcessManager {
	return &ProcessManager{
		client:      client,
		hostname:    hostname,
		passive:     passive,
		sessions:    sessions,
		logger:      logger,
		managed:     make(map[UniqueID]lease.Instance, 8),
		skipped:     make(map[UniqueID]struct{}),
		invocations: make(map[lease.Instance]*Invocation, 8),
	}
}

// Enforce causes the process manager to enforce the given policy set.
func (m *ProcessManager) Enforce(policies policy.Set) error {
	procs, err := Scan(policies)
	if err != nil {
		return err
	}

	scanned := make(map[UniqueID]struct{}, len(procs))

	m.mutex.Lock()
	defer m.mutex.Unlock()

	var pending []ProcessData
	for _, proc := range procs {
		id := proc.UniqueID()
		scanned[id] = struct{}{} // Record the ID in the map of scanned procs

		// Don't manage blacklisted processes
		if Blacklisted(proc) {
			if instance, exists := m.managed[id]; exists {
				// Stop the invocation
				if inv := m.invocations[instance]; inv != nil {
					inv.Stop()
					delete(m.invocations, instance)
					m.log("Stopped management of blacklisted invocation %s (%s)", instance, proc.Name)
				}
				// Remove from managed and add to skipped
				delete(m.managed, id)
				m.skipped[id] = struct{}{}
				m.log("Stopped management of blacklisted process %s (%s)", id, proc.Name)
			} else if _, exists := m.skipped[id]; !exists {
				// Add to skipped
				m.log("Skipped management of blacklisted process %s (%s)", id, proc.Name)
				m.skipped[id] = struct{}{}
			}
			continue
		} else {
			if _, exists := m.skipped[id]; exists {
				// Remove from skipped
				delete(m.skipped, id)
			}
		}

		// Don't re-manage processes that are already managed
		if _, exists := m.managed[id]; exists {
			continue
		}

		// If it matches a policy add it to the pending slice
		matches := policies.Match(Properties(proc, m.hostname))
		if len(matches) > 0 {
			pending = append(pending, proc)
		}
	}

	// Bookkeeping for dead processes
	for id, instance := range m.managed {
		if _, exists := scanned[id]; !exists {
			if inv := m.invocations[instance]; inv != nil {
				// TODO: Ask the invocation whether it's still alive. Check
				// if it has stale policies. Remove if necessary.
				if inv.Done() {
					inv.Stop()
					delete(m.invocations, instance)
					m.debug("Stopped management of invocation %s", instance.ID)
				}
			}
			delete(m.managed, id)
			m.debug("Stopped management of process %s", id)
		}
	}

	for id := range m.skipped {
		if _, exists := scanned[id]; !exists {
			delete(m.skipped, id)
		}
	}

	// Bookkeeping for dead invocations
	for instance, inv := range m.invocations {
		if !inv.Done() {
			continue
		}
		inv.Stop()
		delete(m.invocations, instance)
		m.debug("Stopped management of invocation %s", instance.ID)
	}

	// Exit early if nothing is pending
	if len(pending) == 0 {
		return nil
	}

	// Begin management of newly discovered processes
	if len(pending) == 1 {
		m.debug("Enforcement found 1 new process: %s", pending[0].ID)
	} else {
		ids := make([]string, len(pending))
		for i := range pending {
			ids[i] = pending[i].ID.String()
		}
		m.debug("Enforcement found %d new processes: %s", len(pending), strings.Join(ids, ", "))
	}

	for _, proc := range pending {
		id := proc.UniqueID()

		// Open a reference to the process.
		process, err := NewProcess(proc, m.passive, m.logger)
		if err != nil {
			// TODO: Retry on some interval with backoff so we don't spam the logs
			m.log("Unable to manage process %s: %v", id, err)
			continue
		}

		// Look for an existing invocation that can absorb the process
		absorbed := false
		for instance, inv := range m.invocations {
			absorbed = inv.Absorb(process)
			if absorbed {
				m.managed[id] = instance
				m.debug("Process %s absorbed into instance %s", id, instance.ID)
				break
			}
		}
		if absorbed {
			continue
		}

		// Create a new invocation
		instance := Instance(m.hostname, proc, NewInstanceID(proc))
		m.debug("Started management of invocation %s", instance.ID)
		m.debug("Started management of process %s", id)
		invocation := NewInvocation(m.client, instance, process, m.sessions.Session(SessionID(proc.SessionID)), m.logger)
		m.invocations[instance] = invocation
		m.managed[id] = instance
	}

	return nil
}

// Stop causes the process manager to stop all process management.
func (m *ProcessManager) Stop() {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	for instance, inv := range m.invocations {
		m.log("Stopping management of invocation %s", instance.ID)
		inv.Stop()
		delete(m.invocations, instance)
		m.log("Stopped management of invocation %s", instance.ID)
	}
	for id := range m.managed {
		delete(m.managed, id)
	}
}

func (m *ProcessManager) log(format string, v ...interface{}) {
	if m.logger == nil {
		return
	}
	m.logger.Log(ServiceEvent{
		Msg: fmt.Sprintf(format, v...),
	})
}

func (m *ProcessManager) debug(format string, v ...interface{}) {
	if m.logger == nil {
		return
	}
	m.logger.Log(ServiceEvent{
		Msg:   fmt.Sprintf(format, v...),
		Debug: true,
	})
}
<|endoftext|>


<commit_before>// Application that does metrics analysis as described in the design doc:
// go/ct_metrics_analysis
//
// Can be tested locally with:
// $ go run go/worker_scripts/metrics_analysis/main.go --start_range=1 --num=3 --run_id=rmistry-test1 --benchmark_extra_args="--output-format=csv" --metric_name="loadingMetric" --logtostderr=true --local
//
package main

import (
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"go.skia.org/infra/go/sklog"

	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/ct/go/worker_scripts/worker_common"
	"go.skia.org/infra/go/common"
	skutil "go.skia.org/infra/go/util"
)

const (
	// The number of goroutines that will run in parallel to download traces and
	// run metrics analysis.
	WORKER_POOL_SIZE = 5

	METRICS_BENCHMARK_TIMEOUT_SECS = 300

	TRACE_OUTPUT_BUCKET = "chrome-telemetry-output"
)

var (
	startRange         = flag.Int("start_range", 1, "The number this worker will run metrics analysis from.")
	num                = flag.Int("num", 100, "The total number of traces to run metrics analysis from the start_range.")
	runID              = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
	benchmarkExtraArgs = flag.String("benchmark_extra_args", "", "The extra arguments that are passed to the specified benchmark.")
	metricName         = flag.String("metric_name", "", "The metric to parse the traces with. Eg: loadingMetric")
)

func metricsAnalysis() error {
	defer common.LogPanic()
	worker_common.Init()
	if !*worker_common.Local {
		defer util.CleanTmpDir()
	}
	defer util.TimeTrack(time.Now(), "Metrics Analysis")
	defer sklog.Flush()

	// Validate required arguments.
	if *runID == "" {
		return errors.New("Must specify --run_id")
	}
	if *metricName == "" {
		return errors.New("Must specify --metric_name")
	}

	ctx := context.Background()

	// Instantiate GcsUtil object.
	gs, err := util.NewGcsUtil(nil)
	if err != nil {
		return fmt.Errorf("GcsUtil instantiation failed: %s", err)
	}

	// Download the trace URLs for this run from Google storage.
	tracesFilename := *runID + ".traces.csv"
	tmpDir, err := ioutil.TempDir(util.PagesetsDir, "traces")
	if err != nil {
		return fmt.Errorf("Could not create tmpdir: %s", err)
	}
	defer skutil.RemoveAll(tmpDir)
	remotePatchesDir := filepath.Join(util.BenchmarkRunsDir, *runID)

	// Download traces.
	if _, err := util.DownloadPatch(filepath.Join(tmpDir, tracesFilename), filepath.Join(remotePatchesDir, tracesFilename), gs); err != nil {
		return fmt.Errorf("Could not download %s: %s", tracesFilename, err)
	}
	traces, err := util.GetCustomPages(filepath.Join(tmpDir, tracesFilename))
	if err != nil {
		return fmt.Errorf("Could not read custom traces file %s: %s", tracesFilename, err)
	}
	if len(traces) == 0 {
		return errors.New(fmt.Sprintf("No traces found in %s", tracesFilename))
	}
	traces = util.GetCustomPagesWithinRange(*startRange, *num, traces)
	sklog.Infof("Using %d traces", len(traces))

	// Establish output paths for trace downloads and metrics.
	traceDownloadDir := filepath.Join(util.StorageDir, util.TraceDownloadsDir, *runID)
	skutil.RemoveAll(traceDownloadDir)
	skutil.MkdirAll(traceDownloadDir, 0700)
	defer skutil.RemoveAll(traceDownloadDir)

	localOutputDir := filepath.Join(util.StorageDir, util.BenchmarkRunsDir, *runID)
	skutil.RemoveAll(localOutputDir)
	skutil.MkdirAll(localOutputDir, 0700)
	defer skutil.RemoveAll(localOutputDir)
	remoteDir := filepath.Join(util.BenchmarkRunsDir, *runID)

	sklog.Infof("===== Going to run the task with %d parallel goroutines =====", WORKER_POOL_SIZE)
	// Create channel that contains all trace ULs. This channel will
	// be consumed by the worker pool.
	traceRequests := getClosedChannelOfTraces(traces)

	var wg sync.WaitGroup

	// If not a single benchmark run succeeds then throw at error at the end.
	atleastOneBenchmarkSucceeded := false
	// Gather traceURLs that could not be downloaded.
	erroredTraces := []string{}
	// Mutex to control access to the above slice.
	var erroredTracesMutex sync.Mutex

	// Loop through workers in the worker pool.
	for i := 0; i < WORKER_POOL_SIZE; i++ {
		// Increment the WaitGroup counter.
		wg.Add(1)

		// Create and run a goroutine closure that runs the analysis benchmark.
		go func() {
			// Decrement the WaitGroup counter when the goroutine completes.
			defer wg.Done()

			for t := range traceRequests {
				sklog.Infof("========== Downloading trace %s ==========", t)
				downloadedTrace, err := downloadTrace(t, traceDownloadDir, gs)
				if err != nil {
					sklog.Errorf("Could not download %s: %s", t, err)
					erroredTracesMutex.Lock()
					erroredTraces = append(erroredTraces, t)
					erroredTracesMutex.Unlock()
					continue
				}

				sklog.Infof("========== Processing %s ==========", t)
				if err := runMetricsAnalysisBenchmark(ctx, localOutputDir, downloadedTrace, t); err != nil {
					sklog.Errorf("Error during run_benchmark: %s", err)
				} else {
					atleastOneBenchmarkSucceeded = true
				}
			}
		}()
	}

	// Wait for all spawned goroutines to complete.
	wg.Wait()

	// Summarize errors.
	if len(erroredTraces) > 0 {
		sklog.Error("The following traces could not be downloaded:")
		for _, erroredTrace := range erroredTraces {
			sklog.Errorf("\t%s", erroredTrace)
		}
	}
	if !atleastOneBenchmarkSucceeded {
		return errors.New("Not a single benchmark run was successful. Something is wrong.")
	}

	// If "--output-format=csv" was specified then merge all CSV files and upload.
	if strings.Contains(*benchmarkExtraArgs, "--output-format=csv") {
		// Construct path to CT's python scripts.
		pathToPyFiles := util.GetPathToPyFiles(!*worker_common.Local)
		if err := util.MergeUploadCSVFilesOnWorkers(ctx, localOutputDir, pathToPyFiles, *runID, remoteDir, gs, *startRange, true /* handleStrings */, false /* addRanks */, map[string]map[string]string{} /* pageRankToAdditionalFields */); err != nil {
			return fmt.Errorf("Error while processing withpatch CSV files: %s", err)
		}
	}

	return nil
}

// runMetricsAnalysisBenchmark runs the analysis_metrics_ct benchmark on the provided trace.
func runMetricsAnalysisBenchmark(ctx context.Context, outputPath, downloadedTrace, cloudTraceLink string) error {
	args := []string{
		filepath.Join(util.GetPathToTelemetryBinaries(!*worker_common.Local), util.BINARY_RUN_BENCHMARK),
		util.BENCHMARK_METRICS_ANALYSIS,
		"--local-trace-path", fmt.Sprintf("file://%s", downloadedTrace),
		"--cloud-trace-link", cloudTraceLink,
		"--metric-name", *metricName,
		"--browser", "system", // No browser is brought up but unfortunately this is needed by the framework.
		"--output-dir", filepath.Join(outputPath, getTraceName(downloadedTrace)),
	}
	// Calculate what timeout should be used when executing run_benchmark.
	timeoutSecs := util.GetRunBenchmarkTimeoutValue(*benchmarkExtraArgs, METRICS_BENCHMARK_TIMEOUT_SECS)
	sklog.Infof("Using %d seconds for timeout", timeoutSecs)
	// Remove from benchmarkExtraArgs "special" flags that are recognized by CT but not
	// by the run_benchmark script.
	extraArgs := util.RemoveFlagsFromArgs(*benchmarkExtraArgs, util.RUN_BENCHMARK_TIMEOUT_FLAG, util.MAX_PAGES_PER_BOT)
	// Split extraArgs if not empty and append to args.
	if extraArgs != "" {
		args = append(args, strings.Fields(extraArgs)...)
	}
	// Set the DISPLAY.
	env := []string{
		"DISPLAY=:0",
	}
	// Append the original environment as well.
	for _, e := range os.Environ() {
		env = append(env, e)
	}

	// Create buffer for capturing the stdout and stderr of the benchmark run.
	var b bytes.Buffer
	if _, err := b.WriteString(fmt.Sprintf("========== Stdout and stderr for %s ==========\n", downloadedTrace)); err != nil {
		return fmt.Errorf("Error writing to output buffer: %s", err)
	}
	if err := util.ExecuteCmdWithConfigurableLogging(ctx, "python", args, env, time.Duration(timeoutSecs)*time.Second, &b, &b, false, false); err != nil {
		output, getErr := util.GetRunBenchmarkOutput(b)
		skutil.LogErr(getErr)
		fmt.Println(output)
		return fmt.Errorf("Run benchmark command failed with: %s", err)
	}
	output, err := util.GetRunBenchmarkOutput(b)
	if err != nil {
		return fmt.Errorf("Could not get run benchmark output: %s", err)
	}
	// Print the output and return.
	fmt.Println(output)
	return nil
}

// downloadTrace downloads the traceURL from google storage into the specified
// destDir.
func downloadTrace(traceURL, destDir string, gs *util.GcsUtil) (string, error) {
	traceName := getTraceName(traceURL)
	traceDest := filepath.Join(destDir, traceName)
	if err := gs.DownloadRemoteFileFromBucket(TRACE_OUTPUT_BUCKET, traceName, traceDest); err != nil {
		return "", fmt.Errorf("Error downloading %s from %s to %s: %s", traceName, TRACE_OUTPUT_BUCKET, traceDest, err)
	}
	return traceDest, nil
}

// getTraceName parses the provided traceURI and returns the name of the trace.
// traceURI could be a file path or a URL.
func getTraceName(traceURI string) string {
	traceTokens := strings.Split(traceURI, "/")
	return traceTokens[len(traceTokens)-1]
}

// getClosedChannelOfTraces returns a channel that contains all trace URLs.
func getClosedChannelOfTraces(traces []string) chan string {
	tracesChannel := make(chan string, len(traces))
	for _, t := range traces {
		tracesChannel <- t
	}
	close(tracesChannel)
	return tracesChannel
}

func main() {
	retCode := 0
	if err := metricsAnalysis(); err != nil {
		sklog.Errorf("Error while running metrics analysis: %s", err)
		retCode = 255
	}
	os.Exit(retCode)
}
<commit_msg>[CT] Experiment making metrics analysis serial<commit_after>// Application that does metrics analysis as described in the design doc:
// go/ct_metrics_analysis
//
// Can be tested locally with:
// $ go run go/worker_scripts/metrics_analysis/main.go --start_range=1 --num=3 --run_id=rmistry-test1 --benchmark_extra_args="--output-format=csv" --metric_name="loadingMetric" --logtostderr=true --local
//
package main

import (
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"go.skia.org/infra/go/sklog"

	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/ct/go/worker_scripts/worker_common"
	"go.skia.org/infra/go/common"
	skutil "go.skia.org/infra/go/util"
)

const (
	// The number of goroutines that will run in parallel to download traces and
	// run metrics analysis.
	WORKER_POOL_SIZE = 1

	METRICS_BENCHMARK_TIMEOUT_SECS = 300

	TRACE_OUTPUT_BUCKET = "chrome-telemetry-output"
)

var (
	startRange         = flag.Int("start_range", 1, "The number this worker will run metrics analysis from.")
	num                = flag.Int("num", 100, "The total number of traces to run metrics analysis from the start_range.")
	runID              = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
	benchmarkExtraArgs = flag.String("benchmark_extra_args", "", "The extra arguments that are passed to the specified benchmark.")
	metricName         = flag.String("metric_name", "", "The metric to parse the traces with. Eg: loadingMetric")
)

func metricsAnalysis() error {
	defer common.LogPanic()
	worker_common.Init()
	if !*worker_common.Local {
		defer util.CleanTmpDir()
	}
	defer util.TimeTrack(time.Now(), "Metrics Analysis")
	defer sklog.Flush()

	// Validate required arguments.
	if *runID == "" {
		return errors.New("Must specify --run_id")
	}
	if *metricName == "" {
		return errors.New("Must specify --metric_name")
	}

	ctx := context.Background()

	// Instantiate GcsUtil object.
	gs, err := util.NewGcsUtil(nil)
	if err != nil {
		return fmt.Errorf("GcsUtil instantiation failed: %s", err)
	}

	// Download the trace URLs for this run from Google storage.
	tracesFilename := *runID + ".traces.csv"
	tmpDir, err := ioutil.TempDir(util.PagesetsDir, "traces")
	if err != nil {
		return fmt.Errorf("Could not create tmpdir: %s", err)
	}
	defer skutil.RemoveAll(tmpDir)
	remotePatchesDir := filepath.Join(util.BenchmarkRunsDir, *runID)

	// Download traces.
	if _, err := util.DownloadPatch(filepath.Join(tmpDir, tracesFilename), filepath.Join(remotePatchesDir, tracesFilename), gs); err != nil {
		return fmt.Errorf("Could not download %s: %s", tracesFilename, err)
	}
	traces, err := util.GetCustomPages(filepath.Join(tmpDir, tracesFilename))
	if err != nil {
		return fmt.Errorf("Could not read custom traces file %s: %s", tracesFilename, err)
	}
	if len(traces) == 0 {
		return errors.New(fmt.Sprintf("No traces found in %s", tracesFilename))
	}
	traces = util.GetCustomPagesWithinRange(*startRange, *num, traces)
	sklog.Infof("Using %d traces", len(traces))

	// Establish output paths for trace downloads and metrics.
	traceDownloadDir := filepath.Join(util.StorageDir, util.TraceDownloadsDir, *runID)
	skutil.RemoveAll(traceDownloadDir)
	skutil.MkdirAll(traceDownloadDir, 0700)
	defer skutil.RemoveAll(traceDownloadDir)

	localOutputDir := filepath.Join(util.StorageDir, util.BenchmarkRunsDir, *runID)
	skutil.RemoveAll(localOutputDir)
	skutil.MkdirAll(localOutputDir, 0700)
	defer skutil.RemoveAll(localOutputDir)
	remoteDir := filepath.Join(util.BenchmarkRunsDir, *runID)

	sklog.Infof("===== Going to run the task with %d parallel goroutines =====", WORKER_POOL_SIZE)
	// Create channel that contains all trace ULs. This channel will
	// be consumed by the worker pool.
	traceRequests := getClosedChannelOfTraces(traces)

	var wg sync.WaitGroup

	// If not a single benchmark run succeeds then throw at error at the end.
	atleastOneBenchmarkSucceeded := false
	// Gather traceURLs that could not be downloaded.
	erroredTraces := []string{}
	// Mutex to control access to the above slice.
	var erroredTracesMutex sync.Mutex

	// Loop through workers in the worker pool.
	for i := 0; i < WORKER_POOL_SIZE; i++ {
		// Increment the WaitGroup counter.
		wg.Add(1)

		// Create and run a goroutine closure that runs the analysis benchmark.
		go func() {
			// Decrement the WaitGroup counter when the goroutine completes.
			defer wg.Done()

			for t := range traceRequests {
				sklog.Infof("========== Downloading trace %s ==========", t)
				downloadedTrace, err := downloadTrace(t, traceDownloadDir, gs)
				if err != nil {
					sklog.Errorf("Could not download %s: %s", t, err)
					erroredTracesMutex.Lock()
					erroredTraces = append(erroredTraces, t)
					erroredTracesMutex.Unlock()
					continue
				}

				sklog.Infof("========== Processing %s ==========", t)
				if err := runMetricsAnalysisBenchmark(ctx, localOutputDir, downloadedTrace, t); err != nil {
					sklog.Errorf("Error during run_benchmark: %s", err)
				} else {
					atleastOneBenchmarkSucceeded = true
				}
			}
		}()
	}

	// Wait for all spawned goroutines to complete.
	wg.Wait()

	// Summarize errors.
	if len(erroredTraces) > 0 {
		sklog.Error("The following traces could not be downloaded:")
		for _, erroredTrace := range erroredTraces {
			sklog.Errorf("\t%s", erroredTrace)
		}
	}
	if !atleastOneBenchmarkSucceeded {
		return errors.New("Not a single benchmark run was successful. Something is wrong.")
	}

	// If "--output-format=csv" was specified then merge all CSV files and upload.
	if strings.Contains(*benchmarkExtraArgs, "--output-format=csv") {
		// Construct path to CT's python scripts.
		pathToPyFiles := util.GetPathToPyFiles(!*worker_common.Local)
		if err := util.MergeUploadCSVFilesOnWorkers(ctx, localOutputDir, pathToPyFiles, *runID, remoteDir, gs, *startRange, true /* handleStrings */, false /* addRanks */, map[string]map[string]string{} /* pageRankToAdditionalFields */); err != nil {
			return fmt.Errorf("Error while processing withpatch CSV files: %s", err)
		}
	}

	return nil
}

// runMetricsAnalysisBenchmark runs the analysis_metrics_ct benchmark on the provided trace.
func runMetricsAnalysisBenchmark(ctx context.Context, outputPath, downloadedTrace, cloudTraceLink string) error {
	args := []string{
		filepath.Join(util.GetPathToTelemetryBinaries(!*worker_common.Local), util.BINARY_RUN_BENCHMARK),
		util.BENCHMARK_METRICS_ANALYSIS,
		"--local-trace-path", fmt.Sprintf("file://%s", downloadedTrace),
		"--cloud-trace-link", cloudTraceLink,
		"--metric-name", *metricName,
		"--browser", "system", // No browser is brought up but unfortunately this is needed by the framework.
		"--output-dir", filepath.Join(outputPath, getTraceName(downloadedTrace)),
	}
	// Calculate what timeout should be used when executing run_benchmark.
	timeoutSecs := util.GetRunBenchmarkTimeoutValue(*benchmarkExtraArgs, METRICS_BENCHMARK_TIMEOUT_SECS)
	sklog.Infof("Using %d seconds for timeout", timeoutSecs)
	// Remove from benchmarkExtraArgs "special" flags that are recognized by CT but not
	// by the run_benchmark script.
	extraArgs := util.RemoveFlagsFromArgs(*benchmarkExtraArgs, util.RUN_BENCHMARK_TIMEOUT_FLAG, util.MAX_PAGES_PER_BOT)
	// Split extraArgs if not empty and append to args.
	if extraArgs != "" {
		args = append(args, strings.Fields(extraArgs)...)
	}
	// Set the DISPLAY.
	env := []string{
		"DISPLAY=:0",
	}
	// Append the original environment as well.
	for _, e := range os.Environ() {
		env = append(env, e)
	}

	// Create buffer for capturing the stdout and stderr of the benchmark run.
	var b bytes.Buffer
	if _, err := b.WriteString(fmt.Sprintf("========== Stdout and stderr for %s ==========\n", downloadedTrace)); err != nil {
		return fmt.Errorf("Error writing to output buffer: %s", err)
	}
	if err := util.ExecuteCmdWithConfigurableLogging(ctx, "python", args, env, time.Duration(timeoutSecs)*time.Second, &b, &b, false, false); err != nil {
		output, getErr := util.GetRunBenchmarkOutput(b)
		skutil.LogErr(getErr)
		fmt.Println(output)
		return fmt.Errorf("Run benchmark command failed with: %s", err)
	}
	output, err := util.GetRunBenchmarkOutput(b)
	if err != nil {
		return fmt.Errorf("Could not get run benchmark output: %s", err)
	}
	// Print the output and return.
	fmt.Println(output)
	return nil
}

// downloadTrace downloads the traceURL from google storage into the specified
// destDir.
func downloadTrace(traceURL, destDir string, gs *util.GcsUtil) (string, error) {
	traceName := getTraceName(traceURL)
	traceDest := filepath.Join(destDir, traceName)
	if err := gs.DownloadRemoteFileFromBucket(TRACE_OUTPUT_BUCKET, traceName, traceDest); err != nil {
		return "", fmt.Errorf("Error downloading %s from %s to %s: %s", traceName, TRACE_OUTPUT_BUCKET, traceDest, err)
	}
	return traceDest, nil
}

// getTraceName parses the provided traceURI and returns the name of the trace.
// traceURI could be a file path or a URL.
func getTraceName(traceURI string) string {
	traceTokens := strings.Split(traceURI, "/")
	return traceTokens[len(traceTokens)-1]
}

// getClosedChannelOfTraces returns a channel that contains all trace URLs.
func getClosedChannelOfTraces(traces []string) chan string {
	tracesChannel := make(chan string, len(traces))
	for _, t := range traces {
		tracesChannel <- t
	}
	close(tracesChannel)
	return tracesChannel
}

func main() {
	retCode := 0
	if err := metricsAnalysis(); err != nil {
		sklog.Errorf("Error while running metrics analysis: %s", err)
		retCode = 255
	}
	os.Exit(retCode)
}
<|endoftext|>


<commit_before>package container

import (
	"strings"

	executorpkg "github.com/docker/docker/daemon/cluster/executor"
	clustertypes "github.com/docker/docker/daemon/cluster/provider"
	"github.com/docker/engine-api/types"
	"github.com/docker/engine-api/types/network"
	networktypes "github.com/docker/libnetwork/types"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"golang.org/x/net/context"
)

type executor struct {
	backend executorpkg.Backend
}

// NewExecutor returns an executor from the docker client.
func NewExecutor(b executorpkg.Backend) exec.Executor {
	return &executor{
		backend: b,
	}
}

// Describe returns the underlying node description from the docker client.
func (e *executor) Describe(ctx context.Context) (*api.NodeDescription, error) {
	info, err := e.backend.SystemInfo()
	if err != nil {
		return nil, err
	}

	var plugins []api.PluginDescription
	addPlugins := func(typ string, names []string) {
		for _, name := range names {
			plugins = append(plugins, api.PluginDescription{
				Type: typ,
				Name: name,
			})
		}
	}

	addPlugins("Volume", info.Plugins.Volume)
	// Add builtin driver "overlay" (the only builtin multi-host driver) to
	// the plugin list by default.
	addPlugins("Network", append([]string{"overlay"}, info.Plugins.Network...))
	addPlugins("Authorization", info.Plugins.Authorization)

	// parse []string labels into a map[string]string
	labels := map[string]string{}
	for _, l := range info.Labels {
		stringSlice := strings.SplitN(l, "=", 2)
		// this will take the last value in the list for a given key
		// ideally, one shouldn't assign multiple values to the same key
		if len(stringSlice) > 1 {
			labels[stringSlice[0]] = stringSlice[1]
		}
	}

	description := &api.NodeDescription{
		Hostname: info.Name,
		Platform: &api.Platform{
			Architecture: info.Architecture,
			OS:           info.OSType,
		},
		Engine: &api.EngineDescription{
			EngineVersion: info.ServerVersion,
			Labels:        labels,
			Plugins:       plugins,
		},
		Resources: &api.Resources{
			NanoCPUs:    int64(info.NCPU) * 1e9,
			MemoryBytes: info.MemTotal,
		},
	}

	return description, nil
}

func (e *executor) Configure(ctx context.Context, node *api.Node) error {
	na := node.Attachment
	if na == nil {
		return nil
	}

	options := types.NetworkCreate{
		Driver: na.Network.DriverState.Name,
		IPAM: network.IPAM{
			Driver: na.Network.IPAM.Driver.Name,
		},
		Options:        na.Network.DriverState.Options,
		CheckDuplicate: true,
	}

	for _, ic := range na.Network.IPAM.Configs {
		c := network.IPAMConfig{
			Subnet:  ic.Subnet,
			IPRange: ic.Range,
			Gateway: ic.Gateway,
		}
		options.IPAM.Config = append(options.IPAM.Config, c)
	}

	return e.backend.SetupIngress(clustertypes.NetworkCreateRequest{
		na.Network.ID,
		types.NetworkCreateRequest{
			Name:          na.Network.Spec.Annotations.Name,
			NetworkCreate: options,
		},
	}, na.Addresses[0])
}

// Controller returns a docker container runner.
func (e *executor) Controller(t *api.Task) (exec.Controller, error) {
	ctlr, err := newController(e.backend, t)
	if err != nil {
		return nil, err
	}

	return ctlr, nil
}

func (e *executor) SetNetworkBootstrapKeys(keys []*api.EncryptionKey) error {
	nwKeys := []*networktypes.EncryptionKey{}
	for _, key := range keys {
		nwKey := &networktypes.EncryptionKey{
			Subsystem:   key.Subsystem,
			Algorithm:   int32(key.Algorithm),
			Key:         make([]byte, len(key.Key)),
			LamportTime: key.LamportTime,
		}
		copy(nwKey.Key, key.Key)
		nwKeys = append(nwKeys, nwKey)
	}
	e.backend.SetNetworkBootstrapKeys(nwKeys)

	return nil
}
<commit_msg>Sort plugin names in node description<commit_after>package container

import (
	"sort"
	"strings"

	executorpkg "github.com/docker/docker/daemon/cluster/executor"
	clustertypes "github.com/docker/docker/daemon/cluster/provider"
	"github.com/docker/engine-api/types"
	"github.com/docker/engine-api/types/network"
	networktypes "github.com/docker/libnetwork/types"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"golang.org/x/net/context"
)

type executor struct {
	backend executorpkg.Backend
}

// NewExecutor returns an executor from the docker client.
func NewExecutor(b executorpkg.Backend) exec.Executor {
	return &executor{
		backend: b,
	}
}

// Describe returns the underlying node description from the docker client.
func (e *executor) Describe(ctx context.Context) (*api.NodeDescription, error) {
	info, err := e.backend.SystemInfo()
	if err != nil {
		return nil, err
	}

	var plugins []api.PluginDescription
	addPlugins := func(typ string, names []string) {
		for _, name := range names {
			plugins = append(plugins, api.PluginDescription{
				Type: typ,
				Name: name,
			})
		}
	}

	addPlugins("Volume", info.Plugins.Volume)
	// Add builtin driver "overlay" (the only builtin multi-host driver) to
	// the plugin list by default.
	addPlugins("Network", append([]string{"overlay"}, info.Plugins.Network...))
	addPlugins("Authorization", info.Plugins.Authorization)

	sort.Sort(sortedPlugins(plugins))

	// parse []string labels into a map[string]string
	labels := map[string]string{}
	for _, l := range info.Labels {
		stringSlice := strings.SplitN(l, "=", 2)
		// this will take the last value in the list for a given key
		// ideally, one shouldn't assign multiple values to the same key
		if len(stringSlice) > 1 {
			labels[stringSlice[0]] = stringSlice[1]
		}
	}

	description := &api.NodeDescription{
		Hostname: info.Name,
		Platform: &api.Platform{
			Architecture: info.Architecture,
			OS:           info.OSType,
		},
		Engine: &api.EngineDescription{
			EngineVersion: info.ServerVersion,
			Labels:        labels,
			Plugins:       plugins,
		},
		Resources: &api.Resources{
			NanoCPUs:    int64(info.NCPU) * 1e9,
			MemoryBytes: info.MemTotal,
		},
	}

	return description, nil
}

func (e *executor) Configure(ctx context.Context, node *api.Node) error {
	na := node.Attachment
	if na == nil {
		return nil
	}

	options := types.NetworkCreate{
		Driver: na.Network.DriverState.Name,
		IPAM: network.IPAM{
			Driver: na.Network.IPAM.Driver.Name,
		},
		Options:        na.Network.DriverState.Options,
		CheckDuplicate: true,
	}

	for _, ic := range na.Network.IPAM.Configs {
		c := network.IPAMConfig{
			Subnet:  ic.Subnet,
			IPRange: ic.Range,
			Gateway: ic.Gateway,
		}
		options.IPAM.Config = append(options.IPAM.Config, c)
	}

	return e.backend.SetupIngress(clustertypes.NetworkCreateRequest{
		na.Network.ID,
		types.NetworkCreateRequest{
			Name:          na.Network.Spec.Annotations.Name,
			NetworkCreate: options,
		},
	}, na.Addresses[0])
}

// Controller returns a docker container runner.
func (e *executor) Controller(t *api.Task) (exec.Controller, error) {
	ctlr, err := newController(e.backend, t)
	if err != nil {
		return nil, err
	}

	return ctlr, nil
}

func (e *executor) SetNetworkBootstrapKeys(keys []*api.EncryptionKey) error {
	nwKeys := []*networktypes.EncryptionKey{}
	for _, key := range keys {
		nwKey := &networktypes.EncryptionKey{
			Subsystem:   key.Subsystem,
			Algorithm:   int32(key.Algorithm),
			Key:         make([]byte, len(key.Key)),
			LamportTime: key.LamportTime,
		}
		copy(nwKey.Key, key.Key)
		nwKeys = append(nwKeys, nwKey)
	}
	e.backend.SetNetworkBootstrapKeys(nwKeys)

	return nil
}

type sortedPlugins []api.PluginDescription

func (sp sortedPlugins) Len() int { return len(sp) }

func (sp sortedPlugins) Swap(i, j int) { sp[i], sp[j] = sp[j], sp[i] }

func (sp sortedPlugins) Less(i, j int) bool {
	if sp[i].Type != sp[j].Type {
		return sp[i].Type < sp[j].Type
	}
	return sp[i].Name < sp[j].Name
}
<|endoftext|>


<commit_before>//
// Copyright 2016 Rackspace
//
// 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 poller

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/racker/rackspace-monitoring-poller/config"
	"github.com/rackerlabs/go-connect-tunnel"
)

var (
	metricsConnectionsAttempt = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "poller",
			Subsystem: "connection",
			Name:      "attempts",
			Help:      "Conveys the number of connection attempts per remote address",
		},
		[]string{
			metricLabelAddress,
		},
	)
	metricsConnectionsConnected = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "poller",
			Subsystem: "connection",
			Name:      "connected",
			Help:      "Conveys the number of successful connections per remote address",
		},
		[]string{
			metricLabelAddress,
		},
	)
	metricsConnectionsDisconnected = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "poller",
			Subsystem: "connection",
			Name:      "disconnected",
			Help:      "Conveys the number of disconnects per remote address",
		},
		[]string{
			metricLabelAddress,
		},
	)
)

// EleConnection implements Connection
// See Connection interface for more information
type EleConnection struct {
	checksReconciler ChecksReconciler

	session Session
	conn    io.ReadWriteCloser

	address string
	guid    string

	connectionTimeout time.Duration

	authenticated chan struct{}
}

func init() {
	metricsRegistry.MustRegister(metricsConnectionsAttempt)
	metricsRegistry.MustRegister(metricsConnectionsConnected)
	metricsRegistry.MustRegister(metricsConnectionsDisconnected)
}

// NewConnection instantiates a new EleConnection
// It sets up the address, unique guid, and connection timeout
// for this conneciton
func NewConnection(address string, guid string, checksReconciler ChecksReconciler) Connection {
	return &EleConnection{
		address:           address,
		guid:              guid,
		checksReconciler:  checksReconciler,
		connectionTimeout: time.Duration(10) * time.Second,
		authenticated:     make(chan struct{}, 1),
	}
}

// GetGUID retrieves connection's guid
func (conn *EleConnection) GetGUID() string {
	return conn.guid
}

// GetFarendWriter gets a writer directed towards the endpoint server
func (conn *EleConnection) GetFarendWriter() io.Writer {
	return conn.conn
}

// GetFarendReader gets a reader to consume from the endpoint server
func (conn *EleConnection) GetFarendReader() io.Reader {
	return conn.conn
}

// GetSession retrieves connection's session
func (conn *EleConnection) GetSession() Session {
	return conn.session
}

func (conn *EleConnection) GetClockOffset() int64 {
	return conn.session.GetClockOffset()
}

func (conn *EleConnection) GetLatency() int64 {
	return conn.session.GetLatency()
}

func (conn *EleConnection) HasLatencyMeasurements() bool {
	return conn.session.HasLatencyMeasurements()
}

// SetReadDeadline sets up the read deadline for a socket.
// read deadline - time spent reading response for future calls
func (conn *EleConnection) SetReadDeadline(deadline time.Time) {
	socket, ok := conn.conn.(*net.TCPConn)
	if ok {
		socket.SetReadDeadline(deadline)
	}
}

// SetWriteDeadline sets up the write deadline for a socket.
// write deadline - time spent writing request to socket for future calls
func (conn *EleConnection) SetWriteDeadline(deadline time.Time) {
	socket, ok := conn.conn.(*net.TCPConn)
	if ok {
		socket.SetWriteDeadline(deadline)
	}
}

func (conn *EleConnection) GetLogPrefix() string {
	return conn.address
}

func (conn *EleConnection) String() string {
	return conn.address
}

// Connect sets up a tcp connection with connection defined address
// and passed in tlsConfig
// If context is not set, ErrUndefinedContext is returned
// The end result of this function is a usable connection ready to
// send data.
func (conn *EleConnection) Connect(ctx context.Context, config *config.Config, tlsConfig *tls.Config) error {
	if ctx == nil {
		return ErrUndefinedContext
	}
	log.WithFields(log.Fields{
		"prefix":  conn.GetLogPrefix(),
		"timeout": conn.connectionTimeout,
	}).Info("Connecting to agent/poller endpoint")

	counter, err := metricsConnectionsAttempt.GetMetricWithLabelValues(conn.address)
	if err == nil {
		counter.Inc()
	} else {
		log.WithField("err", err).Debug("Failed to get metricsConnectionsAttempt counter")
	}

	if err = conn.dial(config, tlsConfig); err != nil {
		return err
	}

	conn.session = NewSession(ctx, conn, conn.checksReconciler, config)
	log.WithFields(log.Fields{
		"prefix":         conn.GetLogPrefix(),
		"remote_address": conn.address,
	}).Info("Connected")

	counter, err = metricsConnectionsConnected.GetMetricWithLabelValues(conn.address)
	if err == nil {
		counter.Inc()
	} else {
		log.WithField("err", err).Debug("Failed to get metricsConnectionsConnected counter")
	}

	return nil
}

func (conn *EleConnection) dial(config *config.Config, tlsConfig *tls.Config) (err error) {
	nd := net.Dialer{Timeout: conn.connectionTimeout}

	var proxyConn net.Conn
	if config.ProxyUrl != nil {
		log.WithFields(log.Fields{
			"prefix": conn.GetLogPrefix(),
			"proxy":  config.ProxyUrl,
		}).Info("Connecting to agent/poller endpoint via proxy")
		proxyConn, err = tunnel.DialViaProxy(config.ProxyUrl, conn.address)
		if err != nil {
			return err
		}
	}

	if tlsConfig != nil {
		if proxyConn != nil {
			conn.conn = tls.Client(proxyConn, tlsConfig)
		} else {
			conn.conn, err = tls.DialWithDialer(&nd, "tcp", conn.address, tlsConfig)
			if err != nil {
				return err
			}
		}
	} else {
		if proxyConn != nil {
			conn.conn = proxyConn;
		} else {
			conn.conn, err = nd.Dial("tcp", conn.address)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Close closes the session
func (conn *EleConnection) Close() {
	if conn.conn != nil {
		log.WithFields(log.Fields{
			"prefix":         conn.GetLogPrefix(),
			"remote_address": conn.address,
		}).Info("Disconnected")

		counter, err := metricsConnectionsDisconnected.GetMetricWithLabelValues(conn.address)
		if err == nil {
			counter.Inc()
		} else {
			log.WithField("err", err).Debug("Failed to get metricsConnectionsDisconnected counter")
		}

		conn.conn.Close()
		conn.conn = nil
	}
	conn.GetSession().Close()
}

// Wait returns a channel that is populated when the connection is finished or closed.
func (conn *EleConnection) Done() <-chan struct{} {
	return conn.GetSession().Done()
}

func (conn *EleConnection) SetAuthenticated() {
	close(conn.authenticated)
}

func (conn *EleConnection) Authenticated() <-chan struct{} {
	return conn.authenticated
}
<commit_msg>Don't log proxy credentials (#191)<commit_after>//
// Copyright 2016 Rackspace
//
// 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 poller

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/racker/rackspace-monitoring-poller/config"
	"github.com/rackerlabs/go-connect-tunnel"
)

var (
	metricsConnectionsAttempt = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "poller",
			Subsystem: "connection",
			Name:      "attempts",
			Help:      "Conveys the number of connection attempts per remote address",
		},
		[]string{
			metricLabelAddress,
		},
	)
	metricsConnectionsConnected = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "poller",
			Subsystem: "connection",
			Name:      "connected",
			Help:      "Conveys the number of successful connections per remote address",
		},
		[]string{
			metricLabelAddress,
		},
	)
	metricsConnectionsDisconnected = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "poller",
			Subsystem: "connection",
			Name:      "disconnected",
			Help:      "Conveys the number of disconnects per remote address",
		},
		[]string{
			metricLabelAddress,
		},
	)
)

// EleConnection implements Connection
// See Connection interface for more information
type EleConnection struct {
	checksReconciler ChecksReconciler

	session Session
	conn    io.ReadWriteCloser

	address string
	guid    string

	connectionTimeout time.Duration

	authenticated chan struct{}
}

func init() {
	metricsRegistry.MustRegister(metricsConnectionsAttempt)
	metricsRegistry.MustRegister(metricsConnectionsConnected)
	metricsRegistry.MustRegister(metricsConnectionsDisconnected)
}

// NewConnection instantiates a new EleConnection
// It sets up the address, unique guid, and connection timeout
// for this conneciton
func NewConnection(address string, guid string, checksReconciler ChecksReconciler) Connection {
	return &EleConnection{
		address:           address,
		guid:              guid,
		checksReconciler:  checksReconciler,
		connectionTimeout: time.Duration(10) * time.Second,
		authenticated:     make(chan struct{}, 1),
	}
}

// GetGUID retrieves connection's guid
func (conn *EleConnection) GetGUID() string {
	return conn.guid
}

// GetFarendWriter gets a writer directed towards the endpoint server
func (conn *EleConnection) GetFarendWriter() io.Writer {
	return conn.conn
}

// GetFarendReader gets a reader to consume from the endpoint server
func (conn *EleConnection) GetFarendReader() io.Reader {
	return conn.conn
}

// GetSession retrieves connection's session
func (conn *EleConnection) GetSession() Session {
	return conn.session
}

func (conn *EleConnection) GetClockOffset() int64 {
	return conn.session.GetClockOffset()
}

func (conn *EleConnection) GetLatency() int64 {
	return conn.session.GetLatency()
}

func (conn *EleConnection) HasLatencyMeasurements() bool {
	return conn.session.HasLatencyMeasurements()
}

// SetReadDeadline sets up the read deadline for a socket.
// read deadline - time spent reading response for future calls
func (conn *EleConnection) SetReadDeadline(deadline time.Time) {
	socket, ok := conn.conn.(*net.TCPConn)
	if ok {
		socket.SetReadDeadline(deadline)
	}
}

// SetWriteDeadline sets up the write deadline for a socket.
// write deadline - time spent writing request to socket for future calls
func (conn *EleConnection) SetWriteDeadline(deadline time.Time) {
	socket, ok := conn.conn.(*net.TCPConn)
	if ok {
		socket.SetWriteDeadline(deadline)
	}
}

func (conn *EleConnection) GetLogPrefix() string {
	return conn.address
}

func (conn *EleConnection) String() string {
	return conn.address
}

// Connect sets up a tcp connection with connection defined address
// and passed in tlsConfig
// If context is not set, ErrUndefinedContext is returned
// The end result of this function is a usable connection ready to
// send data.
func (conn *EleConnection) Connect(ctx context.Context, config *config.Config, tlsConfig *tls.Config) error {
	if ctx == nil {
		return ErrUndefinedContext
	}
	log.WithFields(log.Fields{
		"prefix":  conn.GetLogPrefix(),
		"timeout": conn.connectionTimeout,
	}).Info("Connecting to agent/poller endpoint")

	counter, err := metricsConnectionsAttempt.GetMetricWithLabelValues(conn.address)
	if err == nil {
		counter.Inc()
	} else {
		log.WithField("err", err).Debug("Failed to get metricsConnectionsAttempt counter")
	}

	if err = conn.dial(config, tlsConfig); err != nil {
		return err
	}

	conn.session = NewSession(ctx, conn, conn.checksReconciler, config)
	log.WithFields(log.Fields{
		"prefix":         conn.GetLogPrefix(),
		"remote_address": conn.address,
	}).Info("Connected")

	counter, err = metricsConnectionsConnected.GetMetricWithLabelValues(conn.address)
	if err == nil {
		counter.Inc()
	} else {
		log.WithField("err", err).Debug("Failed to get metricsConnectionsConnected counter")
	}

	return nil
}

func (conn *EleConnection) dial(config *config.Config, tlsConfig *tls.Config) (err error) {
	nd := net.Dialer{Timeout: conn.connectionTimeout}

	var proxyConn net.Conn
	if config.ProxyUrl != nil {
		log.WithFields(log.Fields{
			"prefix": conn.GetLogPrefix(),
			"proxy":  config.ProxyUrl.Host,
		}).Info("Connecting to agent/poller endpoint via proxy")
		proxyConn, err = tunnel.DialViaProxy(config.ProxyUrl, conn.address)
		if err != nil {
			return err
		}
	}

	if tlsConfig != nil {
		if proxyConn != nil {
			conn.conn = tls.Client(proxyConn, tlsConfig)
		} else {
			conn.conn, err = tls.DialWithDialer(&nd, "tcp", conn.address, tlsConfig)
			if err != nil {
				return err
			}
		}
	} else {
		if proxyConn != nil {
			conn.conn = proxyConn;
		} else {
			conn.conn, err = nd.Dial("tcp", conn.address)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Close closes the session
func (conn *EleConnection) Close() {
	if conn.conn != nil {
		log.WithFields(log.Fields{
			"prefix":         conn.GetLogPrefix(),
			"remote_address": conn.address,
		}).Info("Disconnected")

		counter, err := metricsConnectionsDisconnected.GetMetricWithLabelValues(conn.address)
		if err == nil {
			counter.Inc()
		} else {
			log.WithField("err", err).Debug("Failed to get metricsConnectionsDisconnected counter")
		}

		conn.conn.Close()
		conn.conn = nil
	}
	conn.GetSession().Close()
}

// Wait returns a channel that is populated when the connection is finished or closed.
func (conn *EleConnection) Done() <-chan struct{} {
	return conn.GetSession().Done()
}

func (conn *EleConnection) SetAuthenticated() {
	close(conn.authenticated)
}

func (conn *EleConnection) Authenticated() <-chan struct{} {
	return conn.authenticated
}
<|endoftext|>