<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package e2e

import (
	"fmt"
	"io/ioutil"
	"net"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
)

const (
	podName = "pfpod"
)

// TODO support other ports besides 80
var portForwardRegexp = regexp.MustCompile("Forwarding from 127.0.0.1:([0-9]+) -> 80")

func pfPod(expectedClientData, chunks, chunkSize, chunkIntervalMillis string) *api.Pod {
	return &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Name:   podName,
			Labels: map[string]string{"name": podName},
		},
		Spec: api.PodSpec{
			Containers: []api.Container{
				{
					Name:  "portforwardtester",
					Image: "gcr.io/google_containers/portforwardtester:1.0",
					Env: []api.EnvVar{
						{
							Name:  "BIND_PORT",
							Value: "80",
						},
						{
							Name:  "EXPECTED_CLIENT_DATA",
							Value: expectedClientData,
						},
						{
							Name:  "CHUNKS",
							Value: chunks,
						},
						{
							Name:  "CHUNK_SIZE",
							Value: chunkSize,
						},
						{
							Name:  "CHUNK_INTERVAL",
							Value: chunkIntervalMillis,
						},
					},
				},
			},
			RestartPolicy: api.RestartPolicyNever,
		},
	}
}

type portForwardCommand struct {
	cmd  *exec.Cmd
	port int
}

// Stop attempts to gracefully stop `kubectl port-forward`, only killing it if necessary.
// This helps avoid spdy goroutine leaks in the Kubelet.
func (c *portForwardCommand) Stop() {
	// SIGINT signals that kubectl port-forward should gracefully terminate
	if err := c.cmd.Process.Signal(syscall.SIGINT); err != nil {
		framework.Logf("error sending SIGINT to kubectl port-forward: %v", err)
	}

	// try to wait for a clean exit
	done := make(chan error)
	go func() {
		done <- c.cmd.Wait()
	}()

	expired := time.NewTimer(wait.ForeverTestTimeout)
	defer expired.Stop()

	select {
	case err := <-done:
		if err == nil {
			// success
			return
		}
		framework.Logf("error waiting for kubectl port-forward to exit: %v", err)
	case <-expired.C:
		framework.Logf("timed out waiting for kubectl port-forward to exit")
	}

	framework.Logf("trying to forcibly kill kubectl port-forward")
	framework.TryKill(c.cmd)
}

func runPortForward(ns, podName string, port int) *portForwardCommand {
	cmd := framework.KubectlCmd("port-forward", fmt.Sprintf("--namespace=%v", ns), podName, fmt.Sprintf(":%d", port))
	// This is somewhat ugly but is the only way to retrieve the port that was picked
	// by the port-forward command. We don't want to hard code the port as we have no
	// way of guaranteeing we can pick one that isn't in use, particularly on Jenkins.
	framework.Logf("starting port-forward command and streaming output")
	stdout, _, err := framework.StartCmdAndStreamOutput(cmd)
	if err != nil {
		framework.Failf("Failed to start port-forward command: %v", err)
	}

	buf := make([]byte, 128)
	var n int
	framework.Logf("reading from `kubectl port-forward` command's stdout")
	if n, err = stdout.Read(buf); err != nil {
		framework.Failf("Failed to read from kubectl port-forward stdout: %v", err)
	}
	portForwardOutput := string(buf[:n])
	match := portForwardRegexp.FindStringSubmatch(portForwardOutput)
	if len(match) != 2 {
		framework.Failf("Failed to parse kubectl port-forward output: %s", portForwardOutput)
	}

	listenPort, err := strconv.Atoi(match[1])
	if err != nil {
		framework.Failf("Error converting %s to an int: %v", match[1], err)
	}

	return &portForwardCommand{
		cmd:  cmd,
		port: listenPort,
	}
}

var _ = framework.KubeDescribe("Port forwarding", func() {
	f := framework.NewDefaultFramework("port-forwarding")

	framework.KubeDescribe("With a server that expects a client request", func() {
		It("should support a client that connects, sends no data, and disconnects [Conformance]", func() {
			By("creating the target pod")
			pod := pfPod("abc", "1", "1", "1")
			if _, err := f.Client.Pods(f.Namespace.Name).Create(pod); err != nil {
				framework.Failf("Couldn't create pod: %v", err)
			}
			if err := f.WaitForPodRunning(pod.Name); err != nil {
				framework.Failf("Pod did not start running: %v", err)
			}

			By("Running 'kubectl port-forward'")
			cmd := runPortForward(f.Namespace.Name, pod.Name, 80)
			defer cmd.Stop()

			By("Dialing the local port")
			conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.port))
			if err != nil {
				framework.Failf("Couldn't connect to port %d: %v", cmd.port, err)
			}

			By("Closing the connection to the local port")
			conn.Close()

			By("Waiting for the target pod to stop running")
			if err := f.WaitForPodNoLongerRunning(pod.Name); err != nil {
				framework.Failf("Pod did not stop running: %v", err)
			}

			By("Retrieving logs from the target pod")
			logOutput, err := framework.GetPodLogs(f.Client, f.Namespace.Name, pod.Name, "portforwardtester")
			if err != nil {
				framework.Failf("Error retrieving logs: %v", err)
			}

			By("Verifying logs")
			verifyLogMessage(logOutput, "Accepted client connection")
			verifyLogMessage(logOutput, "Expected to read 3 bytes from client, but got 0 instead")
		})

		It("should support a client that connects, sends data, and disconnects [Conformance]", func() {
			By("creating the target pod")
			pod := pfPod("abc", "10", "10", "100")
			if _, err := f.Client.Pods(f.Namespace.Name).Create(pod); err != nil {
				framework.Failf("Couldn't create pod: %v", err)
			}
			if err := f.WaitForPodRunning(pod.Name); err != nil {
				framework.Failf("Pod did not start running: %v", err)
			}

			By("Running 'kubectl port-forward'")
			cmd := runPortForward(f.Namespace.Name, pod.Name, 80)
			defer cmd.Stop()

			By("Dialing the local port")
			addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.port))
			if err != nil {
				framework.Failf("Error resolving tcp addr: %v", err)
			}
			conn, err := net.DialTCP("tcp", nil, addr)
			if err != nil {
				framework.Failf("Couldn't connect to port %d: %v", cmd.port, err)
			}
			defer func() {
				By("Closing the connection to the local port")
				conn.Close()
			}()

			By("Sending the expected data to the local port")
			fmt.Fprint(conn, "abc")

			By("Closing the write half of the client's connection")
			conn.CloseWrite()

			By("Reading data from the local port")
			fromServer, err := ioutil.ReadAll(conn)
			if err != nil {
				framework.Failf("Unexpected error reading data from the server: %v", err)
			}

			if e, a := strings.Repeat("x", 100), string(fromServer); e != a {
				framework.Failf("Expected %q from server, got %q", e, a)
			}

			By("Waiting for the target pod to stop running")
			if err := f.WaitForPodNoLongerRunning(pod.Name); err != nil {
				framework.Failf("Pod did not stop running: %v", err)
			}

			By("Retrieving logs from the target pod")
			logOutput, err := framework.GetPodLogs(f.Client, f.Namespace.Name, pod.Name, "portforwardtester")
			if err != nil {
				framework.Failf("Error retrieving logs: %v", err)
			}

			By("Verifying logs")
			verifyLogMessage(logOutput, "^Accepted client connection$")
			verifyLogMessage(logOutput, "^Received expected client data$")
			verifyLogMessage(logOutput, "^Done$")
		})
	})
	framework.KubeDescribe("With a server that expects no client request", func() {
		It("should support a client that connects, sends no data, and disconnects [Conformance]", func() {
			By("creating the target pod")
			pod := pfPod("", "10", "10", "100")
			if _, err := f.Client.Pods(f.Namespace.Name).Create(pod); err != nil {
				framework.Failf("Couldn't create pod: %v", err)
			}
			if err := f.WaitForPodRunning(pod.Name); err != nil {
				framework.Failf("Pod did not start running: %v", err)
			}

			By("Running 'kubectl port-forward'")
			cmd := runPortForward(f.Namespace.Name, pod.Name, 80)
			defer cmd.Stop()

			By("Dialing the local port")
			conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.port))
			if err != nil {
				framework.Failf("Couldn't connect to port %d: %v", cmd.port, err)
			}
			defer func() {
				By("Closing the connection to the local port")
				conn.Close()
			}()

			By("Reading data from the local port")
			fromServer, err := ioutil.ReadAll(conn)
			if err != nil {
				framework.Failf("Unexpected error reading data from the server: %v", err)
			}

			if e, a := strings.Repeat("x", 100), string(fromServer); e != a {
				framework.Failf("Expected %q from server, got %q", e, a)
			}

			By("Waiting for the target pod to stop running")
			if err := f.WaitForPodNoLongerRunning(pod.Name); err != nil {
				framework.Failf("Pod did not stop running: %v", err)
			}

			By("Retrieving logs from the target pod")
			logOutput, err := framework.GetPodLogs(f.Client, f.Namespace.Name, pod.Name, "portforwardtester")
			if err != nil {
				framework.Failf("Error retrieving logs: %v", err)
			}

			By("Verifying logs")
			verifyLogMessage(logOutput, "Accepted client connection")
			verifyLogMessage(logOutput, "Done")
		})
	})
})

func verifyLogMessage(log, expected string) {
	re := regexp.MustCompile(expected)
	lines := strings.Split(log, "\n")
	for i := range lines {
		if re.MatchString(lines[i]) {
			return
		}
	}
	framework.Failf("Missing %q from log: %s", expected, log)
}
<commit_msg>portforwarding e2e: print logs when pod doesn't terminate<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package e2e

import (
	"fmt"
	"io/ioutil"
	"net"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/test/e2e/framework"

	. "github.com/onsi/ginkgo"
)

const (
	podName = "pfpod"
)

// TODO support other ports besides 80
var portForwardRegexp = regexp.MustCompile("Forwarding from 127.0.0.1:([0-9]+) -> 80")

func pfPod(expectedClientData, chunks, chunkSize, chunkIntervalMillis string) *api.Pod {
	return &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Name:   podName,
			Labels: map[string]string{"name": podName},
		},
		Spec: api.PodSpec{
			Containers: []api.Container{
				{
					Name:  "portforwardtester",
					Image: "gcr.io/google_containers/portforwardtester:1.0",
					Env: []api.EnvVar{
						{
							Name:  "BIND_PORT",
							Value: "80",
						},
						{
							Name:  "EXPECTED_CLIENT_DATA",
							Value: expectedClientData,
						},
						{
							Name:  "CHUNKS",
							Value: chunks,
						},
						{
							Name:  "CHUNK_SIZE",
							Value: chunkSize,
						},
						{
							Name:  "CHUNK_INTERVAL",
							Value: chunkIntervalMillis,
						},
					},
				},
			},
			RestartPolicy: api.RestartPolicyNever,
		},
	}
}

type portForwardCommand struct {
	cmd  *exec.Cmd
	port int
}

// Stop attempts to gracefully stop `kubectl port-forward`, only killing it if necessary.
// This helps avoid spdy goroutine leaks in the Kubelet.
func (c *portForwardCommand) Stop() {
	// SIGINT signals that kubectl port-forward should gracefully terminate
	if err := c.cmd.Process.Signal(syscall.SIGINT); err != nil {
		framework.Logf("error sending SIGINT to kubectl port-forward: %v", err)
	}

	// try to wait for a clean exit
	done := make(chan error)
	go func() {
		done <- c.cmd.Wait()
	}()

	expired := time.NewTimer(wait.ForeverTestTimeout)
	defer expired.Stop()

	select {
	case err := <-done:
		if err == nil {
			// success
			return
		}
		framework.Logf("error waiting for kubectl port-forward to exit: %v", err)
	case <-expired.C:
		framework.Logf("timed out waiting for kubectl port-forward to exit")
	}

	framework.Logf("trying to forcibly kill kubectl port-forward")
	framework.TryKill(c.cmd)
}

func runPortForward(ns, podName string, port int) *portForwardCommand {
	cmd := framework.KubectlCmd("port-forward", fmt.Sprintf("--namespace=%v", ns), podName, fmt.Sprintf(":%d", port))
	// This is somewhat ugly but is the only way to retrieve the port that was picked
	// by the port-forward command. We don't want to hard code the port as we have no
	// way of guaranteeing we can pick one that isn't in use, particularly on Jenkins.
	framework.Logf("starting port-forward command and streaming output")
	stdout, _, err := framework.StartCmdAndStreamOutput(cmd)
	if err != nil {
		framework.Failf("Failed to start port-forward command: %v", err)
	}

	buf := make([]byte, 128)
	var n int
	framework.Logf("reading from `kubectl port-forward` command's stdout")
	if n, err = stdout.Read(buf); err != nil {
		framework.Failf("Failed to read from kubectl port-forward stdout: %v", err)
	}
	portForwardOutput := string(buf[:n])
	match := portForwardRegexp.FindStringSubmatch(portForwardOutput)
	if len(match) != 2 {
		framework.Failf("Failed to parse kubectl port-forward output: %s", portForwardOutput)
	}

	listenPort, err := strconv.Atoi(match[1])
	if err != nil {
		framework.Failf("Error converting %s to an int: %v", match[1], err)
	}

	return &portForwardCommand{
		cmd:  cmd,
		port: listenPort,
	}
}

var _ = framework.KubeDescribe("Port forwarding", func() {
	f := framework.NewDefaultFramework("port-forwarding")

	framework.KubeDescribe("With a server that expects a client request", func() {
		It("should support a client that connects, sends no data, and disconnects [Conformance]", func() {
			By("creating the target pod")
			pod := pfPod("abc", "1", "1", "1")
			if _, err := f.Client.Pods(f.Namespace.Name).Create(pod); err != nil {
				framework.Failf("Couldn't create pod: %v", err)
			}
			if err := f.WaitForPodRunning(pod.Name); err != nil {
				framework.Failf("Pod did not start running: %v", err)
			}

			By("Running 'kubectl port-forward'")
			cmd := runPortForward(f.Namespace.Name, pod.Name, 80)
			defer cmd.Stop()

			By("Dialing the local port")
			conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.port))
			if err != nil {
				framework.Failf("Couldn't connect to port %d: %v", cmd.port, err)
			}

			By("Closing the connection to the local port")
			conn.Close()

			By("Waiting for the target pod to stop running")
			waitErr := f.WaitForPodNoLongerRunning(pod.Name)

			By("Retrieving logs from the target pod")
			logOutput, err := framework.GetPodLogs(f.Client, f.Namespace.Name, pod.Name, "portforwardtester")
			if err != nil {
				framework.Failf("Error retrieving logs: %v", err)
			}

			if waitErr != nil {
				framework.Logf("Pod log:\n%s", logOutput)
				framework.Failf("Pod did not stop running: %v", waitErr)
			}

			By("Verifying logs")
			verifyLogMessage(logOutput, "Accepted client connection")
			verifyLogMessage(logOutput, "Expected to read 3 bytes from client, but got 0 instead")
		})

		It("should support a client that connects, sends data, and disconnects [Conformance]", func() {
			By("creating the target pod")
			pod := pfPod("abc", "10", "10", "100")
			if _, err := f.Client.Pods(f.Namespace.Name).Create(pod); err != nil {
				framework.Failf("Couldn't create pod: %v", err)
			}
			if err := f.WaitForPodRunning(pod.Name); err != nil {
				framework.Failf("Pod did not start running: %v", err)
			}

			By("Running 'kubectl port-forward'")
			cmd := runPortForward(f.Namespace.Name, pod.Name, 80)
			defer cmd.Stop()

			By("Dialing the local port")
			addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.port))
			if err != nil {
				framework.Failf("Error resolving tcp addr: %v", err)
			}
			conn, err := net.DialTCP("tcp", nil, addr)
			if err != nil {
				framework.Failf("Couldn't connect to port %d: %v", cmd.port, err)
			}
			defer func() {
				By("Closing the connection to the local port")
				conn.Close()
			}()

			By("Sending the expected data to the local port")
			fmt.Fprint(conn, "abc")

			By("Closing the write half of the client's connection")
			conn.CloseWrite()

			By("Reading data from the local port")
			fromServer, err := ioutil.ReadAll(conn)
			if err != nil {
				framework.Failf("Unexpected error reading data from the server: %v", err)
			}

			if e, a := strings.Repeat("x", 100), string(fromServer); e != a {
				framework.Failf("Expected %q from server, got %q", e, a)
			}

			By("Waiting for the target pod to stop running")
			waitErr := f.WaitForPodNoLongerRunning(pod.Name)

			By("Retrieving logs from the target pod")
			logOutput, err := framework.GetPodLogs(f.Client, f.Namespace.Name, pod.Name, "portforwardtester")
			if err != nil {
				framework.Failf("Error retrieving logs: %v", err)
			}

			if waitErr != nil {
				framework.Logf("Pod log:\n%s", logOutput)
				framework.Failf("Pod did not stop running: %v", waitErr)
			}

			By("Verifying logs")
			verifyLogMessage(logOutput, "^Accepted client connection$")
			verifyLogMessage(logOutput, "^Received expected client data$")
			verifyLogMessage(logOutput, "^Done$")
		})
	})
	framework.KubeDescribe("With a server that expects no client request", func() {
		It("should support a client that connects, sends no data, and disconnects [Conformance]", func() {
			By("creating the target pod")
			pod := pfPod("", "10", "10", "100")
			if _, err := f.Client.Pods(f.Namespace.Name).Create(pod); err != nil {
				framework.Failf("Couldn't create pod: %v", err)
			}
			if err := f.WaitForPodRunning(pod.Name); err != nil {
				framework.Failf("Pod did not start running: %v", err)
			}

			By("Running 'kubectl port-forward'")
			cmd := runPortForward(f.Namespace.Name, pod.Name, 80)
			defer cmd.Stop()

			By("Dialing the local port")
			conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", cmd.port))
			if err != nil {
				framework.Failf("Couldn't connect to port %d: %v", cmd.port, err)
			}
			defer func() {
				By("Closing the connection to the local port")
				conn.Close()
			}()

			By("Reading data from the local port")
			fromServer, err := ioutil.ReadAll(conn)
			if err != nil {
				framework.Failf("Unexpected error reading data from the server: %v", err)
			}

			if e, a := strings.Repeat("x", 100), string(fromServer); e != a {
				framework.Failf("Expected %q from server, got %q", e, a)
			}

			By("Waiting for the target pod to stop running")
			waitErr := f.WaitForPodNoLongerRunning(pod.Name)

			By("Retrieving logs from the target pod")
			logOutput, err := framework.GetPodLogs(f.Client, f.Namespace.Name, pod.Name, "portforwardtester")
			if err != nil {
				framework.Failf("Error retrieving logs: %v", err)
			}

			if waitErr != nil {
				framework.Logf("Pod log:\n%s", logOutput)
				framework.Failf("Pod did not stop running: %v", waitErr)
			}

			By("Verifying logs")
			verifyLogMessage(logOutput, "Accepted client connection")
			verifyLogMessage(logOutput, "Done")
		})
	})
})

func verifyLogMessage(log, expected string) {
	re := regexp.MustCompile(expected)
	lines := strings.Split(log, "\n")
	for i := range lines {
		if re.MatchString(lines[i]) {
			return
		}
	}
	framework.Failf("Missing %q from log: %s", expected, log)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"

	"github.com/andrew-d/go-termutil"
	"github.com/mgutz/ansi"
)

type tFuncAppend func([]string, string, string) []string
type tFuncFilter func(string) bool

// Lltsv is a context for processing LTSV.
type Lltsv struct {
	keys         []string
	ignoreKeyMap map[string]struct{}
	noKey        bool
	filters      []string
	exprs        []string
	funcAppend   tFuncAppend
	funcFilters  map[string]tFuncFilter
	exprRunners  map[string]*ExprRunner
}

func newLltsv(keys []string, ignoreKeys []string, noKey bool, filters []string, exprs []string) *Lltsv {
	ignoreKeyMap := make(map[string]struct{})
	for _, key := range ignoreKeys {
		ignoreKeyMap[key] = struct{}{}
	}
	return &Lltsv{
		keys:         keys,
		ignoreKeyMap: ignoreKeyMap,
		noKey:        noKey,
		filters:      filters,
		exprs:        exprs,
		funcAppend:   getFuncAppend(noKey),
		funcFilters:  getFuncFilters(filters),
		exprRunners:  getExprRunners(exprs),
	}
}

func (lltsv *Lltsv) scanAndWrite(file *os.File) error {
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		lvs, keys := lltsv.parseLtsv(line)
		lltsv.expr(lvs)

		if lltsv.filter(lvs) {
			ltsv := lltsv.restructLtsv(lvs, keys)
			os.Stdout.WriteString(ltsv + "\n")
		}

	}
	return scanner.Err()
}

func (lltsv *Lltsv) filter(lvs map[string]string) bool {
	shouldOutput := true

	for key, funcFilter := range lltsv.funcFilters {
		if !funcFilter(lvs[key]) {
			shouldOutput = false
			break
		}
	}

	return shouldOutput
}

func (lltsv *Lltsv) expr(lvs map[string]string) {
	ctx := &ExprContext{
		vars: lvs,
	}
	for key, runner := range lltsv.exprRunners {
		v, err := evalExpr(runner.expr, ctx)
		if err == nil {
			lvs[key] = v.String()
		}
	}
}

// lvs: label and value pairs
func (lltsv *Lltsv) restructLtsv(lvs map[string]string, keys []string) string {
	// specified keys or all keys
	orders := lltsv.keys
	if len(lltsv.keys) == 0 {
		orders = keys
	}
	// make slice with enough capacity so that append does not newly create object
	// cf. http://golang.org/pkg/builtin/#append
	selected := make([]string, 0, len(orders))
	for _, label := range orders {
		if _, ok := lltsv.ignoreKeyMap[label]; ok {
			continue
		}
		value := lvs[label]
		selected = lltsv.funcAppend(selected, label, value)
	}
	return strings.Join(selected, "\t")
}

func (lltsv *Lltsv) parseLtsv(line string) (map[string]string, []string) {
	columns := strings.Split(line, "\t")
	lvs := make(map[string]string)
	keys := make([]string, 0, len(columns))
	for _, column := range columns {
		lv := strings.SplitN(column, ":", 2)
		if len(lv) < 2 {
			continue
		}
		label, value := lv[0], lv[1]
		lvs[label] = value
		keys = append(keys, label)
	}
	return lvs, keys
}

// Return function pointer to avoid `if` evaluation occurs in each iteration
func getFuncAppend(noKey bool) tFuncAppend {
	if noKey {
		return func(selected []string, label string, value string) []string {
			return append(selected, value)
		}
	}

	if termutil.Isatty(os.Stdout.Fd()) {
		return func(selected []string, label string, value string) []string {
			return append(selected, ansi.Color(label, "green")+":"+ansi.Color(value, "magenta"))
		}
	}

	// if pipe or redirect
	return func(selected []string, label string, value string) []string {
		return append(selected, label+":"+value)
	}
}

func getFuncFilters(filters []string) map[string]tFuncFilter {
	funcFilters := map[string]tFuncFilter{}
	for _, f := range filters {
		token := strings.SplitN(f, " ", 3)
		if len(token) < 3 {
			log.Fatalf("filter expression is invalid: %s\n", f)
		}
		key := token[0]
		switch token[1] {
		case ">", ">=", "<=", "<":
			r, err := strconv.ParseFloat(token[2], 64)
			if err != nil {
				log.Fatal(err)
			}

			funcFilters[key] = func(val string) bool {
				num, err := strconv.ParseFloat(val, 64)
				if err != nil {
					log.Println(err)
					return false
				}
				switch token[1] {
				case ">":
					return num > r
				case ">=":
					return num >= r
				case "<=":
					return num <= r
				case "<":
					return num < r
				default:
					return false
				}
			}
		case "==":
			funcFilters[key] = func(val string) bool {
				return val == token[2]
			}
		case "==*":
			funcFilters[key] = func(val string) bool {
				return strings.ToLower(val) == strings.ToLower(token[2])
			}
		case "!=":
			funcFilters[key] = func(val string) bool {
				return val != token[2]
			}
		case "!=*":
			funcFilters[key] = func(val string) bool {
				return strings.ToLower(val) != strings.ToLower(token[2])
			}
		case "=~", "!~", "=~*", "!~*":
			if token[1] == "=~*" || token[1] == "!~*" {
				token[2] = strings.ToLower(token[2])
			}
			re := regexp.MustCompile(token[2])
			funcFilters[key] = func(val string) bool {
				switch token[1] {
				case "=~":
					return re.MatchString(val)
				case "!~":
					return !re.MatchString(val)
				case "=~*":
					return re.MatchString(strings.ToLower(val))
				case "!~*":
					return !re.MatchString(strings.ToLower(val))
				default:
					return false
				}
			}
		}
	}
	return funcFilters
}

func getExprRunners(exprs []string) map[string]*ExprRunner {
	funcExprs := make(map[string]*ExprRunner, len(exprs))
	for _, f := range exprs {
		token := strings.SplitN(f, "=", 2)
		if len(token) != 2 {
			log.Printf("expression is invalid: %s\n", f)
			continue
		}

		expr, err := parseExpr(token[1])
		if err != nil {
			log.Printf("expression is invalid: %s\n", f)
			continue
		}

		key := strings.Trim(token[0], " ")

		funcExprs[key] = &ExprRunner{
			expr: expr,
		}
	}
	return funcExprs
}
<commit_msg>Use https instead of http in comment.<commit_after>package main

import (
	"bufio"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"

	"github.com/andrew-d/go-termutil"
	"github.com/mgutz/ansi"
)

type tFuncAppend func([]string, string, string) []string
type tFuncFilter func(string) bool

// Lltsv is a context for processing LTSV.
type Lltsv struct {
	keys         []string
	ignoreKeyMap map[string]struct{}
	noKey        bool
	filters      []string
	exprs        []string
	funcAppend   tFuncAppend
	funcFilters  map[string]tFuncFilter
	exprRunners  map[string]*ExprRunner
}

func newLltsv(keys []string, ignoreKeys []string, noKey bool, filters []string, exprs []string) *Lltsv {
	ignoreKeyMap := make(map[string]struct{})
	for _, key := range ignoreKeys {
		ignoreKeyMap[key] = struct{}{}
	}
	return &Lltsv{
		keys:         keys,
		ignoreKeyMap: ignoreKeyMap,
		noKey:        noKey,
		filters:      filters,
		exprs:        exprs,
		funcAppend:   getFuncAppend(noKey),
		funcFilters:  getFuncFilters(filters),
		exprRunners:  getExprRunners(exprs),
	}
}

func (lltsv *Lltsv) scanAndWrite(file *os.File) error {
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		lvs, keys := lltsv.parseLtsv(line)
		lltsv.expr(lvs)

		if lltsv.filter(lvs) {
			ltsv := lltsv.restructLtsv(lvs, keys)
			os.Stdout.WriteString(ltsv + "\n")
		}

	}
	return scanner.Err()
}

func (lltsv *Lltsv) filter(lvs map[string]string) bool {
	shouldOutput := true

	for key, funcFilter := range lltsv.funcFilters {
		if !funcFilter(lvs[key]) {
			shouldOutput = false
			break
		}
	}

	return shouldOutput
}

func (lltsv *Lltsv) expr(lvs map[string]string) {
	ctx := &ExprContext{
		vars: lvs,
	}
	for key, runner := range lltsv.exprRunners {
		v, err := evalExpr(runner.expr, ctx)
		if err == nil {
			lvs[key] = v.String()
		}
	}
}

// lvs: label and value pairs
func (lltsv *Lltsv) restructLtsv(lvs map[string]string, keys []string) string {
	// specified keys or all keys
	orders := lltsv.keys
	if len(lltsv.keys) == 0 {
		orders = keys
	}
	// make slice with enough capacity so that append does not newly create object
	// cf. https://golang.org/pkg/builtin/#append
	selected := make([]string, 0, len(orders))
	for _, label := range orders {
		if _, ok := lltsv.ignoreKeyMap[label]; ok {
			continue
		}
		value := lvs[label]
		selected = lltsv.funcAppend(selected, label, value)
	}
	return strings.Join(selected, "\t")
}

func (lltsv *Lltsv) parseLtsv(line string) (map[string]string, []string) {
	columns := strings.Split(line, "\t")
	lvs := make(map[string]string)
	keys := make([]string, 0, len(columns))
	for _, column := range columns {
		lv := strings.SplitN(column, ":", 2)
		if len(lv) < 2 {
			continue
		}
		label, value := lv[0], lv[1]
		lvs[label] = value
		keys = append(keys, label)
	}
	return lvs, keys
}

// Return function pointer to avoid `if` evaluation occurs in each iteration
func getFuncAppend(noKey bool) tFuncAppend {
	if noKey {
		return func(selected []string, label string, value string) []string {
			return append(selected, value)
		}
	}

	if termutil.Isatty(os.Stdout.Fd()) {
		return func(selected []string, label string, value string) []string {
			return append(selected, ansi.Color(label, "green")+":"+ansi.Color(value, "magenta"))
		}
	}

	// if pipe or redirect
	return func(selected []string, label string, value string) []string {
		return append(selected, label+":"+value)
	}
}

func getFuncFilters(filters []string) map[string]tFuncFilter {
	funcFilters := map[string]tFuncFilter{}
	for _, f := range filters {
		token := strings.SplitN(f, " ", 3)
		if len(token) < 3 {
			log.Fatalf("filter expression is invalid: %s\n", f)
		}
		key := token[0]
		switch token[1] {
		case ">", ">=", "<=", "<":
			r, err := strconv.ParseFloat(token[2], 64)
			if err != nil {
				log.Fatal(err)
			}

			funcFilters[key] = func(val string) bool {
				num, err := strconv.ParseFloat(val, 64)
				if err != nil {
					log.Println(err)
					return false
				}
				switch token[1] {
				case ">":
					return num > r
				case ">=":
					return num >= r
				case "<=":
					return num <= r
				case "<":
					return num < r
				default:
					return false
				}
			}
		case "==":
			funcFilters[key] = func(val string) bool {
				return val == token[2]
			}
		case "==*":
			funcFilters[key] = func(val string) bool {
				return strings.ToLower(val) == strings.ToLower(token[2])
			}
		case "!=":
			funcFilters[key] = func(val string) bool {
				return val != token[2]
			}
		case "!=*":
			funcFilters[key] = func(val string) bool {
				return strings.ToLower(val) != strings.ToLower(token[2])
			}
		case "=~", "!~", "=~*", "!~*":
			if token[1] == "=~*" || token[1] == "!~*" {
				token[2] = strings.ToLower(token[2])
			}
			re := regexp.MustCompile(token[2])
			funcFilters[key] = func(val string) bool {
				switch token[1] {
				case "=~":
					return re.MatchString(val)
				case "!~":
					return !re.MatchString(val)
				case "=~*":
					return re.MatchString(strings.ToLower(val))
				case "!~*":
					return !re.MatchString(strings.ToLower(val))
				default:
					return false
				}
			}
		}
	}
	return funcFilters
}

func getExprRunners(exprs []string) map[string]*ExprRunner {
	funcExprs := make(map[string]*ExprRunner, len(exprs))
	for _, f := range exprs {
		token := strings.SplitN(f, "=", 2)
		if len(token) != 2 {
			log.Printf("expression is invalid: %s\n", f)
			continue
		}

		expr, err := parseExpr(token[1])
		if err != nil {
			log.Printf("expression is invalid: %s\n", f)
			continue
		}

		key := strings.Trim(token[0], " ")

		funcExprs[key] = &ExprRunner{
			expr: expr,
		}
	}
	return funcExprs
}
<|endoftext|>


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

package optimize

import (
	"math"
	"time"

	"github.com/gonum/floats"
)

// Local finds a local minimum of a function using a sequential algorithm.
// In order to maximize a function, multiply the output by -1.
//
// The first argument is of Function type representing the function to be minimized.
// Type switching is used to see if the function implements Gradient, FunctionGradient
// and Statuser.
//
// The second argument is the initial location at which to start the minimization.
// The initial location must be supplied, and must have a length equal to the
// problem dimension.
//
// The third argument contains the settings for the minimization. It is here that
// gradient tolerance, etc. are specified. The DefaultSettings() function
// can be called for a Settings struct with the default values initialized.
// If settings == nil, the default settings are used. Please see the documentation
// for the Settings structure for more information. The optimization Method used
// may also contain settings, see documentation for the appropriate optimizer.
//
// The final argument is the optimization method to use. If method == nil, then
// an appropriate default is chosen based on the properties of the other arguments
// (dimension, gradient-free or gradient-based, etc.). The optimization
// methods in this package are designed such that reasonable defaults occur
// if options are not specified explicitly. For example, the code
//  method := &Bfgs{}
// creates a pointer to a new Bfgs struct. When minimize is called, the settings
// in the method will be populated with default values. The methods are also
// designed such that they can be reused in future calls to method.
//
// Local returns a Result struct and any error that occurred. Please see the
// documentation of Result for more information.
//
// Please be aware that the default behavior of Local is to find the minimum.
// For certain functions and optimization methods, this process can take many
// function evaluations. If you would like to put limits on this, for example
// maximum runtime or maximum function evaluations, please modify the Settings
// input struct.
func Local(f Function, initX []float64, settings *Settings, method Method) (*Result, error) {
	if len(initX) == 0 {
		panic("local: initial X has zero length")
	}

	startTime := time.Now()

	funcs, funcInfo := getFunctionInfo(f)

	if method == nil {
		method = getDefaultMethod(funcInfo)
	}

	if settings == nil {
		settings = DefaultSettings()
	}

	stats := &Stats{}
	location, err := getStartingLocation(f, funcs, funcInfo, initX, stats, settings)
	if err != nil {
		return nil, err
	}

	optLoc := &Location{}
	// update stats (grad norm, function value, etc.) so that things are
	// initialized for the first convergence check
	update(location, optLoc, stats, funcInfo, NoEvaluation, NoIteration, startTime)

	if settings.Recorder != nil {
		err = settings.Recorder.Init(funcInfo)
		if err != nil {
			return &Result{Status: Failure}, err
		}
	}

	// actually perform the minimization
	status, err := minimize(settings, location, method, funcInfo, stats, funcs, optLoc, startTime)

	// cleanup at exit
	if settings.Recorder != nil && err == nil {
		err = settings.Recorder.Record(optLoc, NoEvaluation, PostIteration, stats)
	}
	stats.Runtime = time.Since(startTime)
	return &Result{
		Stats:    *stats,
		Location: *optLoc,
		Status:   status,
	}, err
}

func minimize(settings *Settings, location *Location, method Method, funcInfo *FunctionInfo, stats *Stats, funcs functions, optLoc *Location, startTime time.Time) (status Status, err error) {
	methodStatus, methodIsStatuser := method.(Statuser)
	xNext := make([]float64, len(location.X))

	evalType, iterType, err := method.Init(location, funcInfo, xNext)
	if err != nil {
		return Failure, err
	}
	copyLocation(optLoc, location)

	for {
		if settings.Recorder != nil {
			err = settings.Recorder.Record(location, evalType, iterType, stats)
			if err != nil {
				status = Failure
				return
			}
		}

		status = checkConvergence(location, iterType, stats, settings)
		if status != NotTerminated {
			return
		}

		if funcInfo.IsStatuser {
			status, err = funcs.status.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		if methodIsStatuser {
			status, err = methodStatus.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		// Compute the new function and update the statistics
		err = evaluate(funcs, funcInfo, evalType, xNext, location)
		if err != nil {
			status = Failure
			return
		}
		update(location, optLoc, stats, funcInfo, evalType, iterType, startTime)

		// Find the next location
		evalType, iterType, err = method.Iterate(location, xNext)
		if err != nil {
			status = Failure
			return
		}
	}
	panic("unreachable")
}

func copyLocation(dst, src *Location) {
	dst.X = resize(dst.X, len(src.X))
	copy(dst.X, src.X)

	dst.F = src.F

	dst.Gradient = resize(dst.Gradient, len(src.Gradient))
	copy(dst.Gradient, src.Gradient)
}

func getFunctionInfo(f Function) (functions, *FunctionInfo) {
	// Not sure how/if we want to compute timing to be used with functions
	gradient, isGradient := f.(Gradient)
	gradFunc, isFunGrad := f.(FunctionGradient)

	status, isStatuser := f.(Statuser)

	funcInfo := &FunctionInfo{
		IsGradient:         isGradient,
		IsFunctionGradient: isFunGrad,
		IsStatuser:         isStatuser,
	}
	funcs := functions{
		function: f,
		gradient: gradient,
		gradFunc: gradFunc,
		status:   status,
	}

	return funcs, funcInfo
}

func getDefaultMethod(funcInfo *FunctionInfo) Method {
	if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
		return &BFGS{}
	}
	// TODO: Implement a gradient-free method
	panic("optimize: gradient-free methods not yet coded")
}

// Combine location and stats because maybe in the future we'll add evaluation times
// to functionStats?
func getStartingLocation(f Function, funcs functions, funcInfo *FunctionInfo, initX []float64, stats *Stats, settings *Settings) (*Location, error) {
	l := &Location{
		X: make([]float64, len(initX)),
	}
	copy(l.X, initX)
	if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
		l.Gradient = make([]float64, len(l.X))
	}

	if settings.UseInitialData {
		l.F = settings.InitialFunctionValue
		initG := settings.InitialGradient
		if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
			if len(l.Gradient) != len(initG) {
				panic("local: initial location size mismatch")
			}
			copy(l.Gradient, initG)
		}
	} else {
		// Compute missing information in the initial state.
		if funcInfo.IsFunctionGradient {
			l.F = funcs.gradFunc.FDf(l.X, l.Gradient)
			stats.FunctionGradientEvals++
		} else {
			l.F = funcs.function.F(l.X)
			stats.FunctionEvals++
			if funcInfo.IsGradient {
				funcs.gradient.Df(l.X, l.Gradient)
				stats.GradientEvals++
			}
		}
	}

	if math.IsNaN(l.F) {
		return l, ErrNaN
	}
	// Do we allow Inf initial function value?
	if math.IsInf(l.F, 1) {
		return l, ErrInf
	}
	return l, nil
}

func checkConvergence(loc *Location, itertype IterationType, stats *Stats, settings *Settings) Status {
	if itertype == MajorIteration && loc.Gradient != nil {
		if stats.GradientNorm <= settings.GradientAbsTol {
			return GradientAbsoluteConvergence
		}
	}

	if itertype == MajorIteration && loc.F < settings.FunctionAbsTol {
		return FunctionAbsoluteConvergence
	}

	// Check every step for negative infinity because it could break the
	// linesearches and -inf is the best you can do anyway.
	if math.IsInf(loc.F, -1) {
		return FunctionNegativeInfinity
	}

	if settings.FunctionEvals > 0 {
		totalFun := stats.FunctionEvals + stats.FunctionGradientEvals
		if totalFun >= settings.FunctionEvals {
			return FunctionEvaluationLimit
		}
	}

	if settings.GradientEvals > 0 {
		totalGrad := stats.GradientEvals + stats.FunctionGradientEvals
		if totalGrad >= settings.GradientEvals {
			return GradientEvaluationLimit
		}
	}

	if settings.Runtime > 0 {
		if stats.Runtime >= settings.Runtime {
			return RuntimeLimit
		}
	}

	if itertype == MajorIteration && settings.MajorIterations > 0 {
		if stats.MajorIterations >= settings.MajorIterations {
			return IterationLimit
		}
	}
	return NotTerminated
}

// evaluate evaluates the function and stores the answer in place
func evaluate(funcs functions, funcInfo *FunctionInfo, evalType EvaluationType, xNext []float64, location *Location) error {
	sameX := floats.Equal(location.X, xNext)
	if !sameX {
		copy(location.X, xNext)
	}
	switch evalType {
	case FunctionEval:
		location.F = funcs.function.F(location.X)
		if !sameX {
			for i := range location.Gradient {
				location.Gradient[i] = math.NaN()
			}
		}
		return nil
	case GradientEval:
		if funcInfo.IsGradient {
			if !sameX {
				location.F = math.NaN()
			}
			funcs.gradient.Df(location.X, location.Gradient)
			return nil
		}
		if funcInfo.IsFunctionGradient {
			location.F = funcs.gradFunc.FDf(location.X, location.Gradient)
			return nil
		}
		return ErrMismatch{Type: evalType}
	case FunctionAndGradientEval:
		if funcInfo.IsFunctionGradient {
			location.F = funcs.gradFunc.FDf(location.X, location.Gradient)
			return nil
		}
		if funcInfo.IsGradient {
			location.F = funcs.function.F(location.X)
			funcs.gradient.Df(location.X, location.Gradient)
			return nil
		}
		return ErrMismatch{Type: evalType}
	default:
		panic("unreachable")
	}
}

// update updates the stats given the new evaluation
func update(location *Location, optLoc *Location, stats *Stats, funcInfo *FunctionInfo, evalType EvaluationType, iterType IterationType, startTime time.Time) {
	switch evalType {
	case FunctionEval:
		stats.FunctionEvals++
	case GradientEval:
		if funcInfo.IsGradient {
			stats.GradientEvals++
		} else if funcInfo.IsFunctionGradient {
			stats.FunctionGradientEvals++
		}
	case FunctionAndGradientEval:
		if funcInfo.IsFunctionGradient {
			stats.FunctionGradientEvals++
		} else if funcInfo.IsGradient {
			stats.FunctionEvals++
			stats.GradientEvals++
		}
	}
	if iterType == MajorIteration {
		stats.MajorIterations++
	}
	if location.F < optLoc.F {
		copyLocation(optLoc, location)
	}
	stats.Runtime = time.Since(startTime)
	if location.Gradient != nil {
		stats.GradientNorm = floats.Norm(location.Gradient, math.Inf(1))
	}
}
<commit_msg>Move updating of f/g call stats to evaluate()<commit_after>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package optimize

import (
	"math"
	"time"

	"github.com/gonum/floats"
)

// Local finds a local minimum of a function using a sequential algorithm.
// In order to maximize a function, multiply the output by -1.
//
// The first argument is of Function type representing the function to be minimized.
// Type switching is used to see if the function implements Gradient, FunctionGradient
// and Statuser.
//
// The second argument is the initial location at which to start the minimization.
// The initial location must be supplied, and must have a length equal to the
// problem dimension.
//
// The third argument contains the settings for the minimization. It is here that
// gradient tolerance, etc. are specified. The DefaultSettings() function
// can be called for a Settings struct with the default values initialized.
// If settings == nil, the default settings are used. Please see the documentation
// for the Settings structure for more information. The optimization Method used
// may also contain settings, see documentation for the appropriate optimizer.
//
// The final argument is the optimization method to use. If method == nil, then
// an appropriate default is chosen based on the properties of the other arguments
// (dimension, gradient-free or gradient-based, etc.). The optimization
// methods in this package are designed such that reasonable defaults occur
// if options are not specified explicitly. For example, the code
//  method := &Bfgs{}
// creates a pointer to a new Bfgs struct. When minimize is called, the settings
// in the method will be populated with default values. The methods are also
// designed such that they can be reused in future calls to method.
//
// Local returns a Result struct and any error that occurred. Please see the
// documentation of Result for more information.
//
// Please be aware that the default behavior of Local is to find the minimum.
// For certain functions and optimization methods, this process can take many
// function evaluations. If you would like to put limits on this, for example
// maximum runtime or maximum function evaluations, please modify the Settings
// input struct.
func Local(f Function, initX []float64, settings *Settings, method Method) (*Result, error) {
	if len(initX) == 0 {
		panic("local: initial X has zero length")
	}

	startTime := time.Now()

	funcs, funcInfo := getFunctionInfo(f)

	if method == nil {
		method = getDefaultMethod(funcInfo)
	}

	if settings == nil {
		settings = DefaultSettings()
	}

	stats := &Stats{}
	location, err := getStartingLocation(f, funcs, funcInfo, initX, stats, settings)
	if err != nil {
		return nil, err
	}

	optLoc := &Location{}
	// update stats (grad norm, function value, etc.) so that things are
	// initialized for the first convergence check
	update(location, optLoc, stats, NoIteration, startTime)

	if settings.Recorder != nil {
		err = settings.Recorder.Init(funcInfo)
		if err != nil {
			return &Result{Status: Failure}, err
		}
	}

	// actually perform the minimization
	status, err := minimize(settings, location, method, funcInfo, stats, funcs, optLoc, startTime)

	// cleanup at exit
	if settings.Recorder != nil && err == nil {
		err = settings.Recorder.Record(optLoc, NoEvaluation, PostIteration, stats)
	}
	stats.Runtime = time.Since(startTime)
	return &Result{
		Stats:    *stats,
		Location: *optLoc,
		Status:   status,
	}, err
}

func minimize(settings *Settings, location *Location, method Method, funcInfo *FunctionInfo, stats *Stats, funcs functions, optLoc *Location, startTime time.Time) (status Status, err error) {
	methodStatus, methodIsStatuser := method.(Statuser)
	xNext := make([]float64, len(location.X))

	evalType, iterType, err := method.Init(location, funcInfo, xNext)
	if err != nil {
		return Failure, err
	}
	copyLocation(optLoc, location)

	for {
		if settings.Recorder != nil {
			err = settings.Recorder.Record(location, evalType, iterType, stats)
			if err != nil {
				status = Failure
				return
			}
		}

		status = checkConvergence(location, iterType, stats, settings)
		if status != NotTerminated {
			return
		}

		if funcInfo.IsStatuser {
			status, err = funcs.status.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		if methodIsStatuser {
			status, err = methodStatus.Status()
			if err != nil || status != NotTerminated {
				return
			}
		}

		// Compute the new function and update the statistics
		err = evaluate(funcs, funcInfo, evalType, xNext, location, stats)
		if err != nil {
			status = Failure
			return
		}
		update(location, optLoc, stats, iterType, startTime)

		// Find the next location
		evalType, iterType, err = method.Iterate(location, xNext)
		if err != nil {
			status = Failure
			return
		}
	}
	panic("unreachable")
}

func copyLocation(dst, src *Location) {
	dst.X = resize(dst.X, len(src.X))
	copy(dst.X, src.X)

	dst.F = src.F

	dst.Gradient = resize(dst.Gradient, len(src.Gradient))
	copy(dst.Gradient, src.Gradient)
}

func getFunctionInfo(f Function) (functions, *FunctionInfo) {
	// Not sure how/if we want to compute timing to be used with functions
	gradient, isGradient := f.(Gradient)
	gradFunc, isFunGrad := f.(FunctionGradient)

	status, isStatuser := f.(Statuser)

	funcInfo := &FunctionInfo{
		IsGradient:         isGradient,
		IsFunctionGradient: isFunGrad,
		IsStatuser:         isStatuser,
	}
	funcs := functions{
		function: f,
		gradient: gradient,
		gradFunc: gradFunc,
		status:   status,
	}

	return funcs, funcInfo
}

func getDefaultMethod(funcInfo *FunctionInfo) Method {
	if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
		return &BFGS{}
	}
	// TODO: Implement a gradient-free method
	panic("optimize: gradient-free methods not yet coded")
}

// Combine location and stats because maybe in the future we'll add evaluation times
// to functionStats?
func getStartingLocation(f Function, funcs functions, funcInfo *FunctionInfo, initX []float64, stats *Stats, settings *Settings) (*Location, error) {
	l := &Location{
		X: make([]float64, len(initX)),
	}
	copy(l.X, initX)
	if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
		l.Gradient = make([]float64, len(l.X))
	}

	if settings.UseInitialData {
		l.F = settings.InitialFunctionValue
		initG := settings.InitialGradient
		if funcInfo.IsGradient || funcInfo.IsFunctionGradient {
			if len(l.Gradient) != len(initG) {
				panic("local: initial location size mismatch")
			}
			copy(l.Gradient, initG)
		}
	} else {
		// Compute missing information in the initial state.
		if funcInfo.IsFunctionGradient {
			l.F = funcs.gradFunc.FDf(l.X, l.Gradient)
			stats.FunctionGradientEvals++
		} else {
			l.F = funcs.function.F(l.X)
			stats.FunctionEvals++
			if funcInfo.IsGradient {
				funcs.gradient.Df(l.X, l.Gradient)
				stats.GradientEvals++
			}
		}
	}

	if math.IsNaN(l.F) {
		return l, ErrNaN
	}
	// Do we allow Inf initial function value?
	if math.IsInf(l.F, 1) {
		return l, ErrInf
	}
	return l, nil
}

func checkConvergence(loc *Location, itertype IterationType, stats *Stats, settings *Settings) Status {
	if itertype == MajorIteration && loc.Gradient != nil {
		if stats.GradientNorm <= settings.GradientAbsTol {
			return GradientAbsoluteConvergence
		}
	}

	if itertype == MajorIteration && loc.F < settings.FunctionAbsTol {
		return FunctionAbsoluteConvergence
	}

	// Check every step for negative infinity because it could break the
	// linesearches and -inf is the best you can do anyway.
	if math.IsInf(loc.F, -1) {
		return FunctionNegativeInfinity
	}

	if settings.FunctionEvals > 0 {
		totalFun := stats.FunctionEvals + stats.FunctionGradientEvals
		if totalFun >= settings.FunctionEvals {
			return FunctionEvaluationLimit
		}
	}

	if settings.GradientEvals > 0 {
		totalGrad := stats.GradientEvals + stats.FunctionGradientEvals
		if totalGrad >= settings.GradientEvals {
			return GradientEvaluationLimit
		}
	}

	if settings.Runtime > 0 {
		if stats.Runtime >= settings.Runtime {
			return RuntimeLimit
		}
	}

	if itertype == MajorIteration && settings.MajorIterations > 0 {
		if stats.MajorIterations >= settings.MajorIterations {
			return IterationLimit
		}
	}
	return NotTerminated
}

// evaluate evaluates the function and stores the answer in place
func evaluate(funcs functions, funcInfo *FunctionInfo, evalType EvaluationType, xNext []float64, location *Location, stats *Stats) error {
	sameX := floats.Equal(location.X, xNext)
	if !sameX {
		copy(location.X, xNext)
	}
	switch evalType {
	case FunctionEval:
		location.F = funcs.function.F(location.X)
		stats.FunctionEvals++
		if !sameX {
			for i := range location.Gradient {
				location.Gradient[i] = math.NaN()
			}
		}
		return nil
	case GradientEval:
		if funcInfo.IsGradient {
			funcs.gradient.Df(location.X, location.Gradient)
			stats.GradientEvals++
			if !sameX {
				location.F = math.NaN()
			}
			return nil
		}
		if funcInfo.IsFunctionGradient {
			location.F = funcs.gradFunc.FDf(location.X, location.Gradient)
			stats.FunctionGradientEvals++
			return nil
		}
		return ErrMismatch{Type: evalType}
	case FunctionAndGradientEval:
		if funcInfo.IsFunctionGradient {
			location.F = funcs.gradFunc.FDf(location.X, location.Gradient)
			stats.FunctionGradientEvals++
			return nil
		}
		if funcInfo.IsGradient {
			location.F = funcs.function.F(location.X)
			stats.FunctionEvals++
			funcs.gradient.Df(location.X, location.Gradient)
			stats.GradientEvals++
			return nil
		}
		return ErrMismatch{Type: evalType}
	default:
		panic("unreachable")
	}
}

// update updates the stats given the new evaluation
func update(location *Location, optLoc *Location, stats *Stats, iterType IterationType, startTime time.Time) {
	if iterType == MajorIteration {
		stats.MajorIterations++
	}
	if location.F < optLoc.F {
		copyLocation(optLoc, location)
	}
	stats.Runtime = time.Since(startTime)
	if location.Gradient != nil {
		stats.GradientNorm = floats.Norm(location.Gradient, math.Inf(1))
	}
}
<|endoftext|>


<commit_before>package main

import "log"

import (
	gl "github.com/go-gl/gl"
	glfw "github.com/go-gl/glfw3"
	glh "github.com/go-gl/glh"
)

func keyCallback(w *glfw.Window, key glfw.Key, scancode int, action glfw.Action, mods glfw.ModifierKey) {
	if key == glfw.KeyEscape && action == glfw.Press {
		w.SetShouldClose(true)
	}
}

func main() {
	if !glfw.Init() {
		log.Fatal("Failed to init glfw")
	}

	defer glfw.Terminate()

	window, err := glfw.CreateWindow(400, 400, "loop1", nil, nil)
	if err != nil {
		log.Fatal("Failed to create window")
	}

	window.SetKeyCallback(keyCallback)
	window.MakeContextCurrent()

	glfw.SwapInterval(1)

	if gl.Init() != 0 {
		log.Fatal("failed to init gl")
	}

	for !window.ShouldClose() {
		window.SwapBuffers()
		glfw.PollEvents()
	}

	glh.OpenGLSentinel()
}
<commit_msg>more boilerplate<commit_after>package main

import "log"

import (
	gl "github.com/go-gl/gl"
	glfw "github.com/go-gl/glfw3"
	glh "github.com/go-gl/glh"
)

var (
	width  = 400
	height = 400
)

func errorCallback(err glfw.ErrorCode, desc string) {
	log.Print("Error: ", err, " ", desc)
}

func keyCallback(w *glfw.Window, key glfw.Key, scancode int, action glfw.Action, mods glfw.ModifierKey) {
	if key == glfw.KeyEscape && action == glfw.Press {
		w.SetShouldClose(true)
	}
}

func main() {
	glfw.SetErrorCallback(errorCallback)

	if !glfw.Init() {
		log.Fatal("Failed to init glfw")
	}

	defer glfw.Terminate()

	window, err := glfw.CreateWindow(width, height, "loop1", nil, nil)
	if err != nil {
		log.Fatal("Failed to create window")
	}

	window.SetKeyCallback(keyCallback)
	window.MakeContextCurrent()

	glfw.SwapInterval(1)

	if gl.Init() != 0 {
		log.Fatal("failed to init gl")
	}

	gl.ClearColor(0.2, 0.2, 0.2, 0)

	for !window.ShouldClose() {
		gl.Viewport(0, 0, width, height)
		gl.Clear(gl.COLOR_BUFFER_BIT)

		window.SwapBuffers()
		glfw.PollEvents()
	}

	glh.OpenGLSentinel()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/flynn/flynn-host/client"
	"github.com/flynn/flynn-host/sampi"
	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-dockerclient"
	"github.com/flynn/go-flynn/attempt"
	rpc "github.com/flynn/rpcplus/comborpc"
	"github.com/technoweenie/grohl"
)

// Attempts is the attempt strategy that is used to connect to discoverd.
var Attempts = attempt.Strategy{
	Min:   5,
	Total: 5 * time.Second,
	Delay: 200 * time.Millisecond,
}

// A command line flag to accumulate multiple key-value pairs into Attributes,
// e.g. flynn-host -attribute foo=bar -attribute bar=foo
type AttributeFlag map[string]string

func (a AttributeFlag) Set(val string) error {
	kv := strings.SplitN(val, "=", 2)
	a[kv[0]] = kv[1]
	return nil
}

func (a AttributeFlag) String() string {
	res := make([]string, 0, len(a))
	for k, v := range a {
		res = append(res, k+"="+v)
	}
	return strings.Join(res, ", ")
}

func main() {
	hostname, _ := os.Hostname()
	externalAddr := flag.String("external", "", "external IP of host")
	configFile := flag.String("config", "", "configuration file")
	manifestFile := flag.String("manifest", "", "manifest file")
	hostID := flag.String("id", hostname, "host id")
	force := flag.Bool("force", false, "kill all containers booted by flynn-host before starting")
	attributes := make(AttributeFlag)
	flag.Var(&attributes, "attribute", "key=value pair to add as an attribute")
	flag.Parse()
	grohl.AddContext("app", "lorne")
	grohl.Log(grohl.Data{"at": "start"})
	g := grohl.NewContext(grohl.Data{"fn": "main"})

	dockerc, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		log.Fatal(err)
	}

	if *force {
		if err := killExistingContainers(dockerc); err != nil {
			os.Exit(1)
		}
	}

	state := NewState()
	ports := make(chan int)

	go allocatePorts(ports, 55000, 65535)
	go serveHTTP(&Host{state: state, docker: dockerc}, &attachHandler{state: state, docker: dockerc})
	go streamEvents(dockerc, state)

	processor := &jobProcessor{
		externalAddr: *externalAddr,
		docker:       dockerc,
		state:        state,
		discoverd:    os.Getenv("DISCOVERD"),
	}

	runner := &manifestRunner{
		env:        parseEnviron(),
		externalIP: *externalAddr,
		ports:      ports,
		processor:  processor,
		docker:     dockerc,
	}

	var disc *discoverd.Client
	if *manifestFile != "" {
		f, err := os.Open(*manifestFile)
		if err != nil {
			log.Fatal(err)
		}
		services, err := runner.runManifest(f)
		if err != nil {
			log.Fatal(err)
		}
		f.Close()

		if d, ok := services["discoverd"]; ok {
			processor.discoverd = fmt.Sprintf("%s:%d", d.InternalIP, d.TCPPorts[0])
			var disc *discoverd.Client
			err = Attempts.Run(func() (err error) {
				disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
				return
			})
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	if processor.discoverd == "" && *externalAddr != "" {
		processor.discoverd = *externalAddr + ":1111"
	}
	// HACK: use env as global for discoverd connection in sampic
	os.Setenv("DISCOVERD", processor.discoverd)
	if disc == nil {
		disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
		if err != nil {
			log.Fatal(err)
		}
	}
	sampiStandby, err := disc.RegisterAndStandby("flynn-host", *externalAddr+":1113", map[string]string{"id": *hostID})
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		<-sampiStandby
		rpc.Register(sampi.NewCluster(sampi.NewState()))
	}()

	cluster, err := client.New()
	if err != nil {
		log.Fatal(err)
	}
	g.Log(grohl.Data{"at": "sampi_connected"})

	events := make(chan host.Event)
	state.AddListener("all", events)
	go syncScheduler(cluster, events)

	h := &host.Host{}
	if *configFile != "" {
		h, err = openConfig(*configFile)
		if err != nil {
			log.Fatal(err)
		}
	}
	h.ID = *hostID
	h.Jobs = state.ClusterJobs()

	if h.Attributes == nil {
		h.Attributes = make(map[string]string)
	}

	for k, v := range attributes {
		h.Attributes[k] = v
	}

	jobs := make(chan *host.Job)
	hostErr := cluster.ConnectHost(h, jobs)
	g.Log(grohl.Data{"at": "host_registered"})
	processor.Process(ports, jobs)
	log.Fatal(*hostErr)
}

type jobProcessor struct {
	externalAddr string
	discoverd    string
	docker       interface {
		CreateContainer(*docker.Config) (*docker.Container, error)
		PullImage(docker.PullImageOptions, io.Writer) error
		StartContainer(string, *docker.HostConfig) error
		InspectContainer(string) (*docker.Container, error)
	}
	state *State
}

func killExistingContainers(dc *docker.Client) error {
	g := grohl.NewContext(grohl.Data{"fn": "kill_existing"})
	g.Log(grohl.Data{"at": "start"})
	containers, err := dc.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		g.Log(grohl.Data{"at": "list", "status": "error", "err": err})
		return err
	}
outer:
	for _, c := range containers {
		for _, name := range c.Names {
			if strings.HasPrefix(name, "/flynn-") {
				g.Log(grohl.Data{"at": "kill", "container.id": c.ID, "container.name": name})
				if err := dc.KillContainer(c.ID); err != nil {
					g.Log(grohl.Data{"at": "kill", "container.id": c.ID, "container.name": name, "status": "error", "err": err})
				}
				continue outer
			}
		}
	}
	g.Log(grohl.Data{"at": "finish"})
	return nil
}

func (p *jobProcessor) Process(ports <-chan int, jobs chan *host.Job) {
	for job := range jobs {
		p.processJob(ports, job)
	}
}

func (p *jobProcessor) processJob(ports <-chan int, job *host.Job) (*docker.Container, error) {
	g := grohl.NewContext(grohl.Data{"fn": "process_job", "job.id": job.ID})
	g.Log(grohl.Data{"at": "start", "job.image": job.Config.Image, "job.cmd": job.Config.Cmd, "job.entrypoint": job.Config.Entrypoint})

	for i := 0; i < job.TCPPorts; i++ {
		port := strconv.Itoa(<-ports)
		if i == 0 {
			job.Config.Env = append(job.Config.Env, "PORT="+port)
		}
		job.Config.Env = append(job.Config.Env, fmt.Sprintf("PORT_%d=%s", i, port))
		if job.Config.ExposedPorts == nil {
			job.Config.ExposedPorts = make(map[string]struct{}, job.TCPPorts)
		}
		job.Config.AttachStdout = true
		job.Config.AttachStderr = true
		job.Config.ExposedPorts[port+"/tcp"] = struct{}{}
		if job.HostConfig == nil {
			job.HostConfig = &docker.HostConfig{
				PortBindings:    make(map[string][]docker.PortBinding, job.TCPPorts),
				PublishAllPorts: true,
			}
		}
		job.HostConfig.PortBindings[port+"/tcp"] = []docker.PortBinding{{HostPort: port}}
		if strings.HasPrefix(job.ID, "flynn-") {
			job.Config.Name = job.ID
		} else {
			job.Config.Name = "flynn-" + job.ID
		}
	}
	if p.externalAddr != "" {
		job.Config.Env = appendUnique(job.Config.Env, "EXTERNAL_IP="+p.externalAddr, "SD_HOST="+p.externalAddr, "DISCOVERD="+p.discoverd)
	}
	p.state.AddJob(job)
	g.Log(grohl.Data{"at": "create_container"})
	container, err := p.docker.CreateContainer(job.Config)
	if err == docker.ErrNoSuchImage {
		g.Log(grohl.Data{"at": "pull_image"})
		err = p.docker.PullImage(docker.PullImageOptions{Repository: job.Config.Image}, os.Stdout)
		if err != nil {
			g.Log(grohl.Data{"at": "pull_image", "status": "error", "err": err})
			p.state.SetStatusFailed(job.ID, err)
			return nil, err
		}
		container, err = p.docker.CreateContainer(job.Config)
	}
	if err != nil {
		g.Log(grohl.Data{"at": "create_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetContainerID(job.ID, container.ID)
	p.state.WaitAttach(job.ID)
	g.Log(grohl.Data{"at": "start_container"})
	if err := p.docker.StartContainer(container.ID, job.HostConfig); err != nil {
		g.Log(grohl.Data{"at": "start_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	container, err = p.docker.InspectContainer(container.ID)
	if err != nil {
		g.Log(grohl.Data{"at": "inspect_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetStatusRunning(job.ID, container.Volumes)
	g.Log(grohl.Data{"at": "finish"})
	return container, nil
}

func appendUnique(s []string, vars ...string) []string {
outer:
	for _, v := range vars {
		for _, existing := range s {
			if strings.HasPrefix(existing, strings.SplitN(v, "=", 2)[0]+"=") {
				continue outer
			}
		}
		s = append(s, v)
	}
	return s
}

type sampiSyncClient interface {
	RemoveJobs([]string) error
}

func syncScheduler(scheduler sampiSyncClient, events <-chan host.Event) {
	for event := range events {
		if event.Event != "stop" {
			continue
		}
		grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "job.id": event.JobID})
		if err := scheduler.RemoveJobs([]string{event.JobID}); err != nil {
			grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "status": "error", "err": err, "job.id": event.JobID})
			// TODO: try to reconnect?
		}
	}
}

type dockerStreamClient interface {
	Events() (*docker.EventStream, error)
	InspectContainer(string) (*docker.Container, error)
}

func streamEvents(client dockerStreamClient, state *State) {
	stream, err := client.Events()
	if err != nil {
		log.Fatal(err)
	}
	for event := range stream.Events {
		if event.Status != "die" {
			continue
		}
		container, err := client.InspectContainer(event.ID)
		if err != nil {
			log.Println("inspect container", event.ID, "error:", err)
			// TODO: set job status anyway?
			continue
		}
		state.SetStatusDone(event.ID, container.State.ExitCode)
	}
}

// TODO: fix this, horribly broken

func allocatePorts(ports chan<- int, startPort, endPort int) {
	for i := startPort; i < endPort; i++ {
		ports <- i
	}
	// TODO: handle wrap-around
}
<commit_msg>Fix job attribute assignment<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/flynn/flynn-host/client"
	"github.com/flynn/flynn-host/sampi"
	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-dockerclient"
	"github.com/flynn/go-flynn/attempt"
	rpc "github.com/flynn/rpcplus/comborpc"
	"github.com/technoweenie/grohl"
)

// Attempts is the attempt strategy that is used to connect to discoverd.
var Attempts = attempt.Strategy{
	Min:   5,
	Total: 5 * time.Second,
	Delay: 200 * time.Millisecond,
}

// A command line flag to accumulate multiple key-value pairs into Attributes,
// e.g. flynn-host -attribute foo=bar -attribute bar=foo
type AttributeFlag map[string]string

func (a AttributeFlag) Set(val string) error {
	kv := strings.SplitN(val, "=", 2)
	a[kv[0]] = kv[1]
	return nil
}

func (a AttributeFlag) String() string {
	res := make([]string, 0, len(a))
	for k, v := range a {
		res = append(res, k+"="+v)
	}
	return strings.Join(res, ", ")
}

func main() {
	hostname, _ := os.Hostname()
	externalAddr := flag.String("external", "", "external IP of host")
	configFile := flag.String("config", "", "configuration file")
	manifestFile := flag.String("manifest", "", "manifest file")
	hostID := flag.String("id", hostname, "host id")
	force := flag.Bool("force", false, "kill all containers booted by flynn-host before starting")
	attributes := make(AttributeFlag)
	flag.Var(&attributes, "attribute", "key=value pair to add as an attribute")
	flag.Parse()
	grohl.AddContext("app", "lorne")
	grohl.Log(grohl.Data{"at": "start"})
	g := grohl.NewContext(grohl.Data{"fn": "main"})

	dockerc, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		log.Fatal(err)
	}

	if *force {
		if err := killExistingContainers(dockerc); err != nil {
			os.Exit(1)
		}
	}

	state := NewState()
	ports := make(chan int)

	go allocatePorts(ports, 55000, 65535)
	go serveHTTP(&Host{state: state, docker: dockerc}, &attachHandler{state: state, docker: dockerc})
	go streamEvents(dockerc, state)

	processor := &jobProcessor{
		externalAddr: *externalAddr,
		docker:       dockerc,
		state:        state,
		discoverd:    os.Getenv("DISCOVERD"),
	}

	runner := &manifestRunner{
		env:        parseEnviron(),
		externalIP: *externalAddr,
		ports:      ports,
		processor:  processor,
		docker:     dockerc,
	}

	var disc *discoverd.Client
	if *manifestFile != "" {
		f, err := os.Open(*manifestFile)
		if err != nil {
			log.Fatal(err)
		}
		services, err := runner.runManifest(f)
		if err != nil {
			log.Fatal(err)
		}
		f.Close()

		if d, ok := services["discoverd"]; ok {
			processor.discoverd = fmt.Sprintf("%s:%d", d.InternalIP, d.TCPPorts[0])
			var disc *discoverd.Client
			err = Attempts.Run(func() (err error) {
				disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
				return
			})
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	if processor.discoverd == "" && *externalAddr != "" {
		processor.discoverd = *externalAddr + ":1111"
	}
	// HACK: use env as global for discoverd connection in sampic
	os.Setenv("DISCOVERD", processor.discoverd)
	if disc == nil {
		disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
		if err != nil {
			log.Fatal(err)
		}
	}
	sampiStandby, err := disc.RegisterAndStandby("flynn-host", *externalAddr+":1113", map[string]string{"id": *hostID})
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		<-sampiStandby
		rpc.Register(sampi.NewCluster(sampi.NewState()))
	}()

	cluster, err := client.New()
	if err != nil {
		log.Fatal(err)
	}
	g.Log(grohl.Data{"at": "sampi_connected"})

	events := make(chan host.Event)
	state.AddListener("all", events)
	go syncScheduler(cluster, events)

	h := &host.Host{}
	if *configFile != "" {
		h, err = openConfig(*configFile)
		if err != nil {
			log.Fatal(err)
		}
	}
	h.ID = *hostID
	h.Jobs = state.ClusterJobs()

	if h.Attributes == nil {
		h.Attributes = make(map[string]string)
	}

	for k, v := range attributes {
		h.Attributes[k] = v
	}

	jobs := make(chan *host.Job)
	hostErr := cluster.ConnectHost(h, jobs)
	g.Log(grohl.Data{"at": "host_registered"})
	processor.Process(ports, jobs)
	log.Fatal(*hostErr)
}

type jobProcessor struct {
	externalAddr string
	discoverd    string
	docker       interface {
		CreateContainer(*docker.Config) (*docker.Container, error)
		PullImage(docker.PullImageOptions, io.Writer) error
		StartContainer(string, *docker.HostConfig) error
		InspectContainer(string) (*docker.Container, error)
	}
	state *State
}

func killExistingContainers(dc *docker.Client) error {
	g := grohl.NewContext(grohl.Data{"fn": "kill_existing"})
	g.Log(grohl.Data{"at": "start"})
	containers, err := dc.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		g.Log(grohl.Data{"at": "list", "status": "error", "err": err})
		return err
	}
outer:
	for _, c := range containers {
		for _, name := range c.Names {
			if strings.HasPrefix(name, "/flynn-") {
				g.Log(grohl.Data{"at": "kill", "container.id": c.ID, "container.name": name})
				if err := dc.KillContainer(c.ID); err != nil {
					g.Log(grohl.Data{"at": "kill", "container.id": c.ID, "container.name": name, "status": "error", "err": err})
				}
				continue outer
			}
		}
	}
	g.Log(grohl.Data{"at": "finish"})
	return nil
}

func (p *jobProcessor) Process(ports <-chan int, jobs chan *host.Job) {
	for job := range jobs {
		p.processJob(ports, job)
	}
}

func (p *jobProcessor) processJob(ports <-chan int, job *host.Job) (*docker.Container, error) {
	g := grohl.NewContext(grohl.Data{"fn": "process_job", "job.id": job.ID})
	g.Log(grohl.Data{"at": "start", "job.image": job.Config.Image, "job.cmd": job.Config.Cmd, "job.entrypoint": job.Config.Entrypoint})

	if job.HostConfig == nil {
		job.HostConfig = &docker.HostConfig{
			PortBindings:    make(map[string][]docker.PortBinding, job.TCPPorts),
			PublishAllPorts: true,
		}
	}
	if job.Config.ExposedPorts == nil {
		job.Config.ExposedPorts = make(map[string]struct{}, job.TCPPorts)
	}
	for i := 0; i < job.TCPPorts; i++ {
		port := strconv.Itoa(<-ports)
		if i == 0 {
			job.Config.Env = append(job.Config.Env, "PORT="+port)
		}
		job.Config.Env = append(job.Config.Env, fmt.Sprintf("PORT_%d=%s", i, port))
		job.Config.ExposedPorts[port+"/tcp"] = struct{}{}
		job.HostConfig.PortBindings[port+"/tcp"] = []docker.PortBinding{{HostPort: port}}
	}

	job.Config.AttachStdout = true
	job.Config.AttachStderr = true
	if strings.HasPrefix(job.ID, "flynn-") {
		job.Config.Name = job.ID
	} else {
		job.Config.Name = "flynn-" + job.ID
	}
	if p.externalAddr != "" {
		job.Config.Env = appendUnique(job.Config.Env, "EXTERNAL_IP="+p.externalAddr, "SD_HOST="+p.externalAddr, "DISCOVERD="+p.discoverd)
	}

	p.state.AddJob(job)
	g.Log(grohl.Data{"at": "create_container"})
	container, err := p.docker.CreateContainer(job.Config)
	if err == docker.ErrNoSuchImage {
		g.Log(grohl.Data{"at": "pull_image"})
		err = p.docker.PullImage(docker.PullImageOptions{Repository: job.Config.Image}, os.Stdout)
		if err != nil {
			g.Log(grohl.Data{"at": "pull_image", "status": "error", "err": err})
			p.state.SetStatusFailed(job.ID, err)
			return nil, err
		}
		container, err = p.docker.CreateContainer(job.Config)
	}
	if err != nil {
		g.Log(grohl.Data{"at": "create_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetContainerID(job.ID, container.ID)
	p.state.WaitAttach(job.ID)
	g.Log(grohl.Data{"at": "start_container"})
	if err := p.docker.StartContainer(container.ID, job.HostConfig); err != nil {
		g.Log(grohl.Data{"at": "start_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	container, err = p.docker.InspectContainer(container.ID)
	if err != nil {
		g.Log(grohl.Data{"at": "inspect_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetStatusRunning(job.ID, container.Volumes)
	g.Log(grohl.Data{"at": "finish"})
	return container, nil
}

func appendUnique(s []string, vars ...string) []string {
outer:
	for _, v := range vars {
		for _, existing := range s {
			if strings.HasPrefix(existing, strings.SplitN(v, "=", 2)[0]+"=") {
				continue outer
			}
		}
		s = append(s, v)
	}
	return s
}

type sampiSyncClient interface {
	RemoveJobs([]string) error
}

func syncScheduler(scheduler sampiSyncClient, events <-chan host.Event) {
	for event := range events {
		if event.Event != "stop" {
			continue
		}
		grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "job.id": event.JobID})
		if err := scheduler.RemoveJobs([]string{event.JobID}); err != nil {
			grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "status": "error", "err": err, "job.id": event.JobID})
			// TODO: try to reconnect?
		}
	}
}

type dockerStreamClient interface {
	Events() (*docker.EventStream, error)
	InspectContainer(string) (*docker.Container, error)
}

func streamEvents(client dockerStreamClient, state *State) {
	stream, err := client.Events()
	if err != nil {
		log.Fatal(err)
	}
	for event := range stream.Events {
		if event.Status != "die" {
			continue
		}
		container, err := client.InspectContainer(event.ID)
		if err != nil {
			log.Println("inspect container", event.ID, "error:", err)
			// TODO: set job status anyway?
			continue
		}
		state.SetStatusDone(event.ID, container.State.ExitCode)
	}
}

// TODO: fix this, horribly broken

func allocatePorts(ports chan<- int, startPort, endPort int) {
	for i := startPort; i < endPort; i++ {
		ports <- i
	}
	// TODO: handle wrap-around
}
<|endoftext|>


<commit_before>/*
Copyright 2015 All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package main

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"

	log "github.com/Sirupsen/logrus"
	"github.com/coreos/go-oidc/jose"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

const (
	fakeClientID = "test"
	fakeSecret   = fakeClientID

	fakeAdminRoleURL       = "/admin"
	fakeTestRoleURL        = "/test_role"
	fakeTestAdminRolesURL  = "/test_admin_roles"
	fakeAuthAllURL         = "/auth_all"
	fakeTestWhitelistedURL = fakeAuthAllURL + "/white_listed"
	fakeTestListenOrdered  = fakeAuthAllURL + "/bad_order"

	fakeAdminRole = "role:admin"
	fakeTestRole  = "role:test"
)

func newFakeKeycloakProxyWithResources(t *testing.T, resources []*Resource) *oauthProxy {
	kc := newFakeKeycloakProxy(t)
	kc.config.Resources = resources
	kc.endpoint = &url.URL{
		Host: "127.0.0.1",
	}
	return kc
}

func newFakeKeycloakConfig(t *testing.T) *Config {
	return &Config{
		DiscoveryURL:          "127.0.0.1:8080",
		ClientID:              fakeClientID,
		ClientSecret:          fakeSecret,
		EncryptionKey:         "AgXa7xRcoClDEU0ZDSH4X0XhL5Qy2Z2j",
		SkipTokenVerification: true,
		Scopes:                []string{},
		EnableRefreshTokens:   false,
		SecureCookie:          false,
		CookieAccessName:      "kc-access",
		CookieRefreshName:     "kc-state",
		Resources: []*Resource{
			{
				URL:     fakeAdminRoleURL,
				Methods: []string{"GET"},
				Roles:   []string{fakeAdminRole},
			},
			{
				URL:     fakeTestRoleURL,
				Methods: []string{"GET"},
				Roles:   []string{fakeTestRole},
			},
			{
				URL:     fakeTestAdminRolesURL,
				Methods: []string{"GET"},
				Roles:   []string{fakeAdminRole, fakeTestRole},
			},
			{
				URL:         fakeTestWhitelistedURL,
				WhiteListed: true,
				Methods:     []string{},
				Roles:       []string{},
			},
			{
				URL:     fakeAuthAllURL,
				Methods: []string{"ANY"},
				Roles:   []string{},
			},
			{
				URL:         fakeTestWhitelistedURL,
				WhiteListed: true,
				Methods:     []string{},
				Roles:       []string{},
			},
		},
		CrossOrigin: CORS{},
	}
}

func newFakeKeycloakProxy(t *testing.T) *oauthProxy {
	log.SetOutput(ioutil.Discard)

	kc := &oauthProxy{
		config:   newFakeKeycloakConfig(t),
		upstream: new(fakeReverseProxy),
		endpoint: &url.URL{Host: "127.0.0.1"},
	}
	gin.SetMode(gin.ReleaseMode)
	kc.router = gin.New()
	// step: add the gin routing
	kc.createEndpoints()

	return kc
}

func newTestProxyService(t *testing.T, config *Config) (*oauthProxy, *fakeOAuthServer, string) {
	auth := newFakeOAuthServer(t)
	if config == nil {
		config = newFakeKeycloakConfig(t)
	}
	config.LogRequests = true
	config.SkipTokenVerification = false
	config.DiscoveryURL = auth.getLocation()
	config.Verbose = false
	log.SetOutput(ioutil.Discard)

	proxy, err := newProxy(config)
	if err != nil {
		t.Fatalf("failed to create proxy service, error: %s", err)
	}
	proxy.upstream = new(fakeReverseProxy)
	service := httptest.NewServer(proxy.router)
	config.RedirectionURL = service.URL

	proxy.client, proxy.provider, err = createOpenIDClient(config)
	if err != nil {
		t.Fatalf("failed to recreate the openid client, error: %s", err)
	}

	return proxy, auth, service.URL
}

func TestNewKeycloakProxy(t *testing.T) {
	proxy, err := newProxy(newFakeKeycloakConfig(t))
	assert.NoError(t, err)
	assert.NotNil(t, proxy)
	assert.NotNil(t, proxy.config)
	assert.NotNil(t, proxy.router)
	assert.NotNil(t, proxy.endpoint)
}

func TestRedirectToAuthorization(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	proxy.config.SkipTokenVerification = false
	proxy.redirectToAuthorization(context)
	assert.Equal(t, http.StatusTemporaryRedirect, context.Writer.Status())
}

func TestRedirectToAuthorizationSkipToken(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	proxy.config.SkipTokenVerification = true
	proxy.redirectToAuthorization(context)
	assert.Equal(t, http.StatusForbidden, context.Writer.Status())
}

func TestRedirectToAuthorizationUnauthorized(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)
	proxy.config.SkipTokenVerification = false
	proxy.config.NoRedirects = true

	proxy.redirectToAuthorization(context)
	assert.Equal(t, http.StatusUnauthorized, context.Writer.Status())
}

func TestInitializeReverseProxy(t *testing.T) {
	proxy := newFakeKeycloakProxy(t)

	uri, _ := url.Parse("http://127.0.0.1:8080")
	err := proxy.createUpstream(uri)
	assert.NoError(t, err)
	assert.NotNil(t, proxy.router)
}

func TestRedirectURL(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	if proxy.redirectToURL("http://127.0.0.1", context); context.Writer.Status() != http.StatusTemporaryRedirect {
		t.Errorf("we should have recieved a redirect")
	}

	if !context.IsAborted() {
		t.Errorf("the context should have been aborted")
	}
}

func TestAccessForbidden(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	proxy.config.SkipTokenVerification = false
	if proxy.accessForbidden(context); context.Writer.Status() != http.StatusForbidden {
		t.Errorf("we should have recieved a forbidden access")
	}

	proxy.config.SkipTokenVerification = true
	if proxy.accessForbidden(context); context.Writer.Status() != http.StatusForbidden {
		t.Errorf("we should have recieved a forbidden access")
	}
}

func newFakeResponse() *fakeResponse {
	return &fakeResponse{
		status:  http.StatusOK,
		headers: make(http.Header, 0),
	}
}

func newFakeGinContext(method, uri string) *gin.Context {
	return &gin.Context{
		Request: &http.Request{
			Method:     method,
			Host:       "127.0.0.1",
			RequestURI: uri,
			URL: &url.URL{
				Scheme: "http",
				Host:   "127.0.0.1",
				Path:   uri,
			},
			Header:     make(http.Header, 0),
			RemoteAddr: "127.0.0.1:8989",
		},
		Writer: newFakeResponse(),
	}
}

func newFakeGinContextWithCookies(method, url string, cookies []*http.Cookie) *gin.Context {
	cx := newFakeGinContext(method, url)
	for _, x := range cookies {
		cx.Request.AddCookie(x)
	}

	return cx
}

func newFakeJWTToken(t *testing.T, claims jose.Claims) *jose.JWT {
	token, err := jose.NewJWT(
		jose.JOSEHeader{"alg": "RS256"}, claims,
	)
	if err != nil {
		t.Fatalf("failed to create the jwt token, error: %s", err)
	}

	return &token
}

type fakeReverseProxy struct{}

func (r fakeReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {}

type fakeResponse struct {
	size    int
	status  int
	headers http.Header
	body    bytes.Buffer
	written bool
}

func (r *fakeResponse) Flush()              {}
func (r *fakeResponse) Written() bool       { return r.written }
func (r *fakeResponse) WriteHeaderNow()     {}
func (r *fakeResponse) Size() int           { return r.size }
func (r *fakeResponse) Status() int         { return r.status }
func (r *fakeResponse) Header() http.Header { return r.headers }
func (r *fakeResponse) WriteHeader(code int) {
	r.status = code
	r.written = true
}
func (r *fakeResponse) Write(content []byte) (int, error)            { return len(content), nil }
func (r *fakeResponse) WriteString(s string) (int, error)            { return len(s), nil }
func (r *fakeResponse) Hijack() (net.Conn, *bufio.ReadWriter, error) { return nil, nil, nil }
func (r *fakeResponse) CloseNotify() <-chan bool                     { return make(chan bool, 0) }
<commit_msg>- fixed the typo on the server_test.go<commit_after>/*
Copyright 2015 All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package main

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"

	log "github.com/Sirupsen/logrus"
	"github.com/coreos/go-oidc/jose"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

const (
	fakeClientID = "test"
	fakeSecret   = fakeClientID

	fakeAdminRoleURL       = "/admin"
	fakeTestRoleURL        = "/test_role"
	fakeTestAdminRolesURL  = "/test_admin_roles"
	fakeAuthAllURL         = "/auth_all"
	fakeTestWhitelistedURL = fakeAuthAllURL + "/white_listed"
	fakeTestListenOrdered  = fakeAuthAllURL + "/bad_order"

	fakeAdminRole = "role:admin"
	fakeTestRole  = "role:test"
)

func newFakeKeycloakProxyWithResources(t *testing.T, resources []*Resource) *oauthProxy {
	kc := newFakeKeycloakProxy(t)
	kc.config.Resources = resources
	kc.endpoint = &url.URL{
		Host: "127.0.0.1",
	}
	return kc
}

func newFakeKeycloakConfig(t *testing.T) *Config {
	return &Config{
		DiscoveryURL:          "127.0.0.1:8080",
		ClientID:              fakeClientID,
		ClientSecret:          fakeSecret,
		EncryptionKey:         "AgXa7xRcoClDEU0ZDSH4X0XhL5Qy2Z2j",
		SkipTokenVerification: true,
		Scopes:                []string{},
		EnableRefreshTokens:   false,
		SecureCookie:          false,
		CookieAccessName:      "kc-access",
		CookieRefreshName:     "kc-state",
		Resources: []*Resource{
			{
				URL:     fakeAdminRoleURL,
				Methods: []string{"GET"},
				Roles:   []string{fakeAdminRole},
			},
			{
				URL:     fakeTestRoleURL,
				Methods: []string{"GET"},
				Roles:   []string{fakeTestRole},
			},
			{
				URL:     fakeTestAdminRolesURL,
				Methods: []string{"GET"},
				Roles:   []string{fakeAdminRole, fakeTestRole},
			},
			{
				URL:         fakeTestWhitelistedURL,
				WhiteListed: true,
				Methods:     []string{},
				Roles:       []string{},
			},
			{
				URL:     fakeAuthAllURL,
				Methods: []string{"ANY"},
				Roles:   []string{},
			},
			{
				URL:         fakeTestWhitelistedURL,
				WhiteListed: true,
				Methods:     []string{},
				Roles:       []string{},
			},
		},
		CrossOrigin: CORS{},
	}
}

func newFakeKeycloakProxy(t *testing.T) *oauthProxy {
	log.SetOutput(ioutil.Discard)

	kc := &oauthProxy{
		config:   newFakeKeycloakConfig(t),
		upstream: new(fakeReverseProxy),
		endpoint: &url.URL{Host: "127.0.0.1"},
	}
	gin.SetMode(gin.ReleaseMode)
	kc.router = gin.New()
	// step: add the gin routing
	kc.createEndpoints()

	return kc
}

func newTestProxyService(t *testing.T, config *Config) (*oauthProxy, *fakeOAuthServer, string) {
	auth := newFakeOAuthServer(t)
	if config == nil {
		config = newFakeKeycloakConfig(t)
	}
	config.LogRequests = true
	config.SkipTokenVerification = false
	config.DiscoveryURL = auth.getLocation()
	config.Verbose = false
	log.SetOutput(ioutil.Discard)

	proxy, err := newProxy(config)
	if err != nil {
		t.Fatalf("failed to create proxy service, error: %s", err)
	}
	proxy.upstream = new(fakeReverseProxy)
	service := httptest.NewServer(proxy.router)
	config.RedirectionURL = service.URL

	proxy.client, proxy.provider, err = createOpenIDClient(config)
	if err != nil {
		t.Fatalf("failed to recreate the openid client, error: %s", err)
	}

	return proxy, auth, service.URL
}

func TestNewKeycloakProxy(t *testing.T) {
	proxy, err := newProxy(newFakeKeycloakConfig(t))
	assert.NoError(t, err)
	assert.NotNil(t, proxy)
	assert.NotNil(t, proxy.config)
	assert.NotNil(t, proxy.router)
	assert.NotNil(t, proxy.endpoint)
}

func TestRedirectToAuthorization(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	proxy.config.SkipTokenVerification = false
	proxy.redirectToAuthorization(context)
	assert.Equal(t, http.StatusTemporaryRedirect, context.Writer.Status())
}

func TestRedirectToAuthorizationSkipToken(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	proxy.config.SkipTokenVerification = true
	proxy.redirectToAuthorization(context)
	assert.Equal(t, http.StatusForbidden, context.Writer.Status())
}

func TestRedirectToAuthorizationUnauthorized(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)
	proxy.config.SkipTokenVerification = false
	proxy.config.NoRedirects = true

	proxy.redirectToAuthorization(context)
	assert.Equal(t, http.StatusUnauthorized, context.Writer.Status())
}

func TestInitializeReverseProxy(t *testing.T) {
	proxy := newFakeKeycloakProxy(t)

	uri, _ := url.Parse("http://127.0.0.1:8080")
	err := proxy.createUpstreamProxy(uri)
	assert.NoError(t, err)
	assert.NotNil(t, proxy.router)
}

func TestRedirectURL(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	if proxy.redirectToURL("http://127.0.0.1", context); context.Writer.Status() != http.StatusTemporaryRedirect {
		t.Errorf("we should have recieved a redirect")
	}

	if !context.IsAborted() {
		t.Errorf("the context should have been aborted")
	}
}

func TestAccessForbidden(t *testing.T) {
	context := newFakeGinContext("GET", "/admin")
	proxy := newFakeKeycloakProxy(t)

	proxy.config.SkipTokenVerification = false
	if proxy.accessForbidden(context); context.Writer.Status() != http.StatusForbidden {
		t.Errorf("we should have recieved a forbidden access")
	}

	proxy.config.SkipTokenVerification = true
	if proxy.accessForbidden(context); context.Writer.Status() != http.StatusForbidden {
		t.Errorf("we should have recieved a forbidden access")
	}
}

func newFakeResponse() *fakeResponse {
	return &fakeResponse{
		status:  http.StatusOK,
		headers: make(http.Header, 0),
	}
}

func newFakeGinContext(method, uri string) *gin.Context {
	return &gin.Context{
		Request: &http.Request{
			Method:     method,
			Host:       "127.0.0.1",
			RequestURI: uri,
			URL: &url.URL{
				Scheme: "http",
				Host:   "127.0.0.1",
				Path:   uri,
			},
			Header:     make(http.Header, 0),
			RemoteAddr: "127.0.0.1:8989",
		},
		Writer: newFakeResponse(),
	}
}

func newFakeGinContextWithCookies(method, url string, cookies []*http.Cookie) *gin.Context {
	cx := newFakeGinContext(method, url)
	for _, x := range cookies {
		cx.Request.AddCookie(x)
	}

	return cx
}

func newFakeJWTToken(t *testing.T, claims jose.Claims) *jose.JWT {
	token, err := jose.NewJWT(
		jose.JOSEHeader{"alg": "RS256"}, claims,
	)
	if err != nil {
		t.Fatalf("failed to create the jwt token, error: %s", err)
	}

	return &token
}

type fakeReverseProxy struct{}

func (r fakeReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {}

type fakeResponse struct {
	size    int
	status  int
	headers http.Header
	body    bytes.Buffer
	written bool
}

func (r *fakeResponse) Flush()              {}
func (r *fakeResponse) Written() bool       { return r.written }
func (r *fakeResponse) WriteHeaderNow()     {}
func (r *fakeResponse) Size() int           { return r.size }
func (r *fakeResponse) Status() int         { return r.status }
func (r *fakeResponse) Header() http.Header { return r.headers }
func (r *fakeResponse) WriteHeader(code int) {
	r.status = code
	r.written = true
}
func (r *fakeResponse) Write(content []byte) (int, error)            { return len(content), nil }
func (r *fakeResponse) WriteString(s string) (int, error)            { return len(s), nil }
func (r *fakeResponse) Hijack() (net.Conn, *bufio.ReadWriter, error) { return nil, nil, nil }
func (r *fakeResponse) CloseNotify() <-chan bool                     { return make(chan bool, 0) }
<|endoftext|>


<commit_before>package server

import (
	"io/ioutil"
	"strconv"
	"testing"

	"net/http"

	. "github.com/bborbe/assert"

 	http_client_builder "github.com/bborbe/http/client_builder"
 )

func TestImplementsServer(t *testing.T) {
	addr := ":12345"
	handler := new(helloWorld)
	srv := NewServer(addr, handler)
	var server *Server
	err := AssertThat(srv, Implements(server).Message("should be Server"))
	if err != nil {
		t.Fatal(err)
	}
}

func TestStartStop(t *testing.T) {
	var err error
	port := 12345
	addr := "0.0.0.0:" + strconv.Itoa(port)
	handler := new(helloWorld)
	srv := NewServer(addr, handler)
	err = AssertThat(srv, NotNilValue().Message("create server failed"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Start(), NilValue().Message("not started"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Start(), NotNilValue().Message("already started"))
	if err != nil {
		t.Fatal(err)
	}

	d := http_client.NewNoProxy()
	resp, err := d.Get("http://" + addr)
	if err != nil {
		t.Fatal(err)
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(string(b), Is("hello, world!\n").Message("content false"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Stop(), NilValue().Message("not stopped"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Stop(), NotNilValue().Message("already stopped"))
	if err != nil {
		t.Fatal(err)
	}
}

type helloWorld struct {
}

func (h *helloWorld) GetPath() string {
	return "/"
}

func (h *helloWorld) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte("hello, world!\n"))
}
<commit_msg>remove client adapter<commit_after>package server

import (
	"io/ioutil"
	"strconv"
	"testing"

	"net/http"

	. "github.com/bborbe/assert"
	http_client "github.com/bborbe/http/client"
	http_client_builder "github.com/bborbe/http/client/builder"
 )

func TestImplementsServer(t *testing.T) {
	addr := ":12345"
	handler := new(helloWorld)
	srv := NewServer(addr, handler)
	var server *Server
	err := AssertThat(srv, Implements(server).Message("should be Server"))
	if err != nil {
		t.Fatal(err)
	}
}

func TestStartStop(t *testing.T) {
	var err error
	port := 12345
	addr := "0.0.0.0:" + strconv.Itoa(port)
	handler := new(helloWorld)
	srv := NewServer(addr, handler)
	err = AssertThat(srv, NotNilValue().Message("create server failed"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Start(), NilValue().Message("not started"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Start(), NotNilValue().Message("already started"))
	if err != nil {
		t.Fatal(err)
	}

	httpClientBuilder := http_client_builder.New()
	httpClient := http_client.New(httpClientBuilder.Build())

	resp, err := httpClient.Get("http://" + addr)
	if err != nil {
		t.Fatal(err)
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(string(b), Is("hello, world!\n").Message("content false"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Stop(), NilValue().Message("not stopped"))
	if err != nil {
		t.Fatal(err)
	}

	err = AssertThat(srv.Stop(), NotNilValue().Message("already stopped"))
	if err != nil {
		t.Fatal(err)
	}
}

type helloWorld struct {
}

func (h *helloWorld) GetPath() string {
	return "/"
}

func (h *helloWorld) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte("hello, world!\n"))
}
<|endoftext|>


<commit_before>package service

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

	"github.com/tracerun/tracerun/lg"
	"go.uber.org/zap"
)

const (
	readTimeout = 30
)

// TCPServer to define a TCP server
type TCPServer struct {
	port   uint16
	router map[uint8]RouteFunc
	ln     net.Listener
}

// NewTCPServer to create a TCP server instance
func NewTCPServer(port uint16, router map[uint8]RouteFunc) *TCPServer {
	return &TCPServer{
		port:   port,
		router: router,
	}
}

// Start the server
func (s *TCPServer) Start() error {
	if s.ln != nil {
		return fmt.Errorf("already started")
	}

	port := fmt.Sprintf(":%d", s.port)
	// net.ListenUDP("udp", laddr*net.UDPAddr)
	ln, err := net.Listen("tcp", port)
	if err != nil {
		return err
	}

	s.ln = ln
	lg.L.Info("started to listen socket connections", zap.Uint16("port", s.port))

	for {
		conn, err := ln.Accept()

		if err != nil {
			lg.L.Error("error accept connection", zap.Error(err))
		}
		lg.L.Debug("new connection come")

		go s.handleConn(conn)
	}
}

// Stop the server
func (s *TCPServer) Stop() error {
	if s.ln != nil {
		return s.ln.Close()
	}
	return nil
}

func (s *TCPServer) handleConn(c net.Conn) {
	defer func() {
		if r := recover(); r != nil {
			lg.L.Warn("recovered", zap.Any("error", r), zap.Stack("info"))
		}
	}()

	for {
		// read header
		c.SetReadDeadline(time.Now().Add(readTimeout * time.Second))
		count, route, err := readHeader(c)

		if err != nil {
			if err == io.EOF { // may be closed by client
				break
			}
			if neterr, ok := err.(net.Error); ok && neterr.Timeout() { // timeout
				break
			}
			lg.L.Error("error to read header", zap.Error(err))
		}
		lg.L.Debug("header", zap.Uint8("route", route), zap.Uint16("count", count))

		// read data
		var bytes []byte
		if count > 0 {
			c.SetReadDeadline(time.Now().Add(readTimeout * time.Second))
			bytes, err = readData(c, count)

			if err != nil {
				if err == io.EOF { // may be closed by client
					break
				}
				if neterr, ok := err.(net.Error); ok && neterr.Timeout() { // timeout
					break
				}
				lg.L.Error("error to read data", zap.Error(err))
			}
			lg.L.Debug("data", zap.Binary("data", bytes))
		}

		// get routed function
		fn, ok := s.router[route]
		if !ok {
			lg.L.Warn("not found")
		} else {
			fn(bytes, c)
		}
	}

	if err := c.Close(); err != nil {
		lg.L.Error("error close", zap.Error(err))
	}
	lg.L.Debug("connection closed")
}
<commit_msg>Fix some error detections.<commit_after>package service

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

	"github.com/tracerun/tracerun/lg"
	"go.uber.org/zap"
)

const (
	readTimeout = 30
)

// TCPServer to define a TCP server
type TCPServer struct {
	port   uint16
	router map[uint8]RouteFunc
	ln     net.Listener
}

// NewTCPServer to create a TCP server instance
func NewTCPServer(port uint16, router map[uint8]RouteFunc) *TCPServer {
	return &TCPServer{
		port:   port,
		router: router,
	}
}

// Start the server
func (s *TCPServer) Start() error {
	if s.ln != nil {
		return fmt.Errorf("already started")
	}

	port := fmt.Sprintf(":%d", s.port)
	// net.ListenUDP("udp", laddr*net.UDPAddr)
	ln, err := net.Listen("tcp", port)
	if err != nil {
		return err
	}

	s.ln = ln
	lg.L.Info("started to listen socket connections", zap.Uint16("port", s.port))

	for {
		conn, err := ln.Accept()

		if err != nil {
			lg.L.Error("error accept connection", zap.Error(err))
		}
		lg.L.Debug("new connection come")

		go s.handleConn(conn)
	}
}

// Stop the server
func (s *TCPServer) Stop() error {
	if s.ln != nil {
		return s.ln.Close()
	}
	return nil
}

func (s *TCPServer) handleConn(c net.Conn) {
	defer func() {
		if r := recover(); r != nil {
			lg.L.Warn("recovered", zap.Any("error", r), zap.Stack("info"))
		}
	}()

	for {
		// read header
		c.SetReadDeadline(time.Now().Add(readTimeout * time.Second))
		count, route, err := readHeader(c)

		if err != nil {
			recordConnError(err)
			break
		}
		lg.L.Debug("header", zap.Uint8("route", route), zap.Uint16("count", count))

		// read data
		var bytes []byte
		if count > 0 {
			c.SetReadDeadline(time.Now().Add(readTimeout * time.Second))
			bytes, err = readData(c, count)

			if err != nil {
				recordConnError(err)
				break
			}
			lg.L.Debug("data", zap.Binary("data", bytes))
		}

		// get routed function
		fn, ok := s.router[route]
		if !ok {
			lg.L.Warn("not found")
		} else {
			fn(bytes, c)
		}
	}

	if err := c.Close(); err != nil {
		lg.L.Error("error close", zap.Error(err))
	}
	lg.L.Debug("connection closed")
}

func recordConnError(err error) {
	if err == io.EOF {
		lg.L.Debug("EOF")
		return
	}
	if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
		lg.L.Debug("timeout")
		return
	}
	if _, ok := err.(*net.OpError); ok {
		lg.L.Debug("operror")
		return
	}

	lg.L.Error("error to read data", zap.Error(err))
}
<|endoftext|>