<commit_before>package main

import (
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"koding/databases/redis"
	"koding/db/mongodb"
	"koding/db/mongodb/modelhelper"
	"koding/kites/os/ldapserver"
	"koding/kodingkite"
	"koding/tools/config"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/utils"
	"koding/virt"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	redigo "github.com/garyburd/redigo/redis"
	kitelib "github.com/koding/kite"

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

const OSKITE_NAME = "oskite"

type VMInfo struct {
	vm              *virt.VM
	useCounter      int
	timeout         *time.Timer
	mutex           sync.Mutex
	totalCpuUsage   int
	currentCpus     []string
	currentHostname string

	State               string `json:"state"`
	CpuUsage            int    `json:"cpuUsage"`
	CpuShares           int    `json:"cpuShares"`
	MemoryUsage         int    `json:"memoryUsage"`
	PhysicalMemoryLimit int    `json:"physicalMemoryLimit"`
	TotalMemoryLimit    int    `json:"totalMemoryLimit"`
}

var (
	log         = logger.New(OSKITE_NAME)
	logLevel    logger.Level
	mongodbConn *mongodb.MongoDB
	conf        *config.Config

	flagProfile      = flag.String("c", "", "Configuration profile from file")
	flagRegion       = flag.String("r", "", "Configuration region from file")
	flagDebug        = flag.Bool("d", false, "Debug mode")
	flagTemplates    = flag.String("t", "", "Change template directory")
	flagTimeout      = flag.String("s", "50m", "Shut down timeout for a single VM")
	flagDisableGuest = flag.Bool("noguest", false, "Disable Guest VM creation")
	flagLimit        = flag.Int("limit", 100, "Disable Guest VM creation")

	infos            = make(map[bson.ObjectId]*VMInfo)
	infosMutex       sync.Mutex
	templateDir      = "files/templates" // should be in the same dir as the binary
	firstContainerIP net.IP
	containerSubnet  *net.IPNet
	shuttingDown     = false
	requestWaitGroup sync.WaitGroup

	vmTimeout         = time.Minute * 50
	prepareQueueLimit = 8 + 1 // number of concurrent VM preparations, should be CPU + 1
	prepareQueue      = make(chan func(chan struct{}))
)

func main() {
	flag.Parse()
	if *flagProfile == "" || *flagRegion == "" {
		log.Fatal("Please specify profile via -c and region via -r. Aborting.")
	}

	conf = config.MustConfig(*flagProfile)
	mongodbConn = mongodb.NewMongoDB(conf.Mongo)
	modelhelper.Initialize(conf.Mongo)

	if *flagDebug {
		logLevel = logger.DEBUG
	} else {
		logLevel = logger.GetLoggingLevelFromConfig(OSKITE_NAME, *flagProfile)
	}
	log.SetLevel(logLevel)

	if *flagTemplates != "" {
		templateDir = *flagTemplates
	}

	var newTimeout time.Duration
	var err error
	newTimeout, err = time.ParseDuration(*flagTimeout)
	if err != nil {
		log.Warning("Timeout parameter is not correct: %v", err.Error())
		log.Notice("Using default VM timeout: %s", vmTimeout)
	} else {
		// use our new timeout
		log.Info("Using default VM timeout: %s", vmTimeout)
		vmTimeout = newTimeout
	}

	// set seed for even randomness, needed for randomMinutes() function.
	rand.Seed(time.Now().UnixNano())

	initializeSettings()

	k := prepareOsKite()

	runNewKite(k.ServiceUniqueName)
	handleCurrentVMS(k)   // handle leftover VMs
	startPinnedVMS(k)     // start pinned always-on VMs
	setupSignalHandler(k) // handle SIGUSR1 and other signals.

	// startPrepareWorkers starts multiple workers (based on prepareQueueLimit)
	// that accepts vmPrepare/vmStart functions.
	for i := 0; i < prepareQueueLimit; i++ {
		go prepareWorker()
	}

	// register current client-side methods
	registerVmMethod(k, "vm.start", false, vmStart)
	registerVmMethod(k, "vm.shutdown", false, vmShutdown)
	registerVmMethod(k, "vm.unprepare", false, vmUnprepare)
	registerVmMethod(k, "vm.stop", false, vmStop)
	registerVmMethod(k, "vm.reinitialize", false, vmReinitialize)
	registerVmMethod(k, "vm.info", false, vmInfo)
	registerVmMethod(k, "vm.resizeDisk", false, vmResizeDisk)
	registerVmMethod(k, "vm.createSnapshot", false, vmCreateSnaphost)
	registerVmMethod(k, "spawn", true, spawnFunc)
	registerVmMethod(k, "exec", true, execFunc)

	registerVmMethod(k, "oskite.Info", true, oskiteInfo)
	registerVmMethod(k, "oskite.All", true, oskiteAll)
	registerVmMethod(k, "oskite.All", true, oskiteAll)

	syscall.Umask(0) // don't know why richard calls this
	registerVmMethod(k, "fs.readDirectory", false, fsReadDirectory)
	registerVmMethod(k, "fs.glob", false, fsGlob)
	registerVmMethod(k, "fs.readFile", false, fsReadFile)
	registerVmMethod(k, "fs.writeFile", false, fsWriteFile)
	registerVmMethod(k, "fs.ensureNonexistentPath", false, fsEnsureNonexistentPath)
	registerVmMethod(k, "fs.getInfo", false, fsGetInfo)
	registerVmMethod(k, "fs.setPermissions", false, fsSetPermissions)
	registerVmMethod(k, "fs.remove", false, fsRemove)
	registerVmMethod(k, "fs.rename", false, fsRename)
	registerVmMethod(k, "fs.createDirectory", false, fsCreateDirectory)

	registerVmMethod(k, "app.install", false, appInstall)
	registerVmMethod(k, "app.download", false, appDownload)
	registerVmMethod(k, "app.publish", false, appPublish)
	registerVmMethod(k, "app.skeleton", false, appSkeleton)

	// this method is special cased in oskite.go to allow foreign access
	registerVmMethod(k, "webterm.connect", false, webtermConnect)
	registerVmMethod(k, "webterm.getSessions", false, webtermGetSessions)

	registerVmMethod(k, "s3.store", true, s3Store)
	registerVmMethod(k, "s3.delete", true, s3Delete)

	go oskiteRedis(k.ServiceUniqueName)

	k.Run()<commit_msg>oskite: do not register a method twice<commit_after>package main

import (
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"koding/databases/redis"
	"koding/db/mongodb"
	"koding/db/mongodb/modelhelper"
	"koding/kites/os/ldapserver"
	"koding/kodingkite"
	"koding/tools/config"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/utils"
	"koding/virt"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	redigo "github.com/garyburd/redigo/redis"
	kitelib "github.com/koding/kite"

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

const OSKITE_NAME = "oskite"

type VMInfo struct {
	vm              *virt.VM
	useCounter      int
	timeout         *time.Timer
	mutex           sync.Mutex
	totalCpuUsage   int
	currentCpus     []string
	currentHostname string

	State               string `json:"state"`
	CpuUsage            int    `json:"cpuUsage"`
	CpuShares           int    `json:"cpuShares"`
	MemoryUsage         int    `json:"memoryUsage"`
	PhysicalMemoryLimit int    `json:"physicalMemoryLimit"`
	TotalMemoryLimit    int    `json:"totalMemoryLimit"`
}

var (
	log         = logger.New(OSKITE_NAME)
	logLevel    logger.Level
	mongodbConn *mongodb.MongoDB
	conf        *config.Config

	flagProfile      = flag.String("c", "", "Configuration profile from file")
	flagRegion       = flag.String("r", "", "Configuration region from file")
	flagDebug        = flag.Bool("d", false, "Debug mode")
	flagTemplates    = flag.String("t", "", "Change template directory")
	flagTimeout      = flag.String("s", "50m", "Shut down timeout for a single VM")
	flagDisableGuest = flag.Bool("noguest", false, "Disable Guest VM creation")
	flagLimit        = flag.Int("limit", 100, "Disable Guest VM creation")

	infos            = make(map[bson.ObjectId]*VMInfo)
	infosMutex       sync.Mutex
	templateDir      = "files/templates" // should be in the same dir as the binary
	firstContainerIP net.IP
	containerSubnet  *net.IPNet
	shuttingDown     = false
	requestWaitGroup sync.WaitGroup

	vmTimeout         = time.Minute * 50
	prepareQueueLimit = 8 + 1 // number of concurrent VM preparations, should be CPU + 1
	prepareQueue      = make(chan func(chan struct{}))
)

func main() {
	flag.Parse()
	if *flagProfile == "" || *flagRegion == "" {
		log.Fatal("Please specify profile via -c and region via -r. Aborting.")
	}

	conf = config.MustConfig(*flagProfile)
	mongodbConn = mongodb.NewMongoDB(conf.Mongo)
	modelhelper.Initialize(conf.Mongo)

	if *flagDebug {
		logLevel = logger.DEBUG
	} else {
		logLevel = logger.GetLoggingLevelFromConfig(OSKITE_NAME, *flagProfile)
	}
	log.SetLevel(logLevel)

	if *flagTemplates != "" {
		templateDir = *flagTemplates
	}

	var newTimeout time.Duration
	var err error
	newTimeout, err = time.ParseDuration(*flagTimeout)
	if err != nil {
		log.Warning("Timeout parameter is not correct: %v", err.Error())
		log.Notice("Using default VM timeout: %s", vmTimeout)
	} else {
		// use our new timeout
		log.Info("Using default VM timeout: %s", vmTimeout)
		vmTimeout = newTimeout
	}

	// set seed for even randomness, needed for randomMinutes() function.
	rand.Seed(time.Now().UnixNano())

	initializeSettings()

	k := prepareOsKite()

	runNewKite(k.ServiceUniqueName)
	handleCurrentVMS(k)   // handle leftover VMs
	startPinnedVMS(k)     // start pinned always-on VMs
	setupSignalHandler(k) // handle SIGUSR1 and other signals.

	// startPrepareWorkers starts multiple workers (based on prepareQueueLimit)
	// that accepts vmPrepare/vmStart functions.
	for i := 0; i < prepareQueueLimit; i++ {
		go prepareWorker()
	}

	// register current client-side methods
	registerVmMethod(k, "vm.start", false, vmStart)
	registerVmMethod(k, "vm.shutdown", false, vmShutdown)
	registerVmMethod(k, "vm.unprepare", false, vmUnprepare)
	registerVmMethod(k, "vm.stop", false, vmStop)
	registerVmMethod(k, "vm.reinitialize", false, vmReinitialize)
	registerVmMethod(k, "vm.info", false, vmInfo)
	registerVmMethod(k, "vm.resizeDisk", false, vmResizeDisk)
	registerVmMethod(k, "vm.createSnapshot", false, vmCreateSnaphost)
	registerVmMethod(k, "spawn", true, spawnFunc)
	registerVmMethod(k, "exec", true, execFunc)

	registerVmMethod(k, "oskite.Info", true, oskiteInfo)
	registerVmMethod(k, "oskite.All", true, oskiteAll)

	syscall.Umask(0) // don't know why richard calls this
	registerVmMethod(k, "fs.readDirectory", false, fsReadDirectory)
	registerVmMethod(k, "fs.glob", false, fsGlob)
	registerVmMethod(k, "fs.readFile", false, fsReadFile)
	registerVmMethod(k, "fs.writeFile", false, fsWriteFile)
	registerVmMethod(k, "fs.ensureNonexistentPath", false, fsEnsureNonexistentPath)
	registerVmMethod(k, "fs.getInfo", false, fsGetInfo)
	registerVmMethod(k, "fs.setPermissions", false, fsSetPermissions)
	registerVmMethod(k, "fs.remove", false, fsRemove)
	registerVmMethod(k, "fs.rename", false, fsRename)
	registerVmMethod(k, "fs.createDirectory", false, fsCreateDirectory)

	registerVmMethod(k, "app.install", false, appInstall)
	registerVmMethod(k, "app.download", false, appDownload)
	registerVmMethod(k, "app.publish", false, appPublish)
	registerVmMethod(k, "app.skeleton", false, appSkeleton)

	// this method is special cased in oskite.go to allow foreign access
	registerVmMethod(k, "webterm.connect", false, webtermConnect)
	registerVmMethod(k, "webterm.getSessions", false, webtermGetSessions)

	registerVmMethod(k, "s3.store", true, s3Store)
	registerVmMethod(k, "s3.delete", true, s3Delete)

	go oskiteRedis(k.ServiceUniqueName)

	k.Run()
{'fraction_non_alphanumeric': 0.12649465072372562, 'fraction_numerical': 0.003416344511372831, 'mean_word_length': 4.34550696780394, 'pattern_counts': {'":': 0, '<': 5, '<?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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2123508', 'n_tokens_mistral': 4027, 'n_tokens_neox': 3906, 'n_words': 881}