<commit_before>package worker

import (
	"fmt"
	"github.com/go-redis/redis"
	"github.com/stitchfix/flotilla-os/queue"
	"math/rand"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
)

type statusWorker struct {
	sm           state.Manager
	ee           engine.Engine
	conf         config.Config
	log          flotillaLog.Logger
	pollInterval time.Duration
	t            tomb.Tomb
	engine       *string
	redisClient  *redis.Client
	workerId     string
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, ee engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, engine *string, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = ee
	sw.log = log
	sw.engine = engine
	sw.workerId = fmt.Sprintf("%d", rand.Int())
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker", "engine", *engine)
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.ECSEngine {
				sw.runOnceECS()
				time.Sleep(sw.pollInterval)
			}

			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -3).Format(time.RFC3339),
		},
		"status": {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	//sw.log.Log("message", "recieved runs for processing", "count", len(runs))
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		lock := sw.acquireLock(run, "status", 15*time.Second)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	_ = sw.log.Log("message", "received locked runs for processing", "count", len(lockedRuns))
	for _, run := range lockedRuns {
		sw.processEKSRun(run)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	set, err := sw.redisClient.SetNX(fmt.Sprintf("%s-%s", run.RunID, purpose), sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)

	if err != nil {
		message := fmt.Sprintf("%+v", err)
		_ = sw.log.Log("message", "unable to receive eks runs", "error", message)

		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status {
			_ = sw.log.Log("message", "updating eks run", "run", updatedRun.RunID, "status", updatedRun.Status)
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.ee.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) runOnceECS() {
	runReceipt, err := sw.ee.PollStatus()
	if err != nil {
		sw.log.Log("message", "unable to receive status message", "error", fmt.Sprintf("%+v", err))
		return
	}

	// Ensure update is in the env required, otherwise, ack without taking action
	update := runReceipt.Run
	if update != nil {
		//
		// Relies on the reserved env var, FLOTILLA_SERVER_MODE to ensure update
		// belongs to -this- mode of Flotilla
		//
		var serverMode string
		if update.Env != nil {
			for _, kv := range *update.Env {
				if kv.Name == "FLOTILLA_SERVER_MODE" {
					serverMode = kv.Value
				}
			}
		}

		shouldProcess := len(serverMode) > 0 && serverMode == sw.conf.GetString("flotilla_mode")
		if shouldProcess {
			run, err := sw.findRun(update.TaskArn)
			if err != nil {
				sw.log.Log("message", "unable to find run to apply update to", "error", fmt.Sprintf("%+v", err))
				return
			}

			_, err = sw.sm.UpdateRun(run.RunID, *update)
			if err != nil {
				sw.log.Log("message", "error applying status update", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
				return
			}

			// emit status update event
			sw.logStatusUpdate(*update)
		}

		if sw.conf.GetString("flotilla_mode") != "dev" {
			_ = sw.log.Log("message", "Acking status update", "arn", update.TaskArn)
		}
		if err = runReceipt.Done(); err != nil {
			sw.log.Log("message", "Acking status update failed", "arn", update.TaskArn, "error", fmt.Sprintf("%+v", err))
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<commit_msg>remove log line, reducing verbosity<commit_after>package worker

import (
	"fmt"
	"github.com/go-redis/redis"
	"github.com/stitchfix/flotilla-os/queue"
	"math/rand"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
)

type statusWorker struct {
	sm           state.Manager
	ee           engine.Engine
	conf         config.Config
	log          flotillaLog.Logger
	pollInterval time.Duration
	t            tomb.Tomb
	engine       *string
	redisClient  *redis.Client
	workerId     string
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, ee engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, engine *string, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = ee
	sw.log = log
	sw.engine = engine
	sw.workerId = fmt.Sprintf("%d", rand.Int())
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker", "engine", *engine)
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.ECSEngine {
				sw.runOnceECS()
				time.Sleep(sw.pollInterval)
			}

			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -3).Format(time.RFC3339),
		},
		"status": {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	//sw.log.Log("message", "recieved runs for processing", "count", len(runs))
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		lock := sw.acquireLock(run, "status", 15*time.Second)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	//_ = sw.log.Log("message", "received locked runs for processing", "count", len(lockedRuns))
	for _, run := range lockedRuns {
		sw.processEKSRun(run)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	set, err := sw.redisClient.SetNX(fmt.Sprintf("%s-%s", run.RunID, purpose), sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)

	if err != nil {
		message := fmt.Sprintf("%+v", err)
		_ = sw.log.Log("message", "unable to receive eks runs", "error", message)

		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status {
			_ = sw.log.Log("message", "updating eks run", "run", updatedRun.RunID, "status", updatedRun.Status)
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.ee.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) runOnceECS() {
	runReceipt, err := sw.ee.PollStatus()
	if err != nil {
		sw.log.Log("message", "unable to receive status message", "error", fmt.Sprintf("%+v", err))
		return
	}

	// Ensure update is in the env required, otherwise, ack without taking action
	update := runReceipt.Run
	if update != nil {
		//
		// Relies on the reserved env var, FLOTILLA_SERVER_MODE to ensure update
		// belongs to -this- mode of Flotilla
		//
		var serverMode string
		if update.Env != nil {
			for _, kv := range *update.Env {
				if kv.Name == "FLOTILLA_SERVER_MODE" {
					serverMode = kv.Value
				}
			}
		}

		shouldProcess := len(serverMode) > 0 && serverMode == sw.conf.GetString("flotilla_mode")
		if shouldProcess {
			run, err := sw.findRun(update.TaskArn)
			if err != nil {
				sw.log.Log("message", "unable to find run to apply update to", "error", fmt.Sprintf("%+v", err))
				return
			}

			_, err = sw.sm.UpdateRun(run.RunID, *update)
			if err != nil {
				sw.log.Log("message", "error applying status update", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
				return
			}

			// emit status update event
			sw.logStatusUpdate(*update)
		}

		if sw.conf.GetString("flotilla_mode") != "dev" {
			_ = sw.log.Log("message", "Acking status update", "arn", update.TaskArn)
		}
		if err = runReceipt.Done(); err != nil {
			sw.log.Log("message", "Acking status update failed", "arn", update.TaskArn, "error", fmt.Sprintf("%+v", err))
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<|endoftext|>


<commit_before>package controllers

import (
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/gorilla/mux"
	"github.com/prest/prest/adapters/postgres"
	"github.com/prest/prest/config"
	"github.com/prest/prest/config/router"
	"github.com/prest/prest/middlewares"
)

func TestMain(m *testing.M) {
	os.Setenv("PREST_CONF", "./testdata/prest.toml")
	config.Load()

	createMockScripts(config.PrestConf.QueriesPath)
	writeMockScripts(config.PrestConf.QueriesPath)

	code := m.Run()

	removeMockScripts(config.PrestConf.QueriesPath)
	os.Exit(code)
}

func TestExecuteScriptQuery(t *testing.T) {
	r := mux.NewRouter()
	r.HandleFunc("/testing/script-get/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp, err := ExecuteScriptQuery(r, "fulltable", "get_all")
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}
		w.Write(resp)
	}))

	r.HandleFunc("/testing/script-post/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp, err := ExecuteScriptQuery(r, "fulltable", "write_all")
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}
		w.Write(resp)
	}))

	ts := httptest.NewServer(r)
	defer ts.Close()

	var testCases = []struct {
		description string
		url         string
		method      string
		status      int
	}{
		{"Execute script GET method", "/testing/script-get/?field1=gopher", "GET", http.StatusOK},
		{"Execute script POST method", "/testing/script-post/?field1=gopherzin&field2=pereira", "POST", http.StatusOK},
	}

	for _, tc := range testCases {
		t.Log(tc.description)
		doRequest(t, ts.URL+tc.url, nil, tc.method, tc.status, "ExecuteScriptQuery")
	}
}

func TestExecuteFromScripts(t *testing.T) {
	router := mux.NewRouter()
	router.HandleFunc("/_QUERIES/{queriesLocation}/{script}", ExecuteFromScripts)
	server := httptest.NewServer(router)
	defer server.Close()

	var testCases = []struct {
		description string
		url         string
		method      string
		status      int
	}{
		{"Get results using scripts and funcs by GET method", "/_QUERIES/fulltable/funcs", "GET", http.StatusOK},
		{"Get results using scripts by GET method", "/_QUERIES/fulltable/get_all?field1=gopher", "GET", http.StatusOK},
		{"Get results using scripts by GET method (2)", "/_QUERIES/fulltable/get_header", "GET", http.StatusOK},
		{"Get results using scripts by POST method", "/_QUERIES/fulltable/write_all?field1=gopherzin&field2=pereira", "POST", http.StatusOK},
		{"Get results using scripts by PUT method", "/_QUERIES/fulltable/put_all?field1=trump&field2=pereira", "PUT", http.StatusOK},
		{"Get results using scripts by PATCH method", "/_QUERIES/fulltable/patch_all?field1=temer&field2=trump", "PATCH", http.StatusOK},
		{"Get results using scripts by DELETE method", "/_QUERIES/fulltable/delete_all?field1=trump", "DELETE", http.StatusOK},
		// errors
		{"Get errors using nonexistent folder", "/_QUERIES/fullnon/delete_all?field1=trump", "DELETE", http.StatusBadRequest},
		{"Get errors using nonexistent script", "/_QUERIES/fulltable/some_com_all?field1=trump", "DELETE", http.StatusBadRequest},
		{"Get errors with invalid execution of sql", "/_QUERIES/fulltable/create_table?field1=test7", "POST", http.StatusBadRequest},
	}

	for _, tc := range testCases {
		t.Log(tc.description)
		doRequest(t, server.URL+tc.url, nil, tc.method, tc.status, "ExecuteFromScripts")
	}
}

func TestRenderWithXML(t *testing.T) {
	var testCases = []struct {
		description string
		url         string
		method      string
		status      int
		body        string
	}{
		{"Get schemas with COUNT clause with XML Render", "/schemas?_count=*&_renderer=xml", "GET", 200, "<objects><object><count>4</count></object></objects>"},
	}
	os.Setenv("PREST_DEBUG", "true")
	config.Load()
	postgres.Load()
	n := middlewares.GetApp()
	r := router.Get()

	r.HandleFunc("/schemas", GetSchemas).Methods("GET")
	n.UseHandler(r)
	server := httptest.NewServer(n)
	defer server.Close()

	for _, tc := range testCases {
		t.Log(tc.description)
		doRequest(t, server.URL+tc.url, nil, tc.method, tc.status, "GetSchemas", tc.body)

	}
}
<commit_msg>remove createMockScripts Queries<commit_after>package controllers

import (
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/gorilla/mux"
	"github.com/prest/prest/adapters/postgres"
	"github.com/prest/prest/config"
	"github.com/prest/prest/config/router"
	"github.com/prest/prest/middlewares"
)

func TestMain(m *testing.M) {
	os.Setenv("PREST_CONF", "./testdata/prest.toml")
	config.Load()
	code := m.Run()
	os.Exit(code)
}

func TestExecuteScriptQuery(t *testing.T) {
	r := mux.NewRouter()
	r.HandleFunc("/testing/script-get/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp, err := ExecuteScriptQuery(r, "fulltable", "get_all")
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}
		w.Write(resp)
	}))

	r.HandleFunc("/testing/script-post/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp, err := ExecuteScriptQuery(r, "fulltable", "write_all")
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
		}
		w.Write(resp)
	}))

	ts := httptest.NewServer(r)
	defer ts.Close()

	var testCases = []struct {
		description string
		url         string
		method      string
		status      int
	}{
		{"Execute script GET method", "/testing/script-get/?field1=gopher", "GET", http.StatusOK},
		{"Execute script POST method", "/testing/script-post/?field1=gopherzin&field2=pereira", "POST", http.StatusOK},
	}

	for _, tc := range testCases {
		t.Log(tc.description)
		doRequest(t, ts.URL+tc.url, nil, tc.method, tc.status, "ExecuteScriptQuery")
	}
}

func TestExecuteFromScripts(t *testing.T) {
	router := mux.NewRouter()
	router.HandleFunc("/_QUERIES/{queriesLocation}/{script}", ExecuteFromScripts)
	server := httptest.NewServer(router)
	defer server.Close()

	var testCases = []struct {
		description string
		url         string
		method      string
		status      int
	}{
		{"Get results using scripts and funcs by GET method", "/_QUERIES/fulltable/funcs", "GET", http.StatusOK},
		{"Get results using scripts by GET method", "/_QUERIES/fulltable/get_all?field1=gopher", "GET", http.StatusOK},
		{"Get results using scripts by GET method (2)", "/_QUERIES/fulltable/get_header", "GET", http.StatusOK},
		{"Get results using scripts by POST method", "/_QUERIES/fulltable/write_all?field1=gopherzin&field2=pereira", "POST", http.StatusOK},
		{"Get results using scripts by PUT method", "/_QUERIES/fulltable/put_all?field1=trump&field2=pereira", "PUT", http.StatusOK},
		{"Get results using scripts by PATCH method", "/_QUERIES/fulltable/patch_all?field1=temer&field2=trump", "PATCH", http.StatusOK},
		{"Get results using scripts by DELETE method", "/_QUERIES/fulltable/delete_all?field1=trump", "DELETE", http.StatusOK},
		// errors
		{"Get errors using nonexistent folder", "/_QUERIES/fullnon/delete_all?field1=trump", "DELETE", http.StatusBadRequest},
		{"Get errors using nonexistent script", "/_QUERIES/fulltable/some_com_all?field1=trump", "DELETE", http.StatusBadRequest},
		{"Get errors with invalid execution of sql", "/_QUERIES/fulltable/create_table?field1=test7", "POST", http.StatusBadRequest},
	}

	for _, tc := range testCases {
		t.Log(tc.description)
		doRequest(t, server.URL+tc.url, nil, tc.method, tc.status, "ExecuteFromScripts")
	}
}

func TestRenderWithXML(t *testing.T) {
	var testCases = []struct {
		description string
		url         string
		method      string
		status      int
		body        string
	}{
		{"Get schemas with COUNT clause with XML Render", "/schemas?_count=*&_renderer=xml", "GET", 200, "<objects><object><count>4</count></object></objects>"},
	}
	os.Setenv("PREST_DEBUG", "true")
	config.Load()
	postgres.Load()
	n := middlewares.GetApp()
	r := router.Get()

	r.HandleFunc("/schemas", GetSchemas).Methods("GET")
	n.UseHandler(r)
	server := httptest.NewServer(n)
	defer server.Close()

	for _, tc := range testCases {
		t.Log(tc.description)
		doRequest(t, server.URL+tc.url, nil, tc.method, tc.status, "GetSchemas", tc.body)

	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/tools/logger"
	"os/exec"
	"strings"
	"time"

	"github.com/fatih/set"
	"github.com/op/go-logging"
)

var (
	log      = logger.New("rbd-cleaner")
	interval = time.Minute * 30
)

func main() {
	logging.SetLevel(logging.INFO, "rbd-cleaner")
	log.Info("starting rbd-cleaner")
	log.Info("interval set to %s", interval)

	// first start
	err := rbdCleaner()
	if err != nil {
		log.Error(err.Error())
	}

	// and then for every other intervals
	for _ = range time.Tick(interval) {
		err := rbdCleaner()
		if err != nil {
			log.Error(err.Error())
		}
	}
}

func rbdCleaner() error {
	s, err := rbdList()
	if err != nil {
		return err
	}

	log.Info("cleaner started. total images: %d", s.Size())
	cleanCount := 0
	errCount := 0
	s.Each(func(rbdName interface{}) bool {
		vmName := rbdName.(string)
		vmId := strings.TrimPrefix(vmName, "vm-")
		vm := new(models.VM)

		err := mongodb.One("jVMs", vmId, vm)
		if err == nil {
			return true //  rbd image does exists in jVMS, don't touch it and continue
		}

		log.Info("removing image '%s' from rbd", vmName)
		res, err := removeRBDImage(vmName)
		if err != nil {
			errCount++
			log.Error("err removing '%s'\n", err)
			return true
		}

		cleanCount++
		log.Info(res)
		return true
	})
	log.Info("cleaner ended. images cleaned: %d images with err: %d", cleanCount, errCount)

	return nil
}

func removeRBDImage(imageName string) (string, error) {
	out, err := exec.Command("/usr/bin/rbd", "rm", "vms/"+imageName).CombinedOutput()
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(out)), nil
}

func rbdList() (*set.Set, error) {
	s := set.New()
	out, err := exec.Command("/usr/bin/rbd", "ls", "vms").CombinedOutput()
	if err != nil {
		return nil, err
	}

	rbds := string(bytes.TrimSpace(out))
	for _, r := range strings.Split(rbds, "\n") {
		s.Add(r)
	}
	return s, nil
}
<commit_msg>rbdCleaner: fix log<commit_after>package main

import (
	"bytes"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/tools/logger"
	"os/exec"
	"strings"
	"time"

	"github.com/fatih/set"
	"github.com/op/go-logging"
)

var (
	log      = logger.New("rbd-cleaner")
	interval = time.Minute * 30
)

func main() {
	logging.SetLevel(logging.INFO, "rbd-cleaner")
	log.Info("starting rbd-cleaner")
	log.Info("interval set to %s", interval)

	// first start
	err := rbdCleaner()
	if err != nil {
		log.Error(err.Error())
	}

	// and then for every other intervals
	for _ = range time.Tick(interval) {
		err := rbdCleaner()
		if err != nil {
			log.Error(err.Error())
		}
	}
}

func rbdCleaner() error {
	s, err := rbdList()
	if err != nil {
		return err
	}

	log.Info("cleaner started. total images in rbd pool: %d", s.Size())
	cleanCount := 0
	errCount := 0
	s.Each(func(rbdName interface{}) bool {
		vmName := rbdName.(string)
		vmId := strings.TrimPrefix(vmName, "vm-")
		vm := new(models.VM)

		err := mongodb.One("jVMs", vmId, vm)
		if err == nil {
			return true //  rbd image does exists in jVMS, don't touch it and continue
		}

		log.Info("Removing image '%s' from rbd", vmName)
		res, err := removeRBDImage(vmName)
		if err != nil {
			errCount++
			log.Error("err removing '%s'\n", err)
			return true
		}

		cleanCount++
		log.Info("%s", res)
		return true
	})
	log.Info("cleaner ended. images cleaned: %d images with err: %d", cleanCount, errCount)

	return nil
}

func removeRBDImage(imageName string) (string, error) {
	out, err := exec.Command("/usr/bin/rbd", "rm", "vms/"+imageName).CombinedOutput()
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(out)), nil
}

func rbdList() (*set.Set, error) {
	s := set.New()
	out, err := exec.Command("/usr/bin/rbd", "ls", "vms").CombinedOutput()
	if err != nil {
		return nil, err
	}

	rbds := string(bytes.TrimSpace(out))
	for _, r := range strings.Split(rbds, "\n") {
		s.Add(r)
	}
	return s, nil
}
<|endoftext|>


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

package corridor

import (
	"math/rand"
	"time"

	"github.com/gonum/matrix/mat64"
)

func Mvnrnd(mu, sigma *mat64.Dense) (out *mat64.Dense) {

	// initialize vector slices
	o := make([]float64, 2)
	n := make([]float64, 2)

	// generate random numbers from normal distribution
	rand.Seed(time.Now().Unix())
	for i := 0; i <= 1; i++ {
		n[i] = rand.NormFloat64()
	}

	// convert to matrix type
	rnd := mat64.NewDense(1, 2, n)
	output := mat64.NewDense(1, 2, o)

	// perform cholesky decomposition on covariance matrix
	cholFactor := mat64.Cholesky(sigma)
	lower := cholFactor.L

	// compute output
	output.Mul(rnd, lower)
	output.Add(output, mu)

	//return final output
	return output
}

func Fixrnd(rnd *mat64.Dense) (out *mat64.Dense) {

	// initialize vector slice
	o := make([]float64, 2)

	// write up down movement direction
	if rnd.At(0, 0) > 1 {
		o[0] = 1
	} else if rnd.At(0, 0) >= -1 && rnd.At(0, 0) <= 1 {
		o[0] = 0
	} else if rnd.At(0, 0) < -1 {
		o[0] = -1
	}

	// write left right movement direction
	if rnd.At(0, 1) > 1 {
		o[1] = 1
	} else if rnd.At(0, 1) >= -1 && rnd.At(0, 0) <= 1 {
		o[1] = 0
	} else if rnd.At(0, 1) < -1 {
		o[1] = -1
	}

	// convert to matrix type
	output := mat64.NewDense(1, 2, o)

	// return final output
	return output
}
<commit_msg>Update random routines to include multivariate random number generation.<commit_after>// Copyright ©2015 The corridor Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package corridor

import (
	"math/rand"
	"time"

	"github.com/gonum/matrix/mat64"
)

func Mvnrnd(mu, sigma *mat64.Dense) (rndsmp *mat64.Dense) {

	// initialize vector slices
	o := make([]float64, 2)
	n := make([]float64, 2)

	// generate random numbers from normal distribution
	rand.Seed(time.Now().Unix())
	for i := 0; i <= 1; i++ {
		n[i] = rand.NormFloat64()
	}

	// convert to matrix type
	rnd := mat64.NewDense(1, 2, n)
	output := mat64.NewDense(1, 2, o)

	// perform cholesky decomposition on covariance matrix
	cholFactor := mat64.Cholesky(sigma)
	lower := cholFactor.L

	// compute output
	output.Mul(rnd, lower)
	output.Add(output, mu)

	//return final output
	return output
}

func Fixrnd(rndsmp *mat64.Dense) (fixsmp *mat64.Dense) {

	// initialize vector slice
	o := make([]float64, 2)

	// write up down movement direction
	if rndsmp.At(0, 0) > 1 {
		o[0] = 1
	} else if rndsmp.At(0, 0) >= -1 && rndsmp.At(0, 0) <= 1 {
		o[0] = 0
	} else if rndsmp.At(0, 0) < -1 {
		o[0] = -1
	}

	// write left right movement direction
	if rndsmp.At(0, 1) > 1 {
		o[1] = 1
	} else if rndsmp.At(0, 1) >= -1 && rndsmp.At(0, 0) <= 1 {
		o[1] = 0
	} else if rndsmp.At(0, 1) < -1 {
		o[1] = -1
	}

	// convert to matrix type
	output := mat64.NewDense(1, 2, o)

	// return final output
	return output
}
<|endoftext|>


<commit_before>package queue

import (
	"sync/atomic"

	log "github.com/Cepave/open-falcon-backend/common/logruslog"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/model"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/rdb"
)

var logger = log.NewDefaultLogger("INFO")

type Queue struct {
	q       *commonQueue.Queue
	c       *commonQueue.Config
	cnt     uint64 // counter for the dequeued elements
	running bool
	flush   chan struct{}
	done    chan struct{}
}

func New(c *commonQueue.Config) *Queue {
	return &Queue{
		q:     commonQueue.New(),
		c:     c,
		flush: make(chan struct{}),
		done:  make(chan struct{}),
	}
}

func (q *Queue) Count() uint64 {
	return q.cnt
}

func (q *Queue) Start() {
	if q.running {
		return
	}
	q.running = true
	go q.drain()
}

func (q *Queue) drain() {
	for {
		select {
		default:
			reqs := q.q.DrainNWithDurationByType(q.c, new(model.NqmAgentHeartbeatRequest)).([]*model.NqmAgentHeartbeatRequest)
			d := uint64(len(reqs))
			atomic.AddUint64(&q.cnt, d)
			logger.Debugf("drained %d NQM agent heartbeat requests from queue\n", d)

			rdb.UpdateNqmAgentHeartbeat(reqs)
		case <-q.flush:
			c := &commonQueue.Config{Num: q.c.Num, Dur: 0}
			for {
				reqs := q.q.DrainNWithDurationByType(c, new(model.NqmAgentHeartbeatRequest)).([]*model.NqmAgentHeartbeatRequest)
				if len(reqs) == 0 {
					close(q.done)
					return
				}
				d := uint64(len(reqs))
				atomic.AddUint64(&q.cnt, d)
				logger.Debugf("flushed %d NQM agent heartbeat requests from queue\n", d)

				rdb.UpdateNqmAgentHeartbeat(reqs)
			}
		}
	}
}

func (q *Queue) Stop() {
	if !q.running {
		return
	}
	q.running = false
	close(q.flush)
	<-q.done
}

func (q *Queue) Put(v interface{}) {
	if q.running {
		q.q.Enqueue(v)
	}
	return
}

func (q *Queue) Len() int {
	return q.q.Len()
}

var NqmQueue *Queue

func InitNqmHeartbeat(c *commonQueue.Config) {
	NqmQueue = New(c)
	NqmQueue.Start()
}
<commit_msg>[OWL-1730][nqm-mng] Unify the coding style<commit_after>package queue

import (
	"sync/atomic"

	log "github.com/Cepave/open-falcon-backend/common/logruslog"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/model"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/rdb"
)

var logger = log.NewDefaultLogger("INFO")

type Queue struct {
	q       *commonQueue.Queue
	c       *commonQueue.Config
	cnt     uint64 // counter for the dequeued elements
	running bool
	flush   chan struct{}
	done    chan struct{}
}

func New(c *commonQueue.Config) *Queue {
	return &Queue{
		q:     commonQueue.New(),
		c:     c,
		flush: make(chan struct{}),
		done:  make(chan struct{}),
	}
}

func (q *Queue) Count() uint64 {
	return q.cnt
}

func (q *Queue) Start() {
	if q.running {
		return
	}
	q.running = true
	go q.drain()
}

func (q *Queue) drain() {
	for {
		select {
		default:
			reqs := q.q.DrainNWithDurationByType(q.c, new(model.NqmAgentHeartbeatRequest)).([]*model.NqmAgentHeartbeatRequest)
			d := uint64(len(reqs))
			atomic.AddUint64(&q.cnt, d)
			logger.Debugf("drained %d NQM agent heartbeat requests from queue\n", d)

			rdb.UpdateNqmAgentHeartbeat(reqs)
		case <-q.flush:
			c := &commonQueue.Config{Num: q.c.Num, Dur: 0}
			for {
				reqs := q.q.DrainNWithDurationByType(c, new(model.NqmAgentHeartbeatRequest)).([]*model.NqmAgentHeartbeatRequest)
				if len(reqs) == 0 {
					close(q.done)
					return
				}
				d := uint64(len(reqs))
				atomic.AddUint64(&q.cnt, d)
				logger.Debugf("flushed %d NQM agent heartbeat requests from queue\n", d)

				rdb.UpdateNqmAgentHeartbeat(reqs)
			}
		}
	}
}

func (q *Queue) Stop() {
	if !q.running {
		return
	}
	q.running = false
	close(q.flush)
	<-q.done
}

func (q *Queue) Put(v interface{}) {
	if !q.running {
		return
	}
	q.q.Enqueue(v)
}

func (q *Queue) Len() int {
	return q.q.Len()
}

var NqmQueue *Queue

func InitNqmHeartbeat(c *commonQueue.Config) {
	NqmQueue = New(c)
	NqmQueue.Start()
}
<|endoftext|>


<commit_before>package postgresrunner

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"syscall"

	"code.cloudfoundry.org/lager/lagertest"

	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/db/encryption"
	"github.com/concourse/concourse/atc/db/migration"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit/ginkgomon"
)

type Runner struct {
	Port int
}

func (runner Runner) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	defer ginkgo.GinkgoRecover()

	pgBase := filepath.Join(os.TempDir(), "concourse-pg-runner")

	err := os.MkdirAll(pgBase, 0755)
	Expect(err).NotTo(HaveOccurred())

	tmpdir, err := ioutil.TempDir(pgBase, "postgres")
	Expect(err).NotTo(HaveOccurred())

	currentUser, err := user.Current()
	Expect(err).NotTo(HaveOccurred())

	initdbPath, err := exec.LookPath("initdb")
	Expect(err).NotTo(HaveOccurred())

	postgresPath, err := exec.LookPath("postgres")
	Expect(err).NotTo(HaveOccurred())

	initCmd := exec.Command(initdbPath, "-U", "postgres", "-D", tmpdir, "-E", "UTF8", "--no-local")
	startCmd := exec.Command(postgresPath, "-k", "/tmp", "-D", tmpdir, "-h", "127.0.0.1", "-p", strconv.Itoa(runner.Port))

	if currentUser.Uid == "0" {
		pgUser, err := user.Lookup("postgres")
		Expect(err).NotTo(HaveOccurred())

		var uid, gid uint32
		_, err = fmt.Sscanf(pgUser.Uid, "%d", &uid)
		Expect(err).NotTo(HaveOccurred())

		_, err = fmt.Sscanf(pgUser.Gid, "%d", &gid)
		Expect(err).NotTo(HaveOccurred())

		err = os.Chown(tmpdir, int(uid), int(gid))
		Expect(err).NotTo(HaveOccurred())

		credential := &syscall.Credential{Uid: uid, Gid: gid}

		initCmd.SysProcAttr = &syscall.SysProcAttr{}
		initCmd.SysProcAttr.Credential = credential

		startCmd.SysProcAttr = &syscall.SysProcAttr{}
		startCmd.SysProcAttr.Credential = credential
	}

	session, err := gexec.Start(
		initCmd,
		gexec.NewPrefixedWriter("[o][initdb] ", ginkgo.GinkgoWriter),
		gexec.NewPrefixedWriter("[e][initdb] ", ginkgo.GinkgoWriter),
	)
	Expect(err).NotTo(HaveOccurred())

	<-session.Exited

	Expect(session).To(gexec.Exit(0))

	ginkgoRunner := &ginkgomon.Runner{
		Name:          "postgres",
		Command:       startCmd,
		AnsiColorCode: "90m",
		StartCheck:    "database system is ready to accept connections",
		Cleanup: func() {
			os.RemoveAll(tmpdir)
		},
	}

	return ginkgoRunner.Run(signals, ready)
}

func (runner *Runner) MigrateToVersion(version int) {
	err := migration.NewOpenHelper(
		"postgres",
		runner.DataSourceName(),
		nil,
		encryption.NewNoEncryption(),
	).MigrateToVersion(version)
	Expect(err).NotTo(HaveOccurred())
}

func (runner *Runner) TryOpenDBAtVersion(version int) (*sql.DB, error) {
	dbConn, err := migration.NewOpenHelper(
		"postgres",
		runner.DataSourceName(),
		nil,
		encryption.NewNoEncryption(),
	).OpenAtVersion(version)

	if err != nil {
		return nil, err
	}

	// only allow one connection so that we can detect any code paths that
	// require more than one, which will deadlock if it's at the limit
	dbConn.SetMaxOpenConns(1)

	return dbConn, nil
}

func (runner *Runner) OpenDBAtVersion(version int) *sql.DB {
	dbConn, err := runner.TryOpenDBAtVersion(version)
	Expect(err).NotTo(HaveOccurred())
	return dbConn
}

func (runner *Runner) OpenDB() *sql.DB {
	dbConn, err := migration.NewOpenHelper(
		"postgres",
		runner.DataSourceName(),
		nil,
		encryption.NewNoEncryption(),
	).Open()
	Expect(err).NotTo(HaveOccurred())

	// only allow one connection so that we can detect any code paths that
	// require more than one, which will deadlock if it's at the limit
	dbConn.SetMaxOpenConns(1)

	return dbConn
}

func (runner *Runner) OpenConn() db.Conn {
	dbConn, err := db.Open(
		lagertest.NewTestLogger("postgres-runner"),
		"postgres",
		runner.DataSourceName(),
		nil,
		nil,
		"postgresrunner",
		nil,
	)
	Expect(err).NotTo(HaveOccurred())

	// only allow one connection so that we can detect any code paths that
	// require more than one, which will deadlock if it's at the limit
	dbConn.SetMaxOpenConns(1)

	return dbConn
}

func (runner *Runner) OpenSingleton() *sql.DB {
	dbConn, err := sql.Open("postgres", runner.DataSourceName())
	Expect(err).NotTo(HaveOccurred())

	// only allow one connection, period. this matches production code use case,
	// as this is used for advisory locks.
	dbConn.SetMaxIdleConns(1)
	dbConn.SetMaxOpenConns(1)
	dbConn.SetConnMaxLifetime(0)

	return dbConn
}

func (runner *Runner) DataSourceName() string {
	return fmt.Sprintf("host=/tmp user=postgres dbname=testdb sslmode=disable port=%d", runner.Port)
}

func (runner *Runner) CreateTestDB() {
	createdb := exec.Command("createdb", "-h", "/tmp", "-U", "postgres", "-p", strconv.Itoa(runner.Port), "testdb")

	createS, err := gexec.Start(createdb, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	<-createS.Exited

	if createS.ExitCode() != 0 {
		runner.DropTestDB()

		createdb := exec.Command("createdb", "-h", "/tmp", "-U", "postgres", "-p", strconv.Itoa(runner.Port), "testdb")
		createS, err = gexec.Start(createdb, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())
	}

	<-createS.Exited

	Expect(createS).To(gexec.Exit(0))
}

func (runner *Runner) DropTestDB() {
	dropdb := exec.Command("dropdb", "-h", "/tmp", "-U", "postgres", "-p", strconv.Itoa(runner.Port), "testdb")
	dropS, err := gexec.Start(dropdb, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	<-dropS.Exited

	Expect(dropS).To(gexec.Exit(0))
}

func (runner *Runner) Truncate() {
	truncate := exec.Command(
		"psql",
		"-U", "postgres",
		"-p", strconv.Itoa(runner.Port),
		"testdb",
		"-c", `
			SET client_min_messages TO WARNING;

			CREATE OR REPLACE FUNCTION truncate_tables() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT tablename FROM pg_tables
							WHERE schemaname = 'public' AND tablename != 'migrations_history';
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'TRUNCATE TABLE ' || quote_ident(stmt.tablename) || ' RESTART IDENTITY CASCADE;';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			CREATE OR REPLACE FUNCTION drop_ephemeral_sequences() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname LIKE 'build_event_id_seq_%';
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'DROP SEQUENCE ' || quote_ident(stmt.relname) || ';';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			CREATE OR REPLACE FUNCTION drop_ephemeral_tables() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname LIKE 'pipeline_build_events_%'
							AND relkind = 'r';
					team_statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname LIKE 'team_build_events_%'
							AND relkind = 'r';
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'DROP TABLE ' || quote_ident(stmt.relname) || ';';
					END LOOP;
					FOR stmt IN team_statements LOOP
							EXECUTE 'DROP TABLE ' || quote_ident(stmt.relname) || ';';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			CREATE OR REPLACE FUNCTION reset_global_sequences() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname IN ('one_off_name', 'config_version_seq');
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'ALTER SEQUENCE ' || quote_ident(stmt.relname) || ' RESTART WITH 1;';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			SELECT truncate_tables();
			SELECT drop_ephemeral_sequences();
			SELECT drop_ephemeral_tables();
			SELECT reset_global_sequences();
		`,
	)

	truncateS, err := gexec.Start(truncate, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	<-truncateS.Exited

	Expect(truncateS).To(gexec.Exit(0))
}
<commit_msg>use -h /tmp on Truncate also<commit_after>package postgresrunner

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"syscall"

	"code.cloudfoundry.org/lager/lagertest"

	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/db/encryption"
	"github.com/concourse/concourse/atc/db/migration"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/tedsuo/ifrit/ginkgomon"
)

type Runner struct {
	Port int
}

func (runner Runner) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	defer ginkgo.GinkgoRecover()

	pgBase := filepath.Join(os.TempDir(), "concourse-pg-runner")

	err := os.MkdirAll(pgBase, 0755)
	Expect(err).NotTo(HaveOccurred())

	tmpdir, err := ioutil.TempDir(pgBase, "postgres")
	Expect(err).NotTo(HaveOccurred())

	currentUser, err := user.Current()
	Expect(err).NotTo(HaveOccurred())

	initdbPath, err := exec.LookPath("initdb")
	Expect(err).NotTo(HaveOccurred())

	postgresPath, err := exec.LookPath("postgres")
	Expect(err).NotTo(HaveOccurred())

	initCmd := exec.Command(initdbPath, "-U", "postgres", "-D", tmpdir, "-E", "UTF8", "--no-local")
	startCmd := exec.Command(postgresPath, "-k", "/tmp", "-D", tmpdir, "-h", "127.0.0.1", "-p", strconv.Itoa(runner.Port))

	if currentUser.Uid == "0" {
		pgUser, err := user.Lookup("postgres")
		Expect(err).NotTo(HaveOccurred())

		var uid, gid uint32
		_, err = fmt.Sscanf(pgUser.Uid, "%d", &uid)
		Expect(err).NotTo(HaveOccurred())

		_, err = fmt.Sscanf(pgUser.Gid, "%d", &gid)
		Expect(err).NotTo(HaveOccurred())

		err = os.Chown(tmpdir, int(uid), int(gid))
		Expect(err).NotTo(HaveOccurred())

		credential := &syscall.Credential{Uid: uid, Gid: gid}

		initCmd.SysProcAttr = &syscall.SysProcAttr{}
		initCmd.SysProcAttr.Credential = credential

		startCmd.SysProcAttr = &syscall.SysProcAttr{}
		startCmd.SysProcAttr.Credential = credential
	}

	session, err := gexec.Start(
		initCmd,
		gexec.NewPrefixedWriter("[o][initdb] ", ginkgo.GinkgoWriter),
		gexec.NewPrefixedWriter("[e][initdb] ", ginkgo.GinkgoWriter),
	)
	Expect(err).NotTo(HaveOccurred())

	<-session.Exited

	Expect(session).To(gexec.Exit(0))

	ginkgoRunner := &ginkgomon.Runner{
		Name:          "postgres",
		Command:       startCmd,
		AnsiColorCode: "90m",
		StartCheck:    "database system is ready to accept connections",
		Cleanup: func() {
			os.RemoveAll(tmpdir)
		},
	}

	return ginkgoRunner.Run(signals, ready)
}

func (runner *Runner) MigrateToVersion(version int) {
	err := migration.NewOpenHelper(
		"postgres",
		runner.DataSourceName(),
		nil,
		encryption.NewNoEncryption(),
	).MigrateToVersion(version)
	Expect(err).NotTo(HaveOccurred())
}

func (runner *Runner) TryOpenDBAtVersion(version int) (*sql.DB, error) {
	dbConn, err := migration.NewOpenHelper(
		"postgres",
		runner.DataSourceName(),
		nil,
		encryption.NewNoEncryption(),
	).OpenAtVersion(version)

	if err != nil {
		return nil, err
	}

	// only allow one connection so that we can detect any code paths that
	// require more than one, which will deadlock if it's at the limit
	dbConn.SetMaxOpenConns(1)

	return dbConn, nil
}

func (runner *Runner) OpenDBAtVersion(version int) *sql.DB {
	dbConn, err := runner.TryOpenDBAtVersion(version)
	Expect(err).NotTo(HaveOccurred())
	return dbConn
}

func (runner *Runner) OpenDB() *sql.DB {
	dbConn, err := migration.NewOpenHelper(
		"postgres",
		runner.DataSourceName(),
		nil,
		encryption.NewNoEncryption(),
	).Open()
	Expect(err).NotTo(HaveOccurred())

	// only allow one connection so that we can detect any code paths that
	// require more than one, which will deadlock if it's at the limit
	dbConn.SetMaxOpenConns(1)

	return dbConn
}

func (runner *Runner) OpenConn() db.Conn {
	dbConn, err := db.Open(
		lagertest.NewTestLogger("postgres-runner"),
		"postgres",
		runner.DataSourceName(),
		nil,
		nil,
		"postgresrunner",
		nil,
	)
	Expect(err).NotTo(HaveOccurred())

	// only allow one connection so that we can detect any code paths that
	// require more than one, which will deadlock if it's at the limit
	dbConn.SetMaxOpenConns(1)

	return dbConn
}

func (runner *Runner) OpenSingleton() *sql.DB {
	dbConn, err := sql.Open("postgres", runner.DataSourceName())
	Expect(err).NotTo(HaveOccurred())

	// only allow one connection, period. this matches production code use case,
	// as this is used for advisory locks.
	dbConn.SetMaxIdleConns(1)
	dbConn.SetMaxOpenConns(1)
	dbConn.SetConnMaxLifetime(0)

	return dbConn
}

func (runner *Runner) DataSourceName() string {
	return fmt.Sprintf("host=/tmp user=postgres dbname=testdb sslmode=disable port=%d", runner.Port)
}

func (runner *Runner) CreateTestDB() {
	createdb := exec.Command("createdb", "-h", "/tmp", "-U", "postgres", "-p", strconv.Itoa(runner.Port), "testdb")

	createS, err := gexec.Start(createdb, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	<-createS.Exited

	if createS.ExitCode() != 0 {
		runner.DropTestDB()

		createdb := exec.Command("createdb", "-h", "/tmp", "-U", "postgres", "-p", strconv.Itoa(runner.Port), "testdb")
		createS, err = gexec.Start(createdb, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())
	}

	<-createS.Exited

	Expect(createS).To(gexec.Exit(0))
}

func (runner *Runner) DropTestDB() {
	dropdb := exec.Command("dropdb", "-h", "/tmp", "-U", "postgres", "-p", strconv.Itoa(runner.Port), "testdb")
	dropS, err := gexec.Start(dropdb, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	<-dropS.Exited

	Expect(dropS).To(gexec.Exit(0))
}

func (runner *Runner) Truncate() {
	truncate := exec.Command(
		"psql",
		"-h", "/tmp",
		"-U", "postgres",
		"-p", strconv.Itoa(runner.Port),
		"testdb",
		"-c", `
			SET client_min_messages TO WARNING;

			CREATE OR REPLACE FUNCTION truncate_tables() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT tablename FROM pg_tables
							WHERE schemaname = 'public' AND tablename != 'migrations_history';
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'TRUNCATE TABLE ' || quote_ident(stmt.tablename) || ' RESTART IDENTITY CASCADE;';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			CREATE OR REPLACE FUNCTION drop_ephemeral_sequences() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname LIKE 'build_event_id_seq_%';
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'DROP SEQUENCE ' || quote_ident(stmt.relname) || ';';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			CREATE OR REPLACE FUNCTION drop_ephemeral_tables() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname LIKE 'pipeline_build_events_%'
							AND relkind = 'r';
					team_statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname LIKE 'team_build_events_%'
							AND relkind = 'r';
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'DROP TABLE ' || quote_ident(stmt.relname) || ';';
					END LOOP;
					FOR stmt IN team_statements LOOP
							EXECUTE 'DROP TABLE ' || quote_ident(stmt.relname) || ';';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			CREATE OR REPLACE FUNCTION reset_global_sequences() RETURNS void AS $$
			DECLARE
					statements CURSOR FOR
							SELECT relname FROM pg_class
							WHERE relname IN ('one_off_name', 'config_version_seq');
			BEGIN
					FOR stmt IN statements LOOP
							EXECUTE 'ALTER SEQUENCE ' || quote_ident(stmt.relname) || ' RESTART WITH 1;';
					END LOOP;
			END;
			$$ LANGUAGE plpgsql;

			SELECT truncate_tables();
			SELECT drop_ephemeral_sequences();
			SELECT drop_ephemeral_tables();
			SELECT reset_global_sequences();
		`,
	)

	truncateS, err := gexec.Start(truncate, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	<-truncateS.Exited

	Expect(truncateS).To(gexec.Exit(0))
}
<|endoftext|>


<commit_before>package emitter_test

import (
	"net"
	"sync"

	"github.com/cloudfoundry/dropsonde/control"
	"github.com/cloudfoundry/dropsonde/emitter"
	"github.com/cloudfoundry/dropsonde/factories"
	"github.com/gogo/protobuf/proto"
	uuid "github.com/nu7hatch/gouuid"

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

var _ = Describe("UdpEmitter", func() {
	var testData = []byte("hello")

	Describe("Close()", func() {
		It("closes the UDP connection", func() {

			udpEmitter, _ := emitter.NewUdpEmitter("localhost:42420")

			udpEmitter.Close()

			err := udpEmitter.Emit(testData)
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(Equal("use of closed network connection"))
		})
	})

	Describe("Emit()", func() {
		var udpEmitter emitter.ByteEmitter

		Context("when the agent is listening", func() {

			var agentListener net.PacketConn

			BeforeEach(func() {
				var err error
				agentListener, err = net.ListenPacket("udp4", "")
				Expect(err).ToNot(HaveOccurred())

				udpEmitter, err = emitter.NewUdpEmitter(agentListener.LocalAddr().String())
				Expect(err).ToNot(HaveOccurred())
			})

			AfterEach(func() {
				agentListener.Close()
			})

			It("should send the data", func(done Done) {
				err := udpEmitter.Emit(testData)
				Expect(err).ToNot(HaveOccurred())

				buffer := make([]byte, 4096)
				readCount, _, err := agentListener.ReadFrom(buffer)
				Expect(err).ToNot(HaveOccurred())
				Expect(buffer[:readCount]).To(Equal(testData))

				close(done)
			})
		})

		Context("when the agent is not listening", func() {
			BeforeEach(func() {
				udpEmitter, _ = emitter.NewUdpEmitter("localhost:12345")
			})

			It("should attempt to send the data", func() {
				err := udpEmitter.Emit(testData)
				Expect(err).ToNot(HaveOccurred())
			})

			Context("then the agent starts Listening", func() {
				It("should eventually send data", func(done Done) {
					err := udpEmitter.Emit(testData)
					Expect(err).ToNot(HaveOccurred())

					agentListener, err := net.ListenPacket("udp4", ":12345")
					Expect(err).ToNot(HaveOccurred())

					err = udpEmitter.Emit(testData)
					Expect(err).ToNot(HaveOccurred())

					buffer := make([]byte, 4096)
					readCount, _, err := agentListener.ReadFrom(buffer)
					Expect(err).ToNot(HaveOccurred())
					Expect(buffer[:readCount]).To(Equal(testData))

					close(done)
				})
			})
		})
	})

	Describe("NewUdpEmitter()", func() {
		Context("when ResolveUDPAddr fails", func() {
			It("returns an error", func() {
				emitter, err := emitter.NewUdpEmitter("invalid-address:")
				Expect(emitter).To(BeNil())
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when all is good", func() {
			It("creates an emitter", func() {
				emitter, err := emitter.NewUdpEmitter("localhost:123")
				Expect(emitter).ToNot(BeNil())
				Expect(err).ToNot(HaveOccurred())
			})
		})
	})

	Describe("ListenForHeartbeatRequest", func() {
		var timesCalled int
		var lastControlMessage *control.ControlMessage
		var lock sync.Mutex

		var fakeResponder = func(controlMessage *control.ControlMessage) {
			lock.Lock()
			defer lock.Unlock()
			timesCalled++
			lastControlMessage = controlMessage
		}

		var getTimesCalled = func() int {
			lock.Lock()
			defer lock.Unlock()
			return timesCalled
		}

		var getReceivedHeartbeatRequest = func() *control.ControlMessage {
			lock.Lock()
			defer lock.Unlock()

			return lastControlMessage
		}

		BeforeEach(func() {
			lock = sync.Mutex{}
			timesCalled = 0
		})

		It("calls responder with the correct heartbeat request when when heartbeat is requested", func() {
			emitter, _ := emitter.NewUdpEmitter("localhost:123")
			go emitter.ListenForHeartbeatRequest(fakeResponder)

			Expect(timesCalled).To(BeZero())

			heartbeatRequest := newHeartbeatRequest()

			sendHeartbeatRequest(emitter.Address(), heartbeatRequest)
			Eventually(getTimesCalled).Should(Equal(1))
			Expect(getReceivedHeartbeatRequest()).To(Equal(heartbeatRequest))
		})

		It("responds to multiple heartbeat requests", func() {
			emitter, _ := emitter.NewUdpEmitter("localhost:123")
			go emitter.ListenForHeartbeatRequest(fakeResponder)
			sendHeartbeatRequest(emitter.Address(), newHeartbeatRequest())
			sendHeartbeatRequest(emitter.Address(), newHeartbeatRequest())

			Eventually(getTimesCalled).Should(Equal(2))
		})

		It("returns an error if listening on the UDP port fails", func() {
			emitter, _ := emitter.NewUdpEmitter("localhost:123")
			emitter.Close()
			err := emitter.ListenForHeartbeatRequest(fakeResponder)
			Expect(err).To(HaveOccurred())
			Expect(err).To(MatchError("use of closed network connection"))
		})
	})
})

func sendHeartbeatRequest(addr net.Addr, message *control.ControlMessage) {
	encodedMessage, _ := proto.Marshal(message)
	conn, _ := net.ListenPacket("udp4", "")
	conn.WriteTo(encodedMessage, addr)
}

func newHeartbeatRequest() *control.ControlMessage {
	id, _ := uuid.NewV4()

	return &control.ControlMessage{
		Origin:      proto.String("test"),
		Identifier:  factories.NewControlUUID(id),
		Timestamp:   proto.Int64(0),
		ControlType: control.ControlMessage_HeartbeatRequest.Enum(),
	}
}
<commit_msg>Fix udp_emitter_test to work with Go tip (1.5)<commit_after>package emitter_test

import (
	"net"
	"sync"

	"github.com/cloudfoundry/dropsonde/control"
	"github.com/cloudfoundry/dropsonde/emitter"
	"github.com/cloudfoundry/dropsonde/factories"
	"github.com/gogo/protobuf/proto"
	uuid "github.com/nu7hatch/gouuid"

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

var _ = Describe("UdpEmitter", func() {
	var testData = []byte("hello")

	Describe("Close()", func() {
		It("closes the UDP connection", func() {

			udpEmitter, _ := emitter.NewUdpEmitter("localhost:42420")

			udpEmitter.Close()

			err := udpEmitter.Emit(testData)
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("use of closed network connection"))
		})
	})

	Describe("Emit()", func() {
		var udpEmitter emitter.ByteEmitter

		Context("when the agent is listening", func() {

			var agentListener net.PacketConn

			BeforeEach(func() {
				var err error
				agentListener, err = net.ListenPacket("udp4", "")
				Expect(err).ToNot(HaveOccurred())

				udpEmitter, err = emitter.NewUdpEmitter(agentListener.LocalAddr().String())
				Expect(err).ToNot(HaveOccurred())
			})

			AfterEach(func() {
				agentListener.Close()
			})

			It("should send the data", func(done Done) {
				err := udpEmitter.Emit(testData)
				Expect(err).ToNot(HaveOccurred())

				buffer := make([]byte, 4096)
				readCount, _, err := agentListener.ReadFrom(buffer)
				Expect(err).ToNot(HaveOccurred())
				Expect(buffer[:readCount]).To(Equal(testData))

				close(done)
			})
		})

		Context("when the agent is not listening", func() {
			BeforeEach(func() {
				udpEmitter, _ = emitter.NewUdpEmitter("localhost:12345")
			})

			It("should attempt to send the data", func() {
				err := udpEmitter.Emit(testData)
				Expect(err).ToNot(HaveOccurred())
			})

			Context("then the agent starts Listening", func() {
				It("should eventually send data", func(done Done) {
					err := udpEmitter.Emit(testData)
					Expect(err).ToNot(HaveOccurred())

					agentListener, err := net.ListenPacket("udp4", ":12345")
					Expect(err).ToNot(HaveOccurred())

					err = udpEmitter.Emit(testData)
					Expect(err).ToNot(HaveOccurred())

					buffer := make([]byte, 4096)
					readCount, _, err := agentListener.ReadFrom(buffer)
					Expect(err).ToNot(HaveOccurred())
					Expect(buffer[:readCount]).To(Equal(testData))

					close(done)
				})
			})
		})
	})

	Describe("NewUdpEmitter()", func() {
		Context("when ResolveUDPAddr fails", func() {
			It("returns an error", func() {
				emitter, err := emitter.NewUdpEmitter("invalid-address:")
				Expect(emitter).To(BeNil())
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when all is good", func() {
			It("creates an emitter", func() {
				emitter, err := emitter.NewUdpEmitter("localhost:123")
				Expect(emitter).ToNot(BeNil())
				Expect(err).ToNot(HaveOccurred())
			})
		})
	})

	Describe("ListenForHeartbeatRequest", func() {
		var timesCalled int
		var lastControlMessage *control.ControlMessage
		var lock sync.Mutex

		var fakeResponder = func(controlMessage *control.ControlMessage) {
			lock.Lock()
			defer lock.Unlock()
			timesCalled++
			lastControlMessage = controlMessage
		}

		var getTimesCalled = func() int {
			lock.Lock()
			defer lock.Unlock()
			return timesCalled
		}

		var getReceivedHeartbeatRequest = func() *control.ControlMessage {
			lock.Lock()
			defer lock.Unlock()

			return lastControlMessage
		}

		BeforeEach(func() {
			lock = sync.Mutex{}
			timesCalled = 0
		})

		It("calls responder with the correct heartbeat request when when heartbeat is requested", func() {
			emitter, _ := emitter.NewUdpEmitter("localhost:123")
			go emitter.ListenForHeartbeatRequest(fakeResponder)

			Expect(timesCalled).To(BeZero())

			heartbeatRequest := newHeartbeatRequest()

			sendHeartbeatRequest(emitter.Address(), heartbeatRequest)
			Eventually(getTimesCalled).Should(Equal(1))
			Expect(getReceivedHeartbeatRequest()).To(Equal(heartbeatRequest))
		})

		It("responds to multiple heartbeat requests", func() {
			emitter, _ := emitter.NewUdpEmitter("localhost:123")
			go emitter.ListenForHeartbeatRequest(fakeResponder)
			sendHeartbeatRequest(emitter.Address(), newHeartbeatRequest())
			sendHeartbeatRequest(emitter.Address(), newHeartbeatRequest())

			Eventually(getTimesCalled).Should(Equal(2))
		})

		It("returns an error if listening on the UDP port fails", func() {
			emitter, _ := emitter.NewUdpEmitter("localhost:123")
			emitter.Close()
			err := emitter.ListenForHeartbeatRequest(fakeResponder)
			Expect(err).To(HaveOccurred())
			Expect(err).To(MatchError("use of closed network connection"))
		})
	})
})

func sendHeartbeatRequest(addr net.Addr, message *control.ControlMessage) {
	encodedMessage, _ := proto.Marshal(message)
	conn, _ := net.ListenPacket("udp4", "")
	conn.WriteTo(encodedMessage, addr)
}

func newHeartbeatRequest() *control.ControlMessage {
	id, _ := uuid.NewV4()

	return &control.ControlMessage{
		Origin:      proto.String("test"),
		Identifier:  factories.NewControlUUID(id),
		Timestamp:   proto.Int64(0),
		ControlType: control.ControlMessage_HeartbeatRequest.Enum(),
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

// Package tailer provides a class that is responsible for tailing log files
// and extracting new log lines to be passed into the virtual machines.
package tailer

// For regular files, mtail gets notified on modifications (i.e. appends) to
// log files that are being watched, in order to read the new lines. Log files
// can also be rotated, so mtail is also notified of creates in the log file
// directory.

import (
	"errors"
	"expvar"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"sync"
	"syscall"
	"unicode/utf8"

	"github.com/golang/glog"

	"github.com/google/mtail/watcher"

	"github.com/spf13/afero"
)

var (
	logCount     = expvar.NewInt("log_count")
	logErrors    = expvar.NewMap("log_errors_total")
	logRotations = expvar.NewMap("log_rotations_total")
)

// Tailer receives notification of changes from a Watcher and extracts new log
// lines from files. It also handles new log file creation events and log
// rotations.
type Tailer struct {
	w watcher.Watcher

	watched     map[string]struct{}   // Names of logs being watched.
	watchedLock sync.RWMutex          // protects `watched'
	lines       chan<- string         // Logfile lines being emitted.
	files       map[string]afero.File // File handles for each pathname.
	filesLock   sync.Mutex            // protects `files'
	partials    map[string]string     // Accumulator for the currently read line for each pathname.

	fs afero.Fs // mockable filesystem interface
}

// Options configures a Tailer
type Options struct {
	Lines chan<- string
	W     watcher.Watcher // Not required, will use watcher.LogWatcher if it is zero.
	FS    afero.Fs        // Not required, will use afero.OsFs if it is zero.
}

// New returns a new Tailer, configured with the supplied Options
func New(o Options) (*Tailer, error) {
	if o.Lines == nil {
		return nil, errors.New("tailer needs lines")
	}
	fs := o.FS
	if fs == nil {
		fs = &afero.OsFs{}
	}
	w := o.W
	if w == nil {
		var err error
		w, err = watcher.NewLogWatcher()
		if err != nil {
			return nil, fmt.Errorf("Couldn't create a watcher for tailer: %s", err)
		}
	}
	t := &Tailer{
		w:        w,
		watched:  make(map[string]struct{}),
		lines:    o.Lines,
		files:    make(map[string]afero.File),
		partials: make(map[string]string),
		fs:       fs,
	}
	go t.run()
	return t, nil
}

func (t *Tailer) addWatched(path string) {
	t.watchedLock.Lock()
	defer t.watchedLock.Unlock()
	t.watched[path] = struct{}{}
}

func (t *Tailer) isWatching(path string) bool {
	t.watchedLock.RLock()
	defer t.watchedLock.RUnlock()
	_, ok := t.watched[path]
	return ok
}

// Tail registers a file path to be tailed.
func (t *Tailer) Tail(pathname string) {
	fullpath, err := filepath.Abs(pathname)
	if err != nil {
		glog.Infof("Failed to find absolute path for %q: %s\n", pathname, err)
		return
	}
	if !t.isWatching(fullpath) {
		t.addWatched(fullpath)
		logCount.Add(1)
		t.openLogPath(fullpath, false)
	}
}

// TailFile registers a file descriptor to be tailed.
func (t *Tailer) TailFile(f afero.File) error {
	logCount.Add(1)
	return t.startNewFile(f, false)
}

// handleLogUpdate reads all available bytes from an already opened file
// identified by pathname, and sends them to be processed on the lines channel.
func (t *Tailer) handleLogUpdate(pathname string) {
	f, ok := t.files[pathname]
	if !ok {
		glog.Infof("No file found for %q", pathname)
		return
	}
	t.readFile(f)
}

func (t *Tailer) readFile(f afero.File) {
	for {
		partial, err := t.readPartial(f, t.partials[f.Name()])
		t.partials[f.Name()] = partial
		if err != nil {
			if err == io.EOF {
				// end of file for now, return
				return
			}
			glog.Infof("Failed to read updates from %q: %s", f.Name(), err)
			return
		}
	}
}

// readPartial reads blocks of 4096 bytes from the File, sending lines to the
// channel as it encounters newlines.  If EOF is encountered, the partial line
// is returned to be concatenated with on the next call.
func (t *Tailer) readPartial(f afero.File, partial string) (string, error) {
	b := make([]byte, 4096)
	n, err := f.Read(b)
	if err != nil {
		return partial, err // TODO(jaq): handle this path better
	}
	for i, width := 0, 0; i < len(b) && i < n; i += width {
		var rune rune
		rune, width = utf8.DecodeRune(b[i:])
		switch {
		case rune != '\n':
			partial += string(rune)
		default:
			// send off line for processing
			t.lines <- partial
			// reset accumulator
			partial = ""
		}
	}
	return partial, nil
}

// inode returns the inode number of a file, or 0 if the file has no underlying Sys implementation.
func inode(f os.FileInfo) uint64 {
	s := f.Sys()
	if s == nil {
		return 0
	}
	switch s := s.(type) {
	case *syscall.Stat_t:
		return s.Ino
	default:
		return 0
	}
}

// handleLogCreate handles both new and rotated log files.
func (t *Tailer) handleLogCreate(pathname string) {
	if !t.isWatching(pathname) {
		glog.V(1).Info("Not watching path %q, ignoring.\n", pathname)
		return
	}

	t.filesLock.Lock()
	fd, ok := t.files[pathname]
	t.filesLock.Unlock()
	if ok {
		s1, err := fd.Stat()
		if err != nil {
			glog.Infof("Stat failed on %q: %s", t.files[pathname].Name(), err)
			return
		}
		s2, err := t.fs.Stat(pathname)
		if err != nil {
			glog.Infof("Stat failed on %q: %s", pathname, err)
			return
		}
		if inode(s1) != inode(s2) {
			logRotations.Add(pathname, 1)
			// flush the old log, pathname is still an index into t.files with the old inode.
			t.handleLogUpdate(pathname)
			fd.Close()
			err := t.w.Remove(pathname)
			if err != nil {
				glog.Info("Failed removing watches on", pathname)
			}
			// Always seek to start on log rotation.
			glog.Infof("Seek to start on %s", pathname)
			t.openLogPath(pathname, true)
		} else {
			glog.Infof("Path %s already being watched, and inode not changed.",
				pathname)
		}
	} else {
		// Freshly opened log file, never seen before.
		t.openLogPath(pathname, true)
	}
}

// openLogPath opens a new log file at pathname, and optionally seeks to the
// start or end of the file. Rotated logs should read from the start, but logs
// opened for the first time read from the end.
func (t *Tailer) openLogPath(pathname string, seekStart bool) {
	if !t.isWatching(pathname) {
		glog.Infof("Not watching %q, ignoring.", pathname)
		return
	}

	d := path.Dir(pathname)
	if !t.isWatching(d) {
		err := t.w.Add(d)
		if err != nil {
			glog.Infof("Adding a create watch failed on %q: %s", d, err)
		}
		t.addWatched(d)
	}

	f, err := t.fs.Open(pathname)
	if err != nil {
		// Doesn't exist yet. We're watching the directory, so we'll pick it up
		// again on create; return successfully.
		if os.IsNotExist(err) {
			return
		}
		glog.Infof("Failed to open %q for reading: %s", pathname, err)
		logErrors.Add(pathname, 1)
		return
	}
	err = t.startNewFile(f, seekStart)
	if err != nil {
		glog.Info(err)
	}
}

func (t *Tailer) startNewFile(f afero.File, seekStart bool) error {
	fi, err := f.Stat()
	if err != nil {
		// Stat failed, log error and return.
		logErrors.Add(f.Name(), 1)
		return fmt.Errorf("Failed to stat %q: %s", f.Name(), err)
	}
	glog.Infof("Mode: %q", fi.Mode())
	switch m := fi.Mode(); {
	case m&os.ModeType == 0:
		if seekStart {
			f.Seek(0, os.SEEK_SET)
		} else {
			f.Seek(0, os.SEEK_END)
		}
		err = t.w.Add(f.Name())
		if err != nil {
			return fmt.Errorf("Adding a change watch failed on %q: %s", f.Name(), err)
		}
	case m&os.ModeType == os.ModeNamedPipe:
		return fmt.Errorf("named pipe")
	default:
		return fmt.Errorf("Can't open files with mode %v: %s", m&os.ModeType, f.Name())
	}
	t.filesLock.Lock()
	t.files[f.Name()] = f
	t.filesLock.Unlock()
	glog.Infof("Tailing %s", f.Name())

	// In case the new log has been written to already, attempt to read the first lines.
	t.readFile(f)
	return nil
}

// start is the main event loop for the Tailer.
// It receives notification of log file changes from the watcher channel, and
// handles them.
func (t *Tailer) run() {
	for e := range t.w.Events() {
		switch e := e.(type) {
		case watcher.UpdateEvent:
			t.handleLogUpdate(e.Pathname)
		case watcher.CreateEvent:
			t.handleLogCreate(e.Pathname)
		default:
			glog.Infof("Unexpected event %q", e)
		}
	}
	glog.Infof("Shutting down tailer.")
	close(t.lines)
}

// Close signals termination to the watcher.
func (t *Tailer) Close() {
	t.w.Close()
}
<commit_msg>Fork off a goroutine for reading from pipes.<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

// Package tailer provides a class that is responsible for tailing log files
// and extracting new log lines to be passed into the virtual machines.
package tailer

// For regular files, mtail gets notified on modifications (i.e. appends) to
// log files that are being watched, in order to read the new lines. Log files
// can also be rotated, so mtail is also notified of creates in the log file
// directory.

import (
	"errors"
	"expvar"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"sync"
	"syscall"
	"unicode/utf8"

	"github.com/golang/glog"

	"github.com/google/mtail/watcher"

	"github.com/spf13/afero"
)

var (
	logCount     = expvar.NewInt("log_count")
	logErrors    = expvar.NewMap("log_errors_total")
	logRotations = expvar.NewMap("log_rotations_total")
)

// Tailer receives notification of changes from a Watcher and extracts new log
// lines from files. It also handles new log file creation events and log
// rotations.
type Tailer struct {
	w watcher.Watcher

	pipeGroup sync.WaitGroup
	pipeQuit  chan struct{}

	watched     map[string]struct{}   // Names of logs being watched.
	watchedLock sync.RWMutex          // protects `watched'
	lines       chan<- string         // Logfile lines being emitted.
	files       map[string]afero.File // File handles for each pathname.
	filesLock   sync.Mutex            // protects `files'
	partials    map[string]string     // Accumulator for the currently read line for each pathname.

	fs afero.Fs // mockable filesystem interface
}

// Options configures a Tailer
type Options struct {
	Lines chan<- string
	W     watcher.Watcher // Not required, will use watcher.LogWatcher if it is zero.
	FS    afero.Fs        // Not required, will use afero.OsFs if it is zero.
}

// New returns a new Tailer, configured with the supplied Options
func New(o Options) (*Tailer, error) {
	if o.Lines == nil {
		return nil, errors.New("tailer needs lines")
	}
	fs := o.FS
	if fs == nil {
		fs = &afero.OsFs{}
	}
	w := o.W
	if w == nil {
		var err error
		w, err = watcher.NewLogWatcher()
		if err != nil {
			return nil, fmt.Errorf("Couldn't create a watcher for tailer: %s", err)
		}
	}
	t := &Tailer{
		w:        w,
		watched:  make(map[string]struct{}),
		lines:    o.Lines,
		files:    make(map[string]afero.File),
		partials: make(map[string]string),
		fs:       fs,
	}
	go t.run()
	return t, nil
}

func (t *Tailer) addWatched(path string) {
	t.watchedLock.Lock()
	defer t.watchedLock.Unlock()
	t.watched[path] = struct{}{}
}

func (t *Tailer) isWatching(path string) bool {
	t.watchedLock.RLock()
	defer t.watchedLock.RUnlock()
	_, ok := t.watched[path]
	return ok
}

// Tail registers a file path to be tailed.
func (t *Tailer) Tail(pathname string) {
	fullpath, err := filepath.Abs(pathname)
	if err != nil {
		glog.Infof("Failed to find absolute path for %q: %s\n", pathname, err)
		return
	}
	if !t.isWatching(fullpath) {
		t.addWatched(fullpath)
		logCount.Add(1)
		t.openLogPath(fullpath, false)
	}
}

// TailFile registers a file descriptor to be tailed.
func (t *Tailer) TailFile(f afero.File) error {
	logCount.Add(1)
	return t.startNewFile(f, false)
}

// handleLogUpdate reads all available bytes from an already opened file
// identified by pathname, and sends them to be processed on the lines channel.
func (t *Tailer) handleLogUpdate(pathname string) {
	f, ok := t.files[pathname]
	if !ok {
		glog.Infof("No file found for %q", pathname)
		return
	}
	var err error
	t.partials[pathname], err = t.readFile(f, t.partials[pathname])
	if err != nil {
		glog.Info(err)
	}
}

func (t *Tailer) readFile(f afero.File, partialIn string) (partialOut string, err error) {
	for {
		partialOut, err = t.readPartial(f, partialIn)
		if err != nil {
			if err == io.EOF {
				// end of file for now, return
				err = nil
				return
			}
			glog.Infof("Failed to read updates from %q: %s", f.Name(), err)
			return
		}
	}
}

// readPartial reads blocks of 4096 bytes from the File, sending lines to the
// channel as it encounters newlines.  If EOF is encountered, the partial line
// is returned to be concatenated with on the next call.
func (t *Tailer) readPartial(f afero.File, partial string) (string, error) {
	b := make([]byte, 4096)
	n, err := f.Read(b)
	if err != nil {
		return partial, err // TODO(jaq): handle this path better
	}
	for i, width := 0, 0; i < len(b) && i < n; i += width {
		var rune rune
		rune, width = utf8.DecodeRune(b[i:])
		switch {
		case rune != '\n':
			partial += string(rune)
		default:
			// send off line for processing
			t.lines <- partial
			// reset accumulator
			partial = ""
		}
	}
	return partial, nil
}

// inode returns the inode number of a file, or 0 if the file has no underlying Sys implementation.
func inode(f os.FileInfo) uint64 {
	s := f.Sys()
	if s == nil {
		return 0
	}
	switch s := s.(type) {
	case *syscall.Stat_t:
		return s.Ino
	default:
		return 0
	}
}

// handleLogCreate handles both new and rotated log files.
func (t *Tailer) handleLogCreate(pathname string) {
	if !t.isWatching(pathname) {
		glog.V(1).Info("Not watching path %q, ignoring.\n", pathname)
		return
	}

	t.filesLock.Lock()
	fd, ok := t.files[pathname]
	t.filesLock.Unlock()
	if ok {
		s1, err := fd.Stat()
		if err != nil {
			glog.Infof("Stat failed on %q: %s", t.files[pathname].Name(), err)
			return
		}
		s2, err := t.fs.Stat(pathname)
		if err != nil {
			glog.Infof("Stat failed on %q: %s", pathname, err)
			return
		}
		if inode(s1) != inode(s2) {
			logRotations.Add(pathname, 1)
			// flush the old log, pathname is still an index into t.files with the old inode.
			t.handleLogUpdate(pathname)
			fd.Close()
			err := t.w.Remove(pathname)
			if err != nil {
				glog.Info("Failed removing watches on", pathname)
			}
			// Always seek to start on log rotation.
			glog.Infof("Seek to start on %s", pathname)
			t.openLogPath(pathname, true)
		} else {
			glog.Infof("Path %s already being watched, and inode not changed.",
				pathname)
		}
	} else {
		// Freshly opened log file, never seen before.
		t.openLogPath(pathname, true)
	}
}

// openLogPath opens a new log file at pathname, and optionally seeks to the
// start or end of the file. Rotated logs should read from the start, but logs
// opened for the first time read from the end.
func (t *Tailer) openLogPath(pathname string, seekStart bool) {
	if !t.isWatching(pathname) {
		glog.Infof("Not watching %q, ignoring.", pathname)
		return
	}

	d := path.Dir(pathname)
	if !t.isWatching(d) {
		err := t.w.Add(d)
		if err != nil {
			glog.Infof("Adding a create watch failed on %q: %s", d, err)
		}
		t.addWatched(d)
	}

	f, err := t.fs.Open(pathname)
	if err != nil {
		// Doesn't exist yet. We're watching the directory, so we'll pick it up
		// again on create; return successfully.
		if os.IsNotExist(err) {
			return
		}
		glog.Infof("Failed to open %q for reading: %s", pathname, err)
		logErrors.Add(pathname, 1)
		return
	}
	err = t.startNewFile(f, seekStart)
	if err != nil {
		glog.Info(err)
	}
}

func (t *Tailer) startNewFile(f afero.File, seekStart bool) error {
	fi, err := f.Stat()
	if err != nil {
		// Stat failed, log error and return.
		logErrors.Add(f.Name(), 1)
		return fmt.Errorf("Failed to stat %q: %s", f.Name(), err)
	}
	glog.Infof("Mode: %q", fi.Mode())
	switch m := fi.Mode(); {
	case m&os.ModeType == 0:
		if seekStart {
			f.Seek(0, os.SEEK_SET)
		} else {
			f.Seek(0, os.SEEK_END)
		}
		err = t.w.Add(f.Name())
		if err != nil {
			return fmt.Errorf("Adding a change watch failed on %q: %s", f.Name(), err)
		}
		// In case the new log has been written to already, attempt to read the
		// first lines.
		t.partials[f.Name()], err = t.readFile(f, "")
		if err != nil {
			return err
		}
	case m&os.ModeType == os.ModeNamedPipe:
		go t.readForever(f)
	default:
		return fmt.Errorf("Can't open files with mode %v: %s", m&os.ModeType, f.Name())
	}
	t.filesLock.Lock()
	t.files[f.Name()] = f
	t.filesLock.Unlock()
	glog.Infof("Tailing %s", f.Name())

	return nil
}

// start is the main event loop for the Tailer.
// It receives notification of log file changes from the watcher channel, and
// handles them.
func (t *Tailer) run() {
	for e := range t.w.Events() {
		switch e := e.(type) {
		case watcher.UpdateEvent:
			t.handleLogUpdate(e.Pathname)
		case watcher.CreateEvent:
			t.handleLogCreate(e.Pathname)
		default:
			glog.Infof("Unexpected event %q", e)
		}
	}
	glog.Infof("Shutting down tailer.")
	close(t.lines)
}

func (t *Tailer) readForever(f afero.File) {
	t.pipeGroup.Add(1)
	defer t.pipeGroup.Done()

	var err error
	partial := ""
	for {
		select {
		case <-t.pipeQuit:
			return
		default:
			partial, err = t.readFile(f, partial)
			if err != nil {
				glog.Infof("%s: %s", err, f.Name())
				return
			}
		}
	}
}

// Close signals termination to the watcher.
func (t *Tailer) Close() {
	t.w.Close()
	close(t.pipeQuit)
	t.pipeGroup.Wait()
}
<|endoftext|>


<commit_before>package controller

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"os"
	"path/filepath"
	"strings"
	"regexp"

	log "github.com/Sirupsen/logrus"
	"github.com/drone/drone/engine"
	"github.com/drone/drone/model"
	"github.com/drone/drone/remote"
	"github.com/drone/drone/router/middleware/context"
	"github.com/drone/drone/shared/httputil"
	"github.com/drone/drone/shared/token"
	"github.com/drone/drone/store"
	"github.com/drone/drone/yaml"
	"github.com/drone/drone/yaml/matrix"
)

func PostHook(c *gin.Context) {
	remote_ := remote.FromContext(c)

	tmprepo, build, err := remote_.Hook(c.Request)
	if err != nil {
		log.Errorf("failure to parse hook. %s", err)
		c.AbortWithError(400, err)
		return
	}
	if build == nil {
		c.Writer.WriteHeader(200)
		return
	}
	if tmprepo == nil {
		log.Errorf("failure to ascertain repo from hook.")
		c.Writer.WriteHeader(400)
		return
	}

	// skip the build if any case-insensitive combination of the words "skip" and "ci"
	// wrapped in square brackets appear in the commit message
	skipRe := regexp.MustCompile(`\[(?i:ci *skip|skip *ci)\]`)
	skipMatches := skipRe.FindStringSubmatch(build.Message)
	if len(skipMatches) > 0 {
		log.Infof("ignoring hook. %s found in %s", skipMatches[0], build.Commit)
		c.Writer.WriteHeader(204)
		return
	}

	repo, err := store.GetRepoOwnerName(c, tmprepo.Owner, tmprepo.Name)
	if err != nil {
		log.Errorf("failure to find repo %s/%s from hook. %s", tmprepo.Owner, tmprepo.Name, err)
		c.AbortWithError(404, err)
		return
	}

	// get the token and verify the hook is authorized
	parsed, err := token.ParseRequest(c.Request, func(t *token.Token) (string, error) {
		return repo.Hash, nil
	})
	if err != nil {
		log.Errorf("failure to parse token from hook for %s. %s", repo.FullName, err)
		c.AbortWithError(400, err)
		return
	}
	if parsed.Text != repo.FullName {
		log.Errorf("failure to verify token from hook. Expected %s, got %s", repo.FullName, parsed.Text)
		c.AbortWithStatus(403)
		return
	}

	if repo.UserID == 0 {
		log.Warnf("ignoring hook. repo %s has no owner.", repo.FullName)
		c.Writer.WriteHeader(204)
		return
	}
	var skipped = true
	if (build.Event == model.EventPush && repo.AllowPush) ||
		(build.Event == model.EventPull && repo.AllowPull) ||
		(build.Event == model.EventDeploy && repo.AllowDeploy) ||
		(build.Event == model.EventTag && repo.AllowTag) {
		skipped = false
	}

	if skipped {
		log.Infof("ignoring hook. repo %s is disabled for %s events.", repo.FullName, build.Event)
		c.Writer.WriteHeader(204)
		return
	}

	user, err := store.GetUser(c, repo.UserID)
	if err != nil {
		log.Errorf("failure to find repo owner %s. %s", repo.FullName, err)
		c.AbortWithError(500, err)
		return
	}

	// if there is no email address associated with the pull request,
	// we lookup the email address based on the authors github login.
	//
	// my initial hesitation with this code is that it has the ability
	// to expose your email address. At the same time, your email address
	// is already exposed in the public .git log. So while some people will
	// a small number of people will probably be upset by this, I'm not sure
	// it is actually that big of a deal.
	if len(build.Email) == 0 {
		author, err := store.GetUserLogin(c, build.Author)
		if err == nil {
			build.Email = author.Email
		}
	}

	// if the remote has a refresh token, the current access token
	// may be stale. Therefore, we should refresh prior to dispatching
	// the job.
	if refresher, ok := remote_.(remote.Refresher); ok {
		ok, _ := refresher.Refresh(user)
		if ok {
			store.UpdateUser(c, user)
		}
	}

	// fetch the .drone.yml file from the database
	raw, sec, err := remote_.Script(user, repo, build)
	if err != nil {
		log.Errorf("failure to get .drone.yml for %s. %s", repo.FullName, err)
		c.AbortWithError(404, err)
		return
	}

	axes, err := matrix.Parse(string(raw))
	if err != nil {
		log.Errorf("failure to calculate matrix for %s. %s", repo.FullName, err)
		c.AbortWithError(400, err)
		return
	}
	if len(axes) == 0 {
		axes = append(axes, matrix.Axis{})
	}

	netrc, err := remote_.Netrc(user, repo)
	if err != nil {
		log.Errorf("failure to generate netrc for %s. %s", repo.FullName, err)
		c.AbortWithError(500, err)
		return
	}

	key, _ := store.GetKey(c, repo)

	// verify the branches can be built vs skipped
	yconfig, _ := yaml.Parse(string(raw))
	var match = false
	for _, branch := range yconfig.Branches {
		if branch == build.Branch {
			match = true
			break
		}
		match, _ = filepath.Match(branch, build.Branch)
		if match {
			break
		}
	}
	if !match && len(yconfig.Branches) != 0 {
		log.Infof("ignoring hook. yaml file excludes repo and branch %s %s", repo.FullName, build.Branch)
		c.AbortWithStatus(200)
		return
	}

	// update some build fields
	build.Status = model.StatusPending
	build.RepoID = repo.ID

	// and use a transaction
	var jobs []*model.Job
	for num, axis := range axes {
		jobs = append(jobs, &model.Job{
			BuildID:     build.ID,
			Number:      num + 1,
			Status:      model.StatusPending,
			Environment: axis,
		})
	}
	err = store.CreateBuild(c, build, jobs...)
	if err != nil {
		log.Errorf("failure to save commit for %s. %s", repo.FullName, err)
		c.AbortWithError(500, err)
		return
	}

	c.JSON(200, build)

	url := fmt.Sprintf("%s/%s/%d", httputil.GetURL(c.Request), repo.FullName, build.Number)
	err = remote_.Status(user, repo, build, url)
	if err != nil {
		log.Errorf("error setting commit status for %s/%d", repo.FullName, build.Number)
	}

	// get the previous build so taht we can send
	// on status change notifications
	last, _ := store.GetBuildLastBefore(c, repo, build.Branch, build.ID)

	engine_ := context.Engine(c)
	go engine_.Schedule(c.Copy(), &engine.Task{
		User:      user,
		Repo:      repo,
		Build:     build,
		BuildPrev: last,
		Jobs:      jobs,
		Keys:      key,
		Netrc:     netrc,
		Config:    string(raw),
		Secret:    string(sec),
		System: &model.System{
			Link:    httputil.GetURL(c.Request),
			Plugins: strings.Split(os.Getenv("PLUGIN_FILTER"), " "),
			Globals: strings.Split(os.Getenv("PLUGIN_PARAMS"), " "),
		},
	})

}
<commit_msg>move build skipping regular expression variable into global scope<commit_after>package controller

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"os"
	"path/filepath"
	"strings"
	"regexp"

	log "github.com/Sirupsen/logrus"
	"github.com/drone/drone/engine"
	"github.com/drone/drone/model"
	"github.com/drone/drone/remote"
	"github.com/drone/drone/router/middleware/context"
	"github.com/drone/drone/shared/httputil"
	"github.com/drone/drone/shared/token"
	"github.com/drone/drone/store"
	"github.com/drone/drone/yaml"
	"github.com/drone/drone/yaml/matrix"
)

var skipRe = regexp.MustCompile(`\[(?i:ci *skip|skip *ci)\]`)

func PostHook(c *gin.Context) {
	remote_ := remote.FromContext(c)

	tmprepo, build, err := remote_.Hook(c.Request)
	if err != nil {
		log.Errorf("failure to parse hook. %s", err)
		c.AbortWithError(400, err)
		return
	}
	if build == nil {
		c.Writer.WriteHeader(200)
		return
	}
	if tmprepo == nil {
		log.Errorf("failure to ascertain repo from hook.")
		c.Writer.WriteHeader(400)
		return
	}

	// skip the build if any case-insensitive combination of the words "skip" and "ci"
	// wrapped in square brackets appear in the commit message
	skipMatches := skipRe.FindStringSubmatch(build.Message)
	if len(skipMatches) > 0 {
		log.Infof("ignoring hook. %s found in %s", skipMatches[0], build.Commit)
		c.Writer.WriteHeader(204)
		return
	}

	repo, err := store.GetRepoOwnerName(c, tmprepo.Owner, tmprepo.Name)
	if err != nil {
		log.Errorf("failure to find repo %s/%s from hook. %s", tmprepo.Owner, tmprepo.Name, err)
		c.AbortWithError(404, err)
		return
	}

	// get the token and verify the hook is authorized
	parsed, err := token.ParseRequest(c.Request, func(t *token.Token) (string, error) {
		return repo.Hash, nil
	})
	if err != nil {
		log.Errorf("failure to parse token from hook for %s. %s", repo.FullName, err)
		c.AbortWithError(400, err)
		return
	}
	if parsed.Text != repo.FullName {
		log.Errorf("failure to verify token from hook. Expected %s, got %s", repo.FullName, parsed.Text)
		c.AbortWithStatus(403)
		return
	}

	if repo.UserID == 0 {
		log.Warnf("ignoring hook. repo %s has no owner.", repo.FullName)
		c.Writer.WriteHeader(204)
		return
	}
	var skipped = true
	if (build.Event == model.EventPush && repo.AllowPush) ||
		(build.Event == model.EventPull && repo.AllowPull) ||
		(build.Event == model.EventDeploy && repo.AllowDeploy) ||
		(build.Event == model.EventTag && repo.AllowTag) {
		skipped = false
	}

	if skipped {
		log.Infof("ignoring hook. repo %s is disabled for %s events.", repo.FullName, build.Event)
		c.Writer.WriteHeader(204)
		return
	}

	user, err := store.GetUser(c, repo.UserID)
	if err != nil {
		log.Errorf("failure to find repo owner %s. %s", repo.FullName, err)
		c.AbortWithError(500, err)
		return
	}

	// if there is no email address associated with the pull request,
	// we lookup the email address based on the authors github login.
	//
	// my initial hesitation with this code is that it has the ability
	// to expose your email address. At the same time, your email address
	// is already exposed in the public .git log. So while some people will
	// a small number of people will probably be upset by this, I'm not sure
	// it is actually that big of a deal.
	if len(build.Email) == 0 {
		author, err := store.GetUserLogin(c, build.Author)
		if err == nil {
			build.Email = author.Email
		}
	}

	// if the remote has a refresh token, the current access token
	// may be stale. Therefore, we should refresh prior to dispatching
	// the job.
	if refresher, ok := remote_.(remote.Refresher); ok {
		ok, _ := refresher.Refresh(user)
		if ok {
			store.UpdateUser(c, user)
		}
	}

	// fetch the .drone.yml file from the database
	raw, sec, err := remote_.Script(user, repo, build)
	if err != nil {
		log.Errorf("failure to get .drone.yml for %s. %s", repo.FullName, err)
		c.AbortWithError(404, err)
		return
	}

	axes, err := matrix.Parse(string(raw))
	if err != nil {
		log.Errorf("failure to calculate matrix for %s. %s", repo.FullName, err)
		c.AbortWithError(400, err)
		return
	}
	if len(axes) == 0 {
		axes = append(axes, matrix.Axis{})
	}

	netrc, err := remote_.Netrc(user, repo)
	if err != nil {
		log.Errorf("failure to generate netrc for %s. %s", repo.FullName, err)
		c.AbortWithError(500, err)
		return
	}

	key, _ := store.GetKey(c, repo)

	// verify the branches can be built vs skipped
	yconfig, _ := yaml.Parse(string(raw))
	var match = false
	for _, branch := range yconfig.Branches {
		if branch == build.Branch {
			match = true
			break
		}
		match, _ = filepath.Match(branch, build.Branch)
		if match {
			break
		}
	}
	if !match && len(yconfig.Branches) != 0 {
		log.Infof("ignoring hook. yaml file excludes repo and branch %s %s", repo.FullName, build.Branch)
		c.AbortWithStatus(200)
		return
	}

	// update some build fields
	build.Status = model.StatusPending
	build.RepoID = repo.ID

	// and use a transaction
	var jobs []*model.Job
	for num, axis := range axes {
		jobs = append(jobs, &model.Job{
			BuildID:     build.ID,
			Number:      num + 1,
			Status:      model.StatusPending,
			Environment: axis,
		})
	}
	err = store.CreateBuild(c, build, jobs...)
	if err != nil {
		log.Errorf("failure to save commit for %s. %s", repo.FullName, err)
		c.AbortWithError(500, err)
		return
	}

	c.JSON(200, build)

	url := fmt.Sprintf("%s/%s/%d", httputil.GetURL(c.Request), repo.FullName, build.Number)
	err = remote_.Status(user, repo, build, url)
	if err != nil {
		log.Errorf("error setting commit status for %s/%d", repo.FullName, build.Number)
	}

	// get the previous build so taht we can send
	// on status change notifications
	last, _ := store.GetBuildLastBefore(c, repo, build.Branch, build.ID)

	engine_ := context.Engine(c)
	go engine_.Schedule(c.Copy(), &engine.Task{
		User:      user,
		Repo:      repo,
		Build:     build,
		BuildPrev: last,
		Jobs:      jobs,
		Keys:      key,
		Netrc:     netrc,
		Config:    string(raw),
		Secret:    string(sec),
		System: &model.System{
			Link:    httputil.GetURL(c.Request),
			Plugins: strings.Split(os.Getenv("PLUGIN_FILTER"), " "),
			Globals: strings.Split(os.Getenv("PLUGIN_PARAMS"), " "),
		},
	})

}
<|endoftext|>