<commit_before>package net

import (
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
)

var (
	BindLocal = false // TODO(geoah) refactor to remove global
	bindIpv6  = false // TODO(geoah) refactor to remove global
)

func init() {
	BindLocal, _ = strconv.ParseBool(os.Getenv("BIND_LOCAL"))
	bindIpv6, _ = strconv.ParseBool(os.Getenv("BIND_IPV6"))
}

// GetAddresses returns the addresses the transport is listening to
func GetAddresses(protocol string, l net.Listener) []string {
	port := l.Addr().(*net.TCPAddr).Port
	// TODO log errors
	addrs, _ := GetLocalPeerAddresses(port)
	for i, addr := range addrs {
		addrs[i] = fmt.Sprintf("%s:%s", protocol, addr)
	}
	return addrs
}

func fmtAddress(protocol, address string, port int) string {
	return fmt.Sprintf("%s:%s:%d", protocol, address, port)
}

// GetLocalPeerAddresses returns the addresses TCP can listen to on the local machine
func GetLocalPeerAddresses(port int) ([]string, error) {
	// go through all ifs
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	// gather addresses of all ifs
	ips := []net.Addr{}
	for _, iface := range ifaces {
		ifIPs, err := iface.Addrs()
		if err != nil {
			continue
		}
		ips = append(ips, ifIPs...)
	}

	// gather valid addresses
	addrs := []string{}
	for _, ip := range ips {
		cleanIP, valid := isValidIP(ip)
		if valid {
			hostPort := fmt.Sprintf("%s:%d", cleanIP, port)
			addrs = append(addrs, hostPort)
		}
	}
	return addrs, nil
}

func isValidIP(addr net.Addr) (string, bool) {
	var ip net.IP
	switch v := addr.(type) {
	case *net.IPNet:
		ip = v.IP
	case *net.IPAddr:
		ip = v.IP
	}
	if ip == nil {
		return "", false
	}
	if !BindLocal && (ip.IsLoopback() || isPrivate(ip)) {
		return "", false
	}
	if !bindIpv6 && isIPv6(ip.String()) {
		return "", false
	}
	return ip.String(), true
}

func isIPv6(address string) bool {
	return strings.Count(address, ":") >= 2
}

func isPrivate(ip net.IP) bool {
	_, object24, _ := net.ParseCIDR("10.0.0.0/8")
	_, object20, _ := net.ParseCIDR("172.16.0.0/12")
	_, object16, _ := net.ParseCIDR("192.168.0.0/16")
	return object16.Contains(ip) || object20.Contains(ip) || object24.Contains(ip)
}
<commit_msg>feat(net): split binding to local and private ips<commit_after>package net

import (
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
)

var (
	BindLocal   = false // TODO(geoah) refactor to remove global
	BindPrivate = false // TODO(geoah) refactor to remove global
	bindIpv6    = false // TODO(geoah) refactor to remove global
)

func init() {
	BindLocal, _ = strconv.ParseBool(os.Getenv("BIND_LOCAL"))
	BindPrivate, _ = strconv.ParseBool(os.Getenv("BIND_PRIVATE"))
	bindIpv6, _ = strconv.ParseBool(os.Getenv("BIND_IPV6"))
}

// GetAddresses returns the addresses the transport is listening to
func GetAddresses(protocol string, l net.Listener) []string {
	port := l.Addr().(*net.TCPAddr).Port
	// TODO log errors
	addrs, _ := GetLocalPeerAddresses(port)
	for i, addr := range addrs {
		addrs[i] = fmt.Sprintf("%s:%s", protocol, addr)
	}
	return addrs
}

func fmtAddress(protocol, address string, port int) string {
	return fmt.Sprintf("%s:%s:%d", protocol, address, port)
}

// GetLocalPeerAddresses returns the addresses TCP can listen to on the local machine
func GetLocalPeerAddresses(port int) ([]string, error) {
	// go through all ifs
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	// gather addresses of all ifs
	ips := []net.Addr{}
	for _, iface := range ifaces {
		ifIPs, err := iface.Addrs()
		if err != nil {
			continue
		}
		ips = append(ips, ifIPs...)
	}

	// gather valid addresses
	addrs := []string{}
	for _, ip := range ips {
		cleanIP, valid := isValidIP(ip)
		if valid {
			hostPort := fmt.Sprintf("%s:%d", cleanIP, port)
			addrs = append(addrs, hostPort)
		}
	}
	return addrs, nil
}

func isValidIP(addr net.Addr) (string, bool) {
	var ip net.IP
	switch v := addr.(type) {
	case *net.IPNet:
		ip = v.IP
	case *net.IPAddr:
		ip = v.IP
	}
	if ip == nil {
		return "", false
	}
	if !BindLocal && ip.IsLoopback() {
		return "", false
	}
	if !BindPrivate && isPrivate(ip) {
		return "", false
	}
	if !bindIpv6 && isIPv6(ip.String()) {
		return "", false
	}
	return ip.String(), true
}

func isIPv6(address string) bool {
	return strings.Count(address, ":") >= 2
}

func isPrivate(ip net.IP) bool {
	_, object24, _ := net.ParseCIDR("10.0.0.0/8")
	_, object20, _ := net.ParseCIDR("172.16.0.0/12")
	_, object16, _ := net.ParseCIDR("192.168.0.0/16")
	return object16.Contains(ip) || object20.Contains(ip) || object24.Contains(ip)
}
<|endoftext|>


<commit_before>// Package prog provides the entry point to Elvish. Its subpackages correspond
// to subprograms of Elvish.
package prog

// This package sets up the basic environment and calls the appropriate
// "subprogram", one of the daemon, the terminal interface, or the web
// interface.

import (
	"flag"
	"fmt"
	"io"
	"os"
	"runtime/pprof"

	"src.elv.sh/pkg/logutil"
)

// Default port on which the web interface runs. The number is chosen because it
// resembles "elvi".
const defaultWebPort = 3171

// DeprecationLevel is a global flag that controls which deprecations to show.
// If its value is X, Elvish shows deprecations that should be shown for version
// 0.X.
var DeprecationLevel = 15

// SetDeprecationLevel sets ShowDeprecations to the given value, and returns a
// function to restore the old value.
func SetDeprecationLevel(level int) func() {
	save := DeprecationLevel
	DeprecationLevel = level
	return func() { DeprecationLevel = save }
}

// Flags keeps command-line flags.
type Flags struct {
	Log, CPUProfile string

	Help, Version, BuildInfo, JSON bool

	CodeInArg, CompileOnly, NoRc bool

	Web  bool
	Port int

	Daemon bool
	Forked int

	DB, Sock string
}

func newFlagSet(f *Flags) *flag.FlagSet {
	fs := flag.NewFlagSet("elvish", flag.ContinueOnError)
	// Error and usage will be printed explicitly.
	fs.SetOutput(io.Discard)

	fs.StringVar(&f.Log, "log", "", "a file to write debug log to except for the daemon")
	fs.StringVar(&f.CPUProfile, "cpuprofile", "", "write cpu profile to file")

	fs.BoolVar(&f.Help, "help", false, "show usage help and quit")
	fs.BoolVar(&f.Version, "version", false, "show version and quit")
	fs.BoolVar(&f.BuildInfo, "buildinfo", false, "show build info and quit")
	fs.BoolVar(&f.JSON, "json", false, "show output in JSON. Useful with -buildinfo.")

	// The `-i` option is for compatibility with POSIX shells so that programs, such as the `script`
	// command, will work when asked to launch an interactive Elvish shell.
	fs.Bool("i", false, "force interactive mode; currently ignored")
	fs.BoolVar(&f.CodeInArg, "c", false, "take first argument as code to execute")
	fs.BoolVar(&f.CompileOnly, "compileonly", false, "Parse/Compile but do not execute")
	fs.BoolVar(&f.NoRc, "norc", false, "run elvish without invoking rc.elv")

	fs.BoolVar(&f.Web, "web", false, "run backend of web interface")
	fs.IntVar(&f.Port, "port", defaultWebPort, "the port of the web backend")

	fs.BoolVar(&f.Daemon, "daemon", false, "[internal flag] run the storage daemon instead of shell")

	fs.StringVar(&f.DB, "db", "", "[internal flag] path to the database")
	fs.StringVar(&f.Sock, "sock", "", "[internal flag] path to the daemon socket")

	fs.IntVar(&DeprecationLevel, "deprecation-level", DeprecationLevel, "show warnings for all features deprecated as of version 0.X")

	return fs
}

func usage(out io.Writer, f *flag.FlagSet) {
	fmt.Fprintln(out, "Usage: elvish [flags] [script]")
	fmt.Fprintln(out, "Supported flags:")
	f.SetOutput(out)
	f.PrintDefaults()
}

// Run parses command-line flags and runs the first applicable subprogram. It
// returns the exit status of the program.
func Run(fds [3]*os.File, args []string, programs ...Program) int {
	f := &Flags{}
	fs := newFlagSet(f)
	err := fs.Parse(args[1:])
	if err != nil {
		if err == flag.ErrHelp {
			// (*flag.FlagSet).Parse returns ErrHelp when -h or -help was
			// requested but *not* defined. Elvish defines -help, but not -h; so
			// this means that -h has been requested. Handle this by printing
			// the same message as an undefined flag.
			fmt.Fprintln(fds[2], "flag provided but not defined: -h")
		} else {
			fmt.Fprintln(fds[2], err)
		}
		usage(fds[2], fs)
		return 2
	}

	// Handle flags common to all subprograms.
	if f.CPUProfile != "" {
		f, err := os.Create(f.CPUProfile)
		if err != nil {
			fmt.Fprintln(fds[2], "Warning: cannot create CPU profile:", err)
			fmt.Fprintln(fds[2], "Continuing without CPU profiling.")
		} else {
			pprof.StartCPUProfile(f)
			defer pprof.StopCPUProfile()
		}
	}

	if f.Daemon {
		// We expect our stdout file handle is open on a unique log file for the daemon to write its
		// log messages. See daemon.Spawn() in pkg/daemon.
		logutil.SetOutput(fds[1])
	} else if f.Log != "" {
		err = logutil.SetOutputFile(f.Log)
		if err != nil {
			fmt.Fprintln(fds[2], err)
		}
	}

	if f.Help {
		usage(fds[1], fs)
		return 0
	}

	p := findProgram(f, programs)
	if p == nil {
		fmt.Fprintln(fds[2], "program bug: no suitable subprogram")
		return 2
	}

	err = p.Run(fds, f, fs.Args())
	if err == nil {
		return 0
	}
	if msg := err.Error(); msg != "" {
		fmt.Fprintln(fds[2], msg)
	}
	switch err := err.(type) {
	case badUsageError:
		usage(fds[2], fs)
	case exitError:
		return err.exit
	}
	return 2
}

func findProgram(f *Flags, programs []Program) Program {
	for _, program := range programs {
		if program.ShouldRun(f) {
			return program
		}
	}
	return nil
}

// BadUsage returns an error that may be returned by Program.Main, which
// requests the main program to print out a message, the usage information and
// exit with 2.
func BadUsage(msg string) error { return badUsageError{msg} }

type badUsageError struct{ msg string }

func (e badUsageError) Error() string { return e.msg }

// Exit returns an error that may be returned by Program.Main, which requests the
// main program to exit with the given code. If the exit code is 0, it returns nil.
func Exit(exit int) error {
	if exit == 0 {
		return nil
	}
	return exitError{exit}
}

type exitError struct{ exit int }

func (e exitError) Error() string { return "" }

// Program represents a subprogram.
type Program interface {
	// ShouldRun returns whether the subprogram should run.
	ShouldRun(f *Flags) bool
	// Run runs the subprogram.
	Run(fds [3]*os.File, f *Flags, args []string) error
}
<commit_msg>Fix build on Go 1.15.<commit_after>// Package prog provides the entry point to Elvish. Its subpackages correspond
// to subprograms of Elvish.
package prog

// This package sets up the basic environment and calls the appropriate
// "subprogram", one of the daemon, the terminal interface, or the web
// interface.

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"runtime/pprof"

	"src.elv.sh/pkg/logutil"
)

// Default port on which the web interface runs. The number is chosen because it
// resembles "elvi".
const defaultWebPort = 3171

// DeprecationLevel is a global flag that controls which deprecations to show.
// If its value is X, Elvish shows deprecations that should be shown for version
// 0.X.
var DeprecationLevel = 15

// SetDeprecationLevel sets ShowDeprecations to the given value, and returns a
// function to restore the old value.
func SetDeprecationLevel(level int) func() {
	save := DeprecationLevel
	DeprecationLevel = level
	return func() { DeprecationLevel = save }
}

// Flags keeps command-line flags.
type Flags struct {
	Log, CPUProfile string

	Help, Version, BuildInfo, JSON bool

	CodeInArg, CompileOnly, NoRc bool

	Web  bool
	Port int

	Daemon bool
	Forked int

	DB, Sock string
}

func newFlagSet(f *Flags) *flag.FlagSet {
	fs := flag.NewFlagSet("elvish", flag.ContinueOnError)
	// Error and usage will be printed explicitly.
	fs.SetOutput(ioutil.Discard)

	fs.StringVar(&f.Log, "log", "", "a file to write debug log to except for the daemon")
	fs.StringVar(&f.CPUProfile, "cpuprofile", "", "write cpu profile to file")

	fs.BoolVar(&f.Help, "help", false, "show usage help and quit")
	fs.BoolVar(&f.Version, "version", false, "show version and quit")
	fs.BoolVar(&f.BuildInfo, "buildinfo", false, "show build info and quit")
	fs.BoolVar(&f.JSON, "json", false, "show output in JSON. Useful with -buildinfo.")

	// The `-i` option is for compatibility with POSIX shells so that programs, such as the `script`
	// command, will work when asked to launch an interactive Elvish shell.
	fs.Bool("i", false, "force interactive mode; currently ignored")
	fs.BoolVar(&f.CodeInArg, "c", false, "take first argument as code to execute")
	fs.BoolVar(&f.CompileOnly, "compileonly", false, "Parse/Compile but do not execute")
	fs.BoolVar(&f.NoRc, "norc", false, "run elvish without invoking rc.elv")

	fs.BoolVar(&f.Web, "web", false, "run backend of web interface")
	fs.IntVar(&f.Port, "port", defaultWebPort, "the port of the web backend")

	fs.BoolVar(&f.Daemon, "daemon", false, "[internal flag] run the storage daemon instead of shell")

	fs.StringVar(&f.DB, "db", "", "[internal flag] path to the database")
	fs.StringVar(&f.Sock, "sock", "", "[internal flag] path to the daemon socket")

	fs.IntVar(&DeprecationLevel, "deprecation-level", DeprecationLevel, "show warnings for all features deprecated as of version 0.X")

	return fs
}

func usage(out io.Writer, f *flag.FlagSet) {
	fmt.Fprintln(out, "Usage: elvish [flags] [script]")
	fmt.Fprintln(out, "Supported flags:")
	f.SetOutput(out)
	f.PrintDefaults()
}

// Run parses command-line flags and runs the first applicable subprogram. It
// returns the exit status of the program.
func Run(fds [3]*os.File, args []string, programs ...Program) int {
	f := &Flags{}
	fs := newFlagSet(f)
	err := fs.Parse(args[1:])
	if err != nil {
		if err == flag.ErrHelp {
			// (*flag.FlagSet).Parse returns ErrHelp when -h or -help was
			// requested but *not* defined. Elvish defines -help, but not -h; so
			// this means that -h has been requested. Handle this by printing
			// the same message as an undefined flag.
			fmt.Fprintln(fds[2], "flag provided but not defined: -h")
		} else {
			fmt.Fprintln(fds[2], err)
		}
		usage(fds[2], fs)
		return 2
	}

	// Handle flags common to all subprograms.
	if f.CPUProfile != "" {
		f, err := os.Create(f.CPUProfile)
		if err != nil {
			fmt.Fprintln(fds[2], "Warning: cannot create CPU profile:", err)
			fmt.Fprintln(fds[2], "Continuing without CPU profiling.")
		} else {
			pprof.StartCPUProfile(f)
			defer pprof.StopCPUProfile()
		}
	}

	if f.Daemon {
		// We expect our stdout file handle is open on a unique log file for the daemon to write its
		// log messages. See daemon.Spawn() in pkg/daemon.
		logutil.SetOutput(fds[1])
	} else if f.Log != "" {
		err = logutil.SetOutputFile(f.Log)
		if err != nil {
			fmt.Fprintln(fds[2], err)
		}
	}

	if f.Help {
		usage(fds[1], fs)
		return 0
	}

	p := findProgram(f, programs)
	if p == nil {
		fmt.Fprintln(fds[2], "program bug: no suitable subprogram")
		return 2
	}

	err = p.Run(fds, f, fs.Args())
	if err == nil {
		return 0
	}
	if msg := err.Error(); msg != "" {
		fmt.Fprintln(fds[2], msg)
	}
	switch err := err.(type) {
	case badUsageError:
		usage(fds[2], fs)
	case exitError:
		return err.exit
	}
	return 2
}

func findProgram(f *Flags, programs []Program) Program {
	for _, program := range programs {
		if program.ShouldRun(f) {
			return program
		}
	}
	return nil
}

// BadUsage returns an error that may be returned by Program.Main, which
// requests the main program to print out a message, the usage information and
// exit with 2.
func BadUsage(msg string) error { return badUsageError{msg} }

type badUsageError struct{ msg string }

func (e badUsageError) Error() string { return e.msg }

// Exit returns an error that may be returned by Program.Main, which requests the
// main program to exit with the given code. If the exit code is 0, it returns nil.
func Exit(exit int) error {
	if exit == 0 {
		return nil
	}
	return exitError{exit}
}

type exitError struct{ exit int }

func (e exitError) Error() string { return "" }

// Program represents a subprogram.
type Program interface {
	// ShouldRun returns whether the subprogram should run.
	ShouldRun(f *Flags) bool
	// Run runs the subprogram.
	Run(fds [3]*os.File, f *Flags, args []string) error
}
<|endoftext|>


<commit_before>package util

import (
	"bufio"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

const (
	OvfEnvPath = "/var/lib/waagent/ovf-env.xml"
)

// ParseINI basic INI config file format into a map.
// Example expected format:
//     KEY=VAL
//     KEY2=VAL2
func ParseINI(s string) (map[string]string, error) {
	m := make(map[string]string)
	sc := bufio.NewScanner(strings.NewReader(s))

	for sc.Scan() {
		l := sc.Text() // format: K=V
		p := strings.Split(l, "=")
		if len(p) != 2 {
			return nil, fmt.Errorf("Unexpected config line: %q", l)
		}
		m[p[0]] = p[1]
	}
	if err := sc.Err(); err != nil {
		return nil, fmt.Errorf("Could not scan config file: %v", err)
	}
	return m, nil
}

// ScriptDir returns the absolute path of the running process.
func ScriptDir() (string, error) {
	p, err := filepath.Abs(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Dir(p), nil
}

// GetAzureUser returns the username provided at VM provisioning time to Azure.
func GetAzureUser() (string, error) {
	b, err := ioutil.ReadFile(OvfEnvPath)
	if err != nil {
		return "", err
	}

	var v struct {
		XMLName  xml.Name `xml:"Environment"`
		UserName string   `xml:"ProvisioningSection>LinuxProvisioningConfigurationSet>UserName"`
	}
	if err := xml.Unmarshal(b, &v); err != nil {
		return "", err
	}
	return v.UserName, nil
}

// PathExists checks if a path is a directory or file on the
// filesystem.
func PathExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, fmt.Errorf("util: error checking path %s: %v", path, err)
}

<commit_msg>Remove empty line<commit_after>package util

import (
	"bufio"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

const (
	OvfEnvPath = "/var/lib/waagent/ovf-env.xml"
)

// ParseINI basic INI config file format into a map.
// Example expected format:
//     KEY=VAL
//     KEY2=VAL2
func ParseINI(s string) (map[string]string, error) {
	m := make(map[string]string)
	sc := bufio.NewScanner(strings.NewReader(s))

	for sc.Scan() {
		l := sc.Text() // format: K=V
		p := strings.Split(l, "=")
		if len(p) != 2 {
			return nil, fmt.Errorf("Unexpected config line: %q", l)
		}
		m[p[0]] = p[1]
	}
	if err := sc.Err(); err != nil {
		return nil, fmt.Errorf("Could not scan config file: %v", err)
	}
	return m, nil
}

// ScriptDir returns the absolute path of the running process.
func ScriptDir() (string, error) {
	p, err := filepath.Abs(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Dir(p), nil
}

// GetAzureUser returns the username provided at VM provisioning time to Azure.
func GetAzureUser() (string, error) {
	b, err := ioutil.ReadFile(OvfEnvPath)
	if err != nil {
		return "", err
	}

	var v struct {
		XMLName  xml.Name `xml:"Environment"`
		UserName string   `xml:"ProvisioningSection>LinuxProvisioningConfigurationSet>UserName"`
	}
	if err := xml.Unmarshal(b, &v); err != nil {
		return "", err
	}
	return v.UserName, nil
}

// PathExists checks if a path is a directory or file on the
// filesystem.
func PathExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, fmt.Errorf("util: error checking path %s: %v", path, err)
}
<|endoftext|>


<commit_before>// goroar is an implementation of Roaring Bitmaps in Golang.
//Roaring bitmaps is a new form of compressed bitmaps, proposed by Daniel Lemire et. al., which often offers better compression and fast access than other compressed bitmap approaches.
package goroar

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strconv"
	"text/template"
)

type entry struct {
	key       uint16
	container container
}

type RoaringBitmap struct {
	containers []entry
}

// New creates a new RoaringBitmap
func New() *RoaringBitmap {
	containers := make([]entry, 0, 4)
	return &RoaringBitmap{containers}
}

// BitmapOf generates a new bitmap with the specified values set to true. The provided values don't have to be in sorted order, but it may be preferable to sort them from a performance point of view.
func BitmapOf(values ...uint32) *RoaringBitmap {
	rb := New()
	for _, value := range values {
		rb.Add(value)
	}
	return rb
}

// Add adds a uint32 value to the RoaringBitmap
func (rb *RoaringBitmap) Add(x uint32) {
	hb, lb := highlowbits(x)

	pos := rb.containerIndex(hb)
	if pos >= 0 {
		container := rb.containers[pos].container
		rb.containers[pos].container = container.add(lb)
	} else {
		ac := newArrayContainer()
		ac.add(lb)
		rb.increaseCapacity()

		loc := -pos - 1

		// insertion : shift the elements > x by one position to
		// the right and put x in it's appropriate place
		rb.containers = rb.containers[:len(rb.containers)+1]
		copy(rb.containers[loc+1:], rb.containers[loc:])
		e := entry{hb, ac}
		rb.containers[loc] = e
	}
}

// Contains checks whether the value in included, which is equivalent to checking
// if the corresponding bit is set (get in BitSet class).
func (rb *RoaringBitmap) Contains(i uint32) bool {
	pos := rb.containerIndex(highbits(i))
	if pos < 0 {
		return false
	}
	return rb.containers[pos].container.contains(lowbits(i))
}

// Cardinality returns the number of distinct integers (uint32) in the bitmap.
func (rb *RoaringBitmap) Cardinality() int {
	var cardinality int
	for _, entry := range rb.containers {
		cardinality = cardinality + entry.container.getCardinality()
	}
	return cardinality
}

// And computes the bitwise AND operation.
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) And(other *RoaringBitmap) {
	pos1 := 0
	pos2 := 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

Main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				rb.removeAtIndex(pos1)
				length1 = length1 - 1
				if pos1 == length1 {
					break Main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				pos2 = pos2 + 1
				if pos2 == length2 {
					break Main
				}
				s2 = other.keyAtIndex(pos2)
			} else {
				c := rb.containers[pos1].container.and(other.containers[pos2].container)

				if c.getCardinality() > 0 {
					rb.containers[pos1].container = c
					pos1 = pos1 + 1
				} else {
					rb.removeAtIndex(pos1)
					length1 = length1 - 1
				}
				pos2 = pos2 + 1
				if (pos1 == length1) || (pos2 == length2) {
					break Main
				}
				s1 = rb.keyAtIndex(pos1)
				s2 = other.keyAtIndex(pos2)
			}
		}
	}
	rb.resize(pos1)
}

// Or computes the bitwise OR operation.
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) Or(other *RoaringBitmap) {
	pos1, pos2 := 0, 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				pos1++
				if pos1 == length1 {
					break main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				rb.insertAt(pos1, s2, other.containers[pos2].container)
				pos1++
				length1++
				pos2++
				if pos2 == length2 {
					break main
				}
				s2 = other.containers[pos2].key
			} else {
				rb.containers[pos1].container = rb.containers[pos1].container.or(other.containers[pos2].container)
				pos1++
				pos2++
				if pos1 == length1 || pos2 == length2 {
					break main
				}
				s1 = rb.containers[pos1].key
				s2 = other.containers[pos2].key
			}
		}
	}
	if pos1 == length1 {
		rb.containers = append(rb.containers, other.containers[pos2:length2]...)
	}
}

// Xor computes the bitwise XOR operation.
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) Xor(other *RoaringBitmap) {
	pos1, pos2 := 0, 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				pos1++
				if pos1 == length1 {
					break main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				rb.insertAt(pos1, s2, other.containers[pos2].container)
				pos1++
				length1++
				pos2++
				if pos2 == length2 {
					break main
				}
				s2 = other.containers[pos2].key
			} else {
				c := rb.containers[pos1].container.xor(other.containers[pos2].container)
				if c.getCardinality() > 0 {
					rb.containers[pos1].container = c
					pos1++
				} else {
					rb.removeAtIndex(pos1)
					length1--
				}
				pos2++
				if pos1 == length1 || pos2 == length2 {
					break main
				}
				s1 = rb.containers[pos1].key
				s2 = other.containers[pos2].key
			}
		}
	}
	if pos1 == length1 {
		rb.containers = append(rb.containers, other.containers[pos2:length2]...)
	}
}

// AndNot computes the bitwise andNot operation (difference)
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) AndNot(other *RoaringBitmap) {
	pos1, pos2 := 0, 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				pos1++
				if pos1 == length1 {
					break main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				pos2++
				if pos2 == length2 {
					break main
				}
				s2 = other.containers[pos2].key
			} else {
				c := rb.containers[pos1].container.andNot(other.containers[pos2].container)
				if c.getCardinality() > 0 {
					rb.containers[pos1].container = c
					pos1++
				} else {
					rb.removeAtIndex(pos1)
					length1--
				}
				pos2++
				if pos1 == length1 || pos2 == length2 {
					break main
				}
				s1 = rb.containers[pos1].key
				s2 = other.containers[pos2].key
			}
		}
	}
}

// Iterator returns an iterator over the RoaringBitmap which can be used with "for range".
func (rb *RoaringBitmap) Iterator() <-chan uint32 {
	ch := make(chan uint32)
	go func() {
		// iterate over data
		for _, entry := range rb.containers {
			hs := uint32(entry.key) << 16
			switch typedContainer := entry.container.(type) {
			case *arrayContainer:
				pos := 0
				for pos < typedContainer.cardinality {
					ls := typedContainer.content[pos]
					pos = pos + 1
					ch <- (hs | uint32(ls))
				}
			case *bitmapContainer:
				i := typedContainer.nextSetBit(0)
				for i >= 0 {
					ch <- (hs | uint32(i))
					i = typedContainer.nextSetBit(i + 1)
				}
			}
		}
		close(ch)
	}()
	return ch
}

// Clone returns a copy of the original RoaringBitmap
func (rb *RoaringBitmap) Clone() *RoaringBitmap {
	containers := make([]entry, len(rb.containers))
	for i, value := range rb.containers {
		containers[i] = entry{value.key, value.container.clone()}
	}
	// copy(containers, rb.containers[0:])
	return &RoaringBitmap{containers}
}

func (rb *RoaringBitmap) String() string {
	var buffer bytes.Buffer
	name := []byte("RoaringBitmap[")

	buffer.Write(name)
	for val := range rb.Iterator() {
		buffer.WriteString(strconv.Itoa(int(val)))
		buffer.WriteString(", ")
	}
	if buffer.Len() > len(name) {
		buffer.Truncate(buffer.Len() - 2) // removes the last ", "
	}
	buffer.WriteString("]")
	return buffer.String()
}

// Stats prints statistics about the Roaring Bitmap's internals.
func (rb *RoaringBitmap) Stats() {
	const output = `* Roaring Bitmap Stats *
Cardinality: {{.Cardinality}}
Size uncompressed: {{.UncompressedSize}} bytes
Size compressed: {{.CompressedSize}} bytes
Number of containers: {{.TotalContainers}}
    {{.TotalAC}} ArrayContainers
    {{.TotalBC}} BitmapContainers
Average entries per ArrayContainer: {{.AverageAC}}
Max entries per ArrayContainer: {{.MaxAC}}
`
	type stats struct {
		Cardinality, TotalContainers, TotalAC, TotalBC int
		AverageAC, MaxAC                               string
		CompressedSize, UncompressedSize               int
	}

	var totalAC, totalBC, totalCardinalityAC int
	var maxAC int

	for _, c := range rb.containers {
		switch typedContainer := c.container.(type) {
		case *arrayContainer:
			if typedContainer.cardinality > maxAC {
				maxAC = typedContainer.cardinality
			}
			totalCardinalityAC += typedContainer.cardinality
			totalAC++
		case *bitmapContainer:
			totalBC++
		default:
		}
	}

	s := new(stats)
	s.Cardinality = rb.Cardinality()
	s.TotalContainers = len(rb.containers)
	s.TotalAC = totalAC
	s.TotalBC = totalBC
	s.CompressedSize = rb.SizeInBytes()
	s.UncompressedSize = rb.Cardinality() * 4

	if totalCardinalityAC > 0 {
		s.AverageAC = fmt.Sprintf("%6.2f", float32(totalCardinalityAC)/float32(totalAC))
		s.MaxAC = fmt.Sprintf("%d", maxAC)
	} else {
		s.AverageAC = "--"
		s.MaxAC = "--"
	}

	t := template.Must(template.New("stats").Parse(output))
	if err := t.Execute(os.Stdout, s); err != nil {
		log.Println("RoaringBitmap stats: ", err)
	}
}

func (rb *RoaringBitmap) SizeInBytes() int {
	size := 12 // size of RoaringBitmap struct
	for _, c := range rb.containers {
		size += 12 + c.container.sizeInBytes()
	}
	return size
}

func (rb *RoaringBitmap) resize(newLength int) {
	for i := newLength; i < len(rb.containers); i++ {
		rb.containers[i] = entry{}
	}
	rb.containers = rb.containers[:newLength]
}

func (rb *RoaringBitmap) keyAtIndex(pos int) uint16 {
	return rb.containers[pos].key
}

func (rb *RoaringBitmap) removeAtIndex(i int) {
	copy(rb.containers[i:], rb.containers[i+1:])
	rb.containers[len(rb.containers)-1] = entry{}
	rb.containers = rb.containers[:len(rb.containers)-1]
}

func (rb *RoaringBitmap) insertAt(i int, key uint16, c container) {
	rb.containers = append(rb.containers, entry{})
	copy(rb.containers[i+1:], rb.containers[i:])
	rb.containers[i] = entry{key, c}
}

func (rb *RoaringBitmap) containerIndex(key uint16) int {
	length := len(rb.containers)

	if length == 0 || rb.containers[length-1].key == key {
		return length - 1
	}

	return searchContainer(rb.containers, length, key)
}

func searchContainer(containers []entry, length int, key uint16) int {
	low := 0
	high := length - 1

	for low <= high {
		middleIndex := (low + high) >> 1
		middleValue := containers[middleIndex].key

		switch {
		case middleValue < key:
			low = middleIndex + 1
		case middleValue > key:
			high = middleIndex - 1
		default:
			return middleIndex
		}
	}
	return -(low + 1)
}

// increaseCapacity increases the slice capacity keeping the same length.
func (rb *RoaringBitmap) increaseCapacity() {
	length := len(rb.containers)
	if length+1 > cap(rb.containers) {
		var newCapacity int
		if length < 1024 {
			newCapacity = 2 * (length + 1)
		} else {
			newCapacity = 5 * (length + 1) / 4
		}

		newSlice := make([]entry, length, newCapacity)
		copy(newSlice, rb.containers)

		// increasing the length by 1
		rb.containers = newSlice
	}
}

// And computes the bitwise AND operation on two RoaringBitmaps.
// The input bitmaps are not modified.
// func And(x1, x2 *RoaringBitmap) *RoaringBitmap {
// 	panic("Not implemented")
// }
<commit_msg>Compression rate in Stats()<commit_after>// goroar is an implementation of Roaring Bitmaps in Golang.
//Roaring bitmaps is a new form of compressed bitmaps, proposed by Daniel Lemire et. al., which often offers better compression and fast access than other compressed bitmap approaches.
package goroar

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strconv"
	"text/template"
)

type entry struct {
	key       uint16
	container container
}

type RoaringBitmap struct {
	containers []entry
}

// New creates a new RoaringBitmap
func New() *RoaringBitmap {
	containers := make([]entry, 0, 4)
	return &RoaringBitmap{containers}
}

// BitmapOf generates a new bitmap with the specified values set to true. The provided values don't have to be in sorted order, but it may be preferable to sort them from a performance point of view.
func BitmapOf(values ...uint32) *RoaringBitmap {
	rb := New()
	for _, value := range values {
		rb.Add(value)
	}
	return rb
}

// Add adds a uint32 value to the RoaringBitmap
func (rb *RoaringBitmap) Add(x uint32) {
	hb, lb := highlowbits(x)

	pos := rb.containerIndex(hb)
	if pos >= 0 {
		container := rb.containers[pos].container
		rb.containers[pos].container = container.add(lb)
	} else {
		ac := newArrayContainer()
		ac.add(lb)
		rb.increaseCapacity()

		loc := -pos - 1

		// insertion : shift the elements > x by one position to
		// the right and put x in it's appropriate place
		rb.containers = rb.containers[:len(rb.containers)+1]
		copy(rb.containers[loc+1:], rb.containers[loc:])
		e := entry{hb, ac}
		rb.containers[loc] = e
	}
}

// Contains checks whether the value in included, which is equivalent to checking
// if the corresponding bit is set (get in BitSet class).
func (rb *RoaringBitmap) Contains(i uint32) bool {
	pos := rb.containerIndex(highbits(i))
	if pos < 0 {
		return false
	}
	return rb.containers[pos].container.contains(lowbits(i))
}

// Cardinality returns the number of distinct integers (uint32) in the bitmap.
func (rb *RoaringBitmap) Cardinality() int {
	var cardinality int
	for _, entry := range rb.containers {
		cardinality = cardinality + entry.container.getCardinality()
	}
	return cardinality
}

// And computes the bitwise AND operation.
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) And(other *RoaringBitmap) {
	pos1 := 0
	pos2 := 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

Main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				rb.removeAtIndex(pos1)
				length1 = length1 - 1
				if pos1 == length1 {
					break Main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				pos2 = pos2 + 1
				if pos2 == length2 {
					break Main
				}
				s2 = other.keyAtIndex(pos2)
			} else {
				c := rb.containers[pos1].container.and(other.containers[pos2].container)

				if c.getCardinality() > 0 {
					rb.containers[pos1].container = c
					pos1 = pos1 + 1
				} else {
					rb.removeAtIndex(pos1)
					length1 = length1 - 1
				}
				pos2 = pos2 + 1
				if (pos1 == length1) || (pos2 == length2) {
					break Main
				}
				s1 = rb.keyAtIndex(pos1)
				s2 = other.keyAtIndex(pos2)
			}
		}
	}
	rb.resize(pos1)
}

// Or computes the bitwise OR operation.
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) Or(other *RoaringBitmap) {
	pos1, pos2 := 0, 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				pos1++
				if pos1 == length1 {
					break main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				rb.insertAt(pos1, s2, other.containers[pos2].container)
				pos1++
				length1++
				pos2++
				if pos2 == length2 {
					break main
				}
				s2 = other.containers[pos2].key
			} else {
				rb.containers[pos1].container = rb.containers[pos1].container.or(other.containers[pos2].container)
				pos1++
				pos2++
				if pos1 == length1 || pos2 == length2 {
					break main
				}
				s1 = rb.containers[pos1].key
				s2 = other.containers[pos2].key
			}
		}
	}
	if pos1 == length1 {
		rb.containers = append(rb.containers, other.containers[pos2:length2]...)
	}
}

// Xor computes the bitwise XOR operation.
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) Xor(other *RoaringBitmap) {
	pos1, pos2 := 0, 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				pos1++
				if pos1 == length1 {
					break main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				rb.insertAt(pos1, s2, other.containers[pos2].container)
				pos1++
				length1++
				pos2++
				if pos2 == length2 {
					break main
				}
				s2 = other.containers[pos2].key
			} else {
				c := rb.containers[pos1].container.xor(other.containers[pos2].container)
				if c.getCardinality() > 0 {
					rb.containers[pos1].container = c
					pos1++
				} else {
					rb.removeAtIndex(pos1)
					length1--
				}
				pos2++
				if pos1 == length1 || pos2 == length2 {
					break main
				}
				s1 = rb.containers[pos1].key
				s2 = other.containers[pos2].key
			}
		}
	}
	if pos1 == length1 {
		rb.containers = append(rb.containers, other.containers[pos2:length2]...)
	}
}

// AndNot computes the bitwise andNot operation (difference)
// The receiving RoaringBitmap is modified - the input one is not.
func (rb *RoaringBitmap) AndNot(other *RoaringBitmap) {
	pos1, pos2 := 0, 0
	length1 := len(rb.containers)
	length2 := len(other.containers)

main:
	for pos1 < length1 && pos2 < length2 {
		s1 := rb.keyAtIndex(pos1)
		s2 := other.keyAtIndex(pos2)
		for {
			if s1 < s2 {
				pos1++
				if pos1 == length1 {
					break main
				}
				s1 = rb.keyAtIndex(pos1)
			} else if s1 > s2 {
				pos2++
				if pos2 == length2 {
					break main
				}
				s2 = other.containers[pos2].key
			} else {
				c := rb.containers[pos1].container.andNot(other.containers[pos2].container)
				if c.getCardinality() > 0 {
					rb.containers[pos1].container = c
					pos1++
				} else {
					rb.removeAtIndex(pos1)
					length1--
				}
				pos2++
				if pos1 == length1 || pos2 == length2 {
					break main
				}
				s1 = rb.containers[pos1].key
				s2 = other.containers[pos2].key
			}
		}
	}
}

// Iterator returns an iterator over the RoaringBitmap which can be used with "for range".
func (rb *RoaringBitmap) Iterator() <-chan uint32 {
	ch := make(chan uint32)
	go func() {
		// iterate over data
		for _, entry := range rb.containers {
			hs := uint32(entry.key) << 16
			switch typedContainer := entry.container.(type) {
			case *arrayContainer:
				pos := 0
				for pos < typedContainer.cardinality {
					ls := typedContainer.content[pos]
					pos = pos + 1
					ch <- (hs | uint32(ls))
				}
			case *bitmapContainer:
				i := typedContainer.nextSetBit(0)
				for i >= 0 {
					ch <- (hs | uint32(i))
					i = typedContainer.nextSetBit(i + 1)
				}
			}
		}
		close(ch)
	}()
	return ch
}

// Clone returns a copy of the original RoaringBitmap
func (rb *RoaringBitmap) Clone() *RoaringBitmap {
	containers := make([]entry, len(rb.containers))
	for i, value := range rb.containers {
		containers[i] = entry{value.key, value.container.clone()}
	}
	// copy(containers, rb.containers[0:])
	return &RoaringBitmap{containers}
}

func (rb *RoaringBitmap) String() string {
	var buffer bytes.Buffer
	name := []byte("RoaringBitmap[")

	buffer.Write(name)
	for val := range rb.Iterator() {
		buffer.WriteString(strconv.Itoa(int(val)))
		buffer.WriteString(", ")
	}
	if buffer.Len() > len(name) {
		buffer.Truncate(buffer.Len() - 2) // removes the last ", "
	}
	buffer.WriteString("]")
	return buffer.String()
}

// Stats prints statistics about the Roaring Bitmap's internals.
func (rb *RoaringBitmap) Stats() {
	const output = `* Roaring Bitmap Stats *
Cardinality: {{.Cardinality}}
Size uncompressed: {{.UncompressedSize}} bytes
Size compressed: {{.CompressedSize}} bytes ({{.CompressionRate}}%)
Number of containers: {{.TotalContainers}}
    {{.TotalAC}} ArrayContainers
    {{.TotalBC}} BitmapContainers
Average entries per ArrayContainer: {{.AverageAC}}
Max entries per ArrayContainer: {{.MaxAC}}
`
	type stats struct {
		Cardinality, TotalContainers, TotalAC, TotalBC int
		AverageAC, MaxAC                               string
		CompressedSize, UncompressedSize               int
		CompressionRate                                string
	}

	var totalAC, totalBC, totalCardinalityAC int
	var maxAC int

	for _, c := range rb.containers {
		switch typedContainer := c.container.(type) {
		case *arrayContainer:
			if typedContainer.cardinality > maxAC {
				maxAC = typedContainer.cardinality
			}
			totalCardinalityAC += typedContainer.cardinality
			totalAC++
		case *bitmapContainer:
			totalBC++
		default:
		}
	}

	s := new(stats)
	s.Cardinality = rb.Cardinality()
	s.TotalContainers = len(rb.containers)
	s.TotalAC = totalAC
	s.TotalBC = totalBC
	s.CompressedSize = rb.SizeInBytes()
	s.UncompressedSize = rb.Cardinality() * 4
	s.CompressionRate = fmt.Sprintf("%3.1f",
		float32(s.CompressedSize)/float32(s.UncompressedSize)*100.0)

	if totalCardinalityAC > 0 {
		s.AverageAC = fmt.Sprintf("%6.2f", float32(totalCardinalityAC)/float32(totalAC))
		s.MaxAC = fmt.Sprintf("%d", maxAC)
	} else {
		s.AverageAC = "--"
		s.MaxAC = "--"
	}

	t := template.Must(template.New("stats").Parse(output))
	if err := t.Execute(os.Stdout, s); err != nil {
		log.Println("RoaringBitmap stats: ", err)
	}
}

func (rb *RoaringBitmap) SizeInBytes() int {
	size := 12 // size of RoaringBitmap struct
	for _, c := range rb.containers {
		size += 12 + c.container.sizeInBytes()
	}
	return size
}

func (rb *RoaringBitmap) resize(newLength int) {
	for i := newLength; i < len(rb.containers); i++ {
		rb.containers[i] = entry{}
	}
	rb.containers = rb.containers[:newLength]
}

func (rb *RoaringBitmap) keyAtIndex(pos int) uint16 {
	return rb.containers[pos].key
}

func (rb *RoaringBitmap) removeAtIndex(i int) {
	copy(rb.containers[i:], rb.containers[i+1:])
	rb.containers[len(rb.containers)-1] = entry{}
	rb.containers = rb.containers[:len(rb.containers)-1]
}

func (rb *RoaringBitmap) insertAt(i int, key uint16, c container) {
	rb.containers = append(rb.containers, entry{})
	copy(rb.containers[i+1:], rb.containers[i:])
	rb.containers[i] = entry{key, c}
}

func (rb *RoaringBitmap) containerIndex(key uint16) int {
	length := len(rb.containers)

	if length == 0 || rb.containers[length-1].key == key {
		return length - 1
	}

	return searchContainer(rb.containers, length, key)
}

func searchContainer(containers []entry, length int, key uint16) int {
	low := 0
	high := length - 1

	for low <= high {
		middleIndex := (low + high) >> 1
		middleValue := containers[middleIndex].key

		switch {
		case middleValue < key:
			low = middleIndex + 1
		case middleValue > key:
			high = middleIndex - 1
		default:
			return middleIndex
		}
	}
	return -(low + 1)
}

// increaseCapacity increases the slice capacity keeping the same length.
func (rb *RoaringBitmap) increaseCapacity() {
	length := len(rb.containers)
	if length+1 > cap(rb.containers) {
		var newCapacity int
		if length < 1024 {
			newCapacity = 2 * (length + 1)
		} else {
			newCapacity = 5 * (length + 1) / 4
		}

		newSlice := make([]entry, length, newCapacity)
		copy(newSlice, rb.containers)

		// increasing the length by 1
		rb.containers = newSlice
	}
}

// And computes the bitwise AND operation on two RoaringBitmaps.
// The input bitmaps are not modified.
// func And(x1, x2 *RoaringBitmap) *RoaringBitmap {
// 	panic("Not implemented")
// }
<|endoftext|>


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

// +build ignore
//
// This build tag means that "go install golang.org/x/exp/shiny/..." doesn't
// install this example program. Use "go run main.go draw.go xy.go" to run it.

package main

import (
	"image"
	"image/color"
	_ "image/jpeg"
	_ "image/png"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/image/draw"
)

const maxBoard = 21 // Maximum board size we can handle.

var ZP = image.ZP // For brevity.

type stoneColor uint8

const (
	blank stoneColor = iota // Unused
	black
	white
)

// Piece represents a stone on the board. A nil Piece is "blank".
// The delta records pixel offset from the central dot.
type Piece struct {
	stone *Stone
	ij    IJ
	delta image.Point
	color stoneColor
}

type Board struct {
	Dims
	pieces         []*Piece // The board. Dimensions are 1-indexed. 1, 1 is the lower left corner.
	image          *image.RGBA
	stone          []Stone // All the black stones, followed by all the white stones.
	numBlackStones int
	numWhiteStones int
}

type Stone struct {
	originalImage *image.RGBA
	originalMask  *image.Alpha
	image         *image.RGBA
	mask          *image.Alpha
}

func NewBoard(dim, percent int) *Board {
	switch dim {
	case 9, 13, 19, 21:
	default:
		return nil
	}
	boardTexture := get("goboard.jpg", 0)
	b := new(Board)
	b.Dims.Init(dim, 100)
	b.pieces = make([]*Piece, maxBoard*maxBoard)
	b.image = image.NewRGBA(boardTexture.Bounds())
	draw.Draw(b.image, b.image.Bounds(), boardTexture, ZP, draw.Src)
	dir, err := os.Open("asset")
	if err != nil {
		log.Fatal(err)
	}
	defer dir.Close()
	names, err := dir.Readdirnames(0)
	if err != nil {
		log.Fatal(err)
	}
	circleMask := makeCircle()
	// Blackstones go first
	for _, name := range names {
		if strings.HasPrefix(name, "blackstone") {
			s, m := makeStone(name, circleMask)
			b.stone = append(b.stone, Stone{s, m, nil, nil})
			b.numBlackStones++
		}
	}
	for _, name := range names {
		if strings.HasPrefix(name, "whitestone") {
			s, m := makeStone(name, circleMask)
			b.stone = append(b.stone, Stone{s, m, nil, nil})
			b.numWhiteStones++
		}
	}
	b.Resize(percent) // TODO
	return b
}

func (b *Board) Resize(percent int) {
	b.Dims.Resize(percent)
	for i := range b.stone {
		stone := &b.stone[i]
		stone.image = resizeRGBA(stone.originalImage, b.stoneDiam)
		stone.mask = resizeAlpha(stone.originalMask, b.stoneDiam)
	}
}

func resizeRGBA(src *image.RGBA, size int) *image.RGBA {
	dst := image.NewRGBA(image.Rect(0, 0, size, size))
	draw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Src, nil)
	return dst
}

func resizeAlpha(src *image.Alpha, size int) *image.Alpha {
	dst := image.NewAlpha(image.Rect(0, 0, size, size))
	draw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Src, nil)
	return dst
}

func (b *Board) piece(ij IJ) *Piece {
	return b.pieces[(ij.j-1)*b.Dims.dim+ij.i-1]
}

func jitter() int {
	max := 25 * *scale / 100
	if max&1 == 0 {
		max++
	}
	return rand.Intn(max) - max/2
}

func (b *Board) putPiece(ij IJ, piece *Piece) {
	b.pieces[(ij.j-1)*b.Dims.dim+ij.i-1] = piece
	if piece != nil {
		piece.ij = ij
		piece.delta = image.Point{jitter(), jitter()}
	}
}

func (b *Board) selectBlackPiece() *Piece {
	return &Piece{
		stone: &b.stone[rand.Intn(b.numBlackStones)],
		color: black,
	}
}

func (b *Board) selectWhitePiece() *Piece {
	return &Piece{
		stone: &b.stone[b.numBlackStones+rand.Intn(b.numWhiteStones)],
		color: white,
	}
}

func makeStone(name string, circleMask *image.Alpha) (*image.RGBA, *image.Alpha) {
	stone := get(name, stoneSize0)
	dst := image.NewRGBA(stone.Bounds())
	// Make the whole area black, for the shadow.
	draw.Draw(dst, dst.Bounds(), image.Black, ZP, draw.Src)
	// Lay in the stone within the circle so it shows up inside the shadow.
	draw.DrawMask(dst, dst.Bounds(), stone, ZP, circleMask, ZP, draw.Over)
	return dst, makeShadowMask(stone)
}

func get(name string, size int) image.Image {
	f, err := os.Open(filepath.Join("asset", name))
	if err != nil {
		log.Fatal(err)
	}
	i, _, err := image.Decode(f)
	if err != nil {
		log.Fatal(err)
	}
	f.Close()
	if size != 0 {
		r := i.Bounds()
		if r.Dx() != size || r.Dy() != size {
			log.Fatalf("bad stone size %s for %s; must be %d[2]×%d[2]", r, name, size)
		}
	}
	return i
}

func makeCircle() *image.Alpha {
	mask := image.NewAlpha(image.Rect(0, 0, stoneSize0, stoneSize0))
	// Make alpha work on stone.
	// Shade gives shape, to be applied with black.
	for y := 0; y < stoneSize0; y++ {
		y2 := stoneSize0/2 - y
		y2 *= y2
		for x := 0; x < stoneSize0; x++ {
			x2 := stoneSize0/2 - x
			x2 *= x2
			if x2+y2 <= stoneRad2 {
				mask.SetAlpha(x, y, color.Alpha{255})
			}
		}
	}
	return mask
}

func makeShadowMask(stone image.Image) *image.Alpha {
	mask := image.NewAlpha(stone.Bounds())
	// Make alpha work on stone.
	// Shade gives shape, to be applied with black.
	const size = 256
	const diam = 225
	for y := 0; y < size; y++ {
		y2 := size/2 - y
		y2 *= y2
		for x := 0; x < size; x++ {
			x2 := size/2 - x
			x2 *= x2
			if x2+y2 > stoneRad2 {
				red, _, _, _ := stone.At(x, y).RGBA()
				mask.SetAlpha(x, y, color.Alpha{255 - uint8(red>>8)})
			} else {
				mask.SetAlpha(x, y, color.Alpha{255})
			}
		}
	}
	return mask
}

func (b *Board) Draw(m *image.RGBA) {
	r := b.image.Bounds()
	draw.Draw(m, r, b.image, ZP, draw.Src)
	// Vertical lines.
	x := b.xInset + b.squareWidth/2
	y := b.yInset + b.squareHeight/2
	wid := b.lineWidth
	for i := 0; i < b.dim; i++ {
		r := image.Rect(x, y, x+wid, y+(b.dim-1)*b.squareHeight)
		draw.Draw(m, r, image.Black, ZP, draw.Src)
		x += b.squareWidth
	}
	// Horizontal lines.
	x = b.xInset + b.squareWidth/2
	for i := 0; i < b.dim; i++ {
		r := image.Rect(x, y, x+(b.dim-1)*b.squareWidth+wid, y+wid)
		draw.Draw(m, r, image.Black, ZP, draw.Src)
		y += b.squareHeight
	}
	// Points.
	spot := 4
	if b.dim < 13 {
		spot = 3
	}
	points := []IJ{
		{spot, spot},
		{spot, (b.dim + 1) / 2},
		{spot, b.dim + 1 - spot},
		{(b.dim + 1) / 2, spot},
		{(b.dim + 1) / 2, (b.dim + 1) / 2},
		{(b.dim + 1) / 2, b.dim + 1 - spot},
		{b.dim + 1 - spot, spot},
		{b.dim + 1 - spot, (b.dim + 1) / 2},
		{b.dim + 1 - spot, b.dim + 1 - spot},
	}
	for _, ij := range points {
		b.drawPoint(m, ij)
	}
	// Pieces.
	for i := 1; i <= b.dim; i++ {
		for j := 1; j <= b.dim; j++ {
			ij := IJ{i, j}
			if p := b.piece(ij); p != nil {
				b.drawPiece(m, ij, p)
			}
		}
	}
}

func (b *Board) drawPoint(m *image.RGBA, ij IJ) {
	pt := ij.XYCenter(&b.Dims)
	wid := b.lineWidth
	sz := wid * 3 / 2
	r := image.Rect(pt.x-sz, pt.y-sz, pt.x+wid+sz, pt.y+wid+sz)
	draw.Draw(m, r, image.Black, ZP, draw.Src)
}

func (b *Board) drawPiece(m *image.RGBA, ij IJ, piece *Piece) {
	xy := ij.XYStone(&b.Dims)
	xy = xy.Add(piece.delta)
	draw.DrawMask(m, xy, piece.stone.image, ZP, piece.stone.mask, ZP, draw.Over)
}

func (b *Board) click(m *image.RGBA, x, y, button int) {
	ij, ok := XY{x, y}.IJ(&b.Dims)
	if !ok {
		return
	}
	switch button {
	case 1:
		b.putPiece(ij, b.selectBlackPiece())
	case 2:
		b.putPiece(ij, b.selectWhitePiece())
	case 3:
		b.putPiece(ij, nil)
	}
	render(m, b) // TODO: Connect this to paint events.
}
<commit_msg>shiny/example/goban: Fix comment<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build ignore
//
// This build tag means that "go install golang.org/x/exp/shiny/..." doesn't
// install this example program. Use "go run main.go board.go xy.go" to run it.

package main

import (
	"image"
	"image/color"
	_ "image/jpeg"
	_ "image/png"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/image/draw"
)

const maxBoard = 21 // Maximum board size we can handle.

var ZP = image.ZP // For brevity.

type stoneColor uint8

const (
	blank stoneColor = iota // Unused
	black
	white
)

// Piece represents a stone on the board. A nil Piece is "blank".
// The delta records pixel offset from the central dot.
type Piece struct {
	stone *Stone
	ij    IJ
	delta image.Point
	color stoneColor
}

type Board struct {
	Dims
	pieces         []*Piece // The board. Dimensions are 1-indexed. 1, 1 is the lower left corner.
	image          *image.RGBA
	stone          []Stone // All the black stones, followed by all the white stones.
	numBlackStones int
	numWhiteStones int
}

type Stone struct {
	originalImage *image.RGBA
	originalMask  *image.Alpha
	image         *image.RGBA
	mask          *image.Alpha
}

func NewBoard(dim, percent int) *Board {
	switch dim {
	case 9, 13, 19, 21:
	default:
		return nil
	}
	boardTexture := get("goboard.jpg", 0)
	b := new(Board)
	b.Dims.Init(dim, 100)
	b.pieces = make([]*Piece, maxBoard*maxBoard)
	b.image = image.NewRGBA(boardTexture.Bounds())
	draw.Draw(b.image, b.image.Bounds(), boardTexture, ZP, draw.Src)
	dir, err := os.Open("asset")
	if err != nil {
		log.Fatal(err)
	}
	defer dir.Close()
	names, err := dir.Readdirnames(0)
	if err != nil {
		log.Fatal(err)
	}
	circleMask := makeCircle()
	// Blackstones go first
	for _, name := range names {
		if strings.HasPrefix(name, "blackstone") {
			s, m := makeStone(name, circleMask)
			b.stone = append(b.stone, Stone{s, m, nil, nil})
			b.numBlackStones++
		}
	}
	for _, name := range names {
		if strings.HasPrefix(name, "whitestone") {
			s, m := makeStone(name, circleMask)
			b.stone = append(b.stone, Stone{s, m, nil, nil})
			b.numWhiteStones++
		}
	}
	b.Resize(percent) // TODO
	return b
}

func (b *Board) Resize(percent int) {
	b.Dims.Resize(percent)
	for i := range b.stone {
		stone := &b.stone[i]
		stone.image = resizeRGBA(stone.originalImage, b.stoneDiam)
		stone.mask = resizeAlpha(stone.originalMask, b.stoneDiam)
	}
}

func resizeRGBA(src *image.RGBA, size int) *image.RGBA {
	dst := image.NewRGBA(image.Rect(0, 0, size, size))
	draw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Src, nil)
	return dst
}

func resizeAlpha(src *image.Alpha, size int) *image.Alpha {
	dst := image.NewAlpha(image.Rect(0, 0, size, size))
	draw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Src, nil)
	return dst
}

func (b *Board) piece(ij IJ) *Piece {
	return b.pieces[(ij.j-1)*b.Dims.dim+ij.i-1]
}

func jitter() int {
	max := 25 * *scale / 100
	if max&1 == 0 {
		max++
	}
	return rand.Intn(max) - max/2
}

func (b *Board) putPiece(ij IJ, piece *Piece) {
	b.pieces[(ij.j-1)*b.Dims.dim+ij.i-1] = piece
	if piece != nil {
		piece.ij = ij
		piece.delta = image.Point{jitter(), jitter()}
	}
}

func (b *Board) selectBlackPiece() *Piece {
	return &Piece{
		stone: &b.stone[rand.Intn(b.numBlackStones)],
		color: black,
	}
}

func (b *Board) selectWhitePiece() *Piece {
	return &Piece{
		stone: &b.stone[b.numBlackStones+rand.Intn(b.numWhiteStones)],
		color: white,
	}
}

func makeStone(name string, circleMask *image.Alpha) (*image.RGBA, *image.Alpha) {
	stone := get(name, stoneSize0)
	dst := image.NewRGBA(stone.Bounds())
	// Make the whole area black, for the shadow.
	draw.Draw(dst, dst.Bounds(), image.Black, ZP, draw.Src)
	// Lay in the stone within the circle so it shows up inside the shadow.
	draw.DrawMask(dst, dst.Bounds(), stone, ZP, circleMask, ZP, draw.Over)
	return dst, makeShadowMask(stone)
}

func get(name string, size int) image.Image {
	f, err := os.Open(filepath.Join("asset", name))
	if err != nil {
		log.Fatal(err)
	}
	i, _, err := image.Decode(f)
	if err != nil {
		log.Fatal(err)
	}
	f.Close()
	if size != 0 {
		r := i.Bounds()
		if r.Dx() != size || r.Dy() != size {
			log.Fatalf("bad stone size %s for %s; must be %d[2]×%d[2]", r, name, size)
		}
	}
	return i
}

func makeCircle() *image.Alpha {
	mask := image.NewAlpha(image.Rect(0, 0, stoneSize0, stoneSize0))
	// Make alpha work on stone.
	// Shade gives shape, to be applied with black.
	for y := 0; y < stoneSize0; y++ {
		y2 := stoneSize0/2 - y
		y2 *= y2
		for x := 0; x < stoneSize0; x++ {
			x2 := stoneSize0/2 - x
			x2 *= x2
			if x2+y2 <= stoneRad2 {
				mask.SetAlpha(x, y, color.Alpha{255})
			}
		}
	}
	return mask
}

func makeShadowMask(stone image.Image) *image.Alpha {
	mask := image.NewAlpha(stone.Bounds())
	// Make alpha work on stone.
	// Shade gives shape, to be applied with black.
	const size = 256
	const diam = 225
	for y := 0; y < size; y++ {
		y2 := size/2 - y
		y2 *= y2
		for x := 0; x < size; x++ {
			x2 := size/2 - x
			x2 *= x2
			if x2+y2 > stoneRad2 {
				red, _, _, _ := stone.At(x, y).RGBA()
				mask.SetAlpha(x, y, color.Alpha{255 - uint8(red>>8)})
			} else {
				mask.SetAlpha(x, y, color.Alpha{255})
			}
		}
	}
	return mask
}

func (b *Board) Draw(m *image.RGBA) {
	r := b.image.Bounds()
	draw.Draw(m, r, b.image, ZP, draw.Src)
	// Vertical lines.
	x := b.xInset + b.squareWidth/2
	y := b.yInset + b.squareHeight/2
	wid := b.lineWidth
	for i := 0; i < b.dim; i++ {
		r := image.Rect(x, y, x+wid, y+(b.dim-1)*b.squareHeight)
		draw.Draw(m, r, image.Black, ZP, draw.Src)
		x += b.squareWidth
	}
	// Horizontal lines.
	x = b.xInset + b.squareWidth/2
	for i := 0; i < b.dim; i++ {
		r := image.Rect(x, y, x+(b.dim-1)*b.squareWidth+wid, y+wid)
		draw.Draw(m, r, image.Black, ZP, draw.Src)
		y += b.squareHeight
	}
	// Points.
	spot := 4
	if b.dim < 13 {
		spot = 3
	}
	points := []IJ{
		{spot, spot},
		{spot, (b.dim + 1) / 2},
		{spot, b.dim + 1 - spot},
		{(b.dim + 1) / 2, spot},
		{(b.dim + 1) / 2, (b.dim + 1) / 2},
		{(b.dim + 1) / 2, b.dim + 1 - spot},
		{b.dim + 1 - spot, spot},
		{b.dim + 1 - spot, (b.dim + 1) / 2},
		{b.dim + 1 - spot, b.dim + 1 - spot},
	}
	for _, ij := range points {
		b.drawPoint(m, ij)
	}
	// Pieces.
	for i := 1; i <= b.dim; i++ {
		for j := 1; j <= b.dim; j++ {
			ij := IJ{i, j}
			if p := b.piece(ij); p != nil {
				b.drawPiece(m, ij, p)
			}
		}
	}
}

func (b *Board) drawPoint(m *image.RGBA, ij IJ) {
	pt := ij.XYCenter(&b.Dims)
	wid := b.lineWidth
	sz := wid * 3 / 2
	r := image.Rect(pt.x-sz, pt.y-sz, pt.x+wid+sz, pt.y+wid+sz)
	draw.Draw(m, r, image.Black, ZP, draw.Src)
}

func (b *Board) drawPiece(m *image.RGBA, ij IJ, piece *Piece) {
	xy := ij.XYStone(&b.Dims)
	xy = xy.Add(piece.delta)
	draw.DrawMask(m, xy, piece.stone.image, ZP, piece.stone.mask, ZP, draw.Over)
}

func (b *Board) click(m *image.RGBA, x, y, button int) {
	ij, ok := XY{x, y}.IJ(&b.Dims)
	if !ok {
		return
	}
	switch button {
	case 1:
		b.putPiece(ij, b.selectBlackPiece())
	case 2:
		b.putPiece(ij, b.selectWhitePiece())
	case 3:
		b.putPiece(ij, nil)
	}
	render(m, b) // TODO: Connect this to paint events.
}
<|endoftext|>


<commit_before>package jobs

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/joyent/containerpilot/commands"
	"github.com/joyent/containerpilot/discovery"
	"github.com/joyent/containerpilot/events"
	log "github.com/sirupsen/logrus"
)

// Some magic numbers used internally by restart limits
const (
	unlimited       = -1
	eventBufferSize = 1000
)

// JobStatus is an enum of job health status
type JobStatus int

// JobStatus enum
const (
	statusIdle JobStatus = iota // will be default value before starting
	statusUnknown
	statusHealthy
	statusUnhealthy
	statusMaintenance
	statusAlwaysHealthy
)

func (i JobStatus) String() string {
	switch i {
	case 2:
		return "healthy"
	case 3:
		return "unhealthy"
	case 4:
		return "maintenance"
	case 5:
		// for hardcoded "always healthy" jobs
		return "healthy"
	default:
		// both idle and unknown return unknown for purposes of serialization
		return "unknown"
	}
}

// Job manages the state of a job and its start/stop conditions
type Job struct {
	Name string
	exec *commands.Command

	// service health and discovery
	Status          JobStatus
	statusLock      *sync.RWMutex
	Service         *discovery.ServiceDefinition
	healthCheckExec *commands.Command
	healthCheckName string

	// starting events
	startEvent        events.Event
	startTimeout      time.Duration
	startsRemain      int
	startTimeoutEvent events.Event

	// stopping events
	stoppingWaitEvent events.Event
	stoppingTimeout   time.Duration

	// timing and restarts
	heartbeat      time.Duration
	restartLimit   int
	restartsRemain int
	frequency      time.Duration

	events.EventHandler // Event handling
}

// NewJob creates a new Job from a Config
func NewJob(cfg *Config) *Job {
	job := &Job{
		Name:              cfg.Name,
		exec:              cfg.exec,
		heartbeat:         cfg.heartbeatInterval,
		Service:           cfg.serviceDefinition,
		healthCheckExec:   cfg.healthCheckExec,
		startEvent:        cfg.whenEvent,
		startTimeout:      cfg.whenTimeout,
		startsRemain:      cfg.whenStartsLimit,
		stoppingWaitEvent: cfg.stoppingWaitEvent,
		stoppingTimeout:   cfg.stoppingTimeout,
		restartLimit:      cfg.restartLimit,
		restartsRemain:    cfg.restartLimit,
		frequency:         cfg.freqInterval,
	}
	job.Rx = make(chan events.Event, eventBufferSize)
	job.statusLock = &sync.RWMutex{}
	if job.Name == "containerpilot" {
		// right now this hardcodes the telemetry service to
		// be always "healthy", but maybe we want to have it verify itself
		// before heartbeating in the future?
		job.setStatus(statusAlwaysHealthy)
	}
	return job
}

// FromConfigs creates Jobs from a slice of validated Configs
func FromConfigs(cfgs []*Config) []*Job {
	jobs := []*Job{}
	for _, cfg := range cfgs {
		job := NewJob(cfg)
		jobs = append(jobs, job)
	}
	return jobs
}

// SendHeartbeat sends a heartbeat for this Job's service
func (job *Job) SendHeartbeat() {
	if job.Service != nil {
		job.Service.SendHeartbeat()
	}
}

// GetStatus returns the current health status of the Job
func (job *Job) GetStatus() JobStatus {
	job.statusLock.RLock()
	defer job.statusLock.RUnlock()
	return job.Status
}

func (job *Job) setStatus(status JobStatus) {
	job.statusLock.Lock()
	defer job.statusLock.Unlock()
	if job.Status != statusAlwaysHealthy {
		job.Status = status
	}
}

// MarkForMaintenance marks this Job's service for maintenance
func (job *Job) MarkForMaintenance() {
	job.setStatus(statusMaintenance)
	if job.Service != nil {
		job.Service.MarkForMaintenance()
	}
}

// Deregister will deregister this instance of Job's service
func (job *Job) Deregister() {
	if job.Service != nil {
		job.Service.Deregister()
	}
}

// HealthCheck runs the Job's health check executable
func (job *Job) HealthCheck(ctx context.Context) {
	if job.healthCheckExec != nil {
		job.healthCheckExec.Run(ctx, job.Bus)
	}
}

// StartJob runs the Job's executable
func (job *Job) StartJob(ctx context.Context) {
	job.startTimeoutEvent = events.NonEvent
	job.setStatus(statusUnknown)
	if job.exec != nil {
		job.exec.Run(ctx, job.Bus)
	}
}

// Kill sends SIGTERM to the Job's executable, if any
func (job *Job) Kill() {
	if job.exec != nil {
		job.exec.Kill()
	}
}

// Run executes the event loop for the Job
func (job *Job) Run() {
	ctx, cancel := context.WithCancel(context.Background())

	if job.frequency > 0 {
		events.NewEventTimer(ctx, job.Rx, job.frequency,
			fmt.Sprintf("%s.run-every", job.Name))
	}
	if job.heartbeat > 0 {
		events.NewEventTimer(ctx, job.Rx, job.heartbeat,
			fmt.Sprintf("%s.heartbeat", job.Name))
	}
	if job.startTimeout > 0 {
		timeoutName := fmt.Sprintf("%s.wait-timeout", job.Name)
		events.NewEventTimeout(ctx, job.Rx, job.startTimeout, timeoutName)
		job.startTimeoutEvent = events.Event{events.TimerExpired, timeoutName}
	} else {
		job.startTimeoutEvent = events.NonEvent
	}

	go func() {
		defer job.cleanup(ctx, cancel)
		for {
			select {
			case event, ok := <-job.Rx:
				if !ok {
					return
				}
				if job.processEvent(ctx, event) {
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()
}

func (job *Job) processEvent(ctx context.Context, event events.Event) bool {
	runEverySource := fmt.Sprintf("%s.run-every", job.Name)
	heartbeatSource := fmt.Sprintf("%s.heartbeat", job.Name)
	var healthCheckName string
	if job.healthCheckExec != nil {
		healthCheckName = job.healthCheckExec.Name
	}

	switch event {
	case events.Event{events.TimerExpired, heartbeatSource}:
		status := job.GetStatus()
		if status != statusMaintenance && status != statusIdle {
			if job.healthCheckExec != nil {
				job.HealthCheck(ctx)
			} else if job.Service != nil {
				// this is the case for non-checked but advertised
				// services like the telemetry endpoint
				job.SendHeartbeat()
			}
		}
	case job.startTimeoutEvent:
		job.Bus.Publish(events.Event{
			Code: events.TimerExpired, Source: job.Name})
		job.Rx <- events.Event{Code: events.Quit, Source: job.Name}
	case events.Event{events.TimerExpired, runEverySource}:
		if !job.restartPermitted() {
			log.Debugf("interval expired but restart not permitted: %v",
				job.Name)
			job.startEvent = events.NonEvent
			return true
		}
		job.restartsRemain--
		job.StartJob(ctx)
	case events.Event{events.ExitFailed, healthCheckName}:
		if job.GetStatus() != statusMaintenance {
			job.setStatus(statusUnhealthy)
			job.Bus.Publish(events.Event{events.StatusUnhealthy, job.Name})
		}
	case events.Event{events.ExitSuccess, healthCheckName}:
		if job.GetStatus() != statusMaintenance {
			job.setStatus(statusHealthy)
			job.Bus.Publish(events.Event{events.StatusHealthy, job.Name})
			job.SendHeartbeat()
		}
	case
		events.Event{events.Quit, job.Name},
		events.QuitByClose,
		events.GlobalShutdown:
		if (job.startEvent.Code == events.Stopping ||
			job.startEvent.Code == events.Stopped) &&
			job.exec != nil {
			// "pre-stop" and "post-stop" style jobs ignore the global
			// shutdown and return on their ExitSuccess/ExitFailed.
			// if the stop timeout on the global shutdown is exceeded
			// the whole process gets SIGKILL
			break
		}
		job.startEvent = events.NonEvent
		return true
	case events.GlobalEnterMaintenance:
		job.MarkForMaintenance()
	case events.GlobalExitMaintenance:
		job.setStatus(statusUnknown)
	case
		events.Event{events.ExitSuccess, job.Name},
		events.Event{events.ExitFailed, job.Name}:
		if job.frequency > 0 {
			break // periodic jobs ignore previous events
		}
		if job.restartPermitted() {
			job.restartsRemain--
			job.StartJob(ctx)
			break
		}
		if job.startsRemain != 0 {
			break
		}
		log.Debugf("job exited but restart not permitted: %v", job.Name)
		job.startEvent = events.NonEvent
		job.setStatus(statusUnknown)
		return true
	case job.startEvent:
		if job.startsRemain == 0 {
			job.startEvent = events.NonEvent
			return true
		}
		if job.startsRemain != unlimited {
			// if we have unlimited restarts we want to make sure we don't
			// decrement forever and then wrap-around
			job.startsRemain--
			if job.startsRemain == 0 || job.restartsRemain == 0 {
				// prevent ourselves from receiving the start event again
				// if it fires while we're still running the job's exec
				job.startEvent = events.NonEvent
			}
		}
		job.StartJob(ctx)
	}
	return false
}

func (job *Job) restartPermitted() bool {
	if job.restartLimit == unlimited || job.restartsRemain > 0 {
		return true
	}
	return false
}

// cleanup fires the Stopping event and will wait to receive a stoppingWaitEvent
// if one is configured. cleans up registration to event bus and closes all
// channels and contexts when done.
func (job *Job) cleanup(ctx context.Context, cancel context.CancelFunc) {
	stoppingTimeout := fmt.Sprintf("%s.stopping-timeout", job.Name)
	job.Bus.Publish(events.Event{Code: events.Stopping, Source: job.Name})
	if job.stoppingWaitEvent != events.NonEvent {
		if job.stoppingTimeout > 0 {
			// not having this set is a programmer error not a runtime error
			events.NewEventTimeout(ctx, job.Rx,
				job.stoppingTimeout, stoppingTimeout)
		}
	loop:
		for {
			event := <-job.Rx
			switch event {
			case job.stoppingWaitEvent:
				break loop
			case events.Event{events.Stopping, stoppingTimeout}:
				break loop
			}
		}
	}
	cancel()
	job.Deregister()         // deregister from Consul
	job.Unsubscribe(job.Bus) // deregister from events
	job.Bus.Publish(events.Event{Code: events.Stopped, Source: job.Name})
}

// String implements the stdlib fmt.Stringer interface for pretty-printing
func (job *Job) String() string {
	return "jobs.Job[" + job.Name + "]"
}
<commit_msg>make stopping/stopped jobs complete exit on reload (#466)<commit_after>package jobs

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/joyent/containerpilot/commands"
	"github.com/joyent/containerpilot/discovery"
	"github.com/joyent/containerpilot/events"
	log "github.com/sirupsen/logrus"
)

// Some magic numbers used internally by restart limits
const (
	unlimited       = -1
	eventBufferSize = 1000
)

// JobStatus is an enum of job health status
type JobStatus int

// JobStatus enum
const (
	statusIdle JobStatus = iota // will be default value before starting
	statusUnknown
	statusHealthy
	statusUnhealthy
	statusMaintenance
	statusAlwaysHealthy
)

func (i JobStatus) String() string {
	switch i {
	case 2:
		return "healthy"
	case 3:
		return "unhealthy"
	case 4:
		return "maintenance"
	case 5:
		// for hardcoded "always healthy" jobs
		return "healthy"
	default:
		// both idle and unknown return unknown for purposes of serialization
		return "unknown"
	}
}

// Job manages the state of a job and its start/stop conditions
type Job struct {
	Name string
	exec *commands.Command

	// service health and discovery
	Status          JobStatus
	statusLock      *sync.RWMutex
	Service         *discovery.ServiceDefinition
	healthCheckExec *commands.Command
	healthCheckName string

	// starting events
	startEvent        events.Event
	startTimeout      time.Duration
	startsRemain      int
	startTimeoutEvent events.Event

	// stopping events
	stoppingWaitEvent events.Event
	stoppingTimeout   time.Duration

	// timing and restarts
	heartbeat      time.Duration
	restartLimit   int
	restartsRemain int
	frequency      time.Duration

	events.EventHandler // Event handling
}

// NewJob creates a new Job from a Config
func NewJob(cfg *Config) *Job {
	job := &Job{
		Name:              cfg.Name,
		exec:              cfg.exec,
		heartbeat:         cfg.heartbeatInterval,
		Service:           cfg.serviceDefinition,
		healthCheckExec:   cfg.healthCheckExec,
		startEvent:        cfg.whenEvent,
		startTimeout:      cfg.whenTimeout,
		startsRemain:      cfg.whenStartsLimit,
		stoppingWaitEvent: cfg.stoppingWaitEvent,
		stoppingTimeout:   cfg.stoppingTimeout,
		restartLimit:      cfg.restartLimit,
		restartsRemain:    cfg.restartLimit,
		frequency:         cfg.freqInterval,
	}
	job.Rx = make(chan events.Event, eventBufferSize)
	job.statusLock = &sync.RWMutex{}
	if job.Name == "containerpilot" {
		// right now this hardcodes the telemetry service to
		// be always "healthy", but maybe we want to have it verify itself
		// before heartbeating in the future?
		job.setStatus(statusAlwaysHealthy)
	}
	return job
}

// FromConfigs creates Jobs from a slice of validated Configs
func FromConfigs(cfgs []*Config) []*Job {
	jobs := []*Job{}
	for _, cfg := range cfgs {
		job := NewJob(cfg)
		jobs = append(jobs, job)
	}
	return jobs
}

// SendHeartbeat sends a heartbeat for this Job's service
func (job *Job) SendHeartbeat() {
	if job.Service != nil {
		job.Service.SendHeartbeat()
	}
}

// GetStatus returns the current health status of the Job
func (job *Job) GetStatus() JobStatus {
	job.statusLock.RLock()
	defer job.statusLock.RUnlock()
	return job.Status
}

func (job *Job) setStatus(status JobStatus) {
	job.statusLock.Lock()
	defer job.statusLock.Unlock()
	if job.Status != statusAlwaysHealthy {
		job.Status = status
	}
}

// MarkForMaintenance marks this Job's service for maintenance
func (job *Job) MarkForMaintenance() {
	job.setStatus(statusMaintenance)
	if job.Service != nil {
		job.Service.MarkForMaintenance()
	}
}

// Deregister will deregister this instance of Job's service
func (job *Job) Deregister() {
	if job.Service != nil {
		job.Service.Deregister()
	}
}

// HealthCheck runs the Job's health check executable
func (job *Job) HealthCheck(ctx context.Context) {
	if job.healthCheckExec != nil {
		job.healthCheckExec.Run(ctx, job.Bus)
	}
}

// StartJob runs the Job's executable
func (job *Job) StartJob(ctx context.Context) {
	job.startTimeoutEvent = events.NonEvent
	job.setStatus(statusUnknown)
	if job.exec != nil {
		job.exec.Run(ctx, job.Bus)
	}
}

// Kill sends SIGTERM to the Job's executable, if any
func (job *Job) Kill() {
	if job.exec != nil {
		job.exec.Kill()
	}
}

// Run executes the event loop for the Job
func (job *Job) Run() {
	ctx, cancel := context.WithCancel(context.Background())

	if job.frequency > 0 {
		events.NewEventTimer(ctx, job.Rx, job.frequency,
			fmt.Sprintf("%s.run-every", job.Name))
	}
	if job.heartbeat > 0 {
		events.NewEventTimer(ctx, job.Rx, job.heartbeat,
			fmt.Sprintf("%s.heartbeat", job.Name))
	}
	if job.startTimeout > 0 {
		timeoutName := fmt.Sprintf("%s.wait-timeout", job.Name)
		events.NewEventTimeout(ctx, job.Rx, job.startTimeout, timeoutName)
		job.startTimeoutEvent = events.Event{events.TimerExpired, timeoutName}
	} else {
		job.startTimeoutEvent = events.NonEvent
	}

	go func() {
		defer job.cleanup(ctx, cancel)
		for {
			select {
			case event, ok := <-job.Rx:
				if !ok {
					return
				}
				if job.processEvent(ctx, event) {
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()
}

func (job *Job) processEvent(ctx context.Context, event events.Event) bool {
	runEverySource := fmt.Sprintf("%s.run-every", job.Name)
	heartbeatSource := fmt.Sprintf("%s.heartbeat", job.Name)
	var healthCheckName string
	if job.healthCheckExec != nil {
		healthCheckName = job.healthCheckExec.Name
	}

	switch event {
	case events.Event{events.TimerExpired, heartbeatSource}:
		status := job.GetStatus()
		if status != statusMaintenance && status != statusIdle {
			if job.healthCheckExec != nil {
				job.HealthCheck(ctx)
			} else if job.Service != nil {
				// this is the case for non-checked but advertised
				// services like the telemetry endpoint
				job.SendHeartbeat()
			}
		}
	case job.startTimeoutEvent:
		job.Bus.Publish(events.Event{
			Code: events.TimerExpired, Source: job.Name})
		job.Rx <- events.Event{Code: events.Quit, Source: job.Name}
	case events.Event{events.TimerExpired, runEverySource}:
		if !job.restartPermitted() {
			log.Debugf("interval expired but restart not permitted: %v",
				job.Name)
			job.startEvent = events.NonEvent
			return true
		}
		job.restartsRemain--
		job.StartJob(ctx)
	case events.Event{events.ExitFailed, healthCheckName}:
		if job.GetStatus() != statusMaintenance {
			job.setStatus(statusUnhealthy)
			job.Bus.Publish(events.Event{events.StatusUnhealthy, job.Name})
		}
	case events.Event{events.ExitSuccess, healthCheckName}:
		if job.GetStatus() != statusMaintenance {
			job.setStatus(statusHealthy)
			job.Bus.Publish(events.Event{events.StatusHealthy, job.Name})
			job.SendHeartbeat()
		}
	case
		events.Event{events.Quit, job.Name},
		events.QuitByClose,
		events.GlobalShutdown:
		job.restartsRemain = 0 // no more restarts
		if (job.startEvent.Code == events.Stopping ||
			job.startEvent.Code == events.Stopped) &&
			job.exec != nil {
			// "pre-stop" and "post-stop" style jobs ignore the global
			// shutdown and return on their ExitSuccess/ExitFailed.
			// if the stop timeout on the global shutdown is exceeded
			// the whole process gets SIGKILL
			if job.startsRemain == unlimited {
				job.startsRemain = 1
			}
			break
		}
		job.startsRemain = 0
		job.startEvent = events.NonEvent
		return true
	case events.GlobalEnterMaintenance:
		job.MarkForMaintenance()
	case events.GlobalExitMaintenance:
		job.setStatus(statusUnknown)
	case
		events.Event{events.ExitSuccess, job.Name},
		events.Event{events.ExitFailed, job.Name}:
		if job.frequency > 0 {
			break // periodic jobs ignore previous events
		}
		if job.restartPermitted() {
			job.restartsRemain--
			job.StartJob(ctx)
			break
		}
		if job.startsRemain != 0 {
			break
		}
		log.Debugf("job exited but restart not permitted: %v", job.Name)
		job.startEvent = events.NonEvent
		job.setStatus(statusUnknown)
		return true
	case job.startEvent:
		if job.startsRemain == 0 {
			job.startEvent = events.NonEvent
			return true
		}
		if job.startsRemain != unlimited {
			// if we have unlimited restarts we want to make sure we don't
			// decrement forever and then wrap-around
			job.startsRemain--
			if job.startsRemain == 0 || job.restartsRemain == 0 {
				// prevent ourselves from receiving the start event again
				// if it fires while we're still running the job's exec
				job.startEvent = events.NonEvent
			}
		}
		job.StartJob(ctx)
	}
	return false
}

func (job *Job) restartPermitted() bool {
	if job.restartLimit == unlimited || job.restartsRemain > 0 {
		return true
	}
	return false
}

// cleanup fires the Stopping event and will wait to receive a stoppingWaitEvent
// if one is configured. cleans up registration to event bus and closes all
// channels and contexts when done.
func (job *Job) cleanup(ctx context.Context, cancel context.CancelFunc) {
	stoppingTimeout := fmt.Sprintf("%s.stopping-timeout", job.Name)
	job.Bus.Publish(events.Event{Code: events.Stopping, Source: job.Name})
	if job.stoppingWaitEvent != events.NonEvent {
		if job.stoppingTimeout > 0 {
			// not having this set is a programmer error not a runtime error
			events.NewEventTimeout(ctx, job.Rx,
				job.stoppingTimeout, stoppingTimeout)
		}
	loop:
		for {
			event := <-job.Rx
			switch event {
			case job.stoppingWaitEvent:
				break loop
			case events.Event{events.Stopping, stoppingTimeout}:
				break loop
			}
		}
	}
	cancel()
	job.Deregister()         // deregister from Consul
	job.Unsubscribe(job.Bus) // deregister from events
	job.Bus.Publish(events.Event{Code: events.Stopped, Source: job.Name})
}

// String implements the stdlib fmt.Stringer interface for pretty-printing
func (job *Job) String() string {
	return "jobs.Job[" + job.Name + "]"
}
<|endoftext|>


<commit_before>package router

import (
	"encoding/json"
	"errors"
	"log"
	"net/http"

	"github.com/oursky/ourd/oderr"
)

// Handler specifies the function signature of a request handler function
type Handler func(*Payload, *Response)

// pipeline encapsulates a transformation which a request will come throught
// from preprocessors to the actual handler. (and postprocessor later)
type pipeline struct {
	Action        string
	Preprocessors []Processor
	Handler
}

// Router to dispatch HTTP request to respective handler
type Router struct {
	actions map[string]pipeline
}

// Processor specifies the function signature for a Preprocessor
type Processor func(*Payload, *Response) int

// NewRouter is factory for Router
func NewRouter() *Router {
	return &Router{actions: make(map[string]pipeline)}
}

// Map to register action to handle mapping
func (r *Router) Map(action string, handler Handler, preprocessors ...Processor) {
	r.actions[action] = pipeline{
		Action:        action,
		Preprocessors: preprocessors,
		Handler:       handler,
	}
}

func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var (
		httpStatus = http.StatusOK
		reqJSON    map[string]interface{}
		resp       Response
	)
	defer func() {
		w.WriteHeader(httpStatus)
		if err := json.NewEncoder(w).Encode(resp); err != nil {
			panic(err)
		}
	}()

	if err := json.NewDecoder(req.Body).Decode(&reqJSON); err != nil {
		httpStatus = http.StatusBadRequest
		resp.Err = oderr.NewFmt(oderr.RequestInvalidErr, err.Error())
		return
	}

	payload := Payload{
		Meta: map[string]interface{}{},
		Data: reqJSON,
	}

	if pipeline, ok := r.actions[payload.RouteAction()]; ok {
		for _, p := range pipeline.Preprocessors {
			httpStatus = p(&payload, &resp)
			if resp.Err != nil {
				if httpStatus == 200 {
					httpStatus = 500
				}
				if _, ok := resp.Err.(oderr.Error); !ok {
					resp.Err = oderr.New(oderr.UnknownErr, resp.Err.Error())
				}
				return
			}
		}
		pipeline.Handler(&payload, &resp)
	} else {
		httpStatus = http.StatusNotFound
		resp.Err = oderr.New(oderr.RequestInvalidErr, "Unmatched Route")
	}
}

// CheckAuth will check on the AccessToken, attach DB/RequestID to the response
// This is a no-op if the request action belong to "auth:" group
func CheckAuth(payload *Payload, response *Response) (status int, err error) {
	log.Println("CheckAuth")

	token := payload.AccessToken()

	if token == "validToken" {
		log.Println("CheckAuth -> validToken, ", token)
		return http.StatusOK, nil
	}
	log.Println("CheckAuth -> inValidToken, ", token)
	return http.StatusUnauthorized, errors.New("Unauthorized request")
}
<commit_msg>Treat empty request as empty map in router, #23<commit_after>package router

import (
	"encoding/json"
	"errors"
	"io"
	"log"
	"net/http"

	"github.com/oursky/ourd/oderr"
)

// Handler specifies the function signature of a request handler function
type Handler func(*Payload, *Response)

// pipeline encapsulates a transformation which a request will come throught
// from preprocessors to the actual handler. (and postprocessor later)
type pipeline struct {
	Action        string
	Preprocessors []Processor
	Handler
}

// Router to dispatch HTTP request to respective handler
type Router struct {
	actions map[string]pipeline
}

// Processor specifies the function signature for a Preprocessor
type Processor func(*Payload, *Response) int

// NewRouter is factory for Router
func NewRouter() *Router {
	return &Router{actions: make(map[string]pipeline)}
}

// Map to register action to handle mapping
func (r *Router) Map(action string, handler Handler, preprocessors ...Processor) {
	r.actions[action] = pipeline{
		Action:        action,
		Preprocessors: preprocessors,
		Handler:       handler,
	}
}

func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var (
		httpStatus = http.StatusOK
		reqJSON    map[string]interface{}
		resp       Response
	)
	defer func() {
		w.WriteHeader(httpStatus)
		if err := json.NewEncoder(w).Encode(resp); err != nil {
			panic(err)
		}
	}()

	if err := json.NewDecoder(req.Body).Decode(&reqJSON); err != nil {
		if err == io.EOF {
			reqJSON = map[string]interface{}{}
		} else {
			httpStatus = http.StatusBadRequest
			resp.Err = oderr.NewFmt(oderr.RequestInvalidErr, err.Error())
			return
		}
	}

	payload := Payload{
		Meta: map[string]interface{}{},
		Data: reqJSON,
	}

	if pipeline, ok := r.actions[payload.RouteAction()]; ok {
		for _, p := range pipeline.Preprocessors {
			httpStatus = p(&payload, &resp)
			if resp.Err != nil {
				if httpStatus == 200 {
					httpStatus = 500
				}
				if _, ok := resp.Err.(oderr.Error); !ok {
					resp.Err = oderr.New(oderr.UnknownErr, resp.Err.Error())
				}
				return
			}
		}
		pipeline.Handler(&payload, &resp)
	} else {
		httpStatus = http.StatusNotFound
		resp.Err = oderr.New(oderr.RequestInvalidErr, "Unmatched Route")
	}
}

// CheckAuth will check on the AccessToken, attach DB/RequestID to the response
// This is a no-op if the request action belong to "auth:" group
func CheckAuth(payload *Payload, response *Response) (status int, err error) {
	log.Println("CheckAuth")

	token := payload.AccessToken()

	if token == "validToken" {
		log.Println("CheckAuth -> validToken, ", token)
		return http.StatusOK, nil
	}
	log.Println("CheckAuth -> inValidToken, ", token)
	return http.StatusUnauthorized, errors.New("Unauthorized request")
}
<|endoftext|>