<commit_before>// Copyright 2013 Matthew Baird
// 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 elastigo

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"
	"time"
)

const (
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer = 16384
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs = 100
	// Max delay before forcing a flush to Elasticearch
	BulkDelaySeconds = 5
	// maximum wait shutdown seconds
	MAX_SHUTDOWN_SECS = 5
)

type ErrorBuffer struct {
	Err error
	Buf *bytes.Buffer
}

// A bulk indexer creates goroutines, and channels for connecting and sending data
// to elasticsearch in bulk, using buffers.
type BulkIndexer struct {
	conn *Conn

	// We are creating a variable defining the func responsible for sending
	// to allow a mock sendor for test purposes
	Sender func(*bytes.Buffer) error

	// If we encounter an error in sending, we are going to retry for this long
	// before returning an error
	// if 0 it will not retry
	RetryForSeconds int

	// channel for getting errors
	ErrorChannel chan *ErrorBuffer

	// channel for sending to background indexer
	bulkChannel chan []byte

	// numErrors is a running total of errors seen
	numErrors uint64

	// shutdown channel
	shutdownChan chan chan struct{}
	// Channel to shutdown http send go-routines
	httpDoneChan chan bool
	// channel to shutdown timer
	timerDoneChan chan bool
	// channel to shutdown doc go-routines
	docDoneChan chan bool

	// Channel to send a complete byte.Buffer to the http sendor
	sendBuf chan *bytes.Buffer
	// byte buffer for docs that have been converted to bytes, but not yet sent
	buf *bytes.Buffer
	// Buffer for Max number of time before forcing flush
	BufferDelayMax time.Duration
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer int // 1048576
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs int // 100

	// Number of documents we have send through so far on this session
	docCt int
	// Max number of http conns in flight at one time
	maxConns int
	// If we are indexing enough docs per bufferdelaymax, we won't need to do time
	// based eviction, else we do.
	needsTimeBasedFlush bool
	// Lock for document writes/operations
	mu sync.Mutex
	// Wait Group for the http sends
	sendWg *sync.WaitGroup
}

func (b *BulkIndexer) NumErrors() uint64 {
	return b.numErrors
}

func (c *Conn) NewBulkIndexer(maxConns int) *BulkIndexer {
	b := BulkIndexer{conn: c, sendBuf: make(chan *bytes.Buffer, maxConns)}
	b.needsTimeBasedFlush = true
	b.buf = new(bytes.Buffer)
	b.maxConns = maxConns
	b.BulkMaxBuffer = BulkMaxBuffer
	b.BulkMaxDocs = BulkMaxDocs
	b.BufferDelayMax = time.Duration(BulkDelaySeconds) * time.Second
	b.bulkChannel = make(chan []byte, 100)
	b.sendWg = new(sync.WaitGroup)
	b.docDoneChan = make(chan bool)
	b.timerDoneChan = make(chan bool)
	b.httpDoneChan = make(chan bool)
	return &b
}

// A bulk indexer with more control over error handling
//    @maxConns is the max number of in flight http requests
//    @retrySeconds is # of seconds to wait before retrying falied requests
//
//   done := make(chan bool)
//   BulkIndexerGlobalRun(100, done)
func (c *Conn) NewBulkIndexerErrors(maxConns, retrySeconds int) *BulkIndexer {
	b := c.NewBulkIndexer(maxConns)
	b.RetryForSeconds = retrySeconds
	b.ErrorChannel = make(chan *ErrorBuffer, 20)
	return b
}

// Starts this bulk Indexer running, this Run opens a go routine so is
// Non blocking
func (b *BulkIndexer) Start() {
	b.shutdownChan = make(chan chan struct{})

	go func() {
		// XXX(j): Refactor this stuff to use an interface.
		if b.Sender == nil {
			b.Sender = b.Send
		}
		// Backwards compatibility
		b.startHttpSender()
		b.startDocChannel()
		b.startTimer()
		ch := <-b.shutdownChan
		b.Flush()
		b.shutdown()
		ch <- struct{}{}
		close(ch)
	}()
}

// Stop stops the bulk indexer, blocking the caller until it is complete.
func (b *BulkIndexer) Stop() {
	ch := make(chan struct{})
	b.shutdownChan <- ch
	<-ch
	close(b.shutdownChan)
}

// Make a channel that will close when the given WaitGroup is done.
func wgChan(wg *sync.WaitGroup) <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		wg.Wait()
		close(ch)
	}()
	return ch
}

func (b *BulkIndexer) PendingDocuments() int {
	return b.docCt
}

// Flush all current documents to ElasticSearch
func (b *BulkIndexer) Flush() {
	b.mu.Lock()
	if b.docCt > 0 {
		b.send(b.buf)
	}
	b.mu.Unlock()
	for {
		select {
		case <-wgChan(b.sendWg):
			// done
			return
		case <-time.After(time.Second * time.Duration(MAX_SHUTDOWN_SECS)):
			// timeout!
			return
		}
	}
}

func (b *BulkIndexer) startHttpSender() {

	// this sends http requests to elasticsearch it uses maxConns to open up that
	// many goroutines, each of which will synchronously call ElasticSearch
	// in theory, the whole set will cause a backup all the way to IndexBulk if
	// we have consumed all maxConns
	for i := 0; i < b.maxConns; i++ {
		go func() {
			for {
				select {
				case buf := <-b.sendBuf:
					b.sendWg.Add(1)
					// Copy for the potential re-send.
					bufCopy := bytes.NewBuffer(buf.Bytes())
					err := b.Sender(buf)

					// Perhaps a b.FailureStrategy(err)  ??  with different types of strategies
					//  1.  Retry, then panic
					//  2.  Retry then return error and let runner decide
					//  3.  Retry, then log to disk?   retry later?
					if err != nil {
						if b.RetryForSeconds > 0 {
							time.Sleep(time.Second * time.Duration(b.RetryForSeconds))
							err = b.Sender(bufCopy)
							if err == nil {
								// Successfully re-sent with no error
								b.sendWg.Done()
								continue
							}
						}
						if b.ErrorChannel != nil {
							b.ErrorChannel <- &ErrorBuffer{err, buf}
						}
					}
					b.sendWg.Done()
				case <-b.httpDoneChan:
					// shutdown this go routine
					return
				}

			}
		}()
	}
}

// start a timer for checking back and forcing flush ever BulkDelaySeconds seconds
// even if we haven't hit max messages/size
func (b *BulkIndexer) startTimer() {
	ticker := time.NewTicker(b.BufferDelayMax)
	go func() {
		for {
			select {
			case <-ticker.C:
				b.mu.Lock()
				// don't send unless last sendor was the time,
				// otherwise an indication of other thresholds being hit
				// where time isn't needed
				if b.buf.Len() > 0 && b.needsTimeBasedFlush {
					b.needsTimeBasedFlush = true
					b.send(b.buf)
				} else if b.buf.Len() > 0 {
					b.needsTimeBasedFlush = true
				}
				b.mu.Unlock()
			case <-b.timerDoneChan:
				// shutdown this go routine
				return
			}

		}
	}()
}

func (b *BulkIndexer) startDocChannel() {
	// This goroutine accepts incoming byte arrays from the IndexBulk function and
	// writes to buffer
	go func() {
		for {
			select {
			case docBytes := <-b.bulkChannel:
				b.mu.Lock()
				b.docCt += 1
				b.buf.Write(docBytes)
				if b.buf.Len() >= b.BulkMaxBuffer || b.docCt >= b.BulkMaxDocs {
					b.needsTimeBasedFlush = false
					//log.Printf("Send due to size:  docs=%d  bufsize=%d", b.docCt, b.buf.Len())
					b.send(b.buf)
				}
				b.mu.Unlock()
			case <-b.docDoneChan:
				// shutdown this go routine
				return
			}
		}
	}()
}

func (b *BulkIndexer) send(buf *bytes.Buffer) {
	//b2 := *b.buf
	b.sendBuf <- buf
	b.buf = new(bytes.Buffer)
	//	b.buf.Reset()
	b.docCt = 0
}

func (b *BulkIndexer) shutdown() {
	// This must be called After flush
	b.docDoneChan <- true
	b.timerDoneChan <- true
	for i := 0; i < b.maxConns; i++ {
		b.httpDoneChan <- true
	}
}

// The index bulk API adds or updates a typed JSON document to a specific index, making it searchable.
// it operates by buffering requests, and ocassionally flushing to elasticsearch
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func (b *BulkIndexer) Index(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("index", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Update(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("update", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Delete(index, _type, id string, refresh bool) {
	queryLine := fmt.Sprintf("{\"delete\":{\"_index\":%q,\"_type\":%q,\"_id\":%q,\"refresh\":%t}}\n", index, _type, id, refresh)
	b.bulkChannel <- []byte(queryLine)
	return
}

func (b *BulkIndexer) UpdateWithPartialDoc(index string, _type string, id, ttl string, date *time.Time, partialDoc interface{}, upsert bool, refresh bool) error {

	var data map[string]interface{} = make(map[string]interface{})

	data["doc"] = partialDoc
	if upsert {
		data["doc_as_upsert"] = true
	}
	return b.Update(index, _type, id, ttl, date, data, refresh)
}

// This does the actual send of a buffer, which has already been formatted
// into bytes of ES formatted bulk data
func (b *BulkIndexer) Send(buf *bytes.Buffer) error {
	type responseStruct struct {
		Took   int64                    `json:"took"`
		Errors bool                     `json:"errors"`
		Items  []map[string]interface{} `json:"items"`
	}

	response := responseStruct{}

	body, err := b.conn.DoCommand("POST", "/_bulk", nil, buf)

	if err != nil {
		b.numErrors += 1
		return err
	}
	// check for response errors, bulk insert will give 200 OK but then include errors in response
	jsonErr := json.Unmarshal(body, &response)
	if jsonErr == nil {
		if response.Errors {
			b.numErrors += uint64(len(response.Items))
			return fmt.Errorf("Bulk Insertion Error. Failed item count [%d]", len(response.Items))
		}
	}
	return nil
}

// Given a set of arguments for index, type, id, data create a set of bytes that is formatted for bulkd index
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func WriteBulkBytes(op string, index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) ([]byte, error) {
	// only index and update are currently supported
	if op != "index" && op != "update" {
		return nil, errors.New(fmt.Sprintf("Operation '%s' is not yet supported", op))
	}

	// First line
	buf := bytes.Buffer{}
	buf.WriteString(fmt.Sprintf(`{"%s":{"_index":"`, op))
	buf.WriteString(index)
	buf.WriteString(`","_type":"`)
	buf.WriteString(_type)
	buf.WriteString(`"`)
	if len(id) > 0 {
		buf.WriteString(`,"_id":"`)
		buf.WriteString(id)
		buf.WriteString(`"`)
	}

	if op == "update" {
		buf.WriteString(`,"retry_on_conflict":3`)
	}

	if len(ttl) > 0 {
		buf.WriteString(`,"ttl":"`)
		buf.WriteString(ttl)
		buf.WriteString(`"`)
	}
	if date != nil {
		buf.WriteString(`,"_timestamp":"`)
		buf.WriteString(strconv.FormatInt(date.UnixNano()/1e6, 10))
		buf.WriteString(`"`)
	}
	if refresh {
		buf.WriteString(`,"refresh":true`)
	}
	buf.WriteString(`}}`)
	buf.WriteRune('\n')
	//buf.WriteByte('\n')
	switch v := data.(type) {
	case *bytes.Buffer:
		io.Copy(&buf, v)
	case []byte:
		buf.Write(v)
	case string:
		buf.WriteString(v)
	default:
		body, jsonErr := json.Marshal(data)
		if jsonErr != nil {
			return nil, jsonErr
		}
		buf.Write(body)
	}
	buf.WriteRune('\n')
	return buf.Bytes(), nil
}
<commit_msg>Support reporting stats from BulkUpdater<commit_after>// Copyright 2013 Matthew Baird
// 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 elastigo

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"sync"
	"time"
)

const (
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer = 16384
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs = 100
	// Max delay before forcing a flush to Elasticearch
	BulkDelaySeconds = 5
	// maximum wait shutdown seconds
	MAX_SHUTDOWN_SECS = 5
)

type ErrorBuffer struct {
	Err error
	Buf *bytes.Buffer
}

// A bulk indexer creates goroutines, and channels for connecting and sending data
// to elasticsearch in bulk, using buffers.
type BulkIndexer struct {
	conn *Conn

	// We are creating a variable defining the func responsible for sending
	// to allow a mock sendor for test purposes
	Sender func(*bytes.Buffer) error

	// If we encounter an error in sending, we are going to retry for this long
	// before returning an error
	// if 0 it will not retry
	RetryForSeconds int

	// channel for getting errors
	ErrorChannel chan *ErrorBuffer

	// channel for sending to background indexer
	bulkChannel chan []byte

	// numErrors is a running total of errors seen
	numErrors uint64

	// shutdown channel
	shutdownChan chan chan struct{}
	// Channel to shutdown http send go-routines
	httpDoneChan chan bool
	// channel to shutdown timer
	timerDoneChan chan bool
	// channel to shutdown doc go-routines
	docDoneChan chan bool

	// Channel to send a complete byte.Buffer to the http sendor
	sendBuf chan *bytes.Buffer
	// byte buffer for docs that have been converted to bytes, but not yet sent
	buf *bytes.Buffer
	// Buffer for Max number of time before forcing flush
	BufferDelayMax time.Duration
	// Max buffer size in bytes before flushing to elasticsearch
	BulkMaxBuffer int // 1048576
	// Max number of Docs to hold in buffer before forcing flush
	BulkMaxDocs int // 100

	// Number of documents we have send through so far on this session
	docCt int
	// Max number of http conns in flight at one time
	maxConns int
	// If we are indexing enough docs per bufferdelaymax, we won't need to do time
	// based eviction, else we do.
	needsTimeBasedFlush bool
	// Lock for document writes/operations
	mu sync.Mutex
	// Wait Group for the http sends
	sendWg *sync.WaitGroup
	// Stats channel reporting on timings of requests to ElasticSearch
	StatsCh chan *BulkIndexStats
}

type BulkIndexStats struct {
	NumErrors int
	Duration  time.Duration
}

func (b *BulkIndexer) NumErrors() uint64 {
	return b.numErrors
}

func (c *Conn) NewBulkIndexer(maxConns int, reportStats bool) *BulkIndexer {
	b := BulkIndexer{conn: c, sendBuf: make(chan *bytes.Buffer, maxConns)}
	b.needsTimeBasedFlush = true
	b.buf = new(bytes.Buffer)
	b.maxConns = maxConns
	b.BulkMaxBuffer = BulkMaxBuffer
	b.BulkMaxDocs = BulkMaxDocs
	b.BufferDelayMax = time.Duration(BulkDelaySeconds) * time.Second
	b.bulkChannel = make(chan []byte, 100)
	b.sendWg = new(sync.WaitGroup)
	b.docDoneChan = make(chan bool)
	b.timerDoneChan = make(chan bool)
	b.httpDoneChan = make(chan bool)
	if reportStats {
		b.StatsCh = make(chan *BulkIndexStats, 100) // Limit stats we send for proof of concept
	}
	return &b
}

// A bulk indexer with more control over error handling
//    @maxConns is the max number of in flight http requests
//    @retrySeconds is # of seconds to wait before retrying falied requests
//
//   done := make(chan bool)
//   BulkIndexerGlobalRun(100, done)
func (c *Conn) NewBulkIndexerErrors(maxConns, retrySeconds int) *BulkIndexer {
	b := c.NewBulkIndexer(maxConns, false)
	b.RetryForSeconds = retrySeconds
	b.ErrorChannel = make(chan *ErrorBuffer, 20)
	return b
}

// Starts this bulk Indexer running, this Run opens a go routine so is
// Non blocking
func (b *BulkIndexer) Start() {
	b.shutdownChan = make(chan chan struct{})

	go func() {
		// XXX(j): Refactor this stuff to use an interface.
		if b.Sender == nil {
			b.Sender = b.Send
		}
		// Backwards compatibility
		b.startHttpSender()
		b.startDocChannel()
		b.startTimer()
		ch := <-b.shutdownChan
		b.Flush()
		b.shutdown()
		ch <- struct{}{}
		close(ch)
	}()
}

// Stop stops the bulk indexer, blocking the caller until it is complete.
func (b *BulkIndexer) Stop() {
	ch := make(chan struct{})
	b.shutdownChan <- ch
	<-ch
	close(b.shutdownChan)
}

// Make a channel that will close when the given WaitGroup is done.
func wgChan(wg *sync.WaitGroup) <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		wg.Wait()
		close(ch)
	}()
	return ch
}

func (b *BulkIndexer) PendingDocuments() int {
	return b.docCt
}

// Flush all current documents to ElasticSearch
func (b *BulkIndexer) Flush() {
	b.mu.Lock()
	if b.docCt > 0 {
		b.send(b.buf)
	}
	b.mu.Unlock()
	for {
		select {
		case <-wgChan(b.sendWg):
			// done
			return
		case <-time.After(time.Second * time.Duration(MAX_SHUTDOWN_SECS)):
			// timeout!
			return
		}
	}
}

func (b *BulkIndexer) startHttpSender() {

	// this sends http requests to elasticsearch it uses maxConns to open up that
	// many goroutines, each of which will synchronously call ElasticSearch
	// in theory, the whole set will cause a backup all the way to IndexBulk if
	// we have consumed all maxConns
	for i := 0; i < b.maxConns; i++ {
		go func() {
			for {
				select {
				case buf := <-b.sendBuf:
					b.sendWg.Add(1)
					// Copy for the potential re-send.
					bufCopy := bytes.NewBuffer(buf.Bytes())
					err := b.Sender(buf)

					// Perhaps a b.FailureStrategy(err)  ??  with different types of strategies
					//  1.  Retry, then panic
					//  2.  Retry then return error and let runner decide
					//  3.  Retry, then log to disk?   retry later?
					if err != nil {
						if b.RetryForSeconds > 0 {
							time.Sleep(time.Second * time.Duration(b.RetryForSeconds))
							err = b.Sender(bufCopy)
							if err == nil {
								// Successfully re-sent with no error
								b.sendWg.Done()
								continue
							}
						}
						if b.ErrorChannel != nil {
							b.ErrorChannel <- &ErrorBuffer{err, buf}
						}
					}
					b.sendWg.Done()
				case <-b.httpDoneChan:
					// shutdown this go routine
					return
				}

			}
		}()
	}
}

// start a timer for checking back and forcing flush ever BulkDelaySeconds seconds
// even if we haven't hit max messages/size
func (b *BulkIndexer) startTimer() {
	ticker := time.NewTicker(b.BufferDelayMax)
	go func() {
		for {
			select {
			case <-ticker.C:
				b.mu.Lock()
				// don't send unless last sendor was the time,
				// otherwise an indication of other thresholds being hit
				// where time isn't needed
				if b.buf.Len() > 0 && b.needsTimeBasedFlush {
					b.needsTimeBasedFlush = true
					b.send(b.buf)
				} else if b.buf.Len() > 0 {
					b.needsTimeBasedFlush = true
				}
				b.mu.Unlock()
			case <-b.timerDoneChan:
				// shutdown this go routine
				return
			}

		}
	}()
}

func (b *BulkIndexer) startDocChannel() {
	// This goroutine accepts incoming byte arrays from the IndexBulk function and
	// writes to buffer
	go func() {
		for {
			select {
			case docBytes := <-b.bulkChannel:
				b.mu.Lock()
				b.docCt += 1
				b.buf.Write(docBytes)
				if b.buf.Len() >= b.BulkMaxBuffer || b.docCt >= b.BulkMaxDocs {
					b.needsTimeBasedFlush = false
					//log.Printf("Send due to size:  docs=%d  bufsize=%d", b.docCt, b.buf.Len())
					b.send(b.buf)
				}
				b.mu.Unlock()
			case <-b.docDoneChan:
				// shutdown this go routine
				return
			}
		}
	}()
}

func (b *BulkIndexer) send(buf *bytes.Buffer) {
	//b2 := *b.buf
	b.sendBuf <- buf
	b.buf = new(bytes.Buffer)
	//	b.buf.Reset()
	b.docCt = 0
}

func (b *BulkIndexer) shutdown() {
	// This must be called After flush
	b.docDoneChan <- true
	b.timerDoneChan <- true
	for i := 0; i < b.maxConns; i++ {
		b.httpDoneChan <- true
	}
}

// The index bulk API adds or updates a typed JSON document to a specific index, making it searchable.
// it operates by buffering requests, and ocassionally flushing to elasticsearch
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func (b *BulkIndexer) Index(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("index", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Update(index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) error {
	//{ "index" : { "_index" : "test", "_type" : "type1", "_id" : "1" } }
	by, err := WriteBulkBytes("update", index, _type, id, ttl, date, data, refresh)
	if err != nil {
		return err
	}
	b.bulkChannel <- by
	return nil
}

func (b *BulkIndexer) Delete(index, _type, id string, refresh bool) {
	queryLine := fmt.Sprintf("{\"delete\":{\"_index\":%q,\"_type\":%q,\"_id\":%q,\"refresh\":%t}}\n", index, _type, id, refresh)
	b.bulkChannel <- []byte(queryLine)
	return
}

func (b *BulkIndexer) UpdateWithPartialDoc(index string, _type string, id, ttl string, date *time.Time, partialDoc interface{}, upsert bool, refresh bool) error {

	var data map[string]interface{} = make(map[string]interface{})

	data["doc"] = partialDoc
	if upsert {
		data["doc_as_upsert"] = true
	}
	return b.Update(index, _type, id, ttl, date, data, refresh)
}

// This does the actual send of a buffer, which has already been formatted
// into bytes of ES formatted bulk data
func (b *BulkIndexer) Send(buf *bytes.Buffer) error {
	var err error
	type responseStruct struct {
		Took   int64                    `json:"took"`
		Errors bool                     `json:"errors"`
		Items  []map[string]interface{} `json:"items"`
	}

	numErrors := 0

	response := responseStruct{}

	start := time.Now()
	body, err := b.conn.DoCommand("POST", "/_bulk", nil, buf)

	if err != nil {
		b.numErrors += 1
		return err
	}
	// check for response errors, bulk insert will give 200 OK but then include errors in response
	jsonErr := json.Unmarshal(body, &response)
	if jsonErr == nil {
		if response.Errors {
			numErrors = len(response.Items)
			b.numErrors += uint64(len(response.Items))
			err = fmt.Errorf("Bulk Insertion Error. Failed item count [%d]", len(response.Items))
		}
	}
	b.StatsCh <- &BulkIndexStats{NumErrors: numErrors, Duration: time.Now().Sub(start)}
	return nil
}

// Given a set of arguments for index, type, id, data create a set of bytes that is formatted for bulkd index
// http://www.elasticsearch.org/guide/reference/api/bulk.html
func WriteBulkBytes(op string, index string, _type string, id, ttl string, date *time.Time, data interface{}, refresh bool) ([]byte, error) {
	// only index and update are currently supported
	if op != "index" && op != "update" {
		return nil, errors.New(fmt.Sprintf("Operation '%s' is not yet supported", op))
	}

	// First line
	buf := bytes.Buffer{}
	buf.WriteString(fmt.Sprintf(`{"%s":{"_index":"`, op))
	buf.WriteString(index)
	buf.WriteString(`","_type":"`)
	buf.WriteString(_type)
	buf.WriteString(`"`)
	if len(id) > 0 {
		buf.WriteString(`,"_id":"`)
		buf.WriteString(id)
		buf.WriteString(`"`)
	}

	if op == "update" {
		buf.WriteString(`,"retry_on_conflict":3`)
	}

	if len(ttl) > 0 {
		buf.WriteString(`,"ttl":"`)
		buf.WriteString(ttl)
		buf.WriteString(`"`)
	}
	if date != nil {
		buf.WriteString(`,"_timestamp":"`)
		buf.WriteString(strconv.FormatInt(date.UnixNano()/1e6, 10))
		buf.WriteString(`"`)
	}
	if refresh {
		buf.WriteString(`,"refresh":true`)
	}
	buf.WriteString(`}}`)
	buf.WriteRune('\n')
	//buf.WriteByte('\n')
	switch v := data.(type) {
	case *bytes.Buffer:
		io.Copy(&buf, v)
	case []byte:
		buf.Write(v)
	case string:
		buf.WriteString(v)
	default:
		body, jsonErr := json.Marshal(data)
		if jsonErr != nil {
			return nil, jsonErr
		}
		buf.Write(body)
	}
	buf.WriteRune('\n')
	return buf.Bytes(), nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSDynamoDbTable(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSDynamoDbTableDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSDynamoDbConfigInitialState,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckInitialAWSDynamoDbTableExists("aws_dynamodb_table.basic-dynamodb-table"),
				),
			},
			resource.TestStep{
				Config: testAccAWSDynamoDbConfigAddSecondaryGSI,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckDynamoDbTableWasUpdated("aws_dynamodb_table.basic-dynamodb-table"),
				),
			},
		},
	})
}

func testAccCheckAWSDynamoDbTableDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).dynamodbconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_dynamodb_table" {
			continue
		}

		fmt.Printf("[DEBUG] Checking if DynamoDB table %s exists", rs.Primary.ID)
		// Check if queue exists by checking for its attributes
		params := &dynamodb.DescribeTableInput{
			TableName: aws.String(rs.Primary.ID),
		}
		_, err := conn.DescribeTable(params)
		if err == nil {
			return fmt.Errorf("DynamoDB table %s still exists. Failing!", rs.Primary.ID)
		}

		// Verify the error is what we want
		_, ok := err.(awserr.Error)
		if !ok {
			return err
		}
	}

	return nil
}

func testAccCheckInitialAWSDynamoDbTableExists(n string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		fmt.Printf("[DEBUG] Trying to create initial table state!")
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No DynamoDB table name specified!")
		}

		conn := testAccProvider.Meta().(*AWSClient).dynamodbconn

		params := &dynamodb.DescribeTableInput{
			TableName: aws.String(rs.Primary.ID),
		}

		resp, err := conn.DescribeTable(params)

		if err != nil {
			fmt.Printf("[ERROR] Problem describing table '%s': %s", rs.Primary.ID, err)
			return err
		}

		table := resp.Table

		fmt.Printf("[DEBUG] Checking on table %s", rs.Primary.ID)

		if *table.ProvisionedThroughput.WriteCapacityUnits != 20 {
			return fmt.Errorf("Provisioned write capacity was %d, not 20!", table.ProvisionedThroughput.WriteCapacityUnits)
		}

		if *table.ProvisionedThroughput.ReadCapacityUnits != 10 {
			return fmt.Errorf("Provisioned read capacity was %d, not 10!", table.ProvisionedThroughput.ReadCapacityUnits)
		}

		attrCount := len(table.AttributeDefinitions)
		gsiCount := len(table.GlobalSecondaryIndexes)
		lsiCount := len(table.LocalSecondaryIndexes)

		if attrCount != 4 {
			return fmt.Errorf("There were %d attributes, not 4 like there should have been!", attrCount)
		}

		if gsiCount != 1 {
			return fmt.Errorf("There were %d GSIs, not 1 like there should have been!", gsiCount)
		}

		if lsiCount != 1 {
			return fmt.Errorf("There were %d LSIs, not 1 like there should have been!", lsiCount)
		}

		attrmap := dynamoDbAttributesToMap(&table.AttributeDefinitions)
		if attrmap["TestTableHashKey"] != "S" {
			return fmt.Errorf("Test table hash key was of type %s instead of S!", attrmap["TestTableHashKey"])
		}
		if attrmap["TestTableRangeKey"] != "S" {
			return fmt.Errorf("Test table range key was of type %s instead of S!", attrmap["TestTableRangeKey"])
		}
		if attrmap["TestLSIRangeKey"] != "N" {
			return fmt.Errorf("Test table LSI range key was of type %s instead of N!", attrmap["TestLSIRangeKey"])
		}
		if attrmap["TestGSIRangeKey"] != "S" {
			return fmt.Errorf("Test table GSI range key was of type %s instead of S!", attrmap["TestGSIRangeKey"])
		}

		return nil
	}
}

func testAccCheckDynamoDbTableWasUpdated(n string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No DynamoDB table name specified!")
		}

		conn := testAccProvider.Meta().(*AWSClient).dynamodbconn

		params := &dynamodb.DescribeTableInput{
			TableName: aws.String(rs.Primary.ID),
		}
		resp, err := conn.DescribeTable(params)
		table := resp.Table

		if err != nil {
			return err
		}

		attrCount := len(table.AttributeDefinitions)
		gsiCount := len(table.GlobalSecondaryIndexes)
		lsiCount := len(table.LocalSecondaryIndexes)

		if attrCount != 4 {
			return fmt.Errorf("There were %d attributes, not 4 like there should have been!", attrCount)
		}

		if gsiCount != 1 {
			return fmt.Errorf("There were %d GSIs, not 1 like there should have been!", gsiCount)
		}

		if lsiCount != 1 {
			return fmt.Errorf("There were %d LSIs, not 1 like there should have been!", lsiCount)
		}

		if dynamoDbGetGSIIndex(&table.GlobalSecondaryIndexes, "ReplacementTestTableGSI") == -1 {
			return fmt.Errorf("Could not find GSI named 'ReplacementTestTableGSI' in the table!")
		}

		if dynamoDbGetGSIIndex(&table.GlobalSecondaryIndexes, "InitialTestTableGSI") != -1 {
			return fmt.Errorf("Should have removed 'InitialTestTableGSI' but it still exists!")
		}

		attrmap := dynamoDbAttributesToMap(&table.AttributeDefinitions)
		if attrmap["TestTableHashKey"] != "S" {
			return fmt.Errorf("Test table hash key was of type %s instead of S!", attrmap["TestTableHashKey"])
		}
		if attrmap["TestTableRangeKey"] != "S" {
			return fmt.Errorf("Test table range key was of type %s instead of S!", attrmap["TestTableRangeKey"])
		}
		if attrmap["TestLSIRangeKey"] != "N" {
			return fmt.Errorf("Test table LSI range key was of type %s instead of N!", attrmap["TestLSIRangeKey"])
		}
		if attrmap["ReplacementGSIRangeKey"] != "N" {
			return fmt.Errorf("Test table replacement GSI range key was of type %s instead of N!", attrmap["ReplacementGSIRangeKey"])
		}

		return nil
	}
}

func dynamoDbGetGSIIndex(gsiList *[]*dynamodb.GlobalSecondaryIndexDescription, target string) int {
	for idx, gsiObject := range *gsiList {
		if *gsiObject.IndexName == target {
			return idx
		}
	}

	return -1
}

func dynamoDbAttributesToMap(attributes *[]*dynamodb.AttributeDefinition) map[string]string {
	attrmap := make(map[string]string)

	for _, attrdef := range *attributes {
		attrmap[*(attrdef.AttributeName)] = *(attrdef.AttributeType)
	}

	return attrmap
}

const testAccAWSDynamoDbConfigInitialState = `
resource "aws_dynamodb_table" "basic-dynamodb-table" {
    name = "TerraformTestTable"
		read_capacity = 10
		write_capacity = 20
		hash_key = "TestTableHashKey"
		range_key = "TestTableRangeKey"
		attribute {
			name = "TestTableHashKey"
			type = "S"
		}
		attribute {
			name = "TestTableRangeKey"
			type = "S"
		}
		attribute {
			name = "TestLSIRangeKey"
			type = "N"
		}
		attribute {
			name = "TestGSIRangeKey"
			type = "S"
		}
		local_secondary_index {
			name = "TestTableLSI"
			range_key = "TestLSIRangeKey"
			projection_type = "ALL"
		}
		global_secondary_index {
			name = "InitialTestTableGSI"
			hash_key = "TestTableHashKey"
			range_key = "TestGSIRangeKey"
			write_capacity = 10
			read_capacity = 10
			projection_type = "ALL"
		}
}
`

const testAccAWSDynamoDbConfigAddSecondaryGSI = `
resource "aws_dynamodb_table" "basic-dynamodb-table" {
    name = "TerraformTestTable"
		read_capacity = 20
		write_capacity = 20
		hash_key = "TestTableHashKey"
		range_key = "TestTableRangeKey"
		attribute {
			name = "TestTableHashKey"
			type = "S"
		}
		attribute {
			name = "TestTableRangeKey"
			type = "S"
		}
		attribute {
			name = "TestLSIRangeKey"
			type = "N"
		}
		attribute {
			name = "ReplacementGSIRangeKey"
			type = "N"
		}
		local_secondary_index {
			name = "TestTableLSI"
			range_key = "TestLSIRangeKey"
			projection_type = "ALL"
		}
		global_secondary_index {
			name = "ReplacementTestTableGSI"
			hash_key = "TestTableHashKey"
			range_key = "ReplacementGSIRangeKey"
			write_capacity = 5
			read_capacity = 5
			projection_type = "ALL"
		}
}
`
<commit_msg>simple acc test covering more projection types<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSDynamoDbTable(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSDynamoDbTableDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSDynamoDbConfigInitialState,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckInitialAWSDynamoDbTableExists("aws_dynamodb_table.basic-dynamodb-table"),
				),
			},
			resource.TestStep{
				Config: testAccAWSDynamoDbConfigAddSecondaryGSI,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckDynamoDbTableWasUpdated("aws_dynamodb_table.basic-dynamodb-table"),
				),
			},
		},
	})
}

func testAccCheckAWSDynamoDbTableDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).dynamodbconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_dynamodb_table" {
			continue
		}

		fmt.Printf("[DEBUG] Checking if DynamoDB table %s exists", rs.Primary.ID)
		// Check if queue exists by checking for its attributes
		params := &dynamodb.DescribeTableInput{
			TableName: aws.String(rs.Primary.ID),
		}
		_, err := conn.DescribeTable(params)
		if err == nil {
			return fmt.Errorf("DynamoDB table %s still exists. Failing!", rs.Primary.ID)
		}

		// Verify the error is what we want
		_, ok := err.(awserr.Error)
		if !ok {
			return err
		}
	}

	return nil
}

func testAccCheckInitialAWSDynamoDbTableExists(n string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		fmt.Printf("[DEBUG] Trying to create initial table state!")
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No DynamoDB table name specified!")
		}

		conn := testAccProvider.Meta().(*AWSClient).dynamodbconn

		params := &dynamodb.DescribeTableInput{
			TableName: aws.String(rs.Primary.ID),
		}

		resp, err := conn.DescribeTable(params)

		if err != nil {
			fmt.Printf("[ERROR] Problem describing table '%s': %s", rs.Primary.ID, err)
			return err
		}

		table := resp.Table

		fmt.Printf("[DEBUG] Checking on table %s", rs.Primary.ID)

		if *table.ProvisionedThroughput.WriteCapacityUnits != 20 {
			return fmt.Errorf("Provisioned write capacity was %d, not 20!", table.ProvisionedThroughput.WriteCapacityUnits)
		}

		if *table.ProvisionedThroughput.ReadCapacityUnits != 10 {
			return fmt.Errorf("Provisioned read capacity was %d, not 10!", table.ProvisionedThroughput.ReadCapacityUnits)
		}

		attrCount := len(table.AttributeDefinitions)
		gsiCount := len(table.GlobalSecondaryIndexes)
		lsiCount := len(table.LocalSecondaryIndexes)

		if attrCount != 4 {
			return fmt.Errorf("There were %d attributes, not 4 like there should have been!", attrCount)
		}

		if gsiCount != 1 {
			return fmt.Errorf("There were %d GSIs, not 1 like there should have been!", gsiCount)
		}

		if lsiCount != 1 {
			return fmt.Errorf("There were %d LSIs, not 1 like there should have been!", lsiCount)
		}

		attrmap := dynamoDbAttributesToMap(&table.AttributeDefinitions)
		if attrmap["TestTableHashKey"] != "S" {
			return fmt.Errorf("Test table hash key was of type %s instead of S!", attrmap["TestTableHashKey"])
		}
		if attrmap["TestTableRangeKey"] != "S" {
			return fmt.Errorf("Test table range key was of type %s instead of S!", attrmap["TestTableRangeKey"])
		}
		if attrmap["TestLSIRangeKey"] != "N" {
			return fmt.Errorf("Test table LSI range key was of type %s instead of N!", attrmap["TestLSIRangeKey"])
		}
		if attrmap["TestGSIRangeKey"] != "S" {
			return fmt.Errorf("Test table GSI range key was of type %s instead of S!", attrmap["TestGSIRangeKey"])
		}

		return nil
	}
}

func testAccCheckDynamoDbTableWasUpdated(n string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No DynamoDB table name specified!")
		}

		conn := testAccProvider.Meta().(*AWSClient).dynamodbconn

		params := &dynamodb.DescribeTableInput{
			TableName: aws.String(rs.Primary.ID),
		}
		resp, err := conn.DescribeTable(params)
		table := resp.Table

		if err != nil {
			return err
		}

		attrCount := len(table.AttributeDefinitions)
		gsiCount := len(table.GlobalSecondaryIndexes)
		lsiCount := len(table.LocalSecondaryIndexes)

		if attrCount != 4 {
			return fmt.Errorf("There were %d attributes, not 4 like there should have been!", attrCount)
		}

		if gsiCount != 1 {
			return fmt.Errorf("There were %d GSIs, not 1 like there should have been!", gsiCount)
		}

		if lsiCount != 1 {
			return fmt.Errorf("There were %d LSIs, not 1 like there should have been!", lsiCount)
		}

		if dynamoDbGetGSIIndex(&table.GlobalSecondaryIndexes, "ReplacementTestTableGSI") == -1 {
			return fmt.Errorf("Could not find GSI named 'ReplacementTestTableGSI' in the table!")
		}

		if dynamoDbGetGSIIndex(&table.GlobalSecondaryIndexes, "InitialTestTableGSI") != -1 {
			return fmt.Errorf("Should have removed 'InitialTestTableGSI' but it still exists!")
		}

		attrmap := dynamoDbAttributesToMap(&table.AttributeDefinitions)
		if attrmap["TestTableHashKey"] != "S" {
			return fmt.Errorf("Test table hash key was of type %s instead of S!", attrmap["TestTableHashKey"])
		}
		if attrmap["TestTableRangeKey"] != "S" {
			return fmt.Errorf("Test table range key was of type %s instead of S!", attrmap["TestTableRangeKey"])
		}
		if attrmap["TestLSIRangeKey"] != "N" {
			return fmt.Errorf("Test table LSI range key was of type %s instead of N!", attrmap["TestLSIRangeKey"])
		}
		if attrmap["ReplacementGSIRangeKey"] != "N" {
			return fmt.Errorf("Test table replacement GSI range key was of type %s instead of N!", attrmap["ReplacementGSIRangeKey"])
		}

		return nil
	}
}

func dynamoDbGetGSIIndex(gsiList *[]*dynamodb.GlobalSecondaryIndexDescription, target string) int {
	for idx, gsiObject := range *gsiList {
		if *gsiObject.IndexName == target {
			return idx
		}
	}

	return -1
}

func dynamoDbAttributesToMap(attributes *[]*dynamodb.AttributeDefinition) map[string]string {
	attrmap := make(map[string]string)

	for _, attrdef := range *attributes {
		attrmap[*(attrdef.AttributeName)] = *(attrdef.AttributeType)
	}

	return attrmap
}

const testAccAWSDynamoDbConfigInitialState = `
resource "aws_dynamodb_table" "basic-dynamodb-table" {
    name = "TerraformTestTable"
		read_capacity = 10
		write_capacity = 20
		hash_key = "TestTableHashKey"
		range_key = "TestTableRangeKey"
		attribute {
			name = "TestTableHashKey"
			type = "S"
		}
		attribute {
			name = "TestTableRangeKey"
			type = "S"
		}
		attribute {
			name = "TestLSIRangeKey"
			type = "N"
		}
		attribute {
			name = "TestGSIRangeKey"
			type = "S"
		}
		local_secondary_index {
			name = "TestTableLSI"
			range_key = "TestLSIRangeKey"
			projection_type = "ALL"
		}
		global_secondary_index {
			name = "InitialTestTableGSI"
			hash_key = "TestTableHashKey"
			range_key = "TestGSIRangeKey"
			write_capacity = 10
			read_capacity = 10
			projection_type = "ALL"
		}
}
`

const testAccAWSDynamoDbConfigAddSecondaryGSI = `
resource "aws_dynamodb_table" "basic-dynamodb-table" {
    name = "TerraformTestTable"
		read_capacity = 20
		write_capacity = 20
		hash_key = "TestTableHashKey"
		range_key = "TestTableRangeKey"
		attribute {
			name = "TestTableHashKey"
			type = "S"
		}
		attribute {
			name = "TestTableRangeKey"
			type = "S"
		}
		attribute {
			name = "TestLSIRangeKey"
			type = "N"
		}
		attribute {
			name = "ReplacementGSIRangeKey"
			type = "N"
		}
		attribute {
			name = "TestNonKeyAttribute"
			type = "S"
		}
		local_secondary_index {
			name = "TestTableLSI"
			range_key = "TestLSIRangeKey"
			projection_type = "ALL"
		}
		global_secondary_index {
			name = "ReplacementTestTableGSI"
			hash_key = "TestTableHashKey"
			range_key = "ReplacementGSIRangeKey"
			write_capacity = 5
			read_capacity = 5
			projection_type = "INCLUDE"
			non_key_attributes = ["TestNonKeyAttribute"]
		}
}
`
<|endoftext|>


<commit_before>package version

import (
	"fmt"

	"go.pedge.io/proto/version"
)

const (
	// MajorVersion is the current major version for pachyderm.
	MajorVersion = 1
	// MinorVersion is the current minor version for pachyderm.
	MinorVersion = 0
	// MicroVersion is the patch number for pachyderm.
	MicroVersion = 1
)

var (
	// AdditionalVersion is the string provided at release time
	// The value is passed to the linker at build time
	AdditionalVersion string
	// Version is the current version for pachyderm.
	Version = &protoversion.Version{
		Major:      MajorVersion,
		Minor:      MinorVersion,
		Micro:      MicroVersion,
		Additional: AdditionalVersion,
	}
)

func PrettyPrintVersion(version *protoversion.Version) string {
	result := fmt.Sprintf("%d.%d.%d", version.Major, version.Minor, version.Micro)
	if version.Additional != "" {
		result += fmt.Sprintf("-%s", version.Additional)
	}
	return result
}
<commit_msg>Resolve merge conflict<commit_after>package version

import (
	"fmt"

	"go.pedge.io/proto/version"
)

const (
	// MajorVersion is the current major version for pachyderm.
	MajorVersion = 1
	// MinorVersion is the current minor version for pachyderm.
	MinorVersion = 0
	// MicroVersion is the patch number for pachyderm.
	MicroVersion = 1
)

var (
	// AdditionalVersion is the string provided at release time
	// The value is passed to the linker at build time
	AdditionalVersion string
	// Version is the current version for pachyderm.
	Version = &protoversion.Version{
		Major:      MajorVersion,
		Minor:      MinorVersion,
		Micro:      MicroVersion,
		Additional: AdditionalVersion,
	}
	// AdditionalVersion is the string provided at release time
	AdditionalVersion string
)

func PrettyPrintVersion(version *protoversion.Version) string {
	result := fmt.Sprintf("%d.%d.%d", version.Major, version.Minor, version.Micro)
	if version.Additional != "" {
		result += fmt.Sprintf("-%s", version.Additional)
	}
	return result
}
<<<<<<< HEAD
||||||| merged common ancestors

func getAdditionalVersion() string {
	value := os.Getenv("VERSION_ADDITIONAL")
	if value == "" {
		out, err := exec.Command("git", "log", "--pretty=format:%H").Output()
		if err != nil {
			panic(err)
		}
		lines := strings.SplitAfterN(string(out), "\n", 2)
		if len(lines) < 2 {
			panic("Couldn't determine current commit hash")
		}
		value = strings.TrimSpace(lines[0])
	}
	return value
}
=======

func getAdditionalVersion() string {
	value := AdditionalVersion
	if value == "" {
		_, err := exec.LookPath("git")
		if err != nil {
			return "dirty"
		}
		out, err := exec.Command("git", "log", "--pretty=format:%H").Output()
		if err != nil {
			panic(err)
		}
		lines := strings.SplitAfterN(string(out), "\n", 2)
		if len(lines) < 2 {
			panic("Couldn't determine current commit hash")
		}
		value = strings.TrimSpace(lines[0])
	}
	return value
}
>>>>>>> 800dd5ec21d1eab30ab9dcb101744e46c1887ae5
<|endoftext|>


<commit_before>package aws

import (
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsEcsTaskDefinition() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsEcsTaskDefinitionCreate,
		Read:   resourceAwsEcsTaskDefinitionRead,
		Update: resourceAwsEcsTaskDefinitionUpdate,
		Delete: resourceAwsEcsTaskDefinitionDelete,

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"family": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"revision": &schema.Schema{
				Type:     schema.TypeInt,
				Computed: true,
			},

			"container_definitions": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				StateFunc: func(v interface{}) string {
					hash := sha1.Sum([]byte(v.(string)))
					return hex.EncodeToString(hash[:])
				},
			},

			"volume": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"host_path": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
				Set: resourceAwsEcsTaskDefinitionVolumeHash,
			},
		},
	}
}

func resourceAwsEcsTaskDefinitionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	rawDefinitions := d.Get("container_definitions").(string)
	definitions, err := expandEcsContainerDefinitions(rawDefinitions)
	if err != nil {
		return err
	}

	input := ecs.RegisterTaskDefinitionInput{
		ContainerDefinitions: definitions,
		Family:               aws.String(d.Get("family").(string)),
	}

	if v, ok := d.GetOk("volume"); ok {
		volumes, err := expandEcsVolumes(v.(*schema.Set).List())
		if err != nil {
			return err
		}
		input.Volumes = volumes
	}

	log.Printf("[DEBUG] Registering ECS task definition: %s", input)
	out, err := conn.RegisterTaskDefinition(&input)
	if err != nil {
		return err
	}

	taskDefinition := *out.TaskDefinition

	log.Printf("[DEBUG] ECS task definition registered: %q (rev. %d)",
		*taskDefinition.TaskDefinitionArn, *taskDefinition.Revision)

	d.SetId(*taskDefinition.Family)
	d.Set("arn", *taskDefinition.TaskDefinitionArn)

	return resourceAwsEcsTaskDefinitionRead(d, meta)
}

func resourceAwsEcsTaskDefinitionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	log.Printf("[DEBUG] Reading task definition %s", d.Id())
	out, err := conn.DescribeTaskDefinition(&ecs.DescribeTaskDefinitionInput{
		TaskDefinition: aws.String(d.Get("arn").(string)),
	})
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Received task definition %s", out)

	taskDefinition := out.TaskDefinition

	d.SetId(*taskDefinition.Family)
	d.Set("arn", *taskDefinition.TaskDefinitionArn)
	d.Set("family", *taskDefinition.Family)
	d.Set("revision", *taskDefinition.Revision)
	d.Set("container_definitions", taskDefinition.ContainerDefinitions)
	d.Set("volumes", flattenEcsVolumes(taskDefinition.Volumes))

	return nil
}

func resourceAwsEcsTaskDefinitionUpdate(d *schema.ResourceData, meta interface{}) error {
	oldArn := d.Get("arn").(string)

	log.Printf("[DEBUG] Creating new revision of task definition %q", d.Id())
	err := resourceAwsEcsTaskDefinitionCreate(d, meta)
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] New revision of %q created: %q", d.Id(), d.Get("arn").(string))

	log.Printf("[DEBUG] Deregistering old revision of task definition %q: %q", d.Id(), oldArn)
	conn := meta.(*AWSClient).ecsconn
	_, err = conn.DeregisterTaskDefinition(&ecs.DeregisterTaskDefinitionInput{
		TaskDefinition: aws.String(oldArn),
	})
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Old revision of task definition deregistered: %q", oldArn)

	return resourceAwsEcsTaskDefinitionRead(d, meta)
}

func resourceAwsEcsTaskDefinitionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	_, err := conn.DeregisterTaskDefinition(&ecs.DeregisterTaskDefinitionInput{
		TaskDefinition: aws.String(d.Get("arn").(string)),
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Task definition %q deregistered.", d.Get("arn").(string))

	return nil
}

func resourceAwsEcsTaskDefinitionVolumeHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["name"].(string)))
	buf.WriteString(fmt.Sprintf("%s-", m["host_path"].(string)))

	return hashcode.String(buf.String())
}
<commit_msg>provider/aws: Make all fields in ecs_task_definition ForceNew  - fixes https://github.com/hashicorp/terraform/issues/2694<commit_after>package aws

import (
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsEcsTaskDefinition() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsEcsTaskDefinitionCreate,
		Read:   resourceAwsEcsTaskDefinitionRead,
		Delete: resourceAwsEcsTaskDefinitionDelete,

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},

			"family": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"revision": &schema.Schema{
				Type:     schema.TypeInt,
				Computed: true,
			},

			"container_definitions": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
				StateFunc: func(v interface{}) string {
					hash := sha1.Sum([]byte(v.(string)))
					return hex.EncodeToString(hash[:])
				},
			},

			"volume": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"host_path": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
				Set: resourceAwsEcsTaskDefinitionVolumeHash,
			},
		},
	}
}

func resourceAwsEcsTaskDefinitionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	rawDefinitions := d.Get("container_definitions").(string)
	definitions, err := expandEcsContainerDefinitions(rawDefinitions)
	if err != nil {
		return err
	}

	input := ecs.RegisterTaskDefinitionInput{
		ContainerDefinitions: definitions,
		Family:               aws.String(d.Get("family").(string)),
	}

	if v, ok := d.GetOk("volume"); ok {
		volumes, err := expandEcsVolumes(v.(*schema.Set).List())
		if err != nil {
			return err
		}
		input.Volumes = volumes
	}

	log.Printf("[DEBUG] Registering ECS task definition: %s", input)
	out, err := conn.RegisterTaskDefinition(&input)
	if err != nil {
		return err
	}

	taskDefinition := *out.TaskDefinition

	log.Printf("[DEBUG] ECS task definition registered: %q (rev. %d)",
		*taskDefinition.TaskDefinitionArn, *taskDefinition.Revision)

	d.SetId(*taskDefinition.Family)
	d.Set("arn", *taskDefinition.TaskDefinitionArn)

	return resourceAwsEcsTaskDefinitionRead(d, meta)
}

func resourceAwsEcsTaskDefinitionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	log.Printf("[DEBUG] Reading task definition %s", d.Id())
	out, err := conn.DescribeTaskDefinition(&ecs.DescribeTaskDefinitionInput{
		TaskDefinition: aws.String(d.Get("arn").(string)),
	})
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] Received task definition %s", out)

	taskDefinition := out.TaskDefinition

	d.SetId(*taskDefinition.Family)
	d.Set("arn", *taskDefinition.TaskDefinitionArn)
	d.Set("family", *taskDefinition.Family)
	d.Set("revision", *taskDefinition.Revision)
	d.Set("container_definitions", taskDefinition.ContainerDefinitions)
	d.Set("volumes", flattenEcsVolumes(taskDefinition.Volumes))

	return nil
}

func resourceAwsEcsTaskDefinitionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	_, err := conn.DeregisterTaskDefinition(&ecs.DeregisterTaskDefinitionInput{
		TaskDefinition: aws.String(d.Get("arn").(string)),
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Task definition %q deregistered.", d.Get("arn").(string))

	return nil
}

func resourceAwsEcsTaskDefinitionVolumeHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["name"].(string)))
	buf.WriteString(fmt.Sprintf("%s-", m["host_path"].(string)))

	return hashcode.String(buf.String())
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"strings"
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/route53"
)

func TestAccRoute53Record(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckRoute53RecordDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccRoute53RecordConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckRoute53RecordExists("aws_route53_record.default"),
				),
			},
		},
	})
}

func testAccCheckRoute53RecordDestroy(s *terraform.State) error {
	conn := testAccProvider.route53
	for _, rs := range s.Resources {
		if rs.Type != "aws_route53_record" {
			continue
		}

		parts := strings.Split(rs.ID, "_")
		zone := parts[0]
		name := parts[1]
		rType := parts[2]

		lopts := &route53.ListOpts{Name: name, Type: rType}
		resp, err := conn.ListResourceRecordSets(zone, lopts)
		if err != nil {
			return err
		}
		if len(resp.Records) == 0 {
			return nil
		}
		rec := resp.Records[0]
		if route53.FQDN(rec.Name) == route53.FQDN(name) && rec.Type == rType {
			return fmt.Errorf("Record still exists: %#v", rec)
		}
	}
	return nil
}

func testAccCheckRoute53RecordExists(n string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.route53
		rs, ok := s.Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.ID == "" {
			return fmt.Errorf("No hosted zone ID is set")
		}

		parts := strings.Split(rs.ID, "_")
		zone := parts[0]
		name := parts[1]
		rType := parts[2]

		lopts := &route53.ListOpts{Name: name, Type: rType}
		resp, err := conn.ListResourceRecordSets(zone, lopts)
		if err != nil {
			return err
		}
		if len(resp.Records) == 0 {
			return fmt.Errorf("Record does not exist")
		}
		rec := resp.Records[0]
		if route53.FQDN(rec.Name) == route53.FQDN(name) && rec.Type == rType {
			return nil
		}
		return fmt.Errorf("Record does not exist: %#v", rec)
	}
}

const testAccRoute53RecordConfig = `
resource "aws_route53_zone" "main" {
	name = "hashicorp.com"
}

resource "aws_route53_record" "default" {
	zone_id = "${aws_route53_zone.main.zone_id}"
	name = "www.hashicorp.com"
	type = "A"
	ttl = "30"
	records = ["127.0.0.1"]
}
`
<commit_msg>providers/aws: use example.com for tests<commit_after>package aws

import (
	"fmt"
	"strings"
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/route53"
)

func TestAccRoute53Record(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckRoute53RecordDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccRoute53RecordConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckRoute53RecordExists("aws_route53_record.default"),
				),
			},
		},
	})
}

func testAccCheckRoute53RecordDestroy(s *terraform.State) error {
	conn := testAccProvider.route53
	for _, rs := range s.Resources {
		if rs.Type != "aws_route53_record" {
			continue
		}

		parts := strings.Split(rs.ID, "_")
		zone := parts[0]
		name := parts[1]
		rType := parts[2]

		lopts := &route53.ListOpts{Name: name, Type: rType}
		resp, err := conn.ListResourceRecordSets(zone, lopts)
		if err != nil {
			return err
		}
		if len(resp.Records) == 0 {
			return nil
		}
		rec := resp.Records[0]
		if route53.FQDN(rec.Name) == route53.FQDN(name) && rec.Type == rType {
			return fmt.Errorf("Record still exists: %#v", rec)
		}
	}
	return nil
}

func testAccCheckRoute53RecordExists(n string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.route53
		rs, ok := s.Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.ID == "" {
			return fmt.Errorf("No hosted zone ID is set")
		}

		parts := strings.Split(rs.ID, "_")
		zone := parts[0]
		name := parts[1]
		rType := parts[2]

		lopts := &route53.ListOpts{Name: name, Type: rType}
		resp, err := conn.ListResourceRecordSets(zone, lopts)
		if err != nil {
			return err
		}
		if len(resp.Records) == 0 {
			return fmt.Errorf("Record does not exist")
		}
		rec := resp.Records[0]
		if route53.FQDN(rec.Name) == route53.FQDN(name) && rec.Type == rType {
			return nil
		}
		return fmt.Errorf("Record does not exist: %#v", rec)
	}
}

const testAccRoute53RecordConfig = `
resource "aws_route53_zone" "main" {
	name = "example.com"
}

resource "aws_route53_record" "default" {
	zone_id = "${aws_route53_zone.main.zone_id}"
	name = "www.example.com"
	type = "A"
	ttl = "30"
	records = ["127.0.0.1"]
}
`
<|endoftext|>


<commit_before>package rpc

import (
	"os"
	"testing"

	"github.com/thomasylee/GoRaft/global"
	"github.com/thomasylee/GoRaft/state"
)

const port string = "8000"

func TestMain(m *testing.M) {
	global.SetUpLogger()
	global.SetLogLevel("debug")

	global.TimeoutChannel = make(chan bool, 1)

	go RunServer(port)
	global.Log.Debug("gRPC server for tests has been started.")

	os.Exit(m.Run())
}

func resetTestEnvironment() {
	close(global.TimeoutChannel)
	global.TimeoutChannel = make(chan bool, 1)

	state.Node = state.NewNodeState(
		state.NewMemoryDataStore(),
		state.NewMemoryDataStore())
}

func Test_AppendEntries_WhenRequestHasNoEntries_ReturnsSuccessTrue(t *testing.T) {
	resetTestEnvironment()

	request := &AppendEntriesRequest{
		Term:         0,
		LeaderId:     "123",
		PrevLogIndex: 0,
		PrevLogTerm:  0,
		// Heartbeat request since entries is empty.
		Entries:      []*AppendEntriesRequest_Entry{},
		LeaderCommit: 0,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if !response.Success {
		t.Error("Success was false")
	}
}

func Test_AppendEntries_WhenTermIsOlderThanCurrentTerm_ReturnsSuccessFalse(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetCurrentTerm(1)

	request := &AppendEntriesRequest{
		// Request should fail since CurrentTerm is 1, which is > 0.
		Term:         0,
		LeaderId:     "123",
		PrevLogIndex: 0,
		PrevLogTerm:  0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "a",
				Value: "A",
			},
		},
		LeaderCommit: 0,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if response.Success {
		t.Error("Success was true")
	}

	if response.Term != 1 {
		t.Error("Term was not 1:", response.Term)
	}
}

func Test_AppendEntries_WhenLogAtPrevLogIndexDoesNotMatch_ReturnsSuccessFalse(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetCurrentTerm(1)

	entry := state.LogEntry{Term: 1, Key: "a", Value: "A"}
	state.Node.SetLogEntry(1, entry)

	request := &AppendEntriesRequest{
		Term:         1,
		LeaderId:     "123",
		PrevLogIndex: 1,
		// Request should fail since term for entry 1 is 1, not 0.
		PrevLogTerm: 0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "b",
				Value: "B",
			},
		},
		LeaderCommit: 0,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if response.Success {
		t.Error("Success was true")
	}

	if response.Term != 1 {
		t.Error("Term was not 1:", response.Term)
	}
}

func Test_AppendEntries_WhenLeaderCommitIsGreaterThanCommitIndex_IncreasesCommitIndex(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetCurrentTerm(0)

	entry := state.LogEntry{Term: 0, Key: "a", Value: "A"}
	state.Node.SetLogEntry(1, entry)

	request := &AppendEntriesRequest{
		Term:         1,
		LeaderId:     "123",
		PrevLogIndex: 1,
		// Request should fail since term for entry 1 is 1, not 0.
		PrevLogTerm: 0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "b",
				Value: "B",
			},
		},
		LeaderCommit: 2,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if !response.Success {
		t.Error("Success was true")
	}

	if response.Term != 1 {
		t.Error("Term was not 1:", response.Term)
	}

	if state.Node.CommitIndex != 2 {
		t.Error("CommitIndex was not 2:", state.Node.CommitIndex)
	}
}

func Test_RequestVote(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetVotedFor("1")

	request := &AppendEntriesRequest{
		Term:         1,
		LeaderId:     "123",
		PrevLogIndex: 0,
		PrevLogTerm:  0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "a",
				Value: "A",
			},
		},
		LeaderCommit: 0,
	}

	_, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	var tests = []struct {
		request  RequestVoteRequest
		response RequestVoteResponse
	}{
		{
			// Term is too old.
			RequestVoteRequest{
				Term:         0,
				CandidateId:  "1",
				LastLogIndex: 1,
				LastLogTerm:  0,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// Voted for a different candidate.
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "2",
				LastLogIndex: 1,
				LastLogTerm:  1,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// The log is out of date.
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "1",
				LastLogIndex: 0,
				LastLogTerm:  0,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// The log is in conflict with this node's log.
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "1",
				LastLogIndex: 1,
				LastLogTerm:  2,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// Everything is good!
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "1",
				LastLogIndex: 1,
				LastLogTerm:  1,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: true,
			},
		},
	}

	for _, test := range tests {
		response, err := SendRequestVote("127.0.0.1:"+port, &test.request)
		if err != nil {
			t.Error(err)
			continue
		}

		if *response != test.response {
			t.Errorf("RequestVote test case failed %v. Expected %v but found %v",
				test.request, test.response, *response)
		}
	}
}
<commit_msg>Give test server time to start before testing<commit_after>package rpc

import (
	"os"
	"testing"
	"time"

	"github.com/thomasylee/GoRaft/global"
	"github.com/thomasylee/GoRaft/state"
)

const port string = "8000"

func TestMain(m *testing.M) {
	global.SetUpLogger()
	global.SetLogLevel("debug")

	global.TimeoutChannel = make(chan bool, 1)

	go RunServer(port)
	// Give the server a few seconds to start.
	time.Sleep(3 * time.Second)
	global.Log.Debug("gRPC server for tests has been started.")

	os.Exit(m.Run())
}

func resetTestEnvironment() {
	close(global.TimeoutChannel)
	global.TimeoutChannel = make(chan bool, 1)

	state.Node = state.NewNodeState(
		state.NewMemoryDataStore(),
		state.NewMemoryDataStore())
}

func Test_AppendEntries_WhenRequestHasNoEntries_ReturnsSuccessTrue(t *testing.T) {
	resetTestEnvironment()

	request := &AppendEntriesRequest{
		Term:         0,
		LeaderId:     "123",
		PrevLogIndex: 0,
		PrevLogTerm:  0,
		// Heartbeat request since entries is empty.
		Entries:      []*AppendEntriesRequest_Entry{},
		LeaderCommit: 0,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if !response.Success {
		t.Error("Success was false")
	}
}

func Test_AppendEntries_WhenTermIsOlderThanCurrentTerm_ReturnsSuccessFalse(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetCurrentTerm(1)

	request := &AppendEntriesRequest{
		// Request should fail since CurrentTerm is 1, which is > 0.
		Term:         0,
		LeaderId:     "123",
		PrevLogIndex: 0,
		PrevLogTerm:  0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "a",
				Value: "A",
			},
		},
		LeaderCommit: 0,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if response.Success {
		t.Error("Success was true")
	}

	if response.Term != 1 {
		t.Error("Term was not 1:", response.Term)
	}
}

func Test_AppendEntries_WhenLogAtPrevLogIndexDoesNotMatch_ReturnsSuccessFalse(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetCurrentTerm(1)

	entry := state.LogEntry{Term: 1, Key: "a", Value: "A"}
	state.Node.SetLogEntry(1, entry)

	request := &AppendEntriesRequest{
		Term:         1,
		LeaderId:     "123",
		PrevLogIndex: 1,
		// Request should fail since term for entry 1 is 1, not 0.
		PrevLogTerm: 0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "b",
				Value: "B",
			},
		},
		LeaderCommit: 0,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if response.Success {
		t.Error("Success was true")
	}

	if response.Term != 1 {
		t.Error("Term was not 1:", response.Term)
	}
}

func Test_AppendEntries_WhenLeaderCommitIsGreaterThanCommitIndex_IncreasesCommitIndex(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetCurrentTerm(0)

	entry := state.LogEntry{Term: 0, Key: "a", Value: "A"}
	state.Node.SetLogEntry(1, entry)

	request := &AppendEntriesRequest{
		Term:         1,
		LeaderId:     "123",
		PrevLogIndex: 1,
		// Request should fail since term for entry 1 is 1, not 0.
		PrevLogTerm: 0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "b",
				Value: "B",
			},
		},
		LeaderCommit: 2,
	}

	response, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	if !response.Success {
		t.Error("Success was true")
	}

	if response.Term != 1 {
		t.Error("Term was not 1:", response.Term)
	}

	if state.Node.CommitIndex != 2 {
		t.Error("CommitIndex was not 2:", state.Node.CommitIndex)
	}
}

func Test_RequestVote(t *testing.T) {
	resetTestEnvironment()

	state.Node.SetVotedFor("1")

	request := &AppendEntriesRequest{
		Term:         1,
		LeaderId:     "123",
		PrevLogIndex: 0,
		PrevLogTerm:  0,
		Entries: []*AppendEntriesRequest_Entry{
			&AppendEntriesRequest_Entry{
				Key:   "a",
				Value: "A",
			},
		},
		LeaderCommit: 0,
	}

	_, err := SendAppendEntries("127.0.0.1:"+port, request)
	if err != nil {
		t.Fatal(err)
	}

	var tests = []struct {
		request  RequestVoteRequest
		response RequestVoteResponse
	}{
		{
			// Term is too old.
			RequestVoteRequest{
				Term:         0,
				CandidateId:  "1",
				LastLogIndex: 1,
				LastLogTerm:  0,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// Voted for a different candidate.
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "2",
				LastLogIndex: 1,
				LastLogTerm:  1,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// The log is out of date.
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "1",
				LastLogIndex: 0,
				LastLogTerm:  0,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// The log is in conflict with this node's log.
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "1",
				LastLogIndex: 1,
				LastLogTerm:  2,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: false,
			},
		},
		{
			// Everything is good!
			RequestVoteRequest{
				Term:         1,
				CandidateId:  "1",
				LastLogIndex: 1,
				LastLogTerm:  1,
			},
			RequestVoteResponse{
				Term:        1,
				VoteGranted: true,
			},
		},
	}

	for _, test := range tests {
		response, err := SendRequestVote("127.0.0.1:"+port, &test.request)
		if err != nil {
			t.Error(err)
			continue
		}

		if *response != test.response {
			t.Errorf("RequestVote test case failed %v. Expected %v but found %v",
				test.request, test.response, *response)
		}
	}
}
<|endoftext|>


<commit_before>package rpcfs

import (
	"bytes"
	"io/ioutil"
	"log"
	"net"
	"os"
	"rpc"
	"sort"
)

type Master struct {
	cache             *DiskFileCache
	fileServer        *FsServer
	fileServerAddress string
	secret            []byte
	workServers       []*rpc.Client
	masterRun         *LocalMaster
}

func NewMaster(cacheDir string, workers []string, secret []byte, excluded []string) *Master {
	c := NewDiskFileCache(cacheDir)
	me := &Master{
		cache:      c,
		fileServer: NewFsServer("/", c, excluded),
		secret:     secret,
	}
	me.masterRun = &LocalMaster{me}
	me.setupWorkers(workers)
	me.secret = secret
	return me
}

func (me *Master) Start(myAddress, mySocket string) {
	me.fileServerAddress = myAddress
	go me.startServer(me.fileServer, myAddress)
	me.startLocalServer(mySocket)
}

func (me *Master) startLocalServer(sock string) {
	listener, err := net.Listen("unix", sock)
	if err != nil {
		log.Fatal("startLocalServer", err)
	}
	err = os.Chmod(sock, 0700)
	if err != nil {
		log.Fatal("sock chmod", err)
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		rpcServer := rpc.NewServer()
		err = rpcServer.Register(me.masterRun)
		if err != nil {
			log.Fatal("could not register server", err)
		}
		go rpcServer.ServeConn(conn)
	}
}

func (me *Master) setupWorkers(addresses []string) {
	var out []*rpc.Client
	for _, addr := range addresses {
		conn, err := SetupClient(addr, me.secret)
		if err != nil {
			log.Println("Failed setting up connection with: ", addr, err)
			continue
		}
		out = append(out, rpc.NewClient(conn))
	}
	if len(out) == 0 {
		log.Fatal("No workers available.")
	}
	me.workServers = out
}

// StartServer starts the connection listener.  Should be invoked in a coroutine.
func (me *Master) startServer(server interface{}, addr string) {
	out := make(chan net.Conn)
	go SetupServer(addr, me.secret, out)
	for {
		conn := <-out
		rpcServer := rpc.NewServer()
		err := rpcServer.Register(server)
		if err != nil {
			log.Fatal("could not register server", err)
		}
		log.Println("Server started...")
		go rpcServer.ServeConn(conn)
	}
}

func (me *Master) run(req *WorkRequest, rep *WorkReply) os.Error {
	// TODO: random pick.
	worker := me.workServers[0]

	req.FileServer = me.fileServerAddress
	err := worker.Call("WorkerDaemon.Run", &req, &rep)
	if err != nil {
		return err
	}
	me.replayFileModifications(worker, rep.Files)
	rep.Files = nil
	return err
}

func (me *Master) replayFileModifications(worker *rpc.Client, infos []FileInfo) {
	entries := make(map[string]*FileInfo)
	names := []string{}
	for i, info := range infos {
		names = append(names, info.Path)
		entries[info.Path] = &infos[i]
	}

	// Sort so we get parents before children.
	sort.SortStrings(names)
	for _, name := range names {
		info := entries[name]
		var err os.Error

		if info.FileInfo.IsDirectory() {
			if name == "" {
				name = "/"
			}
			_, err = os.Lstat(name)
			if err != nil {
				log.Println("Replay mkdir:", name)
				err = os.Mkdir(name, info.FileInfo.Mode&07777)
			}
		}
		if info.Hash != nil {
			log.Printf("Replay file content %s %x", name, info.Hash)
			c := FetchFromContentServer(
				worker, "WorkerDaemon.FileContent", info.FileInfo.Size, info.Hash)
			hash := me.cache.Save(c)
			if bytes.Compare(info.Hash, hash) != 0 {
				log.Fatal("Hash mismatch.")
			}
			err = ioutil.WriteFile(info.Path, c, info.FileInfo.Mode&07777)
		}
		if info.LinkContent != "" {
			log.Println("Replay symlink:", name)
			err = os.Symlink(info.LinkContent, info.Path)
		}
		if info.Delete {
			log.Println("Replay delete:", name)
			err = os.Remove(info.Path)
		}
		if err != nil {
			log.Fatal("Replay error", info.Path, err)
		}
	}
}

////////////////

// Expose functionality for the local tool to use.
type LocalMaster struct {
	master *Master
}

func (me *LocalMaster) Run(req *WorkRequest, rep *WorkReply) os.Error {
	return me.master.run(req, rep)
}

func (me *LocalMaster) Quit(ignoreIn *int, ignoreOut *int) os.Error {
	// TODO - make this actually do something.
	return nil
}
<commit_msg>Pick a random server.<commit_after>package rpcfs

import (
	"bytes"
	"io/ioutil"
	"log"
	"net"
	"os"
	"rand"
	"rpc"
	"sort"
)

type Master struct {
	cache             *DiskFileCache
	fileServer        *FsServer
	fileServerAddress string
	secret            []byte
	workServers       []*rpc.Client
	masterRun         *LocalMaster
}

func NewMaster(cacheDir string, workers []string, secret []byte, excluded []string) *Master {
	c := NewDiskFileCache(cacheDir)
	me := &Master{
		cache:      c,
		fileServer: NewFsServer("/", c, excluded),
		secret:     secret,
	}
	me.masterRun = &LocalMaster{me}
	me.setupWorkers(workers)
	me.secret = secret
	return me
}

func (me *Master) Start(myAddress, mySocket string) {
	me.fileServerAddress = myAddress
	go me.startServer(me.fileServer, myAddress)
	me.startLocalServer(mySocket)
}

func (me *Master) startLocalServer(sock string) {
	listener, err := net.Listen("unix", sock)
	if err != nil {
		log.Fatal("startLocalServer", err)
	}
	err = os.Chmod(sock, 0700)
	if err != nil {
		log.Fatal("sock chmod", err)
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		rpcServer := rpc.NewServer()
		err = rpcServer.Register(me.masterRun)
		if err != nil {
			log.Fatal("could not register server", err)
		}
		go rpcServer.ServeConn(conn)
	}
}

func (me *Master) setupWorkers(addresses []string) {
	var out []*rpc.Client
	for _, addr := range addresses {
		conn, err := SetupClient(addr, me.secret)
		if err != nil {
			log.Println("Failed setting up connection with: ", addr, err)
			continue
		}
		out = append(out, rpc.NewClient(conn))
	}
	if len(out) == 0 {
		log.Fatal("No workers available.")
	}
	me.workServers = out
}

// StartServer starts the connection listener.  Should be invoked in a coroutine.
func (me *Master) startServer(server interface{}, addr string) {
	out := make(chan net.Conn)
	go SetupServer(addr, me.secret, out)
	for {
		conn := <-out
		rpcServer := rpc.NewServer()
		err := rpcServer.Register(server)
		if err != nil {
			log.Fatal("could not register server", err)
		}
		log.Println("Server started...")
		go rpcServer.ServeConn(conn)
	}
}

func (me *Master) run(req *WorkRequest, rep *WorkReply) os.Error {
	idx := rand.Intn(len(me.workServers))
	worker := me.workServers[idx]

	req.FileServer = me.fileServerAddress
	err := worker.Call("WorkerDaemon.Run", &req, &rep)
	if err != nil {
		return err
	}
	me.replayFileModifications(worker, rep.Files)
	rep.Files = nil
	return err
}

func (me *Master) replayFileModifications(worker *rpc.Client, infos []FileInfo) {
	entries := make(map[string]*FileInfo)
	names := []string{}
	for i, info := range infos {
		names = append(names, info.Path)
		entries[info.Path] = &infos[i]
	}

	// Sort so we get parents before children.
	sort.SortStrings(names)
	for _, name := range names {
		info := entries[name]
		var err os.Error

		if info.FileInfo.IsDirectory() {
			if name == "" {
				name = "/"
			}
			_, err = os.Lstat(name)
			if err != nil {
				log.Println("Replay mkdir:", name)
				err = os.Mkdir(name, info.FileInfo.Mode&07777)
			}
		}
		if info.Hash != nil {
			log.Printf("Replay file content %s %x", name, info.Hash)
			c := FetchFromContentServer(
				worker, "WorkerDaemon.FileContent", info.FileInfo.Size, info.Hash)
			hash := me.cache.Save(c)
			if bytes.Compare(info.Hash, hash) != 0 {
				log.Fatal("Hash mismatch.")
			}
			err = ioutil.WriteFile(info.Path, c, info.FileInfo.Mode&07777)
		}
		if info.LinkContent != "" {
			log.Println("Replay symlink:", name)
			err = os.Symlink(info.LinkContent, info.Path)
		}
		if info.Delete {
			log.Println("Replay delete:", name)
			err = os.Remove(info.Path)
		}
		if err != nil {
			log.Fatal("Replay error", info.Path, err)
		}
	}
}

////////////////

// Expose functionality for the local tool to use.
type LocalMaster struct {
	master *Master
}

func (me *LocalMaster) Run(req *WorkRequest, rep *WorkReply) os.Error {
	return me.master.run(req, rep)
}

func (me *LocalMaster) Quit(ignoreIn *int, ignoreOut *int) os.Error {
	// TODO - make this actually do something.
	return nil
}
<|endoftext|>


<commit_before>package provision

import (
	"bytes"
	"fmt"
	"text/template"

	"github.com/docker/machine/drivers"
	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/provision/pkgaction"
	"github.com/docker/machine/libmachine/swarm"
	"github.com/docker/machine/log"
	"github.com/docker/machine/utils"
)

const (
	// TODO: eventually the RPM install process will be integrated
	// into the get.docker.com install script; for now
	// we install via vendored RPMs
	dockerRPMPath = "https://docker-mcn.s3.amazonaws.com/public/redhat/rpms/docker-engine-1.6.1-0.0.20150511.171646.git1b47f9f.el7.centos.x86_64.rpm"
)

func init() {
	Register("RedHat", &RegisteredProvisioner{
		New: NewRedHatProvisioner,
	})
}

func NewRedHatProvisioner(d drivers.Driver) Provisioner {
	return &RedHatProvisioner{
		GenericProvisioner{
			DockerOptionsDir:  "/etc/docker",
			DaemonOptionsFile: "/lib/systemd/system/docker.service",
			OsReleaseId:       "rhel",
			Packages: []string{
				"curl",
			},
			Driver: d,
		},
	}
}

type RedHatProvisioner struct {
	GenericProvisioner
}

func (provisioner *RedHatProvisioner) Service(name string, action pkgaction.ServiceAction) error {
	reloadDaemon := false
	switch action {
	case pkgaction.Start, pkgaction.Restart:
		reloadDaemon = true
	}

	// systemd needs reloaded when config changes on disk; we cannot
	// be sure exactly when it changes from the provisioner so
	// we call a reload on every restart to be safe
	if reloadDaemon {
		if _, err := provisioner.SSHCommand("sudo systemctl daemon-reload"); err != nil {
			return err
		}
	}

	command := fmt.Sprintf("sudo systemctl %s %s", action.String(), name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) Package(name string, action pkgaction.PackageAction) error {
	var packageAction string

	switch action {
	case pkgaction.Install:
		packageAction = "install"
	case pkgaction.Remove:
		packageAction = "remove"
	case pkgaction.Upgrade:
		packageAction = "upgrade"
	}

	command := fmt.Sprintf("sudo -E yum %s -y %s", packageAction, name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) isAWS() bool {
	if _, err := provisioner.SSHCommand("curl -s http://169.254.169.254/latest/meta-data/ami-id"); err != nil {
		return false
	}

	return true
}

func (provisioner *RedHatProvisioner) configureRepos() error {
	log.Debug("configuring extra repo")
	repoCmd := "subscription-manager repos --enable=rhel-7-server-extras-rpms"
	if provisioner.isAWS() {
		repoCmd = "yum-config-manager --enable rhui-REGION-rhel-server-extras"
	}

	if _, err := provisioner.SSHCommand(fmt.Sprintf("sudo %s", repoCmd)); err != nil {
		return err
	}

	return nil
}

func installDocker(provisioner *RedHatProvisioner) error {
	if err := provisioner.installOfficialDocker(); err != nil {
		return err
	}

	if err := provisioner.Service("docker", pkgaction.Restart); err != nil {
		return err
	}

	if err := provisioner.Service("docker", pkgaction.Enable); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) installOfficialDocker() error {
	log.Debug("installing docker")

	if _, err := provisioner.SSHCommand(fmt.Sprintf("sudo yum install -y --nogpgcheck  %s", dockerRPMPath)); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) dockerDaemonResponding() bool {
	if _, err := provisioner.SSHCommand("sudo docker version"); err != nil {
		log.Warn("Error getting SSH command to check if the daemon is up: %s", err)
		return false
	}

	// The daemon is up if the command worked.  Carry on.
	return true
}

func (provisioner *RedHatProvisioner) Provision(swarmOptions swarm.SwarmOptions, authOptions auth.AuthOptions, engineOptions engine.EngineOptions) error {
	provisioner.SwarmOptions = swarmOptions
	provisioner.AuthOptions = authOptions
	provisioner.EngineOptions = engineOptions

	// set default storage driver for redhat
	if provisioner.EngineOptions.StorageDriver == "" {
		provisioner.EngineOptions.StorageDriver = "devicemapper"
	}

	if err := provisioner.SetHostname(provisioner.Driver.GetMachineName()); err != nil {
		return err
	}

	// setup extras repo
	if err := provisioner.configureRepos(); err != nil {
		return err
	}

	for _, pkg := range provisioner.Packages {
		if err := provisioner.Package(pkg, pkgaction.Install); err != nil {
			return err
		}
	}

	// install docker
	if err := installDocker(provisioner); err != nil {
		return err
	}

	if err := utils.WaitFor(provisioner.dockerDaemonResponding); err != nil {
		return err
	}

	if err := makeDockerOptionsDir(provisioner); err != nil {
		return err
	}

	provisioner.AuthOptions = setRemoteAuthOptions(provisioner)

	if err := ConfigureAuth(provisioner); err != nil {
		return err
	}

	if err := configureSwarm(provisioner, swarmOptions); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) GenerateDockerOptions(dockerPort int) (*DockerOptions, error) {
	var (
		engineCfg  bytes.Buffer
		configPath = provisioner.DaemonOptionsFile
	)

	// remove existing
	if _, err := provisioner.SSHCommand(fmt.Sprintf("sudo rm %s", configPath)); err != nil {
		return nil, err
	}

	driverNameLabel := fmt.Sprintf("provider=%s", provisioner.Driver.DriverName())
	provisioner.EngineOptions.Labels = append(provisioner.EngineOptions.Labels, driverNameLabel)

	// systemd / redhat will not load options if they are on newlines
	// instead, it just continues with a different set of options; yeah...
	engineConfigTmpl := `[Unit]
Description=Docker Application Container Engine
Documentation=http://docs.docker.com
After=network.target docker.socket
Required=docker.socket

[Service]
ExecStart=/usr/bin/docker -d -H tcp://0.0.0.0:{{.DockerPort}} -H unix:///var/run/docker.sock --storage-driver {{.EngineOptions.StorageDriver}} --tlsverify --tlscacert {{.AuthOptions.CaCertRemotePath}} --tlscert {{.AuthOptions.ServerCertRemotePath}} --tlskey {{.AuthOptions.ServerKeyRemotePath}} {{ range .EngineOptions.Labels }}--label {{.}} {{ end }}{{ range .EngineOptions.InsecureRegistry }}--insecure-registry {{.}} {{ end }}{{ range .EngineOptions.RegistryMirror }}--registry-mirror {{.}} {{ end }}{{ range .EngineOptions.ArbitraryFlags }}--{{.}} {{ end }}
MountFlags=slave
LimitNOFILE=1048576
LimitNPROC=1048576
LimitCORE=infinity

[Install]
WantedBy=multi-user.target
`
	t, err := template.New("engineConfig").Parse(engineConfigTmpl)
	if err != nil {
		return nil, err
	}

	engineConfigContext := EngineConfigContext{
		DockerPort:       dockerPort,
		AuthOptions:      provisioner.AuthOptions,
		EngineOptions:    provisioner.EngineOptions,
		DockerOptionsDir: provisioner.DockerOptionsDir,
	}

	t.Execute(&engineCfg, engineConfigContext)

	daemonOptsDir := configPath
	return &DockerOptions{
		EngineOptions:     engineCfg.String(),
		EngineOptionsPath: daemonOptsDir,
	}, nil
}
<commit_msg>redhat: remove configure repo as we use our own rpms now<commit_after>package provision

import (
	"bytes"
	"fmt"
	"text/template"

	"github.com/docker/machine/drivers"
	"github.com/docker/machine/libmachine/auth"
	"github.com/docker/machine/libmachine/engine"
	"github.com/docker/machine/libmachine/provision/pkgaction"
	"github.com/docker/machine/libmachine/swarm"
	"github.com/docker/machine/log"
	"github.com/docker/machine/utils"
)

const (
	// TODO: eventually the RPM install process will be integrated
	// into the get.docker.com install script; for now
	// we install via vendored RPMs
	dockerRPMPath = "https://docker-mcn.s3.amazonaws.com/public/redhat/rpms/docker-engine-1.6.1-0.0.20150511.171646.git1b47f9f.el7.centos.x86_64.rpm"
)

func init() {
	Register("RedHat", &RegisteredProvisioner{
		New: NewRedHatProvisioner,
	})
}

func NewRedHatProvisioner(d drivers.Driver) Provisioner {
	return &RedHatProvisioner{
		GenericProvisioner{
			DockerOptionsDir:  "/etc/docker",
			DaemonOptionsFile: "/lib/systemd/system/docker.service",
			OsReleaseId:       "rhel",
			Packages: []string{
				"curl",
			},
			Driver: d,
		},
	}
}

type RedHatProvisioner struct {
	GenericProvisioner
}

func (provisioner *RedHatProvisioner) Service(name string, action pkgaction.ServiceAction) error {
	reloadDaemon := false
	switch action {
	case pkgaction.Start, pkgaction.Restart:
		reloadDaemon = true
	}

	// systemd needs reloaded when config changes on disk; we cannot
	// be sure exactly when it changes from the provisioner so
	// we call a reload on every restart to be safe
	if reloadDaemon {
		if _, err := provisioner.SSHCommand("sudo systemctl daemon-reload"); err != nil {
			return err
		}
	}

	command := fmt.Sprintf("sudo systemctl %s %s", action.String(), name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) Package(name string, action pkgaction.PackageAction) error {
	var packageAction string

	switch action {
	case pkgaction.Install:
		packageAction = "install"
	case pkgaction.Remove:
		packageAction = "remove"
	case pkgaction.Upgrade:
		packageAction = "upgrade"
	}

	command := fmt.Sprintf("sudo -E yum %s -y %s", packageAction, name)

	if _, err := provisioner.SSHCommand(command); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) isAWS() bool {
	if _, err := provisioner.SSHCommand("curl -s http://169.254.169.254/latest/meta-data/ami-id"); err != nil {
		return false
	}

	return true
}

func installDocker(provisioner *RedHatProvisioner) error {
	if err := provisioner.installOfficialDocker(); err != nil {
		return err
	}

	if err := provisioner.Service("docker", pkgaction.Restart); err != nil {
		return err
	}

	if err := provisioner.Service("docker", pkgaction.Enable); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) installOfficialDocker() error {
	log.Debug("installing docker")

	if _, err := provisioner.SSHCommand(fmt.Sprintf("sudo yum install -y --nogpgcheck  %s", dockerRPMPath)); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) dockerDaemonResponding() bool {
	if _, err := provisioner.SSHCommand("sudo docker version"); err != nil {
		log.Warn("Error getting SSH command to check if the daemon is up: %s", err)
		return false
	}

	// The daemon is up if the command worked.  Carry on.
	return true
}

func (provisioner *RedHatProvisioner) Provision(swarmOptions swarm.SwarmOptions, authOptions auth.AuthOptions, engineOptions engine.EngineOptions) error {
	provisioner.SwarmOptions = swarmOptions
	provisioner.AuthOptions = authOptions
	provisioner.EngineOptions = engineOptions

	// set default storage driver for redhat
	if provisioner.EngineOptions.StorageDriver == "" {
		provisioner.EngineOptions.StorageDriver = "devicemapper"
	}

	if err := provisioner.SetHostname(provisioner.Driver.GetMachineName()); err != nil {
		return err
	}

	for _, pkg := range provisioner.Packages {
		if err := provisioner.Package(pkg, pkgaction.Install); err != nil {
			return err
		}
	}

	// install docker
	if err := installDocker(provisioner); err != nil {
		return err
	}

	if err := utils.WaitFor(provisioner.dockerDaemonResponding); err != nil {
		return err
	}

	if err := makeDockerOptionsDir(provisioner); err != nil {
		return err
	}

	provisioner.AuthOptions = setRemoteAuthOptions(provisioner)

	if err := ConfigureAuth(provisioner); err != nil {
		return err
	}

	if err := configureSwarm(provisioner, swarmOptions); err != nil {
		return err
	}

	return nil
}

func (provisioner *RedHatProvisioner) GenerateDockerOptions(dockerPort int) (*DockerOptions, error) {
	var (
		engineCfg  bytes.Buffer
		configPath = provisioner.DaemonOptionsFile
	)

	// remove existing
	if _, err := provisioner.SSHCommand(fmt.Sprintf("sudo rm %s", configPath)); err != nil {
		return nil, err
	}

	driverNameLabel := fmt.Sprintf("provider=%s", provisioner.Driver.DriverName())
	provisioner.EngineOptions.Labels = append(provisioner.EngineOptions.Labels, driverNameLabel)

	// systemd / redhat will not load options if they are on newlines
	// instead, it just continues with a different set of options; yeah...
	engineConfigTmpl := `[Unit]
Description=Docker Application Container Engine
Documentation=http://docs.docker.com
After=network.target docker.socket
Required=docker.socket

[Service]
ExecStart=/usr/bin/docker -d -H tcp://0.0.0.0:{{.DockerPort}} -H unix:///var/run/docker.sock --storage-driver {{.EngineOptions.StorageDriver}} --tlsverify --tlscacert {{.AuthOptions.CaCertRemotePath}} --tlscert {{.AuthOptions.ServerCertRemotePath}} --tlskey {{.AuthOptions.ServerKeyRemotePath}} {{ range .EngineOptions.Labels }}--label {{.}} {{ end }}{{ range .EngineOptions.InsecureRegistry }}--insecure-registry {{.}} {{ end }}{{ range .EngineOptions.RegistryMirror }}--registry-mirror {{.}} {{ end }}{{ range .EngineOptions.ArbitraryFlags }}--{{.}} {{ end }}
MountFlags=slave
LimitNOFILE=1048576
LimitNPROC=1048576
LimitCORE=infinity

[Install]
WantedBy=multi-user.target
`
	t, err := template.New("engineConfig").Parse(engineConfigTmpl)
	if err != nil {
		return nil, err
	}

	engineConfigContext := EngineConfigContext{
		DockerPort:       dockerPort,
		AuthOptions:      provisioner.AuthOptions,
		EngineOptions:    provisioner.EngineOptions,
		DockerOptionsDir: provisioner.DockerOptionsDir,
	}

	t.Execute(&engineCfg, engineConfigContext)

	daemonOptsDir := configPath
	return &DockerOptions{
		EngineOptions:     engineCfg.String(),
		EngineOptionsPath: daemonOptsDir,
	}, nil
}
<|endoftext|>