<commit_before>package dominoes

// Source: exercism/problem-specifications
// Commit: b5bc74d dominoes: Apply new "input" policy
// Problem Specifications Version: 2.1.0

var testCases = []struct {
	description string
	dominoes    []Domino
	valid       bool // true => can chain, false => cannot chain
}{
	{
		"empty input = empty output",
		[]Domino{},
		true,
	},
	{
		"singleton input = singleton output",
		[]Domino{{1, 1}},
		true,
	},
	{
		"singleton that can't be chained",
		[]Domino{{1, 2}},
		false,
	},
	{
		"three elements",
		[]Domino{{1, 2}, {3, 1}, {2, 3}},
		true,
	},
	{
		"can reverse dominoes",
		[]Domino{{1, 2}, {1, 3}, {2, 3}},
		true,
	},
	{
		"can't be chained",
		[]Domino{{1, 2}, {4, 1}, {2, 3}},
		false,
	},
	{
		"disconnected - simple",
		// This meets the requirement of being possibly-Euclidean.
		// All vertices have even degree.
		// Nevertheless, there is no chain here, as there's no way to get from 1 to 2.
		// This test (and the two following) prevent solutions from using the even-degree test as the sole criterion,
		// as that is not a sufficient condition.
		[]Domino{{1, 1}, {2, 2}},
		false,
	},
	{
		"disconnected - double loop",
		[]Domino{{1, 2}, {2, 1}, {3, 4}, {4, 3}},
		false,
	},
	{
		"disconnected - single isolated",
		[]Domino{{1, 2}, {2, 3}, {3, 1}, {4, 4}},
		false,
	},
	// {
	// 	"need backtrack",
	// 	// Some solutions may make a chain out of (1, 2), (2, 3), (3, 1)
	// 	// then decide that since there are no more dominoes containing a 1,
	// 	// there is no chain possible.
	// 	// There is indeed a chain here, so this test checks for this line of reasoning.
	// 	// You need to place the (2, 4) after the (1, 2) rather than the (2, 3).
	// 	[]Domino{{1, 2}, {2, 3}, {3, 1}, {2, 4}, {2, 4}},
	// 	true,
	// },
	// {
	// 	"separate loops",
	// 	[]Domino{{1, 2}, {2, 3}, {3, 1}, {1, 1}, {2, 2}, {3, 3}},
	// 	true,
	// },
	// {
	// 	"nine elements",
	// 	[]Domino{{1, 2}, {5, 3}, {3, 1}, {1, 2}, {2, 4}, {1, 6}, {2, 3}, {3, 4}, {5, 6}},
	// 	true,
	// },
}
<commit_msg>Solve all but last test case<commit_after>package dominoes

// Source: exercism/problem-specifications
// Commit: b5bc74d dominoes: Apply new "input" policy
// Problem Specifications Version: 2.1.0

var testCases = []struct {
	description string
	dominoes    []Domino
	valid       bool // true => can chain, false => cannot chain
}{
	{
		"empty input = empty output",
		[]Domino{},
		true,
	},
	{
		"singleton input = singleton output",
		[]Domino{{1, 1}},
		true,
	},
	{
		"singleton that can't be chained",
		[]Domino{{1, 2}},
		false,
	},
	{
		"three elements",
		[]Domino{{1, 2}, {3, 1}, {2, 3}},
		true,
	},
	{
		"can reverse dominoes",
		[]Domino{{1, 2}, {1, 3}, {2, 3}},
		true,
	},
	{
		"can't be chained",
		[]Domino{{1, 2}, {4, 1}, {2, 3}},
		false,
	},
	{
		"disconnected - simple",
		// This meets the requirement of being possibly-Euclidean.
		// All vertices have even degree.
		// Nevertheless, there is no chain here, as there's no way to get from 1 to 2.
		// This test (and the two following) prevent solutions from using the even-degree test as the sole criterion,
		// as that is not a sufficient condition.
		[]Domino{{1, 1}, {2, 2}},
		false,
	},
	{
		"disconnected - double loop",
		[]Domino{{1, 2}, {2, 1}, {3, 4}, {4, 3}},
		false,
	},
	{
		"disconnected - single isolated",
		[]Domino{{1, 2}, {2, 3}, {3, 1}, {4, 4}},
		false,
	},
	{
		"need backtrack",
		// Some solutions may make a chain out of (1, 2), (2, 3), (3, 1)
		// then decide that since there are no more dominoes containing a 1,
		// there is no chain possible.
		// There is indeed a chain here, so this test checks for this line of reasoning.
		// You need to place the (2, 4) after the (1, 2) rather than the (2, 3).
		[]Domino{{1, 2}, {2, 3}, {3, 1}, {2, 4}, {2, 4}},
		true,
	},
	{
		"separate loops",
		[]Domino{{1, 2}, {2, 3}, {3, 1}, {1, 1}, {2, 2}, {3, 3}},
		true,
	},
	// {
	// 	"nine elements",
	// 	[]Domino{{1, 2}, {5, 3}, {3, 1}, {1, 2}, {2, 4}, {1, 6}, {2, 3}, {3, 4}, {5, 6}},
	// 	true,
	// },
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"runtime"

	"github.com/codegangsta/cli"
	"github.com/exercism/cli/api"
	"github.com/exercism/cli/cmd"
	"github.com/exercism/cli/debug"
)

const (
	// Version is the current release of the command-line app.
	// We try to follow Semantic Versioning (http://semver.org),
	// but with the http://exercism.io app being a prototype, a
	// lot of things get out of hand.
	Version = "2.2.6"

	descConfigure = "Writes config values to a JSON file."
	descDebug     = "Outputs useful debug information."
	descDownload  = "Downloads a solution given the ID of the latest iteration."
	descFetch     = "Fetches the next unsubmitted problem in each track."
	descList      = "Lists the available problems for a language track, given its ID."
	descOpen      = "Opens exercism.io to your most recent iteration of a problem given the track ID and problem slug."
	descRestore   = "Downloads the most recent iteration for each of your solutions on exercism.io."
	descSkip      = "Skips a problem given a track ID and problem slug."
	descStatus    = "Fetches information about your progress with a given language track."
	descSubmit    = "Submits a new iteration to a problem on exercism.io."
	descTracks    = "Lists the available language tracks."
	descUpgrade   = "Upgrades the CLI to the latest released version."

	descLongDownload = "The submission ID is the last part of the URL when looking at a solution on exercism.io."
	descLongRestore  = "Restore will pull the latest revisions of exercises that have already been submitted. It will *not* overwrite existing files. If you have made changes to a file and have not submitted it, and you're trying to restore the last submitted version, first move that file out of the way, then call restore."
)

func main() {
	api.UserAgent = fmt.Sprintf("github.com/exercism/cli v%s (%s/%s)", Version, runtime.GOOS, runtime.GOARCH)

	app := cli.NewApp()
	app.Name = "exercism"
	app.Usage = "A command line tool to interact with http://exercism.io"
	app.Version = Version
	app.Before = func(ctx *cli.Context) error {
		debug.Verbose = ctx.GlobalBool("verbose")
		debug.Println("verbose logging enabled")

		return nil
	}
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "config, c",
			Usage:  "path to config file",
			EnvVar: "EXERCISM_CONFIG_FILE,XDG_CONFIG_HOME",
		},
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "turn on verbose logging",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "configure",
			Usage: descConfigure,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "dir, d",
					Usage: "path to exercises directory",
				},
				cli.StringFlag{
					Name:  "host, u",
					Usage: "exercism api host",
				},
				cli.StringFlag{
					Name:  "key, k",
					Usage: "exercism.io API key (see http://exercism.io/account/key)",
				},
				cli.StringFlag{
					Name:  "api, a",
					Usage: "exercism xapi host",
				},
			},
			Action: cmd.Configure,
		},
		{
			Name:   "debug",
			Usage:  descDebug,
			Action: cmd.Debug,
		},
		{
			Name:        "download",
			ShortName:   "dl",
			Usage:       descDownload,
			Description: descLongDownload,
			Action:      cmd.Download,
		},
		{
			Name:      "fetch",
			ShortName: "f",
			Usage:     descFetch,
			Action:    cmd.Fetch,
		},
		{
			Name:      "list",
			ShortName: "li",
			Usage:     descList,
			Action:    cmd.List,
		},
		{
			Name:      "open",
			ShortName: "op",
			Usage:     descOpen,
			Action:    cmd.Open,
		},
		{
			Name:        "restore",
			ShortName:   "r",
			Usage:       descRestore,
			Description: descLongRestore,
			Action:      cmd.Restore,
		},
		{
			Name:   "skip",
			Usage:  descSkip,
			Action: cmd.Skip,
		},
		{
			Name:      "status",
			ShortName: "st",
			Usage:     descStatus,
			Action:    cmd.Status,
		},
		{
			Name:      "submit",
			ShortName: "s",
			Usage:     descSubmit,
			Action:    cmd.Submit,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "test",
					Usage: "allow submission of test files",
				},
				cli.StringFlag{
					Name:  "comment",
					Usage: "includes a comment with the submission",
				},
			},
		},
		{
			Name:      "tracks",
			ShortName: "t",
			Usage:     descTracks,
			Action:    cmd.Tracks,
		},
		{
			Name:      "unsubmit",
			ShortName: "u",
			Usage:     "REMOVED",
			Action: func(*cli.Context) {
				fmt.Println("For security reasons, this command is no longer in use.\nYou can delete iterations in the web interface.")
			},
		},
		{
			Name:   "upgrade",
			Usage:  descUpgrade,
			Action: cmd.Upgrade,
		},
	}
	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Improve usage/help text of submit command. Fix #330<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"runtime"

	"github.com/codegangsta/cli"
	"github.com/exercism/cli/api"
	"github.com/exercism/cli/cmd"
	"github.com/exercism/cli/debug"
)

const (
	// Version is the current release of the command-line app.
	// We try to follow Semantic Versioning (http://semver.org),
	// but with the http://exercism.io app being a prototype, a
	// lot of things get out of hand.
	Version = "2.2.6"

	descConfigure = "Writes config values to a JSON file."
	descDebug     = "Outputs useful debug information."
	descDownload  = "Downloads a solution given the ID of the latest iteration."
	descFetch     = "Fetches the next unsubmitted problem in each track."
	descList      = "Lists the available problems for a language track, given its ID."
	descOpen      = "Opens exercism.io to your most recent iteration of a problem given the track ID and problem slug."
	descRestore   = "Downloads the most recent iteration for each of your solutions on exercism.io."
	descSkip      = "Skips a problem given a track ID and problem slug."
	descStatus    = "Fetches information about your progress with a given language track."
	descSubmit    = "Submits a new iteration to a problem on exercism.io."
	descTracks    = "Lists the available language tracks."
	descUpgrade   = "Upgrades the CLI to the latest released version."

	descLongDownload = "The submission ID is the last part of the URL when looking at a solution on exercism.io."
	descLongRestore  = "Restore will pull the latest revisions of exercises that have already been submitted. It will *not* overwrite existing files. If you have made changes to a file and have not submitted it, and you're trying to restore the last submitted version, first move that file out of the way, then call restore."
)

func main() {
	api.UserAgent = fmt.Sprintf("github.com/exercism/cli v%s (%s/%s)", Version, runtime.GOOS, runtime.GOARCH)

	app := cli.NewApp()
	app.Name = "exercism"
	app.Usage = "A command line tool to interact with http://exercism.io"
	app.Version = Version
	app.Before = func(ctx *cli.Context) error {
		debug.Verbose = ctx.GlobalBool("verbose")
		debug.Println("verbose logging enabled")

		return nil
	}
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "config, c",
			Usage:  "path to config file",
			EnvVar: "EXERCISM_CONFIG_FILE,XDG_CONFIG_HOME",
		},
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "turn on verbose logging",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "configure",
			Usage: descConfigure,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "dir, d",
					Usage: "path to exercises directory",
				},
				cli.StringFlag{
					Name:  "host, u",
					Usage: "exercism api host",
				},
				cli.StringFlag{
					Name:  "key, k",
					Usage: "exercism.io API key (see http://exercism.io/account/key)",
				},
				cli.StringFlag{
					Name:  "api, a",
					Usage: "exercism xapi host",
				},
			},
			Action: cmd.Configure,
		},
		{
			Name:   "debug",
			Usage:  descDebug,
			Action: cmd.Debug,
		},
		{
			Name:        "download",
			ShortName:   "dl",
			Usage:       descDownload,
			Description: descLongDownload,
			Action:      cmd.Download,
		},
		{
			Name:      "fetch",
			ShortName: "f",
			Usage:     descFetch,
			Action:    cmd.Fetch,
		},
		{
			Name:      "list",
			ShortName: "li",
			Usage:     descList,
			Action:    cmd.List,
		},
		{
			Name:      "open",
			ShortName: "op",
			Usage:     descOpen,
			Action:    cmd.Open,
		},
		{
			Name:        "restore",
			ShortName:   "r",
			Usage:       descRestore,
			Description: descLongRestore,
			Action:      cmd.Restore,
		},
		{
			Name:   "skip",
			Usage:  descSkip,
			Action: cmd.Skip,
		},
		{
			Name:      "status",
			ShortName: "st",
			Usage:     descStatus,
			Action:    cmd.Status,
		},
		{
			Name:      "submit",
			ShortName: "s",
			Usage:     descSubmit,
			Action:    cmd.Submit,
			ArgsUsage: "file1 [file2, etc...]",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "test",
					Usage: "allow submission of test files",
				},
				cli.StringFlag{
					Name:  "comment",
					Usage: "includes a comment with the submission",
				},
			},
		},
		{
			Name:      "tracks",
			ShortName: "t",
			Usage:     descTracks,
			Action:    cmd.Tracks,
		},
		{
			Name:      "unsubmit",
			ShortName: "u",
			Usage:     "REMOVED",
			Action: func(*cli.Context) {
				fmt.Println("For security reasons, this command is no longer in use.\nYou can delete iterations in the web interface.")
			},
		},
		{
			Name:   "upgrade",
			Usage:  descUpgrade,
			Action: cmd.Upgrade,
		},
	}
	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

/*
#cgo CFLAGS: -I/usr/include
#cgo LDFLAGS: -framework IOKit -framework CoreFoundation
#include <stdlib.h>
#include <IOKit/pwr_mgt/IOPMLib.h>

// kIOPMAssertionTypeNoDisplaySleep prevents display sleep,
// kIOPMAssertionTypeNoIdleSleep prevents idle sleep

// reasonForActivity is a descriptive string used by the system whenever it needs
// to tell the user why the system is not sleeping. For example,
// "Mail Compacting Mailboxes" would be a useful string.

long PreventSleep()
{
	// NOTE: IOPMAssertionCreateWithName limits the string to 128 characters.
	CFStringRef reasonForActivity= CFSTR("Encoding");

	IOPMAssertionID assertionID;
	IOReturn success = IOPMAssertionCreateWithName(kIOPMAssertionTypeNoIdleSleep,
	                                    kIOPMAssertionLevelOn, reasonForActivity, &assertionID);
	if (success == kIOReturnSuccess)
	{
		return (int)assertionID;
	}
	return -1;
}

void AllowSleep(long assert)
{
	IOPMAssertionID assertionID = assert;
	IOPMAssertionRelease(assertionID);
}
*/
import "C"

import (
	"errors"
	"fmt"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"regexp"
	"runtime"
	"syscall"
)

func createPidFile(filename string) (*os.File, error) {

	file, err := os.OpenFile(filename,
		os.O_WRONLY|os.O_CREATE|os.O_EXCL|os.O_SYNC,
		syscall.S_IRUSR|syscall.S_IWUSR)

	if err != nil {
		return nil, err
	}

	err = syscall.Flock(int(file.Fd()), syscall.F_WRLCK)
	if err != nil {
		file.Close()
		return nil, err
	}

	_, err = file.WriteString(fmt.Sprintf("%d", os.Getpid()))
	if err != nil {
		file.Close()
		return nil, err
	}

	return file, nil
}

func removePidFile(file *os.File) {

	err := syscall.Flock(int(file.Fd()), syscall.F_UNLCK)
	if err != nil {
		return
	}

	file.Close()
	os.Remove(file.Name())
}

type TranscodeQueue struct {
	Requests chan string
}

func NewTranscodeQueue() *TranscodeQueue {
	return &TranscodeQueue{
		Requests: make(chan string, 100),
	}
}

func (this *TranscodeQueue) DoTranscode(file string, reply *bool) error {
	log.Println("Asked to transcode:", file)
	this.Requests <- file
	*reply = true
	return nil
}

func runCommand(cmd *exec.Cmd, die <-chan os.Signal) (bool, error) {

	isDone := make(chan error)
	if err := cmd.Start(); err != nil {
		return false, err
	}

	go func() {
		if err := cmd.Wait(); err != nil {
			isDone <- err
		}
		isDone <- nil
	}()

	select {
	case err := <-isDone:
		return true, err
	case <-die:
		log.Println("Killing process")
		if err := cmd.Process.Kill(); err != nil {
			return false, err
		}
		return false, nil
	}

}

var wasCanceled error = errors.New("Was Canceled")

func performTranscode(infile, outfile string, isDone chan<- error, sigc <-chan os.Signal) {
	//infile = fmt.Sprintf("%q", infile)
	//outfile = fmt.Sprintf("%q", outfile)

	//transcodeCommand := exec.Command("/bin/sleep", "10")
	transcodeCommand := exec.Command(
		"/usr/local/bin/HandBrakeCLI",
		"-O", "-I",
		"-f", "mp4",
		"--encoder", "x264",
		"--x264-preset", "faster",
		"--x264-tune", "film",
		"--h264-profile", "auto",
		"--h264-level", "auto",
		"--quality", "20",
		"--large-file",
		"--aencoder", "ca_aac,copy:ac3",
		"-B", "160",
		"--mixdown", "dpl2",
		"--aname", "English",
		"--loose-anamorphic",
		"--decomb",
		"--modulus", "2",
		"-i", infile,
		"-o", outfile)

	//transferCommand := exec.Command("/bin/sleep", "5")
	transferCommand := exec.Command(
		"/usr/bin/scp", "-B", "-C", "-q", outfile,
		"guy@mediaserver.local:/srv/Media/Movies/0\\ -\\ Inbox/TV")

	log.Println("Performing transcode of", infile)
	log.Println("Writing file", outfile)

	done, err := runCommand(transcodeCommand, sigc)

	if err != nil {
		log.Println(err.Error())
		isDone <- err
		return
	}
	if !done {
		isDone <- wasCanceled
		return
	}

	log.Println("Done transcoding file")
	log.Println("Transfering file to server")

	done, err = runCommand(transferCommand, sigc)

	if err == nil && done {
		os.Remove(infile)
		//os.Remove(outfile)
		// leave commented out until im sure i understand the scp problem
	} else {
		if err != nil {
			log.Println(err.Error())
			isDone <- err

		} else {
			isDone <- wasCanceled
		}
		return
	}

	log.Println("Done transfering file to server")
	isDone <- nil
}

func requestTranscode(infile, unixSocket string) (reply bool, err error) {

	client, err := rpc.Dial("unix", unixSocket)
	if err != nil {
		return
	}
	err = client.Call("TranscodeQueue.DoTranscode", infile, &reply)
	return
}

func createTranscodeServer(pidfile *os.File, unixSocket string) error {

	sigc := make(chan os.Signal, 10)
	signal.Notify(sigc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	defer signal.Stop(sigc)

	q := NewTranscodeQueue()
	server := rpc.NewServer()
	server.Register(q)

	l, err := net.Listen("unix", unixSocket)
	if err != nil {
		return err
	}

	defer l.Close()

	go server.Accept(l)

	isDone := make(chan error)
	rxp, err := regexp.Compile(`(.*)(\.mpg)$`)

	if err != nil {
		return err
	}

	q.Requests <- os.Args[1]

	for {
		select {
		case req := <-q.Requests:

			outfile := rxp.ReplaceAllString(req, "$1.m4v")
			go performTranscode(req, outfile, isDone, sigc)

			select {
			case e := <-isDone:
				if e != nil {
					return e
				}
			}
		default:
			return nil
		}
	}
}

func main() {

	sleepId := C.PreventSleep()
	defer C.AllowSleep(sleepId)

	if len(os.Args) != 2 {
		return
	}

	runtime.GOMAXPROCS(runtime.NumCPU())

	log.Println("Starting...")

	toTranscode := os.Args[1]
	pidFileName := path.Join(os.TempDir(), "transcodequeue.pid")
	unixSocket := path.Join(os.TempDir(), "transcodequeue.sock")

	pidFile, err := createPidFile(pidFileName)

	if err != nil {
		log.Println("Sending server request to transcode", toTranscode)

		reply, err := requestTranscode(toTranscode, unixSocket)
		if err != nil {
			log.Panicln(err.Error())
		}

		if reply {
			log.Println("Successfully requested transcode")
		}
	} else {
		log.Println("Creating Server")

		defer removePidFile(pidFile)

		err = createTranscodeServer(pidFile, unixSocket)
		if err != nil {
			log.Println(err.Error())
		}
	}

	log.Println("Finished")
}
<commit_msg>Delete input file on sucess<commit_after>package main

/*
#cgo CFLAGS: -I/usr/include
#cgo LDFLAGS: -framework IOKit -framework CoreFoundation
#include <stdlib.h>
#include <IOKit/pwr_mgt/IOPMLib.h>

// kIOPMAssertionTypeNoDisplaySleep prevents display sleep,
// kIOPMAssertionTypeNoIdleSleep prevents idle sleep

// reasonForActivity is a descriptive string used by the system whenever it needs
// to tell the user why the system is not sleeping. For example,
// "Mail Compacting Mailboxes" would be a useful string.

long PreventSleep()
{
	// NOTE: IOPMAssertionCreateWithName limits the string to 128 characters.
	CFStringRef reasonForActivity= CFSTR("Encoding");

	IOPMAssertionID assertionID;
	IOReturn success = IOPMAssertionCreateWithName(kIOPMAssertionTypeNoIdleSleep,
	                                    kIOPMAssertionLevelOn, reasonForActivity, &assertionID);
	if (success == kIOReturnSuccess)
	{
		return (int)assertionID;
	}
	return -1;
}

void AllowSleep(long assert)
{
	IOPMAssertionID assertionID = assert;
	IOPMAssertionRelease(assertionID);
}
*/
import "C"

import (
	"errors"
	"fmt"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"regexp"
	"runtime"
	"syscall"
)

func createPidFile(filename string) (*os.File, error) {

	file, err := os.OpenFile(filename,
		os.O_WRONLY|os.O_CREATE|os.O_EXCL|os.O_SYNC,
		syscall.S_IRUSR|syscall.S_IWUSR)

	if err != nil {
		return nil, err
	}

	err = syscall.Flock(int(file.Fd()), syscall.F_WRLCK)
	if err != nil {
		file.Close()
		return nil, err
	}

	_, err = file.WriteString(fmt.Sprintf("%d", os.Getpid()))
	if err != nil {
		file.Close()
		return nil, err
	}

	return file, nil
}

func removePidFile(file *os.File) {

	err := syscall.Flock(int(file.Fd()), syscall.F_UNLCK)
	if err != nil {
		return
	}

	file.Close()
	os.Remove(file.Name())
}

type TranscodeQueue struct {
	Requests chan string
}

func NewTranscodeQueue() *TranscodeQueue {
	return &TranscodeQueue{
		Requests: make(chan string, 100),
	}
}

func (this *TranscodeQueue) DoTranscode(file string, reply *bool) error {
	log.Println("Asked to transcode:", file)
	this.Requests <- file
	*reply = true
	return nil
}

func runCommand(cmd *exec.Cmd, die <-chan os.Signal) (bool, error) {

	isDone := make(chan error)
	if err := cmd.Start(); err != nil {
		return false, err
	}

	go func() {
		if err := cmd.Wait(); err != nil {
			isDone <- err
		}
		isDone <- nil
	}()

	select {
	case err := <-isDone:
		return true, err
	case <-die:
		log.Println("Killing process")
		if err := cmd.Process.Kill(); err != nil {
			return false, err
		}
		return false, nil
	}

}

var wasCanceled error = errors.New("Was Canceled")

func performTranscode(infile, outfile string, isDone chan<- error, sigc <-chan os.Signal) {
	//infile = fmt.Sprintf("%q", infile)
	//outfile = fmt.Sprintf("%q", outfile)

	//transcodeCommand := exec.Command("/bin/sleep", "10")
	transcodeCommand := exec.Command(
		"/usr/local/bin/HandBrakeCLI",
		"-O", "-I",
		"-f", "mp4",
		"--encoder", "x264",
		"--x264-preset", "faster",
		"--x264-tune", "film",
		"--h264-profile", "auto",
		"--h264-level", "auto",
		"--quality", "20",
		"--large-file",
		"--aencoder", "ca_aac,copy:ac3",
		"-B", "160",
		"--mixdown", "dpl2",
		"--aname", "English",
		"--loose-anamorphic",
		"--decomb",
		"--modulus", "2",
		"-i", infile,
		"-o", outfile)

	//transferCommand := exec.Command("/bin/sleep", "5")
	transferCommand := exec.Command(
		"/usr/bin/scp", "-B", "-C", "-q", outfile,
		"guy@mediaserver.local:/srv/Media/Movies/0\\ -\\ Inbox/TV")

	log.Println("Performing transcode of", infile)
	log.Println("Writing file", outfile)

	done, err := runCommand(transcodeCommand, sigc)

	if err != nil {
		log.Println(err.Error())
		isDone <- err
		return
	}
	if !done {
		isDone <- wasCanceled
		return
	}

	log.Println("Done transcoding file")
	log.Println("Transfering file to server")

	done, err = runCommand(transferCommand, sigc)

	if err == nil && done {
		os.Remove(infile)
		os.Remove(outfile)
		// leave commented out until im sure i understand the scp problem
	} else {
		if err != nil {
			log.Println(err.Error())
			isDone <- err

		} else {
			isDone <- wasCanceled
		}
		return
	}

	log.Println("Done transfering file to server")
	isDone <- nil
}

func requestTranscode(infile, unixSocket string) (reply bool, err error) {

	client, err := rpc.Dial("unix", unixSocket)
	if err != nil {
		return
	}
	err = client.Call("TranscodeQueue.DoTranscode", infile, &reply)
	return
}

func createTranscodeServer(pidfile *os.File, unixSocket string) error {

	sigc := make(chan os.Signal, 10)
	signal.Notify(sigc,
		syscall.SIGHUP,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)
	defer signal.Stop(sigc)

	q := NewTranscodeQueue()
	server := rpc.NewServer()
	server.Register(q)

	l, err := net.Listen("unix", unixSocket)
	if err != nil {
		return err
	}

	defer l.Close()

	go server.Accept(l)

	isDone := make(chan error)
	rxp, err := regexp.Compile(`(.*)(\.mpg)$`)

	if err != nil {
		return err
	}

	q.Requests <- os.Args[1]

	for {
		select {
		case req := <-q.Requests:

			outfile := rxp.ReplaceAllString(req, "$1.m4v")
			go performTranscode(req, outfile, isDone, sigc)

			select {
			case e := <-isDone:
				if e != nil {
					return e
				}
			}
		default:
			return nil
		}
	}
}

func main() {

	sleepId := C.PreventSleep()
	defer C.AllowSleep(sleepId)

	if len(os.Args) != 2 {
		return
	}

	runtime.GOMAXPROCS(runtime.NumCPU())

	log.Println("Starting...")

	toTranscode := os.Args[1]
	pidFileName := path.Join(os.TempDir(), "transcodequeue.pid")
	unixSocket := path.Join(os.TempDir(), "transcodequeue.sock")

	pidFile, err := createPidFile(pidFileName)

	if err != nil {
		log.Println("Sending server request to transcode", toTranscode)

		reply, err := requestTranscode(toTranscode, unixSocket)
		if err != nil {
			log.Panicln(err.Error())
		}

		if reply {
			log.Println("Successfully requested transcode")
		}
	} else {
		log.Println("Creating Server")

		defer removePidFile(pidFile)

		err = createTranscodeServer(pidFile, unixSocket)
		if err != nil {
			log.Println(err.Error())
		}
	}

	log.Println("Finished")
}
<|endoftext|>


<commit_before>package exec

import (
	"bytes"
	"context"
	"reflect"
	"strconv"
	"sync"

	"github.com/neelance/graphql-go/errors"
	"github.com/neelance/graphql-go/internal/common"
	"github.com/neelance/graphql-go/internal/exec/resolvable"
	"github.com/neelance/graphql-go/internal/exec/selected"
	"github.com/neelance/graphql-go/internal/query"
	"github.com/neelance/graphql-go/internal/schema"
	"github.com/neelance/graphql-go/log"
	"github.com/neelance/graphql-go/trace"
)

type Request struct {
	selected.Request
	Limiter chan struct{}
	Tracer  trace.Tracer
	Logger  log.Logger
}

type fieldResult struct {
	name  string
	value []byte
}

func (r *Request) handlePanic(ctx context.Context) {
	if value := recover(); value != nil {
		r.Logger.LogPanic(ctx, value)
		r.AddError(makePanicError(value))
	}
}

func makePanicError(value interface{}) *errors.QueryError {
	return errors.Errorf("graphql: panic occurred: %v", value)
}

func (r *Request) Execute(ctx context.Context, s *resolvable.Schema, op *query.Operation) ([]byte, []*errors.QueryError) {
	var out bytes.Buffer
	func() {
		defer r.handlePanic(ctx)
		sels := selected.ApplyOperation(&r.Request, s, op)
		r.execSelections(ctx, sels, s.Resolver, &out, op.Type == query.Mutation)
	}()

	if err := ctx.Err(); err != nil {
		return nil, []*errors.QueryError{errors.Errorf("%s", err)}
	}

	return out.Bytes(), r.Errs
}

type fieldWithResolver struct {
	field    *selected.SchemaField
	resolver reflect.Value
	out      bytes.Buffer
}

func (r *Request) execSelections(ctx context.Context, sels []selected.Selection, resolver reflect.Value, out *bytes.Buffer, serially bool) {
	async := !serially && selected.HasAsyncSel(sels)

	var fields []*fieldWithResolver
	collectFieldsToResolve(sels, resolver, &fields)

	if async {
		var wg sync.WaitGroup
		wg.Add(len(fields))
		for _, f := range fields {
			go func(f *fieldWithResolver) {
				defer r.handlePanic(ctx)
				r.execFieldSelection(ctx, f.field, f.resolver, &f.out, false)
				wg.Done()
			}(f)
		}
		wg.Wait()
	}

	out.WriteByte('{')
	for i, f := range fields {
		if i > 0 {
			out.WriteByte(',')
		}
		out.WriteByte('"')
		out.WriteString(f.field.Alias)
		out.WriteByte('"')
		out.WriteByte(':')
		if async {
			out.Write(f.out.Bytes())
			continue
		}
		r.execFieldSelection(ctx, f.field, f.resolver, out, false)
	}
	out.WriteByte('}')
}

func collectFieldsToResolve(sels []selected.Selection, resolver reflect.Value, fields *[]*fieldWithResolver) {
	for _, sel := range sels {
		switch sel := sel.(type) {
		case *selected.SchemaField:
			*fields = append(*fields, &fieldWithResolver{field: sel, resolver: resolver})

		case *selected.TypenameField:
			sf := &selected.SchemaField{
				Field:       resolvable.MetaFieldTypename,
				Alias:       sel.Alias,
				FixedResult: reflect.ValueOf(typeOf(sel, resolver)),
			}
			*fields = append(*fields, &fieldWithResolver{field: sf, resolver: resolver})

		case *selected.TypeAssertion:
			out := resolver.Method(sel.MethodIndex).Call(nil)
			if !out[1].Bool() {
				continue
			}
			collectFieldsToResolve(sel.Sels, out[0], fields)

		default:
			panic("unreachable")
		}
	}
}

func typeOf(tf *selected.TypenameField, resolver reflect.Value) string {
	if len(tf.TypeAssertions) == 0 {
		return tf.Name
	}
	for name, a := range tf.TypeAssertions {
		out := resolver.Method(a.MethodIndex).Call(nil)
		if out[1].Bool() {
			return name
		}
	}
	return ""
}

func (r *Request) execFieldSelection(ctx context.Context, field *selected.SchemaField, resolver reflect.Value, out *bytes.Buffer, applyLimiter bool) {
	if applyLimiter {
		r.Limiter <- struct{}{}
	}

	var result reflect.Value
	var err *errors.QueryError

	traceCtx, finish := r.Tracer.TraceField(ctx, field.TraceLabel, field.TypeName, field.Name, !field.Async, field.Args)
	defer func() {
		finish(err)
	}()

	err = func() (err *errors.QueryError) {
		defer func() {
			if panicValue := recover(); panicValue != nil {
				r.Logger.LogPanic(ctx, panicValue)
				err = makePanicError(panicValue)
			}
		}()

		if field.FixedResult.IsValid() {
			result = field.FixedResult
			return nil
		}

		if err := traceCtx.Err(); err != nil {
			return errors.Errorf("%s", err) // don't execute any more resolvers if context got cancelled
		}

		var in []reflect.Value
		if field.HasContext {
			in = append(in, reflect.ValueOf(traceCtx))
		}
		if field.ArgsPacker != nil {
			in = append(in, field.PackedArgs)
		}
		callOut := resolver.Method(field.MethodIndex).Call(in)
		result = callOut[0]
		if field.HasError && !callOut[1].IsNil() {
			resolverErr := callOut[1].Interface().(error)
			err := errors.Errorf("%s", resolverErr)
			err.ResolverError = resolverErr
			return err
		}
		return nil
	}()

	if applyLimiter {
		<-r.Limiter
	}

	if err != nil {
		r.AddError(err)
		out.WriteString("null") // TODO handle non-nil
		return
	}

	r.execSelectionSet(traceCtx, field.Sels, field.Type, result, out)
}

func (r *Request) execSelectionSet(ctx context.Context, sels []selected.Selection, typ common.Type, resolver reflect.Value, out *bytes.Buffer) {
	t, nonNull := unwrapNonNull(typ)
	switch t := t.(type) {
	case *schema.Object, *schema.Interface, *schema.Union:
		if resolver.IsNil() {
			if nonNull {
				panic(errors.Errorf("got nil for non-null %q", t))
			}
			out.WriteString("null")
			return
		}

		r.execSelections(ctx, sels, resolver, out, false)
		return
	}

	if !nonNull {
		if resolver.IsNil() {
			out.WriteString("null")
			return
		}
		resolver = resolver.Elem()
	}

	switch t := t.(type) {
	case *common.List:
		l := resolver.Len()

		if selected.HasAsyncSel(sels) {
			var wg sync.WaitGroup
			wg.Add(l)
			entryouts := make([]bytes.Buffer, l)
			for i := 0; i < l; i++ {
				go func(i int) {
					defer r.handlePanic(ctx)
					r.execSelectionSet(ctx, sels, t.OfType, resolver.Index(i), &entryouts[i])
					wg.Done()
				}(i)
			}
			wg.Wait()

			out.WriteByte('[')
			for i, entryout := range entryouts {
				if i > 0 {
					out.WriteByte(',')
				}
				out.Write(entryout.Bytes())
			}
			out.WriteByte(']')
			return
		}

		out.WriteByte('[')
		for i := 0; i < l; i++ {
			if i > 0 {
				out.WriteByte(',')
			}
			r.execSelectionSet(ctx, sels, t.OfType, resolver.Index(i), out)
		}
		out.WriteByte(']')

	case *schema.Scalar:
		var b []byte // TODO use scratch
		switch t.Name {
		case "Int":
			out.Write(strconv.AppendInt(b, resolver.Int(), 10))
		case "Float":
			out.Write(strconv.AppendFloat(b, resolver.Float(), 'f', -1, 64))
		case "String":
			out.Write(strconv.AppendQuote(b, resolver.String()))
		case "Boolean":
			out.Write(strconv.AppendBool(b, resolver.Bool()))
		default:
			v := resolver.Interface().(marshaler)
			data, err := v.MarshalJSON()
			if err != nil {
				panic(errors.Errorf("could not marshal %v", v))
			}
			out.Write(data)
		}

	case *schema.Enum:
		out.WriteByte('"')
		out.WriteString(resolver.String())
		out.WriteByte('"')

	default:
		panic("unreachable")
	}
}

func unwrapNonNull(t common.Type) (common.Type, bool) {
	if nn, ok := t.(*common.NonNull); ok {
		return nn.OfType, true
	}
	return t, false
}

type marshaler interface {
	MarshalJSON() ([]byte, error)
}
<commit_msg>use encoding/json to encode scalars<commit_after>package exec

import (
	"bytes"
	"context"
	"encoding/json"
	"reflect"
	"sync"

	"github.com/neelance/graphql-go/errors"
	"github.com/neelance/graphql-go/internal/common"
	"github.com/neelance/graphql-go/internal/exec/resolvable"
	"github.com/neelance/graphql-go/internal/exec/selected"
	"github.com/neelance/graphql-go/internal/query"
	"github.com/neelance/graphql-go/internal/schema"
	"github.com/neelance/graphql-go/log"
	"github.com/neelance/graphql-go/trace"
)

type Request struct {
	selected.Request
	Limiter chan struct{}
	Tracer  trace.Tracer
	Logger  log.Logger
}

type fieldResult struct {
	name  string
	value []byte
}

func (r *Request) handlePanic(ctx context.Context) {
	if value := recover(); value != nil {
		r.Logger.LogPanic(ctx, value)
		r.AddError(makePanicError(value))
	}
}

func makePanicError(value interface{}) *errors.QueryError {
	return errors.Errorf("graphql: panic occurred: %v", value)
}

func (r *Request) Execute(ctx context.Context, s *resolvable.Schema, op *query.Operation) ([]byte, []*errors.QueryError) {
	var out bytes.Buffer
	func() {
		defer r.handlePanic(ctx)
		sels := selected.ApplyOperation(&r.Request, s, op)
		r.execSelections(ctx, sels, s.Resolver, &out, op.Type == query.Mutation)
	}()

	if err := ctx.Err(); err != nil {
		return nil, []*errors.QueryError{errors.Errorf("%s", err)}
	}

	return out.Bytes(), r.Errs
}

type fieldWithResolver struct {
	field    *selected.SchemaField
	resolver reflect.Value
	out      bytes.Buffer
}

func (r *Request) execSelections(ctx context.Context, sels []selected.Selection, resolver reflect.Value, out *bytes.Buffer, serially bool) {
	async := !serially && selected.HasAsyncSel(sels)

	var fields []*fieldWithResolver
	collectFieldsToResolve(sels, resolver, &fields)

	if async {
		var wg sync.WaitGroup
		wg.Add(len(fields))
		for _, f := range fields {
			go func(f *fieldWithResolver) {
				defer r.handlePanic(ctx)
				r.execFieldSelection(ctx, f.field, f.resolver, &f.out, false)
				wg.Done()
			}(f)
		}
		wg.Wait()
	}

	out.WriteByte('{')
	for i, f := range fields {
		if i > 0 {
			out.WriteByte(',')
		}
		out.WriteByte('"')
		out.WriteString(f.field.Alias)
		out.WriteByte('"')
		out.WriteByte(':')
		if async {
			out.Write(f.out.Bytes())
			continue
		}
		r.execFieldSelection(ctx, f.field, f.resolver, out, false)
	}
	out.WriteByte('}')
}

func collectFieldsToResolve(sels []selected.Selection, resolver reflect.Value, fields *[]*fieldWithResolver) {
	for _, sel := range sels {
		switch sel := sel.(type) {
		case *selected.SchemaField:
			*fields = append(*fields, &fieldWithResolver{field: sel, resolver: resolver})

		case *selected.TypenameField:
			sf := &selected.SchemaField{
				Field:       resolvable.MetaFieldTypename,
				Alias:       sel.Alias,
				FixedResult: reflect.ValueOf(typeOf(sel, resolver)),
			}
			*fields = append(*fields, &fieldWithResolver{field: sf, resolver: resolver})

		case *selected.TypeAssertion:
			out := resolver.Method(sel.MethodIndex).Call(nil)
			if !out[1].Bool() {
				continue
			}
			collectFieldsToResolve(sel.Sels, out[0], fields)

		default:
			panic("unreachable")
		}
	}
}

func typeOf(tf *selected.TypenameField, resolver reflect.Value) string {
	if len(tf.TypeAssertions) == 0 {
		return tf.Name
	}
	for name, a := range tf.TypeAssertions {
		out := resolver.Method(a.MethodIndex).Call(nil)
		if out[1].Bool() {
			return name
		}
	}
	return ""
}

func (r *Request) execFieldSelection(ctx context.Context, field *selected.SchemaField, resolver reflect.Value, out *bytes.Buffer, applyLimiter bool) {
	if applyLimiter {
		r.Limiter <- struct{}{}
	}

	var result reflect.Value
	var err *errors.QueryError

	traceCtx, finish := r.Tracer.TraceField(ctx, field.TraceLabel, field.TypeName, field.Name, !field.Async, field.Args)
	defer func() {
		finish(err)
	}()

	err = func() (err *errors.QueryError) {
		defer func() {
			if panicValue := recover(); panicValue != nil {
				r.Logger.LogPanic(ctx, panicValue)
				err = makePanicError(panicValue)
			}
		}()

		if field.FixedResult.IsValid() {
			result = field.FixedResult
			return nil
		}

		if err := traceCtx.Err(); err != nil {
			return errors.Errorf("%s", err) // don't execute any more resolvers if context got cancelled
		}

		var in []reflect.Value
		if field.HasContext {
			in = append(in, reflect.ValueOf(traceCtx))
		}
		if field.ArgsPacker != nil {
			in = append(in, field.PackedArgs)
		}
		callOut := resolver.Method(field.MethodIndex).Call(in)
		result = callOut[0]
		if field.HasError && !callOut[1].IsNil() {
			resolverErr := callOut[1].Interface().(error)
			err := errors.Errorf("%s", resolverErr)
			err.ResolverError = resolverErr
			return err
		}
		return nil
	}()

	if applyLimiter {
		<-r.Limiter
	}

	if err != nil {
		r.AddError(err)
		out.WriteString("null") // TODO handle non-nil
		return
	}

	r.execSelectionSet(traceCtx, field.Sels, field.Type, result, out)
}

func (r *Request) execSelectionSet(ctx context.Context, sels []selected.Selection, typ common.Type, resolver reflect.Value, out *bytes.Buffer) {
	t, nonNull := unwrapNonNull(typ)
	switch t := t.(type) {
	case *schema.Object, *schema.Interface, *schema.Union:
		if resolver.IsNil() {
			if nonNull {
				panic(errors.Errorf("got nil for non-null %q", t))
			}
			out.WriteString("null")
			return
		}

		r.execSelections(ctx, sels, resolver, out, false)
		return
	}

	if !nonNull {
		if resolver.IsNil() {
			out.WriteString("null")
			return
		}
		resolver = resolver.Elem()
	}

	switch t := t.(type) {
	case *common.List:
		l := resolver.Len()

		if selected.HasAsyncSel(sels) {
			var wg sync.WaitGroup
			wg.Add(l)
			entryouts := make([]bytes.Buffer, l)
			for i := 0; i < l; i++ {
				go func(i int) {
					defer r.handlePanic(ctx)
					r.execSelectionSet(ctx, sels, t.OfType, resolver.Index(i), &entryouts[i])
					wg.Done()
				}(i)
			}
			wg.Wait()

			out.WriteByte('[')
			for i, entryout := range entryouts {
				if i > 0 {
					out.WriteByte(',')
				}
				out.Write(entryout.Bytes())
			}
			out.WriteByte(']')
			return
		}

		out.WriteByte('[')
		for i := 0; i < l; i++ {
			if i > 0 {
				out.WriteByte(',')
			}
			r.execSelectionSet(ctx, sels, t.OfType, resolver.Index(i), out)
		}
		out.WriteByte(']')

	case *schema.Scalar:
		v := resolver.Interface()
		data, err := json.Marshal(v)
		if err != nil {
			panic(errors.Errorf("could not marshal %v", v))
		}
		out.Write(data)

	case *schema.Enum:
		out.WriteByte('"')
		out.WriteString(resolver.String())
		out.WriteByte('"')

	default:
		panic("unreachable")
	}
}

func unwrapNonNull(t common.Type) (common.Type, bool) {
	if nn, ok := t.(*common.NonNull); ok {
		return nn.OfType, true
	}
	return t, false
}

type marshaler interface {
	MarshalJSON() ([]byte, error)
}
<|endoftext|>


<commit_before>package exec

import (
	"bytes"
	"context"
	"encoding/json"
	"reflect"
	"sync"

	"github.com/graph-gophers/graphql-go/errors"
	"github.com/graph-gophers/graphql-go/internal/common"
	"github.com/graph-gophers/graphql-go/internal/exec/resolvable"
	"github.com/graph-gophers/graphql-go/internal/exec/selected"
	"github.com/graph-gophers/graphql-go/internal/query"
	"github.com/graph-gophers/graphql-go/internal/schema"
	"github.com/graph-gophers/graphql-go/log"
	"github.com/graph-gophers/graphql-go/trace"
)

type Request struct {
	selected.Request
	Limiter chan struct{}
	Tracer  trace.Tracer
	Logger  log.Logger
}

func (r *Request) handlePanic(ctx context.Context) {
	if value := recover(); value != nil {
		r.Logger.LogPanic(ctx, value)
		r.AddError(makePanicError(value))
	}
}

func makePanicError(value interface{}) *errors.QueryError {
	return errors.Errorf("graphql: panic occurred: %v", value)
}

func (r *Request) Execute(ctx context.Context, s *resolvable.Schema, op *query.Operation) ([]byte, []*errors.QueryError) {
	var out bytes.Buffer
	func() {
		defer r.handlePanic(ctx)
		sels := selected.ApplyOperation(&r.Request, s, op)
		r.execSelections(ctx, sels, nil, s.Resolver, &out, op.Type == query.Mutation)
	}()

	if err := ctx.Err(); err != nil {
		return nil, []*errors.QueryError{errors.Errorf("%s", err)}
	}

	return out.Bytes(), r.Errs
}

type fieldToExec struct {
	field    *selected.SchemaField
	sels     []selected.Selection
	resolver reflect.Value
	out      *bytes.Buffer
}

func (r *Request) execSelections(ctx context.Context, sels []selected.Selection, path *pathSegment, resolver reflect.Value, out *bytes.Buffer, serially bool) {
	async := !serially && selected.HasAsyncSel(sels)

	var fields []*fieldToExec
	collectFieldsToResolve(sels, resolver, &fields, make(map[string]*fieldToExec))

	if async {
		var wg sync.WaitGroup
		wg.Add(len(fields))
		for _, f := range fields {
			go func(f *fieldToExec) {
				defer wg.Done()
				defer r.handlePanic(ctx)
				f.out = new(bytes.Buffer)
				execFieldSelection(ctx, r, f, &pathSegment{path, f.field.Alias}, true)
			}(f)
		}
		wg.Wait()
	}

	out.WriteByte('{')
	for i, f := range fields {
		if i > 0 {
			out.WriteByte(',')
		}
		out.WriteByte('"')
		out.WriteString(f.field.Alias)
		out.WriteByte('"')
		out.WriteByte(':')
		if async {
			out.Write(f.out.Bytes())
			continue
		}
		f.out = out
		execFieldSelection(ctx, r, f, &pathSegment{path, f.field.Alias}, false)
	}
	out.WriteByte('}')
}

func collectFieldsToResolve(sels []selected.Selection, resolver reflect.Value, fields *[]*fieldToExec, fieldByAlias map[string]*fieldToExec) {
	for _, sel := range sels {
		switch sel := sel.(type) {
		case *selected.SchemaField:
			field, ok := fieldByAlias[sel.Alias]
			if !ok { // validation already checked for conflict (TODO)
				field = &fieldToExec{field: sel, resolver: resolver}
				fieldByAlias[sel.Alias] = field
				*fields = append(*fields, field)
			}
			field.sels = append(field.sels, sel.Sels...)

		case *selected.TypenameField:
			sf := &selected.SchemaField{
				Field:       resolvable.MetaFieldTypename,
				Alias:       sel.Alias,
				FixedResult: reflect.ValueOf(typeOf(sel, resolver)),
			}
			*fields = append(*fields, &fieldToExec{field: sf, resolver: resolver})

		case *selected.TypeAssertion:
			out := resolver.Method(sel.MethodIndex).Call(nil)
			if !out[1].Bool() {
				continue
			}
			collectFieldsToResolve(sel.Sels, out[0], fields, fieldByAlias)

		default:
			panic("unreachable")
		}
	}
}

func typeOf(tf *selected.TypenameField, resolver reflect.Value) string {
	if len(tf.TypeAssertions) == 0 {
		return tf.Name
	}
	for name, a := range tf.TypeAssertions {
		out := resolver.Method(a.MethodIndex).Call(nil)
		if out[1].Bool() {
			return name
		}
	}
	return ""
}

func execFieldSelection(ctx context.Context, r *Request, f *fieldToExec, path *pathSegment, applyLimiter bool) {
	if applyLimiter {
		r.Limiter <- struct{}{}
	}

	var result reflect.Value
	var err *errors.QueryError

	traceCtx, finish := r.Tracer.TraceField(ctx, f.field.TraceLabel, f.field.TypeName, f.field.Name, !f.field.Async, f.field.Args)
	defer func() {
		finish(err)
	}()

	err = func() (err *errors.QueryError) {
		defer func() {
			if panicValue := recover(); panicValue != nil {
				r.Logger.LogPanic(ctx, panicValue)
				err = makePanicError(panicValue)
				err.Path = path.toSlice()
			}
		}()

		if f.field.FixedResult.IsValid() {
			result = f.field.FixedResult
			return nil
		}

		if err := traceCtx.Err(); err != nil {
			return errors.Errorf("%s", err) // don't execute any more resolvers if context got cancelled
		}

		var in []reflect.Value
		if f.field.HasContext {
			in = append(in, reflect.ValueOf(traceCtx))
		}
		if f.field.ArgsPacker != nil {
			in = append(in, f.field.PackedArgs)
		}
		callOut := f.resolver.Method(f.field.MethodIndex).Call(in)
		result = callOut[0]
		if f.field.HasError && !callOut[1].IsNil() {
			resolverErr := callOut[1].Interface().(error)
			err := errors.Errorf("%s", resolverErr)
			err.Path = path.toSlice()
			err.ResolverError = resolverErr
			return err
		}
		return nil
	}()

	if applyLimiter {
		<-r.Limiter
	}

	if err != nil {
		r.AddError(err)
		f.out.WriteString("null") // TODO handle non-nil
		return
	}

	r.execSelectionSet(traceCtx, f.sels, f.field.Type, path, result, f.out)
}

func (r *Request) execSelectionSet(ctx context.Context, sels []selected.Selection, typ common.Type, path *pathSegment, resolver reflect.Value, out *bytes.Buffer) {
	t, nonNull := unwrapNonNull(typ)
	switch t := t.(type) {
	case *schema.Object, *schema.Interface, *schema.Union:
		if resolver.Kind() == reflect.Ptr && resolver.IsNil() {
			if nonNull {
				panic(errors.Errorf("got nil for non-null %q", t))
			}
			out.WriteString("null")
			return
		}

		r.execSelections(ctx, sels, path, resolver, out, false)
		return
	}

	if !nonNull {
		if resolver.IsNil() {
			out.WriteString("null")
			return
		}
		resolver = resolver.Elem()
	}

	switch t := t.(type) {
	case *common.List:
		l := resolver.Len()

		if selected.HasAsyncSel(sels) {
			var wg sync.WaitGroup
			wg.Add(l)
			entryouts := make([]bytes.Buffer, l)
			for i := 0; i < l; i++ {
				go func(i int) {
					defer wg.Done()
					defer r.handlePanic(ctx)
					r.execSelectionSet(ctx, sels, t.OfType, &pathSegment{path, i}, resolver.Index(i), &entryouts[i])
				}(i)
			}
			wg.Wait()

			out.WriteByte('[')
			for i, entryout := range entryouts {
				if i > 0 {
					out.WriteByte(',')
				}
				out.Write(entryout.Bytes())
			}
			out.WriteByte(']')
			return
		}

		out.WriteByte('[')
		for i := 0; i < l; i++ {
			if i > 0 {
				out.WriteByte(',')
			}
			r.execSelectionSet(ctx, sels, t.OfType, &pathSegment{path, i}, resolver.Index(i), out)
		}
		out.WriteByte(']')

	case *schema.Scalar:
		v := resolver.Interface()
		data, err := json.Marshal(v)
		if err != nil {
			panic(errors.Errorf("could not marshal %v", v))
		}
		out.Write(data)

	case *schema.Enum:
		out.WriteByte('"')
		out.WriteString(resolver.String())
		out.WriteByte('"')

	default:
		panic("unreachable")
	}
}

func unwrapNonNull(t common.Type) (common.Type, bool) {
	if nn, ok := t.(*common.NonNull); ok {
		return nn.OfType, true
	}
	return t, false
}

type pathSegment struct {
	parent *pathSegment
	value  interface{}
}

func (p *pathSegment) toSlice() []interface{} {
	if p == nil {
		return nil
	}
	return append(p.parent.toSlice(), p.value)
}
<commit_msg>Pass back JSON marshaling error<commit_after>package exec

import (
	"bytes"
	"context"
	"encoding/json"
	"reflect"
	"sync"

	"github.com/graph-gophers/graphql-go/errors"
	"github.com/graph-gophers/graphql-go/internal/common"
	"github.com/graph-gophers/graphql-go/internal/exec/resolvable"
	"github.com/graph-gophers/graphql-go/internal/exec/selected"
	"github.com/graph-gophers/graphql-go/internal/query"
	"github.com/graph-gophers/graphql-go/internal/schema"
	"github.com/graph-gophers/graphql-go/log"
	"github.com/graph-gophers/graphql-go/trace"
)

type Request struct {
	selected.Request
	Limiter chan struct{}
	Tracer  trace.Tracer
	Logger  log.Logger
}

func (r *Request) handlePanic(ctx context.Context) {
	if value := recover(); value != nil {
		r.Logger.LogPanic(ctx, value)
		r.AddError(makePanicError(value))
	}
}

func makePanicError(value interface{}) *errors.QueryError {
	return errors.Errorf("graphql: panic occurred: %v", value)
}

func (r *Request) Execute(ctx context.Context, s *resolvable.Schema, op *query.Operation) ([]byte, []*errors.QueryError) {
	var out bytes.Buffer
	func() {
		defer r.handlePanic(ctx)
		sels := selected.ApplyOperation(&r.Request, s, op)
		r.execSelections(ctx, sels, nil, s.Resolver, &out, op.Type == query.Mutation)
	}()

	if err := ctx.Err(); err != nil {
		return nil, []*errors.QueryError{errors.Errorf("%s", err)}
	}

	return out.Bytes(), r.Errs
}

type fieldToExec struct {
	field    *selected.SchemaField
	sels     []selected.Selection
	resolver reflect.Value
	out      *bytes.Buffer
}

func (r *Request) execSelections(ctx context.Context, sels []selected.Selection, path *pathSegment, resolver reflect.Value, out *bytes.Buffer, serially bool) {
	async := !serially && selected.HasAsyncSel(sels)

	var fields []*fieldToExec
	collectFieldsToResolve(sels, resolver, &fields, make(map[string]*fieldToExec))

	if async {
		var wg sync.WaitGroup
		wg.Add(len(fields))
		for _, f := range fields {
			go func(f *fieldToExec) {
				defer wg.Done()
				defer r.handlePanic(ctx)
				f.out = new(bytes.Buffer)
				execFieldSelection(ctx, r, f, &pathSegment{path, f.field.Alias}, true)
			}(f)
		}
		wg.Wait()
	}

	out.WriteByte('{')
	for i, f := range fields {
		if i > 0 {
			out.WriteByte(',')
		}
		out.WriteByte('"')
		out.WriteString(f.field.Alias)
		out.WriteByte('"')
		out.WriteByte(':')
		if async {
			out.Write(f.out.Bytes())
			continue
		}
		f.out = out
		execFieldSelection(ctx, r, f, &pathSegment{path, f.field.Alias}, false)
	}
	out.WriteByte('}')
}

func collectFieldsToResolve(sels []selected.Selection, resolver reflect.Value, fields *[]*fieldToExec, fieldByAlias map[string]*fieldToExec) {
	for _, sel := range sels {
		switch sel := sel.(type) {
		case *selected.SchemaField:
			field, ok := fieldByAlias[sel.Alias]
			if !ok { // validation already checked for conflict (TODO)
				field = &fieldToExec{field: sel, resolver: resolver}
				fieldByAlias[sel.Alias] = field
				*fields = append(*fields, field)
			}
			field.sels = append(field.sels, sel.Sels...)

		case *selected.TypenameField:
			sf := &selected.SchemaField{
				Field:       resolvable.MetaFieldTypename,
				Alias:       sel.Alias,
				FixedResult: reflect.ValueOf(typeOf(sel, resolver)),
			}
			*fields = append(*fields, &fieldToExec{field: sf, resolver: resolver})

		case *selected.TypeAssertion:
			out := resolver.Method(sel.MethodIndex).Call(nil)
			if !out[1].Bool() {
				continue
			}
			collectFieldsToResolve(sel.Sels, out[0], fields, fieldByAlias)

		default:
			panic("unreachable")
		}
	}
}

func typeOf(tf *selected.TypenameField, resolver reflect.Value) string {
	if len(tf.TypeAssertions) == 0 {
		return tf.Name
	}
	for name, a := range tf.TypeAssertions {
		out := resolver.Method(a.MethodIndex).Call(nil)
		if out[1].Bool() {
			return name
		}
	}
	return ""
}

func execFieldSelection(ctx context.Context, r *Request, f *fieldToExec, path *pathSegment, applyLimiter bool) {
	if applyLimiter {
		r.Limiter <- struct{}{}
	}

	var result reflect.Value
	var err *errors.QueryError

	traceCtx, finish := r.Tracer.TraceField(ctx, f.field.TraceLabel, f.field.TypeName, f.field.Name, !f.field.Async, f.field.Args)
	defer func() {
		finish(err)
	}()

	err = func() (err *errors.QueryError) {
		defer func() {
			if panicValue := recover(); panicValue != nil {
				r.Logger.LogPanic(ctx, panicValue)
				err = makePanicError(panicValue)
				err.Path = path.toSlice()
			}
		}()

		if f.field.FixedResult.IsValid() {
			result = f.field.FixedResult
			return nil
		}

		if err := traceCtx.Err(); err != nil {
			return errors.Errorf("%s", err) // don't execute any more resolvers if context got cancelled
		}

		var in []reflect.Value
		if f.field.HasContext {
			in = append(in, reflect.ValueOf(traceCtx))
		}
		if f.field.ArgsPacker != nil {
			in = append(in, f.field.PackedArgs)
		}
		callOut := f.resolver.Method(f.field.MethodIndex).Call(in)
		result = callOut[0]
		if f.field.HasError && !callOut[1].IsNil() {
			resolverErr := callOut[1].Interface().(error)
			err := errors.Errorf("%s", resolverErr)
			err.Path = path.toSlice()
			err.ResolverError = resolverErr
			return err
		}
		return nil
	}()

	if applyLimiter {
		<-r.Limiter
	}

	if err != nil {
		r.AddError(err)
		f.out.WriteString("null") // TODO handle non-nil
		return
	}

	r.execSelectionSet(traceCtx, f.sels, f.field.Type, path, result, f.out)
}

func (r *Request) execSelectionSet(ctx context.Context, sels []selected.Selection, typ common.Type, path *pathSegment, resolver reflect.Value, out *bytes.Buffer) {
	t, nonNull := unwrapNonNull(typ)
	switch t := t.(type) {
	case *schema.Object, *schema.Interface, *schema.Union:
		if resolver.Kind() == reflect.Ptr && resolver.IsNil() {
			if nonNull {
				panic(errors.Errorf("got nil for non-null %q", t))
			}
			out.WriteString("null")
			return
		}

		r.execSelections(ctx, sels, path, resolver, out, false)
		return
	}

	if !nonNull {
		if resolver.IsNil() {
			out.WriteString("null")
			return
		}
		resolver = resolver.Elem()
	}

	switch t := t.(type) {
	case *common.List:
		l := resolver.Len()

		if selected.HasAsyncSel(sels) {
			var wg sync.WaitGroup
			wg.Add(l)
			entryouts := make([]bytes.Buffer, l)
			for i := 0; i < l; i++ {
				go func(i int) {
					defer wg.Done()
					defer r.handlePanic(ctx)
					r.execSelectionSet(ctx, sels, t.OfType, &pathSegment{path, i}, resolver.Index(i), &entryouts[i])
				}(i)
			}
			wg.Wait()

			out.WriteByte('[')
			for i, entryout := range entryouts {
				if i > 0 {
					out.WriteByte(',')
				}
				out.Write(entryout.Bytes())
			}
			out.WriteByte(']')
			return
		}

		out.WriteByte('[')
		for i := 0; i < l; i++ {
			if i > 0 {
				out.WriteByte(',')
			}
			r.execSelectionSet(ctx, sels, t.OfType, &pathSegment{path, i}, resolver.Index(i), out)
		}
		out.WriteByte(']')

	case *schema.Scalar:
		v := resolver.Interface()
		data, err := json.Marshal(v)
		if err != nil {
			panic(errors.Errorf("could not marshal %v: %s", v, err))
		}
		out.Write(data)

	case *schema.Enum:
		out.WriteByte('"')
		out.WriteString(resolver.String())
		out.WriteByte('"')

	default:
		panic("unreachable")
	}
}

func unwrapNonNull(t common.Type) (common.Type, bool) {
	if nn, ok := t.(*common.NonNull); ok {
		return nn.OfType, true
	}
	return t, false
}

type pathSegment struct {
	parent *pathSegment
	value  interface{}
}

func (p *pathSegment) toSlice() []interface{} {
	if p == nil {
		return nil
	}
	return append(p.parent.toSlice(), p.value)
}
<|endoftext|>


<commit_before>package pool

import (
	"context"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-redis/redis/internal"
)

var ErrClosed = errors.New("redis: client is closed")
var ErrPoolTimeout = errors.New("redis: connection pool timeout")

var timers = sync.Pool{
	New: func() interface{} {
		t := time.NewTimer(time.Hour)
		t.Stop()
		return t
	},
}

// Stats contains pool state information and accumulated stats.
type Stats struct {
	Hits     uint32 // number of times free connection was found in the pool
	Misses   uint32 // number of times free connection was NOT found in the pool
	Timeouts uint32 // number of times a wait timeout occurred

	TotalConns uint32 // number of total connections in the pool
	IdleConns  uint32 // number of idle connections in the pool
	StaleConns uint32 // number of stale connections removed from the pool
}

type Pooler interface {
	NewConn(context.Context) (*Conn, error)
	CloseConn(*Conn) error

	Get(context.Context) (*Conn, error)
	Put(*Conn)
	Remove(*Conn)

	Len() int
	IdleLen() int
	Stats() *Stats

	Close() error
}

type Options struct {
	Dialer  func(c context.Context) (net.Conn, error)
	OnClose func(*Conn) error

	PoolSize           int
	MinIdleConns       int
	MaxConnAge         time.Duration
	PoolTimeout        time.Duration
	IdleTimeout        time.Duration
	IdleCheckFrequency time.Duration
}

type ConnPool struct {
	opt *Options

	dialErrorsNum uint32 // atomic

	lastDialErrorMu sync.RWMutex
	lastDialError   error

	queue chan struct{}

	connsMu      sync.Mutex
	conns        []*Conn
	idleConns    []*Conn
	poolSize     int
	idleConnsLen int

	stats Stats

	_closed uint32 // atomic
}

var _ Pooler = (*ConnPool)(nil)

func NewConnPool(opt *Options) *ConnPool {
	p := &ConnPool{
		opt: opt,

		queue:     make(chan struct{}, opt.PoolSize),
		conns:     make([]*Conn, 0, opt.PoolSize),
		idleConns: make([]*Conn, 0, opt.PoolSize),
	}

	for i := 0; i < opt.MinIdleConns; i++ {
		p.checkMinIdleConns()
	}

	if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 {
		go p.reaper(opt.IdleCheckFrequency)
	}

	return p
}

func (p *ConnPool) checkMinIdleConns() {
	if p.opt.MinIdleConns == 0 {
		return
	}
	if p.poolSize < p.opt.PoolSize && p.idleConnsLen < p.opt.MinIdleConns {
		p.poolSize++
		p.idleConnsLen++
		go p.addIdleConn()
	}
}

func (p *ConnPool) addIdleConn() {
	cn, err := p.newConn(context.TODO(), true)
	if err != nil {
		return
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	p.idleConns = append(p.idleConns, cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) NewConn(ctx context.Context) (*Conn, error) {
	return p._NewConn(ctx, false)
}

func (p *ConnPool) _NewConn(ctx context.Context, pooled bool) (*Conn, error) {
	cn, err := p.newConn(ctx, pooled)
	if err != nil {
		return nil, err
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	if pooled {
		// If pool is full remove the cn on next Put.
		if p.poolSize >= p.opt.PoolSize {
			cn.pooled = false
		} else {
			p.poolSize++
		}
	}
	p.connsMu.Unlock()
	return cn, nil
}

func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	if atomic.LoadUint32(&p.dialErrorsNum) >= uint32(p.opt.PoolSize) {
		return nil, p.getLastDialError()
	}

	netConn, err := p.opt.Dialer(ctx)
	if err != nil {
		p.setLastDialError(err)
		if atomic.AddUint32(&p.dialErrorsNum, 1) == uint32(p.opt.PoolSize) {
			go p.tryDial()
		}
		return nil, err
	}

	cn := NewConn(netConn)
	cn.pooled = pooled
	return cn, nil
}

func (p *ConnPool) tryDial() {
	for {
		if p.closed() {
			return
		}

		conn, err := p.opt.Dialer(context.Background())
		if err != nil {
			p.setLastDialError(err)
			time.Sleep(time.Second)
			continue
		}

		atomic.StoreUint32(&p.dialErrorsNum, 0)
		_ = conn.Close()
		return
	}
}

func (p *ConnPool) setLastDialError(err error) {
	p.lastDialErrorMu.Lock()
	p.lastDialError = err
	p.lastDialErrorMu.Unlock()
}

func (p *ConnPool) getLastDialError() error {
	p.lastDialErrorMu.RLock()
	err := p.lastDialError
	p.lastDialErrorMu.RUnlock()
	return err
}

// Get returns existed connection from the pool or creates a new one.
func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	err := p.waitTurn(ctx)
	if err != nil {
		return nil, err
	}

	for {
		p.connsMu.Lock()
		cn := p.popIdle()
		p.connsMu.Unlock()

		if cn == nil {
			break
		}

		if p.isStaleConn(cn) {
			_ = p.CloseConn(cn)
			continue
		}

		atomic.AddUint32(&p.stats.Hits, 1)
		return cn, nil
	}

	atomic.AddUint32(&p.stats.Misses, 1)

	newcn, err := p._NewConn(ctx, true)
	if err != nil {
		p.freeTurn()
		return nil, err
	}

	return newcn, nil
}

func (p *ConnPool) getTurn() {
	p.queue <- struct{}{}
}

func (p *ConnPool) waitTurn(ctx context.Context) error {
	var done <-chan struct{}
	if ctx != nil {
		done = ctx.Done()
	}

	select {
	case <-done:
		return ctx.Err()
	case p.queue <- struct{}{}:
		return nil
	default:
		timer := timers.Get().(*time.Timer)
		timer.Reset(p.opt.PoolTimeout)

		select {
		case <-done:
			if !timer.Stop() {
				<-timer.C
			}
			timers.Put(timer)
			return ctx.Err()
		case p.queue <- struct{}{}:
			if !timer.Stop() {
				<-timer.C
			}
			timers.Put(timer)
			return nil
		case <-timer.C:
			timers.Put(timer)
			atomic.AddUint32(&p.stats.Timeouts, 1)
			return ErrPoolTimeout
		}
	}
}

func (p *ConnPool) freeTurn() {
	<-p.queue
}

func (p *ConnPool) popIdle() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	idx := len(p.idleConns) - 1
	cn := p.idleConns[idx]
	p.idleConns = p.idleConns[:idx]
	p.idleConnsLen--
	p.checkMinIdleConns()
	return cn
}

func (p *ConnPool) Put(cn *Conn) {
	if !cn.pooled {
		p.Remove(cn)
		return
	}

	p.connsMu.Lock()
	p.idleConns = append(p.idleConns, cn)
	p.idleConnsLen++
	p.connsMu.Unlock()
	p.freeTurn()
}

func (p *ConnPool) Remove(cn *Conn) {
	p.removeConnWithLock(cn)
	p.freeTurn()
	_ = p.closeConn(cn)
}

func (p *ConnPool) CloseConn(cn *Conn) error {
	p.removeConnWithLock(cn)
	return p.closeConn(cn)
}

func (p *ConnPool) removeConnWithLock(cn *Conn) {
	p.connsMu.Lock()
	p.removeConn(cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) removeConn(cn *Conn) {
	for i, c := range p.conns {
		if c == cn {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			if cn.pooled {
				p.poolSize--
				p.checkMinIdleConns()
			}
			return
		}
	}
}

func (p *ConnPool) closeConn(cn *Conn) error {
	if p.opt.OnClose != nil {
		_ = p.opt.OnClose(cn)
	}
	return cn.Close()
}

// Len returns total number of connections.
func (p *ConnPool) Len() int {
	p.connsMu.Lock()
	n := len(p.conns)
	p.connsMu.Unlock()
	return n
}

// IdleLen returns number of idle connections.
func (p *ConnPool) IdleLen() int {
	p.connsMu.Lock()
	n := p.idleConnsLen
	p.connsMu.Unlock()
	return n
}

func (p *ConnPool) Stats() *Stats {
	idleLen := p.IdleLen()
	return &Stats{
		Hits:     atomic.LoadUint32(&p.stats.Hits),
		Misses:   atomic.LoadUint32(&p.stats.Misses),
		Timeouts: atomic.LoadUint32(&p.stats.Timeouts),

		TotalConns: uint32(p.Len()),
		IdleConns:  uint32(idleLen),
		StaleConns: atomic.LoadUint32(&p.stats.StaleConns),
	}
}

func (p *ConnPool) closed() bool {
	return atomic.LoadUint32(&p._closed) == 1
}

func (p *ConnPool) Filter(fn func(*Conn) bool) error {
	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if fn(cn) {
			if err := p.closeConn(cn); err != nil && firstErr == nil {
				firstErr = err
			}
		}
	}
	p.connsMu.Unlock()
	return firstErr
}

func (p *ConnPool) Close() error {
	if !atomic.CompareAndSwapUint32(&p._closed, 0, 1) {
		return ErrClosed
	}

	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if err := p.closeConn(cn); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	p.conns = nil
	p.poolSize = 0
	p.idleConns = nil
	p.idleConnsLen = 0
	p.connsMu.Unlock()

	return firstErr
}

func (p *ConnPool) reaper(frequency time.Duration) {
	ticker := time.NewTicker(frequency)
	defer ticker.Stop()

	for range ticker.C {
		if p.closed() {
			break
		}
		n, err := p.ReapStaleConns()
		if err != nil {
			internal.Logger.Printf("ReapStaleConns failed: %s", err)
			continue
		}
		atomic.AddUint32(&p.stats.StaleConns, uint32(n))
	}
}

func (p *ConnPool) ReapStaleConns() (int, error) {
	var n int
	for {
		p.getTurn()

		p.connsMu.Lock()
		cn := p.reapStaleConn()
		p.connsMu.Unlock()
		p.freeTurn()

		if cn != nil {
			p.closeConn(cn)
			n++
		} else {
			break
		}
	}
	return n, nil
}

func (p *ConnPool) reapStaleConn() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	cn := p.idleConns[0]
	if !p.isStaleConn(cn) {
		return nil
	}

	p.idleConns = append(p.idleConns[:0], p.idleConns[1:]...)
	p.idleConnsLen--
	p.removeConn(cn)

	return cn
}

func (p *ConnPool) isStaleConn(cn *Conn) bool {
	if p.opt.IdleTimeout == 0 && p.opt.MaxConnAge == 0 {
		return false
	}

	now := time.Now()
	if p.opt.IdleTimeout > 0 && now.Sub(cn.UsedAt()) >= p.opt.IdleTimeout {
		return true
	}
	if p.opt.MaxConnAge > 0 && now.Sub(cn.createdAt) >= p.opt.MaxConnAge {
		return true
	}

	return false
}
<commit_msg>Inc StaleConns in ReapStaleConns<commit_after>package pool

import (
	"context"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-redis/redis/internal"
)

var ErrClosed = errors.New("redis: client is closed")
var ErrPoolTimeout = errors.New("redis: connection pool timeout")

var timers = sync.Pool{
	New: func() interface{} {
		t := time.NewTimer(time.Hour)
		t.Stop()
		return t
	},
}

// Stats contains pool state information and accumulated stats.
type Stats struct {
	Hits     uint32 // number of times free connection was found in the pool
	Misses   uint32 // number of times free connection was NOT found in the pool
	Timeouts uint32 // number of times a wait timeout occurred

	TotalConns uint32 // number of total connections in the pool
	IdleConns  uint32 // number of idle connections in the pool
	StaleConns uint32 // number of stale connections removed from the pool
}

type Pooler interface {
	NewConn(context.Context) (*Conn, error)
	CloseConn(*Conn) error

	Get(context.Context) (*Conn, error)
	Put(*Conn)
	Remove(*Conn)

	Len() int
	IdleLen() int
	Stats() *Stats

	Close() error
}

type Options struct {
	Dialer  func(c context.Context) (net.Conn, error)
	OnClose func(*Conn) error

	PoolSize           int
	MinIdleConns       int
	MaxConnAge         time.Duration
	PoolTimeout        time.Duration
	IdleTimeout        time.Duration
	IdleCheckFrequency time.Duration
}

type ConnPool struct {
	opt *Options

	dialErrorsNum uint32 // atomic

	lastDialErrorMu sync.RWMutex
	lastDialError   error

	queue chan struct{}

	connsMu      sync.Mutex
	conns        []*Conn
	idleConns    []*Conn
	poolSize     int
	idleConnsLen int

	stats Stats

	_closed uint32 // atomic
}

var _ Pooler = (*ConnPool)(nil)

func NewConnPool(opt *Options) *ConnPool {
	p := &ConnPool{
		opt: opt,

		queue:     make(chan struct{}, opt.PoolSize),
		conns:     make([]*Conn, 0, opt.PoolSize),
		idleConns: make([]*Conn, 0, opt.PoolSize),
	}

	for i := 0; i < opt.MinIdleConns; i++ {
		p.checkMinIdleConns()
	}

	if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 {
		go p.reaper(opt.IdleCheckFrequency)
	}

	return p
}

func (p *ConnPool) checkMinIdleConns() {
	if p.opt.MinIdleConns == 0 {
		return
	}
	if p.poolSize < p.opt.PoolSize && p.idleConnsLen < p.opt.MinIdleConns {
		p.poolSize++
		p.idleConnsLen++
		go p.addIdleConn()
	}
}

func (p *ConnPool) addIdleConn() {
	cn, err := p.newConn(context.TODO(), true)
	if err != nil {
		return
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	p.idleConns = append(p.idleConns, cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) NewConn(ctx context.Context) (*Conn, error) {
	return p._NewConn(ctx, false)
}

func (p *ConnPool) _NewConn(ctx context.Context, pooled bool) (*Conn, error) {
	cn, err := p.newConn(ctx, pooled)
	if err != nil {
		return nil, err
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	if pooled {
		// If pool is full remove the cn on next Put.
		if p.poolSize >= p.opt.PoolSize {
			cn.pooled = false
		} else {
			p.poolSize++
		}
	}
	p.connsMu.Unlock()
	return cn, nil
}

func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	if atomic.LoadUint32(&p.dialErrorsNum) >= uint32(p.opt.PoolSize) {
		return nil, p.getLastDialError()
	}

	netConn, err := p.opt.Dialer(ctx)
	if err != nil {
		p.setLastDialError(err)
		if atomic.AddUint32(&p.dialErrorsNum, 1) == uint32(p.opt.PoolSize) {
			go p.tryDial()
		}
		return nil, err
	}

	cn := NewConn(netConn)
	cn.pooled = pooled
	return cn, nil
}

func (p *ConnPool) tryDial() {
	for {
		if p.closed() {
			return
		}

		conn, err := p.opt.Dialer(context.Background())
		if err != nil {
			p.setLastDialError(err)
			time.Sleep(time.Second)
			continue
		}

		atomic.StoreUint32(&p.dialErrorsNum, 0)
		_ = conn.Close()
		return
	}
}

func (p *ConnPool) setLastDialError(err error) {
	p.lastDialErrorMu.Lock()
	p.lastDialError = err
	p.lastDialErrorMu.Unlock()
}

func (p *ConnPool) getLastDialError() error {
	p.lastDialErrorMu.RLock()
	err := p.lastDialError
	p.lastDialErrorMu.RUnlock()
	return err
}

// Get returns existed connection from the pool or creates a new one.
func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	err := p.waitTurn(ctx)
	if err != nil {
		return nil, err
	}

	for {
		p.connsMu.Lock()
		cn := p.popIdle()
		p.connsMu.Unlock()

		if cn == nil {
			break
		}

		if p.isStaleConn(cn) {
			_ = p.CloseConn(cn)
			continue
		}

		atomic.AddUint32(&p.stats.Hits, 1)
		return cn, nil
	}

	atomic.AddUint32(&p.stats.Misses, 1)

	newcn, err := p._NewConn(ctx, true)
	if err != nil {
		p.freeTurn()
		return nil, err
	}

	return newcn, nil
}

func (p *ConnPool) getTurn() {
	p.queue <- struct{}{}
}

func (p *ConnPool) waitTurn(ctx context.Context) error {
	var done <-chan struct{}
	if ctx != nil {
		done = ctx.Done()
	}

	select {
	case <-done:
		return ctx.Err()
	case p.queue <- struct{}{}:
		return nil
	default:
		timer := timers.Get().(*time.Timer)
		timer.Reset(p.opt.PoolTimeout)

		select {
		case <-done:
			if !timer.Stop() {
				<-timer.C
			}
			timers.Put(timer)
			return ctx.Err()
		case p.queue <- struct{}{}:
			if !timer.Stop() {
				<-timer.C
			}
			timers.Put(timer)
			return nil
		case <-timer.C:
			timers.Put(timer)
			atomic.AddUint32(&p.stats.Timeouts, 1)
			return ErrPoolTimeout
		}
	}
}

func (p *ConnPool) freeTurn() {
	<-p.queue
}

func (p *ConnPool) popIdle() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	idx := len(p.idleConns) - 1
	cn := p.idleConns[idx]
	p.idleConns = p.idleConns[:idx]
	p.idleConnsLen--
	p.checkMinIdleConns()
	return cn
}

func (p *ConnPool) Put(cn *Conn) {
	if !cn.pooled {
		p.Remove(cn)
		return
	}

	p.connsMu.Lock()
	p.idleConns = append(p.idleConns, cn)
	p.idleConnsLen++
	p.connsMu.Unlock()
	p.freeTurn()
}

func (p *ConnPool) Remove(cn *Conn) {
	p.removeConnWithLock(cn)
	p.freeTurn()
	_ = p.closeConn(cn)
}

func (p *ConnPool) CloseConn(cn *Conn) error {
	p.removeConnWithLock(cn)
	return p.closeConn(cn)
}

func (p *ConnPool) removeConnWithLock(cn *Conn) {
	p.connsMu.Lock()
	p.removeConn(cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) removeConn(cn *Conn) {
	for i, c := range p.conns {
		if c == cn {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			if cn.pooled {
				p.poolSize--
				p.checkMinIdleConns()
			}
			return
		}
	}
}

func (p *ConnPool) closeConn(cn *Conn) error {
	if p.opt.OnClose != nil {
		_ = p.opt.OnClose(cn)
	}
	return cn.Close()
}

// Len returns total number of connections.
func (p *ConnPool) Len() int {
	p.connsMu.Lock()
	n := len(p.conns)
	p.connsMu.Unlock()
	return n
}

// IdleLen returns number of idle connections.
func (p *ConnPool) IdleLen() int {
	p.connsMu.Lock()
	n := p.idleConnsLen
	p.connsMu.Unlock()
	return n
}

func (p *ConnPool) Stats() *Stats {
	idleLen := p.IdleLen()
	return &Stats{
		Hits:     atomic.LoadUint32(&p.stats.Hits),
		Misses:   atomic.LoadUint32(&p.stats.Misses),
		Timeouts: atomic.LoadUint32(&p.stats.Timeouts),

		TotalConns: uint32(p.Len()),
		IdleConns:  uint32(idleLen),
		StaleConns: atomic.LoadUint32(&p.stats.StaleConns),
	}
}

func (p *ConnPool) closed() bool {
	return atomic.LoadUint32(&p._closed) == 1
}

func (p *ConnPool) Filter(fn func(*Conn) bool) error {
	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if fn(cn) {
			if err := p.closeConn(cn); err != nil && firstErr == nil {
				firstErr = err
			}
		}
	}
	p.connsMu.Unlock()
	return firstErr
}

func (p *ConnPool) Close() error {
	if !atomic.CompareAndSwapUint32(&p._closed, 0, 1) {
		return ErrClosed
	}

	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if err := p.closeConn(cn); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	p.conns = nil
	p.poolSize = 0
	p.idleConns = nil
	p.idleConnsLen = 0
	p.connsMu.Unlock()

	return firstErr
}

func (p *ConnPool) reaper(frequency time.Duration) {
	ticker := time.NewTicker(frequency)
	defer ticker.Stop()

	for range ticker.C {
		if p.closed() {
			break
		}
		_, err := p.ReapStaleConns()
		if err != nil {
			internal.Logger.Printf("ReapStaleConns failed: %s", err)
			continue
		}
	}
}

func (p *ConnPool) ReapStaleConns() (int, error) {
	var n int
	for {
		p.getTurn()

		p.connsMu.Lock()
		cn := p.reapStaleConn()
		p.connsMu.Unlock()
		p.freeTurn()

		if cn != nil {
			p.closeConn(cn)
			n++
		} else {
			break
		}
	}
	atomic.AddUint32(&p.stats.StaleConns, uint32(n))
	return n, nil
}

func (p *ConnPool) reapStaleConn() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	cn := p.idleConns[0]
	if !p.isStaleConn(cn) {
		return nil
	}

	p.idleConns = append(p.idleConns[:0], p.idleConns[1:]...)
	p.idleConnsLen--
	p.removeConn(cn)

	return cn
}

func (p *ConnPool) isStaleConn(cn *Conn) bool {
	if p.opt.IdleTimeout == 0 && p.opt.MaxConnAge == 0 {
		return false
	}

	now := time.Now()
	if p.opt.IdleTimeout > 0 && now.Sub(cn.UsedAt()) >= p.opt.IdleTimeout {
		return true
	}
	if p.opt.MaxConnAge > 0 && now.Sub(cn.createdAt) >= p.opt.MaxConnAge {
		return true
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

// mtail programs may be updated while mtail is running, and they will be
// reloaded without having to restart the mtail process. Programs can be
// created and deleted as well, and some configuration systems do an atomic
// rename of the program when it is installed, so mtail is also aware of file
// moves.  The Master Control Program is responsible for managing the lifetime
// of mtail programs.

import (
	"context"
	"expvar"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"go.opencensus.io/trace"

	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/watcher"
)

var (
	// LineCount counts the number of lines received by the program loader.
	LineCount = expvar.NewInt("lines_total")
	// ProgLoads counts the number of program load events.
	ProgLoads = expvar.NewMap("prog_loads_total")
	// ProgLoadErrors counts the number of program load errors.
	ProgLoadErrors    = expvar.NewMap("prog_load_errors_total")
	progRuntimeErrors = expvar.NewMap("prog_runtime_errors_total")
)

const (
	fileExt = ".mtail"
)

// LoadAllPrograms loads all programs in a directory and starts watching the
// directory for filesystem changes.  Any compile errors are stored for later retrieival.
// This function returns an error if an internal error occurs.
func (l *Loader) LoadAllPrograms() error {
	s, err := os.Stat(l.programPath)
	if err != nil {
		return errors.Wrapf(err, "failed to stat %q", l.programPath)
	}
	switch {
	case s.IsDir():
		fis, rerr := ioutil.ReadDir(l.programPath)
		if rerr != nil {
			return errors.Wrapf(rerr, "Failed to list programs in %q", l.programPath)
		}

		for _, fi := range fis {
			if fi.IsDir() {
				continue
			}
			err = l.LoadProgram(path.Join(l.programPath, fi.Name()))
			if err != nil {
				if l.errorsAbort {
					return err
				}
				glog.Warning(err)
			}
		}
	default:
		err = l.LoadProgram(l.programPath)
		if err != nil {
			if l.errorsAbort {
				return err
			}
			glog.Warning(err)
		}
	}
	return nil
}

// LoadProgram loads or reloads a program from the full pathname programPath.  The name of
// the program is the basename of the file.
func (l *Loader) LoadProgram(programPath string) error {
	name := filepath.Base(programPath)
	if strings.HasPrefix(name, ".") {
		glog.V(2).Infof("Skipping %s because it is a hidden file.", programPath)
		return nil
	}
	if filepath.Ext(name) != fileExt {
		glog.V(2).Infof("Skipping %s due to file extension.", programPath)
		return nil
	}
	f, err := os.OpenFile(programPath, os.O_RDONLY, 0600)
	if err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "Failed to read program %q", programPath)
	}
	defer func() {
		if err := f.Close(); err != nil {
			glog.Warning(err)
		}
	}()
	l.programErrorMu.Lock()
	defer l.programErrorMu.Unlock()
	l.programErrors[name] = l.CompileAndRun(name, f)
	if l.programErrors[name] != nil {
		if l.errorsAbort {
			return l.programErrors[name]
		}
		glog.Infof("Compile errors for %s:\n%s", name, l.programErrors[name])
	}
	return nil
}

const loaderTemplate = `

Program Loader

program name
errors
load errors
load successes
runtime errors
last runtime error

{{range $name, $errors := $.Errors}}
{{$name}}

{{if $errors}}
{{$errors}}
{{else}}
No compile errors
{{end}}

{{index $.Loaderrors $name}}
{{index $.Loadsuccess $name}}
{{index $.RuntimeErrors $name}}

{{index $.RuntimeErrorString $name}}

{{end}}

`

// WriteStatusHTML writes the current state of the loader as HTML to the given writer w.
func (l *Loader) WriteStatusHTML(w io.Writer) error {
	t, err := template.New("loader").Parse(loaderTemplate)
	if err != nil {
		return err
	}
	l.programErrorMu.RLock()
	defer l.programErrorMu.RUnlock()
	data := struct {
		Errors             map[string]error
		Loaderrors         map[string]string
		Loadsuccess        map[string]string
		RuntimeErrors      map[string]string
		RuntimeErrorString map[string]string
	}{
		l.programErrors,
		make(map[string]string),
		make(map[string]string),
		make(map[string]string),
		make(map[string]string),
	}
	for name := range l.programErrors {
		if ProgLoadErrors.Get(name) != nil {
			data.Loaderrors[name] = ProgLoadErrors.Get(name).String()
		}
		if ProgLoads.Get(name) != nil {
			data.Loadsuccess[name] = ProgLoads.Get(name).String()
		}
		if progRuntimeErrors.Get(name) != nil {
			data.RuntimeErrors[name] = progRuntimeErrors.Get(name).String()
		}
		data.RuntimeErrorString[name] = l.handles[name].RuntimeErrorString()
	}
	return t.Execute(w, data)
}

// CompileAndRun compiles a program read from the input, starting execution if
// it succeeds.  If an existing virtual machine of the same name already
// exists, the previous virtual machine is terminated and the new loaded over
// it.  If the new program fails to compile, any existing virtual machine with
// the same name remains running.
func (l *Loader) CompileAndRun(name string, input io.Reader) error {
	glog.V(2).Infof("CompileAndRun %s", name)
	v, errs := Compile(name, input, l.dumpAst, l.dumpAstTypes, l.syslogUseCurrentYear, l.overrideLocation)
	if errs != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("compile failed for %s:\n%s", name, errs)
	}
	if v == nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("Internal error: Compilation failed for %s: No program returned, but no errors.", name)
	}

	if l.dumpBytecode {
		glog.Info("Dumping program objects and bytecode\n", v.DumpByteCode())
	}

	// Load the metrics from the compilation into the global metric storage for export.
	for _, m := range v.m {
		if !m.Hidden {
			if l.omitMetricSource {
				m.Source = ""
			}
			err := l.ms.Add(m)
			if err != nil {
				return err
			}
		}
	}

	ProgLoads.Add(name, 1)
	glog.Infof("Loaded program %s", name)

	if l.compileOnly {
		return nil
	}

	l.handleMu.Lock()
	defer l.handleMu.Unlock()

	l.handles[name] = v
	return nil
}

// Loader handles the lifecycle of programs and virtual machines, by watching
// the configured program source directory, compiling changes to programs, and
// managing the virtual machines.
type Loader struct {
	ms          *metrics.Store        // pointer to metrics.Store to pass to compiler
	w           watcher.Watcher       // watches for program changes
	reg         prometheus.Registerer // plce to reg metrics
	programPath string                // Path that contains mtail programs.

	handleMu sync.RWMutex   // guards accesses to handles
	handles  map[string]*VM // map of program names to virtual machines

	programErrorMu sync.RWMutex     // guards access to programErrors
	programErrors  map[string]error // errors from the last compile attempt of the program

	overrideLocation     *time.Location // Instructs the vm to override the timezone with the specified zone.
	compileOnly          bool           // Only compile programs and report errors, do not load VMs.
	errorsAbort          bool           // Compiler errors abort the loader.
	dumpAst              bool           // print the AST after parse
	dumpAstTypes         bool           // print the AST after type check
	dumpBytecode         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	syslogUseCurrentYear bool           // Instructs the VM to overwrite zero years with the current year in a strptime instruction.
	omitMetricSource     bool

	signalQuit chan struct{} // When closed stops the signal handler goroutine.
}

// OverrideLocation sets the timezone location for the VM.
func OverrideLocation(loc *time.Location) func(*Loader) error {
	return func(l *Loader) error {
		l.overrideLocation = loc
		return nil
	}
}

// CompileOnly sets the Loader to compile programs only, without executing them.
func CompileOnly(l *Loader) error {
	l.compileOnly = true
	return ErrorsAbort(l)
}

// ErrorsAbort sets the Loader to abort the Loader on compile errors.
func ErrorsAbort(l *Loader) error {
	l.errorsAbort = true
	return nil
}

// DumpAst instructs the Loader to print the AST after program compilation.
func DumpAst(l *Loader) error {
	l.dumpAst = true
	return nil
}

// DumpAstTypes instructs the Loader to print the AST after type checking.
func DumpAstTypes(l *Loader) error {
	l.dumpAstTypes = true
	return nil
}

// DumpBytecode instructs the loader to print the compiled bytecode after code generation.
func DumpBytecode(l *Loader) error {
	l.dumpBytecode = true
	return nil
}

// SyslogUseCurrentYear instructs the VM to annotate yearless timestamps with the current year.
func SyslogUseCurrentYear(l *Loader) error {
	l.syslogUseCurrentYear = true
	return nil
}

// OmitMetricSource instructs the Loader to not annotate metrics with their program source when added to the metric store.
func OmitMetricSource(l *Loader) error {
	l.omitMetricSource = true
	return nil
}

// PrometheusRegisterer passes in a registry for setting up exported metrics.
func PrometheusRegisterer(reg prometheus.Registerer) func(l *Loader) error {
	return func(l *Loader) error {
		l.reg = reg
		return nil
	}
}

// NewLoader creates a new program loader that reads programs from programPath.
func NewLoader(programPath string, store *metrics.Store, options ...func(*Loader) error) (*Loader, error) {
	if store == nil {
		return nil, errors.New("loader needs a store")
	}
	l := &Loader{
		ms:            store,
		programPath:   programPath,
		handles:       make(map[string]*VM),
		programErrors: make(map[string]error),
		signalQuit:    make(chan struct{}),
	}
	if err := l.SetOption(options...); err != nil {
		return nil, err
	}
	if l.reg != nil {
		l.reg.MustRegister(lineProcessingDurations)
	}
	go func() {
		n := make(chan os.Signal, 1)
		signal.Notify(n, syscall.SIGHUP)
		defer signal.Stop(n)
		for {
			select {
			case <-n:
				l.LoadAllPrograms()
			case <-l.signalQuit:
				return
			}
		}
	}()
	return l, nil
}

// SetOption takes one or more option functions and applies them in order to Loader.
func (l *Loader) SetOption(options ...func(*Loader) error) error {
	for _, option := range options {
		if err := option(l); err != nil {
			return err
		}
	}
	return nil
}

func (l *Loader) Close() {
	glog.Info("Shutting down loader.")
	// Shut down the signal handler before deleting the programs.
	// A more complete shutdown would wait on the signal handler
	// shutdown before proceeding, but the risk of a SIGHUP showing up right
	// now is mitigated by the lock below -- at worst we might reload all
	// programmes before deleting them all.
	close(l.signalQuit)
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	for prog := range l.handles {
		delete(l.handles, prog)
	}
}

// ProcessLogLine satisfies the LogLine.Processor interface.
func (l *Loader) ProcessLogLine(ctx context.Context, ll *logline.LogLine) {
	ctx, span := trace.StartSpan(ctx, "Loader.ProcessLogLine")
	defer span.End()
	LineCount.Add(1)
	l.handleMu.RLock()
	defer l.handleMu.RUnlock()
	for prog := range l.handles {
		l.handles[prog].ProcessLogLine(ctx, ll)
	}
}

// UnloadProgram removes the named program, any currently running VM goroutine.
func (l *Loader) UnloadProgram(pathname string) {
	name := filepath.Base(pathname)
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	if _, ok := l.handles[name]; ok {
		delete(l.handles, name)
	}
}

func (l *Loader) ProgzHandler(w http.ResponseWriter, r *http.Request) {
	prog := r.URL.Query().Get("prog")
	if prog != "" {
		l.handleMu.RLock()
		v, ok := l.handles[prog]
		l.handleMu.RUnlock()
		if !ok {
			http.Error(w, "No program found", http.StatusNotFound)
			return
		}
		fmt.Fprintf(w, v.DumpByteCode())
		fmt.Fprintf(w, "\nLast runtime error:\n%s", v.RuntimeErrorString())
		return
	}
	l.handleMu.RLock()
	defer l.handleMu.RUnlock()
	w.Header().Add("Content-type", "text/html")
	fmt.Fprintf(w, "
")
	for prog := range l.handles {
		fmt.Fprintf(w, "
%s
", prog, prog)
	}
	fmt.Fprintf(w, "
")
}
<commit_msg>Handle the error when calling LoadAllPrograms in the loader poll loop.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

// mtail programs may be updated while mtail is running, and they will be
// reloaded without having to restart the mtail process. Programs can be
// created and deleted as well, and some configuration systems do an atomic
// rename of the program when it is installed, so mtail is also aware of file
// moves.  The Master Control Program is responsible for managing the lifetime
// of mtail programs.

import (
	"context"
	"expvar"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"go.opencensus.io/trace"

	"github.com/google/mtail/internal/logline"
	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/watcher"
)

var (
	// LineCount counts the number of lines received by the program loader.
	LineCount = expvar.NewInt("lines_total")
	// ProgLoads counts the number of program load events.
	ProgLoads = expvar.NewMap("prog_loads_total")
	// ProgLoadErrors counts the number of program load errors.
	ProgLoadErrors    = expvar.NewMap("prog_load_errors_total")
	progRuntimeErrors = expvar.NewMap("prog_runtime_errors_total")
)

const (
	fileExt = ".mtail"
)

// LoadAllPrograms loads all programs in a directory and starts watching the
// directory for filesystem changes.  Any compile errors are stored for later retrieival.
// This function returns an error if an internal error occurs.
func (l *Loader) LoadAllPrograms() error {
	s, err := os.Stat(l.programPath)
	if err != nil {
		return errors.Wrapf(err, "failed to stat %q", l.programPath)
	}
	switch {
	case s.IsDir():
		fis, rerr := ioutil.ReadDir(l.programPath)
		if rerr != nil {
			return errors.Wrapf(rerr, "Failed to list programs in %q", l.programPath)
		}

		for _, fi := range fis {
			if fi.IsDir() {
				continue
			}
			err = l.LoadProgram(path.Join(l.programPath, fi.Name()))
			if err != nil {
				if l.errorsAbort {
					return err
				}
				glog.Warning(err)
			}
		}
	default:
		err = l.LoadProgram(l.programPath)
		if err != nil {
			if l.errorsAbort {
				return err
			}
			glog.Warning(err)
		}
	}
	return nil
}

// LoadProgram loads or reloads a program from the full pathname programPath.  The name of
// the program is the basename of the file.
func (l *Loader) LoadProgram(programPath string) error {
	name := filepath.Base(programPath)
	if strings.HasPrefix(name, ".") {
		glog.V(2).Infof("Skipping %s because it is a hidden file.", programPath)
		return nil
	}
	if filepath.Ext(name) != fileExt {
		glog.V(2).Infof("Skipping %s due to file extension.", programPath)
		return nil
	}
	f, err := os.OpenFile(programPath, os.O_RDONLY, 0600)
	if err != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Wrapf(err, "Failed to read program %q", programPath)
	}
	defer func() {
		if err := f.Close(); err != nil {
			glog.Warning(err)
		}
	}()
	l.programErrorMu.Lock()
	defer l.programErrorMu.Unlock()
	l.programErrors[name] = l.CompileAndRun(name, f)
	if l.programErrors[name] != nil {
		if l.errorsAbort {
			return l.programErrors[name]
		}
		glog.Infof("Compile errors for %s:\n%s", name, l.programErrors[name])
	}
	return nil
}

const loaderTemplate = `

Program Loader

program name
errors
load errors
load successes
runtime errors
last runtime error

{{range $name, $errors := $.Errors}}
{{$name}}

{{if $errors}}
{{$errors}}
{{else}}
No compile errors
{{end}}

{{index $.Loaderrors $name}}
{{index $.Loadsuccess $name}}
{{index $.RuntimeErrors $name}}

{{index $.RuntimeErrorString $name}}

{{end}}

`

// WriteStatusHTML writes the current state of the loader as HTML to the given writer w.
func (l *Loader) WriteStatusHTML(w io.Writer) error {
	t, err := template.New("loader").Parse(loaderTemplate)
	if err != nil {
		return err
	}
	l.programErrorMu.RLock()
	defer l.programErrorMu.RUnlock()
	data := struct {
		Errors             map[string]error
		Loaderrors         map[string]string
		Loadsuccess        map[string]string
		RuntimeErrors      map[string]string
		RuntimeErrorString map[string]string
	}{
		l.programErrors,
		make(map[string]string),
		make(map[string]string),
		make(map[string]string),
		make(map[string]string),
	}
	for name := range l.programErrors {
		if ProgLoadErrors.Get(name) != nil {
			data.Loaderrors[name] = ProgLoadErrors.Get(name).String()
		}
		if ProgLoads.Get(name) != nil {
			data.Loadsuccess[name] = ProgLoads.Get(name).String()
		}
		if progRuntimeErrors.Get(name) != nil {
			data.RuntimeErrors[name] = progRuntimeErrors.Get(name).String()
		}
		data.RuntimeErrorString[name] = l.handles[name].RuntimeErrorString()
	}
	return t.Execute(w, data)
}

// CompileAndRun compiles a program read from the input, starting execution if
// it succeeds.  If an existing virtual machine of the same name already
// exists, the previous virtual machine is terminated and the new loaded over
// it.  If the new program fails to compile, any existing virtual machine with
// the same name remains running.
func (l *Loader) CompileAndRun(name string, input io.Reader) error {
	glog.V(2).Infof("CompileAndRun %s", name)
	v, errs := Compile(name, input, l.dumpAst, l.dumpAstTypes, l.syslogUseCurrentYear, l.overrideLocation)
	if errs != nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("compile failed for %s:\n%s", name, errs)
	}
	if v == nil {
		ProgLoadErrors.Add(name, 1)
		return errors.Errorf("Internal error: Compilation failed for %s: No program returned, but no errors.", name)
	}

	if l.dumpBytecode {
		glog.Info("Dumping program objects and bytecode\n", v.DumpByteCode())
	}

	// Load the metrics from the compilation into the global metric storage for export.
	for _, m := range v.m {
		if !m.Hidden {
			if l.omitMetricSource {
				m.Source = ""
			}
			err := l.ms.Add(m)
			if err != nil {
				return err
			}
		}
	}

	ProgLoads.Add(name, 1)
	glog.Infof("Loaded program %s", name)

	if l.compileOnly {
		return nil
	}

	l.handleMu.Lock()
	defer l.handleMu.Unlock()

	l.handles[name] = v
	return nil
}

// Loader handles the lifecycle of programs and virtual machines, by watching
// the configured program source directory, compiling changes to programs, and
// managing the virtual machines.
type Loader struct {
	ms          *metrics.Store        // pointer to metrics.Store to pass to compiler
	w           watcher.Watcher       // watches for program changes
	reg         prometheus.Registerer // plce to reg metrics
	programPath string                // Path that contains mtail programs.

	handleMu sync.RWMutex   // guards accesses to handles
	handles  map[string]*VM // map of program names to virtual machines

	programErrorMu sync.RWMutex     // guards access to programErrors
	programErrors  map[string]error // errors from the last compile attempt of the program

	overrideLocation     *time.Location // Instructs the vm to override the timezone with the specified zone.
	compileOnly          bool           // Only compile programs and report errors, do not load VMs.
	errorsAbort          bool           // Compiler errors abort the loader.
	dumpAst              bool           // print the AST after parse
	dumpAstTypes         bool           // print the AST after type check
	dumpBytecode         bool           // Instructs the loader to dump to stdout the compiled program after compilation.
	syslogUseCurrentYear bool           // Instructs the VM to overwrite zero years with the current year in a strptime instruction.
	omitMetricSource     bool

	signalQuit chan struct{} // When closed stops the signal handler goroutine.
}

// OverrideLocation sets the timezone location for the VM.
func OverrideLocation(loc *time.Location) func(*Loader) error {
	return func(l *Loader) error {
		l.overrideLocation = loc
		return nil
	}
}

// CompileOnly sets the Loader to compile programs only, without executing them.
func CompileOnly(l *Loader) error {
	l.compileOnly = true
	return ErrorsAbort(l)
}

// ErrorsAbort sets the Loader to abort the Loader on compile errors.
func ErrorsAbort(l *Loader) error {
	l.errorsAbort = true
	return nil
}

// DumpAst instructs the Loader to print the AST after program compilation.
func DumpAst(l *Loader) error {
	l.dumpAst = true
	return nil
}

// DumpAstTypes instructs the Loader to print the AST after type checking.
func DumpAstTypes(l *Loader) error {
	l.dumpAstTypes = true
	return nil
}

// DumpBytecode instructs the loader to print the compiled bytecode after code generation.
func DumpBytecode(l *Loader) error {
	l.dumpBytecode = true
	return nil
}

// SyslogUseCurrentYear instructs the VM to annotate yearless timestamps with the current year.
func SyslogUseCurrentYear(l *Loader) error {
	l.syslogUseCurrentYear = true
	return nil
}

// OmitMetricSource instructs the Loader to not annotate metrics with their program source when added to the metric store.
func OmitMetricSource(l *Loader) error {
	l.omitMetricSource = true
	return nil
}

// PrometheusRegisterer passes in a registry for setting up exported metrics.
func PrometheusRegisterer(reg prometheus.Registerer) func(l *Loader) error {
	return func(l *Loader) error {
		l.reg = reg
		return nil
	}
}

// NewLoader creates a new program loader that reads programs from programPath.
func NewLoader(programPath string, store *metrics.Store, options ...func(*Loader) error) (*Loader, error) {
	if store == nil {
		return nil, errors.New("loader needs a store")
	}
	l := &Loader{
		ms:            store,
		programPath:   programPath,
		handles:       make(map[string]*VM),
		programErrors: make(map[string]error),
		signalQuit:    make(chan struct{}),
	}
	if err := l.SetOption(options...); err != nil {
		return nil, err
	}
	if l.reg != nil {
		l.reg.MustRegister(lineProcessingDurations)
	}
	go func() {
		n := make(chan os.Signal, 1)
		signal.Notify(n, syscall.SIGHUP)
		defer signal.Stop(n)
		for {
			select {
			case <-n:
				if err := l.LoadAllPrograms(); err != nil {
					glog.Info(err)
				}
			case <-l.signalQuit:
				return
			}
		}
	}()
	return l, nil
}

// SetOption takes one or more option functions and applies them in order to Loader.
func (l *Loader) SetOption(options ...func(*Loader) error) error {
	for _, option := range options {
		if err := option(l); err != nil {
			return err
		}
	}
	return nil
}

func (l *Loader) Close() {
	glog.Info("Shutting down loader.")
	// Shut down the signal handler before deleting the programs.
	// A more complete shutdown would wait on the signal handler
	// shutdown before proceeding, but the risk of a SIGHUP showing up right
	// now is mitigated by the lock below -- at worst we might reload all
	// programmes before deleting them all.
	close(l.signalQuit)
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	for prog := range l.handles {
		delete(l.handles, prog)
	}
}

// ProcessLogLine satisfies the LogLine.Processor interface.
func (l *Loader) ProcessLogLine(ctx context.Context, ll *logline.LogLine) {
	ctx, span := trace.StartSpan(ctx, "Loader.ProcessLogLine")
	defer span.End()
	LineCount.Add(1)
	l.handleMu.RLock()
	defer l.handleMu.RUnlock()
	for prog := range l.handles {
		l.handles[prog].ProcessLogLine(ctx, ll)
	}
}

// UnloadProgram removes the named program, any currently running VM goroutine.
func (l *Loader) UnloadProgram(pathname string) {
	name := filepath.Base(pathname)
	l.handleMu.Lock()
	defer l.handleMu.Unlock()
	if _, ok := l.handles[name]; ok {
		delete(l.handles, name)
	}
}

func (l *Loader) ProgzHandler(w http.ResponseWriter, r *http.Request) {
	prog := r.URL.Query().Get("prog")
	if prog != "" {
		l.handleMu.RLock()
		v, ok := l.handles[prog]
		l.handleMu.RUnlock()
		if !ok {
			http.Error(w, "No program found", http.StatusNotFound)
			return
		}
		fmt.Fprintf(w, v.DumpByteCode())
		fmt.Fprintf(w, "\nLast runtime error:\n%s", v.RuntimeErrorString())
		return
	}
	l.handleMu.RLock()
	defer l.handleMu.RUnlock()
	w.Header().Add("Content-type", "text/html")
	fmt.Fprintf(w, "
")
	for prog := range l.handles {
		fmt.Fprintf(w, "
%s
", prog, prog)
	}
	fmt.Fprintf(w, "
")
}
<|endoftext|>