<commit_before>package util

import (
	"fmt"
	"github.com/codegangsta/cli"
)

func (u *SomaUtil) GetCliArgumentCount(c *cli.Context) int {
	a := c.Args()
	if !a.Present() {
		return 0
	}
	return len(a.Tail()) + 1
}

func (u *SomaUtil) ValidateCliArgument(c *cli.Context, pos uint8, s string) {
	a := c.Args()
	if a.Get(int(pos)-1) != s {
		u.Abort(fmt.Sprintf("Syntax error, missing keyword: ", s))
	}
}

func (u *SomaUtil) ValidateCliArgumentCount(c *cli.Context, i uint8) {
	a := c.Args()
	if i == 0 {
		if a.Present() {
			u.Abort("Syntax error, command takes no arguments")
		}
	} else {
		if !a.Present() || len(a.Tail()) != (int(i)-1) {
			u.Abort("Syntax error")
		}
	}
}

func (u *SomaUtil) GetFullArgumentSlice(c *cli.Context) []string {
	sl := []string{c.Args().First()}
	sl = append(sl, c.Args().Tail()...)
	return sl
}

func (u *SomaUtil) ParseVariableArguments(keys []string, rKeys []string, args []string) (map[string]string, []string) {
	// return map of the parse result
	result := make(map[string]string)
	// map to test which required keys were found
	argumentCheck := make(map[string]bool)
	// return slice which optional keys were found
	optionalKeys := make([]string, 0)
	// no required keys is valid
	if len(rKeys) > 0 {
		for _, key := range rKeys {
			argumentCheck[key] = false
		}
	}
	skipNext := false

	for pos, val := range args {
		// skip current argument if last argument was a keyword
		if skipNext {
			skipNext = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check back-to-back keywords
			u.CheckStringNotAKeyword(args[pos+1], keys)
			result[val] = args[pos+1]
			argumentCheck[val] = true
			skipNext = true
			if !u.SliceContainsString(val, rKeys) {
				optionalKeys = append(optionalKeys, val)
			}
			continue
		}
		// keywords trigger continue, arguments are skipped over.
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check we managed to collect all required keywords
	for _, v := range argumentCheck {
		if !v {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", v))
		}
	}

	return result, optionalKeys
}

/*
 * This function parses whitespace separated argument lists of
 * keyword/value pairs were keywords can be specified multiple
 * times, some keywords are required and some only allowed once.
 * Sequence of multiple keywords are detected and lead to abort
 *
 * multKeys => [ "port", "transport" ]
 * uniqKeys => [ "team" ]
 * reqKeys  => [ "team" ]
 * args     => [ "port", "53", "transport", "tcp", "transport",
 *               "udp", "team", "ITOMI" ]
 *
 * result => result["team"] = [ "ITOMI" ]
 *           result["port"] = [ "53" ]
 *           result["transport"] = [ "tcp", "udp" ]
 */
func (u *SomaUtil) ParseVariadicArguments(
	multKeys []string, // keys that may appear multiple times
	uniqKeys []string, // keys that are allowed at most once
	reqKeys []string, // keys that are required at least one
	args []string, // arguments to parse
) map[string][]string {
	// returns a map of slices of string
	result := make(map[string][]string)

	// merge key slices
	keys := append(multKeys, uniqKeys...)

	// helper to skip over next value in args slice
	skip := false

	for pos, val := range args {
		// skip current arg if last argument was a keyword
		if skip {
			skip = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check for back-to-back keyswords
			u.CheckStringNotAKeyword(args[pos+1], keys)

			// append value of current keyword into result map
			result[val] = append(result[val], args[pos+1])
			skip = true
			continue
		}
		// keywords trigger continue before this
		// values after keywords are skip'ed
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if we managed to collect all required keywords
	for _, key := range reqKeys {
		_, ok := result[key] // ok is false if slice is nil
		if !ok {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specified once
	for _, key := range uniqKeys {
		sl, ok := result[key]
		if ok && (len(sl) > 1) {
			u.Abort(fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

	return result
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Improve argument validation error msg<commit_after>package util

import (
	"fmt"
	"github.com/codegangsta/cli"
)

func (u *SomaUtil) GetCliArgumentCount(c *cli.Context) int {
	a := c.Args()
	if !a.Present() {
		return 0
	}
	return len(a.Tail()) + 1
}

func (u *SomaUtil) ValidateCliArgument(c *cli.Context, pos uint8, s string) {
	a := c.Args()
	if a.Get(int(pos)-1) != s {
		u.Abort(fmt.Sprintf("Syntax error, missing keyword: ", s))
	}
}

func (u *SomaUtil) ValidateCliArgumentCount(c *cli.Context, i uint8) {
	a := c.Args()
	if i == 0 {
		if a.Present() {
			u.Abort("Syntax error, command takes no arguments")
		}
	} else {
		if !a.Present() || len(a.Tail()) != (int(i)-1) {
			u.Abort("Syntax error, incorrect argument count")
		}
	}
}

func (u *SomaUtil) GetFullArgumentSlice(c *cli.Context) []string {
	sl := []string{c.Args().First()}
	sl = append(sl, c.Args().Tail()...)
	return sl
}

func (u *SomaUtil) ParseVariableArguments(keys []string, rKeys []string, args []string) (map[string]string, []string) {
	// return map of the parse result
	result := make(map[string]string)
	// map to test which required keys were found
	argumentCheck := make(map[string]bool)
	// return slice which optional keys were found
	optionalKeys := make([]string, 0)
	// no required keys is valid
	if len(rKeys) > 0 {
		for _, key := range rKeys {
			argumentCheck[key] = false
		}
	}
	skipNext := false

	for pos, val := range args {
		// skip current argument if last argument was a keyword
		if skipNext {
			skipNext = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check back-to-back keywords
			u.CheckStringNotAKeyword(args[pos+1], keys)
			result[val] = args[pos+1]
			argumentCheck[val] = true
			skipNext = true
			if !u.SliceContainsString(val, rKeys) {
				optionalKeys = append(optionalKeys, val)
			}
			continue
		}
		// keywords trigger continue, arguments are skipped over.
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check we managed to collect all required keywords
	for _, v := range argumentCheck {
		if !v {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", v))
		}
	}

	return result, optionalKeys
}

/*
 * This function parses whitespace separated argument lists of
 * keyword/value pairs were keywords can be specified multiple
 * times, some keywords are required and some only allowed once.
 * Sequence of multiple keywords are detected and lead to abort
 *
 * multKeys => [ "port", "transport" ]
 * uniqKeys => [ "team" ]
 * reqKeys  => [ "team" ]
 * args     => [ "port", "53", "transport", "tcp", "transport",
 *               "udp", "team", "ITOMI" ]
 *
 * result => result["team"] = [ "ITOMI" ]
 *           result["port"] = [ "53" ]
 *           result["transport"] = [ "tcp", "udp" ]
 */
func (u *SomaUtil) ParseVariadicArguments(
	multKeys []string, // keys that may appear multiple times
	uniqKeys []string, // keys that are allowed at most once
	reqKeys []string, // keys that are required at least one
	args []string, // arguments to parse
) map[string][]string {
	// returns a map of slices of string
	result := make(map[string][]string)

	// merge key slices
	keys := append(multKeys, uniqKeys...)

	// helper to skip over next value in args slice
	skip := false

	for pos, val := range args {
		// skip current arg if last argument was a keyword
		if skip {
			skip = false
			continue
		}

		if u.SliceContainsString(val, keys) {
			// check for back-to-back keyswords
			u.CheckStringNotAKeyword(args[pos+1], keys)

			// append value of current keyword into result map
			result[val] = append(result[val], args[pos+1])
			skip = true
			continue
		}
		// keywords trigger continue before this
		// values after keywords are skip'ed
		// reaching this is an error
		u.Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if we managed to collect all required keywords
	for _, key := range reqKeys {
		_, ok := result[key] // ok is false if slice is nil
		if !ok {
			u.Abort(fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specified once
	for _, key := range uniqKeys {
		sl, ok := result[key]
		if ok && (len(sl) > 1) {
			u.Abort(fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

	return result
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package lock

import (
	"fmt"
	"hash/crc32"
	"strconv"
	"strings"
	"sync"

	"code.cloudfoundry.org/lager"
)

const (
	LockTypeResourceConfigChecking = iota
	LockTypeResourceTypeChecking
	LockTypeBuildTracking
	LockTypePipelineScheduling
	LockTypeResourceCheckingForJob
	LockTypeBatch
	LockTypeVolumeCreating
	LockTypeContainerCreating
)

func NewBuildTrackingLockID(buildID int) LockID {
	return LockID{LockTypeBuildTracking, buildID}
}

func NewResourceConfigCheckingLockID(resourceConfigID int) LockID {
	return LockID{LockTypeResourceConfigChecking, resourceConfigID}
}

func NewResourceTypeCheckingLockID(resourceTypeID int) LockID {
	return LockID{LockTypeResourceTypeChecking, resourceTypeID}
}

func NewPipelineSchedulingLockLockID(pipelineID int) LockID {
	return LockID{LockTypePipelineScheduling, pipelineID}
}

func NewResourceCheckingForJobLockID(jobID int) LockID {
	return LockID{LockTypeResourceCheckingForJob, jobID}
}

func NewTaskLockID(taskName string) LockID {
	return LockID{LockTypeBatch, lockIDFromString(taskName)}
}

func NewVolumeCreatingLockID(volumeID int) LockID {
	return LockID{LockTypeVolumeCreating, volumeID}
}

func NewContainerCreatingLockID(containerID int) LockID {
	return LockID{LockTypeContainerCreating, containerID}
}

//go:generate counterfeiter . LockFactory

type LockFactory interface {
	NewLock(logger lager.Logger, ids LockID) Lock
}

type lockFactory struct {
	db           LockDB
	locks        lockRepo
	acquireMutex *sync.Mutex
}

func NewLockFactory(conn *RetryableConn) LockFactory {
	return &lockFactory{
		db: &lockDB{
			conn:  conn,
			mutex: &sync.Mutex{},
		},
		locks: lockRepo{
			locks: map[string]bool{},
			mutex: &sync.Mutex{},
		},
		acquireMutex: &sync.Mutex{},
	}
}

func NewTestLockFactory(db LockDB) LockFactory {
	return &lockFactory{
		db: db,
		locks: lockRepo{
			locks: map[string]bool{},
			mutex: &sync.Mutex{},
		},
		acquireMutex: &sync.Mutex{},
	}
}

func (f *lockFactory) NewLock(logger lager.Logger, id LockID) Lock {
	return &lock{
		logger:       logger,
		db:           f.db,
		id:           id,
		locks:        f.locks,
		acquireMutex: f.acquireMutex,
	}
}

//go:generate counterfeiter . Lock

type Lock interface {
	Acquire() (bool, error)
	Release() error
}

//go:generate counterfeiter . LockDB

type LockDB interface {
	Acquire(id LockID) (bool, error)
	Release(id LockID) error
}

type lock struct {
	id LockID

	logger       lager.Logger
	db           LockDB
	locks        lockRepo
	acquireMutex *sync.Mutex
}

func (l *lock) Acquire() (bool, error) {
	l.acquireMutex.Lock()
	defer l.acquireMutex.Unlock()

	logger := l.logger.Session("acquire", lager.Data{"id": l.id})

	if l.locks.IsRegistered(l.id) {
		logger.Debug("not-acquired-already-held-locally")
		return false, nil
	}

	acquired, err := l.db.Acquire(l.id)
	if err != nil {
		logger.Error("failed-to-register-in-db", err)
		return false, err
	}

	if !acquired {
		logger.Debug("not-acquired-already-held-in-db")
		return false, err
	}

	l.locks.Register(l.id)

	return true, nil
}

func (l *lock) Release() error {
	logger := l.logger.Session("release", lager.Data{"id": l.id})

	err := l.db.Release(l.id)
	if err != nil {
		logger.Error("failed-to-release-in-db-but-continuing-anyway", err)
	}

	l.locks.Unregister(l.id)

	return nil
}

type lockDB struct {
	conn  *RetryableConn
	mutex *sync.Mutex
}

func (db *lockDB) Acquire(id LockID) (bool, error) {
	db.mutex.Lock()
	defer db.mutex.Unlock()

	var acquired bool
	err := db.conn.QueryRow(`SELECT pg_try_advisory_lock(`+id.toDBParams()+`)`, id.toDBArgs()...).Scan(&acquired)
	if err != nil {
		return false, err
	}

	return acquired, nil
}

func (db *lockDB) Release(id LockID) error {
	db.mutex.Lock()
	defer db.mutex.Unlock()

	_, err := db.conn.Exec(`SELECT pg_advisory_unlock(`+id.toDBParams()+`)`, id.toDBArgs()...)
	return err
}

type lockRepo struct {
	locks map[string]bool
	mutex *sync.Mutex
}

func (lr lockRepo) IsRegistered(id LockID) bool {
	lr.mutex.Lock()
	defer lr.mutex.Unlock()

	if _, ok := lr.locks[id.toKey()]; ok {
		return true
	}
	return false
}

func (lr lockRepo) Register(id LockID) {
	lr.mutex.Lock()
	defer lr.mutex.Unlock()

	lr.locks[id.toKey()] = true
}

func (lr lockRepo) Unregister(id LockID) {
	lr.mutex.Lock()
	defer lr.mutex.Unlock()

	delete(lr.locks, id.toKey())
}

type LockID []int

func (l LockID) toKey() string {
	s := []string{}
	for i := range l {
		s = append(s, strconv.Itoa(l[i]))
	}
	return strings.Join(s, "+")
}

func (l LockID) toDBParams() string {
	s := []string{}
	for i := range l {
		s = append(s, fmt.Sprintf("$%d", i+1))
	}

	return strings.Join(s, ",")
}

func (l LockID) toDBArgs() []interface{} {
	result := []interface{}{}
	for i := range l {
		result = append(result, l[i])
	}

	return result
}

func lockIDFromString(taskName string) int {
	return int(int32(crc32.ChecksumIEEE([]byte(taskName))))
}
<commit_msg>re-add simple logs at least<commit_after>package lock

import (
	"fmt"
	"hash/crc32"
	"strconv"
	"strings"
	"sync"

	"code.cloudfoundry.org/lager"
)

const (
	LockTypeResourceConfigChecking = iota
	LockTypeResourceTypeChecking
	LockTypeBuildTracking
	LockTypePipelineScheduling
	LockTypeResourceCheckingForJob
	LockTypeBatch
	LockTypeVolumeCreating
	LockTypeContainerCreating
)

func NewBuildTrackingLockID(buildID int) LockID {
	return LockID{LockTypeBuildTracking, buildID}
}

func NewResourceConfigCheckingLockID(resourceConfigID int) LockID {
	return LockID{LockTypeResourceConfigChecking, resourceConfigID}
}

func NewResourceTypeCheckingLockID(resourceTypeID int) LockID {
	return LockID{LockTypeResourceTypeChecking, resourceTypeID}
}

func NewPipelineSchedulingLockLockID(pipelineID int) LockID {
	return LockID{LockTypePipelineScheduling, pipelineID}
}

func NewResourceCheckingForJobLockID(jobID int) LockID {
	return LockID{LockTypeResourceCheckingForJob, jobID}
}

func NewTaskLockID(taskName string) LockID {
	return LockID{LockTypeBatch, lockIDFromString(taskName)}
}

func NewVolumeCreatingLockID(volumeID int) LockID {
	return LockID{LockTypeVolumeCreating, volumeID}
}

func NewContainerCreatingLockID(containerID int) LockID {
	return LockID{LockTypeContainerCreating, containerID}
}

//go:generate counterfeiter . LockFactory

type LockFactory interface {
	NewLock(logger lager.Logger, ids LockID) Lock
}

type lockFactory struct {
	db           LockDB
	locks        lockRepo
	acquireMutex *sync.Mutex
}

func NewLockFactory(conn *RetryableConn) LockFactory {
	return &lockFactory{
		db: &lockDB{
			conn:  conn,
			mutex: &sync.Mutex{},
		},
		locks: lockRepo{
			locks: map[string]bool{},
			mutex: &sync.Mutex{},
		},
		acquireMutex: &sync.Mutex{},
	}
}

func NewTestLockFactory(db LockDB) LockFactory {
	return &lockFactory{
		db: db,
		locks: lockRepo{
			locks: map[string]bool{},
			mutex: &sync.Mutex{},
		},
		acquireMutex: &sync.Mutex{},
	}
}

func (f *lockFactory) NewLock(logger lager.Logger, id LockID) Lock {
	return &lock{
		logger:       logger,
		db:           f.db,
		id:           id,
		locks:        f.locks,
		acquireMutex: f.acquireMutex,
	}
}

//go:generate counterfeiter . Lock

type Lock interface {
	Acquire() (bool, error)
	Release() error
}

//go:generate counterfeiter . LockDB

type LockDB interface {
	Acquire(id LockID) (bool, error)
	Release(id LockID) error
}

type lock struct {
	id LockID

	logger       lager.Logger
	db           LockDB
	locks        lockRepo
	acquireMutex *sync.Mutex
}

func (l *lock) Acquire() (bool, error) {
	l.acquireMutex.Lock()
	defer l.acquireMutex.Unlock()

	logger := l.logger.Session("acquire", lager.Data{"id": l.id})

	if l.locks.IsRegistered(l.id) {
		logger.Debug("not-acquired-already-held-locally")
		return false, nil
	}

	acquired, err := l.db.Acquire(l.id)
	if err != nil {
		logger.Error("failed-to-register-in-db", err)
		return false, err
	}

	if !acquired {
		logger.Debug("not-acquired-already-held-in-db")
		return false, err
	}

	l.locks.Register(l.id)

	logger.Debug("acquired")

	return true, nil
}

func (l *lock) Release() error {
	logger := l.logger.Session("release", lager.Data{"id": l.id})

	err := l.db.Release(l.id)
	if err != nil {
		logger.Error("failed-to-release-in-db-but-continuing-anyway", err)
	}

	l.locks.Unregister(l.id)

	logger.Debug("released")

	return nil
}

type lockDB struct {
	conn  *RetryableConn
	mutex *sync.Mutex
}

func (db *lockDB) Acquire(id LockID) (bool, error) {
	db.mutex.Lock()
	defer db.mutex.Unlock()

	var acquired bool
	err := db.conn.QueryRow(`SELECT pg_try_advisory_lock(`+id.toDBParams()+`)`, id.toDBArgs()...).Scan(&acquired)
	if err != nil {
		return false, err
	}

	return acquired, nil
}

func (db *lockDB) Release(id LockID) error {
	db.mutex.Lock()
	defer db.mutex.Unlock()

	_, err := db.conn.Exec(`SELECT pg_advisory_unlock(`+id.toDBParams()+`)`, id.toDBArgs()...)
	return err
}

type lockRepo struct {
	locks map[string]bool
	mutex *sync.Mutex
}

func (lr lockRepo) IsRegistered(id LockID) bool {
	lr.mutex.Lock()
	defer lr.mutex.Unlock()

	if _, ok := lr.locks[id.toKey()]; ok {
		return true
	}
	return false
}

func (lr lockRepo) Register(id LockID) {
	lr.mutex.Lock()
	defer lr.mutex.Unlock()

	lr.locks[id.toKey()] = true
}

func (lr lockRepo) Unregister(id LockID) {
	lr.mutex.Lock()
	defer lr.mutex.Unlock()

	delete(lr.locks, id.toKey())
}

type LockID []int

func (l LockID) toKey() string {
	s := []string{}
	for i := range l {
		s = append(s, strconv.Itoa(l[i]))
	}
	return strings.Join(s, "+")
}

func (l LockID) toDBParams() string {
	s := []string{}
	for i := range l {
		s = append(s, fmt.Sprintf("$%d", i+1))
	}

	return strings.Join(s, ",")
}

func (l LockID) toDBArgs() []interface{} {
	result := []interface{}{}
	for i := range l {
		result = append(result, l[i])
	}

	return result
}

func lockIDFromString(taskName string) int {
	return int(int32(crc32.ChecksumIEEE([]byte(taskName))))
}
<|endoftext|>


<commit_before>package db

import (
	"errors"
	_ "github.com/lib/pq"
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestMain(t *testing.T) {

	Convey("db.Results", t, func() {
		query := &mockQuery{2}

		results, err := Results(query)

		So(err, ShouldBeNil)
		So(len(results), ShouldEqual, 2)
	})

	Convey("db.First", t, func() {
		Convey("returns the first record", func() {
			query := &mockQuery{2}
			output, err := First(query)
			So(err, ShouldBeNil)
			So(output.(mockResult), ShouldResemble, mockResult{0})
		})

		Convey("Missing records returns nil", func() {
			query := &mockQuery{0}
			output, err := First(query)
			So(err, ShouldBeNil)
			So(output, ShouldBeNil)
		})

		Convey("Properly forwards non-RecordNotFound errors", func() {
			query := &BrokenQuery{errors.New("Some error")}
			_, err := First(query)

			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldEqual, "Some error")
		})
	})
}
<commit_msg>Add db.First example<commit_after>package db

import (
	"errors"
	_ "github.com/lib/pq"
	. "github.com/smartystreets/goconvey/convey"
	"log"
	"testing"
)

func TestMain(t *testing.T) {

	Convey("db.Results", t, func() {
		query := &mockQuery{2}

		results, err := Results(query)

		So(err, ShouldBeNil)
		So(len(results), ShouldEqual, 2)
	})

	Convey("db.First", t, func() {
		Convey("returns the first record", func() {
			query := &mockQuery{2}
			output, err := First(query)
			So(err, ShouldBeNil)
			So(output.(mockResult), ShouldResemble, mockResult{0})
		})

		Convey("Missing records returns nil", func() {
			query := &mockQuery{0}
			output, err := First(query)
			So(err, ShouldBeNil)
			So(output, ShouldBeNil)
		})

		Convey("Properly forwards non-RecordNotFound errors", func() {
			query := &BrokenQuery{errors.New("Some error")}
			_, err := First(query)

			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldEqual, "Some error")
		})
	})
}

func ExampleFirst() {
	db := OpenStellarCoreTestDatabase()
	q := CoreAccountByAddressQuery{
		GormQuery{&db},
		"gspbxqXqEUZkiCCEFFCN9Vu4FLucdjLLdLcsV6E82Qc1T7ehsTC",
	}

	record, err := First(q)

	if err != nil {
		panic(err)
	}

	account := record.(CoreAccountRecord)
	log.Println(account.Accountid)

}
<|endoftext|>


<commit_before>package hypervisor

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/golang/glog"
	"github.com/hyperhq/runv/api"
	hyperstartapi "github.com/hyperhq/runv/hyperstart/api/json"
	"github.com/hyperhq/runv/hypervisor/types"
	"github.com/hyperhq/runv/lib/utils"
)

const CURRENT_PERSIST_VERSION = 20170224

type PersistContainerInfo struct {
	Description *api.ContainerDescription
	Root        *PersistVolumeInfo
	Process     *hyperstartapi.Process
	Fsmap       []*hyperstartapi.FsmapDescriptor
	VmVolumes   []*hyperstartapi.VolumeDescriptor
}

type PersistVolumeInfo struct {
	Name       string
	Filename   string
	Format     string
	Fstype     string
	DeviceName string
	ScsiId     int
	Containers []int
	MontPoints []string
}

type PersistNetworkInfo struct {
	Id         string
	Index      int
	PciAddr    int
	DeviceName string
	IpAddr     string
}

type PersistInfo struct {
	PersistVersion int
	Id             string
	DriverInfo     map[string]interface{}
	VmSpec         *hyperstartapi.Pod
	HwStat         *VmHwStatus
	VolumeList     []*PersistVolumeInfo
	NetworkList    []*PersistNetworkInfo
	PortList       []*api.PortDescription
	ContainerList  []*PersistContainerInfo
}

func (ctx *VmContext) dump() (*PersistInfo, error) {
	dr, err := ctx.DCtx.Dump()
	if err != nil {
		return nil, err
	}

	nc := ctx.networks
	info := &PersistInfo{
		PersistVersion: CURRENT_PERSIST_VERSION,
		Id:             ctx.Id,
		DriverInfo:     dr,
		VmSpec:         ctx.networks.sandboxInfo(),
		HwStat:         ctx.dumpHwInfo(),
		VolumeList:     make([]*PersistVolumeInfo, len(ctx.volumes)),
		NetworkList:    make([]*PersistNetworkInfo, len(nc.eth)+len(nc.lo)),
		PortList:       make([]*api.PortDescription, len(nc.ports)),
		ContainerList:  make([]*PersistContainerInfo, len(ctx.containers)),
	}

	vid := 0
	for _, vol := range ctx.volumes {
		info.VolumeList[vid] = vol.dump()
		vid++
	}

	for i, p := range nc.ports {
		info.PortList[i] = &api.PortDescription{
			HostPort:      p.HostPort,
			ContainerPort: p.ContainerPort,
			Protocol:      p.Protocol,
		}
	}
	nid := 0
	for _, nic := range nc.lo {
		info.NetworkList[nid] = &PersistNetworkInfo{
			Id:         nic.Id,
			Index:      nic.Index,
			PciAddr:    nic.PCIAddr,
			DeviceName: nic.DeviceName,
			IpAddr:     nic.IpAddr,
		}
		nid++
	}
	nc.slotLock.RLock()
	for _, nic := range nc.eth {
		info.NetworkList[nid] = &PersistNetworkInfo{
			Id:         nic.Id,
			Index:      nic.Index,
			PciAddr:    nic.PCIAddr,
			DeviceName: nic.DeviceName,
			IpAddr:     nic.IpAddr,
		}
		nid++
	}
	defer nc.slotLock.RUnlock()

	cid := 0
	for _, c := range ctx.containers {
		info.ContainerList[cid] = c.dump()
		cid++
	}

	return info, nil
}

func (ctx *VmContext) dumpHwInfo() *VmHwStatus {
	return &VmHwStatus{
		PciAddr:  ctx.pciAddr,
		ScsiId:   ctx.scsiId,
		AttachId: ctx.hyperstart.LastStreamSeq(),
		GuestCid: ctx.GuestCid,
	}
}

func (ctx *VmContext) loadHwStatus(pinfo *PersistInfo) error {
	ctx.pciAddr = pinfo.HwStat.PciAddr
	ctx.scsiId = pinfo.HwStat.ScsiId
	ctx.GuestCid = pinfo.HwStat.GuestCid
	if ctx.GuestCid != 0 {
		if !VsockCidManager.MarkCidInuse(ctx.GuestCid) {
			return fmt.Errorf("conflicting vsock guest cid %d: already in use", ctx.GuestCid)
		}
		ctx.Boot.EnableVsock = true
	}
	return nil
}

func (blk *DiskDescriptor) dump() *PersistVolumeInfo {
	return &PersistVolumeInfo{
		Name:       blk.Name,
		Filename:   blk.Filename,
		Format:     blk.Format,
		Fstype:     blk.Fstype,
		DeviceName: blk.DeviceName,
		ScsiId:     blk.ScsiId,
	}
}

func (vol *PersistVolumeInfo) blockInfo() *DiskDescriptor {
	return &DiskDescriptor{
		Name:       vol.Name,
		Filename:   vol.Filename,
		Format:     vol.Format,
		Fstype:     vol.Fstype,
		DeviceName: vol.DeviceName,
		ScsiId:     vol.ScsiId,
	}
}

func (nc *NetworkContext) load(pinfo *PersistInfo) {
	nc.SandboxConfig = &api.SandboxConfig{
		Hostname: pinfo.VmSpec.Hostname,
		Dns:      pinfo.VmSpec.Dns,
	}
	portWhilteList := pinfo.VmSpec.PortmappingWhiteLists
	if portWhilteList != nil {
		nc.Neighbors = &api.NeighborNetworks{
			InternalNetworks: portWhilteList.InternalNetworks,
			ExternalNetworks: portWhilteList.ExternalNetworks,
		}
	}

	for i, p := range pinfo.PortList {
		nc.ports[i] = p
	}
	for _, pi := range pinfo.NetworkList {
		ifc := &InterfaceCreated{
			Id:         pi.Id,
			Index:      pi.Index,
			PCIAddr:    pi.PciAddr,
			DeviceName: pi.DeviceName,
			IpAddr:     pi.IpAddr,
		}
		// if empty, may be old data, generate one for compatibility.
		if ifc.Id == "" {
			ifc.Id = utils.RandStr(8, "alpha")
		}
		// use device name distinguish from lo and eth
		if ifc.DeviceName == DEFAULT_LO_DEVICE_NAME {
			nc.lo[pi.IpAddr] = ifc
		} else {
			nc.eth[pi.Index] = ifc
		}
		nc.idMap[pi.Id] = ifc
	}
}

func (cc *ContainerContext) dump() *PersistContainerInfo {
	return &PersistContainerInfo{
		Description: cc.ContainerDescription,
		Root:        cc.root.dump(),
		Process:     cc.process,
		Fsmap:       cc.fsmap,
		VmVolumes:   cc.vmVolumes,
	}
}

func vmDeserialize(s []byte) (*PersistInfo, error) {
	info := &PersistInfo{}
	err := json.Unmarshal(s, info)
	return info, err
}

func (pinfo *PersistInfo) serialize() ([]byte, error) {
	return json.Marshal(pinfo)
}

func (pinfo *PersistInfo) vmContext(hub chan VmEvent, client chan *types.VmResponse) (*VmContext, error) {
	oldVersion := pinfo.PersistVersion < CURRENT_PERSIST_VERSION

	dc, err := HDriver.LoadContext(pinfo.DriverInfo)
	if err != nil {
		glog.Error("cannot load driver context: ", err.Error())
		return nil, err
	}

	ctx, err := InitContext(pinfo.Id, hub, client, dc, &BootConfig{})
	if err != nil {
		return nil, err
	}

	err = ctx.loadHwStatus(pinfo)
	if err != nil {
		return nil, err
	}

	ctx.networks.load(pinfo)

	oldVersionVolume := false
	if len(pinfo.VolumeList) > 0 && (len(pinfo.VolumeList[0].Containers) > 0 || len(pinfo.VolumeList[0].MontPoints) > 0) {
		oldVersionVolume = true
	}
	imageMap := make(map[string]*DiskDescriptor)
	for _, vol := range pinfo.VolumeList {
		binfo := vol.blockInfo()
		if oldVersionVolume {
			if len(vol.Containers) != len(vol.MontPoints) {
				return nil, fmt.Error("persistent data corrupt, volume info mismatch")
			}
			if len(vol.MontPoints) == 1 && vol.MontPoints[0] == "/" {
				imageMap[vol.Name] = binfo
				continue
			}
		}
		ctx.volumes[binfo.Name] = &DiskContext{
			DiskDescriptor: binfo,
			sandbox:        ctx,
			observers:      make(map[string]*sync.WaitGroup),
			lock:           &sync.RWMutex{},
			// FIXME: is there any trouble if we set it as ready when restoring from persistence
			ready: true,
		}
	}

	for _, pc := range pinfo.ContainerList {
		c := pc.Description
		cc := &ContainerContext{
			ContainerDescription: c,
			fsmap:                pc.Fsmap,
			vmVolumes:            pc.VmVolumes,
			process:              pc.Process,
			sandbox:              ctx,
			logPrefix:            fmt.Sprintf("SB[%s] Con[%s] ", ctx.Id, c.Id),
			root: &DiskContext{
				DiskDescriptor: pc.Root.blockInfo(),
				sandbox:        ctx,
				isRootVol:      true,
				ready:          true,
				observers:      make(map[string]*sync.WaitGroup),
				lock:           &sync.RWMutex{},
			},
		}
		// restore wg for volumes attached to container
		wgDisk := &sync.WaitGroup{}
		for vn := range c.Volumes {
			entry, ok := ctx.volumes[vn]
			if !ok {
				cc.Log(ERROR, "restoring container volume does not exist in volume map, skip")
				continue
			}
			entry.wait(c.Id, wgDisk)
		}
		cc.root.wait(c.Id, wgDisk)

		ctx.containers[c.Id] = cc
	}

	return ctx, nil
}
<commit_msg>make vmContext load/dump compatible with old data<commit_after>package hypervisor

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/golang/glog"
	"github.com/hyperhq/runv/api"
	hyperstartapi "github.com/hyperhq/runv/hyperstart/api/json"
	"github.com/hyperhq/runv/hypervisor/types"
	"github.com/hyperhq/runv/lib/utils"
)

const CURRENT_PERSIST_VERSION = 20170224

type PersistVolumeInfo struct {
	Name         string
	Filename     string
	Format       string
	Fstype       string
	DeviceName   string
	ScsiId       int
	ContainerIds []string
	IsRootVol    bool
	Containers   []int // deprecated
	MontPoints   []string
}

type PersistNetworkInfo struct {
	Id         string
	Index      int
	PciAddr    int
	DeviceName string
	IpAddr     string
}

type PersistInfo struct {
	PersistVersion int
	Id             string
	DriverInfo     map[string]interface{}
	VmSpec         *hyperstartapi.Pod
	HwStat         *VmHwStatus
	VolumeList     []*PersistVolumeInfo
	NetworkList    []*PersistNetworkInfo
	PortList       []*api.PortDescription
}

func (ctx *VmContext) dump() (*PersistInfo, error) {
	dr, err := ctx.DCtx.Dump()
	if err != nil {
		return nil, err
	}

	nc := ctx.networks
	info := &PersistInfo{
		PersistVersion: CURRENT_PERSIST_VERSION,
		Id:             ctx.Id,
		DriverInfo:     dr,
		VmSpec:         ctx.networks.sandboxInfo(),
		HwStat:         ctx.dumpHwInfo(),
		VolumeList:     make([]*PersistVolumeInfo, len(ctx.volumes)+len(ctx.containers)),
		NetworkList:    make([]*PersistNetworkInfo, len(nc.eth)+len(nc.lo)),
		PortList:       make([]*api.PortDescription, len(nc.ports)),
	}

	vid := 0
	for _, vol := range ctx.volumes {
		v := vol.dump()
		for id := range vol.observers {
			v.ContainerIds = append(v.ContainerIds, id)
		}
		info.VolumeList[vid] = v
		vid++
	}

	for i, p := range nc.ports {
		info.PortList[i] = &api.PortDescription{
			HostPort:      p.HostPort,
			ContainerPort: p.ContainerPort,
			Protocol:      p.Protocol,
		}
	}
	nid := 0
	for _, nic := range nc.lo {
		info.NetworkList[nid] = &PersistNetworkInfo{
			Id:         nic.Id,
			Index:      nic.Index,
			PciAddr:    nic.PCIAddr,
			DeviceName: nic.DeviceName,
			IpAddr:     nic.IpAddr,
		}
		nid++
	}
	nc.slotLock.RLock()
	for _, nic := range nc.eth {
		info.NetworkList[nid] = &PersistNetworkInfo{
			Id:         nic.Id,
			Index:      nic.Index,
			PciAddr:    nic.PCIAddr,
			DeviceName: nic.DeviceName,
			IpAddr:     nic.IpAddr,
		}
		nid++
	}
	defer nc.slotLock.RUnlock()

	cid := 0
	info.VmSpec.DeprecatedContainers = make([]hyperstartapi.Container, len(ctx.containers))
	for _, c := range ctx.containers {
		info.VmSpec.DeprecatedContainers[cid] = *c.VmSpec()
		rootVolume := c.root.dump()
		rootVolume.ContainerIds = []string{c.Id}
		rootVolume.IsRootVol = true
		info.VolumeList[vid] = rootVolume
		vid++
		cid++
	}

	return info, nil
}

func (ctx *VmContext) dumpHwInfo() *VmHwStatus {
	return &VmHwStatus{
		PciAddr:  ctx.pciAddr,
		ScsiId:   ctx.scsiId,
		AttachId: ctx.hyperstart.LastStreamSeq(),
		GuestCid: ctx.GuestCid,
	}
}

func (ctx *VmContext) loadHwStatus(pinfo *PersistInfo) error {
	ctx.pciAddr = pinfo.HwStat.PciAddr
	ctx.scsiId = pinfo.HwStat.ScsiId
	ctx.GuestCid = pinfo.HwStat.GuestCid
	if ctx.GuestCid != 0 {
		if !VsockCidManager.MarkCidInuse(ctx.GuestCid) {
			return fmt.Errorf("conflicting vsock guest cid %d: already in use", ctx.GuestCid)
		}
		ctx.Boot.EnableVsock = true
	}
	return nil
}

func (blk *DiskDescriptor) dump() *PersistVolumeInfo {
	return &PersistVolumeInfo{
		Name:       blk.Name,
		Filename:   blk.Filename,
		Format:     blk.Format,
		Fstype:     blk.Fstype,
		DeviceName: blk.DeviceName,
		ScsiId:     blk.ScsiId,
	}
}

func (vol *PersistVolumeInfo) blockInfo() *DiskDescriptor {
	return &DiskDescriptor{
		Name:       vol.Name,
		Filename:   vol.Filename,
		Format:     vol.Format,
		Fstype:     vol.Fstype,
		DeviceName: vol.DeviceName,
		ScsiId:     vol.ScsiId,
	}
}

func (nc *NetworkContext) load(pinfo *PersistInfo) {
	nc.SandboxConfig = &api.SandboxConfig{
		Hostname: pinfo.VmSpec.Hostname,
		Dns:      pinfo.VmSpec.Dns,
	}
	portWhilteList := pinfo.VmSpec.PortmappingWhiteLists
	if portWhilteList != nil {
		nc.Neighbors = &api.NeighborNetworks{
			InternalNetworks: portWhilteList.InternalNetworks,
			ExternalNetworks: portWhilteList.ExternalNetworks,
		}
	}

	for i, p := range pinfo.PortList {
		nc.ports[i] = p
	}
	for _, pi := range pinfo.NetworkList {
		ifc := &InterfaceCreated{
			Id:         pi.Id,
			Index:      pi.Index,
			PCIAddr:    pi.PciAddr,
			DeviceName: pi.DeviceName,
			IpAddr:     pi.IpAddr,
		}
		// if empty, may be old data, generate one for compatibility.
		if ifc.Id == "" {
			ifc.Id = utils.RandStr(8, "alpha")
		}
		// use device name distinguish from lo and eth
		if ifc.DeviceName == DEFAULT_LO_DEVICE_NAME {
			nc.lo[pi.IpAddr] = ifc
		} else {
			nc.eth[pi.Index] = ifc
		}
		nc.idMap[pi.Id] = ifc
	}
}

func vmDeserialize(s []byte) (*PersistInfo, error) {
	info := &PersistInfo{}
	err := json.Unmarshal(s, info)
	return info, err
}

func (pinfo *PersistInfo) serialize() ([]byte, error) {
	return json.Marshal(pinfo)
}

func (pinfo *PersistInfo) vmContext(hub chan VmEvent, client chan *types.VmResponse) (*VmContext, error) {
	oldVersion := pinfo.PersistVersion < CURRENT_PERSIST_VERSION

	dc, err := HDriver.LoadContext(pinfo.DriverInfo)
	if err != nil {
		glog.Error("cannot load driver context: ", err.Error())
		return nil, err
	}

	ctx, err := InitContext(pinfo.Id, hub, client, dc, &BootConfig{})
	if err != nil {
		return nil, err
	}

	err = ctx.loadHwStatus(pinfo)
	if err != nil {
		return nil, err
	}

	ctx.networks.load(pinfo)

	// map container id to image DiskDescriptor
	imageMap := make(map[string]*DiskDescriptor)
	// map container id to volume DiskContext list
	volumeMap := make(map[string][]*DiskContext)
	pcList := pinfo.VmSpec.DeprecatedContainers
	for _, vol := range pinfo.VolumeList {
		binfo := vol.blockInfo()
		if oldVersion {
			if len(vol.Containers) != len(vol.MontPoints) {
				return nil, fmt.Errorf("persistent data corrupt, volume info mismatch")
			}
			if len(vol.MontPoints) == 1 && vol.MontPoints[0] == "/" {
				imageMap[pcList[vol.Containers[0]].Id] = binfo
				continue
			}
		} else {
			if vol.IsRootVol {
				if len(vol.ContainerIds) != 1 {
					return nil, fmt.Errorf("persistent data corrupt, root volume mismatch")
				}
				imageMap[vol.ContainerIds[0]] = binfo
				continue
			}
		}
		ctx.volumes[binfo.Name] = &DiskContext{
			DiskDescriptor: binfo,
			sandbox:        ctx,
			observers:      make(map[string]*sync.WaitGroup),
			lock:           &sync.RWMutex{},
			// FIXME: is there any trouble if we set it as ready when restoring from persistence
			ready: true,
		}
		if oldVersion {
			for _, idx := range vol.Containers {
				volumeMap[pcList[idx].Id] = append(volumeMap[pcList[idx].Id], ctx.volumes[binfo.Name])
			}
		} else {
			for _, id := range vol.ContainerIds {
				volumeMap[id] = append(volumeMap[id], ctx.volumes[binfo.Name])
			}
		}
	}

	for _, pc := range pcList {
		bInfo, ok := imageMap[pc.Id]
		if !ok {
			return nil, fmt.Errorf("persistent data corrupt, lack of container root volume")
		}
		cc := &ContainerContext{
			ContainerDescription: &api.ContainerDescription{
				Id:         pc.Id,
				RootPath:   pc.Rootfs,
				Initialize: pc.Initialize,
				Sysctl:     pc.Sysctl,
				RootVolume: &api.VolumeDescription{
					Name:         bInfo.Name,
					Source:       bInfo.Filename,
					Format:       bInfo.Format,
					Fstype:       bInfo.Fstype,
					DockerVolume: bInfo.DockerVolume,
				},
			},
			fsmap:     pc.Fsmap,
			process:   pc.Process,
			vmVolumes: pc.Volumes,
			sandbox:   ctx,
			logPrefix: fmt.Sprintf("SB[%s] Con[%s] ", ctx.Id, pc.Id),
			root: &DiskContext{
				DiskDescriptor: bInfo,
				sandbox:        ctx,
				isRootVol:      true,
				ready:          true,
				observers:      make(map[string]*sync.WaitGroup),
				lock:           &sync.RWMutex{},
			},
		}
		if cc.process.Id == "" {
			cc.process.Id = "init"
		}
		// restore wg for volumes attached to container
		wgDisk := &sync.WaitGroup{}
		volList, ok := volumeMap[pc.Id]
		if ok {
			cc.Volumes = make(map[string]*api.VolumeReference, len(volList))
			for _, vol := range volList {
				// for unwait attached volumes when removing container
				cc.Volumes[vol.Name] = &api.VolumeReference{}
				vol.wait(cc.Id, wgDisk)
			}
		}
		cc.root.wait(cc.Id, wgDisk)

		ctx.containers[cc.Id] = cc
	}

	return ctx, nil
}
<|endoftext|>


<commit_before>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
	"sync"
)

// AWSSQSBroker represents a AWS SQS broker
type AWSSQSBroker struct {
	Broker
	processingWG      sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal
	receivingWG       sync.WaitGroup
	stopReceivingChan chan int
	sess              *session.Session
	service           sqsiface.SQSAPI
}

// NewAWSSQSBroker creates new Broker instance
func NewAWSSQSBroker(cnf *config.Config) Interface {
	b := &AWSSQSBroker{Broker: New(cnf)}
	b.sess = session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	b.service = sqs.New(b.sess)

	return b
}

// GetPendingTasks returns a slice of task.Signatures waiting in the queue
func (b *AWSSQSBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
	return nil, errors.New("Not implemented")
}

// StartConsuming enters a loop and waits for incoming messages
func (b *AWSSQSBroker) StartConsuming(consumerTag string, concurrency int, taskProcessor TaskProcessor) (bool, error) {
	b.startConsuming(consumerTag, taskProcessor)
	qURL := b.defaultQueueURL()
	deliveries := make(chan *sqs.ReceiveMessageOutput)

	b.stopReceivingChan = make(chan int)
	b.receivingWG.Add(1)

	go func() {
		defer b.receivingWG.Done()

		log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

		for {
			select {
			// A way to stop this goroutine from b.StopConsuming
			case <-b.stopReceivingChan:
				return
			default:
				output, err := b.receiveMessage(qURL)
				if err != nil {
					log.ERROR.Printf("Queue consume error: %s", err)
					continue
				}
				if len(output.Messages) == 0 {
					continue
				}

				deliveries <- output
			}

			whetherContinue, err := b.continueReceivingMessages(qURL, deliveries)
			if err != nil {
				log.ERROR.Printf("Error when receiving messages. Error: %v", err)
			}
			if whetherContinue == false {
				return
			}
		}
	}()

	if err := b.consume(deliveries, concurrency, taskProcessor); err != nil {
		return b.retry, err
	}

	return b.retry, nil
}

// StopConsuming quits the loop
func (b *AWSSQSBroker) StopConsuming() {
	b.stopConsuming()

	b.stopReceiving()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	// Waiting for the receiving goroutine to have stopped
	b.receivingWG.Wait()
}

// Publish places a new message on the default queue
func (b *AWSSQSBroker) Publish(signature *tasks.Signature) error {

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check that signature.RoutingKey is set, if not switch to DefaultQueue
	b.AdjustRoutingKey(signature)

	MsgInput := &sqs.SendMessageInput{
		MessageBody: aws.String(string(msg)),
		QueueUrl:    aws.String(b.cnf.Broker + "/" + signature.RoutingKey),
	}

	// if this is a fifo queue, there needs to be some addtional parameters.
	if strings.HasSuffix(signature.RoutingKey, ".fifo") {
		// Use Machinery's signature Task UUID as SQS Message Group ID.
		MsgDedupID := signature.UUID
		MsgInput.MessageDeduplicationId = aws.String(MsgDedupID)

		// Use Machinery's signature Group UUID as SQS Message Group ID.
		MsgGroupID := signature.GroupUUID
		MsgInput.MessageGroupId = aws.String(MsgGroupID)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// and is not a fifo queue, set a delay in seconds for the task.
	if signature.ETA != nil && !strings.HasSuffix(signature.RoutingKey, ".fifo") {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			MsgInput.DelaySeconds = aws.Int64(signature.ETA.Unix() - now.Unix())
		}
	}

	result, err := b.service.SendMessage(MsgInput)

	if err != nil {
		log.INFO.Println("Error", err)
		return err

	}
	log.INFO.Println("Success", *result.MessageId)
	return nil

}

// TODO: Add GetPendingTasks() & add AssignWorker(), refer to
// RichardKnop's broker.

func (b *AWSSQSBroker) consume(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		b.initializePool(pool, concurrency)
	}()

	errorsChan := make(chan error)

	for {
		err := b.consumeDeliveries(deliveries, concurrency, taskProcessor, pool, errorsChan)
		if err != nil {
			return err
		}
	}
}

func (b *AWSSQSBroker) consumeOne(delivery *sqs.ReceiveMessageOutput, taskProcessor TaskProcessor) error {
	sig := new(tasks.Signature)
	if len(delivery.Messages) == 0 {
		log.ERROR.Printf("received an empty message, the delivery was %v", delivery)
		return errors.New("received empty message, the delivery is " + delivery.GoString())
	}
	msg := delivery.Messages[0].Body
	if err := json.Unmarshal([]byte(*msg), sig); err != nil {
		log.ERROR.Printf("unmarshal error. the delivery is %v", delivery)
		return err
	}

	// Delete message after consuming successfully
	err := b.deleteOne(delivery)
	if err != nil {
		log.ERROR.Printf("error when deleting the delivery. the delivery is %v", delivery)
	}

	// If the task is not registered, we requeue it,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(sig.Name) {
		err := b.Publish(sig)
		if err != nil {
			return err
		}
		log.INFO.Printf("requeue a task to default queue: %v", sig)
		return nil
	}
	return taskProcessor.Process(sig)
}

func (b *AWSSQSBroker) deleteOne(delivery *sqs.ReceiveMessageOutput) error {
	qURL := b.defaultQueueURL()
	_, err := b.service.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      qURL,
		ReceiptHandle: delivery.Messages[0].ReceiptHandle,
	})

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

func (b *AWSSQSBroker) defaultQueueURL() *string {
	return aws.String(b.cnf.Broker + "/" + b.cnf.DefaultQueue)
}

func (b *AWSSQSBroker) receiveMessage(qURL *string) (*sqs.ReceiveMessageOutput, error) {
	result, err := b.service.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            qURL,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(int64(b.cnf.ResultsExpireIn)), // 10 hours
		WaitTimeSeconds:     aws.Int64(0),
	})
	if err != nil {
		return nil, err
	}
	return result, err
}

func (b *AWSSQSBroker) initializePool(pool chan struct{}, concurrency int) {
	for i := 0; i < concurrency; i++ {
		pool <- struct{}{}
	}
}

func (b *AWSSQSBroker) consumeDeliveries(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor, pool chan struct{}, errorsChan chan error) error {
	select {
	case err := <-errorsChan:
		return err
	case d := <-deliveries:
		if concurrency > 0 {
			// get worker from pool (blocks until one is available)
			<-pool
		}

		b.processingWG.Add(1)

		// Consume the task inside a goroutine so multiple tasks
		// can be processed concurrently
		go func() {

			if err := b.consumeOne(d, taskProcessor); err != nil {
				errorsChan <- err
			}

			b.processingWG.Done()

			if concurrency > 0 {
				// give worker back to pool
				pool <- struct{}{}
			}
		}()
	case <-b.stopChan:
		return nil
	}
	return nil
}

func (b *AWSSQSBroker) continueReceivingMessages(qURL *string, deliveries chan *sqs.ReceiveMessageOutput) (bool, error) {
	select {
	// A way to stop this goroutine from b.StopConsuming
	case <-b.stopReceivingChan:
		return false, nil
	default:
		output, err := b.receiveMessage(qURL)
		if err != nil {
			return true, err
		}
		if len(output.Messages) == 0 {
			return true, nil
		}
		go func() { deliveries <- output }()
	}
	return true, nil
}

func (b *AWSSQSBroker) stopReceiving() {
	// Stop the receiving goroutine
	b.stopReceivingChan <- 1
}
<commit_msg>fix: resolve issue with sqs routing check<commit_after>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
	"sync"
)

// AWSSQSBroker represents a AWS SQS broker
type AWSSQSBroker struct {
	Broker
	processingWG      sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal
	receivingWG       sync.WaitGroup
	stopReceivingChan chan int
	sess              *session.Session
	service           sqsiface.SQSAPI
}

// NewAWSSQSBroker creates new Broker instance
func NewAWSSQSBroker(cnf *config.Config) Interface {
	b := &AWSSQSBroker{Broker: New(cnf)}
	b.sess = session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	b.service = sqs.New(b.sess)

	return b
}

// GetPendingTasks returns a slice of task.Signatures waiting in the queue
func (b *AWSSQSBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
	return nil, errors.New("Not implemented")
}

// StartConsuming enters a loop and waits for incoming messages
func (b *AWSSQSBroker) StartConsuming(consumerTag string, concurrency int, taskProcessor TaskProcessor) (bool, error) {
	b.startConsuming(consumerTag, taskProcessor)
	qURL := b.defaultQueueURL()
	deliveries := make(chan *sqs.ReceiveMessageOutput)

	b.stopReceivingChan = make(chan int)
	b.receivingWG.Add(1)

	go func() {
		defer b.receivingWG.Done()

		log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

		for {
			select {
			// A way to stop this goroutine from b.StopConsuming
			case <-b.stopReceivingChan:
				return
			default:
				output, err := b.receiveMessage(qURL)
				if err != nil {
					log.ERROR.Printf("Queue consume error: %s", err)
					continue
				}
				if len(output.Messages) == 0 {
					continue
				}

				deliveries <- output
			}

			whetherContinue, err := b.continueReceivingMessages(qURL, deliveries)
			if err != nil {
				log.ERROR.Printf("Error when receiving messages. Error: %v", err)
			}
			if whetherContinue == false {
				return
			}
		}
	}()

	if err := b.consume(deliveries, concurrency, taskProcessor); err != nil {
		return b.retry, err
	}

	return b.retry, nil
}

// StopConsuming quits the loop
func (b *AWSSQSBroker) StopConsuming() {
	b.stopConsuming()

	b.stopReceiving()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	// Waiting for the receiving goroutine to have stopped
	b.receivingWG.Wait()
}

// Publish places a new message on the default queue
func (b *AWSSQSBroker) Publish(signature *tasks.Signature) error {

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check that signature.RoutingKey is set, if not switch to DefaultQueue
	AdjustRoutingKey(b, signature)

	MsgInput := &sqs.SendMessageInput{
		MessageBody: aws.String(string(msg)),
		QueueUrl:    aws.String(b.cnf.Broker + "/" + signature.RoutingKey),
	}

	// if this is a fifo queue, there needs to be some addtional parameters.
	if strings.HasSuffix(signature.RoutingKey, ".fifo") {
		// Use Machinery's signature Task UUID as SQS Message Group ID.
		MsgDedupID := signature.UUID
		MsgInput.MessageDeduplicationId = aws.String(MsgDedupID)

		// Use Machinery's signature Group UUID as SQS Message Group ID.
		MsgGroupID := signature.GroupUUID
		MsgInput.MessageGroupId = aws.String(MsgGroupID)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// and is not a fifo queue, set a delay in seconds for the task.
	if signature.ETA != nil && !strings.HasSuffix(signature.RoutingKey, ".fifo") {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			MsgInput.DelaySeconds = aws.Int64(signature.ETA.Unix() - now.Unix())
		}
	}

	result, err := b.service.SendMessage(MsgInput)

	if err != nil {
		log.INFO.Println("Error", err)
		return err

	}
	log.INFO.Println("Success", *result.MessageId)
	return nil

}

// TODO: Add GetPendingTasks() & add AssignWorker(), refer to
// RichardKnop's broker.

func (b *AWSSQSBroker) consume(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		b.initializePool(pool, concurrency)
	}()

	errorsChan := make(chan error)

	for {
		err := b.consumeDeliveries(deliveries, concurrency, taskProcessor, pool, errorsChan)
		if err != nil {
			return err
		}
	}
}

func (b *AWSSQSBroker) consumeOne(delivery *sqs.ReceiveMessageOutput, taskProcessor TaskProcessor) error {
	sig := new(tasks.Signature)
	if len(delivery.Messages) == 0 {
		log.ERROR.Printf("received an empty message, the delivery was %v", delivery)
		return errors.New("received empty message, the delivery is " + delivery.GoString())
	}
	msg := delivery.Messages[0].Body
	if err := json.Unmarshal([]byte(*msg), sig); err != nil {
		log.ERROR.Printf("unmarshal error. the delivery is %v", delivery)
		return err
	}

	// Delete message after consuming successfully
	err := b.deleteOne(delivery)
	if err != nil {
		log.ERROR.Printf("error when deleting the delivery. the delivery is %v", delivery)
	}

	// If the task is not registered, we requeue it,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(sig.Name) {
		err := b.Publish(sig)
		if err != nil {
			return err
		}
		log.INFO.Printf("requeue a task to default queue: %v", sig)
		return nil
	}
	return taskProcessor.Process(sig)
}

func (b *AWSSQSBroker) deleteOne(delivery *sqs.ReceiveMessageOutput) error {
	qURL := b.defaultQueueURL()
	_, err := b.service.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      qURL,
		ReceiptHandle: delivery.Messages[0].ReceiptHandle,
	})

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

func (b *AWSSQSBroker) defaultQueueURL() *string {
	return aws.String(b.cnf.Broker + "/" + b.cnf.DefaultQueue)
}

func (b *AWSSQSBroker) receiveMessage(qURL *string) (*sqs.ReceiveMessageOutput, error) {
	result, err := b.service.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            qURL,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(int64(b.cnf.ResultsExpireIn)), // 10 hours
		WaitTimeSeconds:     aws.Int64(0),
	})
	if err != nil {
		return nil, err
	}
	return result, err
}

func (b *AWSSQSBroker) initializePool(pool chan struct{}, concurrency int) {
	for i := 0; i < concurrency; i++ {
		pool <- struct{}{}
	}
}

func (b *AWSSQSBroker) consumeDeliveries(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor, pool chan struct{}, errorsChan chan error) error {
	select {
	case err := <-errorsChan:
		return err
	case d := <-deliveries:
		if concurrency > 0 {
			// get worker from pool (blocks until one is available)
			<-pool
		}

		b.processingWG.Add(1)

		// Consume the task inside a goroutine so multiple tasks
		// can be processed concurrently
		go func() {

			if err := b.consumeOne(d, taskProcessor); err != nil {
				errorsChan <- err
			}

			b.processingWG.Done()

			if concurrency > 0 {
				// give worker back to pool
				pool <- struct{}{}
			}
		}()
	case <-b.stopChan:
		return nil
	}
	return nil
}

func (b *AWSSQSBroker) continueReceivingMessages(qURL *string, deliveries chan *sqs.ReceiveMessageOutput) (bool, error) {
	select {
	// A way to stop this goroutine from b.StopConsuming
	case <-b.stopReceivingChan:
		return false, nil
	default:
		output, err := b.receiveMessage(qURL)
		if err != nil {
			return true, err
		}
		if len(output.Messages) == 0 {
			return true, nil
		}
		go func() { deliveries <- output }()
	}
	return true, nil
}

func (b *AWSSQSBroker) stopReceiving() {
	// Stop the receiving goroutine
	b.stopReceivingChan <- 1
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"time"

	"github.com/TheJumpCloud/jcapi"
)

const (
	COMMAND_NAME_RANDOM_PART int = 16

	MAX_COMMAND_LEN int = 16384 // Maximum size of a command to send to JumpCloud

	RESULT_POLL_TIME     int = 5  // Check every RESULT_POLL_TIME seconds for command results
	RESULT_MAX_POLL_TIME int = 70 // Stop checking for results after RESULT_MAX_POLL_TIME_SECONDS (must be a minimum of 60 seconds)

	URL_BASE string = "https://console.jumpcloud.com/api"
)

func makeImmediateCommand(name, command, commandType, shell, user string) (cmd jcapi.JCCommand) {
	cmd = jcapi.JCCommand{
		Name:        name,
		Command:     command,
		CommandType: commandType,
		User:        user,
		LaunchType:  "manual",
		Schedule:    "immediate",
		Timeout:     "0", // No timeout
		ListensTo:   "",
		Trigger:     "",
		Sudo:        false,
		Shell:       shell,
		Skip:        0,
		Limit:       10,
	}

	return
}

func deleteCommandResultsByName(jc jcapi.JCAPI, commandName string) (err error) {
	results, err := jc.GetCommandResultsByName(commandName)
	if err != nil {
		err = fmt.Errorf("Could not find the command result for '%s', err='%s'", commandName, err.Error())
		return
	}

	errBuf := ""

	for _, result := range results {
		errBuf += jc.DeleteCommandResult(result.Id).Error()
	}

	if err != nil {
		err = fmt.Errorf("One or more deletes failed, err='%s'", err.Error())
		return
	}

	return
}

func findSystemsByOSType(systems []jcapi.JCSystem, osTypeRegEx string) (indices []int, err error) {
	r, err := regexp.Compile(osTypeRegEx)
	if err != nil {
		err = fmt.Errorf("Could not compile regex for '%s', err='%s'", osTypeRegEx, err.Error())
		return
	}

	for idx, system := range systems {
		if r.Match([]byte(system.Os)) {
			indices = append(indices, idx)
		}
	}

	return
}

func getSystemNameMap(jc jcapi.JCAPI, idList []string) (systemNameMap map[string]jcapi.JCSystem, err error) {
	systemNameMap = make(map[string]jcapi.JCSystem)

	//
	// Walk the IDs on which we ran the command, and create a map of host name to
	// JCSystem object. This allows us to determine which command results go with
	// which commands.
	//
	for _, id := range idList {
		system, err2 := jc.GetSystemById(id, false)
		if err2 != nil {
			err = fmt.Errorf("Could not get system object for id '%s', err='%s'", id, err.Error())
			return
		}

		systemNameMap[system.Hostname] = system
	}

	return
}

func waitForAndProcessResults(jc jcapi.JCAPI, commandObj jcapi.JCCommand) (outputBuffer string, err error) {
	systemsFound := make(map[string]jcapi.JCResponse)

	systemNameMap, err := getSystemNameMap(jc, commandObj.Systems)
	if err != nil {
		err = fmt.Errorf("Could not map system names to systems, err='%s'", err.Error())
		return
	}

	pollTime := RESULT_POLL_TIME

	fmt.Printf("\nWaiting for results...")

	for i := 0; i < RESULT_MAX_POLL_TIME; i += pollTime {
		time.Sleep(time.Duration(pollTime) * time.Second)

		fmt.Printf(".")

		results, err2 := jc.GetCommandResultsByName(commandObj.Name)
		if err2 != nil {
			err = fmt.Errorf("Could not find the command result for '%s', err='%s'", commandObj.Name, err2.Error())
			return
		}

		// Walk the results and add their exit code to the map (maps system name to the result data)
		for _, result := range results {
			systemsFound[result.System] = result.Response
		}

		if len(results) == len(commandObj.Systems) {
			break
		}
	}

	fmt.Printf("done.\n\n")

	// Output the results for all systems that were successful
	headerBuffer := fmt.Sprintf("Systems with a zero exit code\n")
	headerBuffer += fmt.Sprintf("=============================\n")

	dataBuffer := ""

	for systemId, result := range systemsFound {
		if result.Data.ExitCode == 0 && result.Error == "" {
			dataBuffer += fmt.Sprintf("ID %s: Completed, exit code=%d, output=[%s]\n", systemId, result.Data.ExitCode, result.Data.Output)
		}
	}

	if dataBuffer != "" {
		outputBuffer = headerBuffer + dataBuffer
	}

	// Output the results for all systems that returned a non-zero exit code
	headerBuffer = fmt.Sprintf("\nSystems with non-zero exit code\n")
	headerBuffer += fmt.Sprintf("===============================\n")

	dataBuffer = ""

	for systemId, result := range systemsFound {
		if result.Data.ExitCode > 0 || result.Error != "" {
			dataBuffer += fmt.Sprintf("ID %s: Failed, exit code=%d, error=[%s]\n", systemId, result.Data.ExitCode, result.Error)
		}
	}

	if dataBuffer != "" {
		outputBuffer = headerBuffer + dataBuffer
	}

	// Output the results for all systems that did not return any result
	headerBuffer = fmt.Sprintf("\nSystems with no command result\n")
	headerBuffer += fmt.Sprintf("===============================\n")

	dataBuffer = ""

	for hostName, _ := range systemNameMap {
		if _, exists := systemsFound[hostName]; !exists {
			dataBuffer += fmt.Sprintf("ID %s (%s): Received no command result\n", systemNameMap[hostName].Id, hostName)
		}
	}

	if dataBuffer != "" {
		outputBuffer += headerBuffer + dataBuffer
	}

	return
}

//
// This example accepts an API key and the name of the file to execute commands against. You can then
// specify a regular expression for the JCSystem.Os value of the systems you want the command to run
// against.
//
// It will write the results of the commands to stdout when exitCode != 0
//
func main() {
	apiKey := flag.String("api-key", "none", "Your JumpCloud Administrator API Key")
	commandFile := flag.String("command-file", "myCommand.txt", "The name of a file containing a command to run")
	commandType := flag.String("command-type", "windows", "OS type of the command, 'linux', 'windows', or 'mac'")
	shell := flag.String("shell", "powershell", "Shell (Windows-only) (powershell/cmd)")
	osType := flag.String("os-type", "Windows.*", "A regular expression to match your systems for OS type")
	deleteFlag := flag.Bool("delete-after-run", false, "When true, delete commands and results at completion.")

	flag.Parse()

	jc := jcapi.NewJCAPI(*apiKey, URL_BASE)

	// Generate a randomized command name
	commandName := "CMD " + makeRandomString(COMMAND_NAME_RANDOM_PART)

	// Read the input file
	commandData, err := readFile(*commandFile, MAX_COMMAND_LEN)
	if err != nil {
		log.Fatalf("Could not read command, err='%s'", err.Error())
	}

	// Make the command object
	commandObj := makeImmediateCommand(commandName, commandData, *commandType, *shell, jcapi.COMMAND_ROOT_USER)

	//
	// Get the list of matching servers and add them to the command
	//
	systems, err := jc.GetSystems(false)
	if err != nil {
		log.Fatalf("Could not get a list of all systems, err='%s'")
	}

	indices, err := findSystemsByOSType(systems, *osType)
	if err != nil {
		log.Fatalf("Could not search a list of systems for OS type matching '%s', err='%s'", *osType, err.Error())
	}

	if len(indices) == 0 {
		log.Fatalf("No systems match '%s' on your JumpCloud account\n", *osType)
	}

	fmt.Printf("Executing Command on the Following Systems\n")
	fmt.Printf("------------------------------------------\n")

	for _, index := range indices {
		fmt.Printf("%s\t%s\n", systems[index].Id, systems[index].Hostname)

		commandObj.Systems = append(commandObj.Systems, systems[index].Id)
	}

	//
	// Add the command object to the JumpCloud account
	//
	commandObj, err = jc.AddUpdateCommand(jcapi.Insert, commandObj)
	if err != nil {
		log.Fatalf("Could not insert a new command, err='%s'", err.Error())
	}
	if *deleteFlag {
		defer jc.DeleteCommand(commandObj)
	}

	//
	// Run the command
	//
	err = jc.RunCommand(commandObj)
	if err != nil {
		log.Fatalf("Could not run the command '%s', err='%s'", commandObj.ToString(), err.Error())
	}
	if *deleteFlag {
		defer deleteCommandResultsByName(jc, commandName)
	}

	//
	// Wait for results to come back from each command executed, but don't wait forever if some of the results
	// don't come back. Write to stdout as results are returned, and after a specified timeout, output the list
	// of failed commands and commands that did not run within the specified timeout.
	//
	outputBuffer, err := waitForAndProcessResults(jc, commandObj)
	if err != nil {
		log.Fatalf("Could not wait for and process results, err='%s'", err.Error())
	}

	fmt.Printf(outputBuffer)

	os.Exit(0)
}
<commit_msg>Added check for active systems only<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"time"

	"github.com/TheJumpCloud/jcapi"
)

const (
	COMMAND_NAME_RANDOM_PART int = 16

	MAX_COMMAND_LEN int = 16384 // Maximum size of a command to send to JumpCloud

	RESULT_POLL_TIME     int = 5  // Check every RESULT_POLL_TIME seconds for command results
	RESULT_MAX_POLL_TIME int = 70 // Stop checking for results after RESULT_MAX_POLL_TIME_SECONDS (must be a minimum of 60 seconds)

	URL_BASE string = "https://console.jumpcloud.com/api"
)

func makeImmediateCommand(name, command, commandType, shell, user string) (cmd jcapi.JCCommand) {
	cmd = jcapi.JCCommand{
		Name:        name,
		Command:     command,
		CommandType: commandType,
		User:        user,
		LaunchType:  "manual",
		Schedule:    "immediate",
		Timeout:     "0", // No timeout
		ListensTo:   "",
		Trigger:     "",
		Sudo:        false,
		Shell:       shell,
		Skip:        0,
		Limit:       10,
	}

	return
}

func deleteCommandResultsByName(jc jcapi.JCAPI, commandName string) (err error) {
	results, err := jc.GetCommandResultsByName(commandName)
	if err != nil {
		err = fmt.Errorf("Could not find the command result for '%s', err='%s'", commandName, err.Error())
		return
	}

	errBuf := ""

	for _, result := range results {
		errBuf += jc.DeleteCommandResult(result.Id).Error()
	}

	if err != nil {
		err = fmt.Errorf("One or more deletes failed, err='%s'", err.Error())
		return
	}

	return
}

func findSystemsByOSType(systems []jcapi.JCSystem, osTypeRegEx string) (indices []int, err error) {
	r, err := regexp.Compile(osTypeRegEx)
	if err != nil {
		err = fmt.Errorf("Could not compile regex for '%s', err='%s'", osTypeRegEx, err.Error())
		return
	}

	for idx, system := range systems {
		if r.Match([]byte(system.Os)) && system.Active {
			indices = append(indices, idx)
		}
	}

	return
}

func getSystemNameMap(jc jcapi.JCAPI, idList []string) (systemNameMap map[string]jcapi.JCSystem, err error) {
	systemNameMap = make(map[string]jcapi.JCSystem)

	//
	// Walk the IDs on which we ran the command, and create a map of host name to
	// JCSystem object. This allows us to determine which command results go with
	// which commands.
	//
	for _, id := range idList {
		system, err2 := jc.GetSystemById(id, false)
		if err2 != nil {
			err = fmt.Errorf("Could not get system object for id '%s', err='%s'", id, err.Error())
			return
		}

		systemNameMap[system.Hostname] = system
	}

	return
}

func waitForAndProcessResults(jc jcapi.JCAPI, commandObj jcapi.JCCommand) (outputBuffer string, err error) {
	systemsFound := make(map[string]jcapi.JCResponse)

	systemNameMap, err := getSystemNameMap(jc, commandObj.Systems)
	if err != nil {
		err = fmt.Errorf("Could not map system names to systems, err='%s'", err.Error())
		return
	}

	pollTime := RESULT_POLL_TIME

	fmt.Printf("\nWaiting for results...")

	for i := 0; i < RESULT_MAX_POLL_TIME; i += pollTime {
		time.Sleep(time.Duration(pollTime) * time.Second)

		fmt.Printf(".")

		results, err2 := jc.GetCommandResultsByName(commandObj.Name)
		if err2 != nil {
			err = fmt.Errorf("Could not find the command result for '%s', err='%s'", commandObj.Name, err2.Error())
			return
		}

		// Walk the results and add their exit code to the map (maps system name to the result data)
		for _, result := range results {
			systemsFound[result.System] = result.Response
		}

		if len(results) == len(commandObj.Systems) {
			break
		}
	}

	fmt.Printf("done.\n\n")

	// Output the results for all systems that were successful
	headerBuffer := fmt.Sprintf("Systems with a zero exit code\n")
	headerBuffer += fmt.Sprintf("=============================\n")

	dataBuffer := ""

	for systemId, result := range systemsFound {
		if result.Data.ExitCode == 0 && result.Error == "" {
			dataBuffer += fmt.Sprintf("ID %s: Completed, exit code=%d, output=[%s]\n", systemId, result.Data.ExitCode, result.Data.Output)
		}
	}

	if dataBuffer != "" {
		outputBuffer = headerBuffer + dataBuffer
	}

	// Output the results for all systems that returned a non-zero exit code
	headerBuffer = fmt.Sprintf("\nSystems with non-zero exit code\n")
	headerBuffer += fmt.Sprintf("===============================\n")

	dataBuffer = ""

	for systemId, result := range systemsFound {
		if result.Data.ExitCode > 0 || result.Error != "" {
			dataBuffer += fmt.Sprintf("ID %s: Failed, exit code=%d, error=[%s]\n", systemId, result.Data.ExitCode, result.Error)
		}
	}

	if dataBuffer != "" {
		outputBuffer = headerBuffer + dataBuffer
	}

	// Output the results for all systems that did not return any result
	headerBuffer = fmt.Sprintf("\nSystems with no command result\n")
	headerBuffer += fmt.Sprintf("===============================\n")

	dataBuffer = ""

	for hostName, _ := range systemNameMap {
		if _, exists := systemsFound[hostName]; !exists {
			dataBuffer += fmt.Sprintf("ID %s (%s): Received no command result\n", systemNameMap[hostName].Id, hostName)
		}
	}

	if dataBuffer != "" {
		outputBuffer += headerBuffer + dataBuffer
	}

	return
}

//
// This example accepts an API key and the name of the file to execute commands against. You can then
// specify a regular expression for the JCSystem.Os value of the systems you want the command to run
// against.
//
// It will write the results of the commands to stdout when exitCode != 0
//
func main() {
	apiKey := flag.String("api-key", "none", "Your JumpCloud Administrator API Key")
	commandFile := flag.String("command-file", "myCommand.txt", "The name of a file containing a command to run")
	commandType := flag.String("command-type", "windows", "OS type of the command, 'linux', 'windows', or 'mac'")
	shell := flag.String("shell", "powershell", "Shell (Windows-only) (powershell/cmd)")
	osType := flag.String("os-type", "Windows.*", "A regular expression to match your systems for OS type")
	deleteFlag := flag.Bool("delete-after-run", false, "When true, delete commands and results at completion.")

	flag.Parse()

	jc := jcapi.NewJCAPI(*apiKey, URL_BASE)

	// Generate a randomized command name
	commandName := "CMD " + makeRandomString(COMMAND_NAME_RANDOM_PART)

	// Read the input file
	commandData, err := readFile(*commandFile, MAX_COMMAND_LEN)
	if err != nil {
		log.Fatalf("Could not read command, err='%s'", err.Error())
	}

	// Make the command object
	commandObj := makeImmediateCommand(commandName, commandData, *commandType, *shell, jcapi.COMMAND_ROOT_USER)

	//
	// Get the list of matching servers and add them to the command
	//
	systems, err := jc.GetSystems(false)
	if err != nil {
		log.Fatalf("Could not get a list of all systems, err='%s'")
	}

	indices, err := findSystemsByOSType(systems, *osType)
	if err != nil {
		log.Fatalf("Could not search a list of systems for OS type matching '%s', err='%s'", *osType, err.Error())
	}

	if len(indices) == 0 {
		log.Fatalf("No systems match '%s' on your JumpCloud account\n", *osType)
	}

	fmt.Printf("Executing Command on the Following Systems\n")
	fmt.Printf("------------------------------------------\n")

	for _, index := range indices {
		fmt.Printf("%s\t%s\n", systems[index].Id, systems[index].Hostname)

		commandObj.Systems = append(commandObj.Systems, systems[index].Id)
	}

	//
	// Add the command object to the JumpCloud account
	//
	commandObj, err = jc.AddUpdateCommand(jcapi.Insert, commandObj)
	if err != nil {
		log.Fatalf("Could not insert a new command, err='%s'", err.Error())
	}
	if *deleteFlag {
		defer jc.DeleteCommand(commandObj)
	}

	//
	// Run the command
	//
	err = jc.RunCommand(commandObj)
	if err != nil {
		log.Fatalf("Could not run the command '%s', err='%s'", commandObj.ToString(), err.Error())
	}
	if *deleteFlag {
		defer deleteCommandResultsByName(jc, commandName)
	}

	//
	// Wait for results to come back from each command executed, but don't wait forever if some of the results
	// don't come back. Write to stdout as results are returned, and after a specified timeout, output the list
	// of failed commands and commands that did not run within the specified timeout.
	//
	outputBuffer, err := waitForAndProcessResults(jc, commandObj)
	if err != nil {
		log.Fatalf("Could not wait for and process results, err='%s'", err.Error())
	}

	fmt.Printf(outputBuffer)

	os.Exit(0)
}
<|endoftext|>


<commit_before>package service

import (
	"io"
	"sync"
	"text/template"

	"github.com/builderscon/octav/octav/assets"
	"github.com/builderscon/octav/octav/gettext"
	pdebug "github.com/lestrrat/go-pdebug"
)

var templateSvc *TemplateSvc
var templateOnce sync.Once

func Template() *TemplateSvc {
	templateOnce.Do(templateSvc.Init)
	return templateSvc
}

func (v *TemplateSvc) Init() {
	if pdebug.Enabled {
		g := pdebug.Marker("service.Template.Init")
		defer g.End()
	}

	var t *template.Template

	var parsed int
	for _, n := range assets.AssetNames() {
		b, err := assets.Asset(n)
		if err != nil {
			panic(err.Error())
		}

		if t == nil {
			t = template.New(n).Funcs(map[string]interface{}{
				"gettext": gettext.Get,
			})
		}

		var tmpl *template.Template
		if n == t.Name() {
			tmpl = t
		} else {
			tmpl = t.New(n)
		}

		if pdebug.Enabled {
			pdebug.Printf("Parsing template %s", n)
		}
		if _, err := tmpl.Parse(string(b)); err != nil {
			panic(err.Error())
		}
		parsed++
	}

	if pdebug.Enabled {
		pdebug.Printf("Parsed %d templates", parsed)
	}

	*v = TemplateSvc{
		template: t,
	}
}

func (v *TemplateSvc) Execute(dst io.Writer, name string, vars interface{}) (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("service.Template.Execute").BindError(&err)
		defer g.End()
	}

	return v.template.ExecuteTemplate(dst, name, vars)
}
<commit_msg>see if this works<commit_after>package service

import (
	"io"
	"sync"
	"text/template"

	"github.com/builderscon/octav/octav/assets"
	"github.com/builderscon/octav/octav/gettext"
	pdebug "github.com/lestrrat/go-pdebug"
)

var templateSvc TemplateSvc
var templateOnce sync.Once

func Template() *TemplateSvc {
	templateOnce.Do(templateSvc.Init)
	return &templateSvc
}

func (v *TemplateSvc) Init() {
	if pdebug.Enabled {
		g := pdebug.Marker("service.Template.Init")
		defer g.End()
	}

	var t *template.Template

	var parsed int
	for _, n := range assets.AssetNames() {
		b, err := assets.Asset(n)
		if err != nil {
			panic(err.Error())
		}

		if t == nil {
			t = template.New(n).Funcs(map[string]interface{}{
				"gettext": gettext.Get,
			})
		}

		var tmpl *template.Template
		if n == t.Name() {
			tmpl = t
		} else {
			tmpl = t.New(n)
		}

		if pdebug.Enabled {
			pdebug.Printf("Parsing template %s", n)
		}
		if _, err := tmpl.Parse(string(b)); err != nil {
			panic(err.Error())
		}
		parsed++
	}

	if pdebug.Enabled {
		pdebug.Printf("Parsed %d templates", parsed)
	}

	v.template = t
}

func (v *TemplateSvc) Execute(dst io.Writer, name string, vars interface{}) (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("service.Template.Execute").BindError(&err)
		defer g.End()
	}

	return v.template.ExecuteTemplate(dst, name, vars)
}
<|endoftext|>


<commit_before>package db

import (
	"bytes"
	"encoding/binary"

	wirefs "bazil.org/bazil/fs/wire"
	"bazil.org/fuse"
	"github.com/boltdb/bolt"
	"github.com/golang/protobuf/proto"
)

type Dirs struct {
	b *bolt.Bucket
}

func dirKey(parentInode uint64, name string) []byte {
	buf := make([]byte, 8+len(name))
	binary.BigEndian.PutUint64(buf, parentInode)
	copy(buf[8:], name)
	return buf
}

func basename(dirKey []byte) []byte {
	return dirKey[8:]
}

// Get the entry in parent directory with the given name.
//
// Returned value is valid after the transaction.
func (b *Dirs) Get(parentInode uint64, name string) (*wirefs.Dirent, error) {
	key := dirKey(parentInode, name)
	buf := b.b.Get(key)
	if buf == nil {
		return nil, fuse.ENOENT
	}

	var de wirefs.Dirent
	if err := proto.Unmarshal(buf, &de); err != nil {
		return nil, err
	}
	return &de, nil
}

// Put an entry in parent directory with the given name.
func (b *Dirs) Put(parentInode uint64, name string, de *wirefs.Dirent) error {
	buf, err := proto.Marshal(de)
	if err != nil {
		return err
	}
	key := dirKey(parentInode, name)
	if err := b.b.Put(key, buf); err != nil {
		return err
	}
	return nil
}

// Delete the entry in parent directory with the given name.
//
// Returns fuse.ENOENT if an entry does not exist.
func (b *Dirs) Delete(parentInode uint64, name string) error {
	key := dirKey(parentInode, name)
	if b.b.Get(key) == nil {
		return fuse.ENOENT
	}
	if err := b.b.Delete(key); err != nil {
		return err
	}
	return nil
}

// TombstoneCreate marks the entry in parent directory with the given
// name as removed, but leaves a hint of its existence. It creates a
// new entry if one does not exist yet.
func (b *Dirs) TombstoneCreate(parentInode uint64, name string) error {
	de := &wirefs.Dirent{Tombstone: &wirefs.Tombstone{}}
	if err := b.Put(parentInode, name, de); err != nil {
		return err
	}
	return nil
}

// Rename renames an entry in the parent directory from oldName to
// newName.
//
// Returns the overwritten entry, or nil.
func (b *Dirs) Rename(parentInode uint64, oldName string, newName string) (*DirEntry, error) {
	keyOld := dirKey(parentInode, oldName)
	keyNew := dirKey(parentInode, newName)

	bufOld := b.b.Get(keyOld)
	if bufOld == nil {
		return nil, fuse.ENOENT
	}

	// the file getting overwritten
	var loser *DirEntry
	if buf := b.b.Get(keyNew); buf != nil {
		// overwriting
		loser = &DirEntry{
			name: basename(keyNew),
			data: buf,
		}
	}

	if err := b.b.Put(keyNew, bufOld); err != nil {
		return nil, err
	}
	tombDE := &wirefs.Dirent{Tombstone: &wirefs.Tombstone{}}
	tombBuf, err := proto.Marshal(tombDE)
	if err != nil {
		return nil, err
	}
	if err := b.b.Put(keyOld, tombBuf); err != nil {
		return nil, err
	}

	return loser, nil
}

func (b *Dirs) List(inode uint64) *DirsCursor {
	c := b.b.Cursor()
	prefix := dirKey(inode, "")
	return &DirsCursor{
		inode:  inode,
		prefix: prefix,
		c:      c,
	}
}

type DirsCursor struct {
	inode  uint64
	prefix []byte
	c      *bolt.Cursor
}

func (c *DirsCursor) First() *DirEntry {
	return c.item(c.c.Seek(c.prefix))
}

func (c *DirsCursor) Next() *DirEntry {
	return c.item(c.c.Next())
}

// Seek to first name equal to name, or the next one if exact match is
// not found.
//
// Passing an empty name seeks to the beginning of the directory.
func (c *DirsCursor) Seek(name string) *DirEntry {
	k := make([]byte, 0, len(c.prefix)+len(name))
	k = append(k, c.prefix...)
	k = append(k, name...)
	return c.item(c.c.Seek(k))
}

func (c *DirsCursor) item(k, v []byte) *DirEntry {
	if !bytes.HasPrefix(k, c.prefix) {
		// past the end of the directory
		return nil
	}
	name := k[len(c.prefix):]
	return &DirEntry{name: name, data: v}
}

type DirEntry struct {
	name []byte
	data []byte
}

// Name returns the basename of this directory entry.
//
// name is valid after the transaction.
func (e *DirEntry) Name() string {
	return string(e.name)
}

// Unmarshal the directory entry to out.
//
// out is valid after the transaction.
func (e *DirEntry) Unmarshal(out *wirefs.Dirent) error {
	if err := proto.Unmarshal(e.data, out); err != nil {
		return err
	}
	return nil
}
<commit_msg>db: Support marking dir entries as tombstones<commit_after>package db

import (
	"bytes"
	"encoding/binary"

	wirefs "bazil.org/bazil/fs/wire"
	"bazil.org/fuse"
	"github.com/boltdb/bolt"
	"github.com/golang/protobuf/proto"
)

type Dirs struct {
	b *bolt.Bucket
}

func dirKey(parentInode uint64, name string) []byte {
	buf := make([]byte, 8+len(name))
	binary.BigEndian.PutUint64(buf, parentInode)
	copy(buf[8:], name)
	return buf
}

func basename(dirKey []byte) []byte {
	return dirKey[8:]
}

// Get the entry in parent directory with the given name.
//
// Returned value is valid after the transaction.
func (b *Dirs) Get(parentInode uint64, name string) (*wirefs.Dirent, error) {
	key := dirKey(parentInode, name)
	buf := b.b.Get(key)
	if buf == nil {
		return nil, fuse.ENOENT
	}

	var de wirefs.Dirent
	if err := proto.Unmarshal(buf, &de); err != nil {
		return nil, err
	}
	return &de, nil
}

// Put an entry in parent directory with the given name.
func (b *Dirs) Put(parentInode uint64, name string, de *wirefs.Dirent) error {
	buf, err := proto.Marshal(de)
	if err != nil {
		return err
	}
	key := dirKey(parentInode, name)
	if err := b.b.Put(key, buf); err != nil {
		return err
	}
	return nil
}

// Delete the entry in parent directory with the given name.
//
// Returns fuse.ENOENT if an entry does not exist.
func (b *Dirs) Delete(parentInode uint64, name string) error {
	key := dirKey(parentInode, name)
	if b.b.Get(key) == nil {
		return fuse.ENOENT
	}
	if err := b.b.Delete(key); err != nil {
		return err
	}
	return nil
}

// Tombstone marks the entry in parent directory with the given name
// as removed, but leaves a hint of its existence.
//
// Returns fuse.ENOENT if an entry does not exist.
func (b *Dirs) Tombstone(parentInode uint64, name string) error {
	key := dirKey(parentInode, name)
	if b.b.Get(key) == nil {
		return fuse.ENOENT
	}
	de := &wirefs.Dirent{Tombstone: &wirefs.Tombstone{}}
	if err := b.Put(parentInode, name, de); err != nil {
		return err
	}
	return nil
}

// TombstoneCreate marks the entry in parent directory with the given
// name as removed, but leaves a hint of its existence. It creates a
// new entry if one does not exist yet.
func (b *Dirs) TombstoneCreate(parentInode uint64, name string) error {
	de := &wirefs.Dirent{Tombstone: &wirefs.Tombstone{}}
	if err := b.Put(parentInode, name, de); err != nil {
		return err
	}
	return nil
}

// Rename renames an entry in the parent directory from oldName to
// newName.
//
// Returns the overwritten entry, or nil.
func (b *Dirs) Rename(parentInode uint64, oldName string, newName string) (*DirEntry, error) {
	keyOld := dirKey(parentInode, oldName)
	keyNew := dirKey(parentInode, newName)

	bufOld := b.b.Get(keyOld)
	if bufOld == nil {
		return nil, fuse.ENOENT
	}

	// the file getting overwritten
	var loser *DirEntry
	if buf := b.b.Get(keyNew); buf != nil {
		// overwriting
		loser = &DirEntry{
			name: basename(keyNew),
			data: buf,
		}
	}

	if err := b.b.Put(keyNew, bufOld); err != nil {
		return nil, err
	}
	tombDE := &wirefs.Dirent{Tombstone: &wirefs.Tombstone{}}
	tombBuf, err := proto.Marshal(tombDE)
	if err != nil {
		return nil, err
	}
	if err := b.b.Put(keyOld, tombBuf); err != nil {
		return nil, err
	}

	return loser, nil
}

func (b *Dirs) List(inode uint64) *DirsCursor {
	c := b.b.Cursor()
	prefix := dirKey(inode, "")
	return &DirsCursor{
		inode:  inode,
		prefix: prefix,
		c:      c,
	}
}

type DirsCursor struct {
	inode  uint64
	prefix []byte
	c      *bolt.Cursor
}

func (c *DirsCursor) First() *DirEntry {
	return c.item(c.c.Seek(c.prefix))
}

func (c *DirsCursor) Next() *DirEntry {
	return c.item(c.c.Next())
}

// Seek to first name equal to name, or the next one if exact match is
// not found.
//
// Passing an empty name seeks to the beginning of the directory.
func (c *DirsCursor) Seek(name string) *DirEntry {
	k := make([]byte, 0, len(c.prefix)+len(name))
	k = append(k, c.prefix...)
	k = append(k, name...)
	return c.item(c.c.Seek(k))
}

func (c *DirsCursor) item(k, v []byte) *DirEntry {
	if !bytes.HasPrefix(k, c.prefix) {
		// past the end of the directory
		return nil
	}
	name := k[len(c.prefix):]
	return &DirEntry{name: name, data: v}
}

type DirEntry struct {
	name []byte
	data []byte
}

// Name returns the basename of this directory entry.
//
// name is valid after the transaction.
func (e *DirEntry) Name() string {
	return string(e.name)
}

// Unmarshal the directory entry to out.
//
// out is valid after the transaction.
func (e *DirEntry) Unmarshal(out *wirefs.Dirent) error {
	if err := proto.Unmarshal(e.data, out); err != nil {
		return err
	}
	return nil
}
<|endoftext|>