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

package web

import (
	"context"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/prometheus/storage/tsdb"
	libtsdb "github.com/prometheus/tsdb"
)

func TestGlobalURL(t *testing.T) {
	opts := &Options{
		ListenAddress: ":9090",
		ExternalURL: &url.URL{
			Scheme: "https",
			Host:   "externalhost:80",
			Path:   "/path/prefix",
		},
	}

	tests := []struct {
		inURL  string
		outURL string
	}{
		{
			// Nothing should change if the input URL is not on localhost, even if the port is our listening port.
			inURL:  "http://somehost:9090/metrics",
			outURL: "http://somehost:9090/metrics",
		},
		{
			// Port and host should change if target is on localhost and port is our listening port.
			inURL:  "http://localhost:9090/metrics",
			outURL: "https://externalhost:80/metrics",
		},
		{
			// Only the host should change if the port is not our listening port, but the host is localhost.
			inURL:  "http://localhost:8000/metrics",
			outURL: "http://externalhost:8000/metrics",
		},
		{
			// Alternative localhost representations should also work.
			inURL:  "http://127.0.0.1:9090/metrics",
			outURL: "https://externalhost:80/metrics",
		},
	}

	for i, test := range tests {
		inURL, err := url.Parse(test.inURL)
		if err != nil {
			t.Fatalf("%d. Error parsing input URL: %s", i, err)
		}
		globalURL := tmplFuncs("", opts)["globalURL"].(func(u *url.URL) *url.URL)
		outURL := globalURL(inURL)

		if outURL.String() != test.outURL {
			t.Fatalf("%d. got %s, want %s", i, outURL.String(), test.outURL)
		}
	}
}

func TestReadyAndHealthy(t *testing.T) {
	t.Parallel()
	dbDir, err := ioutil.TempDir("", "tsdb-ready")
	if err != nil {
		t.Fatalf("Unexpected error creating a tmpDir: %s", err)
	}
	defer os.RemoveAll(dbDir)
	db, err := libtsdb.Open(dbDir, nil, nil, nil)
	if err != nil {
		t.Fatalf("Unexpected error opening empty dir: %s", err)
	}

	opts := &Options{
		ListenAddress:  ":9090",
		ReadTimeout:    30 * time.Second,
		MaxConnections: 512,
		Context:        nil,
		Storage:        &tsdb.ReadyStorage{},
		QueryEngine:    nil,
		TargetManager:  nil,
		RuleManager:    nil,
		Notifier:       nil,
		RoutePrefix:    "/",
		MetricsPath:    "/metrics/",
		EnableAdminAPI: true,
		TSDB:           func() *libtsdb.DB { return db },
	}

	opts.Flags = map[string]string{}

	webHandler := New(nil, opts)
	go webHandler.Run(context.Background())

	// Give some time for the web goroutine to run since we need the server
	// to be up before starting tests.
	time.Sleep(5 * time.Second)

	resp, err := http.Get("http://localhost:9090/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /-/healthy with server unready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /-/ready with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /version with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", resp.Status)
	}

	// Set to ready.
	webHandler.Ready()

	resp, err = http.Get("http://localhost:9090/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /-/healthy with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /-/ready with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /version with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", resp.Status)
	}
}

func TestRoutePrefix(t *testing.T) {
	t.Parallel()
	dbDir, err := ioutil.TempDir("", "tsdb-ready")
	if err != nil {
		t.Fatalf("Unexpected error creating a tmpDir: %s", err)
	}
	defer os.RemoveAll(dbDir)
	db, err := libtsdb.Open(dbDir, nil, nil, nil)
	if err != nil {
		t.Fatalf("Unexpected error opening empty dir: %s", err)
	}

	opts := &Options{
		ListenAddress:  ":9091",
		ReadTimeout:    30 * time.Second,
		MaxConnections: 512,
		Context:        nil,
		Storage:        &tsdb.ReadyStorage{},
		QueryEngine:    nil,
		TargetManager:  nil,
		RuleManager:    nil,
		Notifier:       nil,
		RoutePrefix:    "/prometheus",
		MetricsPath:    "/prometheus/metrics",
		EnableAdminAPI: true,
		TSDB:           func() *libtsdb.DB { return db },
	}

	opts.Flags = map[string]string{}

	webHandler := New(nil, opts)
	go func() {
		err := webHandler.Run(context.Background())
		if err != nil {
			t.Fatalf("Can't start webhandler error %s", err)
		}
	}()

	// Give some time for the web goroutine to run since we need the server
	// to be up before starting tests.
	time.Sleep(5 * time.Second)

	resp, err := http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path %s/-/healthy with server unready test, Expected status 200 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/-/ready with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/version with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	// Set to ready.
	webHandler.Ready()

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path "+opts.RoutePrefix+"/-/healthy with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path "+opts.RoutePrefix+"/-/ready with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path "+opts.RoutePrefix+"/version with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path %s/api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path %s/api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}
}

func TestDebugHandler(t *testing.T) {
	for _, tc := range []struct {
		prefix, url string
		code        int
	}{
		{"/", "/debug/pprof/cmdline", 200},
		{"/foo", "/foo/debug/pprof/cmdline", 200},

		{"/", "/debug/pprof/goroutine", 200},
		{"/foo", "/foo/debug/pprof/goroutine", 200},

		{"/", "/debug/pprof/foo", 404},
		{"/foo", "/bar/debug/pprof/goroutine", 404},
	} {
		opts := &Options{
			RoutePrefix: tc.prefix,
			MetricsPath: "/metrics",
		}
		handler := New(nil, opts)
		handler.Ready()

		w := httptest.NewRecorder()
		req, err := http.NewRequest("GET", tc.url, nil)
		if err != nil {
			t.Fatalf("Unexpected error %s", err)
		}

		handler.router.ServeHTTP(w, req)
		if w.Code != tc.code {
			t.Fatalf("Unexpected status code %d: %s", w.Code, w.Body.String())
		}
	}
}
<commit_msg>Use panic in goroutine in test, not t.Fatalf<commit_after>// Copyright 2016 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package web

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/prometheus/storage/tsdb"
	libtsdb "github.com/prometheus/tsdb"
)

func TestGlobalURL(t *testing.T) {
	opts := &Options{
		ListenAddress: ":9090",
		ExternalURL: &url.URL{
			Scheme: "https",
			Host:   "externalhost:80",
			Path:   "/path/prefix",
		},
	}

	tests := []struct {
		inURL  string
		outURL string
	}{
		{
			// Nothing should change if the input URL is not on localhost, even if the port is our listening port.
			inURL:  "http://somehost:9090/metrics",
			outURL: "http://somehost:9090/metrics",
		},
		{
			// Port and host should change if target is on localhost and port is our listening port.
			inURL:  "http://localhost:9090/metrics",
			outURL: "https://externalhost:80/metrics",
		},
		{
			// Only the host should change if the port is not our listening port, but the host is localhost.
			inURL:  "http://localhost:8000/metrics",
			outURL: "http://externalhost:8000/metrics",
		},
		{
			// Alternative localhost representations should also work.
			inURL:  "http://127.0.0.1:9090/metrics",
			outURL: "https://externalhost:80/metrics",
		},
	}

	for i, test := range tests {
		inURL, err := url.Parse(test.inURL)
		if err != nil {
			t.Fatalf("%d. Error parsing input URL: %s", i, err)
		}
		globalURL := tmplFuncs("", opts)["globalURL"].(func(u *url.URL) *url.URL)
		outURL := globalURL(inURL)

		if outURL.String() != test.outURL {
			t.Fatalf("%d. got %s, want %s", i, outURL.String(), test.outURL)
		}
	}
}

func TestReadyAndHealthy(t *testing.T) {
	t.Parallel()
	dbDir, err := ioutil.TempDir("", "tsdb-ready")
	if err != nil {
		t.Fatalf("Unexpected error creating a tmpDir: %s", err)
	}
	defer os.RemoveAll(dbDir)
	db, err := libtsdb.Open(dbDir, nil, nil, nil)
	if err != nil {
		t.Fatalf("Unexpected error opening empty dir: %s", err)
	}

	opts := &Options{
		ListenAddress:  ":9090",
		ReadTimeout:    30 * time.Second,
		MaxConnections: 512,
		Context:        nil,
		Storage:        &tsdb.ReadyStorage{},
		QueryEngine:    nil,
		TargetManager:  nil,
		RuleManager:    nil,
		Notifier:       nil,
		RoutePrefix:    "/",
		MetricsPath:    "/metrics/",
		EnableAdminAPI: true,
		TSDB:           func() *libtsdb.DB { return db },
	}

	opts.Flags = map[string]string{}

	webHandler := New(nil, opts)
	go webHandler.Run(context.Background())

	// Give some time for the web goroutine to run since we need the server
	// to be up before starting tests.
	time.Sleep(5 * time.Second)

	resp, err := http.Get("http://localhost:9090/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /-/healthy with server unready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /-/ready with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /version with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path /api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", resp.Status)
	}

	// Set to ready.
	webHandler.Ready()

	resp, err = http.Get("http://localhost:9090/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /-/healthy with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /-/ready with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9090/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /version with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9090/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path /api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", resp.Status)
	}
}

func TestRoutePrefix(t *testing.T) {
	t.Parallel()
	dbDir, err := ioutil.TempDir("", "tsdb-ready")
	if err != nil {
		t.Fatalf("Unexpected error creating a tmpDir: %s", err)
	}
	defer os.RemoveAll(dbDir)
	db, err := libtsdb.Open(dbDir, nil, nil, nil)
	if err != nil {
		t.Fatalf("Unexpected error opening empty dir: %s", err)
	}

	opts := &Options{
		ListenAddress:  ":9091",
		ReadTimeout:    30 * time.Second,
		MaxConnections: 512,
		Context:        nil,
		Storage:        &tsdb.ReadyStorage{},
		QueryEngine:    nil,
		TargetManager:  nil,
		RuleManager:    nil,
		Notifier:       nil,
		RoutePrefix:    "/prometheus",
		MetricsPath:    "/prometheus/metrics",
		EnableAdminAPI: true,
		TSDB:           func() *libtsdb.DB { return db },
	}

	opts.Flags = map[string]string{}

	webHandler := New(nil, opts)
	go func() {
		err := webHandler.Run(context.Background())
		if err != nil {
			panic(fmt.Sprintf("Can't start webhandler error %s", err))
		}
	}()

	// Give some time for the web goroutine to run since we need the server
	// to be up before starting tests.
	time.Sleep(5 * time.Second)

	resp, err := http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path %s/-/healthy with server unready test, Expected status 200 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/-/ready with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/version with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("Path %s/api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	// Set to ready.
	webHandler.Ready()

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/healthy")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path "+opts.RoutePrefix+"/-/healthy with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/-/ready")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path "+opts.RoutePrefix+"/-/ready with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Get("http://localhost:9091" + opts.RoutePrefix + "/version")
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path "+opts.RoutePrefix+"/version with server ready test, Expected status 200 got: %s", resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/snapshot", "", strings.NewReader(""))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path %s/api/v2/admin/tsdb/snapshot with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}

	resp, err = http.Post("http://localhost:9091"+opts.RoutePrefix+"/api/v2/admin/tsdb/delete_series", "", strings.NewReader("{}"))
	if err != nil {
		t.Fatalf("Unexpected HTTP error %s", err)
	}
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("Path %s/api/v2/admin/tsdb/delete_series with server unready test, Expected status 503 got: %s", opts.RoutePrefix, resp.Status)
	}
}

func TestDebugHandler(t *testing.T) {
	for _, tc := range []struct {
		prefix, url string
		code        int
	}{
		{"/", "/debug/pprof/cmdline", 200},
		{"/foo", "/foo/debug/pprof/cmdline", 200},

		{"/", "/debug/pprof/goroutine", 200},
		{"/foo", "/foo/debug/pprof/goroutine", 200},

		{"/", "/debug/pprof/foo", 404},
		{"/foo", "/bar/debug/pprof/goroutine", 404},
	} {
		opts := &Options{
			RoutePrefix: tc.prefix,
			MetricsPath: "/metrics",
		}
		handler := New(nil, opts)
		handler.Ready()

		w := httptest.NewRecorder()
		req, err := http.NewRequest("GET", tc.url, nil)
		if err != nil {
			t.Fatalf("Unexpected error %s", err)
		}

		handler.router.ServeHTTP(w, req)
		if w.Code != tc.code {
			t.Fatalf("Unexpected status code %d: %s", w.Code, w.Body.String())
		}
	}
}
<|endoftext|>


<commit_before>package tesTaskEngineWorker

import (
	"fmt"
	proto "github.com/golang/protobuf/proto"
	"os"
	"path"
	"path/filepath"
	"strings"
	pbe "tes/ga4gh"
)

// FileMapper is responsible for mapping paths into a working directory on the
// worker's host file system.
//
// Every job needs it's own directory to work in. When a file is downloaded for
// a job, it needs to be stored in the job's working directory. Similar for job
// outputs, uploads, stdin/out/err, etc. FileMapper helps the worker engine
// manage all these paths.
type FileMapper struct {
	Volumes []Volume
	Inputs  []*pbe.TaskParameter
	Outputs []*pbe.TaskParameter
	dir     string
}

// Volume represents a volume mounted into a docker container.
// This includes a HostPath, the path on the host file system,
// and a ContainerPath, the path on the container file system,
// and the mode ("rw" = read-only, "ro" = read-write).
type Volume struct {
	// The path in tes worker.
	HostPath string
	// The path in Docker.
	ContainerPath string
	Mode          string
}

// NewFileMapper returns a new FileMapper configured to map files for a job.
//
// The following example will return a FileMapper that maps into the
// "/path/to/workdir/123/" directory on the host file system.
//     NewJobFileMapper("123", "/path/to/workdir")
func NewJobFileMapper(jobID string, baseDir string) *FileMapper {
	dir := path.Join(baseDir, jobID)
	// TODO error handling
	dir, _ = filepath.Abs(dir)
	return &FileMapper{dir: dir}
}

// AddVolume adds a mapped volume to the mapper. A corresponding Volume record
// is added to mapper.Volumes.
//
// Currently, volumes are hard-coded to "rw" (read-write).
//
// If the volume paths are invalid or can't be mapped, an error is returned.
func (mapper *FileMapper) AddVolume(source string, mountPoint string) error {
	if source != "" {
		return fmt.Errorf("Could not create a volume: 'source' is not supported for %s", source)
	}
	if mountPoint == "" {
		return fmt.Errorf("Could not create a volume: 'mountPoint' is required for %s", mountPoint)
	}

	hostPath, err := mapper.HostPath(mountPoint)
	if err != nil {
		return err
	}

	v := Volume{
		HostPath:      hostPath,
		ContainerPath: mountPoint,
		// TODO should be read only?
		Mode: "rw",
	}

  // Ensure that the volume directory exists on the host
	perr := ensurePath(hostPath)
	if perr != nil { return err }

	mapper.Volumes = append(mapper.Volumes, v)
	return nil
}

// HostPath returns a mapped path.
//
// The path is concatenated to the mapper's base dir.
// e.g. If the mapper is configured with a base dir of "/tmp/mapped_files", then
// mapper.HostPath("/home/ubuntu/myfile") will return "/tmp/mapped_files/home/ubuntu/myfile".
//
// The mapped path is required to be a subpath of the mapper's base directory.
// e.g. mapper.HostPath("../../foo") should fail with an error.
func (mapper *FileMapper) HostPath(src string) (string, error) {
	p := path.Join(mapper.dir, src)
	p = path.Clean(p)
	if !mapper.IsSubpath(p, mapper.dir) {
		return "", fmt.Errorf("Invalid path: %s is not a valid subpath of %s", p, mapper.dir)
	}
	return p, nil
}

// OpenHostFile opens a file on the host file system at a mapped path.
// "src" is an unmapped path. This function will handle mapping the path.
//
// This function calls os.Open
//
// If the path can't be mapped or the file can't be opened, an error is returned.
func (mapper *FileMapper) OpenHostFile(src string) (*os.File, error) {
	p, perr := mapper.HostPath(src)
	if perr != nil {
		return nil, perr
	}
	f, oerr := os.Open(p)
	if oerr != nil {
		return nil, oerr
	}
	return f, nil
}

// CreateHostFile creates a file on the host file system at a mapped path.
// "src" is an unmapped path. This function will handle mapping the path.
//
// This function calls os.Create
//
// If the path can't be mapped or the file can't be created, an error is returned.
func (mapper *FileMapper) CreateHostFile(src string) (*os.File, error) {
	p, perr := mapper.HostPath(src)
	if perr != nil {
		return nil, perr
	}
	err := ensurePath(p)
	if err != nil {
		return nil, err
	}
	f, oerr := os.Create(p)
	if oerr != nil {
		return nil, oerr
	}
	return f, nil
}

// AppInput adds an input to the mapped files for the given TaskParameter.
// A copy of the TaskParameter will be added to mapper.Inputs, with the
// "Path" field updated to the mapped host path.
//
// If the path can't be mapped, or the path is not in an existing volume,
// an error is returned.
func (mapper *FileMapper) AddInput(input *pbe.TaskParameter) error {
	p, err := mapper.HostPath(input.Path)
	if err != nil {
		return err
	}
	// Require that the path be in a defined volume
	if !mapper.IsInVolume(p) {
		return fmt.Errorf("Input path is required to be in a volume: %s", input.Path)
	}

	perr := ensurePath(p)
	if perr != nil { return err }

	// Create a TaskParameter for the input with a path mapped to the host
	hostIn := proto.Clone(input).(*pbe.TaskParameter)
	hostIn.Path = p
	mapper.Inputs = append(mapper.Inputs, hostIn)
	return nil
}

// AddOutput adds an output to the mapped files for the given TaskParameter.
// A copy of the TaskParameter will be added to mapper.Outputs, with the
// "Path" field updated to the mapped host path.
//
// If the Create flag is set on the TaskParameter, the file will be created
// on the host file system.
//
// If the path can't be mapped, or the path is not in an existing volume,
// an error is returned.
func (mapper *FileMapper) AddOutput(output *pbe.TaskParameter) error {
	p, err := mapper.HostPath(output.Path)
	if err != nil {
		return err
	}
	// Require that the path be in a defined volume
	if !mapper.IsInVolume(p) {
		return fmt.Errorf("Output path is required to be in a volume: %s", output.Path)
	}
	// Create the file if needed, as per the TES spec
	if output.Create {
		err := ensureFile(p, output.Class)
		if err != nil {
			return err
		}
	}
	// Create a TaskParameter for the out with a path mapped to the host
	hostOut := proto.Clone(output).(*pbe.TaskParameter)
	hostOut.Path = p
	mapper.Outputs = append(mapper.Outputs, hostOut)
	return nil
}

// IsSubpath returns true if the given path "p" is a subpath of "base".
func (mapper *FileMapper) IsSubpath(p string, base string) bool {
	return strings.HasPrefix(p, base)
}

// InInVolume checks whether a given path is in a mapped volume.
func (mapper *FileMapper) IsInVolume(p string) bool {
	for _, vol := range mapper.Volumes {
		if mapper.IsSubpath(p, vol.HostPath) {
			return true
		}
	}
	return false
}
<commit_msg>Return the correct err obj<commit_after>package tesTaskEngineWorker

import (
	"fmt"
	proto "github.com/golang/protobuf/proto"
	"os"
	"path"
	"path/filepath"
	"strings"
	pbe "tes/ga4gh"
)

// FileMapper is responsible for mapping paths into a working directory on the
// worker's host file system.
//
// Every job needs it's own directory to work in. When a file is downloaded for
// a job, it needs to be stored in the job's working directory. Similar for job
// outputs, uploads, stdin/out/err, etc. FileMapper helps the worker engine
// manage all these paths.
type FileMapper struct {
	Volumes []Volume
	Inputs  []*pbe.TaskParameter
	Outputs []*pbe.TaskParameter
	dir     string
}

// Volume represents a volume mounted into a docker container.
// This includes a HostPath, the path on the host file system,
// and a ContainerPath, the path on the container file system,
// and the mode ("rw" = read-only, "ro" = read-write).
type Volume struct {
	// The path in tes worker.
	HostPath string
	// The path in Docker.
	ContainerPath string
	Mode          string
}

// NewFileMapper returns a new FileMapper configured to map files for a job.
//
// The following example will return a FileMapper that maps into the
// "/path/to/workdir/123/" directory on the host file system.
//     NewJobFileMapper("123", "/path/to/workdir")
func NewJobFileMapper(jobID string, baseDir string) *FileMapper {
	dir := path.Join(baseDir, jobID)
	// TODO error handling
	dir, _ = filepath.Abs(dir)
	return &FileMapper{dir: dir}
}

// AddVolume adds a mapped volume to the mapper. A corresponding Volume record
// is added to mapper.Volumes.
//
// Currently, volumes are hard-coded to "rw" (read-write).
//
// If the volume paths are invalid or can't be mapped, an error is returned.
func (mapper *FileMapper) AddVolume(source string, mountPoint string) error {
	if source != "" {
		return fmt.Errorf("Could not create a volume: 'source' is not supported for %s", source)
	}
	if mountPoint == "" {
		return fmt.Errorf("Could not create a volume: 'mountPoint' is required for %s", mountPoint)
	}

	hostPath, err := mapper.HostPath(mountPoint)
	if err != nil {
		return err
	}

	v := Volume{
		HostPath:      hostPath,
		ContainerPath: mountPoint,
		// TODO should be read only?
		Mode: "rw",
	}

  // Ensure that the volume directory exists on the host
	perr := ensurePath(hostPath)
	if perr != nil { return perr }

	mapper.Volumes = append(mapper.Volumes, v)
	return nil
}

// HostPath returns a mapped path.
//
// The path is concatenated to the mapper's base dir.
// e.g. If the mapper is configured with a base dir of "/tmp/mapped_files", then
// mapper.HostPath("/home/ubuntu/myfile") will return "/tmp/mapped_files/home/ubuntu/myfile".
//
// The mapped path is required to be a subpath of the mapper's base directory.
// e.g. mapper.HostPath("../../foo") should fail with an error.
func (mapper *FileMapper) HostPath(src string) (string, error) {
	p := path.Join(mapper.dir, src)
	p = path.Clean(p)
	if !mapper.IsSubpath(p, mapper.dir) {
		return "", fmt.Errorf("Invalid path: %s is not a valid subpath of %s", p, mapper.dir)
	}
	return p, nil
}

// OpenHostFile opens a file on the host file system at a mapped path.
// "src" is an unmapped path. This function will handle mapping the path.
//
// This function calls os.Open
//
// If the path can't be mapped or the file can't be opened, an error is returned.
func (mapper *FileMapper) OpenHostFile(src string) (*os.File, error) {
	p, perr := mapper.HostPath(src)
	if perr != nil {
		return nil, perr
	}
	f, oerr := os.Open(p)
	if oerr != nil {
		return nil, oerr
	}
	return f, nil
}

// CreateHostFile creates a file on the host file system at a mapped path.
// "src" is an unmapped path. This function will handle mapping the path.
//
// This function calls os.Create
//
// If the path can't be mapped or the file can't be created, an error is returned.
func (mapper *FileMapper) CreateHostFile(src string) (*os.File, error) {
	p, perr := mapper.HostPath(src)
	if perr != nil {
		return nil, perr
	}
	err := ensurePath(p)
	if err != nil {
		return nil, err
	}
	f, oerr := os.Create(p)
	if oerr != nil {
		return nil, oerr
	}
	return f, nil
}

// AppInput adds an input to the mapped files for the given TaskParameter.
// A copy of the TaskParameter will be added to mapper.Inputs, with the
// "Path" field updated to the mapped host path.
//
// If the path can't be mapped, or the path is not in an existing volume,
// an error is returned.
func (mapper *FileMapper) AddInput(input *pbe.TaskParameter) error {
	p, err := mapper.HostPath(input.Path)
	if err != nil {
		return err
	}
	// Require that the path be in a defined volume
	if !mapper.IsInVolume(p) {
		return fmt.Errorf("Input path is required to be in a volume: %s", input.Path)
	}

	perr := ensurePath(p)
	if perr != nil { return perr }

	// Create a TaskParameter for the input with a path mapped to the host
	hostIn := proto.Clone(input).(*pbe.TaskParameter)
	hostIn.Path = p
	mapper.Inputs = append(mapper.Inputs, hostIn)
	return nil
}

// AddOutput adds an output to the mapped files for the given TaskParameter.
// A copy of the TaskParameter will be added to mapper.Outputs, with the
// "Path" field updated to the mapped host path.
//
// If the Create flag is set on the TaskParameter, the file will be created
// on the host file system.
//
// If the path can't be mapped, or the path is not in an existing volume,
// an error is returned.
func (mapper *FileMapper) AddOutput(output *pbe.TaskParameter) error {
	p, err := mapper.HostPath(output.Path)
	if err != nil {
		return err
	}
	// Require that the path be in a defined volume
	if !mapper.IsInVolume(p) {
		return fmt.Errorf("Output path is required to be in a volume: %s", output.Path)
	}
	// Create the file if needed, as per the TES spec
	if output.Create {
		err := ensureFile(p, output.Class)
		if err != nil {
			return err
		}
	}
	// Create a TaskParameter for the out with a path mapped to the host
	hostOut := proto.Clone(output).(*pbe.TaskParameter)
	hostOut.Path = p
	mapper.Outputs = append(mapper.Outputs, hostOut)
	return nil
}

// IsSubpath returns true if the given path "p" is a subpath of "base".
func (mapper *FileMapper) IsSubpath(p string, base string) bool {
	return strings.HasPrefix(p, base)
}

// InInVolume checks whether a given path is in a mapped volume.
func (mapper *FileMapper) IsInVolume(p string) bool {
	for _, vol := range mapper.Volumes {
		if mapper.IsSubpath(p, vol.HostPath) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>//
// Package shorthand provides shorthand definition and expansion.
//
// shorthand.go - A simple definition and expansion notation to use
// as shorthand when a template language is too much.
//
// @author R. S. Doiel, <rsdoiel@gmail.com>
// copyright (c) 2015 all rights reserved.
// Released under the BSD 2-Clause license
// See: http://opensource.org/licenses/BSD-2-Clause
//
package shorthand

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/russross/blackfriday"
)

// The version nummber of library and utility
const Version = "v0.0.4"

//
// An Op is built from a multi character symbol
// Each element in the symbol has meaning
// " :" is the start of a glyph band the end ": " is a colon and trailing space
// = the source value is next, this is basic assignment of a string value to a symbol
// < is input from a file
// { expand (resolved label values)
// } assign a statement (i.e. label, op, value)
// ! input from a shell expression
// [ is a markdown expansion
// > is to output an to a file
// @ operate on whole symbol table
//

// SourceMap holds the source and value of an assignment
type SourceMap struct {
	Label    string // Label is the symbol to be replace based on Op and Source
	Op       string // Op is the type of assignment being made (if is an empty string if not an assignment)
	Source   string // Source is argument to the right of Op
	Expanded string // Expanded is the value calculated based on Label, Op and Source
	LineNo   int
}

// SymbolTable holds the exressions, values and other errata of parsing assignments making expansions
type SymbolTable struct {
	entries []SourceMap
	labels  map[string]int
}

// OperatorMap is a map of operator testings and their related functions
type OperatorMap map[string]func(SourceMap) (SourceMap, error)

// Assignment Ops
const (
	AssignString             string = " :=: "
	AssignInclude            string = " :=<: "
	IncludeAssignments       string = " :}<: "
	AssignExpansion          string = " :{: "
	AssignExpandExpansion    string = " :{{: "
	IncludeExpansion         string = " :{<: "
	AssignShell              string = " :!: "
	AssignExpandShell        string = " :{!: "
	AssignMarkdown           string = " :[: "
	AssignExpandMarkdown     string = " :{[: "
	IncludeMarkdown          string = " :[<: "
	IncludeExpandMarkdown    string = " :{[<: "
	OutputAssignedExpansion  string = " :>: "
	OutputAssignedExpansions string = " :@>: "
	OutputAssignment         string = " :}>: "
	OutputAssignments        string = " :@}>: "
)

var ops = []string{
	AssignString,
	AssignInclude,
	IncludeAssignments,
	AssignExpansion,
	AssignExpandExpansion,
	IncludeExpansion,
	AssignShell,
	AssignExpandShell,
	AssignMarkdown,
	AssignExpandMarkdown,
	IncludeMarkdown,
	IncludeExpandMarkdown,
	OutputAssignedExpansion,
	OutputAssignedExpansions,
	OutputAssignment,
	OutputAssignments,
}

func warning(msgs string) {
	fmt.Fprintln(os.Stderr, msgs)
}

// IsAssignment checks to see if a string contains an assignment (e.g. has a ' := ' in the string.)
func IsAssignment(text string) bool {
	for _, op := range ops {
		if strings.Index(text, op) != -1 {
			return true
		}
	}
	return false
}

// HasAssignment checks to see if a shortcut has already been assigned.
func HasAssignment(table *SymbolTable, label string) bool {
	_, ok := table.labels[label]
	if ok == true {
		return true
	}
	return false
}

// Parse returns SourceMap object and bool with true if assignment or false if not
func Parse(s string, lineNo int) (SourceMap, bool) {
	for _, op := range ops {
		if strings.Index(s, op) != -1 {
			parts := strings.SplitN(strings.TrimSpace(s), op, 2)
			return SourceMap{Label: parts[0], Op: op, Source: parts[1], LineNo: lineNo, Expanded: ""}, true
		}
	}
	return SourceMap{Label: "", Op: "", Source: "", LineNo: lineNo, Expanded: s}, false
}

// Expand takes some text and expands all labels to their values
func Expand(table *SymbolTable, text string) string {
	// labels hash should also point at the last known state of
	// the label
	result := text
	for _, i := range table.labels {
		sm := table.entries[i]
		if strings.Contains(text, sm.Label) {
			tmp := strings.Replace(result, sm.Label, sm.Expanded, -1)
			result = tmp
		}
	}
	return result
}

// Assign stores a shorthand and its expansion or writes and assignment or assignments
func Assign(table *SymbolTable, s string, lineNo int) bool {
	sm, ok := Parse(s, lineNo)
	if ok == false {
		return ok
	}

	// These functions do not change the symbol table
	if sm.Op == OutputAssignedExpansion {
		return WriteAssignment(sm.Source, sm.Label, table, false)
	}

	if sm.Op == OutputAssignedExpansions {
		return WriteAssignments(sm.Source, table, false)
	}

	if sm.Op == OutputAssignment {
		return WriteAssignment(sm.Source, sm.Label, table, true)
	}

	if sm.Op == OutputAssignments {
		return WriteAssignments(sm.Source, table, true)
	}

	// These functions change the symbol table
	if sm.Op == AssignString {
		sm.Expanded = sm.Source
	} else if sm.Op == AssignInclude {
		buf, err := ioutil.ReadFile(sm.Source)
		if err != nil {
			warning(fmt.Sprintf("Cannot read %s: %s\n", sm.Source, err))
			return false
		}
		sm.Expanded = string(buf)
	} else if sm.Op == IncludeExpansion {
		buf, err := ioutil.ReadFile(sm.Source)
		if err != nil {
			warning(fmt.Sprintf("Cannot read %s: %s\n", sm.Source, err))
			return false
		}
		sm.Expanded = Expand(table, string(buf))
	} else if sm.Op == AssignShell {
		buf, err := exec.Command("bash", "-c", sm.Source).Output()
		if err != nil {
			warning(fmt.Sprintf("Shell command returned error: %s\n", err))
			return false
		}
		sm.Expanded = string(buf)
	} else if sm.Op == AssignExpandShell {
		buf, err := exec.Command("bash", "-c", Expand(table, sm.Source)).Output()
		if err != nil {
			warning(fmt.Sprintf("Shell command returned error: %s\n", err))
			return false
		}
		sm.Expanded = string(buf)
	} else if sm.Op == AssignExpansion {
		sm.Expanded = Expand(table, sm.Source)
	} else if sm.Op == AssignExpandExpansion {
		tmp := Expand(table, sm.Source)
		sm.Expanded = Expand(table, tmp)
	} else if sm.Op == IncludeAssignments {
		err := ReadAssignments(sm.Source, table)
		if err != nil {
			warning(fmt.Sprintf("Error processing %s: %s\n", sm.Source, err))
			return false
		}
	} else if sm.Op == AssignMarkdown {
		sm.Expanded = strings.TrimSpace(string(blackfriday.MarkdownCommon([]byte(sm.Source))))
	} else if sm.Op == AssignExpandMarkdown {
		sm.Expanded = strings.TrimSpace(string(blackfriday.MarkdownCommon([]byte(Expand(table, sm.Source)))))
	} else if sm.Op == IncludeMarkdown {
		sm.Expanded = ReadMarkdown(sm.Source)
	} else if sm.Op == IncludeExpandMarkdown {
		buf, err := ioutil.ReadFile(sm.Source)
		if err != nil {
			warning(fmt.Sprintf("Cannot read %s: %s\n", sm.Source, err))
			return false
		}
		sm.Expanded = strings.TrimSpace(string(blackfriday.MarkdownCommon([]byte(Expand(table, string(buf))))))
	}
	if sm.Label != "" {
		table.entries = append(table.entries, sm)
		if table.labels == nil {
			table.labels = make(map[string]int)
		}
		table.labels[sm.Label] = len(table.entries) - 1
	}
	return ok
}

// WriteAssignment writes a single assignment statement to filename
func WriteAssignment(fname string, label string, table *SymbolTable, writeSourceCode bool) bool {
	fp, err := os.Create(fname)
	if err != nil {
		warning(fmt.Sprintf("%s", err))
		return false
	}
	defer fp.Close()

	i, ok := table.labels[label]
	if ok == true {
		sm := table.entries[i]
		if writeSourceCode == true {
			fmt.Fprintf(fp, "%s%s%s", sm.Label, sm.Op, sm.Source)
		} else {
			fmt.Fprintf(fp, "%s", sm.Expanded)
		}
		return true
	}
	return false
}

// WriteAssignments write all assignment statements to filename
func WriteAssignments(fname string, table *SymbolTable, writeSourceCode bool) bool {
	fp, err := os.Create(fname)
	if err != nil {
		warning(fmt.Sprintf("Cannot write to %s, error: %s\n", fname, err))
		return false
	}
	defer fp.Close()
	for _, sm := range table.entries {
		if writeSourceCode == true {
			fmt.Fprintf(fp, "%s%s%s\n", sm.Label, sm.Op, sm.Source)
		} else {
			fmt.Fprintf(fp, "%s\n", sm.Expanded)
		}
	}
	return true
}

// ReadAssignments read in all of the lines of fname and
// add any assignment statements found to Abbreviations
// and expand any assignments and return as a string.
func ReadAssignments(fname string, table *SymbolTable) error {
	buf, err := ioutil.ReadFile(fname)
	if err != nil {
		return fmt.Errorf("Cannot read %s: %s\n", fname, err)
	}
	for i, text := range strings.Split(string(buf), "\n") {
		if IsAssignment(text) {
			ok := Assign(table, text, i)
			if ok == false {
				return fmt.Errorf("Error at line %d in %s\n", i+1, fname)
			}
		}
	}
	return nil
}

// ReadMarkdown reads in a file and processes it with Blackfriday MarkdownCommon
func ReadMarkdown(fname string) string {
	buf, err := ioutil.ReadFile(fname)
	if err != nil {
		warning(fmt.Sprintf("Cannot read %s: %s\n", fname, err))
		return ""
	}
	return string(blackfriday.MarkdownCommon(buf))
}

//
// An Op is built from a multi character symbol
// Each element in the symbol has meaning
// " :" is the start of a glyph band the end ": " is a colon and trailing space
// = the source value is next, this is basic assignment of a string value to a symbol
// < is input from a file
// { expand (resolved label values)
// } assign a statement (i.e. label, op, value)
// ! input from a shell expression
// [ is a markdown expansion
// > is to output an to a file
// @ operate on whole symbol table
//

type VirtualMachine struct {
	Symbols   *SymbolTable
	Operators OperatorMap
}

func New() *VirtualMachine {
	vm := new(VirtualMachine)
	vm.Symbols = new(SymbolTable)
	vm.Operators = make(OperatorMap)
	// Now register the built-in operators
	vm.RegisterOp(" :=: ", AssignStringCallback)
	vm.RegisterOp(" :=<: ", AssignIncludeCallback)
	vm.RegisterOp(" :}<: ", IncludeAssignmentsCallback)
	vm.RegisterOp(" :{: ", AssignExpansionCallback)
	vm.RegisterOp(" :{{: ", AssignExpandExpansionCallback)
	vm.RegisterOp(" :{<: ", IncludeExpansionCallback)
	vm.RegisterOp(" :!: ", AssignShellCallback)
	vm.RegisterOp(" :{!: ", AssignExpandShellCallback)
	vm.RegisterOp(" :[: ", AssignMarkdownCallback)
	vm.RegisterOp(" :{[: ", AssignExpandMarkdownCallback)
	vm.RegisterOp(" :[<: ", IncludeMarkdownCallback)
	vm.RegisterOp(" :{[<: ", IncludeExpandMarkdownCallback)
	vm.RegisterOp(" :>: ", OutputAssignedExpansionCallback)
	vm.RegisterOp(" :@>: ", OutputAssignedExpansionsCallback)
	vm.RegisterOp(" :}>: ", OutputAssignmentCallback)
	vm.RegisterOp(" :@}>: ", OutputAssignmentsCallback)
	vm.RegisterOp(" :exit: ", ExitShorthand)
	vm.RegisterOp(" :quit: ", ExitShorthand)
	return vm
}

// GetOps returns a list of registered Shorthand operators.
func (vm *VirtualMachine) GetOps() []string {
	var ops []string

	for op := range vm.Operators {
		ops = append(ops, op)
	}
	return ops
}

// RegisterOp associate a operations with a function
func (vm *VirtualMachine) RegisterOp(op string, callback func(SourceMap) (SourceMap, error)) error {
	_, ok := vm.Operators[op]
	if ok == true {
		return fmt.Errorf("Cannot redefine function %s\n", op)
	}
	vm.Operators[op] = callback
	return nil
}

// Eval stores a shorthand assignment or expands and writes the content to stdout
func (vm *VirtualMachine) Eval(s string, lineNo int) (string, error) {
	sm, isAssignment := Parse(s, lineNo)
	if isAssignment == false {
		return fmt.Sprintf("%s", Expand(vm.Symbols, s)), nil
	}

	callback, ok := vm.Operators[sm.Op]
	if ok == false {
		return "", fmt.Errorf("ERROR (%d): %s is not an expansion or valid assignment.\n", lineNo, s)
	}

	newSM, err := callback(sm)
	if err != nil {
		return "", fmt.Errorf("ERROR (%d): %s failed %s\n", lineNo, s, err)
	}

	if newSM.Label != "" {
		vm.Symbols.entries = append(vm.Symbols.entries, newSM)
		if vm.Symbols.labels == nil {
			vm.Symbols.labels = make(map[string]int)
		}
		vm.Symbols.labels[newSM.Label] = len(vm.Symbols.entries) - 1
	}
	return "", nil
}

// Run takes a reader (e.g. os.Stdin), and two writers (e.g. os.Stdout and os.Stderr)
// It reads until EOF, :exit:, or :quit: operation is encountered
// returns the number of lines processed.
func (vm *VirtualMachine) Run(in *bufio.Reader) int {
	lineNo := 0
	for {
		src, rErr := in.ReadString('\n')
		if rErr != nil {
			break
		}
		lineNo += 1
		if strings.Contains(src, ":exit:") || strings.Contains(src, ":quit:") {
			break
		}
		out, err := vm.Eval(src, lineNo)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR (%d): %s\n", lineNo, err)
		}
		if out != "" {
			fmt.Fprintf(os.Stdout, out)
		}
	}
	return lineNo
}
<commit_msg>Sketch code refactor<commit_after>//
// Package shorthand provides shorthand definition and expansion.
//
// shorthand.go - A simple definition and expansion notation to use
// as shorthand when a template language is too much.
//
// @author R. S. Doiel, <rsdoiel@gmail.com>
// copyright (c) 2015 all rights reserved.
// Released under the BSD 2-Clause license
// See: http://opensource.org/licenses/BSD-2-Clause
//
package shorthand

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/russross/blackfriday"
)

// The version nummber of library and utility
const Version = "v0.0.4-next"

//
// An Op is built from a multi character symbol
// Each element in the symbol has meaning
// " :" is the start of a glyph band the end ": " is a colon and trailing space
// = the source value is next, this is basic assignment of a string value to a symbol
// < is input from a file
// { expand (resolved label values)
// } assign a statement (i.e. label, op, value)
// ! input from a shell expression
// [ is a markdown expansion
// > is to output an to a file
// @ operate on whole symbol table
//

// SourceMap holds the source and value of an assignment
type SourceMap struct {
	Label    string // Label is the symbol to be replace based on Op and Source
	Op       string // Op is the type of assignment being made (if is an empty string if not an assignment)
	Source   string // Source is argument to the right of Op
	Expanded string // Expanded is the value calculated based on Label, Op and Source
	LineNo   int
}

// SymbolTable holds the exressions, values and other errata of parsing assignments making expansions
type SymbolTable struct {
	entries []SourceMap
	labels  map[string]int
}

// OperatorMap is a map of operator testings and their related functions
type OperatorMap map[string]func(SourceMap) (SourceMap, error)

// Assignment Ops
const (
	AssignString             string = " :=: "
	AssignInclude            string = " :=<: "
	IncludeAssignments       string = " :}<: "
	AssignExpansion          string = " :{: "
	AssignExpandExpansion    string = " :{{: "
	IncludeExpansion         string = " :{<: "
	AssignShell              string = " :!: "
	AssignExpandShell        string = " :{!: "
	AssignMarkdown           string = " :[: "
	AssignExpandMarkdown     string = " :{[: "
	IncludeMarkdown          string = " :[<: "
	IncludeExpandMarkdown    string = " :{[<: "
	OutputAssignedExpansion  string = " :>: "
	OutputAssignedExpansions string = " :@>: "
	OutputAssignment         string = " :}>: "
	OutputAssignments        string = " :@}>: "
)

var ops = []string{
	AssignString,
	AssignInclude,
	IncludeAssignments,
	AssignExpansion,
	AssignExpandExpansion,
	IncludeExpansion,
	AssignShell,
	AssignExpandShell,
	AssignMarkdown,
	AssignExpandMarkdown,
	IncludeMarkdown,
	IncludeExpandMarkdown,
	OutputAssignedExpansion,
	OutputAssignedExpansions,
	OutputAssignment,
	OutputAssignments,
}

func warning(msgs string) {
	fmt.Fprintln(os.Stderr, msgs)
}

// IsAssignment checks to see if a string contains an assignment (e.g. has a ' := ' in the string.)
func IsAssignment(text string) bool {
	for _, op := range ops {
		if strings.Index(text, op) != -1 {
			return true
		}
	}
	return false
}

// HasAssignment checks to see if a shortcut has already been assigned.
func HasAssignment(table *SymbolTable, label string) bool {
	_, ok := table.labels[label]
	if ok == true {
		return true
	}
	return false
}

// Parse returns SourceMap object and bool with true if assignment or false if not
func Parse(s string, lineNo int) (SourceMap, bool) {
	for _, op := range ops {
		if strings.Index(s, op) != -1 {
			parts := strings.SplitN(strings.TrimSpace(s), op, 2)
			return SourceMap{Label: parts[0], Op: op, Source: parts[1], LineNo: lineNo, Expanded: ""}, true
		}
	}
	return SourceMap{Label: "", Op: "", Source: "", LineNo: lineNo, Expanded: s}, false
}

// Expand takes some text and expands all labels to their values
func Expand(table *SymbolTable, text string) string {
	// labels hash should also point at the last known state of
	// the label
	result := text
	for _, i := range table.labels {
		sm := table.entries[i]
		if strings.Contains(text, sm.Label) {
			tmp := strings.Replace(result, sm.Label, sm.Expanded, -1)
			result = tmp
		}
	}
	return result
}

// Assign stores a shorthand and its expansion or writes and assignment or assignments
func Assign(table *SymbolTable, s string, lineNo int) bool {
	sm, ok := Parse(s, lineNo)
	if ok == false {
		return ok
	}

	// These functions do not change the symbol table
	if sm.Op == OutputAssignedExpansion {
		return WriteAssignment(sm.Source, sm.Label, table, false)
	}

	if sm.Op == OutputAssignedExpansions {
		return WriteAssignments(sm.Source, table, false)
	}

	if sm.Op == OutputAssignment {
		return WriteAssignment(sm.Source, sm.Label, table, true)
	}

	if sm.Op == OutputAssignments {
		return WriteAssignments(sm.Source, table, true)
	}

	// These functions change the symbol table
	if sm.Op == AssignString {
		sm.Expanded = sm.Source
	} else if sm.Op == AssignInclude {
		buf, err := ioutil.ReadFile(sm.Source)
		if err != nil {
			warning(fmt.Sprintf("Cannot read %s: %s\n", sm.Source, err))
			return false
		}
		sm.Expanded = string(buf)
	} else if sm.Op == IncludeExpansion {
		buf, err := ioutil.ReadFile(sm.Source)
		if err != nil {
			warning(fmt.Sprintf("Cannot read %s: %s\n", sm.Source, err))
			return false
		}
		sm.Expanded = Expand(table, string(buf))
	} else if sm.Op == AssignShell {
		buf, err := exec.Command("bash", "-c", sm.Source).Output()
		if err != nil {
			warning(fmt.Sprintf("Shell command returned error: %s\n", err))
			return false
		}
		sm.Expanded = string(buf)
	} else if sm.Op == AssignExpandShell {
		buf, err := exec.Command("bash", "-c", Expand(table, sm.Source)).Output()
		if err != nil {
			warning(fmt.Sprintf("Shell command returned error: %s\n", err))
			return false
		}
		sm.Expanded = string(buf)
	} else if sm.Op == AssignExpansion {
		sm.Expanded = Expand(table, sm.Source)
	} else if sm.Op == AssignExpandExpansion {
		tmp := Expand(table, sm.Source)
		sm.Expanded = Expand(table, tmp)
	} else if sm.Op == IncludeAssignments {
		err := ReadAssignments(sm.Source, table)
		if err != nil {
			warning(fmt.Sprintf("Error processing %s: %s\n", sm.Source, err))
			return false
		}
	} else if sm.Op == AssignMarkdown {
		sm.Expanded = strings.TrimSpace(string(blackfriday.MarkdownCommon([]byte(sm.Source))))
	} else if sm.Op == AssignExpandMarkdown {
		sm.Expanded = strings.TrimSpace(string(blackfriday.MarkdownCommon([]byte(Expand(table, sm.Source)))))
	} else if sm.Op == IncludeMarkdown {
		sm.Expanded = ReadMarkdown(sm.Source)
	} else if sm.Op == IncludeExpandMarkdown {
		buf, err := ioutil.ReadFile(sm.Source)
		if err != nil {
			warning(fmt.Sprintf("Cannot read %s: %s\n", sm.Source, err))
			return false
		}
		sm.Expanded = strings.TrimSpace(string(blackfriday.MarkdownCommon([]byte(Expand(table, string(buf))))))
	}
	if sm.Label != "" {
		table.entries = append(table.entries, sm)
		if table.labels == nil {
			table.labels = make(map[string]int)
		}
		table.labels[sm.Label] = len(table.entries) - 1
	}
	return ok
}

// WriteAssignment writes a single assignment statement to filename
func WriteAssignment(fname string, label string, table *SymbolTable, writeSourceCode bool) bool {
	fp, err := os.Create(fname)
	if err != nil {
		warning(fmt.Sprintf("%s", err))
		return false
	}
	defer fp.Close()

	i, ok := table.labels[label]
	if ok == true {
		sm := table.entries[i]
		if writeSourceCode == true {
			fmt.Fprintf(fp, "%s%s%s", sm.Label, sm.Op, sm.Source)
		} else {
			fmt.Fprintf(fp, "%s", sm.Expanded)
		}
		return true
	}
	return false
}

// WriteAssignments write all assignment statements to filename
func WriteAssignments(fname string, table *SymbolTable, writeSourceCode bool) bool {
	fp, err := os.Create(fname)
	if err != nil {
		warning(fmt.Sprintf("Cannot write to %s, error: %s\n", fname, err))
		return false
	}
	defer fp.Close()
	for _, sm := range table.entries {
		if writeSourceCode == true {
			fmt.Fprintf(fp, "%s%s%s\n", sm.Label, sm.Op, sm.Source)
		} else {
			fmt.Fprintf(fp, "%s\n", sm.Expanded)
		}
	}
	return true
}

// ReadAssignments read in all of the lines of fname and
// add any assignment statements found to Abbreviations
// and expand any assignments and return as a string.
func ReadAssignments(fname string, table *SymbolTable) error {
	buf, err := ioutil.ReadFile(fname)
	if err != nil {
		return fmt.Errorf("Cannot read %s: %s\n", fname, err)
	}
	for i, text := range strings.Split(string(buf), "\n") {
		if IsAssignment(text) {
			ok := Assign(table, text, i)
			if ok == false {
				return fmt.Errorf("Error at line %d in %s\n", i+1, fname)
			}
		}
	}
	return nil
}

// ReadMarkdown reads in a file and processes it with Blackfriday MarkdownCommon
func ReadMarkdown(fname string) string {
	buf, err := ioutil.ReadFile(fname)
	if err != nil {
		warning(fmt.Sprintf("Cannot read %s: %s\n", fname, err))
		return ""
	}
	return string(blackfriday.MarkdownCommon(buf))
}

//
// An Op is built from a multi character symbol
// Each element in the symbol has meaning
// " :" is the start of a glyph band the end ": " is a colon and trailing space
// = the source value is next, this is basic assignment of a string value to a symbol
// < is input from a file
// { expand (resolved label values)
// } assign a statement (i.e. label, op, value)
// ! input from a shell expression
// [ is a markdown expansion
// > is to output an to a file
// @ operate on whole symbol table
//

type VirtualMachine struct {
	Symbols   *SymbolTable
	Operators OperatorMap
}

func New() *VirtualMachine {
	vm := new(VirtualMachine)
	vm.Symbols = new(SymbolTable)
	vm.Operators = make(OperatorMap)
	// Now register the built-in operators
	vm.RegisterOp(" :=: ", AssignStringCallback)
	vm.RegisterOp(" :=<: ", AssignIncludeCallback)
	vm.RegisterOp(" :}<: ", IncludeAssignmentsCallback)
	vm.RegisterOp(" :{: ", AssignExpansionCallback)
	vm.RegisterOp(" :{{: ", AssignExpandExpansionCallback)
	vm.RegisterOp(" :{<: ", IncludeExpansionCallback)
	vm.RegisterOp(" :!: ", AssignShellCallback)
	vm.RegisterOp(" :{!: ", AssignExpandShellCallback)
	vm.RegisterOp(" :[: ", AssignMarkdownCallback)
	vm.RegisterOp(" :{[: ", AssignExpandMarkdownCallback)
	vm.RegisterOp(" :[<: ", IncludeMarkdownCallback)
	vm.RegisterOp(" :{[<: ", IncludeExpandMarkdownCallback)
	vm.RegisterOp(" :>: ", OutputAssignedExpansionCallback)
	vm.RegisterOp(" :@>: ", OutputAssignedExpansionsCallback)
	vm.RegisterOp(" :}>: ", OutputAssignmentCallback)
	vm.RegisterOp(" :@}>: ", OutputAssignmentsCallback)
	vm.RegisterOp(" :exit: ", ExitShorthand)
	vm.RegisterOp(" :quit: ", ExitShorthand)
	return vm
}

// GetOps returns a list of registered Shorthand operators.
func (vm *VirtualMachine) GetOps() []string {
	var ops []string

	for op := range vm.Operators {
		ops = append(ops, op)
	}
	return ops
}

// RegisterOp associate a operations with a function
func (vm *VirtualMachine) RegisterOp(op string, callback func(SourceMap) (SourceMap, error)) error {
	_, ok := vm.Operators[op]
	if ok == true {
		return fmt.Errorf("Cannot redefine function %s\n", op)
	}
	vm.Operators[op] = callback
	return nil
}

// Eval stores a shorthand assignment or expands and writes the content to stdout
func (vm *VirtualMachine) Eval(s string, lineNo int) (string, error) {
	sm, isAssignment := Parse(s, lineNo)
	if isAssignment == false {
		return fmt.Sprintf("%s", Expand(vm.Symbols, s)), nil
	}

	callback, ok := vm.Operators[sm.Op]
	if ok == false {
		return "", fmt.Errorf("ERROR (%d): %s is not an expansion or valid assignment.\n", lineNo, s)
	}

	newSM, err := callback(sm)
	if err != nil {
		return "", fmt.Errorf("ERROR (%d): %s failed %s\n", lineNo, s, err)
	}

	if newSM.Label != "" {
		vm.Symbols.entries = append(vm.Symbols.entries, newSM)
		if vm.Symbols.labels == nil {
			vm.Symbols.labels = make(map[string]int)
		}
		vm.Symbols.labels[newSM.Label] = len(vm.Symbols.entries) - 1
	}
	return "", nil
}

// Run takes a reader (e.g. os.Stdin), and two writers (e.g. os.Stdout and os.Stderr)
// It reads until EOF, :exit:, or :quit: operation is encountered
// returns the number of lines processed.
func (vm *VirtualMachine) Run(in *bufio.Reader) int {
	lineNo := 0
	for {
		src, rErr := in.ReadString('\n')
		if rErr != nil {
			break
		}
		lineNo += 1
		if strings.Contains(src, ":exit:") || strings.Contains(src, ":quit:") {
			break
		}
		out, err := vm.Eval(src, lineNo)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR (%d): %s\n", lineNo, err)
		}
		if out != "" {
			fmt.Fprintf(os.Stdout, out)
		}
	}
	return lineNo
}
<|endoftext|>


<commit_before>package sim

import (
	"fmt"
	"github.com/Queens-Hacks/Propagate/sandbox"
)

type TileType int

const (
	DirtTile TileType = iota
	AirTile
	PlantTile
)

type growthRoot struct {
	PlantId string
	Loc     Location
	node    *sandbox.Node
}

type Location struct {
	X int `json:"x"`
	Y int `json:"y"`
}

type plantInfo struct {
	PlantId string   `json:"plantId"`
	Parent  Location `json:"parent"`
	IsRoot  bool     `json:"isRoot"`
}

type Tile struct {
	T     TileType   `json:"tileType"`
	Plant *plantInfo `json:"plant"`
}

type Plant struct {
	Color  int    `json:"color"`
	Source string `json:"source"`
	Author string `json:"author"`
	refCnt int
}

type gameState struct {
	World    [][]*Tile         `json:"world"`
	Plants   map[string]*Plant `json:"plants"`
	roots    []*growthRoot
	maxPlant int
}

type tileDiff struct {
	Loc  Location `json:"loc"`
	Tile Tile     `json:"tile"`
}

type diff struct {
	TileDiffs     []tileDiff        `json:"tileDiff"`
	NewPlants     map[string]*Plant `json:"newPlants"`
	RemovedPlants []string          `json:"removedPlants"`
}

func (d *diff) isEmpty() bool {
	return len(d.TileDiffs) == 0 && len(d.NewPlants) == 0 && len(d.RemovedPlants) == 0
}

type State struct {
	state gameState
	diff  diff
}

// Records a reference to a plant, causing the plant to be kept in the structure
func (s *State) plantAddRef(plantId string) {
	s.GetPlant(plantId).refCnt++
}

// Records a reference to a plant, causing the plant to be removed from the structure
func (s *State) plantRelease(plantId string) {
	plant := s.GetPlant(plantId)
	plant.refCnt--

	// If the reference count has reached zero, remove the plant from the thing
	if plant.refCnt <= 0 {
		s.diff.RemovedPlants = append(s.diff.RemovedPlants, plantId)
		delete(s.state.Plants, plantId)
	}
}

func (s *State) Width() int {
	return len(s.state.World[0])
}

func (s *State) Height() int {
	return len(s.state.World)
}

func (s *State) GetPlant(plantId string) *Plant {
	return s.state.Plants[plantId]
}

// Adds a species to the stateAndDiff, and returns the string key for the plant
// This plant is created with a refCnt of zero, but will not be dropped until
// its reference count hits zero again.
func (s *State) AddSpecies(color int, source string, author string) string {
	p := Plant{color, source, author, 0}
	s.state.maxPlant += 1
	key := fmt.Sprintf("%d", s.state.maxPlant)
	s.diff.NewPlants[key] = &p
	s.state.Plants[key] = &p
	return key
}

func (s *State) GetTile(loc Location) *Tile {
	return s.state.World[loc.Y][loc.X]
}

// Set the tile at a location to a new tile
func (s *State) SetTile(loc Location, new Tile) {
	// Manage the addref and releases
	if new.Plant != nil {
		s.plantAddRef(new.Plant.PlantId)
	}
	old := s.GetTile(loc)
	if old.Plant != nil {
		s.plantRelease(old.Plant.PlantId)
	}

	// Actually update the tile and record the tilediffs
	*old = new
	s.diff.TileDiffs = append(s.diff.TileDiffs, tileDiff{loc, new})
}

func (s *State) AddPlant(loc Location, id string) *growthRoot {
	plant := s.GetPlant(id)

	// Create the sandbox node for the plant object
	node := sandbox.AddNode(plant.Source)

	// Create the root node for the object, and append it to the roots list
	root := growthRoot{id, loc, node}
	s.state.roots = append(s.state.roots, &root)

	isUnderground := false
	if s.GetTile(loc).T == DirtTile {
		isUnderground = true
	}

	// Set the tile at the base of the plant to a plant tile
	s.SetTile(loc, Tile{PlantTile, &plantInfo{id, loc, isUnderground}})

	// Return a reference to the root node we previously appended
	return &root
}
<commit_msg>start growing from ground<commit_after>package sim

import (
	"fmt"

	"github.com/Queens-Hacks/Propagate/sandbox"
)

type TileType int

const (
	DirtTile TileType = iota
	AirTile
	PlantTile
)

type growthRoot struct {
	PlantId string
	Loc     Location
	node    *sandbox.Node
}

type Location struct {
	X int `json:"x"`
	Y int `json:"y"`
}

type plantInfo struct {
	PlantId string   `json:"plantId"`
	Parent  Location `json:"parent"`
	IsRoot  bool     `json:"isRoot"`
}

type Tile struct {
	T     TileType   `json:"tileType"`
	Plant *plantInfo `json:"plant"`
}

type Plant struct {
	Color  int    `json:"color"`
	Source string `json:"source"`
	Author string `json:"author"`
	refCnt int
}

type gameState struct {
	World    [][]*Tile         `json:"world"`
	Plants   map[string]*Plant `json:"plants"`
	roots    []*growthRoot
	maxPlant int
}

type tileDiff struct {
	Loc  Location `json:"loc"`
	Tile Tile     `json:"tile"`
}

type diff struct {
	TileDiffs     []tileDiff        `json:"tileDiff"`
	NewPlants     map[string]*Plant `json:"newPlants"`
	RemovedPlants []string          `json:"removedPlants"`
}

func (d *diff) isEmpty() bool {
	return len(d.TileDiffs) == 0 && len(d.NewPlants) == 0 && len(d.RemovedPlants) == 0
}

type State struct {
	state gameState
	diff  diff
}

// Records a reference to a plant, causing the plant to be kept in the structure
func (s *State) plantAddRef(plantId string) {
	s.GetPlant(plantId).refCnt++
}

// Records a reference to a plant, causing the plant to be removed from the structure
func (s *State) plantRelease(plantId string) {
	plant := s.GetPlant(plantId)
	plant.refCnt--

	// If the reference count has reached zero, remove the plant from the thing
	if plant.refCnt <= 0 {
		s.diff.RemovedPlants = append(s.diff.RemovedPlants, plantId)
		delete(s.state.Plants, plantId)
	}
}

func (s *State) Width() int {
	return len(s.state.World[0])
}

func (s *State) Height() int {
	return len(s.state.World)
}

func (s *State) GetPlant(plantId string) *Plant {
	return s.state.Plants[plantId]
}

// Adds a species to the stateAndDiff, and returns the string key for the plant
// This plant is created with a refCnt of zero, but will not be dropped until
// its reference count hits zero again.
func (s *State) AddSpecies(color int, source string, author string) string {
	p := Plant{color, source, author, 0}
	s.state.maxPlant += 1
	key := fmt.Sprintf("%d", s.state.maxPlant)
	s.diff.NewPlants[key] = &p
	s.state.Plants[key] = &p
	return key
}

func (s *State) GetTile(loc Location) *Tile {
	return s.state.World[loc.Y][loc.X]
}

// Set the tile at a location to a new tile
func (s *State) SetTile(loc Location, new Tile) {
	// Manage the addref and releases
	if new.Plant != nil {
		s.plantAddRef(new.Plant.PlantId)
	}
	old := s.GetTile(loc)
	if old.Plant != nil {
		s.plantRelease(old.Plant.PlantId)
	}

	// Actually update the tile and record the tilediffs
	*old = new
	s.diff.TileDiffs = append(s.diff.TileDiffs, tileDiff{loc, new})
}

func (s *State) AddPlant(loc Location, id string) *growthRoot {
	plant := s.GetPlant(id)

	// Create the sandbox node for the plant object
	node := sandbox.AddNode(plant.Source)

	s.lowerToDirt(&loc)

	// Create the root node for the object, and append it to the roots list
	root := growthRoot{id, loc, node}
	s.state.roots = append(s.state.roots, &root)

	isUnderground := false
	if s.GetTile(loc).T == DirtTile {
		isUnderground = true
	}

	// Set the tile at the base of the plant to a plant tile
	s.SetTile(loc, Tile{PlantTile, &plantInfo{id, loc, isUnderground}})

	// Return a reference to the root node we previously appended
	return &root
}

func (s *State) lowerToDirt(loc *Location) {
	var base int
	for y := 0; y < s.Height(); y++ {
		t := s.GetTile(Location{loc.X, y})
		if t.T == DirtTile {
			base = y - 1
			break
		}
	}
	loc.Y = base
}
<|endoftext|>


<commit_before>package scheduler

import (
	"net/http"
	"strings"

	"github.com/Dataman-Cloud/swan/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/mesosproto/sched"
	"github.com/Sirupsen/logrus"
)

func (s *Scheduler) status(status *mesos.TaskStatus) {
	// send ack
	if status.GetUuid() != nil {
		call := &sched.Call{
			FrameworkId: s.framework.GetId(),
			Type:        sched.Call_ACKNOWLEDGE.Enum(),
			Acknowledge: &sched.Call_Acknowledge{
				AgentId: status.GetAgentId(),
				TaskId:  status.GetTaskId(),
				Uuid:    status.GetUuid(),
			},
		}

		// send call
		resp, err := s.send(call)
		if err != nil {
			logrus.Error("Unable to send Acknowledge Call: ", err)
			return
		}
		if resp.StatusCode != http.StatusAccepted {
			logrus.Errorf("Acknowledge call returned unexpected status: %d", resp.StatusCode)
			return
		}
	}

	ID := status.TaskId.GetValue()
	state := status.GetState()

	taskId := strings.Split(ID, "-")[1]
	appId := strings.Split(taskId, ".")[1]

	var STATUS string

	switch state {
	case mesos.TaskState_TASK_STAGING:
		STATUS = "STAGING"
		if err := s.store.UpdateTaskStatus(taskId, "STAGING"); err != nil {
			logrus.Errorf("updating task %s status to STAGING failed", taskId)
		}
	case mesos.TaskState_TASK_STARTING:
		STATUS = "STARTING"
		if err := s.store.UpdateTaskStatus(taskId, "STARTING"); err != nil {
			logrus.Errorf("updating task %s status to STARTING failed", taskId)
		}
	case mesos.TaskState_TASK_RUNNING:
		STATUS = "RUNNING"
		if err := s.store.IncreaseApplicationRunningInstances(appId); err != nil {
			logrus.Errorf("Updating application got error: %s", err.Error())
		}

		if err := s.store.UpdateTaskStatus(taskId, "RUNNING"); err != nil {
			logrus.Errorf("updating task %s status to RUNNING failed: %s", taskId, err.Error())
			return
		}

		app, err := s.store.FetchApplication(appId)
		if err != nil {
			break
		}

		if app.RunningInstances == app.Instances {
			if err := s.store.UpdateApplicationStatus(appId, "RUNNING"); err != nil {
				logrus.Errorf("Updating application got error: %s", err.Error())
			}
		}

	case mesos.TaskState_TASK_FINISHED:
		logrus.Info("Task Finished")
		STATUS = "RESCHEDULING"
	case mesos.TaskState_TASK_FAILED:
		logrus.Info("Task Failed")
		STATUS = "RESCHEDULING"
	case mesos.TaskState_TASK_KILLED:
		logrus.Info("Task Killed")
		STATUS = "KILLED"
	case mesos.TaskState_TASK_LOST:
		logrus.Info("Task Lost")
		STATUS = "RESCHEDULING"
	}

	task, err := s.store.FetchTask(taskId)
	if err != nil {
		logrus.Errorf("Fetch task %s failed: %s", taskId, err.Error())
		return
	}

	app, err := s.store.FetchApplication(appId)
	if err != nil {
		logrus.Errorf("Fetch application %s failed: %s", appId, err.Error())
		return
	}

	if STATUS == "RESCHEDULING" &&
		len(task.HealthChecks) == 0 &&
		task.Status != "RESCHEDULING" &&
		app.Status != "UPDATING" &&
		app.Status != "ROLLINGBACK" {
		if err := s.store.UpdateTaskStatus(taskId, "RESCHEDULING"); err != nil {
			logrus.Errorf("updating task status to RESCHEDULING failed: %s", taskId)
		}

		s.Status = "busy"

		resources := s.BuildResources(task.Cpus, task.Mem, task.Disk)
		offers, err := s.RequestOffers(resources)
		if err != nil {
			logrus.Errorf("Request offers failed: %s for rescheduling", err.Error())
			return
		}

		var choosedOffer *mesos.Offer
		for _, offer := range offers {
			cpus, mem, disk := s.OfferedResources(offer)
			if cpus >= task.Cpus && mem >= task.Mem && disk >= task.Disk {
				choosedOffer = offer
				break
			}
		}

		var taskInfos []*mesos.TaskInfo
		taskInfo := s.BuildTaskInfo(choosedOffer, resources, task)
		taskInfos = append(taskInfos, taskInfo)

		resp, err := s.LaunchTasks(choosedOffer, taskInfos)
		if err != nil {
			logrus.Errorf("Launchs task failed: %s for rescheduling", err.Error())
		}

		if resp != nil && resp.StatusCode != http.StatusAccepted {
			logrus.Errorf("Launchs task failed: status code %d for rescheduling", resp.StatusCode)
		}
	}
}
<commit_msg>add task failure message<commit_after>package scheduler

import (
	"net/http"
	"strings"

	"github.com/Dataman-Cloud/swan/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/mesosproto/sched"
	"github.com/Sirupsen/logrus"
)

func (s *Scheduler) status(status *mesos.TaskStatus) {
	// send ack
	if status.GetUuid() != nil {
		call := &sched.Call{
			FrameworkId: s.framework.GetId(),
			Type:        sched.Call_ACKNOWLEDGE.Enum(),
			Acknowledge: &sched.Call_Acknowledge{
				AgentId: status.GetAgentId(),
				TaskId:  status.GetTaskId(),
				Uuid:    status.GetUuid(),
			},
		}

		// send call
		resp, err := s.send(call)
		if err != nil {
			logrus.Error("Unable to send Acknowledge Call: ", err)
			return
		}
		if resp.StatusCode != http.StatusAccepted {
			logrus.Errorf("Acknowledge call returned unexpected status: %d", resp.StatusCode)
			return
		}
	}

	ID := status.TaskId.GetValue()
	state := status.GetState()

	taskId := strings.Split(ID, "-")[1]
	appId := strings.Split(taskId, ".")[1]

	var STATUS string

	switch state {
	case mesos.TaskState_TASK_STAGING:
		STATUS = "STAGING"
		if err := s.store.UpdateTaskStatus(taskId, "STAGING"); err != nil {
			logrus.Errorf("updating task %s status to STAGING failed", taskId)
		}
	case mesos.TaskState_TASK_STARTING:
		STATUS = "STARTING"
		if err := s.store.UpdateTaskStatus(taskId, "STARTING"); err != nil {
			logrus.Errorf("updating task %s status to STARTING failed", taskId)
		}
	case mesos.TaskState_TASK_RUNNING:
		STATUS = "RUNNING"
		if err := s.store.IncreaseApplicationRunningInstances(appId); err != nil {
			logrus.Errorf("Updating application got error: %s", err.Error())
		}

		if err := s.store.UpdateTaskStatus(taskId, "RUNNING"); err != nil {
			logrus.Errorf("updating task %s status to RUNNING failed: %s", taskId, err.Error())
			return
		}

		app, err := s.store.FetchApplication(appId)
		if err != nil {
			break
		}

		if app.RunningInstances == app.Instances {
			if err := s.store.UpdateApplicationStatus(appId, "RUNNING"); err != nil {
				logrus.Errorf("Updating application got error: %s", err.Error())
			}
		}

	case mesos.TaskState_TASK_FINISHED:
		logrus.Infof("Task Finished, message: %s", status.GetMessage())
		STATUS = "RESCHEDULING"
	case mesos.TaskState_TASK_FAILED:
		logrus.Infof("Task Failed, message: %s", status.GetMessage())
		STATUS = "RESCHEDULING"
	case mesos.TaskState_TASK_KILLED:
		logrus.Infof("Task Killed, message: %s", status.GetMessage())
		STATUS = "KILLED"
	case mesos.TaskState_TASK_LOST:
		logrus.Infof("Task Lost, message: %s", status.GetMessage())
		STATUS = "RESCHEDULING"
	}

	task, err := s.store.FetchTask(taskId)
	if err != nil {
		logrus.Errorf("Fetch task %s failed: %s", taskId, err.Error())
		return
	}

	app, err := s.store.FetchApplication(appId)
	if err != nil {
		logrus.Errorf("Fetch application %s failed: %s", appId, err.Error())
		return
	}

	if STATUS == "RESCHEDULING" &&
		len(task.HealthChecks) == 0 &&
		task.Status != "RESCHEDULING" &&
		app.Status != "UPDATING" &&
		app.Status != "ROLLINGBACK" {
		if err := s.store.UpdateTaskStatus(taskId, "RESCHEDULING"); err != nil {
			logrus.Errorf("updating task status to RESCHEDULING failed: %s", taskId)
		}

		s.Status = "busy"

		resources := s.BuildResources(task.Cpus, task.Mem, task.Disk)
		offers, err := s.RequestOffers(resources)
		if err != nil {
			logrus.Errorf("Request offers failed: %s for rescheduling", err.Error())
			return
		}

		var choosedOffer *mesos.Offer
		for _, offer := range offers {
			cpus, mem, disk := s.OfferedResources(offer)
			if cpus >= task.Cpus && mem >= task.Mem && disk >= task.Disk {
				choosedOffer = offer
				break
			}
		}

		var taskInfos []*mesos.TaskInfo
		taskInfo := s.BuildTaskInfo(choosedOffer, resources, task)
		taskInfos = append(taskInfos, taskInfo)

		resp, err := s.LaunchTasks(choosedOffer, taskInfos)
		if err != nil {
			logrus.Errorf("Launchs task failed: %s for rescheduling", err.Error())
		}

		if resp != nil && resp.StatusCode != http.StatusAccepted {
			logrus.Errorf("Launchs task failed: status code %d for rescheduling", resp.StatusCode)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package addressupdater

import (
	stderrors "errors"
	"fmt"
	"math"
	"sync"
	"sync/atomic"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	coretesting "launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/testing/testbase"
)

var _ = gc.Suite(&machineSuite{})

type machineSuite struct {
	testbase.LoggingSuite
}

var testAddrs = []instance.Address{instance.NewAddress("127.0.0.1")}

func (*machineSuite) TestSetsAddressInitially(c *gc.C) {
	context := &testMachineContext{
		getAddresses: addressesGetter(c, "i1234", testAddrs, nil),
		dyingc:       make(chan struct{}),
	}
	m := &testMachine{
		id:         "99",
		instanceId: "i1234",
		refresh:    func() error { return nil },
		life:       state.Alive,
	}
	died := make(chan machine)
	// Change the poll intervals to be short, so that we know
	// that we've polled (probably) at least a few times.
	defer testbase.PatchValue(&ShortPoll, coretesting.ShortWait/10).Restore()
	defer testbase.PatchValue(&LongPoll, coretesting.ShortWait/10).Restore()

	go runMachine(context, m, nil, died)
	time.Sleep(coretesting.ShortWait)

	killMachineLoop(c, m, context.dyingc, died)
	c.Assert(context.killAllErr, gc.Equals, nil)
	c.Assert(m.addresses, gc.DeepEquals, testAddrs)
	c.Assert(m.setAddressCount, gc.Equals, 1)
}

func (*machineSuite) TestShortPollIntervalWhenNoAddress(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, 1*time.Millisecond).Restore()
	defer testbase.PatchValue(&LongPoll, coretesting.LongWait).Restore()
	count := countPolls(c, nil)
	c.Assert(count, jc.GreaterThan, 2)
}

func (*machineSuite) TestShortPollIntervalExponent(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, 1*time.Microsecond).Restore()
	defer testbase.PatchValue(&LongPoll, coretesting.LongWait).Restore()
	defer testbase.PatchValue(&ShortPollBackoff, 2.0).Restore()

	// With an exponent of 2, the maximum number of polls that can
	// occur within the given interval ShortWait is log to the base
	// base ShortPollBackoff of ShortWait/ShortPoll, given that sleep will
	// sleep for at least the requested interval.
	maxCount := int(math.Log(float64(coretesting.ShortWait)/float64(ShortPoll))/math.Log(ShortPollBackoff) + 1)
	count := countPolls(c, nil)
	c.Assert(count, jc.GreaterThan, 2)
	c.Assert(count, jc.LessThan, maxCount)
	c.Logf("actual count: %v; max %v", count, maxCount)
}

func (*machineSuite) TestLongPollIntervalWhenHasAddress(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, coretesting.LongWait).Restore()
	defer testbase.PatchValue(&LongPoll, 1*time.Millisecond).Restore()
	count := countPolls(c, testAddrs)
	c.Assert(count, jc.GreaterThan, 2)
}

// countPolls sets up a machine loop with the given
// addresses to be returned from getAddresses,
// waits for coretesting.ShortWait, and returns the
// number of times the addresses are polled.
func countPolls(c *gc.C, addrs []instance.Address) int {
	count := int32(0)
	getAddresses := func(id instance.Id) ([]instance.Address, error) {
		c.Check(id, gc.Equals, instance.Id("i1234"))
		atomic.AddInt32(&count, 1)
		if addrs == nil {
			return nil, fmt.Errorf("no instance addresses available")
		}
		return addrs, nil
	}
	context := &testMachineContext{
		getAddresses: getAddresses,
		dyingc:       make(chan struct{}),
	}
	m := &testMachine{
		instanceId: "i1234",
		refresh:    func() error { return nil },
		addresses:  addrs,
		life:       state.Alive,
	}
	died := make(chan machine)

	go runMachine(context, m, nil, died)

	time.Sleep(coretesting.ShortWait)
	killMachineLoop(c, m, context.dyingc, died)
	c.Assert(context.killAllErr, gc.Equals, nil)
	return int(count)
}

func (*machineSuite) TestSinglePollWhenAddressesUnimplemented(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, 1*time.Millisecond).Restore()
	defer testbase.PatchValue(&LongPoll, 1*time.Millisecond).Restore()
	count := int32(0)
	getAddresses := func(id instance.Id) ([]instance.Address, error) {
		c.Check(id, gc.Equals, instance.Id("i1234"))
		atomic.AddInt32(&count, 1)
		return nil, errors.NewNotImplementedError("instance address")
	}
	context := &testMachineContext{
		getAddresses: getAddresses,
		dyingc:       make(chan struct{}),
	}
	m := &testMachine{
		instanceId: "i1234",
		refresh:    func() error { return nil },
		life:       state.Alive,
	}
	died := make(chan machine)

	go runMachine(context, m, nil, died)

	time.Sleep(coretesting.ShortWait)
	killMachineLoop(c, m, context.dyingc, died)
	c.Assert(context.killAllErr, gc.Equals, nil)
	c.Assert(count, gc.Equals, int32(1))
}

func (*machineSuite) TestChangedRefreshes(c *gc.C) {
	context := &testMachineContext{
		getAddresses: addressesGetter(c, "i1234", testAddrs, nil),
		dyingc:       make(chan struct{}),
	}
	refreshc := make(chan struct{})
	m := &testMachine{
		instanceId: "i1234",
		refresh: func() error {
			refreshc <- struct{}{}
			return nil
		},
		addresses: testAddrs,
		life:      state.Dead,
	}
	died := make(chan machine)
	changed := make(chan struct{})
	go runMachine(context, m, changed, died)
	select {
	case <-died:
		c.Fatalf("machine died prematurely")
	case <-time.After(coretesting.ShortWait):
	}

	// Notify the machine that it has changed; it should
	// refresh, and publish the fact that it no longer has
	// an address.
	changed <- struct{}{}

	select {
	case <-refreshc:
	case <-time.After(coretesting.LongWait):
		c.Fatalf("timed out waiting for refresh")
	}
	select {
	case <-died:
	case <-time.After(coretesting.LongWait):
		c.Fatalf("expected death after life set to dying")
	}
	// The machine addresses should remain the same even
	// after death.
	c.Assert(m.addresses, gc.DeepEquals, testAddrs)
}

var terminatingErrorsTests = []struct {
	about  string
	mutate func(m *testMachine, err error)
}{{
	about: "set addresses",
	mutate: func(m *testMachine, err error) {
		m.setAddressesErr = err
	},
}, {
	about: "refresh",
	mutate: func(m *testMachine, err error) {
		m.refresh = func() error {
			return err
		}
	},
}, {
	about: "instance id",
	mutate: func(m *testMachine, err error) {
		m.instanceIdErr = err
	},
}}

func (*machineSuite) TestTerminatingErrors(c *gc.C) {
	for i, test := range terminatingErrorsTests {
		c.Logf("test %d: %s", i, test.about)
		testTerminatingErrors(c, test.mutate)
	}
}

// testTerminatingErrors checks that when a testMachine is
// changed with the given mutate function, the machine goroutine
// will die having called its context's killAll function with the
// given error.  The test is cunningly structured so that it in the normal course
// of things it will go through all possible places that can return an error.
func testTerminatingErrors(c *gc.C, mutate func(m *testMachine, err error)) {
	context := &testMachineContext{
		getAddresses: addressesGetter(c, "i1234", testAddrs, nil),
		dyingc:       make(chan struct{}),
	}
	expectErr := stderrors.New("a very unusual error")
	m := &testMachine{
		id:         "99",
		instanceId: "i1234",
		refresh:    func() error { return nil },
		life:       state.Alive,
	}
	mutate(m, expectErr)
	died := make(chan machine)
	changed := make(chan struct{}, 1)
	go runMachine(context, m, changed, died)
	changed <- struct{}{}
	select {
	case <-died:
	case <-time.After(coretesting.LongWait):
		c.Fatalf("timed out waiting for machine to die")
	}
	c.Assert(context.killAllErr, gc.ErrorMatches, ".*"+expectErr.Error())
}

func killMachineLoop(c *gc.C, m machine, dying chan struct{}, died <-chan machine) {
	close(dying)
	select {
	case diedm := <-died:
		c.Assert(diedm, gc.Equals, m)
	case <-time.After(coretesting.LongWait):
		c.Fatalf("updater did not die after dying channel was closed")
	}
}

func addressesGetter(c *gc.C, expectId instance.Id, addrs []instance.Address, err error) func(id instance.Id) ([]instance.Address, error) {
	return func(id instance.Id) ([]instance.Address, error) {
		c.Check(id, gc.Equals, expectId)
		return addrs, err
	}
}

type testMachineContext struct {
	killAllErr   error
	getAddresses func(instance.Id) ([]instance.Address, error)
	dyingc       chan struct{}
}

func (context *testMachineContext) killAll(err error) {
	if err == nil {
		panic("killAll with nil error")
	}
	context.killAllErr = err
}

func (context *testMachineContext) addresses(id instance.Id) ([]instance.Address, error) {
	return context.getAddresses(id)
}

func (context *testMachineContext) dying() <-chan struct{} {
	return context.dyingc
}

type testMachine struct {
	instanceId      instance.Id
	instanceIdErr   error
	id              string
	refresh         func() error
	setAddressesErr error
	// mu protects the following fields.
	mu              sync.Mutex
	life            state.Life
	addresses       []instance.Address
	setAddressCount int
}

func (m *testMachine) Id() string {
	if m.id == "" {
		panic("Id called but not set")
	}
	return m.id
}

func (m *testMachine) Addresses() []instance.Address {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.addresses
}

func (m *testMachine) InstanceId() (instance.Id, error) {
	return m.instanceId, m.instanceIdErr
}

func (m *testMachine) SetAddresses(addrs []instance.Address) error {
	if m.setAddressesErr != nil {
		return m.setAddressesErr
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.addresses = append(m.addresses[:0], addrs...)
	m.setAddressCount++
	return nil
}

func (m *testMachine) String() string {
	return m.id
}

func (m *testMachine) Refresh() error {
	return m.refresh()
}

func (m *testMachine) Life() state.Life {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.life
}

func (m *testMachine) setLife(life state.Life) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.life = life
}
<commit_msg>worker/addressupdater: fix comment<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package addressupdater

import (
	stderrors "errors"
	"fmt"
	"math"
	"sync"
	"sync/atomic"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	coretesting "launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/testing/testbase"
)

var _ = gc.Suite(&machineSuite{})

type machineSuite struct {
	testbase.LoggingSuite
}

var testAddrs = []instance.Address{instance.NewAddress("127.0.0.1")}

func (*machineSuite) TestSetsAddressInitially(c *gc.C) {
	context := &testMachineContext{
		getAddresses: addressesGetter(c, "i1234", testAddrs, nil),
		dyingc:       make(chan struct{}),
	}
	m := &testMachine{
		id:         "99",
		instanceId: "i1234",
		refresh:    func() error { return nil },
		life:       state.Alive,
	}
	died := make(chan machine)
	// Change the poll intervals to be short, so that we know
	// that we've polled (probably) at least a few times.
	defer testbase.PatchValue(&ShortPoll, coretesting.ShortWait/10).Restore()
	defer testbase.PatchValue(&LongPoll, coretesting.ShortWait/10).Restore()

	go runMachine(context, m, nil, died)
	time.Sleep(coretesting.ShortWait)

	killMachineLoop(c, m, context.dyingc, died)
	c.Assert(context.killAllErr, gc.Equals, nil)
	c.Assert(m.addresses, gc.DeepEquals, testAddrs)
	c.Assert(m.setAddressCount, gc.Equals, 1)
}

func (*machineSuite) TestShortPollIntervalWhenNoAddress(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, 1*time.Millisecond).Restore()
	defer testbase.PatchValue(&LongPoll, coretesting.LongWait).Restore()
	count := countPolls(c, nil)
	c.Assert(count, jc.GreaterThan, 2)
}

func (*machineSuite) TestShortPollIntervalExponent(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, 1*time.Microsecond).Restore()
	defer testbase.PatchValue(&LongPoll, coretesting.LongWait).Restore()
	defer testbase.PatchValue(&ShortPollBackoff, 2.0).Restore()

	// With an exponent of 2, the maximum number of polls that can
	// occur within the given interval ShortWait is log to the base
	// ShortPollBackoff of ShortWait/ShortPoll, given that sleep will
	// sleep for at least the requested interval.
	maxCount := int(math.Log(float64(coretesting.ShortWait)/float64(ShortPoll))/math.Log(ShortPollBackoff) + 1)
	count := countPolls(c, nil)
	c.Assert(count, jc.GreaterThan, 2)
	c.Assert(count, jc.LessThan, maxCount)
	c.Logf("actual count: %v; max %v", count, maxCount)
}

func (*machineSuite) TestLongPollIntervalWhenHasAddress(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, coretesting.LongWait).Restore()
	defer testbase.PatchValue(&LongPoll, 1*time.Millisecond).Restore()
	count := countPolls(c, testAddrs)
	c.Assert(count, jc.GreaterThan, 2)
}

// countPolls sets up a machine loop with the given
// addresses to be returned from getAddresses,
// waits for coretesting.ShortWait, and returns the
// number of times the addresses are polled.
func countPolls(c *gc.C, addrs []instance.Address) int {
	count := int32(0)
	getAddresses := func(id instance.Id) ([]instance.Address, error) {
		c.Check(id, gc.Equals, instance.Id("i1234"))
		atomic.AddInt32(&count, 1)
		if addrs == nil {
			return nil, fmt.Errorf("no instance addresses available")
		}
		return addrs, nil
	}
	context := &testMachineContext{
		getAddresses: getAddresses,
		dyingc:       make(chan struct{}),
	}
	m := &testMachine{
		instanceId: "i1234",
		refresh:    func() error { return nil },
		addresses:  addrs,
		life:       state.Alive,
	}
	died := make(chan machine)

	go runMachine(context, m, nil, died)

	time.Sleep(coretesting.ShortWait)
	killMachineLoop(c, m, context.dyingc, died)
	c.Assert(context.killAllErr, gc.Equals, nil)
	return int(count)
}

func (*machineSuite) TestSinglePollWhenAddressesUnimplemented(c *gc.C) {
	defer testbase.PatchValue(&ShortPoll, 1*time.Millisecond).Restore()
	defer testbase.PatchValue(&LongPoll, 1*time.Millisecond).Restore()
	count := int32(0)
	getAddresses := func(id instance.Id) ([]instance.Address, error) {
		c.Check(id, gc.Equals, instance.Id("i1234"))
		atomic.AddInt32(&count, 1)
		return nil, errors.NewNotImplementedError("instance address")
	}
	context := &testMachineContext{
		getAddresses: getAddresses,
		dyingc:       make(chan struct{}),
	}
	m := &testMachine{
		instanceId: "i1234",
		refresh:    func() error { return nil },
		life:       state.Alive,
	}
	died := make(chan machine)

	go runMachine(context, m, nil, died)

	time.Sleep(coretesting.ShortWait)
	killMachineLoop(c, m, context.dyingc, died)
	c.Assert(context.killAllErr, gc.Equals, nil)
	c.Assert(count, gc.Equals, int32(1))
}

func (*machineSuite) TestChangedRefreshes(c *gc.C) {
	context := &testMachineContext{
		getAddresses: addressesGetter(c, "i1234", testAddrs, nil),
		dyingc:       make(chan struct{}),
	}
	refreshc := make(chan struct{})
	m := &testMachine{
		instanceId: "i1234",
		refresh: func() error {
			refreshc <- struct{}{}
			return nil
		},
		addresses: testAddrs,
		life:      state.Dead,
	}
	died := make(chan machine)
	changed := make(chan struct{})
	go runMachine(context, m, changed, died)
	select {
	case <-died:
		c.Fatalf("machine died prematurely")
	case <-time.After(coretesting.ShortWait):
	}

	// Notify the machine that it has changed; it should
	// refresh, and publish the fact that it no longer has
	// an address.
	changed <- struct{}{}

	select {
	case <-refreshc:
	case <-time.After(coretesting.LongWait):
		c.Fatalf("timed out waiting for refresh")
	}
	select {
	case <-died:
	case <-time.After(coretesting.LongWait):
		c.Fatalf("expected death after life set to dying")
	}
	// The machine addresses should remain the same even
	// after death.
	c.Assert(m.addresses, gc.DeepEquals, testAddrs)
}

var terminatingErrorsTests = []struct {
	about  string
	mutate func(m *testMachine, err error)
}{{
	about: "set addresses",
	mutate: func(m *testMachine, err error) {
		m.setAddressesErr = err
	},
}, {
	about: "refresh",
	mutate: func(m *testMachine, err error) {
		m.refresh = func() error {
			return err
		}
	},
}, {
	about: "instance id",
	mutate: func(m *testMachine, err error) {
		m.instanceIdErr = err
	},
}}

func (*machineSuite) TestTerminatingErrors(c *gc.C) {
	for i, test := range terminatingErrorsTests {
		c.Logf("test %d: %s", i, test.about)
		testTerminatingErrors(c, test.mutate)
	}
}

// testTerminatingErrors checks that when a testMachine is
// changed with the given mutate function, the machine goroutine
// will die having called its context's killAll function with the
// given error.  The test is cunningly structured so that it in the normal course
// of things it will go through all possible places that can return an error.
func testTerminatingErrors(c *gc.C, mutate func(m *testMachine, err error)) {
	context := &testMachineContext{
		getAddresses: addressesGetter(c, "i1234", testAddrs, nil),
		dyingc:       make(chan struct{}),
	}
	expectErr := stderrors.New("a very unusual error")
	m := &testMachine{
		id:         "99",
		instanceId: "i1234",
		refresh:    func() error { return nil },
		life:       state.Alive,
	}
	mutate(m, expectErr)
	died := make(chan machine)
	changed := make(chan struct{}, 1)
	go runMachine(context, m, changed, died)
	changed <- struct{}{}
	select {
	case <-died:
	case <-time.After(coretesting.LongWait):
		c.Fatalf("timed out waiting for machine to die")
	}
	c.Assert(context.killAllErr, gc.ErrorMatches, ".*"+expectErr.Error())
}

func killMachineLoop(c *gc.C, m machine, dying chan struct{}, died <-chan machine) {
	close(dying)
	select {
	case diedm := <-died:
		c.Assert(diedm, gc.Equals, m)
	case <-time.After(coretesting.LongWait):
		c.Fatalf("updater did not die after dying channel was closed")
	}
}

func addressesGetter(c *gc.C, expectId instance.Id, addrs []instance.Address, err error) func(id instance.Id) ([]instance.Address, error) {
	return func(id instance.Id) ([]instance.Address, error) {
		c.Check(id, gc.Equals, expectId)
		return addrs, err
	}
}

type testMachineContext struct {
	killAllErr   error
	getAddresses func(instance.Id) ([]instance.Address, error)
	dyingc       chan struct{}
}

func (context *testMachineContext) killAll(err error) {
	if err == nil {
		panic("killAll with nil error")
	}
	context.killAllErr = err
}

func (context *testMachineContext) addresses(id instance.Id) ([]instance.Address, error) {
	return context.getAddresses(id)
}

func (context *testMachineContext) dying() <-chan struct{} {
	return context.dyingc
}

type testMachine struct {
	instanceId      instance.Id
	instanceIdErr   error
	id              string
	refresh         func() error
	setAddressesErr error
	// mu protects the following fields.
	mu              sync.Mutex
	life            state.Life
	addresses       []instance.Address
	setAddressCount int
}

func (m *testMachine) Id() string {
	if m.id == "" {
		panic("Id called but not set")
	}
	return m.id
}

func (m *testMachine) Addresses() []instance.Address {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.addresses
}

func (m *testMachine) InstanceId() (instance.Id, error) {
	return m.instanceId, m.instanceIdErr
}

func (m *testMachine) SetAddresses(addrs []instance.Address) error {
	if m.setAddressesErr != nil {
		return m.setAddressesErr
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.addresses = append(m.addresses[:0], addrs...)
	m.setAddressCount++
	return nil
}

func (m *testMachine) String() string {
	return m.id
}

func (m *testMachine) Refresh() error {
	return m.refresh()
}

func (m *testMachine) Life() state.Life {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.life
}

func (m *testMachine) setLife(life state.Life) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.life = life
}
<|endoftext|>