<commit_before>package net

import (
	"log"
	"strconv"
	"net"
	"io"
	"sync"
	"time"
	"runtime"
	"math/rand"
	"sync/atomic"
	"GoOnchain/common"
)

// The node state
const (
	INIT = 0
	HANDSHAKEING = 1
	HANDSHAKED = 2
	ESTABLISH = 3
	INACTIVITY = 4
)

// The node capability flag
const (
	RELAY  = 0x01
	SERVER = 0x02
	NODESERVICES = 0x01
)

type node struct {
	state		uint		// node status
	id		string		// The nodes's id, MAC or IP?
	addr		string 		// The address of the node
	conn		net.Conn	// Connect socket with the peer node
	nonce		uint32		// Random number to identify different entity from the same IP
	cap		uint32  	// The node capability set
	version		uint32		// The network protocol the node used
	services	uint64		// The services the node supplied
	port		uint16		// The server port of the node
	relay		bool		// The relay capability of the node (merge into capbility flag)
	handshakeRetry  uint32		// Handshake retry times
	handshakeTime	time.Time	// Last Handshake trigger time
	height		uint64		// The node latest block height
	time		time.Time	// The latest time the node activity
	// TODO does this channel should be a buffer channel
	chF		chan func()	// Channel used to operate the node without lock
	private		*uint		// Reserver for future using
}

type nodeMap struct {
	node *node
	lock sync.RWMutex
	list map[string]*node
}

var nodes nodeMap

func newNode() (*node) {
	node := node{
		state: INIT,
		chF: make(chan func()),
	}

	// Update nonce
	runtime.SetFinalizer(&node, rmNode)
	go node.backend()
	return &node
}

func rmNode(node *node) {
	log.Printf("Remove node %s", node.addr)
}

// TODO pass pointer to method only need modify it
func (node *node) backend() {
	common.Trace()
	for f := range node.chF {
		f()
	}
}

func (node node) getID() string {
	return node.id
}

func (node node) getState() uint {
	return node.state
}

func (node node) getConn() net.Conn {
	return node.conn
}

func (node *node) setState(state uint) {
	node.state = state
}

func (node node) getHandshakeTime() (time.Time) {
	return node.handshakeTime
}

func (node *node) setHandshakeTime(t time.Time) {
	node.handshakeTime = t
}

func (node node) getHandshakeRetry() uint32 {
	return atomic.LoadUint32(&(node.handshakeRetry))
}

func (node *node) setHandshakeRetry(r uint32) {
	node.handshakeRetry = r
	atomic.StoreUint32(&(node.handshakeRetry), r)
}

func (node node) getHeight() uint64 {
	return node.height
}

func (node *node) updateTime(t time.Time) {
	node.time = t
}

func (node *node) rx() error {
	conn := node.getConn()
	from := conn.RemoteAddr().String()

	for {
		buf := make([]byte, MAXBUFLEN)
		len, err := conn.Read(buf[0:(MAXBUFLEN - 1)])
		buf[MAXBUFLEN - 1] = 0 //Prevent overflow

		switch err {
		case nil:
			go handleNodeMsg(node, buf, uint64(len))
			break
		case io.EOF:
			//log.Println("Reading EOF of network conn")
			break
		default:
			log.Printf("read error", err)
			goto DISCONNECT
		}
	}

DISCONNECT:
	err := conn.Close()
	node.setState(INACTIVITY)
	log.Printf("Close connection", from)
	return err
}

// Init the server port, should be run in another thread
func (node *node) initRx () {
	listener, err := net.Listen("tcp", "localhost:" + strconv.Itoa(NODETESTPORT))
	if err != nil {
		log.Println("Error listening", err.Error())
		return
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Error accepting", err.Error())
			return
		}
		node := newNode()
		// Currently we use the address as the ID
		node.id = conn.RemoteAddr().String()
		node.addr = conn.RemoteAddr().String()
		log.Println("Remote node %s connect with %s",
			conn.RemoteAddr(), conn.LocalAddr())
		node.conn = conn
		// TOOD close the conn when erro happened
		// TODO lock the node and assign the connection to Node.
		nodes.add(node)
		go node.rx()
	}
	//TODO When to free the net listen resouce?
}

func (node *node) connect(nodeAddr string)  {
	node.chF <- func() {
		common.Trace()
		conn, err := net.Dial("tcp", nodeAddr)
		if err != nil {
			log.Println("Error dialing", err.Error())
			return
		}

		node := newNode()
		node.conn = conn
		node.id = conn.RemoteAddr().String()
		node.addr = conn.RemoteAddr().String()
		// FixMe Only for testing
		node.height = 1000

		log.Printf("Connect node %s connect with %s with %s",
			conn.LocalAddr().String(), conn.RemoteAddr().String(),
			conn.RemoteAddr().Network())
		// TODO Need lock
		nodes.add(node)
		go node.rx()
	}
}

// TODO construct a TX channel and other application just drop the message to the channel
func (node node) tx(buf []byte) {
	node.chF <- func() {
		common.Trace()
		_, err := node.conn.Write(buf)
		if err != nil {
			log.Println("Error sending messge to peer node", err.Error())
		}
		return
	}
}

func (nodes *nodeMap) broadcast(buf []byte) {
	// TODO lock the map
	// TODO Check whether the node existed or not
	for _, node := range nodes.list {
		if node.state == ESTABLISH {
			go node.tx(buf)
		}
	}
}

func (nodes *nodeMap) add(node *node) {
	//TODO lock the node Map
	// TODO check whether the node existed or not
	// TODO dupicate IP address nodes issue
	nodes.list[node.id] = node
	// Unlock the map
}

func (nodes *nodeMap) delNode(node *node) {
	//TODO lock the node Map
	delete(nodes.list, node.id)
	// Unlock the map
}

func InitNodes() {
	// TODO write lock
	n := newNode()

	n.version = PROTOCOLVERSION
	n.services = NODESERVICES
	n.port = NODETESTPORT
	n.relay = true
	rand.Seed(time.Now().UTC().UnixNano())
	n.nonce = rand.Uint32()

	nodes.node = n
	nodes.list = make(map[string]*node)
}
<commit_msg>Add the Relay interface to transmit packet for other modules<commit_after>package net

import (
	"log"
	"strconv"
	"net"
	"io"
	"sync"
	"time"
	"runtime"
	"math/rand"
	"sync/atomic"
	"GoOnchain/common"
)

// The node state
const (
	INIT = 0
	HANDSHAKEING = 1
	HANDSHAKED = 2
	ESTABLISH = 3
	INACTIVITY = 4
)

// The node capability flag
const (
	RELAY  = 0x01
	SERVER = 0x02
	NODESERVICES = 0x01
)

type node struct {
	state		uint		// node status
	id		string		// The nodes's id, MAC or IP?
	addr		string 		// The address of the node
	conn		net.Conn	// Connect socket with the peer node
	nonce		uint32		// Random number to identify different entity from the same IP
	cap		uint32  	// The node capability set
	version		uint32		// The network protocol the node used
	services	uint64		// The services the node supplied
	port		uint16		// The server port of the node
	relay		bool		// The relay capability of the node (merge into capbility flag)
	handshakeRetry  uint32		// Handshake retry times
	handshakeTime	time.Time	// Last Handshake trigger time
	height		uint64		// The node latest block height
	time		time.Time	// The latest time the node activity
	// TODO does this channel should be a buffer channel
	chF		chan func()	// Channel used to operate the node without lock
	private		*uint		// Reserver for future using
}

type nodeMap struct {
	node *node
	lock sync.RWMutex
	list map[string]*node
}

var nodes nodeMap

func newNode() (*node) {
	node := node{
		state: INIT,
		chF: make(chan func()),
	}

	// Update nonce
	runtime.SetFinalizer(&node, rmNode)
	go node.backend()
	return &node
}

func rmNode(node *node) {
	log.Printf("Remove node %s", node.addr)
}

// TODO pass pointer to method only need modify it
func (node *node) backend() {
	common.Trace()
	for f := range node.chF {
		f()
	}
}

func (node node) getID() string {
	return node.id
}

func (node node) getState() uint {
	return node.state
}

func (node node) getConn() net.Conn {
	return node.conn
}

func (node *node) setState(state uint) {
	node.state = state
}

func (node node) getHandshakeTime() (time.Time) {
	return node.handshakeTime
}

func (node *node) setHandshakeTime(t time.Time) {
	node.handshakeTime = t
}

func (node node) getHandshakeRetry() uint32 {
	return atomic.LoadUint32(&(node.handshakeRetry))
}

func (node *node) setHandshakeRetry(r uint32) {
	node.handshakeRetry = r
	atomic.StoreUint32(&(node.handshakeRetry), r)
}

func (node node) getHeight() uint64 {
	return node.height
}

func (node *node) updateTime(t time.Time) {
	node.time = t
}

func (node *node) rx() error {
	conn := node.getConn()
	from := conn.RemoteAddr().String()

	for {
		buf := make([]byte, MAXBUFLEN)
		len, err := conn.Read(buf[0:(MAXBUFLEN - 1)])
		buf[MAXBUFLEN - 1] = 0 //Prevent overflow

		switch err {
		case nil:
			go handleNodeMsg(node, buf, uint64(len))
			break
		case io.EOF:
			//log.Println("Reading EOF of network conn")
			break
		default:
			log.Printf("read error", err)
			goto DISCONNECT
		}
	}

DISCONNECT:
	err := conn.Close()
	node.setState(INACTIVITY)
	log.Printf("Close connection", from)
	return err
}

// Init the server port, should be run in another thread
func (node *node) initRx () {
	listener, err := net.Listen("tcp", "localhost:" + strconv.Itoa(NODETESTPORT))
	if err != nil {
		log.Println("Error listening", err.Error())
		return
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Error accepting", err.Error())
			return
		}
		node := newNode()
		// Currently we use the address as the ID
		node.id = conn.RemoteAddr().String()
		node.addr = conn.RemoteAddr().String()
		log.Println("Remote node %s connect with %s",
			conn.RemoteAddr(), conn.LocalAddr())
		node.conn = conn
		// TOOD close the conn when erro happened
		// TODO lock the node and assign the connection to Node.
		nodes.add(node)
		go node.rx()
	}
	//TODO When to free the net listen resouce?
}

func (node *node) connect(nodeAddr string)  {
	node.chF <- func() {
		common.Trace()
		conn, err := net.Dial("tcp", nodeAddr)
		if err != nil {
			log.Println("Error dialing", err.Error())
			return
		}

		node := newNode()
		node.conn = conn
		node.id = conn.RemoteAddr().String()
		node.addr = conn.RemoteAddr().String()
		// FixMe Only for testing
		node.height = 1000

		log.Printf("Connect node %s connect with %s with %s",
			conn.LocalAddr().String(), conn.RemoteAddr().String(),
			conn.RemoteAddr().Network())
		// TODO Need lock
		nodes.add(node)
		go node.rx()
	}
}

// TODO construct a TX channel and other application just drop the message to the channel
func (node node) tx(buf []byte) {
	node.chF <- func() {
		common.Trace()
		_, err := node.conn.Write(buf)
		if err != nil {
			log.Println("Error sending messge to peer node", err.Error())
		}
		return
	}
}

func (nodes *nodeMap) broadcast(buf []byte) {
	// TODO lock the map
	// TODO Check whether the node existed or not
	for _, node := range nodes.list {
		if node.state == ESTABLISH {
			go node.tx(buf)
		}
	}
}

func (nodes *nodeMap) add(node *node) {
	//TODO lock the node Map
	// TODO check whether the node existed or not
	// TODO dupicate IP address nodes issue
	nodes.list[node.id] = node
	// Unlock the map
}

func (nodes *nodeMap) delNode(node *node) {
	//TODO lock the node Map
	delete(nodes.list, node.id)
	// Unlock the map
}

func InitNodes() {
	// TODO write lock
	n := newNode()

	n.version = PROTOCOLVERSION
	n.services = NODESERVICES
	n.port = NODETESTPORT
	n.relay = true
	rand.Seed(time.Now().UTC().UnixNano())
	n.nonce = rand.Uint32()

	nodes.node = n
	nodes.list = make(map[string]*node)
}

func Relay(msgType string, msg interface{}) {
	// TODO Unicast or broadcast the message based on the type
	//node.tx()
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcchain

import (
	"fmt"
)

// NotificationType represents the type of a notification message.
type NotificationType int

// Constants for the type of a notification message.
const (
	// NTOrphanBlock indicates an orphan block was processed and the
	// associated block hash is the root of all known orphans which should
	// be used to request the missing blocks.
	NTOrphanBlock NotificationType = iota

	// NTBlockAccepted indicates the associated block was accepted into
	// the block chain.  Note that this does not necessarily mean it was
	// added to the main chain.  For that, use NTBlockConnected.
	NTBlockAccepted

	// NTBlockConnected indicates the associated block was connected to the
	// main chain.
	NTBlockConnected

	// NTBlockDisconnected indicates the associated block was disconnected
	// from the main chain.
	NTBlockDisconnected
)

// notificationTypeStrings is a map of notification types back to their constant
// names for pretty printing.
var notificationTypeStrings = map[NotificationType]string{
	NTOrphanBlock:       "NTOrphanBlock",
	NTBlockAccepted:     "NTBlockAccepted",
	NTBlockConnected:    "NTBlockConnected",
	NTBlockDisconnected: "NTBlockDisconnected",
}

// String returns the NotificationType in human-readable form.
func (n NotificationType) String() string {
	if s, ok := notificationTypeStrings[n]; ok {
		return s
	}
	return fmt.Sprintf("Unknown Notification Type (%d)", int(n))
}

// Notification defines an asynchronous notification that is sent to the caller
// over the notification channel provided during the call to New and consists
// of a notification type as well as associated data that depends on the type as
// follows:
// 	- NTOrphanBlock:       *btcwire.ShaHash
// 	- NTBlockAccepted:     *btcutil.Block
// 	- NTBlockConnected:    *btcutil.Block
// 	- NTBlockDisconnected: *btcutil.Block
type Notification struct {
	Type NotificationType
	Data interface{}
}

// sendNotification sends a notification with the passed type and data if the
// caller requested notifications by providing a channel in the call to New.
func (b *BlockChain) sendNotification(typ NotificationType, data interface{}) {
	// Ignore it if the caller didn't request notifications.
	if b.notifications == nil {
		return
	}

	// Generate and send the notification.
	n := Notification{Type: typ, Data: data}
	b.notifications <- &n
}
<commit_msg>Send notification in their own goroutine.<commit_after>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcchain

import (
	"fmt"
)

// NotificationType represents the type of a notification message.
type NotificationType int

// Constants for the type of a notification message.
const (
	// NTOrphanBlock indicates an orphan block was processed and the
	// associated block hash is the root of all known orphans which should
	// be used to request the missing blocks.
	NTOrphanBlock NotificationType = iota

	// NTBlockAccepted indicates the associated block was accepted into
	// the block chain.  Note that this does not necessarily mean it was
	// added to the main chain.  For that, use NTBlockConnected.
	NTBlockAccepted

	// NTBlockConnected indicates the associated block was connected to the
	// main chain.
	NTBlockConnected

	// NTBlockDisconnected indicates the associated block was disconnected
	// from the main chain.
	NTBlockDisconnected
)

// notificationTypeStrings is a map of notification types back to their constant
// names for pretty printing.
var notificationTypeStrings = map[NotificationType]string{
	NTOrphanBlock:       "NTOrphanBlock",
	NTBlockAccepted:     "NTBlockAccepted",
	NTBlockConnected:    "NTBlockConnected",
	NTBlockDisconnected: "NTBlockDisconnected",
}

// String returns the NotificationType in human-readable form.
func (n NotificationType) String() string {
	if s, ok := notificationTypeStrings[n]; ok {
		return s
	}
	return fmt.Sprintf("Unknown Notification Type (%d)", int(n))
}

// Notification defines an asynchronous notification that is sent to the caller
// over the notification channel provided during the call to New and consists
// of a notification type as well as associated data that depends on the type as
// follows:
// 	- NTOrphanBlock:       *btcwire.ShaHash
// 	- NTBlockAccepted:     *btcutil.Block
// 	- NTBlockConnected:    *btcutil.Block
// 	- NTBlockDisconnected: *btcutil.Block
type Notification struct {
	Type NotificationType
	Data interface{}
}

// sendNotification sends a notification with the passed type and data if the
// caller requested notifications by providing a channel in the call to New.
func (b *BlockChain) sendNotification(typ NotificationType, data interface{}) {
	// Ignore it if the caller didn't request notifications.
	if b.notifications == nil {
		return
	}

	// Generate and send the notification asynchronously.
	go func() {
		n := Notification{Type: typ, Data: data}
		b.notifications <- &n
	}()
}
<|endoftext|>


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

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"time"

	"github.com/golang/protobuf/jsonpb"
	"github.com/maruel/subcommands"
	"google.golang.org/genproto/protobuf/field_mask"

	"go.chromium.org/luci/auth"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/cli"
)

func cmdCollect(defaultAuthOpts auth.Options) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: `collect [flags] <BUILD> [<BUILD>...]`,
		ShortDesc: "wait for builds to end",
		LongDesc: `Wait for builds to end.

Argument BUILD can be an int64 build id or a string
<project>/<bucket>/<builder>/<build_number>, e.g. chromium/ci/linux-rel/1

Optionally writes build details into an output file as JSON array of
bulidbucket.v2.Build proto messages.`,
		CommandRun: func() subcommands.CommandRun {
			r := &collectRun{}
			r.RegisterGlobalFlags(defaultAuthOpts)
			r.Flags.DurationVar(
				&r.intervalArg, "interval", time.Minute,
				"duration to wait between requests")
			r.Flags.StringVar(&r.outputArg, "json-output", "", "path to the file "+
				"into which build details are written; nothing is written if "+
				"unspecified")
			return r
		},
	}
}

type collectRun struct {
	baseCommandRun
	intervalArg time.Duration
	outputArg   string
}

var getRequestFieldMask = &field_mask.FieldMask{
	Paths: []string{
		// Include all output properties, which are omitted by default.
		"output",

		// Since we add output props above, we must list default fields too.
		"builder",
		"create_time",
		"created_by",
		"end_time",
		"id",
		"input.experimental",
		"input.gerrit_changes",
		"input.gitiles_commit",
		"number",
		"start_time",
		"status",
		"update_time",
	},
}

func collectBuildDetails(ctx context.Context, client buildbucketpb.BuildsClient, buildIDs []int64, sleep func()) (map[int64]*buildbucketpb.Build, error) {
	unfinishedBuilds := buildIDs
	endedBuildDetails := make(map[int64]*buildbucketpb.Build, len(buildIDs))
	for {
		breq := &buildbucketpb.BatchRequest{}
		for _, buildID := range unfinishedBuilds {
			getReq := &buildbucketpb.GetBuildRequest{
				Id:     buildID,
				Fields: getRequestFieldMask,
			}
			breq.Requests = append(breq.Requests, &buildbucketpb.BatchRequest_Request{
				Request: &buildbucketpb.BatchRequest_Request_GetBuild{
					GetBuild: getReq,
				},
			})
		}

		fmt.Printf("Checking build statuses... ")
		bresp, err := client.Batch(ctx, breq)
		if err != nil {
			return nil, err
		}

		unfinishedBuilds = nil
		for _, resp := range bresp.Responses {
			build := resp.GetGetBuild()
			if build.Status&buildbucketpb.Status_ENDED_MASK != 0 {
				endedBuildDetails[build.Id] = build
			} else {
				unfinishedBuilds = append(unfinishedBuilds, build.Id)
			}
		}

		if len(unfinishedBuilds) == 0 {
			fmt.Printf("all ended\n")
			break
		}

		fmt.Printf(
			"%d still running/pending: %v\n", len(unfinishedBuilds), unfinishedBuilds)
		sleep()
	}

	return endedBuildDetails, nil
}

func writeBuildDetails(buildIDs []int64, buildDetails map[int64]*buildbucketpb.Build, outputPath string) error {
	outputFile, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer outputFile.Close()

	rawMessages := make([]json.RawMessage, len(buildIDs))
	m := &jsonpb.Marshaler{}
	buf := bytes.Buffer{}
	for i, buildID := range buildIDs {
		buf.Reset()
		if err := m.Marshal(&buf, buildDetails[buildID]); err != nil {
			return err
		}
		rawMessages[i] = append(json.RawMessage{}, buf.Bytes()...)
	}

	return json.NewEncoder(outputFile).Encode(rawMessages)
}

func (r *collectRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	ctx := cli.GetContext(a, r, env)
	if err := r.initClients(ctx); err != nil {
		return r.done(ctx, err)
	}

	buildIDs, err := r.retrieveBuildIDs(ctx, args)
	if err != nil {
		return r.done(ctx, err)
	}

	buildDetails, err := collectBuildDetails(ctx, r.client, buildIDs, func() {
		fmt.Printf("Waiting %s before trying again...\n", r.intervalArg)
		time.Sleep(r.intervalArg)
	})
	if err != nil {
		return r.done(ctx, err)
	}

	if r.outputArg != "" {
		if err := writeBuildDetails(buildIDs, buildDetails, r.outputArg); err != nil {
			return r.done(ctx, err)
		}
	}

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

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"time"

	"github.com/golang/protobuf/jsonpb"
	"github.com/maruel/subcommands"
	"google.golang.org/genproto/protobuf/field_mask"
	"google.golang.org/grpc/codes"

	"go.chromium.org/luci/auth"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/grpc/grpcutil"
)

func cmdCollect(defaultAuthOpts auth.Options) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: `collect [flags] <BUILD> [<BUILD>...]`,
		ShortDesc: "wait for builds to end",
		LongDesc: `Wait for builds to end.

Argument BUILD can be an int64 build id or a string
<project>/<bucket>/<builder>/<build_number>, e.g. chromium/ci/linux-rel/1

Optionally writes build details into an output file as JSON array of
bulidbucket.v2.Build proto messages.`,
		CommandRun: func() subcommands.CommandRun {
			r := &collectRun{}
			r.RegisterGlobalFlags(defaultAuthOpts)
			r.Flags.DurationVar(
				&r.intervalArg, "interval", time.Minute,
				"duration to wait between requests")
			r.Flags.StringVar(&r.outputArg, "json-output", "", "path to the file "+
				"into which build details are written; nothing is written if "+
				"unspecified")
			return r
		},
	}
}

type collectRun struct {
	baseCommandRun
	intervalArg time.Duration
	outputArg   string
}

var getRequestFieldMask = &field_mask.FieldMask{
	Paths: []string{
		// Include all output properties, which are omitted by default.
		"output",

		// Since we add output props above, we must list default fields too.
		"builder",
		"create_time",
		"created_by",
		"end_time",
		"id",
		"input.experimental",
		"input.gerrit_changes",
		"input.gitiles_commit",
		"number",
		"start_time",
		"status",
		"update_time",
	},
}

// dedupInt64s dedups int64s. Reuses the underlying array.
func dedupInt64s(nums []int64) []int64 {
	seen := make(map[int64]struct{}, len(nums))
	res := make([]int64, 0, len(nums))
	for _, x := range nums {
		if _, ok := seen[x]; !ok {
			seen[x] = struct{}{}
			res = append(res, x)
		}
	}
	return res
}

func collectBuildDetails(ctx context.Context, client buildbucketpb.BuildsClient, buildIDs []int64, sleep func()) (map[int64]*buildbucketpb.Build, error) {
	buildIDs = dedupInt64s(buildIDs)
	ended := make(map[int64]*buildbucketpb.Build, len(buildIDs))
	for {
		breq := &buildbucketpb.BatchRequest{}
		for _, id := range buildIDs {
			if _, ok := ended[id]; ok {
				continue
			}

			getReq := &buildbucketpb.GetBuildRequest{
				Id:     id,
				Fields: getRequestFieldMask,
			}
			breq.Requests = append(breq.Requests, &buildbucketpb.BatchRequest_Request{
				Request: &buildbucketpb.BatchRequest_Request_GetBuild{
					GetBuild: getReq,
				},
			})
		}

		logging.Infof(ctx, "checking build statuses...")
		bresp, err := client.Batch(ctx, breq)
		if err != nil {
			return nil, err
		}

		for _, resp := range bresp.Responses {
			error := resp.GetError()
			code := codes.Code(error.GetCode())
			build := resp.GetGetBuild()
			switch {
			case grpcutil.IsTransientCode(code):
				logging.Warningf(ctx, "transient %s error: %s", code, error.Message)

			case code != codes.OK:
				return nil, fmt.Errorf("RPC error %s: %s", code, error.Message)

			case build.Status&buildbucketpb.Status_ENDED_MASK != 0:
				ended[build.Id] = build
				logging.Infof(ctx, "%d ended: %s", build.Id, build.Status)
			}
		}

		logging.Infof(ctx, "%d are running/pending", len(buildIDs)-len(ended))
		if len(buildIDs) == len(ended) {
			break
		}
		sleep()
	}

	return ended, nil
}

func writeBuildDetails(buildIDs []int64, buildDetails map[int64]*buildbucketpb.Build, outputPath string) error {
	outputFile, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer outputFile.Close()

	rawMessages := make([]json.RawMessage, len(buildIDs))
	m := &jsonpb.Marshaler{}
	buf := bytes.Buffer{}
	for i, buildID := range buildIDs {
		buf.Reset()
		if err := m.Marshal(&buf, buildDetails[buildID]); err != nil {
			return err
		}
		rawMessages[i] = append(json.RawMessage{}, buf.Bytes()...)
	}

	return json.NewEncoder(outputFile).Encode(rawMessages)
}

func (r *collectRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	ctx := cli.GetContext(a, r, env)
	if err := r.initClients(ctx); err != nil {
		return r.done(ctx, err)
	}

	buildIDs, err := r.retrieveBuildIDs(ctx, args)
	if err != nil {
		return r.done(ctx, err)
	}

	buildDetails, err := collectBuildDetails(ctx, r.client, buildIDs, func() {
		logging.Infof(ctx, "waiting %s before trying again...", r.intervalArg)
		time.Sleep(r.intervalArg)
	})
	if err != nil {
		return r.done(ctx, err)
	}

	if r.outputArg != "" {
		if err := writeBuildDetails(buildIDs, buildDetails, r.outputArg); err != nil {
			return r.done(ctx, err)
		}
	}

	return r.done(ctx, nil)
}
<|endoftext|>


<commit_before>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type Config

// The amazonebs package contains a packer.Builder implementation that
// builds AMIs for Amazon EC2.
//
// In general, there are two types of AMIs that can be created: ebs-backed or
// instance-store. This builder _only_ builds ebs-backed images.
package ebs

import (
	"context"
	"fmt"

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/iam"
	awscommon "github.com/hashicorp/packer/builder/amazon/common"
	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
)

// The unique ID for this builder
const BuilderId = "mitchellh.amazonebs"

type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`
	awscommon.AMIConfig    `mapstructure:",squash"`
	awscommon.RunConfig    `mapstructure:",squash"`
	// Add one or more block device mappings to the AMI. These will be attached
	// when booting a new instance from your AMI. To add a block device during
	// the Packer build see `launch_block_device_mappings` below. Your options
	// here may vary depending on the type of VM you use. See the
	// [BlockDevices](#block-devices-configuration) documentation for fields.
	AMIMappings awscommon.BlockDevices `mapstructure:"ami_block_device_mappings" required:"false"`
	// Add one or more block devices before the Packer build starts. If you add
	// instance store volumes or EBS volumes in addition to the root device
	// volume, the created AMI will contain block device mapping information
	// for those volumes. Amazon creates snapshots of the source instance's
	// root volume and any other EBS volumes described here. When you launch an
	// instance from this new AMI, the instance automatically launches with
	// these additional volumes, and will restore them from snapshots taken
	// from the source instance. See the
	// [BlockDevices](#block-devices-configuration) documentation for fields.
	LaunchMappings awscommon.BlockDevices `mapstructure:"launch_block_device_mappings" required:"false"`
	// Tags to apply to the volumes that are *launched* to create the AMI.
	// These tags are *not* applied to the resulting AMI unless they're
	// duplicated in `tags`. This is a [template
	// engine](/docs/templates/engine.html), see [Build template
	// data](#build-template-data) for more information.
	VolumeRunTags awscommon.TagMap `mapstructure:"run_volume_tags"`
	// Relevant only to Windows guests: If you set this flag, we'll add clauses
	// to the launch_block_device_mappings that make sure ephemeral drives
	// don't show up in the EC2 console. If you launched from the EC2 console,
	// you'd get this automatically, but the SDK does not provide this service.
	// For more information, see
	// https://docs.aws.amazon.com/AWSEC2/latest/WindowsGuide/InstanceStorage.html.
	// Because we don't validate the OS type of your guest, it is up to you to
	// make sure you don't set this for *nix guests; behavior may be
	// unpredictable.
	NoEphemeral bool `mapstructure:"no_ephemeral" required:"false"`

	ctx interpolate.Context
}

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	b.config.ctx.Funcs = awscommon.TemplateFuncs
	err := config.Decode(&b.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &b.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"ami_description",
				"run_tags",
				"run_volume_tags",
				"spot_tags",
				"snapshot_tags",
				"tags",
			},
		},
	}, raws...)
	if err != nil {
		return nil, err
	}

	if b.config.PackerConfig.PackerForce {
		b.config.AMIForceDeregister = true
	}

	// Accumulate any errors
	var errs *packer.MultiError
	var warns []string

	errs = packer.MultiErrorAppend(errs, b.config.AccessConfig.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs,
		b.config.AMIConfig.Prepare(&b.config.AccessConfig, &b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.AMIMappings.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.LaunchMappings.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.RunConfig.Prepare(&b.config.ctx)...)

	if b.config.IsSpotInstance() && (b.config.AMIENASupport.True() || b.config.AMISriovNetSupport) {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("Spot instances do not support modification, which is required "+
				"when either `ena_support` or `sriov_support` are set. Please ensure "+
				"you use an AMI that already has either SR-IOV or ENA enabled."))
	}

	if b.config.RunConfig.SpotPriceAutoProduct != "" {
		warns = append(warns, "spot_price_auto_product is deprecated and no "+
			"longer necessary for Packer builds. In future versions of "+
			"Packer, inclusion of spot_price_auto_product will error your "+
			"builds. Please take a look at our current documentation to "+
			"understand how Packer requests Spot instances.")
	}

	if errs != nil && len(errs.Errors) > 0 {
		return warns, errs
	}

	packer.LogSecretFilter.Set(b.config.AccessKey, b.config.SecretKey, b.config.Token)
	return warns, nil
}

func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (packer.Artifact, error) {

	session, err := b.config.Session()
	if err != nil {
		return nil, err
	}

	ec2conn := ec2.New(session)
	iam := iam.New(session)
	// Setup the state bag and initial state for the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", &b.config)
	state.Put("access_config", &b.config.AccessConfig)
	state.Put("ami_config", &b.config.AMIConfig)
	state.Put("ec2", ec2conn)
	state.Put("iam", iam)
	state.Put("awsSession", session)
	state.Put("hook", hook)
	state.Put("ui", ui)

	var instanceStep multistep.Step

	if b.config.IsSpotInstance() {
		instanceStep = &awscommon.StepRunSpotInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			LaunchMappings:                    b.config.LaunchMappings,
			BlockDurationMinutes:              b.config.BlockDurationMinutes,
			Ctx:                               b.config.ctx,
			Comm:                              &b.config.RunConfig.Comm,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			ExpectedRootDevice:                "ebs",
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			SourceAMI:                         b.config.SourceAmi,
			SpotPrice:                         b.config.SpotPrice,
			SpotTags:                          b.config.SpotTags,
			Tags:                              b.config.RunTags,
			SpotInstanceTypes:                 b.config.SpotInstanceTypes,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
		}
	} else {
		instanceStep = &awscommon.StepRunSourceInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			LaunchMappings:                    b.config.LaunchMappings,
			Comm:                              &b.config.RunConfig.Comm,
			Ctx:                               b.config.ctx,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			EnableT2Unlimited:                 b.config.EnableT2Unlimited,
			ExpectedRootDevice:                "ebs",
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			IsRestricted:                      b.config.IsChinaCloud() || b.config.IsGovCloud(),
			SourceAMI:                         b.config.SourceAmi,
			Tags:                              b.config.RunTags,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
			NoEphemeral:                       b.config.NoEphemeral,
		}
	}

	// Build the steps
	steps := []multistep.Step{
		&awscommon.StepPreValidate{
			DestAmiName:        b.config.AMIName,
			ForceDeregister:    b.config.AMIForceDeregister,
			AMISkipBuildRegion: b.config.AMISkipBuildRegion,
			VpcId:              b.config.VpcId,
			SubnetId:           b.config.SubnetId,
		},
		&awscommon.StepSourceAMIInfo{
			SourceAmi:                b.config.SourceAmi,
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
			AmiFilters:               b.config.SourceAmiFilter,
			AMIVirtType:              b.config.AMIVirtType,
		},
		&awscommon.StepNetworkInfo{
			VpcId:               b.config.VpcId,
			VpcFilter:           b.config.VpcFilter,
			SecurityGroupIds:    b.config.SecurityGroupIds,
			SecurityGroupFilter: b.config.SecurityGroupFilter,
			SubnetId:            b.config.SubnetId,
			SubnetFilter:        b.config.SubnetFilter,
			AvailabilityZone:    b.config.AvailabilityZone,
		},
		&awscommon.StepKeyPair{
			Debug:        b.config.PackerDebug,
			Comm:         &b.config.RunConfig.Comm,
			DebugKeyPath: fmt.Sprintf("ec2_%s.pem", b.config.PackerBuildName),
		},
		&awscommon.StepSecurityGroup{
			SecurityGroupFilter:    b.config.SecurityGroupFilter,
			SecurityGroupIds:       b.config.SecurityGroupIds,
			CommConfig:             &b.config.RunConfig.Comm,
			TemporarySGSourceCidrs: b.config.TemporarySGSourceCidrs,
		},
		&awscommon.StepIamInstanceProfile{
			IamInstanceProfile:                        b.config.IamInstanceProfile,
			TemporaryIamInstanceProfilePolicyDocument: b.config.TemporaryIamInstanceProfilePolicyDocument,
		},
		&awscommon.StepCleanupVolumes{
			LaunchMappings: b.config.LaunchMappings,
		},
		instanceStep,
		&awscommon.StepGetPassword{
			Debug:     b.config.PackerDebug,
			Comm:      &b.config.RunConfig.Comm,
			Timeout:   b.config.WindowsPasswordTimeout,
			BuildName: b.config.PackerBuildName,
		},
		&communicator.StepConnect{
			Config: &b.config.RunConfig.Comm,
			Host: awscommon.SSHHost(
				ec2conn,
				b.config.SSHInterface,
				b.config.Comm.SSHHost,
			),
			SSHConfig: b.config.RunConfig.Comm.SSHConfigFunc(),
		},
		&common.StepProvision{},
		&common.StepCleanupTempKeys{
			Comm: &b.config.RunConfig.Comm,
		},
		&awscommon.StepStopEBSBackedInstance{
			Skip:                b.config.IsSpotInstance(),
			DisableStopInstance: b.config.DisableStopInstance,
		},
		&awscommon.StepModifyEBSBackedInstance{
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
		},
		&awscommon.StepDeregisterAMI{
			AccessConfig:        &b.config.AccessConfig,
			ForceDeregister:     b.config.AMIForceDeregister,
			ForceDeleteSnapshot: b.config.AMIForceDeleteSnapshot,
			AMIName:             b.config.AMIName,
			Regions:             b.config.AMIRegions,
		},
		&stepCreateAMI{
			AMISkipBuildRegion: b.config.AMISkipBuildRegion,
		},
		&awscommon.StepAMIRegionCopy{
			AccessConfig:       &b.config.AccessConfig,
			Regions:            b.config.AMIRegions,
			AMIKmsKeyId:        b.config.AMIKmsKeyId,
			RegionKeyIds:       b.config.AMIRegionKMSKeyIDs,
			EncryptBootVolume:  b.config.AMIEncryptBootVolume,
			Name:               b.config.AMIName,
			OriginalRegion:     *ec2conn.Config.Region,
			AMISkipBuildRegion: b.config.AMISkipBuildRegion,
		},
		&awscommon.StepModifyAMIAttributes{
			Description:    b.config.AMIDescription,
			Users:          b.config.AMIUsers,
			Groups:         b.config.AMIGroups,
			ProductCodes:   b.config.AMIProductCodes,
			SnapshotUsers:  b.config.SnapshotUsers,
			SnapshotGroups: b.config.SnapshotGroups,
			Ctx:            b.config.ctx,
		},
		&awscommon.StepCreateTags{
			Tags:         b.config.AMITags,
			SnapshotTags: b.config.SnapshotTags,
			Ctx:          b.config.ctx,
		},
	}

	// Run!
	b.runner = common.NewRunner(steps, b.config.PackerConfig, ui)
	b.runner.Run(ctx, state)
	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// If there are no AMIs, then just return
	if _, ok := state.GetOk("amis"); !ok {
		return nil, nil
	}

	// Build the artifact and return it
	artifact := &awscommon.Artifact{
		Amis:           state.Get("amis").(map[string]string),
		BuilderIdValue: BuilderId,
		Session:        session,
	}

	return artifact, nil
}
<commit_msg>send to step run spot instance<commit_after>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type Config

// The amazonebs package contains a packer.Builder implementation that
// builds AMIs for Amazon EC2.
//
// In general, there are two types of AMIs that can be created: ebs-backed or
// instance-store. This builder _only_ builds ebs-backed images.
package ebs

import (
	"context"
	"fmt"

	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/iam"
	awscommon "github.com/hashicorp/packer/builder/amazon/common"
	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
)

// The unique ID for this builder
const BuilderId = "mitchellh.amazonebs"

type Config struct {
	common.PackerConfig    `mapstructure:",squash"`
	awscommon.AccessConfig `mapstructure:",squash"`
	awscommon.AMIConfig    `mapstructure:",squash"`
	awscommon.RunConfig    `mapstructure:",squash"`
	// Add one or more block device mappings to the AMI. These will be attached
	// when booting a new instance from your AMI. To add a block device during
	// the Packer build see `launch_block_device_mappings` below. Your options
	// here may vary depending on the type of VM you use. See the
	// [BlockDevices](#block-devices-configuration) documentation for fields.
	AMIMappings awscommon.BlockDevices `mapstructure:"ami_block_device_mappings" required:"false"`
	// Add one or more block devices before the Packer build starts. If you add
	// instance store volumes or EBS volumes in addition to the root device
	// volume, the created AMI will contain block device mapping information
	// for those volumes. Amazon creates snapshots of the source instance's
	// root volume and any other EBS volumes described here. When you launch an
	// instance from this new AMI, the instance automatically launches with
	// these additional volumes, and will restore them from snapshots taken
	// from the source instance. See the
	// [BlockDevices](#block-devices-configuration) documentation for fields.
	LaunchMappings awscommon.BlockDevices `mapstructure:"launch_block_device_mappings" required:"false"`
	// Tags to apply to the volumes that are *launched* to create the AMI.
	// These tags are *not* applied to the resulting AMI unless they're
	// duplicated in `tags`. This is a [template
	// engine](/docs/templates/engine.html), see [Build template
	// data](#build-template-data) for more information.
	VolumeRunTags awscommon.TagMap `mapstructure:"run_volume_tags"`
	// Relevant only to Windows guests: If you set this flag, we'll add clauses
	// to the launch_block_device_mappings that make sure ephemeral drives
	// don't show up in the EC2 console. If you launched from the EC2 console,
	// you'd get this automatically, but the SDK does not provide this service.
	// For more information, see
	// https://docs.aws.amazon.com/AWSEC2/latest/WindowsGuide/InstanceStorage.html.
	// Because we don't validate the OS type of your guest, it is up to you to
	// make sure you don't set this for *nix guests; behavior may be
	// unpredictable.
	NoEphemeral bool `mapstructure:"no_ephemeral" required:"false"`

	ctx interpolate.Context
}

type Builder struct {
	config Config
	runner multistep.Runner
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	b.config.ctx.Funcs = awscommon.TemplateFuncs
	err := config.Decode(&b.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &b.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"ami_description",
				"run_tags",
				"run_volume_tags",
				"spot_tags",
				"snapshot_tags",
				"tags",
			},
		},
	}, raws...)
	if err != nil {
		return nil, err
	}

	if b.config.PackerConfig.PackerForce {
		b.config.AMIForceDeregister = true
	}

	// Accumulate any errors
	var errs *packer.MultiError
	var warns []string

	errs = packer.MultiErrorAppend(errs, b.config.AccessConfig.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs,
		b.config.AMIConfig.Prepare(&b.config.AccessConfig, &b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.AMIMappings.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.LaunchMappings.Prepare(&b.config.ctx)...)
	errs = packer.MultiErrorAppend(errs, b.config.RunConfig.Prepare(&b.config.ctx)...)

	if b.config.IsSpotInstance() && (b.config.AMIENASupport.True() || b.config.AMISriovNetSupport) {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("Spot instances do not support modification, which is required "+
				"when either `ena_support` or `sriov_support` are set. Please ensure "+
				"you use an AMI that already has either SR-IOV or ENA enabled."))
	}

	if b.config.RunConfig.SpotPriceAutoProduct != "" {
		warns = append(warns, "spot_price_auto_product is deprecated and no "+
			"longer necessary for Packer builds. In future versions of "+
			"Packer, inclusion of spot_price_auto_product will error your "+
			"builds. Please take a look at our current documentation to "+
			"understand how Packer requests Spot instances.")
	}

	if errs != nil && len(errs.Errors) > 0 {
		return warns, errs
	}

	packer.LogSecretFilter.Set(b.config.AccessKey, b.config.SecretKey, b.config.Token)
	return warns, nil
}

func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (packer.Artifact, error) {

	session, err := b.config.Session()
	if err != nil {
		return nil, err
	}

	ec2conn := ec2.New(session)
	iam := iam.New(session)
	// Setup the state bag and initial state for the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", &b.config)
	state.Put("access_config", &b.config.AccessConfig)
	state.Put("ami_config", &b.config.AMIConfig)
	state.Put("ec2", ec2conn)
	state.Put("iam", iam)
	state.Put("awsSession", session)
	state.Put("hook", hook)
	state.Put("ui", ui)

	var instanceStep multistep.Step

	if b.config.IsSpotInstance() {
		instanceStep = &awscommon.StepRunSpotInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			LaunchMappings:                    b.config.LaunchMappings,
			BlockDurationMinutes:              b.config.BlockDurationMinutes,
			Ctx:                               b.config.ctx,
			Comm:                              &b.config.RunConfig.Comm,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			ExpectedRootDevice:                "ebs",
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			SourceAMI:                         b.config.SourceAmi,
			SpotPrice:                         b.config.SpotPrice,
			SpotTags:                          b.config.SpotTags,
			Tags:                              b.config.RunTags,
			SpotInstanceTypes:                 b.config.SpotInstanceTypes,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
			NoEphemeral:                       b.config.NoEphemeral,
		}
	} else {
		instanceStep = &awscommon.StepRunSourceInstance{
			AssociatePublicIpAddress:          b.config.AssociatePublicIpAddress,
			LaunchMappings:                    b.config.LaunchMappings,
			Comm:                              &b.config.RunConfig.Comm,
			Ctx:                               b.config.ctx,
			Debug:                             b.config.PackerDebug,
			EbsOptimized:                      b.config.EbsOptimized,
			EnableT2Unlimited:                 b.config.EnableT2Unlimited,
			ExpectedRootDevice:                "ebs",
			InstanceInitiatedShutdownBehavior: b.config.InstanceInitiatedShutdownBehavior,
			InstanceType:                      b.config.InstanceType,
			IsRestricted:                      b.config.IsChinaCloud() || b.config.IsGovCloud(),
			SourceAMI:                         b.config.SourceAmi,
			Tags:                              b.config.RunTags,
			UserData:                          b.config.UserData,
			UserDataFile:                      b.config.UserDataFile,
			VolumeTags:                        b.config.VolumeRunTags,
			NoEphemeral:                       b.config.NoEphemeral,
		}
	}

	// Build the steps
	steps := []multistep.Step{
		&awscommon.StepPreValidate{
			DestAmiName:        b.config.AMIName,
			ForceDeregister:    b.config.AMIForceDeregister,
			AMISkipBuildRegion: b.config.AMISkipBuildRegion,
			VpcId:              b.config.VpcId,
			SubnetId:           b.config.SubnetId,
		},
		&awscommon.StepSourceAMIInfo{
			SourceAmi:                b.config.SourceAmi,
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
			AmiFilters:               b.config.SourceAmiFilter,
			AMIVirtType:              b.config.AMIVirtType,
		},
		&awscommon.StepNetworkInfo{
			VpcId:               b.config.VpcId,
			VpcFilter:           b.config.VpcFilter,
			SecurityGroupIds:    b.config.SecurityGroupIds,
			SecurityGroupFilter: b.config.SecurityGroupFilter,
			SubnetId:            b.config.SubnetId,
			SubnetFilter:        b.config.SubnetFilter,
			AvailabilityZone:    b.config.AvailabilityZone,
		},
		&awscommon.StepKeyPair{
			Debug:        b.config.PackerDebug,
			Comm:         &b.config.RunConfig.Comm,
			DebugKeyPath: fmt.Sprintf("ec2_%s.pem", b.config.PackerBuildName),
		},
		&awscommon.StepSecurityGroup{
			SecurityGroupFilter:    b.config.SecurityGroupFilter,
			SecurityGroupIds:       b.config.SecurityGroupIds,
			CommConfig:             &b.config.RunConfig.Comm,
			TemporarySGSourceCidrs: b.config.TemporarySGSourceCidrs,
		},
		&awscommon.StepIamInstanceProfile{
			IamInstanceProfile:                        b.config.IamInstanceProfile,
			TemporaryIamInstanceProfilePolicyDocument: b.config.TemporaryIamInstanceProfilePolicyDocument,
		},
		&awscommon.StepCleanupVolumes{
			LaunchMappings: b.config.LaunchMappings,
		},
		instanceStep,
		&awscommon.StepGetPassword{
			Debug:     b.config.PackerDebug,
			Comm:      &b.config.RunConfig.Comm,
			Timeout:   b.config.WindowsPasswordTimeout,
			BuildName: b.config.PackerBuildName,
		},
		&communicator.StepConnect{
			Config: &b.config.RunConfig.Comm,
			Host: awscommon.SSHHost(
				ec2conn,
				b.config.SSHInterface,
				b.config.Comm.SSHHost,
			),
			SSHConfig: b.config.RunConfig.Comm.SSHConfigFunc(),
		},
		&common.StepProvision{},
		&common.StepCleanupTempKeys{
			Comm: &b.config.RunConfig.Comm,
		},
		&awscommon.StepStopEBSBackedInstance{
			Skip:                b.config.IsSpotInstance(),
			DisableStopInstance: b.config.DisableStopInstance,
		},
		&awscommon.StepModifyEBSBackedInstance{
			EnableAMISriovNetSupport: b.config.AMISriovNetSupport,
			EnableAMIENASupport:      b.config.AMIENASupport,
		},
		&awscommon.StepDeregisterAMI{
			AccessConfig:        &b.config.AccessConfig,
			ForceDeregister:     b.config.AMIForceDeregister,
			ForceDeleteSnapshot: b.config.AMIForceDeleteSnapshot,
			AMIName:             b.config.AMIName,
			Regions:             b.config.AMIRegions,
		},
		&stepCreateAMI{
			AMISkipBuildRegion: b.config.AMISkipBuildRegion,
		},
		&awscommon.StepAMIRegionCopy{
			AccessConfig:       &b.config.AccessConfig,
			Regions:            b.config.AMIRegions,
			AMIKmsKeyId:        b.config.AMIKmsKeyId,
			RegionKeyIds:       b.config.AMIRegionKMSKeyIDs,
			EncryptBootVolume:  b.config.AMIEncryptBootVolume,
			Name:               b.config.AMIName,
			OriginalRegion:     *ec2conn.Config.Region,
			AMISkipBuildRegion: b.config.AMISkipBuildRegion,
		},
		&awscommon.StepModifyAMIAttributes{
			Description:    b.config.AMIDescription,
			Users:          b.config.AMIUsers,
			Groups:         b.config.AMIGroups,
			ProductCodes:   b.config.AMIProductCodes,
			SnapshotUsers:  b.config.SnapshotUsers,
			SnapshotGroups: b.config.SnapshotGroups,
			Ctx:            b.config.ctx,
		},
		&awscommon.StepCreateTags{
			Tags:         b.config.AMITags,
			SnapshotTags: b.config.SnapshotTags,
			Ctx:          b.config.ctx,
		},
	}

	// Run!
	b.runner = common.NewRunner(steps, b.config.PackerConfig, ui)
	b.runner.Run(ctx, state)
	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// If there are no AMIs, then just return
	if _, ok := state.GetOk("amis"); !ok {
		return nil, nil
	}

	// Build the artifact and return it
	artifact := &awscommon.Artifact{
		Amis:           state.Get("amis").(map[string]string),
		BuilderIdValue: BuilderId,
		Session:        session,
	}

	return artifact, nil
}
<|endoftext|>


<commit_before>package softlayer

import (
	"fmt"
	"log"
)

// Artifact represents a Softlayer image as the result of a Packer build.
type Artifact struct {
	imageName      string
	imageId        string
	datacenterName string
	client         *SoftlayerClient
}

// BuilderId returns the builder Id.
func (*Artifact) BuilderId() string {
	return BuilderId
}

// Destroy destroys the Softlayer image represented by the artifact.
func (self *Artifact) Destroy() error {
	log.Printf("Destroying image: %s", self.String())
	err := self.client.destroyImage(self.imageId)
	return err
}

// Files returns the files represented by the artifact.
func (*Artifact) Files() []string {
	return nil
}

// Id returns the Softlayer image ID.
func (self *Artifact) Id() string {
	return self.imageId
}

// String returns the string representation of the artifact.
func (self *Artifact) String() string {
	return fmt.Sprintf("%s::%s (%s)", self.datacenterName, self.imageId, self.imageName)
}
<commit_msg>Add State to Artifact struct to support the new packer api (v0.7.2)<commit_after>package softlayer

import (
	"fmt"
	"log"
)

// Artifact represents a Softlayer image as the result of a Packer build.
type Artifact struct {
	imageName      string
	imageId        string
	datacenterName string
	client         *SoftlayerClient
}

// BuilderId returns the builder Id.
func (*Artifact) BuilderId() string {
	return BuilderId
}

// Destroy destroys the Softlayer image represented by the artifact.
func (self *Artifact) Destroy() error {
	log.Printf("Destroying image: %s", self.String())
	err := self.client.destroyImage(self.imageId)
	return err
}

// Files returns the files represented by the artifact.
func (*Artifact) Files() []string {
	return nil
}

// Id returns the Softlayer image ID.
func (self *Artifact) Id() string {
	return self.imageId
}

func (self *Artifact) State(name string) interface{} {
	return nil
}

// String returns the string representation of the artifact.
func (self *Artifact) String() string {
	return fmt.Sprintf("%s::%s (%s)", self.datacenterName, self.imageId, self.imageName)
}
<|endoftext|>


<commit_before>package file

import (
	"fmt"
	"os"
	"sync"

	"github.com/hashicorp/vault/audit"
	"github.com/hashicorp/vault/logical"
)

func Factory(conf map[string]string) (audit.Backend, error) {
	path, ok := conf["path"]
	if !ok {
		return nil, fmt.Errorf("path is required")
	}

	return &Backend{Path: path}, nil
}

// Backend is the audit backend for the file-based audit store.
//
// NOTE: This audit backend is currently very simple: it appends to a file.
// It doesn't do anything more at the moment to assist with rotation
// or reset the write cursor, this should be done in the future.
type Backend struct {
	Path string

	once sync.Once
	f    *os.File
}

func (b *Backend) LogRequest(auth *logical.Auth, req *logical.Request) error {
	if err := b.open(); err != nil {
		return err
	}

	// TODO
	return nil
}

func (b *Backend) LogResponse(
	auth *logical.Auth,
	req *logical.Request,
	resp *logical.Response,
	err error) error {
	if err := b.open(); err != nil {
		return err
	}

	// TODO
	return nil
}

func (b *Backend) open() error {
	if b.f != nil {
		return nil
	}

	var err error
	b.f, err = os.Create(b.Path)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>audit/file: use JSON formatter to write output<commit_after>package file

import (
	"fmt"
	"os"
	"sync"

	"github.com/hashicorp/vault/audit"
	"github.com/hashicorp/vault/logical"
)

func Factory(conf map[string]string) (audit.Backend, error) {
	path, ok := conf["path"]
	if !ok {
		return nil, fmt.Errorf("path is required")
	}

	return &Backend{Path: path}, nil
}

// Backend is the audit backend for the file-based audit store.
//
// NOTE: This audit backend is currently very simple: it appends to a file.
// It doesn't do anything more at the moment to assist with rotation
// or reset the write cursor, this should be done in the future.
type Backend struct {
	Path string

	once sync.Once
	f    *os.File
}

func (b *Backend) LogRequest(auth *logical.Auth, req *logical.Request) error {
	if err := b.open(); err != nil {
		return err
	}

	var format audit.FormatJSON
	return format.FormatRequest(b.f, auth, req)
}

func (b *Backend) LogResponse(
	auth *logical.Auth,
	req *logical.Request,
	resp *logical.Response,
	err error) error {
	if err := b.open(); err != nil {
		return err
	}

	var format audit.FormatJSON
	return format.FormatResponse(b.f, auth, req, resp, err)
}

func (b *Backend) open() error {
	if b.f != nil {
		return nil
	}

	var err error
	b.f, err = os.Create(b.Path)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package replicaset

import (
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	// MaxPeers defines the maximum number of peers that mongo supports.
	MaxPeers = 7

	// maxInitiateAttempts is the maximum number of times to attempt
	// replSetInitiate for each call to Initiate.
	maxInitiateAttempts = 10

	// initiateAttemptDelay is the amount of time to sleep between failed
	// attempts to replSetInitiate.
	initiateAttemptDelay = 100 * time.Millisecond

	// rsMembersUnreachableError is the error message returned from mongo
	// when it thinks that replicaset members are unreachable. This can
	// occur if replSetInitiate is executed shortly after starting up mongo.
	rsMembersUnreachableError = "all members and seeds must be reachable to initiate set"
)

var logger = loggo.GetLogger("juju.replicaset")

// Initiate sets up a replica set with the given replica set name with the
// single given member.  It need be called only once for a given mongo replica
// set.  The tags specified will be added as tags on the member that is created
// in the replica set.
//
// Note that you must set DialWithInfo and set Direct = true when dialing into a
// specific non-initiated mongo server.
//
// See http://docs.mongodb.org/manual/reference/method/rs.initiate/ for more
// details.
func Initiate(session *mgo.Session, address, name string, tags map[string]string) error {
	monotonicSession := session.Clone()
	monotonicSession.SetMode(mgo.Monotonic, true)
	cfg := Config{
		Name:    name,
		Version: 1,
		Members: []Member{{
			Id:      1,
			Address: address,
			Tags:    tags,
		}},
	}
	logger.Infof("Initiating replicaset with config %#v", cfg)
	var err error
	for i := 0; i < maxInitiateAttempts; i++ {
		err = monotonicSession.Run(bson.D{{"replSetInitiate", cfg}}, nil)
		if err != nil && err.Error() == rsMembersUnreachableError {
			time.Sleep(initiateAttemptDelay)
			continue
		}
		break
	}
	return err
}

// Member holds configuration information for a replica set member.
//
// See http://docs.mongodb.org/manual/reference/replica-configuration/
// for more details
type Member struct {
	// Id is a unique id for a member in a set.
	Id int `bson:"_id"`

	// Address holds the network address of the member,
	// in the form hostname:port.
	Address string `bson:"host"`

	// Arbiter holds whether the member is an arbiter only.
	// This value is optional; it defaults to false.
	Arbiter *bool `bson:"arbiterOnly,omitempty"`

	// BuildIndexes determines whether the mongod builds indexes on this member.
	// This value is optional; it defaults to true.
	BuildIndexes *bool `bson:"buildIndexes,omitempty"`

	// Hidden determines whether the replica set hides this member from
	// the output of IsMaster.
	// This value is optional; it defaults to false.
	Hidden *bool `bson:"hidden,omitempty"`

	// Priority determines eligibility of a member to become primary.
	// This value is optional; it defaults to 1.
	Priority *float64 `bson:"priority,omitempty"`

	// Tags store additional information about a replica member, often used for
	// customizing read preferences and write concern.
	Tags map[string]string `bson:"tags,omitempty"`

	// SlaveDelay describes the number of seconds behind the master that this
	// replica set member should lag rounded up to the nearest second.
	// This value is optional; it defaults to 0.
	SlaveDelay *time.Duration `bson:"slaveDelay,omitempty"`

	// Votes controls the number of votes a server has in a replica set election.
	// This value is optional; it defaults to 1.
	Votes *int `bson:"votes,omitempty"`
}

func fmtConfigForLog(config *Config) string {
	memberInfo := make([]string, len(config.Members))
	for i, member := range config.Members {
		memberInfo[i] = fmt.Sprintf("Member{%d %q %v}", member.Id, member.Address, member.Tags)

	}
	return fmt.Sprintf("{Name: %s, Version: %d, Members: {%s}}",
		config.Name, config.Version, strings.Join(memberInfo, ", "))
}

// applyRelSetConfig applies the new config to the mongo session. It also logs
// what the changes are. It checks if the replica set changes cause the DB
// connection to be dropped. If so, it Refreshes the session and tries to Ping
// again.
func applyRelSetConfig(cmd string, session *mgo.Session, oldconfig, newconfig *Config) error {
	logger.Debugf("%s() changing replica set\nfrom %s\n  to %s",
		cmd, fmtConfigForLog(oldconfig), fmtConfigForLog(newconfig))
	err := session.Run(bson.D{{"replSetReconfig", newconfig}}, nil)
	// We will only try to Ping 2 times
	for i := 0; i < 2; i++ {
		if err == io.EOF {
			// If the primary changes due to replSetReconfig, then all
			// current connections are dropped.
			// Refreshing should fix us up.
			logger.Debugf("got EOF while running %s(), calling session.Refresh()", cmd)
			session.Refresh()
		} else if err != nil {
			// For all errors that aren't EOF, return immediately
			return err
		}
		// err is either nil or EOF and we called Refresh, so Ping to
		// make sure we're actually connected
		err = session.Ping()
		// Change the command because it is the new command we ran
		cmd = "Ping"
	}
	return err
}

// Add adds the given members to the session's replica set.  Duplicates of
// existing replicas will be ignored.
//
// Members will have their Ids set automatically if they are not already > 0
func Add(session *mgo.Session, members ...Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	oldconfig := *config
	config.Version++
	max := 0
	for _, member := range config.Members {
		if member.Id > max {
			max = member.Id
		}
	}

outerLoop:
	for _, newMember := range members {
		for _, member := range config.Members {
			if member.Address == newMember.Address {
				// already exists, skip it
				continue outerLoop
			}
		}
		// let the caller specify an id if they want, treat zero as unspecified
		if newMember.Id < 1 {
			max++
			newMember.Id = max
		}
		config.Members = append(config.Members, newMember)
	}
	return applyRelSetConfig("Add", session, &oldconfig, config)
}

// Remove removes members with the given addresses from the replica set. It is
// not an error to remove addresses of non-existent replica set members.
func Remove(session *mgo.Session, addrs ...string) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}
	oldconfig := *config
	config.Version++
	for _, rem := range addrs {
		for n, repl := range config.Members {
			if repl.Address == rem {
				config.Members = append(config.Members[:n], config.Members[n+1:]...)
				break
			}
		}
	}
	return applyRelSetConfig("Remove", session, &oldconfig, config)
}

// Set changes the current set of replica set members.  Members will have their
// ids set automatically if their ids are not already > 0.
func Set(session *mgo.Session, members []Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	// Copy the current configuration for logging
	oldconfig := *config
	config.Version++

	// Assign ids to members that did not previously exist, starting above the
	// value of the highest id that already existed
	ids := map[string]int{}
	max := 0
	for _, m := range config.Members {
		ids[m.Address] = m.Id
		if m.Id > max {
			max = m.Id
		}
	}

	for x, m := range members {
		if id, ok := ids[m.Address]; ok {
			m.Id = id
		} else if m.Id < 1 {
			max++
			m.Id = max
		}
		members[x] = m
	}

	config.Members = members

	return applyRelSetConfig("Set", session, &oldconfig, config)
}

// Config reports information about the configuration of a given mongo node
type IsMasterResults struct {
	// The following fields hold information about the specific mongodb node.
	IsMaster  bool      `bson:"ismaster"`
	Secondary bool      `bson:"secondary"`
	Arbiter   bool      `bson:"arbiterOnly"`
	Address   string    `bson:"me"`
	LocalTime time.Time `bson:"localTime"`

	// The following fields hold information about the replica set.
	ReplicaSetName string   `bson:"setName"`
	Addresses      []string `bson:"hosts"`
	Arbiters       []string `bson:"arbiters"`
	PrimaryAddress string   `bson:"primary"`
}

// IsMaster returns information about the configuration of the node that
// the given session is connected to.
func IsMaster(session *mgo.Session) (*IsMasterResults, error) {
	results := &IsMasterResults{}
	err := session.Run("isMaster", results)
	if err != nil {
		return nil, err
	}
	return results, nil
}

var ErrMasterNotConfigured = fmt.Errorf("mongo master not configured")

// MasterHostPort returns the "address:port" string for the primary
// mongo server in the replicaset. It returns ErrMasterNotConfigured if
// the replica set has not yet been initiated.
func MasterHostPort(session *mgo.Session) (string, error) {
	results, err := IsMaster(session)
	if err != nil {
		return "", err
	}
	if results.PrimaryAddress == "" {
		return "", ErrMasterNotConfigured
	}
	return results.PrimaryAddress, nil
}

// CurrentMembers returns the current members of the replica set.
func CurrentMembers(session *mgo.Session) ([]Member, error) {
	cfg, err := CurrentConfig(session)
	if err != nil {
		return nil, err
	}
	return cfg.Members, nil
}

// CurrentConfig returns the Config for the given session's replica set.  If
// there is no current config, the error returned will be mgo.ErrNotFound.
func CurrentConfig(session *mgo.Session) (*Config, error) {
	cfg := &Config{}
	monotonicSession := session.Clone()
	monotonicSession.SetMode(mgo.Monotonic, true)
	err := monotonicSession.DB("local").C("system.replset").Find(nil).One(cfg)
	if err == mgo.ErrNotFound {
		return nil, err
	}
	if err != nil {
		return nil, fmt.Errorf("cannot get replset config: %s", err.Error())
	}
	return cfg, nil
}

// Config is the document stored in mongodb that defines the servers in the
// replica set
type Config struct {
	Name    string   `bson:"_id"`
	Version int      `bson:"version"`
	Members []Member `bson:"members"`
}

// CurrentStatus returns the status of the replica set for the given session.
func CurrentStatus(session *mgo.Session) (*Status, error) {
	status := &Status{}
	err := session.Run("replSetGetStatus", status)
	if err != nil {
		return nil, fmt.Errorf("cannot get replica set status: %v", err)
	}
	return status, nil
}

// Status holds data about the status of members of the replica set returned
// from replSetGetStatus
//
// See http://docs.mongodb.org/manual/reference/command/replSetGetStatus/#dbcmd.replSetGetStatus
type Status struct {
	Name    string         `bson:"set"`
	Members []MemberStatus `bson:"members"`
}

// Status holds the status of a replica set member returned from
// replSetGetStatus.
type MemberStatus struct {
	// Id holds the replica set id of the member that the status is describing.
	Id int `bson:"_id"`

	// Address holds address of the member that the status is describing.
	Address string `bson:"name"`

	// Self holds whether this is the status for the member that
	// the session is connected to.
	Self bool `bson:"self"`

	// ErrMsg holds the most recent error or status message received
	// from the member.
	ErrMsg string `bson:"errmsg"`

	// Healthy reports whether the member is up. It is true for the
	// member that the request was made to.
	Healthy bool `bson:"health"`

	// State describes the current state of the member.
	State MemberState `bson:"state"`

	// Uptime describes how long the member has been online.
	Uptime time.Duration `bson:"uptime"`

	// Ping describes the length of time a round-trip packet takes to travel
	// between the remote member and the local instance.  It is zero for the
	// member that the session is connected to.
	Ping time.Duration `bson:"pingMS"`
}

// MemberState represents the state of a replica set member.
// See http://docs.mongodb.org/manual/reference/replica-states/
type MemberState int

const (
	StartupState = iota
	PrimaryState
	SecondaryState
	RecoveringState
	FatalState
	Startup2State
	UnknownState
	ArbiterState
	DownState
	RollbackState
	ShunnedState
)

var memberStateStrings = []string{
	StartupState:    "STARTUP",
	PrimaryState:    "PRIMARY",
	SecondaryState:  "SECONDARY",
	RecoveringState: "RECOVERING",
	FatalState:      "FATAL",
	Startup2State:   "STARTUP2",
	UnknownState:    "UNKNOWN",
	ArbiterState:    "ARBITER",
	DownState:       "DOWN",
	RollbackState:   "ROLLBACK",
	ShunnedState:    "SHUNNED",
}

// String returns a string describing the state.
func (state MemberState) String() string {
	if state < 0 || int(state) >= len(memberStateStrings) {
		return "INVALID_MEMBER_STATE"
	}
	return memberStateStrings[state]
}
<commit_msg>[r=axwalk] replicaset: close cloned sessions<commit_after>package replicaset

import (
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	// MaxPeers defines the maximum number of peers that mongo supports.
	MaxPeers = 7

	// maxInitiateAttempts is the maximum number of times to attempt
	// replSetInitiate for each call to Initiate.
	maxInitiateAttempts = 10

	// initiateAttemptDelay is the amount of time to sleep between failed
	// attempts to replSetInitiate.
	initiateAttemptDelay = 100 * time.Millisecond

	// rsMembersUnreachableError is the error message returned from mongo
	// when it thinks that replicaset members are unreachable. This can
	// occur if replSetInitiate is executed shortly after starting up mongo.
	rsMembersUnreachableError = "all members and seeds must be reachable to initiate set"
)

var logger = loggo.GetLogger("juju.replicaset")

// Initiate sets up a replica set with the given replica set name with the
// single given member.  It need be called only once for a given mongo replica
// set.  The tags specified will be added as tags on the member that is created
// in the replica set.
//
// Note that you must set DialWithInfo and set Direct = true when dialing into a
// specific non-initiated mongo server.
//
// See http://docs.mongodb.org/manual/reference/method/rs.initiate/ for more
// details.
func Initiate(session *mgo.Session, address, name string, tags map[string]string) error {
	monotonicSession := session.Clone()
	defer monotonicSession.Close()
	monotonicSession.SetMode(mgo.Monotonic, true)
	cfg := Config{
		Name:    name,
		Version: 1,
		Members: []Member{{
			Id:      1,
			Address: address,
			Tags:    tags,
		}},
	}
	logger.Infof("Initiating replicaset with config %#v", cfg)
	var err error
	for i := 0; i < maxInitiateAttempts; i++ {
		err = monotonicSession.Run(bson.D{{"replSetInitiate", cfg}}, nil)
		if err != nil && err.Error() == rsMembersUnreachableError {
			time.Sleep(initiateAttemptDelay)
			continue
		}
		break
	}
	return err
}

// Member holds configuration information for a replica set member.
//
// See http://docs.mongodb.org/manual/reference/replica-configuration/
// for more details
type Member struct {
	// Id is a unique id for a member in a set.
	Id int `bson:"_id"`

	// Address holds the network address of the member,
	// in the form hostname:port.
	Address string `bson:"host"`

	// Arbiter holds whether the member is an arbiter only.
	// This value is optional; it defaults to false.
	Arbiter *bool `bson:"arbiterOnly,omitempty"`

	// BuildIndexes determines whether the mongod builds indexes on this member.
	// This value is optional; it defaults to true.
	BuildIndexes *bool `bson:"buildIndexes,omitempty"`

	// Hidden determines whether the replica set hides this member from
	// the output of IsMaster.
	// This value is optional; it defaults to false.
	Hidden *bool `bson:"hidden,omitempty"`

	// Priority determines eligibility of a member to become primary.
	// This value is optional; it defaults to 1.
	Priority *float64 `bson:"priority,omitempty"`

	// Tags store additional information about a replica member, often used for
	// customizing read preferences and write concern.
	Tags map[string]string `bson:"tags,omitempty"`

	// SlaveDelay describes the number of seconds behind the master that this
	// replica set member should lag rounded up to the nearest second.
	// This value is optional; it defaults to 0.
	SlaveDelay *time.Duration `bson:"slaveDelay,omitempty"`

	// Votes controls the number of votes a server has in a replica set election.
	// This value is optional; it defaults to 1.
	Votes *int `bson:"votes,omitempty"`
}

func fmtConfigForLog(config *Config) string {
	memberInfo := make([]string, len(config.Members))
	for i, member := range config.Members {
		memberInfo[i] = fmt.Sprintf("Member{%d %q %v}", member.Id, member.Address, member.Tags)

	}
	return fmt.Sprintf("{Name: %s, Version: %d, Members: {%s}}",
		config.Name, config.Version, strings.Join(memberInfo, ", "))
}

// applyRelSetConfig applies the new config to the mongo session. It also logs
// what the changes are. It checks if the replica set changes cause the DB
// connection to be dropped. If so, it Refreshes the session and tries to Ping
// again.
func applyRelSetConfig(cmd string, session *mgo.Session, oldconfig, newconfig *Config) error {
	logger.Debugf("%s() changing replica set\nfrom %s\n  to %s",
		cmd, fmtConfigForLog(oldconfig), fmtConfigForLog(newconfig))
	err := session.Run(bson.D{{"replSetReconfig", newconfig}}, nil)
	// We will only try to Ping 2 times
	for i := 0; i < 2; i++ {
		if err == io.EOF {
			// If the primary changes due to replSetReconfig, then all
			// current connections are dropped.
			// Refreshing should fix us up.
			logger.Debugf("got EOF while running %s(), calling session.Refresh()", cmd)
			session.Refresh()
		} else if err != nil {
			// For all errors that aren't EOF, return immediately
			return err
		}
		// err is either nil or EOF and we called Refresh, so Ping to
		// make sure we're actually connected
		err = session.Ping()
		// Change the command because it is the new command we ran
		cmd = "Ping"
	}
	return err
}

// Add adds the given members to the session's replica set.  Duplicates of
// existing replicas will be ignored.
//
// Members will have their Ids set automatically if they are not already > 0
func Add(session *mgo.Session, members ...Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	oldconfig := *config
	config.Version++
	max := 0
	for _, member := range config.Members {
		if member.Id > max {
			max = member.Id
		}
	}

outerLoop:
	for _, newMember := range members {
		for _, member := range config.Members {
			if member.Address == newMember.Address {
				// already exists, skip it
				continue outerLoop
			}
		}
		// let the caller specify an id if they want, treat zero as unspecified
		if newMember.Id < 1 {
			max++
			newMember.Id = max
		}
		config.Members = append(config.Members, newMember)
	}
	return applyRelSetConfig("Add", session, &oldconfig, config)
}

// Remove removes members with the given addresses from the replica set. It is
// not an error to remove addresses of non-existent replica set members.
func Remove(session *mgo.Session, addrs ...string) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}
	oldconfig := *config
	config.Version++
	for _, rem := range addrs {
		for n, repl := range config.Members {
			if repl.Address == rem {
				config.Members = append(config.Members[:n], config.Members[n+1:]...)
				break
			}
		}
	}
	return applyRelSetConfig("Remove", session, &oldconfig, config)
}

// Set changes the current set of replica set members.  Members will have their
// ids set automatically if their ids are not already > 0.
func Set(session *mgo.Session, members []Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	// Copy the current configuration for logging
	oldconfig := *config
	config.Version++

	// Assign ids to members that did not previously exist, starting above the
	// value of the highest id that already existed
	ids := map[string]int{}
	max := 0
	for _, m := range config.Members {
		ids[m.Address] = m.Id
		if m.Id > max {
			max = m.Id
		}
	}

	for x, m := range members {
		if id, ok := ids[m.Address]; ok {
			m.Id = id
		} else if m.Id < 1 {
			max++
			m.Id = max
		}
		members[x] = m
	}

	config.Members = members

	return applyRelSetConfig("Set", session, &oldconfig, config)
}

// Config reports information about the configuration of a given mongo node
type IsMasterResults struct {
	// The following fields hold information about the specific mongodb node.
	IsMaster  bool      `bson:"ismaster"`
	Secondary bool      `bson:"secondary"`
	Arbiter   bool      `bson:"arbiterOnly"`
	Address   string    `bson:"me"`
	LocalTime time.Time `bson:"localTime"`

	// The following fields hold information about the replica set.
	ReplicaSetName string   `bson:"setName"`
	Addresses      []string `bson:"hosts"`
	Arbiters       []string `bson:"arbiters"`
	PrimaryAddress string   `bson:"primary"`
}

// IsMaster returns information about the configuration of the node that
// the given session is connected to.
func IsMaster(session *mgo.Session) (*IsMasterResults, error) {
	results := &IsMasterResults{}
	err := session.Run("isMaster", results)
	if err != nil {
		return nil, err
	}
	return results, nil
}

var ErrMasterNotConfigured = fmt.Errorf("mongo master not configured")

// MasterHostPort returns the "address:port" string for the primary
// mongo server in the replicaset. It returns ErrMasterNotConfigured if
// the replica set has not yet been initiated.
func MasterHostPort(session *mgo.Session) (string, error) {
	results, err := IsMaster(session)
	if err != nil {
		return "", err
	}
	if results.PrimaryAddress == "" {
		return "", ErrMasterNotConfigured
	}
	return results.PrimaryAddress, nil
}

// CurrentMembers returns the current members of the replica set.
func CurrentMembers(session *mgo.Session) ([]Member, error) {
	cfg, err := CurrentConfig(session)
	if err != nil {
		return nil, err
	}
	return cfg.Members, nil
}

// CurrentConfig returns the Config for the given session's replica set.  If
// there is no current config, the error returned will be mgo.ErrNotFound.
func CurrentConfig(session *mgo.Session) (*Config, error) {
	cfg := &Config{}
	monotonicSession := session.Clone()
	defer monotonicSession.Close()
	monotonicSession.SetMode(mgo.Monotonic, true)
	err := monotonicSession.DB("local").C("system.replset").Find(nil).One(cfg)
	if err == mgo.ErrNotFound {
		return nil, err
	}
	if err != nil {
		return nil, fmt.Errorf("cannot get replset config: %s", err.Error())
	}
	return cfg, nil
}

// Config is the document stored in mongodb that defines the servers in the
// replica set
type Config struct {
	Name    string   `bson:"_id"`
	Version int      `bson:"version"`
	Members []Member `bson:"members"`
}

// CurrentStatus returns the status of the replica set for the given session.
func CurrentStatus(session *mgo.Session) (*Status, error) {
	status := &Status{}
	err := session.Run("replSetGetStatus", status)
	if err != nil {
		return nil, fmt.Errorf("cannot get replica set status: %v", err)
	}
	return status, nil
}

// Status holds data about the status of members of the replica set returned
// from replSetGetStatus
//
// See http://docs.mongodb.org/manual/reference/command/replSetGetStatus/#dbcmd.replSetGetStatus
type Status struct {
	Name    string         `bson:"set"`
	Members []MemberStatus `bson:"members"`
}

// Status holds the status of a replica set member returned from
// replSetGetStatus.
type MemberStatus struct {
	// Id holds the replica set id of the member that the status is describing.
	Id int `bson:"_id"`

	// Address holds address of the member that the status is describing.
	Address string `bson:"name"`

	// Self holds whether this is the status for the member that
	// the session is connected to.
	Self bool `bson:"self"`

	// ErrMsg holds the most recent error or status message received
	// from the member.
	ErrMsg string `bson:"errmsg"`

	// Healthy reports whether the member is up. It is true for the
	// member that the request was made to.
	Healthy bool `bson:"health"`

	// State describes the current state of the member.
	State MemberState `bson:"state"`

	// Uptime describes how long the member has been online.
	Uptime time.Duration `bson:"uptime"`

	// Ping describes the length of time a round-trip packet takes to travel
	// between the remote member and the local instance.  It is zero for the
	// member that the session is connected to.
	Ping time.Duration `bson:"pingMS"`
}

// MemberState represents the state of a replica set member.
// See http://docs.mongodb.org/manual/reference/replica-states/
type MemberState int

const (
	StartupState = iota
	PrimaryState
	SecondaryState
	RecoveringState
	FatalState
	Startup2State
	UnknownState
	ArbiterState
	DownState
	RollbackState
	ShunnedState
)

var memberStateStrings = []string{
	StartupState:    "STARTUP",
	PrimaryState:    "PRIMARY",
	SecondaryState:  "SECONDARY",
	RecoveringState: "RECOVERING",
	FatalState:      "FATAL",
	Startup2State:   "STARTUP2",
	UnknownState:    "UNKNOWN",
	ArbiterState:    "ARBITER",
	DownState:       "DOWN",
	RollbackState:   "ROLLBACK",
	ShunnedState:    "SHUNNED",
}

// String returns a string describing the state.
func (state MemberState) String() string {
	if state < 0 || int(state) >= len(memberStateStrings) {
		return "INVALID_MEMBER_STATE"
	}
	return memberStateStrings[state]
}
<|endoftext|>