<commit_before>package agent

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/hashicorp/memberlist"
	"github.com/hashicorp/serf/serf"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"sync"
)

// Agent starts and manages a Serf instance, adding some niceties
// on top of Serf such as storing logs that you can later retrieve,
// and invoking EventHandlers when events occur.
type Agent struct {
	// Stores the serf configuration
	conf *serf.Config

	// Stores the agent configuration
	agentConf *Config

	// eventCh is used for Serf to deliver events on
	eventCh chan serf.Event

	// eventHandlers is the registered handlers for events
	eventHandlers     map[EventHandler]struct{}
	eventHandlersLock sync.Mutex

	// logger instance wraps the logOutput
	logger *log.Logger

	// This is the underlying Serf we are wrapping
	serf *serf.Serf

	// shutdownCh is used for shutdowns
	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// Start creates a new agent, potentially returning an error
func Create(agentConf *Config, conf *serf.Config, logOutput io.Writer) (*Agent, error) {
	// Ensure we have a log sink
	if logOutput == nil {
		logOutput = os.Stderr
	}

	// Setup the underlying loggers
	conf.MemberlistConfig.LogOutput = logOutput
	conf.LogOutput = logOutput

	// Create a channel to listen for events from Serf
	eventCh := make(chan serf.Event, 64)
	conf.EventCh = eventCh

	// Setup the agent
	agent := &Agent{
		conf:          conf,
		agentConf:     agentConf,
		eventCh:       eventCh,
		eventHandlers: make(map[EventHandler]struct{}),
		logger:        log.New(logOutput, "", log.LstdFlags),
		shutdownCh:    make(chan struct{}),
	}

	// Restore agent tags from a tags file
	if agentConf.TagsFile != "" {
		if err := agent.loadTagsFile(agentConf.TagsFile); err != nil {
			return nil, err
		}
	}

	// Load in a keyring file if provided
	if agentConf.KeyringFile != "" {
		if err := agent.loadKeyringFile(agentConf.KeyringFile); err != nil {
			return nil, err
		}
	}

	return agent, nil
}

// Start is used to initiate the event listeners. It is seperate from
// create so that there isn't a race condition between creating the
// agent and registering handlers
func (a *Agent) Start() error {
	a.logger.Printf("[INFO] agent: Serf agent starting")

	// Create serf first
	serf, err := serf.Create(a.conf)
	if err != nil {
		return fmt.Errorf("Error creating Serf: %s", err)
	}
	a.serf = serf

	// Start event loop
	go a.eventLoop()
	return nil
}

// Leave prepares for a graceful shutdown of the agent and its processes
func (a *Agent) Leave() error {
	if a.serf == nil {
		return nil
	}

	a.logger.Println("[INFO] agent: requesting graceful leave from Serf")
	return a.serf.Leave()
}

// Shutdown closes this agent and all of its processes. Should be preceeded
// by a Leave for a graceful shutdown.
func (a *Agent) Shutdown() error {
	a.shutdownLock.Lock()
	defer a.shutdownLock.Unlock()

	if a.shutdown {
		return nil
	}

	if a.serf == nil {
		goto EXIT
	}

	a.logger.Println("[INFO] agent: requesting serf shutdown")
	if err := a.serf.Shutdown(); err != nil {
		return err
	}

EXIT:
	a.logger.Println("[INFO] agent: shutdown complete")
	a.shutdown = true
	close(a.shutdownCh)
	return nil
}

// ShutdownCh returns a channel that can be selected to wait
// for the agent to perform a shutdown.
func (a *Agent) ShutdownCh() <-chan struct{} {
	return a.shutdownCh
}

// Returns the Serf agent of the running Agent.
func (a *Agent) Serf() *serf.Serf {
	return a.serf
}

// Returns the Serf config of the running Agent.
func (a *Agent) SerfConfig() *serf.Config {
	return a.conf
}

// Join asks the Serf instance to join. See the Serf.Join function.
func (a *Agent) Join(addrs []string, replay bool) (n int, err error) {
	a.logger.Printf("[INFO] agent: joining: %v replay: %v", addrs, replay)
	ignoreOld := !replay
	n, err = a.serf.Join(addrs, ignoreOld)
	if n > 0 {
		a.logger.Printf("[INFO] agent: joined: %d nodes", n)
	}
	if err != nil {
		a.logger.Printf("[WARN] agent: error joining: %v", err)
	}
	return
}

// ForceLeave is used to eject a failed node from the cluster
func (a *Agent) ForceLeave(node string) error {
	a.logger.Printf("[INFO] agent: Force leaving node: %s", node)
	err := a.serf.RemoveFailedNode(node)
	if err != nil {
		a.logger.Printf("[WARN] agent: failed to remove node: %v", err)
	}
	return err
}

// UserEvent sends a UserEvent on Serf, see Serf.UserEvent.
func (a *Agent) UserEvent(name string, payload []byte, coalesce bool) error {
	a.logger.Printf("[DEBUG] agent: Requesting user event send: %s. Coalesced: %#v. Payload: %#v",
		name, coalesce, string(payload))
	err := a.serf.UserEvent(name, payload, coalesce)
	if err != nil {
		a.logger.Printf("[WARN] agent: failed to send user event: %v", err)
	}
	return err
}

// Query sends a Query on Serf, see Serf.Query.
func (a *Agent) Query(name string, payload []byte, params *serf.QueryParam) (*serf.QueryResponse, error) {
	// Prevent the use of the internal prefix
	if strings.HasPrefix(name, serf.InternalQueryPrefix) {
		// Allow the special "ping" query
		if name != serf.InternalQueryPrefix+"ping" || payload != nil {
			return nil, fmt.Errorf("Queries cannot contain the '%s' prefix", serf.InternalQueryPrefix)
		}
	}
	a.logger.Printf("[DEBUG] agent: Requesting query send: %s. Payload: %#v",
		name, string(payload))
	resp, err := a.serf.Query(name, payload, params)
	if err != nil {
		a.logger.Printf("[WARN] agent: failed to start user query: %v", err)
	}
	return resp, err
}

// RegisterEventHandler adds an event handler to recieve event notifications
func (a *Agent) RegisterEventHandler(eh EventHandler) {
	a.eventHandlersLock.Lock()
	defer a.eventHandlersLock.Unlock()
	a.eventHandlers[eh] = struct{}{}
}

// DeregisterEventHandler removes an EventHandler and prevents more invocations
func (a *Agent) DeregisterEventHandler(eh EventHandler) {
	a.eventHandlersLock.Lock()
	defer a.eventHandlersLock.Unlock()
	delete(a.eventHandlers, eh)
}

// eventLoop listens to events from Serf and fans out to event handlers
func (a *Agent) eventLoop() {
	serfShutdownCh := a.serf.ShutdownCh()
	for {
		select {
		case e := <-a.eventCh:
			a.logger.Printf("[INFO] agent: Received event: %s", e.String())
			a.eventHandlersLock.Lock()
			for eh, _ := range a.eventHandlers {
				eh.HandleEvent(e)
			}
			a.eventHandlersLock.Unlock()

		case <-serfShutdownCh:
			a.logger.Printf("[WARN] agent: Serf shutdown detected, quitting")
			a.Shutdown()
			return

		case <-a.shutdownCh:
			return
		}
	}
}

// InstallKey initiates a query to install a new key on all members
func (a *Agent) InstallKey(key string) (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating key installation")
	manager := a.serf.KeyManager()
	return manager.InstallKey(key)
}

// UseKey sends a query instructing all members to switch primary keys
func (a *Agent) UseKey(key string) (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating primary key change")
	manager := a.serf.KeyManager()
	return manager.UseKey(key)
}

// RemoveKey sends a query to all members to remove a key from the keyring
func (a *Agent) RemoveKey(key string) (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating key removal")
	manager := a.serf.KeyManager()
	return manager.RemoveKey(key)
}

// ListKeys sends a query to all members to return a list of their keys
func (a *Agent) ListKeys() (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating key listing")
	manager := a.serf.KeyManager()
	return manager.ListKeys()
}

// SetTags is used to update the tags. The agent will make sure to
// persist tags if necessary before gossiping to the cluster.
func (a *Agent) SetTags(tags map[string]string) error {
	// Update the tags file if we have one
	if a.agentConf.TagsFile != "" {
		if err := a.writeTagsFile(tags); err != nil {
			a.logger.Printf("[ERR] agent: %s", err)
			return err
		}
	}

	// Set the tags in Serf, start gossiping out
	return a.serf.SetTags(tags)
}

// loadTagsFile will load agent tags out of a file and set them in the
// current serf configuration.
func (a *Agent) loadTagsFile(tagsFile string) error {
	// Avoid passing tags and using a tags file at the same time
	if len(a.agentConf.Tags) > 0 {
		return fmt.Errorf("Tags config not allowed while using tag files")
	}

	if _, err := os.Stat(tagsFile); err == nil {
		tagData, err := ioutil.ReadFile(tagsFile)
		if err != nil {
			return fmt.Errorf("Failed to read tags file: %s", err)
		}
		if err := json.Unmarshal(tagData, &a.conf.Tags); err != nil {
			return fmt.Errorf("Failed to decode tags file: %s", err)
		}
		a.logger.Printf("[INFO] agent: Restored %d tag(s) from %s",
			len(a.conf.Tags), tagsFile)
	}

	// Success!
	return nil
}

// writeTagsFile will write the current tags to the configured tags file.
func (a *Agent) writeTagsFile(tags map[string]string) error {
	encoded, err := json.MarshalIndent(tags, "", "  ")
	if err != nil {
		return fmt.Errorf("Failed to encode tags: %s", err)
	}

	// Use 0600 for permissions, in case tag data is sensitive
	if err = ioutil.WriteFile(a.agentConf.TagsFile, encoded, 0600); err != nil {
		return fmt.Errorf("Failed to write tags file: %s", err)
	}

	// Success!
	return nil
}

// MarshalTags is a utility function which takes a map of tag key/value pairs
// and returns the same tags as strings in 'key=value' format.
func MarshalTags(tags map[string]string) []string {
	var result []string
	for name, value := range tags {
		result = append(result, fmt.Sprintf("%s=%s", name, value))
	}
	return result
}

// UnmarshalTags is a utility function which takes a slice of strings in
// key=value format and returns them as a tag mapping.
func UnmarshalTags(tags []string) (map[string]string, error) {
	result := make(map[string]string)
	for _, tag := range tags {
		parts := strings.SplitN(tag, "=", 2)
		if len(parts) != 2 || len(parts[0]) == 0 {
			return nil, fmt.Errorf("Invalid tag: '%s'", tag)
		}
		result[parts[0]] = parts[1]
	}
	return result, nil
}

// loadKeyringFile will load a keyring out of a file
func (a *Agent) loadKeyringFile(keyringFile string) error {
	// Avoid passing an encryption key and a keyring file at the same time
	if len(a.agentConf.EncryptKey) > 0 {
		return fmt.Errorf("Encryption key not allowed while using a keyring")
	}

	if _, err := os.Stat(keyringFile); err != nil {
		return err
	}

	// Read in the keyring file data
	keyringData, err := ioutil.ReadFile(keyringFile)
	if err != nil {
		return fmt.Errorf("Failed to read keyring file: %s", err)
	}

	// Decode keyring JSON
	keys := make([]string, 0)
	if err := json.Unmarshal(keyringData, &keys); err != nil {
		return fmt.Errorf("Failed to decode keyring file: %s", err)
	}

	// Decode base64 values
	keysDecoded := make([][]byte, len(keys))
	for i, key := range keys {
		keyBytes, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			return fmt.Errorf("Failed to decode key from keyring: %s", err)
		}
		keysDecoded[i] = keyBytes
	}

	// Create the keyring
	keyring, err := memberlist.NewKeyring(keysDecoded, keysDecoded[0])
	if err != nil {
		return fmt.Errorf("Failed to restore keyring: %s", err)
	}
	a.conf.MemberlistConfig.Keyring = keyring
	a.logger.Printf("[INFO] agent: Restored keyring with %d keys from %s",
		len(keys), keyringFile)

	// Success!
	return nil
}
<commit_msg>agent: Adding Stats method<commit_after>package agent

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/hashicorp/memberlist"
	"github.com/hashicorp/serf/serf"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"sync"
)

// Agent starts and manages a Serf instance, adding some niceties
// on top of Serf such as storing logs that you can later retrieve,
// and invoking EventHandlers when events occur.
type Agent struct {
	// Stores the serf configuration
	conf *serf.Config

	// Stores the agent configuration
	agentConf *Config

	// eventCh is used for Serf to deliver events on
	eventCh chan serf.Event

	// eventHandlers is the registered handlers for events
	eventHandlers     map[EventHandler]struct{}
	eventHandlersLock sync.Mutex

	// logger instance wraps the logOutput
	logger *log.Logger

	// This is the underlying Serf we are wrapping
	serf *serf.Serf

	// shutdownCh is used for shutdowns
	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// Start creates a new agent, potentially returning an error
func Create(agentConf *Config, conf *serf.Config, logOutput io.Writer) (*Agent, error) {
	// Ensure we have a log sink
	if logOutput == nil {
		logOutput = os.Stderr
	}

	// Setup the underlying loggers
	conf.MemberlistConfig.LogOutput = logOutput
	conf.LogOutput = logOutput

	// Create a channel to listen for events from Serf
	eventCh := make(chan serf.Event, 64)
	conf.EventCh = eventCh

	// Setup the agent
	agent := &Agent{
		conf:          conf,
		agentConf:     agentConf,
		eventCh:       eventCh,
		eventHandlers: make(map[EventHandler]struct{}),
		logger:        log.New(logOutput, "", log.LstdFlags),
		shutdownCh:    make(chan struct{}),
	}

	// Restore agent tags from a tags file
	if agentConf.TagsFile != "" {
		if err := agent.loadTagsFile(agentConf.TagsFile); err != nil {
			return nil, err
		}
	}

	// Load in a keyring file if provided
	if agentConf.KeyringFile != "" {
		if err := agent.loadKeyringFile(agentConf.KeyringFile); err != nil {
			return nil, err
		}
	}

	return agent, nil
}

// Start is used to initiate the event listeners. It is seperate from
// create so that there isn't a race condition between creating the
// agent and registering handlers
func (a *Agent) Start() error {
	a.logger.Printf("[INFO] agent: Serf agent starting")

	// Create serf first
	serf, err := serf.Create(a.conf)
	if err != nil {
		return fmt.Errorf("Error creating Serf: %s", err)
	}
	a.serf = serf

	// Start event loop
	go a.eventLoop()
	return nil
}

// Leave prepares for a graceful shutdown of the agent and its processes
func (a *Agent) Leave() error {
	if a.serf == nil {
		return nil
	}

	a.logger.Println("[INFO] agent: requesting graceful leave from Serf")
	return a.serf.Leave()
}

// Shutdown closes this agent and all of its processes. Should be preceeded
// by a Leave for a graceful shutdown.
func (a *Agent) Shutdown() error {
	a.shutdownLock.Lock()
	defer a.shutdownLock.Unlock()

	if a.shutdown {
		return nil
	}

	if a.serf == nil {
		goto EXIT
	}

	a.logger.Println("[INFO] agent: requesting serf shutdown")
	if err := a.serf.Shutdown(); err != nil {
		return err
	}

EXIT:
	a.logger.Println("[INFO] agent: shutdown complete")
	a.shutdown = true
	close(a.shutdownCh)
	return nil
}

// ShutdownCh returns a channel that can be selected to wait
// for the agent to perform a shutdown.
func (a *Agent) ShutdownCh() <-chan struct{} {
	return a.shutdownCh
}

// Returns the Serf agent of the running Agent.
func (a *Agent) Serf() *serf.Serf {
	return a.serf
}

// Returns the Serf config of the running Agent.
func (a *Agent) SerfConfig() *serf.Config {
	return a.conf
}

// Join asks the Serf instance to join. See the Serf.Join function.
func (a *Agent) Join(addrs []string, replay bool) (n int, err error) {
	a.logger.Printf("[INFO] agent: joining: %v replay: %v", addrs, replay)
	ignoreOld := !replay
	n, err = a.serf.Join(addrs, ignoreOld)
	if n > 0 {
		a.logger.Printf("[INFO] agent: joined: %d nodes", n)
	}
	if err != nil {
		a.logger.Printf("[WARN] agent: error joining: %v", err)
	}
	return
}

// ForceLeave is used to eject a failed node from the cluster
func (a *Agent) ForceLeave(node string) error {
	a.logger.Printf("[INFO] agent: Force leaving node: %s", node)
	err := a.serf.RemoveFailedNode(node)
	if err != nil {
		a.logger.Printf("[WARN] agent: failed to remove node: %v", err)
	}
	return err
}

// UserEvent sends a UserEvent on Serf, see Serf.UserEvent.
func (a *Agent) UserEvent(name string, payload []byte, coalesce bool) error {
	a.logger.Printf("[DEBUG] agent: Requesting user event send: %s. Coalesced: %#v. Payload: %#v",
		name, coalesce, string(payload))
	err := a.serf.UserEvent(name, payload, coalesce)
	if err != nil {
		a.logger.Printf("[WARN] agent: failed to send user event: %v", err)
	}
	return err
}

// Query sends a Query on Serf, see Serf.Query.
func (a *Agent) Query(name string, payload []byte, params *serf.QueryParam) (*serf.QueryResponse, error) {
	// Prevent the use of the internal prefix
	if strings.HasPrefix(name, serf.InternalQueryPrefix) {
		// Allow the special "ping" query
		if name != serf.InternalQueryPrefix+"ping" || payload != nil {
			return nil, fmt.Errorf("Queries cannot contain the '%s' prefix", serf.InternalQueryPrefix)
		}
	}
	a.logger.Printf("[DEBUG] agent: Requesting query send: %s. Payload: %#v",
		name, string(payload))
	resp, err := a.serf.Query(name, payload, params)
	if err != nil {
		a.logger.Printf("[WARN] agent: failed to start user query: %v", err)
	}
	return resp, err
}

// RegisterEventHandler adds an event handler to recieve event notifications
func (a *Agent) RegisterEventHandler(eh EventHandler) {
	a.eventHandlersLock.Lock()
	defer a.eventHandlersLock.Unlock()
	a.eventHandlers[eh] = struct{}{}
}

// DeregisterEventHandler removes an EventHandler and prevents more invocations
func (a *Agent) DeregisterEventHandler(eh EventHandler) {
	a.eventHandlersLock.Lock()
	defer a.eventHandlersLock.Unlock()
	delete(a.eventHandlers, eh)
}

// eventLoop listens to events from Serf and fans out to event handlers
func (a *Agent) eventLoop() {
	serfShutdownCh := a.serf.ShutdownCh()
	for {
		select {
		case e := <-a.eventCh:
			a.logger.Printf("[INFO] agent: Received event: %s", e.String())
			a.eventHandlersLock.Lock()
			for eh, _ := range a.eventHandlers {
				eh.HandleEvent(e)
			}
			a.eventHandlersLock.Unlock()

		case <-serfShutdownCh:
			a.logger.Printf("[WARN] agent: Serf shutdown detected, quitting")
			a.Shutdown()
			return

		case <-a.shutdownCh:
			return
		}
	}
}

// InstallKey initiates a query to install a new key on all members
func (a *Agent) InstallKey(key string) (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating key installation")
	manager := a.serf.KeyManager()
	return manager.InstallKey(key)
}

// UseKey sends a query instructing all members to switch primary keys
func (a *Agent) UseKey(key string) (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating primary key change")
	manager := a.serf.KeyManager()
	return manager.UseKey(key)
}

// RemoveKey sends a query to all members to remove a key from the keyring
func (a *Agent) RemoveKey(key string) (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating key removal")
	manager := a.serf.KeyManager()
	return manager.RemoveKey(key)
}

// ListKeys sends a query to all members to return a list of their keys
func (a *Agent) ListKeys() (*serf.KeyResponse, error) {
	a.logger.Print("[INFO] agent: Initiating key listing")
	manager := a.serf.KeyManager()
	return manager.ListKeys()
}

// SetTags is used to update the tags. The agent will make sure to
// persist tags if necessary before gossiping to the cluster.
func (a *Agent) SetTags(tags map[string]string) error {
	// Update the tags file if we have one
	if a.agentConf.TagsFile != "" {
		if err := a.writeTagsFile(tags); err != nil {
			a.logger.Printf("[ERR] agent: %s", err)
			return err
		}
	}

	// Set the tags in Serf, start gossiping out
	return a.serf.SetTags(tags)
}

// loadTagsFile will load agent tags out of a file and set them in the
// current serf configuration.
func (a *Agent) loadTagsFile(tagsFile string) error {
	// Avoid passing tags and using a tags file at the same time
	if len(a.agentConf.Tags) > 0 {
		return fmt.Errorf("Tags config not allowed while using tag files")
	}

	if _, err := os.Stat(tagsFile); err == nil {
		tagData, err := ioutil.ReadFile(tagsFile)
		if err != nil {
			return fmt.Errorf("Failed to read tags file: %s", err)
		}
		if err := json.Unmarshal(tagData, &a.conf.Tags); err != nil {
			return fmt.Errorf("Failed to decode tags file: %s", err)
		}
		a.logger.Printf("[INFO] agent: Restored %d tag(s) from %s",
			len(a.conf.Tags), tagsFile)
	}

	// Success!
	return nil
}

// writeTagsFile will write the current tags to the configured tags file.
func (a *Agent) writeTagsFile(tags map[string]string) error {
	encoded, err := json.MarshalIndent(tags, "", "  ")
	if err != nil {
		return fmt.Errorf("Failed to encode tags: %s", err)
	}

	// Use 0600 for permissions, in case tag data is sensitive
	if err = ioutil.WriteFile(a.agentConf.TagsFile, encoded, 0600); err != nil {
		return fmt.Errorf("Failed to write tags file: %s", err)
	}

	// Success!
	return nil
}

// MarshalTags is a utility function which takes a map of tag key/value pairs
// and returns the same tags as strings in 'key=value' format.
func MarshalTags(tags map[string]string) []string {
	var result []string
	for name, value := range tags {
		result = append(result, fmt.Sprintf("%s=%s", name, value))
	}
	return result
}

// UnmarshalTags is a utility function which takes a slice of strings in
// key=value format and returns them as a tag mapping.
func UnmarshalTags(tags []string) (map[string]string, error) {
	result := make(map[string]string)
	for _, tag := range tags {
		parts := strings.SplitN(tag, "=", 2)
		if len(parts) != 2 || len(parts[0]) == 0 {
			return nil, fmt.Errorf("Invalid tag: '%s'", tag)
		}
		result[parts[0]] = parts[1]
	}
	return result, nil
}

// loadKeyringFile will load a keyring out of a file
func (a *Agent) loadKeyringFile(keyringFile string) error {
	// Avoid passing an encryption key and a keyring file at the same time
	if len(a.agentConf.EncryptKey) > 0 {
		return fmt.Errorf("Encryption key not allowed while using a keyring")
	}

	if _, err := os.Stat(keyringFile); err != nil {
		return err
	}

	// Read in the keyring file data
	keyringData, err := ioutil.ReadFile(keyringFile)
	if err != nil {
		return fmt.Errorf("Failed to read keyring file: %s", err)
	}

	// Decode keyring JSON
	keys := make([]string, 0)
	if err := json.Unmarshal(keyringData, &keys); err != nil {
		return fmt.Errorf("Failed to decode keyring file: %s", err)
	}

	// Decode base64 values
	keysDecoded := make([][]byte, len(keys))
	for i, key := range keys {
		keyBytes, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			return fmt.Errorf("Failed to decode key from keyring: %s", err)
		}
		keysDecoded[i] = keyBytes
	}

	// Create the keyring
	keyring, err := memberlist.NewKeyring(keysDecoded, keysDecoded[0])
	if err != nil {
		return fmt.Errorf("Failed to restore keyring: %s", err)
	}
	a.conf.MemberlistConfig.Keyring = keyring
	a.logger.Printf("[INFO] agent: Restored keyring with %d keys from %s",
		len(keys), keyringFile)

	// Success!
	return nil
}

// Stats is used to get various runtime information and stats
func (a *Agent) Stats() map[string]map[string]string {
	local := a.serf.LocalMember()
	output := map[string]map[string]string{
		"agent": map[string]string{
			"name": local.Name,
		},
		"runtime": runtimeStats(),
		"serf":    a.serf.Stats(),
		"tags":    local.Tags,
	}
	return output
}
<|endoftext|>


<commit_before>package command

import (
	"os"
	"fmt"
	"errors"

	"github.com/coreos/etcdctl/third_party/github.com/codegangsta/cli"
	"github.com/coreos/etcdctl/third_party/github.com/coreos/go-etcd/etcd"
)

// NewGetCommand returns the CLI command for "get".
func NewGetCommand() cli.Command {
	return cli.Command{
		Name:	"get",
		Usage:	"retrieve the value of a key",
		Flags: []cli.Flag{
			cli.BoolFlag{"sort", "returns result in sorted order"},
			cli.BoolFlag{"consistent", "send request to the leader, thereby guranteeing that any earlier writes will be seen by the read"},
		},
		Action: func(c *cli.Context) {
			handleGet(c, getCommandFunc)
		},
	}
}

// handleGet handles a request that intends to do get-like operations.
func handleGet(c *cli.Context, fn handlerFunc) {
	handlePrint(c, fn, printGet)
}

// printGet writes error message when getting the value of a directory.
func printGet(resp *etcd.Response, format string) {
	if resp.Node.Dir {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("Cannot get the value [%s: Is a directory]\nPlease use ls to list a directory", resp.Node.Key))
		os.Exit(1)
	}

	printKey(resp, format)
}

// getCommandFunc executes the "get" command.
func getCommandFunc(c *cli.Context, client *etcd.Client) (*etcd.Response, error) {
	if len(c.Args()) == 0 {
		return nil, errors.New("Key required")
	}
	key := c.Args()[0]
	consistent := c.Bool("consistent")
	sorted := c.Bool("sort")

	// Setup consistency on the client.
	if consistent {
		client.SetConsistency(etcd.STRONG_CONSISTENCY)
	} else {
		client.SetConsistency(etcd.WEAK_CONSISTENCY)
	}

	// Retrieve the value from the server.
	return client.Get(key, sorted, false)
}
<commit_msg>fix(command): use unix style error messages<commit_after>package command

import (
	"os"
	"fmt"
	"errors"

	"github.com/coreos/etcdctl/third_party/github.com/codegangsta/cli"
	"github.com/coreos/etcdctl/third_party/github.com/coreos/go-etcd/etcd"
)

// NewGetCommand returns the CLI command for "get".
func NewGetCommand() cli.Command {
	return cli.Command{
		Name:	"get",
		Usage:	"retrieve the value of a key",
		Flags: []cli.Flag{
			cli.BoolFlag{"sort", "returns result in sorted order"},
			cli.BoolFlag{"consistent", "send request to the leader, thereby guranteeing that any earlier writes will be seen by the read"},
		},
		Action: func(c *cli.Context) {
			handleGet(c, getCommandFunc)
		},
	}
}

// handleGet handles a request that intends to do get-like operations.
func handleGet(c *cli.Context, fn handlerFunc) {
	handlePrint(c, fn, printGet)
}

// printGet writes error message when getting the value of a directory.
func printGet(resp *etcd.Response, format string) {
	if resp.Node.Dir {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("%s: is a directory", resp.Node.Key))
		os.Exit(1)
	}

	printKey(resp, format)
}

// getCommandFunc executes the "get" command.
func getCommandFunc(c *cli.Context, client *etcd.Client) (*etcd.Response, error) {
	if len(c.Args()) == 0 {
		return nil, errors.New("Key required")
	}
	key := c.Args()[0]
	consistent := c.Bool("consistent")
	sorted := c.Bool("sort")

	// Setup consistency on the client.
	if consistent {
		client.SetConsistency(etcd.STRONG_CONSISTENCY)
	} else {
		client.SetConsistency(etcd.WEAK_CONSISTENCY)
	}

	// Retrieve the value from the server.
	return client.Get(key, sorted, false)
}
<|endoftext|>


<commit_before>/*
 *  Copyright 2016 Adobe Systems Incorporated. All rights reserved.
 *  This file is licensed to you 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 REPRESENTATIONS
 *  OF ANY KIND, either express or implied. See the License for the specific language
 *  governing permissions and limitations under the License.
 */
package build

import (
	"os"

	"github.com/adobe-platform/porter/conf"
	"github.com/adobe-platform/porter/constants"
	"github.com/adobe-platform/porter/hook"
	"github.com/adobe-platform/porter/logger"
	"github.com/adobe-platform/porter/provision"

	"github.com/phylake/go-cli"
)

type PackCmd struct{}

func (recv *PackCmd) Name() string {
	return "pack"
}

func (recv *PackCmd) ShortHelp() string {
	return "Create a service payload"
}

func (recv *PackCmd) LongHelp() string {
	return ""
}

func (recv *PackCmd) SubCommands() []cli.Command {
	return nil
}

func (recv *PackCmd) Execute(args []string) bool {

	log := logger.CLI("cmd", "pack")

	commandSuccess := hook.Execute(log, constants.HookPrePack, "", nil, true)

	if commandSuccess {
		config, success := conf.GetConfig(log, true)
		if !success {
			os.Exit(1)
		}

		if os.Getenv(constants.EnvConfig) != "" {
			config.Print()
		}

		commandSuccess = provision.Package(log, config)
	}

	commandSuccess = hook.Execute(log, constants.HookPostPack, "", nil, commandSuccess)

	if !commandSuccess {
		os.Exit(1)
	}

	log.Info("Packaged service", "FilePath", constants.PayloadPath)

	return true
}
<commit_msg>update pack hook exec to be consistent with other commands<commit_after>/*
 *  Copyright 2016 Adobe Systems Incorporated. All rights reserved.
 *  This file is licensed to you 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 REPRESENTATIONS
 *  OF ANY KIND, either express or implied. See the License for the specific language
 *  governing permissions and limitations under the License.
 */
package build

import (
	"os"

	"github.com/adobe-platform/porter/conf"
	"github.com/adobe-platform/porter/constants"
	"github.com/adobe-platform/porter/hook"
	"github.com/adobe-platform/porter/logger"
	"github.com/adobe-platform/porter/provision"

	"github.com/phylake/go-cli"
)

type PackCmd struct{}

func (recv *PackCmd) Name() string {
	return "pack"
}

func (recv *PackCmd) ShortHelp() string {
	return "Create a service payload"
}

func (recv *PackCmd) LongHelp() string {
	return ""
}

func (recv *PackCmd) SubCommands() []cli.Command {
	return nil
}

func (recv *PackCmd) Execute(args []string) bool {

	if !doPack() {
		os.Exit(1)
	}

	return true
}

func doPack() (success bool) {
	log := logger.CLI("cmd", "pack")

	defer func() {

		postHookSuccess := hook.Execute(log, constants.HookPostPack, "", nil, success)

		success = success && postHookSuccess
	}()

	if !hook.Execute(log, constants.HookPrePack, "", nil, true) {
		return
	}

	config, getConfigSuccess := conf.GetConfig(log, true)
	if !getConfigSuccess {
		return
	}

	if os.Getenv(constants.EnvConfig) != "" {
		config.Print()
	}

	success = provision.Package(log, config)

	if success {
		log.Info("Packaged service", "FilePath", constants.PayloadPath)
	}

	success = true
	return
}
<|endoftext|>


<commit_before>package agent

import (
	"bytes"
	"github.com/hashicorp/consul/consul/structs"
	"io"
	"net/http"
	"strconv"
	"strings"
)

func (s *HTTPServer) KVSEndpoint(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	// Set default DC
	args := structs.KeyRequest{}
	if done := s.parse(resp, req, &args.Datacenter, &args.BlockingQuery); done {
		return nil, nil
	}

	// Pull out the key name, validation left to each sub-handler
	args.Key = strings.TrimPrefix(req.URL.Path, "/v1/kv/")

	// Switch on the method
	switch req.Method {
	case "GET":
		return s.KVSGet(resp, req, &args)
	case "PUT":
		return s.KVSPut(resp, req, &args)
	case "DELETE":
		return s.KVSDelete(resp, req, &args)
	default:
		resp.WriteHeader(405)
		return nil, nil
	}
	return nil, nil
}

// KVSGet handles a GET request
func (s *HTTPServer) KVSGet(resp http.ResponseWriter, req *http.Request, args *structs.KeyRequest) (interface{}, error) {
	// Check for recurse
	method := "KVS.Get"
	params := req.URL.Query()
	if _, ok := params["recurse"]; ok {
		method = "KVS.List"
	} else if missingKey(resp, args) {
		return nil, nil
	}

	// Make the RPC
	var out structs.IndexedDirEntries
	if err := s.agent.RPC(method, &args, &out); err != nil {
		return nil, err
	}
	setIndex(resp, out.Index)

	// Check if we get a not found
	if len(out.Entries) == 0 {
		resp.WriteHeader(404)
		return nil, nil
	}
	return out.Entries, nil
}

// KVSPut handles a PUT request
func (s *HTTPServer) KVSPut(resp http.ResponseWriter, req *http.Request, args *structs.KeyRequest) (interface{}, error) {
	if missingKey(resp, args) {
		return nil, nil
	}
	applyReq := structs.KVSRequest{
		Datacenter: args.Datacenter,
		Op:         structs.KVSSet,
		DirEnt: structs.DirEntry{
			Key:   args.Key,
			Flags: 0,
			Value: nil,
		},
	}

	// Check for flags
	params := req.URL.Query()
	if _, ok := params["flags"]; ok {
		flagVal, err := strconv.ParseUint(params.Get("flags"), 10, 64)
		if err != nil {
			return nil, err
		}
		applyReq.DirEnt.Flags = flagVal
	}

	// Check for cas value
	if _, ok := params["cas"]; ok {
		casVal, err := strconv.ParseUint(params.Get("flags"), 10, 64)
		if err != nil {
			return nil, err
		}
		applyReq.DirEnt.ModifyIndex = casVal
		applyReq.Op = structs.KVSCAS
	}

	// Copy the value
	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, req.Body); err != nil {
		return nil, err
	}
	applyReq.DirEnt.Value = buf.Bytes()

	// Make the RPC
	var out bool
	if err := s.agent.RPC("KVS.Apply", &applyReq, &out); err != nil {
		return nil, err
	}

	// Only use the out value if this was a CAS
	if applyReq.Op == structs.KVSSet {
		return true, nil
	} else {
		return out, nil
	}
}

// KVSPut handles a DELETE request
func (s *HTTPServer) KVSDelete(resp http.ResponseWriter, req *http.Request, args *structs.KeyRequest) (interface{}, error) {
	applyReq := structs.KVSRequest{
		Datacenter: args.Datacenter,
		Op:         structs.KVSDelete,
		DirEnt: structs.DirEntry{
			Key: args.Key,
		},
	}

	// Check for recurse
	params := req.URL.Query()
	if _, ok := params["recurse"]; ok {
		applyReq.Op = structs.KVSDeleteTree
	} else if missingKey(resp, args) {
		return nil, nil
	}

	// Make the RPC
	var out bool
	if err := s.agent.RPC("KVS.Apply", &applyReq, &out); err != nil {
		return nil, err
	}
	return nil, nil
}

// missingKey checks if the key is missing
func missingKey(resp http.ResponseWriter, args *structs.KeyRequest) bool {
	if args.Key == "" {
		resp.WriteHeader(400)
		resp.Write([]byte("Missing key name"))
		return true
	}
	return false
}
<commit_msg>agnet: Fix parsing of cas flag<commit_after>package agent

import (
	"bytes"
	"github.com/hashicorp/consul/consul/structs"
	"io"
	"net/http"
	"strconv"
	"strings"
)

func (s *HTTPServer) KVSEndpoint(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
	// Set default DC
	args := structs.KeyRequest{}
	if done := s.parse(resp, req, &args.Datacenter, &args.BlockingQuery); done {
		return nil, nil
	}

	// Pull out the key name, validation left to each sub-handler
	args.Key = strings.TrimPrefix(req.URL.Path, "/v1/kv/")

	// Switch on the method
	switch req.Method {
	case "GET":
		return s.KVSGet(resp, req, &args)
	case "PUT":
		return s.KVSPut(resp, req, &args)
	case "DELETE":
		return s.KVSDelete(resp, req, &args)
	default:
		resp.WriteHeader(405)
		return nil, nil
	}
	return nil, nil
}

// KVSGet handles a GET request
func (s *HTTPServer) KVSGet(resp http.ResponseWriter, req *http.Request, args *structs.KeyRequest) (interface{}, error) {
	// Check for recurse
	method := "KVS.Get"
	params := req.URL.Query()
	if _, ok := params["recurse"]; ok {
		method = "KVS.List"
	} else if missingKey(resp, args) {
		return nil, nil
	}

	// Make the RPC
	var out structs.IndexedDirEntries
	if err := s.agent.RPC(method, &args, &out); err != nil {
		return nil, err
	}
	setIndex(resp, out.Index)

	// Check if we get a not found
	if len(out.Entries) == 0 {
		resp.WriteHeader(404)
		return nil, nil
	}
	return out.Entries, nil
}

// KVSPut handles a PUT request
func (s *HTTPServer) KVSPut(resp http.ResponseWriter, req *http.Request, args *structs.KeyRequest) (interface{}, error) {
	if missingKey(resp, args) {
		return nil, nil
	}
	applyReq := structs.KVSRequest{
		Datacenter: args.Datacenter,
		Op:         structs.KVSSet,
		DirEnt: structs.DirEntry{
			Key:   args.Key,
			Flags: 0,
			Value: nil,
		},
	}

	// Check for flags
	params := req.URL.Query()
	if _, ok := params["flags"]; ok {
		flagVal, err := strconv.ParseUint(params.Get("flags"), 10, 64)
		if err != nil {
			return nil, err
		}
		applyReq.DirEnt.Flags = flagVal
	}

	// Check for cas value
	if _, ok := params["cas"]; ok {
		casVal, err := strconv.ParseUint(params.Get("cas"), 10, 64)
		if err != nil {
			return nil, err
		}
		applyReq.DirEnt.ModifyIndex = casVal
		applyReq.Op = structs.KVSCAS
	}

	// Copy the value
	buf := bytes.NewBuffer(nil)
	if _, err := io.Copy(buf, req.Body); err != nil {
		return nil, err
	}
	applyReq.DirEnt.Value = buf.Bytes()

	// Make the RPC
	var out bool
	if err := s.agent.RPC("KVS.Apply", &applyReq, &out); err != nil {
		return nil, err
	}

	// Only use the out value if this was a CAS
	if applyReq.Op == structs.KVSSet {
		return true, nil
	} else {
		return out, nil
	}
}

// KVSPut handles a DELETE request
func (s *HTTPServer) KVSDelete(resp http.ResponseWriter, req *http.Request, args *structs.KeyRequest) (interface{}, error) {
	applyReq := structs.KVSRequest{
		Datacenter: args.Datacenter,
		Op:         structs.KVSDelete,
		DirEnt: structs.DirEntry{
			Key: args.Key,
		},
	}

	// Check for recurse
	params := req.URL.Query()
	if _, ok := params["recurse"]; ok {
		applyReq.Op = structs.KVSDeleteTree
	} else if missingKey(resp, args) {
		return nil, nil
	}

	// Make the RPC
	var out bool
	if err := s.agent.RPC("KVS.Apply", &applyReq, &out); err != nil {
		return nil, err
	}
	return nil, nil
}

// missingKey checks if the key is missing
func missingKey(resp http.ResponseWriter, args *structs.KeyRequest) bool {
	if args.Key == "" {
		resp.WriteHeader(400)
		resp.Write([]byte("Missing key name"))
		return true
	}
	return false
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package provisioner

import (
	"launchpad.net/golxc"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/container/lxc"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	"launchpad.net/loggo"
)

var lxcLogger = loggo.GetLogger("juju.provisioner.lxc")

var _ Broker = (*lxcBroker)(nil)

func NewLxcBroker(factory golxc.ContainerFactory, config *config.Config, tools *state.Tools) Broker {
	return &lxcBroker{
		golxc:   factory,
		manager: lxc.NewContainerManager(factory, "juju"),
		config:  config,
		tools:   tools,
	}
}

type lxcBroker struct {
	golxc   golxc.ContainerFactory
	manager lxc.ContainerManager
	config  *config.Config
	tools   *state.Tools
}

func (broker *lxcBroker) StartInstance(machineId, machineNonce string, series string, cons constraints.Value, info *state.Info, apiInfo *api.Info) (instance.Instance, error) {
	lxcLogger.Infof("starting lxc container for machineId: %s", machineId)

	inst, err := broker.manager.StartContainer(machineId, series, machineNonce, broker.tools, broker.config, info, apiInfo)
	if err != nil {
		lxcLogger.Errorf("failed to start container: %v", err)
		return nil, err
	}
	return inst, nil
}

// StopInstances shuts down the given instances.
func (broker *lxcBroker) StopInstances(instances []instance.Instance) error {
	// TODO: potentially parallelise.
	for _, instance := range instances {
		lxcLogger.Infof("stopping lxc container for instance: %s", instance.Id())
		if err := broker.manager.StopContainer(instance); err != nil {
			lxcLogger.Errorf("container did not stop: %v", err)
			return err
		}
	}
	return nil
}

// AllInstances only returns running containers.
func (broker *lxcBroker) AllInstances() (result []instance.Instance, err error) {
	return broker.manager.ListContainers()
}
<commit_msg>Update implementation to match interface changes.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package provisioner

import (
	"launchpad.net/golxc"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/container/lxc"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	"launchpad.net/loggo"
)

var lxcLogger = loggo.GetLogger("juju.provisioner.lxc")

var _ Broker = (*lxcBroker)(nil)

func NewLxcBroker(factory golxc.ContainerFactory, config *config.Config, tools *state.Tools) Broker {
	return &lxcBroker{
		golxc:   factory,
		manager: lxc.NewContainerManager(factory, "juju"),
		config:  config,
		tools:   tools,
	}
}

type lxcBroker struct {
	golxc   golxc.ContainerFactory
	manager lxc.ContainerManager
	config  *config.Config
	tools   *state.Tools
}

func (broker *lxcBroker) StartInstance(machineId, machineNonce string, series string, cons constraints.Value, info *state.Info, apiInfo *api.Info) (instance.Instance, *instance.HardwareCharacteristics, error) {
	lxcLogger.Infof("starting lxc container for machineId: %s", machineId)

	inst, err := broker.manager.StartContainer(machineId, series, machineNonce, broker.tools, broker.config, info, apiInfo)
	if err != nil {
		lxcLogger.Errorf("failed to start container: %v", err)
		return nil, nil, err
	}
	return inst, nil, nil
}

// StopInstances shuts down the given instances.
func (broker *lxcBroker) StopInstances(instances []instance.Instance) error {
	// TODO: potentially parallelise.
	for _, instance := range instances {
		lxcLogger.Infof("stopping lxc container for instance: %s", instance.Id())
		if err := broker.manager.StopContainer(instance); err != nil {
			lxcLogger.Errorf("container did not stop: %v", err)
			return err
		}
	}
	return nil
}

// AllInstances only returns running containers.
func (broker *lxcBroker) AllInstances() (result []instance.Instance, err error) {
	return broker.manager.ListContainers()
}
<|endoftext|>


<commit_before>package webcam

import (
	"bytes"
	"encoding/binary"
	"unsafe"

	"github.com/blackjack/webcam/ioctl"
	"golang.org/x/sys/unix"
)

const (
	V4L2_CAP_VIDEO_CAPTURE      uint32 = 0x00000001
	V4L2_CAP_STREAMING          uint32 = 0x04000000
	V4L2_BUF_TYPE_VIDEO_CAPTURE uint32 = 1
	V4L2_MEMORY_MMAP            uint32 = 1
	V4L2_FIELD_ANY              uint32 = 0
)

const (
	V4L2_FRMSIZE_TYPE_DISCRETE   uint32 = 1
	V4L2_FRMSIZE_TYPE_CONTINUOUS uint32 = 2
	V4L2_FRMSIZE_TYPE_STEPWISE   uint32 = 3
)

var (
	VIDIOC_QUERYCAP = ioctl.IoR(uintptr('V'), 0, unsafe.Sizeof(v4l2_capability{}))
	VIDIOC_ENUM_FMT = ioctl.IoRW(uintptr('V'), 2, unsafe.Sizeof(v4l2_fmtdesc{}))
	VIDIOC_S_FMT    = ioctl.IoRW(uintptr('V'), 5, unsafe.Sizeof(v4l2_format{}))
	VIDIOC_REQBUFS  = ioctl.IoRW(uintptr('V'), 8, unsafe.Sizeof(v4l2_requestbuffers{}))
	VIDIOC_QUERYBUF = ioctl.IoRW(uintptr('V'), 9, unsafe.Sizeof(v4l2_buffer{}))
	VIDIOC_QBUF     = ioctl.IoRW(uintptr('V'), 15, unsafe.Sizeof(v4l2_buffer{}))
	VIDIOC_DQBUF    = ioctl.IoRW(uintptr('V'), 17, unsafe.Sizeof(v4l2_buffer{}))
	//sizeof int32
	VIDIOC_STREAMON        = ioctl.IoW(uintptr('V'), 18, 4)
	VIDIOC_ENUM_FRAMESIZES = ioctl.IoRW(uintptr('V'), 74, unsafe.Sizeof(v4l2_frmsizeenum{}))
)

type v4l2_capability struct {
	driver       [16]uint8
	card         [32]uint8
	bus_info     [32]uint8
	version      uint32
	capabilities uint32
	device_caps  uint32
	reserved     [3]uint32
}

type v4l2_fmtdesc struct {
	index       uint32
	_type       uint32
	flags       uint32
	description [32]uint8
	pixelformat uint32
	reserved    [4]uint32
}

type v4l2_frmsizeenum struct {
	index        uint32
	pixel_format uint32
	_type        uint32
	union        [24]uint8
	reserved     [2]uint32
}

type v4l2_frmsize_discrete struct {
	Width  uint32
	Height uint32
}

type v4l2_frmsize_stepwise struct {
	Min_width   uint32
	Max_width   uint32
	Step_width  uint32
	Min_height  uint32
	Max_height  uint32
	Step_height uint32
}

type v4l2_format struct {
	_type uint32
	union [204]uint8
}

type v4l2_pix_format struct {
	Width        uint32
	Height       uint32
	Pixelformat  uint32
	Field        uint32
	Bytesperline uint32
	Sizeimage    uint32
	Colorspace   uint32
	Priv         uint32
	Flags        uint32
	Ycbcr_enc    uint32
	Quantization uint32
	Xfer_func    uint32
}

type v4l2_requestbuffers struct {
	count    uint32
	_type    uint32
	memory   uint32
	reserved [2]uint32
}

type v4l2_buffer struct {
	index     uint32
	_type     uint32
	bytesused uint32
	flags     uint32
	field     uint32
	timestamp unix.Timeval
	timecode  v4l2_timecode
	sequence  uint32
	memory    uint32
	union     [8]uint8
	length    uint32
	reserved2 uint32
	reserved  uint32
}

type v4l2_timecode struct {
	_type    uint32
	flags    uint32
	frames   uint8
	seconds  uint8
	minutes  uint8
	hours    uint8
	userbits [4]uint8
}

func checkCapabilities(fd uintptr) (supportsVideoCapture bool, supportsVideoStreaming bool, err error) {

	caps := &v4l2_capability{}

	err = ioctl.Ioctl(fd, VIDIOC_QUERYCAP, uintptr(unsafe.Pointer(caps)))

	if err != nil {
		return
	}

	supportsVideoCapture = (caps.capabilities & V4L2_CAP_VIDEO_CAPTURE) != 0
	supportsVideoStreaming = (caps.capabilities & V4L2_CAP_STREAMING) != 0
	return

}

func getPixelFormat(fd uintptr, index uint32) (code uint32, description string, err error) {

	fmtdesc := &v4l2_fmtdesc{}

	fmtdesc.index = index
	fmtdesc._type = V4L2_BUF_TYPE_VIDEO_CAPTURE

	err = ioctl.Ioctl(fd, VIDIOC_ENUM_FMT, uintptr(unsafe.Pointer(fmtdesc)))

	if err != nil {
		return
	}

	code = fmtdesc.pixelformat
	description = CToGoString(fmtdesc.description[:])

	return
}

func getFrameSize(fd uintptr, index uint32, code uint32) (frameSize FrameSize, err error) {

	frmsizeenum := &v4l2_frmsizeenum{}
	frmsizeenum.index = index
	frmsizeenum.pixel_format = code

	err = ioctl.Ioctl(fd, VIDIOC_ENUM_FRAMESIZES, uintptr(unsafe.Pointer(frmsizeenum)))

	if err != nil {
		return
	}

	switch frmsizeenum._type {

	case V4L2_FRMSIZE_TYPE_DISCRETE:
		discrete := &v4l2_frmsize_discrete{}
		err = binary.Read(bytes.NewBuffer(frmsizeenum.union[:]), binary.LittleEndian, discrete)

		if err != nil {
			return
		}

		frameSize.MinWidth = discrete.Width
		frameSize.MaxWidth = discrete.Width
		frameSize.StepWidth = 0
		frameSize.MinHeight = discrete.Height
		frameSize.MaxHeight = discrete.Height
		frameSize.StepHeight = 0

	case V4L2_FRMSIZE_TYPE_CONTINUOUS:

	case V4L2_FRMSIZE_TYPE_STEPWISE:
		stepwise := &v4l2_frmsize_stepwise{}
		err = binary.Read(bytes.NewBuffer(frmsizeenum.union[:]), binary.LittleEndian, stepwise)

		if err != nil {
			return
		}

		frameSize.MinWidth = stepwise.Min_width
		frameSize.MaxWidth = stepwise.Max_width
		frameSize.StepWidth = stepwise.Step_width
		frameSize.MinHeight = stepwise.Min_height
		frameSize.MaxHeight = stepwise.Max_height
		frameSize.StepHeight = stepwise.Step_height
	}

	return
}

func setImageFormat(fd uintptr, formatcode *uint32, width *uint32, height *uint32) (err error) {

	format := &v4l2_format{
		_type: V4L2_BUF_TYPE_VIDEO_CAPTURE,
	}

	pix := v4l2_pix_format{
		Width:       *width,
		Height:      *height,
		Pixelformat: *formatcode,
		Field:       V4L2_FIELD_ANY,
	}

	pixbytes := &bytes.Buffer{}
	err = binary.Write(pixbytes, binary.LittleEndian, pix)

	if err != nil {
		return
	}

	copy(format.union[:], pixbytes.Bytes())

	err = ioctl.Ioctl(fd, VIDIOC_S_FMT, uintptr(unsafe.Pointer(format)))

	if err != nil {
		return
	}

	pixReverse := &v4l2_pix_format{}
	err = binary.Read(bytes.NewBuffer(format.union[:]), binary.LittleEndian, pixReverse)

	if err != nil {
		return
	}

	*width = pixReverse.Width
	*height = pixReverse.Height
	*formatcode = pixReverse.Pixelformat

	return

}

func mmapRequestBuffers(fd uintptr, buf_count *uint32) (err error) {

	req := &v4l2_requestbuffers{}
	req.count = *buf_count
	req._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	req.memory = V4L2_MEMORY_MMAP

	err = ioctl.Ioctl(fd, VIDIOC_REQBUFS, uintptr(unsafe.Pointer(req)))

	if err != nil {
		return
	}

	*buf_count = req.count

	return

}

func mmapQueryBuffer(fd uintptr, index uint32, length *uint32) (start unsafe.Pointer, err error) {

	buffer := &v4l2_buffer{}

	buffer._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	buffer.memory = V4L2_MEMORY_MMAP
	buffer.index = index

	err = ioctl.Ioctl(fd, VIDIOC_QUERYBUF, uintptr(unsafe.Pointer(buffer)))

	if err != nil {
		return
	}

	var offset uint32
	err = binary.Read(bytes.NewBuffer(buffer.union[:]), binary.LittleEndian, offset)

	if err != nil {
		return
	}

	*length = buffer.length

	returnPointer, _, errno := unix.Syscall6(unix.SYS_MMAP, 0, uintptr(buffer.length), unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED, fd, uintptr(offset))

	start = unsafe.Pointer(returnPointer)

	if errno != 0 {
		err = errno
	}

	return

}

func mmapDequeueBuffer(fd uintptr, index *uint32, length *uint32) (err error) {

	buffer := &v4l2_buffer{}

	buffer._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	buffer.memory = V4L2_MEMORY_MMAP

	err = ioctl.Ioctl(fd, VIDIOC_DQBUF, uintptr(unsafe.Pointer(buffer)))

	if err != nil {
		return
	}

	*index = buffer.index
	*length = buffer.bytesused

	return

}

func mmapEnqueueBuffer(fd uintptr, index uint32) (err error) {

	buffer := &v4l2_buffer{}

	buffer._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	buffer.memory = V4L2_MEMORY_MMAP
	buffer.index = index

	err = ioctl.Ioctl(fd, VIDIOC_QBUF, uintptr(unsafe.Pointer(buffer)))
	return

}

func mmapReleaseBuffer(start unsafe.Pointer, length uint32) (err error) {

	_, _, errno := unix.Syscall(unix.SYS_MUNMAP, uintptr(start), uintptr(length), uintptr(0))

	if errno != 0 {
		err = errno
	}

	return

}

func startStreaming(fd uintptr) (err error) {

	var uintPointer uint32 = V4L2_BUF_TYPE_VIDEO_CAPTURE
	err = ioctl.Ioctl(fd, VIDIOC_STREAMON, uintptr(unsafe.Pointer(&uintPointer)))
	return

}

func waitForFrame(fd uintptr, timeout uint32) (count int, err error) {

	for {

		fds := &FdSet{}
		fds.Set(fd)

		tv := &unix.Timeval{}
		tv.Sec = int64(timeout)

		countReturn, _, errno := unix.Syscall6(unix.SYS_SELECT, uintptr(fd+1), uintptr(unsafe.Pointer(fds)), uintptr(0), uintptr(0), uintptr(unsafe.Pointer(tv)), 0)

		count = int(countReturn)

		if errno != 0 {
			err = errno
		}

		if count < 0 {

			if err == unix.EINTR {
				continue
			}

			return

		}

		return

	}

}

func CToGoString(c []byte) string {
	n := -1
	for i, b := range c {
		if b == 0 {
			break
		}
		n = i
	}
	return string(c[:n+1])
}

//unix.Fdset produces wrong results (visible in strace)

const (
	FD_BITS    = uintptr(unsafe.Sizeof(0) * 8)
	FD_SETSIZE = 1024
)

type FdSet struct {
	bits [FD_SETSIZE / FD_BITS]int32
}

func (fds *FdSet) Set(fd uintptr) {
	fdmask := uint(1) << (uint32(fd) % uint32(FD_BITS))
	fds.bits[fd/FD_BITS] |= int32(fdmask)
}

func (fds *FdSet) IsSet(fd uintptr) bool {
	fdmask := uint32(1) << (uint32(fd) % uint32(FD_BITS))
	return (fds.bits[fd/FD_BITS] & int32(fdmask)) != 0
}
<commit_msg>Add fix padding of v4l2_format (need more investigation)<commit_after>package webcam

import (
	"bytes"
	"encoding/binary"
	"unsafe"

	"github.com/blackjack/webcam/ioctl"
	"golang.org/x/sys/unix"
)

const (
	V4L2_CAP_VIDEO_CAPTURE      uint32 = 0x00000001
	V4L2_CAP_STREAMING          uint32 = 0x04000000
	V4L2_BUF_TYPE_VIDEO_CAPTURE uint32 = 1
	V4L2_MEMORY_MMAP            uint32 = 1
	V4L2_FIELD_ANY              uint32 = 0
)

const (
	V4L2_FRMSIZE_TYPE_DISCRETE   uint32 = 1
	V4L2_FRMSIZE_TYPE_CONTINUOUS uint32 = 2
	V4L2_FRMSIZE_TYPE_STEPWISE   uint32 = 3
)

var (
	VIDIOC_QUERYCAP = ioctl.IoR(uintptr('V'), 0, unsafe.Sizeof(v4l2_capability{}))
	VIDIOC_ENUM_FMT = ioctl.IoRW(uintptr('V'), 2, unsafe.Sizeof(v4l2_fmtdesc{}))
	VIDIOC_S_FMT    = ioctl.IoRW(uintptr('V'), 5, unsafe.Sizeof(v4l2_format{}))
	VIDIOC_REQBUFS  = ioctl.IoRW(uintptr('V'), 8, unsafe.Sizeof(v4l2_requestbuffers{}))
	VIDIOC_QUERYBUF = ioctl.IoRW(uintptr('V'), 9, unsafe.Sizeof(v4l2_buffer{}))
	VIDIOC_QBUF     = ioctl.IoRW(uintptr('V'), 15, unsafe.Sizeof(v4l2_buffer{}))
	VIDIOC_DQBUF    = ioctl.IoRW(uintptr('V'), 17, unsafe.Sizeof(v4l2_buffer{}))
	//sizeof int32
	VIDIOC_STREAMON        = ioctl.IoW(uintptr('V'), 18, 4)
	VIDIOC_ENUM_FRAMESIZES = ioctl.IoRW(uintptr('V'), 74, unsafe.Sizeof(v4l2_frmsizeenum{}))
)

type v4l2_capability struct {
	driver       [16]uint8
	card         [32]uint8
	bus_info     [32]uint8
	version      uint32
	capabilities uint32
	device_caps  uint32
	reserved     [3]uint32
}

type v4l2_fmtdesc struct {
	index       uint32
	_type       uint32
	flags       uint32
	description [32]uint8
	pixelformat uint32
	reserved    [4]uint32
}

type v4l2_frmsizeenum struct {
	index        uint32
	pixel_format uint32
	_type        uint32
	union        [24]uint8
	reserved     [2]uint32
}

type v4l2_frmsize_discrete struct {
	Width  uint32
	Height uint32
}

type v4l2_frmsize_stepwise struct {
	Min_width   uint32
	Max_width   uint32
	Step_width  uint32
	Min_height  uint32
	Max_height  uint32
	Step_height uint32
}

type v4l2_format struct {
	_type uint32
	_     [4]byte
	union [200]uint8
}

type v4l2_pix_format struct {
	Width        uint32
	Height       uint32
	Pixelformat  uint32
	Field        uint32
	Bytesperline uint32
	Sizeimage    uint32
	Colorspace   uint32
	Priv         uint32
	Flags        uint32
	Ycbcr_enc    uint32
	Quantization uint32
	Xfer_func    uint32
}

type v4l2_requestbuffers struct {
	count    uint32
	_type    uint32
	memory   uint32
	reserved [2]uint32
}

type v4l2_buffer struct {
	index     uint32
	_type     uint32
	bytesused uint32
	flags     uint32
	field     uint32
	timestamp unix.Timeval
	timecode  v4l2_timecode
	sequence  uint32
	memory    uint32
	union     [8]uint8
	length    uint32
	reserved2 uint32
	reserved  uint32
}

type v4l2_timecode struct {
	_type    uint32
	flags    uint32
	frames   uint8
	seconds  uint8
	minutes  uint8
	hours    uint8
	userbits [4]uint8
}

func checkCapabilities(fd uintptr) (supportsVideoCapture bool, supportsVideoStreaming bool, err error) {

	caps := &v4l2_capability{}

	err = ioctl.Ioctl(fd, VIDIOC_QUERYCAP, uintptr(unsafe.Pointer(caps)))

	if err != nil {
		return
	}

	supportsVideoCapture = (caps.capabilities & V4L2_CAP_VIDEO_CAPTURE) != 0
	supportsVideoStreaming = (caps.capabilities & V4L2_CAP_STREAMING) != 0
	return

}

func getPixelFormat(fd uintptr, index uint32) (code uint32, description string, err error) {

	fmtdesc := &v4l2_fmtdesc{}

	fmtdesc.index = index
	fmtdesc._type = V4L2_BUF_TYPE_VIDEO_CAPTURE

	err = ioctl.Ioctl(fd, VIDIOC_ENUM_FMT, uintptr(unsafe.Pointer(fmtdesc)))

	if err != nil {
		return
	}

	code = fmtdesc.pixelformat
	description = CToGoString(fmtdesc.description[:])

	return
}

func getFrameSize(fd uintptr, index uint32, code uint32) (frameSize FrameSize, err error) {

	frmsizeenum := &v4l2_frmsizeenum{}
	frmsizeenum.index = index
	frmsizeenum.pixel_format = code

	err = ioctl.Ioctl(fd, VIDIOC_ENUM_FRAMESIZES, uintptr(unsafe.Pointer(frmsizeenum)))

	if err != nil {
		return
	}

	switch frmsizeenum._type {

	case V4L2_FRMSIZE_TYPE_DISCRETE:
		discrete := &v4l2_frmsize_discrete{}
		err = binary.Read(bytes.NewBuffer(frmsizeenum.union[:]), binary.LittleEndian, discrete)

		if err != nil {
			return
		}

		frameSize.MinWidth = discrete.Width
		frameSize.MaxWidth = discrete.Width
		frameSize.StepWidth = 0
		frameSize.MinHeight = discrete.Height
		frameSize.MaxHeight = discrete.Height
		frameSize.StepHeight = 0

	case V4L2_FRMSIZE_TYPE_CONTINUOUS:

	case V4L2_FRMSIZE_TYPE_STEPWISE:
		stepwise := &v4l2_frmsize_stepwise{}
		err = binary.Read(bytes.NewBuffer(frmsizeenum.union[:]), binary.LittleEndian, stepwise)

		if err != nil {
			return
		}

		frameSize.MinWidth = stepwise.Min_width
		frameSize.MaxWidth = stepwise.Max_width
		frameSize.StepWidth = stepwise.Step_width
		frameSize.MinHeight = stepwise.Min_height
		frameSize.MaxHeight = stepwise.Max_height
		frameSize.StepHeight = stepwise.Step_height
	}

	return
}

func setImageFormat(fd uintptr, formatcode *uint32, width *uint32, height *uint32) (err error) {

	format := &v4l2_format{
		_type: V4L2_BUF_TYPE_VIDEO_CAPTURE,
	}

	pix := v4l2_pix_format{
		Width:       *width,
		Height:      *height,
		Pixelformat: *formatcode,
		Field:       V4L2_FIELD_ANY,
	}

	pixbytes := &bytes.Buffer{}
	err = binary.Write(pixbytes, binary.LittleEndian, pix)

	if err != nil {
		return
	}

	copy(format.union[:], pixbytes.Bytes())

	err = ioctl.Ioctl(fd, VIDIOC_S_FMT, uintptr(unsafe.Pointer(format)))

	if err != nil {
		return
	}

	pixReverse := &v4l2_pix_format{}
	err = binary.Read(bytes.NewBuffer(format.union[:]), binary.LittleEndian, pixReverse)

	if err != nil {
		return
	}

	*width = pixReverse.Width
	*height = pixReverse.Height
	*formatcode = pixReverse.Pixelformat

	return

}

func mmapRequestBuffers(fd uintptr, buf_count *uint32) (err error) {

	req := &v4l2_requestbuffers{}
	req.count = *buf_count
	req._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	req.memory = V4L2_MEMORY_MMAP

	err = ioctl.Ioctl(fd, VIDIOC_REQBUFS, uintptr(unsafe.Pointer(req)))

	if err != nil {
		return
	}

	*buf_count = req.count

	return

}

func mmapQueryBuffer(fd uintptr, index uint32, length *uint32) (start unsafe.Pointer, err error) {

	buffer := &v4l2_buffer{}

	buffer._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	buffer.memory = V4L2_MEMORY_MMAP
	buffer.index = index

	err = ioctl.Ioctl(fd, VIDIOC_QUERYBUF, uintptr(unsafe.Pointer(buffer)))

	if err != nil {
		return
	}

	var offset uint32
	err = binary.Read(bytes.NewBuffer(buffer.union[:]), binary.LittleEndian, offset)

	if err != nil {
		return
	}

	*length = buffer.length

	returnPointer, _, errno := unix.Syscall6(unix.SYS_MMAP, 0, uintptr(buffer.length), unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED, fd, uintptr(offset))

	start = unsafe.Pointer(returnPointer)

	if errno != 0 {
		err = errno
	}

	return

}

func mmapDequeueBuffer(fd uintptr, index *uint32, length *uint32) (err error) {

	buffer := &v4l2_buffer{}

	buffer._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	buffer.memory = V4L2_MEMORY_MMAP

	err = ioctl.Ioctl(fd, VIDIOC_DQBUF, uintptr(unsafe.Pointer(buffer)))

	if err != nil {
		return
	}

	*index = buffer.index
	*length = buffer.bytesused

	return

}

func mmapEnqueueBuffer(fd uintptr, index uint32) (err error) {

	buffer := &v4l2_buffer{}

	buffer._type = V4L2_BUF_TYPE_VIDEO_CAPTURE
	buffer.memory = V4L2_MEMORY_MMAP
	buffer.index = index

	err = ioctl.Ioctl(fd, VIDIOC_QBUF, uintptr(unsafe.Pointer(buffer)))
	return

}

func mmapReleaseBuffer(start unsafe.Pointer, length uint32) (err error) {

	_, _, errno := unix.Syscall(unix.SYS_MUNMAP, uintptr(start), uintptr(length), uintptr(0))

	if errno != 0 {
		err = errno
	}

	return

}

func startStreaming(fd uintptr) (err error) {

	var uintPointer uint32 = V4L2_BUF_TYPE_VIDEO_CAPTURE
	err = ioctl.Ioctl(fd, VIDIOC_STREAMON, uintptr(unsafe.Pointer(&uintPointer)))
	return

}

func waitForFrame(fd uintptr, timeout uint32) (count int, err error) {

	for {

		fds := &FdSet{}
		fds.Set(fd)

		tv := &unix.Timeval{}
		tv.Sec = int64(timeout)

		countReturn, _, errno := unix.Syscall6(unix.SYS_SELECT, uintptr(fd+1), uintptr(unsafe.Pointer(fds)), uintptr(0), uintptr(0), uintptr(unsafe.Pointer(tv)), 0)

		count = int(countReturn)

		if errno != 0 {
			err = errno
		}

		if count < 0 {

			if err == unix.EINTR {
				continue
			}

			return

		}

		return

	}

}

func CToGoString(c []byte) string {
	n := -1
	for i, b := range c {
		if b == 0 {
			break
		}
		n = i
	}
	return string(c[:n+1])
}

//unix.Fdset produces wrong results (visible in strace)

const (
	FD_BITS    = uintptr(unsafe.Sizeof(0) * 8)
	FD_SETSIZE = 1024
)

type FdSet struct {
	bits [FD_SETSIZE / FD_BITS]int32
}

func (fds *FdSet) Set(fd uintptr) {
	fdmask := uint(1) << (uint32(fd) % uint32(FD_BITS))
	fds.bits[fd/FD_BITS] |= int32(fdmask)
}

func (fds *FdSet) IsSet(fd uintptr) bool {
	fdmask := uint32(1) << (uint32(fd) % uint32(FD_BITS))
	return (fds.bits[fd/FD_BITS] & int32(fdmask)) != 0
}
<|endoftext|>


<commit_before><commit_msg>Remove log line from Go unit test<commit_after><|endoftext|>


<commit_before><commit_msg>fixing test<commit_after><|endoftext|>


<commit_before>package gunfish

import (
	"time"
)

// Version
const (
	Version = "v0.1.1"
)

// Limit values
const (
	MaxWorkerNum           = 119   // Maximum of worker number
	MinWorkerNum           = 1     // Minimum of worker number
	MaxSenderNum           = 150   // Maximum of sender number
	MinSenderNum           = 1     // Minimum of sender number
	MaxQueueSize           = 40960 // Maximum queue size.
	MinQueueSize           = 128   // Minimum Queue size.
	MaxRequestSize         = 5000  // Maximum of requset count.
	MinRequestSize         = 1     // Minimum of request size.
	LimitApnsTokenByteSize = 100   // Payload byte size.
)

// Default values
const (
	// SendRetryCount is the threashold which is resend count.
	SendRetryCount = 10
	// RetryWaitTime is periodical time to retrieve notifications from retry queue to resend
	RetryWaitTime = time.Millisecond * 500
	// RetryOnceCount is the number of sending notification at once.
	RetryOnceCount = 1000
	// Default multiplicity of sending notifications to apns. If not configures
	// at file, this value is set.
	DefaultApnsSenderNum = 20
	// Default array size of posted data. If not configures at file, this value is set.
	DefaultRequestQueueSize = 2000
	// Default port number of provider server
	DefaultPort = 8003
	// Default supervisor's queue size. If not configures at file, this value is set.
	DefaultQueueSize = 1000
	// About the average time of response from apns. That value is not accurate
	// because that is defined heuristically in Japan.
	AverageResponseTime = time.Millisecond * 150
	// Minimum RetryAfter time (seconds).
	RetryAfterSecond = time.Second * 10
	// Gunfish returns RetryAfter header based on 'Exponential Backoff'. Therefore,
	// that defines the wait time threshold so as not to wait too long.
	ResetRetryAfterSecond = time.Second * 60
	// FlowRateInterval is the designed value to enable to delivery notifications
	// for that value seconds. Gunfish is designed as to ensure to delivery
	// notifications for 10 seconds.
	FlowRateInterval = time.Second * 10
	// Default flow rate as notification requests per sec (req/sec).
	DefaultFlowRatePerSec = 2000
	// Wait millisecond interval when to shutdown.
	ShutdownWaitTime = time.Millisecond * 10
	// That is the count while request counter is 0 in the 'ShutdownWaitTime' period.
	RestartWaitCount = 50

	// That is GCM client timeout
	GCMClientTimeout = time.Second * 10
)

// Apns endpoints
const (
	APNsDevURL  = "https://api.development.push.apple.com"
	APNsProdURL = "https://api.push.apple.com"
	APNsMockURL = "https://localhost:2195"
	GCMProdURL  = "https://gcm-http.googleapis.com/gcm/send"
)

// Supports Content-Type
const (
	ApplicationJSON              = "application/json"
	ApplicationXW3FormURLEncoded = "application/x-www-form-urlencoded"
)

// Environment struct
type Environment int

// Executed environment
const (
	Production Environment = iota
	Development
	Test
	Disable
)

// Alert fields mapping
var (
	AlertKeyToField = map[string]string{
		"title":          "Title",
		"body":           "Body",
		"title-loc-key":  "TitleLocKey",
		"title-loc-args": "TitleLocArgs",
		"action-loc-key": "ActionLocKey",
		"loc-key":        "LocKey",
		"loc-args":       "LocArgs",
		"launch-image":   "LaunchImage",
	}
)
<commit_msg>this change should be modified other branch<commit_after>package gunfish

import (
	"time"
)

// Version
const (
	Version = "v0.1.1"
)

// Limit values
const (
	MaxWorkerNum           = 119   // Maximum of worker number
	MinWorkerNum           = 1     // Minimum of worker number
	MaxSenderNum           = 150   // Maximum of sender number
	MinSenderNum           = 1     // Minimum of sender number
	MaxQueueSize           = 40960 // Maximum queue size.
	MinQueueSize           = 128   // Minimum Queue size.
	MaxRequestSize         = 5000  // Maximum of requset count.
	MinRequestSize         = 1     // Minimum of request size.
	LimitApnsTokenByteSize = 100   // Payload byte size.
)

// Default values
const (
	// SendRetryCount is the threashold which is resend count.
	SendRetryCount = 10
	// RetryWaitTime is periodical time to retrieve notifications from retry queue to resend
	RetryWaitTime = time.Millisecond * 500
	// RetryOnceCount is the number of sending notification at once.
	RetryOnceCount = 1000
	// Default multiplicity of sending notifications to apns. If not configures
	// at file, this value is set.
	DefaultApnsSenderNum = 20
	// Default array size of posted data. If not configures at file, this value is set.
	DefaultRequestQueueSize = 2000
	// Default port number of provider server
	DefaultPort = 8003
	// Default supervisor's queue size. If not configures at file, this value is set.
	DefaultQueueSize = 1000
	// About the average time of response from apns. That value is not accurate
	// because that is defined heuristically in Japan.
	AverageResponseTime = time.Millisecond * 150
	// Minimum RetryAfter time (seconds).
	RetryAfterSecond = time.Second * 10
	// Gunfish returns RetryAfter header based on 'Exponential Backoff'. Therefore,
	// that defines the wait time threshold so as not to wait too long.
	ResetRetryAfterSecond = time.Second * 60
	// FlowRateInterval is the designed value to enable to delivery notifications
	// for that value seconds. Gunfish is designed as to ensure to delivery
	// notifications for 10 seconds.
	FlowRateInterval = time.Second * 10
	// Default flow rate as notification requests per sec (req/sec).
	DefaultFlowRatePerSec = 2000
	// Wait millisecond interval when to shutdown.
	ShutdownWaitTime = time.Millisecond * 10
	// That is the count while request counter is 0 in the 'ShutdownWaitTime' period.
	RestartWaitCount = 50
)

// Apns endpoints
const (
	DevServer  = "https://api.development.push.apple.com"
	ProdServer = "https://api.push.apple.com"
	MockServer = "https://localhost:2195"
)

// Supports Content-Type
const (
	ApplicationJSON              = "application/json"
	ApplicationXW3FormURLEncoded = "application/x-www-form-urlencoded"
)

// Environment struct
type Environment int

// Executed environment
const (
	Production Environment = iota
	Development
	Test
	Disable
)

// Alert fields mapping
var (
	AlertKeyToField = map[string]string{
		"title":          "Title",
		"body":           "Body",
		"title-loc-key":  "TitleLocKey",
		"title-loc-args": "TitleLocArgs",
		"action-loc-key": "ActionLocKey",
		"loc-key":        "LocKey",
		"loc-args":       "LocArgs",
		"launch-image":   "LaunchImage",
	}
)
<|endoftext|>


<commit_before>package pool

import (
	"context"
	"io/ioutil"
	"net/http"
	"time"
)

func crawl(t Task) Task {
	startTime := time.Now()
	client := &http.Client{
		Timeout: timeout,
	}
	if t.Proxy.Host != "" {
		client.Transport = &http.Transport{
			Proxy:             http.ProxyURL(t.Proxy),
			DisableKeepAlives: true,
		}
	}
	req, err := http.NewRequest("GET", t.Hostname, nil)
	if err != nil {
		t.Error = err
		return t
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	req = req.WithContext(ctx)
	resp, err := client.Do(req)
	if err != nil {
		t.Error = err
		return t
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Error = err
		return t
	}
	err = resp.Body.Close()
	if err != nil {
		t.Error = err
		return t
	}
	t.Body = body
	t.Response = resp
	t.ResponceTime = time.Since(startTime)
	return t
}
<commit_msg>move to task<commit_after><|endoftext|>


<commit_before>package query

import (
	"fmt"
	"time"

	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/lib/series"
	"github.com/yuuki/diamondb/lib/storage"
)

type UnsupportedFunctionError struct {
	funcName string
}

func (e *UnsupportedFunctionError) Error() string {
	return fmt.Sprintf("Unsupported function %s", e.funcName)
}

type funcArg struct {
	expr        Expr
	seriesSlice series.SeriesSlice
}

type funcArgs []*funcArg

func EvalTarget(fetcher storage.Fetcher, target string, startTime, endTime time.Time) (series.SeriesSlice, error) {
	expr, err := ParseTarget(target)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to ParseTarget %s", target)
	}
	ss, err := invokeExpr(fetcher, expr, startTime, endTime)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to invokeExpr %v %d %d", expr, startTime, endTime)
	}
	return ss, err
}

func invokeExpr(fetcher storage.Fetcher, expr Expr, startTime, endTime time.Time) (series.SeriesSlice, error) {
	switch e := expr.(type) {
	case SeriesListExpr:
		ss, err := fetcher.FetchSeriesSlice(e.Literal, startTime, endTime)
		if err != nil {
			return nil, errors.Wrapf(err,
				"Failed to FetchSeriesSlice %s %d %d",
				e.Literal, startTime.Unix(), endTime.Unix(),
			)
		}
		return ss, nil
	case FuncExpr:
		args := funcArgs{}
		for _, expr := range e.SubExprs {
			switch expr.(type) {
			case BoolExpr:
				args = append(args, &funcArg{expr: expr})
			case NumberExpr:
				args = append(args, &funcArg{expr: expr})
			case StringExpr:
				args = append(args, &funcArg{expr: expr})
			case SeriesListExpr, FuncExpr:
				ss, err := invokeExpr(fetcher, expr, startTime, endTime)
				if err != nil {
					return nil, errors.Wrapf(err, "Failed to invokeExpr %v %d %d", expr, startTime, endTime)
				}
				args = append(args, &funcArg{expr: expr, seriesSlice: ss})
			default:
				return nil, errors.Errorf("Unknown expression %+v", expr)
			}
		}
		switch e.Name {
		case "alias":
			ss, err := doAlias(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "averageSeries", "avg":
			ss, err := doAverageSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "sumSeries", "sum":
			ss, err := doSumSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "minSeries":
			ss, err := doMinSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "maxSeries":
			ss, err := doMaxSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "multiplySeries":
			ss, err := doMultiplySeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "divideSeries":
			ss, err := doDivideSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		default:
			return nil, &UnsupportedFunctionError{funcName: e.Name}
		}
	default:
		return nil, errors.Errorf("Unknown expression %+v", expr)
	}
}
<commit_msg>Fix bug that the argument validation of nested function's arguments is not working<commit_after>package query

import (
	"fmt"
	"time"

	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/lib/series"
	"github.com/yuuki/diamondb/lib/storage"
)

type UnsupportedFunctionError struct {
	funcName string
}

func (e *UnsupportedFunctionError) Error() string {
	return fmt.Sprintf("Unsupported function %s", e.funcName)
}

type funcArg struct {
	expr        Expr
	seriesSlice series.SeriesSlice
}

type funcArgs []*funcArg

func EvalTarget(fetcher storage.Fetcher, target string, startTime, endTime time.Time) (series.SeriesSlice, error) {
	expr, err := ParseTarget(target)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to ParseTarget %s", target)
	}
	ss, err := invokeExpr(fetcher, expr, startTime, endTime)
	if err != nil {
		return nil, errors.Wrapf(err, "Failed to invokeExpr %v %d %d", expr, startTime, endTime)
	}
	return ss, err
}

func invokeExpr(fetcher storage.Fetcher, expr Expr, startTime, endTime time.Time) (series.SeriesSlice, error) {
	switch e := expr.(type) {
	case SeriesListExpr:
		ss, err := fetcher.FetchSeriesSlice(e.Literal, startTime, endTime)
		if err != nil {
			return nil, errors.Wrapf(err,
				"Failed to FetchSeriesSlice %s %d %d",
				e.Literal, startTime.Unix(), endTime.Unix(),
			)
		}
		return ss, nil
	case FuncExpr:
		args := funcArgs{}
		for _, expr := range e.SubExprs {
			switch e2 := expr.(type) {
			case BoolExpr:
				args = append(args, &funcArg{expr: expr})
			case NumberExpr:
				args = append(args, &funcArg{expr: expr})
			case StringExpr:
				args = append(args, &funcArg{expr: expr})
			case SeriesListExpr:
				ss, err := invokeExpr(fetcher, expr, startTime, endTime)
				if err != nil {
					return nil, errors.Wrapf(err, "Failed to invokeExpr %v %d %d", expr, startTime, endTime)
				}
				ex := SeriesListExpr{Literal: ss.FormatedName()}
				args = append(args, &funcArg{expr: ex, seriesSlice: ss})
			case FuncExpr:
				ss, err := invokeExpr(fetcher, expr, startTime, endTime)
				if err != nil {
					return nil, errors.Wrapf(err, "Failed to invokeExpr %v %d %d", expr, startTime, endTime)
				}
				// Regard FuncExpr as SeriesListExpr after process function
				ex := SeriesListExpr{Literal: fmt.Sprintf("%s(%s)", e2.Name, ss.FormatedName())}
				args = append(args, &funcArg{expr: ex, seriesSlice: ss})
			default:
				return nil, errors.Errorf("Unknown expression %+v", expr)
			}
		}
		switch e.Name {
		case "alias":
			ss, err := doAlias(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "averageSeries", "avg":
			ss, err := doAverageSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "sumSeries", "sum":
			ss, err := doSumSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "minSeries":
			ss, err := doMinSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "maxSeries":
			ss, err := doMaxSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "multiplySeries":
			ss, err := doMultiplySeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		case "divideSeries":
			ss, err := doDivideSeries(args)
			if err != nil {
				return nil, errors.WithStack(err)
			}
			return ss, err
		default:
			return nil, &UnsupportedFunctionError{funcName: e.Name}
		}
	default:
		return nil, errors.Errorf("Unknown expression %+v", expr)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	lua "github.com/anaminus/gopher-lua"
	"github.com/anaminus/rbxmk"
	"github.com/anaminus/rbxmk/rtypes"
	"github.com/anaminus/snek"
	"github.com/robloxapi/types"
)

const testdata = "testdata"

// Replace scriptArguments[x] with test script.
const replaceIndex = 3

var scriptArguments = [...]string{
	"rbxmk_test",
	"run",
	"--debug",
	"-",
	"true",
	"false",
	"nil",
	"42",
	"3.141592653589793",
	"-1e-8",
	"hello, world!",
	"hello\000world!",
}

type dummyFile struct {
	r    io.Reader
	info *dummyInfo
}

func (d *dummyFile) Name() string               { return "test" }
func (d *dummyFile) Stat() (os.FileInfo, error) { return d.info, nil }
func (d *dummyFile) Read(b []byte) (int, error) { return d.r.Read(b) }
func (d *dummyFile) Write([]byte) (int, error)  { return 0, nil }

type dummyInfo struct {
	name  string
	size  int64
	mode  os.FileMode
	time  time.Time
	isdir bool
}

func (d *dummyInfo) Name() string       { return d.name }
func (d *dummyInfo) Size() int64        { return d.size }
func (d *dummyInfo) Mode() os.FileMode  { return d.mode }
func (d *dummyInfo) ModTime() time.Time { return d.time }
func (d *dummyInfo) IsDir() bool        { return d.isdir }
func (d *dummyInfo) Sys() interface{}   { return d }

func initMain(s rbxmk.State, t *testing.T) {
	T := s.L.CreateTable(0, 2)

	// Pass makes a positive assertion. If the first argument is a non-function,
	// then an error is emitted if the value is falsy. Otherwise, the function
	// is called. If the call errors or returns a falsy value, then an error is
	// emitted. Returning no value counts as truthy. The error is emitted to
	// testing.T, but does not cause a Lua error to be thrown. The second
	// optional argument is a string that describes the assertion, which is
	// included with an emitted error.
	T.RawSetString("Pass", s.WrapFunc(func(s rbxmk.State) int {
		v := s.CheckAny(1)
		msg := s.OptString(2, "expected pass")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err != nil {
				t.Errorf("%s: %s", msg, err.Error())
				return 0
			}
			if s.Count() > n {
				if !s.L.ToBool(n + 1) {
					if m := s.L.ToString(n + 2); m != "" {
						t.Errorf("%s: %s", msg, m)
					} else {
						t.Errorf(msg)
					}
					return 0
				}
			}
		default:
			if !lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// Fail makes a negative assertion. If the first argument is a non-function,
	// then an error is emitted if the value is truthy. Otherwise, the function
	// is called. If the call does not error or returns a truthy value, then an
	// error is emitted. Returning no value counts as falsy. The error is
	// emitted to testing.T, but does not cause a Lua error to be thrown. The
	// second optional argument is a string that describes the assertion, which
	// is included with an emitted error.
	T.RawSetString("Fail", s.WrapFunc(func(s rbxmk.State) int {
		v := s.CheckAny(1)
		msg := s.OptString(2, "expected fail")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err == nil {
				if s.Count() > n {
					if s.L.ToBool(n + 1) {
						if m := s.L.ToString(n + 2); m != "" {
							t.Errorf("%s: %s", msg, m)
						} else {
							t.Errorf(msg)
						}
					}
					return 0
				}
				t.Errorf(msg)
				return 0
			} else if lua.LVAsBool(s.L.GetGlobal("SHOW_ERRORS")) {
				t.Logf("ERROR: %s\n", err)
			}
		default:
			if lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// GC runs the garbage collector.
	T.RawSetString("GC", s.L.NewFunction(func(l *lua.LState) int {
		runtime.GC()
		return 0
	}))

	// UserDataCacheLen returns the number of cached userdata values.
	T.RawSetString("UserDataCacheLen", s.WrapFunc(func(s rbxmk.State) int {
		return s.Push(types.Int(s.UserDataCacheLen()))
	}))

	// UserDataID returns a string that identifies a userdata value.
	T.RawSetString("UserDataID", s.WrapFunc(func(s rbxmk.State) int {
		return s.Push(types.String(fmt.Sprintf("%p", s.CheckUserData(1))))
	}))

	T.RawSetString("DummySymbol", s.UserDataOf(rtypes.Symbol{Name: "DummySymbol"}, "Symbol"))

	s.L.SetGlobal("T", T)
}

// TestScripts runs each .lua file in testdata as a Lua script. If the first
// line starts with a comment that contains "fail", then the script is expected
// to throw an error. All scripts receive the arguments from scriptArguments.
func TestScripts(t *testing.T) {
	program := snek.NewProgram("", scriptArguments[:])
	program.Register(snek.Def{
		Name: "run",
		New: func() snek.Command {
			return &RunCommand{Init: func(s rbxmk.State) { initMain(s, t) }}
		},
	})

	var files []string
	wd, _ := os.Getwd()
	for _, arg := range os.Args[2:] {
		if strings.HasPrefix(arg, "-test.") {
			continue
		}
		rel, err := filepath.Rel(wd, arg)
		if err != nil {
			rel = arg
		}
		files = append(files, rel)
	}
	if len(files) == 0 {
		err := filepath.Walk(testdata, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() &&
				filepath.Ext(info.Name()) == ".lua" &&
				!strings.HasPrefix(filepath.Base(info.Name()), "_") {
				files = append(files, path)
			}
			return nil
		})
		if err != nil {
			t.Fatalf("error walking testdata: %s", err)
		}
	}
	for _, file := range files {
		t.Run(filepath.ToSlash(file), func(t *testing.T) {
			args := make([]string, len(scriptArguments))
			copy(args, scriptArguments[:])
			args[replaceIndex] = file
			err := program.RunWithInput("run", snek.Input{
				Program:   args[0],
				Arguments: args[2:],
			})
			if err != nil {
				t.Errorf("script %s: %s", file, err)
			}
		})
	}
}
<commit_msg>Add T.Equal function for comparing values.<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	lua "github.com/anaminus/gopher-lua"
	"github.com/anaminus/rbxmk"
	"github.com/anaminus/rbxmk/rtypes"
	"github.com/anaminus/snek"
	"github.com/robloxapi/types"
)

const testdata = "testdata"

// Replace scriptArguments[x] with test script.
const replaceIndex = 3

var scriptArguments = [...]string{
	"rbxmk_test",
	"run",
	"--debug",
	"-",
	"true",
	"false",
	"nil",
	"42",
	"3.141592653589793",
	"-1e-8",
	"hello, world!",
	"hello\000world!",
}

type dummyFile struct {
	r    io.Reader
	info *dummyInfo
}

func (d *dummyFile) Name() string               { return "test" }
func (d *dummyFile) Stat() (os.FileInfo, error) { return d.info, nil }
func (d *dummyFile) Read(b []byte) (int, error) { return d.r.Read(b) }
func (d *dummyFile) Write([]byte) (int, error)  { return 0, nil }

type dummyInfo struct {
	name  string
	size  int64
	mode  os.FileMode
	time  time.Time
	isdir bool
}

func (d *dummyInfo) Name() string       { return d.name }
func (d *dummyInfo) Size() int64        { return d.size }
func (d *dummyInfo) Mode() os.FileMode  { return d.mode }
func (d *dummyInfo) ModTime() time.Time { return d.time }
func (d *dummyInfo) IsDir() bool        { return d.isdir }
func (d *dummyInfo) Sys() interface{}   { return d }

func deepeq(t *testing.T, s rbxmk.State, msg string, a, b lua.LValue) bool {
	if b, ok := b.(*lua.LTable); ok {
		if a, ok := a.(*lua.LTable); ok {
			visited := map[lua.LValue]struct{}{}
			b.ForEach(func(k, v lua.LValue) error {
				switch k := k.(type) {
				case lua.LNumber:
					deepeq(t, s, msg+"["+k.String()+"]", a.RawGetInt(int(k)), v)
				case lua.LString:
					deepeq(t, s, msg+"."+k.String(), a.RawGetString(string(k)), v)
				default:
					deepeq(t, s, msg+"["+k.String()+"]", a.RawGet(k), v)
				}
				visited[k] = struct{}{}
				return nil
			})
			a.ForEach(func(k, v lua.LValue) error {
				if _, ok := visited[k]; ok {
					return nil
				}
				switch k := k.(type) {
				case lua.LNumber:
					deepeq(t, s, msg+"["+k.String()+"]", v, b.RawGetInt(int(k)))
				case lua.LString:
					deepeq(t, s, msg+"."+k.String(), v, b.RawGetString(string(k)))
				default:
					deepeq(t, s, msg+"["+k.String()+"]", v, b.RawGet(k))
				}
				return nil
			})
			return true
		}
	}
	if !s.L.Equal(a, b) {
		if a.Type() != b.Type() {
			t.Errorf("%s: expected type %s, got %s", msg, b.Type(), a.Type())
			return false
		}
		if a.Type() == lua.LTString {
			t.Errorf("%s: expected %q, got %q", msg, b.String(), a.String())
			return false
		}
		t.Errorf("%s: expected %s, got %s", msg, b.String(), a.String())
		return false
	}
	return true
}

func initMain(s rbxmk.State, t *testing.T) {
	T := s.L.CreateTable(0, 2)

	// Pass makes a positive assertion. If the first argument is a non-function,
	// then an error is emitted if the value is falsy. Otherwise, the function
	// is called. If the call errors or returns a falsy value, then an error is
	// emitted. Returning no value counts as truthy. The error is emitted to
	// testing.T, but does not cause a Lua error to be thrown. The second
	// optional argument is a string that describes the assertion, which is
	// included with an emitted error.
	T.RawSetString("Pass", s.WrapFunc(func(s rbxmk.State) int {
		v := s.CheckAny(1)
		msg := s.OptString(2, "expected pass")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err != nil {
				t.Errorf("%s: %s", msg, err.Error())
				return 0
			}
			if s.Count() > n {
				if !s.L.ToBool(n + 1) {
					if m := s.L.ToString(n + 2); m != "" {
						t.Errorf("%s: %s", msg, m)
					} else {
						t.Errorf(msg)
					}
					return 0
				}
			}
		default:
			if !lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// Fail makes a negative assertion. If the first argument is a non-function,
	// then an error is emitted if the value is truthy. Otherwise, the function
	// is called. If the call does not error or returns a truthy value, then an
	// error is emitted. Returning no value counts as falsy. The error is
	// emitted to testing.T, but does not cause a Lua error to be thrown. The
	// second optional argument is a string that describes the assertion, which
	// is included with an emitted error.
	T.RawSetString("Fail", s.WrapFunc(func(s rbxmk.State) int {
		v := s.CheckAny(1)
		msg := s.OptString(2, "expected fail")
		switch v := v.(type) {
		case *lua.LFunction:
			n := s.Count()
			s.L.Push(v)
			if err := s.L.PCall(0, lua.MultRet, nil); err == nil {
				if s.Count() > n {
					if s.L.ToBool(n + 1) {
						if m := s.L.ToString(n + 2); m != "" {
							t.Errorf("%s: %s", msg, m)
						} else {
							t.Errorf(msg)
						}
					}
					return 0
				}
				t.Errorf(msg)
				return 0
			} else if lua.LVAsBool(s.L.GetGlobal("SHOW_ERRORS")) {
				t.Logf("ERROR: %s\n", err)
			}
		default:
			if lua.LVAsBool(v) {
				t.Errorf(msg)
				return 0
			}
		}
		return 0
	}))

	// Equal asserts whether two values are deeply equal.
	T.RawSetString("Equal", s.L.NewFunction(func(l *lua.LState) int {
		msg := s.CheckString(1)
		a := s.CheckAny(2)
		b := s.CheckAny(3)
		deepeq(t, s, msg, a, b)
		return 0
	}))

	// GC runs the garbage collector.
	T.RawSetString("GC", s.L.NewFunction(func(l *lua.LState) int {
		runtime.GC()
		return 0
	}))

	// UserDataCacheLen returns the number of cached userdata values.
	T.RawSetString("UserDataCacheLen", s.WrapFunc(func(s rbxmk.State) int {
		return s.Push(types.Int(s.UserDataCacheLen()))
	}))

	// UserDataID returns a string that identifies a userdata value.
	T.RawSetString("UserDataID", s.WrapFunc(func(s rbxmk.State) int {
		return s.Push(types.String(fmt.Sprintf("%p", s.CheckUserData(1))))
	}))

	T.RawSetString("DummySymbol", s.UserDataOf(rtypes.Symbol{Name: "DummySymbol"}, "Symbol"))

	s.L.SetGlobal("T", T)
}

// TestScripts runs each .lua file in testdata as a Lua script. If the first
// line starts with a comment that contains "fail", then the script is expected
// to throw an error. All scripts receive the arguments from scriptArguments.
func TestScripts(t *testing.T) {
	program := snek.NewProgram("", scriptArguments[:])
	program.Register(snek.Def{
		Name: "run",
		New: func() snek.Command {
			return &RunCommand{Init: func(s rbxmk.State) { initMain(s, t) }}
		},
	})

	var files []string
	wd, _ := os.Getwd()
	for _, arg := range os.Args[2:] {
		if strings.HasPrefix(arg, "-test.") {
			continue
		}
		rel, err := filepath.Rel(wd, arg)
		if err != nil {
			rel = arg
		}
		files = append(files, rel)
	}
	if len(files) == 0 {
		err := filepath.Walk(testdata, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() &&
				filepath.Ext(info.Name()) == ".lua" &&
				!strings.HasPrefix(filepath.Base(info.Name()), "_") {
				files = append(files, path)
			}
			return nil
		})
		if err != nil {
			t.Fatalf("error walking testdata: %s", err)
		}
	}
	for _, file := range files {
		t.Run(filepath.ToSlash(file), func(t *testing.T) {
			args := make([]string, len(scriptArguments))
			copy(args, scriptArguments[:])
			args[replaceIndex] = file
			err := program.RunWithInput("run", snek.Input{
				Program:   args[0],
				Arguments: args[2:],
			})
			if err != nil {
				t.Errorf("script %s: %s", file, err)
			}
		})
	}
}
<|endoftext|>