<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"bitbucket.org/kardianos/osext"
	"code.google.com/p/gcfg"
	"github.com/influxdb/influxdb/client"
	"github.com/soniah/gosnmp"
)

const layout = "2006-01-02 15:04:05"

type SnmpConfig struct {
	Host      string `gcfg:"host"`
	Public    string `gcfg:"community"`
	Port      int    `gcfg:"port"`
	Retries   int    `gcfg:"retries"`
	Timeout   int    `gcfg:"timeout"`
	Repeat    int    `gcfg:"repeat"`
	Freq      int    `gcfg:"freq"`
	PortFile  string `gcfg:"portfile"`
	Config    string `gcfg:"config"`
	labels    map[string]string
	asName    map[string]string
	asOID     map[string]string
	oids      []string
	mib       *MibConfig
	Influx    *InfluxConfig
	LastError time.Time
	Requests  int64
	Gets      int64
	Errors    int64
	debugging chan bool
	enabled   chan chan bool
}

type InfluxConfig struct {
	Host      string `gcfg:"host"`
	Port      int    `gcfg:"port"`
	DB        string `gcfg:"db"`
	User      string `gcfg:"user"`
	Password  string `gcfg:"password"`
	Retention string `gcfg:"retention"`
	iChan     chan *client.BatchPoints
	conn      *client.Client
	Sent      int64
	Errors    int64
}

type HTTPConfig struct {
	Port int `gcfg:"port"`
}

type GeneralConfig struct {
	LogDir  string `gcfg:"logdir"`
	OidFile string `gcfg:"oidfile"`
}

type MibConfig struct {
	Scalers bool     `gcfg:"scalers"`
	Name    string   `gcfg:"name"`
	Columns []string `gcfg:"column"`
}

var (
	quit          = make(chan struct{})
	verbose       bool
	startTime     = time.Now()
	testing       bool
	snmpNames     bool
	repeat        = 0
	freq          = 30
	httpPort      = 8080
	oidToName     = make(map[string]string)
	nameToOid     = make(map[string]string)
	appdir, _     = osext.ExecutableFolder()
	logDir        = filepath.Join(appdir, "log")
	oidFile       = filepath.Join(appdir, "oids.txt")
	configFile    = filepath.Join(appdir, "config.gcfg")
	errorLog      *os.File
	errorDuration = time.Duration(10 * time.Minute)
	errorPeriod   = errorDuration.String()
	errorMax      = 100
	errorName     string

	cfg = struct {
		Snmp    map[string]*SnmpConfig
		Mibs    map[string]*MibConfig
		Influx  map[string]*InfluxConfig
		HTTP    HTTPConfig
		General GeneralConfig
	}{}
)

func fatal(v ...interface{}) {
	log.SetOutput(os.Stderr)
	log.Fatalln(v...)
}

func (c *SnmpConfig) DebugAction() string {
	debug := make(chan bool)
	c.enabled <- debug
	if <-debug {
		return "disable"
	}
	return "enable"
}

func (c *SnmpConfig) LoadPorts() {
	c.labels = make(map[string]string)
	if len(c.PortFile) == 0 {
		return
	}
	data, err := ioutil.ReadFile(filepath.Join(appdir, c.PortFile))
	if err != nil {
		log.Fatal(err)
	}
	for _, line := range strings.Split(string(data), "\n") {
		// strip comments
		comment := strings.Index(line, "#")
		if comment >= 0 {
			line = line[:comment]
		}
		f := strings.Fields(line)
		if len(f) < 2 {
			continue
		}
		c.labels[f[0]] = f[1]
	}
}

func (c *SnmpConfig) incRequests() {
	atomic.AddInt64(&c.Requests, 1)
}

func (c *SnmpConfig) incGets() {
	atomic.AddInt64(&c.Gets, 1)
}

func (c *SnmpConfig) incErrors() {
	atomic.AddInt64(&c.Errors, 1)
}

func (c *InfluxConfig) incErrors() {
	atomic.AddInt64(&c.Errors, 1)
}

func (c *InfluxConfig) incSent() {
	atomic.AddInt64(&c.Sent, 1)
}

// loads [last_octet]name for device
func (c *SnmpConfig) Translate() {
	client, err := snmpClient(c)
	if err != nil {
		fatal("Client connect error:", err)
	}
	defer client.Conn.Close()
	spew("Looking up column names for:", c.Host)
	pdus, err := client.BulkWalkAll(nameOid)
	if err != nil {
		fatal("SNMP bulkwalk error", err)
	}
	c.asName = make(map[string]string)
	c.asOID = make(map[string]string)
	for _, pdu := range pdus {
		switch pdu.Type {
		case gosnmp.OctetString:
			i := strings.LastIndex(pdu.Name, ".")
			suffix := pdu.Name[i+1:]
			name := string(pdu.Value.([]byte))
			_, ok := c.labels[name]
			if len(c.PortFile) == 0 || ok {
				c.asName[name] = suffix
				c.asOID[suffix] = name
			}
		}
	}
	// make sure we got everything
	for k := range c.labels {
		if _, ok := c.asName[k]; !ok {
			fatal("No OID found for:", k)
		}
	}

}

func spew(x ...interface{}) {
	if verbose {
		fmt.Println(x...)
	}
}

func (c *SnmpConfig) OIDs() {
	if c.mib == nil {
		fatal("NO MIB!")
	}
	c.oids = []string{}
	for _, col := range c.mib.Columns {
		base, ok := nameToOid[col]
		if !ok {
			fatal("no oid for col:", col)
		}
		// just named columns
		if len(c.PortFile) > 0 {
			for k := range c.asOID {
				c.oids = append(c.oids, base+"."+k)
			}
		} else if c.mib.Scalers {
			// or plain old scaler instances
			c.oids = append(c.oids, base+".0")
		} else {
			c.oids = append(c.oids, base)
		}
	}
	if len(c.mib.Columns) > 0 {
		spew("COLUMNS", c.mib.Columns)
		spew(c.oids)
	}
}

// load oid lookup data
func init() {
	data, err := ioutil.ReadFile(oidFile)
	if err != nil {
		log.Fatal(err)
	}
	for _, line := range strings.Split(string(data), "\n") {
		f := strings.Fields(line)
		if len(f) < 2 {
			continue
		}
		nameToOid[f[0]] = f[1]
		oidToName[f[1]] = f[0]
	}
}

func flags() *flag.FlagSet {
	var f flag.FlagSet
	f.BoolVar(&testing, "testing", testing, "print data w/o saving")
	f.BoolVar(&snmpNames, "names", snmpNames, "print column names and exit")
	f.StringVar(&configFile, "config", configFile, "config file")
	f.BoolVar(&verbose, "verbose", verbose, "verbose mode")
	f.IntVar(&repeat, "repeat", repeat, "number of times to repeat")
	f.IntVar(&freq, "freq", freq, "delay (in seconds)")
	f.IntVar(&httpPort, "http", httpPort, "http port")
	f.StringVar(&logDir, "logs", logDir, "log directory")
	f.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		f.VisitAll(func(flag *flag.Flag) {
			format := "%10s: %s\n"
			fmt.Fprintf(os.Stderr, format, "-"+flag.Name, flag.Usage)
		})
		fmt.Fprintf(os.Stderr, "\nAll settings can be set in config file: %s\n", configFile)
		os.Exit(1)

	}
	return &f
}

func init() {
	// parse first time to see if config file is being specified
	f := flags()
	f.Parse(os.Args[1:])
	// now load up config settings
	if _, err := os.Stat(configFile); err != nil {
		log.Fatal(err)
	} else {
		data, err := ioutil.ReadFile(configFile)
		if err != nil {
			log.Fatal(err)
		}
		err = gcfg.ReadStringInto(&cfg, string(data))
		if err != nil {
			log.Fatalf("Failed to parse gcfg data: %s", err)
		}
		httpPort = cfg.HTTP.Port
	}

	if len(cfg.General.LogDir) > 0 {
		logDir = cfg.General.LogDir
	}
	if len(cfg.General.OidFile) > 0 {
		oidFile = cfg.General.OidFile
	}

	for _, s := range cfg.Snmp {
		s.LoadPorts()
		s.debugging = make(chan bool)
		s.enabled = make(chan chan bool)
	}
	var ok bool
	for name, c := range cfg.Snmp {
		if c.mib, ok = cfg.Mibs[name]; !ok {
			if c.mib, ok = cfg.Mibs["*"]; !ok {
				fatal("No mib data found for config:", name)
			}
		}
		c.Translate()
		c.OIDs()
		if c.Freq == 0 {
			c.Freq = freq
		}
	}

	// only run when one needs to see the interface names of the device
	if snmpNames {
		for _, c := range cfg.Snmp {
			fmt.Println("\nSNMP host:", c.Host)
			fmt.Println("=========================================")
			printSnmpNames(c)
		}
		os.Exit(0)
	}

	// re-read cmd line args to override as indicated
	f = flags()
	f.Parse(os.Args[1:])
	os.Mkdir(logDir, 0755)

	// now make sure each snmp device has a db
	for name, c := range cfg.Snmp {
		// default is to use name of snmp config, but it can be overridden
		if len(c.Config) > 0 {
			name = c.Config
		}
		if c.Influx, ok = cfg.Influx[name]; !ok {
			if c.Influx, ok = cfg.Influx["*"]; !ok {
				fatal("No influx config for snmp device:", name)
			}
		}
		c.Influx.Init()
	}

	var ferr error
	errorName = fmt.Sprintf("error.%d.log", cfg.HTTP.Port)
	errorPath := filepath.Join(logDir, errorName)
	errorLog, ferr = os.OpenFile(errorPath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0664)
	if ferr != nil {
		log.Fatal("Can't open error log:", ferr)
	}
}

func errLog(msg string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, msg, args...)
	fmt.Fprintf(errorLog, msg, args...)
}

func errMsg(msg string, err error) {
	now := time.Now()
	errLog("%s\t%s: %s\n", now.Format(layout), msg, err)
}

func main() {
	var wg sync.WaitGroup
	defer func() {
		errorLog.Close()
	}()
	for _, c := range cfg.Snmp {
		wg.Add(1)
		go c.Gather(repeat, &wg)
	}
	if repeat > 0 {
		wg.Wait()
	} else {
		if httpPort > 0 {
			webServer(httpPort)
		} else {
			<-quit
		}
	}
}
<commit_msg>Fixed import warnings<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/kardianos/osext"
	"gopkg.in/gcfg.v1"
	"github.com/influxdb/influxdb/client"
	"github.com/soniah/gosnmp"
)

const layout = "2006-01-02 15:04:05"

type SnmpConfig struct {
	Host      string `gcfg:"host"`
	Public    string `gcfg:"community"`
	Port      int    `gcfg:"port"`
	Retries   int    `gcfg:"retries"`
	Timeout   int    `gcfg:"timeout"`
	Repeat    int    `gcfg:"repeat"`
	Freq      int    `gcfg:"freq"`
	PortFile  string `gcfg:"portfile"`
	Config    string `gcfg:"config"`
	labels    map[string]string
	asName    map[string]string
	asOID     map[string]string
	oids      []string
	mib       *MibConfig
	Influx    *InfluxConfig
	LastError time.Time
	Requests  int64
	Gets      int64
	Errors    int64
	debugging chan bool
	enabled   chan chan bool
}

type InfluxConfig struct {
	Host      string `gcfg:"host"`
	Port      int    `gcfg:"port"`
	DB        string `gcfg:"db"`
	User      string `gcfg:"user"`
	Password  string `gcfg:"password"`
	Retention string `gcfg:"retention"`
	iChan     chan *client.BatchPoints
	conn      *client.Client
	Sent      int64
	Errors    int64
}

type HTTPConfig struct {
	Port int `gcfg:"port"`
}

type GeneralConfig struct {
	LogDir  string `gcfg:"logdir"`
	OidFile string `gcfg:"oidfile"`
}

type MibConfig struct {
	Scalers bool     `gcfg:"scalers"`
	Name    string   `gcfg:"name"`
	Columns []string `gcfg:"column"`
}

var (
	quit          = make(chan struct{})
	verbose       bool
	startTime     = time.Now()
	testing       bool
	snmpNames     bool
	repeat        = 0
	freq          = 30
	httpPort      = 8080
	oidToName     = make(map[string]string)
	nameToOid     = make(map[string]string)
	appdir, _     = osext.ExecutableFolder()
	logDir        = filepath.Join(appdir, "log")
	oidFile       = filepath.Join(appdir, "oids.txt")
	configFile    = filepath.Join(appdir, "config.gcfg")
	errorLog      *os.File
	errorDuration = time.Duration(10 * time.Minute)
	errorPeriod   = errorDuration.String()
	errorMax      = 100
	errorName     string

	cfg = struct {
		Snmp    map[string]*SnmpConfig
		Mibs    map[string]*MibConfig
		Influx  map[string]*InfluxConfig
		HTTP    HTTPConfig
		General GeneralConfig
	}{}
)

func fatal(v ...interface{}) {
	log.SetOutput(os.Stderr)
	log.Fatalln(v...)
}

func (c *SnmpConfig) DebugAction() string {
	debug := make(chan bool)
	c.enabled <- debug
	if <-debug {
		return "disable"
	}
	return "enable"
}

func (c *SnmpConfig) LoadPorts() {
	c.labels = make(map[string]string)
	if len(c.PortFile) == 0 {
		return
	}
	data, err := ioutil.ReadFile(filepath.Join(appdir, c.PortFile))
	if err != nil {
		log.Fatal(err)
	}
	for _, line := range strings.Split(string(data), "\n") {
		// strip comments
		comment := strings.Index(line, "#")
		if comment >= 0 {
			line = line[:comment]
		}
		f := strings.Fields(line)
		if len(f) < 2 {
			continue
		}
		c.labels[f[0]] = f[1]
	}
}

func (c *SnmpConfig) incRequests() {
	atomic.AddInt64(&c.Requests, 1)
}

func (c *SnmpConfig) incGets() {
	atomic.AddInt64(&c.Gets, 1)
}

func (c *SnmpConfig) incErrors() {
	atomic.AddInt64(&c.Errors, 1)
}

func (c *InfluxConfig) incErrors() {
	atomic.AddInt64(&c.Errors, 1)
}

func (c *InfluxConfig) incSent() {
	atomic.AddInt64(&c.Sent, 1)
}

// loads [last_octet]name for device
func (c *SnmpConfig) Translate() {
	client, err := snmpClient(c)
	if err != nil {
		fatal("Client connect error:", err)
	}
	defer client.Conn.Close()
	spew("Looking up column names for:", c.Host)
	pdus, err := client.BulkWalkAll(nameOid)
	if err != nil {
		fatal("SNMP bulkwalk error", err)
	}
	c.asName = make(map[string]string)
	c.asOID = make(map[string]string)
	for _, pdu := range pdus {
		switch pdu.Type {
		case gosnmp.OctetString:
			i := strings.LastIndex(pdu.Name, ".")
			suffix := pdu.Name[i+1:]
			name := string(pdu.Value.([]byte))
			_, ok := c.labels[name]
			if len(c.PortFile) == 0 || ok {
				c.asName[name] = suffix
				c.asOID[suffix] = name
			}
		}
	}
	// make sure we got everything
	for k := range c.labels {
		if _, ok := c.asName[k]; !ok {
			fatal("No OID found for:", k)
		}
	}

}

func spew(x ...interface{}) {
	if verbose {
		fmt.Println(x...)
	}
}

func (c *SnmpConfig) OIDs() {
	if c.mib == nil {
		fatal("NO MIB!")
	}
	c.oids = []string{}
	for _, col := range c.mib.Columns {
		base, ok := nameToOid[col]
		if !ok {
			fatal("no oid for col:", col)
		}
		// just named columns
		if len(c.PortFile) > 0 {
			for k := range c.asOID {
				c.oids = append(c.oids, base+"."+k)
			}
		} else if c.mib.Scalers {
			// or plain old scaler instances
			c.oids = append(c.oids, base+".0")
		} else {
			c.oids = append(c.oids, base)
		}
	}
	if len(c.mib.Columns) > 0 {
		spew("COLUMNS", c.mib.Columns)
		spew(c.oids)
	}
}

// load oid lookup data
func init() {
	data, err := ioutil.ReadFile(oidFile)
	if err != nil {
		log.Fatal(err)
	}
	for _, line := range strings.Split(string(data), "\n") {
		f := strings.Fields(line)
		if len(f) < 2 {
			continue
		}
		nameToOid[f[0]] = f[1]
		oidToName[f[1]] = f[0]
	}
}

func flags() *flag.FlagSet {
	var f flag.FlagSet
	f.BoolVar(&testing, "testing", testing, "print data w/o saving")
	f.BoolVar(&snmpNames, "names", snmpNames, "print column names and exit")
	f.StringVar(&configFile, "config", configFile, "config file")
	f.BoolVar(&verbose, "verbose", verbose, "verbose mode")
	f.IntVar(&repeat, "repeat", repeat, "number of times to repeat")
	f.IntVar(&freq, "freq", freq, "delay (in seconds)")
	f.IntVar(&httpPort, "http", httpPort, "http port")
	f.StringVar(&logDir, "logs", logDir, "log directory")
	f.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		f.VisitAll(func(flag *flag.Flag) {
			format := "%10s: %s\n"
			fmt.Fprintf(os.Stderr, format, "-"+flag.Name, flag.Usage)
		})
		fmt.Fprintf(os.Stderr, "\nAll settings can be set in config file: %s\n", configFile)
		os.Exit(1)

	}
	return &f
}

func init() {
	// parse first time to see if config file is being specified
	f := flags()
	f.Parse(os.Args[1:])
	// now load up config settings
	if _, err := os.Stat(configFile); err != nil {
		log.Fatal(err)
	} else {
		data, err := ioutil.ReadFile(configFile)
		if err != nil {
			log.Fatal(err)
		}
		err = gcfg.ReadStringInto(&cfg, string(data))
		if err != nil {
			log.Fatalf("Failed to parse gcfg data: %s", err)
		}
		httpPort = cfg.HTTP.Port
	}

	if len(cfg.General.LogDir) > 0 {
		logDir = cfg.General.LogDir
	}
	if len(cfg.General.OidFile) > 0 {
		oidFile = cfg.General.OidFile
	}

	for _, s := range cfg.Snmp {
		s.LoadPorts()
		s.debugging = make(chan bool)
		s.enabled = make(chan chan bool)
	}
	var ok bool
	for name, c := range cfg.Snmp {
		if c.mib, ok = cfg.Mibs[name]; !ok {
			if c.mib, ok = cfg.Mibs["*"]; !ok {
				fatal("No mib data found for config:", name)
			}
		}
		c.Translate()
		c.OIDs()
		if c.Freq == 0 {
			c.Freq = freq
		}
	}

	// only run when one needs to see the interface names of the device
	if snmpNames {
		for _, c := range cfg.Snmp {
			fmt.Println("\nSNMP host:", c.Host)
			fmt.Println("=========================================")
			printSnmpNames(c)
		}
		os.Exit(0)
	}

	// re-read cmd line args to override as indicated
	f = flags()
	f.Parse(os.Args[1:])
	os.Mkdir(logDir, 0755)

	// now make sure each snmp device has a db
	for name, c := range cfg.Snmp {
		// default is to use name of snmp config, but it can be overridden
		if len(c.Config) > 0 {
			name = c.Config
		}
		if c.Influx, ok = cfg.Influx[name]; !ok {
			if c.Influx, ok = cfg.Influx["*"]; !ok {
				fatal("No influx config for snmp device:", name)
			}
		}
		c.Influx.Init()
	}

	var ferr error
	errorName = fmt.Sprintf("error.%d.log", cfg.HTTP.Port)
	errorPath := filepath.Join(logDir, errorName)
	errorLog, ferr = os.OpenFile(errorPath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0664)
	if ferr != nil {
		log.Fatal("Can't open error log:", ferr)
	}
}

func errLog(msg string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, msg, args...)
	fmt.Fprintf(errorLog, msg, args...)
}

func errMsg(msg string, err error) {
	now := time.Now()
	errLog("%s\t%s: %s\n", now.Format(layout), msg, err)
}

func main() {
	var wg sync.WaitGroup
	defer func() {
		errorLog.Close()
	}()
	for _, c := range cfg.Snmp {
		wg.Add(1)
		go c.Gather(repeat, &wg)
	}
	if repeat > 0 {
		wg.Wait()
	} else {
		if httpPort > 0 {
			webServer(httpPort)
		} else {
			<-quit
		}
	}
}
<|endoftext|>


<commit_before>package cache

import (
	"container/list"
	"sync"
	"time"
)

type BurnStrategy int

const (
	BurnStrategyRandom    = BurnStrategy(1)
	BurnStrategyOldest    = BurnStrategy(2)
	BurnStrategyOldestLRU = BurnStrategy(3)
)

type Cache struct {
	l        *list.List
	contents map[string]*list.Element
	mutex    *sync.Mutex
	dead     bool
	hits     int64
	misses   int64
	options  CacheOptions
}

type CachedItem struct {
	key   string
	value interface{}
}

type CacheOptions struct {
	MaxEntries     int //If this is set to 0, you must have an expiration time set.
	Upper          int
	BurnStrategy   BurnStrategy
	ExpirationTime time.Duration
	JobInvertal    time.Duration
}

func NewCache(c CacheOptions) *Cache {
	newc := &Cache{}
	newc.options = c
	newc.contents = make(map[string]*list.Element)
	newc.mutex = &sync.Mutex{}
	newc.l = list.New()
	if newc.options.MaxEntries > newc.options.Upper && newc.options.Upper > 0 {
		newc.options.Upper = newc.options.MaxEntries
	}
	return newc
}

func (c *Cache) Set(key string, value interface{}) {
	c.lock()
	defer c.unlock()
	if value == nil {
		c.deleteItem(c.contents[key])
		return
	}
	if c.l.Len()+1 >= c.options.Upper && c.options.Upper > 0 {
		c.burnEntryByStrategy()
	}
	newitem := &CachedItem{}
	newitem.value = value
	newitem.key = key
	e := c.l.PushFront(newitem)
	c.contents[key] = e
	if c.options.ExpirationTime > 0 {
		go c.expireIn(c.options.ExpirationTime, e)
	}
}

func (c *Cache) Get(key string) interface{} {
	c.lock()
	defer c.unlock()
	k, ok := c.contents[key]
	if ok {
		c.hits++
		if c.options.BurnStrategy == BurnStrategyOldestLRU {
			c.l.MoveToFront(k)
		}
		return k.Value.(*CachedItem).value
	} else {
		c.misses++
		return nil
	}
}

func (c *Cache) Start() {
	if c.options.JobInvertal > 0 {
		go c.runner()
	}
}

func (c *Cache) Stop() {
	c.dead = true
}

func (c *Cache) Hits() int64 {
	return c.hits
}

func (c *Cache) Misses() int64 {
	return c.misses
}

func (c *Cache) RemoveItem(key string) {
	c.lock()
	defer c.unlock()
	c.deleteItem(c.contents[key])
}

func (c *Cache) Trim(num int) {
	c.lock()
	defer c.unlock()
	for num > 0 && c.l.Len() > 0 {
		num--
		c.burnEntryByStrategy()
	}
}

func (c *Cache) Bump(key string) {
	k, ok := c.contents[key]
	if ok {
		c.l.MoveToFront(k)
	}
}

//For now we, can't use this internally since it is mutex'd
func (c *Cache) Len() int {
	c.lock()
	defer c.unlock()
	return c.l.Len()
}

//private functions
func (c *Cache) burnEntryByStrategy() {
	if c.options.BurnStrategy == BurnStrategyOldest || c.options.BurnStrategy == BurnStrategyOldestLRU {
		c.burnEntryByOldest()
	} else {
		c.burnEntryByRandom()
	}
}

func (c *Cache) burnEntryByRandom() {
	for _, a := range c.contents {
		c.deleteItem(a)
		break
	}
}

func (c *Cache) burnEntryByOldest() {
	i := c.l.Back()
	if i != nil {
		c.deleteItem(i)
	}
}

func (c *Cache) lock() {
	c.mutex.Lock()
}

func (c *Cache) unlock() {
	c.mutex.Unlock()
}

func (c *Cache) runner() {
	for {
		time.Sleep(c.options.JobInvertal)
		if c.dead {
			break
		}
		c.lock()
		if c.options.MaxEntries > 0 {
			for len(c.contents) > c.options.MaxEntries {
				c.burnEntryByStrategy()
			}
		}
		c.unlock()
	}
}

func (c *Cache) expireIn(t time.Duration, i *list.Element) {
	time.Sleep(t)
	if c.dead {
		return
	}
	c.lock()
	defer c.unlock()
	c.deleteItem(i)
}

func (c *Cache) deleteItem(i *list.Element) {
	c.l.Remove(i)
	delete(c.contents, i.Value.(*CachedItem).key)
}
<commit_msg>Add safe range<commit_after>package cache

import (
	"container/list"
	"sync"
	"time"
)

type BurnStrategy int

const (
	BurnStrategyRandom    = BurnStrategy(1)
	BurnStrategyOldest    = BurnStrategy(2)
	BurnStrategyOldestLRU = BurnStrategy(3)
)

type Cache struct {
	l        *list.List
	contents map[string]*list.Element
	mutex    *sync.Mutex
	dead     bool
	hits     int64
	misses   int64
	options  CacheOptions
}

type CachedItem struct {
	key   string
	value interface{}
}

type CacheOptions struct {
	MaxEntries     int //If this is set to 0, you must have an expiration time set.
	Upper          int
	BurnStrategy   BurnStrategy
	ExpirationTime time.Duration
	JobInvertal    time.Duration
	SafeRange      int
}

func NewCache(c CacheOptions) *Cache {
	newc := &Cache{}
	newc.options = c
	newc.contents = make(map[string]*list.Element)
	newc.mutex = &sync.Mutex{}
	newc.l = list.New()
	if newc.options.MaxEntries > newc.options.Upper && newc.options.Upper > 0 {
		newc.options.Upper = newc.options.MaxEntries
	}
	return newc
}

func (c *Cache) Set(key string, value interface{}) {
	c.lock()
	defer c.unlock()
	if value == nil {
		c.deleteItem(c.contents[key])
		return
	}
	if c.l.Len()+1 >= c.options.Upper && c.options.Upper > 0 {
		c.burnEntryByStrategy()
	}
	newitem := &CachedItem{}
	newitem.value = value
	newitem.key = key
	e := c.l.PushFront(newitem)
	c.contents[key] = e
	if c.options.ExpirationTime > 0 {
		go c.expireIn(c.options.ExpirationTime, e)
	}
}

func (c *Cache) Get(key string) interface{} {
	c.lock()
	defer c.unlock()
	k, ok := c.contents[key]
	if ok {
		c.hits++
		if c.options.BurnStrategy == BurnStrategyOldestLRU {
			c.l.MoveToFront(k)
		}
		return k.Value.(*CachedItem).value
	} else {
		c.misses++
		return nil
	}
}

func (c *Cache) Start() {
	if c.options.JobInvertal > 0 {
		go c.runner()
	}
}

func (c *Cache) Stop() {
	c.dead = true
}

func (c *Cache) Hits() int64 {
	return c.hits
}

func (c *Cache) Misses() int64 {
	return c.misses
}

func (c *Cache) RemoveItem(key string) {
	c.lock()
	defer c.unlock()
	c.deleteItem(c.contents[key])
}

func (c *Cache) Trim(num int) {
	c.lock()
	defer c.unlock()
	for num > 0 && c.l.Len() > 0 {
		num--
		c.burnEntryByStrategy()
	}
}

func (c *Cache) Bump(key string) {
	k, ok := c.contents[key]
	if ok {
		c.l.MoveToFront(k)
	}
}

//For now we, can't use this internally since it is mutex'd
func (c *Cache) Len() int {
	c.lock()
	defer c.unlock()
	return c.l.Len()
}

//private functions
func (c *Cache) burnEntryByStrategy() {
	if c.options.BurnStrategy == BurnStrategyOldest || c.options.BurnStrategy == BurnStrategyOldestLRU {
		c.burnEntryByOldest()
	} else {
		c.burnEntryByRandom()
	}
}

func (c *Cache) burnEntryByRandom() {
	for _, a := range c.contents {
		c.deleteItem(a)
		break
	}
}

func (c *Cache) burnEntryByOldest() {
	i := c.l.Back()
	if i != nil {
		c.deleteItem(i)
	}
}

func (c *Cache) lock() {
	c.mutex.Lock()
}

func (c *Cache) unlock() {
	c.mutex.Unlock()
}

func (c *Cache) runner() {
	for {
		time.Sleep(c.options.JobInvertal)
		if c.dead {
			break
		}
		c.lock()
		if c.options.MaxEntries > 0 {
			for len(c.contents) > c.options.MaxEntries-c.options.SafeRange {
				c.burnEntryByStrategy()
			}
		}
		c.unlock()
	}
}

func (c *Cache) expireIn(t time.Duration, i *list.Element) {
	time.Sleep(t)
	if c.dead {
		return
	}
	c.lock()
	defer c.unlock()
	c.deleteItem(i)
}

func (c *Cache) deleteItem(i *list.Element) {
	c.l.Remove(i)
	delete(c.contents, i.Value.(*CachedItem).key)
}
<|endoftext|>


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

import (
	"encoding/json"
	"fmt"
	"github.com/confluentinc/confluent-kafka-go/kafka"
	"gopkg.in/alecthomas/kingpin.v2"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	verbosity     = 1
	exit_eof      = false
	eof_cnt       = 0
	partition_cnt = 0
	key_delim     = ""
	sigs          chan os.Signal
)

func send(name string, msg map[string]interface{}) {
	if msg == nil {
		msg = make(map[string]interface{})
	}
	msg["name"] = name
	msg["_time"] = time.Now().Unix()
	b, err := json.Marshal(msg)
	if err != nil {
		panic(err)
	}
	fmt.Println(string(b))
}

func partitions_to_map(partitions []kafka.TopicPartition) []map[string]interface{} {
	parts := make([]map[string]interface{}, len(partitions))
	for i, tp := range partitions {
		parts[i] = map[string]interface{}{"topic": *tp.Topic, "partition": tp.Partition}
	}
	return parts
}

func send_partitions(name string, partitions []kafka.TopicPartition) {

	msg := make(map[string]interface{})
	msg["partitions"] = partitions_to_map(partitions)

	send(name, msg)
}

type assigned_partition struct {
	tp            kafka.TopicPartition
	consumed_msgs int
	min_offset    int64
	max_offset    int64
}

func assignment_key(tp kafka.TopicPartition) string {
	return fmt.Sprintf("%s-%d", *tp.Topic, tp.Partition)
}

func find_assignment(tp kafka.TopicPartition) *assigned_partition {
	a, ok := state.curr_assignment[assignment_key(tp)]
	if !ok {
		return nil
	}
	return a
}

func add_assignment(tp kafka.TopicPartition) {
	state.curr_assignment[assignment_key(tp)] = &assigned_partition{tp: tp, min_offset: -1, max_offset: -1}
}

func clear_curr_assignment() {
	state.curr_assignment = make(map[string]*assigned_partition)
}

type comm_state struct {
	consumed_msgs                int
	consumed_msgs_last_reported  int
	consumed_msgs_at_last_commit int
	curr_assignment              map[string]*assigned_partition
	max_messages                 int
	auto_commit                  bool
	async_commit                 bool
	c                            *kafka.Consumer
}

var state comm_state

func send_records_consumed(immediate bool) {
	if len(state.curr_assignment) == 0 ||
		(!immediate && state.consumed_msgs_last_reported+1000 > state.consumed_msgs) {
		return
	}

	msg := map[string]interface{}{}
	msg["count"] = state.consumed_msgs - state.consumed_msgs_last_reported
	parts := make([]map[string]interface{}, len(state.curr_assignment))
	i := 0
	for _, a := range state.curr_assignment {
		if a.min_offset == -1 {
			// Skip partitions that havent had any messages since last time.
			// This is to circumvent some minOffset checks in kafkatest.
			continue
		}
		parts[i] = map[string]interface{}{"topic": *a.tp.Topic,
			"partition":     a.tp.Partition,
			"consumed_msgs": a.consumed_msgs,
			"minOffset":     a.min_offset,
			"maxOffset":     a.max_offset}
		a.min_offset = -1
		i += 1
	}
	msg["partitions"] = parts[0:i]

	send("records_consumed", msg)

	state.consumed_msgs_last_reported = state.consumed_msgs
}

// do_commit commits every 1000 messages or whenever there is a consume timeout, or when immediate==true
func do_commit(immediate bool, async bool) {
	if state.auto_commit || !immediate ||
		state.consumed_msgs_at_last_commit+1000 > state.consumed_msgs {
		return
	}

	// Make sure we report consumption before commit,
	// otherwise tests may fail because of commit > consumed
	if state.consumed_msgs_at_last_commit < state.consumed_msgs {
		send_records_consumed(true)
	}

	async = state.async_commit

	fmt.Fprintf(os.Stderr, "%% Committing %d messages (async=%v)\n",
		state.consumed_msgs-state.consumed_msgs_at_last_commit, async)

	err := state.c.Commit(async)
	if err != nil {
		kerr, ok := err.(kafka.KafkaError)
		if ok && kerr.Code() == kafka.ERR__NO_OFFSET {
			fmt.Fprintf(os.Stderr, "%% No offsets to commit\n")
		} else {
			panic(fmt.Sprintf("Commit failed: %s", err))
		}
	}

	state.consumed_msgs_at_last_commit = state.consumed_msgs
}

// returns false when consumer should terminate, else true to keep running.
func handle_msg(m *kafka.Message) bool {
	if verbosity >= 2 {
		fmt.Fprintf(os.Stderr, "%% Message receved: %v:\n", m.TopicPartition)
	}

	a := find_assignment(m.TopicPartition)
	if a == nil {
		fmt.Fprintf(os.Stderr, "%% Received message on unassigned partition: %v\n", m.TopicPartition)
		return true
	}

	a.consumed_msgs += 1
	offset := int64(m.TopicPartition.Offset)
	if a.min_offset == -1 {
		a.min_offset = offset
	}
	if a.max_offset < offset {
		a.max_offset = offset
	}

	state.consumed_msgs += 1

	send_records_consumed(false)
	do_commit(false, state.async_commit)

	if state.max_messages > 0 && state.consumed_msgs >= state.max_messages {
		// ignore extra messages
		return false
	}

	return true

}

func run_consumer(config *kafka.ConfigMap, topic string) {
	c, err := kafka.NewConsumer(config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create consumer: %s\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(os.Stderr, "%% Created Consumer %v\n", c)
	state.c = c

	c.Subscribe(topic, nil)

	send("startup_complete", nil)
	run := true

	for run == true {
		select {

		case sig := <-sigs:
			fmt.Fprintf(os.Stderr, "%% Terminating on signal %v\n", sig)
			run = false

		case ev := <-c.Events:
			switch e := ev.(type) {
			case kafka.AssignedPartitions:
				if len(state.curr_assignment) > 0 {
					panic(fmt.Sprintf("Assign: curr_assignment should have been empty: %v", state.curr_assignment))
				}
				state.curr_assignment = make(map[string]*assigned_partition)
				for _, tp := range e.Partitions {
					add_assignment(tp)
				}
				send_partitions("partitions_assigned", e.Partitions)
				c.Assign(e.Partitions)

			case kafka.RevokedPartitions:
				send_records_consumed(true)
				do_commit(true, false)
				send_partitions("partitions_revoked", e.Partitions)
				clear_curr_assignment()
				c.Unassign()

			case *kafka.Message:
				run = handle_msg(e)

			case kafka.KafkaError:
				fmt.Fprintf(os.Stderr, "%% Error: %v\n", e)
				run = false
			default:
				fmt.Fprintf(os.Stderr, "%% Unhandled event %T ignored: %v\n", e, e)
			}

		case _ = <-time.After(1 * time.Second):
			// Report consumed messages
			send_records_consumed(true)
			// Commit on timeout as well (not just every 1000 messages)
			do_commit(true, state.async_commit)
		}
	}

	fmt.Fprintf(os.Stderr, "%% Consumer shutting down\n")

	send_records_consumed(true)

	if !state.auto_commit {
		do_commit(true, false)
	}

	fmt.Fprintf(os.Stderr, "%% Closing consumer\n")

	c.Close()

	send("shutdown_complete", nil)
}

func main() {
	sigs = make(chan os.Signal)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	// Default config
	conf := kafka.ConfigMap{"default.topic.config": kafka.ConfigMap{"auto.offset.reset": "earliest"}}

	/* Required options */
	group := kingpin.Flag("group-id", "Consumer group").Required().String()
	topic := kingpin.Flag("topic", "Topic to consume").Required().String()
	brokers := kingpin.Flag("broker-list", "Bootstrap broker(s)").Required().String()
	session_timeout := kingpin.Flag("session-timeout", "Session timeout").Required().Int()

	/* Optionals */
	enable_autocommit := kingpin.Flag("enable-autocommit", "Enable auto-commit").Default("true").Bool()
	max_messages := kingpin.Flag("max-messages", "Max messages to consume").Default("10000000").Int()
	java_assignment_strategy := kingpin.Flag("assignment-strategy", "Assignment strategy (Java class name)").String()
	config_file := kingpin.Flag("consumer.config", "Config file").File()
	debug := kingpin.Flag("debug", "Debug flags").String()

	kingpin.Parse()

	conf["bootstrap.servers"] = *brokers
	conf["group.id"] = *group
	conf["session.timeout.ms"] = *session_timeout
	conf["enable.auto.commit"] = *enable_autocommit
	fmt.Println("Config: ", conf)

	if len(*debug) > 0 {
		conf["debug"] = *debug
	}

	/* Convert Java assignment strategy to librdkafka one.
	 * "[java.class.path.]Strategy[Assignor]" -> "strategy" */
	if java_assignment_strategy != nil && len(*java_assignment_strategy) > 0 {
		s := strings.Split(*java_assignment_strategy, ".")
		strategy := strings.ToLower(strings.TrimSuffix(s[len(s)-1], "Assignor"))
		conf["partition.assignment.strategy"] = strategy
		fmt.Fprintf(os.Stderr, "%% Mapped %s -> %s\n",
			*java_assignment_strategy, conf["partition.assignment.strategy"])
	}

	if *config_file != nil {
		fmt.Fprintf(os.Stderr, "%% Ignoring config file %s\n", *config_file)
	}

	conf["go.events.channel.enable"] = true
	conf["go.application.rebalance.enable"] = true

	state.auto_commit = *enable_autocommit
	state.max_messages = *max_messages
	run_consumer((*kafka.ConfigMap)(&conf), *topic)

}
<commit_msg>go_verifiable_consumer: commit improvements<commit_after>/**
 * Copyright 2016 Confluent Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package main

import (
	"encoding/json"
	"fmt"
	"github.com/confluentinc/confluent-kafka-go/kafka"
	"gopkg.in/alecthomas/kingpin.v2"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	verbosity     = 1
	exit_eof      = false
	eof_cnt       = 0
	partition_cnt = 0
	key_delim     = ""
	sigs          chan os.Signal
)

func send(name string, msg map[string]interface{}) {
	if msg == nil {
		msg = make(map[string]interface{})
	}
	msg["name"] = name
	msg["_time"] = time.Now().Unix()
	b, err := json.Marshal(msg)
	if err != nil {
		panic(err)
	}
	fmt.Println(string(b))
}

func partitions_to_map(partitions []kafka.TopicPartition) []map[string]interface{} {
	parts := make([]map[string]interface{}, len(partitions))
	for i, tp := range partitions {
		parts[i] = map[string]interface{}{"topic": *tp.Topic, "partition": tp.Partition, "offset": tp.Offset}
	}
	return parts
}

func send_offsets_committed(offsets []kafka.TopicPartition, err error) {
	if len(state.curr_assignment) == 0 {
		// Dont emit offsets_committed if there is no current assignment
		// This happens when auto_commit is enabled since we also
		// force a manual commit on rebalance to make sure
		// offsets_committed is emitted prior to partitions_revoked,
		// so the builtin auto committer will also kick in and post
		// this later OffsetsCommitted event which we simply ignore..
		fmt.Fprintf(os.Stderr, "%% Ignore OffsetsCommitted(%v) without a valid assignment\n", err)
		return
	}
	msg := make(map[string]interface{})

	if err != nil {
		msg["success"] = false
		msg["error"] = fmt.Sprintf("%v", err)

		kerr, ok := err.(kafka.KafkaError)
		if ok && kerr.Code() == kafka.ERR__NO_OFFSET {
			fmt.Fprintf(os.Stderr, "%% No offsets to commit\n")
			return
		}

		fmt.Fprintf(os.Stderr, "%% Commit failed: %v", msg["error"])
	} else {
		msg["success"] = true

	}

	if offsets != nil {
		msg["offsets"] = partitions_to_map(offsets)
	}

	// Make sure we report consumption before commit,
	// otherwise tests may fail because of commit > consumed
	send_records_consumed(true)

	send("offsets_committed", msg)
}

func send_partitions(name string, partitions []kafka.TopicPartition) {

	msg := make(map[string]interface{})
	msg["partitions"] = partitions_to_map(partitions)

	send(name, msg)
}

type assigned_partition struct {
	tp            kafka.TopicPartition
	consumed_msgs int
	min_offset    int64
	max_offset    int64
}

func assignment_key(tp kafka.TopicPartition) string {
	return fmt.Sprintf("%s-%d", *tp.Topic, tp.Partition)
}

func find_assignment(tp kafka.TopicPartition) *assigned_partition {
	a, ok := state.curr_assignment[assignment_key(tp)]
	if !ok {
		return nil
	}
	return a
}

func add_assignment(tp kafka.TopicPartition) {
	state.curr_assignment[assignment_key(tp)] = &assigned_partition{tp: tp, min_offset: -1, max_offset: -1}
}

func clear_curr_assignment() {
	state.curr_assignment = make(map[string]*assigned_partition)
}

type comm_state struct {
	consumed_msgs                int
	consumed_msgs_last_reported  int
	consumed_msgs_at_last_commit int
	curr_assignment              map[string]*assigned_partition
	max_messages                 int
	auto_commit                  bool
	async_commit                 bool
	c                            *kafka.Consumer
}

var state comm_state

func send_records_consumed(immediate bool) {
	if len(state.curr_assignment) == 0 ||
		(!immediate && state.consumed_msgs_last_reported+1000 > state.consumed_msgs) {
		return
	}

	msg := map[string]interface{}{}
	msg["count"] = state.consumed_msgs - state.consumed_msgs_last_reported
	parts := make([]map[string]interface{}, len(state.curr_assignment))
	i := 0
	for _, a := range state.curr_assignment {
		if a.min_offset == -1 {
			// Skip partitions that havent had any messages since last time.
			// This is to circumvent some minOffset checks in kafkatest.
			continue
		}
		parts[i] = map[string]interface{}{"topic": *a.tp.Topic,
			"partition":     a.tp.Partition,
			"consumed_msgs": a.consumed_msgs,
			"minOffset":     a.min_offset,
			"maxOffset":     a.max_offset}
		a.min_offset = -1
		i += 1
	}
	msg["partitions"] = parts[0:i]

	send("records_consumed", msg)

	state.consumed_msgs_last_reported = state.consumed_msgs
}

// do_commit commits every 1000 messages or whenever there is a consume timeout, or when immediate==true
func do_commit(immediate bool, async bool) {
	if !immediate &&
		(state.auto_commit ||
			state.consumed_msgs_at_last_commit+1000 > state.consumed_msgs) {
		return
	}

	async = state.async_commit

	fmt.Fprintf(os.Stderr, "%% Committing %d messages (async=%v)\n",
		state.consumed_msgs-state.consumed_msgs_at_last_commit, async)

	state.consumed_msgs_at_last_commit = state.consumed_msgs

	var wait_committed chan bool

	if !async {
		wait_committed = make(chan bool)
	}

	go func() {
		offsets, err := state.c.Commit()

		send_offsets_committed(offsets, err)

		if !async {
			close(wait_committed)
		}
	}()

	if !async {
		_, _ = <-wait_committed
	}
}

// returns false when consumer should terminate, else true to keep running.
func handle_msg(m *kafka.Message) bool {
	if verbosity >= 2 {
		fmt.Fprintf(os.Stderr, "%% Message receved: %v:\n", m.TopicPartition)
	}

	a := find_assignment(m.TopicPartition)
	if a == nil {
		fmt.Fprintf(os.Stderr, "%% Received message on unassigned partition: %v\n", m.TopicPartition)
		return true
	}

	a.consumed_msgs += 1
	offset := int64(m.TopicPartition.Offset)
	if a.min_offset == -1 {
		a.min_offset = offset
	}
	if a.max_offset < offset {
		a.max_offset = offset
	}

	state.consumed_msgs += 1

	send_records_consumed(false)
	do_commit(false, state.async_commit)

	if state.max_messages > 0 && state.consumed_msgs >= state.max_messages {
		// ignore extra messages
		return false
	}

	return true

}

func run_consumer(config *kafka.ConfigMap, topic string) {
	c, err := kafka.NewConsumer(config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create consumer: %s\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(os.Stderr, "%% Created Consumer %v\n", c)
	state.c = c

	c.Subscribe(topic, nil)

	send("startup_complete", nil)
	run := true

	for run == true {
		select {

		case sig := <-sigs:
			fmt.Fprintf(os.Stderr, "%% Terminating on signal %v\n", sig)
			run = false

		case ev := <-c.Events:
			switch e := ev.(type) {
			case kafka.AssignedPartitions:
				if len(state.curr_assignment) > 0 {
					panic(fmt.Sprintf("Assign: curr_assignment should have been empty: %v", state.curr_assignment))
				}
				state.curr_assignment = make(map[string]*assigned_partition)
				for _, tp := range e.Partitions {
					add_assignment(tp)
				}
				send_partitions("partitions_assigned", e.Partitions)
				c.Assign(e.Partitions)

			case kafka.RevokedPartitions:
				send_records_consumed(true)
				do_commit(true, false)
				send_partitions("partitions_revoked", e.Partitions)
				clear_curr_assignment()
				c.Unassign()

			case *kafka.Message:
				run = handle_msg(e)

			case kafka.KafkaError:
				fmt.Fprintf(os.Stderr, "%% Error: %v\n", e)
				run = false
			default:
				fmt.Fprintf(os.Stderr, "%% Unhandled event %T ignored: %v\n", e, e)
			}

		case _ = <-time.After(1 * time.Second):
			// Report consumed messages
			send_records_consumed(true)
			// Commit on timeout as well (not just every 1000 messages)
			do_commit(true, state.async_commit)
		}
	}

	fmt.Fprintf(os.Stderr, "%% Consumer shutting down\n")

	send_records_consumed(true)

	if !state.auto_commit {
		do_commit(true, false)
	}

	fmt.Fprintf(os.Stderr, "%% Closing consumer\n")

	c.Close()

	send("shutdown_complete", nil)
}

func main() {
	sigs = make(chan os.Signal)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	// Default config
	conf := kafka.ConfigMap{"default.topic.config": kafka.ConfigMap{"auto.offset.reset": "earliest"}}

	/* Required options */
	group := kingpin.Flag("group-id", "Consumer group").Required().String()
	topic := kingpin.Flag("topic", "Topic to consume").Required().String()
	brokers := kingpin.Flag("broker-list", "Bootstrap broker(s)").Required().String()
	session_timeout := kingpin.Flag("session-timeout", "Session timeout").Required().Int()

	/* Optionals */
	enable_autocommit := kingpin.Flag("enable-autocommit", "Enable auto-commit").Default("true").Bool()
	max_messages := kingpin.Flag("max-messages", "Max messages to consume").Default("10000000").Int()
	java_assignment_strategy := kingpin.Flag("assignment-strategy", "Assignment strategy (Java class name)").String()
	config_file := kingpin.Flag("consumer.config", "Config file").File()
	debug := kingpin.Flag("debug", "Debug flags").String()

	kingpin.Parse()

	conf["bootstrap.servers"] = *brokers
	conf["group.id"] = *group
	conf["session.timeout.ms"] = *session_timeout
	conf["enable.auto.commit"] = *enable_autocommit
	fmt.Println("Config: ", conf)

	if len(*debug) > 0 {
		conf["debug"] = *debug
	}

	/* Convert Java assignment strategy to librdkafka one.
	 * "[java.class.path.]Strategy[Assignor]" -> "strategy" */
	if java_assignment_strategy != nil && len(*java_assignment_strategy) > 0 {
		s := strings.Split(*java_assignment_strategy, ".")
		strategy := strings.ToLower(strings.TrimSuffix(s[len(s)-1], "Assignor"))
		conf["partition.assignment.strategy"] = strategy
		fmt.Fprintf(os.Stderr, "%% Mapped %s -> %s\n",
			*java_assignment_strategy, conf["partition.assignment.strategy"])
	}

	if *config_file != nil {
		fmt.Fprintf(os.Stderr, "%% Ignoring config file %s\n", *config_file)
	}

	conf["go.events.channel.enable"] = true
	conf["go.application.rebalance.enable"] = true

	state.auto_commit = *enable_autocommit
	state.max_messages = *max_messages
	run_consumer((*kafka.ConfigMap)(&conf), *topic)

}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"database/sql"
	"fmt"
	"html/template"
	"log"
	mathRand "math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/daaku/go.httpgzip"
	"golang.org/x/crypto/acme/autocert"
)

//go:generate esc -private -o bindata.go -prefix static ./views ./static/
var templates = template.New("").Funcs(template.FuncMap{"add": func(a, b int) int { return a + b }})
var templateSources = []string{
	"/views/main.tmpl",
	"/views/_util.tmpl",
	"/views/geekhack.tmpl",
}

var hostname_whitelist = []string{
	"www.sadbox.org", "sadbox.org", "mail.sadbox.org",
	"www.sadbox.es", "sadbox.es",
	"www.geekwhack.org", "geekwhack.org",
}

var sadboxDB *sql.DB

var channels []channel

type channel struct {
	LinkName    string
	ChannelName string
}

type WebsiteName struct {
	Title, Brand string
}

type Main struct {
	Channels []channel
}

type TemplateContext struct {
	Geekhack *Geekhack
	Webname  *WebsiteName
	Main     *Main
}

func NewContext(r *http.Request) *TemplateContext {
	title := "sadbox \u00B7 org"
	host := "sadbox.org"
	for _, v := range hostname_whitelist {
		if v == r.Host {
			trimmed := strings.TrimPrefix(r.Host, "www.")
			title = strings.Replace(trimmed, ".", " \u00B7 ", -1)
			host = trimmed
			break
		}
	}
	return &TemplateContext{
		Webname: &WebsiteName{
			title,
			host,
		},
	}
}

func CatchPanic(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered from panic: %v", r)
				http.Error(w, "Something went wrong!", http.StatusInternalServerError)
			}
		}()

		handler.ServeHTTP(w, r)
	})
}

func SendToHTTPS(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Connection", "close")
	http.Redirect(w, r, "https://"+r.Host+r.RequestURI, http.StatusMovedPermanently)
}

func RedirectToHTTPS(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			SendToHTTPS(w, r)
			return
		}

		ip := net.ParseIP(host)
		if ip == nil {
			SendToHTTPS(w, r)
			return
		}

		if !ip.IsLoopback() {
			SendToHTTPS(w, r)
			return
		}

		handler.ServeHTTP(w, r)
	})
}

func AddHeaders(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "max-age=120")
		w.Header().Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains; preload")
		w.Header().Set("Content-Security-Policy", "default-src 'self'; object-src 'none'; frame-ancestors 'none'")
		w.Header().Set("X-Frame-Options", "DENY")
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("Referrer-Policy", "no-referrer, same-origin")
		w.Header().Set("X-XSS-Protection", "1; mode=block")
		handler.ServeHTTP(w, r)
	})
}

func Log(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		remoteHost := r.Header.Get("X-Forwarded-For")
		if remoteHost == "" {
			remoteHost = r.RemoteAddr
		}
		log.Printf("%s %s %s", remoteHost, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}

type sessionKeys struct {
	keys [][32]byte
	buf  []byte
	tls  *tls.Config
}

func NewSessionKeys(tls *tls.Config) *sessionKeys {
	sk := &sessionKeys{
		buf: make([]byte, 32),
		tls: tls,
	}
	sk.addNewKey()
	return sk
}

func (sk *sessionKeys) addNewKey() {
	rand.Read(sk.buf)
	newKey := [32]byte{}
	copy(newKey[:], sk.buf)
	sk.keys = append([][32]byte{newKey}, sk.keys...)
	if len(sk.keys) > 4 {
		sk.keys = sk.keys[:5]
	}
	log.Println("Rotating Session Keys")
	sk.tls.SetSessionTicketKeys(sk.keys)
}

func (sk *sessionKeys) Spin() {
	for range time.Tick(24 * time.Hour) {
		sk.addNewKey()
	}
}

func main() {
	log.Println("Starting sadbox.org")

	for _, filename := range templateSources {
		template.Must(templates.Parse(_escFSMustString(false, filename)))
	}

	var err error
	sadboxDB, err = sql.Open("mysql", os.Getenv("GEEKHACK_DB"))
	if err != nil {
		log.Fatal(err)
	}
	defer sadboxDB.Close()

	err = sadboxDB.Ping()
	if err != nil {
		log.Fatal(err)
	}

	staticFileServer := http.FileServer(_escFS(false))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/" {
			ctx := NewContext(r)
			ctx.Main = &Main{channels}
			if err := templates.ExecuteTemplate(w, "main", ctx); err != nil {
				log.Println(err)
			}
		} else {
			staticFileServer.ServeHTTP(w, r)
		}
	})

	mathRand.Seed(time.Now().UnixNano())

	rows, err := sadboxDB.Query(`SELECT Channel from Channels;`)
	if err != nil {
		log.Fatal(err)
	}
	for rows.Next() {
		var channelName string
		err := rows.Scan(&channelName)
		if err != nil {
			log.Fatal(err)
		}

		chanInfo := channel{fmt.Sprintf("/%s/", strings.Trim(channelName, "#")), channelName}

		ircChanHandler, err := NewIRCChannel(chanInfo)
		if err != nil {
			log.Fatal(err)
		}
		http.Handle(chanInfo.LinkName, ircChanHandler)

		channels = append(channels, chanInfo)
	}

	// Redirects to the right URL so I don't break old links
	http.Handle("/ghstats", http.RedirectHandler("/geekhack/", http.StatusMovedPermanently))
	http.Handle("/geekhack", http.RedirectHandler("/geekhack/", http.StatusMovedPermanently))

	// This will redirect people to the gmail page
	http.Handle("mail.sadbox.org/", http.RedirectHandler("https://mail.google.com/a/sadbox.org", http.StatusFound))

	localhost_znc, err := url.Parse("http://127.0.0.1:6698")
	if err != nil {
		log.Fatal(err)
	}
	http.Handle("/znc/", httputil.NewSingleHostReverseProxy(localhost_znc))
	http.Handle("/znc", http.RedirectHandler("/znc/", http.StatusMovedPermanently))

	servemux := httpgzip.NewHandler(
		CatchPanic(
			Log(
				AddHeaders(http.DefaultServeMux))))

	httpSrv := &http.Server{
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
		Handler:      RedirectToHTTPS(servemux),
	}
	go func() { log.Fatal(httpSrv.ListenAndServe()) }()

	m := autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		Cache:      autocert.DirCache("/home/sadbox-web/cert-cache"),
		HostPolicy: autocert.HostWhitelist(hostname_whitelist...),
		Email:      "blue6249@gmail.com",
		ForceRSA:   true,
	}

	var certs []tls.Certificate
	sadboxCert, err := tls.LoadX509KeyPair("/home/sadbox-web/cert-cache/sadbox.org",
		"/home/sadbox-web/cert-cache/sadbox.org")

	if err == nil {
		certs = append(certs, sadboxCert)
	}

	tlsconfig := &tls.Config{
		PreferServerCipherSuites: true,
		GetCertificate:           m.GetCertificate,
		Certificates:             certs,
	}

	go NewSessionKeys(tlsconfig).Spin()

	server := &http.Server{
		Addr:         ":https",
		Handler:      servemux,
		TLSConfig:    tlsconfig,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
	}
	log.Fatal(server.ListenAndServeTLS("", ""))
}
<commit_msg>Move to NYTimes/gziphandler since the other one is unmaintained<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"database/sql"
	"fmt"
	"html/template"
	"log"
	mathRand "math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/NYTimes/gziphandler"
	"golang.org/x/crypto/acme/autocert"
)

//go:generate esc -private -o bindata.go -prefix static ./views ./static/
var templates = template.New("").Funcs(template.FuncMap{"add": func(a, b int) int { return a + b }})
var templateSources = []string{
	"/views/main.tmpl",
	"/views/_util.tmpl",
	"/views/geekhack.tmpl",
}

var hostname_whitelist = []string{
	"www.sadbox.org", "sadbox.org", "mail.sadbox.org",
	"www.sadbox.es", "sadbox.es",
	"www.geekwhack.org", "geekwhack.org",
}

var sadboxDB *sql.DB

var channels []channel

type channel struct {
	LinkName    string
	ChannelName string
}

type WebsiteName struct {
	Title, Brand string
}

type Main struct {
	Channels []channel
}

type TemplateContext struct {
	Geekhack *Geekhack
	Webname  *WebsiteName
	Main     *Main
}

func NewContext(r *http.Request) *TemplateContext {
	title := "sadbox \u00B7 org"
	host := "sadbox.org"
	for _, v := range hostname_whitelist {
		if v == r.Host {
			trimmed := strings.TrimPrefix(r.Host, "www.")
			title = strings.Replace(trimmed, ".", " \u00B7 ", -1)
			host = trimmed
			break
		}
	}
	return &TemplateContext{
		Webname: &WebsiteName{
			title,
			host,
		},
	}
}

func CatchPanic(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered from panic: %v", r)
				http.Error(w, "Something went wrong!", http.StatusInternalServerError)
			}
		}()

		handler.ServeHTTP(w, r)
	})
}

func SendToHTTPS(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Connection", "close")
	http.Redirect(w, r, "https://"+r.Host+r.RequestURI, http.StatusMovedPermanently)
}

func RedirectToHTTPS(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			SendToHTTPS(w, r)
			return
		}

		ip := net.ParseIP(host)
		if ip == nil {
			SendToHTTPS(w, r)
			return
		}

		if !ip.IsLoopback() {
			SendToHTTPS(w, r)
			return
		}

		handler.ServeHTTP(w, r)
	})
}

func AddHeaders(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "max-age=120")
		w.Header().Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains; preload")
		w.Header().Set("Content-Security-Policy", "default-src 'self'; object-src 'none'; frame-ancestors 'none'")
		w.Header().Set("X-Frame-Options", "DENY")
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("Referrer-Policy", "no-referrer, same-origin")
		w.Header().Set("X-XSS-Protection", "1; mode=block")
		handler.ServeHTTP(w, r)
	})
}

func Log(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		remoteHost := r.Header.Get("X-Forwarded-For")
		if remoteHost == "" {
			remoteHost = r.RemoteAddr
		}
		log.Printf("%s %s %s", remoteHost, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}

type sessionKeys struct {
	keys [][32]byte
	buf  []byte
	tls  *tls.Config
}

func NewSessionKeys(tls *tls.Config) *sessionKeys {
	sk := &sessionKeys{
		buf: make([]byte, 32),
		tls: tls,
	}
	sk.addNewKey()
	return sk
}

func (sk *sessionKeys) addNewKey() {
	rand.Read(sk.buf)
	newKey := [32]byte{}
	copy(newKey[:], sk.buf)
	sk.keys = append([][32]byte{newKey}, sk.keys...)
	if len(sk.keys) > 4 {
		sk.keys = sk.keys[:5]
	}
	log.Println("Rotating Session Keys")
	sk.tls.SetSessionTicketKeys(sk.keys)
}

func (sk *sessionKeys) Spin() {
	for range time.Tick(24 * time.Hour) {
		sk.addNewKey()
	}
}

func main() {
	log.Println("Starting sadbox.org")

	for _, filename := range templateSources {
		template.Must(templates.Parse(_escFSMustString(false, filename)))
	}

	var err error
	sadboxDB, err = sql.Open("mysql", os.Getenv("GEEKHACK_DB"))
	if err != nil {
		log.Fatal(err)
	}
	defer sadboxDB.Close()

	err = sadboxDB.Ping()
	if err != nil {
		log.Fatal(err)
	}

	staticFileServer := http.FileServer(_escFS(false))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/" {
			ctx := NewContext(r)
			ctx.Main = &Main{channels}
			if err := templates.ExecuteTemplate(w, "main", ctx); err != nil {
				log.Println(err)
			}
		} else {
			staticFileServer.ServeHTTP(w, r)
		}
	})

	mathRand.Seed(time.Now().UnixNano())

	rows, err := sadboxDB.Query(`SELECT Channel from Channels;`)
	if err != nil {
		log.Fatal(err)
	}
	for rows.Next() {
		var channelName string
		err := rows.Scan(&channelName)
		if err != nil {
			log.Fatal(err)
		}

		chanInfo := channel{fmt.Sprintf("/%s/", strings.Trim(channelName, "#")), channelName}

		ircChanHandler, err := NewIRCChannel(chanInfo)
		if err != nil {
			log.Fatal(err)
		}
		http.Handle(chanInfo.LinkName, ircChanHandler)

		channels = append(channels, chanInfo)
	}

	// Redirects to the right URL so I don't break old links
	http.Handle("/ghstats", http.RedirectHandler("/geekhack/", http.StatusMovedPermanently))
	http.Handle("/geekhack", http.RedirectHandler("/geekhack/", http.StatusMovedPermanently))

	// This will redirect people to the gmail page
	http.Handle("mail.sadbox.org/", http.RedirectHandler("https://mail.google.com/a/sadbox.org", http.StatusFound))

	localhost_znc, err := url.Parse("http://127.0.0.1:6698")
	if err != nil {
		log.Fatal(err)
	}
	http.Handle("/znc/", httputil.NewSingleHostReverseProxy(localhost_znc))
	http.Handle("/znc", http.RedirectHandler("/znc/", http.StatusMovedPermanently))

	servemux := gziphandler.GzipHandler(
		CatchPanic(
			Log(
				AddHeaders(http.DefaultServeMux))))

	httpSrv := &http.Server{
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
		Handler:      RedirectToHTTPS(servemux),
	}
	go func() { log.Fatal(httpSrv.ListenAndServe()) }()

	m := autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		Cache:      autocert.DirCache("/home/sadbox-web/cert-cache"),
		HostPolicy: autocert.HostWhitelist(hostname_whitelist...),
		Email:      "blue6249@gmail.com",
		ForceRSA:   true,
	}

	var certs []tls.Certificate
	sadboxCert, err := tls.LoadX509KeyPair("/home/sadbox-web/cert-cache/sadbox.org",
		"/home/sadbox-web/cert-cache/sadbox.org")

	if err == nil {
		certs = append(certs, sadboxCert)
	}

	tlsconfig := &tls.Config{
		PreferServerCipherSuites: true,
		GetCertificate:           m.GetCertificate,
		Certificates:             certs,
	}

	go NewSessionKeys(tlsconfig).Spin()

	server := &http.Server{
		Addr:         ":https",
		Handler:      servemux,
		TLSConfig:    tlsconfig,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 5 * time.Second,
	}
	log.Fatal(server.ListenAndServeTLS("", ""))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"time"
)

var streams []stream
var commands []*exec.Cmd

type stream struct {
	Name   string
	Stream string
	Image  string
}

func main() {
	//Read in urls of webcams from configuration file
	data, err := ioutil.ReadFile("streams.json")
	if err != nil {
		panic(err)
	}

	if err := json.Unmarshal(data, &streams); err != nil {
		panic(err)
	}

	//Startup webserver to listen to commands to execute video player
	http.HandleFunc("/", serveIndex)
	http.HandleFunc("/all", serveAll)
	http.HandleFunc("/pick", serveOne)
	http.HandleFunc("/stop", serveStop)
	// showAll()
	log.Println("server started...")
	log.Fatal(http.ListenAndServe(":2000", nil))
}

func renderWebsite(w http.ResponseWriter) {
	const tpl = `
<!DOCTYPE html>

		<meta charset="UTF-8">
		<title>Zoo Cam Viewer</title>

All

Stop

		{{range .Streams}}
{{ .Name }}

{{else}}
no streams
{{end}}

`
	t, err := template.New("webpage").Parse(tpl)
	if err != nil {
		panic(err)
	}
	t.Execute(w, struct{ Streams []stream }{Streams: streams})
}

func serveIndex(w http.ResponseWriter, r *http.Request) {
	//Serve up basic website with icons to choose from with "all" option
	renderWebsite(w)
}

func serveAll(w http.ResponseWriter, r *http.Request) {
	//Close all existing processes, and fire up all the videos
	showAll()
	renderWebsite(w)
}

func serveOne(w http.ResponseWriter, r *http.Request) {
	//Close all existing processes, and fire up the single video passed in
	name := r.FormValue("name")
	log.Println("got stream name of " + name)
	for _, stream := range streams {
		if stream.Name == name {
			log.Println("Starting stream " + stream.Name)
			showOne(stream)
		}
	}
	renderWebsite(w)
}

func serveStop(w http.ResponseWriter, r *http.Request) {
	killAll()
	renderWebsite(w)
}

func showAll() {
	killAll()
	width := 1920
	height := 1080
	streamCount := len(streams)

	//Determine how many streams we have to make even boxed grids
	boxes := 1
	for ; boxes*boxes < streamCount; boxes++ {
	}

	startWidth := 0
	startHeight := 0
	widthStep := width / boxes
	heightStep := height / boxes
	//We now have a box X box width screen (say 3x3), so split the screen appropriately
	for index, s := range streams {
		endWidth := startWidth + ((index + 1) * widthStep)
		endHeight := startHeight + ((index + 1) * heightStep)
		log.Printf("end width is %v and end height is %v\n", endWidth, endHeight)
		log.Printf("dimensions of window: %v,%v,%v,%v", startWidth, startHeight, endWidth, endHeight)
		cmd := exec.Command("omxplayer", "--win", fmt.Sprintf("%v,%v,%v,%v", startWidth, startHeight, endWidth, endHeight), s.Stream)
		// cmd := exec.Command("mplayer", s.Stream)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		commands = append(commands, cmd)
		cmd.Start()
		time.Sleep(5 * time.Second)
		startWidth = endWidth
		startHeight = endHeight
	}
}

func showOne(s stream) {
	killAll()
	//Startup in fullscreen
	cmd := exec.Command("omxplayer", "-b", s.Stream)
	// cmd := exec.Command("mplayer", s.Stream)
	cmd.Start()
	commands = append(commands, cmd)
}

func killAll() {
	log.Println("killing all existing streams")
	for _, proc := range commands {
		proc.Process.Kill()
	}
}
<commit_msg>trying to increase timeout to give pi more time to process<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"time"
)

var streams []stream
var commands []*exec.Cmd

type stream struct {
	Name   string
	Stream string
	Image  string
}

func main() {
	//Read in urls of webcams from configuration file
	data, err := ioutil.ReadFile("streams.json")
	if err != nil {
		panic(err)
	}

	if err := json.Unmarshal(data, &streams); err != nil {
		panic(err)
	}

	//Startup webserver to listen to commands to execute video player
	http.HandleFunc("/", serveIndex)
	http.HandleFunc("/all", serveAll)
	http.HandleFunc("/pick", serveOne)
	http.HandleFunc("/stop", serveStop)
	// showAll()
	log.Println("server started...")
	log.Fatal(http.ListenAndServe(":2000", nil))
}

func renderWebsite(w http.ResponseWriter) {
	const tpl = `
<!DOCTYPE html>

		<meta charset="UTF-8">
		<title>Zoo Cam Viewer</title>

All

Stop

		{{range .Streams}}
{{ .Name }}

{{else}}
no streams
{{end}}

`
	t, err := template.New("webpage").Parse(tpl)
	if err != nil {
		panic(err)
	}
	t.Execute(w, struct{ Streams []stream }{Streams: streams})
}

func serveIndex(w http.ResponseWriter, r *http.Request) {
	//Serve up basic website with icons to choose from with "all" option
	renderWebsite(w)
}

func serveAll(w http.ResponseWriter, r *http.Request) {
	//Close all existing processes, and fire up all the videos
	showAll()
	renderWebsite(w)
}

func serveOne(w http.ResponseWriter, r *http.Request) {
	//Close all existing processes, and fire up the single video passed in
	name := r.FormValue("name")
	log.Println("got stream name of " + name)
	for _, stream := range streams {
		if stream.Name == name {
			log.Println("Starting stream " + stream.Name)
			showOne(stream)
		}
	}
	renderWebsite(w)
}

func serveStop(w http.ResponseWriter, r *http.Request) {
	killAll()
	renderWebsite(w)
}

func showAll() {
	killAll()
	width := 1920
	height := 1080
	streamCount := len(streams)

	//Determine how many streams we have to make even boxed grids
	boxes := 1
	for ; boxes*boxes < streamCount; boxes++ {
	}

	startWidth := 0
	startHeight := 0
	widthStep := width / boxes
	heightStep := height / boxes
	//We now have a box X box width screen (say 3x3), so split the screen appropriately
	for index, s := range streams {
		endWidth := startWidth + ((index + 1) * widthStep)
		endHeight := startHeight + ((index + 1) * heightStep)
		log.Printf("end width is %v and end height is %v\n", endWidth, endHeight)
		log.Printf("dimensions of window: %v,%v,%v,%v", startWidth, startHeight, endWidth, endHeight)
		cmd := exec.Command("omxplayer", "--win", fmt.Sprintf("%v,%v,%v,%v", startWidth, startHeight, endWidth, endHeight), s.Stream)
		// cmd := exec.Command("mplayer", s.Stream)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		commands = append(commands, cmd)
		cmd.Start()
		time.Sleep(7 * time.Second)
		startWidth += widthStep
		startHeight += heightStep
	}
}

func showOne(s stream) {
	killAll()
	//Startup in fullscreen
	cmd := exec.Command("omxplayer", "-b", s.Stream)
	// cmd := exec.Command("mplayer", s.Stream)
	cmd.Start()
	commands = append(commands, cmd)
}

func killAll() {
	log.Println("killing all existing streams")
	for _, proc := range commands {
		proc.Process.Kill()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"golang.org/x/image/bmp"
	"image"
	"image/color"
	"os"
)

func main() {
	i := flag.Uint("size", 3, "Set the map size in multiples of 2 (total will be 2^n + 1)")
	colorValue := flag.Int("color", 32725, "Color to make outputted gif")
	path := flag.String("path", "test.bmp", "Filename to write to")
	dry := flag.Bool("dry", false, "Dry run (don't create image, just print size)")
	verbose := flag.Bool("v", false, "verbose")
	veryverbose := flag.Bool("vv", false, "very verbose")
	flag.Parse()

	if *colorValue > 65535 {
		fmt.Println("Color must be less than or equal to 65535")
		return
	}

	switch {
	case *veryverbose:
		log.SetLevel(log.DebugLevel)
	case *verbose:
		log.SetLevel(log.InfoLevel)
	default:
		log.SetLevel(log.WarnLevel)
	}

	t := InitializeTerrain(uint16(*i), 65535)

	fmt.Println("Generating terrain of size", t.max)
	t.Generate(0.3, 32767, 32767, 32767, 32767)

	if *dry {
		return
	}

	intMax := int(t.max)

	r := image.Rect(0, 0, intMax, intMax)
	m := image.NewGray16(r)

	for x := 0; x < intMax; x++ {
		for y := 0; y < intMax; y++ {
			n, _ := t.GetHeight(uint16(x), uint16(y))
			m.SetGray16(x, y, color.Gray16{Y: n})
		}
	}

	fi, err := os.Create(*path)
	defer fi.Close()
	if err != nil {
		panic(err)
	}
	bmp.Encode(fi, m)
}
<commit_msg>Updating argument syntax<commit_after>package main

import (
	"flag"
	log "github.com/Sirupsen/logrus"
	"golang.org/x/image/bmp"
	"image"
	"image/color"
	"os"
)

func main() {
	i := flag.Uint("size", 3, "Set the map size in multiples of 2 (total will be 2^n + 1)")
	path := flag.String("path", "test.bmp", "Filename to write to")
	dry := flag.Bool("dry", false, "Dry run (don't create image, just print size)")
	loglevel := flag.String("loglevel", "Info", "Log level. Can be Debug, Info, Warn, Error, Fatal, or Panic")
	flag.Parse()

	switch *loglevel {
	case "Debug":
		log.SetLevel(log.DebugLevel)
	case "Info":
		log.SetLevel(log.InfoLevel)
	case "Warn":
		log.SetLevel(log.WarnLevel)
	case "Error":
		log.SetLevel(log.ErrorLevel)
	case "Fatal":
		log.SetLevel(log.FatalLevel)
	case "Panic":
		log.SetLevel(log.PanicLevel)
	default:
		log.Fatal("Unknown log level. Valid values are Debug, Info, Warn, Error, Fatal, or Panic. Supplied value: ", *loglevel)
		return // Unnecessary, but do so just in case the former one didn't
	}

	t := InitializeTerrain(uint16(*i), 65535)

	log.Info("Generating terrain of size ", t.max)
	t.Generate(0.3, 32767, 32767, 32767, 32767)

	if *dry {
		return
	}

	intMax := int(t.max)

	r := image.Rect(0, 0, intMax, intMax)
	m := image.NewGray16(r)

	for x := 0; x < intMax; x++ {
		for y := 0; y < intMax; y++ {
			n, _ := t.GetHeight(uint16(x), uint16(y))
			m.SetGray16(x, y, color.Gray16{Y: n})
		}
	}

	fi, err := os.Create(*path)
	defer fi.Close()
	if err != nil {
		panic(err)
	}
	bmp.Encode(fi, m)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"crypto/subtle"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"time"

	"github.com/codegangsta/cli"
	"golang.org/x/net/websocket"
)

func main() {
	app := cli.NewApp()
	app.Usage = "turn webhooks into websockets"

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "bind, b",
			Value: ":8080",
			Usage: "address to listen on",
		},
	}

	app.Action = ActionMain

	app.RunAndExitOnError()
}

func ActionMain(c *cli.Context) {
	var (
		key           = os.Getenv("HOOKBOT_KEY")
		github_secret = os.Getenv("HOOKBOT_GITHUB_SECRET")
	)

	if key == "" || github_secret == "" {
		log.Fatalln("Error: HOOKBOT_KEY or HOOKBOT_GITHUB_SECRET not set")
	}

	hookbot := NewHookbot(key, github_secret)
	http.Handle("/", hookbot)
	http.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "OK")
	})

	log.Println("Listening on", c.String("bind"))
	err := http.ListenAndServe(c.String("bind"), nil)
	if err != nil {
		log.Fatal(err)
	}
}

type Message struct {
	Topic string
	Body  []byte
}

type Listener struct {
	Topic string
	c     chan []byte
}

type Hookbot struct {
	key, github_secret string

	http.Handler

	message                  chan Message
	addListener, delListener chan Listener
}

func NewHookbot(key, github_secret string) *Hookbot {
	h := &Hookbot{
		key: key, github_secret: github_secret,

		message:     make(chan Message, 1),
		addListener: make(chan Listener, 1),
		delListener: make(chan Listener, 1),
	}

	mux := http.NewServeMux()
	mux.Handle("/sub/", websocket.Handler(h.ServeSubscribe))
	mux.HandleFunc("/pub/", h.ServePublish)

	// Middlewares
	h.Handler = mux
	h.Handler = h.KeyChecker(h.Handler)

	go h.Loop()

	return h
}

var timeout = 1 * time.Second

func TimeoutSend(c chan []byte, m []byte) {
	select {
	case c <- m:
	case <-time.After(timeout):
	}
}

// Manage fanout from h.message onto listeners
func (h *Hookbot) Loop() {
	listeners := map[Listener]struct{}{}
	for {
		select {
		case m := <-h.message:
			for listener := range listeners {
				if listener.Topic == m.Topic {
					go TimeoutSend(listener.c, m.Body)
				}
			}

		case l := <-h.addListener:
			listeners[l] = struct{}{}
		case l := <-h.delListener:
			delete(listeners, l)
		}
	}
}

func (h *Hookbot) Add(topic string) Listener {
	l := Listener{Topic: topic, c: make(chan []byte)}
	h.addListener <- l
	return l
}

func (h *Hookbot) Del(l Listener) {
	h.delListener <- l
}

func SecureEqual(x, y string) bool {
	if subtle.ConstantTimeCompare([]byte(x), []byte(y)) == 1 {
		return true
	}
	return false
}

func (h *Hookbot) IsGithubKeyOK(w http.ResponseWriter, r *http.Request) bool {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Not Authorized", http.StatusUnauthorized)
	}

	r.Body = ioutil.NopCloser(bytes.NewReader(body))

	mac := hmac.New(sha1.New, []byte(h.github_secret))
	mac.Reset()
	mac.Write(body)

	signature := fmt.Sprintf("sha1=%x", mac.Sum(nil))

	return SecureEqual(r.Header.Get("X-Hub-Signature"), signature)
}

func (h *Hookbot) IsKeyOK(w http.ResponseWriter, r *http.Request) bool {

	if _, ok := r.Header["X-Hub-Signature"]; ok {
		if !h.IsGithubKeyOK(w, r) {
			http.NotFound(w, r)
			return false
		}
		return true
	}

	lhs := r.Header.Get("Authorization")
	rhs := fmt.Sprintf("Bearer %v", h.key)

	if !SecureEqual(lhs, rhs) {
		http.NotFound(w, r)
		return false
	}

	return true
}

func (h *Hookbot) KeyChecker(wrapped http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !h.IsKeyOK(w, r) {
			return
		}

		wrapped.ServeHTTP(w, r)
	}
}

// The topic is everything after the "/pub/" or "/sub/"
var TopicRE *regexp.Regexp = regexp.MustCompile("/[^/]+/(.*)")

func Topic(url *url.URL) string {
	m := TopicRE.FindStringSubmatch(url.Path)
	if m == nil {
		return ""
	}
	return m[1]
}

func (h *Hookbot) ServePublish(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		log.Printf("Error serving %v: %v", r.URL, err)
		return
	}

	topic := Topic(r.URL)
	h.message <- Message{Topic: topic, Body: body}
}

func (h *Hookbot) ServeSubscribe(conn *websocket.Conn) {

	topic := Topic(conn.Request().URL)

	listener := h.Add(topic)
	defer h.Del(listener)

	closed := make(chan struct{})

	go func() {
		defer close(closed)
		_, _ = io.Copy(ioutil.Discard, conn)
	}()

	var message []byte

	for {
		select {
		case message = <-listener.c:
		case <-closed:
			log.Printf("Client disconnected")
			return
		}

		conn.SetWriteDeadline(time.Now().Add(90 * time.Second))
		n, err := conn.Write(message)
		switch {
		case n != len(message):
			log.Printf("Short write %d != %d", n, len(message))
			return // short write
		case err == io.EOF:
			return // done
		case err != nil:
			log.Printf("Error in conn.Write: %v", err)
			return // unknown error
		}
	}
}
<commit_msg>Make authentication token URL-specific<commit_after>package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"golang.org/x/net/websocket"
)

func main() {
	app := cli.NewApp()
	app.Usage = "turn webhooks into websockets"

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "bind, b",
			Value: ":8080",
			Usage: "address to listen on",
		},
	}

	app.Action = ActionMain

	app.RunAndExitOnError()
}

func ActionMain(c *cli.Context) {
	var (
		key           = os.Getenv("HOOKBOT_KEY")
		github_secret = os.Getenv("HOOKBOT_GITHUB_SECRET")
	)

	if key == "" || github_secret == "" {
		log.Fatalln("Error: HOOKBOT_KEY or HOOKBOT_GITHUB_SECRET not set")
	}

	hookbot := NewHookbot(key, github_secret)
	http.Handle("/", hookbot)
	http.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "OK")
	})

	log.Println("Listening on", c.String("bind"))
	err := http.ListenAndServe(c.String("bind"), nil)
	if err != nil {
		log.Fatal(err)
	}
}

type Message struct {
	Topic string
	Body  []byte
}

type Listener struct {
	Topic string
	c     chan []byte
}

type Hookbot struct {
	key, github_secret string

	http.Handler

	message                  chan Message
	addListener, delListener chan Listener
}

func NewHookbot(key, github_secret string) *Hookbot {
	h := &Hookbot{
		key: key, github_secret: github_secret,

		message:     make(chan Message, 1),
		addListener: make(chan Listener, 1),
		delListener: make(chan Listener, 1),
	}

	mux := http.NewServeMux()
	mux.Handle("/sub/", websocket.Handler(h.ServeSubscribe))
	mux.HandleFunc("/pub/", h.ServePublish)

	// Middlewares
	h.Handler = mux
	h.Handler = h.KeyChecker(h.Handler)

	go h.Loop()

	return h
}

var timeout = 1 * time.Second

func TimeoutSend(c chan []byte, m []byte) {
	select {
	case c <- m:
	case <-time.After(timeout):
	}
}

// Manage fanout from h.message onto listeners
func (h *Hookbot) Loop() {
	listeners := map[Listener]struct{}{}
	for {
		select {
		case m := <-h.message:
			for listener := range listeners {
				if listener.Topic == m.Topic {
					go TimeoutSend(listener.c, m.Body)
				}
			}

		case l := <-h.addListener:
			listeners[l] = struct{}{}
		case l := <-h.delListener:
			delete(listeners, l)
		}
	}
}

func (h *Hookbot) Add(topic string) Listener {
	l := Listener{Topic: topic, c: make(chan []byte)}
	h.addListener <- l
	return l
}

func (h *Hookbot) Del(l Listener) {
	h.delListener <- l
}

func SecureEqual(x, y string) bool {
	if subtle.ConstantTimeCompare([]byte(x), []byte(y)) == 1 {
		return true
	}
	return false
}

func (h *Hookbot) IsGithubKeyOK(w http.ResponseWriter, r *http.Request) bool {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Not Authorized", http.StatusUnauthorized)
	}

	r.Body = ioutil.NopCloser(bytes.NewReader(body))

	expected := fmt.Sprintf("sha1=%v", Sha1HMAC(h.github_secret, body))

	return SecureEqual(r.Header.Get("X-Hub-Signature"), expected)
}

func Sha1HMAC(key, payload string) string {
	mac := hmac.New(sha1.New, []byte(key))
	_, _ = mac.Write([]byte(payload))
	return fmt.Sprintf("%x", mac.Sum(nil))
}

func (h *Hookbot) IsKeyOK(w http.ResponseWriter, r *http.Request) bool {

	if _, ok := r.Header["X-Hub-Signature"]; ok {
		if !h.IsGithubKeyOK(w, r) {
			return false
		}
		return true
	}

	authorization := r.Header.Get("Authorization")
	fields := strings.Fields(authorization)

	if len(fields) != 2 {
		return false
	}

	authType, givenKey := fields[0], fields[1]

	var givenMac string

	switch strings.ToLower(authType) {
	default:
		return false // Not understood
	case "basic":
		givenMacBytes, err := base64.StdEncoding.DecodeString(givenKey)
		if err != nil {
			return false
		}
		// Remove trailing right colon, since it should be blank.
		givenMac = strings.TrimRight(string(givenMacBytes), ":")

	case "bearer":
		givenMac = givenKey // No processing required
	}

	expectedMac := Sha1HMAC(h.key, r.URL.Path)

	if !SecureEqual(givenMac, expectedMac) {
		return false
	}

	return true
}

func (h *Hookbot) KeyChecker(wrapped http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !h.IsKeyOK(w, r) {
			http.NotFound(w, r)
			return
		}

		wrapped.ServeHTTP(w, r)
	}
}

// The topic is everything after the "/pub/" or "/sub/"
var TopicRE *regexp.Regexp = regexp.MustCompile("/[^/]+/(.*)")

func Topic(url *url.URL) string {
	m := TopicRE.FindStringSubmatch(url.Path)
	if m == nil {
		return ""
	}
	return m[1]
}

func (h *Hookbot) ServePublish(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		log.Printf("Error serving %v: %v", r.URL, err)
		return
	}

	topic := Topic(r.URL)
	h.message <- Message{Topic: topic, Body: body}
}

func (h *Hookbot) ServeSubscribe(conn *websocket.Conn) {

	topic := Topic(conn.Request().URL)

	listener := h.Add(topic)
	defer h.Del(listener)

	closed := make(chan struct{})

	go func() {
		defer close(closed)
		_, _ = io.Copy(ioutil.Discard, conn)
	}()

	var message []byte

	for {
		select {
		case message = <-listener.c:
		case <-closed:
			log.Printf("Client disconnected")
			return
		}

		conn.SetWriteDeadline(time.Now().Add(90 * time.Second))
		n, err := conn.Write(message)
		switch {
		case n != len(message):
			log.Printf("Short write %d != %d", n, len(message))
			return // short write
		case err == io.EOF:
			return // done
		case err != nil:
			log.Printf("Error in conn.Write: %v", err)
			return // unknown error
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strings"
	"syscall"
	"time"

	"github.com/braintree/manners"
	"github.com/codegangsta/cli"
)

func handleSignals(server *manners.GracefulServer) {
	sigChan := make(chan os.Signal, 3)
	go func() {
		for sig := range sigChan {
			if sig == os.Interrupt || sig == os.Kill {
				server.Close()
			}
			if sig == syscall.SIGUSR1 {
				var buf []byte
				var written int
				currLen := 1024
				for written == len(buf) {
					buf = make([]byte, currLen)
					written = runtime.Stack(buf, true)
					currLen *= 2
				}
				log.Print(string(buf[:written]))
			}
			if sig == syscall.SIGUSR2 {
				go func() {
					cpufile, _ := os.OpenFile("./planb_cpu.pprof", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
					memfile, _ := os.OpenFile("./planb_mem.pprof", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
					log.Println("enabling profile...")
					pprof.WriteHeapProfile(memfile)
					memfile.Close()
					pprof.StartCPUProfile(cpufile)
					time.Sleep(60 * time.Second)
					pprof.StopCPUProfile()
					cpufile.Close()
					log.Println("profiling done")
				}()
			}
		}
	}()
	signal.Notify(sigChan, os.Interrupt, os.Kill, syscall.SIGUSR1, syscall.SIGUSR2)
}

func runServer(c *cli.Context) {
	listener, err := net.Listen("tcp", c.String("listen"))
	if err != nil {
		log.Fatal(err)
	}
	router := Router{
		ReadRedisHost:  c.String("read-redis-host"),
		ReadRedisPort:  c.Int("read-redis-port"),
		WriteRedisHost: c.String("write-redis-host"),
		WriteRedisPort: c.Int("write-redis-port"),
		LogPath:        c.String("access-log"),
		RequestTimeout: time.Duration(c.Int("request-timeout")) * time.Second,
		DialTimeout:    time.Duration(c.Int("dial-timeout")) * time.Second,
		DeadBackendTTL: c.Int("dead-backend-time"),
	}
	err = router.Init()
	if err != nil {
		log.Fatal(err)
	}
	s := manners.NewWithServer(&http.Server{Handler: &router})
	handleSignals(s)
	log.Printf("Listening on %v...\n", listener.Addr())
	err = s.Serve(listener)
	router.Stop()
	if err != nil {
		log.Fatal(err)
	}
}

func fixUsage(s string) string {
	parts := strings.Split(s, " ")
	currLen := 0
	lastPart := 0
	var lines []string
	for i := range parts {
		if currLen+len(parts[i])+1 > 55 {
			lines = append(lines, strings.Join(parts[lastPart:i], " "))
			currLen = 0
			lastPart = i
		}
		currLen += len(parts[i]) + 1
	}
	lines = append(lines, strings.Join(parts[lastPart:], " "))
	return strings.Join(lines, "\n\t")
}

func main() {
	app := cli.NewApp()
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "listen, l",
			Value: "0.0.0.0:8989",
			Usage: "Address to listen",
		},
		cli.StringFlag{
			Name:  "read-redis-host",
			Value: "127.0.0.1",
		},
		cli.IntFlag{
			Name:  "read-redis-port",
			Value: 6379,
		},
		cli.StringFlag{
			Name:  "write-redis-host",
			Value: "127.0.0.1",
		},
		cli.IntFlag{
			Name:  "write-redis-port",
			Value: 6379,
		},
		cli.StringFlag{
			Name:  "access-log",
			Value: "./access.log",
			Usage: fixUsage("File path where access log will be written. If value equals `syslog` log will be sent to local syslog."),
		},
		cli.IntFlag{
			Name:  "request-timeout",
			Value: 30,
			Usage: "Total backend request timeout in seconds",
		},
		cli.IntFlag{
			Name:  "dial-timeout",
			Value: 10,
			Usage: "Dial backend request timeout in seconds",
		},
		cli.IntFlag{
			Name:  "dead-backend-time",
			Value: 30,
			Usage: fixUsage("Time in seconds a backend will remain disabled after a network failure."),
		},
	}
	app.Version = "0.1.0"
	app.Name = "planb"
	app.Usage = "http and websockets reverse proxy"
	app.Action = runServer
	app.Author = "tsuru team"
	app.Email = "https://github.com/tsuru/planb"
	app.Run(os.Args)
}
<commit_msg>increment version<commit_after>// Copyright 2015 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strings"
	"syscall"
	"time"

	"github.com/braintree/manners"
	"github.com/codegangsta/cli"
)

func handleSignals(server *manners.GracefulServer) {
	sigChan := make(chan os.Signal, 3)
	go func() {
		for sig := range sigChan {
			if sig == os.Interrupt || sig == os.Kill {
				server.Close()
			}
			if sig == syscall.SIGUSR1 {
				var buf []byte
				var written int
				currLen := 1024
				for written == len(buf) {
					buf = make([]byte, currLen)
					written = runtime.Stack(buf, true)
					currLen *= 2
				}
				log.Print(string(buf[:written]))
			}
			if sig == syscall.SIGUSR2 {
				go func() {
					cpufile, _ := os.OpenFile("./planb_cpu.pprof", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
					memfile, _ := os.OpenFile("./planb_mem.pprof", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
					log.Println("enabling profile...")
					pprof.WriteHeapProfile(memfile)
					memfile.Close()
					pprof.StartCPUProfile(cpufile)
					time.Sleep(60 * time.Second)
					pprof.StopCPUProfile()
					cpufile.Close()
					log.Println("profiling done")
				}()
			}
		}
	}()
	signal.Notify(sigChan, os.Interrupt, os.Kill, syscall.SIGUSR1, syscall.SIGUSR2)
}

func runServer(c *cli.Context) {
	listener, err := net.Listen("tcp", c.String("listen"))
	if err != nil {
		log.Fatal(err)
	}
	router := Router{
		ReadRedisHost:  c.String("read-redis-host"),
		ReadRedisPort:  c.Int("read-redis-port"),
		WriteRedisHost: c.String("write-redis-host"),
		WriteRedisPort: c.Int("write-redis-port"),
		LogPath:        c.String("access-log"),
		RequestTimeout: time.Duration(c.Int("request-timeout")) * time.Second,
		DialTimeout:    time.Duration(c.Int("dial-timeout")) * time.Second,
		DeadBackendTTL: c.Int("dead-backend-time"),
	}
	err = router.Init()
	if err != nil {
		log.Fatal(err)
	}
	s := manners.NewWithServer(&http.Server{Handler: &router})
	handleSignals(s)
	log.Printf("Listening on %v...\n", listener.Addr())
	err = s.Serve(listener)
	router.Stop()
	if err != nil {
		log.Fatal(err)
	}
}

func fixUsage(s string) string {
	parts := strings.Split(s, " ")
	currLen := 0
	lastPart := 0
	var lines []string
	for i := range parts {
		if currLen+len(parts[i])+1 > 55 {
			lines = append(lines, strings.Join(parts[lastPart:i], " "))
			currLen = 0
			lastPart = i
		}
		currLen += len(parts[i]) + 1
	}
	lines = append(lines, strings.Join(parts[lastPart:], " "))
	return strings.Join(lines, "\n\t")
}

func main() {
	app := cli.NewApp()
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "listen, l",
			Value: "0.0.0.0:8989",
			Usage: "Address to listen",
		},
		cli.StringFlag{
			Name:  "read-redis-host",
			Value: "127.0.0.1",
		},
		cli.IntFlag{
			Name:  "read-redis-port",
			Value: 6379,
		},
		cli.StringFlag{
			Name:  "write-redis-host",
			Value: "127.0.0.1",
		},
		cli.IntFlag{
			Name:  "write-redis-port",
			Value: 6379,
		},
		cli.StringFlag{
			Name:  "access-log",
			Value: "./access.log",
			Usage: fixUsage("File path where access log will be written. If value equals `syslog` log will be sent to local syslog."),
		},
		cli.IntFlag{
			Name:  "request-timeout",
			Value: 30,
			Usage: "Total backend request timeout in seconds",
		},
		cli.IntFlag{
			Name:  "dial-timeout",
			Value: 10,
			Usage: "Dial backend request timeout in seconds",
		},
		cli.IntFlag{
			Name:  "dead-backend-time",
			Value: 30,
			Usage: fixUsage("Time in seconds a backend will remain disabled after a network failure."),
		},
	}
	app.Version = "0.1.1"
	app.Name = "planb"
	app.Usage = "http and websockets reverse proxy"
	app.Action = runServer
	app.Author = "tsuru team"
	app.Email = "https://github.com/tsuru/planb"
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/Clever/shorty/db"
	"github.com/Clever/shorty/routes"
	"github.com/gorilla/mux"
	"gopkg.in/Clever/kayvee-go.v3/logger"
)

const (
	pgBackend    = "postgres"
	redisBackend = "redis"
)

var (
	port     = flag.String("port", "80", "port to listen on")
	database = flag.String("db", pgBackend, "datastore option to use, one of: ['postgres', 'redis']")
	readonly = flag.Bool("readonly", false, "set readonly mode (useful for external-facing instance)")
	protocol = flag.String("protocol", "http", "protocol for the short handler - useful to separate for external-facing separate instance")
	domain   = flag.String("domain", "go", "set the domain for the short URL reported to the user")
	lg       = logger.New("shorty")
)

func main() {
	flag.Parse()

	var sdb db.ShortenBackend
	switch *database {
	case pgBackend:
		sdb = db.NewPostgresDB()
	case redisBackend:
		sdb = db.NewRedisDB()
	default:
		lg.CriticalD("missing-backed", logger.M{
			"msg": fmt.Sprintf("'%s' backend is not offered", *database)})
		os.Exit(1)
	}

	// default to ReadOnly mode for POSTs and list of slugs
	deleteHandler := routes.ReadOnlyHandler()
	shortenHandler := routes.ReadOnlyHandler()
	listHandler := routes.ReadOnlyHandler()
	if *readonly == false {
		deleteHandler = routes.DeleteHandler(sdb)
		shortenHandler = routes.ShortenHandler(sdb)
		listHandler = routes.ListHandler(sdb)
	}
	r := mux.NewRouter()
	r.HandleFunc("/delete", deleteHandler).Methods("POST")
	r.HandleFunc("/shorten", shortenHandler).Methods("POST")
	r.HandleFunc("/list", listHandler).Methods("GET")

	// Safe for public consumption no matter what below here
	// Technically someone could scrape the whole slug space to discover
	// all the slugs, but that comes along with the territory
	r.HandleFunc("/meta", routes.MetaHandler(*protocol, *domain)).Methods("GET")
	r.PathPrefix("/Shortener.jsx").Handler(http.FileServer(http.Dir("./static")))
	r.PathPrefix("/favicon.png").Handler(http.FileServer(http.Dir("./static")))
	r.HandleFunc("/{slug}", routes.RedirectHandler(sdb, *domain, "")).Methods("GET")
	r.HandleFunc("/{slug}/{suffix}", routes.RedirectHandler(sdb, *domain, "/")).Methods("GET")
	r.HandleFunc("/health/check", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "STATUS OK")
	})

	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		http.ServeFile(w, r, "./static/index.html")
	}).Methods("GET")
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("./static")))
	http.Handle("/", r)

	lg.InfoD("starting-server", logger.M{"port": *port})
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<commit_msg>Order matters for routes.<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/Clever/shorty/db"
	"github.com/Clever/shorty/routes"
	"github.com/gorilla/mux"
	"gopkg.in/Clever/kayvee-go.v3/logger"
)

const (
	pgBackend    = "postgres"
	redisBackend = "redis"
)

var (
	port     = flag.String("port", "80", "port to listen on")
	database = flag.String("db", pgBackend, "datastore option to use, one of: ['postgres', 'redis']")
	readonly = flag.Bool("readonly", false, "set readonly mode (useful for external-facing instance)")
	protocol = flag.String("protocol", "http", "protocol for the short handler - useful to separate for external-facing separate instance")
	domain   = flag.String("domain", "go", "set the domain for the short URL reported to the user")
	lg       = logger.New("shorty")
)

func main() {
	flag.Parse()

	var sdb db.ShortenBackend
	switch *database {
	case pgBackend:
		sdb = db.NewPostgresDB()
	case redisBackend:
		sdb = db.NewRedisDB()
	default:
		lg.CriticalD("missing-backed", logger.M{
			"msg": fmt.Sprintf("'%s' backend is not offered", *database)})
		os.Exit(1)
	}

	// default to ReadOnly mode for POSTs and list of slugs
	deleteHandler := routes.ReadOnlyHandler()
	shortenHandler := routes.ReadOnlyHandler()
	listHandler := routes.ReadOnlyHandler()
	if *readonly == false {
		deleteHandler = routes.DeleteHandler(sdb)
		shortenHandler = routes.ShortenHandler(sdb)
		listHandler = routes.ListHandler(sdb)
	}
	r := mux.NewRouter()
	r.HandleFunc("/delete", deleteHandler).Methods("POST")
	r.HandleFunc("/shorten", shortenHandler).Methods("POST")
	r.HandleFunc("/list", listHandler).Methods("GET")

	// Safe for public consumption no matter what below here
	// Technically someone could scrape the whole slug space to discover
	// all the slugs, but that comes along with the territory
	r.HandleFunc("/meta", routes.MetaHandler(*protocol, *domain)).Methods("GET")
	r.PathPrefix("/Shortener.jsx").Handler(http.FileServer(http.Dir("./static")))
	r.PathPrefix("/favicon.png").Handler(http.FileServer(http.Dir("./static")))
	r.HandleFunc("/health/check", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "STATUS OK")
	})

	r.HandleFunc("/{slug}", routes.RedirectHandler(sdb, *domain, "")).Methods("GET")
	r.HandleFunc("/{slug}/{suffix}", routes.RedirectHandler(sdb, *domain, "/")).Methods("GET")

	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		http.ServeFile(w, r, "./static/index.html")
	}).Methods("GET")
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("./static")))
	http.Handle("/", r)

	lg.InfoD("starting-server", logger.M{"port": *port})
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

import (
	"code.google.com/p/go.net/websocket"
	"mozilla.org/simplepush"
	storage "mozilla.org/simplepush/storage/mcstorage"
	mozutil "mozilla.org/util"

	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime/pprof"
)

var (
	configFile *string = flag.String("config", "config.ini", "Configuration File")
	profile    *string = flag.String("profile", "", "Profile file output")
	logger     *mozutil.HekaLogger
	store      *storage.Storage
)

// -- main
func main() {
	flag.Parse()
	config := mozutil.MzGetConfig(*configFile)

	config = simplepush.FixConfig(config)
	log.Printf("CurrentHost: %s", config["shard.current_host"])

	if *profile != "" {
		f, err := os.Create(*profile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	logger = mozutil.NewHekaLogger(config)
	store = storage.New(config, logger)

	// Initialize the common server.
	simplepush.InitServer(config, logger)
	handlers := simplepush.NewHandler(config, logger, store)

	// Register the handlers
	// each websocket gets it's own handler.
	http.HandleFunc("/update/", handlers.UpdateHandler)
	http.HandleFunc("/status/", handlers.StatusHandler)
	http.HandleFunc("/realstatus/", handlers.RealStatusHandler)
	http.Handle("/", websocket.Handler(handlers.PushSocketHandler))

	// Config the server
	host := mozutil.MzGet(config, "host", "localhost")
	port := mozutil.MzGet(config, "port", "8080")

	// Hoist the main sail
	logger.Info("main",
		fmt.Sprintf("listening on %s:%s", host, port), nil)
	err := http.ListenAndServe(fmt.Sprintf("%s:%s", host, port), nil)
	if err != nil {
		panic("ListenAndServe: " + err.Error())
	}
}

// 04fs
// vim: set tabstab=4 softtabstop=4 shiftwidth=4 noexpandtab
<commit_msg>Capture sigint so we can gracefully exit.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

import (
	"code.google.com/p/go.net/websocket"
	"mozilla.org/simplepush"
	storage "mozilla.org/simplepush/storage/mcstorage"
	mozutil "mozilla.org/util"

	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"runtime/pprof"
	"syscall"
)

var (
	configFile *string = flag.String("config", "config.ini", "Configuration File")
	profile    *string = flag.String("profile", "", "Profile file output")
	logger     *mozutil.HekaLogger
	store      *storage.Storage
)

const SIGUSR1 = syscall.SIGUSR1

// -- main
func main() {
	flag.Parse()
	config := mozutil.MzGetConfig(*configFile)

	config = simplepush.FixConfig(config)
	log.Printf("CurrentHost: %s", config["shard.current_host"])

	if *profile != "" {
		f, err := os.Create(*profile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	logger = mozutil.NewHekaLogger(config)
	store = storage.New(config, logger)

	// Initialize the common server.
	simplepush.InitServer(config, logger)
	handlers := simplepush.NewHandler(config, logger, store)

	// Register the handlers
	// each websocket gets it's own handler.
	http.HandleFunc("/update/", handlers.UpdateHandler)
	http.HandleFunc("/status/", handlers.StatusHandler)
	http.HandleFunc("/realstatus/", handlers.RealStatusHandler)
	http.Handle("/", websocket.Handler(handlers.PushSocketHandler))

	// Config the server
	host := mozutil.MzGet(config, "host", "localhost")
	port := mozutil.MzGet(config, "port", "8080")

	// Hoist the main sail
	logger.Info("main",
		fmt.Sprintf("listening on %s:%s", host, port), nil)

	// wait for sigint
	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGHUP, SIGUSR1)

	errChan := make(chan error)
	go func() {
		errChan <- http.ListenAndServe(fmt.Sprintf("%s:%s", host, port), nil)
	}()

	select {
	case err := <-errChan:
		if err != nil {
			panic("ListenAndServe: " + err.Error())
		}
	case <-sigChan:
		logger.Info("main", "Recieved signal, shutting down.", nil)
	}
}

// 04fs
// vim: set tabstab=4 softtabstop=4 shiftwidth=4 noexpandtab
<|endoftext|>