<commit_before>package security

import (
	"bytes"
	"errors"
	"fmt"
	"reflect"
	"sync"

	"veyron.io/veyron/veyron/security/serialization"

	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"
)

var errStoreAddMismatch = errors.New("blessing's public key does not match store's public key")

type blessings struct {
	Value       security.WireBlessings
	unmarshaled security.Blessings
}

func (w *blessings) Blessings() security.Blessings {
	if w == nil {
		return nil
	}
	return w.unmarshaled
}

func (w *blessings) Verify() error {
	var err error
	if w.unmarshaled == nil {
		w.unmarshaled, err = security.NewBlessings(w.Value)
	}
	return err
}

func newWireBlessings(b security.Blessings) *blessings {
	return &blessings{Value: security.MarshalBlessings(b), unmarshaled: b}
}

type state struct {
	// Store maps BlessingPatterns to the Blessings object that is to be shared
	// with peers which present blessings of their own that match the pattern.
	//
	// All blessings bind to the same public key.
	Store map[security.BlessingPattern]*blessings
	// Default is the default Blessings to be shared with peers for which
	// no other information is available to select blessings.
	Default *blessings
}

// blessingStore implements security.BlessingStore.
type blessingStore struct {
	publicKey  security.PublicKey
	serializer SerializerReaderWriter
	signer     serialization.Signer
	mu         sync.RWMutex
	state      state // GUARDED_BY(mu)
}

func (bs *blessingStore) Set(blessings security.Blessings, forPeers security.BlessingPattern) (security.Blessings, error) {
	if !forPeers.IsValid() {
		return nil, fmt.Errorf("%q is an invalid BlessingPattern", forPeers)
	}
	if blessings != nil && !reflect.DeepEqual(blessings.PublicKey(), bs.publicKey) {
		return nil, errStoreAddMismatch
	}
	bs.mu.Lock()
	defer bs.mu.Unlock()
	old, hadold := bs.state.Store[forPeers]
	if blessings != nil {
		bs.state.Store[forPeers] = newWireBlessings(blessings)
	} else {
		delete(bs.state.Store, forPeers)
	}
	if err := bs.save(); err != nil {
		if hadold {
			bs.state.Store[forPeers] = old
		} else {
			delete(bs.state.Store, forPeers)
		}
		return nil, err
	}
	return old.Blessings(), nil
}

func (bs *blessingStore) ForPeer(peerBlessings ...string) security.Blessings {
	bs.mu.RLock()
	defer bs.mu.RUnlock()

	var ret security.Blessings
	for pattern, wb := range bs.state.Store {
		if pattern.MatchedBy(peerBlessings...) {
			b := wb.Blessings()
			if union, err := security.UnionOfBlessings(ret, b); err != nil {
				vlog.Errorf("UnionOfBlessings(%v, %v) failed: %v, dropping the latter from BlessingStore.ForPeers(%v)", ret, b, err, peerBlessings)
			} else {
				ret = union
			}
		}
	}
	return ret
}

func (bs *blessingStore) Default() security.Blessings {
	bs.mu.RLock()
	defer bs.mu.RUnlock()
	if bs.state.Default != nil {
		return bs.state.Default.Blessings()
	}
	return bs.ForPeer()
}

func (bs *blessingStore) SetDefault(blessings security.Blessings) error {
	bs.mu.Lock()
	defer bs.mu.Unlock()
	if !reflect.DeepEqual(blessings.PublicKey(), bs.publicKey) {
		return errStoreAddMismatch
	}
	oldDefault := bs.state.Default
	bs.state.Default = newWireBlessings(blessings)
	if err := bs.save(); err != nil {
		bs.state.Default = oldDefault
	}
	return nil
}

func (bs *blessingStore) PublicKey() security.PublicKey {
	return bs.publicKey
}

func (bs *blessingStore) String() string {
	return fmt.Sprintf("{state: %v, publicKey: %v}", bs.state, bs.publicKey)
}

// DebugString return a human-readable string encoding of the store
// in the following format
// Default blessing : <Default blessing of the store>
//
// Peer pattern : Blessings
// <pattern>    : <blessings>
// ...
// <pattern>    : <blessings>
func (bs *blessingStore) DebugString() string {
	const format = "%-30s : %s\n"
	b := bytes.NewBufferString(fmt.Sprintf("Default blessings: %v\n", bs.state.Default.Blessings()))

	b.WriteString(fmt.Sprintf(format, "Peer pattern", "Blessings"))
	for pattern, wb := range bs.state.Store {
		b.WriteString(fmt.Sprintf(format, pattern, wb.Blessings()))
	}
	return b.String()
}

func (bs *blessingStore) save() error {
	if (bs.signer == nil) && (bs.serializer == nil) {
		return nil
	}
	data, signature, err := bs.serializer.Writers()
	if err != nil {
		return err
	}
	return encodeAndStore(bs.state, data, signature, bs.signer)
}

// newInMemoryBlessingStore returns an in-memory security.BlessingStore for a
// principal with the provided PublicKey.
//
// The returned BlessingStore is initialized with an empty set of blessings.
func newInMemoryBlessingStore(publicKey security.PublicKey) security.BlessingStore {
	return &blessingStore{
		publicKey: publicKey,
		state:     state{Store: make(map[security.BlessingPattern]*blessings)},
	}
}

// TODO(ataly, ashankar): Get rid of this struct once we have switched all credentials
// directories to the new serialization format.
type oldState struct {
	Store   map[security.BlessingPattern]security.WireBlessings
	Default security.WireBlessings
}

// TODO(ataly, ashankar): Get rid of this method once we have switched all
// credentials directories to the new serialization format.
func (bs *blessingStore) tryOldFormat() bool {
	var empty security.WireBlessings
	if len(bs.state.Store) == 0 {
		return bs.state.Default == nil || reflect.DeepEqual(bs.state.Default.Value, empty)
	}
	for _, wb := range bs.state.Store {
		if len(wb.Value.CertificateChains) == 0 {
			return true
		}
	}
	return false
}

// TODO(ataly, ashankar): Get rid of this method once we have switched all
// credentials directories to the new serialization format.
func (bs *blessingStore) deserializeOld() error {
	data, signature, err := bs.serializer.Readers()
	if err != nil {
		return err
	}
	if data == nil && signature == nil {
		return nil
	}
	var old oldState
	if err := decodeFromStorage(&old, data, signature, bs.signer.PublicKey()); err != nil {
		return err
	}
	for p, wire := range old.Store {
		bs.state.Store[p] = &blessings{Value: wire}
	}
	bs.state.Default = &blessings{Value: old.Default}
	return nil
}

func (bs *blessingStore) deserialize() error {
	data, signature, err := bs.serializer.Readers()
	if err != nil {
		return err
	}
	if data == nil && signature == nil {
		return nil
	}
	if err := decodeFromStorage(&bs.state, data, signature, bs.signer.PublicKey()); err == nil && !bs.tryOldFormat() {
		return nil
	}
	if err := bs.deserializeOld(); err != nil {
		return err
	}
	return nil
}

// newPersistingBlessingStore returns a security.BlessingStore for a principal
// that is initialized with the persisted data. The returned security.BlessingStore
// also persists any updates to its state.
func newPersistingBlessingStore(serializer SerializerReaderWriter, signer serialization.Signer) (security.BlessingStore, error) {
	verifyBlessings := func(wb *blessings, key security.PublicKey) error {
		if err := wb.Verify(); err != nil {
			return err
		}
		if b := wb.Blessings(); b != nil && !reflect.DeepEqual(b.PublicKey(), key) {
			return fmt.Errorf("read Blessings: %v that are not for provided PublicKey: %v", b, key)
		}
		return nil
	}
	if serializer == nil || signer == nil {
		return nil, errors.New("persisted data or signer is not specified")
	}
	bs := &blessingStore{
		publicKey:  signer.PublicKey(),
		state:      state{Store: make(map[security.BlessingPattern]*blessings)},
		serializer: serializer,
		signer:     signer,
	}
	if err := bs.deserialize(); err != nil {
		return nil, err
	}
	for _, wb := range bs.state.Store {
		if err := verifyBlessings(wb, bs.publicKey); err != nil {
			return nil, err
		}
	}
	if bs.state.Default != nil {
		if err := verifyBlessings(bs.state.Default, bs.publicKey); err != nil {
			return nil, err
		}
	}
	// Save the blessingstore in the new serialization format. This will ensure
	// that all credentials directories in the old format will switch to the new
	// format.
	// TODO(ataly, ashankar): Get rid of this  once we have switched all
	// credentials directories to the new serialization format.
	if err := bs.save(); err != nil {
		return nil, err
	}
	return bs, nil
}
<commit_msg>"veyron/security": BlessingStore bug<commit_after>package security

import (
	"bytes"
	"errors"
	"fmt"
	"reflect"
	"sync"

	"veyron.io/veyron/veyron/security/serialization"

	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"
)

var errStoreAddMismatch = errors.New("blessing's public key does not match store's public key")

type blessings struct {
	Value       security.WireBlessings
	unmarshaled security.Blessings
}

func (w *blessings) Blessings() security.Blessings {
	if w == nil {
		return nil
	}
	return w.unmarshaled
}

func (w *blessings) Verify() error {
	var err error
	if w.unmarshaled == nil {
		w.unmarshaled, err = security.NewBlessings(w.Value)
	}
	return err
}

func newWireBlessings(b security.Blessings) *blessings {
	return &blessings{Value: security.MarshalBlessings(b), unmarshaled: b}
}

type state struct {
	// Store maps BlessingPatterns to the Blessings object that is to be shared
	// with peers which present blessings of their own that match the pattern.
	//
	// All blessings bind to the same public key.
	Store map[security.BlessingPattern]*blessings
	// Default is the default Blessings to be shared with peers for which
	// no other information is available to select blessings.
	Default *blessings
}

// blessingStore implements security.BlessingStore.
type blessingStore struct {
	publicKey  security.PublicKey
	serializer SerializerReaderWriter
	signer     serialization.Signer
	mu         sync.RWMutex
	state      state // GUARDED_BY(mu)
}

func (bs *blessingStore) Set(blessings security.Blessings, forPeers security.BlessingPattern) (security.Blessings, error) {
	if !forPeers.IsValid() {
		return nil, fmt.Errorf("%q is an invalid BlessingPattern", forPeers)
	}
	if blessings != nil && !reflect.DeepEqual(blessings.PublicKey(), bs.publicKey) {
		return nil, errStoreAddMismatch
	}
	bs.mu.Lock()
	defer bs.mu.Unlock()
	old, hadold := bs.state.Store[forPeers]
	if blessings != nil {
		bs.state.Store[forPeers] = newWireBlessings(blessings)
	} else {
		delete(bs.state.Store, forPeers)
	}
	if err := bs.save(); err != nil {
		if hadold {
			bs.state.Store[forPeers] = old
		} else {
			delete(bs.state.Store, forPeers)
		}
		return nil, err
	}
	return old.Blessings(), nil
}

func (bs *blessingStore) ForPeer(peerBlessings ...string) security.Blessings {
	bs.mu.RLock()
	defer bs.mu.RUnlock()

	var ret security.Blessings
	for pattern, wb := range bs.state.Store {
		if pattern.MatchedBy(peerBlessings...) {
			b := wb.Blessings()
			if union, err := security.UnionOfBlessings(ret, b); err != nil {
				vlog.Errorf("UnionOfBlessings(%v, %v) failed: %v, dropping the latter from BlessingStore.ForPeers(%v)", ret, b, err, peerBlessings)
			} else {
				ret = union
			}
		}
	}
	return ret
}

func (bs *blessingStore) Default() security.Blessings {
	bs.mu.RLock()
	defer bs.mu.RUnlock()
	if bs.state.Default != nil {
		return bs.state.Default.Blessings()
	}
	return bs.ForPeer()
}

func (bs *blessingStore) SetDefault(blessings security.Blessings) error {
	bs.mu.Lock()
	defer bs.mu.Unlock()
	if !reflect.DeepEqual(blessings.PublicKey(), bs.publicKey) {
		return errStoreAddMismatch
	}
	oldDefault := bs.state.Default
	bs.state.Default = newWireBlessings(blessings)
	if err := bs.save(); err != nil {
		bs.state.Default = oldDefault
	}
	return nil
}

func (bs *blessingStore) PublicKey() security.PublicKey {
	return bs.publicKey
}

func (bs *blessingStore) String() string {
	return fmt.Sprintf("{state: %v, publicKey: %v}", bs.state, bs.publicKey)
}

// DebugString return a human-readable string encoding of the store
// in the following format
// Default blessing : <Default blessing of the store>
//
// Peer pattern : Blessings
// <pattern>    : <blessings>
// ...
// <pattern>    : <blessings>
func (bs *blessingStore) DebugString() string {
	const format = "%-30s : %s\n"
	b := bytes.NewBufferString(fmt.Sprintf("Default blessings: %v\n", bs.state.Default.Blessings()))

	b.WriteString(fmt.Sprintf(format, "Peer pattern", "Blessings"))
	for pattern, wb := range bs.state.Store {
		b.WriteString(fmt.Sprintf(format, pattern, wb.Blessings()))
	}
	return b.String()
}

func (bs *blessingStore) save() error {
	if (bs.signer == nil) && (bs.serializer == nil) {
		return nil
	}
	data, signature, err := bs.serializer.Writers()
	if err != nil {
		return err
	}
	return encodeAndStore(bs.state, data, signature, bs.signer)
}

// newInMemoryBlessingStore returns an in-memory security.BlessingStore for a
// principal with the provided PublicKey.
//
// The returned BlessingStore is initialized with an empty set of blessings.
func newInMemoryBlessingStore(publicKey security.PublicKey) security.BlessingStore {
	return &blessingStore{
		publicKey: publicKey,
		state:     state{Store: make(map[security.BlessingPattern]*blessings)},
	}
}

// TODO(ataly, ashankar): Get rid of this struct once we have switched all credentials
// directories to the new serialization format.
type oldState struct {
	Store   map[security.BlessingPattern]security.WireBlessings
	Default security.WireBlessings
}

// TODO(ataly, ashankar): Get rid of this method once we have switched all
// credentials directories to the new serialization format.
func (bs *blessingStore) tryOldFormat() bool {
	var empty security.WireBlessings
	if len(bs.state.Store) == 0 {
		return bs.state.Default == nil || reflect.DeepEqual(bs.state.Default.Value, empty)
	}
	for _, wb := range bs.state.Store {
		if len(wb.Value.CertificateChains) == 0 {
			return true
		}
	}
	return false
}

func (bs *blessingStore) verifyState() error {
	verifyBlessings := func(wb *blessings, key security.PublicKey) error {
		if err := wb.Verify(); err != nil {
			return err
		}
		if b := wb.Blessings(); b != nil && !reflect.DeepEqual(b.PublicKey(), key) {
			return fmt.Errorf("read Blessings: %v that are not for provided PublicKey: %v", b, key)
		}
		return nil
	}
	for _, wb := range bs.state.Store {
		if err := verifyBlessings(wb, bs.publicKey); err != nil {
			return err
		}
	}
	if bs.state.Default != nil {
		if err := verifyBlessings(bs.state.Default, bs.publicKey); err != nil {
			return err
		}
	}
	return nil
}

// TODO(ataly, ashankar): Get rid of this method once we have switched all
// credentials directories to the new serialization format.
func (bs *blessingStore) deserializeOld() error {
	data, signature, err := bs.serializer.Readers()
	if err != nil {
		return err
	}
	if data == nil && signature == nil {
		return nil
	}
	var old oldState
	if err := decodeFromStorage(&old, data, signature, bs.signer.PublicKey()); err != nil {
		return err
	}
	for p, wire := range old.Store {
		bs.state.Store[p] = &blessings{Value: wire}
	}
	bs.state.Default = &blessings{Value: old.Default}

	if err := bs.verifyState(); err != nil {
		return err
	}
	// Save the blessingstore in the new serialization format. This will ensure
	// that all credentials directories in the old format will switch to the new
	// format.
	if err := bs.save(); err != nil {
		return err
	}
	return nil
}

func (bs *blessingStore) deserialize() error {
	data, signature, err := bs.serializer.Readers()
	if err != nil {
		return err
	}
	if data == nil && signature == nil {
		return nil
	}
	if err := decodeFromStorage(&bs.state, data, signature, bs.signer.PublicKey()); err == nil && !bs.tryOldFormat() {
		return bs.verifyState()
	}
	if err := bs.deserializeOld(); err != nil {
		return err
	}
	return nil
}

// newPersistingBlessingStore returns a security.BlessingStore for a principal
// that is initialized with the persisted data. The returned security.BlessingStore
// also persists any updates to its state.
func newPersistingBlessingStore(serializer SerializerReaderWriter, signer serialization.Signer) (security.BlessingStore, error) {
	if serializer == nil || signer == nil {
		return nil, errors.New("persisted data or signer is not specified")
	}
	bs := &blessingStore{
		publicKey:  signer.PublicKey(),
		state:      state{Store: make(map[security.BlessingPattern]*blessings)},
		serializer: serializer,
		signer:     signer,
	}
	if err := bs.deserialize(); err != nil {
		return nil, err
	}
	return bs, nil
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, sub to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package leaplib

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path"
	"strconv"
)

/*--------------------------------------------------------------------------------------------------
 */

/*
FileStore - Most basic persistent implementation of DocumentStore. Simple stores each document into
a file within a configured directory.
*/
type FileStore struct {
	config DocumentStoreConfig
}

/*
Store - Store document in its file location.
*/
func (s *FileStore) Store(id string, doc *Document) error {
	file, err := os.Create(path.Join(s.config.StoreDirectory, id))
	if err != nil {
		return err
	}
	defer file.Close()

	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(doc.Title)); err != nil {
		return err
	}
	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(doc.Description)); err != nil {
		return err
	}
	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(doc.Type)); err != nil {
		return err
	}
	serialized, err := SerializeDocumentContent(doc.Type, doc.Content)
	if err != nil {
		return err
	}
	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(serialized)); err != nil {
		return err
	}

	return nil
}

/*
Fetch - Fetch document from its file location.
*/
func (s *FileStore) Fetch(id string) (*Document, error) {
	file, err := os.Open(path.Join(s.config.StoreDirectory, id))
	if err != nil {
		return nil, err
	}
	defer file.Close()

	doc := Document{ID: id}

	scanner := bufio.NewScanner(file)
	if !scanner.Scan() {
		return nil, errors.New("failed to read title from document file")
	}
	doc.Title, err = strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}

	if !scanner.Scan() {
		return nil, errors.New("failed to read description from document file")
	}
	doc.Description, err = strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}

	if !scanner.Scan() {
		return nil, errors.New("failed to read type from document file")
	}
	doc.Type, err = strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}

	if !scanner.Scan() {
		return nil, errors.New("failed to read content from document file")
	}
	unquotedContent, err := strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}
	if doc.Content, err = ParseDocumentContent(doc.Type, unquotedContent); err != nil {
		return nil, err
	}

	return &doc, nil
}

/*
GetFileStore - Just a func that returns a FileStore
*/
func GetFileStore(config DocumentStoreConfig) (DocumentStore, error) {
	if len(config.StoreDirectory) == 0 {
		return nil, errors.New("A file store document configuration requires a valid directory")
	}
	if _, err := os.Stat(config.StoreDirectory); os.IsNotExist(err) {
		if err = os.MkdirAll(config.StoreDirectory, os.ModePerm); err != nil {
			return nil, fmt.Errorf("Cannot create file store for documents: %v", err)
		}
	}
	return &FileStore{config: config}, nil
}

/*--------------------------------------------------------------------------------------------------
 */
<commit_msg>Comments are important<commit_after>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, sub to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package leaplib

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path"
	"strconv"
)

/*--------------------------------------------------------------------------------------------------
 */

/*
FileStore - Most basic persistent implementation of DocumentStore. Simple stores each document into
a file within a configured directory.
*/
type FileStore struct {
	config DocumentStoreConfig
}

/*
Store - Store document in its file location.
*/
func (s *FileStore) Store(id string, doc *Document) error {
	file, err := os.Create(path.Join(s.config.StoreDirectory, id))
	if err != nil {
		return err
	}
	defer file.Close()

	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(doc.Title)); err != nil {
		return err
	}
	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(doc.Description)); err != nil {
		return err
	}
	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(doc.Type)); err != nil {
		return err
	}
	serialized, err := SerializeDocumentContent(doc.Type, doc.Content)
	if err != nil {
		return err
	}
	if _, err = fmt.Fprintln(file, strconv.QuoteToASCII(serialized)); err != nil {
		return err
	}

	return nil
}

/*
Fetch - Fetch document from its file location.
*/
func (s *FileStore) Fetch(id string) (*Document, error) {
	file, err := os.Open(path.Join(s.config.StoreDirectory, id))
	if err != nil {
		return nil, err
	}
	defer file.Close()

	doc := Document{ID: id}

	scanner := bufio.NewScanner(file)

	// Get title
	if !scanner.Scan() {
		return nil, errors.New("failed to read title from document file")
	}
	doc.Title, err = strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}

	// Get description
	if !scanner.Scan() {
		return nil, errors.New("failed to read description from document file")
	}
	doc.Description, err = strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}

	// Get type
	if !scanner.Scan() {
		return nil, errors.New("failed to read type from document file")
	}
	doc.Type, err = strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}

	// Get content
	if !scanner.Scan() {
		return nil, errors.New("failed to read content from document file")
	}
	unquotedContent, err := strconv.Unquote(scanner.Text())
	if err != nil {
		return nil, fmt.Errorf("unquote error: %v", err)
	}
	if doc.Content, err = ParseDocumentContent(doc.Type, unquotedContent); err != nil {
		return nil, err
	}

	return &doc, nil
}

/*
GetFileStore - Just a func that returns a FileStore
*/
func GetFileStore(config DocumentStoreConfig) (DocumentStore, error) {
	if len(config.StoreDirectory) == 0 {
		return nil, errors.New("A file store document configuration requires a valid directory")
	}
	if _, err := os.Stat(config.StoreDirectory); os.IsNotExist(err) {
		if err = os.MkdirAll(config.StoreDirectory, os.ModePerm); err != nil {
			return nil, fmt.Errorf("Cannot create file store for documents: %v", err)
		}
	}
	return &FileStore{config: config}, nil
}

/*--------------------------------------------------------------------------------------------------
 */
<|endoftext|>


<commit_before>/*
                       WARNING WARNING WARNING

                Attention all potential contributors

   This testfile is not in the best state. We've been slowly transitioning
   from the built in "testing" package to using Ginkgo. As you can see, we've
   changed the format, but a lot of the setup, test body, descriptions, etc
   are either hardcoded, completely lacking, or misleading.

   For example:

   Describe("Testing with ginkgo"...)      // This is not a great description
   It("TestDoesSoemthing"...)              // This is a horrible description

   Describe("create-user command"...       // Describe the actual object under test
   It("creates a user when provided ..."   // this is more descriptive

   For good examples of writing Ginkgo tests for the cli, refer to

   src/github.com/cloudfoundry/cli/cf/commands/application/delete_app_test.go
   src/github.com/cloudfoundry/cli/cf/terminal/ui_test.go
   src/github.com/cloudfoundry/loggregator_consumer/consumer_test.go
*/

package application_test

import (
	"time"

	"code.google.com/p/gogoprotobuf/proto"
	. "github.com/cloudfoundry/cli/cf/commands/application"
	"github.com/cloudfoundry/cli/cf/errors"
	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/cli/cf/terminal"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testlogs "github.com/cloudfoundry/cli/testhelpers/logs"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	"github.com/cloudfoundry/loggregatorlib/logmessage"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
)

var _ = Describe("logs command", func() {
	It("fails with usage when called without one argument", func() {
		requirementsFactory, logsRepo := getLogsDependencies()

		ui := callLogs([]string{}, requirementsFactory, logsRepo)
		Expect(ui.FailedWithUsage).To(BeTrue())
	})

	It("fails requirements when not logged in", func() {
		requirementsFactory, logsRepo := getLogsDependencies()
		requirementsFactory.LoginSuccess = false

		callLogs([]string{"my-app"}, requirementsFactory, logsRepo)
		Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
	})

	It("TestLogsOutputsRecentLogs", func() {
		app := models.Application{}
		app.Name = "my-app"
		app.Guid = "my-app-guid"

		currentTime := time.Now()

		recentLogs := []*logmessage.LogMessage{
			testlogs.NewLogMessage("Log Line 1", app.Guid, "DEA", currentTime),
			testlogs.NewLogMessage("Log Line 2", app.Guid, "DEA", currentTime),
		}

		requirementsFactory, logsRepo := getLogsDependencies()
		requirementsFactory.Application = app
		logsRepo.RecentLogs = recentLogs

		ui := callLogs([]string{"--recent", "my-app"}, requirementsFactory, logsRepo)

		Expect(requirementsFactory.ApplicationName).To(Equal("my-app"))
		Expect(app.Guid).To(Equal(logsRepo.AppLoggedGuid))
		Expect(ui.Outputs).To(ContainSubstrings(
			[]string{"Connected, dumping recent logs for app", "my-app", "my-org", "my-space", "my-user"},
			[]string{"Log Line 1"},
			[]string{"Log Line 2"},
		))
	})

	It("TestLogsEscapeFormattingVerbs", func() {
		app := models.Application{}
		app.Name = "my-app"
		app.Guid = "my-app-guid"

		recentLogs := []*logmessage.LogMessage{
			testlogs.NewLogMessage("hello%2Bworld%v", app.Guid, "DEA", time.Now()),
		}

		requirementsFactory, logsRepo := getLogsDependencies()
		requirementsFactory.Application = app
		logsRepo.RecentLogs = recentLogs

		ui := callLogs([]string{"--recent", "my-app"}, requirementsFactory, logsRepo)

		Expect(ui.Outputs).To(ContainSubstrings([]string{"hello%2Bworld%v"}))
	})

	It("TestLogsTailsTheAppLogs", func() {
		app := models.Application{}
		app.Name = "my-app"
		app.Guid = "my-app-guid"

		logs := []*logmessage.LogMessage{
			testlogs.NewLogMessage("Log Line 1", app.Guid, "DEA", time.Now()),
		}

		requirementsFactory, logsRepo := getLogsDependencies()
		requirementsFactory.Application = app
		logsRepo.TailLogMessages = logs

		ui := callLogs([]string{"my-app"}, requirementsFactory, logsRepo)

		Expect(requirementsFactory.ApplicationName).To(Equal("my-app"))
		Expect(app.Guid).To(Equal(logsRepo.AppLoggedGuid))
		Expect(ui.Outputs).To(ContainSubstrings(
			[]string{"Connected, tailing logs for app", "my-app", "my-org", "my-space", "my-user"},
			[]string{"Log Line 1"},
		))
	})

	Context("when the loggregator server has an invalid cert", func() {
		var (
			requirementsFactory *testreq.FakeReqFactory
			logsRepo            *testapi.FakeLogsRepository
		)

		BeforeEach(func() {
			requirementsFactory, logsRepo = getLogsDependencies()
		})

		Context("when the skip-ssl-validation flag is not set", func() {
			It("fails and informs the user about the skip-ssl-validation flag", func() {
				logsRepo.TailLogErr = errors.NewInvalidSSLCert("https://example.com", "it don't work good")
				ui := callLogs([]string{"my-app"}, requirementsFactory, logsRepo)

				Expect(ui.Outputs).To(ContainSubstrings(
					[]string{"Received invalid SSL certificate", "https://example.com"},
					[]string{"TIP"},
				))
			})

			It("informs the user of the error when they include the --recent flag", func() {
				logsRepo.RecentLogErr = errors.NewInvalidSSLCert("https://example.com", "how does SSL work???")
				ui := callLogs([]string{"--recent", "my-app"}, requirementsFactory, logsRepo)

				Expect(ui.Outputs).To(ContainSubstrings(
					[]string{"Received invalid SSL certificate", "https://example.com"},
					[]string{"TIP"},
				))
			})
		})
	})

	Context("when the loggregator server has a valid cert", func() {
		var (
			flags               []string
			requirementsFactory *testreq.FakeReqFactory
			logsRepo            *testapi.FakeLogsRepository
		)

		BeforeEach(func() {
			requirementsFactory, logsRepo = getLogsDependencies()
			flags = []string{"my-app"}
		})

		It("tails logs", func() {
			ui := callLogs(flags, requirementsFactory, logsRepo)

			Expect(ui.Outputs).To(ContainSubstrings(
				[]string{"Connected, tailing logs for app", "my-org", "my-space", "my-user"},
			))
		})
	})

	Describe("Helpers", func() {
		date := time.Date(2014, 4, 4, 11, 39, 20, 5, time.UTC)

		createMessage := func(sourceId string, sourceName string, msgType logmessage.LogMessage_MessageType, date time.Time) *logmessage.LogMessage {
			timestamp := date.UnixNano()
			return &logmessage.LogMessage{
				Message:     []byte("Hello World!\n\r\n\r"),
				AppId:       proto.String("my-app-guid"),
				MessageType: &msgType,
				SourceId:    &sourceId,
				Timestamp:   &timestamp,
				SourceName:  &sourceName,
			}
		}

		Context("when the message comes from an app", func() {
			It("includes the instance index", func() {
				msg := createMessage("4", "App", logmessage.LogMessage_OUT, date)
				Expect(terminal.Decolorize(LogMessageOutput(msg, time.UTC))).To(Equal("2014-04-04T11:39:20.00+0000 [App/4]   OUT Hello World!"))
			})
		})

		Context("when the message comes from a cloudfoundry component", func() {
			It("doesn't include the instance index", func() {
				msg := createMessage("4", "DEA", logmessage.LogMessage_OUT, date)
				Expect(terminal.Decolorize(LogMessageOutput(msg, time.UTC))).To(Equal("2014-04-04T11:39:20.00+0000 [DEA]     OUT Hello World!"))
			})
		})

		Context("when the message was written to stderr", func() {
			It("shows the log type as 'ERR'", func() {
				msg := createMessage("4", "DEA", logmessage.LogMessage_ERR, date)
				Expect(terminal.Decolorize(LogMessageOutput(msg, time.UTC))).To(Equal("2014-04-04T11:39:20.00+0000 [DEA]     ERR Hello World!"))
			})
		})

		It("formats the time in the given time zone", func() {
			msg := createMessage("4", "DEA", logmessage.LogMessage_ERR, date)
			Expect(terminal.Decolorize(LogMessageOutput(msg, time.FixedZone("the-zone", 3*60*60)))).To(Equal("2014-04-04T14:39:20.00+0300 [DEA]     ERR Hello World!"))
		})
	})
})

func getLogsDependencies() (requirementsFactory *testreq.FakeReqFactory, logsRepo *testapi.FakeLogsRepository) {
	logsRepo = &testapi.FakeLogsRepository{}
	requirementsFactory = &testreq.FakeReqFactory{LoginSuccess: true}
	return
}

func callLogs(args []string, requirementsFactory *testreq.FakeReqFactory, logsRepo *testapi.FakeLogsRepository) (ui *testterm.FakeUI) {
	ui = new(testterm.FakeUI)

	configRepo := testconfig.NewRepositoryWithDefaults()
	cmd := NewLogs(ui, configRepo, logsRepo)
	testcmd.RunCommand(cmd, args, requirementsFactory)
	return
}
<commit_msg>Cleanup 'cf logs' tests<commit_after>package application_test

import (
	"time"

	"code.google.com/p/gogoprotobuf/proto"
	"github.com/cloudfoundry/cli/cf/errors"
	"github.com/cloudfoundry/cli/cf/models"
	"github.com/cloudfoundry/cli/cf/terminal"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testlogs "github.com/cloudfoundry/cli/testhelpers/logs"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	"github.com/cloudfoundry/loggregatorlib/logmessage"

	. "github.com/cloudfoundry/cli/cf/commands/application"
	"github.com/cloudfoundry/cli/cf/configuration"
	. "github.com/cloudfoundry/cli/testhelpers/matchers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("logs command", func() {
	var (
		ui                  *testterm.FakeUI
		logsRepo            *testapi.FakeLogsRepository
		requirementsFactory *testreq.FakeReqFactory
		configRepo          configuration.ReadWriter
	)

	BeforeEach(func() {
		ui = &testterm.FakeUI{}
		configRepo = testconfig.NewRepositoryWithDefaults()
		logsRepo = &testapi.FakeLogsRepository{}
		requirementsFactory = &testreq.FakeReqFactory{}
	})

	runCommand := func(args ...string) {
		testcmd.RunCommand(NewLogs(ui, configRepo, logsRepo), args, requirementsFactory)
	}

	Describe("requirements", func() {
		It("fails with usage when called without one argument", func() {
			requirementsFactory.LoginSuccess = true

			runCommand()
			Expect(ui.FailedWithUsage).To(BeTrue())
		})

		It("fails requirements when not logged in", func() {
			runCommand("my-app")
			Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
		})
	})

	Context("when logged in", func() {
		var (
			app models.Application
		)

		BeforeEach(func() {
			requirementsFactory.LoginSuccess = true

			app = models.Application{}
			app.Name = "my-app"
			app.Guid = "my-app-guid"

			currentTime := time.Now()
			recentLogs := []*logmessage.LogMessage{
				testlogs.NewLogMessage("Log Line 1", app.Guid, "DEA", currentTime),
				testlogs.NewLogMessage("Log Line 2", app.Guid, "DEA", currentTime),
			}

			appLogs := []*logmessage.LogMessage{
				testlogs.NewLogMessage("Log Line 1", app.Guid, "DEA", time.Now()),
			}

			requirementsFactory.Application = app
			logsRepo.RecentLogs = recentLogs
			logsRepo.TailLogMessages = appLogs
		})

		It("shows the recent logs when the --recent flag is provided", func() {
			runCommand("--recent", "my-app")

			Expect(requirementsFactory.ApplicationName).To(Equal("my-app"))
			Expect(app.Guid).To(Equal(logsRepo.AppLoggedGuid))
			Expect(ui.Outputs).To(ContainSubstrings(
				[]string{"Connected, dumping recent logs for app", "my-app", "my-org", "my-space", "my-user"},
				[]string{"Log Line 1"},
				[]string{"Log Line 2"},
			))
		})

		Context("when the log messages contain format string identifiers", func() {
			BeforeEach(func() {
				logsRepo.RecentLogs = []*logmessage.LogMessage{
					testlogs.NewLogMessage("hello%2Bworld%v", app.Guid, "DEA", time.Now()),
				}
			})

			It("does not treat them as format strings", func() {
				runCommand("--recent", "my-app")
				Expect(ui.Outputs).To(ContainSubstrings([]string{"hello%2Bworld%v"}))
			})
		})

		It("tails the app's logs when no flags are given", func() {
			runCommand("my-app")

			Expect(requirementsFactory.ApplicationName).To(Equal("my-app"))
			Expect(app.Guid).To(Equal(logsRepo.AppLoggedGuid))
			Expect(ui.Outputs).To(ContainSubstrings(
				[]string{"Connected, tailing logs for app", "my-app", "my-org", "my-space", "my-user"},
				[]string{"Log Line 1"},
			))
		})

		Context("when the loggregator server has an invalid cert", func() {
			Context("when the skip-ssl-validation flag is not set", func() {
				It("fails and informs the user about the skip-ssl-validation flag", func() {
					logsRepo.TailLogErr = errors.NewInvalidSSLCert("https://example.com", "it don't work good")
					runCommand("my-app")

					Expect(ui.Outputs).To(ContainSubstrings(
						[]string{"Received invalid SSL certificate", "https://example.com"},
						[]string{"TIP"},
					))
				})

				It("informs the user of the error when they include the --recent flag", func() {
					logsRepo.RecentLogErr = errors.NewInvalidSSLCert("https://example.com", "how does SSL work???")
					runCommand("--recent", "my-app")

					Expect(ui.Outputs).To(ContainSubstrings(
						[]string{"Received invalid SSL certificate", "https://example.com"},
						[]string{"TIP"},
					))
				})
			})
		})

		Context("when the loggregator server has a valid cert", func() {
			It("tails logs", func() {
				runCommand("my-app")
				Expect(ui.Outputs).To(ContainSubstrings(
					[]string{"Connected, tailing logs for app", "my-org", "my-space", "my-user"},
				))
			})
		})

		Describe("Helpers", func() {
			date := time.Date(2014, 4, 4, 11, 39, 20, 5, time.UTC)

			createMessage := func(sourceId string, sourceName string, msgType logmessage.LogMessage_MessageType, date time.Time) *logmessage.LogMessage {
				timestamp := date.UnixNano()
				return &logmessage.LogMessage{
					Message:     []byte("Hello World!\n\r\n\r"),
					AppId:       proto.String("my-app-guid"),
					MessageType: &msgType,
					SourceId:    &sourceId,
					Timestamp:   &timestamp,
					SourceName:  &sourceName,
				}
			}

			Context("when the message comes from an app", func() {
				It("includes the instance index", func() {
					msg := createMessage("4", "App", logmessage.LogMessage_OUT, date)
					Expect(terminal.Decolorize(LogMessageOutput(msg, time.UTC))).To(Equal("2014-04-04T11:39:20.00+0000 [App/4]   OUT Hello World!"))
				})
			})

			Context("when the message comes from a cloudfoundry component", func() {
				It("doesn't include the instance index", func() {
					msg := createMessage("4", "DEA", logmessage.LogMessage_OUT, date)
					Expect(terminal.Decolorize(LogMessageOutput(msg, time.UTC))).To(Equal("2014-04-04T11:39:20.00+0000 [DEA]     OUT Hello World!"))
				})
			})

			Context("when the message was written to stderr", func() {
				It("shows the log type as 'ERR'", func() {
					msg := createMessage("4", "DEA", logmessage.LogMessage_ERR, date)
					Expect(terminal.Decolorize(LogMessageOutput(msg, time.UTC))).To(Equal("2014-04-04T11:39:20.00+0000 [DEA]     ERR Hello World!"))
				})
			})

			It("formats the time in the given time zone", func() {
				msg := createMessage("4", "DEA", logmessage.LogMessage_ERR, date)
				Expect(terminal.Decolorize(LogMessageOutput(msg, time.FixedZone("the-zone", 3*60*60)))).To(Equal("2014-04-04T14:39:20.00+0300 [DEA]     ERR Hello World!"))
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

package db

import (
	"leveldb/descriptor"
	"leveldb/errors"
	"leveldb/iter"
	"leveldb/log"
	"leveldb/opt"
	"sync"
)

type session struct {
	sync.RWMutex

	desc   descriptor.Descriptor
	o      *iOptions
	cmp    *iComparer
	filter *iFilter
	tops   *tOps

	manifest       *log.Writer
	manifestFile   descriptor.File
	manifestWriter descriptor.Writer

	st struct {
		sync.RWMutex
		version         *version
		versions        []*version
		nextNum         stateNum
		logNum          uint64
		seq             uint64
		compactPointers [kNumLevels]iKey
	}
}

func newSession(desc descriptor.Descriptor, o *opt.Options) *session {
	s := new(session)
	s.desc = desc
	s.o = &iOptions{s, o}
	s.cmp = &iComparer{o.GetComparer()}
	filter := o.GetFilter()
	if filter != nil {
		s.filter = &iFilter{filter}
	}
	s.tops = newTableOps(s, s.o.GetMaxOpenFiles())
	s.setVersion(&version{s: s})
	return s
}

// Create a new database session
func (s *session) create() (err error) {
	// create manifest
	err = s.createManifest(s.allocFileNum(), nil, nil)
	if err != nil {
		return
	}
	return
}

// Recover a database session
func (s *session) recover() (err error) {
	s.manifestFile, err = s.desc.GetMainManifest()
	if err != nil {
		return
	}

	r, err := s.manifestFile.Open()
	if err != nil {
		return
	}

	st := &s.st

	cmpName := s.cmp.cmp.Name()
	staging := st.version.newStaging()
	srec := new(sessionRecord)
	lr := log.NewReader(r, true)
	for lr.Next() {
		rec := new(sessionRecord)
		err = rec.decode(lr.Record())
		if err != nil {
			continue
		}

		if rec.hasComparer && rec.comparer != cmpName {
			return errors.ErrInvalid("invalid comparer, " +
				"want '" + cmpName + "', " +
				"got '" + rec.comparer + "'")
		}

		// save compact pointers
		for _, rp := range rec.compactPointers {
			st.compactPointers[rp.level] = iKey(rp.key)
		}

		// commit record to version staging
		staging.commit(rec)

		if rec.hasLogNum {
			srec.setLogNum(rec.logNum)
		}
		if rec.hasNextNum {
			srec.setNextNum(rec.nextNum)
		}
		if rec.hasSeq {
			srec.setSeq(rec.seq)
		}
	}
	// check for error in log reader
	err = lr.Error()
	if err != nil {
		return
	}

	switch false {
	case srec.hasNextNum:
		err = errors.ErrCorrupt("manifest missing next file number")
	case srec.hasLogNum:
		err = errors.ErrCorrupt("manifest missing log file number")
	case srec.hasSeq:
		err = errors.ErrCorrupt("manifest missing seq number")
	}
	if err != nil {
		return
	}

	s.setVersion(staging.finish())
	s.setFileNum(srec.nextNum)
	s.recordCommited(srec)

	return
}

func (s *session) commit(r *sessionRecord) (err error) {
	// spawn new version based on current version
	nv := s.st.version.spawn(r)

	if s.manifest == nil {
		// manifest log writer not yet created, create one
		err = s.createManifest(s.allocFileNum(), r, nv)
	} else {
		err = s.flushManifest(r)
	}

	// finally, apply new version if no error rise
	if err == nil {
		s.setVersion(nv)
	}

	return
}

func (s *session) needCompaction() bool {
	st := &s.st
	st.RLock()
	defer st.RUnlock()
	v := st.version
	return v.compactionScore >= 1 || v.seekCompactionTable != nil
}

func (s *session) pickCompaction() (c *compaction) {
	st := &s.st

	st.RLock()
	v := st.version
	bySize := v.compactionScore >= 1
	bySeek := v.seekCompactionTable != nil
	st.RUnlock()

	icmp := s.cmp
	ucmp := icmp.cmp

	var level int
	var t0 tFiles
	if bySize {
		level = v.compactionLevel
		cp := s.st.compactPointers[level]
		tt := v.tables[level]
		for _, t := range tt {
			if cp == nil || icmp.Compare(t.max, cp) > 0 {
				t0 = append(t0, t)
				break
			}
		}
		if len(t0) == 0 {
			t0 = append(t0, tt[0])
		}
	} else if bySeek {
		level = v.seekCompactionLevel
		t0 = append(t0, v.seekCompactionTable)
	} else {
		return
	}

	c = &compaction{s: s, version: v, level: level}
	if level == 0 {
		min, max := t0.getRange(icmp)
		t0 = nil
		v.tables[0].getOverlaps(min.ukey(), max.ukey(), &t0, false, ucmp)
	}

	c.tables[0] = t0
	c.expand()
	return
}

func (s *session) getCompactionRange(level int, min, max []byte) (c *compaction) {
	st := &s.st

	st.RLock()
	v := st.version
	st.RUnlock()

	var t0 tFiles
	v.tables[level].getOverlaps(min, max, &t0, level != 0, s.cmp.cmp)
	if len(t0) == 0 {
		return nil
	}

	c = &compaction{s: s, version: v, level: level}
	c.tables[0] = t0
	c.expand()
	return
}

type compaction struct {
	s       *session
	version *version

	level  int
	tables [2]tFiles

	gp              tFiles
	gpidx           int
	seenKey         bool
	overlappedBytes uint64
	min, max        iKey

	tPtrs [kNumLevels]int
}

func (c *compaction) expand() {
	s := c.s
	v := c.version
	icmp := s.cmp
	ucmp := icmp.cmp

	level := c.level
	vt0, vt1 := v.tables[level], v.tables[level+1]

	t0, t1 := c.tables[0], c.tables[1]
	min, max := t0.getRange(icmp)
	vt1.getOverlaps(min.ukey(), max.ukey(), &t1, true, ucmp)

	// Get entire range covered by compaction
	amin, amax := append(t0, t1...).getRange(icmp)

	// See if we can grow the number of inputs in "level" without
	// changing the number of "level+1" files we pick up.
	if len(t1) > 0 {
		var exp0 tFiles
		vt0.getOverlaps(amin.ukey(), amax.ukey(), &exp0, level != 0, ucmp)
		if len(exp0) > len(t0) && t1.size()+exp0.size() < kExpCompactionMaxBytes {
			var exp1 tFiles
			xmin, xmax := exp0.getRange(icmp)
			vt1.getOverlaps(xmin.ukey(), xmax.ukey(), &exp1, true, ucmp)
			if len(exp1) == len(t1) {
				s.printf("Compaction: expanding, level=%d from=`%d+%d (%d+%d bytes)' to=`%d+%d (%d+%d bytes)'",
					level, len(t0), len(t1), t0.size(), t1.size(),
					len(exp0), len(exp1), exp0.size(), exp1.size())
				min, max = xmin, xmax
				t0, t1 = exp0, exp1
				amin, amax = append(t0, t1...).getRange(icmp)
			}
		}
	}

	// Compute the set of grandparent files that overlap this compaction
	// (parent == level+1; grandparent == level+2)
	if level+2 < kNumLevels {
		v.tables[level+2].getOverlaps(amin.ukey(), amax.ukey(), &c.gp, true, ucmp)
	}

	c.tables[0], c.tables[1] = t0, t1
	c.min, c.max = min, max
}

func (c *compaction) trivial() bool {
	return len(c.tables[0]) == 1 && len(c.tables[1]) == 0 && c.gp.size() <= kMaxGrandParentOverlapBytes
}

func (c *compaction) isBaseLevelForKey(key []byte) bool {
	s := c.s
	v := c.version
	ucmp := s.cmp.cmp
	for level, tt := range v.tables[c.level+2:] {
		for c.tPtrs[level] < len(tt) {
			t := tt[c.tPtrs[level]]
			if ucmp.Compare(key, t.max.ukey()) <= 0 {
				// We've advanced far enough
				if ucmp.Compare(key, t.min.ukey()) >= 0 {
					// Key falls in this file's range, so definitely not base level
					return false
				}
				break
			}
			c.tPtrs[level]++
		}
	}
	return true
}

func (c *compaction) shouldStopBefore(key iKey) bool {
	icmp := c.s.cmp
	for ; c.gpidx < len(c.gp); c.gpidx++ {
		gp := c.gp[c.gpidx]
		if icmp.Compare(key, gp.max) <= 0 {
			break
		}
		if c.seenKey {
			c.overlappedBytes += gp.size
		}
	}
	c.seenKey = true

	if c.overlappedBytes > kMaxGrandParentOverlapBytes {
		// Too much overlap for current output; start new output
		c.overlappedBytes = 0
		return true
	}
	return false
}

func (c *compaction) newIterator() iter.Iterator {
	s := c.s
	icmp := s.cmp

	level := c.level
	icap := 2
	if c.level == 0 {
		icap = len(c.tables[0]) + 1
	}
	its := make([]iter.Iterator, 0, icap)

	ro := &opt.ReadOptions{}

	for i, tt := range c.tables {
		if len(tt) == 0 {
			continue
		}

		if level+i == 0 {
			for _, t := range tt {
				its = append(its, s.tops.newIterator(t, ro))
			}
		} else {
			it := iter.NewIndexedIterator(tt.newIndexIterator(s.tops, icmp, ro))
			its = append(its, it)
		}
	}

	return iter.NewMergedIterator(its, icmp)
}
<commit_msg>db: close manifest reader when done; during session recovery<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

package db

import (
	"leveldb/descriptor"
	"leveldb/errors"
	"leveldb/iter"
	"leveldb/log"
	"leveldb/opt"
	"sync"
)

type session struct {
	sync.RWMutex

	desc   descriptor.Descriptor
	o      *iOptions
	cmp    *iComparer
	filter *iFilter
	tops   *tOps

	manifest       *log.Writer
	manifestFile   descriptor.File
	manifestWriter descriptor.Writer

	st struct {
		sync.RWMutex
		version         *version
		versions        []*version
		nextNum         stateNum
		logNum          uint64
		seq             uint64
		compactPointers [kNumLevels]iKey
	}
}

func newSession(desc descriptor.Descriptor, o *opt.Options) *session {
	s := new(session)
	s.desc = desc
	s.o = &iOptions{s, o}
	s.cmp = &iComparer{o.GetComparer()}
	filter := o.GetFilter()
	if filter != nil {
		s.filter = &iFilter{filter}
	}
	s.tops = newTableOps(s, s.o.GetMaxOpenFiles())
	s.setVersion(&version{s: s})
	return s
}

// Create a new database session
func (s *session) create() (err error) {
	// create manifest
	err = s.createManifest(s.allocFileNum(), nil, nil)
	if err != nil {
		return
	}
	return
}

// Recover a database session
func (s *session) recover() (err error) {
	s.manifestFile, err = s.desc.GetMainManifest()
	if err != nil {
		return
	}

	r, err := s.manifestFile.Open()
	if err != nil {
		return
	}
	defer r.Close()

	st := &s.st

	cmpName := s.cmp.cmp.Name()
	staging := st.version.newStaging()
	srec := new(sessionRecord)
	lr := log.NewReader(r, true)
	for lr.Next() {
		rec := new(sessionRecord)
		err = rec.decode(lr.Record())
		if err != nil {
			continue
		}

		if rec.hasComparer && rec.comparer != cmpName {
			return errors.ErrInvalid("invalid comparer, " +
				"want '" + cmpName + "', " +
				"got '" + rec.comparer + "'")
		}

		// save compact pointers
		for _, rp := range rec.compactPointers {
			st.compactPointers[rp.level] = iKey(rp.key)
		}

		// commit record to version staging
		staging.commit(rec)

		if rec.hasLogNum {
			srec.setLogNum(rec.logNum)
		}
		if rec.hasNextNum {
			srec.setNextNum(rec.nextNum)
		}
		if rec.hasSeq {
			srec.setSeq(rec.seq)
		}
	}
	// check for error in log reader
	err = lr.Error()
	if err != nil {
		return
	}

	switch false {
	case srec.hasNextNum:
		err = errors.ErrCorrupt("manifest missing next file number")
	case srec.hasLogNum:
		err = errors.ErrCorrupt("manifest missing log file number")
	case srec.hasSeq:
		err = errors.ErrCorrupt("manifest missing seq number")
	}
	if err != nil {
		return
	}

	s.setVersion(staging.finish())
	s.setFileNum(srec.nextNum)
	s.recordCommited(srec)

	return
}

func (s *session) commit(r *sessionRecord) (err error) {
	// spawn new version based on current version
	nv := s.st.version.spawn(r)

	if s.manifest == nil {
		// manifest log writer not yet created, create one
		err = s.createManifest(s.allocFileNum(), r, nv)
	} else {
		err = s.flushManifest(r)
	}

	// finally, apply new version if no error rise
	if err == nil {
		s.setVersion(nv)
	}

	return
}

func (s *session) needCompaction() bool {
	st := &s.st
	st.RLock()
	defer st.RUnlock()
	v := st.version
	return v.compactionScore >= 1 || v.seekCompactionTable != nil
}

func (s *session) pickCompaction() (c *compaction) {
	st := &s.st

	st.RLock()
	v := st.version
	bySize := v.compactionScore >= 1
	bySeek := v.seekCompactionTable != nil
	st.RUnlock()

	icmp := s.cmp
	ucmp := icmp.cmp

	var level int
	var t0 tFiles
	if bySize {
		level = v.compactionLevel
		cp := s.st.compactPointers[level]
		tt := v.tables[level]
		for _, t := range tt {
			if cp == nil || icmp.Compare(t.max, cp) > 0 {
				t0 = append(t0, t)
				break
			}
		}
		if len(t0) == 0 {
			t0 = append(t0, tt[0])
		}
	} else if bySeek {
		level = v.seekCompactionLevel
		t0 = append(t0, v.seekCompactionTable)
	} else {
		return
	}

	c = &compaction{s: s, version: v, level: level}
	if level == 0 {
		min, max := t0.getRange(icmp)
		t0 = nil
		v.tables[0].getOverlaps(min.ukey(), max.ukey(), &t0, false, ucmp)
	}

	c.tables[0] = t0
	c.expand()
	return
}

func (s *session) getCompactionRange(level int, min, max []byte) (c *compaction) {
	st := &s.st

	st.RLock()
	v := st.version
	st.RUnlock()

	var t0 tFiles
	v.tables[level].getOverlaps(min, max, &t0, level != 0, s.cmp.cmp)
	if len(t0) == 0 {
		return nil
	}

	c = &compaction{s: s, version: v, level: level}
	c.tables[0] = t0
	c.expand()
	return
}

type compaction struct {
	s       *session
	version *version

	level  int
	tables [2]tFiles

	gp              tFiles
	gpidx           int
	seenKey         bool
	overlappedBytes uint64
	min, max        iKey

	tPtrs [kNumLevels]int
}

func (c *compaction) expand() {
	s := c.s
	v := c.version
	icmp := s.cmp
	ucmp := icmp.cmp

	level := c.level
	vt0, vt1 := v.tables[level], v.tables[level+1]

	t0, t1 := c.tables[0], c.tables[1]
	min, max := t0.getRange(icmp)
	vt1.getOverlaps(min.ukey(), max.ukey(), &t1, true, ucmp)

	// Get entire range covered by compaction
	amin, amax := append(t0, t1...).getRange(icmp)

	// See if we can grow the number of inputs in "level" without
	// changing the number of "level+1" files we pick up.
	if len(t1) > 0 {
		var exp0 tFiles
		vt0.getOverlaps(amin.ukey(), amax.ukey(), &exp0, level != 0, ucmp)
		if len(exp0) > len(t0) && t1.size()+exp0.size() < kExpCompactionMaxBytes {
			var exp1 tFiles
			xmin, xmax := exp0.getRange(icmp)
			vt1.getOverlaps(xmin.ukey(), xmax.ukey(), &exp1, true, ucmp)
			if len(exp1) == len(t1) {
				s.printf("Compaction: expanding, level=%d from=`%d+%d (%d+%d bytes)' to=`%d+%d (%d+%d bytes)'",
					level, len(t0), len(t1), t0.size(), t1.size(),
					len(exp0), len(exp1), exp0.size(), exp1.size())
				min, max = xmin, xmax
				t0, t1 = exp0, exp1
				amin, amax = append(t0, t1...).getRange(icmp)
			}
		}
	}

	// Compute the set of grandparent files that overlap this compaction
	// (parent == level+1; grandparent == level+2)
	if level+2 < kNumLevels {
		v.tables[level+2].getOverlaps(amin.ukey(), amax.ukey(), &c.gp, true, ucmp)
	}

	c.tables[0], c.tables[1] = t0, t1
	c.min, c.max = min, max
}

func (c *compaction) trivial() bool {
	return len(c.tables[0]) == 1 && len(c.tables[1]) == 0 && c.gp.size() <= kMaxGrandParentOverlapBytes
}

func (c *compaction) isBaseLevelForKey(key []byte) bool {
	s := c.s
	v := c.version
	ucmp := s.cmp.cmp
	for level, tt := range v.tables[c.level+2:] {
		for c.tPtrs[level] < len(tt) {
			t := tt[c.tPtrs[level]]
			if ucmp.Compare(key, t.max.ukey()) <= 0 {
				// We've advanced far enough
				if ucmp.Compare(key, t.min.ukey()) >= 0 {
					// Key falls in this file's range, so definitely not base level
					return false
				}
				break
			}
			c.tPtrs[level]++
		}
	}
	return true
}

func (c *compaction) shouldStopBefore(key iKey) bool {
	icmp := c.s.cmp
	for ; c.gpidx < len(c.gp); c.gpidx++ {
		gp := c.gp[c.gpidx]
		if icmp.Compare(key, gp.max) <= 0 {
			break
		}
		if c.seenKey {
			c.overlappedBytes += gp.size
		}
	}
	c.seenKey = true

	if c.overlappedBytes > kMaxGrandParentOverlapBytes {
		// Too much overlap for current output; start new output
		c.overlappedBytes = 0
		return true
	}
	return false
}

func (c *compaction) newIterator() iter.Iterator {
	s := c.s
	icmp := s.cmp

	level := c.level
	icap := 2
	if c.level == 0 {
		icap = len(c.tables[0]) + 1
	}
	its := make([]iter.Iterator, 0, icap)

	ro := &opt.ReadOptions{}

	for i, tt := range c.tables {
		if len(tt) == 0 {
			continue
		}

		if level+i == 0 {
			for _, t := range tt {
				its = append(its, s.tops.newIterator(t, ro))
			}
		} else {
			it := iter.NewIndexedIterator(tt.newIndexIterator(s.tops, icmp, ro))
			its = append(its, it)
		}
	}

	return iter.NewMergedIterator(its, icmp)
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/mitchellh/cli"
	"os"
	"fmt"
	"log"
	"encoding/base64"
	"time"
)

const AWS_REGION_FOR_TESTING = "us-east-1"
const AMAZON_LINUX_AMI_ID = "ami-08111162"

const TEST_FILE_PATH = "/home/ec2-user/test-file"
const USER_DATA_TEMPLATE =
`#!/bin/bash
set -e
echo '%s' > "%s"
`
// An integration test that runs an EC2 instance, uses create_command to take a snapshot of it, and then delete_command
// to delete that snapshot.
func TestCreateAndDelete(t *testing.T) {
	t.Parallel()

	logger, ui := createLoggerAndUi("TestCreateAndDelete")
	session := session.New(&aws.Config{Region: aws.String(AWS_REGION_FOR_TESTING)})
	svc := ec2.New(session)

	instance, instanceName := launchInstance(svc, logger, t)
	defer terminateInstance(instance, svc, logger, t)
	waitForInstanceToStart(instance, svc, logger, t)

	snapshotId := takeSnapshotWithVerification(instanceName, *instance.InstanceId, ui, svc, logger, t)
	deleteSnapshotWithVerification(instanceName, snapshotId, ui, svc, logger, t)
}

// An integration test that runs an EC2 instance, uses create_command to take a snapshot of it, and then calls the
// delete_command to delete that snapshot, but setting the older than parmaeter in a way that should prevent any actual
// deletion.
func TestDeleteRespectsOlderThan(t *testing.T) {
	t.Parallel()

	logger, ui := createLoggerAndUi("TestDeleteRespectsOlderThan")
	session := session.New(&aws.Config{Region: aws.String(AWS_REGION_FOR_TESTING)})
	svc := ec2.New(session)

	instance, instanceName := launchInstance(svc, logger, t)
	defer terminateInstance(instance, svc, logger, t)
	waitForInstanceToStart(instance, svc, logger, t)

	snapshotId := takeSnapshotWithVerification(instanceName, *instance.InstanceId, ui, svc, logger, t)
	// Always try to delete the snapshot at the end so the tests don't litter the AWS account with snapshots
	defer deleteSnapshotWithVerification(instanceName, snapshotId, ui, svc, logger, t)

	// Set olderThan to "10h" to ensure the snapshot, which is only a few seconds old, does not get deleted
	deleteSnapshotForInstance(instanceName, "10h", 0, ui, logger, t)
	waitForSnapshotToBeDeleted(snapshotId, svc, logger, t)
	verifySnapshotWorks(snapshotId, svc, logger, t)
}

// An integration test that runs an EC2 instance, uses create_command to take a snapshot of it, and then calls the
// delete_command to delete that snapshot, but setting the at least parameter in a way that should prevent any actual
// deletion.
func TestDeleteRespectsAtLeast(t *testing.T) {
	t.Parallel()

	logger, ui := createLoggerAndUi("TestDeleteRespectsAtLeast")
	session := session.New(&aws.Config{Region: aws.String(AWS_REGION_FOR_TESTING)})
	svc := ec2.New(session)

	instance, instanceName := launchInstance(svc, logger, t)
	defer terminateInstance(instance, svc, logger, t)
	waitForInstanceToStart(instance, svc, logger, t)

	snapshotId := takeSnapshotWithVerification(instanceName, *instance.InstanceId, ui, svc, logger, t)
	// Always try to delete the snapshot at the end so the tests don't litter the AWS account with snapshots
	defer deleteSnapshotWithVerification(instanceName, snapshotId, ui, svc, logger, t)

	// Set atLeast to 1 to ensure the snapshot, which is the only one that exists, does not get deleted
	deleteSnapshotForInstance(instanceName, "0h", 1, ui, logger, t)
	waitForSnapshotToBeDeleted(snapshotId, svc, logger, t)
	verifySnapshotWorks(snapshotId, svc, logger, t)
}

func TestCreateWithInvalidInstanceName(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestCreateWithInvalidInstanceName")
	cmd := CreateCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: "not-a-valid-instance-name",
		AmiName: "this-ami-should-not-be-created",
	}

	_, err := create(cmd)

	if err == nil {
		t.Fatalf("Expected an error when creating a snapshot of an instance name that doesn't exist, but instead got nil")
	}
}

func TestCreateWithInvalidInstanceId(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestCreateWithInvalidInstanceId")
	cmd := CreateCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceId: "not-a-valid-instance-id",
		AmiName: "this-ami-should-not-be-created",
	}

	_, err := create(cmd)

	if err == nil {
		t.Fatalf("Expected an error when creating a snapshot of an instance id that doesn't exist, but instead got nil")
	}
}

func TestDeleteWithInvalidInstanceName(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestDeleteWithInvalidInstanceName")
	cmd := DeleteCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: "not-a-valid-instance-name",
		OlderThan: "0h",
		RequireAtLeast: 0,
	}

	err := deleteSnapshots(cmd)

	if err == nil {
		t.Fatalf("Expected an error when deleting a snapshot of an instance name that doesn't exist, but instead got nil")
	}
}

func TestDeleteWithInvalidInstanceId(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestDeleteWithInvalidInstanceId")
	cmd := DeleteCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceId: "not-a-valid-instance-id",
		OlderThan: "0h",
		RequireAtLeast: 0,
	}

	err := deleteSnapshots(cmd)

	if err == nil {
		t.Fatalf("Expected an error when deleting a snapshot of an instance id that doesn't exist, but instead got nil")
	}
}

func launchInstance(svc *ec2.EC2, logger *log.Logger, t *testing.T) (*ec2.Instance, string) {
	instanceName := fmt.Sprintf("ec2-snapper-unit-test-%s", UniqueId())
	userData := fmt.Sprint(USER_DATA_TEMPLATE, instanceName, TEST_FILE_PATH)

	logger.Printf("Launching EC2 instance in region %s. Its User Data will create a file %s with contents %s.", AWS_REGION_FOR_TESTING, TEST_FILE_PATH, instanceName)

	runResult, err := svc.RunInstances(&ec2.RunInstancesInput{
		ImageId:      aws.String(AMAZON_LINUX_AMI_ID),
		InstanceType: aws.String("t2.micro"),
		MinCount:     aws.Int64(1),
		MaxCount:     aws.Int64(1),
		UserData:     aws.String(base64.StdEncoding.EncodeToString([]byte(userData))),
	})

	if err != nil {
		t.Fatal(err)
	}

	if len(runResult.Instances) != 1 {
		t.Fatalf("Expected to launch 1 instance but got %d", len(runResult.Instances))
	}

	instance := runResult.Instances[0]
	logger.Printf("Launched instance %s", *instance.InstanceId)

	tagInstance(instance, instanceName, svc, logger, t)

	return instance, instanceName
}

func tagInstance(instance *ec2.Instance, instanceName string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Adding tags to instance %s", *instance.InstanceId)

	_ , err := svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{instance.InstanceId},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("Name"),
				Value: aws.String(instanceName),
			},
		},
	})
	if err != nil {
		t.Fatal(err)
	}
}

func waitForInstanceToStart(instance *ec2.Instance, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Waiting for instance %s to start...", *instance.InstanceId)

	if err := svc.WaitUntilInstanceRunning(&ec2.DescribeInstancesInput{InstanceIds: []*string{instance.InstanceId}}); err != nil {
		t.Fatal(err)
	}

	logger.Printf("Instance %s is now running", *instance.InstanceId)
}

func terminateInstance(instance *ec2.Instance, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Terminating instance %s", *instance.InstanceId)
	if _, err := svc.TerminateInstances(&ec2.TerminateInstancesInput{InstanceIds: []*string{instance.InstanceId}}); err != nil {
		t.Fatal("Failed to terminate instance %s", *instance.InstanceId)
	}
}

func takeSnapshot(instanceName string, ui cli.Ui, logger *log.Logger, t *testing.T) string {
	log.Printf("Creating a snapshot with name %s.", instanceName)

	cmd := CreateCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: instanceName,
		AmiName: instanceName,
	}

	snapshotId, err := create(cmd)

	if err != nil {
		t.Fatal(err)
	}

	logger.Printf("Created snasphot %s", snapshotId)
	return snapshotId
}

func verifySnapshotWorks(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Verifying snapshot %s exists", snapshotId)

	snapshots := findSnapshots(snapshotId, svc, logger, t)
	if len(snapshots) != 1 {
		t.Fatalf("Expected to find one snapshot with id %s but found %d", snapshotId, len(snapshots))
	}

	snapshot := snapshots[0]

	if *snapshot.State == ec2.ImageStateAvailable {
		logger.Printf("Found snapshot %s in expected state %s", snapshotId, *snapshot.State)
	} else {
		t.Fatalf("Expected image to be in state %s, but it was in state %s", ec2.ImageStateAvailable, *snapshot.State)
	}

	// TODO: fire up a new EC2 instance with the snapshot, SSH to it, and check the file we wrote is still there
}

func verifySnapshotIsDeleted(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Verifying snapshot %s is deleted", snapshotId)
	snapshots := findSnapshots(snapshotId, svc, logger, t)
	if len(snapshots) != 0 {
		t.Fatalf("Expected to find zero snapshots with id %s but found %d", snapshotId, len(snapshots))
	}
}

func findSnapshots(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) []*ec2.Image {
	resp, err := svc.DescribeImages(&ec2.DescribeImagesInput{ImageIds: []*string{&snapshotId}})
	if err != nil {
		t.Fatal(err)
	}

	return resp.Images
}

func waitForSnapshotToBeAvailable(instanceId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Waiting for snapshot for instance %s to become available", instanceId)

	instanceIdTagFilter := &ec2.Filter{
		Name: aws.String(fmt.Sprintf("tag:%s", EC2_SNAPPER_INSTANCE_ID_TAG)),
		Values: []*string{aws.String(instanceId)},
	}

	if err := svc.WaitUntilImageAvailable(&ec2.DescribeImagesInput{Filters: []*ec2.Filter{instanceIdTagFilter}}); err != nil {
		t.Fatal(err)
	}
}

func waitForSnapshotToBeDeleted(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Waiting for snapshot %s to be deleted", snapshotId)

	// We just do a simple sleep, as there is no built-in API call to wait for this.
	time.Sleep(30 * time.Second)
}

func deleteSnapshotForInstance(instanceName string, olderThan string, requireAtLeast int, ui cli.Ui, logger *log.Logger, t *testing.T) {
	logger.Printf("Deleting snapshot for instance %s", instanceName)

	deleteCmd := DeleteCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: instanceName,
		OlderThan: olderThan,
		RequireAtLeast: requireAtLeast,
	}

	if err := deleteSnapshots(deleteCmd); err != nil {
		t.Fatal(err)
	}
}

func takeSnapshotWithVerification(instanceName string, instanceId string, ui cli.Ui, svc *ec2.EC2, logger *log.Logger, t *testing.T) string {
	snapshotId := takeSnapshot(instanceName, ui, logger, t)

	waitForSnapshotToBeAvailable(instanceId, svc, logger, t)
	verifySnapshotWorks(snapshotId, svc, logger, t)

	return snapshotId
}

func deleteSnapshotWithVerification(instanceName string, snapshotId string, ui cli.Ui, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	deleteSnapshotForInstance(instanceName, "0h", 0, ui, logger, t)
	waitForSnapshotToBeDeleted(snapshotId, svc, logger, t)
	verifySnapshotIsDeleted(snapshotId, svc, logger, t)
}

func createLoggerAndUi(testName string) (*log.Logger, cli.Ui) {
	logger := log.New(os.Stdout, testName + " ", log.LstdFlags)

	basicUi := &cli.BasicUi{
		Reader:      os.Stdin,
		Writer:      os.Stdout,
		ErrorWriter: os.Stderr,
	}

	prefixedUi := &cli.PrefixedUi{
		AskPrefix:		logger.Prefix(),
		AskSecretPrefix:	logger.Prefix(),
		OutputPrefix:		logger.Prefix(),
		InfoPrefix:		logger.Prefix(),
		ErrorPrefix:		logger.Prefix(),
		WarnPrefix:		logger.Prefix(),
		Ui:			basicUi,

	}

	return logger, prefixedUi
}<commit_msg>Wait for EC2 instance to exist before tagging it<commit_after>package main

import (
	"testing"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/mitchellh/cli"
	"os"
	"fmt"
	"log"
	"encoding/base64"
	"time"
)

const AWS_REGION_FOR_TESTING = "us-east-1"
const AMAZON_LINUX_AMI_ID = "ami-08111162"

const TEST_FILE_PATH = "/home/ec2-user/test-file"
const USER_DATA_TEMPLATE =
`#!/bin/bash
set -e
echo '%s' > "%s"
`
// An integration test that runs an EC2 instance, uses create_command to take a snapshot of it, and then delete_command
// to delete that snapshot.
func TestCreateAndDelete(t *testing.T) {
	t.Parallel()

	logger, ui := createLoggerAndUi("TestCreateAndDelete")
	session := session.New(&aws.Config{Region: aws.String(AWS_REGION_FOR_TESTING)})
	svc := ec2.New(session)

	instance, instanceName := launchInstance(svc, logger, t)
	defer terminateInstance(instance, svc, logger, t)
	waitForInstanceToStart(instance, svc, logger, t)

	snapshotId := takeSnapshotWithVerification(instanceName, *instance.InstanceId, ui, svc, logger, t)
	deleteSnapshotWithVerification(instanceName, snapshotId, ui, svc, logger, t)
}

// An integration test that runs an EC2 instance, uses create_command to take a snapshot of it, and then calls the
// delete_command to delete that snapshot, but setting the older than parmaeter in a way that should prevent any actual
// deletion.
func TestDeleteRespectsOlderThan(t *testing.T) {
	t.Parallel()

	logger, ui := createLoggerAndUi("TestDeleteRespectsOlderThan")
	session := session.New(&aws.Config{Region: aws.String(AWS_REGION_FOR_TESTING)})
	svc := ec2.New(session)

	instance, instanceName := launchInstance(svc, logger, t)
	defer terminateInstance(instance, svc, logger, t)
	waitForInstanceToStart(instance, svc, logger, t)

	snapshotId := takeSnapshotWithVerification(instanceName, *instance.InstanceId, ui, svc, logger, t)
	// Always try to delete the snapshot at the end so the tests don't litter the AWS account with snapshots
	defer deleteSnapshotWithVerification(instanceName, snapshotId, ui, svc, logger, t)

	// Set olderThan to "10h" to ensure the snapshot, which is only a few seconds old, does not get deleted
	deleteSnapshotForInstance(instanceName, "10h", 0, ui, logger, t)
	waitForSnapshotToBeDeleted(snapshotId, svc, logger, t)
	verifySnapshotWorks(snapshotId, svc, logger, t)
}

// An integration test that runs an EC2 instance, uses create_command to take a snapshot of it, and then calls the
// delete_command to delete that snapshot, but setting the at least parameter in a way that should prevent any actual
// deletion.
func TestDeleteRespectsAtLeast(t *testing.T) {
	t.Parallel()

	logger, ui := createLoggerAndUi("TestDeleteRespectsAtLeast")
	session := session.New(&aws.Config{Region: aws.String(AWS_REGION_FOR_TESTING)})
	svc := ec2.New(session)

	instance, instanceName := launchInstance(svc, logger, t)
	defer terminateInstance(instance, svc, logger, t)
	waitForInstanceToStart(instance, svc, logger, t)

	snapshotId := takeSnapshotWithVerification(instanceName, *instance.InstanceId, ui, svc, logger, t)
	// Always try to delete the snapshot at the end so the tests don't litter the AWS account with snapshots
	defer deleteSnapshotWithVerification(instanceName, snapshotId, ui, svc, logger, t)

	// Set atLeast to 1 to ensure the snapshot, which is the only one that exists, does not get deleted
	deleteSnapshotForInstance(instanceName, "0h", 1, ui, logger, t)
	waitForSnapshotToBeDeleted(snapshotId, svc, logger, t)
	verifySnapshotWorks(snapshotId, svc, logger, t)
}

func TestCreateWithInvalidInstanceName(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestCreateWithInvalidInstanceName")
	cmd := CreateCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: "not-a-valid-instance-name",
		AmiName: "this-ami-should-not-be-created",
	}

	_, err := create(cmd)

	if err == nil {
		t.Fatalf("Expected an error when creating a snapshot of an instance name that doesn't exist, but instead got nil")
	}
}

func TestCreateWithInvalidInstanceId(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestCreateWithInvalidInstanceId")
	cmd := CreateCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceId: "not-a-valid-instance-id",
		AmiName: "this-ami-should-not-be-created",
	}

	_, err := create(cmd)

	if err == nil {
		t.Fatalf("Expected an error when creating a snapshot of an instance id that doesn't exist, but instead got nil")
	}
}

func TestDeleteWithInvalidInstanceName(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestDeleteWithInvalidInstanceName")
	cmd := DeleteCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: "not-a-valid-instance-name",
		OlderThan: "0h",
		RequireAtLeast: 0,
	}

	err := deleteSnapshots(cmd)

	if err == nil {
		t.Fatalf("Expected an error when deleting a snapshot of an instance name that doesn't exist, but instead got nil")
	}
}

func TestDeleteWithInvalidInstanceId(t *testing.T) {
	t.Parallel()

	_, ui := createLoggerAndUi("TestDeleteWithInvalidInstanceId")
	cmd := DeleteCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceId: "not-a-valid-instance-id",
		OlderThan: "0h",
		RequireAtLeast: 0,
	}

	err := deleteSnapshots(cmd)

	if err == nil {
		t.Fatalf("Expected an error when deleting a snapshot of an instance id that doesn't exist, but instead got nil")
	}
}

func launchInstance(svc *ec2.EC2, logger *log.Logger, t *testing.T) (*ec2.Instance, string) {
	instanceName := fmt.Sprintf("ec2-snapper-unit-test-%s", UniqueId())
	userData := fmt.Sprint(USER_DATA_TEMPLATE, instanceName, TEST_FILE_PATH)

	logger.Printf("Launching EC2 instance in region %s. Its User Data will create a file %s with contents %s.", AWS_REGION_FOR_TESTING, TEST_FILE_PATH, instanceName)

	runResult, err := svc.RunInstances(&ec2.RunInstancesInput{
		ImageId:      aws.String(AMAZON_LINUX_AMI_ID),
		InstanceType: aws.String("t2.micro"),
		MinCount:     aws.Int64(1),
		MaxCount:     aws.Int64(1),
		UserData:     aws.String(base64.StdEncoding.EncodeToString([]byte(userData))),
	})

	if err != nil {
		t.Fatal(err)
	}

	if len(runResult.Instances) != 1 {
		t.Fatalf("Expected to launch 1 instance but got %d", len(runResult.Instances))
	}

	instance := runResult.Instances[0]
	logger.Printf("Launched instance %s", *instance.InstanceId)

	err = svc.WaitUntilInstanceExists(&ec2.DescribeInstancesInput{InstanceIds: []*string{instance.InstanceId}})
	if err != nil {
		t.Fatal(err)
	}

	tagInstance(instance, instanceName, svc, logger, t)

	return instance, instanceName
}

func tagInstance(instance *ec2.Instance, instanceName string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Adding tags to instance %s", *instance.InstanceId)

	_ , err := svc.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{instance.InstanceId},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("Name"),
				Value: aws.String(instanceName),
			},
		},
	})
	if err != nil {
		t.Fatal(err)
	}
}

func waitForInstanceToStart(instance *ec2.Instance, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Waiting for instance %s to start...", *instance.InstanceId)

	if err := svc.WaitUntilInstanceRunning(&ec2.DescribeInstancesInput{InstanceIds: []*string{instance.InstanceId}}); err != nil {
		t.Fatal(err)
	}

	logger.Printf("Instance %s is now running", *instance.InstanceId)
}

func terminateInstance(instance *ec2.Instance, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Terminating instance %s", *instance.InstanceId)
	if _, err := svc.TerminateInstances(&ec2.TerminateInstancesInput{InstanceIds: []*string{instance.InstanceId}}); err != nil {
		t.Fatal("Failed to terminate instance %s", *instance.InstanceId)
	}
}

func takeSnapshot(instanceName string, ui cli.Ui, logger *log.Logger, t *testing.T) string {
	log.Printf("Creating a snapshot with name %s.", instanceName)

	cmd := CreateCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: instanceName,
		AmiName: instanceName,
	}

	snapshotId, err := create(cmd)

	if err != nil {
		t.Fatal(err)
	}

	logger.Printf("Created snasphot %s", snapshotId)
	return snapshotId
}

func verifySnapshotWorks(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Verifying snapshot %s exists", snapshotId)

	snapshots := findSnapshots(snapshotId, svc, logger, t)
	if len(snapshots) != 1 {
		t.Fatalf("Expected to find one snapshot with id %s but found %d", snapshotId, len(snapshots))
	}

	snapshot := snapshots[0]

	if *snapshot.State == ec2.ImageStateAvailable {
		logger.Printf("Found snapshot %s in expected state %s", snapshotId, *snapshot.State)
	} else {
		t.Fatalf("Expected image to be in state %s, but it was in state %s", ec2.ImageStateAvailable, *snapshot.State)
	}

	// TODO: fire up a new EC2 instance with the snapshot, SSH to it, and check the file we wrote is still there
}

func verifySnapshotIsDeleted(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Verifying snapshot %s is deleted", snapshotId)
	snapshots := findSnapshots(snapshotId, svc, logger, t)
	if len(snapshots) != 0 {
		t.Fatalf("Expected to find zero snapshots with id %s but found %d", snapshotId, len(snapshots))
	}
}

func findSnapshots(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) []*ec2.Image {
	resp, err := svc.DescribeImages(&ec2.DescribeImagesInput{ImageIds: []*string{&snapshotId}})
	if err != nil {
		t.Fatal(err)
	}

	return resp.Images
}

func waitForSnapshotToBeAvailable(instanceId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Waiting for snapshot for instance %s to become available", instanceId)

	instanceIdTagFilter := &ec2.Filter{
		Name: aws.String(fmt.Sprintf("tag:%s", EC2_SNAPPER_INSTANCE_ID_TAG)),
		Values: []*string{aws.String(instanceId)},
	}

	if err := svc.WaitUntilImageAvailable(&ec2.DescribeImagesInput{Filters: []*ec2.Filter{instanceIdTagFilter}}); err != nil {
		t.Fatal(err)
	}
}

func waitForSnapshotToBeDeleted(snapshotId string, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	logger.Printf("Waiting for snapshot %s to be deleted", snapshotId)

	// We just do a simple sleep, as there is no built-in API call to wait for this.
	time.Sleep(30 * time.Second)
}

func deleteSnapshotForInstance(instanceName string, olderThan string, requireAtLeast int, ui cli.Ui, logger *log.Logger, t *testing.T) {
	logger.Printf("Deleting snapshot for instance %s", instanceName)

	deleteCmd := DeleteCommand{
		Ui: ui,
		AwsRegion: AWS_REGION_FOR_TESTING,
		InstanceName: instanceName,
		OlderThan: olderThan,
		RequireAtLeast: requireAtLeast,
	}

	if err := deleteSnapshots(deleteCmd); err != nil {
		t.Fatal(err)
	}
}

func takeSnapshotWithVerification(instanceName string, instanceId string, ui cli.Ui, svc *ec2.EC2, logger *log.Logger, t *testing.T) string {
	snapshotId := takeSnapshot(instanceName, ui, logger, t)

	waitForSnapshotToBeAvailable(instanceId, svc, logger, t)
	verifySnapshotWorks(snapshotId, svc, logger, t)

	return snapshotId
}

func deleteSnapshotWithVerification(instanceName string, snapshotId string, ui cli.Ui, svc *ec2.EC2, logger *log.Logger, t *testing.T) {
	deleteSnapshotForInstance(instanceName, "0h", 0, ui, logger, t)
	waitForSnapshotToBeDeleted(snapshotId, svc, logger, t)
	verifySnapshotIsDeleted(snapshotId, svc, logger, t)
}

func createLoggerAndUi(testName string) (*log.Logger, cli.Ui) {
	logger := log.New(os.Stdout, testName + " ", log.LstdFlags)

	basicUi := &cli.BasicUi{
		Reader:      os.Stdin,
		Writer:      os.Stdout,
		ErrorWriter: os.Stderr,
	}

	prefixedUi := &cli.PrefixedUi{
		AskPrefix:		logger.Prefix(),
		AskSecretPrefix:	logger.Prefix(),
		OutputPrefix:		logger.Prefix(),
		InfoPrefix:		logger.Prefix(),
		ErrorPrefix:		logger.Prefix(),
		WarnPrefix:		logger.Prefix(),
		Ui:			basicUi,

	}

	return logger, prefixedUi
}<|endoftext|>


<commit_before>package lfs

import (
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

const (
	batchSize = 100
)

type Transferable interface {
	Check() (*objectResource, error)
	Transfer(CopyCallback) error
	Object() *objectResource
	Oid() string
	Size() int64
	Name() string
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	retrying      uint32
	meter         *ProgressMeter
	workers       int // Number of transfer workers to spawn
	transferKind  string
	errors        []error
	transferables map[string]Transferable
	retries       []Transferable
	batcher       *Batcher
	apic          chan Transferable // Channel for processing individual API requests
	transferc     chan Transferable // Channel for processing transfers
	retriesc      chan Transferable // Channel for processing retries
	errorc        chan error        // Channel for processing errors
	watchers      []chan string
	errorwait     sync.WaitGroup
	retrywait     sync.WaitGroup
	wait          sync.WaitGroup
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(files int, size int64, dryRun bool) *TransferQueue {
	q := &TransferQueue{
		meter:         NewProgressMeter(files, size, dryRun),
		apic:          make(chan Transferable, batchSize),
		transferc:     make(chan Transferable, batchSize),
		retriesc:      make(chan Transferable, batchSize),
		errorc:        make(chan error),
		workers:       Config.ConcurrentTransfers(),
		transferables: make(map[string]Transferable),
	}

	q.errorwait.Add(1)
	q.retrywait.Add(1)

	q.run()

	return q
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	q.wait.Add(1)
	q.transferables[t.Oid()] = t

	if q.batcher != nil {
		q.batcher.Add(t)
		return
	}

	q.apic <- t
}

// Wait waits for the queue to finish processing all transfers. Once Wait is
// called, Add will no longer add transferables to the queue. Any failed
// transfers will be automatically retried once.
func (q *TransferQueue) Wait() {
	if q.batcher != nil {
		q.batcher.Exit()
	}

	q.wait.Wait()

	// Handle any retries
	close(q.retriesc)
	q.retrywait.Wait()
	atomic.StoreUint32(&q.retrying, 1)

	if len(q.retries) > 0 && q.batcher != nil {
		tracerx.Printf("tq: retrying %d failed transfers", len(q.retries))
		for _, t := range q.retries {
			q.Add(t)
		}
		q.batcher.Exit()
		q.wait.Wait()
	}

	atomic.StoreUint32(&q.retrying, 0)

	close(q.apic)
	close(q.transferc)
	close(q.errorc)

	for _, watcher := range q.watchers {
		close(watcher)
	}

	q.meter.Finish()
	q.errorwait.Wait()
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, batchSize)
	q.watchers = append(q.watchers, c)
	return c
}

// individualApiRoutine processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) individualApiRoutine(apiWaiter chan interface{}) {
	for t := range q.apic {
		obj, err := t.Check()
		if err != nil {
			if q.canRetry(err) {
				q.retry(t)
			} else {
				q.errorc <- err
			}
			q.wait.Done()
			continue
		}

		if apiWaiter != nil { // Signal to launch more individual api workers
			q.meter.Start()
			select {
			case apiWaiter <- 1:
			default:
			}
		}

		if obj != nil {
			t.SetObject(obj)
			q.meter.Add(t.Name())
			q.transferc <- t
		} else {
			q.meter.Skip(t.Size())
			q.wait.Done()
		}
	}
}

// legacyFallback is used when a batch request is made to a server that does
// not support the batch endpoint. When this happens, the Transferables are
// fed from the batcher into apic to be processed individually.
func (q *TransferQueue) legacyFallback(failedBatch []Transferable) {
	tracerx.Printf("tq: batch api not implemented, falling back to individual")

	q.launchIndividualApiRoutines()

	for _, t := range failedBatch {
		q.apic <- t
	}

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		for _, t := range batch {
			q.apic <- t
		}
	}
}

// batchApiRoutine processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) batchApiRoutine() {
	var startProgress sync.Once

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		tracerx.Printf("tq: sending batch of size %d", len(batch))

		transfers := make([]*objectResource, 0, len(batch))
		for _, t := range batch {
			transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
		}

		objects, err := Batch(transfers, q.transferKind)
		if err != nil {
			if IsNotImplementedError(err) {
				git.Config.SetLocal("", "lfs.batch", "false")

				go q.legacyFallback(batch)
				return
			}

			if q.canRetry(err) {
				for _, t := range batch {
					q.retry(t)
				}
			} else {
				q.errorc <- err
			}

			q.wait.Add(-len(transfers))
			continue
		}

		startProgress.Do(q.meter.Start)

		for _, o := range objects {
			if o.Error != nil {
				q.errorc <- Error(o.Error)
				q.meter.Skip(o.Size)
				q.wait.Done()
				continue
			}

			if _, ok := o.Rel(q.transferKind); ok {
				// This object needs to be transferred
				if transfer, ok := q.transferables[o.Oid]; ok {
					transfer.SetObject(o)
					q.meter.Add(transfer.Name())
					q.transferc <- transfer
				} else {
					q.meter.Skip(transfer.Size())
					q.wait.Done()
				}
			} else {
				q.meter.Skip(o.Size)
				q.wait.Done()
			}
		}
	}
}

// This goroutine collects errors returned from transfers
func (q *TransferQueue) errorCollector() {
	for err := range q.errorc {
		q.errors = append(q.errors, err)
	}
	q.errorwait.Done()
}

func (q *TransferQueue) retryCollector() {
	for t := range q.retriesc {
		q.retries = append(q.retries, t)
	}
	q.retrywait.Done()
}

func (q *TransferQueue) transferWorker() {
	for transfer := range q.transferc {
		cb := func(total, read int64, current int) error {
			q.meter.TransferBytes(q.transferKind, transfer.Name(), read, total, current)
			return nil
		}

		if err := transfer.Transfer(cb); err != nil {
			if q.canRetry(err) {
				tracerx.Printf("tq: retrying object %s", transfer.Oid())
				q.retry(transfer)
			} else {
				q.errorc <- err
			}
		} else {
			oid := transfer.Oid()
			for _, c := range q.watchers {
				c <- oid
			}
		}

		q.meter.FinishTransfer(transfer.Name())

		q.wait.Done()
	}
}

// launchIndividualApiRoutines first launches a single api worker. When it
// receives the first successful api request it launches workers - 1 more
// workers. This prevents being prompted for credentials multiple times at once
// when they're needed.
func (q *TransferQueue) launchIndividualApiRoutines() {
	go func() {
		apiWaiter := make(chan interface{})
		go q.individualApiRoutine(apiWaiter)

		<-apiWaiter

		for i := 0; i < q.workers-1; i++ {
			go q.individualApiRoutine(nil)
		}
	}()
}

// run starts the transfer queue, doing individual or batch transfers depending
// on the Config.BatchTransfer() value. run will transfer files sequentially or
// concurrently depending on the Config.ConcurrentTransfers() value.
func (q *TransferQueue) run() {
	go q.errorCollector()
	go q.retryCollector()

	tracerx.Printf("tq: starting %d transfer workers", q.workers)
	for i := 0; i < q.workers; i++ {
		go q.transferWorker()
	}

	if Config.BatchTransfer() {
		tracerx.Printf("tq: running as batched queue, batch size of %d", batchSize)
		q.batcher = NewBatcher(batchSize)
		go q.batchApiRoutine()
	} else {
		tracerx.Printf("tq: running as individual queue")
		q.launchIndividualApiRoutines()
	}
}

func (q *TransferQueue) retry(t Transferable) {
	q.retriesc <- t
}

func (q *TransferQueue) canRetry(err error) bool {
	if !IsRetriableError(err) || atomic.LoadUint32(&q.retrying) == 1 {
		return false
	}

	return true
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []error {
	return q.errors
}
<commit_msg>アアー アアアア アー<commit_after>package lfs

import (
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

const (
	batchSize = 100
)

type Transferable interface {
	Check() (*objectResource, error)
	Transfer(CopyCallback) error
	Object() *objectResource
	Oid() string
	Size() int64
	Name() string
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	retrying      uint32
	meter         *ProgressMeter
	workers       int // Number of transfer workers to spawn
	transferKind  string
	errors        []error
	transferables map[string]Transferable
	retries       []Transferable
	batcher       *Batcher
	apic          chan Transferable // Channel for processing individual API requests
	transferc     chan Transferable // Channel for processing transfers
	retriesc      chan Transferable // Channel for processing retries
	errorc        chan error        // Channel for processing errors
	watchers      []chan string
	errorwait     sync.WaitGroup
	retrywait     sync.WaitGroup
	wait          sync.WaitGroup
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(files int, size int64, dryRun bool) *TransferQueue {
	q := &TransferQueue{
		meter:         NewProgressMeter(files, size, dryRun),
		apic:          make(chan Transferable, batchSize),
		transferc:     make(chan Transferable, batchSize),
		retriesc:      make(chan Transferable, batchSize),
		errorc:        make(chan error),
		workers:       Config.ConcurrentTransfers(),
		transferables: make(map[string]Transferable),
	}

	q.errorwait.Add(1)
	q.retrywait.Add(1)

	q.run()

	return q
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	q.wait.Add(1)
	q.transferables[t.Oid()] = t

	if q.batcher != nil {
		q.batcher.Add(t)
		return
	}

	q.apic <- t
}

// Wait waits for the queue to finish processing all transfers. Once Wait is
// called, Add will no longer add transferables to the queue. Any failed
// transfers will be automatically retried once.
func (q *TransferQueue) Wait() {
	if q.batcher != nil {
		q.batcher.Exit()
	}

	q.wait.Wait()

	// Handle any retries
	close(q.retriesc)
	q.retrywait.Wait()
	atomic.StoreUint32(&q.retrying, 1)

	if len(q.retries) > 0 {
		tracerx.Printf("tq: retrying %d failed transfers", len(q.retries))
		for _, t := range q.retries {
			q.Add(t)
		}
		if q.batcher != nil {
			q.batcher.Exit()
		}
		q.wait.Wait()
	}

	atomic.StoreUint32(&q.retrying, 0)

	close(q.apic)
	close(q.transferc)
	close(q.errorc)

	for _, watcher := range q.watchers {
		close(watcher)
	}

	q.meter.Finish()
	q.errorwait.Wait()
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, batchSize)
	q.watchers = append(q.watchers, c)
	return c
}

// individualApiRoutine processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) individualApiRoutine(apiWaiter chan interface{}) {
	for t := range q.apic {
		obj, err := t.Check()
		if err != nil {
			if q.canRetry(err) {
				q.retry(t)
			} else {
				q.errorc <- err
			}
			q.wait.Done()
			continue
		}

		if apiWaiter != nil { // Signal to launch more individual api workers
			q.meter.Start()
			select {
			case apiWaiter <- 1:
			default:
			}
		}

		if obj != nil {
			t.SetObject(obj)
			q.meter.Add(t.Name())
			q.transferc <- t
		} else {
			q.meter.Skip(t.Size())
			q.wait.Done()
		}
	}
}

// legacyFallback is used when a batch request is made to a server that does
// not support the batch endpoint. When this happens, the Transferables are
// fed from the batcher into apic to be processed individually.
func (q *TransferQueue) legacyFallback(failedBatch []Transferable) {
	tracerx.Printf("tq: batch api not implemented, falling back to individual")

	q.launchIndividualApiRoutines()

	for _, t := range failedBatch {
		q.apic <- t
	}

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		for _, t := range batch {
			q.apic <- t
		}
	}
}

// batchApiRoutine processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) batchApiRoutine() {
	var startProgress sync.Once

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		tracerx.Printf("tq: sending batch of size %d", len(batch))

		transfers := make([]*objectResource, 0, len(batch))
		for _, t := range batch {
			transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
		}

		objects, err := Batch(transfers, q.transferKind)
		if err != nil {
			if IsNotImplementedError(err) {
				git.Config.SetLocal("", "lfs.batch", "false")

				go q.legacyFallback(batch)
				return
			}

			if q.canRetry(err) {
				for _, t := range batch {
					q.retry(t)
				}
			} else {
				q.errorc <- err
			}

			q.wait.Add(-len(transfers))
			continue
		}

		startProgress.Do(q.meter.Start)

		for _, o := range objects {
			if o.Error != nil {
				q.errorc <- Error(o.Error)
				q.meter.Skip(o.Size)
				q.wait.Done()
				continue
			}

			if _, ok := o.Rel(q.transferKind); ok {
				// This object needs to be transferred
				if transfer, ok := q.transferables[o.Oid]; ok {
					transfer.SetObject(o)
					q.meter.Add(transfer.Name())
					q.transferc <- transfer
				} else {
					q.meter.Skip(transfer.Size())
					q.wait.Done()
				}
			} else {
				q.meter.Skip(o.Size)
				q.wait.Done()
			}
		}
	}
}

// This goroutine collects errors returned from transfers
func (q *TransferQueue) errorCollector() {
	for err := range q.errorc {
		q.errors = append(q.errors, err)
	}
	q.errorwait.Done()
}

func (q *TransferQueue) retryCollector() {
	for t := range q.retriesc {
		q.retries = append(q.retries, t)
	}
	q.retrywait.Done()
}

func (q *TransferQueue) transferWorker() {
	for transfer := range q.transferc {
		cb := func(total, read int64, current int) error {
			q.meter.TransferBytes(q.transferKind, transfer.Name(), read, total, current)
			return nil
		}

		if err := transfer.Transfer(cb); err != nil {
			if q.canRetry(err) {
				tracerx.Printf("tq: retrying object %s", transfer.Oid())
				q.retry(transfer)
			} else {
				q.errorc <- err
			}
		} else {
			oid := transfer.Oid()
			for _, c := range q.watchers {
				c <- oid
			}
		}

		q.meter.FinishTransfer(transfer.Name())

		q.wait.Done()
	}
}

// launchIndividualApiRoutines first launches a single api worker. When it
// receives the first successful api request it launches workers - 1 more
// workers. This prevents being prompted for credentials multiple times at once
// when they're needed.
func (q *TransferQueue) launchIndividualApiRoutines() {
	go func() {
		apiWaiter := make(chan interface{})
		go q.individualApiRoutine(apiWaiter)

		<-apiWaiter

		for i := 0; i < q.workers-1; i++ {
			go q.individualApiRoutine(nil)
		}
	}()
}

// run starts the transfer queue, doing individual or batch transfers depending
// on the Config.BatchTransfer() value. run will transfer files sequentially or
// concurrently depending on the Config.ConcurrentTransfers() value.
func (q *TransferQueue) run() {
	go q.errorCollector()
	go q.retryCollector()

	tracerx.Printf("tq: starting %d transfer workers", q.workers)
	for i := 0; i < q.workers; i++ {
		go q.transferWorker()
	}

	if Config.BatchTransfer() {
		tracerx.Printf("tq: running as batched queue, batch size of %d", batchSize)
		q.batcher = NewBatcher(batchSize)
		go q.batchApiRoutine()
	} else {
		tracerx.Printf("tq: running as individual queue")
		q.launchIndividualApiRoutines()
	}
}

func (q *TransferQueue) retry(t Transferable) {
	q.retriesc <- t
}

func (q *TransferQueue) canRetry(err error) bool {
	if !IsRetriableError(err) || atomic.LoadUint32(&q.retrying) == 1 {
		return false
	}

	return true
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []error {
	return q.errors
}
<|endoftext|>