<commit_before>/*
Copyright 2014 Workiva, LLC

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 queue

import (
	"sync"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestPut(t *testing.T) {
	q := New(10)

	q.Put(`test`)
	assert.Equal(t, int64(1), q.Len())

	results, err := q.Get(1)
	assert.Nil(t, err)

	result := results[0]
	assert.Equal(t, `test`, result)
	assert.True(t, q.Empty())

	q.Put(`test2`)
	assert.Equal(t, int64(1), q.Len())

	results, err = q.Get(1)
	assert.Nil(t, err)

	result = results[0]
	assert.Equal(t, `test2`, result)
	assert.True(t, q.Empty())
}

func TestGet(t *testing.T) {
	q := New(10)

	q.Put(`test`)
	result, err := q.Get(2)
	if !assert.Nil(t, err) {
		return
	}

	assert.Len(t, result, 1)
	assert.Equal(t, `test`, result[0])
	assert.Equal(t, int64(0), q.Len())

	q.Put(`1`)
	q.Put(`2`)

	result, err = q.Get(1)
	if !assert.Nil(t, err) {
		return
	}

	assert.Len(t, result, 1)
	assert.Equal(t, `1`, result[0])
	assert.Equal(t, int64(1), q.Len())

	result, err = q.Get(2)
	if !assert.Nil(t, err) {
		return
	}

	assert.Equal(t, `2`, result[0])
}

func TestAddEmptyPut(t *testing.T) {
	q := New(10)

	q.Put()

	if q.Len() != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, q.Len())
	}
}

func TestGetNonPositiveNumber(t *testing.T) {
	q := New(10)

	q.Put(`test`)
	result, err := q.Get(0)
	if !assert.Nil(t, err) {
		return
	}

	if len(result) != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, len(result))
	}
}

func TestEmpty(t *testing.T) {
	q := New(10)

	if !q.Empty() {
		t.Errorf(`Expected empty queue.`)
	}

	q.Put(`test`)
	if q.Empty() {
		t.Errorf(`Expected non-empty queue.`)
	}
}

func TestGetEmpty(t *testing.T) {
	q := New(10)

	go func() {
		q.Put(`a`)
	}()

	result, err := q.Get(2)
	if !assert.Nil(t, err) {
		return
	}

	assert.Len(t, result, 1)
	assert.Equal(t, `a`, result[0])
}

func TestMultipleGetEmpty(t *testing.T) {
	q := New(10)
	var wg sync.WaitGroup
	wg.Add(2)
	results := make([][]interface{}, 2)

	go func() {
		wg.Done()
		local, err := q.Get(1)
		assert.Nil(t, err)
		results[0] = local
		wg.Done()
	}()

	go func() {
		wg.Done()
		local, err := q.Get(1)
		assert.Nil(t, err)
		results[1] = local
		wg.Done()
	}()

	wg.Wait()
	wg.Add(2)

	q.Put(`a`, `b`, `c`)
	wg.Wait()

	if assert.Len(t, results[0], 1) && assert.Len(t, results[1], 1) {
		assert.True(t, (results[0][0] == `a` && results[1][0] == `b`) ||
			(results[0][0] == `b` && results[1][0] == `a`),
			`The array should be a, b or b, a`)
	}
}

func TestEmptyGetWithDispose(t *testing.T) {
	q := New(10)
	var wg sync.WaitGroup
	wg.Add(1)

	var err error

	go func() {
		wg.Done()
		_, err = q.Get(1)
		wg.Done()
	}()

	wg.Wait()
	wg.Add(1)

	q.Dispose()

	wg.Wait()

	assert.IsType(t, DisposedError{}, err)
}

func TestGetPutDisposed(t *testing.T) {
	q := New(10)

	q.Dispose()

	_, err := q.Get(1)
	assert.IsType(t, DisposedError{}, err)

	err = q.Put(`a`)
	assert.IsType(t, DisposedError{}, err)
}

func BenchmarkQueue(b *testing.B) {
	q := New(int64(b.N))
	var wg sync.WaitGroup
	wg.Add(1)
	i := 0

	go func() {
		for {
			q.Get(1)
			i++
			if i == b.N {
				wg.Done()
				break
			}
		}
	}()

	for i := 0; i < b.N; i++ {
		q.Put(`a`)
	}

	wg.Wait()
}

func BenchmarkChannel(b *testing.B) {
	ch := make(chan interface{}, 1)
	var wg sync.WaitGroup
	wg.Add(1)
	i := 0

	go func() {
		for {
			<-ch
			i++
			if i == b.N {
				wg.Done()
				break
			}
		}
	}()

	for i := 0; i < b.N; i++ {
		ch <- `a`
	}

	wg.Wait()
}

func TestTakeUntil(t *testing.T) {
	q := New(10)
	q.Put(`a`, `b`, `c`)
	result, err := q.TakeUntil(func(item interface{}) bool {
		return item != `c`
	})

	if !assert.Nil(t, err) {
		return
	}

	expected := []interface{}{`a`, `b`}
	assert.Equal(t, expected, result)
}

func TestTakeUntilEmptyQueue(t *testing.T) {
	q := New(10)
	result, err := q.TakeUntil(func(item interface{}) bool {
		return item != `c`
	})

	if !assert.Nil(t, err) {
		return
	}

	expected := []interface{}{}
	assert.Equal(t, expected, result)
}

func TestTakeUntilOnDisposedQueue(t *testing.T) {
	q := New(10)
	q.Dispose()
	result, err := q.TakeUntil(func(item interface{}) bool {
		return true
	})

	assert.Nil(t, result)
	assert.IsType(t, DisposedError{}, err)
}

func TestExecuteInParallel(t *testing.T) {
	q := New(10)
	for i := 0; i < 10; i++ {
		q.Put(i)
	}

	numCalls := uint64(0)

	ExecuteInParallel(q, func(interface{}) {
		atomic.AddUint64(&numCalls, 1)
	})

	assert.Equal(t, 10, numCalls)
	assert.True(t, q.Disposed())
}

func TestExecuteInParallelEmptyQueue(t *testing.T) {
	q := New(1)

	// basically just ensuring we don't deadlock here
	ExecuteInParallel(q, func(interface{}) {
		t.Fail()
	})
}

func BenchmarkQueuePut(b *testing.B) {
	numItems := int64(1000)

	qs := make([]*Queue, 0, b.N)

	for i := 0; i < b.N; i++ {
		q := New(10)
		qs = append(qs, q)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		q := qs[i]
		for j := int64(0); j < numItems; j++ {
			q.Put(j)
		}
	}
}

func BenchmarkQueueGet(b *testing.B) {
	numItems := int64(1000)

	qs := make([]*Queue, 0, b.N)

	for i := 0; i < b.N; i++ {
		q := New(numItems)
		for j := int64(0); j < numItems; j++ {
			q.Put(j)
		}
		qs = append(qs, q)
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		q := qs[i]
		for j := int64(0); j < numItems; j++ {
			q.Get(1)
		}
	}
}
<commit_msg>Added benchmark around execute in parallel.<commit_after>/*
Copyright 2014 Workiva, LLC

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 queue

import (
	"sync"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestPut(t *testing.T) {
	q := New(10)

	q.Put(`test`)
	assert.Equal(t, int64(1), q.Len())

	results, err := q.Get(1)
	assert.Nil(t, err)

	result := results[0]
	assert.Equal(t, `test`, result)
	assert.True(t, q.Empty())

	q.Put(`test2`)
	assert.Equal(t, int64(1), q.Len())

	results, err = q.Get(1)
	assert.Nil(t, err)

	result = results[0]
	assert.Equal(t, `test2`, result)
	assert.True(t, q.Empty())
}

func TestGet(t *testing.T) {
	q := New(10)

	q.Put(`test`)
	result, err := q.Get(2)
	if !assert.Nil(t, err) {
		return
	}

	assert.Len(t, result, 1)
	assert.Equal(t, `test`, result[0])
	assert.Equal(t, int64(0), q.Len())

	q.Put(`1`)
	q.Put(`2`)

	result, err = q.Get(1)
	if !assert.Nil(t, err) {
		return
	}

	assert.Len(t, result, 1)
	assert.Equal(t, `1`, result[0])
	assert.Equal(t, int64(1), q.Len())

	result, err = q.Get(2)
	if !assert.Nil(t, err) {
		return
	}

	assert.Equal(t, `2`, result[0])
}

func TestAddEmptyPut(t *testing.T) {
	q := New(10)

	q.Put()

	if q.Len() != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, q.Len())
	}
}

func TestGetNonPositiveNumber(t *testing.T) {
	q := New(10)

	q.Put(`test`)
	result, err := q.Get(0)
	if !assert.Nil(t, err) {
		return
	}

	if len(result) != 0 {
		t.Errorf(`Expected len: %d, received: %d`, 0, len(result))
	}
}

func TestEmpty(t *testing.T) {
	q := New(10)

	if !q.Empty() {
		t.Errorf(`Expected empty queue.`)
	}

	q.Put(`test`)
	if q.Empty() {
		t.Errorf(`Expected non-empty queue.`)
	}
}

func TestGetEmpty(t *testing.T) {
	q := New(10)

	go func() {
		q.Put(`a`)
	}()

	result, err := q.Get(2)
	if !assert.Nil(t, err) {
		return
	}

	assert.Len(t, result, 1)
	assert.Equal(t, `a`, result[0])
}

func TestMultipleGetEmpty(t *testing.T) {
	q := New(10)
	var wg sync.WaitGroup
	wg.Add(2)
	results := make([][]interface{}, 2)

	go func() {
		wg.Done()
		local, err := q.Get(1)
		assert.Nil(t, err)
		results[0] = local
		wg.Done()
	}()

	go func() {
		wg.Done()
		local, err := q.Get(1)
		assert.Nil(t, err)
		results[1] = local
		wg.Done()
	}()

	wg.Wait()
	wg.Add(2)

	q.Put(`a`, `b`, `c`)
	wg.Wait()

	if assert.Len(t, results[0], 1) && assert.Len(t, results[1], 1) {
		assert.True(t, (results[0][0] == `a` && results[1][0] == `b`) ||
			(results[0][0] == `b` && results[1][0] == `a`),
			`The array should be a, b or b, a`)
	}
}

func TestEmptyGetWithDispose(t *testing.T) {
	q := New(10)
	var wg sync.WaitGroup
	wg.Add(1)

	var err error

	go func() {
		wg.Done()
		_, err = q.Get(1)
		wg.Done()
	}()

	wg.Wait()
	wg.Add(1)

	q.Dispose()

	wg.Wait()

	assert.IsType(t, DisposedError{}, err)
}

func TestGetPutDisposed(t *testing.T) {
	q := New(10)

	q.Dispose()

	_, err := q.Get(1)
	assert.IsType(t, DisposedError{}, err)

	err = q.Put(`a`)
	assert.IsType(t, DisposedError{}, err)
}

func BenchmarkQueue(b *testing.B) {
	q := New(int64(b.N))
	var wg sync.WaitGroup
	wg.Add(1)
	i := 0

	go func() {
		for {
			q.Get(1)
			i++
			if i == b.N {
				wg.Done()
				break
			}
		}
	}()

	for i := 0; i < b.N; i++ {
		q.Put(`a`)
	}

	wg.Wait()
}

func BenchmarkChannel(b *testing.B) {
	ch := make(chan interface{}, 1)
	var wg sync.WaitGroup
	wg.Add(1)
	i := 0

	go func() {
		for {
			<-ch
			i++
			if i == b.N {
				wg.Done()
				break
			}
		}
	}()

	for i := 0; i < b.N; i++ {
		ch <- `a`
	}

	wg.Wait()
}

func TestTakeUntil(t *testing.T) {
	q := New(10)
	q.Put(`a`, `b`, `c`)
	result, err := q.TakeUntil(func(item interface{}) bool {
		return item != `c`
	})

	if !assert.Nil(t, err) {
		return
	}

	expected := []interface{}{`a`, `b`}
	assert.Equal(t, expected, result)
}

func TestTakeUntilEmptyQueue(t *testing.T) {
	q := New(10)
	result, err := q.TakeUntil(func(item interface{}) bool {
		return item != `c`
	})

	if !assert.Nil(t, err) {
		return
	}

	expected := []interface{}{}
	assert.Equal(t, expected, result)
}

func TestTakeUntilOnDisposedQueue(t *testing.T) {
	q := New(10)
	q.Dispose()
	result, err := q.TakeUntil(func(item interface{}) bool {
		return true
	})

	assert.Nil(t, result)
	assert.IsType(t, DisposedError{}, err)
}

func TestExecuteInParallel(t *testing.T) {
	q := New(10)
	for i := 0; i < 10; i++ {
		q.Put(i)
	}

	numCalls := uint64(0)

	ExecuteInParallel(q, func(interface{}) {
		atomic.AddUint64(&numCalls, 1)
	})

	assert.Equal(t, 10, numCalls)
	assert.True(t, q.Disposed())
}

func TestExecuteInParallelEmptyQueue(t *testing.T) {
	q := New(1)

	// basically just ensuring we don't deadlock here
	ExecuteInParallel(q, func(interface{}) {
		t.Fail()
	})
}

func BenchmarkQueuePut(b *testing.B) {
	numItems := int64(1000)

	qs := make([]*Queue, 0, b.N)

	for i := 0; i < b.N; i++ {
		q := New(10)
		qs = append(qs, q)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		q := qs[i]
		for j := int64(0); j < numItems; j++ {
			q.Put(j)
		}
	}
}

func BenchmarkQueueGet(b *testing.B) {
	numItems := int64(1000)

	qs := make([]*Queue, 0, b.N)

	for i := 0; i < b.N; i++ {
		q := New(numItems)
		for j := int64(0); j < numItems; j++ {
			q.Put(j)
		}
		qs = append(qs, q)
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		q := qs[i]
		for j := int64(0); j < numItems; j++ {
			q.Get(1)
		}
	}
}

func BenchmarkExecuteInParallel(b *testing.B) {
	numItems := int64(1000)

	qs := make([]*Queue, 0, b.N)

	for i := 0; i < b.N; i++ {
		q := New(numItems)
		for j := int64(0); j < numItems; j++ {
			q.Put(j)
		}
		qs = append(qs, q)
	}

	var counter int64
	fn := func(ifc interface{}) {
		c := ifc.(int64)
		atomic.AddInt64(&counter, c)
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		q := qs[i]
		ExecuteInParallel(q, fn)
	}
}
<|endoftext|>


<commit_before>package queue

import (
	"github.com/nathan-osman/go-cannon/email"
	"github.com/nathan-osman/go-cannon/util"

	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net/smtp"
	"os"
	"sync"
	"syscall"
	"time"
)

// Persistent connection to an SMTP host.
type Host struct {
	sync.Mutex
	host         string
	directory    string
	lastActivity time.Time
	newEmail     *util.NonBlockingChan
	stop         chan bool
}

// Log the specified message.
func (h *Host) log(msg string) {
	log.Printf("[%s] %s", h.host, msg)
}

// Receive the next email in the queue.
func (h *Host) receiveEmail() *email.Email {

	// The host queue is considered "inactive" while waiting for new emails to
	// arrive - record the current time before entering the select{} block so
	// that the Idle() method can calculate the idle time
	h.Lock()
	h.lastActivity = time.Now()
	h.Unlock()

	// When the function exits, reset the inactive timer to 0
	defer func() {
		h.Lock()
		h.lastActivity = time.Time{}
		h.Unlock()
	}()

	// Either receive a new email or stop the queue
	select {
	case i := <-h.newEmail.Recv:
		return i.(*email.Email)
	case <-h.stop:
		return nil
	}
}

// Attempt to connect to the specified server.
func (h *Host) tryMailServer(server string) (*smtp.Client, error) {

	var (
		c    *smtp.Client
		err  error
		done = make(chan bool)
	)

	// Because Dial() is a blocking function, it must be run in a separate
	// goroutine so that it can be aborted immediately
	go func() {
		c, err = smtp.Dial(fmt.Sprintf("%s:25", server))
		close(done)
	}()

	// Wait for either the goroutine to complete or a stop request
	select {
	case <-done:
	case <-h.stop:
		return nil, nil
	}

	// Attempt to establish a TCP connection to port 25
	if err == nil {

		// Obtain this machine's hostname and say HELO
		if hostname, err := os.Hostname(); err == nil {
			c.Hello(hostname)
		}

		// If the server advertises TLS, attempt to use it - if the server
		// advertises TLS but it doesn't work, that's an error
		if ok, _ := c.Extension("STARTTLS"); ok {
			if err := c.StartTLS(&tls.Config{}); err != nil {
				return nil, err
			}
		}

		return c, nil

	} else {
		return nil, err
	}
}

// Attempt to connect to a mail server until the timeout is reached or until
// stopped.
func (h *Host) connectToMailServer() (*smtp.Client, error) {

	// Obtain the list of mail servers to try
	servers := util.FindMailServers(h.host)

	// RFC 5321 (4.5.4) describes the process for retrying connections to a
	// mail server after failure. The recommended strategy is to retry twice
	// with 30 minute intervals and continue at 120 minute intervals until four
	// days have elapsed. That's *roughly* what is done here.
	for i := 0; i < 50; i++ {

		// Try each of the servers in order
		for _, s := range servers {
			if c, err := h.tryMailServer(s); err == nil {
				return c, nil
			}
		}

		// None of the servers could be reached - wait a few minutes before
		// trying to connect again
		var d time.Duration
		if i < 2 {
			d = 30 * time.Minute
		} else {
			d = 2 * time.Hour
		}

		select {
		case <-time.After(d):
		case <-h.stop:
			return nil, nil
		}
	}

	// All attempts have failed, let the caller know we tried :)
	return nil, errors.New("unable to connect to a mail server")
}

// Attempt to deliver the specified email to the specifed client
func (h *Host) deliverToMailServer(c *smtp.Client, e *email.Email) error {

	// Specify the sender of the emails
	if err := c.Mail(e.From); err != nil {
		return err
	}

	// Add each of the recipients
	for _, t := range e.To {
		if err := c.Rcpt(t); err != nil {
			return err
		}
	}

	// Obtain a writer for writing the actual email
	w, err := c.Data()
	if err != nil {
		return err
	}
	defer w.Close()

	// Write the email
	if _, err = w.Write(e.Message); err != nil {
		return err
	}

	return nil
}

// Receive emails and deliver them to their recipients.
func (h *Host) run() {

	// Close the stop channel when the goroutine exits
	defer close(h.stop)

	// The client must be declared here so that it can be closed after the loop
	var (
		c   *smtp.Client
		e   *email.Email
		err error
	)

	for {

		// Receive the next email from the queue if one hasn't already been
		// retrieved
		if e == nil {
			h.log("waiting for email in queue...")
			if e = h.receiveEmail(); e == nil {
				break
			}
			h.log("email received in queue")
		}

		// Connect to the server if a connection does not already exist
		if c == nil {
			h.log("connecting to mail server...")
			if c, err = h.connectToMailServer(); c == nil {
				if err == nil {
					break
				} else {
					h.log(err.Error())
					continue
				}
			}
			h.log("connection established")
		}

		// Attempt delivery of the message
		if err = h.deliverToMailServer(c, e); err != nil {

			// If the type of error has anything to do with a syscall, assume
			// that the connection was broken and try reconnecting - otherwise,
			// discard the email - either way, reset the error
			if _, ok := err.(syscall.Errno); ok {
				h.log("connection to server lost")
				c, err = nil, nil
				continue
			} else {
				h.log(err.Error())
				e, err = nil, nil
			}
		} else {
			h.log("email successfully delivered")
		}

		// Delete the email
		e.Delete(h.directory)
	}

	// Close the connection if it is still open
	if c != nil {
		c.Quit()
	}

	h.log("shutting down queue")
}

// Create a new host connection.
func NewHost(host, directory string) *Host {

	// Create the host, including the channel used for delivering new emails
	h := &Host{
		host:      host,
		directory: directory,
		newEmail:  util.NewNonBlockingChan(),
		stop:      make(chan bool),
	}

	// Start a goroutine to manage the lifecycle of the host connection
	go h.run()

	return h
}

// Attempt to deliver an email to the host.
func (h *Host) Deliver(e *email.Email) {
	h.newEmail.Send <- e
}

// Retrieve the connection idle time.
func (h *Host) Idle() time.Duration {
	h.Lock()
	defer h.Unlock()
	if h.lastActivity.IsZero() {
		return 0
	} else {
		return time.Since(h.lastActivity)
	}
}

// Close the connection to the host.
func (h *Host) Stop() {

	// Send on the channel to stop it and wait for it to be closed
	h.stop <- true
	<-h.stop
}
<commit_msg>Added line for logging failed server connection attempts.<commit_after>package queue

import (
	"github.com/nathan-osman/go-cannon/email"
	"github.com/nathan-osman/go-cannon/util"

	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net/smtp"
	"os"
	"sync"
	"syscall"
	"time"
)

// Persistent connection to an SMTP host.
type Host struct {
	sync.Mutex
	host         string
	directory    string
	lastActivity time.Time
	newEmail     *util.NonBlockingChan
	stop         chan bool
}

// Log the specified message.
func (h *Host) log(msg string) {
	log.Printf("[%s] %s", h.host, msg)
}

// Receive the next email in the queue.
func (h *Host) receiveEmail() *email.Email {

	// The host queue is considered "inactive" while waiting for new emails to
	// arrive - record the current time before entering the select{} block so
	// that the Idle() method can calculate the idle time
	h.Lock()
	h.lastActivity = time.Now()
	h.Unlock()

	// When the function exits, reset the inactive timer to 0
	defer func() {
		h.Lock()
		h.lastActivity = time.Time{}
		h.Unlock()
	}()

	// Either receive a new email or stop the queue
	select {
	case i := <-h.newEmail.Recv:
		return i.(*email.Email)
	case <-h.stop:
		return nil
	}
}

// Attempt to connect to the specified server.
func (h *Host) tryMailServer(server string) (*smtp.Client, error) {

	var (
		c    *smtp.Client
		err  error
		done = make(chan bool)
	)

	// Because Dial() is a blocking function, it must be run in a separate
	// goroutine so that it can be aborted immediately
	go func() {
		c, err = smtp.Dial(fmt.Sprintf("%s:25", server))
		close(done)
	}()

	// Wait for either the goroutine to complete or a stop request
	select {
	case <-done:
	case <-h.stop:
		return nil, nil
	}

	// Attempt to establish a TCP connection to port 25
	if err == nil {

		// Obtain this machine's hostname and say HELO
		if hostname, err := os.Hostname(); err == nil {
			c.Hello(hostname)
		}

		// If the server advertises TLS, attempt to use it - if the server
		// advertises TLS but it doesn't work, that's an error
		if ok, _ := c.Extension("STARTTLS"); ok {
			if err := c.StartTLS(&tls.Config{}); err != nil {
				return nil, err
			}
		}

		return c, nil

	} else {
		return nil, err
	}
}

// Attempt to connect to a mail server until the timeout is reached or until
// stopped.
func (h *Host) connectToMailServer() (*smtp.Client, error) {

	// Obtain the list of mail servers to try
	servers := util.FindMailServers(h.host)

	// RFC 5321 (4.5.4) describes the process for retrying connections to a
	// mail server after failure. The recommended strategy is to retry twice
	// with 30 minute intervals and continue at 120 minute intervals until four
	// days have elapsed. That's *roughly* what is done here.
	for i := 0; i < 50; i++ {

		// Try each of the servers in order
		for _, s := range servers {
			if c, err := h.tryMailServer(s); err == nil {
				return c, nil
			} else {
				h.log(fmt.Sprintf("unable to connect to %s", s))
			}
		}

		// None of the servers could be reached - wait a few minutes before
		// trying to connect again
		var d time.Duration
		if i < 2 {
			d = 30 * time.Minute
		} else {
			d = 2 * time.Hour
		}

		select {
		case <-time.After(d):
		case <-h.stop:
			return nil, nil
		}
	}

	// All attempts have failed, let the caller know we tried :)
	return nil, errors.New("unable to connect to a mail server")
}

// Attempt to deliver the specified email to the specifed client
func (h *Host) deliverToMailServer(c *smtp.Client, e *email.Email) error {

	// Specify the sender of the emails
	if err := c.Mail(e.From); err != nil {
		return err
	}

	// Add each of the recipients
	for _, t := range e.To {
		if err := c.Rcpt(t); err != nil {
			return err
		}
	}

	// Obtain a writer for writing the actual email
	w, err := c.Data()
	if err != nil {
		return err
	}
	defer w.Close()

	// Write the email
	if _, err = w.Write(e.Message); err != nil {
		return err
	}

	return nil
}

// Receive emails and deliver them to their recipients.
func (h *Host) run() {

	// Close the stop channel when the goroutine exits
	defer close(h.stop)

	// The client must be declared here so that it can be closed after the loop
	var (
		c   *smtp.Client
		e   *email.Email
		err error
	)

	for {

		// Receive the next email from the queue if one hasn't already been
		// retrieved
		if e == nil {
			h.log("waiting for email in queue...")
			if e = h.receiveEmail(); e == nil {
				break
			}
			h.log("email received in queue")
		}

		// Connect to the server if a connection does not already exist
		if c == nil {
			h.log("connecting to mail server...")
			if c, err = h.connectToMailServer(); c == nil {
				if err == nil {
					break
				} else {
					h.log(err.Error())
					continue
				}
			}
			h.log("connection established")
		}

		// Attempt delivery of the message
		if err = h.deliverToMailServer(c, e); err != nil {

			// If the type of error has anything to do with a syscall, assume
			// that the connection was broken and try reconnecting - otherwise,
			// discard the email - either way, reset the error
			if _, ok := err.(syscall.Errno); ok {
				h.log("connection to server lost")
				c, err = nil, nil
				continue
			} else {
				h.log(err.Error())
				e, err = nil, nil
			}
		} else {
			h.log("email successfully delivered")
		}

		// Delete the email
		e.Delete(h.directory)
	}

	// Close the connection if it is still open
	if c != nil {
		c.Quit()
	}

	h.log("shutting down queue")
}

// Create a new host connection.
func NewHost(host, directory string) *Host {

	// Create the host, including the channel used for delivering new emails
	h := &Host{
		host:      host,
		directory: directory,
		newEmail:  util.NewNonBlockingChan(),
		stop:      make(chan bool),
	}

	// Start a goroutine to manage the lifecycle of the host connection
	go h.run()

	return h
}

// Attempt to deliver an email to the host.
func (h *Host) Deliver(e *email.Email) {
	h.newEmail.Send <- e
}

// Retrieve the connection idle time.
func (h *Host) Idle() time.Duration {
	h.Lock()
	defer h.Unlock()
	if h.lastActivity.IsZero() {
		return 0
	} else {
		return time.Since(h.lastActivity)
	}
}

// Close the connection to the host.
func (h *Host) Stop() {

	// Send on the channel to stop it and wait for it to be closed
	h.stop <- true
	<-h.stop
}
<|endoftext|>


<commit_before><commit_msg>Added RemoveAny() tests<commit_after><|endoftext|>


<commit_before><commit_msg>refactoring<commit_after><|endoftext|>


<commit_before><commit_msg>Add host<commit_after><|endoftext|>


<commit_before>package qapi

import (
	"time"
)

// Account represents an account associated with the user on behalf
// of which the API client is authorized.
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts
type Account struct {
	// Type of the account (e.g., "Cash", "Margin").
	Type string `json":"type""`

	// Eight-digit account number (e.g., "26598145")
	// Stored as a string, it's used for making account-related API calls
	Number string `json:"number"`

	// Status of the account (e.g., Active).
	Status string `json:"status"`

	// Whether this is a primary account for the holder.
	IsPrimary bool `json:"isPrimary"`

	// Whether this account is one that gets billed for various expenses such as inactivity fees, market data, etc.
	IsBilling bool `json:"isBilling"`

	// Type of client holding the account (e.g., "Individual").
	ClientAccountType string `json:"clientAccountType"`
}

// Position belonging to an account
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-positions
type Position struct {
	// Position symbol.
	Symbol string `json:"symbol"`

	// Internal symbol identifier
	SymbolID int `json:"symbolId"`

	// Position quantity remaining open.
	OpenQuantity float32 `json:"openQuantity"`

	// Portion of the position that was closed today.
	ClosedQuantity float32 `json:"closedQuantity"`

	// Market value of the position (quantity x price).
	CurrentMarketValue float32 `json:"currentMarketValue"`

	// Current price of the position symbol.
	CurrentPrice float32 `json:"currentPrice"`

	// Average price paid for all executions constituting the position.
	AverageEntryPrice float32 `json:"averageEntryPrice"`

	// Realized profit/loss on this position.
	ClosedPnL float32 `json:"closedPnL"`

	// Unrealized profit/loss on this position.
	OpenPnL float32 `json:"openPnL"`

	// Total cost of the position.
	TotalCost float32 `json:"totalCost"`

	// Designates whether real-time quote was used to compute PnL.
	IsRealTime bool `json:"isRealTime"`

	// Designates whether a symbol is currently undergoing a reorg.
	IsUnderReorg bool `json:"isUnderReorg"`
}

// Balance belonging to an Account
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-balances
type Balance struct {

	// Currency of the balance figure(e.g., "USD" or "CAD").
	Currency string `json:"currency"`

	// Balance amount.
	Cash float32 `json:"cash"`

	// Market value of all securities in the account in a given currency.
	MarketValue float32 `json:"marketValue"`

	// Equity as a difference between cash and marketValue properties.
	TotalEquity float32 `json:"totalEquity"`

	// Buying power for that particular currency side of the account.
	BuyingPower float32 `json:"buyingPower"`

	// Maintenance excess for that particular side of the account.
	MaintenanceExcess float32 `json:"maintenanceExcess"`

	// Whether real-time data was used to calculate the above values.
	IsRealTime bool `json:"isRealTime"`
}

// AccountBalances represents per-currency and combined balances for a specified account.
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-balances
type AccountBalances struct {
	PerCurrencyBalances    []Balance `json:"perCurrencyBalances"`
	CombinedBalances       []Balance `json:"combinedBalances"`
	SODPerCurrencyBalances []Balance `json:"sodPerCurrencyBalances"`
	SODCombinedBalances    []Balance `json:"sodCombinedBalances"`
}

// Execution belonging to an Account
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-executions
type Execution struct {
	// Execution symbol.
	Symbol string `json:"symbol"`

	// Internal symbol identifier
	SymbolID int `json:"symbolId"`

	// Execution quantity.
	Quantity int `json:"quantity"`

	// Client side of the order to which execution belongs.
	Side string `json:"side"`

	// Execution price.
	Price float32 `json:"price"`

	// Internal identifier of the execution.
	ID int `json:"id"`

	// Internal identifier of the order to which the execution belongs.
	OrderID int `json:"orderId"`

	// Internal identifier of the order chain to which the execution belongs.
	OrderChainID int `json:"orderChainId"`

	// Identifier of the execution at the market where it originated.
	ExchangeExecID string `json:"exchangeExecId"`

	// Execution timestamp.
	Timestamp time.Time `json:"timestamp"`

	// Manual notes that may have been entered by Trade Desk staff
	Notes string `json:"notes"`

	// Trading venue where execution originated.
	Venue string `json:"venue"`

	// Execution cost (price x quantity).
	TotalCost float32 `json:"totalCost"`

	// Questrade commission for orders placed with Trade Desk.
	OrderPlacementCommission float32 `json:"orderPlacementCommission"`

	// Questrade commission.
	Commission float32 `json:"commission"`

	// Liquidity fee charged by execution venue.
	ExecutionFee float32 `json:"executionFee"`

	// SEC fee charged on all sales of US securities.
	SecFee float32 `json:"secFee"`

	// Additional execution fee charged by TSX (if applicable).
	CanadianExecutionFee int `json:"canadianExecutionFee"`

	// Internal identifierof the parent order.
	ParentID int `json:"parentId"`
}
<commit_msg>Fixed syntax error<commit_after>package qapi

import (
	"time"
)

// Account represents an account associated with the user on behalf
// of which the API client is authorized.
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts
type Account struct {
	// Type of the account (e.g., "Cash", "Margin").
	Type string `json:"type"`

	// Eight-digit account number (e.g., "26598145")
	// Stored as a string, it's used for making account-related API calls
	Number string `json:"number"`

	// Status of the account (e.g., Active).
	Status string `json:"status"`

	// Whether this is a primary account for the holder.
	IsPrimary bool `json:"isPrimary"`

	// Whether this account is one that gets billed for various expenses such as inactivity fees, market data, etc.
	IsBilling bool `json:"isBilling"`

	// Type of client holding the account (e.g., "Individual").
	ClientAccountType string `json:"clientAccountType"`
}

// Position belonging to an account
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-positions
type Position struct {
	// Position symbol.
	Symbol string `json:"symbol"`

	// Internal symbol identifier
	SymbolID int `json:"symbolId"`

	// Position quantity remaining open.
	OpenQuantity float32 `json:"openQuantity"`

	// Portion of the position that was closed today.
	ClosedQuantity float32 `json:"closedQuantity"`

	// Market value of the position (quantity x price).
	CurrentMarketValue float32 `json:"currentMarketValue"`

	// Current price of the position symbol.
	CurrentPrice float32 `json:"currentPrice"`

	// Average price paid for all executions constituting the position.
	AverageEntryPrice float32 `json:"averageEntryPrice"`

	// Realized profit/loss on this position.
	ClosedPnL float32 `json:"closedPnL"`

	// Unrealized profit/loss on this position.
	OpenPnL float32 `json:"openPnL"`

	// Total cost of the position.
	TotalCost float32 `json:"totalCost"`

	// Designates whether real-time quote was used to compute PnL.
	IsRealTime bool `json:"isRealTime"`

	// Designates whether a symbol is currently undergoing a reorg.
	IsUnderReorg bool `json:"isUnderReorg"`
}

// Balance belonging to an Account
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-balances
type Balance struct {

	// Currency of the balance figure(e.g., "USD" or "CAD").
	Currency string `json:"currency"`

	// Balance amount.
	Cash float32 `json:"cash"`

	// Market value of all securities in the account in a given currency.
	MarketValue float32 `json:"marketValue"`

	// Equity as a difference between cash and marketValue properties.
	TotalEquity float32 `json:"totalEquity"`

	// Buying power for that particular currency side of the account.
	BuyingPower float32 `json:"buyingPower"`

	// Maintenance excess for that particular side of the account.
	MaintenanceExcess float32 `json:"maintenanceExcess"`

	// Whether real-time data was used to calculate the above values.
	IsRealTime bool `json:"isRealTime"`
}

// AccountBalances represents per-currency and combined balances for a specified account.
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-balances
type AccountBalances struct {
	PerCurrencyBalances    []Balance `json:"perCurrencyBalances"`
	CombinedBalances       []Balance `json:"combinedBalances"`
	SODPerCurrencyBalances []Balance `json:"sodPerCurrencyBalances"`
	SODCombinedBalances    []Balance `json:"sodCombinedBalances"`
}

// Execution belonging to an Account
//
// Ref: http://www.questrade.com/api/documentation/rest-operations/account-calls/accounts-id-executions
type Execution struct {
	// Execution symbol.
	Symbol string `json:"symbol"`

	// Internal symbol identifier
	SymbolID int `json:"symbolId"`

	// Execution quantity.
	Quantity int `json:"quantity"`

	// Client side of the order to which execution belongs.
	Side string `json:"side"`

	// Execution price.
	Price float32 `json:"price"`

	// Internal identifier of the execution.
	ID int `json:"id"`

	// Internal identifier of the order to which the execution belongs.
	OrderID int `json:"orderId"`

	// Internal identifier of the order chain to which the execution belongs.
	OrderChainID int `json:"orderChainId"`

	// Identifier of the execution at the market where it originated.
	ExchangeExecID string `json:"exchangeExecId"`

	// Execution timestamp.
	Timestamp time.Time `json:"timestamp"`

	// Manual notes that may have been entered by Trade Desk staff
	Notes string `json:"notes"`

	// Trading venue where execution originated.
	Venue string `json:"venue"`

	// Execution cost (price x quantity).
	TotalCost float32 `json:"totalCost"`

	// Questrade commission for orders placed with Trade Desk.
	OrderPlacementCommission float32 `json:"orderPlacementCommission"`

	// Questrade commission.
	Commission float32 `json:"commission"`

	// Liquidity fee charged by execution venue.
	ExecutionFee float32 `json:"executionFee"`

	// SEC fee charged on all sales of US securities.
	SecFee float32 `json:"secFee"`

	// Additional execution fee charged by TSX (if applicable).
	CanadianExecutionFee int `json:"canadianExecutionFee"`

	// Internal identifierof the parent order.
	ParentID int `json:"parentId"`
}
<|endoftext|>


<commit_before>package engine

// Directory Index

import (
	"bytes"
	"net/http"
	"path/filepath"
	"strings"

	log "github.com/sirupsen/logrus"
	"github.com/xyproto/algernon/themes"
	"github.com/xyproto/algernon/utils"
	"gopkg.in/gcfg.v1"
)

var (
	// List of filenames that should be displayed instead of a directory listing
	indexFilenames = []string{"index.lua", "index.html", "index.md", "index.txt", "index.pongo2", "index.tmpl", "index.po2", "index.amber", "index.happ", "index.hyper", "index.hyper.js", "index.hyper.jsx"}

	doubleP  = utils.Pathsep + utils.Pathsep /* // */
	dotSlash = "." + utils.Pathsep           /* ./ */
)

const (
	dirconfFilename = ".algernon"
)

// directory configuration file structure
type DirConfig struct {
	Main struct {
		Title string
		Theme string
	}
}

// DirectoryListing serves the given directory as a web page with links the the contents
func (ac *Config) DirectoryListing(w http.ResponseWriter, req *http.Request, rootdir, dirname, theme string) {
	var (
		buf          bytes.Buffer
		fullFilename string
		URLpath      string
		title        = dirname
	)

	// Fill the coming HTML body with a list of all the filenames in `dirname`
	for _, filename := range utils.GetFilenames(dirname) {

		if filename == dirconfFilename {
			// Skip
			continue
		}

		// Find the full name
		fullFilename = filepath.Join(dirname, filename)

		// Remove the root directory from the link path
		URLpath = fullFilename[len(rootdir)+1:]

		// Output different entries for files and directories
		buf.WriteString(themes.HTMLLink(filename, URLpath, ac.fs.IsDir(fullFilename)))
	}

	// Read directory configuration, if present
	fullDirConfFilename := filepath.Join(dirname, dirconfFilename)
	if ac.fs.Exists(fullDirConfFilename) {
		var dirConf DirConfig
		if err := gcfg.ReadFileInto(&dirConf, fullDirConfFilename); err == nil {
			if dirConf.Main.Title != "" {
				title = dirConf.Main.Title
			}
			if dirConf.Main.Theme != "" {
				theme = dirConf.Main.Theme
			}
		}
	} else {
		// Strip the leading "./" from the current directory
		title = strings.TrimPrefix(title, dotSlash)

		// Strip double "/" at the end, just keep one
		if strings.Contains(title, doubleP) {
			// Replace "//" with just "/"
			title = strings.Replace(title, doubleP, utils.Pathsep, utils.EveryInstance)
		}
	}

	// Check if the current page contents are empty
	if buf.Len() == 0 {
		buf.WriteString("Empty directory")
	}

	htmldata := themes.MessagePageBytes(title, buf.Bytes(), theme)

	// If the auto-refresh feature has been enabled
	if ac.autoRefreshMode {
		// Insert JavaScript for refreshing the page into the generated HTML
		htmldata = ac.InsertAutoRefresh(req, htmldata)
	}

	// Serve the page
	w.Header().Add("Content-Type", "text/html; charset=utf-8")
	DataToClient(w, req, dirname, htmldata)
}

// DirPage serves a directory, using index.* files, if present.
// The directory must exist.
// rootdir is the base directory (can be ".")
// dirname is the specific directory that is to be served (should never be ".")
func (ac *Config) DirPage(w http.ResponseWriter, req *http.Request, rootdir, dirname, theme string) {

	// Check if we are instructed to quit after serving the first file
	if ac.quitAfterFirstRequest {
		go ac.quitSoon("Quit after first request", defaultSoonDuration)
	}

	// If the URL does not end with a slash, redirect to an URL that does
	if !strings.HasSuffix(req.URL.Path, "/") {
		if req.Method == "POST" {
			log.Warn("Redirecting a POST request: " + req.URL.Path + " -> " + req.URL.Path + "/.")
			log.Warn("Header data may be lost! Please add the missing slash.")
		}
		http.Redirect(w, req, req.URL.Path+"/", http.StatusMovedPermanently)
		return
	}

	// Handle the serving of index files, if needed
	var filename string
	for _, indexfile := range indexFilenames {
		filename = filepath.Join(dirname, indexfile)
		if ac.fs.Exists(filename) {
			ac.FilePage(w, req, filename, ac.defaultLuaDataFilename)
			return
		}
	}

	// Serve a directory listing of no index file is found
	ac.DirectoryListing(w, req, rootdir, dirname, theme)
}
<commit_msg>Revert use of filepath.Join<commit_after>package engine

// Directory Index

import (
	"bytes"
	"net/http"
	"path/filepath"
	"strings"

	log "github.com/sirupsen/logrus"
	"github.com/xyproto/algernon/themes"
	"github.com/xyproto/algernon/utils"
	"gopkg.in/gcfg.v1"
)

var (
	// List of filenames that should be displayed instead of a directory listing
	indexFilenames = []string{"index.lua", "index.html", "index.md", "index.txt", "index.pongo2", "index.tmpl", "index.po2", "index.amber", "index.happ", "index.hyper", "index.hyper.js", "index.hyper.jsx"}

	doubleP  = utils.Pathsep + utils.Pathsep /* // */
	dotSlash = "." + utils.Pathsep           /* ./ */
)

const (
	dirconfFilename = ".algernon"
)

// directory configuration file structure
type DirConfig struct {
	Main struct {
		Title string
		Theme string
	}
}

// DirectoryListing serves the given directory as a web page with links the the contents
func (ac *Config) DirectoryListing(w http.ResponseWriter, req *http.Request, rootdir, dirname, theme string) {
	var (
		buf          bytes.Buffer
		fullFilename string
		URLpath      string
		title        = dirname
	)

	// Fill the coming HTML body with a list of all the filenames in `dirname`
	for _, filename := range utils.GetFilenames(dirname) {

		if filename == dirconfFilename {
			// Skip
			continue
		}

		// Find the full name
		fullFilename = dirname

		// Add a "/" after the directory name, if missing
		if !strings.HasSuffix(fullFilename, utils.Pathsep) {
			fullFilename += utils.Pathsep
		}

		// Add the filename at the end
		fullFilename += filename

		// Remove the root directory from the link path
		URLpath = fullFilename[len(rootdir)+1:]

		// Output different entries for files and directories
		buf.WriteString(themes.HTMLLink(filename, URLpath, ac.fs.IsDir(fullFilename)))
	}

	// Read directory configuration, if present
	fullDirConfFilename := filepath.Join(dirname, dirconfFilename)
	if ac.fs.Exists(fullDirConfFilename) {
		var dirConf DirConfig
		if err := gcfg.ReadFileInto(&dirConf, fullDirConfFilename); err == nil {
			if dirConf.Main.Title != "" {
				title = dirConf.Main.Title
			}
			if dirConf.Main.Theme != "" {
				theme = dirConf.Main.Theme
			}
		}
	} else {
		// Strip the leading "./" from the current directory
		title = strings.TrimPrefix(title, dotSlash)

		// Strip double "/" at the end, just keep one
		if strings.Contains(title, doubleP) {
			// Replace "//" with just "/"
			title = strings.Replace(title, doubleP, utils.Pathsep, utils.EveryInstance)
		}
	}

	// Check if the current page contents are empty
	if buf.Len() == 0 {
		buf.WriteString("Empty directory")
	}

	htmldata := themes.MessagePageBytes(title, buf.Bytes(), theme)

	// If the auto-refresh feature has been enabled
	if ac.autoRefreshMode {
		// Insert JavaScript for refreshing the page into the generated HTML
		htmldata = ac.InsertAutoRefresh(req, htmldata)
	}

	// Serve the page
	w.Header().Add("Content-Type", "text/html; charset=utf-8")
	DataToClient(w, req, dirname, htmldata)
}

// DirPage serves a directory, using index.* files, if present.
// The directory must exist.
// rootdir is the base directory (can be ".")
// dirname is the specific directory that is to be served (should never be ".")
func (ac *Config) DirPage(w http.ResponseWriter, req *http.Request, rootdir, dirname, theme string) {

	// Check if we are instructed to quit after serving the first file
	if ac.quitAfterFirstRequest {
		go ac.quitSoon("Quit after first request", defaultSoonDuration)
	}

	// If the URL does not end with a slash, redirect to an URL that does
	if !strings.HasSuffix(req.URL.Path, "/") {
		if req.Method == "POST" {
			log.Warn("Redirecting a POST request: " + req.URL.Path + " -> " + req.URL.Path + "/.")
			log.Warn("Header data may be lost! Please add the missing slash.")
		}
		http.Redirect(w, req, req.URL.Path+"/", http.StatusMovedPermanently)
		return
	}

	// Handle the serving of index files, if needed
	var filename string
	for _, indexfile := range indexFilenames {
		filename = filepath.Join(dirname, indexfile)
		if ac.fs.Exists(filename) {
			ac.FilePage(w, req, filename, ac.defaultLuaDataFilename)
			return
		}
	}

	// Serve a directory listing of no index file is found
	ac.DirectoryListing(w, req, rootdir, dirname, theme)
}
<|endoftext|>


<commit_before>package engine

import (
	"sort"
	"sync"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/unit"
)

const (
	partitionSize = 5
)

type countbyname map[string]int

func (cbn countbyname) inc(key string) {
	v, ok := cbn[key]
	if !ok {
		cbn[key] = 0
	}
	cbn[key] = v + 1
}

func (cbn countbyname) dec(key string) {
	v, ok := cbn[key]
	if !ok {
		return
	}
	cbn[key] = v - 1
}

type namedCount struct {
	name  string
	count int
}

type namedCounts []*namedCount

func (a namedCounts) Len() int           { return len(a) }
func (a namedCounts) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a namedCounts) Less(i, j int) bool { return a[i].count < a[j].count }

type cluster struct {
	uptodate bool
	mutex    *sync.Mutex

	// some statistics about what goes on in the cluster
	// this will get more sophisticated as needed

	// where are jobs running
	jobsToMachines map[string]string
	// how many jobs on a machine
	machineJobCount countbyname
}

func newCluster() *cluster {
	return &cluster{false, new(sync.Mutex), make(map[string]string), make(map[string]int)}
}

// jobScheduled handles the job scheduled event
func (clust *cluster) jobScheduled(jobName string, mst *machine.MachineState) {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	if !clust.uptodate {
		return
	}

	clust.populateJob(jobName, mst.BootId)
}

// jobStopped handles the job stopped event
func (clust *cluster) jobStopped(jobName string) {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	if !clust.uptodate {
		return
	}

	clust.deleteJob(jobName)
}

// isUptodate returns whether the cluster has been told what happened before it got created
func (clust *cluster) isUptodate() bool {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()
	return clust.uptodate
}

// refreshFrom refreshes from a specified cluster
func (clust *cluster) refreshFrom(cu *cluster) {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	if clust.uptodate {
		return
	}

	clust.jobsToMachines = cu.jobsToMachines
	clust.machineJobCount = cu.machineJobCount
	clust.uptodate = true
}

func (clust *cluster) populateJob(jobName string, machineBootID string) {
	clust.jobsToMachines[jobName] = machineBootID
	clust.machineJobCount.inc(machineBootID)
}

func (clust *cluster) deleteJob(jobName string) {
	machineBootID, ok := clust.jobsToMachines[jobName]

	// TODO(uwedeportivo): this might be a signal that a refresh is needed
	if !ok {
		return
	}

	clust.machineJobCount.dec(machineBootID)
}

func (clust *cluster) kLeastLoaded(k int) []string {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	mas := make(namedCounts, len(clust.machineJobCount))
	cursor := 0
	for k, v := range clust.machineJobCount {
		mas[cursor] = &namedCount{k, v}
		cursor++
	}
	sort.Sort(mas)

	l := k
	if l > len(mas) {
		l = len(mas)
	}

	mbis := make([]string, l)
	for i, nc := range mas {
		mbis[i] = nc.name
	}
	return mbis
}

func (eg *Engine) refreshCluster(force bool) {
	if !force && eg.clust.isUptodate() {
		return
	}

	cu := newCluster()

	jobs := eg.registry.GetAllJobs()
	for _, j := range jobs {
		mst := eg.registry.GetJobTarget(j.Name)
		cu.populateJob(j.Name, mst.BootId)
	}

	eg.clust.refreshFrom(cu)
}

// requiresMachine returns whether specified job requires a specific machine.
func (eg *Engine) requiresMachine(j *job.Job) ([]string, bool) {
	requirements := j.Requirements()

	bootID, ok := requirements[unit.FleetXConditionMachineBootID]
	return bootID, ok && len(bootID) > 0
}

// partitionCluster returns a slice of bootids from a subset of active machines
// that should be considered for scheduling the specified job.
// The returned slice is sorted by ascending lexicographical string value of machine boot id.
func (eg *Engine) partitionCluster(j *job.Job) ([]string, error) {
	if bootID, ok := eg.requiresMachine(j); ok {
		return bootID, nil
	}

	// TODO(uwedeportivo): for now punt on jobs with requirements and offer to all machines
	// because agents are decoding the requirements
	if len(j.Requirements()) > 0 {
		machines := eg.registry.GetActiveMachines()

		machineBootIds := make([]string, len(machines))
		for i, mach := range machines {
			machineBootIds[i] = mach.BootId
		}
		sort.Strings(machineBootIds)
		return machineBootIds, nil
	}

	// this is usually a cheap no-op
	eg.refreshCluster(false)

	// as an initial heuristic, choose the k least loaded, with k = partitionSize
	machineBootIds := eg.clust.kLeastLoaded(partitionSize)

	sort.Strings(machineBootIds)
	return machineBootIds, nil
}
<commit_msg>feat(job offer): restrict job offers to a subset of the machines<commit_after>package engine

import (
	"sort"
	"sync"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/unit"
)

const (
	partitionSize = 5
)

type countbyname map[string]int

func (cbn countbyname) inc(key string) {
	v, ok := cbn[key]
	if !ok {
		cbn[key] = 0
	}
	cbn[key] = v + 1
}

func (cbn countbyname) dec(key string) {
	v, ok := cbn[key]
	if !ok {
		return
	}
	cbn[key] = v - 1
}

type namedCount struct {
	name  string
	count int
}

type namedCounts []*namedCount

func (a namedCounts) Len() int           { return len(a) }
func (a namedCounts) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a namedCounts) Less(i, j int) bool { return a[i].count < a[j].count }

type cluster struct {
	uptodate bool
	mutex    *sync.Mutex

	// some statistics about what goes on in the cluster
	// this will get more sophisticated as needed

	// where are jobs running
	jobsToMachines map[string]string
	// how many jobs on a machine
	machineJobCount countbyname
}

func newCluster() *cluster {
	return &cluster{false, new(sync.Mutex), make(map[string]string), make(map[string]int)}
}

// jobScheduled handles the job scheduled event
func (clust *cluster) jobScheduled(jobName string, mst *machine.MachineState) {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	if !clust.uptodate {
		return
	}

	clust.populateJob(jobName, mst.BootId)
}

// jobStopped handles the job stopped event
func (clust *cluster) jobStopped(jobName string) {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	if !clust.uptodate {
		return
	}

	clust.deleteJob(jobName)
}

// isUptodate returns whether the cluster has been told what happened before it got created
func (clust *cluster) isUptodate() bool {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()
	return clust.uptodate
}

// refreshFrom refreshes from a specified cluster
func (clust *cluster) refreshFrom(cu *cluster, force bool) {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	if !force && clust.uptodate {
		return
	}

	clust.jobsToMachines = cu.jobsToMachines
	clust.machineJobCount = cu.machineJobCount
	clust.uptodate = true
}

func (clust *cluster) populateJob(jobName string, machineBootID string) {
	clust.jobsToMachines[jobName] = machineBootID
	clust.machineJobCount.inc(machineBootID)
}

func (clust *cluster) deleteJob(jobName string) {
	machineBootID, ok := clust.jobsToMachines[jobName]

	// TODO(uwedeportivo): this might be a signal that a refresh is needed
	if !ok {
		return
	}

	clust.machineJobCount.dec(machineBootID)
}

func (clust *cluster) kLeastLoaded(k int) []string {
	clust.mutex.Lock()
	defer clust.mutex.Unlock()

	mas := make(namedCounts, len(clust.machineJobCount))
	cursor := 0
	for k, v := range clust.machineJobCount {
		mas[cursor] = &namedCount{k, v}
		cursor++
	}
	sort.Sort(mas)

	l := k
	if l > len(mas) {
		l = len(mas)
	}

	mbis := make([]string, l)
	for i, nc := range mas {
		mbis[i] = nc.name
	}
	return mbis
}

func (eg *Engine) refreshCluster(force bool) {
	if !force && eg.clust.isUptodate() {
		return
	}

	cu := newCluster()

	jobs := eg.registry.GetAllJobs()
	for _, j := range jobs {
		mst := eg.registry.GetJobTarget(j.Name)
		cu.populateJob(j.Name, mst.BootId)
	}

	eg.clust.refreshFrom(cu, force)
}

// requiresMachine returns whether specified job requires a specific machine.
func (eg *Engine) requiresMachine(j *job.Job) ([]string, bool) {
	requirements := j.Requirements()

	bootID, ok := requirements[unit.FleetXConditionMachineBootID]
	return bootID, ok && len(bootID) > 0
}

// partitionCluster returns a slice of bootids from a subset of active machines
// that should be considered for scheduling the specified job.
// The returned slice is sorted by ascending lexicographical string value of machine boot id.
func (eg *Engine) partitionCluster(j *job.Job) ([]string, error) {
	if bootID, ok := eg.requiresMachine(j); ok {
		return bootID, nil
	}

	// TODO(uwedeportivo): for now punt on jobs with requirements and offer to all machines
	// because agents are decoding the requirements
	if len(j.Requirements()) > 0 {
		machines := eg.registry.GetActiveMachines()

		machineBootIds := make([]string, len(machines))
		for i, mach := range machines {
			machineBootIds[i] = mach.BootId
		}
		sort.Strings(machineBootIds)
		return machineBootIds, nil
	}

	// this is usually a cheap no-op
	eg.refreshCluster(false)

	// as an initial heuristic, choose the k least loaded, with k = partitionSize
	machineBootIds := eg.clust.kLeastLoaded(partitionSize)

	sort.Strings(machineBootIds)
	return machineBootIds, nil
}
<|endoftext|>


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

// The stage0 command looks up the buildlet's URL from its environment
// (GCE metadata service, scaleway, etc), downloads it, and runs
// it. If not on GCE, such as when in a Linux Docker container being
// developed and tested locally, the stage0 instead looks for the
// META_BUILDLET_BINARY_URL environment to have a URL to the buildlet
// binary.
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"cloud.google.com/go/compute/metadata"
	"golang.org/x/build/internal/httpdl"
	"golang.org/x/build/internal/untar"
)

// This lets us be lazy and put the stage0 start-up in rc.local where
// it might race with the network coming up, rather than write proper
// upstart+systemd+init scripts:
var networkWait = flag.Duration("network-wait", 0, "if zero, a default is used if needed")

const osArch = runtime.GOOS + "/" + runtime.GOARCH

const attr = "buildlet-binary-url"

// untar helper, for the Windows image prep script.
var (
	untarFile    = flag.String("untar-file", "", "if non-empty, tar.gz to untar to --untar-dest-dir")
	untarDestDir = flag.String("untar-dest-dir", "", "destination directory to untar --untar-file to")
)

// configureSerialLogOutput is set non-nil on some platforms to configure
// log output to go to the serial console.
var configureSerialLogOutput func()

func main() {
	if configureSerialLogOutput != nil {
		configureSerialLogOutput()
	}
	log.SetPrefix("stage0: ")
	log.Printf("bootstrap binary running.")
	flag.Parse()

	if *untarFile != "" {
		untarMode()
		return
	}

	switch osArch {
	case "linux/arm":
		switch env := os.Getenv("GO_BUILDER_ENV"); env {
		case "linux-arm-arm5spacemonkey", "host-linux-arm-scaleway":
			// No setup currently.
		default:
			panic(fmt.Sprintf("unknown/unspecified $GO_BUILDER_ENV value %q", env))
		}
	case "linux/arm64":
		switch env := os.Getenv("GO_BUILDER_ENV"); env {
		case "host-linux-arm64-packet", "host-linux-arm64-linaro":
			// No special setup.
		default:
			panic(fmt.Sprintf("unknown/unspecified $GO_BUILDER_ENV value %q", env))
		}
	case "linux/ppc64":
		initOregonStatePPC64()
	case "linux/ppc64le":
		initOregonStatePPC64le()
	}

	if !awaitNetwork() {
		sleepFatalf("network didn't become reachable")
	}

	// Note: we name it ".exe" for Windows, but the name also
	// works fine on Linux, etc.
	target := filepath.FromSlash("./buildlet.exe")
	if err := download(target, buildletURL()); err != nil {
		sleepFatalf("Downloading %s: %v", buildletURL, err)
	}

	if runtime.GOOS != "windows" {
		if err := os.Chmod(target, 0755); err != nil {
			log.Fatal(err)
		}
	}

	env := os.Environ()
	if isUnix() && os.Getuid() == 0 {
		if os.Getenv("USER") == "" {
			env = append(env, "USER=root")
		}
		if os.Getenv("HOME") == "" {
			env = append(env, "HOME=/root")
		}
	}

	cmd := exec.Command(target)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = env
	switch buildenv := os.Getenv("GO_BUILDER_ENV"); buildenv {
	case "linux-arm-arm5spacemonkey":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs(buildenv)...)
	case "host-linux-arm-scaleway":
		scalewayArgs := append(
			legacyReverseBuildletArgs(buildenv),
			"--hostname="+os.Getenv("HOSTNAME"),
		)
		cmd.Args = append(cmd.Args,
			scalewayArgs...,
		)
	}
	switch osArch {
	case "linux/s390x":
		cmd.Args = append(cmd.Args, "--workdir=/data/golang/workdir")
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("linux-s390x-ibm")...)
	case "linux/arm64":
		switch v := os.Getenv("GO_BUILDER_ENV"); v {
		case "host-linux-arm64-packet", "host-linux-arm64-linaro":
			hostname := os.Getenv("HOSTNAME") // if empty, docker container name is used
			cmd.Args = append(cmd.Args,
				"--reverse-type="+v,
				"--workdir=/workdir",
				"--hostname="+hostname,
				"--halt=false",
				"--reboot=false",
				"--coordinator=farmer.golang.org:443",
			)
		default:
			panic(fmt.Sprintf("unknown/unspecified $GO_BUILDER_ENV value %q", env))
		}
	case "linux/ppc64":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("linux-ppc64-buildlet")...)
	case "linux/ppc64le":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("linux-ppc64le-buildlet")...)
	case "solaris/amd64":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("solaris-amd64-smartosbuildlet")...)
	}
	if err := cmd.Run(); err != nil {
		sleepFatalf("Error running buildlet: %v", err)
	}
}

// legacyReverseBuildletArgs passes builder as the deprecated --reverse flag.
// New code should use --reverse-type instead.
func legacyReverseBuildletArgs(builder string) []string {
	return []string{
		"--halt=false",
		"--reverse=" + builder,
		"--coordinator=farmer.golang.org:443",
	}
}

// awaitNetwork reports whether the network came up within 30 seconds,
// determined somewhat arbitrarily via a DNS lookup for google.com.
func awaitNetwork() bool {
	timeout := 30 * time.Second
	if runtime.GOOS == "windows" {
		timeout = 5 * time.Minute // empirically slower sometimes?
	}
	if *networkWait != 0 {
		timeout = *networkWait
	}
	deadline := time.Now().Add(timeout)
	var lastSpam time.Time
	log.Printf("waiting for network.")
	for time.Now().Before(deadline) {
		t0 := time.Now()
		if isNetworkUp() {
			log.Printf("network is up.")
			return true
		}
		failAfter := time.Since(t0)
		if now := time.Now(); now.After(lastSpam.Add(5 * time.Second)) {
			log.Printf("network still down; probe failure took %v", failAfter.Round(time.Second/10))
			lastSpam = now
		}
		time.Sleep(1 * time.Second)
	}
	log.Printf("gave up waiting for network")
	return false
}

func isNetworkUp() bool {
	const probeURL = "http://farmer.golang.org/netcheck" // 404 is fine.
	c := &http.Client{
		Timeout: 5 * time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	res, err := c.Get(probeURL)
	if err != nil {
		return false
	}
	res.Body.Close()
	return true
}

func buildletURL() string {
	if os.Getenv("GO_BUILDER_ENV") == "linux-arm-arm5spacemonkey" {
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-arm-arm5"
	}
	switch osArch {
	case "linux/s390x":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-s390x"
	case "linux/arm64":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-arm64"
	case "linux/ppc64":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-ppc64"
	case "linux/ppc64le":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-ppc64le"
	case "solaris/amd64":
		return "https://storage.googleapis.com/go-builder-data/buildlet.solaris-amd64"
	}
	// The buildlet download URL is located in an env var
	// when the buildlet is not running on GCE, or is running
	// on Kubernetes.
	if !metadata.OnGCE() || os.Getenv("IN_KUBERNETES") == "1" {
		if v := os.Getenv("META_BUILDLET_BINARY_URL"); v != "" {
			return v
		}
		sleepFatalf("Not on GCE, and no META_BUILDLET_BINARY_URL specified.")
	}
	v, err := metadata.InstanceAttributeValue(attr)
	if err != nil {
		sleepFatalf("Failed to look up %q attribute value: %v", attr, err)
	}
	return v
}

func sleepFatalf(format string, args ...interface{}) {
	log.Printf(format, args...)
	if runtime.GOOS == "windows" {
		log.Printf("(sleeping for 1 minute before failing)")
		time.Sleep(time.Minute) // so user has time to see it in cmd.exe, maybe
	}
	os.Exit(1)
}

func download(file, url string) error {
	log.Printf("downloading %s to %s ...\n", url, file)
	const maxTry = 3
	var lastErr error
	for try := 1; try <= maxTry; try++ {
		if try > 1 {
			// network should be up by now per awaitNetwork, so just retry
			// shortly a few time on errors.
			time.Sleep(2)
		}
		err := httpdl.Download(file, url)
		if err == nil {
			fi, err := os.Stat(file)
			if err != nil {
				return err
			}
			log.Printf("downloaded %s (%d bytes)", file, fi.Size())
			return nil
		}
		lastErr = err
		log.Printf("try %d/%d download failure: %v", try, maxTry, err)
	}
	return lastErr
}

func aptGetInstall(pkgs ...string) {
	args := append([]string{"--yes", "install"}, pkgs...)
	cmd := exec.Command("apt-get", args...)
	if out, err := cmd.CombinedOutput(); err != nil {
		log.Fatalf("error running apt-get install: %s", out)
	}
}

func initBootstrapDir(destDir, tgzCache string) {
	if err := os.MkdirAll(destDir, 0755); err != nil {
		log.Fatal(err)
	}
	// TODO(bradfitz): rewrite this to use Go instead of curl+tar
	// if this ever gets used on platforms besides Unix. For
	// Windows and Plan 9 we bake in the bootstrap tarball into
	// the image anyway. So this works for now. Solaris might require
	// tweaking to use gtar instead or something.
	latestURL := fmt.Sprintf("https://storage.googleapis.com/go-builder-data/gobootstrap-%s-%s.tar.gz",
		runtime.GOOS, runtime.GOARCH)
	curl := exec.Command("/usr/bin/curl", "-R", "-o", tgzCache, "-z", tgzCache, latestURL)
	out, err := curl.CombinedOutput()
	if err != nil {
		log.Fatalf("curl error fetching %s to %s: %s", latestURL, out, err)
	}
	tar := exec.Command("tar", "zxf", tgzCache)
	tar.Dir = destDir
	out, err = tar.CombinedOutput()
	if err != nil {
		log.Fatalf("error untarring %s to %s: %s", tgzCache, destDir, out)
	}
}

func initOregonStatePPC64() {
	aptGetInstall("gcc", "strace", "libc6-dev", "gdb")
	initBootstrapDir("/usr/local/go-bootstrap", "/usr/local/go-bootstrap.tar.gz")
}

func initOregonStatePPC64le() {
	aptGetInstall("gcc", "strace", "libc6-dev", "gdb")
	initBootstrapDir("/usr/local/go-bootstrap", "/usr/local/go-bootstrap.tar.gz")
}

func isUnix() bool {
	switch runtime.GOOS {
	case "plan9", "windows":
		return false
	}
	return true
}

func untarMode() {
	if *untarDestDir == "" {
		log.Fatal("--untar-dest-dir must not be empty")
	}
	if fi, err := os.Stat(*untarDestDir); err != nil || !fi.IsDir() {
		if err != nil {
			log.Fatalf("--untar-dest-dir %q: %v", *untarDestDir, err)
		}
		log.Fatalf("--untar-dest-dir %q not a directory.", *untarDestDir)
	}
	f, err := os.Open(*untarFile)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	if err := untar.Untar(f, *untarDestDir); err != nil {
		log.Fatalf("Untarring %q to %q: %v", *untarFile, *untarDestDir, err)
	}
}
<commit_msg>cmd/buildlet/stage0: support building with Go 1.8 again<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The stage0 command looks up the buildlet's URL from its environment
// (GCE metadata service, scaleway, etc), downloads it, and runs
// it. If not on GCE, such as when in a Linux Docker container being
// developed and tested locally, the stage0 instead looks for the
// META_BUILDLET_BINARY_URL environment to have a URL to the buildlet
// binary.
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"cloud.google.com/go/compute/metadata"
	"golang.org/x/build/internal/httpdl"
	"golang.org/x/build/internal/untar"
)

// This lets us be lazy and put the stage0 start-up in rc.local where
// it might race with the network coming up, rather than write proper
// upstart+systemd+init scripts:
var networkWait = flag.Duration("network-wait", 0, "if zero, a default is used if needed")

const osArch = runtime.GOOS + "/" + runtime.GOARCH

const attr = "buildlet-binary-url"

// untar helper, for the Windows image prep script.
var (
	untarFile    = flag.String("untar-file", "", "if non-empty, tar.gz to untar to --untar-dest-dir")
	untarDestDir = flag.String("untar-dest-dir", "", "destination directory to untar --untar-file to")
)

// configureSerialLogOutput is set non-nil on some platforms to configure
// log output to go to the serial console.
var configureSerialLogOutput func()

func main() {
	if configureSerialLogOutput != nil {
		configureSerialLogOutput()
	}
	log.SetPrefix("stage0: ")
	log.Printf("bootstrap binary running.")
	flag.Parse()

	if *untarFile != "" {
		untarMode()
		return
	}

	switch osArch {
	case "linux/arm":
		switch env := os.Getenv("GO_BUILDER_ENV"); env {
		case "linux-arm-arm5spacemonkey", "host-linux-arm-scaleway":
			// No setup currently.
		default:
			panic(fmt.Sprintf("unknown/unspecified $GO_BUILDER_ENV value %q", env))
		}
	case "linux/arm64":
		switch env := os.Getenv("GO_BUILDER_ENV"); env {
		case "host-linux-arm64-packet", "host-linux-arm64-linaro":
			// No special setup.
		default:
			panic(fmt.Sprintf("unknown/unspecified $GO_BUILDER_ENV value %q", env))
		}
	case "linux/ppc64":
		initOregonStatePPC64()
	case "linux/ppc64le":
		initOregonStatePPC64le()
	}

	if !awaitNetwork() {
		sleepFatalf("network didn't become reachable")
	}

	// Note: we name it ".exe" for Windows, but the name also
	// works fine on Linux, etc.
	target := filepath.FromSlash("./buildlet.exe")
	if err := download(target, buildletURL()); err != nil {
		sleepFatalf("Downloading %s: %v", buildletURL, err)
	}

	if runtime.GOOS != "windows" {
		if err := os.Chmod(target, 0755); err != nil {
			log.Fatal(err)
		}
	}

	env := os.Environ()
	if isUnix() && os.Getuid() == 0 {
		if os.Getenv("USER") == "" {
			env = append(env, "USER=root")
		}
		if os.Getenv("HOME") == "" {
			env = append(env, "HOME=/root")
		}
	}

	cmd := exec.Command(target)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = env
	switch buildenv := os.Getenv("GO_BUILDER_ENV"); buildenv {
	case "linux-arm-arm5spacemonkey":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs(buildenv)...)
	case "host-linux-arm-scaleway":
		scalewayArgs := append(
			legacyReverseBuildletArgs(buildenv),
			"--hostname="+os.Getenv("HOSTNAME"),
		)
		cmd.Args = append(cmd.Args,
			scalewayArgs...,
		)
	}
	switch osArch {
	case "linux/s390x":
		cmd.Args = append(cmd.Args, "--workdir=/data/golang/workdir")
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("linux-s390x-ibm")...)
	case "linux/arm64":
		switch v := os.Getenv("GO_BUILDER_ENV"); v {
		case "host-linux-arm64-packet", "host-linux-arm64-linaro":
			hostname := os.Getenv("HOSTNAME") // if empty, docker container name is used
			cmd.Args = append(cmd.Args,
				"--reverse-type="+v,
				"--workdir=/workdir",
				"--hostname="+hostname,
				"--halt=false",
				"--reboot=false",
				"--coordinator=farmer.golang.org:443",
			)
		default:
			panic(fmt.Sprintf("unknown/unspecified $GO_BUILDER_ENV value %q", env))
		}
	case "linux/ppc64":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("linux-ppc64-buildlet")...)
	case "linux/ppc64le":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("linux-ppc64le-buildlet")...)
	case "solaris/amd64":
		cmd.Args = append(cmd.Args, legacyReverseBuildletArgs("solaris-amd64-smartosbuildlet")...)
	}
	if err := cmd.Run(); err != nil {
		sleepFatalf("Error running buildlet: %v", err)
	}
}

// legacyReverseBuildletArgs passes builder as the deprecated --reverse flag.
// New code should use --reverse-type instead.
func legacyReverseBuildletArgs(builder string) []string {
	return []string{
		"--halt=false",
		"--reverse=" + builder,
		"--coordinator=farmer.golang.org:443",
	}
}

// awaitNetwork reports whether the network came up within 30 seconds,
// determined somewhat arbitrarily via a DNS lookup for google.com.
func awaitNetwork() bool {
	timeout := 30 * time.Second
	if runtime.GOOS == "windows" {
		timeout = 5 * time.Minute // empirically slower sometimes?
	}
	if *networkWait != 0 {
		timeout = *networkWait
	}
	deadline := time.Now().Add(timeout)
	var lastSpam time.Time
	log.Printf("waiting for network.")
	for time.Now().Before(deadline) {
		t0 := time.Now()
		if isNetworkUp() {
			log.Printf("network is up.")
			return true
		}
		failAfter := time.Since(t0)
		if now := time.Now(); now.After(lastSpam.Add(5 * time.Second)) {
			const round = time.Second / 10
			log.Printf("network still down; probe failure took %v", failAfter/round*round)
			lastSpam = now
		}
		time.Sleep(1 * time.Second)
	}
	log.Printf("gave up waiting for network")
	return false
}

func isNetworkUp() bool {
	const probeURL = "http://farmer.golang.org/netcheck" // 404 is fine.
	c := &http.Client{
		Timeout: 5 * time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	res, err := c.Get(probeURL)
	if err != nil {
		return false
	}
	res.Body.Close()
	return true
}

func buildletURL() string {
	if os.Getenv("GO_BUILDER_ENV") == "linux-arm-arm5spacemonkey" {
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-arm-arm5"
	}
	switch osArch {
	case "linux/s390x":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-s390x"
	case "linux/arm64":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-arm64"
	case "linux/ppc64":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-ppc64"
	case "linux/ppc64le":
		return "https://storage.googleapis.com/go-builder-data/buildlet.linux-ppc64le"
	case "solaris/amd64":
		return "https://storage.googleapis.com/go-builder-data/buildlet.solaris-amd64"
	}
	// The buildlet download URL is located in an env var
	// when the buildlet is not running on GCE, or is running
	// on Kubernetes.
	if !metadata.OnGCE() || os.Getenv("IN_KUBERNETES") == "1" {
		if v := os.Getenv("META_BUILDLET_BINARY_URL"); v != "" {
			return v
		}
		sleepFatalf("Not on GCE, and no META_BUILDLET_BINARY_URL specified.")
	}
	v, err := metadata.InstanceAttributeValue(attr)
	if err != nil {
		sleepFatalf("Failed to look up %q attribute value: %v", attr, err)
	}
	return v
}

func sleepFatalf(format string, args ...interface{}) {
	log.Printf(format, args...)
	if runtime.GOOS == "windows" {
		log.Printf("(sleeping for 1 minute before failing)")
		time.Sleep(time.Minute) // so user has time to see it in cmd.exe, maybe
	}
	os.Exit(1)
}

func download(file, url string) error {
	log.Printf("downloading %s to %s ...\n", url, file)
	const maxTry = 3
	var lastErr error
	for try := 1; try <= maxTry; try++ {
		if try > 1 {
			// network should be up by now per awaitNetwork, so just retry
			// shortly a few time on errors.
			time.Sleep(2)
		}
		err := httpdl.Download(file, url)
		if err == nil {
			fi, err := os.Stat(file)
			if err != nil {
				return err
			}
			log.Printf("downloaded %s (%d bytes)", file, fi.Size())
			return nil
		}
		lastErr = err
		log.Printf("try %d/%d download failure: %v", try, maxTry, err)
	}
	return lastErr
}

func aptGetInstall(pkgs ...string) {
	args := append([]string{"--yes", "install"}, pkgs...)
	cmd := exec.Command("apt-get", args...)
	if out, err := cmd.CombinedOutput(); err != nil {
		log.Fatalf("error running apt-get install: %s", out)
	}
}

func initBootstrapDir(destDir, tgzCache string) {
	if err := os.MkdirAll(destDir, 0755); err != nil {
		log.Fatal(err)
	}
	// TODO(bradfitz): rewrite this to use Go instead of curl+tar
	// if this ever gets used on platforms besides Unix. For
	// Windows and Plan 9 we bake in the bootstrap tarball into
	// the image anyway. So this works for now. Solaris might require
	// tweaking to use gtar instead or something.
	latestURL := fmt.Sprintf("https://storage.googleapis.com/go-builder-data/gobootstrap-%s-%s.tar.gz",
		runtime.GOOS, runtime.GOARCH)
	curl := exec.Command("/usr/bin/curl", "-R", "-o", tgzCache, "-z", tgzCache, latestURL)
	out, err := curl.CombinedOutput()
	if err != nil {
		log.Fatalf("curl error fetching %s to %s: %s", latestURL, out, err)
	}
	tar := exec.Command("tar", "zxf", tgzCache)
	tar.Dir = destDir
	out, err = tar.CombinedOutput()
	if err != nil {
		log.Fatalf("error untarring %s to %s: %s", tgzCache, destDir, out)
	}
}

func initOregonStatePPC64() {
	aptGetInstall("gcc", "strace", "libc6-dev", "gdb")
	initBootstrapDir("/usr/local/go-bootstrap", "/usr/local/go-bootstrap.tar.gz")
}

func initOregonStatePPC64le() {
	aptGetInstall("gcc", "strace", "libc6-dev", "gdb")
	initBootstrapDir("/usr/local/go-bootstrap", "/usr/local/go-bootstrap.tar.gz")
}

func isUnix() bool {
	switch runtime.GOOS {
	case "plan9", "windows":
		return false
	}
	return true
}

func untarMode() {
	if *untarDestDir == "" {
		log.Fatal("--untar-dest-dir must not be empty")
	}
	if fi, err := os.Stat(*untarDestDir); err != nil || !fi.IsDir() {
		if err != nil {
			log.Fatalf("--untar-dest-dir %q: %v", *untarDestDir, err)
		}
		log.Fatalf("--untar-dest-dir %q not a directory.", *untarDestDir)
	}
	f, err := os.Open(*untarFile)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	if err := untar.Untar(f, *untarDestDir); err != nil {
		log.Fatalf("Untarring %q to %q: %v", *untarFile, *untarDestDir, err)
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"testing"

	"encoding/json"

	"github.com/drud/ddev/pkg/exec"
	"github.com/drud/ddev/pkg/plugins/platform"
	"github.com/drud/ddev/pkg/testcommon"
	"github.com/drud/ddev/pkg/util"
	log "github.com/sirupsen/logrus"
	asrt "github.com/stretchr/testify/assert"
)

// TestDescribeBadArgs ensures the binary behaves as expected when used with invalid arguments or working directories.
func TestDescribeBadArgs(t *testing.T) {
	assert := asrt.New(t)

	// Create a temporary directory and switch to it for the duration of this test.
	tmpdir := testcommon.CreateTmpDir("badargs")
	defer testcommon.CleanupDir(tmpdir)
	defer testcommon.Chdir(tmpdir)()

	// Ensure it fails if we run the vanilla describe outside of an application root.
	args := []string{"describe"}
	out, err := exec.RunCommand(DdevBin, args)
	assert.Error(err)
	assert.Contains(string(out), "Please specify a site name or change directories")

	// Ensure we get a failure if we run a describe on a named application which does not exist.
	args = []string{"describe", util.RandString(16)}
	out, err = exec.RunCommand(DdevBin, args)
	assert.Error(err)
	assert.Contains(string(out), "Unable to find any active site")

	// Ensure we get a failure if using too many arguments.
	args = []string{"describe", util.RandString(16), util.RandString(16)}
	out, err = exec.RunCommand(DdevBin, args)
	assert.Error(err)
	assert.Contains(string(out), "Too many arguments provided")

}

// TestDescribe tests that the describe command works properly when using the binary.
func TestDescribe(t *testing.T) {
	assert := asrt.New(t)

	for _, v := range DevTestSites {
		// First, try to do a describe from another directory.
		tmpdir := testcommon.CreateTmpDir("")
		cleanup := testcommon.Chdir(tmpdir)
		defer testcommon.CleanupDir(tmpdir)

		args := []string{"describe", v.Name}
		out, err := exec.RunCommand(DdevBin, args)
		assert.NoError(err)
		assert.Contains(string(out), "NAME")
		assert.Contains(string(out), "LOCATION")
		assert.Contains(string(out), v.Name)
		assert.Contains(string(out), "running")

		cleanup()

		cleanup = v.Chdir()

		args = []string{"describe"}
		out, err = exec.RunCommand(DdevBin, args)
		assert.NoError(err)
		assert.Contains(string(out), "NAME")
		assert.Contains(string(out), "LOCATION")
		assert.Contains(string(out), v.Name)
		assert.Contains(string(out), "running")

		// Test describe in current directory with json flag
		args = []string{"describe", "-j"}
		out, err = exec.RunCommand(DdevBin, args)
		assert.NoError(err)
		// Unmarshall the json results results. The describe function only has 4 fields to output
		data := make(log.Fields, 4)
		err = json.Unmarshal([]byte(out), &data)
		assert.NoError(err)
		raw, ok := data["raw"].(map[string]interface{})
		assert.True(ok)
		assert.EqualValues(raw["status"], "running")
		assert.EqualValues(raw["name"], v.Name)
		assert.EqualValues(raw["shortroot"].(string), platform.RenderHomeRootedDir(v.Dir))
		assert.EqualValues(raw["approot"].(string), v.Dir)
		cleanup()
	}
}

// TestDescribeAppFunction performs unit tests on the describeApp function from the working directory.
func TestDescribeAppFunction(t *testing.T) {
	assert := asrt.New(t)
	for _, v := range DevTestSites {
		cleanup := v.Chdir()

		app, err := platform.GetActiveApp("")
		assert.NoError(err)

		desc, err := app.Describe()
		assert.NoError(err)
		assert.EqualValues(desc["status"], platform.SiteRunning)
		assert.EqualValues(app.GetName(), desc["name"])
		assert.EqualValues(platform.RenderHomeRootedDir(v.Dir), desc["shortroot"].(string))
		assert.EqualValues(v.Dir, desc["approot"].(string))

		out, _ := json.Marshal(desc)
		assert.Contains(string(out), app.URL())
		assert.Contains(string(out), app.GetName())
		assert.Contains(string(out), "\"router_status\":\"healthy\"")
		assert.Contains(string(out), platform.RenderHomeRootedDir(v.Dir))

		// Stop the router using docker and then check the describe
		_, err = exec.RunCommand("docker", []string{"stop", "ddev-router"})
		assert.NoError(err)
		desc, err = app.Describe()
		assert.NoError(err)
		out, _ = json.Marshal(desc)
		assert.NoError(err)
		assert.Contains(string(out), "router_status\":\"exited")
		_, err = exec.RunCommand("docker", []string{"start", "ddev-router"})
		assert.NoError(err)

		cleanup()
	}
}

// TestDescribeAppUsingSitename performs unit tests on the describeApp function using the sitename as an argument.
func TestDescribeAppUsingSitename(t *testing.T) {
	assert := asrt.New(t)

	// Create a temporary directory and switch to it for the duration of this test.
	tmpdir := testcommon.CreateTmpDir("describeAppUsingSitename")
	defer testcommon.CleanupDir(tmpdir)
	defer testcommon.Chdir(tmpdir)()

	for _, v := range DevTestSites {
		app, err := platform.GetActiveApp(v.Name)
		assert.NoError(err)
		desc, err := app.Describe()
		assert.NoError(err)
		assert.EqualValues(desc["status"], platform.SiteRunning)
		assert.EqualValues(app.GetName(), desc["name"])
		assert.EqualValues(platform.RenderHomeRootedDir(v.Dir), desc["approot"].(string))

		out, _ := json.Marshal(desc)
		assert.NoError(err)
		assert.Contains(string(out), "running")
		assert.Contains(string(out), platform.RenderHomeRootedDir(v.Dir))
	}
}

// TestDescribeAppWithInvalidParams performs unit tests on the describeApp function using a variety of invalid parameters.
func TestDescribeAppWithInvalidParams(t *testing.T) {
	assert := asrt.New(t)

	// Create a temporary directory and switch to it for the duration of this test.
	tmpdir := testcommon.CreateTmpDir("TestDescribeAppWithInvalidParams")
	defer testcommon.CleanupDir(tmpdir)
	defer testcommon.Chdir(tmpdir)()

	// Ensure describeApp fails from an invalid working directory.
	_, err := platform.GetActiveApp("")
	assert.Error(err)

	// Ensure describeApp fails with invalid site-names.
	_, err = platform.GetActiveApp(util.RandString(16))
	assert.Error(err)

	// Change to a site's working directory and ensure a failure still occurs with a invalid site name.
	cleanup := DevTestSites[0].Chdir()
	_, err = platform.GetActiveApp(util.RandString(16))
	assert.Error(err)
	cleanup()
}
<commit_msg>Fix missed shortroot conversion in test, broke windows test only somehow (#546)<commit_after>package cmd

import (
	"testing"

	"encoding/json"

	"github.com/drud/ddev/pkg/exec"
	"github.com/drud/ddev/pkg/plugins/platform"
	"github.com/drud/ddev/pkg/testcommon"
	"github.com/drud/ddev/pkg/util"
	log "github.com/sirupsen/logrus"
	asrt "github.com/stretchr/testify/assert"
)

// TestDescribeBadArgs ensures the binary behaves as expected when used with invalid arguments or working directories.
func TestDescribeBadArgs(t *testing.T) {
	assert := asrt.New(t)

	// Create a temporary directory and switch to it for the duration of this test.
	tmpdir := testcommon.CreateTmpDir("badargs")
	defer testcommon.CleanupDir(tmpdir)
	defer testcommon.Chdir(tmpdir)()

	// Ensure it fails if we run the vanilla describe outside of an application root.
	args := []string{"describe"}
	out, err := exec.RunCommand(DdevBin, args)
	assert.Error(err)
	assert.Contains(string(out), "Please specify a site name or change directories")

	// Ensure we get a failure if we run a describe on a named application which does not exist.
	args = []string{"describe", util.RandString(16)}
	out, err = exec.RunCommand(DdevBin, args)
	assert.Error(err)
	assert.Contains(string(out), "Unable to find any active site")

	// Ensure we get a failure if using too many arguments.
	args = []string{"describe", util.RandString(16), util.RandString(16)}
	out, err = exec.RunCommand(DdevBin, args)
	assert.Error(err)
	assert.Contains(string(out), "Too many arguments provided")

}

// TestDescribe tests that the describe command works properly when using the binary.
func TestDescribe(t *testing.T) {
	assert := asrt.New(t)

	for _, v := range DevTestSites {
		// First, try to do a describe from another directory.
		tmpdir := testcommon.CreateTmpDir("")
		cleanup := testcommon.Chdir(tmpdir)
		defer testcommon.CleanupDir(tmpdir)

		args := []string{"describe", v.Name}
		out, err := exec.RunCommand(DdevBin, args)
		assert.NoError(err)
		assert.Contains(string(out), "NAME")
		assert.Contains(string(out), "LOCATION")
		assert.Contains(string(out), v.Name)
		assert.Contains(string(out), "running")

		cleanup()

		cleanup = v.Chdir()

		args = []string{"describe"}
		out, err = exec.RunCommand(DdevBin, args)
		assert.NoError(err)
		assert.Contains(string(out), "NAME")
		assert.Contains(string(out), "LOCATION")
		assert.Contains(string(out), v.Name)
		assert.Contains(string(out), "running")

		// Test describe in current directory with json flag
		args = []string{"describe", "-j"}
		out, err = exec.RunCommand(DdevBin, args)
		assert.NoError(err)
		// Unmarshall the json results results. The describe function only has 4 fields to output
		data := make(log.Fields, 4)
		err = json.Unmarshal([]byte(out), &data)
		assert.NoError(err)
		raw, ok := data["raw"].(map[string]interface{})
		assert.True(ok)
		assert.EqualValues(raw["status"], "running")
		assert.EqualValues(raw["name"], v.Name)
		assert.EqualValues(raw["shortroot"].(string), platform.RenderHomeRootedDir(v.Dir))
		assert.EqualValues(raw["approot"].(string), v.Dir)
		cleanup()
	}
}

// TestDescribeAppFunction performs unit tests on the describeApp function from the working directory.
func TestDescribeAppFunction(t *testing.T) {
	assert := asrt.New(t)
	for _, v := range DevTestSites {
		cleanup := v.Chdir()

		app, err := platform.GetActiveApp("")
		assert.NoError(err)

		desc, err := app.Describe()
		assert.NoError(err)
		assert.EqualValues(desc["status"], platform.SiteRunning)
		assert.EqualValues(app.GetName(), desc["name"])
		assert.EqualValues(platform.RenderHomeRootedDir(v.Dir), desc["shortroot"].(string))
		assert.EqualValues(v.Dir, desc["approot"].(string))

		out, _ := json.Marshal(desc)
		assert.Contains(string(out), app.URL())
		assert.Contains(string(out), app.GetName())
		assert.Contains(string(out), "\"router_status\":\"healthy\"")
		assert.Contains(string(out), platform.RenderHomeRootedDir(v.Dir))

		// Stop the router using docker and then check the describe
		_, err = exec.RunCommand("docker", []string{"stop", "ddev-router"})
		assert.NoError(err)
		desc, err = app.Describe()
		assert.NoError(err)
		out, _ = json.Marshal(desc)
		assert.NoError(err)
		assert.Contains(string(out), "router_status\":\"exited")
		_, err = exec.RunCommand("docker", []string{"start", "ddev-router"})
		assert.NoError(err)

		cleanup()
	}
}

// TestDescribeAppUsingSitename performs unit tests on the describeApp function using the sitename as an argument.
func TestDescribeAppUsingSitename(t *testing.T) {
	assert := asrt.New(t)

	// Create a temporary directory and switch to it for the duration of this test.
	tmpdir := testcommon.CreateTmpDir("describeAppUsingSitename")
	defer testcommon.CleanupDir(tmpdir)
	defer testcommon.Chdir(tmpdir)()

	for _, v := range DevTestSites {
		app, err := platform.GetActiveApp(v.Name)
		assert.NoError(err)
		desc, err := app.Describe()
		assert.NoError(err)
		assert.EqualValues(desc["status"], platform.SiteRunning)
		assert.EqualValues(app.GetName(), desc["name"])
		assert.EqualValues(platform.RenderHomeRootedDir(v.Dir), desc["shortroot"].(string))
		assert.EqualValues(v.Dir, desc["approot"].(string))

		out, _ := json.Marshal(desc)
		assert.NoError(err)
		assert.Contains(string(out), "running")
		assert.Contains(string(out), platform.RenderHomeRootedDir(v.Dir))
	}
}

// TestDescribeAppWithInvalidParams performs unit tests on the describeApp function using a variety of invalid parameters.
func TestDescribeAppWithInvalidParams(t *testing.T) {
	assert := asrt.New(t)

	// Create a temporary directory and switch to it for the duration of this test.
	tmpdir := testcommon.CreateTmpDir("TestDescribeAppWithInvalidParams")
	defer testcommon.CleanupDir(tmpdir)
	defer testcommon.Chdir(tmpdir)()

	// Ensure describeApp fails from an invalid working directory.
	_, err := platform.GetActiveApp("")
	assert.Error(err)

	// Ensure describeApp fails with invalid site-names.
	_, err = platform.GetActiveApp(util.RandString(16))
	assert.Error(err)

	// Change to a site's working directory and ensure a failure still occurs with a invalid site name.
	cleanup := DevTestSites[0].Chdir()
	_, err = platform.GetActiveApp(util.RandString(16))
	assert.Error(err)
	cleanup()
}
<|endoftext|>


<commit_before><commit_msg>make grade fix<commit_after><|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"compress/gzip"
	"errors"
	"fmt"
	"github.com/Symantec/Dominator/imageserver/client"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filesystem/scanner"
	"github.com/Symantec/Dominator/lib/filesystem/untar"
	"github.com/Symantec/Dominator/lib/filter"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/image"
	"github.com/Symantec/Dominator/lib/mbr"
	objectclient "github.com/Symantec/Dominator/lib/objectserver/client"
	"github.com/Symantec/Dominator/lib/srpc"
	"io"
	"os"
	"os/exec"
	"strings"
	"time"
)

func addImagefileSubcommand(args []string) {
	imageSClient, objectClient := getClients()
	err := addImagefile(imageSClient, objectClient, args[0], args[1], args[2],
		args[3])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error adding image: \"%s\"\t%s\n", args[0], err)
		os.Exit(1)
	}
	os.Exit(0)
}

func addImagefile(imageSClient *srpc.Client,
	objectClient *objectclient.ObjectClient,
	name, imageFilename, filterFilename, triggersFilename string) error {
	imageExists, err := client.CheckImage(imageSClient, name)
	if err != nil {
		return errors.New("error checking for image existance: " + err.Error())
	}
	if imageExists {
		return errors.New("image exists")
	}
	newImage := new(image.Image)
	if err := loadImageFiles(newImage, objectClient, filterFilename,
		triggersFilename); err != nil {
		return err
	}
	newImage.FileSystem, err = buildImage(imageSClient, newImage.Filter,
		imageFilename)
	if err != nil {
		return errors.New("error building image: " + err.Error())
	}
	if err := spliceComputedFiles(newImage.FileSystem); err != nil {
		return err
	}
	return addImage(imageSClient, name, newImage)
}

func addImage(imageSClient *srpc.Client, name string, img *image.Image) error {
	if *expiresIn > 0 {
		img.ExpiresAt = time.Now().Add(*expiresIn)
	}
	if err := img.Verify(); err != nil {
		return err
	}
	if err := img.VerifyRequiredPaths(requiredPaths); err != nil {
		return err
	}
	if err := client.AddImage(imageSClient, name, img); err != nil {
		return errors.New("remote error: " + err.Error())
	}
	return nil
}

type hasher struct {
	objQ *objectclient.ObjectAdderQueue
}

func (h *hasher) Hash(reader io.Reader, length uint64) (
	hash.Hash, error) {
	hash, err := h.objQ.Add(reader, length)
	if err != nil {
		return hash, errors.New("error sending image data: " + err.Error())
	}
	return hash, nil
}

func buildImage(imageSClient *srpc.Client, filter *filter.Filter,
	imageFilename string) (*filesystem.FileSystem, error) {
	var h hasher
	var err error
	h.objQ, err = objectclient.NewObjectAdderQueue(imageSClient)
	if err != nil {
		return nil, err
	}
	fs, err := buildImageWithHasher(imageSClient, filter, imageFilename, &h)
	if err != nil {
		h.objQ.Close()
		return nil, err
	}
	err = h.objQ.Close()
	if err != nil {
		return nil, err
	}
	return fs, nil
}

func buildImageWithHasher(imageSClient *srpc.Client, filter *filter.Filter,
	imageFilename string, h *hasher) (*filesystem.FileSystem, error) {
	fi, err := os.Lstat(imageFilename)
	if err != nil {
		return nil, err
	}
	if fi.IsDir() {
		sfs, err := scanner.ScanFileSystem(imageFilename, nil, filter, nil, h,
			nil)
		if err != nil {
			return nil, err
		}
		return &sfs.FileSystem, nil
	}
	imageFile, err := os.Open(imageFilename)
	if err != nil {
		return nil, errors.New("error opening image file: " + err.Error())
	}
	defer imageFile.Close()
	if partitionTable, err := mbr.Decode(imageFile); err != nil {
		return nil, err
	} else if partitionTable != nil {
		return buildImageFromRaw(imageSClient, filter, imageFile,
			partitionTable, h)
	}
	var imageReader io.Reader
	if strings.HasSuffix(imageFilename, ".tar") {
		imageReader = imageFile
	} else if strings.HasSuffix(imageFilename, ".tar.gz") ||
		strings.HasSuffix(imageFilename, ".tgz") {
		gzipReader, err := gzip.NewReader(imageFile)
		if err != nil {
			return nil, errors.New(
				"error creating gzip reader: " + err.Error())
		}
		defer gzipReader.Close()
		imageReader = gzipReader
	} else {
		return nil, errors.New("unrecognised image type")
	}
	tarReader := tar.NewReader(imageReader)
	fs, err := untar.Decode(tarReader, h, filter)
	if err != nil {
		return nil, errors.New("error building image: " + err.Error())
	}
	return fs, nil
}

func buildImageFromRaw(imageSClient *srpc.Client, filter *filter.Filter,
	imageFile *os.File, partitionTable *mbr.Mbr,
	h *hasher) (*filesystem.FileSystem, error) {
	var index uint
	var offsetOfLargest, sizeOfLargest uint64
	numPartitions := partitionTable.GetNumPartitions()
	for index = 0; index < numPartitions; index++ {
		offset := partitionTable.GetPartitionOffset(index)
		size := partitionTable.GetPartitionSize(index)
		if size > sizeOfLargest {
			offsetOfLargest = offset
			sizeOfLargest = size
		}
	}
	if sizeOfLargest < 1 {
		return nil, errors.New("unable to find largest partition")
	}
	cmd := exec.Command("mount", "-o",
		fmt.Sprintf("loop,offset=%d", offsetOfLargest), imageFile.Name(),
		"/mnt")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	fs, err := buildImageWithHasher(imageSClient, filter, "/mnt", h)
	cmd = exec.Command("umount", "/mnt")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	return fs, err
}
<commit_msg>Use private namespace in imagetool when mounting RAW files.<commit_after>package main

import (
	"archive/tar"
	"compress/gzip"
	"errors"
	"fmt"
	"github.com/Symantec/Dominator/imageserver/client"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filesystem/scanner"
	"github.com/Symantec/Dominator/lib/filesystem/untar"
	"github.com/Symantec/Dominator/lib/filter"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/image"
	"github.com/Symantec/Dominator/lib/mbr"
	objectclient "github.com/Symantec/Dominator/lib/objectserver/client"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/lib/wsyscall"
	"io"
	"os"
	"os/exec"
	"strings"
	"time"
)

func addImagefileSubcommand(args []string) {
	imageSClient, objectClient := getClients()
	err := addImagefile(imageSClient, objectClient, args[0], args[1], args[2],
		args[3])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error adding image: \"%s\": %s\n", args[0], err)
		os.Exit(1)
	}
	os.Exit(0)
}

func addImagefile(imageSClient *srpc.Client,
	objectClient *objectclient.ObjectClient,
	name, imageFilename, filterFilename, triggersFilename string) error {
	imageExists, err := client.CheckImage(imageSClient, name)
	if err != nil {
		return errors.New("error checking for image existance: " + err.Error())
	}
	if imageExists {
		return errors.New("image exists")
	}
	newImage := new(image.Image)
	if err := loadImageFiles(newImage, objectClient, filterFilename,
		triggersFilename); err != nil {
		return err
	}
	newImage.FileSystem, err = buildImage(imageSClient, newImage.Filter,
		imageFilename)
	if err != nil {
		return errors.New("error building image: " + err.Error())
	}
	if err := spliceComputedFiles(newImage.FileSystem); err != nil {
		return err
	}
	return addImage(imageSClient, name, newImage)
}

func addImage(imageSClient *srpc.Client, name string, img *image.Image) error {
	if *expiresIn > 0 {
		img.ExpiresAt = time.Now().Add(*expiresIn)
	}
	if err := img.Verify(); err != nil {
		return err
	}
	if err := img.VerifyRequiredPaths(requiredPaths); err != nil {
		return err
	}
	if err := client.AddImage(imageSClient, name, img); err != nil {
		return errors.New("remote error: " + err.Error())
	}
	return nil
}

type hasher struct {
	objQ *objectclient.ObjectAdderQueue
}

func (h *hasher) Hash(reader io.Reader, length uint64) (
	hash.Hash, error) {
	hash, err := h.objQ.Add(reader, length)
	if err != nil {
		return hash, errors.New("error sending image data: " + err.Error())
	}
	return hash, nil
}

func buildImage(imageSClient *srpc.Client, filter *filter.Filter,
	imageFilename string) (*filesystem.FileSystem, error) {
	var h hasher
	var err error
	h.objQ, err = objectclient.NewObjectAdderQueue(imageSClient)
	if err != nil {
		return nil, err
	}
	fs, err := buildImageWithHasher(imageSClient, filter, imageFilename, &h)
	if err != nil {
		h.objQ.Close()
		return nil, err
	}
	err = h.objQ.Close()
	if err != nil {
		return nil, err
	}
	return fs, nil
}

func buildImageWithHasher(imageSClient *srpc.Client, filter *filter.Filter,
	imageFilename string, h *hasher) (*filesystem.FileSystem, error) {
	fi, err := os.Lstat(imageFilename)
	if err != nil {
		return nil, err
	}
	if fi.IsDir() {
		sfs, err := scanner.ScanFileSystem(imageFilename, nil, filter, nil, h,
			nil)
		if err != nil {
			return nil, err
		}
		return &sfs.FileSystem, nil
	}
	imageFile, err := os.Open(imageFilename)
	if err != nil {
		return nil, errors.New("error opening image file: " + err.Error())
	}
	defer imageFile.Close()
	if partitionTable, err := mbr.Decode(imageFile); err != nil {
		return nil, err
	} else if partitionTable != nil {
		return buildImageFromRaw(imageSClient, filter, imageFile,
			partitionTable, h)
	}
	var imageReader io.Reader
	if strings.HasSuffix(imageFilename, ".tar") {
		imageReader = imageFile
	} else if strings.HasSuffix(imageFilename, ".tar.gz") ||
		strings.HasSuffix(imageFilename, ".tgz") {
		gzipReader, err := gzip.NewReader(imageFile)
		if err != nil {
			return nil, errors.New(
				"error creating gzip reader: " + err.Error())
		}
		defer gzipReader.Close()
		imageReader = gzipReader
	} else {
		return nil, errors.New("unrecognised image type")
	}
	tarReader := tar.NewReader(imageReader)
	fs, err := untar.Decode(tarReader, h, filter)
	if err != nil {
		return nil, errors.New("error building image: " + err.Error())
	}
	return fs, nil
}

func buildImageFromRaw(imageSClient *srpc.Client, filter *filter.Filter,
	imageFile *os.File, partitionTable *mbr.Mbr,
	h *hasher) (*filesystem.FileSystem, error) {
	var index uint
	var offsetOfLargest, sizeOfLargest uint64
	numPartitions := partitionTable.GetNumPartitions()
	for index = 0; index < numPartitions; index++ {
		offset := partitionTable.GetPartitionOffset(index)
		size := partitionTable.GetPartitionSize(index)
		if size > sizeOfLargest {
			offsetOfLargest = offset
			sizeOfLargest = size
		}
	}
	if sizeOfLargest < 1 {
		return nil, errors.New("unable to find largest partition")
	}
	if err := wsyscall.UnshareMountNamespace(); err != nil {
		return nil, errors.New(
			"error unsharing mount namespace: " + err.Error())
	}
	cmd := exec.Command("mount", "-o",
		fmt.Sprintf("loop,offset=%d", offsetOfLargest), imageFile.Name(),
		"/mnt")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	return buildImageWithHasher(imageSClient, filter, "/mnt", h)
}
<|endoftext|>