<commit_before>package main

import (
	"errors"
	"fmt" //for debugging //TODO Remove
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/AIS/storage" // use local import?
	ais "github.com/andmarios/aislib"
)

//The Archive stores the information about the ships (and works as a temp. solution for the RTree concurrency)
type Archive struct {
	rt *storage.RTree //Stores the points
	rw *sync.RWMutex  //works as a lock for the RTree (#TODO: RTree should be improved to handle concurrency on its own)

	si *storage.ShipInfo //Contains tracklog and other info for each ship
}

//Returns a pointer to the new Archive
func NewArchive() *Archive {
	return &Archive{
		rt: storage.NewRTree(),
		rw: &sync.RWMutex{},
		si: storage.NewShipInfo(),
	}
}

// Stores the information recieved form the channel
func (a *Archive) Save(msg chan *Message) {
	counter := 0 //TODO Remove
	for {
		select {
		case m := <-msg:
			var err error
			ps := (*ais.PositionReport)(nil)
			switch m.Type {
			case 1, 2, 3: // class A position report (longest)
				cApr, e := ais.DecodeClassAPositionReport(m.ArmoredPayload())
				ps = &cApr.PositionReport
				if e != nil {
					continue //TODO
				}
				err = a.updatePos(ps)
			case 5: // static voyage data
				svd, e := ais.DecodeStaticVoyageData(m.ArmoredPayload())
				if e != nil && svd.MMSI <= 0 {
					continue //TODO
				}
				err = a.si.UpdateSVD(svd.MMSI, svd.Callsign, svd.Destination, svd.VesselName, svd.ToBow, svd.ToStern)
			case 18: // basic class B position report (shorter)
				cBpr, e := ais.DecodeClassBPositionReport(m.ArmoredPayload())
				ps = &cBpr.PositionReport
				if e != nil {
					continue //TODO
				}
				err = a.updatePos(ps)
			}
			if err != nil {
				//fmt.Printf("Had an error saving to Archive... %v\n", err)
				continue //TODO do something...
			}
			counter++              //TODO Remove
			if counter%1000 == 0 { //TODO Remove
				fmt.Printf("Number of boats: %d\n", a.rt.NumOfBoats())
				fmt.Println(a.FindWithin(59.0, 5.54, 59.15, 5.8))
				//fmt.Println(a.FindAll())
			}
		}
	}
}

//Updates the ships position in the structures (message type 1,2,3,18)
func (a *Archive) updatePos(ps *ais.PositionReport) error {
	mmsi := ps.MMSI
	if !okCoords(ps.Lat, ps.Lon) || mmsi <= 0 { //This happends quite frequently (coordinates are set to 91,181)
		return errors.New(fmt.Sprintf("Cannot update position... MMSI: %d, lat: %f, long %f", mmsi, ps.Lat, ps.Lon))
	}
	//Check if it is a known ship
	if a.si.IsKnown(mmsi) {
		oldLat, oldLong := a.si.GetCoords(mmsi) //get the previous coordinates
		a.rw.Lock()
		a.rt.Update(mmsi, oldLat, oldLong, ps.Lat, ps.Lon) //update the position in the R*Tree
		a.rw.Unlock()
	} else {
		a.rw.Lock()
		a.rt.InsertData(ps.Lat, ps.Lon, mmsi) //insert a new ship into the R*Tree
		a.rw.Unlock()
	} //TODO check for error?
	err := a.si.AddCheckpoint(ps.MMSI, ps.Lat, ps.Lon, time.Now(), ps.Heading) //Adds the position to the ships tracklog
	return err
}

// Returns a GeoJSON FeatureCollection containing all the known ships
func (a *Archive) FindAll() string {
	a.rw.RLock()
	geoJsonFC := a.FindWithin(-79.999999, -179.999999, 79.999999, 179.999999)
	a.rw.RUnlock()
	return geoJsonFC
}

/*
Public func for finding all known boats that overlaps a given rectangle of the map [13], [14]
	input:
		minLatitude, minLongitude, maxLatitude, maxLongitude	float64
	output:
		string	-	All matching ships in GeoJSON FeatureCollection

*/
func (a *Archive) FindWithin(minLat, minLong, maxLat, maxLong float64) string {
	r, err := storage.NewRectangle(minLat, minLong, maxLat, maxLong)
	if err != nil {
		return "ERROR, invalid rectangle coordinates"
	}
	a.rw.RLock()
	matchingShips := a.rt.FindWithin(r)
	a.rw.RUnlock()
	features := []string{}
	var name string
	var length, heading uint16
	for _, s := range *matchingShips {
		name, length, heading = a.si.GetFeatures(s.MMSI)
		f := `{
				"type": "Feature",
				"id": ` + strconv.Itoa(int(s.MMSI)) + `,
				"geometry": {
					"type": "Point",
					"coordinates": ` + "[" + strconv.FormatFloat(s.Long, 'f', 6, 64) + ", " + strconv.FormatFloat(s.Lat, 'f', 6, 64) + "]" + `},
				"properties": {
					"name": "` + name + `" ,
					"length": ` + strconv.Itoa(int(length)) + `,
					"heading": ` + strconv.Itoa(int(heading)) + `
				}

			}`
		features = append(features, f)
	}
	return "{ \"type\": \"FeatureCollection\", \"features\": [" + strings.Join(features, ", ") + "]}"
}

// Check if the coordinates are ok.	(<91, 181> seems to be a fallback value for the coordinates)
func okCoords(lat, long float64) bool {
	if lat <= 90 && long <= 180 && lat >= -90 && long >= -180 {
		return true
	}
	return false
}

/*
TODO:
	- Fix rStarTree so that it handles concurrency by itself?
		- Archive controls the concurrency of the RTree at the moment...
				- need not be much point using a RWMutex for the rtree... there are a lot more writes than reads atm ... could use a normal mutex, and thereby save some overhead..
		- This could be improved in the future by modifying the rtree structure

References:
	[1]	http://geojsonlint.com/
	[2]	http://stackoverflow.com/questions/7933460/how-do-you-write-multiline-strings-in-go#7933487
*/
<commit_msg>import path<commit_after>package main

import (
	"errors"
	"fmt" //for debugging //TODO Remove
	"strconv"
	"strings"
	"sync"
	"time"

	ais "github.com/andmarios/aislib"
	"github.com/tormol/AIS/storage"
)

//The Archive stores the information about the ships (and works as a temp. solution for the RTree concurrency)
type Archive struct {
	rt *storage.RTree //Stores the points
	rw *sync.RWMutex  //works as a lock for the RTree (#TODO: RTree should be improved to handle concurrency on its own)

	si *storage.ShipInfo //Contains tracklog and other info for each ship
}

//Returns a pointer to the new Archive
func NewArchive() *Archive {
	return &Archive{
		rt: storage.NewRTree(),
		rw: &sync.RWMutex{},
		si: storage.NewShipInfo(),
	}
}

// Stores the information recieved form the channel
func (a *Archive) Save(msg chan *Message) {
	counter := 0 //TODO Remove
	for {
		select {
		case m := <-msg:
			var err error
			ps := (*ais.PositionReport)(nil)
			switch m.Type {
			case 1, 2, 3: // class A position report (longest)
				cApr, e := ais.DecodeClassAPositionReport(m.ArmoredPayload())
				ps = &cApr.PositionReport
				if e != nil {
					continue //TODO
				}
				err = a.updatePos(ps)
			case 5: // static voyage data
				svd, e := ais.DecodeStaticVoyageData(m.ArmoredPayload())
				if e != nil && svd.MMSI <= 0 {
					continue //TODO
				}
				err = a.si.UpdateSVD(svd.MMSI, svd.Callsign, svd.Destination, svd.VesselName, svd.ToBow, svd.ToStern)
			case 18: // basic class B position report (shorter)
				cBpr, e := ais.DecodeClassBPositionReport(m.ArmoredPayload())
				ps = &cBpr.PositionReport
				if e != nil {
					continue //TODO
				}
				err = a.updatePos(ps)
			}
			if err != nil {
				//fmt.Printf("Had an error saving to Archive... %v\n", err)
				continue //TODO do something...
			}
			counter++              //TODO Remove
			if counter%1000 == 0 { //TODO Remove
				fmt.Printf("Number of boats: %d\n", a.rt.NumOfBoats())
				fmt.Println(a.FindWithin(59.0, 5.54, 59.15, 5.8))
				//fmt.Println(a.FindAll())
			}
		}
	}
}

//Updates the ships position in the structures (message type 1,2,3,18)
func (a *Archive) updatePos(ps *ais.PositionReport) error {
	mmsi := ps.MMSI
	if !okCoords(ps.Lat, ps.Lon) || mmsi <= 0 { //This happends quite frequently (coordinates are set to 91,181)
		return errors.New(fmt.Sprintf("Cannot update position... MMSI: %d, lat: %f, long %f", mmsi, ps.Lat, ps.Lon))
	}
	//Check if it is a known ship
	if a.si.IsKnown(mmsi) {
		oldLat, oldLong := a.si.GetCoords(mmsi) //get the previous coordinates
		a.rw.Lock()
		a.rt.Update(mmsi, oldLat, oldLong, ps.Lat, ps.Lon) //update the position in the R*Tree
		a.rw.Unlock()
	} else {
		a.rw.Lock()
		a.rt.InsertData(ps.Lat, ps.Lon, mmsi) //insert a new ship into the R*Tree
		a.rw.Unlock()
	} //TODO check for error?
	err := a.si.AddCheckpoint(ps.MMSI, ps.Lat, ps.Lon, time.Now(), ps.Heading) //Adds the position to the ships tracklog
	return err
}

// Returns a GeoJSON FeatureCollection containing all the known ships
func (a *Archive) FindAll() string {
	a.rw.RLock()
	geoJsonFC := a.FindWithin(-79.999999, -179.999999, 79.999999, 179.999999)
	a.rw.RUnlock()
	return geoJsonFC
}

/*
Public func for finding all known boats that overlaps a given rectangle of the map [13], [14]
	input:
		minLatitude, minLongitude, maxLatitude, maxLongitude	float64
	output:
		string	-	All matching ships in GeoJSON FeatureCollection

*/
func (a *Archive) FindWithin(minLat, minLong, maxLat, maxLong float64) string {
	r, err := storage.NewRectangle(minLat, minLong, maxLat, maxLong)
	if err != nil {
		return "ERROR, invalid rectangle coordinates"
	}
	a.rw.RLock()
	matchingShips := a.rt.FindWithin(r)
	a.rw.RUnlock()
	features := []string{}
	var name string
	var length, heading uint16
	for _, s := range *matchingShips {
		name, length, heading = a.si.GetFeatures(s.MMSI)
		f := `{
				"type": "Feature",
				"id": ` + strconv.Itoa(int(s.MMSI)) + `,
				"geometry": {
					"type": "Point",
					"coordinates": ` + "[" + strconv.FormatFloat(s.Long, 'f', 6, 64) + ", " + strconv.FormatFloat(s.Lat, 'f', 6, 64) + "]" + `},
				"properties": {
					"name": "` + name + `" ,
					"length": ` + strconv.Itoa(int(length)) + `,
					"heading": ` + strconv.Itoa(int(heading)) + `
				}

			}`
		features = append(features, f)
	}
	return "{ \"type\": \"FeatureCollection\", \"features\": [" + strings.Join(features, ", ") + "]}"
}

// Check if the coordinates are ok.	(<91, 181> seems to be a fallback value for the coordinates)
func okCoords(lat, long float64) bool {
	if lat <= 90 && long <= 180 && lat >= -90 && long >= -180 {
		return true
	}
	return false
}

/*
TODO:
	- Fix rStarTree so that it handles concurrency by itself?
		- Archive controls the concurrency of the RTree at the moment...
				- need not be much point using a RWMutex for the rtree... there are a lot more writes than reads atm ... could use a normal mutex, and thereby save some overhead..
		- This could be improved in the future by modifying the rtree structure

References:
	[1]	http://geojsonlint.com/
	[2]	http://stackoverflow.com/questions/7933460/how-do-you-write-multiline-strings-in-go#7933487
*/
<|endoftext|>


<commit_before>package evedata

import (
	"encoding/gob"
	"os"

	"log"
	"net/http"
	"time"

	"github.com/antihax/evedata/appContext"
	"github.com/antihax/evedata/config"
	"github.com/antihax/evedata/discord"
	"github.com/antihax/evedata/esi"
	"github.com/antihax/evedata/eveConsumer"
	"github.com/antihax/evedata/eveapi"
	"github.com/antihax/evedata/models"

	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/context"
	"github.com/gregjones/httpcache"
	httpredis "github.com/gregjones/httpcache/redis"
	"golang.org/x/oauth2"
	gsr "gopkg.in/boj/redistore.v1"
)

var ctx appContext.AppContext

func GetContext() *appContext.AppContext {
	return &ctx
}

// Load the server
func GoServer() {
	var err error

	// Make a new app context.

	// Read configuation.
	if ctx.Conf, err = config.ReadConfig(); err != nil {
		log.Fatalf("Error reading configuration: %v", err)
	}

	// Build the redis pool
	ctx.Cache = &redis.Pool{
		MaxIdle:     50,
		MaxActive:   0,
		Wait:        false,
		IdleTimeout: 90 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", ctx.Conf.Redis.Address)
			if err != nil {
				return nil, err
			}
			if ctx.Conf.Redis.Password != "" {
				if _, err := c.Do("AUTH", ctx.Conf.Redis.Password); err != nil {
					c.Close()
					return nil, err
				}
			}
			return c, err
		},
	}

	// Create a Redis http client for the CCP APIs.
	transportCache := httpcache.NewTransport(httpredis.NewWithClient(ctx.Cache.Get()))

	// Attach a basic transport with our chained custom transport.
	transportCache.Transport = &transport{&http.Transport{}, &ctx}

	// Build a HTTP Client pool this client will be shared with APIs for:
	//   - ESI
	//   - ZKillboard
	//   - EVE SSO
	//   - EVE CREST and XML

	ctx.HTTPClient = &http.Client{Transport: transportCache}
	if ctx.HTTPClient == nil {
		panic("http client is null")
	}

	// Build Connection Pool
	if ctx.Db, err = models.SetupDatabase(ctx.Conf.Database.Driver, ctx.Conf.Database.Spec); err != nil {
		log.Fatalf("Cannot build database pool: %v", err)
	}

	if len(os.Args) > 1 {
		if os.Args[1] == "dumpdb" {
			err := models.DumpDatabase("./sql/evedata.sql", "evedata")
			if err != nil {
				log.Fatalln(err)
			}
		}
	}

	/*r := ctx.Cache.Get()
	r.Do("FLUSHALL")
	r.Close()
	models.MaintKillMails()*/

	// Setup the SSO authenticator, this is the main login.
	ssoScopes := []string{}

	ctx.SSOAuthenticator = eveapi.NewSSOAuthenticator(
		ctx.HTTPClient,
		ctx.Conf.CREST.SSO.ClientID,
		ctx.Conf.CREST.SSO.SecretKey,
		ctx.Conf.CREST.SSO.RedirectURL,
		ssoScopes)

	// Setup the Token authenticator, this handles sub characters.
	tokenScopes := []string{
		eveapi.ScopeCharacterContractsRead,
		eveapi.ScopeCharacterMarketOrdersRead,
		eveapi.ScopeCharacterResearchRead,
		eveapi.ScopeCharacterWalletRead,
		"esi-assets.read_assets.v1",
		"esi-characters.read_contacts.v1",
		"esi-characters.write_contacts.v1",
		"esi-corporations.read_corporation_membership.v1",
		"esi-location.read_location.v1",
		"esi-location.read_ship_type.v1",
		"esi-planets.manage_planets.v1",
		"esi-search.search_structures.v1",
		"esi-skills.read_skills.v1",
		"esi-ui.open_window.v1",
		"esi-ui.write_waypoint.v1",
		"esi-universe.read_structures.v1",
		"esi-wallet.read_character_wallet.v1",
	}

	ctx.TokenAuthenticator = eveapi.NewSSOAuthenticator(
		ctx.HTTPClient,
		ctx.Conf.CREST.Token.ClientID,
		ctx.Conf.CREST.Token.SecretKey,
		ctx.Conf.CREST.Token.RedirectURL,
		tokenScopes)

	// Setup the EVE ESI Client
	ctx.ESI = esi.NewAPIClient(ctx.HTTPClient, ctx.Conf.UserAgent)

	// Setup the bootstrap authenticator. Needed to update the site main token.
	ctx.ESIBootstrapAuthenticator = eveapi.NewSSOAuthenticator(
		ctx.HTTPClient,
		ctx.Conf.CREST.ESIAccessToken.ClientID,
		ctx.Conf.CREST.ESIAccessToken.SecretKey,
		ctx.Conf.CREST.ESIAccessToken.RedirectURL,
		[]string{"esi-universe.read_structures.v1",
			"esi-search.search_structures.v1",
			"esi-markets.structure_markets.v1"})

	// Get the token from config and build a TokenSource (refreshes the token if needed.)
	token := &eveapi.CRESTToken{
		AccessToken:  ctx.Conf.CREST.ESIAccessToken.AccessToken,
		TokenType:    ctx.Conf.CREST.ESIAccessToken.TokenType,
		RefreshToken: ctx.Conf.CREST.ESIAccessToken.RefreshToken,
		Expiry:       ctx.Conf.CREST.ESIAccessToken.Expiry,
	}

	ctx.ESIPublicToken, err = ctx.ESIBootstrapAuthenticator.TokenSource(token)
	if err != nil {
		log.Fatalf("Error starting bootstrap ESI client: %v", err)
	}

	// Create a memcached session store.
	ctx.Store, err = gsr.NewRediStoreWithPool(ctx.Cache, []byte(ctx.Conf.Store.Key))
	if err != nil {
		log.Fatalf("Cannot build database pool: %v", err)
	}

	ctx.Store.Options.Domain = ctx.Conf.Store.Domain

	// Register structs for storage
	gob.Register(oauth2.Token{})
	gob.Register(eveapi.CRESTToken{})
	gob.Register(eveapi.VerifyResponse{})

	// Anonymous EVE API & Crest Client
	ctx.EVE = eveapi.NewEVEAPIClient(ctx.HTTPClient)

	// Set our logging flags
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	if ctx.Conf.Discord.Enabled {
		go discord.GoDiscordBot(&ctx)
	}

	if ctx.Conf.EVEConsumer.Enabled {
		log.Println("Starting EVE Consumer")
		eC := eveConsumer.NewEVEConsumer(&ctx)
		eC.RunConsumer()
		defer eC.StopConsumer()
	}
	// Allocate the routes
	rtr := NewRouter(&ctx)

	log.Printf("EveData Listening port 3000...\n")
	http.ListenAndServe(":3000", context.ClearHandler(rtr))
	log.Printf("EveData Quitting..\n")
}
<commit_msg>optimize<commit_after>package evedata

import (
	"encoding/gob"
	"net"
	"os"

	"log"
	"net/http"
	"time"

	"github.com/antihax/evedata/appContext"
	"github.com/antihax/evedata/config"
	"github.com/antihax/evedata/discord"
	"github.com/antihax/evedata/esi"
	"github.com/antihax/evedata/eveConsumer"
	"github.com/antihax/evedata/eveapi"
	"github.com/antihax/evedata/models"

	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/context"
	"github.com/gregjones/httpcache"
	httpredis "github.com/gregjones/httpcache/redis"
	"golang.org/x/oauth2"
	gsr "gopkg.in/boj/redistore.v1"
)

var ctx appContext.AppContext

func GetContext() *appContext.AppContext {
	return &ctx
}

// Load the server
func GoServer() {
	var err error

	// Make a new app context.

	// Read configuation.
	if ctx.Conf, err = config.ReadConfig(); err != nil {
		log.Fatalf("Error reading configuration: %v", err)
	}

	// Build the redis pool
	ctx.Cache = &redis.Pool{
		MaxIdle:     50,
		MaxActive:   0,
		Wait:        false,
		IdleTimeout: 90 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", ctx.Conf.Redis.Address)
			if err != nil {
				return nil, err
			}
			if ctx.Conf.Redis.Password != "" {
				if _, err := c.Do("AUTH", ctx.Conf.Redis.Password); err != nil {
					c.Close()
					return nil, err
				}
			}
			return c, err
		},
	}

	// Create a Redis http client for the CCP APIs.
	transportCache := httpcache.NewTransport(httpredis.NewWithClient(ctx.Cache.Get()))

	// Attach a basic transport with our chained custom transport.
	transportCache.Transport = &transport{&http.Transport{
		MaxIdleConns: 0,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		IdleConnTimeout:       60 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 0,
		MaxIdleConnsPerHost:   100,
	}, &ctx}

	// Build a HTTP Client pool this client will be shared with APIs for:
	//   - ESI
	//   - ZKillboard
	//   - EVE SSO
	//   - EVE CREST and XML

	ctx.HTTPClient = &http.Client{Transport: transportCache}
	if ctx.HTTPClient == nil {
		panic("http client is null")
	}

	// Build Connection Pool
	if ctx.Db, err = models.SetupDatabase(ctx.Conf.Database.Driver, ctx.Conf.Database.Spec); err != nil {
		log.Fatalf("Cannot build database pool: %v", err)
	}

	if len(os.Args) > 1 {
		if os.Args[1] == "dumpdb" {
			err := models.DumpDatabase("./sql/evedata.sql", "evedata")
			if err != nil {
				log.Fatalln(err)
			}
		}
	}

	/*r := ctx.Cache.Get()
	r.Do("FLUSHALL")
	r.Close()
	models.MaintKillMails()*/

	// Setup the SSO authenticator, this is the main login.
	ssoScopes := []string{}

	ctx.SSOAuthenticator = eveapi.NewSSOAuthenticator(
		ctx.HTTPClient,
		ctx.Conf.CREST.SSO.ClientID,
		ctx.Conf.CREST.SSO.SecretKey,
		ctx.Conf.CREST.SSO.RedirectURL,
		ssoScopes)

	// Setup the Token authenticator, this handles sub characters.
	tokenScopes := []string{
		eveapi.ScopeCharacterContractsRead,
		eveapi.ScopeCharacterMarketOrdersRead,
		eveapi.ScopeCharacterResearchRead,
		eveapi.ScopeCharacterWalletRead,
		"esi-assets.read_assets.v1",
		"esi-characters.read_contacts.v1",
		"esi-characters.write_contacts.v1",
		"esi-corporations.read_corporation_membership.v1",
		"esi-location.read_location.v1",
		"esi-location.read_ship_type.v1",
		"esi-planets.manage_planets.v1",
		"esi-search.search_structures.v1",
		"esi-skills.read_skills.v1",
		"esi-ui.open_window.v1",
		"esi-ui.write_waypoint.v1",
		"esi-universe.read_structures.v1",
		"esi-wallet.read_character_wallet.v1",
	}

	ctx.TokenAuthenticator = eveapi.NewSSOAuthenticator(
		ctx.HTTPClient,
		ctx.Conf.CREST.Token.ClientID,
		ctx.Conf.CREST.Token.SecretKey,
		ctx.Conf.CREST.Token.RedirectURL,
		tokenScopes)

	// Setup the EVE ESI Client
	ctx.ESI = esi.NewAPIClient(ctx.HTTPClient, ctx.Conf.UserAgent)

	// Setup the bootstrap authenticator. Needed to update the site main token.
	ctx.ESIBootstrapAuthenticator = eveapi.NewSSOAuthenticator(
		ctx.HTTPClient,
		ctx.Conf.CREST.ESIAccessToken.ClientID,
		ctx.Conf.CREST.ESIAccessToken.SecretKey,
		ctx.Conf.CREST.ESIAccessToken.RedirectURL,
		[]string{"esi-universe.read_structures.v1",
			"esi-search.search_structures.v1",
			"esi-markets.structure_markets.v1"})

	// Get the token from config and build a TokenSource (refreshes the token if needed.)
	token := &eveapi.CRESTToken{
		AccessToken:  ctx.Conf.CREST.ESIAccessToken.AccessToken,
		TokenType:    ctx.Conf.CREST.ESIAccessToken.TokenType,
		RefreshToken: ctx.Conf.CREST.ESIAccessToken.RefreshToken,
		Expiry:       ctx.Conf.CREST.ESIAccessToken.Expiry,
	}

	ctx.ESIPublicToken, err = ctx.ESIBootstrapAuthenticator.TokenSource(token)
	if err != nil {
		log.Fatalf("Error starting bootstrap ESI client: %v", err)
	}

	// Create a memcached session store.
	ctx.Store, err = gsr.NewRediStoreWithPool(ctx.Cache, []byte(ctx.Conf.Store.Key))
	if err != nil {
		log.Fatalf("Cannot build database pool: %v", err)
	}

	ctx.Store.Options.Domain = ctx.Conf.Store.Domain

	// Register structs for storage
	gob.Register(oauth2.Token{})
	gob.Register(eveapi.CRESTToken{})
	gob.Register(eveapi.VerifyResponse{})

	// Anonymous EVE API & Crest Client
	ctx.EVE = eveapi.NewEVEAPIClient(ctx.HTTPClient)

	// Set our logging flags
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	if ctx.Conf.Discord.Enabled {
		go discord.GoDiscordBot(&ctx)
	}

	if ctx.Conf.EVEConsumer.Enabled {
		log.Println("Starting EVE Consumer")
		eC := eveConsumer.NewEVEConsumer(&ctx)
		eC.RunConsumer()
		defer eC.StopConsumer()
	}
	// Allocate the routes
	rtr := NewRouter(&ctx)

	log.Printf("EveData Listening port 3000...\n")
	http.ListenAndServe(":3000", context.ClearHandler(rtr))
	log.Printf("EveData Quitting..\n")
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/jingweno/jqplay/config"
	"github.com/jingweno/jqplay/jq"
	"golang.org/x/net/context"
	"gopkg.in/gin-gonic/gin.v1"
)

const (
	JSONPayloadLimit   = JSONPayloadLimitMB * OneMB
	JSONPayloadLimitMB = 10
	OneMB              = 1024000
	JQExecTimeout      = 15 * time.Second
)

type JQHandlerContext struct {
	*config.Config
	JQ string
}

func (c *JQHandlerContext) Asset(path string) string {
	return fmt.Sprintf("%s/%s", c.AssetHost, path)
}

func (c *JQHandlerContext) ShouldInitJQ() bool {
	return c.JQ != ""
}

type JQHandler struct {
	DB     *DB
	Config *config.Config
}

func (h *JQHandler) handleIndex(c *gin.Context) {
	c.HTML(200, "index.tmpl", &JQHandlerContext{Config: h.Config})
}

func (h *JQHandler) checkReqSize(c *gin.Context) error {
	if c.Request.ContentLength > JSONPayloadLimit {
		size := float64(c.Request.ContentLength) / OneMB
		return fmt.Errorf("JSON payload size is %.1fMB, larger than limit %dMB.", size, JSONPayloadLimitMB)
	}

	return nil
}

func (h *JQHandler) handleJqPost(c *gin.Context) {
	if err := h.checkReqSize(c); err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusExpectationFailed, err.Error())
		return
	}

	ctx, cancel := context.WithTimeout(c.Request.Context(), JQExecTimeout)
	defer cancel()

	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, JSONPayloadLimit)

	var jq *jq.JQ
	err := c.BindJSON(&jq)
	if err != nil {
		err = fmt.Errorf("error parsing JSON: %s", err)
		h.logger(c).WithError(err).Infof("error parsing JSON")
		c.String(422, err.Error())
		return
	}

	// Evaling into ResponseWriter sets the status code to 200
	// appending error message in the end if there's any
	err = jq.Eval(ctx, c.Writer)
	if err != nil {
		h.logger(c).WithError(err).WithField("j", jq.J).WithField("q", jq.Q).Infof("jq error")
		fmt.Fprint(c.Writer, err.Error())
	}
}

func (h *JQHandler) handleJqGet(c *gin.Context) {
	jq := &jq.JQ{
		J: c.Query("j"),
		Q: c.Query("q"),
	}

	var jqData string
	if err := jq.Validate(); err == nil {
		d, err := json.Marshal(jq)
		if err == nil {
			jqData = string(d)
		}
	}

	c.HTML(200, "index.tmpl", &JQHandlerContext{Config: h.Config, JQ: jqData})
}

func (h *JQHandler) handleJqSharePost(c *gin.Context) {
	if err := h.checkReqSize(c); err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusExpectationFailed, err.Error())
		return
	}

	var jq *jq.JQ
	if err := c.BindJSON(&jq); err != nil {
		err = fmt.Errorf("error parsing JSON: %s", err)
		h.logger(c).WithError(err)
		c.String(http.StatusUnprocessableEntity, err.Error())
		return
	}

	if err := jq.Validate(); err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusUnprocessableEntity, err.Error())
		return
	}

	id, err := h.DB.UpsertSnippet(FromJQ(jq))
	if err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusUnprocessableEntity, "error sharing snippet")
		return
	}

	c.String(http.StatusCreated, id)
}

func (h *JQHandler) handleJqShareGet(c *gin.Context) {
	id := c.Param("id")

	s, err := h.DB.GetSnippet(id)
	if err != nil {
		h.logger(c).WithError(err).Info("error getting snippet")
		c.Redirect(http.StatusFound, "/")
		return
	}

	var jqData string
	d, err := json.Marshal(ToJQ(s))
	if err == nil {
		jqData = string(d)
	}

	c.HTML(200, "index.tmpl", &JQHandlerContext{
		Config: h.Config,
		JQ:     jqData,
	})
}

func (h *JQHandler) logger(c *gin.Context) *logrus.Entry {
	l, _ := c.Get("logger")
	return l.(*logrus.Entry)
}
<commit_msg>Check body<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/jingweno/jqplay/config"
	"github.com/jingweno/jqplay/jq"
	"golang.org/x/net/context"
	"gopkg.in/gin-gonic/gin.v1"
)

const (
	JSONPayloadLimit   = JSONPayloadLimitMB * OneMB
	JSONPayloadLimitMB = 10
	OneMB              = 1024000
	JQExecTimeout      = 15 * time.Second
)

type JQHandlerContext struct {
	*config.Config
	JQ string
}

func (c *JQHandlerContext) Asset(path string) string {
	return fmt.Sprintf("%s/%s", c.AssetHost, path)
}

func (c *JQHandlerContext) ShouldInitJQ() bool {
	return c.JQ != ""
}

type JQHandler struct {
	DB     *DB
	Config *config.Config
}

func (h *JQHandler) handleIndex(c *gin.Context) {
	c.HTML(200, "index.tmpl", &JQHandlerContext{Config: h.Config})
}

func (h *JQHandler) checkReqSize(c *gin.Context) error {
	if c.Request.ContentLength > JSONPayloadLimit {
		size := float64(c.Request.ContentLength) / OneMB
		return fmt.Errorf("JSON payload size is %.1fMB, larger than limit %dMB.", size, JSONPayloadLimitMB)
	}

	return nil
}

func (h *JQHandler) handleJqPost(c *gin.Context) {
	if err := h.checkReqSize(c); err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusExpectationFailed, err.Error())
		return
	}

	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, JSONPayloadLimit)

	var jq *jq.JQ
	err := c.BindJSON(&jq)
	if err != nil {
		err = fmt.Errorf("error parsing JSON: %s", err)
		h.logger(c).WithError(err).Infof("error parsing JSON")
		c.String(http.StatusUnprocessableEntity, err.Error())
		return
	}

	ctx, cancel := context.WithTimeout(c.Request.Context(), JQExecTimeout)
	defer cancel()

	// Evaling into ResponseWriter sets the status code to 200
	// appending error message in the end if there's any
	err = jq.Eval(ctx, c.Writer)
	if err != nil {
		h.logger(c).WithError(err).WithField("j", jq.J).WithField("q", jq.Q).Infof("jq error")
		fmt.Fprint(c.Writer, err.Error())
	}
}

func (h *JQHandler) handleJqGet(c *gin.Context) {
	jq := &jq.JQ{
		J: c.Query("j"),
		Q: c.Query("q"),
	}

	var jqData string
	if err := jq.Validate(); err == nil {
		d, err := json.Marshal(jq)
		if err == nil {
			jqData = string(d)
		}
	}

	c.HTML(http.StatusOK, "index.tmpl", &JQHandlerContext{Config: h.Config, JQ: jqData})
}

func (h *JQHandler) handleJqSharePost(c *gin.Context) {
	if err := h.checkReqSize(c); err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusExpectationFailed, err.Error())
		return
	}

	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, JSONPayloadLimit)

	var jq *jq.JQ
	if err := c.BindJSON(&jq); err != nil {
		err = fmt.Errorf("error parsing JSON: %s", err)
		h.logger(c).WithError(err)
		c.String(http.StatusUnprocessableEntity, err.Error())
		return
	}

	if err := jq.Validate(); err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusUnprocessableEntity, err.Error())
		return
	}

	id, err := h.DB.UpsertSnippet(FromJQ(jq))
	if err != nil {
		h.logger(c).WithError(err)
		c.String(http.StatusUnprocessableEntity, "error sharing snippet")
		return
	}

	c.String(http.StatusCreated, id)
}

func (h *JQHandler) handleJqShareGet(c *gin.Context) {
	id := c.Param("id")

	s, err := h.DB.GetSnippet(id)
	if err != nil {
		h.logger(c).WithError(err).Info("error getting snippet")
		c.Redirect(http.StatusFound, "/")
		return
	}

	var jqData string
	d, err := json.Marshal(ToJQ(s))
	if err == nil {
		jqData = string(d)
	}

	c.HTML(200, "index.tmpl", &JQHandlerContext{
		Config: h.Config,
		JQ:     jqData,
	})
}

func (h *JQHandler) logger(c *gin.Context) *logrus.Entry {
	l, _ := c.Get("logger")
	return l.(*logrus.Entry)
}
<|endoftext|>


<commit_before>package server

import (
	"bufio"
	"bytes"
	"context"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/docker/client"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/api"
	"github.com/pterodactyl/wings/config"
	"go.uber.org/zap"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sync"
)

// Executes the installation stack for a server process. Bubbles any errors up to the calling
// function which should handle contacting the panel to notify it of the server state.
func (s *Server) Install() error {
	err := s.internalInstall()

	zap.S().Debugw("notifying panel of server install state", zap.String("server", s.Uuid))
	if serr := s.SyncInstallState(err == nil); serr != nil {
		zap.S().Warnw(
			"failed to notify panel of server install state",
			zap.String("server", s.Uuid),
			zap.Bool("was_successful", err == nil),
			zap.Error(serr),
		)
	}

	return err
}

// Reinstalls a server's software by utilizing the install script for the server egg. This
// does not touch any existing files for the server, other than what the script modifies.
func (s *Server) Reinstall() error {
	if s.GetState() != ProcessOfflineState {
		zap.S().Debugw("waiting for server instance to enter a stopped state", zap.String("server", s.Uuid))
		if err := s.Environment.WaitForStop(10, true); err != nil {
			return err
		}
	}

	return s.Install()
}

// Internal installation function used to simplify reporting back to the Panel.
func (s *Server) internalInstall() error {
	script, rerr, err := api.NewRequester().GetInstallationScript(s.Uuid)
	if err != nil || rerr != nil {
		if err != nil {
			return err
		}

		return errors.New(rerr.String())
	}

	p, err := NewInstallationProcess(s, &script)
	if err != nil {
		return errors.WithStack(err)
	}

	zap.S().Infow("beginning installation process for server", zap.String("server", s.Uuid))

	if err := p.Run(); err != nil {
		return err
	}

	zap.S().Infow("completed installation process for server", zap.String("server", s.Uuid))

	return nil
}

type InstallationProcess struct {
	Server *Server
	Script *api.InstallationScript

	client *client.Client
	mutex  *sync.Mutex
}

// Generates a new installation process struct that will be used to create containers,
// and otherwise perform installation commands for a server.
func NewInstallationProcess(s *Server, script *api.InstallationScript) (*InstallationProcess, error) {
	proc := &InstallationProcess{
		Script: script,
		Server: s,
		mutex:  &sync.Mutex{},
	}

	if c, err := client.NewClientWithOpts(client.FromEnv); err != nil {
		return nil, errors.WithStack(err)
	} else {
		proc.client = c
	}

	return proc, nil
}

// Runs the installation process, this is done as a backgrounded thread. This will configure
// the required environment, and then spin up the installation container.
//
// Once the container finishes installing the results will be stored in an installation
// log in the server's configuration directory.
func (ip *InstallationProcess) Run() error {
	installPath, err := ip.BeforeExecute()
	if err != nil {
		return err
	}

	cid, err := ip.Execute(installPath)
	if err != nil {
		return err
	}

	// If this step fails, log a warning but don't exit out of the process. This is completely
	// internal to the daemon's functionality, and does not affect the status of the server itself.
	if err := ip.AfterExecute(cid); err != nil {
		zap.S().Warnw("failed to complete after-execute step of installation process", zap.String("server", ip.Server.Uuid), zap.Error(err))
	}

	return nil
}

// Writes the installation script to a temporary file on the host machine so that it
// can be properly mounted into the installation container and then executed.
func (ip *InstallationProcess) writeScriptToDisk() (string, error) {
	// Make sure the temp directory root exists before trying to make a directory within it. The
	// ioutil.TempDir call expects this base to exist, it won't create it for you.
	if err := os.MkdirAll(path.Join(os.TempDir(), "pterodactyl/"), 0700); err != nil {
		return "", errors.WithStack(err)
	}

	d, err := ioutil.TempDir("", "pterodactyl/")
	if err != nil {
		return "", errors.WithStack(err)
	}

	f, err := os.OpenFile(filepath.Join(d, "install.sh"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return "", errors.WithStack(err)
	}
	defer f.Close()

	w := bufio.NewWriter(f)

	scanner := bufio.NewScanner(bytes.NewReader([]byte(ip.Script.Script)))
	for scanner.Scan() {
		w.WriteString(scanner.Text() + "\n")
	}

	if err := scanner.Err(); err != nil {
		return "", errors.WithStack(err)
	}

	w.Flush()

	return d, nil
}

// Pulls the docker image to be used for the installation container.
func (ip *InstallationProcess) pullInstallationImage() error {
	r, err := ip.client.ImagePull(context.Background(), ip.Script.ContainerImage, types.ImagePullOptions{})
	if err != nil {
		return errors.WithStack(err)
	}

	// Block continuation until the image has been pulled successfully.
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		zap.S().Debugw(scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Runs before the container is executed. This pulls down the required docker container image
// as well as writes the installation script to the disk. This process is executed in an async
// manner, if either one fails the error is returned.
func (ip *InstallationProcess) BeforeExecute() (string, error) {
	wg := sync.WaitGroup{}
	wg.Add(3)

	var e []error
	var fileName string

	go func() {
		defer wg.Done()
		name, err := ip.writeScriptToDisk()
		if err != nil {
			e = append(e, err)
			return
		}

		fileName = name
	}()

	go func() {
		defer wg.Done()
		if err := ip.pullInstallationImage(); err != nil {
			e = append(e, err)
		}
	}()

	go func() {
		defer wg.Done()

		opts := types.ContainerRemoveOptions{
			RemoveVolumes: true,
			Force:         true,
		}

		if err := ip.client.ContainerRemove(context.Background(), ip.Server.Uuid+"_installer", opts); err != nil {
			if !client.IsErrNotFound(err) {
				e = append(e, err)
			}
		}
	}()

	wg.Wait()

	// Maybe a better way to handle this, but if there is at least one error
	// just bail out of the process now.
	if len(e) > 0 {
		return "", errors.WithStack(e[0])
	}

	return fileName, nil
}

// Returns the log path for the installation process.
func (ip *InstallationProcess) GetLogPath() string {
	return filepath.Join(config.Get().System.GetInstallLogPath(), ip.Server.Uuid+".log")
}

// Cleans up after the execution of the installation process. This grabs the logs from the
// process to store in the server configuration directory, and then destroys the associated
// installation container.
func (ip *InstallationProcess) AfterExecute(containerId string) error {
	ctx := context.Background()

	zap.S().Debugw("pulling installation logs for server", zap.String("server", ip.Server.Uuid), zap.String("container_id", containerId))
	reader, err := ip.client.ContainerLogs(ctx, containerId, types.ContainerLogsOptions{
		ShowStdout: true,
		ShowStderr: true,
		Follow:     false,
	})

	if err != nil && !client.IsErrNotFound(err) {
		return errors.WithStack(err)
	}

	f, err := os.OpenFile(ip.GetLogPath(), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return errors.WithStack(err)
	}
	defer f.Close()

	// We write the contents of the container output to a more "permanent" file so that they
	// can be referenced after this container is deleted.
	if _, err := io.Copy(f, reader); err != nil {
		return errors.WithStack(err)
	}

	zap.S().Debugw("removing server installation container", zap.String("server", ip.Server.Uuid), zap.String("container_id", containerId))
	rErr := ip.client.ContainerRemove(ctx, containerId, types.ContainerRemoveOptions{
		RemoveVolumes: true,
		RemoveLinks:   false,
		Force:         true,
	})

	if rErr != nil && !client.IsErrNotFound(rErr) {
		return errors.WithStack(rErr)
	}

	return nil
}

// Executes the installation process inside a specially created docker container.
func (ip *InstallationProcess) Execute(installPath string) (string, error) {
	ctx := context.Background()

	zap.S().Debugw(
		"creating server installer container",
		zap.String("server", ip.Server.Uuid),
		zap.String("script_path", installPath+"/install.sh"),
	)

	conf := &container.Config{
		Hostname:     "installer",
		AttachStdout: true,
		AttachStderr: true,
		AttachStdin:  true,
		OpenStdin:    true,
		Tty:          true,
		Cmd:          []string{ip.Script.Entrypoint, "./mnt/install/install.sh"},
		Image:        ip.Script.ContainerImage,
		Env:          ip.Server.GetEnvironmentVariables(),
		Labels: map[string]string{
			"Service":       "Pterodactyl",
			"ContainerType": "server_installer",
		},
	}

	hostConf := &container.HostConfig{
		Mounts: []mount.Mount{
			{
				Target:   "/mnt/server",
				Source:   ip.Server.Filesystem.Path(),
				Type:     mount.TypeBind,
				ReadOnly: false,
			},
			{
				Target:   "/mnt/install",
				Source:   installPath,
				Type:     mount.TypeBind,
				ReadOnly: false,
			},
		},
		Tmpfs: map[string]string{
			"/tmp": "rw,exec,nosuid,size=50M",
		},
		DNS: []string{"1.1.1.1", "8.8.8.8"},
		LogConfig: container.LogConfig{
			Type: "local",
			Config: map[string]string{
				"max-size": "5m",
				"max-file": "1",
				"compress": "false",
			},
		},
		Privileged:  true,
		NetworkMode: "pterodactyl_nw",
	}

	zap.S().Infow("creating installer container for server process", zap.String("server", ip.Server.Uuid))
	r, err := ip.client.ContainerCreate(ctx, conf, hostConf, nil, ip.Server.Uuid+"_installer")
	if err != nil {
		return "", errors.WithStack(err)
	}

	zap.S().Infow(
		"running installation script for server in container",
		zap.String("server", ip.Server.Uuid),
		zap.String("container_id", r.ID),
	)
	if err := ip.client.ContainerStart(ctx, r.ID, types.ContainerStartOptions{}); err != nil {
		return "", err
	}

	go func(id string) {
		ip.Server.Events().Publish(DaemonMessageEvent, "Starting installation process, this could take a few minutes...")
		if err := ip.StreamOutput(id); err != nil {
			zap.S().Errorw(
				"error handling streaming output for server install process",
				zap.String("container_id", id),
				zap.Error(err),
			)
		}
		ip.Server.Events().Publish(DaemonMessageEvent, "Installation process completed.")
	}(r.ID)

	sChann, eChann := ip.client.ContainerWait(ctx, r.ID, container.WaitConditionNotRunning)
	select {
	case err := <-eChann:
		if err != nil {
			return "", errors.WithStack(err)
		}
	case <-sChann:
	}

	return r.ID, nil
}

// Streams the output of the installation process to a log file in the server configuration
// directory, as well as to a websocket listener so that the process can be viewed in
// the panel by administrators.
func (ip *InstallationProcess) StreamOutput(id string) error {
	reader, err := ip.client.ContainerLogs(context.Background(), id, types.ContainerLogsOptions{
		ShowStdout: true,
		ShowStderr: true,
		Follow:     true,
	})

	if err != nil {
		return errors.WithStack(err)
	}

	defer reader.Close()

	s := bufio.NewScanner(reader)
	for s.Scan() {
		ip.Server.Events().Publish(InstallOutputEvent, s.Text())
	}

	if err := s.Err(); err != nil {
		zap.S().Warnw(
			"error processing scanner line in installation output for server",
			zap.String("server", ip.Server.Uuid),
			zap.String("container_id", id),
			zap.Error(errors.WithStack(err)),
		)
	}

	return nil
}

// Makes a HTTP request to the Panel instance notifying it that the server has
// completed the installation process, and what the state of the server is. A boolean
// value of "true" means everything was successful, "false" means something went
// wrong and the server must be deleted and re-created.
func (s *Server) SyncInstallState(successful bool) error {
	r := api.NewRequester()

	rerr, err := r.SendInstallationStatus(s.Uuid, successful)
	if rerr != nil || err != nil {
		if err != nil {
			return errors.WithStack(err)
		}

		return errors.New(rerr.String())
	}

	return nil
}
<commit_msg>Always clean up the installer containers, even if there is an error during the process; closes pterodactyl/panel#2015<commit_after>package server

import (
	"bufio"
	"bytes"
	"context"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/docker/client"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/api"
	"github.com/pterodactyl/wings/config"
	"go.uber.org/zap"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sync"
)

// Executes the installation stack for a server process. Bubbles any errors up to the calling
// function which should handle contacting the panel to notify it of the server state.
func (s *Server) Install() error {
	err := s.internalInstall()

	zap.S().Debugw("notifying panel of server install state", zap.String("server", s.Uuid))
	if serr := s.SyncInstallState(err == nil); serr != nil {
		zap.S().Warnw(
			"failed to notify panel of server install state",
			zap.String("server", s.Uuid),
			zap.Bool("was_successful", err == nil),
			zap.Error(serr),
		)
	}

	return err
}

// Reinstalls a server's software by utilizing the install script for the server egg. This
// does not touch any existing files for the server, other than what the script modifies.
func (s *Server) Reinstall() error {
	if s.GetState() != ProcessOfflineState {
		zap.S().Debugw("waiting for server instance to enter a stopped state", zap.String("server", s.Uuid))
		if err := s.Environment.WaitForStop(10, true); err != nil {
			return err
		}
	}

	return s.Install()
}

// Internal installation function used to simplify reporting back to the Panel.
func (s *Server) internalInstall() error {
	script, rerr, err := api.NewRequester().GetInstallationScript(s.Uuid)
	if err != nil || rerr != nil {
		if err != nil {
			return err
		}

		return errors.New(rerr.String())
	}

	p, err := NewInstallationProcess(s, &script)
	if err != nil {
		return errors.WithStack(err)
	}

	zap.S().Infow("beginning installation process for server", zap.String("server", s.Uuid))

	if err := p.Run(); err != nil {
		return err
	}

	zap.S().Infow("completed installation process for server", zap.String("server", s.Uuid))

	return nil
}

type InstallationProcess struct {
	Server *Server
	Script *api.InstallationScript

	client *client.Client
	mutex  *sync.Mutex
}

// Generates a new installation process struct that will be used to create containers,
// and otherwise perform installation commands for a server.
func NewInstallationProcess(s *Server, script *api.InstallationScript) (*InstallationProcess, error) {
	proc := &InstallationProcess{
		Script: script,
		Server: s,
		mutex:  &sync.Mutex{},
	}

	if c, err := client.NewClientWithOpts(client.FromEnv); err != nil {
		return nil, errors.WithStack(err)
	} else {
		proc.client = c
	}

	return proc, nil
}

// Removes the installer container for the server.
func (ip *InstallationProcess) RemoveContainer() {
	err := ip.client.ContainerRemove(context.Background(), ip.Server.Uuid + "_installer", types.ContainerRemoveOptions{
		RemoveVolumes: true,
		Force:         true,
	})

	if err != nil && !client.IsErrNotFound(err) {
		zap.S().Warnw("failed to delete server installer container", zap.String("server", ip.Server.Uuid), zap.Error(errors.WithStack(err)))
	}
}

// Runs the installation process, this is done as a backgrounded thread. This will configure
// the required environment, and then spin up the installation container.
//
// Once the container finishes installing the results will be stored in an installation
// log in the server's configuration directory.
func (ip *InstallationProcess) Run() error {
	installPath, err := ip.BeforeExecute()
	if err != nil {
		return err
	}

	cid, err := ip.Execute(installPath)
	if err != nil {
		ip.RemoveContainer()

		return err
	}

	// If this step fails, log a warning but don't exit out of the process. This is completely
	// internal to the daemon's functionality, and does not affect the status of the server itself.
	if err := ip.AfterExecute(cid); err != nil {
		zap.S().Warnw("failed to complete after-execute step of installation process", zap.String("server", ip.Server.Uuid), zap.Error(err))
	}

	return nil
}

// Writes the installation script to a temporary file on the host machine so that it
// can be properly mounted into the installation container and then executed.
func (ip *InstallationProcess) writeScriptToDisk() (string, error) {
	// Make sure the temp directory root exists before trying to make a directory within it. The
	// ioutil.TempDir call expects this base to exist, it won't create it for you.
	if err := os.MkdirAll(path.Join(os.TempDir(), "pterodactyl/"), 0700); err != nil {
		return "", errors.WithStack(err)
	}

	d, err := ioutil.TempDir("", "pterodactyl/")
	if err != nil {
		return "", errors.WithStack(err)
	}

	f, err := os.OpenFile(filepath.Join(d, "install.sh"), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return "", errors.WithStack(err)
	}
	defer f.Close()

	w := bufio.NewWriter(f)

	scanner := bufio.NewScanner(bytes.NewReader([]byte(ip.Script.Script)))
	for scanner.Scan() {
		w.WriteString(scanner.Text() + "\n")
	}

	if err := scanner.Err(); err != nil {
		return "", errors.WithStack(err)
	}

	w.Flush()

	return d, nil
}

// Pulls the docker image to be used for the installation container.
func (ip *InstallationProcess) pullInstallationImage() error {
	r, err := ip.client.ImagePull(context.Background(), ip.Script.ContainerImage, types.ImagePullOptions{})
	if err != nil {
		return errors.WithStack(err)
	}

	// Block continuation until the image has been pulled successfully.
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		zap.S().Debugw(scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Runs before the container is executed. This pulls down the required docker container image
// as well as writes the installation script to the disk. This process is executed in an async
// manner, if either one fails the error is returned.
func (ip *InstallationProcess) BeforeExecute() (string, error) {
	wg := sync.WaitGroup{}
	wg.Add(3)

	var e []error
	var fileName string

	go func() {
		defer wg.Done()
		name, err := ip.writeScriptToDisk()
		if err != nil {
			e = append(e, err)
			return
		}

		fileName = name
	}()

	go func() {
		defer wg.Done()
		if err := ip.pullInstallationImage(); err != nil {
			e = append(e, err)
		}
	}()

	go func() {
		defer wg.Done()

		opts := types.ContainerRemoveOptions{
			RemoveVolumes: true,
			Force:         true,
		}

		if err := ip.client.ContainerRemove(context.Background(), ip.Server.Uuid+"_installer", opts); err != nil {
			if !client.IsErrNotFound(err) {
				e = append(e, err)
			}
		}
	}()

	wg.Wait()

	// Maybe a better way to handle this, but if there is at least one error
	// just bail out of the process now.
	if len(e) > 0 {
		return "", errors.WithStack(e[0])
	}

	return fileName, nil
}

// Returns the log path for the installation process.
func (ip *InstallationProcess) GetLogPath() string {
	return filepath.Join(config.Get().System.GetInstallLogPath(), ip.Server.Uuid+".log")
}

// Cleans up after the execution of the installation process. This grabs the logs from the
// process to store in the server configuration directory, and then destroys the associated
// installation container.
func (ip *InstallationProcess) AfterExecute(containerId string) error {
	ctx := context.Background()
	defer ip.RemoveContainer()

	zap.S().Debugw("pulling installation logs for server", zap.String("server", ip.Server.Uuid), zap.String("container_id", containerId))
	reader, err := ip.client.ContainerLogs(ctx, containerId, types.ContainerLogsOptions{
		ShowStdout: true,
		ShowStderr: true,
		Follow:     false,
	})

	if err != nil && !client.IsErrNotFound(err) {
		return errors.WithStack(err)
	}

	f, err := os.OpenFile(ip.GetLogPath(), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return errors.WithStack(err)
	}
	defer f.Close()

	// We write the contents of the container output to a more "permanent" file so that they
	// can be referenced after this container is deleted.
	if _, err := io.Copy(f, reader); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Executes the installation process inside a specially created docker container.
func (ip *InstallationProcess) Execute(installPath string) (string, error) {
	ctx := context.Background()

	zap.S().Debugw(
		"creating server installer container",
		zap.String("server", ip.Server.Uuid),
		zap.String("script_path", installPath+"/install.sh"),
	)

	conf := &container.Config{
		Hostname:     "installer",
		AttachStdout: true,
		AttachStderr: true,
		AttachStdin:  true,
		OpenStdin:    true,
		Tty:          true,
		Cmd:          []string{ip.Script.Entrypoint, "./mnt/install/install.sh"},
		Image:        ip.Script.ContainerImage,
		Env:          ip.Server.GetEnvironmentVariables(),
		Labels: map[string]string{
			"Service":       "Pterodactyl",
			"ContainerType": "server_installer",
		},
	}

	hostConf := &container.HostConfig{
		Mounts: []mount.Mount{
			{
				Target:   "/mnt/server",
				Source:   ip.Server.Filesystem.Path(),
				Type:     mount.TypeBind,
				ReadOnly: false,
			},
			{
				Target:   "/mnt/install",
				Source:   installPath,
				Type:     mount.TypeBind,
				ReadOnly: false,
			},
		},
		Tmpfs: map[string]string{
			"/tmp": "rw,exec,nosuid,size=50M",
		},
		DNS: config.Get().Docker.Network.Dns,
		LogConfig: container.LogConfig{
			Type: "local",
			Config: map[string]string{
				"max-size": "5m",
				"max-file": "1",
				"compress": "false",
			},
		},
		Privileged:  true,
		NetworkMode: container.NetworkMode(config.Get().Docker.Network.Mode),
	}

	zap.S().Infow("creating installer container for server process", zap.String("server", ip.Server.Uuid))
	r, err := ip.client.ContainerCreate(ctx, conf, hostConf, nil, ip.Server.Uuid+"_installer")
	if err != nil {
		return "", errors.WithStack(err)
	}

	zap.S().Infow(
		"running installation script for server in container",
		zap.String("server", ip.Server.Uuid),
		zap.String("container_id", r.ID),
	)
	if err := ip.client.ContainerStart(ctx, r.ID, types.ContainerStartOptions{}); err != nil {
		return "", err
	}

	go func(id string) {
		ip.Server.Events().Publish(DaemonMessageEvent, "Starting installation process, this could take a few minutes...")
		if err := ip.StreamOutput(id); err != nil {
			zap.S().Errorw(
				"error handling streaming output for server install process",
				zap.String("container_id", id),
				zap.Error(err),
			)
		}
		ip.Server.Events().Publish(DaemonMessageEvent, "Installation process completed.")
	}(r.ID)

	sChann, eChann := ip.client.ContainerWait(ctx, r.ID, container.WaitConditionNotRunning)
	select {
	case err := <-eChann:
		if err != nil {
			return "", errors.WithStack(err)
		}
	case <-sChann:
	}

	return r.ID, nil
}

// Streams the output of the installation process to a log file in the server configuration
// directory, as well as to a websocket listener so that the process can be viewed in
// the panel by administrators.
func (ip *InstallationProcess) StreamOutput(id string) error {
	reader, err := ip.client.ContainerLogs(context.Background(), id, types.ContainerLogsOptions{
		ShowStdout: true,
		ShowStderr: true,
		Follow:     true,
	})

	if err != nil {
		return errors.WithStack(err)
	}

	defer reader.Close()

	s := bufio.NewScanner(reader)
	for s.Scan() {
		ip.Server.Events().Publish(InstallOutputEvent, s.Text())
	}

	if err := s.Err(); err != nil {
		zap.S().Warnw(
			"error processing scanner line in installation output for server",
			zap.String("server", ip.Server.Uuid),
			zap.String("container_id", id),
			zap.Error(errors.WithStack(err)),
		)
	}

	return nil
}

// Makes a HTTP request to the Panel instance notifying it that the server has
// completed the installation process, and what the state of the server is. A boolean
// value of "true" means everything was successful, "false" means something went
// wrong and the server must be deleted and re-created.
func (s *Server) SyncInstallState(successful bool) error {
	r := api.NewRequester()

	rerr, err := r.SendInstallationStatus(s.Uuid, successful)
	if rerr != nil || err != nil {
		if err != nil {
			return errors.WithStack(err)
		}

		return errors.New(rerr.String())
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func TestBasicStreamLoad(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that specific files show up
	filesToCheck := []string{
		"README.md", "project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped", "bin",
		"bin/mybinary.so", "filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file.exe", "folder/file1.txt", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s == nil {
			t.Error("File not found in sandbox: %v", p)
		}
	}
}

func TestStreamLoadOnExistingFiles(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	deletemePath := filepath.Join(sandbox1, "deleteme.txt")
	deleteme, err := os.Create(deletemePath)
	if err != nil {
		panic(err)
	}
	defer deleteme.Close()
	_, err = deleteme.Write([]byte("delete this"))
	if err != nil {
		panic(err)
	}
	deleteme.Close()

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1, "-force=true"}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	s, _ := os.Stat(deletemePath)
	if s != nil {
		t.Fail()
	}

	// TODO maybe someday we should back up what was in the sandbox before the first load
	// Check if the deleteme is backed up in the backup folder
	//s, _ = os.Stat(filepath.Join(sandbox1, backupFolder, "deleteme.txt"))
	//if s == nil {
	//	t.Fail()
	//}
}

func TestLoadAndClobberChanges(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Make adds and mods to the files
	filesToModify := []string{
		"README.md", "project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped",
		"filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file1.txt", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
	}
	for _, file := range filesToModify {
		path := filepath.Join(sandbox1, file)
		s, _ := os.Stat(path)
		if s == nil {
			t.Error("File not found in sandbox: %v", path)
		}

		if s.IsDir() {
			deleteMe, err := os.Create(filepath.Join(path, "deleteMe.txt"))
			if err != nil {
				panic(err)
			}
			defer deleteMe.Close()
			_, err = deleteMe.Write([]byte("test contents"))
			if err != nil {
				panic(err)
			}

			err = os.Mkdir(filepath.Join(path, "deleteMe"), 0700)
			if err != nil {
				panic(err)
			}
		} else {
			modFile, err := os.OpenFile(path, os.O_WRONLY, 0)
			if err != nil {
				panic(err)
			}
			defer modFile.Close()

			_, err = modFile.Write([]byte("new contents123"))
			if err != nil {
				panic(err)
			}
		}
	}

	os.Args = []string{"load", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	status, err := scmStatus(sandbox1, NO_COPY)
	if err != nil {
		t.Fatalf("%v", err.Error())
	}
	if !status.unchanged() {
		t.Fail()
	}

	// Check that all of the files and folders made their way into the backup
	for _, file := range filesToModify {
		path := filepath.Join(sandbox1, backupFolder, file)
		s, _ := os.Stat(path)
		if s == nil {
			t.Fatalf("File not found in backup: %v", path)
		}

		if s.IsDir() {
			s, _ := os.Stat(filepath.Join(path, "deleteMe"))
			if s == nil {
				t.Error("File not found in backup: %v", filepath.Join(path, "deleteMe"))
			}

			s, _ = os.Stat(filepath.Join(path, "deleteMe.txt"))
			if s == nil {
				t.Error("File not found in backup: %v", filepath.Join(path, "deleteMe.txt"))
			}
		}
	}
}

func TestAlternateStreamLoad(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-stream=Alternate Stream", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that specific files show up
	filesToCheck := []string{
		"README.md", "project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped", "bin",
		"bin/mybinary.so", "filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file.exe", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
		"alternateFile.txt", "alternateFolder", "alternateFolder/anotherAlternateFolder",
		"alternateFolder/anotherAlternateFile.txt",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s == nil {
			t.Error("File not found in sandbox: %v", p)
		}
	}
}

func TestEmptyStreamLoad(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-stream=Empty Stream", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that only the jazzMeta file is created in the sandbox
	f, err := os.Open(sandbox1)
	if err != nil {
		panic(err)
	}

	names, err := f.Readdirnames(-1)
	if err != nil {
		panic(err)
	}

	if len(names) != 1 {
		t.Fatalf("Expected 1 file but found %v files", len(names))
	}

	if names[0] != metadataFileName {
		t.Fatalf("Expected only the metadata file but found %v", names[0])
	}
}

func TestSwitchStreams(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-stream=Alternate Stream", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that specific files show up
	filesToCheck := []string{
		"README.md", "project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped", "bin",
		"bin/mybinary.so", "filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file.exe", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
		"alternateFile.txt", "alternateFolder", "alternateFolder/anotherAlternateFolder",
		"alternateFolder/anotherAlternateFile.txt",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s == nil {
			t.Error("File not found in sandbox: %v", p)
		}
	}

	// Verify that the file from the other stream no longer exists
	filesToCheck = []string{
		"folder/file1.txt",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s != nil {
			t.Error("File from the wrong stream found in sandbox: %v", p)
		}
	}

	// Verify that there are no backups since nothing was modified
	s, _ := os.Stat(filepath.Join(sandbox1, backupFolder))
	if s != nil {
		t.Fatalf("Found a backup folder even though no changes were made.")
	}
}
<commit_msg>Add a sandbox root delta to the switch stream case.<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func TestBasicStreamLoad(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that specific files show up
	filesToCheck := []string{
		"README.md", "project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped", "bin",
		"bin/mybinary.so", "filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file.exe", "folder/file1.txt", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s == nil {
			t.Error("File not found in sandbox: %v", p)
		}
	}
}

func TestStreamLoadOnExistingFiles(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	deletemePath := filepath.Join(sandbox1, "deleteme.txt")
	deleteme, err := os.Create(deletemePath)
	if err != nil {
		panic(err)
	}
	defer deleteme.Close()
	_, err = deleteme.Write([]byte("delete this"))
	if err != nil {
		panic(err)
	}
	deleteme.Close()

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1, "-force=true"}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	s, _ := os.Stat(deletemePath)
	if s != nil {
		t.Fail()
	}

	// TODO maybe someday we should back up what was in the sandbox before the first load
	// Check if the deleteme is backed up in the backup folder
	//s, _ = os.Stat(filepath.Join(sandbox1, backupFolder, "deleteme.txt"))
	//if s == nil {
	//	t.Fail()
	//}
}

func TestLoadAndClobberChanges(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Make adds and mods to the files
	filesToModify := []string{
		"README.md", "project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped",
		"filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file1.txt", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
	}
	for _, file := range filesToModify {
		path := filepath.Join(sandbox1, file)
		s, _ := os.Stat(path)
		if s == nil {
			t.Error("File not found in sandbox: %v", path)
		}

		if s.IsDir() {
			deleteMe, err := os.Create(filepath.Join(path, "deleteMe.txt"))
			if err != nil {
				panic(err)
			}
			defer deleteMe.Close()
			_, err = deleteMe.Write([]byte("test contents"))
			if err != nil {
				panic(err)
			}

			err = os.Mkdir(filepath.Join(path, "deleteMe"), 0700)
			if err != nil {
				panic(err)
			}
		} else {
			modFile, err := os.OpenFile(path, os.O_WRONLY, 0)
			if err != nil {
				panic(err)
			}
			defer modFile.Close()

			_, err = modFile.Write([]byte("new contents123"))
			if err != nil {
				panic(err)
			}
		}
	}

	os.Args = []string{"load", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	status, err := scmStatus(sandbox1, NO_COPY)
	if err != nil {
		t.Fatalf("%v", err.Error())
	}
	if !status.unchanged() {
		t.Fail()
	}

	// Check that all of the files and folders made their way into the backup
	for _, file := range filesToModify {
		path := filepath.Join(sandbox1, backupFolder, file)
		s, _ := os.Stat(path)
		if s == nil {
			t.Fatalf("File not found in backup: %v", path)
		}

		if s.IsDir() {
			s, _ := os.Stat(filepath.Join(path, "deleteMe"))
			if s == nil {
				t.Error("File not found in backup: %v", filepath.Join(path, "deleteMe"))
			}

			s, _ = os.Stat(filepath.Join(path, "deleteMe.txt"))
			if s == nil {
				t.Error("File not found in backup: %v", filepath.Join(path, "deleteMe.txt"))
			}
		}
	}
}

func TestAlternateStreamLoad(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-stream=Alternate Stream", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that specific files show up
	filesToCheck := []string{
		"project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped", "bin",
		"bin/mybinary.so", "filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file.exe", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
		"alternateFile.txt", "alternateFolder", "alternateFolder/anotherAlternateFolder",
		"alternateFolder/anotherAlternateFile.txt",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s == nil {
			t.Error("File not found in sandbox: %v", p)
		}
	}
}

func TestEmptyStreamLoad(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-stream=Empty Stream", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that only the jazzMeta file is created in the sandbox
	f, err := os.Open(sandbox1)
	if err != nil {
		panic(err)
	}

	names, err := f.Readdirnames(-1)
	if err != nil {
		panic(err)
	}

	if len(names) != 1 {
		t.Fatalf("Expected 1 file but found %v files", len(names))
	}

	if names[0] != metadataFileName {
		t.Fatalf("Expected only the metadata file but found %v", names[0])
	}
}

func TestSwitchStreams(t *testing.T) {
	sandbox1, err := ioutil.TempDir(os.TempDir(), "gojazz-test")
	if err != nil {
		panic(err)
	}

	defer os.RemoveAll(sandbox1)

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	t.Logf("Loading test project into %v\n", sandbox1)
	os.Args = []string{"load", "sirnewton | gojazz-test", "-stream=Alternate Stream", "-sandbox=" + sandbox1}
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	loadOp()

	// Verify that specific files show up
	filesToCheck := []string{
		"project.json", "bigFile.txt", ".jazzignore",
		".cfignore", "folder", "filename(with)[chars$]^that.must-be-escaped", "bin",
		"bin/mybinary.so", "filename(with)[chars$]^that.must-be-escaped/test.java",
		"folder/file.exe", "folder/file2.jsp", "folder/file3.jar",
		"folder/filename(with)[chars$]^that.must-be-escaped",
		"alternateFile.txt", "alternateFolder", "alternateFolder/anotherAlternateFolder",
		"alternateFolder/anotherAlternateFile.txt",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s == nil {
			t.Error("File not found in sandbox: %v", p)
		}
	}

	// Verify that the file from the other stream no longer exists
	filesToCheck = []string{
		"folder/file1.txt", "README.md",
	}
	for _, file := range filesToCheck {
		p := filepath.Join(sandbox1, file)
		s, _ := os.Stat(p)
		if s != nil {
			t.Error("File from the wrong stream found in sandbox: %v", p)
		}
	}

	// Verify that there are no backups since nothing was modified
	s, _ := os.Stat(filepath.Join(sandbox1, backupFolder))
	if s != nil {
		t.Fatalf("Found a backup folder even though no changes were made.")
	}
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/smancke/guble/protocol"

	"fmt"
	"github.com/docker/distribution/health"
	"net/http"
	"reflect"
	"time"
)

// Startable interface for modules which provide a start mechanism
type Startable interface {
	Start() error
}

// Stopable interface for modules which provide a stop mechanism
type Stopable interface {
	Stop() error
}

// Endpoint adds a HTTP handler for the `GetPrefix()` to the webserver
type Endpoint interface {
	http.Handler
	GetPrefix() string
}

// Service is the main class for simple control of a server
type Service struct {
	webServer  *WebServer
	router     Router
	stopables  []Stopable
	startables []Startable
	// The time given to each Module on Stop()
	StopGracePeriod      time.Duration
	healthCheckFrequency time.Duration
	healthCheckThreshold int
}

// NewService registers the Main Router, where other modules can subscribe for messages
func NewService(addr string, router Router) *Service {
	service := &Service{
		stopables:       make([]Stopable, 0, 5),
		webServer:       NewWebServer(addr),
		router:          router,
		StopGracePeriod: time.Second * 2,
	}
	service.Register(service.webServer)
	service.Register(service.router)

	service.webServer.Handle("/health", http.HandlerFunc(health.StatusHandler))

	return service
}

// Register the supplied module on this service.
// This method checks the module for the following interfaces and
// does the expected registrations:
//   Startable,
//   Stopable: notify when the service stops
//   health.Checker:
//   Endpoint: Register the handler function of the Endpoint in the http service at prefix
func (s *Service) Register(module interface{}) {
	name := reflect.TypeOf(module).String()

	if startable, ok := module.(Startable); ok {
		protocol.Info("register %v as StartListener", name)
		s.startables = append(s.startables, startable)
	}

	if stopable, ok := module.(Stopable); ok {
		protocol.Info("register %v as StopListener", name)
		s.stopables = append(s.stopables, stopable)
	}

	if checker, ok := module.(health.Checker); ok {
		protocol.Info("register %v as HealthChecker", name)
		//TODO parameterize / configure frequency and threshold
		health.RegisterPeriodicThresholdFunc(name, s.healthCheckFrequency, s.healthCheckThreshold, health.CheckFunc(checker.Check))
	}

	if endpoint, ok := module.(Endpoint); ok {
		protocol.Info("register %v as Endpoint to %v", name, endpoint.GetPrefix())
		s.webServer.Handle(endpoint.GetPrefix(), endpoint)
	}
}

func (s *Service) Start() error {
	el := protocol.NewErrorList("Errors occured while startup the service: ")

	for _, startable := range s.startables {
		name := reflect.TypeOf(startable).String()

		protocol.Debug("starting module %v", name)
		if err := startable.Start(); err != nil {
			protocol.Err("error on startup module %v", name)
			el.Add(err)
		}
	}
	return el.ErrorOrNil()
}

func (s *Service) Stop() error {
	errors := make(map[string]error)
	for _, stopable := range s.stopables {
		name := reflect.TypeOf(stopable).String()
		stoppedChan := make(chan bool)
		errorChan := make(chan error)
		protocol.Info("stopping %v ...", name)
		go func() {
			err := stopable.Stop()
			if err != nil {
				errorChan <- err
				return
			}
			stoppedChan <- true
		}()
		select {
		case err := <-errorChan:
			protocol.Err("error while stopping %v: %v", name, err.Error)
			errors[name] = err
		case <-stoppedChan:
			protocol.Info("stopped %v", name)
		case <-time.After(s.StopGracePeriod):
			errors[name] = fmt.Errorf("error while stopping %v: not returned after %v seconds", name, s.StopGracePeriod)
			protocol.Err(errors[name].Error())
		}
	}
	if len(errors) > 0 {
		return fmt.Errorf("Errors while stopping modules %q", errors)
	}
	return nil
}

func (s *Service) GetWebServer() *WebServer {
	return s.webServer
}
<commit_msg>minor refactor<commit_after>package server

import (
	"github.com/smancke/guble/protocol"

	"fmt"
	"github.com/docker/distribution/health"
	"net/http"
	"reflect"
	"time"
)

// Startable interface for modules which provide a start mechanism
type Startable interface {
	Start() error
}

// Stopable interface for modules which provide a stop mechanism
type Stopable interface {
	Stop() error
}

// Endpoint adds a HTTP handler for the `GetPrefix()` to the webserver
type Endpoint interface {
	http.Handler
	GetPrefix() string
}

// Service is the main class for simple control of a server
type Service struct {
	webServer  *WebServer
	router     Router
	stopables  []Stopable
	startables []Startable
	// The time given to each Module on Stop()
	StopGracePeriod      time.Duration
	healthCheckFrequency time.Duration
	healthCheckThreshold int
}

// NewService registers the Main Router, where other modules can subscribe for messages
func NewService(addr string, router Router) *Service {
	service := &Service{
		stopables:       make([]Stopable, 0, 5),
		webServer:       NewWebServer(addr),
		router:          router,
		StopGracePeriod: time.Second * 2,
	}
	service.Register(service.webServer)
	service.Register(service.router)

	service.webServer.Handle("/health", http.HandlerFunc(health.StatusHandler))

	return service
}

// Register the supplied module on this service.
// This method checks the module for the following interfaces and
// does the expected registrations:
//   Startable,
//   Stopable: notify when the service stops
//   health.Checker:
//   Endpoint: Register the handler function of the Endpoint in the http service at prefix
func (s *Service) Register(module interface{}) {
	name := reflect.TypeOf(module).String()

	if startable, ok := module.(Startable); ok {
		protocol.Info("register %v as Startable", name)
		s.startables = append(s.startables, startable)
	}

	if stopable, ok := module.(Stopable); ok {
		protocol.Info("register %v as Stopable", name)
		s.stopables = append(s.stopables, stopable)
	}

	if checker, ok := module.(health.Checker); ok {
		protocol.Info("register %v as HealthChecker", name)
		//TODO parameterize / configure frequency and threshold
		health.RegisterPeriodicThresholdFunc(name, s.healthCheckFrequency, s.healthCheckThreshold, health.CheckFunc(checker.Check))
	}

	if endpoint, ok := module.(Endpoint); ok {
		prefix := endpoint.GetPrefix()
		protocol.Info("register %v as Endpoint to %v", name, prefix)
		s.webServer.Handle(prefix, endpoint)
	}
}

func (s *Service) Start() error {
	el := protocol.NewErrorList("Errors occured while startup the service: ")

	for _, startable := range s.startables {
		name := reflect.TypeOf(startable).String()

		protocol.Debug("starting module %v", name)
		if err := startable.Start(); err != nil {
			protocol.Err("error on startup module %v", name)
			el.Add(err)
		}
	}
	return el.ErrorOrNil()
}

func (s *Service) Stop() error {
	errors := make(map[string]error)
	for _, stopable := range s.stopables {
		name := reflect.TypeOf(stopable).String()
		stoppedChan := make(chan bool)
		errorChan := make(chan error)
		protocol.Info("stopping %v ...", name)
		go func() {
			err := stopable.Stop()
			if err != nil {
				errorChan <- err
				return
			}
			stoppedChan <- true
		}()
		select {
		case err := <-errorChan:
			protocol.Err("error while stopping %v: %v", name, err.Error)
			errors[name] = err
		case <-stoppedChan:
			protocol.Info("stopped %v", name)
		case <-time.After(s.StopGracePeriod):
			errors[name] = fmt.Errorf("error while stopping %v: not returned after %v seconds", name, s.StopGracePeriod)
			protocol.Err(errors[name].Error())
		}
	}
	if len(errors) > 0 {
		return fmt.Errorf("Errors while stopping modules %q", errors)
	}
	return nil
}

func (s *Service) GetWebServer() *WebServer {
	return s.webServer
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"
	"time"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/boltdb/bolt"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.ValueGCRunInterval = 10 * time.Hour
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxReaders(math.MaxInt64)
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly|lmdb.NoReadahead, 0777)
	y.Check(err)
	return lmdbEnv
}

func getBoltDB() *bolt.DB {
	opts := bolt.DefaultOptions
	opts.ReadOnly = true
	boltdb, err := bolt.Open(*flagDir+"/bolt/bolt.db", 0777, opts)
	y.Check(err)
	return boltdb
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

type hitCounter struct {
	found    uint64
	notFound uint64
	errored  uint64
}

func (h *hitCounter) Reset() {
	h.found, h.notFound, h.errored = 0, 0, 0
}

func (h *hitCounter) Update(c *hitCounter) {
	atomic.AddUint64(&h.found, c.found)
	atomic.AddUint64(&h.notFound, c.notFound)
	atomic.AddUint64(&h.errored, c.errored)
}

func (h *hitCounter) Print(storeName string, b *testing.B) {
	b.Logf("%s: %d keys had valid values.", storeName, h.found)
	b.Logf("%s: %d keys had no values", storeName, h.notFound)
	b.Logf("%s: %d keys had errors", storeName, h.errored)
	b.Logf("%s: %d total keys looked at", storeName, h.found+h.notFound+h.errored)
	b.Logf("%s: hit rate : %.2f", storeName, float64(h.found)/float64(h.found+h.notFound+h.errored))
}

// A generic read benchmark that runs the doBench func for a specific key value store,
// aggregates the hit counts and prints them out.
func runRandomReadBenchmark(b *testing.B, storeName string, doBench func(*hitCounter, *testing.PB)) {
	counter := &hitCounter{}
	b.Run("read-random"+storeName, func(b *testing.B) {
		counter.Reset()
		b.RunParallel(func(pb *testing.PB) {
			c := &hitCounter{}
			doBench(c, pb)
			counter.Update(c)
		})
	})
	counter.Print(storeName, b)
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	runRandomReadBenchmark(b, "badger", func(c *hitCounter, pb *testing.PB) {
		var item badger.KVItem
		for pb.Next() {
			key := newKey()
			err := bdb.Get(key, &item)
			if err != nil {
				c.errored++
			} else if v := item.Value(); v != nil {
				y.AssertTruef(len(v) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v), *flagValueSize)
				c.found++
			} else {
				c.notFound++
			}
		}
	})
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	runRandomReadBenchmark(b, "rocksdb", func(c *hitCounter, pb *testing.PB) {
		for pb.Next() {
			key := newKey()
			// FIXME lookup API and increment accordingly
			if _, err := rdb.Get(key); err == nil {
				c.found++
			}
		}
	})
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	runRandomReadBenchmark(b, "lmdb", func(c *hitCounter, pb *testing.PB) {
		err := lmdbEnv.View(func(txn *lmdb.Txn) error {
			txn.RawRead = true
			for pb.Next() {
				key := newKey()
				v, err := txn.Get(lmdbDBI, key)
				if lmdb.IsNotFound(err) {
					c.notFound++
					continue
				} else if err != nil {
					c.errored++
					continue
				}
				y.AssertTruef(len(v) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v), *flagValueSize)
				c.found++
			}
			return nil
		})
		if err != nil {
			y.Check(err)
		}
	})
}

func BenchmarkReadRandomBolt(b *testing.B) {
	boltdb := getBoltDB()
	defer boltdb.Close()

	runRandomReadBenchmark(b, "bolt", func(c *hitCounter, pb *testing.PB) {
		err := boltdb.View(func(txn *bolt.Tx) error {
			boltBkt := txn.Bucket([]byte("bench"))
			y.AssertTrue(boltBkt != nil)
			for pb.Next() {
				key := newKey()
				v := boltBkt.Get(key)
				if v == nil {
					c.notFound++
					continue
				}
				y.AssertTruef(len(v) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v), *flagValueSize)
				c.found++
			}
			return nil
		})
		y.Check(err)
	})
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				txn.RawRead = true
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					y.AssertTruef(len(v1) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v1), *flagValueSize)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count >= 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBolt(b *testing.B) {
	boltdb := getBoltDB()
	defer boltdb.Close()

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("boltdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err := boltdb.View(func(txn *bolt.Tx) error {
				boltBkt := txn.Bucket([]byte("bench"))
				y.AssertTrue(boltBkt != nil)
				cur := boltBkt.Cursor()
				for k1, v1 := cur.First(); k1 != nil; k1, v1 = cur.Next() {
					y.AssertTruef(len(v1) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v1), *flagValueSize)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count >= 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				vsz := len(item.Value())
				y.AssertTruef(vsz == *flagValueSize, "Assertion failed. value size is %d, expected %d", vsz, *flagValueSize)
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<commit_msg>Set lmdb.NoReadahead flag conditionally<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"
	"time"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/boltdb/bolt"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.ValueGCRunInterval = 10 * time.Hour
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb(readahead bool) *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxReaders(math.MaxInt64)
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	var flags uint = lmdb.Readonly
	if !readahead {
		flags |= lmdb.NoReadahead
	}
	err = lmdbEnv.Open(*flagDir+"/lmdb", flags, 0777)
	y.Check(err)
	return lmdbEnv
}

func getBoltDB() *bolt.DB {
	opts := bolt.DefaultOptions
	opts.ReadOnly = true
	boltdb, err := bolt.Open(*flagDir+"/bolt/bolt.db", 0777, opts)
	y.Check(err)
	return boltdb
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

type hitCounter struct {
	found    uint64
	notFound uint64
	errored  uint64
}

func (h *hitCounter) Reset() {
	h.found, h.notFound, h.errored = 0, 0, 0
}

func (h *hitCounter) Update(c *hitCounter) {
	atomic.AddUint64(&h.found, c.found)
	atomic.AddUint64(&h.notFound, c.notFound)
	atomic.AddUint64(&h.errored, c.errored)
}

func (h *hitCounter) Print(storeName string, b *testing.B) {
	b.Logf("%s: %d keys had valid values.", storeName, h.found)
	b.Logf("%s: %d keys had no values", storeName, h.notFound)
	b.Logf("%s: %d keys had errors", storeName, h.errored)
	b.Logf("%s: %d total keys looked at", storeName, h.found+h.notFound+h.errored)
	b.Logf("%s: hit rate : %.2f", storeName, float64(h.found)/float64(h.found+h.notFound+h.errored))
}

// A generic read benchmark that runs the doBench func for a specific key value store,
// aggregates the hit counts and prints them out.
func runRandomReadBenchmark(b *testing.B, storeName string, doBench func(*hitCounter, *testing.PB)) {
	counter := &hitCounter{}
	b.Run("read-random"+storeName, func(b *testing.B) {
		counter.Reset()
		b.RunParallel(func(pb *testing.PB) {
			c := &hitCounter{}
			doBench(c, pb)
			counter.Update(c)
		})
	})
	counter.Print(storeName, b)
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	runRandomReadBenchmark(b, "badger", func(c *hitCounter, pb *testing.PB) {
		var item badger.KVItem
		for pb.Next() {
			key := newKey()
			err := bdb.Get(key, &item)
			if err != nil {
				c.errored++
			} else if v := item.Value(); v != nil {
				y.AssertTruef(len(v) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v), *flagValueSize)
				c.found++
			} else {
				c.notFound++
			}
		}
	})
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	runRandomReadBenchmark(b, "rocksdb", func(c *hitCounter, pb *testing.PB) {
		for pb.Next() {
			key := newKey()
			// FIXME lookup API and increment accordingly
			if _, err := rdb.Get(key); err == nil {
				c.found++
			}
		}
	})
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb(false)
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	runRandomReadBenchmark(b, "lmdb", func(c *hitCounter, pb *testing.PB) {
		err := lmdbEnv.View(func(txn *lmdb.Txn) error {
			txn.RawRead = true
			for pb.Next() {
				key := newKey()
				v, err := txn.Get(lmdbDBI, key)
				if lmdb.IsNotFound(err) {
					c.notFound++
					continue
				} else if err != nil {
					c.errored++
					continue
				}
				y.AssertTruef(len(v) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v), *flagValueSize)
				c.found++
			}
			return nil
		})
		if err != nil {
			y.Check(err)
		}
	})
}

func BenchmarkReadRandomBolt(b *testing.B) {
	boltdb := getBoltDB()
	defer boltdb.Close()

	runRandomReadBenchmark(b, "bolt", func(c *hitCounter, pb *testing.PB) {
		err := boltdb.View(func(txn *bolt.Tx) error {
			boltBkt := txn.Bucket([]byte("bench"))
			y.AssertTrue(boltBkt != nil)
			for pb.Next() {
				key := newKey()
				v := boltBkt.Get(key)
				if v == nil {
					c.notFound++
					continue
				}
				y.AssertTruef(len(v) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v), *flagValueSize)
				c.found++
			}
			return nil
		})
		y.Check(err)
	})
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb(true)
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				txn.RawRead = true
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					y.AssertTruef(len(v1) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v1), *flagValueSize)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count >= 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBolt(b *testing.B) {
	boltdb := getBoltDB()
	defer boltdb.Close()

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("boltdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err := boltdb.View(func(txn *bolt.Tx) error {
				boltBkt := txn.Bucket([]byte("bench"))
				y.AssertTrue(boltBkt != nil)
				cur := boltBkt.Cursor()
				for k1, v1 := cur.First(); k1 != nil; k1, v1 = cur.Next() {
					y.AssertTruef(len(v1) == *flagValueSize, "Assertion failed. value size is %d, expected %d", len(v1), *flagValueSize)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count >= 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				vsz := len(item.Value())
				y.AssertTruef(vsz == *flagValueSize, "Assertion failed. value size is %d, expected %d", vsz, *flagValueSize)
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<|endoftext|>