<commit_before>package main

import (
	"html/template"
	"log"
	"net/http"

	"github.com/sterops/tinkerings/dispotrains.webapp/storage"

	"github.com/eknkc/dateformat"
	"github.com/gorilla/mux"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

var (
	homeTmpl    = template.Must(template.ParseFiles("templates/lines.html"))
	lineTmpl    = template.Must(template.ParseFiles("templates/line.html"))
	stationTmpl = template.Must(template.ParseFiles("templates/station.html"))
)

type Line struct {
	Network      string
	ID           string
	GoodStations []*storage.Station `bson:"goodStations"`
	BadStations  []*storage.Station `bson:"badStations"`
}

type LineHolder struct {
	Value Line
}

type LineSlice []LineHolder

type DisplayStation struct {
	storage.Station
	Elevators []*LocElevator
}

type LocElevator storage.Elevator

func (e *LocElevator) LocalStatusDate() string {
	return dateformat.FormatLocale(e.Status.LastUpdate, "ddd D MMM à HH:MM", dateformat.French)
}

func (ls LineSlice) Lines() []Line {
	r := make([]Line, len(ls))
	for i, v := range ls {
		r[i] = v.Value
	}
	return r
}

func HomeHandler(w http.ResponseWriter, req *http.Request) {
	session, err := mgo.Dial("localhost")
	if err != nil {
		panic(err)
	}
	defer session.Close()
	c := session.DB("dispotrains").C("lines")
	var lines LineSlice = make(LineSlice, 0)
	c.Find(nil).Sort("value.network", "value.id").All(&lines)
	homeTmpl.Execute(w, lines.Lines())
}

func LineHandler(w http.ResponseWriter, req *http.Request) {
	session, err := mgo.Dial("localhost")
	if err != nil {
		panic(err)
	}
	defer session.Close()
	c := session.DB("dispotrains").C("lines")

	vars := mux.Vars(req)
	lineId := vars["line"]

	var line LineHolder
	c.Find(bson.M{"_id": lineId}).One(&line)
	if err = lineTmpl.Execute(w, line.Value); err != nil {
		log.Fatal(err)
	}
}

func StationHandler(w http.ResponseWriter, req *http.Request) {
	session, err := mgo.Dial("localhost")
	if err != nil {
		panic(err)
	}
	defer session.Close()
	c := session.DB("dispotrains").C("stations")

	vars := mux.Vars(req)
	stationName := vars["station"]

	var station DisplayStation
	c.Find(bson.M{"name": stationName}).One(&station)
	if err = stationTmpl.Execute(w, station); err != nil {
		log.Fatal(err)
	}
}

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	r.HandleFunc("/ligne/{line}", LineHandler)
	r.HandleFunc("/gare/{station}", StationHandler)
	r.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))
	http.Handle("/", r)
	log.Fatal(http.ListenAndServe(":9000", nil))
}
<commit_msg>Fix display name on station page<commit_after>package main

import (
	"html/template"
	"log"
	"net/http"

	"github.com/sterops/tinkerings/dispotrains.webapp/storage"

	"github.com/eknkc/dateformat"
	"github.com/gorilla/mux"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

var (
	homeTmpl    = template.Must(template.ParseFiles("templates/lines.html"))
	lineTmpl    = template.Must(template.ParseFiles("templates/line.html"))
	stationTmpl = template.Must(template.ParseFiles("templates/station.html"))
)

type Line struct {
	Network      string
	ID           string
	GoodStations []*storage.Station `bson:"goodStations"`
	BadStations  []*storage.Station `bson:"badStations"`
}

type LineHolder struct {
	Value Line
}

type LineSlice []LineHolder

type DisplayStation struct {
	storage.Station
    DisplayName string
	Elevators []*LocElevator
}

type LocElevator storage.Elevator

func (e *LocElevator) LocalStatusDate() string {
	return dateformat.FormatLocale(e.Status.LastUpdate, "ddd D MMM à HH:MM", dateformat.French)
}

func (ls LineSlice) Lines() []Line {
	r := make([]Line, len(ls))
	for i, v := range ls {
		r[i] = v.Value
	}
	return r
}

func HomeHandler(w http.ResponseWriter, req *http.Request) {
	session, err := mgo.Dial("localhost")
	if err != nil {
		panic(err)
	}
	defer session.Close()
	c := session.DB("dispotrains").C("lines")
	var lines LineSlice = make(LineSlice, 0)
	c.Find(nil).Sort("value.network", "value.id").All(&lines)
	homeTmpl.Execute(w, lines.Lines())
}

func LineHandler(w http.ResponseWriter, req *http.Request) {
	session, err := mgo.Dial("localhost")
	if err != nil {
		panic(err)
	}
	defer session.Close()
	c := session.DB("dispotrains").C("lines")

	vars := mux.Vars(req)
	lineId := vars["line"]

	var line LineHolder
	c.Find(bson.M{"_id": lineId}).One(&line)
	if err = lineTmpl.Execute(w, line.Value); err != nil {
		log.Fatal(err)
	}
}

func StationHandler(w http.ResponseWriter, req *http.Request) {
	session, err := mgo.Dial("localhost")
	if err != nil {
		panic(err)
	}
	defer session.Close()
	c := session.DB("dispotrains").C("stations")

	vars := mux.Vars(req)
	stationName := vars["station"]

	var station DisplayStation
	c.Find(bson.M{"name": stationName}).One(&station)
	if err = stationTmpl.Execute(w, station); err != nil {
		log.Fatal(err)
	}
}

func CacheRequest(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.Header().Set("Cache-control", "public, max-age=259200")
        h.ServeHTTP(w, r)
	})
}

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	r.HandleFunc("/ligne/{line}", LineHandler)
	r.HandleFunc("/gare/{station}", StationHandler)
	r.PathPrefix("/static/").Handler(CacheRequest(http.StripPrefix("/static/", http.FileServer(http.Dir("static")))))
	http.Handle("/", r)
	log.Fatal(http.ListenAndServe(":9000", nil))
}
<|endoftext|>


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

package persistence

// TODO(ericsnow) Eliminate the mongo-related imports here.

import (
	"fmt"
	"reflect"

	"github.com/juju/errors"
	"github.com/juju/names"
	jujutxn "github.com/juju/txn"
	"gopkg.in/juju/charm.v5"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/process"
)

// TODO(ericsnow) Implement persistence using a TXN abstraction (used
// in the business logic) with ops factories available from the
// persistence layer.

// PersistenceBase exposes the core persistence functionality needed
// for workload processes.
type PersistenceBase interface {
	// One populates doc with the document corresponding to the given
	// ID. Missing documents result in errors.NotFound.
	One(collName, id string, doc interface{}) error
	// All populates docs with the list of the documents corresponding
	// to the provided query.
	All(collName string, query, docs interface{}) error
	// Run runs the transaction generated by the provided factory
	// function. It may be retried several times.
	Run(transactions jujutxn.TransactionSource) error
}

// Persistence exposes the high-level persistence functionality
// related to workload processes in Juju.
type Persistence struct {
	st    PersistenceBase
	charm names.CharmTag
	unit  names.UnitTag
}

// NewPersistence builds a new Persistence based on the provided info.
func NewPersistence(st PersistenceBase, charm *names.CharmTag, unit *names.UnitTag) *Persistence {
	pp := &Persistence{
		st: st,
	}
	if charm != nil {
		pp.charm = *charm
	}
	if unit != nil {
		pp.unit = *unit
	}
	return pp
}

// EnsureDefinitions checks persistence to see if records for the
// definitions are already there. If not then they are added. If so then
// they are checked to be sure they match those provided. The lists of
// names for those that already exist and that don't match are returned.
func (pp Persistence) EnsureDefinitions(definitions ...charm.Process) ([]string, []string, error) {
	var found []string
	var mismatched []string

	if len(definitions) == 0 {
		return found, mismatched, nil
	}

	var ids []string
	var ops []txn.Op
	for _, definition := range definitions {
		ids = append(ids, pp.definitionID(definition.Name))
		ops = append(ops, pp.newInsertDefinitionOp(definition))
	}
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			// The last attempt aborted so clear out any ops that failed
			// the DocMissing assertion and try again.
			indexed, err := pp.indexDefinitionDocs(ids)
			if err != nil {
				return nil, errors.Trace(err)
			}

			var okOps []txn.Op
			for _, op := range ops {
				if existing, ok := indexed[op.Id]; !ok {
					okOps = append(okOps, op)
				} else { // Otherwise the op is dropped.
					id := fmt.Sprintf("%s", op.Id)
					found = append(found, id)
					definition, ok := op.Insert.(*ProcessDefinitionDoc)
					if !ok {
						return nil, errors.Errorf("inserting invalid type %T", op.Insert)
					}
					if !reflect.DeepEqual(definition, &existing) {
						mismatched = append(mismatched, id)
					}
				}
			}
			if len(okOps) == 0 {
				return nil, jujutxn.ErrNoOperations
			}
			ops = okOps
		}
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return nil, nil, errors.Trace(err)
	}

	return found, mismatched, nil
}

// Insert adds records for the process to persistence. If the process
// is already there then false gets returned (true if inserted).
// Existing records are not checked for consistency.
func (pp Persistence) Insert(info process.Info) (bool, error) {
	var okay bool
	var ops []txn.Op
	// TODO(ericsnow) Add unitPersistence.newEnsureAliveOp(pp.unit)?
	// TODO(ericsnow) Add pp.newEnsureDefinitionOp(info.Process)?
	ops = append(ops, pp.newInsertProcessOps(info)...)
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			// One of the records already exists.
			okay = false
			return nil, jujutxn.ErrNoOperations
		}
		okay = true
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return false, errors.Trace(err)
	}
	return okay, nil
}

// SetStatus updates the raw status for the identified process in
// persistence. The return value corresponds to whether or not the
// record was found in persistence. Any other problem results in
// an error. The process is not checked for inconsistent records.
func (pp Persistence) SetStatus(id string, status process.Status) (bool, error) {
	var found bool
	var ops []txn.Op
	// TODO(ericsnow) Add unitPersistence.newEnsureAliveOp(pp.unit)?
	ops = append(ops, pp.newSetRawStatusOps(id, status)...)
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			_, err := pp.proc(id)
			if errors.IsNotFound(err) {
				found = false
				return nil, jujutxn.ErrNoOperations
			} else if err != nil {
				return nil, errors.Trace(err)
			}
			// We ignore the request since the proc is dying.
			// TODO(ericsnow) Ensure that procDoc.Status != state.Alive?
			return nil, jujutxn.ErrNoOperations
		}
		found = true
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return false, errors.Trace(err)
	}
	return found, nil
}

// List builds the list of processes found in persistence which match
// the provided IDs. The lists of IDs with missing records is also
// returned. Inconsistent records result in errors.NotValid.
func (pp Persistence) List(ids ...string) ([]process.Info, []string, error) {
	var missing []string

	// TODO(ericsnow) Ensure that the unit is Alive?
	// TODO(ericsnow) fix race that exists between the 3 calls
	definitionDocs, err := pp.definitions(ids)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	launchDocs, err := pp.launches(ids)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	procDocs, err := pp.procs(ids)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	var results []process.Info
	for _, id := range ids {
		proc, missingCount := pp.extractProc(id, definitionDocs, launchDocs, procDocs)
		if missingCount > 0 {
			if missingCount < 7 {
				return nil, nil, errors.Errorf("found inconsistent records for process %q", id)
			}
			missing = append(missing, id)
			continue
		}
		results = append(results, *proc)
	}
	return results, missing, nil
}

// ListAll builds the list of all processes found in persistence.
// Inconsistent records result in errors.NotValid.
func (pp Persistence) ListAll() ([]process.Info, error) {
	// TODO(ericsnow) Ensure that the unit is Alive?
	// TODO(ericsnow) fix race that exists between the 3 calls
	definitionDocs, err := pp.allDefinitions()
	if err != nil {
		return nil, errors.Trace(err)
	}
	launchDocs, err := pp.allLaunches()
	if err != nil {
		return nil, errors.Trace(err)
	}
	procDocs, err := pp.allProcs()
	if err != nil {
		return nil, errors.Trace(err)
	}

	if len(launchDocs) > len(procDocs) {
		return nil, errors.Errorf("found inconsistent records (extra launch docs)")
	}

	var results []process.Info
	for id := range procDocs {
		proc, missingCount := pp.extractProc(id, definitionDocs, launchDocs, procDocs)
		if missingCount > 0 {
			return nil, errors.Errorf("found inconsistent records for process %q", id)
		}
		results = append(results, *proc)
	}
	for name, doc := range definitionDocs {
		matched := false
		for _, proc := range results {
			if name == proc.Name {
				matched = true
				break
			}
		}
		if !matched {
			results = append(results, process.Info{
				Process: doc.definition(),
			})
		}
	}
	return results, nil
}

// TODO(ericsnow) Add procs to state/cleanup.go.

// TODO(ericsnow) How to ensure they are completely removed from state?

// Remove removes all records associated with the identified process
// from persistence. Also returned is whether or not the process was
// found. If the records for the process are not consistent then
// errors.NotValid is returned.
func (pp Persistence) Remove(id string) (bool, error) {
	var found bool
	var ops []txn.Op
	// TODO(ericsnow) Add unitPersistence.newEnsureAliveOp(pp.unit)?
	ops = append(ops, pp.newRemoveProcessOps(id)...)
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			okay, err := pp.checkRecords(id)
			if err != nil {
				return nil, errors.Trace(err)
			}
			// If okay is true, it must be dying.
			found = okay
			return nil, jujutxn.ErrNoOperations
		}
		found = true
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return false, errors.Trace(err)
	}
	return found, nil
}
<commit_msg>Set found correctly.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package persistence

// TODO(ericsnow) Eliminate the mongo-related imports here.

import (
	"fmt"
	"reflect"

	"github.com/juju/errors"
	"github.com/juju/names"
	jujutxn "github.com/juju/txn"
	"gopkg.in/juju/charm.v5"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/process"
)

// TODO(ericsnow) Implement persistence using a TXN abstraction (used
// in the business logic) with ops factories available from the
// persistence layer.

// PersistenceBase exposes the core persistence functionality needed
// for workload processes.
type PersistenceBase interface {
	// One populates doc with the document corresponding to the given
	// ID. Missing documents result in errors.NotFound.
	One(collName, id string, doc interface{}) error
	// All populates docs with the list of the documents corresponding
	// to the provided query.
	All(collName string, query, docs interface{}) error
	// Run runs the transaction generated by the provided factory
	// function. It may be retried several times.
	Run(transactions jujutxn.TransactionSource) error
}

// Persistence exposes the high-level persistence functionality
// related to workload processes in Juju.
type Persistence struct {
	st    PersistenceBase
	charm names.CharmTag
	unit  names.UnitTag
}

// NewPersistence builds a new Persistence based on the provided info.
func NewPersistence(st PersistenceBase, charm *names.CharmTag, unit *names.UnitTag) *Persistence {
	pp := &Persistence{
		st: st,
	}
	if charm != nil {
		pp.charm = *charm
	}
	if unit != nil {
		pp.unit = *unit
	}
	return pp
}

// EnsureDefinitions checks persistence to see if records for the
// definitions are already there. If not then they are added. If so then
// they are checked to be sure they match those provided. The lists of
// names for those that already exist and that don't match are returned.
func (pp Persistence) EnsureDefinitions(definitions ...charm.Process) ([]string, []string, error) {
	var found []string
	var mismatched []string

	if len(definitions) == 0 {
		return found, mismatched, nil
	}

	var ids []string
	var ops []txn.Op
	for _, definition := range definitions {
		ids = append(ids, pp.definitionID(definition.Name))
		ops = append(ops, pp.newInsertDefinitionOp(definition))
	}
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			// The last attempt aborted so clear out any ops that failed
			// the DocMissing assertion and try again.
			found = []string{}
			mismatched = []string{}
			indexed, err := pp.indexDefinitionDocs(ids)
			if err != nil {
				return nil, errors.Trace(err)
			}

			var okOps []txn.Op
			for _, op := range ops {
				if existing, ok := indexed[op.Id]; !ok {
					okOps = append(okOps, op)
				} else { // Otherwise the op is dropped.
					id := fmt.Sprintf("%s", op.Id)
					found = append(found, id)
					definition, ok := op.Insert.(*ProcessDefinitionDoc)
					if !ok {
						return nil, errors.Errorf("inserting invalid type %T", op.Insert)
					}
					if !reflect.DeepEqual(definition, &existing) {
						mismatched = append(mismatched, id)
					}
				}
			}
			if len(okOps) == 0 {
				return nil, jujutxn.ErrNoOperations
			}
			ops = okOps
		}
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return nil, nil, errors.Trace(err)
	}

	return found, mismatched, nil
}

// Insert adds records for the process to persistence. If the process
// is already there then false gets returned (true if inserted).
// Existing records are not checked for consistency.
func (pp Persistence) Insert(info process.Info) (bool, error) {
	var okay bool
	var ops []txn.Op
	// TODO(ericsnow) Add unitPersistence.newEnsureAliveOp(pp.unit)?
	// TODO(ericsnow) Add pp.newEnsureDefinitionOp(info.Process)?
	ops = append(ops, pp.newInsertProcessOps(info)...)
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			// One of the records already exists.
			okay = false
			return nil, jujutxn.ErrNoOperations
		}
		okay = true
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return false, errors.Trace(err)
	}
	return okay, nil
}

// SetStatus updates the raw status for the identified process in
// persistence. The return value corresponds to whether or not the
// record was found in persistence. Any other problem results in
// an error. The process is not checked for inconsistent records.
func (pp Persistence) SetStatus(id string, status process.Status) (bool, error) {
	var found bool
	var ops []txn.Op
	// TODO(ericsnow) Add unitPersistence.newEnsureAliveOp(pp.unit)?
	ops = append(ops, pp.newSetRawStatusOps(id, status)...)
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			_, err := pp.proc(id)
			if errors.IsNotFound(err) {
				found = false
				return nil, jujutxn.ErrNoOperations
			} else if err != nil {
				return nil, errors.Trace(err)
			}
			// We ignore the request since the proc is dying.
			// TODO(ericsnow) Ensure that procDoc.Status != state.Alive?
			return nil, jujutxn.ErrNoOperations
		}
		found = true
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return false, errors.Trace(err)
	}
	return found, nil
}

// List builds the list of processes found in persistence which match
// the provided IDs. The lists of IDs with missing records is also
// returned. Inconsistent records result in errors.NotValid.
func (pp Persistence) List(ids ...string) ([]process.Info, []string, error) {
	var missing []string

	// TODO(ericsnow) Ensure that the unit is Alive?
	// TODO(ericsnow) fix race that exists between the 3 calls
	definitionDocs, err := pp.definitions(ids)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	launchDocs, err := pp.launches(ids)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	procDocs, err := pp.procs(ids)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	var results []process.Info
	for _, id := range ids {
		proc, missingCount := pp.extractProc(id, definitionDocs, launchDocs, procDocs)
		if missingCount > 0 {
			if missingCount < 7 {
				return nil, nil, errors.Errorf("found inconsistent records for process %q", id)
			}
			missing = append(missing, id)
			continue
		}
		results = append(results, *proc)
	}
	return results, missing, nil
}

// ListAll builds the list of all processes found in persistence.
// Inconsistent records result in errors.NotValid.
func (pp Persistence) ListAll() ([]process.Info, error) {
	// TODO(ericsnow) Ensure that the unit is Alive?
	// TODO(ericsnow) fix race that exists between the 3 calls
	definitionDocs, err := pp.allDefinitions()
	if err != nil {
		return nil, errors.Trace(err)
	}
	launchDocs, err := pp.allLaunches()
	if err != nil {
		return nil, errors.Trace(err)
	}
	procDocs, err := pp.allProcs()
	if err != nil {
		return nil, errors.Trace(err)
	}

	if len(launchDocs) > len(procDocs) {
		return nil, errors.Errorf("found inconsistent records (extra launch docs)")
	}

	var results []process.Info
	for id := range procDocs {
		proc, missingCount := pp.extractProc(id, definitionDocs, launchDocs, procDocs)
		if missingCount > 0 {
			return nil, errors.Errorf("found inconsistent records for process %q", id)
		}
		results = append(results, *proc)
	}
	for name, doc := range definitionDocs {
		matched := false
		for _, proc := range results {
			if name == proc.Name {
				matched = true
				break
			}
		}
		if !matched {
			results = append(results, process.Info{
				Process: doc.definition(),
			})
		}
	}
	return results, nil
}

// TODO(ericsnow) Add procs to state/cleanup.go.

// TODO(ericsnow) How to ensure they are completely removed from state?

// Remove removes all records associated with the identified process
// from persistence. Also returned is whether or not the process was
// found. If the records for the process are not consistent then
// errors.NotValid is returned.
func (pp Persistence) Remove(id string) (bool, error) {
	var found bool
	var ops []txn.Op
	// TODO(ericsnow) Add unitPersistence.newEnsureAliveOp(pp.unit)?
	ops = append(ops, pp.newRemoveProcessOps(id)...)
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			okay, err := pp.checkRecords(id)
			if err != nil {
				return nil, errors.Trace(err)
			}
			// If okay is true, it must be dying.
			found = okay
			return nil, jujutxn.ErrNoOperations
		}
		found = true
		return ops, nil
	}
	if err := pp.st.Run(buildTxn); err != nil {
		return false, errors.Trace(err)
	}
	return found, nil
}
<|endoftext|>


<commit_before>package provider

import (
	"encoding/base64"
	"errors"
	"time"

	"github.com/NetSys/di/db"
	"github.com/NetSys/di/dsl"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"

	log "github.com/Sirupsen/logrus"
)

const spotPrice = "0.5"

// Ubuntu 15.10, 64-bit hvm-ssd
var amis = map[string]string{
	"us-west-2": "ami-acd63bcc",
	"us-west-1": "ami-af671bcf",
}

type awsSpotCluster struct {
	sessions map[string]*ec2.EC2

	namespace  string
	aclTrigger db.Trigger
}

type awsID struct {
	spotID string
	region string
}

func getSpotIDs(ids []awsID) []string {
	var spotIDs []string
	for _, id := range ids {
		spotIDs = append(spotIDs, id.spotID)
	}

	return spotIDs
}

func groupByRegion(ids []awsID) map[string][]awsID {
	grouped := make(map[string][]awsID)
	for _, id := range ids {
		region := id.region
		if _, ok := grouped[region]; !ok {
			grouped[region] = []awsID{}
		}
		grouped[region] = append(grouped[region], id)
	}

	return grouped
}

func (clst *awsSpotCluster) Start(conn db.Conn, clusterID int, namespace string) error {
	clst.sessions = make(map[string]*ec2.EC2)
	clst.namespace = namespace
	clst.aclTrigger = conn.TriggerTick(60, db.ClusterTable)

	go clst.watchACLs(conn, clusterID)

	return nil
}

func (clst *awsSpotCluster) Disconnect() {
	/* Ideally we'd close clst.ec2 as well, but the API doesn't export that ability
	* apparently. */
	clst.aclTrigger.Stop()
}

func (clst awsSpotCluster) getSession(region string) *ec2.EC2 {
	if _, ok := clst.sessions[region]; ok {
		return clst.sessions[region]
	}

	session := session.New()
	session.Config.Region = aws.String(region)

	newEC2 := ec2.New(session)
	clst.sessions[region] = newEC2

	return newEC2
}

func (clst awsSpotCluster) Boot(bootSet []Machine) error {
	if len(bootSet) <= 0 {
		return nil
	}

	type bootReq struct {
		cfg    string
		size   string
		region string
	}

	bootReqMap := make(map[bootReq]int64) // From boot request to an instance count.
	for _, m := range bootSet {
		br := bootReq{
			cfg:    cloudConfigUbuntu(m.SSHKeys, "wily"),
			size:   m.Size,
			region: m.Region,
		}
		bootReqMap[br] = bootReqMap[br] + 1
	}

	var awsIDs []awsID
	for br, count := range bootReqMap {
		session := clst.getSession(br.region)
		cloudConfig64 := base64.StdEncoding.EncodeToString([]byte(br.cfg))
		resp, err := session.RequestSpotInstances(&ec2.RequestSpotInstancesInput{
			SpotPrice: aws.String(spotPrice),
			LaunchSpecification: &ec2.RequestSpotLaunchSpecification{
				ImageId:        aws.String(amis[br.region]),
				InstanceType:   aws.String(br.size),
				UserData:       &cloudConfig64,
				SecurityGroups: []*string{&clst.namespace},
			},
			InstanceCount: &count,
		})

		if err != nil {
			return err
		}

		for _, request := range resp.SpotInstanceRequests {
			awsIDs = append(awsIDs, awsID{
				spotID: *request.SpotInstanceRequestId,
				region: br.region})
		}
	}

	if err := clst.tagSpotRequests(awsIDs); err != nil {
		return err
	}

	if err := clst.wait(awsIDs, true); err != nil {
		return err
	}

	return nil
}

func (clst awsSpotCluster) Stop(machines []Machine) error {
	var awsIDs []awsID
	for _, m := range machines {
		awsIDs = append(awsIDs, awsID{
			region: m.Region,
			spotID: m.ID,
		})
	}
	for region, ids := range groupByRegion(awsIDs) {
		session := clst.getSession(region)
		spotIDs := getSpotIDs(ids)

		spots, err := session.DescribeSpotInstanceRequests(
			&ec2.DescribeSpotInstanceRequestsInput{
				SpotInstanceRequestIds: aws.StringSlice(spotIDs),
			})
		if err != nil {
			return err
		}

		instIds := []string{}
		for _, spot := range spots.SpotInstanceRequests {
			if spot.InstanceId != nil {
				instIds = append(instIds, *spot.InstanceId)
			}
		}

		if len(instIds) > 0 {
			_, err = session.TerminateInstances(&ec2.TerminateInstancesInput{
				InstanceIds: aws.StringSlice(instIds),
			})
			if err != nil {
				return err
			}
		}

		_, err = session.CancelSpotInstanceRequests(&ec2.CancelSpotInstanceRequestsInput{
			SpotInstanceRequestIds: aws.StringSlice(spotIDs),
		})
		if err != nil {
			return err
		}

		if err := clst.wait(ids, false); err != nil {
			return err
		}
	}

	return nil
}

func (clst awsSpotCluster) Get() ([]Machine, error) {
	machines := []Machine{}
	for region := range amis {
		session := clst.getSession(region)

		spots, err := session.DescribeSpotInstanceRequests(nil)
		if err != nil {
			return nil, err
		}

		insts, err := session.DescribeInstances(&ec2.DescribeInstancesInput{
			Filters: []*ec2.Filter{
				{
					Name:   aws.String("instance.group-name"),
					Values: []*string{aws.String(clst.namespace)},
				},
			},
		})
		if err != nil {
			return nil, err
		}

		instMap := make(map[string]*ec2.Instance)
		for _, res := range insts.Reservations {
			for _, inst := range res.Instances {
				instMap[*inst.InstanceId] = inst
			}
		}

		for _, spot := range spots.SpotInstanceRequests {
			if *spot.State != ec2.SpotInstanceStateActive &&
				*spot.State != ec2.SpotInstanceStateOpen {
				continue
			}

			var inst *ec2.Instance
			if spot.InstanceId != nil {
				inst = instMap[*spot.InstanceId]
			}

			// Due to a race condition in the AWS API, it's possible that spot
			// requests might lose their Tags.  If handled naively, those spot
			// requests would technically be without a namespace, meaning the
			// instances they create would be live forever as zombies.
			//
			// To mitigate this issue, we rely not only on the spot request tags, but
			// additionally on the instance security group.  If a spot request has a
			// running instance in the appropriate security group, it is by
			// definition in our namespace.  Thus, we only check the tags for spot
			// requests without running instances.
			if inst == nil {
				var isOurs bool
				for _, tag := range spot.Tags {
					ns := clst.namespace
					if tag != nil && tag.Key != nil && *tag.Key == ns {
						isOurs = true
						break
					}
				}

				if !isOurs {
					continue
				}
			}

			machine := Machine{
				ID:       *spot.SpotInstanceRequestId,
				Region:   region,
				Provider: db.AmazonSpot,
			}

			if inst != nil {
				if *inst.State.Name != ec2.InstanceStateNamePending &&
					*inst.State.Name != ec2.InstanceStateNameRunning {
					continue
				}

				if inst.PublicIpAddress != nil {
					machine.PublicIP = *inst.PublicIpAddress
				}

				if inst.PrivateIpAddress != nil {
					machine.PrivateIP = *inst.PrivateIpAddress
				}

				if inst.InstanceType != nil {
					machine.Size = *inst.InstanceType
				}
			}

			machines = append(machines, machine)
		}
	}

	return machines, nil
}

func (clst *awsSpotCluster) PickBestSize(ram dsl.Range, cpu dsl.Range, maxPrice float64) string {
	return pickBestSize(awsDescriptions, ram, cpu, maxPrice)
}

func (clst *awsSpotCluster) tagSpotRequests(awsIDs []awsID) error {
OuterLoop:
	for region, ids := range groupByRegion(awsIDs) {
		session := clst.getSession(region)
		spotIDs := getSpotIDs(ids)

		var err error
		for i := 0; i < 30; i++ {
			_, err = session.CreateTags(&ec2.CreateTagsInput{
				Tags: []*ec2.Tag{
					{Key: aws.String(clst.namespace), Value: aws.String("")},
				},
				Resources: aws.StringSlice(spotIDs),
			})
			if err == nil {
				continue OuterLoop
			}
			time.Sleep(5 * time.Second)
		}

		log.Warn("Failed to tag spot requests: ", err)
		session.CancelSpotInstanceRequests(
			&ec2.CancelSpotInstanceRequestsInput{
				SpotInstanceRequestIds: aws.StringSlice(spotIDs),
			})

		return err
	}

	return nil
}

/* Wait for the spot request 'ids' to have booted or terminated depending on the value of
* 'boot' */
func (clst *awsSpotCluster) wait(awsIDs []awsID, boot bool) error {
OuterLoop:
	for i := 0; i < 100; i++ {
		machines, err := clst.Get()
		if err != nil {
			log.WithError(err).Warn("Failed to get machines.")
			time.Sleep(2 * time.Second)
			continue
		}

		exists := make(map[awsID]struct{})
		for _, inst := range machines {
			id := awsID{
				spotID: inst.ID,
				region: inst.Region,
			}

			exists[id] = struct{}{}
		}

		for _, id := range awsIDs {
			if _, ok := exists[id]; ok != boot {
				time.Sleep(2 * time.Second)
				continue OuterLoop
			}
		}

		return nil
	}

	return errors.New("timed out")
}

func (clst *awsSpotCluster) updateSecurityGroups(acls []string) error {
	for _, session := range clst.sessions {
		resp, err := session.DescribeSecurityGroups(
			&ec2.DescribeSecurityGroupsInput{
				Filters: []*ec2.Filter{
					{
						Name:   aws.String("group-name"),
						Values: []*string{aws.String(clst.namespace)},
					},
				},
			})

		if err != nil {
			return err
		}

		ingress := []*ec2.IpPermission{}
		groups := resp.SecurityGroups
		if len(groups) > 1 {
			return errors.New("Multiple Security Groups with the same name: " +
				clst.namespace)
		} else if len(groups) == 0 {
			session.CreateSecurityGroup(&ec2.CreateSecurityGroupInput{
				Description: aws.String("Declarative Infrastructure Group"),
				GroupName:   aws.String(clst.namespace),
			})
		} else {
			/* XXX: Deal with egress rules. */
			ingress = groups[0].IpPermissions
		}

		permMap := make(map[string]bool)
		for _, acl := range acls {
			permMap[acl] = true
		}

		groupIngressExists := false
		for i, p := range ingress {
			if (i > 0 || p.FromPort != nil || p.ToPort != nil ||
				*p.IpProtocol != "-1") && p.UserIdGroupPairs == nil {
				log.Info("Revoke ingress security group: ", *p)
				_, err = session.RevokeSecurityGroupIngress(
					&ec2.RevokeSecurityGroupIngressInput{
						GroupName:     aws.String(clst.namespace),
						IpPermissions: []*ec2.IpPermission{p}})
				if err != nil {
					return err
				}

				continue
			}

			for _, ipr := range p.IpRanges {
				ip := *ipr.CidrIp
				if !permMap[ip] {
					log.Info("Revoke ingress security group: ", ip)
					_, err = session.RevokeSecurityGroupIngress(
						&ec2.RevokeSecurityGroupIngressInput{
							GroupName:  aws.String(clst.namespace),
							CidrIp:     aws.String(ip),
							FromPort:   p.FromPort,
							IpProtocol: p.IpProtocol,
							ToPort:     p.ToPort})
					if err != nil {
						return err
					}
				} else {
					permMap[ip] = false
				}
			}

			if len(groups) > 0 {
				for _, grp := range p.UserIdGroupPairs {
					if *grp.GroupId != *groups[0].GroupId {
						log.Info("Revoke ingress security group GroupID: ",
							*grp.GroupId)
						_, err = session.RevokeSecurityGroupIngress(
							&ec2.RevokeSecurityGroupIngressInput{
								GroupName:               aws.String(clst.namespace),
								SourceSecurityGroupName: grp.GroupName})
						if err != nil {
							return err
						}
					} else {
						groupIngressExists = true
					}
				}
			}
		}

		if !groupIngressExists {
			log.Info("Add intragroup ACL")
			_, err = session.AuthorizeSecurityGroupIngress(
				&ec2.AuthorizeSecurityGroupIngressInput{
					GroupName:               aws.String(clst.namespace),
					SourceSecurityGroupName: aws.String(clst.namespace)})
		}

		for perm, install := range permMap {
			if !install {
				continue
			}

			log.Info("Add ACL: ", perm)
			_, err = session.AuthorizeSecurityGroupIngress(
				&ec2.AuthorizeSecurityGroupIngressInput{
					CidrIp:     aws.String(perm),
					GroupName:  aws.String(clst.namespace),
					IpProtocol: aws.String("-1")})

			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (clst *awsSpotCluster) watchACLs(conn db.Conn, clusterID int) {
	for range clst.aclTrigger.C {
		var acls []string
		conn.Transact(func(view db.Database) error {
			clusters := view.SelectFromCluster(func(c db.Cluster) bool {
				return c.ID == clusterID
			})

			if len(clusters) == 0 {
				log.Warn("Undefined cluster.")
				return nil
			} else if len(clusters) > 1 {
				panic("Duplicate Clusters")
			}

			acls = clusters[0].ACLs
			return nil
		})

		clst.updateSecurityGroups(acls)
	}
}
<commit_msg>awsSpot: Manually set EBS disk size<commit_after>package provider

import (
	"encoding/base64"
	"errors"
	"time"

	"github.com/NetSys/di/db"
	"github.com/NetSys/di/dsl"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"

	log "github.com/Sirupsen/logrus"
)

const spotPrice = "0.5"
const diskSize = 32

// Ubuntu 15.10, 64-bit hvm-ssd
var amis = map[string]string{
	"us-west-2": "ami-acd63bcc",
	"us-west-1": "ami-af671bcf",
}

type awsSpotCluster struct {
	sessions map[string]*ec2.EC2

	namespace  string
	aclTrigger db.Trigger
}

type awsID struct {
	spotID string
	region string
}

func getSpotIDs(ids []awsID) []string {
	var spotIDs []string
	for _, id := range ids {
		spotIDs = append(spotIDs, id.spotID)
	}

	return spotIDs
}

func groupByRegion(ids []awsID) map[string][]awsID {
	grouped := make(map[string][]awsID)
	for _, id := range ids {
		region := id.region
		if _, ok := grouped[region]; !ok {
			grouped[region] = []awsID{}
		}
		grouped[region] = append(grouped[region], id)
	}

	return grouped
}

func (clst *awsSpotCluster) Start(conn db.Conn, clusterID int, namespace string) error {
	clst.sessions = make(map[string]*ec2.EC2)
	clst.namespace = namespace
	clst.aclTrigger = conn.TriggerTick(60, db.ClusterTable)

	go clst.watchACLs(conn, clusterID)

	return nil
}

func (clst *awsSpotCluster) Disconnect() {
	/* Ideally we'd close clst.ec2 as well, but the API doesn't export that ability
	* apparently. */
	clst.aclTrigger.Stop()
}

func (clst awsSpotCluster) getSession(region string) *ec2.EC2 {
	if _, ok := clst.sessions[region]; ok {
		return clst.sessions[region]
	}

	session := session.New()
	session.Config.Region = aws.String(region)

	newEC2 := ec2.New(session)
	clst.sessions[region] = newEC2

	return newEC2
}

func (clst awsSpotCluster) Boot(bootSet []Machine) error {
	if len(bootSet) <= 0 {
		return nil
	}

	type bootReq struct {
		cfg    string
		size   string
		region string
	}

	bootReqMap := make(map[bootReq]int64) // From boot request to an instance count.
	for _, m := range bootSet {
		br := bootReq{
			cfg:    cloudConfigUbuntu(m.SSHKeys, "wily"),
			size:   m.Size,
			region: m.Region,
		}
		bootReqMap[br] = bootReqMap[br] + 1
	}

	var awsIDs []awsID
	for br, count := range bootReqMap {
		bd := &ec2.BlockDeviceMapping{
			Ebs: &ec2.EbsBlockDevice{
				DeleteOnTermination: aws.Bool(true),
				VolumeSize:          aws.Int64(diskSize),
			},
		}

		session := clst.getSession(br.region)
		cloudConfig64 := base64.StdEncoding.EncodeToString([]byte(br.cfg))
		resp, err := session.RequestSpotInstances(&ec2.RequestSpotInstancesInput{
			SpotPrice: aws.String(spotPrice),
			LaunchSpecification: &ec2.RequestSpotLaunchSpecification{
				ImageId:             aws.String(amis[br.region]),
				InstanceType:        aws.String(br.size),
				UserData:            &cloudConfig64,
				SecurityGroups:      []*string{&clst.namespace},
				BlockDeviceMappings: []*ec2.BlockDeviceMapping{bd},
			},
			InstanceCount: &count,
		})

		if err != nil {
			return err
		}

		for _, request := range resp.SpotInstanceRequests {
			awsIDs = append(awsIDs, awsID{
				spotID: *request.SpotInstanceRequestId,
				region: br.region})
		}
	}

	if err := clst.tagSpotRequests(awsIDs); err != nil {
		return err
	}

	if err := clst.wait(awsIDs, true); err != nil {
		return err
	}

	return nil
}

func (clst awsSpotCluster) Stop(machines []Machine) error {
	var awsIDs []awsID
	for _, m := range machines {
		awsIDs = append(awsIDs, awsID{
			region: m.Region,
			spotID: m.ID,
		})
	}
	for region, ids := range groupByRegion(awsIDs) {
		session := clst.getSession(region)
		spotIDs := getSpotIDs(ids)

		spots, err := session.DescribeSpotInstanceRequests(
			&ec2.DescribeSpotInstanceRequestsInput{
				SpotInstanceRequestIds: aws.StringSlice(spotIDs),
			})
		if err != nil {
			return err
		}

		instIds := []string{}
		for _, spot := range spots.SpotInstanceRequests {
			if spot.InstanceId != nil {
				instIds = append(instIds, *spot.InstanceId)
			}
		}

		if len(instIds) > 0 {
			_, err = session.TerminateInstances(&ec2.TerminateInstancesInput{
				InstanceIds: aws.StringSlice(instIds),
			})
			if err != nil {
				return err
			}
		}

		_, err = session.CancelSpotInstanceRequests(&ec2.CancelSpotInstanceRequestsInput{
			SpotInstanceRequestIds: aws.StringSlice(spotIDs),
		})
		if err != nil {
			return err
		}

		if err := clst.wait(ids, false); err != nil {
			return err
		}
	}

	return nil
}

func (clst awsSpotCluster) Get() ([]Machine, error) {
	machines := []Machine{}
	for region := range amis {
		session := clst.getSession(region)

		spots, err := session.DescribeSpotInstanceRequests(nil)
		if err != nil {
			return nil, err
		}

		insts, err := session.DescribeInstances(&ec2.DescribeInstancesInput{
			Filters: []*ec2.Filter{
				{
					Name:   aws.String("instance.group-name"),
					Values: []*string{aws.String(clst.namespace)},
				},
			},
		})
		if err != nil {
			return nil, err
		}

		instMap := make(map[string]*ec2.Instance)
		for _, res := range insts.Reservations {
			for _, inst := range res.Instances {
				instMap[*inst.InstanceId] = inst
			}
		}

		for _, spot := range spots.SpotInstanceRequests {
			if *spot.State != ec2.SpotInstanceStateActive &&
				*spot.State != ec2.SpotInstanceStateOpen {
				continue
			}

			var inst *ec2.Instance
			if spot.InstanceId != nil {
				inst = instMap[*spot.InstanceId]
			}

			// Due to a race condition in the AWS API, it's possible that spot
			// requests might lose their Tags.  If handled naively, those spot
			// requests would technically be without a namespace, meaning the
			// instances they create would be live forever as zombies.
			//
			// To mitigate this issue, we rely not only on the spot request tags, but
			// additionally on the instance security group.  If a spot request has a
			// running instance in the appropriate security group, it is by
			// definition in our namespace.  Thus, we only check the tags for spot
			// requests without running instances.
			if inst == nil {
				var isOurs bool
				for _, tag := range spot.Tags {
					ns := clst.namespace
					if tag != nil && tag.Key != nil && *tag.Key == ns {
						isOurs = true
						break
					}
				}

				if !isOurs {
					continue
				}
			}

			machine := Machine{
				ID:       *spot.SpotInstanceRequestId,
				Region:   region,
				Provider: db.AmazonSpot,
			}

			if inst != nil {
				if *inst.State.Name != ec2.InstanceStateNamePending &&
					*inst.State.Name != ec2.InstanceStateNameRunning {
					continue
				}

				if inst.PublicIpAddress != nil {
					machine.PublicIP = *inst.PublicIpAddress
				}

				if inst.PrivateIpAddress != nil {
					machine.PrivateIP = *inst.PrivateIpAddress
				}

				if inst.InstanceType != nil {
					machine.Size = *inst.InstanceType
				}
			}

			machines = append(machines, machine)
		}
	}

	return machines, nil
}

func (clst *awsSpotCluster) PickBestSize(ram dsl.Range, cpu dsl.Range, maxPrice float64) string {
	return pickBestSize(awsDescriptions, ram, cpu, maxPrice)
}

func (clst *awsSpotCluster) tagSpotRequests(awsIDs []awsID) error {
OuterLoop:
	for region, ids := range groupByRegion(awsIDs) {
		session := clst.getSession(region)
		spotIDs := getSpotIDs(ids)

		var err error
		for i := 0; i < 30; i++ {
			_, err = session.CreateTags(&ec2.CreateTagsInput{
				Tags: []*ec2.Tag{
					{Key: aws.String(clst.namespace), Value: aws.String("")},
				},
				Resources: aws.StringSlice(spotIDs),
			})
			if err == nil {
				continue OuterLoop
			}
			time.Sleep(5 * time.Second)
		}

		log.Warn("Failed to tag spot requests: ", err)
		session.CancelSpotInstanceRequests(
			&ec2.CancelSpotInstanceRequestsInput{
				SpotInstanceRequestIds: aws.StringSlice(spotIDs),
			})

		return err
	}

	return nil
}

/* Wait for the spot request 'ids' to have booted or terminated depending on the value of
* 'boot' */
func (clst *awsSpotCluster) wait(awsIDs []awsID, boot bool) error {
OuterLoop:
	for i := 0; i < 100; i++ {
		machines, err := clst.Get()
		if err != nil {
			log.WithError(err).Warn("Failed to get machines.")
			time.Sleep(2 * time.Second)
			continue
		}

		exists := make(map[awsID]struct{})
		for _, inst := range machines {
			id := awsID{
				spotID: inst.ID,
				region: inst.Region,
			}

			exists[id] = struct{}{}
		}

		for _, id := range awsIDs {
			if _, ok := exists[id]; ok != boot {
				time.Sleep(2 * time.Second)
				continue OuterLoop
			}
		}

		return nil
	}

	return errors.New("timed out")
}

func (clst *awsSpotCluster) updateSecurityGroups(acls []string) error {
	for _, session := range clst.sessions {
		resp, err := session.DescribeSecurityGroups(
			&ec2.DescribeSecurityGroupsInput{
				Filters: []*ec2.Filter{
					{
						Name:   aws.String("group-name"),
						Values: []*string{aws.String(clst.namespace)},
					},
				},
			})

		if err != nil {
			return err
		}

		ingress := []*ec2.IpPermission{}
		groups := resp.SecurityGroups
		if len(groups) > 1 {
			return errors.New("Multiple Security Groups with the same name: " +
				clst.namespace)
		} else if len(groups) == 0 {
			session.CreateSecurityGroup(&ec2.CreateSecurityGroupInput{
				Description: aws.String("Declarative Infrastructure Group"),
				GroupName:   aws.String(clst.namespace),
			})
		} else {
			/* XXX: Deal with egress rules. */
			ingress = groups[0].IpPermissions
		}

		permMap := make(map[string]bool)
		for _, acl := range acls {
			permMap[acl] = true
		}

		groupIngressExists := false
		for i, p := range ingress {
			if (i > 0 || p.FromPort != nil || p.ToPort != nil ||
				*p.IpProtocol != "-1") && p.UserIdGroupPairs == nil {
				log.Info("Revoke ingress security group: ", *p)
				_, err = session.RevokeSecurityGroupIngress(
					&ec2.RevokeSecurityGroupIngressInput{
						GroupName:     aws.String(clst.namespace),
						IpPermissions: []*ec2.IpPermission{p}})
				if err != nil {
					return err
				}

				continue
			}

			for _, ipr := range p.IpRanges {
				ip := *ipr.CidrIp
				if !permMap[ip] {
					log.Info("Revoke ingress security group: ", ip)
					_, err = session.RevokeSecurityGroupIngress(
						&ec2.RevokeSecurityGroupIngressInput{
							GroupName:  aws.String(clst.namespace),
							CidrIp:     aws.String(ip),
							FromPort:   p.FromPort,
							IpProtocol: p.IpProtocol,
							ToPort:     p.ToPort})
					if err != nil {
						return err
					}
				} else {
					permMap[ip] = false
				}
			}

			if len(groups) > 0 {
				for _, grp := range p.UserIdGroupPairs {
					if *grp.GroupId != *groups[0].GroupId {
						log.Info("Revoke ingress security group GroupID: ",
							*grp.GroupId)
						_, err = session.RevokeSecurityGroupIngress(
							&ec2.RevokeSecurityGroupIngressInput{
								GroupName:               aws.String(clst.namespace),
								SourceSecurityGroupName: grp.GroupName})
						if err != nil {
							return err
						}
					} else {
						groupIngressExists = true
					}
				}
			}
		}

		if !groupIngressExists {
			log.Info("Add intragroup ACL")
			_, err = session.AuthorizeSecurityGroupIngress(
				&ec2.AuthorizeSecurityGroupIngressInput{
					GroupName:               aws.String(clst.namespace),
					SourceSecurityGroupName: aws.String(clst.namespace)})
		}

		for perm, install := range permMap {
			if !install {
				continue
			}

			log.Info("Add ACL: ", perm)
			_, err = session.AuthorizeSecurityGroupIngress(
				&ec2.AuthorizeSecurityGroupIngressInput{
					CidrIp:     aws.String(perm),
					GroupName:  aws.String(clst.namespace),
					IpProtocol: aws.String("-1")})

			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (clst *awsSpotCluster) watchACLs(conn db.Conn, clusterID int) {
	for range clst.aclTrigger.C {
		var acls []string
		conn.Transact(func(view db.Database) error {
			clusters := view.SelectFromCluster(func(c db.Cluster) bool {
				return c.ID == clusterID
			})

			if len(clusters) == 0 {
				log.Warn("Undefined cluster.")
				return nil
			} else if len(clusters) > 1 {
				panic("Duplicate Clusters")
			}

			acls = clusters[0].ACLs
			return nil
		})

		clst.updateSecurityGroups(acls)
	}
}
<|endoftext|>


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

// ct_hammer is a stress/load test for a CT log.
package main

import (
	"bytes"
	"compress/gzip"
	"crypto"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/golang/protobuf/ptypes"
	"github.com/google/certificate-transparency-go/fixchain/ratelimiter"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/trillian/ctfe/configpb"
	"github.com/google/certificate-transparency-go/trillian/integration"
	"github.com/google/trillian/monitoring"
	"github.com/google/trillian/monitoring/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	ct "github.com/google/certificate-transparency-go"
)

var (
	banner              = flag.Bool("banner", true, "Display intro")
	httpServers         = flag.String("ct_http_servers", "localhost:8092", "Comma-separated list of (assumed interchangeable) servers, each as address:port")
	testDir             = flag.String("testdata_dir", "testdata", "Name of directory with test data")
	leafNotAfter        = flag.String("leaf_not_after", "", "Not-After date to use for leaf certs, RFC3339/ISO-8601 format (e.g. 2017-11-26T12:29:19Z)")
	metricsEndpoint     = flag.String("metrics_endpoint", "", "Endpoint for serving metrics; if left empty, metrics will not be exposed")
	seed                = flag.Int64("seed", -1, "Seed for random number generation")
	logConfig           = flag.String("log_config", "", "File holding log config in JSON")
	mmd                 = flag.Duration("mmd", 2*time.Minute, "Default MMD for logs")
	operations          = flag.Uint64("operations", ^uint64(0), "Number of operations to perform")
	minGetEntries       = flag.Int("min_get_entries", 1, "Minimum get-entries request size")
	maxGetEntries       = flag.Int("max_get_entries", 500, "Maximum get-entries request size")
	oversizedGetEntries = flag.Bool("oversized_get_entries", false, "Whether get-entries requests can go beyond log size")
	maxParallelChains   = flag.Int("max_parallel_chains", 2, "Maximum number of chains to add in parallel (will always add at least 1 chain)")
	limit               = flag.Int("rate_limit", 0, "Maximum rate of requests to an individual log; 0 for no rate limit")
	ignoreErrors        = flag.Bool("ignore_errors", false, "Whether to ignore errors and retry the operation")
	maxRetry            = flag.Duration("max_retry", 60*time.Second, "How long to keep retrying when ignore_errors is set")
)
var (
	addChainBias          = flag.Int("add_chain", 20, "Bias for add-chain operations")
	addPreChainBias       = flag.Int("add_pre_chain", 20, "Bias for add-pre-chain operations")
	getSTHBias            = flag.Int("get_sth", 2, "Bias for get-sth operations")
	getSTHConsistencyBias = flag.Int("get_sth_consistency", 2, "Bias for get-sth-consistency operations")
	getProofByHashBias    = flag.Int("get_proof_by_hash", 2, "Bias for get-proof-by-hash operations")
	getEntriesBias        = flag.Int("get_entries", 2, "Bias for get-entries operations")
	getRootsBias          = flag.Int("get_roots", 1, "Bias for get-roots operations")
	getEntryAndProofBias  = flag.Int("get_entry_and_proof", 0, "Bias for get-entry-and-proof operations")
	invalidChance         = flag.Int("invalid_chance", 10, "Chance of generating an invalid operation, as the N in 1-in-N (0 for never)")
)

func newLimiter(rate int) integration.Limiter {
	if rate <= 0 {
		return nil
	}
	return ratelimiter.NewLimiter(rate)
}

func main() {
	flag.Parse()
	if *logConfig == "" {
		glog.Exit("Test aborted as no log config provided (via --log_config)")
	}
	if *seed == -1 {
		*seed = time.Now().UTC().UnixNano() & 0xFFFFFFFF
	}
	fmt.Printf("Today's test has been brought to you by the letters C and T and the number %#x\n", *seed)
	rand.Seed(*seed)

	cfg, err := ctfe.LogConfigFromFile(*logConfig)
	if err != nil {
		glog.Exitf("Failed to read log config: %v", err)
	}

	var leafChain []ct.ASN1Cert
	var signer crypto.Signer
	if *testDir != "" {
		// Retrieve the test data.
		leafChain, err = integration.GetChain(*testDir, "leaf01.chain")
		if err != nil {
			glog.Exitf("failed to load certificate: %v", err)
		}
		signer, err = integration.MakeSigner(*testDir)
		if err != nil {
			glog.Exitf("failed to retrieve signer for re-signing: %v", err)
		}
	} else {
		glog.Warningf("Warning: add-[pre-]chain operations disabled as no cert generation method available")
		*addChainBias = 0
		*addPreChainBias = 0
	}

	var notAfterOverride time.Time
	if *leafNotAfter != "" {
		notAfterOverride, err = time.Parse(time.RFC3339, *leafNotAfter)
		if err != nil {
			glog.Exitf("Failed to parse leaf notAfter: %v", err)
		}
	}

	bias := integration.HammerBias{
		Bias: map[ctfe.EntrypointName]int{
			ctfe.AddChainName:          *addChainBias,
			ctfe.AddPreChainName:       *addPreChainBias,
			ctfe.GetSTHName:            *getSTHBias,
			ctfe.GetSTHConsistencyName: *getSTHConsistencyBias,
			ctfe.GetProofByHashName:    *getProofByHashBias,
			ctfe.GetEntriesName:        *getEntriesBias,
			ctfe.GetRootsName:          *getRootsBias,
			ctfe.GetEntryAndProofName:  *getEntryAndProofBias,
		},
		InvalidChance: map[ctfe.EntrypointName]int{
			ctfe.AddChainName:          *invalidChance,
			ctfe.AddPreChainName:       *invalidChance,
			ctfe.GetSTHName:            0,
			ctfe.GetSTHConsistencyName: *invalidChance,
			ctfe.GetProofByHashName:    *invalidChance,
			ctfe.GetEntriesName:        *invalidChance,
			ctfe.GetRootsName:          0,
			ctfe.GetEntryAndProofName:  0,
		},
	}

	var mf monitoring.MetricFactory
	if *metricsEndpoint != "" {
		mf = prometheus.MetricFactory{}
		http.Handle("/metrics", promhttp.Handler())
		server := http.Server{Addr: *metricsEndpoint, Handler: nil}
		glog.Infof("Serving metrics at %v", *metricsEndpoint)
		go func() {
			err := server.ListenAndServe()
			glog.Warningf("Metrics server exited: %v", err)
		}()
	} else {
		mf = monitoring.InertMetricFactory{}
	}

	if *banner {
		fmt.Print("\n\nStop")
		for i := 0; i < 8; i++ {
			time.Sleep(100 * time.Millisecond)
			fmt.Print(".")
		}
		mc := "H4sIAAAAAAAA/4xVPbLzMAjsv1OkU8FI9LqDOAUFDUNBxe2/QXYSS/HLe5SeXZYfsf73+D1KB8D2B2RxZpGw8gcsSoQYeH1ya0fof1BpnhpuUR+P8ijorESq8Yto6WYWqsrMGh4qSkdI/YFZWu8d3AAAkklEHBGTNAYxbpKltWRgRzQ3A3CImDIjVSVCicThbLK0VjsiAGAGIIKbmUcIq/KkqYo4BNZDqtgZMAPNPSJCRISZZ36d5OiTUbqJZAOYIoCHUreImJsCPMobQ20SqjBbLWWbBGRREhHQU2MMUu9TwB12cC7X3SNrs1yPKvv5gD4yn+kzshOfMg69fVknJNbdcsjuDvgNXWPmTXCuEnuvP4NdlSWymIQjfsFWzbERZ5sz730NpbvoOGMOzu7eeBUaW3w8r4z2iRuD4uY6W9wgZ96+YZvpHW7SabvlH7CviKWQyp81EL2zj7Fcbee7MpSuNHzj2z18LdAvAkAr8pr/3cGFUO+apa2n64TK3XouTBpEch2Rf8GnzajAFY438+SzgURfV7sXT+q1FNTJYdLF9WxJzFheAyNmXfKuiel5/mW2QqSx2umlQ+L2GpTPWZBu5tvpXW5/fy4xTYd2ly+vR052dZbjTIh0u4vzyRDF6kPzoRLRfhp2pqnr5wce5eAGP6onaRv8EYdl7gfd5zIId/gxYvr4pWW7KnbjoU6kRL62e25b44ZQz7Oaf4GrTovnqemNsyOdL40Dls11ocMPn29nYeUvmt3S1v8DAAD//wEAAP//TRo+KHEIAAA="
		mcData, _ := base64.StdEncoding.DecodeString(mc)
		b := bytes.NewReader(mcData)
		r, _ := gzip.NewReader(b)
		io.Copy(os.Stdout, r)
		r.Close()
		fmt.Print("\n\nHammer Time\n\n")
	}

	type result struct {
		prefix string
		err    error
	}
	results := make(chan result, len(cfg))
	var wg sync.WaitGroup
	for _, c := range cfg {
		wg.Add(1)
		pool, err := integration.NewRandomPool(*httpServers, c.PublicKey, c.Prefix)
		if err != nil {
			glog.Exitf("Failed to create client pool: %v", err)
		}

		mmd := *mmd
		// Note: Although the (usually lower than MMD) expected merge delay is not
		// a guarantee, it should be OK for testing.
		if emd := c.ExpectedMergeDelaySec; emd != 0 {
			mmd = time.Second * time.Duration(emd)
		}

		notAfter := notAfterOverride
		if notAfter.IsZero() {
			notAfter, err = notAfterForLog(c)
			if err != nil {
				glog.Exitf("Failed to determine notAfter for %s: %v", c.Prefix, err)
			}
		}
		generator, err := integration.NewSyntheticChainGenerator(leafChain, signer, notAfter)
		if err != nil {
			glog.Exitf("Failed to build chain generator: %v", err)
		}

		cfg := integration.HammerConfig{
			LogCfg:              c,
			MetricFactory:       mf,
			MMD:                 mmd,
			ChainGenerator:      generator,
			ClientPool:          pool,
			EPBias:              bias,
			MinGetEntries:       *minGetEntries,
			MaxGetEntries:       *maxGetEntries,
			OversizedGetEntries: *oversizedGetEntries,
			Operations:          *operations,
			Limiter:             newLimiter(*limit),
			MaxParallelChains:   *maxParallelChains,
			IgnoreErrors:        *ignoreErrors,
			MaxRetryDuration:    *maxRetry,
		}
		go func(cfg integration.HammerConfig) {
			defer wg.Done()
			err := integration.HammerCTLog(cfg)
			results <- result{prefix: cfg.LogCfg.Prefix, err: err}
		}(cfg)
	}
	wg.Wait()

	glog.Infof("completed tests on all %d logs:", len(cfg))
	close(results)
	errCount := 0
	for e := range results {
		if e.err != nil {
			errCount++
			glog.Errorf("  %s: failed with %v", e.prefix, e.err)
		}
	}
	if errCount > 0 {
		glog.Exitf("non-zero error count (%d), exiting", errCount)
	}
	glog.Info("  no errors; done")
}

// notAfterForLog returns a notAfter time to be used for certs submitted
// to the given log instance, allowing for any temporal shard configuration.
func notAfterForLog(c *configpb.LogConfig) (time.Time, error) {
	if c.NotAfterStart == nil && c.NotAfterLimit == nil {
		return time.Now().Add(24 * time.Hour), nil
	}

	if c.NotAfterStart != nil {
		start, err := ptypes.Timestamp(c.NotAfterStart)
		if err != nil {
			return time.Time{}, fmt.Errorf("failed to parse NotAfterStart: %v", err)
		}
		if c.NotAfterLimit == nil {
			return start.Add(24 * time.Hour), nil
		}

		limit, err := ptypes.Timestamp(c.NotAfterLimit)
		if err != nil {
			return time.Time{}, fmt.Errorf("failed to parse NotAfterLimit: %v", err)
		}
		return time.Unix(0, (limit.UnixNano()-start.UnixNano())/2+start.UnixNano()), nil
	}

	limit, err := ptypes.Timestamp(c.NotAfterLimit)
	if err != nil {
		return time.Time{}, fmt.Errorf("failed to parse NotAfterLimit: %v", err)
	}
	return limit.Add(-1 * time.Hour), nil
}
<commit_msg>ct_hammer: put chain generator creation into a factory<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// ct_hammer is a stress/load test for a CT log.
package main

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/golang/protobuf/ptypes"
	"github.com/google/certificate-transparency-go/fixchain/ratelimiter"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/trillian/ctfe/configpb"
	"github.com/google/certificate-transparency-go/trillian/integration"
	"github.com/google/trillian/monitoring"
	"github.com/google/trillian/monitoring/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	banner              = flag.Bool("banner", true, "Display intro")
	httpServers         = flag.String("ct_http_servers", "localhost:8092", "Comma-separated list of (assumed interchangeable) servers, each as address:port")
	testDir             = flag.String("testdata_dir", "testdata", "Name of directory with test data")
	leafNotAfter        = flag.String("leaf_not_after", "", "Not-After date to use for leaf certs, RFC3339/ISO-8601 format (e.g. 2017-11-26T12:29:19Z)")
	metricsEndpoint     = flag.String("metrics_endpoint", "", "Endpoint for serving metrics; if left empty, metrics will not be exposed")
	seed                = flag.Int64("seed", -1, "Seed for random number generation")
	logConfig           = flag.String("log_config", "", "File holding log config in JSON")
	mmd                 = flag.Duration("mmd", 2*time.Minute, "Default MMD for logs")
	operations          = flag.Uint64("operations", ^uint64(0), "Number of operations to perform")
	minGetEntries       = flag.Int("min_get_entries", 1, "Minimum get-entries request size")
	maxGetEntries       = flag.Int("max_get_entries", 500, "Maximum get-entries request size")
	oversizedGetEntries = flag.Bool("oversized_get_entries", false, "Whether get-entries requests can go beyond log size")
	maxParallelChains   = flag.Int("max_parallel_chains", 2, "Maximum number of chains to add in parallel (will always add at least 1 chain)")
	limit               = flag.Int("rate_limit", 0, "Maximum rate of requests to an individual log; 0 for no rate limit")
	ignoreErrors        = flag.Bool("ignore_errors", false, "Whether to ignore errors and retry the operation")
	maxRetry            = flag.Duration("max_retry", 60*time.Second, "How long to keep retrying when ignore_errors is set")
)
var (
	addChainBias          = flag.Int("add_chain", 20, "Bias for add-chain operations")
	addPreChainBias       = flag.Int("add_pre_chain", 20, "Bias for add-pre-chain operations")
	getSTHBias            = flag.Int("get_sth", 2, "Bias for get-sth operations")
	getSTHConsistencyBias = flag.Int("get_sth_consistency", 2, "Bias for get-sth-consistency operations")
	getProofByHashBias    = flag.Int("get_proof_by_hash", 2, "Bias for get-proof-by-hash operations")
	getEntriesBias        = flag.Int("get_entries", 2, "Bias for get-entries operations")
	getRootsBias          = flag.Int("get_roots", 1, "Bias for get-roots operations")
	getEntryAndProofBias  = flag.Int("get_entry_and_proof", 0, "Bias for get-entry-and-proof operations")
	invalidChance         = flag.Int("invalid_chance", 10, "Chance of generating an invalid operation, as the N in 1-in-N (0 for never)")
)

func newLimiter(rate int) integration.Limiter {
	if rate <= 0 {
		return nil
	}
	return ratelimiter.NewLimiter(rate)
}

func main() {
	flag.Parse()
	if *logConfig == "" {
		glog.Exit("Test aborted as no log config provided (via --log_config)")
	}
	if *seed == -1 {
		*seed = time.Now().UTC().UnixNano() & 0xFFFFFFFF
	}
	fmt.Printf("Today's test has been brought to you by the letters C and T and the number %#x\n", *seed)
	rand.Seed(*seed)

	cfg, err := ctfe.LogConfigFromFile(*logConfig)
	if err != nil {
		glog.Exitf("Failed to read log config: %v", err)
	}

	var generatorFactory func(c *configpb.LogConfig) (integration.ChainGenerator, error) // Almost Java-esque...
	if *testDir != "" {
		// Retrieve the test data.
		leafChain, err := integration.GetChain(*testDir, "leaf01.chain")
		if err != nil {
			glog.Exitf("failed to load certificate: %v", err)
		}
		signer, err := integration.MakeSigner(*testDir)
		if err != nil {
			glog.Exitf("failed to retrieve signer for re-signing: %v", err)
		}
		var notAfterOverride time.Time
		if *leafNotAfter != "" {
			notAfterOverride, err = time.Parse(time.RFC3339, *leafNotAfter)
			if err != nil {
				glog.Exitf("Failed to parse leaf notAfter: %v", err)
			}
		}
		// Build a synthetic generator for each target log.
		generatorFactory = func(c *configpb.LogConfig) (integration.ChainGenerator, error) {
			notAfter := notAfterOverride
			if notAfter.IsZero() {
				notAfter, err = notAfterForLog(c)
				if err != nil {
					return nil, fmt.Errorf("failed to determine notAfter for %s: %v", c.Prefix, err)
				}
			}
			return integration.NewSyntheticChainGenerator(leafChain, signer, notAfter)
		}
	}

	if generatorFactory == nil {
		glog.Warningf("Warning: add-[pre-]chain operations disabled as no cert generation method available")
		*addChainBias = 0
		*addPreChainBias = 0
		generatorFactory = func(c *configpb.LogConfig) (integration.ChainGenerator, error) {
			return nil, nil
		}
	}

	bias := integration.HammerBias{
		Bias: map[ctfe.EntrypointName]int{
			ctfe.AddChainName:          *addChainBias,
			ctfe.AddPreChainName:       *addPreChainBias,
			ctfe.GetSTHName:            *getSTHBias,
			ctfe.GetSTHConsistencyName: *getSTHConsistencyBias,
			ctfe.GetProofByHashName:    *getProofByHashBias,
			ctfe.GetEntriesName:        *getEntriesBias,
			ctfe.GetRootsName:          *getRootsBias,
			ctfe.GetEntryAndProofName:  *getEntryAndProofBias,
		},
		InvalidChance: map[ctfe.EntrypointName]int{
			ctfe.AddChainName:          *invalidChance,
			ctfe.AddPreChainName:       *invalidChance,
			ctfe.GetSTHName:            0,
			ctfe.GetSTHConsistencyName: *invalidChance,
			ctfe.GetProofByHashName:    *invalidChance,
			ctfe.GetEntriesName:        *invalidChance,
			ctfe.GetRootsName:          0,
			ctfe.GetEntryAndProofName:  0,
		},
	}

	var mf monitoring.MetricFactory
	if *metricsEndpoint != "" {
		mf = prometheus.MetricFactory{}
		http.Handle("/metrics", promhttp.Handler())
		server := http.Server{Addr: *metricsEndpoint, Handler: nil}
		glog.Infof("Serving metrics at %v", *metricsEndpoint)
		go func() {
			err := server.ListenAndServe()
			glog.Warningf("Metrics server exited: %v", err)
		}()
	} else {
		mf = monitoring.InertMetricFactory{}
	}

	if *banner {
		fmt.Print("\n\nStop")
		for i := 0; i < 8; i++ {
			time.Sleep(100 * time.Millisecond)
			fmt.Print(".")
		}
		mc := "H4sIAAAAAAAA/4xVPbLzMAjsv1OkU8FI9LqDOAUFDUNBxe2/QXYSS/HLe5SeXZYfsf73+D1KB8D2B2RxZpGw8gcsSoQYeH1ya0fof1BpnhpuUR+P8ijorESq8Yto6WYWqsrMGh4qSkdI/YFZWu8d3AAAkklEHBGTNAYxbpKltWRgRzQ3A3CImDIjVSVCicThbLK0VjsiAGAGIIKbmUcIq/KkqYo4BNZDqtgZMAPNPSJCRISZZ36d5OiTUbqJZAOYIoCHUreImJsCPMobQ20SqjBbLWWbBGRREhHQU2MMUu9TwB12cC7X3SNrs1yPKvv5gD4yn+kzshOfMg69fVknJNbdcsjuDvgNXWPmTXCuEnuvP4NdlSWymIQjfsFWzbERZ5sz730NpbvoOGMOzu7eeBUaW3w8r4z2iRuD4uY6W9wgZ96+YZvpHW7SabvlH7CviKWQyp81EL2zj7Fcbee7MpSuNHzj2z18LdAvAkAr8pr/3cGFUO+apa2n64TK3XouTBpEch2Rf8GnzajAFY438+SzgURfV7sXT+q1FNTJYdLF9WxJzFheAyNmXfKuiel5/mW2QqSx2umlQ+L2GpTPWZBu5tvpXW5/fy4xTYd2ly+vR052dZbjTIh0u4vzyRDF6kPzoRLRfhp2pqnr5wce5eAGP6onaRv8EYdl7gfd5zIId/gxYvr4pWW7KnbjoU6kRL62e25b44ZQz7Oaf4GrTovnqemNsyOdL40Dls11ocMPn29nYeUvmt3S1v8DAAD//wEAAP//TRo+KHEIAAA="
		mcData, _ := base64.StdEncoding.DecodeString(mc)
		b := bytes.NewReader(mcData)
		r, _ := gzip.NewReader(b)
		io.Copy(os.Stdout, r)
		r.Close()
		fmt.Print("\n\nHammer Time\n\n")
	}

	type result struct {
		prefix string
		err    error
	}
	results := make(chan result, len(cfg))
	var wg sync.WaitGroup
	for _, c := range cfg {
		wg.Add(1)
		pool, err := integration.NewRandomPool(*httpServers, c.PublicKey, c.Prefix)
		if err != nil {
			glog.Exitf("Failed to create client pool: %v", err)
		}

		mmd := *mmd
		// Note: Although the (usually lower than MMD) expected merge delay is not
		// a guarantee, it should be OK for testing.
		if emd := c.ExpectedMergeDelaySec; emd != 0 {
			mmd = time.Second * time.Duration(emd)
		}

		generator, err := generatorFactory(c)
		if err != nil {
			glog.Exitf("Failed to build chain generator: %v", err)
		}

		cfg := integration.HammerConfig{
			LogCfg:              c,
			MetricFactory:       mf,
			MMD:                 mmd,
			ChainGenerator:      generator,
			ClientPool:          pool,
			EPBias:              bias,
			MinGetEntries:       *minGetEntries,
			MaxGetEntries:       *maxGetEntries,
			OversizedGetEntries: *oversizedGetEntries,
			Operations:          *operations,
			Limiter:             newLimiter(*limit),
			MaxParallelChains:   *maxParallelChains,
			IgnoreErrors:        *ignoreErrors,
			MaxRetryDuration:    *maxRetry,
		}
		go func(cfg integration.HammerConfig) {
			defer wg.Done()
			err := integration.HammerCTLog(cfg)
			results <- result{prefix: cfg.LogCfg.Prefix, err: err}
		}(cfg)
	}
	wg.Wait()

	glog.Infof("completed tests on all %d logs:", len(cfg))
	close(results)
	errCount := 0
	for e := range results {
		if e.err != nil {
			errCount++
			glog.Errorf("  %s: failed with %v", e.prefix, e.err)
		}
	}
	if errCount > 0 {
		glog.Exitf("non-zero error count (%d), exiting", errCount)
	}
	glog.Info("  no errors; done")
}

// notAfterForLog returns a notAfter time to be used for certs submitted
// to the given log instance, allowing for any temporal shard configuration.
func notAfterForLog(c *configpb.LogConfig) (time.Time, error) {
	if c.NotAfterStart == nil && c.NotAfterLimit == nil {
		return time.Now().Add(24 * time.Hour), nil
	}

	if c.NotAfterStart != nil {
		start, err := ptypes.Timestamp(c.NotAfterStart)
		if err != nil {
			return time.Time{}, fmt.Errorf("failed to parse NotAfterStart: %v", err)
		}
		if c.NotAfterLimit == nil {
			return start.Add(24 * time.Hour), nil
		}

		limit, err := ptypes.Timestamp(c.NotAfterLimit)
		if err != nil {
			return time.Time{}, fmt.Errorf("failed to parse NotAfterLimit: %v", err)
		}
		return time.Unix(0, (limit.UnixNano()-start.UnixNano())/2+start.UnixNano()), nil
	}

	limit, err := ptypes.Timestamp(c.NotAfterLimit)
	if err != nil {
		return time.Time{}, fmt.Errorf("failed to parse NotAfterLimit: %v", err)
	}
	return limit.Add(-1 * time.Hour), nil
}
<|endoftext|>


<commit_before>package weed_server

import (
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/util/log_buffer"
	"strings"
	"time"

	"github.com/golang/protobuf/proto"

	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (fs *FilerServer) SubscribeMetadata(req *filer_pb.SubscribeMetadataRequest, stream filer_pb.SeaweedFiler_SubscribeMetadataServer) error {

	peerAddress := findClientAddress(stream.Context(), 0)

	clientName := fs.addClient(req.ClientName, peerAddress)

	defer fs.deleteClient(clientName)

	lastReadTime := time.Unix(0, req.SinceNs)
	glog.V(0).Infof(" %v starts to subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

	eachEventNotificationFn := fs.eachEventNotificationFn(req, stream, clientName, req.Signature)

	eachLogEntryFn := eachLogEntryFn(eachEventNotificationFn)

	var processedTsNs int64
	var readPersistedLogErr error
	var readInMemoryLogErr error

	for {

		glog.V(4).Infof("read on disk %v aggregated subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

		processedTsNs, readPersistedLogErr = fs.filer.ReadPersistedLogBuffer(lastReadTime, eachLogEntryFn)
		if readPersistedLogErr != nil {
			return fmt.Errorf("reading from persisted logs: %v", readPersistedLogErr)
		}

		if processedTsNs != 0 {
			lastReadTime = time.Unix(0, processedTsNs)
		} else {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				time.Sleep(1127 * time.Millisecond)
				continue
			}
		}

		glog.V(4).Infof("read in memory %v aggregated subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

		lastReadTime, readInMemoryLogErr = fs.filer.MetaAggregator.MetaLogBuffer.LoopProcessLogData("aggMeta:"+clientName, lastReadTime, func() bool {
			fs.filer.MetaAggregator.ListenersLock.Lock()
			fs.filer.MetaAggregator.ListenersCond.Wait()
			fs.filer.MetaAggregator.ListenersLock.Unlock()
			return true
		}, eachLogEntryFn)
		if readInMemoryLogErr != nil {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				continue
			}
			glog.Errorf("processed to %v: %v", lastReadTime, readInMemoryLogErr)
			if readInMemoryLogErr != log_buffer.ResumeError {
				break
			}
		}

		time.Sleep(1127 * time.Millisecond)
	}

	return readInMemoryLogErr

}

func (fs *FilerServer) SubscribeLocalMetadata(req *filer_pb.SubscribeMetadataRequest, stream filer_pb.SeaweedFiler_SubscribeLocalMetadataServer) error {

	peerAddress := findClientAddress(stream.Context(), 0)

	clientName := fs.addClient(req.ClientName, peerAddress)

	defer fs.deleteClient(clientName)

	lastReadTime := time.Unix(0, req.SinceNs)
	glog.V(0).Infof(" %v local subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

	eachEventNotificationFn := fs.eachEventNotificationFn(req, stream, clientName, req.Signature)

	eachLogEntryFn := eachLogEntryFn(eachEventNotificationFn)

	var processedTsNs int64
	var readPersistedLogErr error
	var readInMemoryLogErr error

	for {
		// println("reading from persisted logs ...")
		glog.V(0).Infof("read on disk %v local subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)
		processedTsNs, readPersistedLogErr = fs.filer.ReadPersistedLogBuffer(lastReadTime, eachLogEntryFn)
		if readPersistedLogErr != nil {
			return fmt.Errorf("reading from persisted logs: %v", readPersistedLogErr)
		}

		if processedTsNs != 0 {
			lastReadTime = time.Unix(0, processedTsNs)
		} else {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				time.Sleep(1127 * time.Millisecond)
				continue
			}
		}

		glog.V(0).Infof("read in memory %v local subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

		lastReadTime, readInMemoryLogErr = fs.filer.LocalMetaLogBuffer.LoopProcessLogData("localMeta:"+clientName, lastReadTime, func() bool {
			fs.listenersLock.Lock()
			fs.listenersCond.Wait()
			fs.listenersLock.Unlock()
			return true
		}, eachLogEntryFn)
		if readInMemoryLogErr != nil {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				continue
			}
			glog.Errorf("processed to %v: %v", lastReadTime, readInMemoryLogErr)
			time.Sleep(1127 * time.Millisecond)
			if readInMemoryLogErr != log_buffer.ResumeError {
				break
			}
		}
	}

	return readInMemoryLogErr

}

func eachLogEntryFn(eachEventNotificationFn func(dirPath string, eventNotification *filer_pb.EventNotification, tsNs int64) error) func(logEntry *filer_pb.LogEntry) error {
	return func(logEntry *filer_pb.LogEntry) error {
		event := &filer_pb.SubscribeMetadataResponse{}
		if err := proto.Unmarshal(logEntry.Data, event); err != nil {
			glog.Errorf("unexpected unmarshal filer_pb.SubscribeMetadataResponse: %v", err)
			return fmt.Errorf("unexpected unmarshal filer_pb.SubscribeMetadataResponse: %v", err)
		}

		if err := eachEventNotificationFn(event.Directory, event.EventNotification, event.TsNs); err != nil {
			return err
		}

		return nil
	}
}

func (fs *FilerServer) eachEventNotificationFn(req *filer_pb.SubscribeMetadataRequest, stream filer_pb.SeaweedFiler_SubscribeMetadataServer, clientName string, clientSignature int32) func(dirPath string, eventNotification *filer_pb.EventNotification, tsNs int64) error {
	return func(dirPath string, eventNotification *filer_pb.EventNotification, tsNs int64) error {

		foundSelf := false
		for _, sig := range eventNotification.Signatures {
			if sig == clientSignature && clientSignature != 0 {
				return nil
			}
			if sig == fs.filer.Signature {
				foundSelf = true
			}
		}
		if !foundSelf {
			eventNotification.Signatures = append(eventNotification.Signatures, fs.filer.Signature)
		}

		// get complete path to the file or directory
		var entryName string
		if eventNotification.OldEntry != nil {
			entryName = eventNotification.OldEntry.Name
		} else if eventNotification.NewEntry != nil {
			entryName = eventNotification.NewEntry.Name
		}

		fullpath := util.Join(dirPath, entryName)

		// skip on filer internal meta logs
		if strings.HasPrefix(fullpath, filer.SystemLogDir) {
			return nil
		}

		if !strings.HasPrefix(fullpath, req.PathPrefix) {
			if eventNotification.NewParentPath != "" {
				newFullPath := util.Join(eventNotification.NewParentPath, entryName)
				if !strings.HasPrefix(newFullPath, req.PathPrefix) {
					return nil
				}
			} else {
				return nil
			}
		}

		message := &filer_pb.SubscribeMetadataResponse{
			Directory:         dirPath,
			EventNotification: eventNotification,
			TsNs:              tsNs,
		}
		// println("sending", dirPath, entryName)
		if err := stream.Send(message); err != nil {
			glog.V(0).Infof("=> client %v: %+v", clientName, err)
			return err
		}
		return nil
	}
}

func (fs *FilerServer) addClient(clientType string, clientAddress string) (clientName string) {
	clientName = clientType + "@" + clientAddress
	glog.V(0).Infof("+ listener %v", clientName)
	return
}

func (fs *FilerServer) deleteClient(clientName string) {
	glog.V(0).Infof("- listener %v", clientName)
}
<commit_msg>sync empty notification with timestamp<commit_after>package weed_server

import (
	"fmt"
	"strings"
	"time"

	"github.com/golang/protobuf/proto"

	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/util/log_buffer"
)

const (
	// MaxUnsyncedEvents send empty notification with timestamp when certain amount of events have been filtered
	MaxUnsyncedEvents = 1e3
)

func (fs *FilerServer) SubscribeMetadata(req *filer_pb.SubscribeMetadataRequest, stream filer_pb.SeaweedFiler_SubscribeMetadataServer) error {

	peerAddress := findClientAddress(stream.Context(), 0)

	clientName := fs.addClient(req.ClientName, peerAddress)

	defer fs.deleteClient(clientName)

	lastReadTime := time.Unix(0, req.SinceNs)
	glog.V(0).Infof(" %v starts to subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

	eachEventNotificationFn := fs.eachEventNotificationFn(req, stream, clientName)

	eachLogEntryFn := eachLogEntryFn(eachEventNotificationFn)

	var processedTsNs int64
	var readPersistedLogErr error
	var readInMemoryLogErr error

	for {

		glog.V(4).Infof("read on disk %v aggregated subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

		processedTsNs, readPersistedLogErr = fs.filer.ReadPersistedLogBuffer(lastReadTime, eachLogEntryFn)
		if readPersistedLogErr != nil {
			return fmt.Errorf("reading from persisted logs: %v", readPersistedLogErr)
		}

		if processedTsNs != 0 {
			lastReadTime = time.Unix(0, processedTsNs)
		} else {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				time.Sleep(1127 * time.Millisecond)
				continue
			}
		}

		glog.V(4).Infof("read in memory %v aggregated subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

		lastReadTime, readInMemoryLogErr = fs.filer.MetaAggregator.MetaLogBuffer.LoopProcessLogData("aggMeta:"+clientName, lastReadTime, func() bool {
			fs.filer.MetaAggregator.ListenersLock.Lock()
			fs.filer.MetaAggregator.ListenersCond.Wait()
			fs.filer.MetaAggregator.ListenersLock.Unlock()
			return true
		}, eachLogEntryFn)
		if readInMemoryLogErr != nil {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				continue
			}
			glog.Errorf("processed to %v: %v", lastReadTime, readInMemoryLogErr)
			if readInMemoryLogErr != log_buffer.ResumeError {
				break
			}
		}

		time.Sleep(1127 * time.Millisecond)
	}

	return readInMemoryLogErr

}

func (fs *FilerServer) SubscribeLocalMetadata(req *filer_pb.SubscribeMetadataRequest, stream filer_pb.SeaweedFiler_SubscribeLocalMetadataServer) error {

	peerAddress := findClientAddress(stream.Context(), 0)

	clientName := fs.addClient(req.ClientName, peerAddress)

	defer fs.deleteClient(clientName)

	lastReadTime := time.Unix(0, req.SinceNs)
	glog.V(0).Infof(" %v local subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

	eachEventNotificationFn := fs.eachEventNotificationFn(req, stream, clientName)

	eachLogEntryFn := eachLogEntryFn(eachEventNotificationFn)

	var processedTsNs int64
	var readPersistedLogErr error
	var readInMemoryLogErr error

	for {
		// println("reading from persisted logs ...")
		glog.V(0).Infof("read on disk %v local subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)
		processedTsNs, readPersistedLogErr = fs.filer.ReadPersistedLogBuffer(lastReadTime, eachLogEntryFn)
		if readPersistedLogErr != nil {
			return fmt.Errorf("reading from persisted logs: %v", readPersistedLogErr)
		}

		if processedTsNs != 0 {
			lastReadTime = time.Unix(0, processedTsNs)
		} else {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				time.Sleep(1127 * time.Millisecond)
				continue
			}
		}

		glog.V(0).Infof("read in memory %v local subscribe %s from %+v", clientName, req.PathPrefix, lastReadTime)

		lastReadTime, readInMemoryLogErr = fs.filer.LocalMetaLogBuffer.LoopProcessLogData("localMeta:"+clientName, lastReadTime, func() bool {
			fs.listenersLock.Lock()
			fs.listenersCond.Wait()
			fs.listenersLock.Unlock()
			return true
		}, eachLogEntryFn)
		if readInMemoryLogErr != nil {
			if readInMemoryLogErr == log_buffer.ResumeFromDiskError {
				continue
			}
			glog.Errorf("processed to %v: %v", lastReadTime, readInMemoryLogErr)
			time.Sleep(1127 * time.Millisecond)
			if readInMemoryLogErr != log_buffer.ResumeError {
				break
			}
		}
	}

	return readInMemoryLogErr

}

func eachLogEntryFn(eachEventNotificationFn func(dirPath string, eventNotification *filer_pb.EventNotification, tsNs int64) error) func(logEntry *filer_pb.LogEntry) error {
	return func(logEntry *filer_pb.LogEntry) error {
		event := &filer_pb.SubscribeMetadataResponse{}
		if err := proto.Unmarshal(logEntry.Data, event); err != nil {
			glog.Errorf("unexpected unmarshal filer_pb.SubscribeMetadataResponse: %v", err)
			return fmt.Errorf("unexpected unmarshal filer_pb.SubscribeMetadataResponse: %v", err)
		}

		if err := eachEventNotificationFn(event.Directory, event.EventNotification, event.TsNs); err != nil {
			return err
		}

		return nil
	}
}

func (fs *FilerServer) eachEventNotificationFn(req *filer_pb.SubscribeMetadataRequest, stream filer_pb.SeaweedFiler_SubscribeMetadataServer, clientName string) func(dirPath string, eventNotification *filer_pb.EventNotification, tsNs int64) error {
	filtered := 0

	return func(dirPath string, eventNotification *filer_pb.EventNotification, tsNs int64) error {
		defer func() {
			if filtered > MaxUnsyncedEvents {
				if err := stream.Send(&filer_pb.SubscribeMetadataResponse{
					EventNotification: &filer_pb.EventNotification{},
					TsNs:              tsNs,
				}); err == nil {
					filtered = 0
				}
			}
		}()

		filtered++
		foundSelf := false
		for _, sig := range eventNotification.Signatures {
			if sig == req.Signature && req.Signature != 0 {
				return nil
			}
			if sig == fs.filer.Signature {
				foundSelf = true
			}
		}
		if !foundSelf {
			eventNotification.Signatures = append(eventNotification.Signatures, fs.filer.Signature)
		}

		// get complete path to the file or directory
		var entryName string
		if eventNotification.OldEntry != nil {
			entryName = eventNotification.OldEntry.Name
		} else if eventNotification.NewEntry != nil {
			entryName = eventNotification.NewEntry.Name
		}

		fullpath := util.Join(dirPath, entryName)

		// skip on filer internal meta logs
		if strings.HasPrefix(fullpath, filer.SystemLogDir) {
			return nil
		}

		if !strings.HasPrefix(fullpath, req.PathPrefix) {
			if eventNotification.NewParentPath != "" {
				newFullPath := util.Join(eventNotification.NewParentPath, entryName)
				if !strings.HasPrefix(newFullPath, req.PathPrefix) {
					return nil
				}
			} else {
				return nil
			}
		}

		message := &filer_pb.SubscribeMetadataResponse{
			Directory:         dirPath,
			EventNotification: eventNotification,
			TsNs:              tsNs,
		}
		// println("sending", dirPath, entryName)
		if err := stream.Send(message); err != nil {
			glog.V(0).Infof("=> client %v: %+v", clientName, err)
			return err
		}
		filtered = 0
		return nil
	}
}

func (fs *FilerServer) addClient(clientType string, clientAddress string) (clientName string) {
	clientName = clientType + "@" + clientAddress
	glog.V(0).Infof("+ listener %v", clientName)
	return
}

func (fs *FilerServer) deleteClient(clientName string) {
	glog.V(0).Infof("- listener %v", clientName)
}
<|endoftext|>


<commit_before>package readline

import (
	"bufio"
	"io"
	"os"
)

/*
   Clear Screen: \u001b[{n}J clears the screen
       n=0 clears from cursor until end of screen,
       n=1 clears from cursor to beginning of screen
       n=2 clears entire screen
   Clear Line: \u001b[{n}K clears the current line
       n=0 clears from cursor to end of line
       n=1 clears from cursor to start of line
       n=2 clears entire line
*/

const (
	ESC = 0x1b // \033
)

type inputTerm struct {
	r        *bufio.Reader // for input
	termFd   int           // for setting the terminal input
	oldState *State        // old state of terminal
}

func (it *inputTerm) init() error {
	it.termFd = int(os.Stdin.Fd())
	oldState, err := MakeRaw(it.termFd)
	if err != nil {
		panic(err)
		return err
	}
	it.oldState = oldState

	it.r = bufio.NewReader(os.Stdin)

	return nil
}

func (it *inputTerm) deinit() {
	restoreTerm(it.termFd, it.oldState)
}

func (it *inputTerm) readByte() (byte, error) {
	return it.r.ReadByte()
}

type inputHandler func(*Instance) (byte, bool)

var inputMap = map[byte]inputHandler{
	CharEOF:       eofHandler,
	CharTab:       tabHandler,
	CharEnter:     enterHandler,
	CharBackspace: backspaceHandler,
	CharInterrupt: interruptHandler,
}

func eofHandler(inst *Instance) (byte, bool) {
	if inst.line.Len() == 0 {
		inst.Printf("\ngot EOF(Ctrl+D)\n")
		return CharEOF, true
	}
	inst.Print("\n")
	inst.resetCmdline()
	return CharEOF, false
}

func interruptHandler(inst *Instance) (byte, bool) {
	inst.Printf("\ngot Interrupt(Ctrl+C)\n")
	inst.resetCmdline()
	return CharInterrupt, false
}

func backspaceHandler(inst *Instance) (byte, bool) {
	inst.line.backspace()
	return CharBackspace, false
}

func enterHandler(inst *Instance) (byte, bool) {
	inst.Print("\n")
	end := inst.execute(inst.line.String(), inst.data)
	if !end {
		inst.resetCmdline()
	}
	return CharEnter, end
}

func tabHandler(inst *Instance) (byte, bool) {
	key := byte(CharTab)
	if inst.lastKey != CharTab { // First tab
		_, candidates, end, err := getCandidates(inst)
		if err != nil {
			inst.Log("1st tab error: %v\n", err)
			return key, false
		}
		switch len(candidates) {
		case 0:
			inst.Log("TODO: can not happen\n")
		case 1:
			completeWhole(inst, candidates[0])
			if end {
				inst.lineAdd(' ')
				key = ' '
			}
		default:
			completePartial(inst, candidates)
		}
	} else { // Second Tab
		cp, candidates, _, err := getCandidates(inst)
		if err != nil {
			inst.Log("2nd tab error: %v\n", err)
			return key, false
		}
		switch len(candidates) {
		case 0, 1:
			inst.Log("TODO: can not happen\n")
		default:
			inst.Log("multi candidates\n")
			printCandidates(inst, cp, candidates)
			inst.Printf("%s%s", inst.view.prompt, inst.line.String())
		}
	}

	return key, false
}

func InputLoop(inst *Instance) {
	inst.line.reset()
	inst.view.printPrompt()
	for {
		c, err := inst.input.readByte()
		if err != nil {
			if err == io.EOF {
				inst.Printf("got EOF\n")
			} else {
				inst.Printf("error: %v\n", err)
			}
			break
		}

		//inst.Log("[%d]", c)
		handler, ok := inputMap[c]
		if ok {
			key, end := handler(inst)
			if end {
				break
			}
			inst.lastKey = key
		} else {
			inst.line.insert(c)
			inst.lastKey = c
		}

		inst.view.clearLine()
		inst.view.Print(inst.view.prompt + inst.line.String())
		inst.view.resetCursor(false)
		inst.view.setCursor(inst.line.curPos)
		inst.view.flush()
	}
}
<commit_msg>remove useless comment<commit_after>package readline

import (
	"bufio"
	"io"
	"os"
)

type inputTerm struct {
	r        *bufio.Reader // for input
	termFd   int           // for setting the terminal input
	oldState *State        // old state of terminal
}

func (it *inputTerm) init() error {
	it.termFd = int(os.Stdin.Fd())
	oldState, err := MakeRaw(it.termFd)
	if err != nil {
		panic(err)
		return err
	}
	it.oldState = oldState

	it.r = bufio.NewReader(os.Stdin)

	return nil
}

func (it *inputTerm) deinit() {
	restoreTerm(it.termFd, it.oldState)
}

func (it *inputTerm) readByte() (byte, error) {
	return it.r.ReadByte()
}

type inputHandler func(*Instance) (byte, bool)

var inputMap = map[byte]inputHandler{
	CharEOF:       eofHandler,
	CharTab:       tabHandler,
	CharEnter:     enterHandler,
	CharBackspace: backspaceHandler,
	CharInterrupt: interruptHandler,
}

func eofHandler(inst *Instance) (byte, bool) {
	if inst.line.Len() == 0 {
		inst.Printf("\ngot EOF(Ctrl+D)\n")
		return CharEOF, true
	}
	inst.Print("\n")
	inst.resetCmdline()
	return CharEOF, false
}

func interruptHandler(inst *Instance) (byte, bool) {
	inst.Printf("\ngot Interrupt(Ctrl+C)\n")
	inst.resetCmdline()
	return CharInterrupt, false
}

func backspaceHandler(inst *Instance) (byte, bool) {
	inst.line.backspace()
	return CharBackspace, false
}

func enterHandler(inst *Instance) (byte, bool) {
	inst.Print("\n")
	end := inst.execute(inst.line.String(), inst.data)
	if !end {
		inst.resetCmdline()
	}
	return CharEnter, end
}

func tabHandler(inst *Instance) (byte, bool) {
	key := byte(CharTab)
	if inst.lastKey != CharTab { // First tab
		_, candidates, end, err := getCandidates(inst)
		if err != nil {
			inst.Log("1st tab error: %v\n", err)
			return key, false
		}
		switch len(candidates) {
		case 0:
			inst.Log("TODO: can not happen\n")
		case 1:
			completeWhole(inst, candidates[0])
			if end {
				inst.lineAdd(' ')
				key = ' '
			}
		default:
			completePartial(inst, candidates)
		}
	} else { // Second Tab
		cp, candidates, _, err := getCandidates(inst)
		if err != nil {
			inst.Log("2nd tab error: %v\n", err)
			return key, false
		}
		switch len(candidates) {
		case 0, 1:
			inst.Log("TODO: can not happen\n")
		default:
			inst.Log("multi candidates\n")
			printCandidates(inst, cp, candidates)
			inst.Printf("%s%s", inst.view.prompt, inst.line.String())
		}
	}

	return key, false
}

func InputLoop(inst *Instance) {
	inst.line.reset()
	inst.view.printPrompt()
	for {
		c, err := inst.input.readByte()
		if err != nil {
			if err == io.EOF {
				inst.Printf("got EOF\n")
			} else {
				inst.Printf("error: %v\n", err)
			}
			break
		}

		//inst.Log("[%d]", c)
		handler, ok := inputMap[c]
		if ok {
			key, end := handler(inst)
			if end {
				break
			}
			inst.lastKey = key
		} else {
			inst.line.insert(c)
			inst.lastKey = c
		}

		inst.view.clearLine()
		inst.view.Print(inst.view.prompt + inst.line.String())
		inst.view.resetCursor(false)
		inst.view.setCursor(inst.line.curPos)
		inst.view.flush()
	}
}
<|endoftext|>