<commit_before>package koding

import (
	"fmt"
	"koding/db/mongodb"
	"strconv"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	aws "koding/kites/kloud/api/amazon"
	"koding/kites/kloud/klient"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/protocol"
	"koding/kites/kloud/provider/amazon"

	"github.com/koding/kite"
	"github.com/koding/logging"
	"github.com/mitchellh/goamz/ec2"
)

// Checker checks various aspects of a machine. It is used for limiting certain
// aspects of a machine, such as the total allowed machine count, storage size
// and etc.
type Checker interface {
	// Total checks whether the user has reached the current plan's limit of
	// having a total number numbers of machines. It returns an error if the
	// limit is reached or an unexplained error happaned.
	Total() error

	// AlwaysOn checks whether the given machine has reached the current plans
	// always on limit
	AlwaysOn() error

	// Timeout checks whether the user has reached the current plan's
	// inactivity timeout.
	Timeout() error

	// Storage checks whether the user has reached the current plan's limit
	// total storage with the supplied wantStorage information. It returns an
	// error if the limit is reached or an unexplained error happaned.
	Storage(wantStorage int) error

	// AllowedInstances checks whether the given machine has the permisison to
	// create the given instance type
	AllowedInstances(wantInstance InstanceType) error
}

type PlanChecker struct {
	Api      *amazon.AmazonClient
	DB       *mongodb.MongoDB
	Machine  *protocol.Machine
	Provider *Provider
	Kite     *kite.Kite
	Username string
	Log      logging.Logger
	Plan     Plan
}

func (p *PlanChecker) AllowedInstances(wantInstance InstanceType) error {
	allowedInstances := p.Plan.Limits().AllowedInstances

	if _, ok := allowedInstances[wantInstance]; ok {
		return nil
	}

	return fmt.Errorf("not allowed to create instance type: %s", wantInstance)
}

func (p *PlanChecker) AlwaysOn() error {
	alwaysOnLimit := p.Plan.Limits().AlwaysOn

	// get all alwaysOn machines that belongs to this user
	alwaysOnMachines := 0
	if err := p.DB.Run("jMachines", func(c *mgo.Collection) error {
		var err error
		alwaysOnMachines, err = c.Find(bson.M{
			"credential":    p.Machine.Username,
			"meta.alwaysOn": true,
		}).Count()

		return err
	}); err != nil && err != mgo.ErrNotFound {
		// if it's something else just return an error, needs to be fixed
		return err
	}

	p.Log.Debug("[%s] checking alwaysOn limit. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.Machine.Id, alwaysOnMachines, alwaysOnLimit, p.Plan)

	// the user has still not reached the limit
	if alwaysOnMachines <= alwaysOnLimit {
		p.Log.Debug("[%s] allowing user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
			p.Machine.Id, p.Username, alwaysOnMachines, alwaysOnLimit, p.Plan)
		return nil // allow user, it didn't reach the limit
	}

	p.Log.Debug("[%s] denying user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.Machine.Id, p.Username, alwaysOnMachines, alwaysOnLimit, p.Plan)
	return fmt.Errorf("total alwaysOn limit has been reached. Current count: %d Plan limit: %d",
		alwaysOnMachines, alwaysOnLimit)
}

func (p *PlanChecker) Timeout() error {
	// Check klient state before rushing to AWS.
	klientRef, err := klient.Connect(p.Kite, p.Machine.QueryString)
	if err == kite.ErrNoKitesAvailable {
		p.Provider.startTimer(p.Machine)
		return err
	}

	// return if it's something else
	if err != nil {
		return err
	}

	// now the klient is connected and we can ping it, stop the timer and
	// remove it from the list of inactive machines if it's still there.
	p.Provider.stopTimer(p.Machine)

	// get the usage directly from the klient, which is the most predictable source
	usg, err := klientRef.Usage()
	if err != nil {
		klientRef.Close()
		return err
	}
	klientRef.Close()

	// replace with the real and authenticated username
	p.Machine.Builder["username"] = klientRef.Username
	p.Username = klientRef.Username

	// get the timeout from the plan in which the user belongs to
	planTimeout := p.Plan.Limits().Timeout

	p.Log.Debug("[%s] machine [%s] is inactive for %s (plan limit: %s, plan: %s).",
		p.Machine.Id, p.Machine.IpAddress, usg.InactiveDuration, planTimeout, p.Plan)

	// It still have plenty of time to work, do not stop it
	if usg.InactiveDuration <= planTimeout {
		return nil
	}

	p.Log.Info("[%s] machine [%s] has reached current plan limit of %s (plan: %s). Shutting down...",<commit_msg>checker: only call it once and also nil it for GC<commit_after>package koding

import (
	"fmt"
	"koding/db/mongodb"
	"strconv"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	aws "koding/kites/kloud/api/amazon"
	"koding/kites/kloud/klient"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/protocol"
	"koding/kites/kloud/provider/amazon"

	"github.com/koding/kite"
	"github.com/koding/logging"
	"github.com/mitchellh/goamz/ec2"
)

// Checker checks various aspects of a machine. It is used for limiting certain
// aspects of a machine, such as the total allowed machine count, storage size
// and etc.
type Checker interface {
	// Total checks whether the user has reached the current plan's limit of
	// having a total number numbers of machines. It returns an error if the
	// limit is reached or an unexplained error happaned.
	Total() error

	// AlwaysOn checks whether the given machine has reached the current plans
	// always on limit
	AlwaysOn() error

	// Timeout checks whether the user has reached the current plan's
	// inactivity timeout.
	Timeout() error

	// Storage checks whether the user has reached the current plan's limit
	// total storage with the supplied wantStorage information. It returns an
	// error if the limit is reached or an unexplained error happaned.
	Storage(wantStorage int) error

	// AllowedInstances checks whether the given machine has the permisison to
	// create the given instance type
	AllowedInstances(wantInstance InstanceType) error
}

type PlanChecker struct {
	Api      *amazon.AmazonClient
	DB       *mongodb.MongoDB
	Machine  *protocol.Machine
	Provider *Provider
	Kite     *kite.Kite
	Username string
	Log      logging.Logger
	Plan     Plan
}

func (p *PlanChecker) AllowedInstances(wantInstance InstanceType) error {
	allowedInstances := p.Plan.Limits().AllowedInstances

	if _, ok := allowedInstances[wantInstance]; ok {
		return nil
	}

	return fmt.Errorf("not allowed to create instance type: %s", wantInstance)
}

func (p *PlanChecker) AlwaysOn() error {
	alwaysOnLimit := p.Plan.Limits().AlwaysOn

	// get all alwaysOn machines that belongs to this user
	alwaysOnMachines := 0
	if err := p.DB.Run("jMachines", func(c *mgo.Collection) error {
		var err error
		alwaysOnMachines, err = c.Find(bson.M{
			"credential":    p.Machine.Username,
			"meta.alwaysOn": true,
		}).Count()

		return err
	}); err != nil && err != mgo.ErrNotFound {
		// if it's something else just return an error, needs to be fixed
		return err
	}

	p.Log.Debug("[%s] checking alwaysOn limit. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.Machine.Id, alwaysOnMachines, alwaysOnLimit, p.Plan)

	// the user has still not reached the limit
	if alwaysOnMachines <= alwaysOnLimit {
		p.Log.Debug("[%s] allowing user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
			p.Machine.Id, p.Username, alwaysOnMachines, alwaysOnLimit, p.Plan)
		return nil // allow user, it didn't reach the limit
	}

	p.Log.Debug("[%s] denying user '%s'. current alwaysOn count: %d (plan limit: %d, plan: %s)",
		p.Machine.Id, p.Username, alwaysOnMachines, alwaysOnLimit, p.Plan)
	return fmt.Errorf("total alwaysOn limit has been reached. Current count: %d Plan limit: %d",
		alwaysOnMachines, alwaysOnLimit)
}

func (p *PlanChecker) Timeout() error {
	// Check klient state before rushing to AWS.
	klientRef, err := klient.Connect(p.Kite, p.Machine.QueryString)
	if err == kite.ErrNoKitesAvailable {
		p.Provider.startTimer(p.Machine)
		return err
	}

	// return if it's something else
	if err != nil {
		return err
	}

	// now the klient is connected and we can ping it, stop the timer and
	// remove it from the list of inactive machines if it's still there.
	p.Provider.stopTimer(p.Machine)

	// replace with the real and authenticated username
	p.Machine.Builder["username"] = klientRef.Username
	p.Username = klientRef.Username

	// get the usage directly from the klient, which is the most predictable source
	usg, err := klientRef.Usage()

	klientRef.Close()
	klientRef = nil

	if err != nil {
		return err
	}

	// get the timeout from the plan in which the user belongs to
	planTimeout := p.Plan.Limits().Timeout

	p.Log.Debug("[%s] machine [%s] is inactive for %s (plan limit: %s, plan: %s).",
		p.Machine.Id, p.Machine.IpAddress, usg.InactiveDuration, planTimeout, p.Plan)

	// It still have plenty of time to work, do not stop it
	if usg.InactiveDuration <= planTimeout {
		return nil
	}

	p.Log.Info("[%s] machine [%s] has reached current plan limit of %s (plan: %s). Shutting down...",
{'fraction_non_alphanumeric': 0.10879911455451023, 'fraction_numerical': 0.0008854454897620365, 'mean_word_length': 3.7633104902477594, 'pattern_counts': {'":': 4, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3737673', 'n_tokens_mistral': 3039, 'n_tokens_neox': 2901, 'n_words': 1022}