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

// Package watchdog is responsible for monitoring the sentry for tasks that may
// potentially be stuck or looping inderterminally causing hard to debug hungs in
// the untrusted app.
//
// It works by periodically querying all tasks to check whether they are in user
// mode (RunUser), kernel mode (RunSys), or blocked in the kernel (OffCPU). Tasks
// that have been running in kernel mode for a long time in the same syscall
// without blocking are considered stuck and are reported.
//
// When a stuck task is detected, the watchdog can take one of the following actions:
//		1. LogWarning: Logs a warning message followed by a stack dump of all goroutines.
//			 If a tasks continues to be stuck, the message will repeat every minute, unless
//			 a new stuck task is detected
//		2. Panic: same as above, followed by panic()
//
package watchdog

import (
	"bytes"
	"fmt"
	"time"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/log"
	"gvisor.dev/gvisor/pkg/metric"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	ktime "gvisor.dev/gvisor/pkg/sentry/kernel/time"
	"gvisor.dev/gvisor/pkg/sync"
)

// Opts configures the watchdog.
type Opts struct {
	// TaskTimeout is the amount of time to allow a task to execute the
	// same syscall without blocking before it's declared stuck.
	TaskTimeout time.Duration

	// TaskTimeoutAction indicates what action to take when a stuck tasks
	// is detected.
	TaskTimeoutAction Action

	// StartupTimeout is the amount of time to allow between watchdog
	// creation and calling watchdog.Start.
	StartupTimeout time.Duration

	// StartupTimeoutAction indicates what action to take when
	// watchdog.Start is not called within the timeout.
	StartupTimeoutAction Action
}

// DefaultOpts is a default set of options for the watchdog.
var DefaultOpts = Opts{
	// Task timeout.
	TaskTimeout:       3 * time.Minute,
	TaskTimeoutAction: LogWarning,

	// Startup timeout.
	StartupTimeout:       30 * time.Second,
	StartupTimeoutAction: LogWarning,
}

// descheduleThreshold is the amount of time scheduling needs to be off before the entire wait period
// is discounted from task's last update time. It's set high enough that small scheduling delays won't
// trigger it.
const descheduleThreshold = 1 * time.Second

var (
	stuckStartup = metric.MustCreateNewUint64Metric("/watchdog/stuck_startup_detected", true /* sync */, "Incremented once on startup watchdog timeout")
	stuckTasks   = metric.MustCreateNewUint64Metric("/watchdog/stuck_tasks_detected", true /* sync */, "Cumulative count of stuck tasks detected")
)

// Amount of time to wait before dumping the stack to the log again when the same task(s) remains stuck.
var stackDumpSameTaskPeriod = time.Minute

// Action defines what action to take when a stuck task is detected.
type Action int

const (
	// LogWarning logs warning message followed by stack trace.
	LogWarning Action = iota

	// Panic will do the same logging as LogWarning and panic().
	Panic
)

// Set implements flag.Value.
func (a *Action) Set(v string) error {
	switch v {
	case "log", "logwarning":
		*a = LogWarning
	case "panic":
		*a = Panic
	default:
		return fmt.Errorf("invalid watchdog action %q", v)
	}
	return nil
}

// Get implements flag.Value.
func (a *Action) Get() interface{} {
	return *a
}

// String returns Action's string representation.
func (a Action) String() string {
	switch a {
	case LogWarning:
		return "logWarning"
	case Panic:
		return "panic"
	default:
		panic(fmt.Sprintf("Invalid watchdog action: %d", a))
	}
}

// Watchdog is the main watchdog class. It controls a goroutine that periodically
// analyses all tasks and reports if any of them appear to be stuck.
type Watchdog struct {
	// Configuration options are embedded.
	Opts

	// period indicates how often to check all tasks. It's calculated based on
	// opts.TaskTimeout.
	period time.Duration

	// k is where the tasks come from.
	k *kernel.Kernel

	// stop is used to notify to watchdog should stop.
	stop chan struct{}

	// done is used to notify when the watchdog has stopped.
	done chan struct{}

	// offenders map contains all tasks that are currently stuck.
	offenders map[*kernel.Task]*offender

	// lastStackDump tracks the last time a stack dump was generated to prevent
	// spamming the log.
	lastStackDump time.Time

	// lastRun is set to the last time the watchdog executed a monitoring loop.
	lastRun ktime.Time

	// mu protects the fields below.
	mu sync.Mutex

	// running is true if the watchdog is running.
	running bool

	// startCalled is true if Start has ever been called. It remains true
	// even if Stop is called.
	startCalled bool
}

type offender struct {
	lastUpdateTime ktime.Time
}

// New creates a new watchdog.
func New(k *kernel.Kernel, opts Opts) *Watchdog {
	// 4 is arbitrary, just don't want to prolong 'TaskTimeout' too much.
	period := opts.TaskTimeout / 4
	w := &Watchdog{
		Opts:      opts,
		k:         k,
		period:    period,
		offenders: make(map[*kernel.Task]*offender),
		stop:      make(chan struct{}),
		done:      make(chan struct{}),
	}

	// Handle StartupTimeout if it exists.
	if w.StartupTimeout > 0 {
		log.Infof("Watchdog waiting %v for startup", w.StartupTimeout)
		go w.waitForStart() // S/R-SAFE: watchdog is stopped buring save and restarted after restore.
	}

	return w
}

// Start starts the watchdog.
func (w *Watchdog) Start() {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.startCalled = true

	if w.running {
		return
	}

	if w.TaskTimeout == 0 {
		log.Infof("Watchdog task timeout disabled")
		return
	}
	w.lastRun = w.k.MonotonicClock().Now()

	log.Infof("Starting watchdog, period: %v, timeout: %v, action: %v", w.period, w.TaskTimeout, w.TaskTimeoutAction)
	go w.loop() // S/R-SAFE: watchdog is stopped during save and restarted after restore.
	w.running = true
}

// Stop requests the watchdog to stop and wait for it.
func (w *Watchdog) Stop() {
	if w.TaskTimeout == 0 {
		return
	}

	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.running {
		return
	}
	log.Infof("Stopping watchdog")
	w.stop <- struct{}{}
	<-w.done
	w.running = false
	log.Infof("Watchdog stopped")
}

// waitForStart waits for Start to be called and takes action if it does not
// happen within the startup timeout.
func (w *Watchdog) waitForStart() {
	<-time.After(w.StartupTimeout)
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.startCalled {
		// We are fine.
		return
	}

	stuckStartup.Increment()
	metric.WeirdnessMetric.Increment("watchdog_stuck_startup")

	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("Watchdog.Start() not called within %s", w.StartupTimeout))
	w.doAction(w.StartupTimeoutAction, false, &buf)
}

// loop is the main watchdog routine. It only returns when 'Stop()' is called.
func (w *Watchdog) loop() {
	// Loop until someone stops it.
	for {
		select {
		case <-w.stop:
			w.done <- struct{}{}
			return
		case <-time.After(w.period):
			w.runTurn()
		}
	}
}

// runTurn runs a single pass over all tasks and reports anything it finds.
func (w *Watchdog) runTurn() {
	// Someone needs to watch the watchdog. The call below can get stuck if there
	// is a deadlock affecting root's PID namespace mutex. Run it in a goroutine
	// and report if it takes too long to return.
	var tasks []*kernel.Task
	done := make(chan struct{})
	go func() { // S/R-SAFE: watchdog is stopped and restarted during S/R.
		tasks = w.k.TaskSet().Root.Tasks()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(w.TaskTimeout):
		// Report if the watchdog is not making progress.
		// No one is watching the watchdog watcher though.
		w.reportStuckWatchdog()
		<-done
	}

	newOffenders := make(map[*kernel.Task]*offender)
	newTaskFound := false
	now := ktime.FromNanoseconds(int64(w.k.CPUClockNow() * uint64(linux.ClockTick)))

	// The process may be running with low CPU limit making tasks appear stuck because
	// are starved of CPU cycles. An estimate is that Tasks could have been starved
	// since the last time the watchdog run. If the watchdog detects that scheduling
	// is off, it will discount the entire duration since last run from 'lastUpdateTime'.
	discount := time.Duration(0)
	if now.Sub(w.lastRun.Add(w.period)) > descheduleThreshold {
		discount = now.Sub(w.lastRun)
	}
	w.lastRun = now

	log.Infof("Watchdog starting loop, tasks: %d, discount: %v", len(tasks), discount)
	for _, t := range tasks {
		tsched := t.TaskGoroutineSchedInfo()

		// An offender is a task running inside the kernel for longer than the specified timeout.
		if tsched.State == kernel.TaskGoroutineRunningSys {
			lastUpdateTime := ktime.FromNanoseconds(int64(tsched.Timestamp * uint64(linux.ClockTick)))
			elapsed := now.Sub(lastUpdateTime) - discount
			if elapsed > w.TaskTimeout {
				tc, ok := w.offenders[t]
				if !ok {
					// New stuck task detected.
					//
					// Note that tasks blocked doing IO may be considered stuck in kernel,
					// unless they are surrounded by
					// Task.UninterruptibleSleepStart/Finish.
					tc = &offender{lastUpdateTime: lastUpdateTime}
					stuckTasks.Increment()
					metric.WeirdnessMetric.Increment("watchdog_stuck_tasks")
					newTaskFound = true
				}
				newOffenders[t] = tc
			}
		}
	}
	if len(newOffenders) > 0 {
		w.report(newOffenders, newTaskFound, now)
	}

	// Remember which tasks have been reported.
	w.offenders = newOffenders
}

// report takes appropriate action when a stuck task is detected.
func (w *Watchdog) report(offenders map[*kernel.Task]*offender, newTaskFound bool, now ktime.Time) {
	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("Sentry detected %d stuck task(s):\n", len(offenders)))
	for t, o := range offenders {
		tid := w.k.TaskSet().Root.IDOfTask(t)
		buf.WriteString(fmt.Sprintf("\tTask tid: %v (goroutine %d), entered RunSys state %v ago.\n", tid, t.GoroutineID(), now.Sub(o.lastUpdateTime)))
	}
	buf.WriteString("Search for 'goroutine <id>' in the stack dump to find the offending goroutine(s)")

	// Force stack dump only if a new task is detected.
	w.doAction(w.TaskTimeoutAction, newTaskFound, &buf)
}

func (w *Watchdog) reportStuckWatchdog() {
	var buf bytes.Buffer
	buf.WriteString("Watchdog goroutine is stuck")
	w.doAction(w.TaskTimeoutAction, false, &buf)
}

// doAction will take the given action. If the action is LogWarning, the stack
// is not always dumped to the log to prevent log flooding. "forceStack"
// guarantees that the stack will be dumped regardless.
func (w *Watchdog) doAction(action Action, forceStack bool, msg *bytes.Buffer) {
	switch action {
	case LogWarning:
		// Dump stack only if forced or sometime has passed since the last time a
		// stack dump was generated.
		if !forceStack && time.Since(w.lastStackDump) < stackDumpSameTaskPeriod {
			msg.WriteString("\n...[stack dump skipped]...")
			log.Warningf(msg.String())
			return
		}
		log.TracebackAll(msg.String())
		w.lastStackDump = time.Now()

	case Panic:
		// Panic will skip over running tasks, which is likely the culprit here. So manually
		// dump all stacks before panic'ing.
		log.TracebackAll(msg.String())

		// Attempt to flush metrics, timeout and move on in case metrics are stuck as well.
		metricsEmitted := make(chan struct{}, 1)
		go func() { // S/R-SAFE: watchdog is stopped during save and restarted after restore.
			// Flush metrics before killing process.
			metric.EmitMetricUpdate()
			metricsEmitted <- struct{}{}
		}()
		select {
		case <-metricsEmitted:
		case <-time.After(1 * time.Second):
		}
		panic(fmt.Sprintf("%s\nStack for running G's are skipped while panicking.", msg.String()))

	default:
		panic(fmt.Sprintf("Unknown watchdog action %v", action))

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

// Package watchdog is responsible for monitoring the sentry for tasks that may
// potentially be stuck or looping inderterminally causing hard to debug hungs in
// the untrusted app.
//
// It works by periodically querying all tasks to check whether they are in user
// mode (RunUser), kernel mode (RunSys), or blocked in the kernel (OffCPU). Tasks
// that have been running in kernel mode for a long time in the same syscall
// without blocking are considered stuck and are reported.
//
// When a stuck task is detected, the watchdog can take one of the following actions:
//		1. LogWarning: Logs a warning message followed by a stack dump of all goroutines.
//			 If a tasks continues to be stuck, the message will repeat every minute, unless
//			 a new stuck task is detected
//		2. Panic: same as above, followed by panic()
//
package watchdog

import (
	"bytes"
	"fmt"
	"time"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/log"
	"gvisor.dev/gvisor/pkg/metric"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	ktime "gvisor.dev/gvisor/pkg/sentry/kernel/time"
	"gvisor.dev/gvisor/pkg/sync"
)

// Opts configures the watchdog.
type Opts struct {
	// TaskTimeout is the amount of time to allow a task to execute the
	// same syscall without blocking before it's declared stuck.
	TaskTimeout time.Duration

	// TaskTimeoutAction indicates what action to take when a stuck tasks
	// is detected.
	TaskTimeoutAction Action

	// StartupTimeout is the amount of time to allow between watchdog
	// creation and calling watchdog.Start.
	StartupTimeout time.Duration

	// StartupTimeoutAction indicates what action to take when
	// watchdog.Start is not called within the timeout.
	StartupTimeoutAction Action
}

// DefaultOpts is a default set of options for the watchdog.
var DefaultOpts = Opts{
	// Task timeout.
	TaskTimeout:       3 * time.Minute,
	TaskTimeoutAction: LogWarning,

	// Startup timeout.
	StartupTimeout:       30 * time.Second,
	StartupTimeoutAction: LogWarning,
}

// descheduleThreshold is the amount of time scheduling needs to be off before the entire wait period
// is discounted from task's last update time. It's set high enough that small scheduling delays won't
// trigger it.
const descheduleThreshold = 1 * time.Second

// Amount of time to wait before dumping the stack to the log again when the same task(s) remains stuck.
var stackDumpSameTaskPeriod = time.Minute

// Action defines what action to take when a stuck task is detected.
type Action int

const (
	// LogWarning logs warning message followed by stack trace.
	LogWarning Action = iota

	// Panic will do the same logging as LogWarning and panic().
	Panic
)

// Set implements flag.Value.
func (a *Action) Set(v string) error {
	switch v {
	case "log", "logwarning":
		*a = LogWarning
	case "panic":
		*a = Panic
	default:
		return fmt.Errorf("invalid watchdog action %q", v)
	}
	return nil
}

// Get implements flag.Value.
func (a *Action) Get() interface{} {
	return *a
}

// String returns Action's string representation.
func (a Action) String() string {
	switch a {
	case LogWarning:
		return "logWarning"
	case Panic:
		return "panic"
	default:
		panic(fmt.Sprintf("Invalid watchdog action: %d", a))
	}
}

// Watchdog is the main watchdog class. It controls a goroutine that periodically
// analyses all tasks and reports if any of them appear to be stuck.
type Watchdog struct {
	// Configuration options are embedded.
	Opts

	// period indicates how often to check all tasks. It's calculated based on
	// opts.TaskTimeout.
	period time.Duration

	// k is where the tasks come from.
	k *kernel.Kernel

	// stop is used to notify to watchdog should stop.
	stop chan struct{}

	// done is used to notify when the watchdog has stopped.
	done chan struct{}

	// offenders map contains all tasks that are currently stuck.
	offenders map[*kernel.Task]*offender

	// lastStackDump tracks the last time a stack dump was generated to prevent
	// spamming the log.
	lastStackDump time.Time

	// lastRun is set to the last time the watchdog executed a monitoring loop.
	lastRun ktime.Time

	// mu protects the fields below.
	mu sync.Mutex

	// running is true if the watchdog is running.
	running bool

	// startCalled is true if Start has ever been called. It remains true
	// even if Stop is called.
	startCalled bool
}

type offender struct {
	lastUpdateTime ktime.Time
}

// New creates a new watchdog.
func New(k *kernel.Kernel, opts Opts) *Watchdog {
	// 4 is arbitrary, just don't want to prolong 'TaskTimeout' too much.
	period := opts.TaskTimeout / 4
	w := &Watchdog{
		Opts:      opts,
		k:         k,
		period:    period,
		offenders: make(map[*kernel.Task]*offender),
		stop:      make(chan struct{}),
		done:      make(chan struct{}),
	}

	// Handle StartupTimeout if it exists.
	if w.StartupTimeout > 0 {
		log.Infof("Watchdog waiting %v for startup", w.StartupTimeout)
		go w.waitForStart() // S/R-SAFE: watchdog is stopped buring save and restarted after restore.
	}

	return w
}

// Start starts the watchdog.
func (w *Watchdog) Start() {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.startCalled = true

	if w.running {
		return
	}

	if w.TaskTimeout == 0 {
		log.Infof("Watchdog task timeout disabled")
		return
	}
	w.lastRun = w.k.MonotonicClock().Now()

	log.Infof("Starting watchdog, period: %v, timeout: %v, action: %v", w.period, w.TaskTimeout, w.TaskTimeoutAction)
	go w.loop() // S/R-SAFE: watchdog is stopped during save and restarted after restore.
	w.running = true
}

// Stop requests the watchdog to stop and wait for it.
func (w *Watchdog) Stop() {
	if w.TaskTimeout == 0 {
		return
	}

	w.mu.Lock()
	defer w.mu.Unlock()
	if !w.running {
		return
	}
	log.Infof("Stopping watchdog")
	w.stop <- struct{}{}
	<-w.done
	w.running = false
	log.Infof("Watchdog stopped")
}

// waitForStart waits for Start to be called and takes action if it does not
// happen within the startup timeout.
func (w *Watchdog) waitForStart() {
	<-time.After(w.StartupTimeout)
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.startCalled {
		// We are fine.
		return
	}

	metric.WeirdnessMetric.Increment("watchdog_stuck_startup")

	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("Watchdog.Start() not called within %s", w.StartupTimeout))
	w.doAction(w.StartupTimeoutAction, false, &buf)
}

// loop is the main watchdog routine. It only returns when 'Stop()' is called.
func (w *Watchdog) loop() {
	// Loop until someone stops it.
	for {
		select {
		case <-w.stop:
			w.done <- struct{}{}
			return
		case <-time.After(w.period):
			w.runTurn()
		}
	}
}

// runTurn runs a single pass over all tasks and reports anything it finds.
func (w *Watchdog) runTurn() {
	// Someone needs to watch the watchdog. The call below can get stuck if there
	// is a deadlock affecting root's PID namespace mutex. Run it in a goroutine
	// and report if it takes too long to return.
	var tasks []*kernel.Task
	done := make(chan struct{})
	go func() { // S/R-SAFE: watchdog is stopped and restarted during S/R.
		tasks = w.k.TaskSet().Root.Tasks()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(w.TaskTimeout):
		// Report if the watchdog is not making progress.
		// No one is watching the watchdog watcher though.
		w.reportStuckWatchdog()
		<-done
	}

	newOffenders := make(map[*kernel.Task]*offender)
	newTaskFound := false
	now := ktime.FromNanoseconds(int64(w.k.CPUClockNow() * uint64(linux.ClockTick)))

	// The process may be running with low CPU limit making tasks appear stuck because
	// are starved of CPU cycles. An estimate is that Tasks could have been starved
	// since the last time the watchdog run. If the watchdog detects that scheduling
	// is off, it will discount the entire duration since last run from 'lastUpdateTime'.
	discount := time.Duration(0)
	if now.Sub(w.lastRun.Add(w.period)) > descheduleThreshold {
		discount = now.Sub(w.lastRun)
	}
	w.lastRun = now

	log.Infof("Watchdog starting loop, tasks: %d, discount: %v", len(tasks), discount)
	for _, t := range tasks {
		tsched := t.TaskGoroutineSchedInfo()

		// An offender is a task running inside the kernel for longer than the specified timeout.
		if tsched.State == kernel.TaskGoroutineRunningSys {
			lastUpdateTime := ktime.FromNanoseconds(int64(tsched.Timestamp * uint64(linux.ClockTick)))
			elapsed := now.Sub(lastUpdateTime) - discount
			if elapsed > w.TaskTimeout {
				tc, ok := w.offenders[t]
				if !ok {
					// New stuck task detected.
					//
					// Note that tasks blocked doing IO may be considered stuck in kernel,
					// unless they are surrounded by
					// Task.UninterruptibleSleepStart/Finish.
					tc = &offender{lastUpdateTime: lastUpdateTime}
					metric.WeirdnessMetric.Increment("watchdog_stuck_tasks")
					newTaskFound = true
				}
				newOffenders[t] = tc
			}
		}
	}
	if len(newOffenders) > 0 {
		w.report(newOffenders, newTaskFound, now)
	}

	// Remember which tasks have been reported.
	w.offenders = newOffenders
}

// report takes appropriate action when a stuck task is detected.
func (w *Watchdog) report(offenders map[*kernel.Task]*offender, newTaskFound bool, now ktime.Time) {
	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("Sentry detected %d stuck task(s):\n", len(offenders)))
	for t, o := range offenders {
		tid := w.k.TaskSet().Root.IDOfTask(t)
		buf.WriteString(fmt.Sprintf("\tTask tid: %v (goroutine %d), entered RunSys state %v ago.\n", tid, t.GoroutineID(), now.Sub(o.lastUpdateTime)))
	}
	buf.WriteString("Search for 'goroutine <id>' in the stack dump to find the offending goroutine(s)")

	// Force stack dump only if a new task is detected.
	w.doAction(w.TaskTimeoutAction, newTaskFound, &buf)
}

func (w *Watchdog) reportStuckWatchdog() {
	var buf bytes.Buffer
	buf.WriteString("Watchdog goroutine is stuck")
	w.doAction(w.TaskTimeoutAction, false, &buf)
}

// doAction will take the given action. If the action is LogWarning, the stack
// is not always dumped to the log to prevent log flooding. "forceStack"
// guarantees that the stack will be dumped regardless.
func (w *Watchdog) doAction(action Action, forceStack bool, msg *bytes.Buffer) {
	switch action {
	case LogWarning:
		// Dump stack only if forced or sometime has passed since the last time a
		// stack dump was generated.
		if !forceStack && time.Since(w.lastStackDump) < stackDumpSameTaskPeriod {
			msg.WriteString("\n...[stack dump skipped]...")
			log.Warningf(msg.String())
			return
		}
		log.TracebackAll(msg.String())
		w.lastStackDump = time.Now()

	case Panic:
		// Panic will skip over running tasks, which is likely the culprit here. So manually
		// dump all stacks before panic'ing.
		log.TracebackAll(msg.String())

		// Attempt to flush metrics, timeout and move on in case metrics are stuck as well.
		metricsEmitted := make(chan struct{}, 1)
		go func() { // S/R-SAFE: watchdog is stopped during save and restarted after restore.
			// Flush metrics before killing process.
			metric.EmitMetricUpdate()
			metricsEmitted <- struct{}{}
		}()
		select {
		case <-metricsEmitted:
		case <-time.After(1 * time.Second):
		}
		panic(fmt.Sprintf("%s\nStack for running G's are skipped while panicking.", msg.String()))

	default:
		panic(fmt.Sprintf("Unknown watchdog action %v", action))

	}
}
<|endoftext|>


<commit_before>package workloads

import (
	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/swan/mocks"
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

// TestMemcachedWithMockedExecutor
func TestMemcachedWithMockedExecutor(t *testing.T) {
	log.SetLevel(log.ErrorLevel)

	const (
		expectedCommand = "test -p 11211 -u memcached -t 4 -m 64 -c 1024"
	)

	mockedExecutor := new(mocks.Executor)
	mockedExecutor.On("Execute", expectedCommand).Return(nil, nil).Once()

	Convey("While using Memcached launcher", t, func() {
		memcachedLauncher := NewMemcached(
			mockedExecutor,
			DefaultMemcachedConfig("test"))

		Convey("Build command should create proper command", func() {
			command := memcachedLauncher.buildCommand()

			So(command, ShouldEqual, expectedCommand)
		})

		Convey("Arguments passed to Executor should be a proper command", func() {
			memcachedLauncher.Launch()
			mockedExecutor.AssertExpectations(t)
		})
	})
}
<commit_msg>Rebased to master.<commit_after>package workloads

import (
	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/swan/pkg/executor/mocks"
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

// TestMemcachedWithMockedExecutor
func TestMemcachedWithMockedExecutor(t *testing.T) {
	log.SetLevel(log.ErrorLevel)

	const (
		expectedCommand = "test -p 11211 -u memcached -t 4 -m 64 -c 1024"
	)

	mockedExecutor := new(mocks.Executor)
	mockedExecutor.On("Execute", expectedCommand).Return(nil, nil).Once()

	Convey("While using Memcached launcher", t, func() {
		memcachedLauncher := NewMemcached(
			mockedExecutor,
			DefaultMemcachedConfig("test"))

		Convey("Build command should create proper command", func() {
			command := memcachedLauncher.buildCommand()

			So(command, ShouldEqual, expectedCommand)
		})

		Convey("Arguments passed to Executor should be a proper command", func() {
			memcachedLauncher.Launch()
			mockedExecutor.AssertExpectations(t)
		})
	})
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"health_monitor/api"
	"health_monitor/disk"
	"health_monitor/live"
	"health_monitor/setup"
	"health_monitor/utils"
	"health_monitor/webui"
)

// Flags holds the health_monitor command line arguments
type Flags struct {
	NoWebUI *bool
	NoCLI   *bool
	Quite   *bool
}

func main() {
	var (
		wg    sync.WaitGroup
		flags Flags
		chans [5]chan bool //Number of modules
	)

	for i := range chans {
		chans[i] = make(chan bool)
	}

	flags.NoWebUI = flag.Bool("nowebui", false, "Disables the web ui")
	flags.NoCLI = flag.Bool("nocli", false, "Disables cli")
	flags.Quite = flag.Bool("quite", false, "Disables all notifications except email")

	flag.Parse()

	go webui.RunServer(setup.ConfigVars.Port)
	if (*flags.NoCLI == true) || (*flags.NoWebUI == false) {
		fmt.Printf("[*] Server is up and running at 127.0.0.1:%s\n", setup.ConfigVars.Port)
	}

	exitChan := make(chan os.Signal, 1)
	signal.Notify(exitChan, syscall.SIGINT, syscall.SIGTERM)
	wg.Add(1)
	go tearDown(exitChan, &wg)

	runModules(chans, &wg)
	controlModule(chans, &wg)
	setup.SaveStatus()
}

func controlModule(chans [5]chan bool, wg *sync.WaitGroup) {
	api.ControlChan = make(chan utils.Status)
	for {
		data := <-api.ControlChan
		switch data.Module {
		case "live":
			if data.Run && !setup.ModulesStatus.Live {
				setup.ModulesStatus.Live = true
				wg.Add(1)
				go live.Live(chans[0], wg)
			} else if setup.ModulesStatus.Live {
				setup.ModulesStatus.Live = false
				chans[0] <- true
			}
		case "target":
			break
		case "disk":
			if data.Run && !setup.ModulesStatus.Disk {
				setup.ModulesStatus.Disk = true
				wg.Add(1)
				go disk.Disk(chans[2], wg)
			} else if setup.ModulesStatus.Disk {
				setup.ModulesStatus.Disk = false
				chans[2] <- true
			}
		}
	}
}

func runModules(chans [5]chan bool, wg *sync.WaitGroup) {
	if setup.ModulesStatus.Live {
		wg.Add(1)
		go live.Live(chans[0], wg)
	}
	if setup.ModulesStatus.Target {

	}
	if setup.ModulesStatus.Disk {
		wg.Add(1)
		go disk.Disk(chans[2], wg)
	}
}

func tearDown(exitChan chan os.Signal, wg *sync.WaitGroup) {
	<-exitChan
	fmt.Println("in teardown")
	setup.MainLogFile.Close()
	setup.DBLogFile.Close()
	setup.Database.Close()
	setup.SaveStatus()

	var module string
	for module, _ = range api.ConfFunc {
		api.ChangeModuleStatus(module, false)
	}
	wg.Done()
	wg.Wait()
	os.Exit(0)
}
<commit_msg>Added logs in main function<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"health_monitor/api"
	"health_monitor/disk"
	"health_monitor/live"
	"health_monitor/setup"
	"health_monitor/utils"
	"health_monitor/webui"
)

// Flags holds the health_monitor command line arguments
type Flags struct {
	NoWebUI *bool
	NoCLI   *bool
	Quite   *bool
}

func main() {
	var (
		wg    sync.WaitGroup
		flags Flags
		chans [5]chan bool //Number of modules
	)

	for i := range chans {
		chans[i] = make(chan bool)
	}

	flags.NoWebUI = flag.Bool("nowebui", false, "Disables the web ui")
	flags.NoCLI = flag.Bool("nocli", false, "Disables cli")
	flags.Quite = flag.Bool("quite", false, "Disables all notifications except email")

	flag.Parse()

	go webui.RunServer(setup.ConfigVars.Port)
	if (*flags.NoCLI == true) || (*flags.NoWebUI == false) {
		fmt.Printf("[*] Server is up and running at http://127.0.0.1:%s\n", setup.ConfigVars.Port)
	}

	exitChan := make(chan os.Signal, 1)
	signal.Notify(exitChan, syscall.SIGINT, syscall.SIGTERM)
	wg.Add(1)
	go tearDown(exitChan, &wg)

	utils.ModuleLogs(setup.MainLogFile, "Running modules from last saved profile")
	runModules(chans, &wg)
	controlModule(chans, &wg)
	setup.SaveStatus()
}

func controlModule(chans [5]chan bool, wg *sync.WaitGroup) {
	api.ControlChan = make(chan utils.Status)
	for {
		data := <-api.ControlChan
		switch data.Module {
		case "live":
			if data.Run && !setup.ModulesStatus.Live {
				setup.ModulesStatus.Live = true
				wg.Add(1)
				utils.ModuleLogs(setup.MainLogFile, "Started live module")
				go live.Live(chans[0], wg)
			} else if setup.ModulesStatus.Live {
				setup.ModulesStatus.Live = false
				utils.ModuleLogs(setup.MainLogFile, "Stopped live module")
				chans[0] <- true
			}
		case "target":
			break
		case "disk":
			if data.Run && !setup.ModulesStatus.Disk {
				setup.ModulesStatus.Disk = true
				wg.Add(1)
				utils.ModuleLogs(setup.MainLogFile, "Started disk module")
				go disk.Disk(chans[2], wg)
			} else if setup.ModulesStatus.Disk {
				setup.ModulesStatus.Disk = false
				utils.ModuleLogs(setup.MainLogFile, "Stopped disk module")
				chans[2] <- true
			}
		}
	}
}

func runModules(chans [5]chan bool, wg *sync.WaitGroup) {
	if setup.ModulesStatus.Live {
		wg.Add(1)
		utils.ModuleLogs(setup.MainLogFile, "Started live module")
		go live.Live(chans[0], wg)
	}
	if setup.ModulesStatus.Target {
		utils.ModuleLogs(setup.MainLogFile, "Started target module")
	}
	if setup.ModulesStatus.Disk {
		wg.Add(1)
		utils.ModuleLogs(setup.MainLogFile, "Started disk module")
		go disk.Disk(chans[2], wg)
	}
}

func tearDown(exitChan chan os.Signal, wg *sync.WaitGroup) {
	<-exitChan
	utils.ModuleLogs(setup.MainLogFile, "Shutdown signal received.")
	setup.Database.Close()
	setup.SaveStatus()
	utils.ModuleLogs(setup.MainLogFile, "Saved all config data. Stopping running modules")

	var module string
	for module, _ = range api.ConfFunc {
		api.ChangeModuleStatus(module, false)
	}

	setup.MainLogFile.Close()
	setup.DBLogFile.Close()
	wg.Done()
	wg.Wait()
	os.Exit(0)
}
<|endoftext|>


<commit_before>// +build !windows,!darwin

package volume

import (
	"errors"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

var useAmixer bool

func init() {
	if _, err := exec.LookPath("pactl"); err != nil {
		useAmixer = true
	}
}

func cmdEnv() []string {
	return []string{"LANG=C", "LC_ALL=C"}
}

func getVolumeCmd() []string {
	if useAmixer {
		return []string{"amixer", "get", "Master"}
	}
	return []string{"pactl", "list", "sinks"}
}

func getPADefaultSink() (string, error) {
	out, err := execCmd([]string{"pactl", "info"})

	if err != nil {
		return "", err
	}

	lines := strings.Split(string(out), "\n")

	defaultSinkStr := "Default Sink: "
	for _, line := range lines {
		s := strings.TrimLeft(line, " \t")
		if strings.HasPrefix(s, defaultSinkStr) {
			return strings.TrimSpace(strings.Replace(s, defaultSinkStr, "", 1)), nil
		}
	}
	return "", errors.New("Could not find PulseAudio Default Sink")
}

func parseVolume(out string) (int, error) {
	sinkName, sinkNameErr := getPADefaultSink()

	paDefaultSinkHooked := false

	if sinkNameErr != nil {
		paDefaultSinkHooked = true
	}

	lines := strings.Split(out, "\n")
	volumePattern := regexp.MustCompile(`\d+%`)

	for _, line := range lines {
		s := strings.TrimLeft(line, " \t")

		if !useAmixer && strings.Contains(s, "Name: "+string(sinkName)) {
			paDefaultSinkHooked = true
		}

		if useAmixer && strings.Contains(s, "Playback") && strings.Contains(s, "%") ||
			!useAmixer && paDefaultSinkHooked && strings.HasPrefix(s, "Volume:") {
			volumeStr := volumePattern.FindString(s)
			return strconv.Atoi(volumeStr[:len(volumeStr)-1])
		}
	}
	return 0, errors.New("no volume found")
}

func setVolumeCmd(volume int) []string {
	if useAmixer {
		return []string{"amixer", "set", "Master", strconv.Itoa(volume) + "%"}
	}
	return []string{"pactl", "set-sink-volume", "@DEFAULT_SINK@", strconv.Itoa(volume) + "%"}
}

func increaseVolumeCmd(diff int) []string {
	var sign string
	if diff >= 0 {
		sign = "+"
	} else if useAmixer {
		diff = -diff
		sign = "-"
	}
	if useAmixer {
		return []string{"amixer", "set", "Master", strconv.Itoa(diff) + "%" + sign}
	}
	return []string{"pactl", "--", "set-sink-volume", "@DEFAULT_SINK@", sign + strconv.Itoa(diff) + "%"}
}

func getMutedCmd() []string {
	if useAmixer {
		return []string{"amixer", "get", "Master"}
	}
	return []string{"pactl", "list", "sinks"}
}

func parseMuted(out string) (bool, error) {
	sinkName, sinkNameErr := getPADefaultSink()

	paDefaultSinkHooked := false

	if sinkNameErr != nil {
		paDefaultSinkHooked = true
	}

	lines := strings.Split(out, "\n")
	for _, line := range lines {
		s := strings.TrimLeft(line, " \t")

		if !useAmixer && strings.Contains(s, "Name: "+string(sinkName)) {
			paDefaultSinkHooked = true
		}

		if useAmixer && strings.Contains(s, "Playback") && strings.Contains(s, "%") ||
			!useAmixer && paDefaultSinkHooked && strings.HasPrefix(s, "Mute: ") {
			if strings.Contains(s, "[off]") || strings.Contains(s, "yes") {
				return true, nil
			} else if strings.Contains(s, "[on]") || strings.Contains(s, "no") {
				return false, nil
			}
		}
	}
	return false, errors.New("no muted information found")
}

func muteCmd() []string {
	if useAmixer {
		return []string{"amixer", "-D", "pulse", "set", "Master", "mute"}
	}
	return []string{"pactl", "set-sink-mute", "@DEFAULT_SINK@", "1"}
}

func unmuteCmd() []string {
	if useAmixer {
		return []string{"amixer", "-D", "pulse", "set", "Master", "unmute"}
	}
	return []string{"pactl", "set-sink-mute", "@DEFAULT_SINK@", "0"}
}
<commit_msg>Revert "PulseAudio: Use local variable instead of package variable"<commit_after>// +build !windows,!darwin

package volume

import (
	"errors"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

var useAmixer bool

func init() {
	if _, err := exec.LookPath("pactl"); err != nil {
		useAmixer = true
	}
}

func cmdEnv() []string {
	return []string{"LANG=C", "LC_ALL=C"}
}

func getVolumeCmd() []string {
	if useAmixer {
		return []string{"amixer", "get", "Master"}
	}
	return []string{"pactl", "list", "sinks"}
}

func getPADefaultSink() (string, error) {
	out, err := execCmd([]string{"pactl", "info"})

	if err != nil {
		return "", err
	}

	lines := strings.Split(string(out), "\n")

	defaultSinkStr := "Default Sink: "
	for _, line := range lines {
		s := strings.TrimLeft(line, " \t")
		if strings.HasPrefix(s, defaultSinkStr) {
			return strings.TrimSpace(strings.Replace(s, defaultSinkStr, "", 1)), nil
		}
	}
	return "", errors.New("Could not find PulseAudio Default Sink")
}

var volumePattern = regexp.MustCompile(`\d+%`)

func parseVolume(out string) (int, error) {
	sinkName, sinkNameErr := getPADefaultSink()

	paDefaultSinkHooked := false

	if sinkNameErr != nil {
		paDefaultSinkHooked = true
	}

	lines := strings.Split(out, "\n")

	for _, line := range lines {
		s := strings.TrimLeft(line, " \t")

		if !useAmixer && strings.Contains(s, "Name: "+string(sinkName)) {
			paDefaultSinkHooked = true
		}

		if useAmixer && strings.Contains(s, "Playback") && strings.Contains(s, "%") ||
			!useAmixer && paDefaultSinkHooked && strings.HasPrefix(s, "Volume:") {
			volumeStr := volumePattern.FindString(s)
			return strconv.Atoi(volumeStr[:len(volumeStr)-1])
		}
	}
	return 0, errors.New("no volume found")
}

func setVolumeCmd(volume int) []string {
	if useAmixer {
		return []string{"amixer", "set", "Master", strconv.Itoa(volume) + "%"}
	}
	return []string{"pactl", "set-sink-volume", "@DEFAULT_SINK@", strconv.Itoa(volume) + "%"}
}

func increaseVolumeCmd(diff int) []string {
	var sign string
	if diff >= 0 {
		sign = "+"
	} else if useAmixer {
		diff = -diff
		sign = "-"
	}
	if useAmixer {
		return []string{"amixer", "set", "Master", strconv.Itoa(diff) + "%" + sign}
	}
	return []string{"pactl", "--", "set-sink-volume", "@DEFAULT_SINK@", sign + strconv.Itoa(diff) + "%"}
}

func getMutedCmd() []string {
	if useAmixer {
		return []string{"amixer", "get", "Master"}
	}
	return []string{"pactl", "list", "sinks"}
}

func parseMuted(out string) (bool, error) {
	sinkName, sinkNameErr := getPADefaultSink()

	paDefaultSinkHooked := false

	if sinkNameErr != nil {
		paDefaultSinkHooked = true
	}

	lines := strings.Split(out, "\n")
	for _, line := range lines {
		s := strings.TrimLeft(line, " \t")

		if !useAmixer && strings.Contains(s, "Name: "+string(sinkName)) {
			paDefaultSinkHooked = true
		}

		if useAmixer && strings.Contains(s, "Playback") && strings.Contains(s, "%") ||
			!useAmixer && paDefaultSinkHooked && strings.HasPrefix(s, "Mute: ") {
			if strings.Contains(s, "[off]") || strings.Contains(s, "yes") {
				return true, nil
			} else if strings.Contains(s, "[on]") || strings.Contains(s, "no") {
				return false, nil
			}
		}
	}
	return false, errors.New("no muted information found")
}

func muteCmd() []string {
	if useAmixer {
		return []string{"amixer", "-D", "pulse", "set", "Master", "mute"}
	}
	return []string{"pactl", "set-sink-mute", "@DEFAULT_SINK@", "1"}
}

func unmuteCmd() []string {
	if useAmixer {
		return []string{"amixer", "-D", "pulse", "set", "Master", "unmute"}
	}
	return []string{"pactl", "set-sink-mute", "@DEFAULT_SINK@", "0"}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/version"
)

func storagePoolUpdate(state *state.State, name, newDescription string, newConfig map[string]string, withDB bool) error {
	// Handle the new logic
	pool, err := storagePools.GetPoolByName(state, name)
	if err != storageDrivers.ErrUnknownDriver {
		if err != nil {
			return err
		}

		return pool.Update(!withDB, newDescription, newConfig, nil)
	}

	// Old logic
	s, err := storagePoolInit(state, name)
	if err != nil {
		return err
	}

	oldWritable := s.GetStoragePoolWritable()
	newWritable := oldWritable

	// Backup the current state
	oldDescription := oldWritable.Description
	oldConfig := map[string]string{}
	err = shared.DeepCopy(&oldWritable.Config, &oldConfig)
	if err != nil {
		return err
	}

	// Define a function which reverts everything.  Defer this function
	// so that it doesn't need to be explicitly called in every failing
	// return path. Track whether or not we want to undo the changes
	// using a closure.
	undoChanges := true
	defer func() {
		if undoChanges {
			s.SetStoragePoolWritable(&oldWritable)
		}
	}()

	changedConfig, userOnly := storagePools.ConfigDiff(oldConfig, newConfig)
	// Apply config changes if there are any
	if len(changedConfig) != 0 {
		newWritable.Description = newDescription
		newWritable.Config = newConfig

		// Update the storage pool
		if !userOnly {
			if shared.StringInSlice("driver", changedConfig) {
				return fmt.Errorf("the \"driver\" property of a storage pool cannot be changed")
			}

			err = s.StoragePoolUpdate(&newWritable, changedConfig)
			if err != nil {
				return err
			}
		}

		// Apply the new configuration
		s.SetStoragePoolWritable(&newWritable)
	}

	// Update the database if something changed and the withDB flag is true
	// (i.e. this is not a clustering notification.
	if withDB && (len(changedConfig) != 0 || newDescription != oldDescription) {
		err = state.Cluster.StoragePoolUpdate(name, newDescription, newConfig)
		if err != nil {
			return err
		}
	}

	// Success, update the closure to mark that the changes should be kept.
	undoChanges = false

	return nil
}

// Report all LXD objects that are currently using the given storage pool.
// Volumes of type "custom" are not reported.
// /1.0/containers/alp1
// /1.0/containers/alp1/snapshots/snap0
// /1.0/images/cedce20b5b236f1071134beba7a5fd2aa923fda49eea4c66454dd559a5d6e906
// /1.0/profiles/default
func storagePoolUsedByGet(state *state.State, project string, poolID int64, poolName string) ([]string, error) {
	// Retrieve all non-custom volumes that exist on this storage pool.
	volumes, err := state.Cluster.StoragePoolNodeVolumesGet(project, poolID, []int{storagePoolVolumeTypeContainer, storagePoolVolumeTypeImage, storagePoolVolumeTypeCustom, storagePoolVolumeTypeVM})
	if err != nil && err != db.ErrNoSuchObject {
		return []string{}, err
	}

	// Retrieve all profiles that exist on this storage pool.
	profiles, err := profilesUsingPoolGetNames(state.Cluster, project, poolName)

	if err != nil {
		return []string{}, err
	}

	slicelen := len(volumes) + len(profiles)
	if slicelen == 0 {
		return []string{}, nil
	}

	// Save some allocation cycles by preallocating the correct len.
	poolUsedBy := make([]string, slicelen)
	for i := 0; i < len(volumes); i++ {
		apiEndpoint, _ := storagePoolVolumeTypeNameToAPIEndpoint(volumes[i].Type)
		switch apiEndpoint {
		case storagePoolVolumeAPIEndpointContainers:
			if strings.Index(volumes[i].Name, shared.SnapshotDelimiter) > 0 {
				parentName, snapOnlyName, _ := shared.InstanceGetParentAndSnapshotName(volumes[i].Name)
				poolUsedBy[i] = fmt.Sprintf("/%s/containers/%s/snapshots/%s", version.APIVersion, parentName, snapOnlyName)
			} else {
				poolUsedBy[i] = fmt.Sprintf("/%s/containers/%s", version.APIVersion, volumes[i].Name)
			}
		case storagePoolVolumeAPIEndpointVMs:
			if strings.Index(volumes[i].Name, shared.SnapshotDelimiter) > 0 {
				parentName, snapOnlyName, _ := shared.InstanceGetParentAndSnapshotName(volumes[i].Name)
				poolUsedBy[i] = fmt.Sprintf("/%s/virtual-machines/%s/snapshots/%s", version.APIVersion, parentName, snapOnlyName)
			} else {
				poolUsedBy[i] = fmt.Sprintf("/%s/virtual-machines/%s", version.APIVersion, volumes[i].Name)
			}
		case storagePoolVolumeAPIEndpointImages:
			poolUsedBy[i] = fmt.Sprintf("/%s/images/%s", version.APIVersion, volumes[i].Name)
		case storagePoolVolumeAPIEndpointCustom:
			poolUsedBy[i] = fmt.Sprintf("/%s/storage-pools/%s/volumes/%s/%s", version.APIVersion, poolName, volumes[i].Type, volumes[i].Name)
		default:
			// If that happens the db is busted, so report an error.
			return []string{}, fmt.Errorf("invalid storage type for storage volume \"%s\"", volumes[i].Name)
		}
	}

	for i := 0; i < len(profiles); i++ {
		poolUsedBy[i+len(volumes)] = fmt.Sprintf("/%s/profiles/%s", version.APIVersion, profiles[i])
	}

	return poolUsedBy, err
}

func profilesUsingPoolGetNames(db *db.Cluster, project string, poolName string) ([]string, error) {
	usedBy := []string{}

	profiles, err := db.Profiles(project)
	if err != nil {
		return usedBy, err
	}

	for _, pName := range profiles {
		_, profile, err := db.ProfileGet(project, pName)
		if err != nil {
			return usedBy, err
		}

		for _, v := range profile.Devices {
			if v["type"] != "disk" {
				continue
			}

			if v["pool"] == poolName {
				usedBy = append(usedBy, pName)
			}
		}
	}

	return usedBy, nil
}

func storagePoolDBCreate(s *state.State, poolName, poolDescription string, driver string, config map[string]string) (int64, error) {
	// Check that the storage pool does not already exist.
	_, err := s.Cluster.StoragePoolGetID(poolName)
	if err == nil {
		return -1, fmt.Errorf("The storage pool already exists")
	}

	// Make sure that we don't pass a nil to the next function.
	if config == nil {
		config = map[string]string{}
	}
	err = storagePoolValidate(poolName, driver, config)
	if err != nil {
		return -1, err
	}

	// Fill in the defaults
	err = storagePoolFillDefault(poolName, driver, config)
	if err != nil {
		return -1, err
	}

	// Create the database entry for the storage pool.
	id, err := dbStoragePoolCreateAndUpdateCache(s, poolName, poolDescription, driver, config)
	if err != nil {
		return -1, fmt.Errorf("Error inserting %s into database: %s", poolName, err)
	}

	return id, nil
}

func storagePoolValidate(poolName string, driverName string, config map[string]string) error {
	// Check if the storage pool name is valid.
	err := storagePools.ValidName(poolName)
	if err != nil {
		return err
	}

	// Validate the requested storage pool configuration.
	err = storagePoolValidateConfig(poolName, driverName, config, nil)
	if err != nil {
		return err
	}

	return nil
}

func storagePoolCreateGlobal(state *state.State, req api.StoragePoolsPost) error {
	// Create the database entry.
	id, err := storagePoolDBCreate(state, req.Name, req.Description, req.Driver, req.Config)
	if err != nil {
		return err
	}

	// Define a function which reverts everything.  Defer this function
	// so that it doesn't need to be explicitly called in every failing
	// return path. Track whether or not we want to undo the changes
	// using a closure.
	tryUndo := true
	defer func() {
		if !tryUndo {
			return
		}

		dbStoragePoolDeleteAndUpdateCache(state, req.Name)
	}()

	_, err = storagePoolCreateLocal(state, id, req, false)
	if err != nil {
		return err
	}

	tryUndo = false
	return nil
}

// This performs all non-db related work needed to create the pool.
func storagePoolCreateLocal(state *state.State, id int64, req api.StoragePoolsPost, isNotification bool) (map[string]string, error) {
	tryUndo := true

	// Make a copy of the req for later diff.
	var updatedConfig map[string]string
	var updatedReq api.StoragePoolsPost
	shared.DeepCopy(&req, &updatedReq)

	// Attempt to create using the new storage pool logic.
	pool, err := storagePools.CreatePool(state, id, &updatedReq, isNotification, nil)
	if err != storageDrivers.ErrUnknownDriver {
		if err != nil {
			return nil, err
		}

		// Mount the pool
		_, err = pool.Mount()
		if err != nil {
			return nil, err
		}

		// Record the updated config.
		updatedConfig = updatedReq.Config

		// Setup revert function.
		defer func() {
			if !tryUndo {
				return
			}

			pool.Delete(isNotification, nil)
		}()
	} else {
		// Load the old storage struct
		s, err := storagePoolInit(state, req.Name)
		if err != nil {
			return nil, err
		}

		// If this is a clustering notification for a ceph storage, we don't
		// want this node to actually create the pool, as it's already been
		// done by the node that triggered this notification. We just need to
		// create the storage pool directory.
		if s, ok := s.(*storageCeph); ok && isNotification {
			volumeMntPoint := storagePools.GetStoragePoolVolumeMountPoint(s.pool.Name, s.volume.Name)
			return nil, os.MkdirAll(volumeMntPoint, 0711)
		}

		// Create the pool
		err = s.StoragePoolCreate()
		if err != nil {
			return nil, err
		}

		updatedConfig = s.GetStoragePoolWritable().Config

		defer func() {
			if !tryUndo {
				return
			}
			s.StoragePoolDelete()
		}()
	}

	// In case the storage pool config was changed during the pool creation,
	// we need to update the database to reflect this change. This can e.g.
	// happen, when we create a loop file image. This means we append ".img"
	// to the path the user gave us and update the config in the storage
	// callback. So diff the config here to see if something like this has
	// happened.
	configDiff, _ := storagePools.ConfigDiff(req.Config, updatedConfig)
	if len(configDiff) > 0 {
		// Create the database entry for the storage pool.
		err = state.Cluster.StoragePoolUpdate(req.Name, req.Description, updatedConfig)
		if err != nil {
			return nil, fmt.Errorf("Error inserting %s into database: %s", req.Name, err)
		}
	}

	// Success, update the closure to mark that the changes should be kept.
	tryUndo = false

	return updatedConfig, nil
}

// Helper around the low-level DB API, which also updates the driver names cache.
func dbStoragePoolCreateAndUpdateCache(s *state.State, poolName string, poolDescription string, poolDriver string, poolConfig map[string]string) (int64, error) {
	id, err := s.Cluster.StoragePoolCreate(poolName, poolDescription, poolDriver, poolConfig)
	if err != nil {
		return id, err
	}

	// Update the storage drivers cache in api_1.0.go.
	storagePoolDriversCacheUpdate(s)

	return id, nil
}

// Helper around the low-level DB API, which also updates the driver names
// cache.
func dbStoragePoolDeleteAndUpdateCache(s *state.State, poolName string) error {
	_, err := s.Cluster.StoragePoolDelete(poolName)
	if err != nil {
		return err
	}

	// Update the storage drivers cache in api_1.0.go.
	storagePoolDriversCacheUpdate(s)

	return err
}
<commit_msg>lxd/storage/pools/utils: Comment on storagePoolDBCreate<commit_after>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/version"
)

func storagePoolUpdate(state *state.State, name, newDescription string, newConfig map[string]string, withDB bool) error {
	// Handle the new logic
	pool, err := storagePools.GetPoolByName(state, name)
	if err != storageDrivers.ErrUnknownDriver {
		if err != nil {
			return err
		}

		return pool.Update(!withDB, newDescription, newConfig, nil)
	}

	// Old logic
	s, err := storagePoolInit(state, name)
	if err != nil {
		return err
	}

	oldWritable := s.GetStoragePoolWritable()
	newWritable := oldWritable

	// Backup the current state
	oldDescription := oldWritable.Description
	oldConfig := map[string]string{}
	err = shared.DeepCopy(&oldWritable.Config, &oldConfig)
	if err != nil {
		return err
	}

	// Define a function which reverts everything.  Defer this function
	// so that it doesn't need to be explicitly called in every failing
	// return path. Track whether or not we want to undo the changes
	// using a closure.
	undoChanges := true
	defer func() {
		if undoChanges {
			s.SetStoragePoolWritable(&oldWritable)
		}
	}()

	changedConfig, userOnly := storagePools.ConfigDiff(oldConfig, newConfig)
	// Apply config changes if there are any
	if len(changedConfig) != 0 {
		newWritable.Description = newDescription
		newWritable.Config = newConfig

		// Update the storage pool
		if !userOnly {
			if shared.StringInSlice("driver", changedConfig) {
				return fmt.Errorf("the \"driver\" property of a storage pool cannot be changed")
			}

			err = s.StoragePoolUpdate(&newWritable, changedConfig)
			if err != nil {
				return err
			}
		}

		// Apply the new configuration
		s.SetStoragePoolWritable(&newWritable)
	}

	// Update the database if something changed and the withDB flag is true
	// (i.e. this is not a clustering notification.
	if withDB && (len(changedConfig) != 0 || newDescription != oldDescription) {
		err = state.Cluster.StoragePoolUpdate(name, newDescription, newConfig)
		if err != nil {
			return err
		}
	}

	// Success, update the closure to mark that the changes should be kept.
	undoChanges = false

	return nil
}

// Report all LXD objects that are currently using the given storage pool.
// Volumes of type "custom" are not reported.
// /1.0/containers/alp1
// /1.0/containers/alp1/snapshots/snap0
// /1.0/images/cedce20b5b236f1071134beba7a5fd2aa923fda49eea4c66454dd559a5d6e906
// /1.0/profiles/default
func storagePoolUsedByGet(state *state.State, project string, poolID int64, poolName string) ([]string, error) {
	// Retrieve all non-custom volumes that exist on this storage pool.
	volumes, err := state.Cluster.StoragePoolNodeVolumesGet(project, poolID, []int{storagePoolVolumeTypeContainer, storagePoolVolumeTypeImage, storagePoolVolumeTypeCustom, storagePoolVolumeTypeVM})
	if err != nil && err != db.ErrNoSuchObject {
		return []string{}, err
	}

	// Retrieve all profiles that exist on this storage pool.
	profiles, err := profilesUsingPoolGetNames(state.Cluster, project, poolName)

	if err != nil {
		return []string{}, err
	}

	slicelen := len(volumes) + len(profiles)
	if slicelen == 0 {
		return []string{}, nil
	}

	// Save some allocation cycles by preallocating the correct len.
	poolUsedBy := make([]string, slicelen)
	for i := 0; i < len(volumes); i++ {
		apiEndpoint, _ := storagePoolVolumeTypeNameToAPIEndpoint(volumes[i].Type)
		switch apiEndpoint {
		case storagePoolVolumeAPIEndpointContainers:
			if strings.Index(volumes[i].Name, shared.SnapshotDelimiter) > 0 {
				parentName, snapOnlyName, _ := shared.InstanceGetParentAndSnapshotName(volumes[i].Name)
				poolUsedBy[i] = fmt.Sprintf("/%s/containers/%s/snapshots/%s", version.APIVersion, parentName, snapOnlyName)
			} else {
				poolUsedBy[i] = fmt.Sprintf("/%s/containers/%s", version.APIVersion, volumes[i].Name)
			}
		case storagePoolVolumeAPIEndpointVMs:
			if strings.Index(volumes[i].Name, shared.SnapshotDelimiter) > 0 {
				parentName, snapOnlyName, _ := shared.InstanceGetParentAndSnapshotName(volumes[i].Name)
				poolUsedBy[i] = fmt.Sprintf("/%s/virtual-machines/%s/snapshots/%s", version.APIVersion, parentName, snapOnlyName)
			} else {
				poolUsedBy[i] = fmt.Sprintf("/%s/virtual-machines/%s", version.APIVersion, volumes[i].Name)
			}
		case storagePoolVolumeAPIEndpointImages:
			poolUsedBy[i] = fmt.Sprintf("/%s/images/%s", version.APIVersion, volumes[i].Name)
		case storagePoolVolumeAPIEndpointCustom:
			poolUsedBy[i] = fmt.Sprintf("/%s/storage-pools/%s/volumes/%s/%s", version.APIVersion, poolName, volumes[i].Type, volumes[i].Name)
		default:
			// If that happens the db is busted, so report an error.
			return []string{}, fmt.Errorf("invalid storage type for storage volume \"%s\"", volumes[i].Name)
		}
	}

	for i := 0; i < len(profiles); i++ {
		poolUsedBy[i+len(volumes)] = fmt.Sprintf("/%s/profiles/%s", version.APIVersion, profiles[i])
	}

	return poolUsedBy, err
}

func profilesUsingPoolGetNames(db *db.Cluster, project string, poolName string) ([]string, error) {
	usedBy := []string{}

	profiles, err := db.Profiles(project)
	if err != nil {
		return usedBy, err
	}

	for _, pName := range profiles {
		_, profile, err := db.ProfileGet(project, pName)
		if err != nil {
			return usedBy, err
		}

		for _, v := range profile.Devices {
			if v["type"] != "disk" {
				continue
			}

			if v["pool"] == poolName {
				usedBy = append(usedBy, pName)
			}
		}
	}

	return usedBy, nil
}

// storagePoolDBCreate creates a storage pool DB entry and returns the created Pool ID.
func storagePoolDBCreate(s *state.State, poolName, poolDescription string, driver string, config map[string]string) (int64, error) {
	// Check that the storage pool does not already exist.
	_, err := s.Cluster.StoragePoolGetID(poolName)
	if err == nil {
		return -1, fmt.Errorf("The storage pool already exists")
	}

	// Make sure that we don't pass a nil to the next function.
	if config == nil {
		config = map[string]string{}
	}
	err = storagePoolValidate(poolName, driver, config)
	if err != nil {
		return -1, err
	}

	// Fill in the defaults
	err = storagePoolFillDefault(poolName, driver, config)
	if err != nil {
		return -1, err
	}

	// Create the database entry for the storage pool.
	id, err := dbStoragePoolCreateAndUpdateCache(s, poolName, poolDescription, driver, config)
	if err != nil {
		return -1, fmt.Errorf("Error inserting %s into database: %s", poolName, err)
	}

	return id, nil
}

func storagePoolValidate(poolName string, driverName string, config map[string]string) error {
	// Check if the storage pool name is valid.
	err := storagePools.ValidName(poolName)
	if err != nil {
		return err
	}

	// Validate the requested storage pool configuration.
	err = storagePoolValidateConfig(poolName, driverName, config, nil)
	if err != nil {
		return err
	}

	return nil
}

func storagePoolCreateGlobal(state *state.State, req api.StoragePoolsPost) error {
	// Create the database entry.
	id, err := storagePoolDBCreate(state, req.Name, req.Description, req.Driver, req.Config)
	if err != nil {
		return err
	}

	// Define a function which reverts everything.  Defer this function
	// so that it doesn't need to be explicitly called in every failing
	// return path. Track whether or not we want to undo the changes
	// using a closure.
	tryUndo := true
	defer func() {
		if !tryUndo {
			return
		}

		dbStoragePoolDeleteAndUpdateCache(state, req.Name)
	}()

	_, err = storagePoolCreateLocal(state, id, req, false)
	if err != nil {
		return err
	}

	tryUndo = false
	return nil
}

// This performs all non-db related work needed to create the pool.
func storagePoolCreateLocal(state *state.State, id int64, req api.StoragePoolsPost, isNotification bool) (map[string]string, error) {
	tryUndo := true

	// Make a copy of the req for later diff.
	var updatedConfig map[string]string
	var updatedReq api.StoragePoolsPost
	shared.DeepCopy(&req, &updatedReq)

	// Attempt to create using the new storage pool logic.
	pool, err := storagePools.CreatePool(state, id, &updatedReq, isNotification, nil)
	if err != storageDrivers.ErrUnknownDriver {
		if err != nil {
			return nil, err
		}

		// Mount the pool
		_, err = pool.Mount()
		if err != nil {
			return nil, err
		}

		// Record the updated config.
		updatedConfig = updatedReq.Config

		// Setup revert function.
		defer func() {
			if !tryUndo {
				return
			}

			pool.Delete(isNotification, nil)
		}()
	} else {
		// Load the old storage struct
		s, err := storagePoolInit(state, req.Name)
		if err != nil {
			return nil, err
		}

		// If this is a clustering notification for a ceph storage, we don't
		// want this node to actually create the pool, as it's already been
		// done by the node that triggered this notification. We just need to
		// create the storage pool directory.
		if s, ok := s.(*storageCeph); ok && isNotification {
			volumeMntPoint := storagePools.GetStoragePoolVolumeMountPoint(s.pool.Name, s.volume.Name)
			return nil, os.MkdirAll(volumeMntPoint, 0711)
		}

		// Create the pool
		err = s.StoragePoolCreate()
		if err != nil {
			return nil, err
		}

		updatedConfig = s.GetStoragePoolWritable().Config

		defer func() {
			if !tryUndo {
				return
			}
			s.StoragePoolDelete()
		}()
	}

	// In case the storage pool config was changed during the pool creation,
	// we need to update the database to reflect this change. This can e.g.
	// happen, when we create a loop file image. This means we append ".img"
	// to the path the user gave us and update the config in the storage
	// callback. So diff the config here to see if something like this has
	// happened.
	configDiff, _ := storagePools.ConfigDiff(req.Config, updatedConfig)
	if len(configDiff) > 0 {
		// Create the database entry for the storage pool.
		err = state.Cluster.StoragePoolUpdate(req.Name, req.Description, updatedConfig)
		if err != nil {
			return nil, fmt.Errorf("Error inserting %s into database: %s", req.Name, err)
		}
	}

	// Success, update the closure to mark that the changes should be kept.
	tryUndo = false

	return updatedConfig, nil
}

// Helper around the low-level DB API, which also updates the driver names cache.
func dbStoragePoolCreateAndUpdateCache(s *state.State, poolName string, poolDescription string, poolDriver string, poolConfig map[string]string) (int64, error) {
	id, err := s.Cluster.StoragePoolCreate(poolName, poolDescription, poolDriver, poolConfig)
	if err != nil {
		return id, err
	}

	// Update the storage drivers cache in api_1.0.go.
	storagePoolDriversCacheUpdate(s)

	return id, nil
}

// Helper around the low-level DB API, which also updates the driver names
// cache.
func dbStoragePoolDeleteAndUpdateCache(s *state.State, poolName string) error {
	_, err := s.Cluster.StoragePoolDelete(poolName)
	if err != nil {
		return err
	}

	// Update the storage drivers cache in api_1.0.go.
	storagePoolDriversCacheUpdate(s)

	return err
}
<|endoftext|>


<commit_before>package m_etcd_test

import (
	"errors"
	"testing"
	"time"

	"github.com/coreos/go-etcd/etcd"
	"github.com/lytics/metafora"
	"github.com/lytics/metafora/m_etcd"
	"github.com/lytics/metafora/m_etcd/testutil"
	"github.com/lytics/metafora/statemachine"
)

// TestAll is an integration test for all of m_etcd's components.
//
// While huge integration tests like this are rarely desirable as they can be
// overly fragile and complex, I found myself manually repeating the tests I've
// automated here over and over. This is far more reliable than expecting
// developers to do adhoc testing of all of the m_etcd package's features.
func TestAll(t *testing.T) {
	etcdc, hosts := testutil.NewEtcdClient(t)
	t.Parallel()

	const recursive = true
	etcdc.Delete("test-a", recursive)
	etcdc.Delete("test-b", recursive)

	h := func(tid string, cmds <-chan statemachine.Message) statemachine.Message {
		cmd := <-cmds
		if tid == "error-test" {
			return statemachine.Message{Code: statemachine.Error, Err: errors.New("error-test")}
		}
		return cmd
	}

	newC := func(name, ns string) *metafora.Consumer {
		coord, hf, bal, err := m_etcd.New(name, ns, hosts, h)
		if err != nil {
			t.Fatalf("Error creating new etcd stack: %v", err)
		}
		cons, err := metafora.NewConsumer(coord, hf, bal)
		if err != nil {
			t.Fatalf("Error creating consumer %s:%s: %v", ns, name, err)
		}
		go cons.Run()
		return cons
	}
	// Start 4 consumers, 2 per namespace
	cons1a := newC("node1", "test-a")
	cons2a := newC("node2", "test-a")
	cons1b := newC("node1", "test-b")
	cons2b := newC("node2", "test-b")

	// Create clients and start some tests
	cliA := m_etcd.NewClient("test-a", hosts)
	cliB := m_etcd.NewClient("test-b", hosts)

	if err := cliA.SubmitTask("task1"); err != nil {
		t.Fatalf("Error submitting task1 to a: %v", err)
	}
	if err := cliB.SubmitTask("task1"); err != nil {
		t.Fatalf("Error submitting task1 to b: %v", err)
	}

	// Give consumers a bit to pick up tasks
	time.Sleep(250 * time.Millisecond)

	assertRunning := func(tid string, cons ...*metafora.Consumer) {
		found := false
		for _, c := range cons {
			tasks := c.Tasks()
			if len(tasks) > 0 && found {
				t.Fatal("Task already found running but another task is running on a different consumer")
			}
			if len(tasks) > 1 {
				t.Fatalf("Expected at most 1 task, but found: %d", len(tasks))
			}
			if len(tasks) == 1 && tasks[0].ID() == tid {
				found = true
			}
		}
		if !found {
			t.Fatalf("Could not find task=%q", tid)
		}
	}

	assertRunning("task1", cons1a, cons2a)
	assertRunning("task1", cons1b, cons2b)

	// Kill task1 in A
	{
		cmdr := m_etcd.NewCommander("test-a", etcdc)
		if err := cmdr.Send("task1", statemachine.Message{Code: statemachine.Kill}); err != nil {
			t.Fatalf("Error sending kill to task1: %v", err)
		}
		time.Sleep(250 * time.Millisecond)

		for _, c := range []*metafora.Consumer{cons1a, cons2a} {
			tasks := c.Tasks()
			if len(tasks) != 0 {
				t.Fatalf("Expected no tasks but found: %d", len(tasks))
			}
		}
	}

	// Submit a bunch of tasks to A
	{
		tasks := []string{"task2", "task3", "task4", "task5", "task6", "task7"}
		for _, tid := range tasks {
			if err := cliA.SubmitTask(tid); err != nil {
				t.Fatalf("Error submitting task=%q to A: %v", tid, err)
			}
		}

		// Give them time to start
		time.Sleep(500 * time.Millisecond)

		// Ensure they're balanced
		if err := cliA.SubmitCommand("node1", metafora.CommandBalance()); err != nil {
			t.Fatalf("Error submitting balance command to cons1a: %v", err)
		}
		time.Sleep(500 * time.Millisecond)
		if err := cliA.SubmitCommand("node2", metafora.CommandBalance()); err != nil {
			t.Fatalf("Error submitting balance command to cons1a: %v", err)
		}

		a1tasks := cons1a.Tasks()
		a2tasks := cons2a.Tasks()
		for _, task := range a1tasks {
			metafora.Debug("A1: ", task.ID(), " - ", task.Stopped().IsZero())
		}
		for _, task := range a2tasks {
			metafora.Debug("A2: ", task.ID(), " - ", task.Stopped().IsZero())
		}
		time.Sleep(500 * time.Millisecond)

		a1tasks = cons1a.Tasks()
		a2tasks = cons2a.Tasks()
		if len(a1tasks) < 2 || len(a1tasks) > 4 || len(a2tasks) < 2 || len(a2tasks) > 4 {
			t.Fatalf("Namespace A isn't fairly balanced: node1: %d; node2: %d", len(a1tasks), len(a2tasks))
		}

		// Shutting down a consumer should migrate all tasks to the other
		cons1a.Shutdown()
		time.Sleep(500 * time.Millisecond)

		a2tasks = cons2a.Tasks()
		if len(a2tasks) != len(tasks) {
			t.Fatalf("Consumer 2a should have received all %d tasks but only has %d.", len(tasks), len(a2tasks))
		}
	}

	// Use Namespace B to check Error state handling
	{
		tasks := []string{"task8", "error-test"}
		for _, tid := range tasks {
			if err := cliB.SubmitTask(tid); err != nil {
				t.Fatalf("Error submitting task=%q to B: %v", tid, err)
			}
		}

		// Give them time to start
		time.Sleep(500 * time.Millisecond)

		n := len(cons1b.Tasks()) + len(cons2b.Tasks())
		if n != 3 {
			t.Fatalf("Expected B to be running 3 tasks but found %d", n)
		}

		// Resuming error-test 8*2 times should cause it to be failed
		cmdr := m_etcd.NewCommander("test-b", etcdc)
		for i := 0; i < statemachine.DefaultErrMax*2; i++ {
			if err := cmdr.Send("error-test", statemachine.Message{Code: statemachine.Run}); err != nil {
				t.Fatalf("Unexpected error resuming error-test in B: %v", err)
			}
			time.Sleep(200 * time.Millisecond)
		}

		n = len(cons1b.Tasks()) + len(cons2b.Tasks())
		if n != 2 {
			t.Fatalf("Expected B to be running 2 tasks but found %d", n)
		}

		// Resubmitting a failed task shouldn't error but also shouldn't run.
		if err := cliB.SubmitTask("error-test"); err != nil {
			t.Fatalf("Error resubmitting error-test task to B: %v", err)
		}

		// Give the statemachine a moment to load the initial state and exit
		time.Sleep(200 * time.Millisecond)

		n = len(cons1b.Tasks()) + len(cons2b.Tasks())
		if n != 2 {
			t.Fatalf("Expected B to be running 2 tasks but found %d", n)
		}
	}

	// Shutdown
	cons2a.Shutdown()
	cons1b.Shutdown()
	cons2b.Shutdown()

	// Make sure everything is cleaned up
	respA, err := etcdc.Get("/test-a/tasks", true, true)
	if err != nil {
		t.Fatalf("Error getting tasks from etcd: %v", err)
	}
	respB, err := etcdc.Get("/test-b/tasks", true, true)
	if err != nil {
		t.Fatalf("Error getting tasks from etcd: %v", err)
	}

	nodes := []*etcd.Node{}
	nodes = append(nodes, respA.Node.Nodes...)
	nodes = append(nodes, respB.Node.Nodes...)
	for _, node := range nodes {
		if len(node.Nodes) > 0 {
			t.Fatalf("%s has %d nodes. First key: %s", node.Key, len(node.Nodes), node.Nodes[0].Key)
		}
	}
}

// TestTaskResurrectionInt ensures that a Claim won't recreate a task that had
// been deleted (marked as done). taskmgr has a non-integration version of this
// test.
func TestTaskResurrectionInt(t *testing.T) {
	etcdc, hosts := testutil.NewEtcdClient(t)
	t.Parallel()

	const recursive = true
	etcdc.Delete("test-resurrect", recursive)

	coord, err := m_etcd.NewEtcdCoordinator("r-node", "test-resurrect", hosts)
	if err != nil {
		t.Fatalf("Error creating coordinator: %v", err)
	}
	if err := coord.Init(nil); err != nil {
		t.Fatalf("Error initializing coordinator: %v", err)
	}

	// Try to claim a nonexistent
	if claimed := coord.Claim("xyz"); claimed {
		t.Fatal("Claiming a nonexistent task should not work but did!")
	}

	// Create a task, mark it as done, and try to claim it again
	client := m_etcd.NewClient("test-resurrect", hosts)
	if err := client.SubmitTask("xyz"); err != nil {
		t.Fatalf("Error submitting task xyz: %v", err)
	}

	if claimed := coord.Claim("xyz"); !claimed {
		t.Fatal("Failed to claim task xyz")
	}

	coord.Done("xyz")

	if claimed := coord.Claim("xyz"); claimed {
		t.Fatal("Reclaimed task that was marked as done.")
	}
}
<commit_msg>adding a test for sleeping<commit_after>package m_etcd_test

import (
	"errors"
	"testing"
	"time"

	"github.com/coreos/go-etcd/etcd"
	"github.com/lytics/metafora"
	"github.com/lytics/metafora/m_etcd"
	"github.com/lytics/metafora/m_etcd/testutil"
	"github.com/lytics/metafora/statemachine"
)

// TestSleepTest is an integration test for all of m_etcd's components.
//
func TestSleepTest(t *testing.T) {
	etcdc, hosts := testutil.NewEtcdClient(t)
	t.Parallel()
	const namespace = "sleeptest-metafora"
	const nonsleepingtask = "task1"
	const sleepingtasks = "sleeping-task1"

	const recursive = true
	etcdc.Delete(namespace, recursive)

	h := func(tid string, cmds <-chan statemachine.Message) statemachine.Message {

		if tid == sleepingtasks {
			sleeptil := 5 * time.Second
			nextstarttime := (time.Now().Add(sleeptil))
			t.Logf("sleeping task:%v sleepfor:%v", tid, nextstarttime)
			return statemachine.Message{Code: statemachine.Sleep, Until: &nextstarttime}
		}

		cmd := <-cmds
		t.Logf("non sleeping task:%v", tid)

		return cmd
	}

	newC := func(name, ns string) *metafora.Consumer {
		coord, hf, bal, err := m_etcd.New(name, ns, hosts, h)
		if err != nil {
			t.Fatalf("Error creating new etcd stack: %v", err)
		}
		cons, err := metafora.NewConsumer(coord, hf, bal)
		if err != nil {
			t.Fatalf("Error creating consumer %s:%s: %v", ns, name, err)
		}
		go cons.Run()
		return cons
	}
	// Start 2 consumers
	cons1 := newC("node1", namespace)
	cons2 := newC("node2", namespace)

	// Create clients and start some tests
	cliA := m_etcd.NewClient(namespace, hosts)

	if err := cliA.SubmitTask(nonsleepingtask); err != nil {
		t.Fatalf("Error submitting task1 to a: %v", err)
	}

	// Give consumers a bit to pick up tasks
	time.Sleep(500 * time.Millisecond)

	assertRunning := func(tid string, cons ...*metafora.Consumer) {
		found := false
		for _, c := range cons {
			tasks := c.Tasks()
			if len(tasks) > 0 && found {
				t.Fatal("Task already found running but another task is running on a different consumer")
			}
			if len(tasks) > 1 {
				t.Fatalf("Expected at most 1 task, but found: %d", len(tasks))
			}
			if len(tasks) == 1 && tasks[0].ID() == tid {
				found = true
			}
		}
		if !found {
			t.Fatalf("Could not find task=%q", tid)
		}
	}

	assertRunning(nonsleepingtask, cons1, cons2)

	// Kill task1 in A
	{
		cmdr := m_etcd.NewCommander(namespace, etcdc)
		if err := cmdr.Send(nonsleepingtask, statemachine.Message{Code: statemachine.Kill}); err != nil {
			t.Fatalf("Error sending kill to task1: %v", err)
		}
		time.Sleep(400 * time.Millisecond)

		for _, c := range []*metafora.Consumer{cons1, cons2} {
			tasks := c.Tasks()
			if len(tasks) != 0 {
				t.Fatalf("Expected no tasks but found: %d", len(tasks))
			}
		}
	}

	if err := cliA.SubmitTask(sleepingtasks); err != nil {
		t.Fatalf("Error submitting task1 to a: %v", err)
	}

	// Give consumers a bit to pick up tasks
	time.Sleep(1 * time.Second)

	assertRunning(sleepingtasks, cons1, cons2)

	wait1 := make(chan bool)
	go func() {
		defer close(wait1)
		// Shutdown
		cons1.Shutdown()
		cons2.Shutdown()
	}()

	timeout := time.NewTimer(5 * time.Second)
	select {
	case <-wait1:
	case <-timeout.C:
		t.Fatalf("failed waiting for shutdown")
	}

	//	make sure all tasks are released
	for _, c := range []*metafora.Consumer{cons1, cons2} {
		tasks := c.Tasks()
		for _, work := range tasks {
			t.Fatalf("work id %v is still running", work)
		}
	}
}

// TestAll is an integration test for all of m_etcd's components.
//
// While huge integration tests like this are rarely desirable as they can be
// overly fragile and complex, I found myself manually repeating the tests I've
// automated here over and over. This is far more reliable than expecting
// developers to do adhoc testing of all of the m_etcd package's features.
func TestAll(t *testing.T) {
	etcdc, hosts := testutil.NewEtcdClient(t)
	t.Parallel()

	const recursive = true
	etcdc.Delete("test-a", recursive)
	etcdc.Delete("test-b", recursive)

	h := func(tid string, cmds <-chan statemachine.Message) statemachine.Message {
		cmd := <-cmds
		if tid == "error-test" {
			return statemachine.Message{Code: statemachine.Error, Err: errors.New("error-test")}
		}
		return cmd
	}

	newC := func(name, ns string) *metafora.Consumer {
		coord, hf, bal, err := m_etcd.New(name, ns, hosts, h)
		if err != nil {
			t.Fatalf("Error creating new etcd stack: %v", err)
		}
		cons, err := metafora.NewConsumer(coord, hf, bal)
		if err != nil {
			t.Fatalf("Error creating consumer %s:%s: %v", ns, name, err)
		}
		go cons.Run()
		return cons
	}
	// Start 4 consumers, 2 per namespace
	cons1a := newC("node1", "test-a")
	cons2a := newC("node2", "test-a")
	cons1b := newC("node1", "test-b")
	cons2b := newC("node2", "test-b")

	// Create clients and start some tests
	cliA := m_etcd.NewClient("test-a", hosts)
	cliB := m_etcd.NewClient("test-b", hosts)

	if err := cliA.SubmitTask("task1"); err != nil {
		t.Fatalf("Error submitting task1 to a: %v", err)
	}
	if err := cliB.SubmitTask("task1"); err != nil {
		t.Fatalf("Error submitting task1 to b: %v", err)
	}

	// Give consumers a bit to pick up tasks
	time.Sleep(250 * time.Millisecond)

	assertRunning := func(tid string, cons ...*metafora.Consumer) {
		found := false
		for _, c := range cons {
			tasks := c.Tasks()
			if len(tasks) > 0 && found {
				t.Fatal("Task already found running but another task is running on a different consumer")
			}
			if len(tasks) > 1 {
				t.Fatalf("Expected at most 1 task, but found: %d", len(tasks))
			}
			if len(tasks) == 1 && tasks[0].ID() == tid {
				found = true
			}
		}
		if !found {
			t.Fatalf("Could not find task=%q", tid)
		}
	}

	assertRunning("task1", cons1a, cons2a)
	assertRunning("task1", cons1b, cons2b)

	// Kill task1 in A
	{
		cmdr := m_etcd.NewCommander("test-a", etcdc)
		if err := cmdr.Send("task1", statemachine.Message{Code: statemachine.Kill}); err != nil {
			t.Fatalf("Error sending kill to task1: %v", err)
		}
		time.Sleep(250 * time.Millisecond)

		for _, c := range []*metafora.Consumer{cons1a, cons2a} {
			tasks := c.Tasks()
			if len(tasks) != 0 {
				t.Fatalf("Expected no tasks but found: %d", len(tasks))
			}
		}
	}

	// Submit a bunch of tasks to A
	{
		tasks := []string{"task2", "task3", "task4", "task5", "task6", "task7"}
		for _, tid := range tasks {
			if err := cliA.SubmitTask(tid); err != nil {
				t.Fatalf("Error submitting task=%q to A: %v", tid, err)
			}
		}

		// Give them time to start
		time.Sleep(800 * time.Millisecond)

		// Ensure they're balanced
		if err := cliA.SubmitCommand("node1", metafora.CommandBalance()); err != nil {
			t.Fatalf("Error submitting balance command to cons1a: %v", err)
		}
		time.Sleep(800 * time.Millisecond)
		if err := cliA.SubmitCommand("node2", metafora.CommandBalance()); err != nil {
			t.Fatalf("Error submitting balance command to cons1a: %v", err)
		}

		a1tasks := cons1a.Tasks()
		a2tasks := cons2a.Tasks()
		for _, task := range a1tasks {
			metafora.Debug("A1: ", task.ID(), " - ", task.Stopped().IsZero())
		}
		for _, task := range a2tasks {
			metafora.Debug("A2: ", task.ID(), " - ", task.Stopped().IsZero())
		}
		time.Sleep(800 * time.Millisecond)

		a1tasks = cons1a.Tasks()
		a2tasks = cons2a.Tasks()
		if len(a1tasks) < 2 || len(a1tasks) > 4 || len(a2tasks) < 2 || len(a2tasks) > 4 {
			t.Fatalf("Namespace A isn't fairly balanced: node1: %d; node2: %d", len(a1tasks), len(a2tasks))
		}

		// Shutting down a consumer should migrate all tasks to the other
		cons1a.Shutdown()
		time.Sleep(800 * time.Millisecond)

		a2tasks = cons2a.Tasks()
		if len(a2tasks) != len(tasks) {
			t.Fatalf("Consumer 2a should have received all %d tasks but only has %d.", len(tasks), len(a2tasks))
		}
	}

	// Use Namespace B to check Error state handling
	{
		tasks := []string{"task8", "error-test"}
		for _, tid := range tasks {
			if err := cliB.SubmitTask(tid); err != nil {
				t.Fatalf("Error submitting task=%q to B: %v", tid, err)
			}
		}

		// Give them time to start
		time.Sleep(500 * time.Millisecond)

		n := len(cons1b.Tasks()) + len(cons2b.Tasks())
		if n != 3 {
			t.Fatalf("Expected B to be running 3 tasks but found %d", n)
		}

		// Resuming error-test 8*2 times should cause it to be failed
		cmdr := m_etcd.NewCommander("test-b", etcdc)
		for i := 0; i < statemachine.DefaultErrMax*2; i++ {
			if err := cmdr.Send("error-test", statemachine.Message{Code: statemachine.Run}); err != nil {
				t.Fatalf("Unexpected error resuming error-test in B: %v", err)
			}
			time.Sleep(200 * time.Millisecond)
		}

		n = len(cons1b.Tasks()) + len(cons2b.Tasks())
		if n != 2 {
			t.Fatalf("Expected B to be running 2 tasks but found %d", n)
		}

		// Resubmitting a failed task shouldn't error but also shouldn't run.
		if err := cliB.SubmitTask("error-test"); err != nil {
			t.Fatalf("Error resubmitting error-test task to B: %v", err)
		}

		// Give the statemachine a moment to load the initial state and exit
		time.Sleep(200 * time.Millisecond)

		n = len(cons1b.Tasks()) + len(cons2b.Tasks())
		if n != 2 {
			t.Fatalf("Expected B to be running 2 tasks but found %d", n)
		}
	}

	// Shutdown
	cons2a.Shutdown()
	cons1b.Shutdown()
	cons2b.Shutdown()

	// Make sure everything is cleaned up
	respA, err := etcdc.Get("/test-a/tasks", true, true)
	if err != nil {
		t.Fatalf("Error getting tasks from etcd: %v", err)
	}
	respB, err := etcdc.Get("/test-b/tasks", true, true)
	if err != nil {
		t.Fatalf("Error getting tasks from etcd: %v", err)
	}

	nodes := []*etcd.Node{}
	nodes = append(nodes, respA.Node.Nodes...)
	nodes = append(nodes, respB.Node.Nodes...)
	for _, node := range nodes {
		if len(node.Nodes) > 0 {
			t.Fatalf("%s has %d nodes. First key: %s", node.Key, len(node.Nodes), node.Nodes[0].Key)
		}
	}
}

// TestTaskResurrectionInt ensures that a Claim won't recreate a task that had
// been deleted (marked as done). taskmgr has a non-integration version of this
// test.
func TestTaskResurrectionInt(t *testing.T) {
	etcdc, hosts := testutil.NewEtcdClient(t)
	t.Parallel()

	const recursive = true
	etcdc.Delete("test-resurrect", recursive)

	coord, err := m_etcd.NewEtcdCoordinator("r-node", "test-resurrect", hosts)
	if err != nil {
		t.Fatalf("Error creating coordinator: %v", err)
	}
	if err := coord.Init(nil); err != nil {
		t.Fatalf("Error initializing coordinator: %v", err)
	}

	// Try to claim a nonexistent
	if claimed := coord.Claim("xyz"); claimed {
		t.Fatal("Claiming a nonexistent task should not work but did!")
	}

	// Create a task, mark it as done, and try to claim it again
	client := m_etcd.NewClient("test-resurrect", hosts)
	if err := client.SubmitTask("xyz"); err != nil {
		t.Fatalf("Error submitting task xyz: %v", err)
	}

	if claimed := coord.Claim("xyz"); !claimed {
		t.Fatal("Failed to claim task xyz")
	}

	coord.Done("xyz")

	if claimed := coord.Claim("xyz"); claimed {
		t.Fatal("Reclaimed task that was marked as done.")
	}
}
<|endoftext|>


<commit_before>package vxlan

import (
	//"fmt"
	//"strings"
	//"time"
	"net"
	"strconv"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/go-plugins-helpers/network"
	//"github.com/samalba/dockerclient"
	//"github.com/davecgh/go-spew/spew"
	"github.com/vishvananda/netlink"
)

type Driver struct {
	network.Driver
	networks map[string]*NetworkState
}

// NetworkState is filled in at network creation time
// it contains state that we wish to keep for each network
type NetworkState struct {
	Bridge   *netlink.Bridge
	VXLan    *netlink.Vxlan
	Gateway  string
	IPv4Data []*network.IPAMData
	IPv6Data []*network.IPAMData
}

func NewDriver() (*Driver, error) {
	d := &Driver{
		networks: make(map[string]*NetworkState),
	}
	return d, nil
}

func (d *Driver) CreateNetwork(r *network.CreateNetworkRequest) error {
	log.Debugf("Create network request: %+v", r)

	netID := r.NetworkID
	var err error

	bridgeName := "br_" + netID[:12]
	vxlanName := "vx_" + netID[:12]

	// get interface names from options first
	for k, v := range r.Options {
		if k == "com.docker.network.generic" {
			if genericOpts, ok := v.(map[string]interface{}); ok {
				for key, val := range genericOpts {
					if key == "vxlanName" {
						vxlanName = val.(string)
					}
					if key == "bridgeName" {
						bridgeName = val.(string)
					}
				}
			}
		}
	}

	// create links
	bridge := &netlink.Bridge{
		LinkAttrs: netlink.LinkAttrs{
			Name: bridgeName,
		},
	}
	vxlan := &netlink.Vxlan{
		LinkAttrs: netlink.LinkAttrs{
			Name: vxlanName,
		},
	}

	// Parse interface options
	for k, v := range r.Options {
		if k == "com.docker.network.generic" {
			if genericOpts, ok := v.(map[string]interface{}); ok {
				for key, val := range genericOpts {
					log.Debugf("Libnetwork Opts Sent: [ %s ] Value: [ %s ]", key, val)
					if key == "vxlanMTU" {
						vxlan.LinkAttrs.MTU, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "bridgeMTU" {
						bridge.LinkAttrs.MTU, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "vxlanHardwareAddr" {
						vxlan.LinkAttrs.HardwareAddr, err = net.ParseMAC(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "bridgeHardwareAddr" {
						bridge.LinkAttrs.HardwareAddr, err = net.ParseMAC(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "vxlanTxQLen" {
						vxlan.LinkAttrs.TxQLen, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "bridgeTxQLen" {
						bridge.LinkAttrs.TxQLen, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "VxlanId" {
						vxlan.VxlanId, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "VtepDev" {
						vtepDev, err := netlink.LinkByName(val.(string))
						if err != nil {
							return err
						}
						vxlan.VtepDevIndex = vtepDev.Attrs().Index
					}
					if key == "SrcAddr" {
						vxlan.SrcAddr = net.ParseIP(val.(string))
					}
					if key == "Group" {
						vxlan.Group = net.ParseIP(val.(string))
					}
					if key == "TTL" {
						vxlan.TTL, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "TOS" {
						vxlan.TOS, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Learning" {
						vxlan.Learning, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Proxy" {
						vxlan.Proxy, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "RSC" {
						vxlan.RSC, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "L2miss" {
						vxlan.L2miss, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "L3miss" {
						vxlan.L3miss, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "NoAge" {
						vxlan.NoAge, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "GBP" {
						vxlan.GBP, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Age" {
						vxlan.Age, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Limit" {
						vxlan.Limit, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Port" {
						vxlan.Port, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "PortLow" {
						vxlan.PortLow, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "PortHigh" {
						vxlan.PortHigh, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
				}
			}
		}
	}
	// Done parsing options

	// delete links if they already exist, don't worry about errors
	netlink.LinkDel(bridge)
	netlink.LinkDel(vxlan)

	// add links
	err = netlink.LinkAdd(bridge)
	if err != nil {
		return err
	}
	err = netlink.LinkAdd(vxlan)
	if err != nil {
		return err
	}

	// add vxlan to bridge
	err = netlink.LinkSetMaster(vxlan, bridge)
	if err != nil {
		return err
	}

	// bring interfaces up
	err = netlink.LinkSetUp(bridge)
	if err != nil {
		return err
	}
	err = netlink.LinkSetUp(vxlan)
	if err != nil {
		return err
	}

	// store interfaces to be used later
	ns := &NetworkState{
		VXLan:    vxlan,
		Bridge:   bridge,
		IPv4Data: r.IPv4Data,
		IPv6Data: r.IPv6Data,
	}

	// Add IPs to interfaces
	// Process IPv6 first. If both are inclued, IPv4 gateway will be the one that
	// remains, because JoinResponse can only include one Gateway
	for i := range r.IPv6Data {
		gatewayIP, err := netlink.ParseAddr(r.IPv6Data[i].Gateway)
		if err != nil {
			return err
		}
		ns.Gateway = gatewayIP.IP.String()
		netlink.AddrAdd(bridge, gatewayIP)
	}
	for i := range r.IPv4Data {
		gatewayIP, err := netlink.ParseAddr(r.IPv4Data[i].Gateway)
		if err != nil {
			return err
		}
		ns.Gateway = gatewayIP.IP.String()
		netlink.AddrAdd(bridge, gatewayIP)
	}

	d.networks[netID] = ns

	return nil
}

func (d *Driver) DeleteNetwork(r *network.DeleteNetworkRequest) error {
	netID := r.NetworkID

	err := netlink.LinkDel(d.networks[netID].VXLan)
	if err != nil {
		return err
	}
	err = netlink.LinkDel(d.networks[netID].Bridge)
	if err != nil {
		return err
	}

	return nil
}

func (d *Driver) CreateEndpoint(r *network.CreateEndpointRequest) error {
	log.Debugf("Create endpoint request: %+v", r)
	return nil
}

func (d *Driver) DeleteEndpoint(r *network.DeleteEndpointRequest) error {
	log.Debugf("Delete endpoint request: %+v", r)
	return nil
}

func (d *Driver) EndpointInfo(r *network.InfoRequest) (*network.InfoResponse, error) {
	res := &network.InfoResponse{
		Value: make(map[string]string),
	}
	return res, nil
}

func (d *Driver) Join(r *network.JoinRequest) (*network.JoinResponse, error) {
	// create and attach local name to the bridge
	veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: "veth_" + r.EndpointID[:5]},
		PeerName:  "ethc" + r.EndpointID[:5],
	}
	if err := netlink.LinkAdd(veth); err != nil {
		log.Errorf("failed to create the veth pair named: [ %v ] error: [ %s ] ", veth, err)
		return nil, err
	}

	// bring up the veth pair
	err := netlink.LinkSetUp(veth)
	if err != nil {
		log.Warnf("Error enabling  Veth local iface: [ %v ]", veth)
		return nil, err
	}

	bridge := d.networks[r.NetworkID].Bridge
	// add veth to bridge
	err = netlink.LinkSetMaster(veth, bridge)
	if err != nil {
		return nil, err
	}

	// SrcName gets renamed to DstPrefix + ID on the container iface
	res := &network.JoinResponse{
		InterfaceName: network.InterfaceName{
			SrcName:   veth.PeerName,
			DstPrefix: "eth",
		},
		Gateway: d.networks[r.NetworkID].Gateway,
	}
	log.Debugf("Join endpoint %s:%s to %s", r.NetworkID, r.EndpointID, r.SandboxKey)
	return res, nil
}

func (d *Driver) Leave(r *network.LeaveRequest) error {
	log.Debugf("Leave request: %+v", r)

	veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: "veth_" + r.EndpointID[:5]},
		PeerName:  "ethc" + r.EndpointID[:5],
	}

	// bring down the veth pair
	err := netlink.LinkSetDown(veth)
	if err != nil {
		log.Warnf("Error bring down Veth local iface: [ %v ]", veth)
		return err
	}

	// remove veth from bridge
	err = netlink.LinkSetNoMaster(veth)
	if err != nil {
		log.Warnf("Error removing veth from bridge")
		return err
	}

	// delete the veth interface
	err = netlink.LinkDel(veth)
	if err != nil {
		log.Warnf("Error removing veth interface")
		return err
	}

	return nil
}
<commit_msg>rename Driver type in prep to add IPAMDriver<commit_after>package vxlan

import (
	//"fmt"
	//"strings"
	//"time"
	"net"
	"strconv"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/go-plugins-helpers/network"
	//"github.com/samalba/dockerclient"
	//"github.com/davecgh/go-spew/spew"
	"github.com/vishvananda/netlink"
)

type NetDriver struct {
	network.Driver
	networks map[string]*NetworkState
}

// NetworkState is filled in at network creation time
// it contains state that we wish to keep for each network
type NetworkState struct {
	Bridge   *netlink.Bridge
	VXLan    *netlink.Vxlan
	Gateway  string
	IPv4Data []*network.IPAMData
	IPv6Data []*network.IPAMData
}

func NewDriver() (*NetDriver, error) {
	d := &NetDriver{
		networks: make(map[string]*NetworkState),
	}
	return d, nil
}

func (d *NetDriver) CreateNetwork(r *network.CreateNetworkRequest) error {
	log.Debugf("Create network request: %+v", r)

	netID := r.NetworkID
	var err error

	bridgeName := "br_" + netID[:12]
	vxlanName := "vx_" + netID[:12]

	// get interface names from options first
	for k, v := range r.Options {
		if k == "com.docker.network.generic" {
			if genericOpts, ok := v.(map[string]interface{}); ok {
				for key, val := range genericOpts {
					if key == "vxlanName" {
						vxlanName = val.(string)
					}
					if key == "bridgeName" {
						bridgeName = val.(string)
					}
				}
			}
		}
	}

	// create links
	bridge := &netlink.Bridge{
		LinkAttrs: netlink.LinkAttrs{
			Name: bridgeName,
		},
	}
	vxlan := &netlink.Vxlan{
		LinkAttrs: netlink.LinkAttrs{
			Name: vxlanName,
		},
	}

	// Parse interface options
	for k, v := range r.Options {
		if k == "com.docker.network.generic" {
			if genericOpts, ok := v.(map[string]interface{}); ok {
				for key, val := range genericOpts {
					log.Debugf("Libnetwork Opts Sent: [ %s ] Value: [ %s ]", key, val)
					if key == "vxlanMTU" {
						vxlan.LinkAttrs.MTU, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "bridgeMTU" {
						bridge.LinkAttrs.MTU, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "vxlanHardwareAddr" {
						vxlan.LinkAttrs.HardwareAddr, err = net.ParseMAC(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "bridgeHardwareAddr" {
						bridge.LinkAttrs.HardwareAddr, err = net.ParseMAC(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "vxlanTxQLen" {
						vxlan.LinkAttrs.TxQLen, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "bridgeTxQLen" {
						bridge.LinkAttrs.TxQLen, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "VxlanId" {
						vxlan.VxlanId, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "VtepDev" {
						vtepDev, err := netlink.LinkByName(val.(string))
						if err != nil {
							return err
						}
						vxlan.VtepDevIndex = vtepDev.Attrs().Index
					}
					if key == "SrcAddr" {
						vxlan.SrcAddr = net.ParseIP(val.(string))
					}
					if key == "Group" {
						vxlan.Group = net.ParseIP(val.(string))
					}
					if key == "TTL" {
						vxlan.TTL, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "TOS" {
						vxlan.TOS, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Learning" {
						vxlan.Learning, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Proxy" {
						vxlan.Proxy, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "RSC" {
						vxlan.RSC, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "L2miss" {
						vxlan.L2miss, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "L3miss" {
						vxlan.L3miss, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "NoAge" {
						vxlan.NoAge, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "GBP" {
						vxlan.GBP, err = strconv.ParseBool(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Age" {
						vxlan.Age, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Limit" {
						vxlan.Limit, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "Port" {
						vxlan.Port, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "PortLow" {
						vxlan.PortLow, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
					if key == "PortHigh" {
						vxlan.PortHigh, err = strconv.Atoi(val.(string))
						if err != nil {
							return err
						}
					}
				}
			}
		}
	}
	// Done parsing options

	// delete links if they already exist, don't worry about errors
	netlink.LinkDel(bridge)
	netlink.LinkDel(vxlan)

	// add links
	err = netlink.LinkAdd(bridge)
	if err != nil {
		return err
	}
	err = netlink.LinkAdd(vxlan)
	if err != nil {
		return err
	}

	// add vxlan to bridge
	err = netlink.LinkSetMaster(vxlan, bridge)
	if err != nil {
		return err
	}

	// bring interfaces up
	err = netlink.LinkSetUp(bridge)
	if err != nil {
		return err
	}
	err = netlink.LinkSetUp(vxlan)
	if err != nil {
		return err
	}

	// store interfaces to be used later
	ns := &NetworkState{
		VXLan:    vxlan,
		Bridge:   bridge,
		IPv4Data: r.IPv4Data,
		IPv6Data: r.IPv6Data,
	}

	// Add IPs to interfaces
	// Process IPv6 first. If both are inclued, IPv4 gateway will be the one that
	// remains, because JoinResponse can only include one Gateway
	for i := range r.IPv6Data {
		gatewayIP, err := netlink.ParseAddr(r.IPv6Data[i].Gateway)
		if err != nil {
			return err
		}
		ns.Gateway = gatewayIP.IP.String()
		netlink.AddrAdd(bridge, gatewayIP)
	}
	for i := range r.IPv4Data {
		gatewayIP, err := netlink.ParseAddr(r.IPv4Data[i].Gateway)
		if err != nil {
			return err
		}
		ns.Gateway = gatewayIP.IP.String()
		netlink.AddrAdd(bridge, gatewayIP)
	}

	d.networks[netID] = ns

	return nil
}

func (d *NetDriver) DeleteNetwork(r *network.DeleteNetworkRequest) error {
	netID := r.NetworkID

	err := netlink.LinkDel(d.networks[netID].VXLan)
	if err != nil {
		return err
	}
	err = netlink.LinkDel(d.networks[netID].Bridge)
	if err != nil {
		return err
	}

	return nil
}

func (d *NetDriver) CreateEndpoint(r *network.CreateEndpointRequest) error {
	log.Debugf("Create endpoint request: %+v", r)
	return nil
}

func (d *NetDriver) DeleteEndpoint(r *network.DeleteEndpointRequest) error {
	log.Debugf("Delete endpoint request: %+v", r)
	return nil
}

func (d *NetDriver) EndpointInfo(r *network.InfoRequest) (*network.InfoResponse, error) {
	res := &network.InfoResponse{
		Value: make(map[string]string),
	}
	return res, nil
}

func (d *NetDriver) Join(r *network.JoinRequest) (*network.JoinResponse, error) {
	// create and attach local name to the bridge
	veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: "veth_" + r.EndpointID[:5]},
		PeerName:  "ethc" + r.EndpointID[:5],
	}
	if err := netlink.LinkAdd(veth); err != nil {
		log.Errorf("failed to create the veth pair named: [ %v ] error: [ %s ] ", veth, err)
		return nil, err
	}

	// bring up the veth pair
	err := netlink.LinkSetUp(veth)
	if err != nil {
		log.Warnf("Error enabling  Veth local iface: [ %v ]", veth)
		return nil, err
	}

	bridge := d.networks[r.NetworkID].Bridge
	// add veth to bridge
	err = netlink.LinkSetMaster(veth, bridge)
	if err != nil {
		return nil, err
	}

	// SrcName gets renamed to DstPrefix + ID on the container iface
	res := &network.JoinResponse{
		InterfaceName: network.InterfaceName{
			SrcName:   veth.PeerName,
			DstPrefix: "eth",
		},
		Gateway: d.networks[r.NetworkID].Gateway,
	}
	log.Debugf("Join endpoint %s:%s to %s", r.NetworkID, r.EndpointID, r.SandboxKey)
	return res, nil
}

func (d *NetDriver) Leave(r *network.LeaveRequest) error {
	log.Debugf("Leave request: %+v", r)

	veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: "veth_" + r.EndpointID[:5]},
		PeerName:  "ethc" + r.EndpointID[:5],
	}

	// bring down the veth pair
	err := netlink.LinkSetDown(veth)
	if err != nil {
		log.Warnf("Error bring down Veth local iface: [ %v ]", veth)
		return err
	}

	// remove veth from bridge
	err = netlink.LinkSetNoMaster(veth)
	if err != nil {
		log.Warnf("Error removing veth from bridge")
		return err
	}

	// delete the veth interface
	err = netlink.LinkDel(veth)
	if err != nil {
		log.Warnf("Error removing veth interface")
		return err
	}

	return nil
}
<|endoftext|>