<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"
)

// Avoid thundering herd problem on remote services used by this command.
// Interval will be 0, if this is not an issue.
func SpreadWait(interval time.Duration) {
	if interval > 0 {
		// Seed random generator with current process ID
		rand.Seed(int64(os.Getpid()))
		// Sleep for random amount of time within interval
		time.Sleep(time.Duration(rand.Int63n(int64(interval))))
	}
}

// Ok states that execution went well. Logs debug output and reports ok to
// monitoring.
func Ok() {
	var message string
	log.Println("OK")
	if firstbytes == nil {
		message = "OK"
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(monitorOk, message)
}

// NotAvailable states that the command could not be started successfully. It
// might not be installed or has other problems.
func NotAvailable(err error) {
	s := fmt.Sprintln("Cannot start command: ", err)
	log.Println("FATAL:", s)
	monitor(monitorUnknown, s)
}

// TimedOut states that the command took too long and reports failure to the
// monitoring.
func TimedOut() {
	s := "execution took too long"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Busy states that the command hangs and reports failure to the monitoring.
// Those tasks should be automatically killed, if it happens often.
func Busy() {
	s := "previous invocation of command still running"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Failed states that the command didn't execute successfully and reports
// failure to the monitoring. Also Logs error output.
func Failed(err error) {
	var message string
	code, s := error2exit(err)
	log.Println("FATAL:", s)
	if firstbytes == nil {
		message = s
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(code, message)
}

// LockError states that we could not get the lock.
func LockError(err error) {
	s := fmt.Sprintln("Failed to get lock: ", err)
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

var firstbytes *CapWriter
var interval, timeout time.Duration
var pipeStderr, pipeStdout bool
var useSyslog bool
var wrapNagios bool
var killRunning bool
var command string

func parseFlags() {
	flag.DurationVar(&interval, "i", -1,
		"set maximum execution start delay for command, e.g. 45s, 2m, 1h30m, default: 1/10 of timeout")
	flag.DurationVar(&timeout, "t", 1*time.Minute,
		"set execution timeout for command, e.g. 45s, 2m, 1h30m, default: 1m")
	flag.BoolVar(&useSyslog, "s", false, "log via syslog")
	flag.BoolVar(&wrapNagios, "n", false, "wrap nagios plugin (pass on return codes, pass first 8KiB of stdout as message)")
	flag.BoolVar(&pipeStderr, "e", true, "pipe stderr to log")
	flag.StringVar(&monitoringEvent, "E", "", "monitoring event (defaults to check_foo for /path/check_foo.sh ")
	flag.BoolVar(&pipeStdout, "o", true, "pipe stdout to log")
	flag.BoolVar(&killRunning, "k", false, "kill already running instance of command")
	flag.Parse()

	if flag.NArg() < 1 {
		log.Fatal("FATAL: no command to execute")
	}

	command = flag.Arg(0)

	if interval >= timeout {
		log.Fatal("FATAL: interval >= timeout, no time left for actual command execution")
	}

	if interval == -1 {
		interval = timeout / 10
	}
}

func main() {
	var cmd *exec.Cmd
	var wg sync.WaitGroup

	log.SetFlags(0)
	parseFlags()

	// default check_foo for /path/check_foo.sh
	if monitoringEvent == "" {
		monitoringEvent = filepath.Base(command)
		if ext := filepath.Ext(command); ext != "" {
			monitoringEvent = monitoringEvent[0 : len(monitoringEvent)-len("."+ext)+1]
		}
	}

	logger, err := getLogger(useSyslog)
	if err != nil {
		log.Fatal("FATAL: cannot contact syslog")
		return
	}

	loadMonitoringCommands()

	// FIXME(nightlyone) try two intervals instead of one?
	timer := time.AfterFunc(timeout, func() {
		TimedOut()
		if cmd != nil && cmd.Process != nil {
			cmd.Process.Kill()
			// FIXME(nightlyone) log the kill
		}
		os.Exit(0)
	})

	SpreadWait(interval)

	lock, err := createLock(killRunning)
	if err != nil {
		timer.Stop()
		LockError(err)
		return
	}
	defer lock.Unlock()

	cmd = exec.Command(command, flag.Args()[1:]...)

	if pipeStdout {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		if wrapNagios {
			firstbytes = NewCapWriter(8192)
			stdout := io.TeeReader(stdout, firstbytes)
			logStream(stdout, logger, &wg)
		} else {
			logStream(stdout, logger, &wg)
		}
	} else if wrapNagios {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		firstbytes = NewCapWriter(8192)
		go io.Copy(firstbytes, stdout)
	}

	if pipeStderr {
		stderr, err := cmd.StderrPipe()
		if err != nil {
			log.Fatal(err)
		}
		logStream(stderr, logger, &wg)
	}

	if err := cmd.Start(); err != nil {
		timer.Stop()
		NotAvailable(err)
		return
	}

	if err := cmd.Wait(); err != nil {
		timer.Stop()
		Failed(err)
	} else {
		timer.Stop()
		Ok()
	}

	wg.Wait()
}
<commit_msg>use better flag package, incompatible changes<commit_after>package main

import (
	"fmt"
	"github.com/jessevdk/go-flags"
	"io"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"
)

// Avoid thundering herd problem on remote services used by this command.
// Interval will be 0, if this is not an issue.
func SpreadWait(interval time.Duration) {
	if interval > 0 {
		// Seed random generator with current process ID
		rand.Seed(int64(os.Getpid()))
		// Sleep for random amount of time within interval
		time.Sleep(time.Duration(rand.Int63n(int64(interval))))
	}
}

// Ok states that execution went well. Logs debug output and reports ok to
// monitoring.
func Ok() {
	var message string
	log.Println("OK")
	if firstbytes == nil {
		message = "OK"
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(monitorOk, message)
}

// NotAvailable states that the command could not be started successfully. It
// might not be installed or has other problems.
func NotAvailable(err error) {
	s := fmt.Sprintln("Cannot start command: ", err)
	log.Println("FATAL:", s)
	monitor(monitorUnknown, s)
}

// TimedOut states that the command took too long and reports failure to the
// monitoring.
func TimedOut() {
	s := "execution took too long"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Busy states that the command hangs and reports failure to the monitoring.
// Those tasks should be automatically killed, if it happens often.
func Busy() {
	s := "previous invocation of command still running"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Failed states that the command didn't execute successfully and reports
// failure to the monitoring. Also Logs error output.
func Failed(err error) {
	var message string
	code, s := error2exit(err)
	log.Println("FATAL:", s)
	if firstbytes == nil {
		message = s
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(code, message)
}

// LockError states that we could not get the lock.
func LockError(err error) {
	s := fmt.Sprintln("Failed to get lock: ", err)
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

var firstbytes *CapWriter

var opts struct {
	Interval         time.Duration `short:"i" long:"max-start-delay" description:"optional maximum execution start delay for command, e.g. 45s, 2m, 1h30m"`
	Timeout          time.Duration `short:"t" long:"timeout" default:"1m" description:"set execution timeout for command, e.g. 45s, 2m, 1h30m"`
	UseSyslog        bool          `short:"s" long:"use-syslog" description:"log via syslog instead of stderr"`
	WrapNagiosPlugin bool          `short:"n" long:"wrap-nagios-plugin" description:"wrap nagios plugin (pass on return codes, pass first 8KiB of stdout as message)"`
	NoPipeStderr     bool          `long:"no-stream-stderr" description:"do not stream stderr to log"`
	NoPipeStdout     bool          `long:"no-stream-stdout" description:"do not stream stdout to log"`
	MonitoringEvent  string        `short:"E" long:"monitor-event" description:"monitoring event (defaults to check_foo for /path/check_foo.sh)"`
	KillRunning      bool          `short:"k" long:"kill-running" description:"kill already running instance of command"`
}

func parseFlags() []string {
	p := flags.NewParser(&opts, flags.Default)

	// display nice usage message
	p.Usage = "[OPTIONS]... COMMAND\n\nSafely wrap execution of COMMAND in e.g. a cron job"

	args, err := p.Parse()
	if err != nil {
		// --help is not an error
		if e, ok := err.(*flags.Error); ok && e.Type == flags.ErrHelp {
			return nil
		} else {
			log.Fatal("FATAL: invalid arguments")
		}
	}

	if len(args) < 1 {
		log.Fatal("FATAL: no command to execute")
	}

	if opts.Interval >= opts.Timeout {
		log.Fatal("FATAL: interval >= timeout, no time left for actual command execution")
	}

	monitoringEvent = opts.MonitoringEvent
	return args
}

func main() {
	var cmd *exec.Cmd
	var wg sync.WaitGroup

	log.SetFlags(0)
	args := parseFlags()

	if len(args) == 0 {
		return
	}

	command := args[0]

	// default check_foo for /path/check_foo.sh
	if monitoringEvent == "" {
		monitoringEvent = filepath.Base(command)
		if ext := filepath.Ext(command); ext != "" {
			monitoringEvent = monitoringEvent[0 : len(monitoringEvent)-len("."+ext)+1]
		}
	}

	logger, err := getLogger(opts.UseSyslog)
	if err != nil {
		log.Fatal("FATAL: cannot contact syslog")
		return
	}

	loadMonitoringCommands()

	// FIXME(nightlyone) try two intervals instead of one?
	timer := time.AfterFunc(opts.Timeout, func() {
		TimedOut()
		if cmd != nil && cmd.Process != nil {
			cmd.Process.Kill()
			// FIXME(nightlyone) log the kill
		}
		os.Exit(0)
	})

	SpreadWait(opts.Interval)

	lock, err := createLock(opts.KillRunning)
	if err != nil {
		timer.Stop()
		LockError(err)
		return
	}
	defer lock.Unlock()

	cmd = exec.Command(command, args[1:]...)

	if !opts.NoPipeStdout {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		if opts.WrapNagiosPlugin {
			firstbytes = NewCapWriter(8192)
			stdout := io.TeeReader(stdout, firstbytes)
			logStream(stdout, logger, &wg)
		} else {
			logStream(stdout, logger, &wg)
		}
	} else if opts.WrapNagiosPlugin {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		firstbytes = NewCapWriter(8192)
		go io.Copy(firstbytes, stdout)
	}

	if !opts.NoPipeStderr {
		stderr, err := cmd.StderrPipe()
		if err != nil {
			log.Fatal(err)
		}
		logStream(stderr, logger, &wg)
	}

	if err := cmd.Start(); err != nil {
		timer.Stop()
		NotAvailable(err)
		return
	}

	if err := cmd.Wait(); err != nil {
		timer.Stop()
		Failed(err)
	} else {
		timer.Stop()
		Ok()
	}

	wg.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/jessevdk/go-flags"
	"io"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"
)

// Avoid thundering herd problem on remote services used by this command.
// maxDelay will be 0 if this is not an issue.
func SpreadWait(maxDelay time.Duration) {
	if maxDelay > 0 {
		// Seed random generator with current process ID
		rand.Seed(int64(os.Getpid()))
		// Sleep for random amount of time within maxDelay
		time.Sleep(time.Duration(rand.Int63n(int64(maxDelay))))
	}
}

// Ok states that execution went well. Logs debug output and reports ok to
// monitoring.
func Ok() {
	var message string
	log.Println("OK")
	if firstbytes == nil {
		message = "OK"
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(monitorOk, message)
}

// NotAvailable states that the command could not be started successfully. It
// might not be installed or has other problems.
func NotAvailable(err error) {
	s := fmt.Sprintln("Cannot start command: ", err)
	log.Println("FATAL:", s)
	monitor(monitorUnknown, s)
}

// TimedOut states that the command took too long and reports failure to the
// monitoring.
func TimedOut() {
	s := "execution took too long"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Busy states that the command hangs and reports failure to the monitoring.
// Those tasks should be automatically killed, if it happens often.
func Busy() {
	s := "previous invocation of command still running"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Failed states that the command didn't execute successfully and reports
// failure to the monitoring. Also Logs error output.
func Failed(err error) {
	var message string
	code, s := error2exit(err)
	log.Println("FATAL:", s)
	if firstbytes == nil {
		message = s
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(code, message)
}

// LockError states that we could not get the lock.
func LockError(err error) {
	s := fmt.Sprintln("Failed to get lock: ", err)
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

var firstbytes *CapWriter

var opts struct {
	MaxDelay         time.Duration `short:"d" long:"max-start-delay" description:"optional maximum execution start delay for command, e.g. 45s, 2m, 1h30m"`
	Timeout          time.Duration `short:"t" long:"timeout" default:"1m" description:"set execution timeout for command, e.g. 45s, 2m, 1h30m"`
	UseSyslog        bool          `short:"s" long:"use-syslog" description:"log via syslog instead of stderr"`
	WrapNagiosPlugin bool          `short:"n" long:"wrap-nagios-plugin" description:"wrap nagios plugin (pass on return codes, pass first 8KiB of stdout as message)"`
	NoPipeStderr     bool          `long:"no-stream-stderr" description:"do not stream stderr to log"`
	NoPipeStdout     bool          `long:"no-stream-stdout" description:"do not stream stdout to log"`
	MonitoringEvent  string        `short:"E" long:"monitor-event" description:"monitoring event (defaults to check_foo for /path/check_foo.sh)"`
	KillRunning      bool          `short:"k" long:"kill-running" description:"kill already running instance of command"`
	NoMonitoring     bool          `long:"no-monitoring" description:"wrap command without sending monitoring events"`
}

func parseFlags() []string {
	p := flags.NewParser(&opts, flags.Default)

	// display nice usage message
	p.Usage = "[OPTIONS]... COMMAND\n\nSafely wrap execution of COMMAND in e.g. a cron job"

	args, err := p.Parse()
	if err != nil {
		// --help is not an error
		if e, ok := err.(*flags.Error); ok && e.Type == flags.ErrHelp {
			return nil
		} else {
			log.Fatal("FATAL: invalid arguments")
		}
	}

	if len(args) < 1 {
		log.Fatal("FATAL: no command to execute")
	}

	if opts.MaxDelay >= opts.Timeout {
		log.Fatal("FATAL: max delay >= timeout, no time left for actual command execution")
	}

	return args
}

func main() {
	var cmd *exec.Cmd
	var wg sync.WaitGroup

	log.SetFlags(0)
	args := parseFlags()

	// exit program after --help
	if len(args) == 0 {
		return
	}

	command := args[0]

	monitoringEvent = opts.MonitoringEvent
	if monitoringEvent == "" {
		// event for command /path/check_foo.sh will be check_foo
		monitoringEvent = filepath.Base(command)
		if ext := filepath.Ext(command); ext != "" {
			monitoringEvent = monitoringEvent[0 : len(monitoringEvent)-len(ext)]
		}
	}

	logger, err := getLogger(opts.UseSyslog)
	if err != nil {
		log.Fatal("FATAL: cannot contact syslog")
		return
	}

	loadMonitoringCommands()

	// FIXME(nightlyone) try two intervals instead of one?
	timer := time.AfterFunc(opts.Timeout, func() {
		TimedOut()
		if cmd != nil && cmd.Process != nil {
			cmd.Process.Kill()
			// FIXME(nightlyone) log the kill
		}
		os.Exit(0)
	})

	SpreadWait(opts.MaxDelay)

	lock, err := createLock(opts.KillRunning)
	if err != nil {
		timer.Stop()
		LockError(err)
		return
	}
	defer lock.Unlock()

	cmd = exec.Command(command, args[1:]...)

	if !opts.NoPipeStdout {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		if opts.WrapNagiosPlugin {
			firstbytes = NewCapWriter(8192)
			stdout := io.TeeReader(stdout, firstbytes)
			logStream(stdout, logger, &wg)
		} else {
			logStream(stdout, logger, &wg)
		}
	} else if opts.WrapNagiosPlugin {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		firstbytes = NewCapWriter(8192)
		logStream(stdout, firstbytes, &wg)
	}

	if !opts.NoPipeStderr {
		stderr, err := cmd.StderrPipe()
		if err != nil {
			log.Fatal(err)
		}
		logStream(stderr, logger, &wg)
	}

	if err := cmd.Start(); err != nil {
		timer.Stop()
		NotAvailable(err)
		return
	}

	if err := cmd.Wait(); err != nil {
		timer.Stop()
		Failed(err)
	} else {
		timer.Stop()
		Ok()
	}

	wg.Wait()
}
<commit_msg>Make sure log streams are done before command returns<commit_after>package main

import (
	"fmt"
	"github.com/jessevdk/go-flags"
	"io"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"
)

// Avoid thundering herd problem on remote services used by this command.
// maxDelay will be 0 if this is not an issue.
func SpreadWait(maxDelay time.Duration) {
	if maxDelay > 0 {
		// Seed random generator with current process ID
		rand.Seed(int64(os.Getpid()))
		// Sleep for random amount of time within maxDelay
		time.Sleep(time.Duration(rand.Int63n(int64(maxDelay))))
	}
}

// Ok states that execution went well. Logs debug output and reports ok to
// monitoring.
func Ok() {
	var message string
	log.Println("OK")
	if firstbytes == nil {
		message = "OK"
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(monitorOk, message)
}

// NotAvailable states that the command could not be started successfully. It
// might not be installed or has other problems.
func NotAvailable(err error) {
	s := fmt.Sprintln("Cannot start command: ", err)
	log.Println("FATAL:", s)
	monitor(monitorUnknown, s)
}

// TimedOut states that the command took too long and reports failure to the
// monitoring.
func TimedOut() {
	s := "execution took too long"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Busy states that the command hangs and reports failure to the monitoring.
// Those tasks should be automatically killed, if it happens often.
func Busy() {
	s := "previous invocation of command still running"
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

// Failed states that the command didn't execute successfully and reports
// failure to the monitoring. Also Logs error output.
func Failed(err error) {
	var message string
	code, s := error2exit(err)
	log.Println("FATAL:", s)
	if firstbytes == nil {
		message = s
	} else {
		message = string(firstbytes.Bytes())
	}
	monitor(code, message)
}

// LockError states that we could not get the lock.
func LockError(err error) {
	s := fmt.Sprintln("Failed to get lock: ", err)
	log.Println("FATAL:", s)
	monitor(monitorCritical, s)
}

var firstbytes *CapWriter

var opts struct {
	MaxDelay         time.Duration `short:"d" long:"max-start-delay" description:"optional maximum execution start delay for command, e.g. 45s, 2m, 1h30m"`
	Timeout          time.Duration `short:"t" long:"timeout" default:"1m" description:"set execution timeout for command, e.g. 45s, 2m, 1h30m"`
	UseSyslog        bool          `short:"s" long:"use-syslog" description:"log via syslog instead of stderr"`
	WrapNagiosPlugin bool          `short:"n" long:"wrap-nagios-plugin" description:"wrap nagios plugin (pass on return codes, pass first 8KiB of stdout as message)"`
	NoPipeStderr     bool          `long:"no-stream-stderr" description:"do not stream stderr to log"`
	NoPipeStdout     bool          `long:"no-stream-stdout" description:"do not stream stdout to log"`
	MonitoringEvent  string        `short:"E" long:"monitor-event" description:"monitoring event (defaults to check_foo for /path/check_foo.sh)"`
	KillRunning      bool          `short:"k" long:"kill-running" description:"kill already running instance of command"`
	NoMonitoring     bool          `long:"no-monitoring" description:"wrap command without sending monitoring events"`
}

func parseFlags() []string {
	p := flags.NewParser(&opts, flags.Default)

	// display nice usage message
	p.Usage = "[OPTIONS]... COMMAND\n\nSafely wrap execution of COMMAND in e.g. a cron job"

	args, err := p.Parse()
	if err != nil {
		// --help is not an error
		if e, ok := err.(*flags.Error); ok && e.Type == flags.ErrHelp {
			return nil
		} else {
			log.Fatal("FATAL: invalid arguments")
		}
	}

	if len(args) < 1 {
		log.Fatal("FATAL: no command to execute")
	}

	if opts.MaxDelay >= opts.Timeout {
		log.Fatal("FATAL: max delay >= timeout, no time left for actual command execution")
	}

	return args
}

func main() {
	var cmd *exec.Cmd
	var wg sync.WaitGroup

	log.SetFlags(0)
	args := parseFlags()

	// exit program after --help
	if len(args) == 0 {
		return
	}

	command := args[0]

	monitoringEvent = opts.MonitoringEvent
	if monitoringEvent == "" {
		// event for command /path/check_foo.sh will be check_foo
		monitoringEvent = filepath.Base(command)
		if ext := filepath.Ext(command); ext != "" {
			monitoringEvent = monitoringEvent[0 : len(monitoringEvent)-len(ext)]
		}
	}

	logger, err := getLogger(opts.UseSyslog)
	if err != nil {
		log.Fatal("FATAL: cannot contact syslog")
		return
	}

	loadMonitoringCommands()

	// FIXME(nightlyone) try two intervals instead of one?
	timer := time.AfterFunc(opts.Timeout, func() {
		TimedOut()
		if cmd != nil && cmd.Process != nil {
			cmd.Process.Kill()
			// FIXME(nightlyone) log the kill
		}
		os.Exit(0)
	})

	SpreadWait(opts.MaxDelay)

	lock, err := createLock(opts.KillRunning)
	if err != nil {
		timer.Stop()
		LockError(err)
		return
	}
	defer lock.Unlock()

	cmd = exec.Command(command, args[1:]...)

	if !opts.NoPipeStdout {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		if opts.WrapNagiosPlugin {
			firstbytes = NewCapWriter(8192)
			stdout := io.TeeReader(stdout, firstbytes)
			logStream(stdout, logger, &wg)
		} else {
			logStream(stdout, logger, &wg)
		}
	} else if opts.WrapNagiosPlugin {
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		firstbytes = NewCapWriter(8192)
		logStream(stdout, firstbytes, &wg)
	}

	if !opts.NoPipeStderr {
		stderr, err := cmd.StderrPipe()
		if err != nil {
			log.Fatal(err)
		}
		logStream(stderr, logger, &wg)
	}

	if err := cmd.Start(); err != nil {
		timer.Stop()
		NotAvailable(err)
		return
	}

	wg.Wait()

	if err := cmd.Wait(); err != nil {
		timer.Stop()
		Failed(err)
	} else {
		timer.Stop()
		Ok()
	}
}
<|endoftext|>


<commit_before>package run

import (
	"errors"
	"fmt"
	"github.com/ohsu-comp-bio/funnel/cmd/client"
	"github.com/ohsu-comp-bio/funnel/logger"
	"github.com/ohsu-comp-bio/funnel/proto/tes"
	"github.com/spf13/cobra"
	"os"
	"regexp"
	"strings"
)

var log = logger.New("run")

// TODO figure out a nice default for name
var name = "Funnel Run Task"
var workdir = "/opt/funnel"
var server = "http://localhost:9090"
var image, project, description, varsFile string
var printTask, preemptible, wait bool
var cliInputs, cliInputDirs, cliOutputs, cliOutputDirs, cliEnvVars, tags, volumes []string

// TODO validate resources
var cpu int
var ram float64
var disk float64
var zones []string

// TODO with input contents, script could be loaded from file
// TODO load vars from file
// TODO is there a case for glob inputs?

var example = `
    funnel run
    'bowtie2 -f $factor -x $other -p1 $pair1 -p2 $pair2 -o $alignments'
    --image opengenomics/bowtie2:latest
    --name 'Bowtie2 test'
    --description 'Testings an example of using 'funnel run' for a bowtie2 command'
    --in pair1=file://~/pair1.fastq
    --in pair2=file://~/pair2.fastq
    --out alignments=gs://bkt/bowtie/alignments.bam
    --env factor=5
    --vol /tmp
    --vol /opt
    --cpu 8
    --ram 32
    --disk 100
`

// Cmd represents the run command
var Cmd = &cobra.Command{
	Use:     "run [flags] --image IMAGE CMD",
	Short:   "Run a task.\n",
	Long:    ``,
	Example: example,
	Run:     run,
}

func init() {
	f := Cmd.Flags()
	f.StringVarP(&name, "name", "n", name, "Task name")
	f.StringVar(&description, "description", description, "Task description")
	f.StringVar(&project, "project", project, "Project")
	f.StringVarP(&workdir, "workdir", "w", workdir, "Set the containter working directory")
	f.StringVarP(&image, "container", "c", image, "Specify the containter image")
	f.StringSliceVarP(&cliInputs, "in", "i", cliInputs, "A key-value map of input files")
	f.StringSliceVarP(&cliInputDirs, "in-dir", "I", cliInputDirs, "A key-value map of input directories")
	f.StringSliceVarP(&cliOutputs, "out", "o", cliOutputs, "A key-value map of output files")
	f.StringSliceVarP(&cliOutputDirs, "out-dir", "O", cliOutputDirs, "A key-value map of output directories")
	f.StringSliceVarP(&cliEnvVars, "env", "e", cliEnvVars, "A key-value map of enviromental variables")
	f.StringSliceVar(&volumes, "vol", volumes, "Volumes to be defined on the container")
	f.StringSliceVar(&tags, "tag", tags, "A key-value map of arbitrary tags")
	f.IntVar(&cpu, "cpu", cpu, "Number of CPUs requested")
	f.Float64Var(&ram, "ram", ram, "Amount of RAM requested (in GB)")
	f.Float64Var(&disk, "disk", disk, "Amount of disk space requested (in GB)")
	f.BoolVar(&preemptible, "preemptible", preemptible, "Allow task to be scheduled on preemptible workers")
	f.StringSliceVar(&zones, "zone", zones, "Require task be scheduled in certain zones")
	// TODO
	//f.StringVar(&varsFile, "vars-file", varsFile, "Read vars from file")
	//f.StringVar(&cmdFile, "cmd-file", cmdFile, "Read cmd template from file")
	f.StringVarP(&server, "server", "S", server, "Address of Funnel server")
	f.BoolVarP(&printTask, "print", "p", printTask, "Print the task, instead of running it")
	f.BoolVar(&wait, "wait", wait, "Wait for the task to finish before exiting")
}

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

	if len(args) < 1 {
		fmt.Printf("ERROR - You must specify a command to run\n\n")
		cmd.Usage()
		os.Exit(2)
		return
	}

	if len(args) > 1 {
		fmt.Printf("ERROR - extra arg(s) %s\n\n", args[1:])
		fmt.Printf("--in, --out and --env args should have the form 'KEY=VALUE' not 'KEY VALUE'\n\n")
		cmd.Usage()
		os.Exit(2)
		return
	}

	if image == "" {
		fmt.Printf("ERROR - You must specify an image to run your command in\n\n")
		cmd.Usage()
		os.Exit(2)
		return
	}

	// TODO validate cmd
	rawcmd := args[0]
	executorCmd := []string{"bash", "-c", rawcmd}
	log.Debug("Command:", "cmd", executorCmd)

	// Get template variables from the command line.
	// TODO variables file.
	inputFileMap, err := parseCliVars(cliInputs)
	checkErr(err)
	inputDirMap, err := parseCliVars(cliInputDirs)
	checkErr(err)
	outputFileMap, err := parseCliVars(cliOutputs)
	checkErr(err)
	outputDirMap, err := parseCliVars(cliOutputDirs)
	checkErr(err)
	envVarMap, err := parseCliVars(cliEnvVars)
	checkErr(err)
	tagsMap, err := parseCliVars(tags)
	checkErr(err)

	// check for key collisions
	err = compareKeys(inputFileMap, inputDirMap, outputFileMap, outputDirMap, envVarMap)
	checkErr(err)

	// Create map of enviromental variables to be passed to the executor
	inputEnvVars, err := fileMapToEnvVars(inputFileMap, "/opt/funnel/inputs/")
	checkErr(err)
	inputDirEnvVars, err := fileMapToEnvVars(inputDirMap, "/opt/funnel/inputs/")
	checkErr(err)
	outputEnvVars, err := fileMapToEnvVars(outputFileMap, "/opt/funnel/outputs/")
	checkErr(err)
	outputDirEnvVars, err := fileMapToEnvVars(outputDirMap, "/opt/funnel/outputs/")
	checkErr(err)
	environ, err := mergeVars(inputEnvVars, inputDirEnvVars, outputEnvVars, outputDirEnvVars, envVarMap)
	checkErr(err)

	// Build task input parameters
	inputs, err := createTaskParams(inputFileMap, "/opt/funnel/inputs/", tes.FileType_FILE)
	checkErr(err)
	inputDirs, err := createTaskParams(inputDirMap, "/opt/funnel/inputs/", tes.FileType_DIRECTORY)
	checkErr(err)
	inputs = append(inputs, inputDirs...)

	// Build task output parameters
	outputs, err := createTaskParams(outputFileMap, "/opt/funnel/outputs/", tes.FileType_FILE)
	checkErr(err)
	outputDirs, err := createTaskParams(outputFileMap, "/opt/funnel/outputs/", tes.FileType_DIRECTORY)
	checkErr(err)
	outputs = append(outputs, outputDirs...)

	// Build the task message
	task := &tes.Task{
		Name:        name,
		Project:     project,
		Description: description,
		Inputs:      inputs,
		Outputs:     outputs,
		Resources: &tes.Resources{
			CpuCores:    uint32(cpu),
			RamGb:       ram,
			SizeGb:      disk,
			Zones:       zones,
			Preemptible: preemptible,
		},
		Executors: []*tes.Executor{
			{
				ImageName: image,
				Cmd:       executorCmd,
				Environ:   environ,
				Workdir:   workdir,
				Stdin:     "",
				Stdout:    "/opt/funnel/outputs/stdout",
				Stderr:    "/opt/funnel/outputs/stderr",
				// TODO no ports
				Ports: nil,
			},
		},
		Volumes: volumes,
		Tags:    tagsMap,
	}
	log.Debug("Task", "taskmsg", task)

	cli := client.NewClient(server)
	// Marshal message to JSON
	taskJSON, merr := cli.Marshaler.MarshalToString(task)
	checkErr(merr)

	if printTask {
		fmt.Println(taskJSON)
		return
	}

	resp, rerr := cli.CreateTask([]byte(taskJSON))
	checkErr(rerr)

	taskID := resp.Id
	fmt.Println(taskID)

	if wait {
		werr := cli.WaitForTask(taskID)
		checkErr(werr)
	}
}

func stripStoragePrefix(url string) (string, error) {
	re := regexp.MustCompile("[a-z3]+://")
	if !re.MatchString(url) {
		err := errors.New("File paths must be prefixed with one of:\n file://\n gs://\n s3://")
		return "", err
	}
	path := re.ReplaceAllString(url, "")
	return strings.TrimPrefix(path, "/"), nil
}

func checkErr(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(3)
	}
}
<commit_msg>remove debug messages<commit_after>package run

import (
	"errors"
	"fmt"
	"github.com/ohsu-comp-bio/funnel/cmd/client"
	"github.com/ohsu-comp-bio/funnel/logger"
	"github.com/ohsu-comp-bio/funnel/proto/tes"
	"github.com/spf13/cobra"
	"os"
	"regexp"
	"strings"
)

var log = logger.New("run")

// TODO figure out a nice default for name
var name = "Funnel Run Task"
var workdir = "/opt/funnel"
var server = "http://localhost:9090"
var image, project, description, varsFile string
var printTask, preemptible, wait bool
var cliInputs, cliInputDirs, cliOutputs, cliOutputDirs, cliEnvVars, tags, volumes []string

// TODO validate resources
var cpu int
var ram float64
var disk float64
var zones []string

// TODO with input contents, script could be loaded from file
// TODO load vars from file
// TODO is there a case for glob inputs?

var example = `
    funnel run
    'bowtie2 -f $factor -x $other -p1 $pair1 -p2 $pair2 -o $alignments'
    --image opengenomics/bowtie2:latest
    --name 'Bowtie2 test'
    --description 'Testings an example of using 'funnel run' for a bowtie2 command'
    --in pair1=file://~/pair1.fastq
    --in pair2=file://~/pair2.fastq
    --out alignments=gs://bkt/bowtie/alignments.bam
    --env factor=5
    --vol /tmp
    --vol /opt
    --cpu 8
    --ram 32
    --disk 100
`

// Cmd represents the run command
var Cmd = &cobra.Command{
	Use:     "run [flags] --image IMAGE CMD",
	Short:   "Run a task.\n",
	Long:    ``,
	Example: example,
	Run:     run,
}

func init() {
	f := Cmd.Flags()
	f.StringVarP(&name, "name", "n", name, "Task name")
	f.StringVar(&description, "description", description, "Task description")
	f.StringVar(&project, "project", project, "Project")
	f.StringVarP(&workdir, "workdir", "w", workdir, "Set the containter working directory")
	f.StringVarP(&image, "container", "c", image, "Specify the containter image")
	f.StringSliceVarP(&cliInputs, "in", "i", cliInputs, "A key-value map of input files")
	f.StringSliceVarP(&cliInputDirs, "in-dir", "I", cliInputDirs, "A key-value map of input directories")
	f.StringSliceVarP(&cliOutputs, "out", "o", cliOutputs, "A key-value map of output files")
	f.StringSliceVarP(&cliOutputDirs, "out-dir", "O", cliOutputDirs, "A key-value map of output directories")
	f.StringSliceVarP(&cliEnvVars, "env", "e", cliEnvVars, "A key-value map of enviromental variables")
	f.StringSliceVar(&volumes, "vol", volumes, "Volumes to be defined on the container")
	f.StringSliceVar(&tags, "tag", tags, "A key-value map of arbitrary tags")
	f.IntVar(&cpu, "cpu", cpu, "Number of CPUs requested")
	f.Float64Var(&ram, "ram", ram, "Amount of RAM requested (in GB)")
	f.Float64Var(&disk, "disk", disk, "Amount of disk space requested (in GB)")
	f.BoolVar(&preemptible, "preemptible", preemptible, "Allow task to be scheduled on preemptible workers")
	f.StringSliceVar(&zones, "zone", zones, "Require task be scheduled in certain zones")
	// TODO
	//f.StringVar(&varsFile, "vars-file", varsFile, "Read vars from file")
	//f.StringVar(&cmdFile, "cmd-file", cmdFile, "Read cmd template from file")
	f.StringVarP(&server, "server", "S", server, "Address of Funnel server")
	f.BoolVarP(&printTask, "print", "p", printTask, "Print the task, instead of running it")
	f.BoolVar(&wait, "wait", wait, "Wait for the task to finish before exiting")
}

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

	if len(args) < 1 {
		fmt.Printf("ERROR - You must specify a command to run\n\n")
		cmd.Usage()
		os.Exit(2)
		return
	}

	if len(args) > 1 {
		fmt.Printf("ERROR - extra arg(s) %s\n\n", args[1:])
		fmt.Printf("--in, --out and --env args should have the form 'KEY=VALUE' not 'KEY VALUE'\n\n")
		cmd.Usage()
		os.Exit(2)
		return
	}

	if image == "" {
		fmt.Printf("ERROR - You must specify an image to run your command in\n\n")
		cmd.Usage()
		os.Exit(2)
		return
	}

	// TODO validate cmd
	rawcmd := args[0]
	executorCmd := []string{"bash", "-c", rawcmd}

	// Get template variables from the command line.
	// TODO variables file.
	inputFileMap, err := parseCliVars(cliInputs)
	checkErr(err)
	inputDirMap, err := parseCliVars(cliInputDirs)
	checkErr(err)
	outputFileMap, err := parseCliVars(cliOutputs)
	checkErr(err)
	outputDirMap, err := parseCliVars(cliOutputDirs)
	checkErr(err)
	envVarMap, err := parseCliVars(cliEnvVars)
	checkErr(err)
	tagsMap, err := parseCliVars(tags)
	checkErr(err)

	// check for key collisions
	err = compareKeys(inputFileMap, inputDirMap, outputFileMap, outputDirMap, envVarMap)
	checkErr(err)

	// Create map of enviromental variables to be passed to the executor
	inputEnvVars, err := fileMapToEnvVars(inputFileMap, "/opt/funnel/inputs/")
	checkErr(err)
	inputDirEnvVars, err := fileMapToEnvVars(inputDirMap, "/opt/funnel/inputs/")
	checkErr(err)
	outputEnvVars, err := fileMapToEnvVars(outputFileMap, "/opt/funnel/outputs/")
	checkErr(err)
	outputDirEnvVars, err := fileMapToEnvVars(outputDirMap, "/opt/funnel/outputs/")
	checkErr(err)
	environ, err := mergeVars(inputEnvVars, inputDirEnvVars, outputEnvVars, outputDirEnvVars, envVarMap)
	checkErr(err)

	// Build task input parameters
	inputs, err := createTaskParams(inputFileMap, "/opt/funnel/inputs/", tes.FileType_FILE)
	checkErr(err)
	inputDirs, err := createTaskParams(inputDirMap, "/opt/funnel/inputs/", tes.FileType_DIRECTORY)
	checkErr(err)
	inputs = append(inputs, inputDirs...)

	// Build task output parameters
	outputs, err := createTaskParams(outputFileMap, "/opt/funnel/outputs/", tes.FileType_FILE)
	checkErr(err)
	outputDirs, err := createTaskParams(outputFileMap, "/opt/funnel/outputs/", tes.FileType_DIRECTORY)
	checkErr(err)
	outputs = append(outputs, outputDirs...)

	// Build the task message
	task := &tes.Task{
		Name:        name,
		Project:     project,
		Description: description,
		Inputs:      inputs,
		Outputs:     outputs,
		Resources: &tes.Resources{
			CpuCores:    uint32(cpu),
			RamGb:       ram,
			SizeGb:      disk,
			Zones:       zones,
			Preemptible: preemptible,
		},
		Executors: []*tes.Executor{
			{
				ImageName: image,
				Cmd:       executorCmd,
				Environ:   environ,
				Workdir:   workdir,
				Stdin:     "",
				Stdout:    "/opt/funnel/outputs/stdout",
				Stderr:    "/opt/funnel/outputs/stderr",
				// TODO no ports
				Ports: nil,
			},
		},
		Volumes: volumes,
		Tags:    tagsMap,
	}

	cli := client.NewClient(server)
	// Marshal message to JSON
	taskJSON, merr := cli.Marshaler.MarshalToString(task)
	checkErr(merr)

	if printTask {
		fmt.Println(taskJSON)
		return
	}

	resp, rerr := cli.CreateTask([]byte(taskJSON))
	checkErr(rerr)

	taskID := resp.Id
	fmt.Println(taskID)

	if wait {
		werr := cli.WaitForTask(taskID)
		checkErr(werr)
	}
}

func stripStoragePrefix(url string) (string, error) {
	re := regexp.MustCompile("[a-z3]+://")
	if !re.MatchString(url) {
		err := errors.New("File paths must be prefixed with one of:\n file://\n gs://\n s3://")
		return "", err
	}
	path := re.ReplaceAllString(url, "")
	return strings.TrimPrefix(path, "/"), nil
}

func checkErr(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(3)
	}
}
<|endoftext|>


<commit_before>package main

const Version string = "0.2.0"
<commit_msg>Release 0.2.1<commit_after>package main

const Version string = "0.2.1"
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
)

var version string

func init() {
	version = "0.0.2"
	RootCmd.AddCommand(versionCmd)
}

var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "Prints the version number",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Printf("pekka v%s\n", version)
	},
}
<commit_msg>corrected the version number<commit_after>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
)

var version string

func init() {
	version = "0.2.0"
	RootCmd.AddCommand(versionCmd)
}

var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "Prints the version number",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Printf("pekka v%s\n", version)
	},
}
<|endoftext|>


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

package cmd

import (
	"github.com/spf13/cobra"
)

// versionCmd represents the version command
var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "Prints out the version",
	Long: ``,
	Run: func(cmd *cobra.Command, args []string) {
    fmt.Println("devenv version v1.2.0")
	},
}

func init() {
	RootCmd.AddCommand(versionCmd)
}
<commit_msg>Importing fmt<commit_after>// Copyright © 2017 Sascha Andres <sascha.andres@outlook.com>
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
)

// versionCmd represents the version command
var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "Prints out the version",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println("devenv version v1.2.0")
	},
}

func init() {
	RootCmd.AddCommand(versionCmd)
}
<|endoftext|>


<commit_before>package codec

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"reflect"

	"code.google.com/p/gogoprotobuf/proto"
)

const sizeOfUint8 = 1

var (
	ErrMessageAlreadyRegistered = errors.New("Message already registered")
	ErrMessageNotRegistered     = errors.New("Message not registered")
	ErrCannotWriteMessage       = errors.New("Cannot write message")
)

// CodecInterface describes the codec interface,
// which encodes/decodes protobuf messages from/to
// an io.Reader/Writer
type CodecInterface interface {
	// Register registers a message so that the
	// codec can identify the message when reading
	// the TCP connection.
	Register(msg proto.Message) error
	// Encode encodes a message to bytes and
	// writes it to the io.Writer
	Encode(msg proto.Message, w io.Writer) error
	// Decode reads bytes from the io.Reader
	// and decodes it to a message.
	Decode(r io.Reader) (proto.Message, error)
}

// ProtobufCodec implements the codec interface.
type ProtobufCodec struct {
	// registeredMessages is a map from message indices
	// to message types. The indices increase monotonically.
	registeredMessages map[uint8]reflect.Type
	// messageIndices is a map from message types
	// to message indices.
	messageIndices map[reflect.Type]uint8
}

// NewProtobufCodec creates and returns a ProtobufCodec.
func NewProtobufCodec() *ProtobufCodec {
	return &ProtobufCodec{
		registeredMessages: make(map[uint8]reflect.Type),
		messageIndices:     make(map[reflect.Type]uint8),
	}
}

// Register registers a message. Note this is not concurrent-safe.
func (pc *ProtobufCodec) Register(msg proto.Message) error {
	mtype := reflect.TypeOf(msg)
	if _, existed := pc.messageIndices[mtype]; existed {
		return ErrMessageAlreadyRegistered
	}
	index := uint8(len(pc.messageIndices))
	pc.messageIndices[mtype] = index
	pc.registeredMessages[index] = mtype
	return nil
}

// Encode encodes a message to bytes and writes it to the io.Writer.
func (pc *ProtobufCodec) Encode(msg proto.Message, w io.Writer) error {
	index, existed := pc.messageIndices[reflect.TypeOf(msg)]
	if !existed {
		return ErrMessageNotRegistered
	}
	b, err := proto.Marshal(msg)
	if err != nil {
		return err
	}
	bw := bufio.NewWriter(w)

	// 1. Write the length.
	if err := binary.Write(bw, binary.LittleEndian, int32(len(b))); err != nil {
		return err
	}
	// 2. Write the type.
	if err := binary.Write(bw, binary.LittleEndian, index); err != nil {
		return err
	}
	// 3. Write the bytes.
	n, err := bw.Write(b)
	if err != nil {
		return err
	}
	if n != len(b) {
		return ErrCannotWriteMessage
	}
	return bw.Flush()
}

// Decode reads bytes from an io.Reader and decode it to a message.
func (pc *ProtobufCodec) Decode(r io.Reader) (proto.Message, error) {
	var length int32
	var index uint8

	// 1. Read the length.
	if err := binary.Read(r, binary.LittleEndian, &length); err != nil {
		return nil, err
	}
	br := bufio.NewReaderSize(r, int(length)+sizeOfUint8)
	// 1. Read the type.
	if err := binary.Read(br, binary.LittleEndian, &index); err != nil {
		return nil, err
	}
	// 3. Read the bytes.
	b := make([]byte, length)
	_, err := io.ReadFull(br, b)
	if err != nil {
		return nil, err
	}
	// Decode.
	mtype, existed := pc.registeredMessages[index]
	if !existed {
		return nil, ErrMessageNotRegistered
	}
	msg := reflect.New(mtype.Elem()).Interface().(proto.Message)
	if err := proto.Unmarshal(b, msg); err != nil {
		return nil, err
	}
	return msg, nil
}
<commit_msg>Fixed codec. Using bytes.Buffer instead of bufio.<commit_after>package codec

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"reflect"

	"code.google.com/p/gogoprotobuf/proto"
)

const (
	sizeOfUint8 = 1
	sizeOfInt32 = 4
)

var (
	ErrMessageAlreadyRegistered = errors.New("Message already registered")
	ErrMessageNotRegistered     = errors.New("Message not registered")
	ErrCannotWriteMessage       = errors.New("Cannot write message")
)

// CodecInterface describes the codec interface,
// which encodes/decodes protobuf messages from/to
// an io.Reader/Writer
type CodecInterface interface {
	// Register registers a message so that the
	// codec can identify the message when reading
	// the TCP connection.
	Register(msg proto.Message) error
	// Encode encodes a message to bytes and
	// writes it to the io.Writer
	Encode(msg proto.Message, w io.Writer) error
	// Decode reads bytes from the io.Reader
	// and decodes it to a message.
	Decode(r io.Reader) (proto.Message, error)
}

// ProtobufCodec implements the codec interface.
type ProtobufCodec struct {
	// registeredMessages is a map from message indices
	// to message types. The indices increase monotonically.
	registeredMessages map[uint8]reflect.Type
	// messageIndices is a map from message types
	// to message indices.
	messageIndices map[reflect.Type]uint8
}

// NewProtobufCodec creates and returns a ProtobufCodec.
func NewProtobufCodec() *ProtobufCodec {
	return &ProtobufCodec{
		registeredMessages: make(map[uint8]reflect.Type),
		messageIndices:     make(map[reflect.Type]uint8),
	}
}

// Register registers a message. Note this is not concurrent-safe.
func (pc *ProtobufCodec) Register(msg proto.Message) error {
	mtype := reflect.TypeOf(msg)
	if _, existed := pc.messageIndices[mtype]; existed {
		return ErrMessageAlreadyRegistered
	}
	index := uint8(len(pc.messageIndices))
	pc.messageIndices[mtype] = index
	pc.registeredMessages[index] = mtype
	return nil
}

// Encode encodes a message to bytes and writes it to the io.Writer.
func (pc *ProtobufCodec) Encode(msg proto.Message, w io.Writer) error {
	index, existed := pc.messageIndices[reflect.TypeOf(msg)]
	if !existed {
		return ErrMessageNotRegistered
	}
	buf := new(bytes.Buffer)

	b, err := proto.Marshal(msg)
	if err != nil {
		return err
	}
	// Write the length.
	if err := binary.Write(buf, binary.LittleEndian, int32(len(b)+sizeOfUint8)); err != nil {
		return err
	}
	// Write the type.
	if err := binary.Write(buf, binary.LittleEndian, index); err != nil {
		return err
	}
	// Write the bytes.
	buf.Write(b)
	if _, err = buf.WriteTo(w); err != nil {
		return err
	}
	return nil
}

// Decode reads bytes from an io.Reader and decode it to a message.
func (pc *ProtobufCodec) Decode(r io.Reader) (proto.Message, error) {
	var length int32

	// Read the length.
	if err := binary.Read(r, binary.LittleEndian, &length); err != nil {
		return nil, err
	}
	b := make([]byte, length)
	// Read the type and bytes.
	if _, err := io.ReadFull(r, b); err != nil {
		return nil, err
	}
	// Get the index.
	index := uint8(b[0])
	// Decode.
	mtype, existed := pc.registeredMessages[index]
	if !existed {
		return nil, ErrMessageNotRegistered
	}
	msg := reflect.New(mtype.Elem()).Interface().(proto.Message)
	if err := proto.Unmarshal(b[1:], msg); err != nil {
		return nil, err
	}
	return msg, nil
}
<|endoftext|>


<commit_before>package comm

import (
    "bufio"
    "bytes"
    "errors"
    "fmt"
    ev "github.com/hhkbp2/rafted/event"
    logging "github.com/hhkbp2/rafted/logging"
    ps "github.com/hhkbp2/rafted/persist"
    "github.com/ugorji/go/codec"
    "io"
    "net"
    "sync"
)

const (
    DefaultTransportBufferSize = 16
)

type TransportChunk struct {
    Data     []byte
    SourceCh chan []byte
}

type MemoryServerTransport struct {
    addr       net.Addr
    ConsumeCh  chan *TransportChunk
    ResponseCh chan []byte
    register   *MemoryTransportRegister
}

func NewMemoryServerTransport(
    addr net.Addr, register *MemoryTransportRegister) *MemoryServerTransport {

    return &MemoryServerTransport{
        addr:      addr,
        ConsumeCh: make(chan *TransportChunk, DefaultTransportBufferSize),
        register:  register,
    }
}

func (self *MemoryServerTransport) Open() error {
    self.register.Register(self.addr.String(), self)
    return nil
}

func (self *MemoryServerTransport) ReadChunk() (*TransportChunk, error) {
    chunk, ok := <-self.ConsumeCh
    if ok {
        return chunk, nil
    }
    return nil, io.EOF
}

func (self *MemoryServerTransport) WriteChunk(chunk *TransportChunk) error {
    self.ConsumeCh <- chunk
    return nil
}

func BytesCopy(dest []byte, src []byte) int {
    i := 0
    for ; (i < len(dest)) && (i < len(src)); i++ {
        dest[i] = src[i]
    }
    return i
}

func (self *MemoryServerTransport) Read(p []byte) (int, error) {
    chunk, err := self.ReadChunk()
    if err != nil {
        return 0, err
    }
    n := BytesCopy(p, chunk.Data)
    self.ResponseCh = chunk.SourceCh
    return n, nil
}

func (self *MemoryServerTransport) Write(p []byte) (int, error) {
    if self.ResponseCh != nil {
        self.ResponseCh <- p
        return len(p), nil
    }
    return 0, errors.New("don't know where to response")
}

func (self *MemoryServerTransport) Close() error {
    close(self.ConsumeCh)
    if self.ResponseCh != nil {
        close(self.ResponseCh)
    }
    return self.register.Unregister(self.addr.String())
}

func (self *MemoryServerTransport) Addr() net.Addr {
    return self.addr
}

type MemoryTransportRegister struct {
    sync.RWMutex
    transports map[string]*MemoryServerTransport
}

func NewMemoryTransportRegister() *MemoryTransportRegister {
    return &MemoryTransportRegister{
        transports: make(map[string]*MemoryServerTransport),
    }
}

func (self *MemoryTransportRegister) Register(
    id string, transport *MemoryServerTransport) {

    self.Lock()
    defer self.Unlock()
    self.transports[id] = transport
}

func (self *MemoryTransportRegister) Unregister(id string) error {
    self.Lock()
    defer self.Unlock()
    if _, ok := self.transports[id]; ok {
        delete(self.transports, id)
        return nil
    }
    return errors.New(fmt.Sprintf("no transport for id: %v", id))
}

func (self *MemoryTransportRegister) Reset() error {
    self.Lock()
    defer self.Unlock()
    self.transports = make(map[string]*MemoryServerTransport)
    return nil
}

func (self *MemoryTransportRegister) Get(
    id string) (transport *MemoryServerTransport, ok bool) {

    self.RLock()
    defer self.RUnlock()
    transport, ok = self.transports[id]
    return transport, ok
}

func (self *MemoryTransportRegister) String() string {
    self.RLock()
    defer self.RUnlock()
    var buf bytes.Buffer
    for k, _ := range self.transports {
        buf.WriteString(k + "\n")
    }
    return buf.String()
}

type MemoryTransport struct {
    addr      net.Addr
    consumeCh chan []byte
    register  *MemoryTransportRegister
    peer      *MemoryServerTransport
}

func NewMemoryTransport(addr net.Addr, register *MemoryTransportRegister) *MemoryTransport {
    return &MemoryTransport{
        addr:      addr,
        consumeCh: make(chan []byte, DefaultTransportBufferSize),
        register:  register,
    }
}

func (self *MemoryTransport) PeerAddr() net.Addr {
    return self.addr
}

func (self *MemoryTransport) Open() error {
    if transport, ok := self.register.Get(self.addr.String()); ok {
        self.peer = transport
        return nil
    }
    return errors.New(
        fmt.Sprintf("no server transport for id: %v", self.addr.String()))
}

func (self *MemoryTransport) Close() error {
    // empty body
    return nil
}

func (self *MemoryTransport) Read(b []byte) (int, error) {
    data := <-self.consumeCh
    n := BytesCopy(b, data)
    return n, nil
}

func (self *MemoryTransport) Write(b []byte) (int, error) {
    chunk := &TransportChunk{
        Data:     b,
        SourceCh: self.consumeCh,
    }
    if err := self.peer.WriteChunk(chunk); err != nil {
        return 0, err
    }
    return len(b), nil
}

type MemoryConnection struct {
    *MemoryTransport
    reader  *bufio.Reader
    writer  *bufio.Writer
    encoder Encoder
    decoder Decoder
}

func NewMemoryConnection(
    addr net.Addr, register *MemoryTransportRegister) *MemoryConnection {

    conn := &MemoryConnection{
        MemoryTransport: NewMemoryTransport(addr, register),
    }
    conn.reader = bufio.NewReader(conn.MemoryTransport)
    conn.writer = bufio.NewWriter(conn.MemoryTransport)
    conn.decoder = codec.NewDecoder(conn.reader, &codec.MsgpackHandle{})
    conn.encoder = codec.NewEncoder(conn.writer, &codec.MsgpackHandle{})
    return conn
}

func (self *MemoryConnection) CallRPC(
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    if err := WriteEvent(self.writer, self.encoder, request); err != nil {
        self.Close()
        return nil, err
    }

    event, err := ReadResponse(self.reader, self.decoder)
    if err != nil {
        self.Close()
        return nil, err
    }
    return event, nil
}

type MemoryClient struct {
    connectionPool     map[string][]*MemoryConnection
    connectionPoolLock sync.Mutex

    poolSize int
    register *MemoryTransportRegister
}

func NewMemoryClient(
    poolSize int, register *MemoryTransportRegister) *MemoryClient {

    return &MemoryClient{
        connectionPool: make(map[string][]*MemoryConnection),
        poolSize:       poolSize,
        register:       register,
    }
}

func (self *MemoryClient) CallRPCTo(
    target net.Addr,
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    connection, err := self.getConnection(target)
    if err != nil {
        return nil, err
    }

    response, err = connection.CallRPC(request)
    if err == nil {
        self.returnConnectionToPool(connection)
    }
    return response, err
}

func (self *MemoryClient) getConnectionFromPool(
    target net.Addr) (*MemoryConnection, error) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := target.String()
    connections, ok := self.connectionPool[key]

    if !ok || len(connections) == 0 {
        return nil, errors.New("no connection for this target")
    }

    connection := connections[len(connections)-1]
    self.connectionPool[key] = connections[:len(connections)-1]
    return connection, nil
}

func (self *MemoryClient) returnConnectionToPool(
    connection *MemoryConnection) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := connection.PeerAddr().String()
    connections, ok := self.connectionPool[key]
    if !ok {
        connections = make([]*MemoryConnection, 0)
        self.connectionPool[key] = connections
    }

    if len(connections) < self.poolSize {
        self.connectionPool[key] = append(connections, connection)
    } else {
        connection.Close()
    }
}

func (self *MemoryClient) getConnection(
    target net.Addr) (*MemoryConnection, error) {

    connection, err := self.getConnectionFromPool(target)
    if connection != nil && err == nil {
        return connection, nil
    }

    connection = NewMemoryConnection(target, self.register)
    if err := connection.Open(); err != nil {
        return nil, err
    }

    return connection, nil
}

func (self *MemoryClient) Close() error {
    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    var err error
    for target, connections := range self.connectionPool {
        for _, connection := range connections {
            err = connection.Close()
        }
        delete(self.connectionPool, target)
    }
    return err
}

type MemoryServer struct {
    register            *MemoryTransportRegister
    transport           *MemoryServerTransport
    acceptedConnections map[chan []byte]*MemoryServerTransport
    eventHandler        func(ev.RaftRequestEvent)
    logger              logging.Logger
}

func NewMemoryServer(
    bindAddr net.Addr,
    eventHandler func(ev.RaftRequestEvent),
    register *MemoryTransportRegister,
    logger logging.Logger) *MemoryServer {

    transport := NewMemoryServerTransport(bindAddr, register)
    transport.Open()
    return &MemoryServer{
        register:            register,
        transport:           transport,
        acceptedConnections: make(map[chan []byte]*MemoryServerTransport),
        eventHandler:        eventHandler,
        logger:              logger,
    }
}

func (self *MemoryServer) Serve() {
    for {
        chunk, err := self.transport.ReadChunk()
        if err != nil {
            self.logger.Error(
                "memory server fails to read first chunk, error: %s", err)
            // server exit
            return
        }
        transport, ok := self.acceptedConnections[chunk.SourceCh]
        if ok {
            // connection already accepted
            transport.WriteChunk(chunk)
            continue
        } else {
            addr := ps.ServerAddr{
                Protocol: self.transport.Addr().Network(),
                IP: fmt.Sprintf(
                    "%s.%#v", self.transport.Addr().String(), chunk.SourceCh),
            }
            transport := NewMemoryServerTransport(&addr, self.register)
            transport.ResponseCh = chunk.SourceCh
            self.acceptedConnections[chunk.SourceCh] = transport
            transport.WriteChunk(chunk)
            go self.handleConn(transport)
        }
    }
}

func (self *MemoryServer) handleConn(
    transport *MemoryServerTransport) {

    defer func() {
        transport.Close()
        delete(self.acceptedConnections, transport.ResponseCh)
    }()
    reader := bufio.NewReader(transport)
    writer := bufio.NewWriter(transport)
    decoder := codec.NewDecoder(reader, &codec.MsgpackHandle{})
    encoder := codec.NewEncoder(writer, &codec.MsgpackHandle{})

    for {
        if err := self.handleCommand(
            reader, writer, decoder, encoder); err != nil {

            if err != io.EOF {
                self.logger.Error(
                    "memory server fails to handle command, error: %s", err)
            }
            return
        }
        if err := writer.Flush(); err != nil {
            self.logger.Error(
                "memory server fails to flush writer, error: %s", err)
            return
        }
    }
}

func (self *MemoryServer) handleCommand(
    reader *bufio.Reader,
    writer *bufio.Writer,
    decoder Decoder,
    encoder Encoder) error {

    event, err := ReadRequest(reader, decoder)
    if err != nil {
        return err
    }

    self.eventHandler(event)
    response := event.RecvResponse()
    if err := WriteEvent(writer, encoder, response); err != nil {
        return err
    }
    return nil
}

func (self *MemoryServer) Close() error {
    return self.transport.Close()
}
<commit_msg>revise exit logic of memory server<commit_after>package comm

import (
    "bufio"
    "bytes"
    "errors"
    "fmt"
    ev "github.com/hhkbp2/rafted/event"
    logging "github.com/hhkbp2/rafted/logging"
    ps "github.com/hhkbp2/rafted/persist"
    "github.com/ugorji/go/codec"
    "io"
    "net"
    "sync"
)

const (
    DefaultTransportBufferSize = 16
)

type TransportChunk struct {
    Data     []byte
    SourceCh chan []byte
}

type MemoryServerTransport struct {
    addr       net.Addr
    ConsumeCh  chan *TransportChunk
    ResponseCh chan []byte
    register   *MemoryTransportRegister
}

func NewMemoryServerTransport(
    addr net.Addr, register *MemoryTransportRegister) *MemoryServerTransport {

    return &MemoryServerTransport{
        addr:      addr,
        ConsumeCh: make(chan *TransportChunk, DefaultTransportBufferSize),
        register:  register,
    }
}

func (self *MemoryServerTransport) Open() error {
    self.register.Register(self.addr.String(), self)
    return nil
}

func (self *MemoryServerTransport) ReadNextMessage() *TransportChunk {
    chunk := <-self.ConsumeCh
    return chunk
}

func (self *MemoryServerTransport) ReadChunk() (*TransportChunk, error) {
    chunk, ok := <-self.ConsumeCh
    if ok {
        return chunk, nil
    }
    return nil, io.EOF
}

func (self *MemoryServerTransport) WriteChunk(chunk *TransportChunk) error {
    self.ConsumeCh <- chunk
    return nil
}

func BytesCopy(dest []byte, src []byte) int {
    i := 0
    for ; (i < len(dest)) && (i < len(src)); i++ {
        dest[i] = src[i]
    }
    return i
}

func (self *MemoryServerTransport) Read(p []byte) (int, error) {
    chunk, err := self.ReadChunk()
    if err != nil {
        return 0, err
    }
    n := BytesCopy(p, chunk.Data)
    self.ResponseCh = chunk.SourceCh
    return n, nil
}

func (self *MemoryServerTransport) Write(p []byte) (int, error) {
    if self.ResponseCh != nil {
        self.ResponseCh <- p
        return len(p), nil
    }
    return 0, errors.New("don't know where to response")
}

func (self *MemoryServerTransport) Close() error {
    close(self.ConsumeCh)
    if self.ResponseCh != nil {
        close(self.ResponseCh)
    }
    return self.register.Unregister(self.addr.String())
}

func (self *MemoryServerTransport) Addr() net.Addr {
    return self.addr
}

type MemoryTransportRegister struct {
    sync.RWMutex
    transports map[string]*MemoryServerTransport
}

func NewMemoryTransportRegister() *MemoryTransportRegister {
    return &MemoryTransportRegister{
        transports: make(map[string]*MemoryServerTransport),
    }
}

func (self *MemoryTransportRegister) Register(
    id string, transport *MemoryServerTransport) {

    self.Lock()
    defer self.Unlock()
    self.transports[id] = transport
}

func (self *MemoryTransportRegister) Unregister(id string) error {
    self.Lock()
    defer self.Unlock()
    if _, ok := self.transports[id]; ok {
        delete(self.transports, id)
        return nil
    }
    return errors.New(fmt.Sprintf("no transport for id: %v", id))
}

func (self *MemoryTransportRegister) Reset() error {
    self.Lock()
    defer self.Unlock()
    self.transports = make(map[string]*MemoryServerTransport)
    return nil
}

func (self *MemoryTransportRegister) Get(
    id string) (transport *MemoryServerTransport, ok bool) {

    self.RLock()
    defer self.RUnlock()
    transport, ok = self.transports[id]
    return transport, ok
}

func (self *MemoryTransportRegister) String() string {
    self.RLock()
    defer self.RUnlock()
    var buf bytes.Buffer
    for k, _ := range self.transports {
        buf.WriteString(k + "\n")
    }
    return buf.String()
}

type MemoryTransport struct {
    addr      net.Addr
    consumeCh chan []byte
    register  *MemoryTransportRegister
    peer      *MemoryServerTransport
}

func NewMemoryTransport(addr net.Addr, register *MemoryTransportRegister) *MemoryTransport {
    return &MemoryTransport{
        addr:      addr,
        consumeCh: make(chan []byte, DefaultTransportBufferSize),
        register:  register,
    }
}

func (self *MemoryTransport) PeerAddr() net.Addr {
    return self.addr
}

func (self *MemoryTransport) Open() error {
    if transport, ok := self.register.Get(self.addr.String()); ok {
        self.peer = transport
        return nil
    }
    return errors.New(
        fmt.Sprintf("no server transport for id: %v", self.addr.String()))
}

func (self *MemoryTransport) Close() error {
    // empty body
    return nil
}

func (self *MemoryTransport) Read(b []byte) (int, error) {
    data := <-self.consumeCh
    n := BytesCopy(b, data)
    return n, nil
}

func (self *MemoryTransport) Write(b []byte) (int, error) {
    chunk := &TransportChunk{
        Data:     b,
        SourceCh: self.consumeCh,
    }
    if err := self.peer.WriteChunk(chunk); err != nil {
        return 0, err
    }
    return len(b), nil
}

type MemoryConnection struct {
    *MemoryTransport
    reader  *bufio.Reader
    writer  *bufio.Writer
    encoder Encoder
    decoder Decoder
}

func NewMemoryConnection(
    addr net.Addr, register *MemoryTransportRegister) *MemoryConnection {

    conn := &MemoryConnection{
        MemoryTransport: NewMemoryTransport(addr, register),
    }
    conn.reader = bufio.NewReader(conn.MemoryTransport)
    conn.writer = bufio.NewWriter(conn.MemoryTransport)
    conn.decoder = codec.NewDecoder(conn.reader, &codec.MsgpackHandle{})
    conn.encoder = codec.NewEncoder(conn.writer, &codec.MsgpackHandle{})
    return conn
}

func (self *MemoryConnection) CallRPC(
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    if err := WriteEvent(self.writer, self.encoder, request); err != nil {
        self.Close()
        return nil, err
    }

    event, err := ReadResponse(self.reader, self.decoder)
    if err != nil {
        self.Close()
        return nil, err
    }
    return event, nil
}

type MemoryClient struct {
    connectionPool     map[string][]*MemoryConnection
    connectionPoolLock sync.Mutex

    poolSize int
    register *MemoryTransportRegister
}

func NewMemoryClient(
    poolSize int, register *MemoryTransportRegister) *MemoryClient {

    return &MemoryClient{
        connectionPool: make(map[string][]*MemoryConnection),
        poolSize:       poolSize,
        register:       register,
    }
}

func (self *MemoryClient) CallRPCTo(
    target net.Addr,
    request ev.RaftEvent) (response ev.RaftEvent, err error) {

    connection, err := self.getConnection(target)
    if err != nil {
        return nil, err
    }

    response, err = connection.CallRPC(request)
    if err == nil {
        self.returnConnectionToPool(connection)
    }
    return response, err
}

func (self *MemoryClient) getConnectionFromPool(
    target net.Addr) (*MemoryConnection, error) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := target.String()
    connections, ok := self.connectionPool[key]

    if !ok || len(connections) == 0 {
        return nil, errors.New("no connection for this target")
    }

    connection := connections[len(connections)-1]
    self.connectionPool[key] = connections[:len(connections)-1]
    return connection, nil
}

func (self *MemoryClient) returnConnectionToPool(
    connection *MemoryConnection) {

    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    key := connection.PeerAddr().String()
    connections, ok := self.connectionPool[key]
    if !ok {
        connections = make([]*MemoryConnection, 0)
        self.connectionPool[key] = connections
    }

    if len(connections) < self.poolSize {
        self.connectionPool[key] = append(connections, connection)
    } else {
        connection.Close()
    }
}

func (self *MemoryClient) getConnection(
    target net.Addr) (*MemoryConnection, error) {

    connection, err := self.getConnectionFromPool(target)
    if connection != nil && err == nil {
        return connection, nil
    }

    connection = NewMemoryConnection(target, self.register)
    if err := connection.Open(); err != nil {
        return nil, err
    }

    return connection, nil
}

func (self *MemoryClient) Close() error {
    self.connectionPoolLock.Lock()
    defer self.connectionPoolLock.Unlock()

    var err error
    for target, connections := range self.connectionPool {
        for _, connection := range connections {
            err = connection.Close()
        }
        delete(self.connectionPool, target)
    }
    return err
}

type MemoryServer struct {
    register            *MemoryTransportRegister
    transport           *MemoryServerTransport
    acceptedConnections map[chan []byte]*MemoryServerTransport
    eventHandler        func(ev.RaftRequestEvent)
    logger              logging.Logger
}

func NewMemoryServer(
    bindAddr net.Addr,
    eventHandler func(ev.RaftRequestEvent),
    register *MemoryTransportRegister,
    logger logging.Logger) *MemoryServer {

    transport := NewMemoryServerTransport(bindAddr, register)
    transport.Open()
    return &MemoryServer{
        register:            register,
        transport:           transport,
        acceptedConnections: make(map[chan []byte]*MemoryServerTransport),
        eventHandler:        eventHandler,
        logger:              logger,
    }
}

func (self *MemoryServer) Serve() {
    for {
        chunk := self.transport.ReadNextMessage()
        if chunk == nil {
            self.logger.Debug("memory server read no more message")
            // server exit
            return
        }
        transport, ok := self.acceptedConnections[chunk.SourceCh]
        if ok {
            // connection already accepted
            transport.WriteChunk(chunk)
            continue
        } else {
            addr := ps.ServerAddr{
                Protocol: self.transport.Addr().Network(),
                IP: fmt.Sprintf(
                    "%s.%#v", self.transport.Addr().String(), chunk.SourceCh),
            }
            transport := NewMemoryServerTransport(&addr, self.register)
            transport.ResponseCh = chunk.SourceCh
            self.acceptedConnections[chunk.SourceCh] = transport
            transport.WriteChunk(chunk)
            go self.handleConn(transport)
        }
    }
}

func (self *MemoryServer) handleConn(
    transport *MemoryServerTransport) {

    defer func() {
        transport.Close()
        delete(self.acceptedConnections, transport.ResponseCh)
    }()
    reader := bufio.NewReader(transport)
    writer := bufio.NewWriter(transport)
    decoder := codec.NewDecoder(reader, &codec.MsgpackHandle{})
    encoder := codec.NewEncoder(writer, &codec.MsgpackHandle{})

    for {
        if err := self.handleCommand(
            reader, writer, decoder, encoder); err != nil {

            if err != io.EOF {
                self.logger.Error(
                    "memory server fails to handle command, error: %s", err)
            }
            return
        }
        if err := writer.Flush(); err != nil {
            self.logger.Error(
                "memory server fails to flush writer, error: %s", err)
            return
        }
    }
}

func (self *MemoryServer) handleCommand(
    reader *bufio.Reader,
    writer *bufio.Writer,
    decoder Decoder,
    encoder Encoder) error {

    event, err := ReadRequest(reader, decoder)
    if err != nil {
        return err
    }

    self.eventHandler(event)
    response := event.RecvResponse()
    if err := WriteEvent(writer, encoder, response); err != nil {
        return err
    }
    return nil
}

func (self *MemoryServer) Close() error {
    return self.transport.Close()
}
<|endoftext|>


<commit_before>// +build go1.8

package avatica

import (
	"database/sql/driver"
	"fmt"
)

type namedValue struct {
	Name    string
	Ordinal int
	Value   driver.Value
}

func driverValueToNamedValue(values []driver.Value) []namedValue {
	list := make([]namedValue, len(values))

	for i, v := range values {
		list[i] = namedValue{
			Ordinal: i + 1,
			Value:   v,
		}
	}

	return list
}

func driverNamedValueToNamedValue(values []driver.NamedValue) ([]namedValue,error ) {
	list := make([]namedValue, len(values))

	for i, nv := range values {
		list[i] = namedValue(nv)

		if nv.Name != ""{
			return list,fmt.Errorf("named paramters are not supported: %s given", nv.Name)
		}
	}

	return list, nil
}

type isoLevel int32

const (
	isolationUseCurrent      isoLevel = -1
	isolationNone            isoLevel = 0
	isolationReadUncommitted isoLevel = 1
	isolationReadComitted    isoLevel = 2
	isolationRepeatableRead  isoLevel = 4
	isolationSerializable    isoLevel = 8
)
<commit_msg>Remove go 1.8 build tag for compatibility structs and constants<commit_after>package avatica

import (
	"database/sql/driver"
	"fmt"
)

type namedValue struct {
	Name    string
	Ordinal int
	Value   driver.Value
}

func driverValueToNamedValue(values []driver.Value) []namedValue {
	list := make([]namedValue, len(values))

	for i, v := range values {
		list[i] = namedValue{
			Ordinal: i + 1,
			Value:   v,
		}
	}

	return list
}

func driverNamedValueToNamedValue(values []driver.NamedValue) ([]namedValue, error) {
	list := make([]namedValue, len(values))

	for i, nv := range values {
		list[i] = namedValue(nv)

		if nv.Name != "" {
			return list, fmt.Errorf("named paramters are not supported: %s given", nv.Name)
		}
	}

	return list, nil
}

type isoLevel int32

const (
	isolationUseCurrent      isoLevel = -1
	isolationNone            isoLevel = 0
	isolationReadUncommitted isoLevel = 1
	isolationReadComitted    isoLevel = 2
	isolationRepeatableRead  isoLevel = 4
	isolationSerializable    isoLevel = 8
)
<|endoftext|>


<commit_before>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"

	"github.com/mackerelio/golib/logging"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
}

var timeout = 100 * time.Millisecond

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator() *CloudGenerator {
	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isEC2() bool {
	// If the OS is Linux, check /sys/hypervisor/uuid file first.
	// ref. http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/identify_ec2_instances.html
	if runtime.GOOS == "linux" {
		data, err := ioutil.ReadFile("/sys/hypervisor/uuid")
		if err != nil {
			// Probably not EC2.
			return false
		}
		// Probably not EC2.
		if !strings.HasPrefix(string(data), "ec2") {
			return false
		}
	}

	cl := httpCli()
	// '/ami-id` is may be aws specific URL
	resp, err := cl.Get(ec2BaseURL.String() + "/ami-id")
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	return resp.StatusCode == 200
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}
<commit_msg>retry a few times within isEC2<commit_after>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"

	"github.com/Songmu/retry"
	"github.com/mackerelio/golib/logging"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
}

var timeout = 100 * time.Millisecond

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator() *CloudGenerator {
	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isEC2() bool {
	// If the OS is Linux, check /sys/hypervisor/uuid file first.
	// ref. http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/identify_ec2_instances.html
	if runtime.GOOS == "linux" {
		data, err := ioutil.ReadFile("/sys/hypervisor/uuid")
		if err != nil {
			// Probably not EC2.
			return false
		}
		// Probably not EC2.
		if !strings.HasPrefix(string(data), "ec2") {
			return false
		}
	}

	res := false
	cl := httpCli()
	err := retry.Retry(3, 2*timeout, func() error {
		// '/ami-id` is probably an AWS specific URL
		resp, err := cl.Get(ec2BaseURL.String() + "/ami-id")
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		res = resp.StatusCode == 200
		return nil
	})

	if err != nil {
		return res
	}
	return false
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}
<|endoftext|>


<commit_before>package template

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"text/template"
)

type Template struct {
	name string
	path string
	vars map[string]string
}

func New(name, template_path string, vars map[string]string) *Template {
	return &Template{
		name: name,
		path: template_path,
		vars: vars,
	}
}

func (t *Template) Create(w io.Writer) (*Template, error) {
	if _, err := os.Stat(t.path); err == nil {
		tmpl_content, err := ioutil.ReadFile(t.path)
		if err != nil {
			return t, err
		}

		var_names, err := find_vars(string(tmpl_content))
		if err != nil {
			return t, err
		}

		var missing_vars []string
		for _, name := range var_names {
			key := name
			if _, ok := t.vars[key]; !ok {
				missing_vars = append(missing_vars, key)
			}
		}

		if len(missing_vars) > 0 {
			return t, fmt.Errorf(
				"Missing template variables: %s",
				missing_vars,
			)
		}

		tmpl, err := template.New(t.name).Option("missingkey=error").Parse(string(tmpl_content))
		if err != nil {
			return t, err
		}

		err = tmpl.Execute(w, t.vars)
		if err != nil {
			return t, err
		}
	}

	return t, nil
}

func find_vars(template_content string) ([]string, error) {
	var vars []string

	regex, err := regexp.Compile(`\s?{{\.[a-zA-Z0-9][a-zA-Z0-9]*}}`)
	if err != nil {
		return vars, err
	}

	matches := regex.FindAllString(template_content, -1)

	regex, err = regexp.Compile(`[a-zA-Z0-0]+`)
	if err != nil {
		return vars, err
	}

	for _, match := range matches {
		var_name := regex.FindString(match)
		vars = append(vars, var_name)
	}

	return vars, nil
}
<commit_msg>Build template missing variable error<commit_after>package template

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"text/template"
)

type Template struct {
	name string
	path string
	vars map[string]string
}

func New(name, template_path string, vars map[string]string) *Template {
	return &Template{
		name: name,
		path: template_path,
		vars: vars,
	}
}

func (t *Template) Create(w io.Writer) (*Template, error) {
	if _, err := os.Stat(t.path); err == nil {
		tmpl_content, err := ioutil.ReadFile(t.path)
		if err != nil {
			return t, err
		}

		var_names, err := find_vars(string(tmpl_content))
		if err != nil {
			return t, err
		}

		var missing_vars []string
		for _, name := range var_names {
			key := name
			if _, ok := t.vars[key]; !ok {
				missing_vars = append(missing_vars, key)
			}
		}

		if len(missing_vars) > 0 {
			var err_msg bytes.Buffer
			for _, missing_var := range missing_vars {
				err_msg.WriteString("{{" + "." + missing_var + "}}\n")
			}

			return t, fmt.Errorf(
				"%s: Missing template variables:\n%s",
				t.path,
				err_msg.String(),
			)
		}

		tmpl, err := template.New(t.name).Option("missingkey=error").Parse(string(tmpl_content))
		if err != nil {
			return t, err
		}

		err = tmpl.Execute(w, t.vars)
		if err != nil {
			return t, err
		}
	}

	return t, nil
}

func find_vars(template_content string) ([]string, error) {
	var vars []string

	regex, err := regexp.Compile(`\s?{{\.[a-zA-Z0-9][a-zA-Z0-9]*}}`)
	if err != nil {
		return vars, err
	}

	matches := regex.FindAllString(template_content, -1)

	regex, err = regexp.Compile(`[a-zA-Z0-0]+`)
	if err != nil {
		return vars, err
	}

	for _, match := range matches {
		var_name := regex.FindString(match)
		vars = append(vars, var_name)
	}

	return vars, nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"github.com/gin-gonic/gin"
	"net/http"
	"sync"

	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"
	e "github.com/eirka/eirka-libs/errors"
)

var (
	sitemap = make(map[string]*SiteData)
	mu      = new(sync.RWMutex)
)

// SiteData holds imageboard settings
type SiteData struct {
	Ib          uint
	Api         string
	Img         string
	Title       string
	Desc        string
	Nsfw        bool
	Style       string
	Logo        string
	Imageboards []Imageboard
}

type Imageboard struct {
	Title   string
	Address string
}

// Details gets the imageboard settings from the request for the page handler variables
func Details() gin.HandlerFunc {
	return func(c *gin.Context) {

		host := c.Request.Host

		c.Error(host)

		mu.RLock()
		// check the sitemap to see if its cached
		site := sitemap[host]
		mu.RUnlock()

		// if not query the database
		if site == nil {

			sitedata := &SiteData{}

			// Get Database handle
			dbase, err := db.GetDb()
			if err != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.GetDb")
				c.Abort()
				return
			}

			// get the info about the imageboard
			err = dbase.QueryRow(`SELECT ib_id,ib_title,ib_description,ib_nsfw,ib_api,ib_img,ib_style,ib_logo FROM imageboards WHERE ib_domain = ?`, host).Scan(&sitedata.Ib, &sitedata.Title, &sitedata.Desc, &sitedata.Nsfw, &sitedata.Api, &sitedata.Img, &sitedata.Style, &sitedata.Logo)
			if err == sql.ErrNoRows {
				c.JSON(e.ErrorMessage(e.ErrNotFound))
				c.Error(err).SetMeta("Details.QueryRow")
				return
			} else if err != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.QueryRow")
				return
			}

			// collect the links to the other imageboards for nav menu
			rows, err := dbase.Query(`SELECT ib_title,ib_domain FROM imageboards WHERE ib_id != ?`, sitedata.Ib)
			if err != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.Query")
				c.Abort()
				return
			}
			defer rows.Close()

			for rows.Next() {

				ib := Imageboard{}

				err := rows.Scan(&ib.Title, &ib.Address)
				if err != nil {
					return
				}

				sitedata.Imageboards = append(sitedata.Imageboards, ib)
			}
			if rows.Err() != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.Query")
				c.Abort()
				return
			}

			mu.Lock()
			sitemap[host] = sitedata
			mu.Unlock()

		}

		c.Next()

	}

}

// IndexController generates pages for angularjs frontend
func IndexController(c *gin.Context) {

	// Get parameters from csrf middleware
	csrf_token := c.MustGet("csrf_token").(string)

	host := c.Request.Host

	mu.RLock()
	site := sitemap[host]
	mu.RUnlock()

	c.HTML(http.StatusOK, "index", gin.H{
		"primjs":      config.Settings.Prim.Js,
		"primcss":     config.Settings.Prim.Css,
		"ib":          site.Ib,
		"apisrv":      site.Api,
		"imgsrv":      site.Img,
		"title":       site.Title,
		"desc":        site.Desc,
		"nsfw":        site.Nsfw,
		"style":       site.Style,
		"logo":        site.Logo,
		"imageboards": site.Imageboards,
		"csrf":        csrf_token,
	})

	return

}

// ErrorController generates pages and a 404 response
func ErrorController(c *gin.Context) {

	// Get parameters from csrf middleware
	csrf_token := c.MustGet("csrf_token").(string)

	host := c.Request.Host

	mu.RLock()
	site := sitemap[host]
	mu.RUnlock()

	c.HTML(http.StatusNotFound, "index", gin.H{
		"primjs":      config.Settings.Prim.Js,
		"primcss":     config.Settings.Prim.Css,
		"ib":          site.Ib,
		"apisrv":      site.Api,
		"imgsrv":      site.Img,
		"title":       site.Title,
		"desc":        site.Desc,
		"nsfw":        site.Nsfw,
		"style":       site.Style,
		"logo":        site.Logo,
		"imageboards": site.Imageboards,
		"csrf":        csrf_token,
	})

	return

}
<commit_msg>debug<commit_after>package main

import (
	"database/sql"
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"sync"

	"github.com/eirka/eirka-libs/config"
	"github.com/eirka/eirka-libs/db"
	e "github.com/eirka/eirka-libs/errors"
)

var (
	sitemap = make(map[string]*SiteData)
	mu      = new(sync.RWMutex)
)

// SiteData holds imageboard settings
type SiteData struct {
	Ib          uint
	Api         string
	Img         string
	Title       string
	Desc        string
	Nsfw        bool
	Style       string
	Logo        string
	Imageboards []Imageboard
}

type Imageboard struct {
	Title   string
	Address string
}

// Details gets the imageboard settings from the request for the page handler variables
func Details() gin.HandlerFunc {
	return func(c *gin.Context) {

		host := c.Request.Host

		fmt.Println(host)

		mu.RLock()
		// check the sitemap to see if its cached
		site := sitemap[host]
		mu.RUnlock()

		// if not query the database
		if site == nil {

			sitedata := &SiteData{}

			// Get Database handle
			dbase, err := db.GetDb()
			if err != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.GetDb")
				c.Abort()
				return
			}

			// get the info about the imageboard
			err = dbase.QueryRow(`SELECT ib_id,ib_title,ib_description,ib_nsfw,ib_api,ib_img,ib_style,ib_logo FROM imageboards WHERE ib_domain = ?`, host).Scan(&sitedata.Ib, &sitedata.Title, &sitedata.Desc, &sitedata.Nsfw, &sitedata.Api, &sitedata.Img, &sitedata.Style, &sitedata.Logo)
			if err == sql.ErrNoRows {
				c.JSON(e.ErrorMessage(e.ErrNotFound))
				c.Error(err).SetMeta("Details.QueryRow")
				return
			} else if err != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.QueryRow")
				return
			}

			// collect the links to the other imageboards for nav menu
			rows, err := dbase.Query(`SELECT ib_title,ib_domain FROM imageboards WHERE ib_id != ?`, sitedata.Ib)
			if err != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.Query")
				c.Abort()
				return
			}
			defer rows.Close()

			for rows.Next() {

				ib := Imageboard{}

				err := rows.Scan(&ib.Title, &ib.Address)
				if err != nil {
					return
				}

				sitedata.Imageboards = append(sitedata.Imageboards, ib)
			}
			if rows.Err() != nil {
				c.JSON(e.ErrorMessage(e.ErrInternalError))
				c.Error(err).SetMeta("Details.Query")
				c.Abort()
				return
			}

			mu.Lock()
			sitemap[host] = sitedata
			mu.Unlock()

		}

		c.Next()

	}

}

// IndexController generates pages for angularjs frontend
func IndexController(c *gin.Context) {

	// Get parameters from csrf middleware
	csrf_token := c.MustGet("csrf_token").(string)

	host := c.Request.Host

	mu.RLock()
	site := sitemap[host]
	mu.RUnlock()

	c.HTML(http.StatusOK, "index", gin.H{
		"primjs":      config.Settings.Prim.Js,
		"primcss":     config.Settings.Prim.Css,
		"ib":          site.Ib,
		"apisrv":      site.Api,
		"imgsrv":      site.Img,
		"title":       site.Title,
		"desc":        site.Desc,
		"nsfw":        site.Nsfw,
		"style":       site.Style,
		"logo":        site.Logo,
		"imageboards": site.Imageboards,
		"csrf":        csrf_token,
	})

	return

}

// ErrorController generates pages and a 404 response
func ErrorController(c *gin.Context) {

	// Get parameters from csrf middleware
	csrf_token := c.MustGet("csrf_token").(string)

	host := c.Request.Host

	mu.RLock()
	site := sitemap[host]
	mu.RUnlock()

	c.HTML(http.StatusNotFound, "index", gin.H{
		"primjs":      config.Settings.Prim.Js,
		"primcss":     config.Settings.Prim.Css,
		"ib":          site.Ib,
		"apisrv":      site.Api,
		"imgsrv":      site.Img,
		"title":       site.Title,
		"desc":        site.Desc,
		"nsfw":        site.Nsfw,
		"style":       site.Style,
		"logo":        site.Logo,
		"imageboards": site.Imageboards,
		"csrf":        csrf_token,
	})

	return

}
<|endoftext|>