<commit_before>package editor

import (
	"sort"
	"strconv"

	"github.com/animenotifier/notify.moe/utils/animediff"

	"github.com/aerogo/aero"
	"github.com/animenotifier/arn"
	"github.com/animenotifier/mal"
	"github.com/animenotifier/notify.moe/components"
	"github.com/animenotifier/notify.moe/utils"
)

const maxCompareMALEntries = 15

// CompareMAL ...
func CompareMAL(ctx *aero.Context) string {
	user := utils.GetUser(ctx)
	year, _ := ctx.GetInt("year")
	status := ctx.Get("status")
	animeType := ctx.Get("type")

	animes := arn.FilterAnime(func(anime *arn.Anime) bool {
		if status != "" && anime.Status != status {
			return false
		}

		if year != 0 && year != anime.StartDateTime().Year() {
			return false
		}

		if animeType != "" && anime.Type != animeType {
			return false
		}

		return anime.GetMapping("myanimelist/anime") != ""
	})

	sort.Slice(animes, func(i, j int) bool {
		a := animes[i]
		b := animes[j]

		aPop := a.Popularity.Total()
		bPop := b.Popularity.Total()

		if aPop == bPop {
			return a.Title.Canonical < b.Title.Canonical
		}

		return aPop > bPop
	})

	comparisons := []*utils.MALComparison{}
	malAnimeCollection := arn.MAL.Collection("Anime")

	for _, anime := range animes {
		malID := anime.GetMapping("myanimelist/anime")
		obj, err := malAnimeCollection.Get(malID)

		if err != nil {
			continue
		}

		malAnime := obj.(*mal.Anime)
		var differences []animediff.Difference

		// Canonical title
		if anime.Title.Canonical != malAnime.Title {
			hash := utils.HashString(malAnime.Title)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "CanonicalTitle", hash) {
				differences = append(differences, &animediff.CanonicalTitle{
					TitleA:      anime.Title.Canonical,
					TitleB:      malAnime.Title,
					NumericHash: hash,
				})
			}
		}

		// Japanese title
		if anime.Title.Japanese != malAnime.JapaneseTitle {
			hash := utils.HashString(malAnime.JapaneseTitle)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "JapaneseTitle", hash) {
				differences = append(differences, &animediff.JapaneseTitle{
					TitleA:      anime.Title.Japanese,
					TitleB:      malAnime.JapaneseTitle,
					NumericHash: hash,
				})
			}
		}

		// Romaji title
		if anime.Title.Romaji != malAnime.Title {
			hash := utils.HashString(malAnime.Title)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "RomajiTitle", hash) {
				differences = append(differences, &animediff.RomajiTitle{
					TitleA:      anime.Title.Romaji,
					TitleB:      malAnime.Title,
					NumericHash: hash,
				})
			}
		}

		// Airing start date
		if anime.StartDate != malAnime.StartDate && malAnime.StartDate != "" {
			hash := utils.HashString(malAnime.StartDate)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "StartDate", hash) {
				differences = append(differences, &animediff.StartDate{
					DateA:       anime.StartDate,
					DateB:       malAnime.StartDate,
					NumericHash: hash,
				})
			}
		}

		// Airing end date
		if anime.EndDate != malAnime.EndDate && malAnime.EndDate != "" {
			hash := utils.HashString(malAnime.EndDate)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "EndDate", hash) {
				differences = append(differences, &animediff.EndDate{
					DateA:       anime.EndDate,
					DateB:       malAnime.EndDate,
					NumericHash: hash,
				})
			}
		}

		// Status
		if anime.Status != malAnime.Status {
			hash := utils.HashString(malAnime.Status)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "Status", hash) {
				differences = append(differences, &animediff.Status{
					StatusA:     anime.Status,
					StatusB:     malAnime.Status,
					NumericHash: hash,
				})
			}
		}

		// EpisodeCount
		if anime.EpisodeCount != malAnime.EpisodeCount {
			hash := uint64(malAnime.EpisodeCount)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "EpisodeCount", hash) {
				differences = append(differences, &animediff.EpisodeCount{
					EpisodesA:   anime.EpisodeCount,
					EpisodesB:   malAnime.EpisodeCount,
					NumericHash: hash,
				})
			}
		}

		// Synopsis
		if len(anime.Summary) < len(malAnime.Synopsis) {
			hash := utils.HashString(malAnime.Synopsis)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "Synopsis", hash) {
				differences = append(differences, &animediff.Synopsis{
					SynopsisA:   anime.Summary,
					SynopsisB:   malAnime.Synopsis,
					NumericHash: hash,
				})
			}
		}

		// Compare genres
		hashA := utils.HashStringsNoOrder(anime.Genres)
		hashB := utils.HashStringsNoOrder(malAnime.Genres)

		if hashA != hashB {
			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "Genres", hashB) {
				differences = append(differences, &animediff.Genres{
					GenresA:     anime.Genres,
					GenresB:     malAnime.Genres,
					NumericHash: hashB,
				})
			}
		}

		// Add if there were any differences
		if len(differences) > 0 {
			comparisons = append(comparisons, &utils.MALComparison{
				Anime:       anime,
				MALAnime:    malAnime,
				Differences: differences,
			})

			if len(comparisons) >= maxCompareMALEntries {
				break
			}
		}
	}

	return ctx.HTML(components.CompareMAL(comparisons, strconv.Itoa(year), status, animeType, ctx.URI(), user))
}
<commit_msg>Improved synopsis filter<commit_after>package editor

import (
	"sort"
	"strconv"

	"github.com/animenotifier/notify.moe/utils/animediff"

	"github.com/aerogo/aero"
	"github.com/animenotifier/arn"
	"github.com/animenotifier/mal"
	"github.com/animenotifier/notify.moe/components"
	"github.com/animenotifier/notify.moe/utils"
)

const maxCompareMALEntries = 15

// CompareMAL ...
func CompareMAL(ctx *aero.Context) string {
	user := utils.GetUser(ctx)
	year, _ := ctx.GetInt("year")
	status := ctx.Get("status")
	animeType := ctx.Get("type")

	animes := arn.FilterAnime(func(anime *arn.Anime) bool {
		if status != "" && anime.Status != status {
			return false
		}

		if year != 0 && year != anime.StartDateTime().Year() {
			return false
		}

		if animeType != "" && anime.Type != animeType {
			return false
		}

		return anime.GetMapping("myanimelist/anime") != ""
	})

	sort.Slice(animes, func(i, j int) bool {
		a := animes[i]
		b := animes[j]

		aPop := a.Popularity.Total()
		bPop := b.Popularity.Total()

		if aPop == bPop {
			return a.Title.Canonical < b.Title.Canonical
		}

		return aPop > bPop
	})

	comparisons := []*utils.MALComparison{}
	malAnimeCollection := arn.MAL.Collection("Anime")

	for _, anime := range animes {
		malID := anime.GetMapping("myanimelist/anime")
		obj, err := malAnimeCollection.Get(malID)

		if err != nil {
			continue
		}

		malAnime := obj.(*mal.Anime)
		var differences []animediff.Difference

		// Canonical title
		if anime.Title.Canonical != malAnime.Title {
			hash := utils.HashString(malAnime.Title)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "CanonicalTitle", hash) {
				differences = append(differences, &animediff.CanonicalTitle{
					TitleA:      anime.Title.Canonical,
					TitleB:      malAnime.Title,
					NumericHash: hash,
				})
			}
		}

		// Japanese title
		if anime.Title.Japanese != malAnime.JapaneseTitle {
			hash := utils.HashString(malAnime.JapaneseTitle)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "JapaneseTitle", hash) {
				differences = append(differences, &animediff.JapaneseTitle{
					TitleA:      anime.Title.Japanese,
					TitleB:      malAnime.JapaneseTitle,
					NumericHash: hash,
				})
			}
		}

		// Romaji title
		if anime.Title.Romaji != malAnime.Title {
			hash := utils.HashString(malAnime.Title)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "RomajiTitle", hash) {
				differences = append(differences, &animediff.RomajiTitle{
					TitleA:      anime.Title.Romaji,
					TitleB:      malAnime.Title,
					NumericHash: hash,
				})
			}
		}

		// Airing start date
		if anime.StartDate != malAnime.StartDate && malAnime.StartDate != "" {
			hash := utils.HashString(malAnime.StartDate)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "StartDate", hash) {
				differences = append(differences, &animediff.StartDate{
					DateA:       anime.StartDate,
					DateB:       malAnime.StartDate,
					NumericHash: hash,
				})
			}
		}

		// Airing end date
		if anime.EndDate != malAnime.EndDate && malAnime.EndDate != "" {
			hash := utils.HashString(malAnime.EndDate)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "EndDate", hash) {
				differences = append(differences, &animediff.EndDate{
					DateA:       anime.EndDate,
					DateB:       malAnime.EndDate,
					NumericHash: hash,
				})
			}
		}

		// Status
		if anime.Status != malAnime.Status {
			hash := utils.HashString(malAnime.Status)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "Status", hash) {
				differences = append(differences, &animediff.Status{
					StatusA:     anime.Status,
					StatusB:     malAnime.Status,
					NumericHash: hash,
				})
			}
		}

		// EpisodeCount
		if anime.EpisodeCount != malAnime.EpisodeCount {
			hash := uint64(malAnime.EpisodeCount)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "EpisodeCount", hash) {
				differences = append(differences, &animediff.EpisodeCount{
					EpisodesA:   anime.EpisodeCount,
					EpisodesB:   malAnime.EpisodeCount,
					NumericHash: hash,
				})
			}
		}

		// Synopsis
		if len(anime.Summary) < 300 && len(anime.Summary)+50 < len(malAnime.Synopsis) {
			hash := utils.HashString(malAnime.Synopsis)

			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "Synopsis", hash) {
				differences = append(differences, &animediff.Synopsis{
					SynopsisA:   anime.Summary,
					SynopsisB:   malAnime.Synopsis,
					NumericHash: hash,
				})
			}
		}

		// Compare genres
		hashA := utils.HashStringsNoOrder(anime.Genres)
		hashB := utils.HashStringsNoOrder(malAnime.Genres)

		if hashA != hashB {
			if !arn.IsAnimeDifferenceIgnored(anime.ID, "mal", malAnime.ID, "Genres", hashB) {
				differences = append(differences, &animediff.Genres{
					GenresA:     anime.Genres,
					GenresB:     malAnime.Genres,
					NumericHash: hashB,
				})
			}
		}

		// Add if there were any differences
		if len(differences) > 0 {
			comparisons = append(comparisons, &utils.MALComparison{
				Anime:       anime,
				MALAnime:    malAnime,
				Differences: differences,
			})

			if len(comparisons) >= maxCompareMALEntries {
				break
			}
		}
	}

	return ctx.HTML(components.CompareMAL(comparisons, strconv.Itoa(year), status, animeType, ctx.URI(), user))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/Stantheman/youtube-gif-go/api"
	"github.com/Stantheman/youtube-gif-go/config"
	"github.com/Stantheman/youtube-gif-go/logger"
	"github.com/Stantheman/youtube-gif-go/redisPool"
	"github.com/garyburd/redigo/redis"
	"os"
	"os/exec"
)

// workerfunc takes a decoded message and does work on it
type workerfunc func(payload api.PubSubMessage, workdir string) error

type worker struct {
	proc workerfunc
	tell string
}

var jobs = map[string]worker{
	"download": {
		proc: download,
		tell: "chop",
	},
	"chop": {
		proc: chop,
		tell: "stitch",
	},
	"stitch": {
		proc: stitch,
		tell: "finalize",
	},
	"finalize": {
		proc: finalize,
		tell: "",
	},
}

var conf = config.Get()
var l = logger.Get()

func main() {

	jobname := get_worker_name()

	// each worker listens to its -queue for work
	queue := jobname + "-queue"

	// get a nicely wrapped pubsub connection
	psc := redis.PubSubConn{redisPool.Pool.Get()}
	defer psc.Close()

	if err := psc.Subscribe(queue); err != nil {
		l.Err("Can't subscribe to the " + queue + ": " + err.Error())
		return
	}

	for {
		switch v := psc.Receive().(type) {
		case redis.Message:
			go process_message(v.Data, jobname)
		case redis.Subscription:
			fmt.Printf("%s: %s %d\n", v.Channel, v.Kind, v.Count)
		case error:
			fmt.Println(v.Error())
			return
		}
	}
}

func get_worker_name() string {
	if len(os.Args) != 2 {
		fmt.Println("Usage: " + os.Args[0] + " worker-name")
		os.Exit(1)
	}
	if _, ok := jobs[os.Args[1]]; ok != true {
		fmt.Println("Usage: " + os.Args[0] + " worker-name")
		os.Exit(1)
	}

	return os.Args[1]
}

func process_message(data []byte, jobname string) {
	// decode json
	payload := api.PubSubMessage{}
	if err := json.Unmarshal(data, &payload); err != nil {
		l.Err(err.Error())
		return
	}
	// do work
	if err := work(payload, jobname); err != nil {
		l.Err(err.Error())
		markFailed(payload.ID, jobname, err.Error())
	}
}

/* work is a wrapper method that does the setup and teardown around jobs */
func work(payload api.PubSubMessage, job string) (err error) {

	// mark active
	if err := markInProgress(payload.ID, job); err != nil {
		markFailed(payload.ID, job, err.Error())
		return err
	}

	workspace := conf.Worker.Dir + "/" + payload.ID + "/" + job + "/"

	// make workspace
	if err := os.MkdirAll(workspace, 0777); err != nil {
		return err
	}

	// do the job
	if err := jobs[job].proc(payload, workspace); err != nil {
		return err
	}

	// update the status
	if err := update(payload, job, workspace); err != nil {
		return err
	}

	return nil
}

func update(payload api.PubSubMessage, job, workspace string) (err error) {

	// get ready to update the image status and continue pubsubbin
	c := redisPool.Pool.Get()
	keyname := "gif:" + payload.ID

	// if there's no one left to tell, we're done processing
	if jobs[job].tell == "" {
		// Add the image to active_image list, set image to avail
		c.Send("MULTI")
		c.Send("RPUSH", "active_images", payload.ID)
		c.Send("HSET", keyname, "status", "available")
		c.Send("PERSIST", keyname)
		if _, err := c.Do("EXEC"); err != nil {
			return err
		}
		return nil
	}
	// otherwise we need to update status and tell people

	// push the new workspace onto the stack
	payload.PrevDir = workspace

	// serialize payload for sending back out
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	c.Send("MULTI")
	c.Send("HSET", keyname, "status", "pending "+jobs[job].tell)
	c.Send("EXPIRE", keyname, 3600)
	c.Send("PUBLISH", jobs[job].tell+"-queue", data)
	if _, err := c.Do("EXEC"); err != nil {
		return err
	}
	return nil
}

func markInProgress(id, job string) error {
	c := redisPool.Pool.Get()
	keyname := "gif:" + id

	c.Send("MULTI")
	c.Send("HSET", keyname, "status", job+"-ifying")
	c.Send("EXPIRE", keyname, 3600)
	if _, err := c.Do("EXEC"); err != nil {
		return err
	}
	return nil
}

func markFailed(id, job, errormsg string) {
	c := redisPool.Pool.Get()
	keyname := "gif:" + id

	c.Send("MULTI")
	c.Send("HSET", keyname, "status", "failed")
	c.Send("HSET", keyname, "description", errormsg)
	c.Send("PERSIST", keyname)
	if _, err := c.Do("EXEC"); err != nil {
		// since this is the failure marker, log it and move on
		l.Err("Failed on trying to update failure for " + id + " (" + job + ". Moving on")
	}
	l.Err("Making " + id + " as failed in job " + job + ".")
}

/* download uses youtube-dl to take a youtube URL and place it on disk */
func download(payload api.PubSubMessage, workspace string) error {

	// 1.youtube
	outFile := workspace + payload.ID + ".youtube"

	// needs config for youtube-dl path
	cmd := exec.Command("/opt/youtube-dl",
		"--no-progress",
		"--max-filesize", "100M",
		"--output", outFile,
		"--write-info-json",
		"--format", "mp4",
		payload.URL,
	)

	if output, err := cmd.CombinedOutput(); err != nil {
		return errors.New(err.Error() + ". " + string(output))
	}
	// verify that the file is actually there
	if _, err := os.Stat(outFile); os.IsNotExist(err) {
		return err
	}

	return nil
}

/* chop takes a workspace directory with an <id>.youtube file
and uses mplayer to create a PNG file for every frame in the video*/
func chop(payload api.PubSubMessage, workspace string) error {

	// https://www.youtube.com/watch?v=dgKGixi8bp8 short video
	var flags = []string{}

	var wscale, hscale, xcoord, ycoord, swidth, sheight string

	// avconv allows for math based on frame size
	// scale the crop size/coords against the real one
	if payload.Cx != "" {
		wscale = "(in_w/" + payload.Vw + ")*" + payload.Cw
		hscale = "(in_h/" + payload.Vh + ")*" + payload.Ch
		xcoord = "(in_w/" + payload.Vw + ")*" + payload.Cx
		ycoord = "(in_h/" + payload.Vh + ")*" + payload.Cy
		swidth = payload.Cw
		sheight = payload.Ch
	} else {
		swidth = payload.Vw
		sheight = payload.Vh
	}

	flags = append(flags,
		"-i", payload.PrevDir+payload.ID+".youtube",
		"-vf", "crop="+wscale+":"+hscale+":"+xcoord+":"+ycoord+",scale="+swidth+":"+sheight,
	)
	if payload.Dur != "" {
		flags = append(flags,
			"-ss", payload.Start,
			"-t", payload.Dur,
		)
	}
	flags = append(flags,
		workspace+"%08d.png",
	)

	cmd := exec.Command("/usr/local/bin/avconv", flags...)

	if output, err := cmd.CombinedOutput(); err != nil {
		return errors.New(err.Error() + ". " + string(output))
	}

	// bash script would rm [0,2,4,6,8].png to make gifmaking quicker, gif smaller
	// think about doing that at some boundary
	if _, err := os.Stat(workspace + "00000001.png"); os.IsNotExist(err) {
		return err
	}

	return nil
}

/* stitch takes a workspace directory of PNG files and transforms them
into a gif on STDOUT using imagemagick, which is piped into
gifsicle for optimization */
func stitch(payload api.PubSubMessage, workspace string) error {
	gifname := workspace + payload.ID + ".gif"

	// -monitor
	cmd := exec.Command("/bin/bash", "-c",
		"/usr/bin/gm convert +repage -fuzz 1.6% -delay 4 -loop 0 "+payload.PrevDir+"*.png "+
			"gif:-|/usr/bin/gifsicle -O3 --colors 256 > "+gifname,
	)

	if output, err := cmd.CombinedOutput(); err != nil {
		return errors.New(err.Error() + ". " + string(output))
	}

	// verify that the file is actually there
	if _, err := os.Stat(gifname); os.IsNotExist(err) {
		return err
	}

	return nil
}

func finalize(payload api.PubSubMessage, workspace string) error {
	gifname := conf.Site.Gif_Dir + "/" + payload.ID + ".gif"

	if err := os.Rename(payload.PrevDir+payload.ID+".gif", gifname); err != nil {
		return err
	}

	if err := os.RemoveAll(conf.Worker.Dir + "/" + payload.ID); err != nil {
		l.Err("Couldn't remove old directory: " + err.Error())
	}

	return nil
}
<commit_msg>have youtube-dl fall back to flv too. this stuff should be in a config file :/<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/Stantheman/youtube-gif-go/api"
	"github.com/Stantheman/youtube-gif-go/config"
	"github.com/Stantheman/youtube-gif-go/logger"
	"github.com/Stantheman/youtube-gif-go/redisPool"
	"github.com/garyburd/redigo/redis"
	"os"
	"os/exec"
)

// workerfunc takes a decoded message and does work on it
type workerfunc func(payload api.PubSubMessage, workdir string) error

type worker struct {
	proc workerfunc
	tell string
}

var jobs = map[string]worker{
	"download": {
		proc: download,
		tell: "chop",
	},
	"chop": {
		proc: chop,
		tell: "stitch",
	},
	"stitch": {
		proc: stitch,
		tell: "finalize",
	},
	"finalize": {
		proc: finalize,
		tell: "",
	},
}

var conf = config.Get()
var l = logger.Get()

func main() {

	jobname := get_worker_name()

	// each worker listens to its -queue for work
	queue := jobname + "-queue"

	// get a nicely wrapped pubsub connection
	psc := redis.PubSubConn{redisPool.Pool.Get()}
	defer psc.Close()

	if err := psc.Subscribe(queue); err != nil {
		l.Err("Can't subscribe to the " + queue + ": " + err.Error())
		return
	}

	for {
		switch v := psc.Receive().(type) {
		case redis.Message:
			go process_message(v.Data, jobname)
		case redis.Subscription:
			fmt.Printf("%s: %s %d\n", v.Channel, v.Kind, v.Count)
		case error:
			fmt.Println(v.Error())
			return
		}
	}
}

func get_worker_name() string {
	if len(os.Args) != 2 {
		fmt.Println("Usage: " + os.Args[0] + " worker-name")
		os.Exit(1)
	}
	if _, ok := jobs[os.Args[1]]; ok != true {
		fmt.Println("Usage: " + os.Args[0] + " worker-name")
		os.Exit(1)
	}

	return os.Args[1]
}

func process_message(data []byte, jobname string) {
	// decode json
	payload := api.PubSubMessage{}
	if err := json.Unmarshal(data, &payload); err != nil {
		l.Err(err.Error())
		return
	}
	// do work
	if err := work(payload, jobname); err != nil {
		l.Err(err.Error())
		markFailed(payload.ID, jobname, err.Error())
	}
}

/* work is a wrapper method that does the setup and teardown around jobs */
func work(payload api.PubSubMessage, job string) (err error) {

	// mark active
	if err := markInProgress(payload.ID, job); err != nil {
		markFailed(payload.ID, job, err.Error())
		return err
	}

	workspace := conf.Worker.Dir + "/" + payload.ID + "/" + job + "/"

	// make workspace
	if err := os.MkdirAll(workspace, 0777); err != nil {
		return err
	}

	// do the job
	if err := jobs[job].proc(payload, workspace); err != nil {
		return err
	}

	// update the status
	if err := update(payload, job, workspace); err != nil {
		return err
	}

	return nil
}

func update(payload api.PubSubMessage, job, workspace string) (err error) {

	// get ready to update the image status and continue pubsubbin
	c := redisPool.Pool.Get()
	keyname := "gif:" + payload.ID

	// if there's no one left to tell, we're done processing
	if jobs[job].tell == "" {
		// Add the image to active_image list, set image to avail
		c.Send("MULTI")
		c.Send("RPUSH", "active_images", payload.ID)
		c.Send("HSET", keyname, "status", "available")
		c.Send("PERSIST", keyname)
		if _, err := c.Do("EXEC"); err != nil {
			return err
		}
		return nil
	}
	// otherwise we need to update status and tell people

	// push the new workspace onto the stack
	payload.PrevDir = workspace

	// serialize payload for sending back out
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	c.Send("MULTI")
	c.Send("HSET", keyname, "status", "pending "+jobs[job].tell)
	c.Send("EXPIRE", keyname, 3600)
	c.Send("PUBLISH", jobs[job].tell+"-queue", data)
	if _, err := c.Do("EXEC"); err != nil {
		return err
	}
	return nil
}

func markInProgress(id, job string) error {
	c := redisPool.Pool.Get()
	keyname := "gif:" + id

	c.Send("MULTI")
	c.Send("HSET", keyname, "status", job+"-ifying")
	c.Send("EXPIRE", keyname, 3600)
	if _, err := c.Do("EXEC"); err != nil {
		return err
	}
	return nil
}

func markFailed(id, job, errormsg string) {
	c := redisPool.Pool.Get()
	keyname := "gif:" + id

	c.Send("MULTI")
	c.Send("HSET", keyname, "status", "failed")
	c.Send("HSET", keyname, "description", errormsg)
	c.Send("PERSIST", keyname)
	if _, err := c.Do("EXEC"); err != nil {
		// since this is the failure marker, log it and move on
		l.Err("Failed on trying to update failure for " + id + " (" + job + ". Moving on")
	}
	l.Err("Making " + id + " as failed in job " + job + ".")
}

/* download uses youtube-dl to take a youtube URL and place it on disk */
func download(payload api.PubSubMessage, workspace string) error {

	// 1.youtube
	outFile := workspace + payload.ID + ".youtube"

	// needs config for youtube-dl path
	cmd := exec.Command("/opt/youtube-dl",
		"--no-progress",
		"--max-filesize", "100M",
		"--output", outFile,
		"--write-info-json",
		"--format", "mp4,flv",
		payload.URL,
	)

	if output, err := cmd.CombinedOutput(); err != nil {
		return errors.New(err.Error() + ". " + string(output))
	}
	// verify that the file is actually there
	if _, err := os.Stat(outFile); os.IsNotExist(err) {
		return err
	}

	return nil
}

/* chop takes a workspace directory with an <id>.youtube file
and uses mplayer to create a PNG file for every frame in the video*/
func chop(payload api.PubSubMessage, workspace string) error {

	// https://www.youtube.com/watch?v=dgKGixi8bp8 short video
	var flags = []string{}

	var wscale, hscale, xcoord, ycoord, swidth, sheight string

	// avconv allows for math based on frame size
	// scale the crop size/coords against the real one
	if payload.Cx != "" {
		wscale = "(in_w/" + payload.Vw + ")*" + payload.Cw
		hscale = "(in_h/" + payload.Vh + ")*" + payload.Ch
		xcoord = "(in_w/" + payload.Vw + ")*" + payload.Cx
		ycoord = "(in_h/" + payload.Vh + ")*" + payload.Cy
		swidth = payload.Cw
		sheight = payload.Ch
	} else {
		swidth = payload.Vw
		sheight = payload.Vh
	}

	flags = append(flags,
		"-i", payload.PrevDir+payload.ID+".youtube",
		"-vf", "crop="+wscale+":"+hscale+":"+xcoord+":"+ycoord+",scale="+swidth+":"+sheight,
	)
	if payload.Dur != "" {
		flags = append(flags,
			"-ss", payload.Start,
			"-t", payload.Dur,
		)
	}
	flags = append(flags,
		workspace+"%08d.png",
	)

	cmd := exec.Command("/usr/local/bin/avconv", flags...)

	if output, err := cmd.CombinedOutput(); err != nil {
		return errors.New(err.Error() + ". " + string(output))
	}

	// bash script would rm [0,2,4,6,8].png to make gifmaking quicker, gif smaller
	// think about doing that at some boundary
	if _, err := os.Stat(workspace + "00000001.png"); os.IsNotExist(err) {
		return err
	}

	return nil
}

/* stitch takes a workspace directory of PNG files and transforms them
into a gif on STDOUT using imagemagick, which is piped into
gifsicle for optimization */
func stitch(payload api.PubSubMessage, workspace string) error {
	gifname := workspace + payload.ID + ".gif"

	// -monitor
	cmd := exec.Command("/bin/bash", "-c",
		"/usr/bin/gm convert +repage -fuzz 1.6% -delay 4 -loop 0 "+payload.PrevDir+"*.png "+
			"gif:-|/usr/bin/gifsicle -O3 --colors 256 > "+gifname,
	)

	if output, err := cmd.CombinedOutput(); err != nil {
		return errors.New(err.Error() + ". " + string(output))
	}

	// verify that the file is actually there
	if _, err := os.Stat(gifname); os.IsNotExist(err) {
		return err
	}

	return nil
}

func finalize(payload api.PubSubMessage, workspace string) error {
	gifname := conf.Site.Gif_Dir + "/" + payload.ID + ".gif"

	if err := os.Rename(payload.PrevDir+payload.ID+".gif", gifname); err != nil {
		return err
	}

	if err := os.RemoveAll(conf.Worker.Dir + "/" + payload.ID); err != nil {
		l.Err("Couldn't remove old directory: " + err.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package worker

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor"
	"github.com/qor/qor/admin"
)

var (
	jobDB         *gorm.DB
	jobId         uint64
	workers       = map[string]*Worker{}
	queuers       = map[string]Queuer{}
	DefaultJobCli = strings.Join(os.Args, " ")
)

func init() {
	flag.Uint64Var(&jobId, "job-id", 0, "qor job id")
}

// SetJobDB will run a auto migration for creating table jobs
func SetJobDB(db *gorm.DB) (err error) {
	err = db.AutoMigrate(&Job{}).Error
	if err != nil {
		return
	}

	jobDB = db

	return
}

// TODO: UNDONE
func SetAdmin(a *admin.Admin) {
	// ws := defaultWorkerSet.Workers
	job := a.NewResource(&Job{})
	job.IndexAttrs("Id", "QueueJobId", "Interval", "StartAt", "Cli", "WokerSetName", "WorkerName", "Status", "PID")

	// defaultWorkerSet = NewWorkerSet(defaultWorkerSet.Name, "/workers", "", admin)
	// admin.RegisterViewPath(os.Getenv("GOPATH") + "/src/github.com/qor/qor/worker/templates")
	// a.GetRouter().Get("/workers", func(c *admin.Context) {
	// 	content := admin.Content{Context: c, Admin: a}
	// 	a.Render("workers", content)
	// })

	// defaultWorkerSet.Workers = ws
}

// Listen will parse an flag named as "job-id". If the job-id is zero, it
// will run as queue listen server. Otherwise, it will run a specific job
// and terminate the process after the job is run.
//
// It must be executed before http.ListenAndServer
func Listen() {
	flag.Parse()

	if jobId > 0 {
		var job Job
		if err := jobDB.Where("id = ?", jobId).Find(&job).Error; err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		var logger io.Writer
		var err error
		if logger, err = job.GetLogger(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			fmt.Fprintln(logger, err)
			os.Exit(1)
		}

		var w *Worker
		if w, err = job.GetWorker(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			fmt.Fprintln(logger, err)
			os.Exit(1)
		}

		if err = w.Run(&job); err != nil {
			fmt.Fprintln(os.Stderr, err)
			fmt.Fprintln(logger, err)
			os.Exit(1)
		}

		os.Exit(0)
	} else {
		for _, queuer := range queuers {
			go func() {
				for {
					jobId, err := queuer.Dequeue()
					fmt.Println("dequeue job", jobId)
					if err != nil {
						fmt.Println("qor.worker.dequeue.error:", err)
					} else {
						go RunJob(jobId)
					}
				}
			}()
		}
	}
}

func New(queuer Queuer, name string, handle func(job *Job) error) (w *Worker) {
	w = &Worker{
		Name:   name,
		Handle: handle,
		Queuer: queuer,
	}

	workers[w.Name] = w
	queuers[queuer.Name()] = queuer

	return
}

type Worker struct {
	Name   string
	Queuer Queuer
	Config *qor.Config

	Handle    func(job *Job) error
	OnKill    func(job *Job) error
	OnStart   func(job *Job) error
	OnSuccess func(job *Job)
	OnFailed  func(job *Job)
}

func (w *Worker) Run(job *Job) (err error) {
	if err = job.SavePID(); err != nil {
		return
	}
	logger, err := job.GetLogger()
	if err != nil {
		return
	}

	fmt.Fprintf(logger, "run job (%d) with pid (%d)\n", job.Id, job.PID)

	if w.OnStart != nil {
		if err = w.OnStart(job); err != nil {
			logger.Write([]byte("worker.onstart: " + err.Error() + "\n"))

			if err = job.UpdateStatus(JobFailed); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}

			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}

	if err = job.UpdateStatus(JobRunning); err != nil {
		fmt.Fprintf(logger, "error: %s\n", err)
		return
	}

	if err = w.Handle(job); err != nil {
		if err = job.UpdateStatus(JobFailed); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		if w.OnFailed != nil {
			w.OnFailed(job)
		}

		logger.Write([]byte("worker.hanlde: " + err.Error() + "\n"))
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	} else if w.OnSuccess != nil {
		if err = job.UpdateStatus(JobRun); err != nil {
			fmt.Fprintln(os.Stderr, err)
		}

		w.OnSuccess(job)
	}

	if err = job.UpdateStatus(JobRun); err != nil {
		fmt.Fprintf(logger, "error: %s\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(logger, "finish job (%d) with pid (%d)\n", job.Id, job.PID)

	return
}

var ErrJobRun = errors.New("job is already run")

func (w *Worker) Kill(job *Job) (err error) {
	if w.OnKill != nil {
		if err = w.OnKill(job); err != nil {
			return
		}
	}

	switch job.Status {
	case JobToRun:
		err = w.Queuer.Purge(job)
	case JobRunning:
		if job.PID == 0 {
			return errors.New("pid is zero")
		}

		var process *os.Process
		process, err = os.FindProcess(job.PID)
		if err != nil {
			return
		}

		err = process.Kill()
	case JobRun:
		return ErrJobRun
	}

	if err == nil {
		err = job.UpdateStatus(JobKilled)
	}

	return
}

func (w *Worker) NewJob(interval uint64, startAt time.Time) (job *Job, err error) {
	return w.NewJobWithCli(interval, startAt, DefaultJobCli)
}

func (w *Worker) NewJobWithCli(interval uint64, startAt time.Time, cli string) (job *Job, err error) {
	job = &Job{
		Interval:   interval,
		StartAt:    startAt,
		WorkerName: w.Name,
		Cli:        cli,
	}
	if err = jobDB.Save(job).Error; err != nil {
		return
	}

	err = w.Queuer.Enqueue(job)

	if job.QueueJobId != "" {
		if err = jobDB.Save(job).Error; err != nil {
			return
		}
	}

	return
}
<commit_msg>worker: remove WorkerSetName from job IndexAttrs<commit_after>package worker

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor"
	"github.com/qor/qor/admin"
)

var (
	jobDB         *gorm.DB
	jobId         uint64
	workers       = map[string]*Worker{}
	queuers       = map[string]Queuer{}
	DefaultJobCli = strings.Join(os.Args, " ")
)

func init() {
	flag.Uint64Var(&jobId, "job-id", 0, "qor job id")
}

// SetJobDB will run a auto migration for creating table jobs
func SetJobDB(db *gorm.DB) (err error) {
	err = db.AutoMigrate(&Job{}).Error
	if err != nil {
		return
	}

	jobDB = db

	return
}

// TODO: UNDONE
func SetAdmin(a *admin.Admin) {
	// ws := defaultWorkerSet.Workers
	job := a.NewResource(&Job{})
	job.IndexAttrs("Id", "QueueJobId", "Interval", "StartAt", "Cli", "WorkerName", "Status", "PID", "RunCounter", "FailCounter", "SuccessCounter", "KillCounter")

	// admin.RegisterViewPath(os.Getenv("GOPATH") + "/src/github.com/qor/qor/worker/templates")
	// a.GetRouter().Get("/workers", func(c *admin.Context) {
	// 	content := admin.Content{Context: c, Admin: a}
	// 	a.Render("workers", content)
	// })
}

// Listen will parse an flag named as "job-id". If the job-id is zero, it
// will run as queue listen server. Otherwise, it will run a specific job
// and terminate the process after the job is run.
//
// It must be executed before http.ListenAndServer
func Listen() {
	flag.Parse()

	if jobId > 0 {
		var job Job
		if err := jobDB.Where("id = ?", jobId).Find(&job).Error; err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		var logger io.Writer
		var err error
		if logger, err = job.GetLogger(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			fmt.Fprintln(logger, err)
			os.Exit(1)
		}

		var w *Worker
		if w, err = job.GetWorker(); err != nil {
			fmt.Fprintln(os.Stderr, err)
			fmt.Fprintln(logger, err)
			os.Exit(1)
		}

		if err = w.Run(&job); err != nil {
			fmt.Fprintln(os.Stderr, err)
			fmt.Fprintln(logger, err)
			os.Exit(1)
		}

		os.Exit(0)
	} else {
		for _, queuer := range queuers {
			go func() {
				for {
					jobId, err := queuer.Dequeue()
					fmt.Println("dequeue job", jobId)
					if err != nil {
						fmt.Println("qor.worker.dequeue.error:", err)
					} else {
						go RunJob(jobId)
					}
				}
			}()
		}
	}
}

func New(queuer Queuer, name string, handle func(job *Job) error) (w *Worker) {
	w = &Worker{
		Name:   name,
		Handle: handle,
		Queuer: queuer,
	}

	workers[w.Name] = w
	queuers[queuer.Name()] = queuer

	return
}

type Worker struct {
	Name   string
	Queuer Queuer
	Config *qor.Config

	Handle    func(job *Job) error
	OnKill    func(job *Job) error
	OnStart   func(job *Job) error
	OnSuccess func(job *Job)
	OnFailed  func(job *Job)
}

func (w *Worker) Run(job *Job) (err error) {
	if err = job.SavePID(); err != nil {
		return
	}
	logger, err := job.GetLogger()
	if err != nil {
		return
	}

	fmt.Fprintf(logger, "run job (%d) with pid (%d)\n", job.Id, job.PID)

	if w.OnStart != nil {
		if err = w.OnStart(job); err != nil {
			logger.Write([]byte("worker.onstart: " + err.Error() + "\n"))

			if err = job.UpdateStatus(JobFailed); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}

			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}

	if err = job.UpdateStatus(JobRunning); err != nil {
		fmt.Fprintf(logger, "error: %s\n", err)
		return
	}

	if err = w.Handle(job); err != nil {
		if err = job.UpdateStatus(JobFailed); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		if w.OnFailed != nil {
			w.OnFailed(job)
		}

		logger.Write([]byte("worker.hanlde: " + err.Error() + "\n"))
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	} else if w.OnSuccess != nil {
		if err = job.UpdateStatus(JobRun); err != nil {
			fmt.Fprintln(os.Stderr, err)
		}

		w.OnSuccess(job)
	}

	if err = job.UpdateStatus(JobRun); err != nil {
		fmt.Fprintf(logger, "error: %s\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(logger, "finish job (%d) with pid (%d)\n", job.Id, job.PID)

	return
}

var ErrJobRun = errors.New("job is already run")

func (w *Worker) Kill(job *Job) (err error) {
	if w.OnKill != nil {
		if err = w.OnKill(job); err != nil {
			return
		}
	}

	switch job.Status {
	case JobToRun:
		err = w.Queuer.Purge(job)
	case JobRunning:
		if job.PID == 0 {
			return errors.New("pid is zero")
		}

		var process *os.Process
		process, err = os.FindProcess(job.PID)
		if err != nil {
			return
		}

		err = process.Kill()
	case JobRun:
		return ErrJobRun
	}

	if err == nil {
		err = job.UpdateStatus(JobKilled)
	}

	return
}

func (w *Worker) NewJob(interval uint64, startAt time.Time) (job *Job, err error) {
	return w.NewJobWithCli(interval, startAt, DefaultJobCli)
}

func (w *Worker) NewJobWithCli(interval uint64, startAt time.Time, cli string) (job *Job, err error) {
	job = &Job{
		Interval:   interval,
		StartAt:    startAt,
		WorkerName: w.Name,
		Cli:        cli,
	}
	if err = jobDB.Save(job).Error; err != nil {
		return
	}

	err = w.Queuer.Enqueue(job)

	if job.QueueJobId != "" {
		if err = jobDB.Save(job).Error; err != nil {
			return
		}
	}

	return
}
<|endoftext|>


<commit_before>package store

import (
	"path"
	"sort"
	"time"

	etcdErr "github.com/coreos/etcd/error"
)

const (
	normal = iota
	removed
)

var Permanent time.Time

// Node is the basic element in the store system.
// A key-value pair will have a string value
// A directory will have a children map
type Node struct {
	Path string

	CreateIndex   uint64
	CreateTerm    uint64
	ModifiedIndex uint64
	ModifiedTerm  uint64

	Parent *Node `json:"-"` // should not encode this field! avoid cyclical dependency.

	ExpireTime time.Time
	ACL        string
	Value      string           // for key-value pair
	Children   map[string]*Node // for directory

	// A reference to the store this node is attached to.
	store *store
}

// newKV creates a Key-Value pair
func newKV(store *store, nodePath string, value string, createIndex uint64,
	createTerm uint64, parent *Node, ACL string, expireTime time.Time) *Node {

	return &Node{
		Path:          nodePath,
		CreateIndex:   createIndex,
		CreateTerm:    createTerm,
		ModifiedIndex: createIndex,
		ModifiedTerm:  createTerm,
		Parent:        parent,
		ACL:           ACL,
		store:         store,
		ExpireTime:    expireTime,
		Value:         value,
	}
}

// newDir creates a directory
func newDir(store *store, nodePath string, createIndex uint64, createTerm uint64,
	parent *Node, ACL string, expireTime time.Time) *Node {

	return &Node{
		Path:        nodePath,
		CreateIndex: createIndex,
		CreateTerm:  createTerm,
		Parent:      parent,
		ACL:         ACL,
		ExpireTime:  expireTime,
		Children:    make(map[string]*Node),
		store:       store,
	}
}

// IsHidden function checks if the node is a hidden node. A hidden node
// will begin with '_'
// A hidden node will not be shown via get command under a directory
// For example if we have /foo/_hidden and /foo/notHidden, get "/foo"
// will only return /foo/notHidden
func (n *Node) IsHidden() bool {
	_, name := path.Split(n.Path)

	return name[0] == '_'
}

// IsPermanent function checks if the node is a permanent one.
func (n *Node) IsPermanent() bool {
	return n.ExpireTime.IsZero()
}

// IsDir function checks whether the node is a directory.
// If the node is a directory, the function will return true.
// Otherwise the function will return false.
func (n *Node) IsDir() bool {
	return !(n.Children == nil)
}

// Read function gets the value of the node.
// If the receiver node is not a key-value pair, a "Not A File" error will be returned.
func (n *Node) Read() (string, *etcdErr.Error) {
	if n.IsDir() {
		return "", etcdErr.NewError(etcdErr.EcodeNotFile, "", UndefIndex, UndefTerm)
	}

	return n.Value, nil
}

// Write function set the value of the node to the given value.
// If the receiver node is a directory, a "Not A File" error will be returned.
func (n *Node) Write(value string, index uint64, term uint64) *etcdErr.Error {
	if n.IsDir() {
		return etcdErr.NewError(etcdErr.EcodeNotFile, "", UndefIndex, UndefTerm)
	}

	n.Value = value
	n.ModifiedIndex = index
	n.ModifiedTerm = term

	return nil
}

func (n *Node) ExpirationAndTTL() (*time.Time, int64) {
	if !n.IsPermanent() {
		return &n.ExpireTime, int64(n.ExpireTime.Sub(time.Now())/time.Second) + 1
	}
	return nil, 0
}

// List function return a slice of nodes under the receiver node.
// If the receiver node is not a directory, a "Not A Directory" error will be returned.
func (n *Node) List() ([]*Node, *etcdErr.Error) {
	if !n.IsDir() {
		return nil, etcdErr.NewError(etcdErr.EcodeNotDir, "", UndefIndex, UndefTerm)
	}

	nodes := make([]*Node, len(n.Children))

	i := 0
	for _, node := range n.Children {
		nodes[i] = node
		i++
	}

	return nodes, nil
}

// GetChild function returns the child node under the directory node.
// On success, it returns the file node
func (n *Node) GetChild(name string) (*Node, *etcdErr.Error) {
	if !n.IsDir() {
		return nil, etcdErr.NewError(etcdErr.EcodeNotDir, n.Path, UndefIndex, UndefTerm)
	}

	child, ok := n.Children[name]

	if ok {
		return child, nil
	}

	return nil, nil
}

// Add function adds a node to the receiver node.
// If the receiver is not a directory, a "Not A Directory" error will be returned.
// If there is a existing node with the same name under the directory, a "Already Exist"
// error will be returned
func (n *Node) Add(child *Node) *etcdErr.Error {
	if !n.IsDir() {
		return etcdErr.NewError(etcdErr.EcodeNotDir, "", UndefIndex, UndefTerm)
	}

	_, name := path.Split(child.Path)

	_, ok := n.Children[name]

	if ok {
		return etcdErr.NewError(etcdErr.EcodeNodeExist, "", UndefIndex, UndefTerm)
	}

	n.Children[name] = child

	return nil
}

// Remove function remove the node.
func (n *Node) Remove(recursive bool, callback func(path string)) *etcdErr.Error {

	if n.IsDir() && !recursive {
		// cannot delete a directory without set recursive to true
		return etcdErr.NewError(etcdErr.EcodeNotFile, "", UndefIndex, UndefTerm)
	}

	if !n.IsDir() { // key-value pair
		_, name := path.Split(n.Path)

		// find its parent and remove the node from the map
		if n.Parent != nil && n.Parent.Children[name] == n {
			delete(n.Parent.Children, name)
		}

		if callback != nil {
			callback(n.Path)
		}

		if !n.IsPermanent() {
			n.store.ttlKeyHeap.remove(n)
		}

		return nil
	}

	for _, child := range n.Children { // delete all children
		child.Remove(true, callback)
	}

	// delete self
	_, name := path.Split(n.Path)
	if n.Parent != nil && n.Parent.Children[name] == n {
		delete(n.Parent.Children, name)

		if callback != nil {
			callback(n.Path)
		}

		if !n.IsPermanent() {
			n.store.ttlKeyHeap.remove(n)
		}

	}

	return nil
}

func (n *Node) Pair(recurisive, sorted bool) KeyValuePair {
	if n.IsDir() {
		pair := KeyValuePair{
			Key: n.Path,
			Dir: true,
		}
		pair.Expiration, pair.TTL = n.ExpirationAndTTL()

		if !recurisive {
			return pair
		}

		children, _ := n.List()
		pair.KVPairs = make([]KeyValuePair, len(children))

		// we do not use the index in the children slice directly
		// we need to skip the hidden one
		i := 0

		for _, child := range children {

			if child.IsHidden() { // get will not list hidden node
				continue
			}

			pair.KVPairs[i] = child.Pair(recurisive, sorted)

			i++
		}

		// eliminate hidden nodes
		pair.KVPairs = pair.KVPairs[:i]
		if sorted {
			sort.Sort(pair.KVPairs)
		}

		return pair
	}

	pair := KeyValuePair{
		Key:   n.Path,
		Value: n.Value,
	}
	pair.Expiration, pair.TTL = n.ExpirationAndTTL()
	return pair
}

func (n *Node) UpdateTTL(expireTime time.Time) {

	if !n.IsPermanent() {
		if expireTime.IsZero() {
			// from ttl to permanent
			// remove from ttl heap
			n.store.ttlKeyHeap.remove(n)
		} else {
			// update ttl
			n.ExpireTime = expireTime
			// update ttl heap
			n.store.ttlKeyHeap.update(n)
		}

	} else {
		if !expireTime.IsZero() {
			// from permanent to ttl
			n.ExpireTime = expireTime
			// push into ttl heap
			n.store.ttlKeyHeap.push(n)
		}
	}
}

// Clone function clone the node recursively and return the new node.
// If the node is a directory, it will clone all the content under this directory.
// If the node is a key-value pair, it will clone the pair.
func (n *Node) Clone() *Node {
	if !n.IsDir() {
		return newKV(n.store, n.Path, n.Value, n.CreateIndex, n.CreateTerm, n.Parent, n.ACL, n.ExpireTime)
	}

	clone := newDir(n.store, n.Path, n.CreateIndex, n.CreateTerm, n.Parent, n.ACL, n.ExpireTime)

	for key, child := range n.Children {
		clone.Children[key] = child.Clone()
	}

	return clone
}

// recoverAndclean function help to do recovery.
// Two things need to be done: 1. recovery structure; 2. delete expired nodes

// If the node is a directory, it will help recover children's parent pointer and recursively
// call this function on its children.
// We check the expire last since we need to recover the whole structure first and add all the
// notifications into the event history.
func (n *Node) recoverAndclean() {
	if n.IsDir() {
		for _, child := range n.Children {
			child.Parent = n
			child.store = n.store
			child.recoverAndclean()
		}
	}

	if !n.ExpireTime.IsZero() {
		n.store.ttlKeyHeap.push(n)
	}

}
<commit_msg>refactor comments on IsPermanent<commit_after>package store

import (
	"path"
	"sort"
	"time"

	etcdErr "github.com/coreos/etcd/error"
)

const (
	normal = iota
	removed
)

var Permanent time.Time

// Node is the basic element in the store system.
// A key-value pair will have a string value
// A directory will have a children map
type Node struct {
	Path string

	CreateIndex   uint64
	CreateTerm    uint64
	ModifiedIndex uint64
	ModifiedTerm  uint64

	Parent *Node `json:"-"` // should not encode this field! avoid cyclical dependency.

	ExpireTime time.Time
	ACL        string
	Value      string           // for key-value pair
	Children   map[string]*Node // for directory

	// A reference to the store this node is attached to.
	store *store
}

// newKV creates a Key-Value pair
func newKV(store *store, nodePath string, value string, createIndex uint64,
	createTerm uint64, parent *Node, ACL string, expireTime time.Time) *Node {

	return &Node{
		Path:          nodePath,
		CreateIndex:   createIndex,
		CreateTerm:    createTerm,
		ModifiedIndex: createIndex,
		ModifiedTerm:  createTerm,
		Parent:        parent,
		ACL:           ACL,
		store:         store,
		ExpireTime:    expireTime,
		Value:         value,
	}
}

// newDir creates a directory
func newDir(store *store, nodePath string, createIndex uint64, createTerm uint64,
	parent *Node, ACL string, expireTime time.Time) *Node {

	return &Node{
		Path:        nodePath,
		CreateIndex: createIndex,
		CreateTerm:  createTerm,
		Parent:      parent,
		ACL:         ACL,
		ExpireTime:  expireTime,
		Children:    make(map[string]*Node),
		store:       store,
	}
}

// IsHidden function checks if the node is a hidden node. A hidden node
// will begin with '_'
// A hidden node will not be shown via get command under a directory
// For example if we have /foo/_hidden and /foo/notHidden, get "/foo"
// will only return /foo/notHidden
func (n *Node) IsHidden() bool {
	_, name := path.Split(n.Path)

	return name[0] == '_'
}

// IsPermanent function checks if the node is a permanent one.
func (n *Node) IsPermanent() bool {
	// we use a uninitialized time.Time to indicate the node is a
	// permanent one.
	// the uninitialized time.Time should equal zero.
	return n.ExpireTime.IsZero()
}

// IsDir function checks whether the node is a directory.
// If the node is a directory, the function will return true.
// Otherwise the function will return false.
func (n *Node) IsDir() bool {
	return !(n.Children == nil)
}

// Read function gets the value of the node.
// If the receiver node is not a key-value pair, a "Not A File" error will be returned.
func (n *Node) Read() (string, *etcdErr.Error) {
	if n.IsDir() {
		return "", etcdErr.NewError(etcdErr.EcodeNotFile, "", UndefIndex, UndefTerm)
	}

	return n.Value, nil
}

// Write function set the value of the node to the given value.
// If the receiver node is a directory, a "Not A File" error will be returned.
func (n *Node) Write(value string, index uint64, term uint64) *etcdErr.Error {
	if n.IsDir() {
		return etcdErr.NewError(etcdErr.EcodeNotFile, "", UndefIndex, UndefTerm)
	}

	n.Value = value
	n.ModifiedIndex = index
	n.ModifiedTerm = term

	return nil
}

func (n *Node) ExpirationAndTTL() (*time.Time, int64) {
	if !n.IsPermanent() {
		return &n.ExpireTime, int64(n.ExpireTime.Sub(time.Now())/time.Second) + 1
	}
	return nil, 0
}

// List function return a slice of nodes under the receiver node.
// If the receiver node is not a directory, a "Not A Directory" error will be returned.
func (n *Node) List() ([]*Node, *etcdErr.Error) {
	if !n.IsDir() {
		return nil, etcdErr.NewError(etcdErr.EcodeNotDir, "", UndefIndex, UndefTerm)
	}

	nodes := make([]*Node, len(n.Children))

	i := 0
	for _, node := range n.Children {
		nodes[i] = node
		i++
	}

	return nodes, nil
}

// GetChild function returns the child node under the directory node.
// On success, it returns the file node
func (n *Node) GetChild(name string) (*Node, *etcdErr.Error) {
	if !n.IsDir() {
		return nil, etcdErr.NewError(etcdErr.EcodeNotDir, n.Path, UndefIndex, UndefTerm)
	}

	child, ok := n.Children[name]

	if ok {
		return child, nil
	}

	return nil, nil
}

// Add function adds a node to the receiver node.
// If the receiver is not a directory, a "Not A Directory" error will be returned.
// If there is a existing node with the same name under the directory, a "Already Exist"
// error will be returned
func (n *Node) Add(child *Node) *etcdErr.Error {
	if !n.IsDir() {
		return etcdErr.NewError(etcdErr.EcodeNotDir, "", UndefIndex, UndefTerm)
	}

	_, name := path.Split(child.Path)

	_, ok := n.Children[name]

	if ok {
		return etcdErr.NewError(etcdErr.EcodeNodeExist, "", UndefIndex, UndefTerm)
	}

	n.Children[name] = child

	return nil
}

// Remove function remove the node.
func (n *Node) Remove(recursive bool, callback func(path string)) *etcdErr.Error {

	if n.IsDir() && !recursive {
		// cannot delete a directory without set recursive to true
		return etcdErr.NewError(etcdErr.EcodeNotFile, "", UndefIndex, UndefTerm)
	}

	if !n.IsDir() { // key-value pair
		_, name := path.Split(n.Path)

		// find its parent and remove the node from the map
		if n.Parent != nil && n.Parent.Children[name] == n {
			delete(n.Parent.Children, name)
		}

		if callback != nil {
			callback(n.Path)
		}

		if !n.IsPermanent() {
			n.store.ttlKeyHeap.remove(n)
		}

		return nil
	}

	for _, child := range n.Children { // delete all children
		child.Remove(true, callback)
	}

	// delete self
	_, name := path.Split(n.Path)
	if n.Parent != nil && n.Parent.Children[name] == n {
		delete(n.Parent.Children, name)

		if callback != nil {
			callback(n.Path)
		}

		if !n.IsPermanent() {
			n.store.ttlKeyHeap.remove(n)
		}

	}

	return nil
}

func (n *Node) Pair(recurisive, sorted bool) KeyValuePair {
	if n.IsDir() {
		pair := KeyValuePair{
			Key: n.Path,
			Dir: true,
		}
		pair.Expiration, pair.TTL = n.ExpirationAndTTL()

		if !recurisive {
			return pair
		}

		children, _ := n.List()
		pair.KVPairs = make([]KeyValuePair, len(children))

		// we do not use the index in the children slice directly
		// we need to skip the hidden one
		i := 0

		for _, child := range children {

			if child.IsHidden() { // get will not list hidden node
				continue
			}

			pair.KVPairs[i] = child.Pair(recurisive, sorted)

			i++
		}

		// eliminate hidden nodes
		pair.KVPairs = pair.KVPairs[:i]
		if sorted {
			sort.Sort(pair.KVPairs)
		}

		return pair
	}

	pair := KeyValuePair{
		Key:   n.Path,
		Value: n.Value,
	}
	pair.Expiration, pair.TTL = n.ExpirationAndTTL()
	return pair
}

func (n *Node) UpdateTTL(expireTime time.Time) {

	if !n.IsPermanent() {
		if expireTime.IsZero() {
			// from ttl to permanent
			// remove from ttl heap
			n.store.ttlKeyHeap.remove(n)
		} else {
			// update ttl
			n.ExpireTime = expireTime
			// update ttl heap
			n.store.ttlKeyHeap.update(n)
		}

	} else {
		if !expireTime.IsZero() {
			// from permanent to ttl
			n.ExpireTime = expireTime
			// push into ttl heap
			n.store.ttlKeyHeap.push(n)
		}
	}
}

// Clone function clone the node recursively and return the new node.
// If the node is a directory, it will clone all the content under this directory.
// If the node is a key-value pair, it will clone the pair.
func (n *Node) Clone() *Node {
	if !n.IsDir() {
		return newKV(n.store, n.Path, n.Value, n.CreateIndex, n.CreateTerm, n.Parent, n.ACL, n.ExpireTime)
	}

	clone := newDir(n.store, n.Path, n.CreateIndex, n.CreateTerm, n.Parent, n.ACL, n.ExpireTime)

	for key, child := range n.Children {
		clone.Children[key] = child.Clone()
	}

	return clone
}

// recoverAndclean function help to do recovery.
// Two things need to be done: 1. recovery structure; 2. delete expired nodes

// If the node is a directory, it will help recover children's parent pointer and recursively
// call this function on its children.
// We check the expire last since we need to recover the whole structure first and add all the
// notifications into the event history.
func (n *Node) recoverAndclean() {
	if n.IsDir() {
		for _, child := range n.Children {
			child.Parent = n
			child.store = n.store
			child.recoverAndclean()
		}
	}

	if !n.ExpireTime.IsZero() {
		n.store.ttlKeyHeap.push(n)
	}

}
<|endoftext|>


<commit_before>package copperd

import (
	"fmt"
	"sync"

	"github.com/snaury/copper"
)

type localEndpointKey struct {
	client   *serverClient
	targetID int64
}

type localEndpoint struct {
	owner    *server
	pub      *serverPublication
	key      localEndpointKey
	active   uint32
	settings PublishSettings
}

var _ endpointReference = &localEndpoint{}

func (endpoint *localEndpoint) getEndpointsLocked() []Endpoint {
	return nil
}

func (endpoint *localEndpoint) handleRequestLocked(client copper.Stream) bool {
	if endpoint.pub != nil && endpoint.active < endpoint.settings.Concurrency {
		return endpoint.passthruRequestLocked(client)
	}
	return false
}

func (endpoint *localEndpoint) passthruRequestLocked(client copper.Stream) bool {
	endpoint.active++
	if endpoint.active == endpoint.settings.Concurrency {
		delete(endpoint.pub.ready, endpoint)
	}
	defer func() {
		if endpoint.pub != nil {
			endpoint.active--
			endpoint.pub.ready[endpoint] = struct{}{}
			endpoint.pub.wakeupWaitersLocked(1)
		}
	}()
	endpoint.owner.lock.Unlock()
	defer endpoint.owner.lock.Lock()
	remote, err := endpoint.key.client.conn.Open(endpoint.key.targetID)
	if err != nil {
		// this client has already disconnected
		return false
	}
	passthruBoth(client, remote)
	return true
}

type serverPublication struct {
	owner    *server
	name     string
	targetID int64

	endpoints     map[localEndpointKey]*localEndpoint
	distances     map[uint32]int
	maxqueuesizes map[uint32]int
	settings      PublishSettings

	ready map[*localEndpoint]struct{}
	queue map[*sync.Cond]struct{}

	subscriptions map[*serverSubscription]struct{}
}

var _ endpointReference = &serverPublication{}

func (pub *serverPublication) waitInQueueLocked(waiter *sync.Cond) {
	pub.queue[waiter] = struct{}{}
	waiter.Wait()
	delete(pub.queue, waiter)
}

func (pub *serverPublication) wakeupWaitersLocked(n int) {
	for waiter := range pub.queue {
		delete(pub.queue, waiter)
		waiter.Signal()
		n--
		if n == 0 {
			break
		}
	}
}

func (pub *serverPublication) getEndpointsLocked() []Endpoint {
	return []Endpoint{
		Endpoint{
			Network:  "",
			Address:  "",
			TargetID: pub.targetID,
		},
	}
}

func (pub *serverPublication) handleRequestLocked(client copper.Stream) bool {
	var waiter *sync.Cond
	for len(pub.endpoints) > 0 {
		for endpoint := range pub.ready {
			if endpoint.handleRequestLocked(client) {
				return true
			}
		}
		if len(pub.queue) >= int(pub.settings.MaxQueueSize) {
			// Queue for this publication is already full
			client.CloseWithError(ErrOverCapacity)
			return true
		}
		if waiter == nil {
			// TODO: support for timeout and cancellation
			waiter = sync.NewCond(&pub.owner.lock)
		}
		pub.waitInQueueLocked(waiter)
	}
	return false
}

func (s *server) publishLocked(name string, key localEndpointKey, settings PublishSettings) (*localEndpoint, error) {
	pubs := s.pubsByName[name]
	if pubs == nil {
		pubs = make(map[uint32]*serverPublication)
		s.pubsByName[name] = pubs
	}

	pub := pubs[settings.Priority]
	if pub == nil {
		pub = &serverPublication{
			owner:    s,
			name:     name,
			targetID: s.allocateTargetID(),

			endpoints:     make(map[localEndpointKey]*localEndpoint),
			distances:     make(map[uint32]int),
			maxqueuesizes: make(map[uint32]int),
			settings:      settings,

			ready: make(map[*localEndpoint]struct{}),
			queue: make(map[*sync.Cond]struct{}),

			subscriptions: make(map[*serverSubscription]struct{}),
		}
		pubs[settings.Priority] = pub
		s.pubByTarget[pub.targetID] = pub
		for sub := range s.subsByName[pub.name] {
			sub.addPublicationLocked(pub)
		}
	}

	if pub.endpoints[key] != nil {
		return nil, fmt.Errorf("endpoint with key [target=%d] is already published", key.targetID)
	}

	endpoint := &localEndpoint{
		owner:    s,
		pub:      pub,
		key:      key,
		active:   0,
		settings: settings,
	}

	pub.endpoints[endpoint.key] = endpoint
	if len(pub.endpoints) != 1 {
		if pub.settings.Distance < endpoint.settings.Distance {
			pub.settings.Distance = endpoint.settings.Distance
		}
		pub.settings.Concurrency += endpoint.settings.Concurrency
		if pub.settings.MaxQueueSize < endpoint.settings.MaxQueueSize {
			pub.settings.MaxQueueSize = endpoint.settings.MaxQueueSize
		}
	}
	pub.distances[endpoint.settings.Distance]++
	pub.maxqueuesizes[endpoint.settings.MaxQueueSize]++
	pub.ready[endpoint] = struct{}{}
	pub.wakeupWaitersLocked(int(endpoint.settings.Concurrency))

	for watcher := range pub.owner.pubWatchers {
		watcher.addChangedLocked(pub)
	}
	return endpoint, nil
}

func (endpoint *localEndpoint) unpublishLocked() error {
	pub := endpoint.pub
	if pub == nil {
		return fmt.Errorf("endpoint is not published")
	}
	endpoint.pub = nil

	delete(pub.ready, endpoint)
	if decrementCounterUint32(pub.distances, endpoint.settings.Distance) {
		pub.settings.Distance = 0
		for distance := range pub.distances {
			if pub.settings.Distance < distance {
				pub.settings.Distance = distance
			}
		}
	}
	pub.settings.Concurrency -= endpoint.settings.Concurrency
	if decrementCounterUint32(pub.maxqueuesizes, endpoint.settings.MaxQueueSize) {
		pub.settings.MaxQueueSize = 0
		for maxqueuesize := range pub.maxqueuesizes {
			if pub.settings.MaxQueueSize < maxqueuesize {
				pub.settings.MaxQueueSize = maxqueuesize
			}
		}
	}
	delete(pub.endpoints, endpoint.key)

	pub.wakeupWaitersLocked(len(pub.queue))
	if len(pub.endpoints) == 0 {
		delete(pub.owner.pubByTarget, pub.targetID)
		if pubs := pub.owner.pubsByName[pub.name]; pubs != nil {
			delete(pubs, pub.settings.Priority)
			if len(pubs) == 0 {
				delete(pub.owner.pubsByName, pub.name)
			}
		}
		for watcher := range pub.owner.pubWatchers {
			watcher.addRemovedLocked(pub)
		}
		for sub := range pub.subscriptions {
			sub.removePublicationLocked(pub)
		}
	} else {
		for watcher := range pub.owner.pubWatchers {
			watcher.addChangedLocked(pub)
		}
	}
	return nil
}
<commit_msg>Sum queue sizes from publications<commit_after>package copperd

import (
	"fmt"
	"sync"

	"github.com/snaury/copper"
)

type localEndpointKey struct {
	client   *serverClient
	targetID int64
}

type localEndpoint struct {
	owner    *server
	pub      *serverPublication
	key      localEndpointKey
	active   uint32
	settings PublishSettings
}

var _ endpointReference = &localEndpoint{}

func (endpoint *localEndpoint) getEndpointsLocked() []Endpoint {
	return nil
}

func (endpoint *localEndpoint) handleRequestLocked(client copper.Stream) bool {
	if endpoint.pub != nil && endpoint.active < endpoint.settings.Concurrency {
		return endpoint.passthruRequestLocked(client)
	}
	return false
}

func (endpoint *localEndpoint) passthruRequestLocked(client copper.Stream) bool {
	endpoint.active++
	if endpoint.active == endpoint.settings.Concurrency {
		delete(endpoint.pub.ready, endpoint)
	}
	defer func() {
		if endpoint.pub != nil {
			endpoint.active--
			endpoint.pub.ready[endpoint] = struct{}{}
			endpoint.pub.wakeupWaitersLocked(1)
		}
	}()
	endpoint.owner.lock.Unlock()
	defer endpoint.owner.lock.Lock()
	remote, err := endpoint.key.client.conn.Open(endpoint.key.targetID)
	if err != nil {
		// this client has already disconnected
		return false
	}
	passthruBoth(client, remote)
	return true
}

type serverPublication struct {
	owner    *server
	name     string
	targetID int64

	endpoints map[localEndpointKey]*localEndpoint
	distances map[uint32]int
	settings  PublishSettings

	ready map[*localEndpoint]struct{}
	queue map[*sync.Cond]struct{}

	subscriptions map[*serverSubscription]struct{}
}

var _ endpointReference = &serverPublication{}

func (pub *serverPublication) waitInQueueLocked(waiter *sync.Cond) {
	pub.queue[waiter] = struct{}{}
	waiter.Wait()
	delete(pub.queue, waiter)
}

func (pub *serverPublication) wakeupWaitersLocked(n int) {
	for waiter := range pub.queue {
		delete(pub.queue, waiter)
		waiter.Signal()
		n--
		if n == 0 {
			break
		}
	}
}

func (pub *serverPublication) getEndpointsLocked() []Endpoint {
	return []Endpoint{
		Endpoint{
			Network:  "",
			Address:  "",
			TargetID: pub.targetID,
		},
	}
}

func (pub *serverPublication) handleRequestLocked(client copper.Stream) bool {
	var waiter *sync.Cond
	for len(pub.endpoints) > 0 {
		for endpoint := range pub.ready {
			if endpoint.handleRequestLocked(client) {
				return true
			}
		}
		if len(pub.queue) >= int(pub.settings.MaxQueueSize) {
			// Queue for this publication is already full
			client.CloseWithError(ErrOverCapacity)
			return true
		}
		if waiter == nil {
			// TODO: support for timeout and cancellation
			waiter = sync.NewCond(&pub.owner.lock)
		}
		pub.waitInQueueLocked(waiter)
	}
	return false
}

func (s *server) publishLocked(name string, key localEndpointKey, settings PublishSettings) (*localEndpoint, error) {
	pubs := s.pubsByName[name]
	if pubs == nil {
		pubs = make(map[uint32]*serverPublication)
		s.pubsByName[name] = pubs
	}

	pub := pubs[settings.Priority]
	if pub == nil {
		pub = &serverPublication{
			owner:    s,
			name:     name,
			targetID: s.allocateTargetID(),

			endpoints: make(map[localEndpointKey]*localEndpoint),
			distances: make(map[uint32]int),
			settings:  settings,

			ready: make(map[*localEndpoint]struct{}),
			queue: make(map[*sync.Cond]struct{}),

			subscriptions: make(map[*serverSubscription]struct{}),
		}
		pubs[settings.Priority] = pub
		s.pubByTarget[pub.targetID] = pub
		for sub := range s.subsByName[pub.name] {
			sub.addPublicationLocked(pub)
		}
	} else {
		if pub.endpoints[key] != nil {
			return nil, fmt.Errorf("new endpoint with target=%d conflicts with a previous publication", key.targetID)
		}
		if pub.settings.Concurrency+settings.Concurrency < pub.settings.Concurrency {
			return nil, fmt.Errorf("new endpoint with target=%d overflows maximum capacity", key.targetID)
		}
		if pub.settings.MaxQueueSize+settings.MaxQueueSize < pub.settings.MaxQueueSize {
			return nil, fmt.Errorf("new endpoint with target=%d overflows maximum queue size", key.targetID)
		}
		if pub.settings.Distance < settings.Distance {
			pub.settings.Distance = settings.Distance
		}
		pub.settings.Concurrency += settings.Concurrency
		pub.settings.MaxQueueSize += settings.MaxQueueSize
	}
	pub.distances[settings.Distance]++

	endpoint := &localEndpoint{
		owner:    s,
		pub:      pub,
		key:      key,
		active:   0,
		settings: settings,
	}

	pub.endpoints[endpoint.key] = endpoint
	pub.ready[endpoint] = struct{}{}

	pub.wakeupWaitersLocked(int(settings.Concurrency))
	for watcher := range pub.owner.pubWatchers {
		watcher.addChangedLocked(pub)
	}

	return endpoint, nil
}

func (endpoint *localEndpoint) unpublishLocked() error {
	pub := endpoint.pub
	if pub == nil {
		return fmt.Errorf("endpoint is not published")
	}
	endpoint.pub = nil

	delete(pub.ready, endpoint)
	if decrementCounterUint32(pub.distances, endpoint.settings.Distance) {
		pub.settings.Distance = 0
		for distance := range pub.distances {
			if pub.settings.Distance < distance {
				pub.settings.Distance = distance
			}
		}
	}
	pub.settings.Concurrency -= endpoint.settings.Concurrency
	pub.settings.MaxQueueSize -= endpoint.settings.MaxQueueSize
	delete(pub.endpoints, endpoint.key)

	if len(pub.endpoints) == 0 {
		delete(pub.owner.pubByTarget, pub.targetID)
		if pubs := pub.owner.pubsByName[pub.name]; pubs != nil {
			delete(pubs, pub.settings.Priority)
			if len(pubs) == 0 {
				delete(pub.owner.pubsByName, pub.name)
			}
		}
		for waiter := range pub.queue {
			waiter.Signal()
		}
		for watcher := range pub.owner.pubWatchers {
			watcher.addRemovedLocked(pub)
		}
		for sub := range pub.subscriptions {
			sub.removePublicationLocked(pub)
		}
	} else {
		if len(pub.queue) > int(pub.settings.MaxQueueSize) {
			pub.wakeupWaitersLocked(int(pub.settings.MaxQueueSize) - len(pub.queue))
		}
		for watcher := range pub.owner.pubWatchers {
			watcher.addChangedLocked(pub)
		}
	}
	return nil
}
<|endoftext|>


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

package cmac

import (
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestMsb(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type MsbTest struct{}

func init() { RegisterTestSuite(&MsbTest{}) }

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *MsbTest) DoesFoo() {
	ExpectEq("TODO", "")
}
<commit_msg>Added test names.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmac

import (
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestMsb(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type MsbTest struct{}

func init() { RegisterTestSuite(&MsbTest{}) }

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *MsbTest) EmptyBuffer() {
	ExpectEq("TODO", "")
}

func (t *MsbTest) MostSignficantIsOne() {
	ExpectEq("TODO", "")
}

func (t *MsbTest) MostSignficantIsZero() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/rpc"
	"strings"
	"time"

	"github.com/ianremmler/bort"
	"github.com/thoj/go-ircevent"
)

var (
	// flags
	flags   Config
	cfgFile string

	con    *irc.Connection
	client *rpc.Client
)

var cfg = Config{
	Nick:       "bort",
	Server:     "irc.freenode.net:6667",
	Address:    bort.DefaultAddress,
	Prefix:     "bort:",
	PollPeriod: 5,
}

type Config struct {
	Nick       string `json:"nick"`
	Server     string `json:"server"`
	Channel    string `json:"channel"`
	Address    string `json:"address"`
	Prefix     string `json:"prefix"`
	PollPeriod uint   `json:"poll_period"`
}

func main() {
	flag.Parse()
	config()
	if len(cfg.Channel) < 2 || !strings.HasPrefix(cfg.Channel, "#") {
		log.Fatalf("'%s' is not a valid channel", cfg.Channel)
	}
	con = irc.IRC(cfg.Nick, cfg.Nick)
	if err := con.Connect(cfg.Server); err != nil {
		log.Fatalln(err)
	}
	con.AddCallback("001", func(e *irc.Event) {
		con.Join(cfg.Channel)
	})
	con.AddCallback("JOIN", func(e *irc.Event) {
		if e.Nick == cfg.Nick {
			con.ClearCallback("JOIN")
			log.Printf("joined %s\n", e.Message())
			connectPlug()
			go func() {
				poll := time.Tick(time.Duration(cfg.PollPeriod) * time.Second)
				for {
					<-poll
					deliverPushes()
				}
			}()
		}
	})
	con.AddCallback("PRIVMSG", handleEvent)
	con.Loop()
}

func handleEvent(evt *irc.Event) {
	msg := newMessage(evt)
	res := &bort.Response{}
	if client == nil {
		if connectPlug() != nil {
			return
		}
	}
	if err := client.Call("Plugin.Process", msg, res); err != nil {
		log.Println(err)
		switch err {
		case rpc.ErrShutdown, io.EOF, io.ErrUnexpectedEOF:
			log.Println("disconnected from bortplug")
			connectPlug()
		}
		return
	}
	if err := send(res); err != nil {
		log.Println(err)
	}
}

func deliverPushes() {
	if client == nil {
		if connectPlug() != nil {
			return
		}
	}
	res := []bort.Response{}
	if err := client.Call("Plugin.Pull", struct{}{}, &res); err != nil {
		log.Println(err)
		switch err {
		case rpc.ErrShutdown, io.EOF, io.ErrUnexpectedEOF:
			connectPlug()
		}
		return
	}
	for i := range res {
		if res[i].Target == "" {
			res[i].Target = cfg.Channel
		}
		if err := send(&res[i]); err != nil {
			log.Println(err)
		}
	}
}

func send(res *bort.Response) error {
	switch res.Type {
	case bort.None:
	case bort.PrivMsg:
		for _, str := range strings.Split(strings.TrimRight(res.Text, "\n"), "\n") {
			con.Privmsg(res.Target, str)
		}
	case bort.Action:
		con.Action(res.Target, strings.SplitN(res.Text, "\n", 2)[0])
	default:
		return fmt.Errorf("unknown response type: %d", res.Type)
	}
	return nil
}

func newMessage(evt *irc.Event) *bort.Message {
	msg := &bort.Message{
		Channel: cfg.Channel,
		Host:    evt.Host,
		Nick:    evt.Nick,
		Raw:     evt.Raw,
		Source:  evt.Source,
		Target:  cfg.Channel,
		Text:    evt.Message(),
		User:    evt.User,
	}
	text := strings.TrimSpace(evt.Message())
	if strings.HasPrefix(text, cfg.Prefix) {
		cmdStr := strings.TrimLeft(text[len(cfg.Prefix):], " ")
		cmdStr += " " // append space to ensure SplitN returns 2 strings
		cmdAndArgs := strings.SplitN(cmdStr, " ", 2)
		if len(cmdAndArgs) == 2 {
			msg.Command = cmdAndArgs[0]
			msg.Text = strings.TrimSpace(cmdAndArgs[1])
		}
	}
	if evt.Arguments[0] != cfg.Channel {
		msg.Target = evt.Nick
	}
	return msg
}

func connectPlug() error {
	var err error
	if client != nil {
		client.Close()
	}
	client, err = rpc.Dial("tcp", cfg.Address)
	if err == nil {
		log.Println("connected to bortplug")
	}
	return err
}

func config() {
	if cfgData, err := bort.LoadConfig(cfgFile); err != nil {
		log.Println(err)
	} else if err := json.Unmarshal(cfgData, &cfg); err != nil {
		log.Println(err)
	}
	flag.Visit(func(f *flag.Flag) {
		switch f.Name {
		case "n":
			cfg.Nick = flags.Nick
		case "s":
			cfg.Server = flags.Server
		case "c":
			cfg.Channel = flags.Channel
		case "a":
			cfg.Address = flags.Address
		case "p":
			cfg.Prefix = flags.Prefix
		case "t":
			cfg.PollPeriod = flags.PollPeriod
		}
	})
}

func init() {
	flag.StringVar(&flags.Nick, "n", cfg.Nick, "nick of the bot")
	flag.StringVar(&flags.Server, "s", cfg.Server, "IRC server")
	flag.StringVar(&flags.Channel, "c", cfg.Channel, "channel")
	flag.StringVar(&flags.Address, "a", cfg.Address, "bortplug address")
	flag.StringVar(&flags.Prefix, "p", cfg.Prefix, "command prefix")
	flag.UintVar(&flags.PollPeriod, "t", cfg.PollPeriod, "poll period in seconds")
	flag.StringVar(&cfgFile, "f", "", "configuration file")
}
<commit_msg>Poll period sanity check.  Saner default outbox size.<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/rpc"
	"strings"
	"time"

	"github.com/ianremmler/bort"
	"github.com/thoj/go-ircevent"
)

var (
	// flags
	flags   Config
	cfgFile string

	con    *irc.Connection
	client *rpc.Client
)

var cfg = Config{
	Nick:       "bort",
	Server:     "irc.freenode.net:6667",
	Address:    bort.DefaultAddress,
	Prefix:     "bort:",
	PollPeriod: 5,
}

type Config struct {
	Nick       string `json:"nick"`
	Server     string `json:"server"`
	Channel    string `json:"channel"`
	Address    string `json:"address"`
	Prefix     string `json:"prefix"`
	PollPeriod uint   `json:"poll_period"`
}

func main() {
	flag.Parse()
	config()
	if len(cfg.Channel) < 2 || !strings.HasPrefix(cfg.Channel, "#") {
		log.Fatalf("'%s' is not a valid channel", cfg.Channel)
	}
	if cfg.PollPeriod < 1 {
		cfg.PollPeriod = 1
	}

	con = irc.IRC(cfg.Nick, cfg.Nick)
	if err := con.Connect(cfg.Server); err != nil {
		log.Fatalln(err)
	}
	con.AddCallback("001", func(e *irc.Event) {
		con.Join(cfg.Channel)
	})
	con.AddCallback("JOIN", func(e *irc.Event) {
		if e.Nick == cfg.Nick {
			con.ClearCallback("JOIN")
			log.Printf("joined %s\n", e.Message())
			connectPlug()
			go func() {
				poll := time.Tick(time.Duration(cfg.PollPeriod) * time.Second)
				for {
					<-poll
					deliverPushes()
				}
			}()
		}
	})
	con.AddCallback("PRIVMSG", handleEvent)
	con.Loop()
}

func handleEvent(evt *irc.Event) {
	msg := newMessage(evt)
	res := &bort.Response{}
	if client == nil {
		if connectPlug() != nil {
			return
		}
	}
	if err := client.Call("Plugin.Process", msg, res); err != nil {
		log.Println(err)
		switch err {
		case rpc.ErrShutdown, io.EOF, io.ErrUnexpectedEOF:
			log.Println("disconnected from bortplug")
			connectPlug()
		}
		return
	}
	if err := send(res); err != nil {
		log.Println(err)
	}
}

func deliverPushes() {
	if client == nil {
		if connectPlug() != nil {
			return
		}
	}
	res := []bort.Response{}
	if err := client.Call("Plugin.Pull", struct{}{}, &res); err != nil {
		log.Println(err)
		switch err {
		case rpc.ErrShutdown, io.EOF, io.ErrUnexpectedEOF:
			connectPlug()
		}
		return
	}
	for i := range res {
		if res[i].Target == "" {
			res[i].Target = cfg.Channel
		}
		if err := send(&res[i]); err != nil {
			log.Println(err)
		}
	}
}

func send(res *bort.Response) error {
	switch res.Type {
	case bort.None:
	case bort.PrivMsg:
		for _, str := range strings.Split(strings.TrimRight(res.Text, "\n"), "\n") {
			con.Privmsg(res.Target, str)
		}
	case bort.Action:
		con.Action(res.Target, strings.SplitN(res.Text, "\n", 2)[0])
	default:
		return fmt.Errorf("unknown response type: %d", res.Type)
	}
	return nil
}

func newMessage(evt *irc.Event) *bort.Message {
	msg := &bort.Message{
		Channel: cfg.Channel,
		Host:    evt.Host,
		Nick:    evt.Nick,
		Raw:     evt.Raw,
		Source:  evt.Source,
		Target:  cfg.Channel,
		Text:    evt.Message(),
		User:    evt.User,
	}
	text := strings.TrimSpace(evt.Message())
	if strings.HasPrefix(text, cfg.Prefix) {
		cmdStr := strings.TrimLeft(text[len(cfg.Prefix):], " ")
		cmdStr += " " // append space to ensure SplitN returns 2 strings
		cmdAndArgs := strings.SplitN(cmdStr, " ", 2)
		if len(cmdAndArgs) == 2 {
			msg.Command = cmdAndArgs[0]
			msg.Text = strings.TrimSpace(cmdAndArgs[1])
		}
	}
	if evt.Arguments[0] != cfg.Channel {
		msg.Target = evt.Nick
	}
	return msg
}

func connectPlug() error {
	var err error
	if client != nil {
		client.Close()
	}
	client, err = rpc.Dial("tcp", cfg.Address)
	if err == nil {
		log.Println("connected to bortplug")
	}
	return err
}

func config() {
	if cfgData, err := bort.LoadConfig(cfgFile); err != nil {
		log.Println(err)
	} else if err := json.Unmarshal(cfgData, &cfg); err != nil {
		log.Println(err)
	}
	flag.Visit(func(f *flag.Flag) {
		switch f.Name {
		case "n":
			cfg.Nick = flags.Nick
		case "s":
			cfg.Server = flags.Server
		case "c":
			cfg.Channel = flags.Channel
		case "a":
			cfg.Address = flags.Address
		case "p":
			cfg.Prefix = flags.Prefix
		case "t":
			cfg.PollPeriod = flags.PollPeriod
		}
	})
}

func init() {
	flag.StringVar(&flags.Nick, "n", cfg.Nick, "nick of the bot")
	flag.StringVar(&flags.Server, "s", cfg.Server, "IRC server")
	flag.StringVar(&flags.Channel, "c", cfg.Channel, "channel")
	flag.StringVar(&flags.Address, "a", cfg.Address, "bortplug address")
	flag.StringVar(&flags.Prefix, "p", cfg.Prefix, "command prefix")
	flag.UintVar(&flags.PollPeriod, "t", cfg.PollPeriod, "poll period in seconds")
	flag.StringVar(&cfgFile, "f", "", "configuration file")
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/gophergala/dnsp"
)

var version = "0.0"

func main() {
	o := dnsp.Options{}
	flag.StringVar(&o.Bind, "bind", ":53", "address to bind to")
	flag.StringVar(&o.Server, "server", "8.8.8.8", "address to proxy to")
	flag.Parse()

	s := dnsp.NewServer(o)

	catch(func(sig os.Signal) int {
		log.Printf("debug=dnsp_shutdown signal=%s", sig)
		s.Shutdown()
		return 0
	}, syscall.SIGINT, syscall.SIGTERM)
	defer s.Shutdown() // in case of normal exit

	log.Printf("debug=dnsp_start bind=%s server=%s", o.Bind, o.Server)
	if err := s.ListenAndServe(); err != nil {
		log.Fatal(err)
	}
}

// catch handles system calls using the given handler function.
func catch(handler func(os.Signal) int, signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	for _, s := range signals {
		signal.Notify(c, s)
	}
	go func() {
		os.Exit(handler(<-c))
	}()
}
<commit_msg>add a command line flag to enable whitelist mode<commit_after>package main

import (
	"flag"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/gophergala/dnsp"
)

var version = "0.0"

func main() {
	o := dnsp.Options{}
	flag.StringVar(&o.Bind, "bind", ":53", "address to bind to")
	flag.StringVar(&o.Server, "server", "8.8.8.8", "address to proxy to")
	flag.BoolVar(&o.White, "white", false, "use a whitelist")
	flag.Parse()

	s := dnsp.NewServer(o)

	catch(func(sig os.Signal) int {
		log.Printf("debug=dnsp_shutdown signal=%s", sig)
		s.Shutdown()
		return 0
	}, syscall.SIGINT, syscall.SIGTERM)
	defer s.Shutdown() // in case of normal exit

	log.Printf("debug=dnsp_start bind=%s server=%s", o.Bind, o.Server)
	if err := s.ListenAndServe(); err != nil {
		log.Fatal(err)
	}
}

// catch handles system calls using the given handler function.
func catch(handler func(os.Signal) int, signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	for _, s := range signals {
		signal.Notify(c, s)
	}
	go func() {
		os.Exit(handler(<-c))
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"os"

	"github.com/viktomas/godu"
)

func main() {
	flag.Parse()
	roots := flag.Args()
	if len(roots) == 0 {
		roots = []string{"."}
	}
	tree := godu.GetSubTree(roots[0], ioutil.ReadDir)
	godu.InteractiveTree(tree, os.Stdout, os.Stdin)
}
<commit_msg>Fixing command to use file pointer<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"os"

	"github.com/viktomas/godu"
)

func main() {
	flag.Parse()
	roots := flag.Args()
	if len(roots) == 0 {
		roots = []string{"."}
	}
	tree := godu.GetSubTree(roots[0], ioutil.ReadDir)
	godu.InteractiveTree(&tree, os.Stdout, os.Stdin)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main // import "k8s.io/helm/cmd/helm"

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"google.golang.org/grpc"
)

const (
	homeEnvVar      = "HELM_HOME"
	hostEnvVar      = "HELM_HOST"
	tillerNamespace = "kube-system"
)

var (
	helmHome   string
	tillerHost string
)

// flagDebug is a signal that the user wants additional output.
var flagDebug bool

var globalUsage = `The Kubernetes package manager

To begin working with Helm, run the 'helm init' command:

	$ helm init

This will install Tiller to your running Kubernetes cluster.
It will also set up any necessary local configuration.

Common actions from this point include:

- helm search:    search for charts
- helm fetch:     download a chart to your local directory to view
- helm install:   upload the chart to Kubernetes
- helm list:      list releases of charts

Environment:
  $HELM_HOME      Set an alternative location for Helm files. By default, these are stored in ~/.helm
  $HELM_HOST      Set an alternative Tiller host. The format is host:port (default ":44134").
  $KUBECONFIG     Set an alternate Kubernetes configuration file (default: "~/.kube/config").
`

func newRootCmd(out io.Writer) *cobra.Command {
	cmd := &cobra.Command{
		Use:          "helm",
		Short:        "The Helm package manager for Kubernetes.",
		Long:         globalUsage,
		SilenceUsage: true,
		PersistentPostRun: func(cmd *cobra.Command, args []string) {
			teardown()
		},
	}
	home := os.Getenv(homeEnvVar)
	if home == "" {
		home = "$HOME/.helm"
	}
	thost := os.Getenv(hostEnvVar)
	p := cmd.PersistentFlags()
	p.StringVar(&helmHome, "home", home, "location of your Helm config. Overrides $HELM_HOME.")
	p.StringVar(&tillerHost, "host", thost, "address of tiller. Overrides $HELM_HOST.")
	p.BoolVarP(&flagDebug, "debug", "", false, "enable verbose output")

	cmd.AddCommand(
		newCreateCmd(out),
		newDeleteCmd(nil, out),
		newGetCmd(nil, out),
		newInitCmd(out),
		newInspectCmd(nil, out),
		newInstallCmd(nil, out),
		newListCmd(nil, out),
		newStatusCmd(nil, out),
		newUpgradeCmd(nil, out),
	)
	return cmd
}

// RootCommand is the top-level command for Helm.
var RootCommand = newRootCmd(os.Stdout)

func main() {
	cmd := RootCommand
	if err := cmd.Execute(); err != nil {
		os.Exit(1)
	}
}

func setupConnection(c *cobra.Command, args []string) error {
	if tillerHost == "" {
		tunnel, err := newTillerPortForwarder(tillerNamespace)
		if err != nil {
			return err
		}

		tillerHost = fmt.Sprintf(":%d", tunnel.Local)
		if flagDebug {
			fmt.Printf("Created tunnel using local port: '%d'\n", tunnel.Local)
		}
	}

	// Set up the gRPC config.
	if flagDebug {
		fmt.Printf("Server: %q\n", tillerHost)
	}
	return nil
}

func teardown() {
	if tunnel != nil {
		tunnel.Close()
	}
}

func checkArgsLength(expectedNum, actualNum int, requiredArgs ...string) error {
	if actualNum != expectedNum {
		arg := "arguments"
		if expectedNum == 1 {
			arg = "argument"
		}
		return fmt.Errorf("This command needs %v %s: %s", expectedNum, arg, strings.Join(requiredArgs, ", "))
	}
	return nil
}

// requireInit is a PreRunE implementation for validating that $HELM_HOME is configured.
func requireInit(cmd *cobra.Command, args []string) error {
	err := requireHome()
	if err != nil {
		return fmt.Errorf("%s (try running 'helm init')", err)
	}
	return nil
}

// prettyError unwraps or rewrites certain errors to make them more user-friendly.
func prettyError(err error) error {
	if err == nil {
		return nil
	}
	// This is ridiculous. Why is 'grpc.rpcError' not exported? The least they
	// could do is throw an interface on the lib that would let us get back
	// the desc. Instead, we have to pass ALL errors through this.
	return errors.New(grpc.ErrorDesc(err))
}
<commit_msg>fix(cmd): remove default port in cmd doc<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main // import "k8s.io/helm/cmd/helm"

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"google.golang.org/grpc"
)

const (
	homeEnvVar      = "HELM_HOME"
	hostEnvVar      = "HELM_HOST"
	tillerNamespace = "kube-system"
)

var (
	helmHome   string
	tillerHost string
)

// flagDebug is a signal that the user wants additional output.
var flagDebug bool

var globalUsage = `The Kubernetes package manager

To begin working with Helm, run the 'helm init' command:

	$ helm init

This will install Tiller to your running Kubernetes cluster.
It will also set up any necessary local configuration.

Common actions from this point include:

- helm search:    search for charts
- helm fetch:     download a chart to your local directory to view
- helm install:   upload the chart to Kubernetes
- helm list:      list releases of charts

Environment:
  $HELM_HOME      Set an alternative location for Helm files. By default, these are stored in ~/.helm
  $HELM_HOST      Set an alternative Tiller host. The format is host:port.
  $KUBECONFIG     Set an alternate Kubernetes configuration file (default: "~/.kube/config").
`

func newRootCmd(out io.Writer) *cobra.Command {
	cmd := &cobra.Command{
		Use:          "helm",
		Short:        "The Helm package manager for Kubernetes.",
		Long:         globalUsage,
		SilenceUsage: true,
		PersistentPostRun: func(cmd *cobra.Command, args []string) {
			teardown()
		},
	}
	home := os.Getenv(homeEnvVar)
	if home == "" {
		home = "$HOME/.helm"
	}
	thost := os.Getenv(hostEnvVar)
	p := cmd.PersistentFlags()
	p.StringVar(&helmHome, "home", home, "location of your Helm config. Overrides $HELM_HOME.")
	p.StringVar(&tillerHost, "host", thost, "address of tiller. Overrides $HELM_HOST.")
	p.BoolVarP(&flagDebug, "debug", "", false, "enable verbose output")

	cmd.AddCommand(
		newCreateCmd(out),
		newDeleteCmd(nil, out),
		newGetCmd(nil, out),
		newInitCmd(out),
		newInspectCmd(nil, out),
		newInstallCmd(nil, out),
		newListCmd(nil, out),
		newStatusCmd(nil, out),
		newUpgradeCmd(nil, out),
	)
	return cmd
}

// RootCommand is the top-level command for Helm.
var RootCommand = newRootCmd(os.Stdout)

func main() {
	cmd := RootCommand
	if err := cmd.Execute(); err != nil {
		os.Exit(1)
	}
}

func setupConnection(c *cobra.Command, args []string) error {
	if tillerHost == "" {
		tunnel, err := newTillerPortForwarder(tillerNamespace)
		if err != nil {
			return err
		}

		tillerHost = fmt.Sprintf(":%d", tunnel.Local)
		if flagDebug {
			fmt.Printf("Created tunnel using local port: '%d'\n", tunnel.Local)
		}
	}

	// Set up the gRPC config.
	if flagDebug {
		fmt.Printf("Server: %q\n", tillerHost)
	}
	return nil
}

func teardown() {
	if tunnel != nil {
		tunnel.Close()
	}
}

func checkArgsLength(expectedNum, actualNum int, requiredArgs ...string) error {
	if actualNum != expectedNum {
		arg := "arguments"
		if expectedNum == 1 {
			arg = "argument"
		}
		return fmt.Errorf("This command needs %v %s: %s", expectedNum, arg, strings.Join(requiredArgs, ", "))
	}
	return nil
}

// requireInit is a PreRunE implementation for validating that $HELM_HOME is configured.
func requireInit(cmd *cobra.Command, args []string) error {
	err := requireHome()
	if err != nil {
		return fmt.Errorf("%s (try running 'helm init')", err)
	}
	return nil
}

// prettyError unwraps or rewrites certain errors to make them more user-friendly.
func prettyError(err error) error {
	if err == nil {
		return nil
	}
	// This is ridiculous. Why is 'grpc.rpcError' not exported? The least they
	// could do is throw an interface on the lib that would let us get back
	// the desc. Instead, we have to pass ALL errors through this.
	return errors.New(grpc.ErrorDesc(err))
}
<|endoftext|>


<commit_before>package main

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

	log "github.com/Sirupsen/logrus"
	"github.com/spf13/viper"

	"github.com/rvolosatovs/systemgo/system"
	"github.com/rvolosatovs/systemgo/systemctl"
)

func main() {
	Boot()
	Serve()
}

var Conf Configuration

// TODO: introduce proper configuration
type Configuration struct {
	// Default target
	Target string `yaml:omitempty`

	// Paths to search for unit files
	Paths []string `yaml:omitempty`

	// Debug
	Debug bool `yaml:omitempty`

	// Port for system daemon to listen on
	Port port `yaml:omitempty`
}

type port int

func (p port) String() string {
	return fmt.Sprintf(":%v", int(p))
}

// Instance of a system
var System *system.Daemon

func init() {
	viper.SetConfigName("systemgo")

	viper.AddConfigPath(".")
	if os.Getenv("XDG_CONFIG_HOME") != "" {
		viper.AddConfigPath("$XDG_CONFIG_HOME/systemgo")
	}
	viper.AddConfigPath("/etc/systemgo")

	if err := viper.ReadInConfig(); err != nil {
		log.WithField("path", viper.ConfigFileUsed()).Errorf("Error reading config: %s", err)
	}
	log.WithField("path", viper.ConfigFileUsed()).Infof("Found configuration file")

	viper.SetDefault("port", 28357)
	viper.SetDefault("target", "default.target")
	viper.SetDefault("paths", []string{
		"/etc/systemd/system", "/lib/systemd/system",
	})
	viper.SetDefault("debug", false)

	if err := viper.Unmarshal(&Conf); err != nil {
		log.Errorf("Error parsing %s: %s", viper.ConfigFileUsed(), err)
	}

	if Conf.Debug {
		log.SetLevel(log.DebugLevel)
	}
}

// Boot initializes the system, sets the default paths, as specified in configuration and attempts to start the default target, falls back to "rescue.target", if it fails
func Boot() {
	// Initialize system
	System = system.New()
	System.SetPaths(Conf.Paths...)

	// Start the default target
	if err := System.Start(Conf.Target); err != nil {
		System.Log.Errorf("Error starting default target %s: %s", Conf.Target, err)
		if err = System.Start("rescue.target"); err != nil {
			System.Log.Errorf("Error starting rescue target %s: %s", "rescue.target", err)
		}
	}
}

// Listen for systemctl requests
func Serve() {
	if err := listenHTTP(Conf.Port.String()); err != nil {
		log.Fatalf("Error starting server on %s: %s", Conf.Port, err)
	}
}

// Handle systemctl requests using HTTP
func listenHTTP(addr string) (err error) {
	server := http.NewServeMux()

	for name, handler := range systemctl.Handlers(System) {
		func(handler systemctl.Handler) {
			server.HandleFunc("/"+name, func(w http.ResponseWriter, req *http.Request) {
				//msg := []systemctl.Response{}
				v := req.URL.Query()

				names, ok := v["unit"]
				if !ok {
					names = []string{}
				}

				var resp []byte
				if resp, err = json.Marshal(handler(names...)); err != nil {
					log.Printf("json.Marshal(result): %s", err)
					return
				}

				if _, err = w.Write(resp); err != nil {
					log.Printf("Write(resp): %s", err)
				}
			})
		}(handler)
	}

	return http.ListenAndServe(addr, server)
}
<commit_msg>init: reimplementation<commit_after>package main

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

	log "github.com/Sirupsen/logrus"

	"github.com/rvolosatovs/systemgo/config"
	"github.com/rvolosatovs/systemgo/system"
	"github.com/rvolosatovs/systemgo/systemctl"
)

// Initializes the system, sets the default paths, as specified in configuration and attempts to start the default target, falls back to "rescue.target", if it fails
func main() {
	// Initialize system
	System = system.New()
	System.SetPaths(config.Paths...)

	go Serve()

	// Start the default target
	if err := System.Start(config.Target); err != nil {
		log.Errorf("Error starting default target %s: %s", config.Target, err)
		if err = System.Start(config.RESCUE_TARGET); err != nil {
			log.Errorf("Error starting rescue target %s: %s", config.RESCUE_TARGET, err)
		}
	}

	select {}
}

// Instance of a system
var System *system.Daemon

// Listen for systemctl requests
func Serve() {
	for {
		if err := listenHTTP(config.Port.String()); err != nil {
			log.Errorf("Error listening on %v: %s", config.Port, err)
		}
		log.Infof("Retrying in %v seconds", config.Retry)
		time.Sleep(config.Retry)
	}
}

// Handle systemctl requests using HTTP
func listenHTTP(addr string) (err error) {
	server := http.NewServeMux()

	for name, handler := range systemctl.Handlers(System) {
		func(handler systemctl.Handler) {
			server.HandleFunc("/"+name, func(w http.ResponseWriter, req *http.Request) {
				//msg := []systemctl.Response{}
				v := req.URL.Query()

				names, ok := v["unit"]
				if !ok {
					names = []string{}
				}

				var resp []byte
				if resp, err = json.Marshal(handler(names...)); err != nil {
					log.Printf("json.Marshal(result): %s", err)
					return
				}

				if _, err = w.Write(resp); err != nil {
					log.Printf("Write(resp): %s", err)
				}
			})
		}(handler)
	}

	log.Infof("Listening on http://localhost%s", config.Port)
	return http.ListenAndServe(addr, server)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"

	"github.com/benjdewan/pachelbel/config"
	"github.com/benjdewan/pachelbel/connection"
	"github.com/benjdewan/pachelbel/runner"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var provisionCmd = &cobra.Command{
	Use:   "provision",
	Short: "Idempotent provisioner of compose deployments",
	Long: `pachelbel provision reads in YAML configuration(s) describing a list of
deployments that should exist in a list of clusters of specified sizes,
and ensures they do.

If the deployments do not exist, they are created. If they exist, but are
the wrong size they are scaled. If they are deployed as specified in the
configuration no actions are taken.`,
	Run: runProvision,
}

func runProvision(cmd *cobra.Command, args []string) {
	assertCanStart(args)

	cxn, err := connection.New(viper.GetString("api-key"),
		viper.GetString("log-file"))
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		if closeErr := cxn.Close(); closeErr != nil {
			panic(closeErr)
		}
	}()

	cfg, err := readConfigs(cxn, args)
	if err != nil {
		log.Fatal(err)
	} else if len(cfg.Runners) == 0 {
		fmt.Println("Nothing to do")
		return
	}

	ctl := runner.NewController(cxn, viper.GetBool("dry-run"))
	if err := ctl.Run(cfg.Runners); err != nil {
		log.Fatal(err)
	}

	writeOutput(cxn, cfg.EndpointMap)
}

func writeOutput(cxn *connection.Connection, endpointMap map[string]string) {
	dst := viper.GetString("output")
	if err := cxn.ConnectionYAML(endpointMap, dst); err != nil {
		log.Fatal(err)
	}
}

func readConfigs(cxn *connection.Connection, paths []string) (*config.Config, error) {
	var err error
	config.Databases, err = cxn.SupportedDatabases()
	if err != nil {
		return nil, err
	}

	config.Clusters, err = cxn.Clusters()
	if err != nil {
		return nil, err
	}

	config.Datacenters, err = cxn.Datacenters()
	if err != nil {
		return nil, err
	}

	config.CXN = cxn

	config.BuildClusterFilter(viper.GetStringSlice("cluster"))
	config.BuildDatacenterFilter(viper.GetStringSlice("datacenter"))

	return config.ReadFiles(paths)
}

func assertCanStart(args []string) {
	if len(args) == 0 {
		log.Fatal("The 'provision' command requires at least one configuration file or directory as input")
	}
}

func init() {
	RootCmd.AddCommand(provisionCmd)
	addClusterFlag()
	addDatacenterFlag()
	addOutputFlag()
}

func addClusterFlag() {
	provisionCmd.Flags().StringSliceP("cluster", "c", []string{},
		`By default pachelbel provision will provision
				 every deployment provided. Use this flag to
				 limit pachelbel to only process deployments
				 to the specified cluster.

				 This flag can be repeated to specify multiple
				 clusters`)
	viper.BindPFlag("cluster", provisionCmd.Flags().Lookup("cluster"))
}

func addDatacenterFlag() {
	provisionCmd.Flags().StringSliceP("datacenter", "d", []string{},
		`By default pachelbel provision will
				 provision every deployment provided. Use this
				 flat to limit pachelbel to only process
				 deployments to the specified datacenter.

				 This flag can be repeated to specify multiple
				 datacenters.`)
	viper.BindPFlag("datacenter", provisionCmd.Flags().Lookup("datacenter"))
}

func addOutputFlag() {
	provisionCmd.Flags().StringP("output", "o", "./connection-info.yml",
		`The file to write connection string
				 information to.`)
	viper.BindPFlag("output", provisionCmd.Flags().Lookup("output"))
}
<commit_msg>Remove an unnecessary assignment<commit_after>package cmd

import (
	"fmt"
	"log"

	"github.com/benjdewan/pachelbel/config"
	"github.com/benjdewan/pachelbel/connection"
	"github.com/benjdewan/pachelbel/runner"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var provisionCmd = &cobra.Command{
	Use:   "provision",
	Short: "Idempotent provisioner of compose deployments",
	Long: `pachelbel provision reads in YAML configuration(s) describing a list of
deployments that should exist in a list of clusters of specified sizes,
and ensures they do.

If the deployments do not exist, they are created. If they exist, but are
the wrong size they are scaled. If they are deployed as specified in the
configuration no actions are taken.`,
	Run: runProvision,
}

func runProvision(cmd *cobra.Command, args []string) {
	assertCanStart(args)

	cxn, err := connection.New(viper.GetString("api-key"),
		viper.GetString("log-file"))
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		if closeErr := cxn.Close(); closeErr != nil {
			panic(closeErr)
		}
	}()

	cfg, err := readConfigs(cxn, args)
	if err != nil {
		log.Fatal(err)
	} else if len(cfg.Runners) == 0 {
		fmt.Println("Nothing to do")
		return
	}

	if err := runner.
		NewController(cxn, viper.GetBool("dry-run")).
		Run(cfg.Runners); err != nil {
		log.Fatal(err)
	}

	writeOutput(cxn, cfg.EndpointMap)
}

func writeOutput(cxn *connection.Connection, endpointMap map[string]string) {
	dst := viper.GetString("output")
	if err := cxn.ConnectionYAML(endpointMap, dst); err != nil {
		log.Fatal(err)
	}
}

func readConfigs(cxn *connection.Connection, paths []string) (*config.Config, error) {
	var err error
	config.Databases, err = cxn.SupportedDatabases()
	if err != nil {
		return nil, err
	}

	config.Clusters, err = cxn.Clusters()
	if err != nil {
		return nil, err
	}

	config.Datacenters, err = cxn.Datacenters()
	if err != nil {
		return nil, err
	}

	config.CXN = cxn

	config.BuildClusterFilter(viper.GetStringSlice("cluster"))
	config.BuildDatacenterFilter(viper.GetStringSlice("datacenter"))

	return config.ReadFiles(paths)
}

func assertCanStart(args []string) {
	if len(args) == 0 {
		log.Fatal("The 'provision' command requires at least one configuration file or directory as input")
	}
}

func init() {
	RootCmd.AddCommand(provisionCmd)
	addClusterFlag()
	addDatacenterFlag()
	addOutputFlag()
}

func addClusterFlag() {
	provisionCmd.Flags().StringSliceP("cluster", "c", []string{},
		`By default pachelbel provision will provision
				 every deployment provided. Use this flag to
				 limit pachelbel to only process deployments
				 to the specified cluster.

				 This flag can be repeated to specify multiple
				 clusters`)
	viper.BindPFlag("cluster", provisionCmd.Flags().Lookup("cluster"))
}

func addDatacenterFlag() {
	provisionCmd.Flags().StringSliceP("datacenter", "d", []string{},
		`By default pachelbel provision will
				 provision every deployment provided. Use this
				 flat to limit pachelbel to only process
				 deployments to the specified datacenter.

				 This flag can be repeated to specify multiple
				 datacenters.`)
	viper.BindPFlag("datacenter", provisionCmd.Flags().Lookup("datacenter"))
}

func addOutputFlag() {
	provisionCmd.Flags().StringP("output", "o", "./connection-info.yml",
		`The file to write connection string
				 information to.`)
	viper.BindPFlag("output", provisionCmd.Flags().Lookup("output"))
}
<|endoftext|>