<commit_before>// Copyright (c) 2017 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package indexers

import (
	"errors"

	"github.com/btcsuite/fastsha256"
	"github.com/roasbeef/btcd/blockchain"
	"github.com/roasbeef/btcd/chaincfg"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcd/database"
	"github.com/roasbeef/btcd/wire"
	"github.com/roasbeef/btcutil"
	"github.com/roasbeef/btcutil/gcs"
	"github.com/roasbeef/btcutil/gcs/builder"
)

const (
	// cfIndexName is the human-readable name for the index.
	cfIndexName = "committed filter index"
)

// Committed filters come in two flavours: basic and extended. They are
// generated and dropped in pairs, and both are indexed by a block's hash.
// Besides holding different content, they also live in different buckets.
var (
	// cfIndexParentBucketKey is the name of the parent bucket used to house
	// the index. The rest of the buckets live below this bucket.
	cfIndexParentBucketKey = []byte("cfindexparentbucket")

	// cfIndexKeys is an array of db bucket names used to house indexes of
	// block hashes to cfilters.
	cfIndexKeys = [][]byte{
		[]byte("cf0byhashidx"),
		[]byte("cf1byhashidx"),
	}

	// cfHeaderKeys is an array of db bucket names used to house indexes of
	// block hashes to cf headers.
	cfHeaderKeys = [][]byte{
		[]byte("cf0headerbyhashidx"),
		[]byte("cf1headerbyhashidx"),
	}

	maxFilterType = uint8(len(cfHeaderKeys) - 1)
)

// dbFetchFilter retrieves a block's basic or extended filter. A filter's
// absence is not considered an error.
func dbFetchFilter(dbTx database.Tx, key []byte, h *chainhash.Hash) ([]byte, error) {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Get(h[:]), nil
}

// dbFetchFilterHeader retrieves a block's basic or extended filter header.
// A filter's absence is not considered an error.
func dbFetchFilterHeader(dbTx database.Tx, key []byte, h *chainhash.Hash) ([]byte, error) {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)

	fh := idx.Get(h[:])
	if len(fh) != fastsha256.Size {
		return nil, errors.New("invalid filter header length")
	}

	return fh, nil
}

// dbStoreFilter stores a block's basic or extended filter.
func dbStoreFilter(dbTx database.Tx, key []byte, h *chainhash.Hash, f []byte) error {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Put(h[:], f)
}

// dbStoreFilterHeader stores a block's basic or extended filter header.
func dbStoreFilterHeader(dbTx database.Tx, key []byte, h *chainhash.Hash, fh []byte) error {
	if len(fh) != fastsha256.Size {
		return errors.New("invalid filter header length")
	}
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Put(h[:], fh)
}

// dbDeleteFilter deletes a filter's basic or extended filter.
func dbDeleteFilter(dbTx database.Tx, key []byte, h *chainhash.Hash) error {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Delete(h[:])
}

// dbDeleteFilterHeader deletes a filter's basic or extended filter header.
func dbDeleteFilterHeader(dbTx database.Tx, key []byte, h *chainhash.Hash) error {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Delete(h[:])
}

// CfIndex implements a committed filter (cf) by hash index.
type CfIndex struct {
	db          database.DB
	chainParams *chaincfg.Params
}

// Ensure the CfIndex type implements the Indexer interface.
var _ Indexer = (*CfIndex)(nil)

// Init initializes the hash-based cf index. This is part of the Indexer
// interface.
func (idx *CfIndex) Init() error {
	return nil // Nothing to do.
}

// Key returns the database key to use for the index as a byte slice. This is
// part of the Indexer interface.
func (idx *CfIndex) Key() []byte {
	return cfIndexParentBucketKey
}

// Name returns the human-readable name of the index. This is part of the
// Indexer interface.
func (idx *CfIndex) Name() string {
	return cfIndexName
}

// Create is invoked when the indexer manager determines the index needs to
// be created for the first time. It creates buckets for the two hash-based cf
// indexes (simple, extended).
func (idx *CfIndex) Create(dbTx database.Tx) error {
	meta := dbTx.Metadata()

	cfIndexParentBucket, err := meta.CreateBucket(cfIndexParentBucketKey)
	if err != nil {
		return err
	}

	for _, bucketName := range cfIndexKeys {
		_, err = cfIndexParentBucket.CreateBucket(bucketName)
		if err != nil {
			return err
		}
	}

	for _, bucketName := range cfHeaderKeys {
		_, err = cfIndexParentBucket.CreateBucket(bucketName)
		if err != nil {
			return err
		}
	}

	firstHeader := make([]byte, chainhash.HashSize)
	err = dbStoreFilterHeader(
		dbTx,
		cfHeaderKeys[wire.GCSFilterRegular],
		&idx.chainParams.GenesisBlock.Header.PrevBlock,
		firstHeader,
	)
	if err != nil {
		return err
	}

	return dbStoreFilterHeader(
		dbTx,
		cfHeaderKeys[wire.GCSFilterExtended],
		&idx.chainParams.GenesisBlock.Header.PrevBlock,
		firstHeader,
	)
}

// storeFilter stores a given filter, and performs the steps needed to
// generate the filter's header.
func storeFilter(dbTx database.Tx, block *btcutil.Block, f *gcs.Filter,
	filterType wire.FilterType) error {
	if uint8(filterType) > maxFilterType {
		return errors.New("unsupported filter type")
	}

	// Figure out which buckets to use.
	fkey := cfIndexKeys[filterType]
	hkey := cfHeaderKeys[filterType]

	// Start by storing the filter.
	h := block.Hash()
	var basicFilterBytes []byte
	if f != nil {
		basicFilterBytes = f.NBytes()
	}
	err := dbStoreFilter(dbTx, fkey, h, basicFilterBytes)
	if err != nil {
		return err
	}

	// Then fetch the previous block's filter header.
	ph := &block.MsgBlock().Header.PrevBlock
	pfh, err := dbFetchFilterHeader(dbTx, hkey, ph)
	if err != nil {
		return err
	}

	// Construct the new block's filter header, and store it.
	prevHeader, err := chainhash.NewHash(pfh)
	if err != nil {
		return err
	}
	fh := builder.MakeHeaderForFilter(f, *prevHeader)
	return dbStoreFilterHeader(dbTx, hkey, h, fh[:])
}

// ConnectBlock is invoked by the index manager when a new block has been
// connected to the main chain. This indexer adds a hash-to-cf mapping for
// every passed block. This is part of the Indexer interface.
func (idx *CfIndex) ConnectBlock(dbTx database.Tx, block *btcutil.Block,
	view *blockchain.UtxoViewpoint) error {

	f, err := builder.BuildBasicFilter(block.MsgBlock())
	if err != nil && err != gcs.ErrNoData {
		return err
	}

	if err := storeFilter(dbTx, block, f,
		wire.GCSFilterRegular); err != nil {
		return err
	}

	f, err = builder.BuildExtFilter(block.MsgBlock())
	if err != nil && err != gcs.ErrNoData {
		return err
	}

	return storeFilter(dbTx, block, f, wire.GCSFilterExtended)
}

// DisconnectBlock is invoked by the index manager when a block has been
// disconnected from the main chain.  This indexer removes the hash-to-cf
// mapping for every passed block. This is part of the Indexer interface.
func (idx *CfIndex) DisconnectBlock(dbTx database.Tx, block *btcutil.Block,
	view *blockchain.UtxoViewpoint) error {

	for _, key := range cfIndexKeys {
		err := dbDeleteFilter(dbTx, key, block.Hash())
		if err != nil {
			return err
		}
	}

	for _, key := range cfHeaderKeys {
		err := dbDeleteFilterHeader(dbTx, key, block.Hash())
		if err != nil {
			return err
		}
	}

	return nil
}

// FilterByBlockHash returns the serialized contents of a block's basic or
// extended committed filter.
func (idx *CfIndex) FilterByBlockHash(h *chainhash.Hash,
	filterType wire.FilterType) ([]byte, error) {
	var f []byte
	err := idx.db.View(func(dbTx database.Tx) error {
		if uint8(filterType) > maxFilterType {
			return errors.New("unsupported filter type")
		}

		var err error
		f, err = dbFetchFilter(dbTx, cfIndexKeys[filterType], h)
		return err
	})
	return f, err
}

// FilterHeaderByBlockHash returns the serialized contents of a block's basic
// or extended committed filter header.
func (idx *CfIndex) FilterHeaderByBlockHash(h *chainhash.Hash,
	filterType wire.FilterType) ([]byte, error) {
	var fh []byte
	err := idx.db.View(func(dbTx database.Tx) error {
		if uint8(filterType) > maxFilterType {
			return errors.New("unsupported filter type")
		}

		var err error
		fh, err = dbFetchFilterHeader(dbTx,
			cfHeaderKeys[filterType], h)
		return err
	})
	return fh, err
}

// NewCfIndex returns a new instance of an indexer that is used to create a
// mapping of the hashes of all blocks in the blockchain to their respective
// committed filters.
//
// It implements the Indexer interface which plugs into the IndexManager that
// in turn is used by the blockchain package. This allows the index to be
// seamlessly maintained along with the chain.
func NewCfIndex(db database.DB, chainParams *chaincfg.Params) *CfIndex {
	return &CfIndex{db: db, chainParams: chainParams}
}

// DropCfIndex drops the CF index from the provided database if exists.
func DropCfIndex(db database.DB, interrupt <-chan struct{}) error {
	return dropIndex(db, cfIndexParentBucketKey, cfIndexName, interrupt)
}
<commit_msg>glide: Upgrade btcutil.<commit_after>// Copyright (c) 2017 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package indexers

import (
	"errors"

	"github.com/btcsuite/fastsha256"
	"github.com/roasbeef/btcd/blockchain"
	"github.com/roasbeef/btcd/chaincfg"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcd/database"
	"github.com/roasbeef/btcd/wire"
	"github.com/roasbeef/btcutil"
	"github.com/roasbeef/btcutil/gcs"
	"github.com/roasbeef/btcutil/gcs/builder"
)

const (
	// cfIndexName is the human-readable name for the index.
	cfIndexName = "committed filter index"
)

// Committed filters come in two flavours: basic and extended. They are
// generated and dropped in pairs, and both are indexed by a block's hash.
// Besides holding different content, they also live in different buckets.
var (
	// cfIndexParentBucketKey is the name of the parent bucket used to house
	// the index. The rest of the buckets live below this bucket.
	cfIndexParentBucketKey = []byte("cfindexparentbucket")

	// cfIndexKeys is an array of db bucket names used to house indexes of
	// block hashes to cfilters.
	cfIndexKeys = [][]byte{
		[]byte("cf0byhashidx"),
		[]byte("cf1byhashidx"),
	}

	// cfHeaderKeys is an array of db bucket names used to house indexes of
	// block hashes to cf headers.
	cfHeaderKeys = [][]byte{
		[]byte("cf0headerbyhashidx"),
		[]byte("cf1headerbyhashidx"),
	}

	maxFilterType = uint8(len(cfHeaderKeys) - 1)
)

// dbFetchFilter retrieves a block's basic or extended filter. A filter's
// absence is not considered an error.
func dbFetchFilter(dbTx database.Tx, key []byte, h *chainhash.Hash) ([]byte, error) {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Get(h[:]), nil
}

// dbFetchFilterHeader retrieves a block's basic or extended filter header.
// A filter's absence is not considered an error.
func dbFetchFilterHeader(dbTx database.Tx, key []byte, h *chainhash.Hash) ([]byte, error) {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)

	fh := idx.Get(h[:])
	if len(fh) != fastsha256.Size {
		return nil, errors.New("invalid filter header length")
	}

	return fh, nil
}

// dbStoreFilter stores a block's basic or extended filter.
func dbStoreFilter(dbTx database.Tx, key []byte, h *chainhash.Hash, f []byte) error {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Put(h[:], f)
}

// dbStoreFilterHeader stores a block's basic or extended filter header.
func dbStoreFilterHeader(dbTx database.Tx, key []byte, h *chainhash.Hash, fh []byte) error {
	if len(fh) != fastsha256.Size {
		return errors.New("invalid filter header length")
	}
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Put(h[:], fh)
}

// dbDeleteFilter deletes a filter's basic or extended filter.
func dbDeleteFilter(dbTx database.Tx, key []byte, h *chainhash.Hash) error {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Delete(h[:])
}

// dbDeleteFilterHeader deletes a filter's basic or extended filter header.
func dbDeleteFilterHeader(dbTx database.Tx, key []byte, h *chainhash.Hash) error {
	idx := dbTx.Metadata().Bucket(cfIndexParentBucketKey).Bucket(key)
	return idx.Delete(h[:])
}

// CfIndex implements a committed filter (cf) by hash index.
type CfIndex struct {
	db          database.DB
	chainParams *chaincfg.Params
}

// Ensure the CfIndex type implements the Indexer interface.
var _ Indexer = (*CfIndex)(nil)

// Init initializes the hash-based cf index. This is part of the Indexer
// interface.
func (idx *CfIndex) Init() error {
	return nil // Nothing to do.
}

// Key returns the database key to use for the index as a byte slice. This is
// part of the Indexer interface.
func (idx *CfIndex) Key() []byte {
	return cfIndexParentBucketKey
}

// Name returns the human-readable name of the index. This is part of the
// Indexer interface.
func (idx *CfIndex) Name() string {
	return cfIndexName
}

// Create is invoked when the indexer manager determines the index needs to
// be created for the first time. It creates buckets for the two hash-based cf
// indexes (simple, extended).
func (idx *CfIndex) Create(dbTx database.Tx) error {
	meta := dbTx.Metadata()

	cfIndexParentBucket, err := meta.CreateBucket(cfIndexParentBucketKey)
	if err != nil {
		return err
	}

	for _, bucketName := range cfIndexKeys {
		_, err = cfIndexParentBucket.CreateBucket(bucketName)
		if err != nil {
			return err
		}
	}

	for _, bucketName := range cfHeaderKeys {
		_, err = cfIndexParentBucket.CreateBucket(bucketName)
		if err != nil {
			return err
		}
	}

	firstHeader := make([]byte, chainhash.HashSize)
	err = dbStoreFilterHeader(
		dbTx,
		cfHeaderKeys[wire.GCSFilterRegular],
		&idx.chainParams.GenesisBlock.Header.PrevBlock,
		firstHeader,
	)
	if err != nil {
		return err
	}

	return dbStoreFilterHeader(
		dbTx,
		cfHeaderKeys[wire.GCSFilterExtended],
		&idx.chainParams.GenesisBlock.Header.PrevBlock,
		firstHeader,
	)
}

// storeFilter stores a given filter, and performs the steps needed to
// generate the filter's header.
func storeFilter(dbTx database.Tx, block *btcutil.Block, f *gcs.Filter,
	filterType wire.FilterType) error {
	if uint8(filterType) > maxFilterType {
		return errors.New("unsupported filter type")
	}

	// Figure out which buckets to use.
	fkey := cfIndexKeys[filterType]
	hkey := cfHeaderKeys[filterType]

	// Start by storing the filter.
	h := block.Hash()
	filterBytes, err := f.NBytes()
	if err != nil {
		return err
	}
	err = dbStoreFilter(dbTx, fkey, h, filterBytes)
	if err != nil {
		return err
	}

	// Then fetch the previous block's filter header.
	ph := &block.MsgBlock().Header.PrevBlock
	pfh, err := dbFetchFilterHeader(dbTx, hkey, ph)
	if err != nil {
		return err
	}

	// Construct the new block's filter header, and store it.
	prevHeader, err := chainhash.NewHash(pfh)
	if err != nil {
		return err
	}
	fh, err := builder.MakeHeaderForFilter(f, *prevHeader)
	if err != nil {
		return err
	}
	return dbStoreFilterHeader(dbTx, hkey, h, fh[:])
}

// ConnectBlock is invoked by the index manager when a new block has been
// connected to the main chain. This indexer adds a hash-to-cf mapping for
// every passed block. This is part of the Indexer interface.
func (idx *CfIndex) ConnectBlock(dbTx database.Tx, block *btcutil.Block,
	view *blockchain.UtxoViewpoint) error {

	f, err := builder.BuildBasicFilter(block.MsgBlock())
	if err != nil {
		return err
	}

	if err := storeFilter(dbTx, block, f,
		wire.GCSFilterRegular); err != nil {
		return err
	}

	f, err = builder.BuildExtFilter(block.MsgBlock())
	if err != nil {
		return err
	}

	return storeFilter(dbTx, block, f, wire.GCSFilterExtended)
}

// DisconnectBlock is invoked by the index manager when a block has been
// disconnected from the main chain.  This indexer removes the hash-to-cf
// mapping for every passed block. This is part of the Indexer interface.
func (idx *CfIndex) DisconnectBlock(dbTx database.Tx, block *btcutil.Block,
	view *blockchain.UtxoViewpoint) error {

	for _, key := range cfIndexKeys {
		err := dbDeleteFilter(dbTx, key, block.Hash())
		if err != nil {
			return err
		}
	}

	for _, key := range cfHeaderKeys {
		err := dbDeleteFilterHeader(dbTx, key, block.Hash())
		if err != nil {
			return err
		}
	}

	return nil
}

// FilterByBlockHash returns the serialized contents of a block's basic or
// extended committed filter.
func (idx *CfIndex) FilterByBlockHash(h *chainhash.Hash,
	filterType wire.FilterType) ([]byte, error) {
	var f []byte
	err := idx.db.View(func(dbTx database.Tx) error {
		if uint8(filterType) > maxFilterType {
			return errors.New("unsupported filter type")
		}

		var err error
		f, err = dbFetchFilter(dbTx, cfIndexKeys[filterType], h)
		return err
	})
	return f, err
}

// FilterHeaderByBlockHash returns the serialized contents of a block's basic
// or extended committed filter header.
func (idx *CfIndex) FilterHeaderByBlockHash(h *chainhash.Hash,
	filterType wire.FilterType) ([]byte, error) {
	var fh []byte
	err := idx.db.View(func(dbTx database.Tx) error {
		if uint8(filterType) > maxFilterType {
			return errors.New("unsupported filter type")
		}

		var err error
		fh, err = dbFetchFilterHeader(dbTx,
			cfHeaderKeys[filterType], h)
		return err
	})
	return fh, err
}

// NewCfIndex returns a new instance of an indexer that is used to create a
// mapping of the hashes of all blocks in the blockchain to their respective
// committed filters.
//
// It implements the Indexer interface which plugs into the IndexManager that
// in turn is used by the blockchain package. This allows the index to be
// seamlessly maintained along with the chain.
func NewCfIndex(db database.DB, chainParams *chaincfg.Params) *CfIndex {
	return &CfIndex{db: db, chainParams: chainParams}
}

// DropCfIndex drops the CF index from the provided database if exists.
func DropCfIndex(db database.DB, interrupt <-chan struct{}) error {
	return dropIndex(db, cfIndexParentBucketKey, cfIndexName, interrupt)
}
<|endoftext|>


<commit_before>package bootstrapping

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	singleinstance "github.com/allan-simon/go-singleinstance"
	"github.com/codegangsta/cli"
	"github.com/ingrammicro/concerto/api/blueprint"
	"github.com/ingrammicro/concerto/api/types"
	"github.com/ingrammicro/concerto/cmd"
	"github.com/ingrammicro/concerto/utils"
	"github.com/ingrammicro/concerto/utils/format"
)

const (
	//DefaultTimingInterval Default period for looping
	DefaultTimingInterval = 600 // 600 seconds = 10 minutes
	DefaultTimingSplay    = 360 // seconds
	DefaultThresholdLines = 10
	ProcessIDFile         = "imco-bootstrapping.pid"
	RetriesNumber         = 5
)

type bootstrappingProcess struct {
	startedAt                    time.Time
	finishedAt                   time.Time
	policyfiles                  []policyfile
	attributes                   attributes
	thresholdLines               int
	directoryPath                string
	appliedPolicyfileRevisionIDs map[string]string
}
type attributes struct {
	revisionID string
	rawData    *json.RawMessage
}

type policyfile types.BootstrappingPolicyfile

var allPolicyfilesSuccessfullyApplied bool

// Handle signals
func handleSysSignals(cancelFunc context.CancelFunc) {
	log.Debug("handleSysSignals")

	gracefulStop := make(chan os.Signal, 1)
	signal.Notify(gracefulStop, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)
	log.Debug("Ending, signal detected:", <-gracefulStop)
	cancelFunc()
}

// Returns the full path to the tmp directory joined with pid management file name
func getProcessIDFilePath() string {
	return strings.Join([]string{os.TempDir(), string(os.PathSeparator), ProcessIDFile}, "")
}

// Returns the full path to the tmp directory
func generateWorkspaceDir() (string, error) {
	dir := filepath.Join(os.TempDir(), "imco")
	dirInfo, err := os.Stat(dir)
	if err != nil {
		err := os.Mkdir(dir, 0777)
		if err != nil {
			return "", err
		}
	} else {
		if !dirInfo.Mode().IsDir() {
			return "", fmt.Errorf("%s exists but is not a directory", dir)
		}
	}
	return dir, nil
}

// Start the bootstrapping process
func start(c *cli.Context) error {
	log.Debug("start")

	// TODO: replace /etc/imco with a directory taken from configuration/that depends on OS
	lockFile, err := singleinstance.CreateLockFile(filepath.Join("/etc/imco", "imco-bootstrapping.lock"))
	if err != nil {
		return err
	}
	defer lockFile.Close()

	formatter := format.GetFormatter()
	if err := utils.SetProcessIdToFile(getProcessIDFilePath()); err != nil {
		formatter.PrintFatal("cannot create the pid file", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go handleSysSignals(cancel)

	timingInterval := c.Int64("interval")
	if !(timingInterval > 0) {
		timingInterval = DefaultTimingInterval
	}

	timingSplay := c.Int64("splay")
	if !(timingSplay > 0) {
		timingSplay = DefaultTimingSplay
	}

	thresholdLines := c.Int("lines")
	if !(thresholdLines > 0) {
		thresholdLines = DefaultThresholdLines
	}
	log.Debug("routine lines threshold: ", thresholdLines)

	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	bootstrappingSvc, formatter := cmd.WireUpBootstrapping(c)
	for {
		applyPolicyfiles(bootstrappingSvc, formatter, thresholdLines)

		// Sleep for a configured amount of time plus a random amount of time (10 minutes plus 0 to 5 minutes, for instance)
		ticker := time.NewTicker(time.Duration(timingInterval+int64(r.Intn(int(timingSplay)))) * time.Second)

		select {
		case <-ticker.C:
			log.Debug("ticker")
		case <-ctx.Done():
			log.Debug(ctx.Err())
			log.Debug("closing bootstrapping")
		}
		ticker.Stop()
		if ctx.Err() != nil {
			break
		}
	}

	return nil
}

// Stop the bootstrapping process
func stop(c *cli.Context) error {
	log.Debug("cmdStop")

	formatter := format.GetFormatter()
	if err := utils.StopProcess(getProcessIDFilePath()); err != nil {
		formatter.PrintFatal("cannot stop the bootstrapping process", err)
	}

	log.Info("Bootstrapping routine successfully stopped")
	return nil
}

// Subsidiary routine for commands processing
func applyPolicyfiles(bootstrappingSvc *blueprint.BootstrappingService, formatter format.Formatter, thresholdLines int) error {
	log.Debug("applyPolicyfiles")

	// Inquire about desired configuration changes to be applied by querying the `GET /blueprint/configuration` endpoint. This will provide a JSON response with the desired configuration changes
	bsConfiguration, status, err := bootstrappingSvc.GetBootstrappingConfiguration()
	if err == nil && status != 200 {
		err = fmt.Errorf("received non-ok %d response", status)
	}
	if err != nil {
		formatter.PrintError("couldn't receive bootstrapping data", err)
		return err
	}
	dir, err := generateWorkspaceDir()
	if err != nil {
		return err
	}
	bsProcess := &bootstrappingProcess{
		startedAt:                    time.Now().UTC(),
		thresholdLines:               thresholdLines,
		directoryPath:                dir,
		appliedPolicyfileRevisionIDs: make(map[string]string),
	}

	// proto structures
	err = initializePrototype(bsConfiguration, bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	// For every policyfile, ensure its tarball (downloadable through their download_url) has been downloaded to the server ...
	err = downloadPolicyfiles(bootstrappingSvc, bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	//... and clean off any tarball that is no longer needed.
	err = cleanObsoletePolicyfiles(bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	// Store the attributes as JSON in a file with name `attrs-<attribute_revision_id>.json`
	err = saveAttributes(bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	// Process tarballs policies
	err = processPolicyfiles(bootstrappingSvc, bsProcess)
	// Finishing time
	bsProcess.finishedAt = time.Now().UTC()

	// Inform the platform of applied changes via a `PUT /blueprint/applied_configuration` request with a JSON payload similar to
	log.Debug("reporting applied policy files")
	reportErr := reportAppliedConfiguration(bootstrappingSvc, bsProcess)
	if reportErr != nil {
		log.Debug(err)
		return err
	}
	if err != nil {
		return err
	}
	return completeBootstrappingSequence(bsProcess)
}

func initializePrototype(bsConfiguration *types.BootstrappingConfiguration, bsProcess *bootstrappingProcess) error {
	log.Debug("initializePrototype")

	// Attributes
	bsProcess.attributes.revisionID = bsConfiguration.AttributeRevisionID
	bsProcess.attributes.rawData = bsConfiguration.Attributes

	// Policies
	for _, bsConfPolicyfile := range bsConfiguration.Policyfiles {
		bsProcess.policyfiles = append(bsProcess.policyfiles, policyfile(bsConfPolicyfile))
	}
	log.Debug(bsProcess)
	return nil
}

// downloadPolicyfiles For every policy file, ensure its tarball (downloadable through their download_url) has been downloaded to the server ...
func downloadPolicyfiles(bootstrappingSvc *blueprint.BootstrappingService, bsProcess *bootstrappingProcess) error {
	log.Debug("downloadPolicyfiles")

	for _, bsPolicyfile := range bsProcess.policyfiles {
		tarballPath := bsPolicyfile.TarballPath(bsProcess.directoryPath)
		log.Debug("downloading: ", tarballPath)
		queryURL, err := bsPolicyfile.QueryURL()
		if err != nil {
			return err
		}
		_, status, err := bootstrappingSvc.DownloadPolicyfile(queryURL, tarballPath)
		if err == nil && status != 200 {
			err = fmt.Errorf("obtained non-ok response when downloading policyfile %s", queryURL)
		}
		if err != nil {
			return err
		}
		if err = utils.Untar(tarballPath, bsPolicyfile.Path(bsProcess.directoryPath)); err != nil {
			return err
		}
	}
	return nil
}

// cleanObsoletePolicyfiles cleans off any tarball that is no longer needed.
func cleanObsoletePolicyfiles(bsProcess *bootstrappingProcess) error {
	log.Debug("cleanObsoletePolicyFiles")

	// evaluates working folder
	deletableFiles, err := ioutil.ReadDir(bsProcess.directoryPath)
	if err != nil {
		return err
	}

	// builds an array of currently processable files at this looping time
	currentlyProcessableFiles := []string{bsProcess.attributes.FileName()} // saved attributes file name
	for _, bsPolicyFile := range bsProcess.policyfiles {
		currentlyProcessableFiles = append(currentlyProcessableFiles, bsPolicyFile.FileName()) // Downloaded tgz file names
		currentlyProcessableFiles = append(currentlyProcessableFiles, bsPolicyFile.Name())     // Uncompressed folder names
	}

	// removes from deletableFiles array the policy files currently applied
	for _, f := range deletableFiles {
		if !utils.Contains(currentlyProcessableFiles, f.Name()) {
			log.Debug("removing: ", f.Name())
			if err := os.RemoveAll(filepath.Join(bsProcess.directoryPath, f.Name())); err != nil {
				return err
			}
		}
	}
	return nil
}

// saveAttributes stores the attributes as JSON in a file with name `attrs-<attribute_revision_id>.json`
func saveAttributes(bsProcess *bootstrappingProcess) error {
	log.Debug("saveAttributes")

	attrs, err := json.Marshal(bsProcess.attributes.rawData)
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(bsProcess.attributes.FilePath(bsProcess.directoryPath), attrs, 0600); err != nil {
		return err
	}
	return nil
}

// processPolicyfiles applies for each policy the required chef commands, reporting in bunches of N lines
func processPolicyfiles(bootstrappingSvc *blueprint.BootstrappingService, bsProcess *bootstrappingProcess) error {
	log.Debug("processPolicyfiles")

	for _, bsPolicyfile := range bsProcess.policyfiles {
		command := fmt.Sprintf("cd %s", bsPolicyfile.Path(bsProcess.directoryPath))
		if runtime.GOOS == "windows" {
			command = fmt.Sprintf("%s\nSET \"PATH=%PATH%;C:\\ruby\\bin;C:\\opscode\\chef\\bin;C:\\opscode\\chef\\embedded\\bin\"", command)
		}
		command = fmt.Sprintf("%s\nchef-client -z -j %s", command, bsProcess.attributes.FilePath(bsProcess.directoryPath))
		log.Debug(command)

		// Custom method for chunks processing
		fn := func(chunk string) error {
			log.Debug("sendChunks")
			err := utils.Retry(RetriesNumber, time.Second, func() error {
				log.Debug("Sending: ", chunk)

				commandIn := map[string]interface{}{
					"stdout": chunk,
				}

				_, statusCode, err := bootstrappingSvc.ReportBootstrappingLog(&commandIn)
				switch {
				// 0<100 error cases??
				case statusCode == 0:
					return fmt.Errorf("communication error %v %v", statusCode, err)
				case statusCode >= 500:
					return fmt.Errorf("server error %v %v", statusCode, err)
				case statusCode >= 400:
					return fmt.Errorf("client error %v %v", statusCode, err)
				default:
					return nil
				}
			})

			if err != nil {
				return fmt.Errorf("cannot send the chunk data, %v", err)
			}
			return nil
		}

		exitCode, err := utils.RunContinuousCmd(fn, command, -1, bsProcess.thresholdLines)
		if err == nil && exitCode != 0 {
			err = fmt.Errorf("policyfile application exited with %d code", exitCode)
		}
		if err != nil {
			return err
		}

		log.Info("completed: ", exitCode)
		bsProcess.appliedPolicyfileRevisionIDs[bsPolicyfile.ID] = bsPolicyfile.RevisionID
	}
	return nil
}

// reportAppliedConfiguration Inform the platform of applied changes
func reportAppliedConfiguration(bootstrappingSvc *blueprint.BootstrappingService, bsProcess *bootstrappingProcess) error {
	log.Debug("reportAppliedConfiguration")

	payload := map[string]interface{}{
		"started_at":              bsProcess.startedAt,
		"finished_at":             bsProcess.finishedAt,
		"policyfile_revision_ids": bsProcess.appliedPolicyfileRevisionIDs,
		"attribute_revision_id":   bsProcess.attributes.revisionID,
	}
	return bootstrappingSvc.ReportBootstrappingAppliedConfiguration(&payload)
}

// completeBootstrappingSequence evaluates if the first iteration of policies was completed; If case, execute the "scripts boot" command.
func completeBootstrappingSequence(bsProcess *bootstrappingProcess) error {
	log.Debug("completeBootstrappingSequence")

	if !allPolicyfilesSuccessfullyApplied {
		log.Debug("run the boot scripts")
		//run the boot scripts for the server by executing the scripts boot sub-command (as an external process).
		if output, exit, _, _ := utils.RunCmd(strings.Join([]string{os.Args[0], "scripts", "boot"}, " ")); exit != 0 {
			return fmt.Errorf("boot scripts run failed with exit code %d and following output: %s", exit, output)
		}
		allPolicyfilesSuccessfullyApplied = true
	}
	return nil
}

func (pf policyfile) Name() string {
	return strings.Join([]string{pf.ID, "-", pf.RevisionID}, "")
}

func (pf *policyfile) FileName() string {
	return strings.Join([]string{pf.Name(), "tgz"}, ".")
}

func (pf *policyfile) QueryURL() (string, error) {
	if pf.DownloadURL == "" {
		return "", fmt.Errorf("obtaining URL query: empty download URL")
	}
	url, err := url.Parse(pf.DownloadURL)
	if err != nil {
		return "", fmt.Errorf("parsing URL to extract query: %v", err)
	}
	return fmt.Sprintf("%s?%s", url.Path, url.RawQuery), nil
}

func (pf *policyfile) TarballPath(dir string) string {
	return filepath.Join(dir, pf.FileName())
}

func (pf *policyfile) Path(dir string) string {
	return filepath.Join(dir, pf.Name())
}

func (a *attributes) FileName() string {
	return fmt.Sprintf("attrs-%s.json", a.revisionID)
}

func (a *attributes) FilePath(dir string) string {
	return filepath.Join(dir, a.FileName())
}
<commit_msg>Refactored single instance management (issue #90)<commit_after>package bootstrapping

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/url"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	singleinstance "github.com/allan-simon/go-singleinstance"
	"github.com/codegangsta/cli"
	"github.com/ingrammicro/concerto/api/blueprint"
	"github.com/ingrammicro/concerto/api/types"
	"github.com/ingrammicro/concerto/cmd"
	"github.com/ingrammicro/concerto/utils"
	"github.com/ingrammicro/concerto/utils/format"
)

const (
	//DefaultTimingInterval Default period for looping
	DefaultTimingInterval = 600 // 600 seconds = 10 minutes
	DefaultTimingSplay    = 360 // seconds
	DefaultThresholdLines = 10
	ProcessLockFile       = "imco-bootstrapping.lock"
	RetriesNumber         = 5
)

type bootstrappingProcess struct {
	startedAt                    time.Time
	finishedAt                   time.Time
	policyfiles                  []policyfile
	attributes                   attributes
	thresholdLines               int
	directoryPath                string
	appliedPolicyfileRevisionIDs map[string]string
}
type attributes struct {
	revisionID string
	rawData    *json.RawMessage
}

var allPolicyfilesSuccessfullyApplied bool

type policyfile types.BootstrappingPolicyfile

func (pf policyfile) Name() string {
	return strings.Join([]string{pf.ID, "-", pf.RevisionID}, "")
}

func (pf *policyfile) FileName() string {
	return strings.Join([]string{pf.Name(), "tgz"}, ".")
}

func (pf *policyfile) QueryURL() (string, error) {
	if pf.DownloadURL == "" {
		return "", fmt.Errorf("obtaining URL query: empty download URL")
	}
	url, err := url.Parse(pf.DownloadURL)
	if err != nil {
		return "", fmt.Errorf("parsing URL to extract query: %v", err)
	}
	return fmt.Sprintf("%s?%s", url.Path, url.RawQuery), nil
}

func (pf *policyfile) TarballPath(dir string) string {
	return filepath.Join(dir, pf.FileName())
}

func (pf *policyfile) Path(dir string) string {
	return filepath.Join(dir, pf.Name())
}

func (a *attributes) FileName() string {
	return fmt.Sprintf("attrs-%s.json", a.revisionID)
}

func (a *attributes) FilePath(dir string) string {
	return filepath.Join(dir, a.FileName())
}

// Handle signals
func handleSysSignals(cancelFunc context.CancelFunc) {
	log.Debug("handleSysSignals")

	gracefulStop := make(chan os.Signal, 1)
	signal.Notify(gracefulStop, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)
	log.Debug("Ending, signal detected:", <-gracefulStop)
	cancelFunc()
}

// Returns the full path to the tmp directory joined with lock management file name
func getProcessLockFilePath() string {
	return filepath.Join(os.TempDir(), string(os.PathSeparator), ProcessLockFile)
}

// Returns the full path to the tmp directory
func generateWorkspaceDir() (string, error) {
	dir := filepath.Join(os.TempDir(), "imco")
	dirInfo, err := os.Stat(dir)
	if err != nil {
		err := os.Mkdir(dir, 0777)
		if err != nil {
			return "", err
		}
	} else {
		if !dirInfo.Mode().IsDir() {
			return "", fmt.Errorf("%s exists but is not a directory", dir)
		}
	}
	return dir, nil
}

// Start the bootstrapping process
func start(c *cli.Context) error {
	log.Debug("start")

	lockFile, err := singleinstance.CreateLockFile(getProcessLockFilePath())
	if err != nil {
		return err
	}
	defer lockFile.Close()

	formatter := format.GetFormatter()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go handleSysSignals(cancel)

	timingInterval := c.Int64("interval")
	if !(timingInterval > 0) {
		timingInterval = DefaultTimingInterval
	}

	timingSplay := c.Int64("splay")
	if !(timingSplay > 0) {
		timingSplay = DefaultTimingSplay
	}

	thresholdLines := c.Int("lines")
	if !(thresholdLines > 0) {
		thresholdLines = DefaultThresholdLines
	}
	log.Debug("routine lines threshold: ", thresholdLines)

	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	bootstrappingSvc, formatter := cmd.WireUpBootstrapping(c)
	for {
		applyPolicyfiles(bootstrappingSvc, formatter, thresholdLines)

		// Sleep for a configured amount of time plus a random amount of time (10 minutes plus 0 to 5 minutes, for instance)
		ticker := time.NewTicker(time.Duration(timingInterval+int64(r.Intn(int(timingSplay)))) * time.Second)

		select {
		case <-ticker.C:
			log.Debug("ticker")
		case <-ctx.Done():
			log.Debug(ctx.Err())
			log.Debug("closing bootstrapping")
		}
		ticker.Stop()
		if ctx.Err() != nil {
			break
		}
	}

	return nil
}

// Stop the bootstrapping process
func stop(c *cli.Context) error {
	log.Debug("cmdStop")

	formatter := format.GetFormatter()
	if err := utils.StopProcess(getProcessLockFilePath()); err != nil {
		formatter.PrintFatal("cannot stop the bootstrapping process", err)
	}

	log.Info("Bootstrapping routine successfully stopped")
	return nil
}

// Subsidiary routine for commands processing
func applyPolicyfiles(bootstrappingSvc *blueprint.BootstrappingService, formatter format.Formatter, thresholdLines int) error {
	log.Debug("applyPolicyfiles")

	// Inquire about desired configuration changes to be applied by querying the `GET /blueprint/configuration` endpoint. This will provide a JSON response with the desired configuration changes
	bsConfiguration, status, err := bootstrappingSvc.GetBootstrappingConfiguration()
	if err == nil && status != 200 {
		err = fmt.Errorf("received non-ok %d response", status)
	}
	if err != nil {
		formatter.PrintError("couldn't receive bootstrapping data", err)
		return err
	}
	dir, err := generateWorkspaceDir()
	if err != nil {
		return err
	}
	bsProcess := &bootstrappingProcess{
		startedAt:                    time.Now().UTC(),
		thresholdLines:               thresholdLines,
		directoryPath:                dir,
		appliedPolicyfileRevisionIDs: make(map[string]string),
	}

	// proto structures
	err = initializePrototype(bsConfiguration, bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	// For every policyfile, ensure its tarball (downloadable through their download_url) has been downloaded to the server ...
	err = downloadPolicyfiles(bootstrappingSvc, bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	//... and clean off any tarball that is no longer needed.
	err = cleanObsoletePolicyfiles(bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	// Store the attributes as JSON in a file with name `attrs-<attribute_revision_id>.json`
	err = saveAttributes(bsProcess)
	if err != nil {
		log.Debug(err)
		return err
	}
	// Process tarballs policies
	err = processPolicyfiles(bootstrappingSvc, bsProcess)
	// Finishing time
	bsProcess.finishedAt = time.Now().UTC()

	// Inform the platform of applied changes via a `PUT /blueprint/applied_configuration` request with a JSON payload similar to
	log.Debug("reporting applied policy files")
	reportErr := reportAppliedConfiguration(bootstrappingSvc, bsProcess)
	if reportErr != nil {
		log.Debug(err)
		return err
	}
	if err != nil {
		return err
	}
	return completeBootstrappingSequence(bsProcess)
}

func initializePrototype(bsConfiguration *types.BootstrappingConfiguration, bsProcess *bootstrappingProcess) error {
	log.Debug("initializePrototype")

	// Attributes
	bsProcess.attributes.revisionID = bsConfiguration.AttributeRevisionID
	bsProcess.attributes.rawData = bsConfiguration.Attributes

	// Policies
	for _, bsConfPolicyfile := range bsConfiguration.Policyfiles {
		bsProcess.policyfiles = append(bsProcess.policyfiles, policyfile(bsConfPolicyfile))
	}
	log.Debug(bsProcess)
	return nil
}

// downloadPolicyfiles For every policy file, ensure its tarball (downloadable through their download_url) has been downloaded to the server ...
func downloadPolicyfiles(bootstrappingSvc *blueprint.BootstrappingService, bsProcess *bootstrappingProcess) error {
	log.Debug("downloadPolicyfiles")

	for _, bsPolicyfile := range bsProcess.policyfiles {
		tarballPath := bsPolicyfile.TarballPath(bsProcess.directoryPath)
		log.Debug("downloading: ", tarballPath)
		queryURL, err := bsPolicyfile.QueryURL()
		if err != nil {
			return err
		}
		_, status, err := bootstrappingSvc.DownloadPolicyfile(queryURL, tarballPath)
		if err == nil && status != 200 {
			err = fmt.Errorf("obtained non-ok response when downloading policyfile %s", queryURL)
		}
		if err != nil {
			return err
		}
		if err = utils.Untar(tarballPath, bsPolicyfile.Path(bsProcess.directoryPath)); err != nil {
			return err
		}
	}
	return nil
}

// cleanObsoletePolicyfiles cleans off any tarball that is no longer needed.
func cleanObsoletePolicyfiles(bsProcess *bootstrappingProcess) error {
	log.Debug("cleanObsoletePolicyfiles")

	// evaluates working folder
	deletableFiles, err := ioutil.ReadDir(bsProcess.directoryPath)
	if err != nil {
		return err
	}

	// builds an array of currently processable files at this looping time
	currentlyProcessableFiles := []string{bsProcess.attributes.FileName()} // saved attributes file name
	for _, bsPolicyFile := range bsProcess.policyfiles {
		currentlyProcessableFiles = append(currentlyProcessableFiles, bsPolicyFile.FileName()) // Downloaded tgz file names
		currentlyProcessableFiles = append(currentlyProcessableFiles, bsPolicyFile.Name())     // Uncompressed folder names
	}

	// removes from deletableFiles array the policy files currently applied
	for _, f := range deletableFiles {
		if !utils.Contains(currentlyProcessableFiles, f.Name()) {
			log.Debug("removing: ", f.Name())
			if err := os.RemoveAll(filepath.Join(bsProcess.directoryPath, f.Name())); err != nil {
				return err
			}
		}
	}
	return nil
}

// saveAttributes stores the attributes as JSON in a file with name `attrs-<attribute_revision_id>.json`
func saveAttributes(bsProcess *bootstrappingProcess) error {
	log.Debug("saveAttributes")

	attrs, err := json.Marshal(bsProcess.attributes.rawData)
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(bsProcess.attributes.FilePath(bsProcess.directoryPath), attrs, 0600); err != nil {
		return err
	}
	return nil
}

// processPolicyfiles applies for each policy the required chef commands, reporting in bunches of N lines
func processPolicyfiles(bootstrappingSvc *blueprint.BootstrappingService, bsProcess *bootstrappingProcess) error {
	log.Debug("processPolicyfiles")

	for _, bsPolicyfile := range bsProcess.policyfiles {
		command := fmt.Sprintf("cd %s", bsPolicyfile.Path(bsProcess.directoryPath))
		if runtime.GOOS == "windows" {
			command = fmt.Sprintf("%s\nSET \"PATH=%PATH%;C:\\ruby\\bin;C:\\opscode\\chef\\bin;C:\\opscode\\chef\\embedded\\bin\"", command)
		}
		command = fmt.Sprintf("%s\nchef-client -z -j %s", command, bsProcess.attributes.FilePath(bsProcess.directoryPath))
		log.Debug(command)

		// Custom method for chunks processing
		fn := func(chunk string) error {
			log.Debug("sendChunks")
			err := utils.Retry(RetriesNumber, time.Second, func() error {
				log.Debug("Sending: ", chunk)

				commandIn := map[string]interface{}{
					"stdout": chunk,
				}

				_, statusCode, err := bootstrappingSvc.ReportBootstrappingLog(&commandIn)
				switch {
				// 0<100 error cases??
				case statusCode == 0:
					return fmt.Errorf("communication error %v %v", statusCode, err)
				case statusCode >= 500:
					return fmt.Errorf("server error %v %v", statusCode, err)
				case statusCode >= 400:
					return fmt.Errorf("client error %v %v", statusCode, err)
				default:
					return nil
				}
			})

			if err != nil {
				return fmt.Errorf("cannot send the chunk data, %v", err)
			}
			return nil
		}

		exitCode, err := utils.RunContinuousCmd(fn, command, -1, bsProcess.thresholdLines)
		if err == nil && exitCode != 0 {
			err = fmt.Errorf("policyfile application exited with %d code", exitCode)
		}
		if err != nil {
			return err
		}

		log.Info("completed: ", exitCode)
		bsProcess.appliedPolicyfileRevisionIDs[bsPolicyfile.ID] = bsPolicyfile.RevisionID
	}
	return nil
}

// reportAppliedConfiguration Inform the platform of applied changes
func reportAppliedConfiguration(bootstrappingSvc *blueprint.BootstrappingService, bsProcess *bootstrappingProcess) error {
	log.Debug("reportAppliedConfiguration")

	payload := map[string]interface{}{
		"started_at":              bsProcess.startedAt,
		"finished_at":             bsProcess.finishedAt,
		"policyfile_revision_ids": bsProcess.appliedPolicyfileRevisionIDs,
		"attribute_revision_id":   bsProcess.attributes.revisionID,
	}
	return bootstrappingSvc.ReportBootstrappingAppliedConfiguration(&payload)
}

// completeBootstrappingSequence evaluates if the first iteration of policies was completed; If case, execute the "scripts boot" command.
func completeBootstrappingSequence(bsProcess *bootstrappingProcess) error {
	log.Debug("completeBootstrappingSequence")

	if !allPolicyfilesSuccessfullyApplied {
		log.Debug("run the boot scripts")
		//run the boot scripts for the server by executing the scripts boot sub-command (as an external process).
		if output, exit, _, _ := utils.RunCmd(strings.Join([]string{os.Args[0], "scripts", "boot"}, " ")); exit != 0 {
			return fmt.Errorf("boot scripts run failed with exit code %d and following output: %s", exit, output)
		}
		allPolicyfilesSuccessfullyApplied = true
	}
	return nil
}<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cloudsigma

import (
	"sync"

	"github.com/altoros/gosigma"
	"github.com/juju/errors"
	"github.com/juju/utils/featureflag"

	"github.com/juju/juju/constraints"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/environs/simplestreams"
	"github.com/juju/juju/feature"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/juju/arch"
	"github.com/juju/juju/juju/osenv"
	"github.com/juju/juju/provider/common"
)

const (
	CloudsigmaCloudImagesURLTemplate = "https://%v.cloudsigma.com/"
)

// This file contains the core of the Environ implementation.
type environ struct {
	common.SupportsUnitPlacementPolicy
	name string

	lock      sync.Mutex
	archMutex sync.Mutex

	ecfg                   *environConfig
	client                 *environClient
	supportedArchitectures []string
}

var _ environs.Environ = (*environ)(nil)
var _ simplestreams.HasRegion = (*environ)(nil)
var _ simplestreams.MetadataValidator = (*environ)(nil)

// Name returns the Environ's name.
func (env environ) Name() string {
	return env.name
}

// Provider returns the EnvironProvider that created this Environ.
func (environ) Provider() environs.EnvironProvider {
	return providerInstance
}

// SetConfig updates the Environ's configuration.
//
// Calls to SetConfig do not affect the configuration of values previously obtained
// from Storage.
func (env *environ) SetConfig(cfg *config.Config) error {
	env.lock.Lock()
	defer env.lock.Unlock()

	ecfg, err := validateConfig(cfg, env.ecfg)
	if err != nil {
		return errors.Trace(err)
	}

	if env.client == nil || env.ecfg == nil || env.ecfg.clientConfigChanged(ecfg) {
		client, err := newClient(ecfg)
		if err != nil {
			return errors.Trace(err)
		}

		env.client = client
	}

	env.ecfg = ecfg

	return nil
}

// Config returns the configuration data with which the Environ was created.
// Note that this is not necessarily current; the canonical location
// for the configuration data is stored in the state.
func (env *environ) Config() *config.Config {
	return env.ecfg.Config
}

// Bootstrap initializes the state for the environment, possibly
// starting one or more instances.  If the configuration's
// AdminSecret is non-empty, the administrator password on the
// newly bootstrapped state will be set to a hash of it (see
// utils.PasswordHash), When first connecting to the
// environment via the juju package, the password hash will be
// automatically replaced by the real password.
//
// The supplied constraints are used to choose the initial instance
// specification, and will be stored in the new environment's state.
//
// Bootstrap is responsible for selecting the appropriate tools,
// and setting the agent-version configuration attribute prior to
// bootstrapping the environment.
func (env *environ) Bootstrap(ctx environs.BootstrapContext, params environs.BootstrapParams) (string, string, environs.BootstrapFinalizer, error) {
	featureflag.SetFlagsFromEnvironment(osenv.JujuFeatureFlagEnvKey)
	if featureflag.Enabled(feature.CloudSigma) {

		return common.Bootstrap(ctx, env, params)
	}
	return "", "", nil, errors.Errorf("Cloudsigma is currently in beta in this version of Juju. To enable anyway use JUJU_DEV_FEATURE_FLAGS=\"cloudsigma\"")
}

func (e *environ) StateServerInstances() ([]instance.Id, error) {
	return e.client.getStateServerIds()
}

// Destroy shuts down all known machines and destroys the
// rest of the environment. Note that on some providers,
// very recently started instances may not be destroyed
// because they are not yet visible.
//
// When Destroy has been called, any Environ referring to the
// same remote environment may become invalid
func (env *environ) Destroy() error {
	// You can probably ignore this method; the common implementation should work.
	return common.Destroy(env)
}

// PrecheckInstance performs a preflight check on the specified
// series and constraints, ensuring that they are possibly valid for
// creating an instance in this environment.
//
// PrecheckInstance is best effort, and not guaranteed to eliminate
// all invalid parameters. If PrecheckInstance returns nil, it is not
// guaranteed that the constraints are valid; if a non-nil error is
// returned, then the constraints are definitely invalid.
func (env *environ) PrecheckInstance(series string, cons constraints.Value, placement string) error {
	return nil
}

// Region is specified in the HasRegion interface.
func (env *environ) Region() (simplestreams.CloudSpec, error) {
	return env.cloudSpec(env.ecfg.region())
}

func (env *environ) cloudSpec(region string) (simplestreams.CloudSpec, error) {
	endpoint := gosigma.ResolveEndpoint(region)
	return simplestreams.CloudSpec{
		Region:   region,
		Endpoint: endpoint,
	}, nil
}

func (env *environ) MetadataLookupParams(region string) (*simplestreams.MetadataLookupParams, error) {
	if region == "" {
		region = gosigma.DefaultRegion
	}

	cloudSpec, err := env.cloudSpec(region)
	if err != nil {
		return nil, err
	}

	return &simplestreams.MetadataLookupParams{
		Region:        cloudSpec.Region,
		Endpoint:      cloudSpec.Endpoint,
		Architectures: arch.AllSupportedArches,
		Series:        config.PreferredSeries(env.ecfg),
	}, nil
}
<commit_msg>Drop feature flag handling for cloudsigma.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cloudsigma

import (
	"sync"

	"github.com/altoros/gosigma"
	"github.com/juju/errors"

	"github.com/juju/juju/constraints"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/environs/simplestreams"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/juju/arch"
	"github.com/juju/juju/provider/common"
)

const (
	CloudsigmaCloudImagesURLTemplate = "https://%v.cloudsigma.com/"
)

// This file contains the core of the Environ implementation.
type environ struct {
	common.SupportsUnitPlacementPolicy
	name string

	lock      sync.Mutex
	archMutex sync.Mutex

	ecfg                   *environConfig
	client                 *environClient
	supportedArchitectures []string
}

var _ environs.Environ = (*environ)(nil)
var _ simplestreams.HasRegion = (*environ)(nil)
var _ simplestreams.MetadataValidator = (*environ)(nil)

// Name returns the Environ's name.
func (env environ) Name() string {
	return env.name
}

// Provider returns the EnvironProvider that created this Environ.
func (environ) Provider() environs.EnvironProvider {
	return providerInstance
}

// SetConfig updates the Environ's configuration.
//
// Calls to SetConfig do not affect the configuration of values previously obtained
// from Storage.
func (env *environ) SetConfig(cfg *config.Config) error {
	env.lock.Lock()
	defer env.lock.Unlock()

	ecfg, err := validateConfig(cfg, env.ecfg)
	if err != nil {
		return errors.Trace(err)
	}

	if env.client == nil || env.ecfg == nil || env.ecfg.clientConfigChanged(ecfg) {
		client, err := newClient(ecfg)
		if err != nil {
			return errors.Trace(err)
		}

		env.client = client
	}

	env.ecfg = ecfg

	return nil
}

// Config returns the configuration data with which the Environ was created.
// Note that this is not necessarily current; the canonical location
// for the configuration data is stored in the state.
func (env *environ) Config() *config.Config {
	return env.ecfg.Config
}

// Bootstrap initializes the state for the environment, possibly
// starting one or more instances.  If the configuration's
// AdminSecret is non-empty, the administrator password on the
// newly bootstrapped state will be set to a hash of it (see
// utils.PasswordHash), When first connecting to the
// environment via the juju package, the password hash will be
// automatically replaced by the real password.
//
// The supplied constraints are used to choose the initial instance
// specification, and will be stored in the new environment's state.
//
// Bootstrap is responsible for selecting the appropriate tools,
// and setting the agent-version configuration attribute prior to
// bootstrapping the environment.
func (env *environ) Bootstrap(ctx environs.BootstrapContext, params environs.BootstrapParams) (string, string, environs.BootstrapFinalizer, error) {
	return common.Bootstrap(ctx, env, params)
}

func (e *environ) StateServerInstances() ([]instance.Id, error) {
	return e.client.getStateServerIds()
}

// Destroy shuts down all known machines and destroys the
// rest of the environment. Note that on some providers,
// very recently started instances may not be destroyed
// because they are not yet visible.
//
// When Destroy has been called, any Environ referring to the
// same remote environment may become invalid
func (env *environ) Destroy() error {
	// You can probably ignore this method; the common implementation should work.
	return common.Destroy(env)
}

// PrecheckInstance performs a preflight check on the specified
// series and constraints, ensuring that they are possibly valid for
// creating an instance in this environment.
//
// PrecheckInstance is best effort, and not guaranteed to eliminate
// all invalid parameters. If PrecheckInstance returns nil, it is not
// guaranteed that the constraints are valid; if a non-nil error is
// returned, then the constraints are definitely invalid.
func (env *environ) PrecheckInstance(series string, cons constraints.Value, placement string) error {
	return nil
}

// Region is specified in the HasRegion interface.
func (env *environ) Region() (simplestreams.CloudSpec, error) {
	return env.cloudSpec(env.ecfg.region())
}

func (env *environ) cloudSpec(region string) (simplestreams.CloudSpec, error) {
	endpoint := gosigma.ResolveEndpoint(region)
	return simplestreams.CloudSpec{
		Region:   region,
		Endpoint: endpoint,
	}, nil
}

func (env *environ) MetadataLookupParams(region string) (*simplestreams.MetadataLookupParams, error) {
	if region == "" {
		region = gosigma.DefaultRegion
	}

	cloudSpec, err := env.cloudSpec(region)
	if err != nil {
		return nil, err
	}

	return &simplestreams.MetadataLookupParams{
		Region:        cloudSpec.Region,
		Endpoint:      cloudSpec.Endpoint,
		Architectures: arch.AllSupportedArches,
		Series:        config.PreferredSeries(env.ecfg),
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxd

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/utils/arch"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/cloudconfig/instancecfg"
	"github.com/juju/juju/cloudconfig/providerinit"
	"github.com/juju/juju/container/lxd/lxdclient"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/provider/common"
	"github.com/juju/juju/state/multiwatcher"
)

func isStateServer(icfg *instancecfg.InstanceConfig) bool {
	return multiwatcher.AnyJobNeedsState(icfg.Jobs...)
}

// MaintainInstance is specified in the InstanceBroker interface.
func (*environ) MaintainInstance(args environs.StartInstanceParams) error {
	return nil
}

// StartInstance implements environs.InstanceBroker.
func (env *environ) StartInstance(args environs.StartInstanceParams) (*environs.StartInstanceResult, error) {
	// Please note that in order to fulfil the demands made of Instances and
	// AllInstances, it is imperative that some environment feature be used to
	// keep track of which instances were actually started by juju.
	env = env.getSnapshot()

	// Start a new instance.

	if args.InstanceConfig.HasNetworks() {
		return nil, errors.New("starting instances with networks is not supported yet")
	}

	series := args.Tools.OneSeries()
	logger.Debugf("StartInstance: %q, %s", args.InstanceConfig.MachineId, series)

	if err := env.finishInstanceConfig(args); err != nil {
		return nil, errors.Trace(err)
	}

	// TODO(ericsnow) Handle constraints?

	raw, err := env.newRawInstance(args)
	if err != nil {
		return nil, errors.Trace(err)
	}
	logger.Infof("started instance %q", raw.Name)
	inst := newInstance(raw, env)

	// Build the result.
	hwc := env.getHardwareCharacteristics(args, inst)
	result := environs.StartInstanceResult{
		Instance: inst,
		Hardware: hwc,
	}
	return &result, nil
}

func (env *environ) finishInstanceConfig(args environs.StartInstanceParams) error {
	args.InstanceConfig.Tools = args.Tools[0]
	logger.Debugf("tools: %#v", args.InstanceConfig.Tools)

	args.InstanceConfig.MachineContainerType = env.ecfg.containerType()

	if err := instancecfg.FinishInstanceConfig(args.InstanceConfig, env.ecfg.Config); err != nil {
		return errors.Trace(err)
	}

	// TODO: evaluate the impact of setting the constraints on the
	// instanceConfig for all machines rather than just state server nodes.
	// This limitation is why the constraints are assigned directly here.
	args.InstanceConfig.Constraints = args.Constraints

	args.InstanceConfig.AgentEnvironment[agent.Namespace] = env.ecfg.namespace()

	return nil
}

// newRawInstance is where the new physical instance is actually
// provisioned, relative to the provided args and spec. Info for that
// low-level instance is returned.
func (env *environ) newRawInstance(args environs.StartInstanceParams) (*lxdclient.Instance, error) {
	machineID := common.MachineFullName(env, args.InstanceConfig.MachineId)

	metadata, err := getMetadata(args)
	if err != nil {
		return nil, errors.Trace(err)
	}
	//tags := []string{
	//	env.globalFirewallName(),
	//	machineID,
	//}
	// TODO(ericsnow) Use the env ID for the network name (instead of default)?
	// TODO(ericsnow) Make the network name configurable?
	// TODO(ericsnow) Support multiple networks?
	// TODO(ericsnow) Use a different net interface name? Configurable?
	instSpec := lxdclient.InstanceSpec{
		Name: machineID,
		//Type:              spec.InstanceType.Name,
		//Disks:             getDisks(spec, args.Constraints),
		//NetworkInterfaces: []string{"ExternalNAT"},
		Metadata: metadata,
		Profiles: []string{
			env.profileName(),
		},
		//Tags:              tags,
		// Network is omitted (left empty).
	}

	inst, err := env.raw.AddInstance(instSpec)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return inst, nil
}

// getMetadata builds the raw "user-defined" metadata for the new
// instance (relative to the provided args) and returns it.
func getMetadata(args environs.StartInstanceParams) (map[string]string, error) {
	renderer := lxdRenderer{}
	uncompressed, err := providerinit.ComposeUserData(args.InstanceConfig, nil, renderer)
	if err != nil {
		return nil, errors.Annotate(err, "cannot make user data")
	}
	logger.Debugf("LXD user data; %d bytes", len(uncompressed))

	// TODO(ericsnow) Looks like LXD does not handle gzipped userdata
	// correctly.  It likely has to do with the HTTP transport, much
	// as we have to b64encode the userdata for GCE.  Until that is
	// resolved we simply pass the plain text.
	//compressed := utils.Gzip(compressed)
	userdata := string(uncompressed)

	metadata := map[string]string{
		metadataKeyIsState: metadataValueFalse,
		// We store a gz snapshop of information that is used by
		// cloud-init and unpacked in to the /var/lib/cloud/instances folder
		// for the instance.
		metadataKeyCloudInit: userdata,
	}
	if isStateServer(args.InstanceConfig) {
		metadata[metadataKeyIsState] = metadataValueTrue
	}

	return metadata, nil
}

// getHardwareCharacteristics compiles hardware-related details about
// the given instance and relative to the provided spec and returns it.
func (env *environ) getHardwareCharacteristics(args environs.StartInstanceParams, inst *environInstance) *instance.HardwareCharacteristics {
	raw := inst.raw.Hardware

	archStr := raw.Architecture
	if archStr == "unknown" || !arch.IsSupportedArch(archStr) {
		// TODO(ericsnow) This special-case should be improved.
		archStr = arch.HostArch()
	}

	hwc, err := instance.ParseHardware(
		"arch="+archStr,
		fmt.Sprintf("cpu-cores=%d", raw.NumCores),
		fmt.Sprintf("mem=%dM", raw.MemoryMB),
		//"root-disk=",
		//"tags=",
	)
	if err != nil {
		logger.Errorf("unexpected problem parsing hardware info: %v", err)
		// Keep moving...
	}
	return &hwc
}

// AllInstances implements environs.InstanceBroker.
func (env *environ) AllInstances() ([]instance.Instance, error) {
	instances, err := getInstances(env)
	return instances, errors.Trace(err)
}

// StopInstances implements environs.InstanceBroker.
func (env *environ) StopInstances(instances ...instance.Id) error {
	env = env.getSnapshot()

	var ids []string
	for _, id := range instances {
		ids = append(ids, string(id))
	}

	prefix := common.MachineFullName(env, "")
	err := env.raw.RemoveInstances(prefix, ids...)
	return errors.Trace(err)
}
<commit_msg>added default profile to instanceSpec<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// +build go1.3

package lxd

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/utils/arch"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/cloudconfig/instancecfg"
	"github.com/juju/juju/cloudconfig/providerinit"
	"github.com/juju/juju/container/lxd/lxdclient"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/provider/common"
	"github.com/juju/juju/state/multiwatcher"
)

func isStateServer(icfg *instancecfg.InstanceConfig) bool {
	return multiwatcher.AnyJobNeedsState(icfg.Jobs...)
}

// MaintainInstance is specified in the InstanceBroker interface.
func (*environ) MaintainInstance(args environs.StartInstanceParams) error {
	return nil
}

// StartInstance implements environs.InstanceBroker.
func (env *environ) StartInstance(args environs.StartInstanceParams) (*environs.StartInstanceResult, error) {
	// Please note that in order to fulfil the demands made of Instances and
	// AllInstances, it is imperative that some environment feature be used to
	// keep track of which instances were actually started by juju.
	env = env.getSnapshot()

	// Start a new instance.

	if args.InstanceConfig.HasNetworks() {
		return nil, errors.New("starting instances with networks is not supported yet")
	}

	series := args.Tools.OneSeries()
	logger.Debugf("StartInstance: %q, %s", args.InstanceConfig.MachineId, series)

	if err := env.finishInstanceConfig(args); err != nil {
		return nil, errors.Trace(err)
	}

	// TODO(ericsnow) Handle constraints?

	raw, err := env.newRawInstance(args)
	if err != nil {
		return nil, errors.Trace(err)
	}
	logger.Infof("started instance %q", raw.Name)
	inst := newInstance(raw, env)

	// Build the result.
	hwc := env.getHardwareCharacteristics(args, inst)
	result := environs.StartInstanceResult{
		Instance: inst,
		Hardware: hwc,
	}
	return &result, nil
}

func (env *environ) finishInstanceConfig(args environs.StartInstanceParams) error {
	args.InstanceConfig.Tools = args.Tools[0]
	logger.Debugf("tools: %#v", args.InstanceConfig.Tools)

	args.InstanceConfig.MachineContainerType = env.ecfg.containerType()

	if err := instancecfg.FinishInstanceConfig(args.InstanceConfig, env.ecfg.Config); err != nil {
		return errors.Trace(err)
	}

	// TODO: evaluate the impact of setting the constraints on the
	// instanceConfig for all machines rather than just state server nodes.
	// This limitation is why the constraints are assigned directly here.
	args.InstanceConfig.Constraints = args.Constraints

	args.InstanceConfig.AgentEnvironment[agent.Namespace] = env.ecfg.namespace()

	return nil
}

// newRawInstance is where the new physical instance is actually
// provisioned, relative to the provided args and spec. Info for that
// low-level instance is returned.
func (env *environ) newRawInstance(args environs.StartInstanceParams) (*lxdclient.Instance, error) {
	machineID := common.MachineFullName(env, args.InstanceConfig.MachineId)

	metadata, err := getMetadata(args)
	if err != nil {
		return nil, errors.Trace(err)
	}
	//tags := []string{
	//	env.globalFirewallName(),
	//	machineID,
	//}
	// TODO(ericsnow) Use the env ID for the network name (instead of default)?
	// TODO(ericsnow) Make the network name configurable?
	// TODO(ericsnow) Support multiple networks?
	// TODO(ericsnow) Use a different net interface name? Configurable?
	instSpec := lxdclient.InstanceSpec{
		Name: machineID,
		//Type:              spec.InstanceType.Name,
		//Disks:             getDisks(spec, args.Constraints),
		//NetworkInterfaces: []string{"ExternalNAT"},
		Metadata: metadata,
		Profiles: []string{
			//TODO(wwitzel3) move this to environments.yaml allowing the user to specify
			// lxc profiles to apply. This allows the user to setup any custom devices order
			// config settings for their environment. Also we must ensure that a device with
			// the parent: lxcbr0 exists in at least one of the profiles.
			"default",
			env.profileName(),
		},
		//Tags:              tags,
		// Network is omitted (left empty).
	}

	inst, err := env.raw.AddInstance(instSpec)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return inst, nil
}

// getMetadata builds the raw "user-defined" metadata for the new
// instance (relative to the provided args) and returns it.
func getMetadata(args environs.StartInstanceParams) (map[string]string, error) {
	renderer := lxdRenderer{}
	uncompressed, err := providerinit.ComposeUserData(args.InstanceConfig, nil, renderer)
	if err != nil {
		return nil, errors.Annotate(err, "cannot make user data")
	}
	logger.Debugf("LXD user data; %d bytes", len(uncompressed))

	// TODO(ericsnow) Looks like LXD does not handle gzipped userdata
	// correctly.  It likely has to do with the HTTP transport, much
	// as we have to b64encode the userdata for GCE.  Until that is
	// resolved we simply pass the plain text.
	//compressed := utils.Gzip(compressed)
	userdata := string(uncompressed)

	metadata := map[string]string{
		metadataKeyIsState: metadataValueFalse,
		// We store a gz snapshop of information that is used by
		// cloud-init and unpacked in to the /var/lib/cloud/instances folder
		// for the instance.
		metadataKeyCloudInit: userdata,
	}
	if isStateServer(args.InstanceConfig) {
		metadata[metadataKeyIsState] = metadataValueTrue
	}

	return metadata, nil
}

// getHardwareCharacteristics compiles hardware-related details about
// the given instance and relative to the provided spec and returns it.
func (env *environ) getHardwareCharacteristics(args environs.StartInstanceParams, inst *environInstance) *instance.HardwareCharacteristics {
	raw := inst.raw.Hardware

	archStr := raw.Architecture
	if archStr == "unknown" || !arch.IsSupportedArch(archStr) {
		// TODO(ericsnow) This special-case should be improved.
		archStr = arch.HostArch()
	}

	hwc, err := instance.ParseHardware(
		"arch="+archStr,
		fmt.Sprintf("cpu-cores=%d", raw.NumCores),
		fmt.Sprintf("mem=%dM", raw.MemoryMB),
		//"root-disk=",
		//"tags=",
	)
	if err != nil {
		logger.Errorf("unexpected problem parsing hardware info: %v", err)
		// Keep moving...
	}
	return &hwc
}

// AllInstances implements environs.InstanceBroker.
func (env *environ) AllInstances() ([]instance.Instance, error) {
	instances, err := getInstances(env)
	return instances, errors.Trace(err)
}

// StopInstances implements environs.InstanceBroker.
func (env *environ) StopInstances(instances ...instance.Id) error {
	env = env.getSnapshot()

	var ids []string
	for _, id := range instances {
		ids = append(ids, string(id))
	}

	prefix := common.MachineFullName(env, "")
	err := env.raw.RemoveInstances(prefix, ids...)
	return errors.Trace(err)
}
<|endoftext|>


<commit_before>package inbound

import (
	"io"
	"sync"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/bufio"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/loader"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/uuid"
	"v2ray.com/core/proxy"
	"v2ray.com/core/proxy/registry"
	"v2ray.com/core/proxy/vmess"
	"v2ray.com/core/proxy/vmess/encoding"
	"v2ray.com/core/transport/internet"
)

type userByEmail struct {
	sync.RWMutex
	cache           map[string]*protocol.User
	defaultLevel    uint32
	defaultAlterIDs uint16
}

func NewUserByEmail(users []*protocol.User, config *DefaultConfig) *userByEmail {
	cache := make(map[string]*protocol.User)
	for _, user := range users {
		cache[user.Email] = user
	}
	return &userByEmail{
		cache:           cache,
		defaultLevel:    config.Level,
		defaultAlterIDs: uint16(config.AlterId),
	}
}

func (v *userByEmail) Get(email string) (*protocol.User, bool) {
	var user *protocol.User
	var found bool
	v.RLock()
	user, found = v.cache[email]
	v.RUnlock()
	if !found {
		v.Lock()
		user, found = v.cache[email]
		if !found {
			account := &vmess.Account{
				Id:      uuid.New().String(),
				AlterId: uint32(v.defaultAlterIDs),
			}
			user = &protocol.User{
				Level:   v.defaultLevel,
				Email:   email,
				Account: loader.NewTypedSettings(account),
			}
			v.cache[email] = user
		}
		v.Unlock()
	}
	return user, found
}

// Inbound connection handler that handles messages in VMess format.
type VMessInboundHandler struct {
	sync.RWMutex
	packetDispatcher      dispatcher.PacketDispatcher
	inboundHandlerManager proxyman.InboundHandlerManager
	clients               protocol.UserValidator
	usersByEmail          *userByEmail
	accepting             bool
	listener              *internet.TCPHub
	detours               *DetourConfig
	meta                  *proxy.InboundHandlerMeta
}

func (v *VMessInboundHandler) Port() v2net.Port {
	return v.meta.Port
}

func (v *VMessInboundHandler) Close() {
	v.accepting = false
	if v.listener != nil {
		v.Lock()
		v.listener.Close()
		v.listener = nil
		v.clients.Release()
		v.clients = nil
		v.Unlock()
	}
}

func (v *VMessInboundHandler) GetUser(email string) *protocol.User {
	v.RLock()
	defer v.RUnlock()

	if !v.accepting {
		return nil
	}

	user, existing := v.usersByEmail.Get(email)
	if !existing {
		v.clients.Add(user)
	}
	return user
}

func (v *VMessInboundHandler) Start() error {
	if v.accepting {
		return nil
	}

	tcpListener, err := internet.ListenTCP(v.meta.Address, v.meta.Port, v.HandleConnection, v.meta.StreamSettings)
	if err != nil {
		log.Error("VMess|Inbound: Unable to listen tcp ", v.meta.Address, ":", v.meta.Port, ": ", err)
		return err
	}
	v.accepting = true
	v.Lock()
	v.listener = tcpListener
	v.Unlock()
	return nil
}

func (v *VMessInboundHandler) HandleConnection(connection internet.Connection) {
	defer connection.Close()

	if !v.accepting {
		return
	}

	connReader := v2net.NewTimeOutReader(8, connection)
	defer connReader.Release()

	reader := bufio.NewReader(connReader)
	defer reader.Release()

	v.RLock()
	if !v.accepting {
		v.RUnlock()
		return
	}
	session := encoding.NewServerSession(v.clients)
	defer session.Release()

	request, err := session.DecodeRequestHeader(reader)
	v.RUnlock()

	connection.SetReusable(request.Option.Has(protocol.RequestOptionConnectionReuse))

	if err != nil {
		if errors.Cause(err) != io.EOF {
			log.Access(connection.RemoteAddr(), "", log.AccessRejected, err)
			log.Info("VMessIn: Invalid request from ", connection.RemoteAddr(), ": ", err)
		}
		connection.SetReusable(false)
		return
	}
	log.Access(connection.RemoteAddr(), request.Destination(), log.AccessAccepted, "")
	log.Info("VMessIn: Received request for ", request.Destination())

	connection.SetReusable(request.Option.Has(protocol.RequestOptionConnectionReuse))

	ray := v.packetDispatcher.DispatchToOutbound(&proxy.SessionInfo{
		Source:      v2net.DestinationFromAddr(connection.RemoteAddr()),
		Destination: request.Destination(),
		User:        request.User,
		Inbound:     v.meta,
	})
	input := ray.InboundInput()
	output := ray.InboundOutput()
	defer input.Close()
	defer output.Release()

	var readFinish sync.Mutex
	readFinish.Lock()

	userSettings := request.User.GetSettings()
	connReader.SetTimeOut(userSettings.PayloadReadTimeout)
	reader.SetCached(false)

	go func() {
		bodyReader := session.DecodeRequestBody(request, reader)
		if err := buf.PipeUntilEOF(bodyReader, input); err != nil {
			connection.SetReusable(false)
		}
		bodyReader.Release()

		input.Close()
		readFinish.Unlock()
	}()

	writer := bufio.NewWriter(connection)
	defer writer.Release()

	response := &protocol.ResponseHeader{
		Command: v.generateCommand(request),
	}

	if connection.Reusable() {
		response.Option.Set(protocol.ResponseOptionConnectionReuse)
	}

	session.EncodeResponseHeader(response, writer)

	bodyWriter := session.EncodeResponseBody(request, writer)

	// Optimize for small response packet
	if data, err := output.Read(); err == nil {
		if err := bodyWriter.Write(data); err != nil {
			connection.SetReusable(false)
		}

		writer.SetCached(false)

		if err := buf.PipeUntilEOF(output, bodyWriter); err != nil {
			connection.SetReusable(false)
		}

	}
	output.Release()
	if request.Option.Has(protocol.RequestOptionChunkStream) {
		if err := bodyWriter.Write(buf.NewLocal(8)); err != nil {
			connection.SetReusable(false)
		}
	}
	writer.Flush()
	bodyWriter.Release()

	readFinish.Lock()
}

type Factory struct{}

func (v *Factory) StreamCapability() v2net.NetworkList {
	return v2net.NetworkList{
		Network: []v2net.Network{v2net.Network_TCP, v2net.Network_KCP, v2net.Network_WebSocket},
	}
}

func (v *Factory) Create(space app.Space, rawConfig interface{}, meta *proxy.InboundHandlerMeta) (proxy.InboundHandler, error) {
	if !space.HasApp(dispatcher.APP_ID) {
		return nil, common.ErrBadConfiguration
	}
	config := rawConfig.(*Config)

	allowedClients := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	for _, user := range config.User {
		allowedClients.Add(user)
	}

	handler := &VMessInboundHandler{
		packetDispatcher: space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher),
		clients:          allowedClients,
		detours:          config.Detour,
		usersByEmail:     NewUserByEmail(config.User, config.GetDefaultValue()),
		meta:             meta,
	}

	if space.HasApp(proxyman.APP_ID_INBOUND_MANAGER) {
		handler.inboundHandlerManager = space.GetApp(proxyman.APP_ID_INBOUND_MANAGER).(proxyman.InboundHandlerManager)
	}

	return handler, nil
}

func init() {
	registry.MustRegisterInboundHandlerCreator(loader.GetType(new(Config)), new(Factory))
}
<commit_msg>remove unnecessary function call<commit_after>package inbound

import (
	"io"
	"sync"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/proxyman"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/bufio"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/loader"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/uuid"
	"v2ray.com/core/proxy"
	"v2ray.com/core/proxy/registry"
	"v2ray.com/core/proxy/vmess"
	"v2ray.com/core/proxy/vmess/encoding"
	"v2ray.com/core/transport/internet"
)

type userByEmail struct {
	sync.RWMutex
	cache           map[string]*protocol.User
	defaultLevel    uint32
	defaultAlterIDs uint16
}

func NewUserByEmail(users []*protocol.User, config *DefaultConfig) *userByEmail {
	cache := make(map[string]*protocol.User)
	for _, user := range users {
		cache[user.Email] = user
	}
	return &userByEmail{
		cache:           cache,
		defaultLevel:    config.Level,
		defaultAlterIDs: uint16(config.AlterId),
	}
}

func (v *userByEmail) Get(email string) (*protocol.User, bool) {
	var user *protocol.User
	var found bool
	v.RLock()
	user, found = v.cache[email]
	v.RUnlock()
	if !found {
		v.Lock()
		user, found = v.cache[email]
		if !found {
			account := &vmess.Account{
				Id:      uuid.New().String(),
				AlterId: uint32(v.defaultAlterIDs),
			}
			user = &protocol.User{
				Level:   v.defaultLevel,
				Email:   email,
				Account: loader.NewTypedSettings(account),
			}
			v.cache[email] = user
		}
		v.Unlock()
	}
	return user, found
}

// Inbound connection handler that handles messages in VMess format.
type VMessInboundHandler struct {
	sync.RWMutex
	packetDispatcher      dispatcher.PacketDispatcher
	inboundHandlerManager proxyman.InboundHandlerManager
	clients               protocol.UserValidator
	usersByEmail          *userByEmail
	accepting             bool
	listener              *internet.TCPHub
	detours               *DetourConfig
	meta                  *proxy.InboundHandlerMeta
}

func (v *VMessInboundHandler) Port() v2net.Port {
	return v.meta.Port
}

func (v *VMessInboundHandler) Close() {
	v.accepting = false
	if v.listener != nil {
		v.Lock()
		v.listener.Close()
		v.listener = nil
		v.clients.Release()
		v.clients = nil
		v.Unlock()
	}
}

func (v *VMessInboundHandler) GetUser(email string) *protocol.User {
	v.RLock()
	defer v.RUnlock()

	if !v.accepting {
		return nil
	}

	user, existing := v.usersByEmail.Get(email)
	if !existing {
		v.clients.Add(user)
	}
	return user
}

func (v *VMessInboundHandler) Start() error {
	if v.accepting {
		return nil
	}

	tcpListener, err := internet.ListenTCP(v.meta.Address, v.meta.Port, v.HandleConnection, v.meta.StreamSettings)
	if err != nil {
		log.Error("VMess|Inbound: Unable to listen tcp ", v.meta.Address, ":", v.meta.Port, ": ", err)
		return err
	}
	v.accepting = true
	v.Lock()
	v.listener = tcpListener
	v.Unlock()
	return nil
}

func (v *VMessInboundHandler) HandleConnection(connection internet.Connection) {
	defer connection.Close()

	if !v.accepting {
		return
	}

	connReader := v2net.NewTimeOutReader(8, connection)
	defer connReader.Release()

	reader := bufio.NewReader(connReader)
	defer reader.Release()

	v.RLock()
	if !v.accepting {
		v.RUnlock()
		return
	}
	session := encoding.NewServerSession(v.clients)
	defer session.Release()

	request, err := session.DecodeRequestHeader(reader)
	v.RUnlock()

	if err != nil {
		if errors.Cause(err) != io.EOF {
			log.Access(connection.RemoteAddr(), "", log.AccessRejected, err)
			log.Info("VMessIn: Invalid request from ", connection.RemoteAddr(), ": ", err)
		}
		connection.SetReusable(false)
		return
	}
	log.Access(connection.RemoteAddr(), request.Destination(), log.AccessAccepted, "")
	log.Info("VMessIn: Received request for ", request.Destination())

	connection.SetReusable(request.Option.Has(protocol.RequestOptionConnectionReuse))

	ray := v.packetDispatcher.DispatchToOutbound(&proxy.SessionInfo{
		Source:      v2net.DestinationFromAddr(connection.RemoteAddr()),
		Destination: request.Destination(),
		User:        request.User,
		Inbound:     v.meta,
	})
	input := ray.InboundInput()
	output := ray.InboundOutput()
	defer input.Close()
	defer output.Release()

	var readFinish sync.Mutex
	readFinish.Lock()

	userSettings := request.User.GetSettings()
	connReader.SetTimeOut(userSettings.PayloadReadTimeout)
	reader.SetCached(false)

	go func() {
		bodyReader := session.DecodeRequestBody(request, reader)
		if err := buf.PipeUntilEOF(bodyReader, input); err != nil {
			connection.SetReusable(false)
		}
		bodyReader.Release()

		input.Close()
		readFinish.Unlock()
	}()

	writer := bufio.NewWriter(connection)
	defer writer.Release()

	response := &protocol.ResponseHeader{
		Command: v.generateCommand(request),
	}

	if connection.Reusable() {
		response.Option.Set(protocol.ResponseOptionConnectionReuse)
	}

	session.EncodeResponseHeader(response, writer)

	bodyWriter := session.EncodeResponseBody(request, writer)

	// Optimize for small response packet
	if data, err := output.Read(); err == nil {
		if err := bodyWriter.Write(data); err != nil {
			connection.SetReusable(false)
		}

		writer.SetCached(false)

		if err := buf.PipeUntilEOF(output, bodyWriter); err != nil {
			connection.SetReusable(false)
		}

	}
	output.Release()
	if request.Option.Has(protocol.RequestOptionChunkStream) {
		if err := bodyWriter.Write(buf.NewLocal(8)); err != nil {
			connection.SetReusable(false)
		}
	}
	writer.Flush()
	bodyWriter.Release()

	readFinish.Lock()
}

type Factory struct{}

func (v *Factory) StreamCapability() v2net.NetworkList {
	return v2net.NetworkList{
		Network: []v2net.Network{v2net.Network_TCP, v2net.Network_KCP, v2net.Network_WebSocket},
	}
}

func (v *Factory) Create(space app.Space, rawConfig interface{}, meta *proxy.InboundHandlerMeta) (proxy.InboundHandler, error) {
	if !space.HasApp(dispatcher.APP_ID) {
		return nil, common.ErrBadConfiguration
	}
	config := rawConfig.(*Config)

	allowedClients := vmess.NewTimedUserValidator(protocol.DefaultIDHash)
	for _, user := range config.User {
		allowedClients.Add(user)
	}

	handler := &VMessInboundHandler{
		packetDispatcher: space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher),
		clients:          allowedClients,
		detours:          config.Detour,
		usersByEmail:     NewUserByEmail(config.User, config.GetDefaultValue()),
		meta:             meta,
	}

	if space.HasApp(proxyman.APP_ID_INBOUND_MANAGER) {
		handler.inboundHandlerManager = space.GetApp(proxyman.APP_ID_INBOUND_MANAGER).(proxyman.InboundHandlerManager)
	}

	return handler, nil
}

func init() {
	registry.MustRegisterInboundHandlerCreator(loader.GetType(new(Config)), new(Factory))
}
<|endoftext|>


<commit_before>// BucketReader gets a list of files awaiting processing in the S3
// intake buckets and adds each item to the queue for processing.
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/APTrust/bagman"
	"github.com/APTrust/bagman/fluctus/client"
	"github.com/diamondap/goamz/aws"
	"github.com/op/go-logging"
	"net/http"
	"os"
	"strings"
	"time"
)

// Send S3 files to queue in batches of 500.
// Wait X milliseconds between batches. The wait time is really
// only necessary when the queue is running on an AWS small EC2
// instance, where the number of open network connections is
// severely restricted.
const (
	batchSize        = 500
	waitMilliseconds = 5000
)

var (
	config        bagman.Config
	messageLog    *logging.Logger
	fluctusClient *client.Client
	statusCache   map[string]*bagman.ProcessStatus
)

func main() {
	err := initialize()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Initialization failed for bucket_reader: %v", err)
		os.Exit(1)
	}
	run()
}

func initialize() (err error) {
	// Load the config or die.
	requestedConfig := flag.String("config", "", "configuration to run")
	flag.Parse()
	config = bagman.LoadRequestedConfig(requestedConfig)
	messageLog = bagman.InitLogger(config)
	messageLog.Info("Bucket reader started")
	fluctusClient, err = client.New(
		config.FluctusURL,
		config.FluctusAPIVersion,
		os.Getenv("FLUCTUS_API_USER"),
		os.Getenv("FLUCTUS_API_KEY"),
		messageLog)
	return err
}

func run() {
	s3Client, err := bagman.NewS3Client(aws.USEast)
	if err != nil {
		messageLog.Error(err.Error())
		return
	}
	bucketSummaries, err := s3Client.CheckAllBuckets(config.Buckets)
	if err != nil {
		messageLog.Error(err.Error())
		return
	}
	loadStatusCache()
	url := fmt.Sprintf("%s/mput?topic=%s", config.NsqdHttpAddress,
		config.BagProcessorTopic)
	messageLog.Debug("Sending S3 file info to %s \n", url)
	s3Files := filterLargeFiles(bucketSummaries)
	messageLog.Debug("%d S3 Files are within our size limit\n",
		len(s3Files))
	filesToProcess := s3Files
	// SkipAlreadyProcessed will almost always be true.
	// The exception is when we want to reprocess items to test new code.
	if config.SkipAlreadyProcessed == true {
		messageLog.Info("Skipping already processed files, because config says so")
		filesToProcess = filterProcessedFiles(s3Files)
	} else {
		messageLog.Info("Reprocessing already processed files, because config says so")
	}
	start := 0
	end := min(len(filesToProcess), batchSize)
	messageLog.Info("%d Unprocessed files\n", len(filesToProcess))
	for start <= end {
		batch := filesToProcess[start:end]
		messageLog.Info("Queuing batch of %d items\n", len(batch))
		enqueue(url, batch)
		start = end + 1
		if start < len(filesToProcess) {
			end = min(len(filesToProcess), start+batchSize)
		}
		// Sleep so we don't max out connections on EC2 small.
		// The utility server is running a lot of other network I/O
		// in addition to our queue.
		time.Sleep(time.Millisecond * waitMilliseconds)
	}
}

// min returns the minimum of x or y. The Math package has this function
// but you have to cast to floats.
func min(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

// filterLargeFiles returns only those S3 files that are not larger
// than config.MaxFileSize. This is useful when running tests and
// demos on your local machine, so that you can limit your test
// runs to files under 100k (or whatever you set in config.json).
// You don't want to pull down lots of multi-gig files when you're
// just running local tests. In production, set maxFileSize to
// zero, or to some huge value to get all files.
func filterLargeFiles(bucketSummaries []*bagman.BucketSummary) (s3Files []*bagman.S3File) {
	for _, bucketSummary := range bucketSummaries {
		for _, key := range bucketSummary.Keys {
			if config.MaxFileSize == 0 || key.Size < config.MaxFileSize {
				s3Files = append(s3Files, &bagman.S3File{
					BucketName: bucketSummary.BucketName,
					Key: key})
			}
		}
	}
	return s3Files
}

// Remove S3 files that have been processed successfully.
// No need to reprocess those!
func filterProcessedFiles(s3Files []*bagman.S3File) (filesToProcess []*bagman.S3File) {
	for _, s3File := range s3Files {
		bagDate, err := time.Parse(bagman.S3DateFormat, s3File.Key.LastModified)
		if err != nil {
			messageLog.Error("Cannot parse S3File mod date '%s'. "+
				"File %s will be re-processed.",
				s3File.Key.LastModified, s3File.Key.Key)
			filesToProcess = append(filesToProcess, s3File)
			continue
		}
		etag := strings.Replace(s3File.Key.ETag, "\"", "", 2)
		status := findInStatusCache(etag, s3File.Key.Key, bagDate)
		if status == nil {
			status, err = fluctusClient.GetBagStatus(etag, s3File.Key.Key, bagDate)
		}
		if err != nil {
			messageLog.Error("Cannot get Fluctus bag status for %s. "+
				"Will re-process bag. Error was %v", s3File.Key.Key, err)
			filesToProcess = append(filesToProcess, s3File)
		} else if status == nil || (status.Status == "Pending" && status.Retry == true) {
			reason := "Bag has not yet been successfully processed."
			if status == nil {
				err = createFluctusRecord(s3File)
				if err != nil {
					// TODO: Notify someone?
					messageLog.Error("Could not create Fluctus ProcessedItem "+
						"for %s: %v", s3File.Key.Key, err)
				}
			}
			messageLog.Info("Will process bag %s: %s", s3File.Key.Key, reason)
			filesToProcess = append(filesToProcess, s3File)
		} else if status.Status != "Failed" {
			messageLog.Debug("Skipping %s: already processed successfully.", s3File.Key.Key)
		} else if status.Retry == false {
			messageLog.Debug("Skipping %s: retry flag is set to false.", s3File.Key.Key)
		}
	}
	return filesToProcess
}

// Loads status of all bags received in the past two hours from fluctus
// in a single call.
func loadStatusCache() {
	twoHoursAgo := time.Now().Add(time.Hour * -2)
	statusRecords, err := fluctusClient.BulkStatusGet(twoHoursAgo)
	if err != nil {
		messageLog.Warning("Could not get bulk status records")
	} else {
		messageLog.Info("Got %d status records from the fluctopus\n", len(statusRecords))
		statusCache = make(map[string]*bagman.ProcessStatus, len(statusRecords))
		for i := range statusRecords {
			record := statusRecords[i]
			key := fmt.Sprintf("%s%s%s", record.ETag, record.Name, record.BagDate)
			statusCache[key] = record
		}
	}
}

// Finds the status of the specified tar bag in the cache that
// we retrieved from Fluctus. The cache can save us hundreds or
// thousands of HTTP calls each time the bucket reader runs.
func findInStatusCache(etag, name string, bagDate time.Time) *bagman.ProcessStatus {
	key := fmt.Sprintf("%s%s%s", etag, name, bagDate)
	item, exists := statusCache[key]
	if exists {
		//messageLog.Debug("Found item in cache: %s\n", name)
		return item
	}
	//messageLog.Debug("Item not in cache. Will have to ask the fluctopus for %s\n", name)
	return nil
}

func createFluctusRecord(s3File *bagman.S3File) (err error) {
	status := &bagman.ProcessStatus{}
	status.Date = time.Now().UTC()
	status.Action = "Ingest"
	status.Name = s3File.Key.Key
	bagDate, _ := time.Parse(bagman.S3DateFormat, s3File.Key.LastModified)
	status.BagDate = bagDate
	status.Bucket = s3File.BucketName
	// Strip the quotes off the ETag
	status.ETag = strings.Replace(s3File.Key.ETag, "\"", "", 2)
	status.Stage = bagman.StageReceive
	status.Status = bagman.StatusSuccess
	status.Note = "Item is in receiving bucket. Processing has not started."
	status.Institution = bagman.OwnerOf(s3File.BucketName)
	status.Outcome = string(status.Status)
	status.Reviewed = false

	// Retry should be true until we know for sure that there
	// is something wrong with the bag.
	status.Retry = true

	err = fluctusClient.UpdateBagStatus(status)
	if err != nil {
		return err
	}
	messageLog.Info("Created Fluctus ProcessedItem for %s\n",
		s3File.Key.Key)
	return nil
}

// enqueue adds a batch of items to the nsqd work queue
func enqueue(url string, s3Files []*bagman.S3File) {
	jsonData := make([]string, len(s3Files))
	for i, s3File := range s3Files {
		json, err := json.Marshal(s3File)
		if err != nil {
			messageLog.Error("Error marshalling s3 file to JSON: %v", err)
		} else {
			jsonData[i] = string(json)
			messageLog.Info("Put %s into fetch queue", s3File.Key.Key)
		}
	}
	batch := strings.Join(jsonData, "\n")
	resp, err := http.Post(url, "application/json", bytes.NewBuffer([]byte(batch)))
	if err != nil {
		messageLog.Error("nsqd returned an error: %v", err)
	}
	if resp == nil {
		msg := "No response from nsqd. Is it running? bucket_reader is quitting."
		messageLog.Error(msg)
		fmt.Println(msg)
		os.Exit(1)
	} else if resp.StatusCode != 200 {
		messageLog.Error("nsqd returned status code %d on last mput", resp.StatusCode)
	}
}
<commit_msg>Set status to pending when processing begins<commit_after>// BucketReader gets a list of files awaiting processing in the S3
// intake buckets and adds each item to the queue for processing.
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/APTrust/bagman"
	"github.com/APTrust/bagman/fluctus/client"
	"github.com/diamondap/goamz/aws"
	"github.com/op/go-logging"
	"net/http"
	"os"
	"strings"
	"time"
)

// Send S3 files to queue in batches of 500.
// Wait X milliseconds between batches. The wait time is really
// only necessary when the queue is running on an AWS small EC2
// instance, where the number of open network connections is
// severely restricted.
const (
	batchSize        = 500
	waitMilliseconds = 5000
)

var (
	config        bagman.Config
	messageLog    *logging.Logger
	fluctusClient *client.Client
	statusCache   map[string]*bagman.ProcessStatus
)

func main() {
	err := initialize()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Initialization failed for bucket_reader: %v", err)
		os.Exit(1)
	}
	run()
}

func initialize() (err error) {
	// Load the config or die.
	requestedConfig := flag.String("config", "", "configuration to run")
	flag.Parse()
	config = bagman.LoadRequestedConfig(requestedConfig)
	messageLog = bagman.InitLogger(config)
	messageLog.Info("Bucket reader started")
	fluctusClient, err = client.New(
		config.FluctusURL,
		config.FluctusAPIVersion,
		os.Getenv("FLUCTUS_API_USER"),
		os.Getenv("FLUCTUS_API_KEY"),
		messageLog)
	return err
}

func run() {
	s3Client, err := bagman.NewS3Client(aws.USEast)
	if err != nil {
		messageLog.Error(err.Error())
		return
	}
	bucketSummaries, err := s3Client.CheckAllBuckets(config.Buckets)
	if err != nil {
		messageLog.Error(err.Error())
		return
	}
	loadStatusCache()
	url := fmt.Sprintf("%s/mput?topic=%s", config.NsqdHttpAddress,
		config.BagProcessorTopic)
	messageLog.Debug("Sending S3 file info to %s \n", url)
	s3Files := filterLargeFiles(bucketSummaries)
	messageLog.Debug("%d S3 Files are within our size limit\n",
		len(s3Files))
	filesToProcess := s3Files
	// SkipAlreadyProcessed will almost always be true.
	// The exception is when we want to reprocess items to test new code.
	if config.SkipAlreadyProcessed == true {
		messageLog.Info("Skipping already processed files, because config says so")
		filesToProcess = filterProcessedFiles(s3Files)
	} else {
		messageLog.Info("Reprocessing already processed files, because config says so")
	}
	start := 0
	end := min(len(filesToProcess), batchSize)
	messageLog.Info("%d Unprocessed files\n", len(filesToProcess))
	for start <= end {
		batch := filesToProcess[start:end]
		messageLog.Info("Queuing batch of %d items\n", len(batch))
		enqueue(url, batch)
		start = end + 1
		if start < len(filesToProcess) {
			end = min(len(filesToProcess), start+batchSize)
		}
		// Sleep so we don't max out connections on EC2 small.
		// The utility server is running a lot of other network I/O
		// in addition to our queue.
		time.Sleep(time.Millisecond * waitMilliseconds)
	}
}

// min returns the minimum of x or y. The Math package has this function
// but you have to cast to floats.
func min(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

// filterLargeFiles returns only those S3 files that are not larger
// than config.MaxFileSize. This is useful when running tests and
// demos on your local machine, so that you can limit your test
// runs to files under 100k (or whatever you set in config.json).
// You don't want to pull down lots of multi-gig files when you're
// just running local tests. In production, set maxFileSize to
// zero, or to some huge value to get all files.
func filterLargeFiles(bucketSummaries []*bagman.BucketSummary) (s3Files []*bagman.S3File) {
	for _, bucketSummary := range bucketSummaries {
		for _, key := range bucketSummary.Keys {
			if config.MaxFileSize == 0 || key.Size < config.MaxFileSize {
				s3Files = append(s3Files, &bagman.S3File{
					BucketName: bucketSummary.BucketName,
					Key: key})
			}
		}
	}
	return s3Files
}

// Remove S3 files that have been processed successfully.
// No need to reprocess those!
func filterProcessedFiles(s3Files []*bagman.S3File) (filesToProcess []*bagman.S3File) {
	for _, s3File := range s3Files {
		bagDate, err := time.Parse(bagman.S3DateFormat, s3File.Key.LastModified)
		if err != nil {
			messageLog.Error("Cannot parse S3File mod date '%s'. "+
				"File %s will be re-processed.",
				s3File.Key.LastModified, s3File.Key.Key)
			filesToProcess = append(filesToProcess, s3File)
			continue
		}
		etag := strings.Replace(s3File.Key.ETag, "\"", "", 2)
		status := findInStatusCache(etag, s3File.Key.Key, bagDate)
		if status == nil {
			status, err = fluctusClient.GetBagStatus(etag, s3File.Key.Key, bagDate)
		}
		if err != nil {
			messageLog.Error("Cannot get Fluctus bag status for %s. "+
				"Will re-process bag. Error was %v", s3File.Key.Key, err)
			filesToProcess = append(filesToProcess, s3File)
		} else if status == nil || (status.Status == "Pending" && status.Retry == true) {
			reason := "Bag has not yet been successfully processed."
			if status == nil {
				err = createFluctusRecord(s3File)
				if err != nil {
					// TODO: Notify someone?
					messageLog.Error("Could not create Fluctus ProcessedItem "+
						"for %s: %v", s3File.Key.Key, err)
				}
			}
			messageLog.Info("Will process bag %s: %s", s3File.Key.Key, reason)
			filesToProcess = append(filesToProcess, s3File)
		} else if status.Status != "Failed" {
			messageLog.Debug("Skipping %s: already processed successfully.", s3File.Key.Key)
		} else if status.Retry == false {
			messageLog.Debug("Skipping %s: retry flag is set to false.", s3File.Key.Key)
		}
	}
	return filesToProcess
}

// Loads status of all bags received in the past two hours from fluctus
// in a single call.
func loadStatusCache() {
	twoHoursAgo := time.Now().Add(time.Hour * -2)
	statusRecords, err := fluctusClient.BulkStatusGet(twoHoursAgo)
	if err != nil {
		messageLog.Warning("Could not get bulk status records")
	} else {
		messageLog.Info("Got %d status records from the fluctopus\n", len(statusRecords))
		statusCache = make(map[string]*bagman.ProcessStatus, len(statusRecords))
		for i := range statusRecords {
			record := statusRecords[i]
			key := fmt.Sprintf("%s%s%s", record.ETag, record.Name, record.BagDate)
			statusCache[key] = record
		}
	}
}

// Finds the status of the specified tar bag in the cache that
// we retrieved from Fluctus. The cache can save us hundreds or
// thousands of HTTP calls each time the bucket reader runs.
func findInStatusCache(etag, name string, bagDate time.Time) *bagman.ProcessStatus {
	key := fmt.Sprintf("%s%s%s", etag, name, bagDate)
	item, exists := statusCache[key]
	if exists {
		//messageLog.Debug("Found item in cache: %s\n", name)
		return item
	}
	//messageLog.Debug("Item not in cache. Will have to ask the fluctopus for %s\n", name)
	return nil
}

func createFluctusRecord(s3File *bagman.S3File) (err error) {
	status := &bagman.ProcessStatus{}
	status.Date = time.Now().UTC()
	status.Action = "Ingest"
	status.Name = s3File.Key.Key
	bagDate, _ := time.Parse(bagman.S3DateFormat, s3File.Key.LastModified)
	status.BagDate = bagDate
	status.Bucket = s3File.BucketName
	// Strip the quotes off the ETag
	status.ETag = strings.Replace(s3File.Key.ETag, "\"", "", 2)
	status.Stage = bagman.StageReceive
	status.Status = bagman.StatusPending
	status.Note = "Item is in receiving bucket. Processing has not started."
	status.Institution = bagman.OwnerOf(s3File.BucketName)
	status.Outcome = string(status.Status)
	status.Reviewed = false

	// Retry should be true until we know for sure that there
	// is something wrong with the bag.
	status.Retry = true

	err = fluctusClient.UpdateBagStatus(status)
	if err != nil {
		return err
	}
	messageLog.Info("Created Fluctus ProcessedItem for %s\n",
		s3File.Key.Key)
	return nil
}

// enqueue adds a batch of items to the nsqd work queue
func enqueue(url string, s3Files []*bagman.S3File) {
	jsonData := make([]string, len(s3Files))
	for i, s3File := range s3Files {
		json, err := json.Marshal(s3File)
		if err != nil {
			messageLog.Error("Error marshalling s3 file to JSON: %v", err)
		} else {
			jsonData[i] = string(json)
			messageLog.Info("Put %s into fetch queue", s3File.Key.Key)
		}
	}
	batch := strings.Join(jsonData, "\n")
	resp, err := http.Post(url, "application/json", bytes.NewBuffer([]byte(batch)))
	if err != nil {
		messageLog.Error("nsqd returned an error: %v", err)
	}
	if resp == nil {
		msg := "No response from nsqd. Is it running? bucket_reader is quitting."
		messageLog.Error(msg)
		fmt.Println(msg)
		os.Exit(1)
	} else if resp.StatusCode != 200 {
		messageLog.Error("nsqd returned status code %d on last mput", resp.StatusCode)
	}
}
<|endoftext|>