<commit_before>package conservator

import (
	"fmt"
	"log"
	"time"

	"github.com/antihax/evedata/internal/datapackages"
	"github.com/antihax/evedata/internal/gobcoder"
	"github.com/antihax/goesi/notification"
	"github.com/bradfitz/slice"
	nsq "github.com/nsqio/go-nsq"
	yaml "gopkg.in/yaml.v2"
)

func init() {
	addHandler("characterNotifications", spawnCharacterNotificationsConsumer)
}

func spawnCharacterNotificationsConsumer(s *Conservator, consumer *nsq.Consumer) {
	consumer.AddHandler(s.wait(nsq.HandlerFunc(s.characterNotificationsHandler)))
}

func (s *Conservator) characterNotificationsHandler(message *nsq.Message) error {
	notifications := datapackages.CharacterNotifications{}
	err := gobcoder.GobDecoder(message.Body, &notifications)
	if err != nil {
		log.Println(err)
		return err
	}

	if len(notifications.Notifications) == 0 {
		return nil
	}

	// sort by time
	slice.Sort(notifications.Notifications[:], func(i, j int) bool {
		return notifications.Notifications[i].Timestamp.Unix() < notifications.Notifications[j].Timestamp.Unix()
	})

	if len(notifications.Notifications) == 0 {
		log.Println("we broke the notifications")
		return nil
	}

	for _, n := range notifications.Notifications {
		// Skip the notification if if is old
		if n.Timestamp.Before(time.Now().UTC().Add(-time.Hour * 12)) {
			continue
		}

		err := s.checkNotification(notifications.TokenCharacterID, n.NotificationId, n.Type_, n.Text, n.Timestamp)
		if err != nil {
			log.Println(err)
			return err
		}
	}
	return nil
}

func (s *Conservator) checkNotification(characterID int32, notificationID int64, notificationType, text string, timestamp time.Time) error {
	switch notificationType {

	case "LocateCharMsg":
		l := notification.LocateCharMsg{}
		if err := yaml.Unmarshal([]byte(text), &l); err != nil {
			log.Println(err)
			return err
		}
		systemName, err := s.getCelestialName(l.TargetLocation.SolarSystem)
		if err != nil {
			log.Println(err)
			return err
		}
		regionName, _ := s.getCelestialName(l.TargetLocation.Region)

		character, _ := s.getEntityName(l.CharacterID)

		stationName, _ := s.getStationName(l.TargetLocation.Station)

		message := fmt.Sprintf(" %s (https://www.evedata.org/character?id=%d) has been located in %s of %s",
			character.Name, l.CharacterID, systemName, regionName)

		if stationName != "" {
			message = fmt.Sprintf("%s docked at %s", message, stationName)
		}

		return s.sendNotificationMessage("locator", characterID, notificationID, message)

	case "AllWarDeclaredMsg", "CorpWarDeclaredMsg":
		l := notification.AllWarDeclaredMsg{}
		if err := yaml.Unmarshal([]byte(text), &l); err != nil {
			log.Println(err)
			return err
		}

		defender, err := s.getEntityName(l.AgainstID)
		if err != nil {
			log.Println(err)
			return err
		}
		attacker, err := s.getEntityName(l.DeclaredByID)
		if err != nil {
			log.Println(err)
			return err
		}

		message := fmt.Sprintf("[%s](https://www.evedata.org/%s?id=%d) just declared war on [%s](https://www.evedata.org/%s?id=%d)\n",
			attacker.Name, attacker.EntityType, l.DeclaredByID, defender.Name, defender.EntityType, l.AgainstID)

		return s.sendNotificationMessage("war", characterID, notificationID, message)

	case "CorpAppNewMsg":
		{
			l := notification.CorpAppNewMsg{}
			yaml.Unmarshal([]byte(text), &l)
			character, err := s.getEntityName(l.CharID)
			if err != nil {
				character.Name = "!! No fricking clue !!"
			}
			corporation, _ := s.getEntityName(l.CorpID)

			message := fmt.Sprintf("New corporation application from [%s](https://www.evedata.org/character?id=%d) to %s. Application Comment: %s\n",
				character.Name, l.CharID, corporation.Name, l.ApplicationText)

			return s.sendNotificationMessage("application", characterID, notificationID, message)
		}

	case "StructureUnderAttack":
		l := notification.StructureUnderAttack{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		attackerName := ""
		if l.AllianceID > 0 {
			attacker = l.AllianceID
			attackerType = "alliance"
			attackerName = l.AllianceName
		} else if l.CharID > 0 { // There is no corporationID? WTF CCPls?
			attacker = l.CharID
			attackerType = "character"
			attackerName = l.CorpName
		}

		systemName, err := s.getCelestialName(l.SolarsystemID) // -.-
		if err != nil {
			log.Println(err)
			return err
		}

		message := fmt.Sprintf("A structure is under attack in %s by [%s](https://www.evedata.org/%s?id=%d) S: %.1f%%  A: %.1f%%  H: %.1f%% \n",
			systemName, attackerName, attackerType, attacker, l.ShieldPercentage*100, l.ArmorPercentage*100, l.HullPercentage*100)

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "OrbitalAttacked":
		l := notification.OrbitalAttacked{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		if l.AggressorAllianceID > 0 {
			attacker = l.AggressorAllianceID
			attackerType = "alliance"
		} else if l.AggressorCorpID > 0 {
			attacker = l.AggressorCorpID
			attackerType = "corporation"
		}

		locationName, err := s.getCelestialName(l.PlanetID)
		if err != nil {
			log.Println(err)
			return err
		}
		systemName, err := s.getCelestialName(l.SolarSystemID)
		if err != nil {
			log.Println(err)
			return err
		}
		structureType, err := s.getTypeName(l.TypeID)
		if err != nil {
			log.Println(err)
			return err
		}
		attackerName, err := s.getEntityName(attacker)
		if err != nil {
			log.Println(err)
			return err
		}

		message := fmt.Sprintf("%s is under attack at %s in %s by [%s](https://www.evedata.org/%s?id=%d) S: %.1f%%\n",
			structureType, locationName, systemName, attackerName.Name, attackerType, attacker, l.ShieldLevel*100)

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "TowerAlertMsg":
		l := notification.TowerAlertMsg{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		if l.AggressorAllianceID > 0 {
			attacker = l.AggressorAllianceID
			attackerType = "alliance"
		} else if l.AggressorCorpID > 0 {
			attacker = l.AggressorCorpID
			attackerType = "corporation"
		}

		locationName, err := s.getCelestialName(l.MoonID)
		if err != nil {
			log.Println(err)
			return err
		}
		systemName, err := s.getCelestialName(l.SolarSystemID)
		if err != nil {
			log.Println(err)
			return err
		}
		structureType, err := s.getTypeName(l.TypeID)
		if err != nil {
			log.Println(err)
			return err
		}
		attackerName, err := s.getEntityName(attacker)
		if err != nil {
			log.Println(err)
			return err
		}
		message := fmt.Sprintf(" %s is under attack at %s in %s by [%s](https://www.evedata.org/%s?id=%d) S: %.1f%%  A: %.1f%%  H: %.1f%% \n",
			structureType, locationName, systemName, attackerName.Name, attackerType, attacker, l.ShieldValue*100, l.ArmorValue*100, l.HullValue*100)

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "OrbitalReinforced":
		l := notification.OrbitalReinforced{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		if l.AggressorAllianceID > 0 {
			attacker = l.AggressorAllianceID
			attackerType = "alliance"
		} else if l.AggressorCorpID > 0 {
			attacker = l.AggressorCorpID
			attackerType = "corporation"
		}

		locationName, err := s.getCelestialName(l.PlanetID)
		if err != nil {
			log.Println(err)
			return err
		}
		systemName, err := s.getCelestialName(l.SolarSystemID)
		if err != nil {
			log.Println(err)
			return err
		}
		structureType, err := s.getTypeName(l.TypeID)
		if err != nil {
			log.Println(err)
			return err
		}
		attackerName, err := s.getEntityName(attacker)
		if err != nil {
			log.Println(err)
			return err
		}

		message := fmt.Sprintf("%s was reinforced at %s in %s by [%s](https://www.evedata.org/%s?id=%d).\n\n Timer expires at %s\n",
			structureType, locationName, systemName, attackerName.Name, attackerType, attacker,
			time.Unix(datapackages.WintoUnixTimestamp(l.ReinforceExitTime), 0).UTC().String())

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "StructureLostShields", "StructureLostArmor":
		l := notification.StructureLostShields{}
		yaml.Unmarshal([]byte(text), &l)

		systemName, err := s.getCelestialName(l.SolarsystemID)
		if err != nil {
			log.Println(err)
			return err
		}

		structureType, err := s.getTypeName(l.StructureTypeID)
		if err != nil {
			log.Println(err)
			return err
		}

		message := fmt.Sprintf("%s entered %s in %s: Timer expires at %s\n",
			structureType, notificationType, systemName,
			time.Unix(datapackages.WintoUnixTimestamp(l.Timestamp), 0).UTC().String())

		return s.sendNotificationMessage("structure", characterID, notificationID, message)
	}
	return nil
}

func (s *Conservator) sendNotificationMessage(messageType string, characterID int32, notificationID int64, message string) error {
	shares, ok := s.notifications[messageType][characterID]
	if !ok {
		return nil
	}
	for _, share := range shares {
		channelData := unpackChannelData(share.Packed)
		for _, channel := range channelData {
			if inSlice(messageType, channel.Services) {
				if s.outQueue.CheckWorkCompleted(fmt.Sprintf("evedata-bot-notification-sent:%s", channel.ChannelID), notificationID) {
					continue
				}

				c, err := s.getChannel(channel.ChannelID)
				if err != nil {
					continue
				}

				// Get the service
				service, err := s.getService(c.IntegrationID)
				if err != nil {
					log.Println(err)
					return err
				}

				if err := service.Server.SendMessageToChannel(channel.ChannelID, message); err != nil {
					log.Println(err)
					return err
				}
				s.outQueue.SetWorkCompleted(fmt.Sprintf("evedata-bot-notification-sent:%s", channel.ChannelID), notificationID)
			}
		}
	}
	return nil
}

type EntityName struct {
	Name       string `db:"name" json:"name"`
	EntityType string `db:"type" json:"type"`
}

// Obtain entity name and type by ID.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getEntityName(id int32) (*EntityName, error) {
	ref := EntityName{}
	if err := s.db.QueryRowx(`
		SELECT name, 'corporation' AS type FROM evedata.corporations WHERE corporationID = ?
		UNION
		SELECT name, 'alliance' AS type FROM evedata.alliances WHERE allianceID = ?
		UNION
		SELECT name, 'character' AS type FROM evedata.characters WHERE characterID = ?
		LIMIT 1`, id, id, id).StructScan(&ref); err != nil {
		return &ref, err
	}
	return &ref, nil
}

// Obtain type name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getTypeName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT typeName FROM invTypes WHERE typeID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}

// Obtain SolarSystem name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getSystemName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT solarSystemName FROM mapSolarSystems WHERE solarSystemID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}

// Obtain Celestial name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getCelestialName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT itemName FROM mapDenormalize WHERE itemID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}

// Obtain Station name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getStationName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT stationName FROM staStations WHERE stationID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}
<commit_msg>don't error, just return blank<commit_after>package conservator

import (
	"fmt"
	"log"
	"time"

	"github.com/antihax/evedata/internal/datapackages"
	"github.com/antihax/evedata/internal/gobcoder"
	"github.com/antihax/goesi/notification"
	"github.com/bradfitz/slice"
	nsq "github.com/nsqio/go-nsq"
	yaml "gopkg.in/yaml.v2"
)

func init() {
	addHandler("characterNotifications", spawnCharacterNotificationsConsumer)
}

func spawnCharacterNotificationsConsumer(s *Conservator, consumer *nsq.Consumer) {
	consumer.AddHandler(s.wait(nsq.HandlerFunc(s.characterNotificationsHandler)))
}

func (s *Conservator) characterNotificationsHandler(message *nsq.Message) error {
	notifications := datapackages.CharacterNotifications{}
	err := gobcoder.GobDecoder(message.Body, &notifications)
	if err != nil {
		log.Println(err)
		return err
	}

	if len(notifications.Notifications) == 0 {
		return nil
	}

	// sort by time
	slice.Sort(notifications.Notifications[:], func(i, j int) bool {
		return notifications.Notifications[i].Timestamp.Unix() < notifications.Notifications[j].Timestamp.Unix()
	})

	if len(notifications.Notifications) == 0 {
		log.Println("we broke the notifications")
		return nil
	}

	for _, n := range notifications.Notifications {
		// Skip the notification if if is old
		if n.Timestamp.Before(time.Now().UTC().Add(-time.Hour * 12)) {
			continue
		}

		err := s.checkNotification(notifications.TokenCharacterID, n.NotificationId, n.Type_, n.Text, n.Timestamp)
		if err != nil {
			log.Println(err)
			return err
		}
	}
	return nil
}

func (s *Conservator) checkNotification(characterID int32, notificationID int64, notificationType, text string, timestamp time.Time) error {
	switch notificationType {

	case "LocateCharMsg":
		l := notification.LocateCharMsg{}
		if err := yaml.Unmarshal([]byte(text), &l); err != nil {
			log.Println(err)
			return err
		}
		systemName, err := s.getCelestialName(l.TargetLocation.SolarSystem)
		if err != nil {
			log.Println(err)
			return err
		}
		regionName, _ := s.getCelestialName(l.TargetLocation.Region)

		character, _ := s.getEntityName(l.CharacterID)

		stationName, _ := s.getStationName(l.TargetLocation.Station)

		message := fmt.Sprintf(" %s (https://www.evedata.org/character?id=%d) has been located in %s of %s",
			character.Name, l.CharacterID, systemName, regionName)

		if stationName != "" {
			message = fmt.Sprintf("%s docked at %s", message, stationName)
		}

		return s.sendNotificationMessage("locator", characterID, notificationID, message)

	case "AllWarDeclaredMsg", "CorpWarDeclaredMsg":
		l := notification.AllWarDeclaredMsg{}
		if err := yaml.Unmarshal([]byte(text), &l); err != nil {
			log.Println(err)
		}

		defender, err := s.getEntityName(l.AgainstID)
		if err != nil {
			log.Println(err)
		}
		attacker, err := s.getEntityName(l.DeclaredByID)
		if err != nil {
			log.Println(err)
		}

		message := fmt.Sprintf("[%s](https://www.evedata.org/%s?id=%d) just declared war on [%s](https://www.evedata.org/%s?id=%d)\n",
			attacker.Name, attacker.EntityType, l.DeclaredByID, defender.Name, defender.EntityType, l.AgainstID)

		return s.sendNotificationMessage("war", characterID, notificationID, message)

	case "CorpAppNewMsg":
		{
			l := notification.CorpAppNewMsg{}
			yaml.Unmarshal([]byte(text), &l)
			character, err := s.getEntityName(l.CharID)
			if err != nil {
				character.Name = "!! No fricking clue !!"
			}
			corporation, _ := s.getEntityName(l.CorpID)

			message := fmt.Sprintf("New corporation application from [%s](https://www.evedata.org/character?id=%d) to %s. Application Comment: %s\n",
				character.Name, l.CharID, corporation.Name, l.ApplicationText)

			return s.sendNotificationMessage("application", characterID, notificationID, message)
		}

	case "StructureUnderAttack":
		l := notification.StructureUnderAttack{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		attackerName := ""
		if l.AllianceID > 0 {
			attacker = l.AllianceID
			attackerType = "alliance"
			attackerName = l.AllianceName
		} else if l.CharID > 0 { // There is no corporationID? WTF CCPls?
			attacker = l.CharID
			attackerType = "character"
			attackerName = l.CorpName
		}

		systemName, err := s.getCelestialName(l.SolarsystemID) // -.-
		if err != nil {
			log.Println(err)
		}

		message := fmt.Sprintf("A structure is under attack in %s by [%s](https://www.evedata.org/%s?id=%d) S: %.1f%%  A: %.1f%%  H: %.1f%% \n",
			systemName, attackerName, attackerType, attacker, l.ShieldPercentage, l.ArmorPercentage, l.HullPercentage)

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "OrbitalAttacked":
		l := notification.OrbitalAttacked{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		if l.AggressorAllianceID > 0 {
			attacker = l.AggressorAllianceID
			attackerType = "alliance"
		} else if l.AggressorCorpID > 0 {
			attacker = l.AggressorCorpID
			attackerType = "corporation"
		}

		locationName, err := s.getCelestialName(l.PlanetID)
		if err != nil {
			log.Println(err)
		}
		systemName, err := s.getCelestialName(l.SolarSystemID)
		if err != nil {
			log.Println(err)
		}
		structureType, err := s.getTypeName(l.TypeID)
		if err != nil {
			log.Println(err)
		}
		attackerName, err := s.getEntityName(attacker)
		if err != nil {
			log.Println(err)
		}

		message := fmt.Sprintf("%s is under attack at %s in %s by [%s](https://www.evedata.org/%s?id=%d) S: %.1f%%\n",
			structureType, locationName, systemName, attackerName.Name, attackerType, attacker, l.ShieldLevel*100)

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "TowerAlertMsg":
		l := notification.TowerAlertMsg{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		if l.AggressorAllianceID > 0 {
			attacker = l.AggressorAllianceID
			attackerType = "alliance"
		} else if l.AggressorCorpID > 0 {
			attacker = l.AggressorCorpID
			attackerType = "corporation"
		}

		locationName, err := s.getCelestialName(l.MoonID)
		if err != nil {
			log.Println(err)
		}
		systemName, err := s.getCelestialName(l.SolarSystemID)
		if err != nil {
			log.Println(err)
		}
		structureType, err := s.getTypeName(l.TypeID)
		if err != nil {
			log.Println(err)
		}
		attackerName, err := s.getEntityName(attacker)
		if err != nil {
			log.Println(err)
		}
		message := fmt.Sprintf(" %s is under attack at %s in %s by [%s](https://www.evedata.org/%s?id=%d) S: %.1f%%  A: %.1f%%  H: %.1f%% \n",
			structureType, locationName, systemName, attackerName.Name, attackerType, attacker, l.ShieldValue*100, l.ArmorValue*100, l.HullValue*100)

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "OrbitalReinforced":
		l := notification.OrbitalReinforced{}
		yaml.Unmarshal([]byte(text), &l)

		attacker := int32(0)
		attackerType := ""
		if l.AggressorAllianceID > 0 {
			attacker = l.AggressorAllianceID
			attackerType = "alliance"
		} else if l.AggressorCorpID > 0 {
			attacker = l.AggressorCorpID
			attackerType = "corporation"
		}

		locationName, err := s.getCelestialName(l.PlanetID)
		if err != nil {
			log.Println(err)
		}
		systemName, err := s.getCelestialName(l.SolarSystemID)
		if err != nil {
			log.Println(err)
		}
		structureType, err := s.getTypeName(l.TypeID)
		if err != nil {
			log.Println(err)
		}
		attackerName, err := s.getEntityName(attacker)
		if err != nil {
			log.Println(err)
		}

		message := fmt.Sprintf("%s was reinforced at %s in %s by [%s](https://www.evedata.org/%s?id=%d).\n\n Timer expires at %s\n",
			structureType, locationName, systemName, attackerName.Name, attackerType, attacker,
			time.Unix(datapackages.WintoUnixTimestamp(l.ReinforceExitTime), 0).UTC().String())

		return s.sendNotificationMessage("structure", characterID, notificationID, message)

	case "StructureLostShields", "StructureLostArmor":
		l := notification.StructureLostShields{}
		yaml.Unmarshal([]byte(text), &l)

		systemName, err := s.getCelestialName(l.SolarsystemID)
		if err != nil {
			log.Println(err)
		}

		structureType, err := s.getTypeName(l.StructureTypeID)
		if err != nil {
			log.Println(err)
		}

		message := fmt.Sprintf("%s entered %s in %s: Timer expires at %s\n",
			structureType, notificationType, systemName,
			time.Unix(datapackages.WintoUnixTimestamp(l.Timestamp), 0).UTC().String())

		return s.sendNotificationMessage("structure", characterID, notificationID, message)
	}
	return nil
}

func (s *Conservator) sendNotificationMessage(messageType string, characterID int32, notificationID int64, message string) error {
	shares, ok := s.notifications[messageType][characterID]
	if !ok {
		return nil
	}
	for _, share := range shares {
		channelData := unpackChannelData(share.Packed)
		for _, channel := range channelData {
			if inSlice(messageType, channel.Services) {
				if s.outQueue.CheckWorkCompleted(fmt.Sprintf("evedata-bot-notification-sent:%s", channel.ChannelID), notificationID) {
					continue
				}

				c, err := s.getChannel(channel.ChannelID)
				if err != nil {
					continue
				}

				// Get the service
				service, err := s.getService(c.IntegrationID)
				if err != nil {
					log.Println(err)
					return err
				}

				if err := service.Server.SendMessageToChannel(channel.ChannelID, message); err != nil {
					log.Println(err)
					return err
				}
				s.outQueue.SetWorkCompleted(fmt.Sprintf("evedata-bot-notification-sent:%s", channel.ChannelID), notificationID)
			}
		}
	}
	return nil
}

type EntityName struct {
	Name       string `db:"name" json:"name"`
	EntityType string `db:"type" json:"type"`
}

// Obtain entity name and type by ID.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getEntityName(id int32) (*EntityName, error) {
	ref := EntityName{}
	if err := s.db.QueryRowx(`
		SELECT name, 'corporation' AS type FROM evedata.corporations WHERE corporationID = ?
		UNION
		SELECT name, 'alliance' AS type FROM evedata.alliances WHERE allianceID = ?
		UNION
		SELECT name, 'character' AS type FROM evedata.characters WHERE characterID = ?
		LIMIT 1`, id, id, id).StructScan(&ref); err != nil {
		return &ref, err
	}
	return &ref, nil
}

// Obtain type name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getTypeName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT typeName FROM invTypes WHERE typeID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}

// Obtain SolarSystem name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getSystemName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT solarSystemName FROM mapSolarSystems WHERE solarSystemID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}

// Obtain Celestial name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getCelestialName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT itemName FROM mapDenormalize WHERE itemID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}

// Obtain Station name.
// [BENCHMARK] 0.000 sec / 0.000 sec
func (s *Conservator) getStationName(id int32) (string, error) {
	ref := ""
	if err := s.db.QueryRowx(`
		SELECT stationName FROM staStations WHERE stationID = ?
		LIMIT 1`, id).Scan(&ref); err != nil {
		return "", err
	}
	return ref, nil
}
<|endoftext|>


<commit_before>package sinmetalcraft

import (
	"net/http"
	"net/url"
	"time"

	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/taskqueue"
	"google.golang.org/appengine/urlfetch"

	"google.golang.org/api/compute/v1"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

func init() {
	api := MinecraftTQApi{}

	http.HandleFunc("/tq/1/minecraft", api.Handler)
}

type MinecraftTQApi struct{}

func CallMinecraftTQ(c context.Context, minecraftKey *datastore.Key, operationID string) (*taskqueue.Task, error) {
	t := taskqueue.NewPOSTTask("/tq/1/minecraft", url.Values{
		"keyStr":      {minecraftKey.Encode()},
		"operationID": {operationID},
	})
	t.Delay = time.Second * 30
	return taskqueue.Add(c, t, "minecraft")
}

// /tq/1/minecraft handler
func (a *MinecraftTQApi) Handler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)

	keyStr := r.FormValue("KeyStr")
	operationID := r.FormValue("operationID")

	log.Infof(ctx, "keyStr = %s, operationID = %s", keyStr, operationID)

	key, err := datastore.DecodeKey(keyStr)
	if err != nil {
		log.Errorf(ctx, "key decode error. keyStr = %s, err = %s", keyStr, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	client := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(ctx, compute.ComputeScope),
			Base:   &urlfetch.Transport{Context: ctx},
		},
	}
	s, err := compute.New(client)
	if err != nil {
		log.Errorf(ctx, "ERROR compute.New: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	nzos := compute.NewZoneOperationsService(s)
	ope, err := nzos.Get(PROJECT_NAME, "asia-east1-b", operationID).Do()
	if err != nil {
		log.Errorf(ctx, "ERROR compute Zone Operation Get Error. zone = %s, operation = %s, error = %s", "asia-east1-b", operationID, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	WriteLog(ctx, "__GET_ZONE_COMPUTE_OPE__", ope)

	status := "exists"
	if ope.OperationType == "delete" {
		status = "not_exists"
	}

	resStatus := http.StatusOK
	if ope.Status == "DONE" {
		err = datastore.RunInTransaction(ctx, func(c context.Context) error {
			var entity Minecraft
			err := datastore.Get(ctx, key, &entity)
			if err != nil {
				return err
			}

			entity.ResourceID = int64(ope.TargetId)
			entity.Status = status
			entity.OperationStatus = ope.Status
			entity.OperationType = ope.OperationType
			entity.UpdatedAt = time.Now()

			_, err = datastore.Put(ctx, key, &entity)
			if err != nil {
				return err
			}

			return nil
		}, nil)
	} else {
		log.Infof(ctx, "Operation Status = %s", ope.Status)
		resStatus = http.StatusRequestTimeout

		err = datastore.RunInTransaction(ctx, func(c context.Context) error {
			var entity Minecraft
			err := datastore.Get(ctx, key, &entity)
			if err != nil {
				return err
			}

			entity.ResourceID = 0
			entity.OperationStatus = ope.Status
			entity.OperationType = ope.OperationType
			entity.UpdatedAt = time.Now()

			_, err = datastore.Put(ctx, key, &entity)
			if err != nil {
				return err
			}

			return nil
		}, nil)
	}
	if err != nil {
		log.Errorf(ctx, "Minecraft Put Error. error = %s", err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(resStatus)
}
<commit_msg>TQ Add時のバリデーションを追加した<commit_after>package sinmetalcraft

import (
	"net/http"
	"net/url"
	"time"

	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/taskqueue"
	"google.golang.org/appengine/urlfetch"

	"google.golang.org/api/compute/v1"

	"errors"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

func init() {
	api := MinecraftTQApi{}

	http.HandleFunc("/tq/1/minecraft", api.Handler)
}

type MinecraftTQApi struct{}

func CallMinecraftTQ(c context.Context, minecraftKey *datastore.Key, operationID string) (*taskqueue.Task, error) {
	log.Infof(c, "Call Minecraft TQ, key = %v, operationID = %s", minecraftKey, operationID)
	if minecraftKey == nil {
		return nil, errors.New("key is required")
	}
	if len(operationID) < 1 {
		return nil, errors.New("operationID is required")
	}

	t := taskqueue.NewPOSTTask("/tq/1/minecraft", url.Values{
		"keyStr":      {minecraftKey.Encode()},
		"operationID": {operationID},
	})
	t.Delay = time.Second * 30
	return taskqueue.Add(c, t, "minecraft")
}

// /tq/1/minecraft handler
func (a *MinecraftTQApi) Handler(w http.ResponseWriter, r *http.Request) {
	ctx := appengine.NewContext(r)

	keyStr := r.FormValue("keyStr")
	operationID := r.FormValue("operationID")

	log.Infof(ctx, "keyStr = %s, operationID = %s", keyStr, operationID)

	key, err := datastore.DecodeKey(keyStr)
	if err != nil {
		log.Errorf(ctx, "key decode error. keyStr = %s, err = %s", keyStr, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	client := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(ctx, compute.ComputeScope),
			Base:   &urlfetch.Transport{Context: ctx},
		},
	}
	s, err := compute.New(client)
	if err != nil {
		log.Errorf(ctx, "ERROR compute.New: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	nzos := compute.NewZoneOperationsService(s)
	ope, err := nzos.Get(PROJECT_NAME, "asia-east1-b", operationID).Do()
	if err != nil {
		log.Errorf(ctx, "ERROR compute Zone Operation Get Error. zone = %s, operation = %s, error = %s", "asia-east1-b", operationID, err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	WriteLog(ctx, "__GET_ZONE_COMPUTE_OPE__", ope)

	status := "exists"
	if ope.OperationType == "delete" {
		status = "not_exists"
	}

	resStatus := http.StatusOK
	if ope.Status == "DONE" {
		err = datastore.RunInTransaction(ctx, func(c context.Context) error {
			var entity Minecraft
			err := datastore.Get(ctx, key, &entity)
			if err != nil {
				return err
			}

			entity.ResourceID = int64(ope.TargetId)
			entity.Status = status
			entity.OperationStatus = ope.Status
			entity.OperationType = ope.OperationType
			entity.UpdatedAt = time.Now()

			_, err = datastore.Put(ctx, key, &entity)
			if err != nil {
				return err
			}

			return nil
		}, nil)
	} else {
		log.Infof(ctx, "Operation Status = %s", ope.Status)
		resStatus = http.StatusRequestTimeout

		err = datastore.RunInTransaction(ctx, func(c context.Context) error {
			var entity Minecraft
			err := datastore.Get(ctx, key, &entity)
			if err != nil {
				return err
			}

			entity.ResourceID = 0
			entity.OperationStatus = ope.Status
			entity.OperationType = ope.OperationType
			entity.UpdatedAt = time.Now()

			_, err = datastore.Put(ctx, key, &entity)
			if err != nil {
				return err
			}

			return nil
		}, nil)
	}
	if err != nil {
		log.Errorf(ctx, "Minecraft Put Error. error = %s", err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.WriteHeader(resStatus)
}
<|endoftext|>


<commit_before>// Copyright 2012 The Plotinum Authors. All rights reserved.
// Use of this source code is governed by an MIT-style license
// that can be found in the LICENSE file.

package plotter

import (
	"code.google.com/p/plotinum/plot"
)

// Line implements the Plotter interface, drawing a line.
type Line struct {
	// XYs is a copy of the points for this line.
	XYs

	// LineStyle is the style of the line connecting
	// the points.
	plot.LineStyle
}

// NewLine returns a Line that uses the
// default line style and does not draw
// glyphs.
func NewLine(xys XYer) *Line {
	return &Line{
		XYs:       CopyXYs(xys),
		LineStyle: DefaultLineStyle,
	}
}

// Plot draws the Line, implementing the plot.Plotter
// interface.
func (pts *Line) Plot(da plot.DrawArea, plt *plot.Plot) {
	trX, trY := plt.Transforms(&da)
	ps := make([]plot.Point, len(pts.XYs))
	for i, p := range pts.XYs {
		ps[i].X = trX(p.X)
		ps[i].Y = trY(p.Y)
	}
	da.StrokeLines(pts.LineStyle, da.ClipLinesXY(ps)...)
}

// DataRange returns the minimum and maximum
// x and y values, implementing the plot.DataRanger
// interface.
func (pts *Line) DataRange() (xmin, xmax, ymin, ymax float64) {
	return XYRange(pts)
}

// Thumbnail the thumbnail for the Line,
// implementing the plot.Thumbnailer interface.
func (pts *Line) Thumbnail(da *plot.DrawArea) {
	y := da.Center().Y
	da.StrokeLine2(pts.LineStyle, da.Min.X, y, da.Max().X, y)
}

// LinePoints is a line with glyphs at each point.
type LinePoints struct {
	*Line
	Points *Scatter
}

// NewLinePoints returns both a Line and a
// Points for the given point data.
func NewLinePoints(xys XYer) (*Line, *Scatter) {
	s := NewScatter(xys)
	l := &Line{
		XYs:       s.XYs,
		LineStyle: DefaultLineStyle,
	}
	return l, s
}
<commit_msg>Remove the unused LinePoints type.<commit_after>// Copyright 2012 The Plotinum Authors. All rights reserved.
// Use of this source code is governed by an MIT-style license
// that can be found in the LICENSE file.

package plotter

import (
	"code.google.com/p/plotinum/plot"
)

// Line implements the Plotter interface, drawing a line.
type Line struct {
	// XYs is a copy of the points for this line.
	XYs

	// LineStyle is the style of the line connecting
	// the points.
	plot.LineStyle
}

// NewLine returns a Line that uses the
// default line style and does not draw
// glyphs.
func NewLine(xys XYer) *Line {
	return &Line{
		XYs:       CopyXYs(xys),
		LineStyle: DefaultLineStyle,
	}
}

// Plot draws the Line, implementing the plot.Plotter
// interface.
func (pts *Line) Plot(da plot.DrawArea, plt *plot.Plot) {
	trX, trY := plt.Transforms(&da)
	ps := make([]plot.Point, len(pts.XYs))
	for i, p := range pts.XYs {
		ps[i].X = trX(p.X)
		ps[i].Y = trY(p.Y)
	}
	da.StrokeLines(pts.LineStyle, da.ClipLinesXY(ps)...)
}

// DataRange returns the minimum and maximum
// x and y values, implementing the plot.DataRanger
// interface.
func (pts *Line) DataRange() (xmin, xmax, ymin, ymax float64) {
	return XYRange(pts)
}

// Thumbnail the thumbnail for the Line,
// implementing the plot.Thumbnailer interface.
func (pts *Line) Thumbnail(da *plot.DrawArea) {
	y := da.Center().Y
	da.StrokeLine2(pts.LineStyle, da.Min.X, y, da.Max().X, y)
}

// NewLinePoints returns both a Line and a
// Points for the given point data.
func NewLinePoints(xys XYer) (*Line, *Scatter) {
	s := NewScatter(xys)
	l := &Line{
		XYs:       s.XYs,
		LineStyle: DefaultLineStyle,
	}
	return l, s
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"

	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gorilla/context"
	"github.com/stretchr/testify/assert"
)

func TestReadFindOneSimple(t *testing.T) {
	req := getRequest("/Project/65")
	w := httptest.NewRecorder()

	//{"results":{"entities":[{"type":"Project","id":63},{"type":"Project","id":65},{"type":"Project","id":66},{"type":"Project","id":71}],"paging_info":{"current_page":1,"page_count":1,"entity_count":4,"entities_per_page":500}}}
	server, client, config := mockShotgun(200,
		`{"results":{"entities":[{"type":"Project","id":65}],"paging_info":{"current_page":1,"page_count":1,"entity_count":1,"entities_per_page":1}}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	assert.Equal(t, http.StatusOK, w.Code)

	type Entity struct {
		Type string `json:"type"`
		Id   int    `json:"id"`
	}

	expected := []byte(`{"type":"Project","id":65}`)
	var jsonExpected Entity
	err := json.Unmarshal(expected, &jsonExpected)
	assert.Nil(t, err)

	var jsonResp Entity
	err = json.Unmarshal(w.Body.Bytes(), &jsonResp)
	assert.Nil(t, err)

	assert.Equal(t, jsonExpected, jsonResp)
}

func TestReadFindOneWithFields(t *testing.T) {
	req := getRequest("/Project/65?fields=name,sg_status")
	w := httptest.NewRecorder()

	//{"results":{"entities":[{"type":"Project","id":63},{"type":"Project","id":65},{"type":"Project","id":66},{"type":"Project","id":71}],"paging_info":{"current_page":1,"page_count":1,"entity_count":4,"entities_per_page":500}}}
	server, client, config := mockShotgun(200,
		`{"results":{"entities":[{"id":65,"name":"Big Buck Bunny","sg_status":"Active","type":"Project"}],"paging_info":{"current_page":1,"page_count":1,"entity_count":1,"entities_per_page":1}}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	assert.Equal(t, http.StatusOK, w.Code)

	type Entity struct {
		Type   string `json:"type"`
		Id     int    `json:"id"`
		Name   string `json:"name"`
		Status string `json:"sg_status"`
	}

	expected := []byte(`{"id":65,"name":"Big Buck Bunny","sg_status":"Active","type":"Project"}`)
	var jsonExpected Entity
	err := json.Unmarshal(expected, &jsonExpected)
	assert.Nil(t, err)

	var jsonResp Entity
	err = json.Unmarshal(w.Body.Bytes(), &jsonResp)
	assert.Nil(t, err)

	assert.Equal(t, jsonExpected, jsonResp)
}

func TestReadFindOneNoneIntId(t *testing.T) {
	req := getRequest("/Project/foo")
	w := httptest.NewRecorder()

	//{"results":{"entities":[{"type":"Project","id":63},{"type":"Project","id":65},{"type":"Project","id":66},{"type":"Project","id":71}],"paging_info":{"current_page":1,"page_count":1,"entity_count":4,"entities_per_page":500}}}
	server, client, config := mockShotgun(200,
		`{"results":{"entities":[{"id":65, "type":"Project"}],"paging_info":{"current_page":1,"page_count":1,"entity_count":1,"entities_per_page":1}}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)

	assert.Equal(t, http.StatusNotFound, w.Code)
}

func TestGetNotFound(t *testing.T) {
	req := getRequest("/Project/64")
	w := httptest.NewRecorder()

	server, client, config := mockShotgun(200, `{"results":{"entities":[]}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	assert.Equal(t, http.StatusNotFound, w.Code)
}

func TestReadNoShotgunClient(t *testing.T) {
	req := getRequest("/Project")
	w := httptest.NewRecorder()

	server, _, config := mockShotgun(200, "200 Not Ok - Because Shotgun.")
	defer server.Close()

	router(config).ServeHTTP(w, req)
	assert.Equal(t, http.StatusBadGateway, w.Code)
}

func TestGetBadResponseJson(t *testing.T) {
	req := getRequest("/Project/64")
	w := httptest.NewRecorder()

	server, client, config := mockShotgun(200, `foo`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	assert.Equal(t, http.StatusBadGateway, w.Code)
}
<commit_msg>converted entity_get_test.go to testify/suite<commit_after>package main

import (
	"encoding/json"

	"net/http"
	"net/http/httptest"
	"testing"

	log "github.com/Sirupsen/logrus"

	"github.com/gorilla/context"
	"github.com/stretchr/testify/suite"
)

// Define the suite, and absorb the built-in basic suite
// functionality from testify - including a T() method which
// returns the current testing context
type EntityGetTestSuite struct {
	suite.Suite
}

// Make sure that VariableThatShouldStartAtFive is set to five
// before each test
func (suite *EntityGetTestSuite) SetupSuite() {
	manager := GetQPManager()
	log.Info(" -- EntityGet Test Suite --\n")
	log.Debug("EntityGetTestSuite.SetupSuite() - setting active parsers to format1, format2, format3")
	manager.SetActiveParsers("format1", "format2", "format3")
}

// In order for 'go test' to run this suite, we need to create
// a normal test function and pass our suite to suite.Run
func TestEntityGetTestSuite(t *testing.T) {
	//ts := new(Format1TestSuite)
	log.Debug("EntityGetTestSuite - Running test suite")
	suite.Run(t, new(EntityGetTestSuite))
	log.Debug("EntityGetTestSuite - Finished test suite")
}

func (suite *EntityGetTestSuite) TestReadFindOneSimple() {
	req := getRequest("/Project/65")
	w := httptest.NewRecorder()

	//{"results":{"entities":[{"type":"Project","id":63},{"type":"Project","id":65},{"type":"Project","id":66},{"type":"Project","id":71}],"paging_info":{"current_page":1,"page_count":1,"entity_count":4,"entities_per_page":500}}}
	server, client, config := mockShotgun(200,
		`{"results":{"entities":[{"type":"Project","id":65}],"paging_info":{"current_page":1,"page_count":1,"entity_count":1,"entities_per_page":1}}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	suite.Equal(http.StatusOK, w.Code)

	type Entity struct {
		Type string `json:"type"`
		Id   int    `json:"id"`
	}

	expected := []byte(`{"type":"Project","id":65}`)
	var jsonExpected Entity
	err := json.Unmarshal(expected, &jsonExpected)
	suite.Nil(err)

	var jsonResp Entity
	err = json.Unmarshal(w.Body.Bytes(), &jsonResp)
	suite.Nil(err)

	suite.Equal(jsonExpected, jsonResp)
}

func (suite *EntityGetTestSuite) TestReadFindOneWithFields() {
	req := getRequest("/Project/65?fields=name,sg_status")
	w := httptest.NewRecorder()

	//{"results":{"entities":[{"type":"Project","id":63},{"type":"Project","id":65},{"type":"Project","id":66},{"type":"Project","id":71}],"paging_info":{"current_page":1,"page_count":1,"entity_count":4,"entities_per_page":500}}}
	server, client, config := mockShotgun(200,
		`{"results":{"entities":[{"id":65,"name":"Big Buck Bunny","sg_status":"Active","type":"Project"}],"paging_info":{"current_page":1,"page_count":1,"entity_count":1,"entities_per_page":1}}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	suite.Equal(http.StatusOK, w.Code)

	type Entity struct {
		Type   string `json:"type"`
		Id     int    `json:"id"`
		Name   string `json:"name"`
		Status string `json:"sg_status"`
	}

	expected := []byte(`{"id":65,"name":"Big Buck Bunny","sg_status":"Active","type":"Project"}`)
	var jsonExpected Entity
	err := json.Unmarshal(expected, &jsonExpected)
	suite.Nil(err)

	var jsonResp Entity
	err = json.Unmarshal(w.Body.Bytes(), &jsonResp)
	suite.Nil(err)

	suite.Equal(jsonExpected, jsonResp)
}

func (suite *EntityGetTestSuite) TestReadFindOneNoneIntId() {
	req := getRequest("/Project/foo")
	w := httptest.NewRecorder()

	//{"results":{"entities":[{"type":"Project","id":63},{"type":"Project","id":65},{"type":"Project","id":66},{"type":"Project","id":71}],"paging_info":{"current_page":1,"page_count":1,"entity_count":4,"entities_per_page":500}}}
	server, client, config := mockShotgun(200,
		`{"results":{"entities":[{"id":65, "type":"Project"}],"paging_info":{"current_page":1,"page_count":1,"entity_count":1,"entities_per_page":1}}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)

	suite.Equal(http.StatusNotFound, w.Code)
}

func (suite *EntityGetTestSuite) TestGetNotFound() {
	req := getRequest("/Project/64")
	w := httptest.NewRecorder()

	server, client, config := mockShotgun(200, `{"results":{"entities":[]}}`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	suite.Equal(http.StatusNotFound, w.Code)
}

func (suite *EntityGetTestSuite) TestReadNoShotgunClient() {
	req := getRequest("/Project")
	w := httptest.NewRecorder()

	server, _, config := mockShotgun(200, "200 Not Ok - Because Shotgun.")
	defer server.Close()

	router(config).ServeHTTP(w, req)
	suite.Equal(http.StatusBadGateway, w.Code)
}

func (suite *EntityGetTestSuite) TestGetBadResponseJson() {
	req := getRequest("/Project/64")
	w := httptest.NewRecorder()

	server, client, config := mockShotgun(200, `foo`)
	defer server.Close()

	context.Set(req, "sgConn", *client)
	router(config).ServeHTTP(w, req)
	suite.Equal(http.StatusBadGateway, w.Code)
}
<|endoftext|>


<commit_before>// Package assert provides convenience ways to test queue.Queue
// implementations.
package assert

import (
	"fmt"

	"github.com/alcortesm/queue"
)

// Test represents a test that can be set to fail with a formated
// message using the Errorf method, for instance, *testing.T from the
// standard library.
type Test interface {
	Errorf(format string, a ...interface{})
}

// Assert represent an association between a queue.Queue and a test.
// The methods from this type can be used form inside a test to
// manipulate a queue and assert that it behaves in particular ways.  If
// not, a test error will be automatically issued with a sensible
// failure message, prefixed with the value of the Prefix field.
type Assert struct {
	q queue.Queue
	t Test
	// The prefix to use in the error messages when an assertion fails.
	Prefix string
}

// New returns a new association between the given queue and the given
// test.  The Prefix field is initialized to the empty string.
func New(t Test, q queue.Queue) *Assert {
	return &Assert{
		q:      q,
		t:      t,
		Prefix: "",
	}
}

// Calls the Errorf method on the test with the given formated
// arguments, prefixed with the receiver's Prefix value.
func (a *Assert) errorf(format string, v ...interface{}) {
	a.t.Errorf("%s%s", a.Prefix, fmt.Sprintf(format, v...))
}

// Len asserts that the length of the associated queue is the expected value.
func (a *Assert) Len(expected int) {
	obtained := a.q.Len()
	if obtained != expected {
		a.errorf("wrong length: expected %d, got %d", expected, obtained)
	}
}

// IsEmpty asserts that calling IsEmpty on the associated queue returns the
// expected value.
func (a *Assert) IsEmpty(expected bool) {
	obtained := a.q.IsEmpty()
	if obtained != expected {
		a.errorf("wrong IsEmpty: expected %t, got %t", expected, obtained)
	}
}

// Enqueue asserts that enqueuing the given numbers on the associated
// queue success.
func (a *Assert) Enqueue(numbers ...int) {
	for i, n := range numbers {
		if err := a.q.Enqueue(n); err != nil {
			a.errorf("wrong #%d Enqueue(%d): unexpected error: %q",
				i, n, err)
		}
	}
}

// EnqueueErrFull asserts that trying to enqueue any element fails with
// queue.ErrFull.
func (a *Assert) EnqueueErrFull() {
	if err := a.q.Enqueue(42); err != queue.ErrFull {
		a.errorf("wrong Enqueue: should have failed with ErrFull, "+
			"but got %q instead", err)
	}
}

// Head asserts that calling Head on the associated queue success and returns
// the expected number.
func (a *Assert) Head(expected int) {
	obtained, err := a.q.Head()
	if err != nil {
		a.errorf("wrong Head: unexpected error: %s", err)
		return
	}
	n, ok := obtained.(int)
	if !ok {
		a.errorf("wrong Head: obtained (%v) cannot be cast to int", obtained)
		return
	}
	if expected != n {
		a.errorf("wrong Head: expected %d, got %d", expected, n)
	}
}

// HeadErrEmpty asserts that calling Head on the associated queue fails with
// the error ErrEmpty.
func (a *Assert) HeadErrEmpty() {
	obtained, err := a.q.Head()
	if err != queue.ErrEmpty {
		a.errorf("wrong Head: should have failed with ErrEmpty, "+
			"but got %v and an %q error instead", obtained, err)
	}
}

// Dequeue asserts that dequeuing len(expected) times from the
// associated queue success and returns the expected values.
func (a *Assert) Dequeue(expected ...int) {
	for i, e := range expected {
		obtained, err := a.q.Dequeue()
		if err != nil {
			a.errorf("wrong #%d Dequeue: unexpected error: %q", i, err)
			return
		}
		n, ok := obtained.(int)
		if !ok {
			a.errorf("wrong #%d Dequeue: obtained (%v) cannot be cast to int",
				i, obtained)
			return
		}
		if n != e {
			a.errorf("wrong #%d Dequeue: expected %d, got %d", i, e, n)
		}
	}
}

// DequeueErrEmpty asserts that calling Dequeue on the associated queue
// fails with the error ErrEmpty.
func (a *Assert) DequeueErrEmpty() {
	obtained, err := a.q.Dequeue()
	if err != queue.ErrEmpty {
		a.errorf("wrong Dequeue: should have failed with ErrEmpty, "+
			"but got %v and an %q error instead", obtained, err)
	}
}

// Seq returns a slice with the natural numbers from min(a,b) (included)
// to max(a,b) (not included), sorted in ascending order.  This method
// helps generating data for Assert.Enqueue and expected values for
// Assert.Dequeue.
func Seq(a, b int) []int {
	if b < a {
		b, a = a, b
	}
	count := b - a
	ret := make([]int, count)
	for i := 0; i < count; i++ {
		ret[i] = a + i
	}
	return ret
}
<commit_msg>assert: inprove documentation<commit_after>// Package assert provides convenience ways to test queues
// implementations.
package assert

import (
	"fmt"

	"github.com/alcortesm/queue"
)

// Test represents a test that can be set to fail with a formated
// message using the Errorf method, for instance, *testing.T from the
// standard library.
type Test interface {
	Errorf(format string, a ...interface{})
}

// Assert represent an association between a queue.Queue and a Test.
// The methods from this type can be used to manipulate a queue and
// assert that it behaves in particular ways.  If not, the test will
// be notified of the error with a sensible message, prefixed
// with the value of the Prefix field.
type Assert struct {
	// The prefix to use in error messages when an assertion fails.
	Prefix string
	q      queue.Queue
	t      Test
}

// New returns a new association between the given queue and the given
// test.  The Prefix field is initialized to the empty string.
func New(t Test, q queue.Queue) *Assert {
	return &Assert{
		q:      q,
		t:      t,
		Prefix: "",
	}
}

// Calls the Errorf method on the test with the given formated
// arguments, prefixed with the receiver's Prefix value.
func (a *Assert) errorf(format string, v ...interface{}) {
	a.t.Errorf("%s%s", a.Prefix, fmt.Sprintf(format, v...))
}

// Len asserts that the length of the associated queue is the expected value.
func (a *Assert) Len(expected int) {
	obtained := a.q.Len()
	if obtained != expected {
		a.errorf("wrong length: expected %d, got %d", expected, obtained)
	}
}

// IsEmpty asserts that calling IsEmpty on the associated queue returns the
// expected value.
func (a *Assert) IsEmpty(expected bool) {
	obtained := a.q.IsEmpty()
	if obtained != expected {
		a.errorf("wrong IsEmpty: expected %t, got %t", expected, obtained)
	}
}

// Enqueue asserts that enqueuing the given numbers on the associated
// queue success.
func (a *Assert) Enqueue(numbers ...int) {
	for i, n := range numbers {
		if err := a.q.Enqueue(n); err != nil {
			a.errorf("wrong #%d Enqueue(%d): unexpected error: %q",
				i, n, err)
		}
	}
}

// EnqueueErrFull asserts that trying to enqueue any element fails with
// queue.ErrFull.
func (a *Assert) EnqueueErrFull() {
	if err := a.q.Enqueue(42); err != queue.ErrFull {
		a.errorf("wrong Enqueue: should have failed with ErrFull, "+
			"but got %q instead", err)
	}
}

// Head asserts that calling Head on the associated queue success and returns
// the expected number.
func (a *Assert) Head(expected int) {
	obtained, err := a.q.Head()
	if err != nil {
		a.errorf("wrong Head: unexpected error: %s", err)
		return
	}
	n, ok := obtained.(int)
	if !ok {
		a.errorf("wrong Head: obtained (%v) cannot be cast to int", obtained)
		return
	}
	if expected != n {
		a.errorf("wrong Head: expected %d, got %d", expected, n)
	}
}

// HeadErrEmpty asserts that calling Head on the associated queue fails with
// the error ErrEmpty.
func (a *Assert) HeadErrEmpty() {
	obtained, err := a.q.Head()
	if err != queue.ErrEmpty {
		a.errorf("wrong Head: should have failed with ErrEmpty, "+
			"but got %v and an %q error instead", obtained, err)
	}
}

// Dequeue asserts that dequeuing len(expected) times from the
// associated queue success and returns the expected values.
func (a *Assert) Dequeue(expected ...int) {
	for i, e := range expected {
		obtained, err := a.q.Dequeue()
		if err != nil {
			a.errorf("wrong #%d Dequeue: unexpected error: %q", i, err)
			return
		}
		n, ok := obtained.(int)
		if !ok {
			a.errorf("wrong #%d Dequeue: obtained (%v) cannot be cast to int",
				i, obtained)
			return
		}
		if n != e {
			a.errorf("wrong #%d Dequeue: expected %d, got %d", i, e, n)
		}
	}
}

// DequeueErrEmpty asserts that calling Dequeue on the associated queue
// fails with the error ErrEmpty.
func (a *Assert) DequeueErrEmpty() {
	obtained, err := a.q.Dequeue()
	if err != queue.ErrEmpty {
		a.errorf("wrong Dequeue: should have failed with ErrEmpty, "+
			"but got %v and an %q error instead", obtained, err)
	}
}

// Seq returns a slice with the natural numbers from min(a,b) (included)
// to max(a,b) (not included), sorted in ascending order.  This method
// helps generating data for Assert.Enqueue and expected values for
// Assert.Dequeue.
func Seq(a, b int) []int {
	if b < a {
		b, a = a, b
	}
	count := b - a
	ret := make([]int, count)
	for i := 0; i < count; i++ {
		ret[i] = a + i
	}
	return ret
}
<|endoftext|>


<commit_before><commit_msg>Update util_test.go<commit_after><|endoftext|>


<commit_before>package plugins

import (
	"strings"
	"ircclient"
	"bytes"
	"net"
	"crypto/hmac"
	"crypto/rand"
	"hash"
	"log"
	"fmt"
)

type EvaluationPlugin struct {
	ic *ircclient.IRCClient
	c net.Conn
	hmac hash.Hash
	done chan bool
}

func (q *EvaluationPlugin) Register(cl *ircclient.IRCClient) {
	q.ic = cl
	q.done = make(chan bool, 1)
	var hmacKey []byte

	// Read key for HMAC from config
	if value := q.ic.GetStringOption("Eval", "key"); value == "" {
		log.Println("WARNING: No HMAC key for evaluation plugin specified")
		hmacKey = make([]byte, 20)
		if n, err := rand.Read(hmacKey); n != 20 || err != nil {
			log.Printf("ERROR: Unable to generate one. Exiting. (%s)", err.String())
			return
		}
		q.ic.SetStringOption("Eval", "key", fmt.Sprintf("%x", hmacKey))
		log.Println("WARNING: Auto-generated one.")
	} else {
		// Hm. This is hackish. Proposals? :-)
		hmacKey = hmacKey[0:0]
		for len(value) > 1 {
			var tmp int
			fmt.Sscanf(value[0:2], "%x", &tmp)
			hmacKey = append(hmacKey, uint8(tmp))
			value = value[2:]
		}
	}
	// Network config
	var hostport string
	if hostport = q.ic.GetStringOption("Eval", "hostport"); hostport == "" {
		log.Println("WARNING: Added default listener for EvaluationPlugin")
		q.ic.SetStringOption("Eval", "hostport", "0.0.0.0:5486")
		hostport = "0.0.0.0:4386"
	}
	// Channel
	var channel string
	if channel = q.ic.GetStringOption("Eval", "channel"); channel == "" {
		log.Println("WARNING: No channel set in evaluation plugin options. Please set one.")
		return
	}

	addr, err := net.ResolveUDPAddr("udp", hostport)
	if err != nil {
		log.Println("ERROR: Internal error in EvaluationPlugin")
		return
	}
	q.c, err = net.ListenUDP("udp", addr)
	if err != nil {
		log.Println("ERROR: Unable to open listener for evaluation plugin")
		return
	}
	q.hmac = hmac.NewSHA1(hmacKey)
	go func() {
		buf := make([]byte, 512)
		for {
			n, err := q.c.Read(buf)
			select {
			case <- q.done:
				return
			default:
			}
			if err != nil {
				log.Println("ERROR: Unable to receive on evaluation notification")
				return
			}
			if n < q.hmac.Size() + 1 {
				log.Println("WARNING: Invalid evaluation packet received")
				continue
			}
			hash := buf[0:20]
			q.hmac.Reset()
			q.hmac.Write(buf[20:n])
			if bytes.Compare(hash, q.hmac.Sum()) != 0 {
				log.Println("WARNING: Wrong HMAC on evaluation packet")
				continue
			}
			payload := string(buf[20:n])
			parts := strings.SplitN(payload, "\t", 3)
			if len(parts) < 3 {
				continue
			}
			q.ic.SendLine("PRIVMSG #" + channel + " :New lecture evaluation has been added in Semester " + parts[0] + ": " + parts[2] + ", " + parts[1])
		}
	}()
}

func (q *EvaluationPlugin) String() string {
	return "evaluation"
}

func (q *EvaluationPlugin) Info() string {
	return "notifies when new evaluations are available"
}

func (q *EvaluationPlugin) Usage(cmd string) string {
	return "This plugin provides no commands"
}

func (q *EvaluationPlugin) ProcessLine(msg *ircclient.IRCMessage) {
}

func (q *EvaluationPlugin) ProcessCommand(cmd *ircclient.IRCCommand) {
	return
}

func (q *EvaluationPlugin) Unregister() {
	q.done <- true
	if q.c != nil {
		q.c.Close()
	}
	return
}
<commit_msg>Fix eval notify<commit_after>package plugins

import (
	"strings"
	"ircclient"
	"bytes"
	"net"
	"crypto/hmac"
	"crypto/rand"
	"hash"
	"log"
	"fmt"
)

type EvaluationPlugin struct {
	ic *ircclient.IRCClient
	c net.Conn
	hmac hash.Hash
	done chan bool
}

func (q *EvaluationPlugin) Register(cl *ircclient.IRCClient) {
	q.ic = cl
	q.done = make(chan bool, 1)
	var hmacKey []byte

	// Read key for HMAC from config
	if value := q.ic.GetStringOption("Eval", "key"); value == "" {
		log.Println("WARNING: No HMAC key for evaluation plugin specified")
		hmacKey = make([]byte, 20)
		if n, err := rand.Read(hmacKey); n != 20 || err != nil {
			log.Printf("ERROR: Unable to generate one. Exiting. (%s)", err.String())
			return
		}
		q.ic.SetStringOption("Eval", "key", fmt.Sprintf("%x", hmacKey))
		log.Println("WARNING: Auto-generated one.")
	} else {
		// Hm. This is hackish. Proposals? :-)
		hmacKey = hmacKey[0:0]
		for len(value) > 1 {
			var tmp int
			fmt.Sscanf(value[0:2], "%x", &tmp)
			hmacKey = append(hmacKey, uint8(tmp))
			value = value[2:]
		}
	}
	// Network config
	var hostport string
	if hostport = q.ic.GetStringOption("Eval", "hostport"); hostport == "" {
		log.Println("WARNING: Added default listener for EvaluationPlugin")
		q.ic.SetStringOption("Eval", "hostport", "0.0.0.0:5486")
		hostport = "0.0.0.0:4386"
	}
	// Channel
	var channel string
	if channel = q.ic.GetStringOption("Eval", "channel"); channel == "" {
		log.Println("WARNING: No channel set in evaluation plugin options. Please set one.")
		return
	}

	addr, err := net.ResolveUDPAddr("udp", hostport)
	if err != nil {
		log.Println("ERROR: Internal error in EvaluationPlugin")
		return
	}
	q.c, err = net.ListenUDP("udp", addr)
	if err != nil {
		log.Println("ERROR: Unable to open listener for evaluation plugin")
		return
	}
	q.hmac = hmac.NewSHA1(hmacKey)
	go func() {
		buf := make([]byte, 512)
		for {
			n, err := q.c.Read(buf)
			select {
			case <- q.done:
				return
			default:
			}
			if err != nil {
				log.Println("ERROR: Unable to receive on evaluation notification")
				return
			}
			if n < q.hmac.Size() + 1 {
				log.Println("WARNING: Invalid evaluation packet received")
				continue
			}
			hash := buf[0:20]
			q.hmac.Reset()
			q.hmac.Write(buf[20:n])
			if bytes.Compare(hash, q.hmac.Sum()) != 0 {
				log.Println("WARNING: Wrong HMAC on evaluation packet")
				continue
			}
			payload := string(buf[20:n])
			parts := strings.SplitN(payload, "\t", 3)
			if len(parts) < 3 {
				continue
			}
			q.ic.SendLine("PRIVMSG #" + channel + " :New lecture evaluation has been added in Semester " + parts[1] + ": " + parts[0] + ", " + parts[2])
		}
	}()
}

func (q *EvaluationPlugin) String() string {
	return "evaluation"
}

func (q *EvaluationPlugin) Info() string {
	return "notifies when new evaluations are available"
}

func (q *EvaluationPlugin) Usage(cmd string) string {
	return "This plugin provides no commands"
}

func (q *EvaluationPlugin) ProcessLine(msg *ircclient.IRCMessage) {
}

func (q *EvaluationPlugin) ProcessCommand(cmd *ircclient.IRCCommand) {
	return
}

func (q *EvaluationPlugin) Unregister() {
	q.done <- true
	if q.c != nil {
		q.c.Close()
	}
	return
}
<|endoftext|>


<commit_before><commit_msg>fix GetVolume for Windows<commit_after><|endoftext|>


<commit_before>package tsa_test

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

	"github.com/cloudfoundry-incubator/garden"
	gfakes "github.com/cloudfoundry-incubator/garden/fakes"
	"github.com/concourse/atc"
	. "github.com/concourse/tsa"
	"github.com/concourse/tsa/fakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/ghttp"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/lagertest"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/tedsuo/rata"
)

var _ = Describe("Heartbeater", func() {
	type registration struct {
		worker       atc.Worker
		ttl          time.Duration
		lastInterval time.Duration
	}

	var (
		logger lager.Logger

		addrToRegister string
		interval       time.Duration
		cprInterval    time.Duration
		resourceTypes  []atc.WorkerResourceType

		expectedWorker     atc.Worker
		fakeTokenGenerator *fakes.FakeTokenGenerator
		fakeGardenClient   *gfakes.FakeClient
		fakeATC            *ghttp.Server

		heartbeater ifrit.Process

		verifyHeartbeat http.HandlerFunc

		registrations <-chan registration
		clientWriter  *gbytes.Buffer
	)

	BeforeEach(func() {
		logger = lagertest.NewTestLogger("test")

		addrToRegister = "1.2.3.4:7777"
		interval = time.Second
		cprInterval = 100 * time.Millisecond
		resourceTypes = []atc.WorkerResourceType{
			{
				Type:  "git",
				Image: "docker:///concourse/git-resource",
			},
		}

		expectedWorker = atc.Worker{
			GardenAddr:       addrToRegister,
			ActiveContainers: 2,
			ResourceTypes:    resourceTypes,
			Platform:         "some-platform",
			Tags:             []string{"some", "tags"},
		}

		fakeATC = ghttp.NewServer()

		registerRoute, found := atc.Routes.FindRouteByName(atc.RegisterWorker)
		Expect(found).To(BeTrue())

		registered := make(chan registration, 100)
		registrations = registered

		lastRequestTime := time.Now()
		verifyHeartbeat = ghttp.CombineHandlers(
			ghttp.VerifyRequest(registerRoute.Method, registerRoute.Path),
			func(w http.ResponseWriter, r *http.Request) {
				var worker atc.Worker
				Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

				err := json.NewDecoder(r.Body).Decode(&worker)
				Expect(err).NotTo(HaveOccurred())

				ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
				Expect(err).NotTo(HaveOccurred())

				requestTime := time.Now()
				lastInterval := requestTime.Sub(lastRequestTime)
				lastRequestTime = requestTime

				registered <- registration{worker, ttl, lastInterval}
			},
		)

		fakeGardenClient = new(gfakes.FakeClient)
		fakeTokenGenerator = new(fakes.FakeTokenGenerator)

		fakeTokenGenerator.GenerateTokenReturns("yo", nil)
		clientWriter = gbytes.NewBuffer()
	})

	JustBeforeEach(func() {
		atcEndpoint := rata.NewRequestGenerator(fakeATC.URL(), atc.Routes)
		heartbeater = ifrit.Invoke(
			NewHeartbeater(
				logger,
				interval,
				cprInterval,
				fakeGardenClient,
				atcEndpoint,
				fakeTokenGenerator,
				atc.Worker{
					GardenAddr:    addrToRegister,
					ResourceTypes: resourceTypes,
					Platform:      "some-platform",
					Tags:          []string{"some", "tags"},
				},
				clientWriter,
			),
		)
	})

	AfterEach(func() {
		ginkgomon.Interrupt(heartbeater)
	})

	Context("when Garden returns containers", func() {
		BeforeEach(func() {
			containers := make(chan []garden.Container, 4)

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			close(containers)

			fakeGardenClient.ContainersStub = func(garden.Properties) ([]garden.Container, error) {
				return <-containers, nil
			}
		})

		Context("when the ATC responds to heartbeat requests", func() {
			BeforeEach(func() {
				fakeATC.AppendHandlers(verifyHeartbeat, verifyHeartbeat)
			})

			It("immediately registers", func() {
				actualRegistration := <-registrations
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", 0, 20*time.Millisecond))
			})

			It("heartbeats", func() {
				<-registrations
				actualRegistration := <-registrations
				expectedWorker.ActiveContainers = 5
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", interval, 20*time.Millisecond))
			})
		})

		Context("when the ATC doesn't respond to the first heartbeat", func() {
			BeforeEach(func() {
				fakeATC.AppendHandlers(
					verifyHeartbeat,
					ghttp.CombineHandlers(
						verifyHeartbeat,
						func(w http.ResponseWriter, r *http.Request) { fakeATC.CloseClientConnections() },
					),
					verifyHeartbeat,
					verifyHeartbeat,
				)
			})

			It("heartbeats faster according to cprInterval", func() {
				<-registrations
				<-registrations
				actualRegistration := <-registrations
				expectedWorker.ActiveContainers = 4
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", cprInterval, 20*time.Millisecond))
			})

			It("goes back to normal after the heartbeat succeeds", func() {
				<-registrations
				<-registrations
				<-registrations
				actualRegistration := <-registrations
				expectedWorker.ActiveContainers = 3
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", interval, 20*time.Millisecond))
			})
		})
	})
})
<commit_msg>bump timeout<commit_after>package tsa_test

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

	"github.com/cloudfoundry-incubator/garden"
	gfakes "github.com/cloudfoundry-incubator/garden/fakes"
	"github.com/concourse/atc"
	. "github.com/concourse/tsa"
	"github.com/concourse/tsa/fakes"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/ghttp"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/lager/lagertest"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/tedsuo/rata"
)

var _ = Describe("Heartbeater", func() {
	type registration struct {
		worker       atc.Worker
		ttl          time.Duration
		lastInterval time.Duration
	}

	var (
		logger lager.Logger

		addrToRegister string
		interval       time.Duration
		cprInterval    time.Duration
		resourceTypes  []atc.WorkerResourceType

		expectedWorker     atc.Worker
		fakeTokenGenerator *fakes.FakeTokenGenerator
		fakeGardenClient   *gfakes.FakeClient
		fakeATC            *ghttp.Server

		heartbeater ifrit.Process

		verifyHeartbeat http.HandlerFunc

		registrations <-chan registration
		clientWriter  *gbytes.Buffer
	)

	BeforeEach(func() {
		logger = lagertest.NewTestLogger("test")

		addrToRegister = "1.2.3.4:7777"
		interval = time.Second
		cprInterval = 100 * time.Millisecond
		resourceTypes = []atc.WorkerResourceType{
			{
				Type:  "git",
				Image: "docker:///concourse/git-resource",
			},
		}

		expectedWorker = atc.Worker{
			GardenAddr:       addrToRegister,
			ActiveContainers: 2,
			ResourceTypes:    resourceTypes,
			Platform:         "some-platform",
			Tags:             []string{"some", "tags"},
		}

		fakeATC = ghttp.NewServer()

		registerRoute, found := atc.Routes.FindRouteByName(atc.RegisterWorker)
		Expect(found).To(BeTrue())

		registered := make(chan registration, 100)
		registrations = registered

		lastRequestTime := time.Now()
		verifyHeartbeat = ghttp.CombineHandlers(
			ghttp.VerifyRequest(registerRoute.Method, registerRoute.Path),
			func(w http.ResponseWriter, r *http.Request) {
				var worker atc.Worker
				Expect(r.Header.Get("Authorization")).To(Equal("Bearer yo"))

				err := json.NewDecoder(r.Body).Decode(&worker)
				Expect(err).NotTo(HaveOccurred())

				ttl, err := time.ParseDuration(r.URL.Query().Get("ttl"))
				Expect(err).NotTo(HaveOccurred())

				requestTime := time.Now()
				lastInterval := requestTime.Sub(lastRequestTime)
				lastRequestTime = requestTime

				registered <- registration{worker, ttl, lastInterval}
			},
		)

		fakeGardenClient = new(gfakes.FakeClient)
		fakeTokenGenerator = new(fakes.FakeTokenGenerator)

		fakeTokenGenerator.GenerateTokenReturns("yo", nil)
		clientWriter = gbytes.NewBuffer()
	})

	JustBeforeEach(func() {
		atcEndpoint := rata.NewRequestGenerator(fakeATC.URL(), atc.Routes)
		heartbeater = ifrit.Invoke(
			NewHeartbeater(
				logger,
				interval,
				cprInterval,
				fakeGardenClient,
				atcEndpoint,
				fakeTokenGenerator,
				atc.Worker{
					GardenAddr:    addrToRegister,
					ResourceTypes: resourceTypes,
					Platform:      "some-platform",
					Tags:          []string{"some", "tags"},
				},
				clientWriter,
			),
		)
	})

	AfterEach(func() {
		ginkgomon.Interrupt(heartbeater)
	})

	Context("when Garden returns containers", func() {
		BeforeEach(func() {
			containers := make(chan []garden.Container, 4)

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			containers <- []garden.Container{
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
				new(gfakes.FakeContainer),
			}

			close(containers)

			fakeGardenClient.ContainersStub = func(garden.Properties) ([]garden.Container, error) {
				return <-containers, nil
			}
		})

		Context("when the ATC responds to heartbeat requests", func() {
			BeforeEach(func() {
				fakeATC.AppendHandlers(verifyHeartbeat, verifyHeartbeat)
			})

			It("immediately registers", func() {
				actualRegistration := <-registrations
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", 0, 100*time.Millisecond))
			})

			It("heartbeats", func() {
				<-registrations
				actualRegistration := <-registrations
				expectedWorker.ActiveContainers = 5
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", interval, 20*time.Millisecond))
			})
		})

		Context("when the ATC doesn't respond to the first heartbeat", func() {
			BeforeEach(func() {
				fakeATC.AppendHandlers(
					verifyHeartbeat,
					ghttp.CombineHandlers(
						verifyHeartbeat,
						func(w http.ResponseWriter, r *http.Request) { fakeATC.CloseClientConnections() },
					),
					verifyHeartbeat,
					verifyHeartbeat,
				)
			})

			It("heartbeats faster according to cprInterval", func() {
				<-registrations
				<-registrations
				actualRegistration := <-registrations
				expectedWorker.ActiveContainers = 4
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", cprInterval, 20*time.Millisecond))
			})

			It("goes back to normal after the heartbeat succeeds", func() {
				<-registrations
				<-registrations
				<-registrations
				actualRegistration := <-registrations
				expectedWorker.ActiveContainers = 3
				Expect(actualRegistration.worker).To(Equal(expectedWorker))
				Expect(actualRegistration.ttl).To(Equal(2 * interval))
				Expect(actualRegistration.lastInterval).To(BeNumerically("~", interval, 20*time.Millisecond))
			})
		})
	})
})
<|endoftext|>


<commit_before>package adapter

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"

	"golang.org/x/oauth2"
)

// Provider is an OAuth provider.
type Provider interface {
	ParseConfig(configFile map[string]interface{}) (ProviderConfig, error)
}

// ProviderConfig is a config for an OAuth provider.
type ProviderConfig interface {
	Config() oauth2.Config
	Info(c *oauth2.Config, t *oauth2.Token) (string, map[string]interface{}, error)
}

// ProviderInfoContext is for support context.Context.
type ProviderInfoContext interface {
	InfoContext(ctx context.Context, c *oauth2.Config, t *oauth2.Token) (string, map[string]interface{}, error)
}

var providers = map[string]Provider{}

// RegisterProvider registers the OAuth provider.
func RegisterProvider(name string, provider Provider) {
	providers[name] = provider
}

func init() {
	RegisterProvider("development", providerDevelopment{})
}

type providerDevelopment struct{}
type providerConfigDevelopment struct {
	listener net.Listener
}

// ParseConfig parses the config for the development provider.
func (providerDevelopment) ParseConfig(configFile map[string]interface{}) (ProviderConfig, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	c := providerConfigDevelopment{
		listener: l,
	}
	go http.Serve(l, c)
	return c, nil
}

func (pc providerConfigDevelopment) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/auth":
		v := url.Values{}
		v.Add("state", r.FormValue("state"))
		v.Add("code", "development-code")
		http.Redirect(w, r, r.FormValue("redirect_uri")+"?"+v.Encode(), http.StatusFound)
	case "/token":
		fmt.Fprintln(w, "{}")
	default:
		http.NotFound(w, r)
	}
}

func (pc providerConfigDevelopment) Config() oauth2.Config {
	l := pc.listener
	return oauth2.Config{
		Endpoint: oauth2.Endpoint{
			AuthURL:  "http://" + l.Addr().String() + "/auth",
			TokenURL: "http://" + l.Addr().String() + "/token",
		},
		ClientID:     "development@example.com",
		ClientSecret: "development-client-secret",
		Scopes:       []string{},
	}
}

func (pc providerConfigDevelopment) Info(c *oauth2.Config, t *oauth2.Token) (string, map[string]interface{}, error) {
	return "developer", map[string]interface{}{}, nil
}
<commit_msg>the development provider returns access_token<commit_after>package adapter

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"

	"golang.org/x/oauth2"
)

// Provider is an OAuth provider.
type Provider interface {
	ParseConfig(configFile map[string]interface{}) (ProviderConfig, error)
}

// ProviderConfig is a config for an OAuth provider.
type ProviderConfig interface {
	Config() oauth2.Config
	Info(c *oauth2.Config, t *oauth2.Token) (string, map[string]interface{}, error)
}

// ProviderInfoContext is for support context.Context.
type ProviderInfoContext interface {
	InfoContext(ctx context.Context, c *oauth2.Config, t *oauth2.Token) (string, map[string]interface{}, error)
}

var providers = map[string]Provider{}

// RegisterProvider registers the OAuth provider.
func RegisterProvider(name string, provider Provider) {
	providers[name] = provider
}

func init() {
	RegisterProvider("development", providerDevelopment{})
}

type providerDevelopment struct{}
type providerConfigDevelopment struct {
	listener net.Listener
}

// ParseConfig parses the config for the development provider.
func (providerDevelopment) ParseConfig(configFile map[string]interface{}) (ProviderConfig, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	c := providerConfigDevelopment{
		listener: l,
	}
	go http.Serve(l, c)
	return c, nil
}

func (pc providerConfigDevelopment) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/auth":
		v := url.Values{}
		v.Add("state", r.FormValue("state"))
		v.Add("code", "development-code")
		http.Redirect(w, r, r.FormValue("redirect_uri")+"?"+v.Encode(), http.StatusFound)
	case "/token":
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintln(w, `{"access_token":"hogehoge"}`)
	default:
		http.NotFound(w, r)
	}
}

func (pc providerConfigDevelopment) Config() oauth2.Config {
	l := pc.listener
	return oauth2.Config{
		Endpoint: oauth2.Endpoint{
			AuthURL:  "http://" + l.Addr().String() + "/auth",
			TokenURL: "http://" + l.Addr().String() + "/token",
		},
		ClientID:     "development@example.com",
		ClientSecret: "development-client-secret",
		Scopes:       []string{},
	}
}

func (pc providerConfigDevelopment) Info(c *oauth2.Config, t *oauth2.Token) (string, map[string]interface{}, error) {
	return "developer", map[string]interface{}{}, nil
}
<|endoftext|>


<commit_before>package pool

import (
	"container/list"
	"fmt"
	"io"
	"log"
	"net"
	"net/rpc"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/lib"
	hclog "github.com/hashicorp/go-hclog"
	msgpackrpc "github.com/hashicorp/net-rpc-msgpackrpc"
	"github.com/hashicorp/nomad/helper/tlsutil"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/yamux"
)

// NewClientCodec returns a new rpc.ClientCodec to be used to make RPC calls.
func NewClientCodec(conn io.ReadWriteCloser) rpc.ClientCodec {
	return msgpackrpc.NewCodecFromHandle(true, true, conn, structs.HashiMsgpackHandle)
}

// NewServerCodec returns a new rpc.ServerCodec to be used to handle RPCs.
func NewServerCodec(conn io.ReadWriteCloser) rpc.ServerCodec {
	return msgpackrpc.NewCodecFromHandle(true, true, conn, structs.HashiMsgpackHandle)
}

// streamClient is used to wrap a stream with an RPC client
type StreamClient struct {
	stream net.Conn
	codec  rpc.ClientCodec
}

func (sc *StreamClient) Close() {
	sc.stream.Close()
	sc.codec.Close()
}

// Conn is a pooled connection to a Nomad server
type Conn struct {
	refCount    int32
	shouldClose int32

	addr     net.Addr
	session  *yamux.Session
	lastUsed time.Time
	version  int

	pool *ConnPool

	clients    *list.List
	clientLock sync.Mutex
}

// markForUse does all the bookkeeping required to ready a connection for use.
func (c *Conn) markForUse() {
	c.lastUsed = time.Now()
	atomic.AddInt32(&c.refCount, 1)
}

func (c *Conn) Close() error {
	return c.session.Close()
}

// getClient is used to get a cached or new client
func (c *Conn) getClient() (*StreamClient, error) {
	// Check for cached client
	c.clientLock.Lock()
	front := c.clients.Front()
	if front != nil {
		c.clients.Remove(front)
	}
	c.clientLock.Unlock()
	if front != nil {
		return front.Value.(*StreamClient), nil
	}

	// Open a new session
	stream, err := c.session.Open()
	if err != nil {
		return nil, err
	}

	// Create a client codec
	codec := NewClientCodec(stream)

	// Return a new stream client
	sc := &StreamClient{
		stream: stream,
		codec:  codec,
	}
	return sc, nil
}

// returnClient is used when done with a stream
// to allow re-use by a future RPC
func (c *Conn) returnClient(client *StreamClient) {
	didSave := false
	c.clientLock.Lock()
	if c.clients.Len() < c.pool.maxStreams && atomic.LoadInt32(&c.shouldClose) == 0 {
		c.clients.PushFront(client)
		didSave = true

		// If this is a Yamux stream, shrink the internal buffers so that
		// we can GC the idle memory
		if ys, ok := client.stream.(*yamux.Stream); ok {
			ys.Shrink()
		}
	}
	c.clientLock.Unlock()
	if !didSave {
		client.Close()
	}
}

// ConnPool is used to maintain a connection pool to other
// Nomad servers. This is used to reduce the latency of
// RPC requests between servers. It is only used to pool
// connections in the rpcNomad mode. Raft connections
// are pooled separately.
type ConnPool struct {
	sync.Mutex

	// logger is the logger to be used
	logger *log.Logger

	// The maximum time to keep a connection open
	maxTime time.Duration

	// The maximum number of open streams to keep
	maxStreams int

	// Pool maps an address to a open connection
	pool map[string]*Conn

	// limiter is used to throttle the number of connect attempts
	// to a given address. The first thread will attempt a connection
	// and put a channel in here, which all other threads will wait
	// on to close.
	limiter map[string]chan struct{}

	// TLS wrapper
	tlsWrap tlsutil.RegionWrapper

	// Used to indicate the pool is shutdown
	shutdown   bool
	shutdownCh chan struct{}

	// connListener is used to notify a potential listener of a new connection
	// being made.
	connListener chan<- *yamux.Session
}

// NewPool is used to make a new connection pool
// Maintain at most one connection per host, for up to maxTime.
// Set maxTime to 0 to disable reaping. maxStreams is used to control
// the number of idle streams allowed.
// If TLS settings are provided outgoing connections use TLS.
func NewPool(logger hclog.Logger, maxTime time.Duration, maxStreams int, tlsWrap tlsutil.RegionWrapper) *ConnPool {
	pool := &ConnPool{
		logger:     logger.StandardLogger(&hclog.StandardLoggerOptions{InferLevels: true}),
		maxTime:    maxTime,
		maxStreams: maxStreams,
		pool:       make(map[string]*Conn),
		limiter:    make(map[string]chan struct{}),
		tlsWrap:    tlsWrap,
		shutdownCh: make(chan struct{}),
	}
	if maxTime > 0 {
		go pool.reap()
	}
	return pool
}

// Shutdown is used to close the connection pool
func (p *ConnPool) Shutdown() error {
	p.Lock()
	defer p.Unlock()

	for _, conn := range p.pool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)

	if p.shutdown {
		return nil
	}

	if p.connListener != nil {
		close(p.connListener)
		p.connListener = nil
	}

	p.shutdown = true
	close(p.shutdownCh)
	return nil
}

// ReloadTLS reloads TLS configuration on the fly
func (p *ConnPool) ReloadTLS(tlsWrap tlsutil.RegionWrapper) {
	p.Lock()
	defer p.Unlock()

	oldPool := p.pool
	for _, conn := range oldPool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)
	p.tlsWrap = tlsWrap
}

// SetConnListener is used to listen to new connections being made. The
// channel will be closed when the conn pool is closed or a new listener is set.
func (p *ConnPool) SetConnListener(l chan<- *yamux.Session) {
	p.Lock()
	defer p.Unlock()

	// Close the old listener
	if p.connListener != nil {
		close(p.connListener)
	}

	// Store the new listener
	p.connListener = l
}

// Acquire is used to get a connection that is
// pooled or to return a new connection
func (p *ConnPool) acquire(region string, addr net.Addr, version int) (*Conn, error) {
	// Check to see if there's a pooled connection available. This is up
	// here since it should the vastly more common case than the rest
	// of the code here.
	p.Lock()
	c := p.pool[addr.String()]
	if c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	// If not (while we are still locked), set up the throttling structure
	// for this address, which will make everyone else wait until our
	// attempt is done.
	var wait chan struct{}
	var ok bool
	if wait, ok = p.limiter[addr.String()]; !ok {
		wait = make(chan struct{})
		p.limiter[addr.String()] = wait
	}
	isLeadThread := !ok
	p.Unlock()

	// If we are the lead thread, make the new connection and then wake
	// everybody else up to see if we got it.
	if isLeadThread {
		c, err := p.getNewConn(region, addr, version)
		p.Lock()
		delete(p.limiter, addr.String())
		close(wait)
		if err != nil {
			p.Unlock()
			return nil, err
		}

		p.pool[addr.String()] = c

		// If there is a connection listener, notify them of the new connection.
		if p.connListener != nil {
			select {
			case p.connListener <- c.session:
			default:
			}
		}

		p.Unlock()
		return c, nil
	}

	// Otherwise, wait for the lead thread to attempt the connection
	// and use what's in the pool at that point.
	select {
	case <-p.shutdownCh:
		return nil, fmt.Errorf("rpc error: shutdown")
	case <-wait:
	}

	// See if the lead thread was able to get us a connection.
	p.Lock()
	if c := p.pool[addr.String()]; c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	p.Unlock()
	return nil, fmt.Errorf("rpc error: lead thread didn't get connection")
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(region string, addr net.Addr, version int) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsWrap != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(RpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		tlsConn, err := p.tlsWrap(region, conn)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

	// Write the multiplex byte to set the mode
	if _, err := conn.Write([]byte{byte(RpcMultiplex)}); err != nil {
		conn.Close()
		return nil, err
	}

	// Setup the logger
	conf := yamux.DefaultConfig()
	conf.LogOutput = nil
	conf.Logger = p.logger

	// Create a multiplexed session
	session, err := yamux.Client(conn, conf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Wrap the connection
	c := &Conn{
		refCount: 1,
		addr:     addr,
		session:  session,
		clients:  list.New(),
		lastUsed: time.Now(),
		version:  version,
		pool:     p,
	}
	return c, nil
}

// clearConn is used to clear any cached connection, potentially in response to
// an error
func (p *ConnPool) clearConn(conn *Conn) {
	// Ensure returned streams are closed
	atomic.StoreInt32(&conn.shouldClose, 1)

	// Clear from the cache
	p.Lock()
	if c, ok := p.pool[conn.addr.String()]; ok && c == conn {
		delete(p.pool, conn.addr.String())
	}
	p.Unlock()

	// Close down immediately if idle
	if refCount := atomic.LoadInt32(&conn.refCount); refCount == 0 {
		conn.Close()
	}
}

// releaseConn is invoked when we are done with a conn to reduce the ref count
func (p *ConnPool) releaseConn(conn *Conn) {
	refCount := atomic.AddInt32(&conn.refCount, -1)
	if refCount == 0 && atomic.LoadInt32(&conn.shouldClose) == 1 {
		conn.Close()
	}
}

// getClient is used to get a usable client for an address and protocol version
func (p *ConnPool) getClient(region string, addr net.Addr, version int) (*Conn, *StreamClient, error) {
	retries := 0
START:
	// Try to get a conn first
	conn, err := p.acquire(region, addr, version)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get conn: %v", err)
	}

	// Get a client
	client, err := conn.getClient()
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)

		// Try to redial, possible that the TCP session closed due to timeout
		if retries == 0 {
			retries++
			goto START
		}
		return nil, nil, fmt.Errorf("failed to start stream: %v", err)
	}
	return conn, client, nil
}

// RPC is used to make an RPC call to a remote host
func (p *ConnPool) RPC(region string, addr net.Addr, version int, method string, args interface{}, reply interface{}) error {
	// Get a usable client
	conn, sc, err := p.getClient(region, addr, version)
	if err != nil {
		return fmt.Errorf("rpc error: %v", err)
	}

	// Make the RPC call
	err = msgpackrpc.CallWithCodec(sc.codec, method, args, reply)
	if err != nil {
		sc.Close()
		p.releaseConn(conn)

		// If we read EOF, the session is toast. Clear it and open a
		// new session next time
		// See https://github.com/hashicorp/consul/blob/v1.6.3/agent/pool/pool.go#L471-L477
		if lib.IsErrEOF(err) {
			p.clearConn(conn)
		}

		// If the error is an RPC Coded error
		// return the coded error without wrapping
		if structs.IsErrRPCCoded(err) {
			return err
		}

		// TODO wrap with RPCCoded error instead
		return fmt.Errorf("rpc error: %v", err)
	}

	// Done with the connection
	conn.returnClient(sc)
	p.releaseConn(conn)
	return nil
}

// Reap is used to close conns open over maxTime
func (p *ConnPool) reap() {
	for {
		// Sleep for a while
		select {
		case <-p.shutdownCh:
			return
		case <-time.After(time.Second):
		}

		// Reap all old conns
		p.Lock()
		var removed []string
		now := time.Now()
		for host, conn := range p.pool {
			// Skip recently used connections
			if now.Sub(conn.lastUsed) < p.maxTime {
				continue
			}

			// Skip connections with active streams
			if atomic.LoadInt32(&conn.refCount) > 0 {
				continue
			}

			// Close the conn
			conn.Close()

			// Remove from pool
			removed = append(removed, host)
		}
		for _, host := range removed {
			delete(p.pool, host)
		}
		p.Unlock()
	}
}
<commit_msg>pool: Clear connection before releasing<commit_after>package pool

import (
	"container/list"
	"fmt"
	"io"
	"log"
	"net"
	"net/rpc"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/lib"
	hclog "github.com/hashicorp/go-hclog"
	msgpackrpc "github.com/hashicorp/net-rpc-msgpackrpc"
	"github.com/hashicorp/nomad/helper/tlsutil"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/yamux"
)

// NewClientCodec returns a new rpc.ClientCodec to be used to make RPC calls.
func NewClientCodec(conn io.ReadWriteCloser) rpc.ClientCodec {
	return msgpackrpc.NewCodecFromHandle(true, true, conn, structs.HashiMsgpackHandle)
}

// NewServerCodec returns a new rpc.ServerCodec to be used to handle RPCs.
func NewServerCodec(conn io.ReadWriteCloser) rpc.ServerCodec {
	return msgpackrpc.NewCodecFromHandle(true, true, conn, structs.HashiMsgpackHandle)
}

// streamClient is used to wrap a stream with an RPC client
type StreamClient struct {
	stream net.Conn
	codec  rpc.ClientCodec
}

func (sc *StreamClient) Close() {
	sc.stream.Close()
	sc.codec.Close()
}

// Conn is a pooled connection to a Nomad server
type Conn struct {
	refCount    int32
	shouldClose int32

	addr     net.Addr
	session  *yamux.Session
	lastUsed time.Time
	version  int

	pool *ConnPool

	clients    *list.List
	clientLock sync.Mutex
}

// markForUse does all the bookkeeping required to ready a connection for use.
func (c *Conn) markForUse() {
	c.lastUsed = time.Now()
	atomic.AddInt32(&c.refCount, 1)
}

func (c *Conn) Close() error {
	return c.session.Close()
}

// getClient is used to get a cached or new client
func (c *Conn) getClient() (*StreamClient, error) {
	// Check for cached client
	c.clientLock.Lock()
	front := c.clients.Front()
	if front != nil {
		c.clients.Remove(front)
	}
	c.clientLock.Unlock()
	if front != nil {
		return front.Value.(*StreamClient), nil
	}

	// Open a new session
	stream, err := c.session.Open()
	if err != nil {
		return nil, err
	}

	// Create a client codec
	codec := NewClientCodec(stream)

	// Return a new stream client
	sc := &StreamClient{
		stream: stream,
		codec:  codec,
	}
	return sc, nil
}

// returnClient is used when done with a stream
// to allow re-use by a future RPC
func (c *Conn) returnClient(client *StreamClient) {
	didSave := false
	c.clientLock.Lock()
	if c.clients.Len() < c.pool.maxStreams && atomic.LoadInt32(&c.shouldClose) == 0 {
		c.clients.PushFront(client)
		didSave = true

		// If this is a Yamux stream, shrink the internal buffers so that
		// we can GC the idle memory
		if ys, ok := client.stream.(*yamux.Stream); ok {
			ys.Shrink()
		}
	}
	c.clientLock.Unlock()
	if !didSave {
		client.Close()
	}
}

// ConnPool is used to maintain a connection pool to other
// Nomad servers. This is used to reduce the latency of
// RPC requests between servers. It is only used to pool
// connections in the rpcNomad mode. Raft connections
// are pooled separately.
type ConnPool struct {
	sync.Mutex

	// logger is the logger to be used
	logger *log.Logger

	// The maximum time to keep a connection open
	maxTime time.Duration

	// The maximum number of open streams to keep
	maxStreams int

	// Pool maps an address to a open connection
	pool map[string]*Conn

	// limiter is used to throttle the number of connect attempts
	// to a given address. The first thread will attempt a connection
	// and put a channel in here, which all other threads will wait
	// on to close.
	limiter map[string]chan struct{}

	// TLS wrapper
	tlsWrap tlsutil.RegionWrapper

	// Used to indicate the pool is shutdown
	shutdown   bool
	shutdownCh chan struct{}

	// connListener is used to notify a potential listener of a new connection
	// being made.
	connListener chan<- *yamux.Session
}

// NewPool is used to make a new connection pool
// Maintain at most one connection per host, for up to maxTime.
// Set maxTime to 0 to disable reaping. maxStreams is used to control
// the number of idle streams allowed.
// If TLS settings are provided outgoing connections use TLS.
func NewPool(logger hclog.Logger, maxTime time.Duration, maxStreams int, tlsWrap tlsutil.RegionWrapper) *ConnPool {
	pool := &ConnPool{
		logger:     logger.StandardLogger(&hclog.StandardLoggerOptions{InferLevels: true}),
		maxTime:    maxTime,
		maxStreams: maxStreams,
		pool:       make(map[string]*Conn),
		limiter:    make(map[string]chan struct{}),
		tlsWrap:    tlsWrap,
		shutdownCh: make(chan struct{}),
	}
	if maxTime > 0 {
		go pool.reap()
	}
	return pool
}

// Shutdown is used to close the connection pool
func (p *ConnPool) Shutdown() error {
	p.Lock()
	defer p.Unlock()

	for _, conn := range p.pool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)

	if p.shutdown {
		return nil
	}

	if p.connListener != nil {
		close(p.connListener)
		p.connListener = nil
	}

	p.shutdown = true
	close(p.shutdownCh)
	return nil
}

// ReloadTLS reloads TLS configuration on the fly
func (p *ConnPool) ReloadTLS(tlsWrap tlsutil.RegionWrapper) {
	p.Lock()
	defer p.Unlock()

	oldPool := p.pool
	for _, conn := range oldPool {
		conn.Close()
	}
	p.pool = make(map[string]*Conn)
	p.tlsWrap = tlsWrap
}

// SetConnListener is used to listen to new connections being made. The
// channel will be closed when the conn pool is closed or a new listener is set.
func (p *ConnPool) SetConnListener(l chan<- *yamux.Session) {
	p.Lock()
	defer p.Unlock()

	// Close the old listener
	if p.connListener != nil {
		close(p.connListener)
	}

	// Store the new listener
	p.connListener = l
}

// Acquire is used to get a connection that is
// pooled or to return a new connection
func (p *ConnPool) acquire(region string, addr net.Addr, version int) (*Conn, error) {
	// Check to see if there's a pooled connection available. This is up
	// here since it should the vastly more common case than the rest
	// of the code here.
	p.Lock()
	c := p.pool[addr.String()]
	if c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	// If not (while we are still locked), set up the throttling structure
	// for this address, which will make everyone else wait until our
	// attempt is done.
	var wait chan struct{}
	var ok bool
	if wait, ok = p.limiter[addr.String()]; !ok {
		wait = make(chan struct{})
		p.limiter[addr.String()] = wait
	}
	isLeadThread := !ok
	p.Unlock()

	// If we are the lead thread, make the new connection and then wake
	// everybody else up to see if we got it.
	if isLeadThread {
		c, err := p.getNewConn(region, addr, version)
		p.Lock()
		delete(p.limiter, addr.String())
		close(wait)
		if err != nil {
			p.Unlock()
			return nil, err
		}

		p.pool[addr.String()] = c

		// If there is a connection listener, notify them of the new connection.
		if p.connListener != nil {
			select {
			case p.connListener <- c.session:
			default:
			}
		}

		p.Unlock()
		return c, nil
	}

	// Otherwise, wait for the lead thread to attempt the connection
	// and use what's in the pool at that point.
	select {
	case <-p.shutdownCh:
		return nil, fmt.Errorf("rpc error: shutdown")
	case <-wait:
	}

	// See if the lead thread was able to get us a connection.
	p.Lock()
	if c := p.pool[addr.String()]; c != nil {
		c.markForUse()
		p.Unlock()
		return c, nil
	}

	p.Unlock()
	return nil, fmt.Errorf("rpc error: lead thread didn't get connection")
}

// getNewConn is used to return a new connection
func (p *ConnPool) getNewConn(region string, addr net.Addr, version int) (*Conn, error) {
	// Try to dial the conn
	conn, err := net.DialTimeout("tcp", addr.String(), 10*time.Second)
	if err != nil {
		return nil, err
	}

	// Cast to TCPConn
	if tcp, ok := conn.(*net.TCPConn); ok {
		tcp.SetKeepAlive(true)
		tcp.SetNoDelay(true)
	}

	// Check if TLS is enabled
	if p.tlsWrap != nil {
		// Switch the connection into TLS mode
		if _, err := conn.Write([]byte{byte(RpcTLS)}); err != nil {
			conn.Close()
			return nil, err
		}

		// Wrap the connection in a TLS client
		tlsConn, err := p.tlsWrap(region, conn)
		if err != nil {
			conn.Close()
			return nil, err
		}
		conn = tlsConn
	}

	// Write the multiplex byte to set the mode
	if _, err := conn.Write([]byte{byte(RpcMultiplex)}); err != nil {
		conn.Close()
		return nil, err
	}

	// Setup the logger
	conf := yamux.DefaultConfig()
	conf.LogOutput = nil
	conf.Logger = p.logger

	// Create a multiplexed session
	session, err := yamux.Client(conn, conf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Wrap the connection
	c := &Conn{
		refCount: 1,
		addr:     addr,
		session:  session,
		clients:  list.New(),
		lastUsed: time.Now(),
		version:  version,
		pool:     p,
	}
	return c, nil
}

// clearConn is used to clear any cached connection, potentially in response to
// an error
func (p *ConnPool) clearConn(conn *Conn) {
	// Ensure returned streams are closed
	atomic.StoreInt32(&conn.shouldClose, 1)

	// Clear from the cache
	p.Lock()
	if c, ok := p.pool[conn.addr.String()]; ok && c == conn {
		delete(p.pool, conn.addr.String())
	}
	p.Unlock()

	// Close down immediately if idle
	if refCount := atomic.LoadInt32(&conn.refCount); refCount == 0 {
		conn.Close()
	}
}

// releaseConn is invoked when we are done with a conn to reduce the ref count
func (p *ConnPool) releaseConn(conn *Conn) {
	refCount := atomic.AddInt32(&conn.refCount, -1)
	if refCount == 0 && atomic.LoadInt32(&conn.shouldClose) == 1 {
		conn.Close()
	}
}

// getClient is used to get a usable client for an address and protocol version
func (p *ConnPool) getClient(region string, addr net.Addr, version int) (*Conn, *StreamClient, error) {
	retries := 0
START:
	// Try to get a conn first
	conn, err := p.acquire(region, addr, version)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get conn: %v", err)
	}

	// Get a client
	client, err := conn.getClient()
	if err != nil {
		p.clearConn(conn)
		p.releaseConn(conn)

		// Try to redial, possible that the TCP session closed due to timeout
		if retries == 0 {
			retries++
			goto START
		}
		return nil, nil, fmt.Errorf("failed to start stream: %v", err)
	}
	return conn, client, nil
}

// RPC is used to make an RPC call to a remote host
func (p *ConnPool) RPC(region string, addr net.Addr, version int, method string, args interface{}, reply interface{}) error {
	// Get a usable client
	conn, sc, err := p.getClient(region, addr, version)
	if err != nil {
		return fmt.Errorf("rpc error: %v", err)
	}

	// Make the RPC call
	err = msgpackrpc.CallWithCodec(sc.codec, method, args, reply)
	if err != nil {
		sc.Close()

		// If we read EOF, the session is toast. Clear it and open a
		// new session next time
		// See https://github.com/hashicorp/consul/blob/v1.6.3/agent/pool/pool.go#L471-L477
		if lib.IsErrEOF(err) {
			p.clearConn(conn)
		}

		p.releaseConn(conn)

		// If the error is an RPC Coded error
		// return the coded error without wrapping
		if structs.IsErrRPCCoded(err) {
			return err
		}

		// TODO wrap with RPCCoded error instead
		return fmt.Errorf("rpc error: %v", err)
	}

	// Done with the connection
	conn.returnClient(sc)
	p.releaseConn(conn)
	return nil
}

// Reap is used to close conns open over maxTime
func (p *ConnPool) reap() {
	for {
		// Sleep for a while
		select {
		case <-p.shutdownCh:
			return
		case <-time.After(time.Second):
		}

		// Reap all old conns
		p.Lock()
		var removed []string
		now := time.Now()
		for host, conn := range p.pool {
			// Skip recently used connections
			if now.Sub(conn.lastUsed) < p.maxTime {
				continue
			}

			// Skip connections with active streams
			if atomic.LoadInt32(&conn.refCount) > 0 {
				continue
			}

			// Close the conn
			conn.Close()

			// Remove from pool
			removed = append(removed, host)
		}
		for _, host := range removed {
			delete(p.pool, host)
		}
		p.Unlock()
	}
}
<|endoftext|>