<commit_before>package sweetiebot

import (
	"fmt"
	"strings"

	"github.com/bwmarrin/discordgo"
)

type GroupsModule struct {
}

func (w *GroupsModule) Name() string {
	return "Groups"
}

func (w *GroupsModule) Register(info *GuildInfo) {}

func (w *GroupsModule) Commands() []Command {
	return []Command{
		&AddGroupCommand{},
		&JoinGroupCommand{},
		&ListGroupCommand{},
		&LeaveGroupCommand{},
		&PingCommand{},
		&PurgeGroupCommand{},
	}
}

func (w *GroupsModule) Description() string {
	return "Contains commands for manipulating groups and pinging them."
}

type AddGroupCommand struct {
}

func (c *AddGroupCommand) Name() string {
	return "AddGroup"
}

func (c *AddGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to name the group!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if ok {
		return "```That group already exists!```", false, nil
	}

	if len(info.config.Basic.Groups) <= 0 {
		info.config.Basic.Groups = make(map[string]map[string]bool)
	}
	group := make(map[string]bool)
	group[msg.Author.ID] = true
	info.config.Basic.Groups[arg] = group
	info.SaveConfig()

	return "```Successfully created the " + arg + " group! Join it using !joingroup " + arg + " and ping it using !ping " + arg + ".```", false, nil
}
func (c *AddGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Creates a new group and automatically adds you to it. Groups are automatically destroyed when everyone in the group leaves.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "name", Desc: "Name of the new group. Should not contain spaces or anything other than letters and numbers.", Optional: false},
		},
	}
}
func (c *AddGroupCommand) UsageShort() string { return "Creates a new group." }

type JoinGroupCommand struct {
}

func (c *JoinGroupCommand) Name() string {
	return "JoinGroup"
}

func (c *JoinGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup to list existing groups.```", false, nil
	}

	info.config.Basic.Groups[arg][msg.Author.ID] = true
	info.SaveConfig()

	return "```Successfully joined the " + arg + " group! Ping it using !ping " + arg + " or leave it using !leavegroup " + arg + ". WARNING: Pinging a group will ping EVERYONE IN THE GROUP.```", false, nil
}
func (c *JoinGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Joins an existing group.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to join (case-insensitive).", Optional: false},
		},
	}
}
func (c *JoinGroupCommand) UsageShort() string { return "Joins an existing group." }

type ListGroupCommand struct {
}

func (c *ListGroupCommand) Name() string {
	return "ListGroup"
}

func (c *ListGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		if len(info.config.Basic.Groups) <= 0 {
			return "```No groups to list!```", false, nil
		}
		keys := make([]string, len(info.config.Basic.Groups))

		i := 0
		for k := range info.config.Basic.Groups {
			keys[i] = k
			i++
		}

		return "```\n" + strings.Join(keys, ", ") + "```", false, nil
	}

	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup with no arguments to list existing groups.```", false, nil
	}

	pings := make([]string, len(info.config.Basic.Groups[arg]))

	i := 0
	for k := range info.config.Basic.Groups[arg] {
		m, _, _, _ := sb.db.GetUser(SBatoi(k))
		if m != nil {
			pings[i] = m.Username
		}
		i++
	}

	return "```\n" + strings.Join(pings, ", ") + "```", false, nil
}
func (c *ListGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Lists all current groups, or lists all the members of a group.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to display. If omitted, will display all groups instead.", Optional: true},
		},
	}
}
func (c *ListGroupCommand) UsageShort() string { return "Lists all groups." }

type LeaveGroupCommand struct {
}

func (c *LeaveGroupCommand) Name() string {
	return "LeaveGroup"
}

func (c *LeaveGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup to list existing groups.```", false, nil
	}

	_, ok = info.config.Basic.Groups[arg][msg.Author.ID]
	if !ok {
		return "```You aren't in that group!```", false, nil
	}

	delete(info.config.Basic.Groups[arg], msg.Author.ID)

	if len(info.config.Basic.Groups[arg]) <= 0 {
		delete(info.config.Basic.Groups, arg)
	}

	info.SaveConfig()

	return "```You have been removed from " + arg + "```", false, nil
}
func (c *LeaveGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Removes you from the given group, if you are a member of it.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to leave.", Optional: false},
		},
	}
}
func (c *LeaveGroupCommand) UsageShort() string { return "Removes you from a group." }

func getGroupPings(groups []string, info *GuildInfo) string {
	if len(groups) == 0 {
		return ""
	}
	union := make(map[string]bool)
	for _, group := range groups {
		for k, v := range info.config.Basic.Groups[group] {
			union[k] = v
		}
	}
	pings := make([]string, len(union), len(union))

	i := 0
	for k := range union {
		pings[i] = SBitoa(SBatoi(k)) // We convert to integers and then back to strings to prevent bloons from fucking with the bot
		i++
	}

	return "<@" + strings.Join(pings, "> <@") + ">"
}

type PingCommand struct {
}

func (c *PingCommand) Name() string {
	return "Ping"
}

func (c *PingCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	nargs := strings.SplitN(args[0], "\n", 2)
	args = append(nargs, args[1:]...)
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	m := ""
	if len(indices) > 1 {
		m = msg.Content[indices[1]:]
	}

	if !ok {
		groups := strings.Split(arg, "+")
		for _, v := range groups {
			_, ok = info.config.Basic.Groups[v]
			if !ok {
				return fmt.Sprintf("```The %s group doesn't exist! Use !listgroup to list existing groups.```", v), false, nil
			}
			_, ok = info.config.Basic.Groups[v][msg.Author.ID]
			if !ok {
				return fmt.Sprintf("```You aren't a member of %s. You can only ping groups you are a member of.```", v), false, nil
			}
		}
		sb.dg.ChannelMessageSend(msg.ChannelID, arg+": "+getGroupPings(groups, info)+" "+info.SanitizeOutput(m))

	} else {
		_, ok = info.config.Basic.Groups[arg][msg.Author.ID]
		if !ok {
			return "```You can only ping groups you are a member of.```", false, nil
		}
		sb.dg.ChannelMessageSend(msg.ChannelID, arg+": "+getGroupPings([]string{arg}, info)+" "+info.SanitizeOutput(m))
	}
	return "", false, nil
}
func (c *PingCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Pings everyone in a group with the given message, but only if you are a member of the group.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to ping. You can ping multiple groups at the same time by using `group1+group2`", Optional: false},
			CommandUsageParam{Name: "arbitrary string", Desc: "String for Sweetiebot to echo to the group, no spaces required.", Optional: false},
		},
	}
}
func (c *PingCommand) UsageShort() string { return "Pings a group." }

type PurgeGroupCommand struct {
}

func (c *PurgeGroupCommand) Name() string {
	return "PurgeGroup"
}

func (c *PurgeGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup to list existing groups.```", false, nil
	}

	delete(info.config.Basic.Groups, arg)
	info.SaveConfig()

	return "```Deleted " + arg + "```", false, nil
}
func (c *PurgeGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Deletes the group, if it exists.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to delete.", Optional: false},
		},
	}
}
func (c *PurgeGroupCommand) UsageShort() string { return "Deletes a group." }
<commit_msg>Satisfy OCD<commit_after>package sweetiebot

import (
	"fmt"
	"strings"

	"github.com/bwmarrin/discordgo"
)

type GroupsModule struct {
}

func (w *GroupsModule) Name() string {
	return "Groups"
}

func (w *GroupsModule) Register(info *GuildInfo) {}

func (w *GroupsModule) Commands() []Command {
	return []Command{
		&AddGroupCommand{},
		&JoinGroupCommand{},
		&ListGroupCommand{},
		&LeaveGroupCommand{},
		&PingCommand{},
		&PurgeGroupCommand{},
	}
}

func (w *GroupsModule) Description() string {
	return "Contains commands for manipulating groups and pinging them."
}

type AddGroupCommand struct {
}

func (c *AddGroupCommand) Name() string {
	return "AddGroup"
}

func (c *AddGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to name the group!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if ok {
		return "```That group already exists!```", false, nil
	}

	if len(info.config.Basic.Groups) <= 0 {
		info.config.Basic.Groups = make(map[string]map[string]bool)
	}
	group := make(map[string]bool)
	group[msg.Author.ID] = true
	info.config.Basic.Groups[arg] = group
	info.SaveConfig()

	return "```Successfully created the " + arg + " group! Join it using !joingroup " + arg + " and ping it using !ping " + arg + ".```", false, nil
}
func (c *AddGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Creates a new group and automatically adds you to it. Groups are automatically destroyed when everyone in the group leaves.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "name", Desc: "Name of the new group. Should not contain spaces or anything other than letters and numbers.", Optional: false},
		},
	}
}
func (c *AddGroupCommand) UsageShort() string { return "Creates a new group." }

type JoinGroupCommand struct {
}

func (c *JoinGroupCommand) Name() string {
	return "JoinGroup"
}

func (c *JoinGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup to list existing groups.```", false, nil
	}

	info.config.Basic.Groups[arg][msg.Author.ID] = true
	info.SaveConfig()

	return "```Successfully joined the " + arg + " group! Ping it using !ping " + arg + " or leave it using !leavegroup " + arg + ". WARNING: Pinging a group will ping EVERYONE IN THE GROUP.```", false, nil
}
func (c *JoinGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Joins an existing group.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to join (case-insensitive).", Optional: false},
		},
	}
}
func (c *JoinGroupCommand) UsageShort() string { return "Joins an existing group." }

type ListGroupCommand struct {
}

func (c *ListGroupCommand) Name() string {
	return "ListGroup"
}

func (c *ListGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		if len(info.config.Basic.Groups) <= 0 {
			return "```No groups to list!```", false, nil
		}
		keys := make([]string, len(info.config.Basic.Groups))

		i := 0
		for k := range info.config.Basic.Groups {
			keys[i] = k
			i++
		}

		return "```\n" + strings.Join(keys, ", ") + "```", false, nil
	}

	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup with no arguments to list existing groups.```", false, nil
	}

	pings := make([]string, len(info.config.Basic.Groups[arg]))

	i := 0
	for k := range info.config.Basic.Groups[arg] {
		m, _, _, _ := sb.db.GetUser(SBatoi(k))
		if m != nil {
			pings[i] = m.Username
		}
		i++
	}

	return "```\n" + strings.Join(pings, ", ") + "```", false, nil
}
func (c *ListGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Lists all current groups, or lists all the members of a group.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to display. If omitted, will display all groups instead.", Optional: true},
		},
	}
}
func (c *ListGroupCommand) UsageShort() string { return "Lists all groups." }

type LeaveGroupCommand struct {
}

func (c *LeaveGroupCommand) Name() string {
	return "LeaveGroup"
}

func (c *LeaveGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup to list existing groups.```", false, nil
	}

	_, ok = info.config.Basic.Groups[arg][msg.Author.ID]
	if !ok {
		return "```You aren't in that group!```", false, nil
	}

	delete(info.config.Basic.Groups[arg], msg.Author.ID)

	if len(info.config.Basic.Groups[arg]) <= 0 {
		delete(info.config.Basic.Groups, arg)
	}

	info.SaveConfig()

	return "```You have been removed from " + arg + "```", false, nil
}
func (c *LeaveGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Removes you from the given group, if you are a member of it.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to leave.", Optional: false},
		},
	}
}
func (c *LeaveGroupCommand) UsageShort() string { return "Removes you from a group." }

func getGroupPings(groups []string, info *GuildInfo) string {
	if len(groups) == 0 {
		return ""
	}
	union := make(map[string]bool)
	for _, group := range groups {
		for k, v := range info.config.Basic.Groups[group] {
			union[k] = v
		}
	}
	pings := make([]string, len(union), len(union))

	i := 0
	for k := range union {
		pings[i] = SBitoa(SBatoi(k)) // We convert to integers and then back to strings to prevent bloons from fucking with the bot
		i++
	}

	return "<@" + strings.Join(pings, "> <@") + ">"
}

type PingCommand struct {
}

func (c *PingCommand) Name() string {
	return "Ping"
}

func (c *PingCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	nargs := strings.SplitN(args[0], "\n", 2)
	args = append(nargs, args[1:]...)
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	m := ""
	if len(indices) > 1 {
		m = msg.Content[indices[1]:]
	}

	var groups []string
	if !ok {
		groups = strings.Split(arg, "+")
	} else {
		groups = []string{arg}
	}

	for _, v := range groups {
		_, ok = info.config.Basic.Groups[v]
		if !ok {
			return fmt.Sprintf("```The %s group doesn't exist! Use !listgroup to list existing groups.```", v), false, nil
		}
		_, ok = info.config.Basic.Groups[v][msg.Author.ID]
		if !ok {
			return fmt.Sprintf("```You aren't a member of %s. You can only ping groups you are a member of.```", v), false, nil
		}
	}
	sb.dg.ChannelMessageSend(msg.ChannelID, arg+": "+getGroupPings(groups, info)+" "+info.SanitizeOutput(m))

	return "", false, nil
}
func (c *PingCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Pings everyone in a group with the given message, but only if you are a member of the group.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to ping. You can ping multiple groups at the same time by using `group1+group2`", Optional: false},
			CommandUsageParam{Name: "arbitrary string", Desc: "String for Sweetiebot to echo to the group, no spaces required.", Optional: false},
		},
	}
}
func (c *PingCommand) UsageShort() string { return "Pings a group." }

type PurgeGroupCommand struct {
}

func (c *PurgeGroupCommand) Name() string {
	return "PurgeGroup"
}

func (c *PurgeGroupCommand) Process(args []string, msg *discordgo.Message, indices []int, info *GuildInfo) (string, bool, *discordgo.MessageEmbed) {
	if len(args) < 1 {
		return "```You have to provide a group name!```", false, nil
	}
	arg := strings.TrimSpace(strings.ToLower(args[0]))
	_, ok := info.config.Basic.Groups[arg]
	if !ok {
		return "```That group doesn't exist! Use !listgroup to list existing groups.```", false, nil
	}

	delete(info.config.Basic.Groups, arg)
	info.SaveConfig()

	return "```Deleted " + arg + "```", false, nil
}
func (c *PurgeGroupCommand) Usage(info *GuildInfo) *CommandUsage {
	return &CommandUsage{
		Desc: "Deletes the group, if it exists.",
		Params: []CommandUsageParam{
			CommandUsageParam{Name: "group", Desc: "Name of the group to delete.", Optional: false},
		},
	}
}
func (c *PurgeGroupCommand) UsageShort() string { return "Deletes a group." }
<|endoftext|>


<commit_before>package statemachine

import (
	"fmt"
	"sync"
	"time"

	"github.com/lytics/metafora"
)

// StateCode is the actual state key. The State struct adds additional metadata
// related to certain StateCodes.
type StateCode string

const (
	Runnable  StateCode = "runnable"  // Scheduled
	Sleeping  StateCode = "sleeping"  // Scheduled, not running until time has elapsed
	Completed StateCode = "completed" // Terminal, not scheduled
	Killed    StateCode = "killed"    // Terminal, not scheduled
	Failed    StateCode = "failed"    // Terminal, not scheduled
	Fault     StateCode = "fault"     // Scheduled, in error handling / retry logic
	Paused    StateCode = "paused"    // Scheduled, not running
)

// Terminal states will never run and cannot transition to a non-terminal
// state.
func (s StateCode) Terminal() bool {
	switch s {
	case Runnable, Sleeping, Paused, Fault:
		return false
	case Completed, Killed, Failed:
		return true
	default:
		panic("unknown state " + s)
	}
}

// State represents the current state of a stateful handler. See StateCode for
// details. Until and Errors are extra state used by the Sleeping and Fault
// states respectively.
type State struct {
	Code   StateCode  `json:"state"`
	Until  *time.Time `json:"until,omitempty"`
	Errors []Err      `json:"errors,omitempty"`
}

// copy state so mutations to Until and Errors aren't shared.
func (s *State) copy() *State {
	ns := &State{Code: s.Code}
	if s.Until != nil {
		until := *s.Until
		ns.Until = &until
	}
	for i := range s.Errors {
		ns.Errors = append(ns.Errors, s.Errors[i])
	}
	return ns
}

func (s State) String() string {
	switch s.Code {
	case Sleeping:
		return fmt.Sprintf("%s until %s", s.Code, s.Until)
	case Fault:
		return fmt.Sprintf("%s (%d errors)", s.Code, len(s.Errors))
	default:
		return string(s.Code)
	}
}

// Messages are events that cause state transitions. Until and Err are used by
// the Sleep and Error messages respectively.
type Message struct {
	Code MessageCode `json:"message"`

	// Until is when the statemachine should transition from sleeping to runnable
	Until *time.Time `json:"until,omitempty"`

	// Err is the error that caused this Error message
	Err error `json:"error,omitempty"`
}

// Valid returns true if the Message is valid. Invalid messages sent as
// commands are discarded by the state machine.
func (m Message) Valid() bool {
	switch m.Code {
	case Run, Pause, Release, Checkpoint, Complete, Kill:
		return true
	case Sleep:
		return m.Until != nil
	case Error:
		return m.Err != nil
	default:
		return false
	}
}

func (m Message) String() string {
	switch m.Code {
	case Sleep:
		if m.Until != nil {
			return fmt.Sprintf("%s until %s", m.Code, m.Until)
		}
	case Error:
		if m.Err != nil {
			return fmt.Sprintf("%s: %s", m.Code, m.Err.Error())
		}
	}
	return string(m.Code)
}

// MessageCode is the symbolic name of a state transition.
type MessageCode string

const (
	Run        MessageCode = "run"
	Sleep      MessageCode = "sleep"
	Pause      MessageCode = "pause"
	Kill       MessageCode = "kill"
	Error      MessageCode = "error"
	Complete   MessageCode = "complete"
	Checkpoint MessageCode = "checkpoint"

	// Special event which triggers state machine to exit without transitioning
	// between states.
	Release MessageCode = "release"
)

// Transitions represent a state machine transition from one state to another
// given an event message.
type Transition struct {
	Event MessageCode
	From  StateCode
	To    StateCode
}

func (t Transition) String() string {
	return fmt.Sprintf("%v---%v--->%v", t.From, t.Event, t.To)
}

var (
	// Rules is the state transition table.
	Rules = [...]Transition{
		// Runnable can transition to anything
		{Event: Checkpoint, From: Runnable, To: Runnable},
		{Event: Release, From: Runnable, To: Runnable},
		{Event: Sleep, From: Runnable, To: Sleeping},
		{Event: Complete, From: Runnable, To: Completed},
		{Event: Kill, From: Runnable, To: Killed},
		{Event: Error, From: Runnable, To: Fault},
		{Event: Pause, From: Runnable, To: Paused},

		// Sleeping can return to Runnable or be Killed/Paused
		{Event: Checkpoint, From: Sleeping, To: Sleeping},
		{Event: Release, From: Sleeping, To: Sleeping},
		{Event: Sleep, From: Sleeping, To: Sleeping},
		{Event: Run, From: Sleeping, To: Runnable},
		{Event: Kill, From: Sleeping, To: Killed},
		{Event: Pause, From: Sleeping, To: Paused},
		{Event: Error, From: Sleeping, To: Fault},

		// The error state transitions to either sleeping, failed, or released (to
		// allow custom error handlers to workaround localitly related errors).
		{Event: Sleep, From: Fault, To: Sleeping},
		{Event: Error, From: Fault, To: Failed},

		// Paused can return to Runnable, be put to Sleep, or Killed
		{Event: Checkpoint, From: Paused, To: Paused},
		{Event: Release, From: Paused, To: Paused},
		{Event: Run, From: Paused, To: Runnable},
		{Event: Sleep, From: Paused, To: Sleeping},
		{Event: Kill, From: Paused, To: Killed},
		{Event: Pause, From: Paused, To: Paused},

		// Completed, Failed, and Killed are terminal states that cannot transition
		// to anything.
	}
)

// StatefulHandler is the function signature that the state machine is able to
// run. Instead of metafora.Handler's Stop method, StatefulHandlers receive
// Messages via the commands chan and return their exit status via a Message.
//
// Normally StatefulHandlers simply return a Message as soon as it's received
// on the commands chan. However, it's also acceptable for a handler to return
// a different Message. For example if it encounters an error during shutdown,
// it may choose to return that error as an Error Message as opposed to the
// original command.
type StatefulHandler func(taskID string, commands <-chan Message) Message

type stateMachine struct {
	taskID     string
	h          StatefulHandler
	ss         StateStore
	cl         CommandListener
	cmds       chan Message
	errHandler ErrHandler

	mu    *sync.RWMutex
	state *State
	ts    time.Time

	stopL   *sync.Mutex
	stopped chan bool
}

// New handler that creates a state machine and exposes state transitions to
// the given handler by calling its Transition method. It should be created in
// the HandlerFunc you use with metafora's Consumer.
//
// If ErrHandler is nil DefaultErrHandler will be used.
func New(tid string, h StatefulHandler, ss StateStore, cl CommandListener, e ErrHandler) metafora.Handler {
	if e == nil {
		e = DefaultErrHandler
	}
	return &stateMachine{
		taskID:     tid,
		h:          h,
		ss:         ss,
		cl:         cl,
		errHandler: e,
		mu:         &sync.RWMutex{},
		ts:         time.Now(),
		stopL:      &sync.Mutex{},
		stopped:    make(chan bool),
	}
}

// State returns the current state the state machine is in and what time it
// entered that state. The State may be nil if Run() has yet to be called.
func (s *stateMachine) State() (*State, time.Time) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.state, s.ts
}

func (s *stateMachine) setState(state *State) {
	s.mu.Lock()
	s.state = state.copy()
	s.ts = time.Now()
	s.mu.Unlock()
}

// Run the state machine enabled handler. Loads the initial state and passes
// control to the internal stateful handler passing commands from the command
// listener into the handler's commands chan.
func (s *stateMachine) Run() (done bool) {
	// Multiplex external (Stop) messages and internal ones
	s.cmds = make(chan Message)
	go func() {
		for {
			select {
			case m := <-s.cl.Receive():
				if !m.Valid() {
					metafora.Warnf("Ignoring invalid command: %q", m)
					continue
				}
				select {
				case s.cmds <- m:
				case <-s.stopped:
					return
				}
			case <-s.stopped:
				return
			}
		}
	}()

	// Stop the command listener and internal message multiplexer when Run exits
	defer func() {
		s.cl.Stop()
		s.stop()
	}()

	// Load the initial state
	state, err := s.ss.Load(s.taskID)
	if err != nil {
		// A failure to load the state for a task is *fatal* - the task will be
		// unscheduled and requires operator intervention to reschedule.
		metafora.Errorf("task=%q could not load initial state. Marking done! Error: %v", s.taskID, err)
		return true
	}
	if state.Code.Terminal() {
		metafora.Warnf("task=%q in terminal state %s - exiting.", s.taskID, state.Code)
		return true
	}
	s.setState(state)

	// Main Statemachine Loop
	done = false
	for {
		// Enter State
		metafora.Debugf("task=%q in state %s", s.taskID, state.Code)
		msg := s.exec(state)

		// Apply Message
		newstate, ok := apply(state, msg)
		if !ok {
			metafora.Warnf("task=%q Invalid state transition=%q returned by task. Old state=%q", s.taskID, msg.Code, state.Code)
			msg = Message{Code: Error, Err: err}
			if newstate, ok = apply(state, msg); !ok {
				metafora.Errorf("task=%q Unable to transition to error state! Exiting with state=%q", s.taskID, state.Code)
				return state.Code.Terminal()
			}
		}

		metafora.Infof("task=%q transitioning %s --> %s --> %s", s.taskID, state, msg, newstate)

		// Save state
		if err := s.ss.Store(s.taskID, newstate); err != nil {
			metafora.Errorf("task=%q Unable to persist state=%q. Unscheduling.", s.taskID, newstate.Code)
			return true
		}

		// Set next state and loop if non-terminal
		state = newstate

		// Expose the state for introspection
		s.setState(state)

		// Exit and unschedule task on terminal state.
		if state.Code.Terminal() {
			return true
		}

		// Release messages indicate the task should exit but not unschedule.
		if msg.Code == Release {
			return false
		}

		// Alternatively Stop() may have been called but the handler may not have
		// returned the Release message. Always exit if we've been told to Stop()
		// even if the handler has returned a different Message.
		select {
		case <-s.stopped:
			return false
		default:
		}
	}
}

// execute non-terminal states
func (s *stateMachine) exec(state *State) Message {
	switch state.Code {
	case Runnable:
		// Runnable passes control to the stateful handler
		return run(s.h, s.taskID, s.cmds)
	case Paused:
		// Paused until a message arrives
		return <-s.cmds
	case Sleeping:
		// Sleeping until the specified time (or a message)
		if state.Until == nil {
			metafora.Warnf("task=%q told to sleep without a time. Resuming.", s.taskID)
			return Message{Code: Run}
		}
		dur := state.Until.Sub(time.Now())
		metafora.Infof("task=%q sleeping for %s", s.taskID, dur)
		timer := time.NewTimer(dur)
		select {
		case <-timer.C:
			return Message{Code: Run}
		case msg := <-s.cmds:
			timer.Stop()
			// Checkpoint & Release are special cases that shouldn't affect sleep
			// time, so maintain it across the state transition
			if msg.Code == Checkpoint || msg.Code == Release {
				msg.Until = state.Until
			}
			return msg
		}
	case Fault:
		// Special case where we potentially trim the current state to keep
		// errors from growing without bound.
		var msg Message
		msg, state.Errors = s.errHandler(s.taskID, state.Errors)
		return msg
	default:
		panic("invalid state: " + state.String())
	}
}

func run(f StatefulHandler, tid string, cmd <-chan Message) (m Message) {
	defer func() {
		if r := recover(); r != nil {
			metafora.Errorf("task=%q Run method panic()d! Applying Error message. Panic: %v", tid, r)
			m = Message{Code: Error, Err: fmt.Errorf("panic: %v", r)}
		}
	}()

	// Defensive code to give handlers a *copy* of the command chan. That way if
	// a handler keeps receiving on the command chan in a goroutine past the
	// handler's lifetime it doesn't intercept commands intended for the
	// statemachine.
	internalcmd := make(chan Message)
	stopped := make(chan struct{})
	go func() {
		for {
			select {
			case c := <-cmd:
				internalcmd <- c
			case <-stopped:
				return
			}
		}
	}()
	defer close(stopped)

	return f(tid, internalcmd)
}

// Stop sends a Release message to the state machine through the command chan.
func (s *stateMachine) Stop() {
	s.cmds <- Message{Code: Release}

	// Also inform the state machine it should exit since the internal handler
	// may override the release message causing the task to be unreleaseable.
	s.stop()
}

func (s *stateMachine) stop() {
	s.stopL.Lock()
	defer s.stopL.Unlock()
	select {
	case <-s.stopped:
		return
	default:
		close(s.stopped)
	}
}

// apply a message to cause a state transition. Returns false if the state
// transition is invalid.
func apply(cur *State, m Message) (*State, bool) {
	//XXX Is a linear scan of all rules really the best option here?
	for _, trans := range Rules {
		if trans.Event == m.Code && trans.From == cur.Code {
			metafora.Debugf("Transitioned %s", trans)
			if m.Err != nil {
				// Append errors from message
				cur.Errors = append(cur.Errors, Err{Time: time.Now(), Err: m.Err.Error()})
			}

			// New State + Message's Until + Combined Errors
			return &State{Code: trans.To, Until: m.Until, Errors: cur.Errors}, true
		}
	}
	return cur, false
}
<commit_msg>Add some helper functions for commonly used messages<commit_after>package statemachine

import (
	"fmt"
	"sync"
	"time"

	"github.com/lytics/metafora"
)

// StateCode is the actual state key. The State struct adds additional metadata
// related to certain StateCodes.
type StateCode string

const (
	Runnable  StateCode = "runnable"  // Scheduled
	Sleeping  StateCode = "sleeping"  // Scheduled, not running until time has elapsed
	Completed StateCode = "completed" // Terminal, not scheduled
	Killed    StateCode = "killed"    // Terminal, not scheduled
	Failed    StateCode = "failed"    // Terminal, not scheduled
	Fault     StateCode = "fault"     // Scheduled, in error handling / retry logic
	Paused    StateCode = "paused"    // Scheduled, not running
)

// Terminal states will never run and cannot transition to a non-terminal
// state.
func (s StateCode) Terminal() bool {
	switch s {
	case Runnable, Sleeping, Paused, Fault:
		return false
	case Completed, Killed, Failed:
		return true
	default:
		panic("unknown state " + s)
	}
}

// State represents the current state of a stateful handler. See StateCode for
// details. Until and Errors are extra state used by the Sleeping and Fault
// states respectively.
type State struct {
	Code   StateCode  `json:"state"`
	Until  *time.Time `json:"until,omitempty"`
	Errors []Err      `json:"errors,omitempty"`
}

// copy state so mutations to Until and Errors aren't shared.
func (s *State) copy() *State {
	ns := &State{Code: s.Code}
	if s.Until != nil {
		until := *s.Until
		ns.Until = &until
	}
	for i := range s.Errors {
		ns.Errors = append(ns.Errors, s.Errors[i])
	}
	return ns
}

func (s *State) String() string {
	switch s.Code {
	case Sleeping:
		return fmt.Sprintf("%s until %s", s.Code, s.Until)
	case Fault:
		return fmt.Sprintf("%s (%d errors)", s.Code, len(s.Errors))
	default:
		return string(s.Code)
	}
}

// Messages are events that cause state transitions. Until and Err are used by
// the Sleep and Error messages respectively.
type Message struct {
	Code MessageCode `json:"message"`

	// Until is when the statemachine should transition from sleeping to runnable
	Until *time.Time `json:"until,omitempty"`

	// Err is the error that caused this Error message
	Err error `json:"error,omitempty"`
}

// ErrorMessage is a simpler helper for creating error messages from an error.
func ErrorMessage(err error) Message {
	return Message{Code: Error, Err: err}
}

// SleepMessage is a simpler helper for creating sleep messages from a time.
func SleepMessage(t time.Time) Message {
	return Message{Code: Sleep, Until: &t}
}

// Valid returns true if the Message is valid. Invalid messages sent as
// commands are discarded by the state machine.
func (m *Message) Valid() bool {
	switch m.Code {
	case Run, Pause, Release, Checkpoint, Complete, Kill:
		return true
	case Sleep:
		return m.Until != nil
	case Error:
		return m.Err != nil
	default:
		return false
	}
}

func (m *Message) String() string {
	switch m.Code {
	case Sleep:
		if m.Until != nil {
			return fmt.Sprintf("%s until %s", m.Code, m.Until)
		}
	case Error:
		if m.Err != nil {
			return fmt.Sprintf("%s: %s", m.Code, m.Err.Error())
		}
	}
	return string(m.Code)
}

// MessageCode is the symbolic name of a state transition.
type MessageCode string

const (
	Run        MessageCode = "run"
	Sleep      MessageCode = "sleep"
	Pause      MessageCode = "pause"
	Kill       MessageCode = "kill"
	Error      MessageCode = "error"
	Complete   MessageCode = "complete"
	Checkpoint MessageCode = "checkpoint"

	// Special event which triggers state machine to exit without transitioning
	// between states.
	Release MessageCode = "release"
)

// Transitions represent a state machine transition from one state to another
// given an event message.
type Transition struct {
	Event MessageCode
	From  StateCode
	To    StateCode
}

func (t Transition) String() string {
	return fmt.Sprintf("%v---%v--->%v", t.From, t.Event, t.To)
}

var (
	// Rules is the state transition table.
	Rules = [...]Transition{
		// Runnable can transition to anything
		{Event: Checkpoint, From: Runnable, To: Runnable},
		{Event: Release, From: Runnable, To: Runnable},
		{Event: Sleep, From: Runnable, To: Sleeping},
		{Event: Complete, From: Runnable, To: Completed},
		{Event: Kill, From: Runnable, To: Killed},
		{Event: Error, From: Runnable, To: Fault},
		{Event: Pause, From: Runnable, To: Paused},

		// Sleeping can return to Runnable or be Killed/Paused
		{Event: Checkpoint, From: Sleeping, To: Sleeping},
		{Event: Release, From: Sleeping, To: Sleeping},
		{Event: Sleep, From: Sleeping, To: Sleeping},
		{Event: Run, From: Sleeping, To: Runnable},
		{Event: Kill, From: Sleeping, To: Killed},
		{Event: Pause, From: Sleeping, To: Paused},
		{Event: Error, From: Sleeping, To: Fault},

		// The error state transitions to either sleeping, failed, or released (to
		// allow custom error handlers to workaround localitly related errors).
		{Event: Sleep, From: Fault, To: Sleeping},
		{Event: Error, From: Fault, To: Failed},

		// Paused can return to Runnable, be put to Sleep, or Killed
		{Event: Checkpoint, From: Paused, To: Paused},
		{Event: Release, From: Paused, To: Paused},
		{Event: Run, From: Paused, To: Runnable},
		{Event: Sleep, From: Paused, To: Sleeping},
		{Event: Kill, From: Paused, To: Killed},
		{Event: Pause, From: Paused, To: Paused},

		// Completed, Failed, and Killed are terminal states that cannot transition
		// to anything.
	}
)

// StatefulHandler is the function signature that the state machine is able to
// run. Instead of metafora.Handler's Stop method, StatefulHandlers receive
// Messages via the commands chan and return their exit status via a Message.
//
// Normally StatefulHandlers simply return a Message as soon as it's received
// on the commands chan. However, it's also acceptable for a handler to return
// a different Message. For example if it encounters an error during shutdown,
// it may choose to return that error as an Error Message as opposed to the
// original command.
type StatefulHandler func(taskID string, commands <-chan Message) Message

type stateMachine struct {
	taskID     string
	h          StatefulHandler
	ss         StateStore
	cl         CommandListener
	cmds       chan Message
	errHandler ErrHandler

	mu    *sync.RWMutex
	state *State
	ts    time.Time

	stopL   *sync.Mutex
	stopped chan bool
}

// New handler that creates a state machine and exposes state transitions to
// the given handler by calling its Transition method. It should be created in
// the HandlerFunc you use with metafora's Consumer.
//
// If ErrHandler is nil DefaultErrHandler will be used.
func New(tid string, h StatefulHandler, ss StateStore, cl CommandListener, e ErrHandler) metafora.Handler {
	if e == nil {
		e = DefaultErrHandler
	}
	return &stateMachine{
		taskID:     tid,
		h:          h,
		ss:         ss,
		cl:         cl,
		errHandler: e,
		mu:         &sync.RWMutex{},
		ts:         time.Now(),
		stopL:      &sync.Mutex{},
		stopped:    make(chan bool),
	}
}

// State returns the current state the state machine is in and what time it
// entered that state. The State may be nil if Run() has yet to be called.
func (s *stateMachine) State() (*State, time.Time) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.state, s.ts
}

func (s *stateMachine) setState(state *State) {
	s.mu.Lock()
	s.state = state.copy()
	s.ts = time.Now()
	s.mu.Unlock()
}

// Run the state machine enabled handler. Loads the initial state and passes
// control to the internal stateful handler passing commands from the command
// listener into the handler's commands chan.
func (s *stateMachine) Run() (done bool) {
	// Multiplex external (Stop) messages and internal ones
	s.cmds = make(chan Message)
	go func() {
		for {
			select {
			case m := <-s.cl.Receive():
				if !m.Valid() {
					metafora.Warnf("Ignoring invalid command: %q", m)
					continue
				}
				select {
				case s.cmds <- m:
				case <-s.stopped:
					return
				}
			case <-s.stopped:
				return
			}
		}
	}()

	// Stop the command listener and internal message multiplexer when Run exits
	defer func() {
		s.cl.Stop()
		s.stop()
	}()

	// Load the initial state
	state, err := s.ss.Load(s.taskID)
	if err != nil {
		// A failure to load the state for a task is *fatal* - the task will be
		// unscheduled and requires operator intervention to reschedule.
		metafora.Errorf("task=%q could not load initial state. Marking done! Error: %v", s.taskID, err)
		return true
	}
	if state.Code.Terminal() {
		metafora.Warnf("task=%q in terminal state %s - exiting.", s.taskID, state.Code)
		return true
	}
	s.setState(state)

	// Main Statemachine Loop
	done = false
	for {
		// Enter State
		metafora.Debugf("task=%q in state %s", s.taskID, state.Code)
		msg := s.exec(state)

		// Apply Message
		newstate, ok := apply(state, msg)
		if !ok {
			metafora.Warnf("task=%q Invalid state transition=%q returned by task. Old state=%q", s.taskID, msg.Code, state.Code)
			msg = Message{Code: Error, Err: err}
			if newstate, ok = apply(state, msg); !ok {
				metafora.Errorf("task=%q Unable to transition to error state! Exiting with state=%q", s.taskID, state.Code)
				return state.Code.Terminal()
			}
		}

		metafora.Infof("task=%q transitioning %s --> %s --> %s", s.taskID, state, msg, newstate)

		// Save state
		if err := s.ss.Store(s.taskID, newstate); err != nil {
			metafora.Errorf("task=%q Unable to persist state=%q. Unscheduling.", s.taskID, newstate.Code)
			return true
		}

		// Set next state and loop if non-terminal
		state = newstate

		// Expose the state for introspection
		s.setState(state)

		// Exit and unschedule task on terminal state.
		if state.Code.Terminal() {
			return true
		}

		// Release messages indicate the task should exit but not unschedule.
		if msg.Code == Release {
			return false
		}

		// Alternatively Stop() may have been called but the handler may not have
		// returned the Release message. Always exit if we've been told to Stop()
		// even if the handler has returned a different Message.
		select {
		case <-s.stopped:
			return false
		default:
		}
	}
}

// execute non-terminal states
func (s *stateMachine) exec(state *State) Message {
	switch state.Code {
	case Runnable:
		// Runnable passes control to the stateful handler
		return run(s.h, s.taskID, s.cmds)
	case Paused:
		// Paused until a message arrives
		return <-s.cmds
	case Sleeping:
		// Sleeping until the specified time (or a message)
		if state.Until == nil {
			metafora.Warnf("task=%q told to sleep without a time. Resuming.", s.taskID)
			return Message{Code: Run}
		}
		dur := state.Until.Sub(time.Now())
		metafora.Infof("task=%q sleeping for %s", s.taskID, dur)
		timer := time.NewTimer(dur)
		select {
		case <-timer.C:
			return Message{Code: Run}
		case msg := <-s.cmds:
			timer.Stop()
			// Checkpoint & Release are special cases that shouldn't affect sleep
			// time, so maintain it across the state transition
			if msg.Code == Checkpoint || msg.Code == Release {
				msg.Until = state.Until
			}
			return msg
		}
	case Fault:
		// Special case where we potentially trim the current state to keep
		// errors from growing without bound.
		var msg Message
		msg, state.Errors = s.errHandler(s.taskID, state.Errors)
		return msg
	default:
		panic("invalid state: " + state.String())
	}
}

func run(f StatefulHandler, tid string, cmd <-chan Message) (m Message) {
	defer func() {
		if r := recover(); r != nil {
			metafora.Errorf("task=%q Run method panic()d! Applying Error message. Panic: %v", tid, r)
			m = Message{Code: Error, Err: fmt.Errorf("panic: %v", r)}
		}
	}()

	// Defensive code to give handlers a *copy* of the command chan. That way if
	// a handler keeps receiving on the command chan in a goroutine past the
	// handler's lifetime it doesn't intercept commands intended for the
	// statemachine.
	internalcmd := make(chan Message)
	stopped := make(chan struct{})
	go func() {
		for {
			select {
			case c := <-cmd:
				internalcmd <- c
			case <-stopped:
				return
			}
		}
	}()
	defer close(stopped)

	return f(tid, internalcmd)
}

// Stop sends a Release message to the state machine through the command chan.
func (s *stateMachine) Stop() {
	s.cmds <- Message{Code: Release}

	// Also inform the state machine it should exit since the internal handler
	// may override the release message causing the task to be unreleaseable.
	s.stop()
}

func (s *stateMachine) stop() {
	s.stopL.Lock()
	defer s.stopL.Unlock()
	select {
	case <-s.stopped:
		return
	default:
		close(s.stopped)
	}
}

// apply a message to cause a state transition. Returns false if the state
// transition is invalid.
func apply(cur *State, m Message) (*State, bool) {
	//XXX Is a linear scan of all rules really the best option here?
	for _, trans := range Rules {
		if trans.Event == m.Code && trans.From == cur.Code {
			metafora.Debugf("Transitioned %s", trans)
			if m.Err != nil {
				// Append errors from message
				cur.Errors = append(cur.Errors, Err{Time: time.Now(), Err: m.Err.Error()})
			}

			// New State + Message's Until + Combined Errors
			return &State{Code: trans.To, Until: m.Until, Errors: cur.Errors}, true
		}
	}
	return cur, false
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	collectors map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

// Collect sends the collected metrics from each of the collectors to
// prometheus. Collect could be called several times concurrently
// and thus its run is protected by a single mutex.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			execute(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
}

func filterAvailableCollectors(collectors string) string {
	var availableCollectors []string
	for _, c := range strings.Split(collectors, ",") {
		_, ok := collector.Factories[c]
		if ok {
			availableCollectors = append(availableCollectors, c)
		}
	}
	return strings.Join(availableCollectors, ",")
}

func execute(name string, c collector.Collector, ch chan<- prometheus.Metric) {
	begin := time.Now()
	err := c.Collect(ch)
	duration := time.Since(begin)
	var success float64

	if err != nil {
		log.Errorf("ERROR: %s collector failed after %fs: %s", name, duration.Seconds(), err)
		success = 0
	} else {
		log.Debugf("OK: %s collector succeeded after %fs.", name, duration.Seconds())
		success = 1
	}
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration.Seconds(),
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeSuccessDesc,
		prometheus.GaugeValue,
		success,
		name,
	)
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s, _ := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func init() {
	prometheus.MustRegister(version.NewCollector("wmi_exporter"))
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.SWbemServicesClient = s
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nNote: If executing from Powershell, the flags need to quoted. For example:\n%s\n",
		"\twmi_exporter \"-collectors.enabled\" iis")
}

func main() {
	var (
		showVersion       = flag.Bool("version", false, "Print version information.")
		listenAddress     = flag.String("telemetry.addr", ":9182", "host:port for WMI exporter.")
		metricsPath       = flag.String("telemetry.path", "/metrics", "URL path for surfacing collected metrics.")
		enabledCollectors = flag.String("collectors.enabled", filterAvailableCollectors(defaultCollectors), "Comma-separated list of collectors to use. Use '[default]' as a placeholder for all the collectors enabled by default")
		printCollectors   = flag.Bool("collectors.print", false, "If true, print available collectors and exit.")
	)
	flag.Usage = usage
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("wmi_exporter"))
		os.Exit(0)
	}

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	nodeCollector := WmiCollector{collectors: collectors}
	prometheus.MustRegister(nodeCollector)

	http.Handle(*metricsPath, promhttp.Handler())
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}
<commit_msg>correct help message, fixes #118. typo was introduced in a66f0b54751e95a83c5c2d50cee375b5411d36fb<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/windows/svc"

	"github.com/StackExchange/wmi"
	"github.com/martinlindhe/wmi_exporter/collector"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

// WmiCollector implements the prometheus.Collector interface.
type WmiCollector struct {
	collectors map[string]collector.Collector
}

const (
	defaultCollectors            = "cpu,cs,logical_disk,net,os,service,system"
	defaultCollectorsPlaceholder = "[defaults]"
	serviceName                  = "wmi_exporter"
)

var (
	scrapeDurationDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_duration_seconds"),
		"wmi_exporter: Duration of a collection.",
		[]string{"collector"},
		nil,
	)
	scrapeSuccessDesc = prometheus.NewDesc(
		prometheus.BuildFQName(collector.Namespace, "exporter", "collector_success"),
		"wmi_exporter: Whether the collector was successful.",
		[]string{"collector"},
		nil,
	)
)

// Describe sends all the descriptors of the collectors included to
// the provided channel.
func (coll WmiCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- scrapeDurationDesc
	ch <- scrapeSuccessDesc
}

// Collect sends the collected metrics from each of the collectors to
// prometheus. Collect could be called several times concurrently
// and thus its run is protected by a single mutex.
func (coll WmiCollector) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(coll.collectors))
	for name, c := range coll.collectors {
		go func(name string, c collector.Collector) {
			execute(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
}

func filterAvailableCollectors(collectors string) string {
	var availableCollectors []string
	for _, c := range strings.Split(collectors, ",") {
		_, ok := collector.Factories[c]
		if ok {
			availableCollectors = append(availableCollectors, c)
		}
	}
	return strings.Join(availableCollectors, ",")
}

func execute(name string, c collector.Collector, ch chan<- prometheus.Metric) {
	begin := time.Now()
	err := c.Collect(ch)
	duration := time.Since(begin)
	var success float64

	if err != nil {
		log.Errorf("ERROR: %s collector failed after %fs: %s", name, duration.Seconds(), err)
		success = 0
	} else {
		log.Debugf("OK: %s collector succeeded after %fs.", name, duration.Seconds())
		success = 1
	}
	ch <- prometheus.MustNewConstMetric(
		scrapeDurationDesc,
		prometheus.GaugeValue,
		duration.Seconds(),
		name,
	)
	ch <- prometheus.MustNewConstMetric(
		scrapeSuccessDesc,
		prometheus.GaugeValue,
		success,
		name,
	)
}

func expandEnabledCollectors(enabled string) []string {
	expanded := strings.Replace(enabled, defaultCollectorsPlaceholder, defaultCollectors, -1)
	separated := strings.Split(expanded, ",")
	unique := map[string]bool{}
	for _, s := range separated {
		if s != "" {
			unique[s] = true
		}
	}
	result := make([]string, 0, len(unique))
	for s, _ := range unique {
		result = append(result, s)
	}
	return result
}

func loadCollectors(list string) (map[string]collector.Collector, error) {
	collectors := map[string]collector.Collector{}
	enabled := expandEnabledCollectors(list)

	for _, name := range enabled {
		fn, ok := collector.Factories[name]
		if !ok {
			return nil, fmt.Errorf("collector '%s' not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		collectors[name] = c
	}
	return collectors, nil
}

func init() {
	prometheus.MustRegister(version.NewCollector("wmi_exporter"))
}

func initWbem() {
	// This initialization prevents a memory leak on WMF 5+. See
	// https://github.com/martinlindhe/wmi_exporter/issues/77 and linked issues
	// for details.
	log.Debugf("Initializing SWbemServices")
	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
	if err != nil {
		log.Fatal(err)
	}
	wmi.DefaultClient.SWbemServicesClient = s
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nNote: If executing from Powershell, the flags need to quoted. For example:\n%s\n",
		"\twmi_exporter \"-collectors.enabled\" iis")
}

func main() {
	var (
		showVersion       = flag.Bool("version", false, "Print version information.")
		listenAddress     = flag.String("telemetry.addr", ":9182", "host:port for WMI exporter.")
		metricsPath       = flag.String("telemetry.path", "/metrics", "URL path for surfacing collected metrics.")
		enabledCollectors = flag.String("collectors.enabled", filterAvailableCollectors(defaultCollectors), "Comma-separated list of collectors to use. Use '[defaults]' as a placeholder for all the collectors enabled by default")
		printCollectors   = flag.Bool("collectors.print", false, "If true, print available collectors and exit.")
	)
	flag.Usage = usage
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("wmi_exporter"))
		os.Exit(0)
	}

	if *printCollectors {
		collectorNames := make(sort.StringSlice, 0, len(collector.Factories))
		for n := range collector.Factories {
			collectorNames = append(collectorNames, n)
		}
		collectorNames.Sort()
		fmt.Printf("Available collectors:\n")
		for _, n := range collectorNames {
			fmt.Printf(" - %s\n", n)
		}
		return
	}

	initWbem()

	isInteractive, err := svc.IsAnInteractiveSession()
	if err != nil {
		log.Fatal(err)
	}

	stopCh := make(chan bool)
	if !isInteractive {
		go svc.Run(serviceName, &wmiExporterService{stopCh: stopCh})
	}

	collectors, err := loadCollectors(*enabledCollectors)
	if err != nil {
		log.Fatalf("Couldn't load collectors: %s", err)
	}

	log.Infof("Enabled collectors: %v", strings.Join(keys(collectors), ", "))

	nodeCollector := WmiCollector{collectors: collectors}
	prometheus.MustRegister(nodeCollector)

	http.Handle(*metricsPath, promhttp.Handler())
	http.HandleFunc("/health", healthCheck)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, *metricsPath, http.StatusMovedPermanently)
	})

	log.Infoln("Starting WMI exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	go func() {
		log.Infoln("Starting server on", *listenAddress)
		log.Fatalf("cannot start WMI exporter: %s", http.ListenAndServe(*listenAddress, nil))
	}()

	for {
		if <-stopCh {
			log.Info("Shutting down WMI exporter")
			break
		}
	}
}

func healthCheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	io.WriteString(w, `{"status":"ok"}`)
}

func keys(m map[string]collector.Collector) []string {
	ret := make([]string, 0, len(m))
	for key := range m {
		ret = append(ret, key)
	}
	return ret
}

type wmiExporterService struct {
	stopCh chan<- bool
}

func (s *wmiExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
	changes <- svc.Status{State: svc.StartPending}
	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
loop:
	for {
		select {
		case c := <-r:
			switch c.Cmd {
			case svc.Interrogate:
				changes <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s.stopCh <- true
				break loop
			default:
				log.Error(fmt.Sprintf("unexpected control request #%d", c))
			}
		}
	}
	changes <- svc.Status{State: svc.StopPending}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"time"
)

const (
	stamp        = "20060102150405"
	dirPerms     = 0755
	filePerms    = 0644
	sqlExtension = ".sql"
)

// startDump copies creation statements for tables, procedures, functions, triggers and views to a file/directory structure at the path location that trite uses in client mode to restore tables.
func startDump(dir string, dbi *mysqlCredentials) {
	dumpdir := path.Join(dir, dbi.host+"_dump"+time.Now().Format(stamp))
	fmt.Println("Dumping to:", dumpdir)
	fmt.Println()

	// Return a database connection
	db, err := dbi.connect()
	defer db.Close()

	// Problem connecting to database
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Turn off idle connections
	db.SetMaxIdleConns(0)

	// Get a list of schemas in the target database
	db.SetMaxIdleConns(1)
	schemas := schemaList(db)

	// Create dump directory
	err = os.MkdirAll(dumpdir, dirPerms)
	checkErr(err)

	// Schema loop
	count := 0
	total := 0
	fmt.Println()
	for _, schema := range schemas {
		total++
		// Dump schema create
		fmt.Print(schema, ": ")
		dumpSchema(db, dumpdir, schema)

		// Dump table creation statements
		count = dumpTables(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " tables, ")

		// Dump procedure creation statements
		count = dumpProcs(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " procedures, ")

		// Dump function creation statements
		count = dumpFuncs(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " functions, ")

		// Dump trigger creation statements
		count = dumpTriggers(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " triggers, ")

		// Dump view creation statements
		count = dumpViews(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " views\n")
	}

	fmt.Println()
	fmt.Println(total, "total objects dumped")
}

// schemaList returns a string slice of schemas to process. MySQL specific schemas like mysql, information_schema and performance_schema are omitted.
func schemaList(db *sql.DB) []string {
	rows, err := db.Query("show databases")
	checkErr(err)

	// Get schema list
	schemas := make([]string, 0)
	var database string
	for rows.Next() {
		err = rows.Scan(&database)
		checkErr(err)

		if database != "mysql" && database != "information_schema" && database != "performance_schema" {
			schemas = append(schemas, database)
		}
	}

	return schemas
}

// dumpSchema creates a file with the schema creation statement.
func dumpSchema(db *sql.DB, dumpdir string, schema string) {
	dir := path.Join(dumpdir, schema)
	var err error

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var ignore string
	var stmt string
	err = db.QueryRow("show create schema "+addQuotes(schema)).Scan(&ignore, &stmt)
	checkErr(err)

	file := path.Join(dir, schema+sqlExtension)
	err = ioutil.WriteFile(file, []byte(stmt+";\n"), filePerms)
	checkErr(err)
}

// dumpTables creates files containing table creation statments. It processes all tables for the schema passed to it. The /tables directory is hardcoded and expected by trite client code.
func dumpTables(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "tables")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select table_name from information_schema.tables where table_schema='" + schema + "' and table_type = 'BASE TABLE'")
	checkErr(err)

	var tableName string
	var ignore string
	var stmt string
	for rows.Next() {
		err = rows.Scan(&tableName)
		checkErr(err)

		err = db.QueryRow("show create table "+addQuotes(schema)+"."+addQuotes(tableName)).Scan(&ignore, &stmt)
		checkErr(err)

		file := path.Join(dir, tableName+sqlExtension)
		err = ioutil.WriteFile(file, []byte(stmt+";\n"), filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpProcs creates files containing procedure creation statments. It processes all procedures for the schema passed to it. The /procedures directory is hardcoded and expected by trite client code.
func dumpProcs(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "procedures")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select routine_name from information_schema.routines where routine_schema='" + schema + "' and routine_type = 'PROCEDURE'")
	checkErr(err)

	var procName string
	for rows.Next() {
		err = rows.Scan(&procName)
		checkErr(err)

		var procInfo createInfoStruct
		err = db.QueryRow("show create procedure "+addQuotes(schema)+"."+addQuotes(procName)).Scan(&procInfo.Name, &procInfo.SqlMode, &procInfo.Create, &procInfo.CharsetClient, &procInfo.Collation, &procInfo.DbCollation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(procInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, procName+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpFuncs creates files containing function creation statments. It processes all functions for the schema passed to it. The /functions directory is hardcoded and expected by trite client code.
func dumpFuncs(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "functions")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select routine_name from information_schema.routines where routine_schema='" + schema + "' and routine_type = 'FUNCTION'")
	checkErr(err)

	var funcName string
	for rows.Next() {
		err = rows.Scan(&funcName)
		checkErr(err)

		var funcInfo createInfoStruct
		err = db.QueryRow("show create function "+addQuotes(schema)+"."+addQuotes(funcName)).Scan(&funcInfo.Name, &funcInfo.SqlMode, &funcInfo.Create, &funcInfo.CharsetClient, &funcInfo.Collation, &funcInfo.DbCollation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(funcInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, funcName+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpTriggers creates files containing trigger creation statments. It processes all triggers for the schema passed to it. The /triggers directory is hardcoded and expected by trite client code.
func dumpTriggers(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "triggers")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select trigger_name from information_schema.triggers where trigger_schema='" + schema + "'")
	checkErr(err)

	var trigName string
	for rows.Next() {
		err = rows.Scan(&trigName)
		checkErr(err)

		var trigInfo createInfoStruct
		err = db.QueryRow("show create trigger "+addQuotes(schema)+"."+addQuotes(trigName)).Scan(&trigInfo.Name, &trigInfo.SqlMode, &trigInfo.Create, &trigInfo.CharsetClient, &trigInfo.Collation, &trigInfo.DbCollation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(trigInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, trigName+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpViews creates files containing view creation statments. It processes all views for the schema passed to it. The /views directory is hardcoded and expected by trite client code.
func dumpViews(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "views")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select table_name from information_schema.tables where table_schema='" + schema + "' and table_type = 'VIEW'")
	checkErr(err)

	var view string
	for rows.Next() {
		err = rows.Scan(&view)
		checkErr(err)

		var viewInfo createInfoStruct
		err = db.QueryRow("show create view "+addQuotes(schema)+"."+addQuotes(view)).Scan(&viewInfo.Name, &viewInfo.Create, &viewInfo.CharsetClient, &viewInfo.Collation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(viewInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, view+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}
<commit_msg>Spelling correction<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"time"
)

const (
	stamp        = "20060102150405"
	dirPerms     = 0755
	filePerms    = 0644
	sqlExtension = ".sql"
)

// startDump copies creation statements for tables, procedures, functions, triggers and views to a file/directory structure at the path location that trite uses in client mode to restore tables.
func startDump(dir string, dbi *mysqlCredentials) {
	dumpdir := path.Join(dir, dbi.host+"_dump"+time.Now().Format(stamp))
	fmt.Println("Dumping to:", dumpdir)
	fmt.Println()

	// Return a database connection
	db, err := dbi.connect()
	defer db.Close()

	// Problem connecting to database
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Turn off idle connections
	db.SetMaxIdleConns(0)

	// Get a list of schemas in the target database
	db.SetMaxIdleConns(1)
	schemas := schemaList(db)

	// Create dump directory
	err = os.MkdirAll(dumpdir, dirPerms)
	checkErr(err)

	// Schema loop
	count := 0
	total := 0
	fmt.Println()
	for _, schema := range schemas {
		total++
		// Dump schema create
		fmt.Print(schema, ": ")
		dumpSchema(db, dumpdir, schema)

		// Dump table creation statements
		count = dumpTables(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " tables, ")

		// Dump procedure creation statements
		count = dumpProcs(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " procedures, ")

		// Dump function creation statements
		count = dumpFuncs(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " functions, ")

		// Dump trigger creation statements
		count = dumpTriggers(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " triggers, ")

		// Dump view creation statements
		count = dumpViews(db, dumpdir, schema)
		total = total + count
		fmt.Print(count, " views\n")
	}

	fmt.Println()
	fmt.Println(total, "total objects dumped")
}

// schemaList returns a string slice of schemas to process. MySQL specific schemas like mysql, information_schema and performance_schema are omitted.
func schemaList(db *sql.DB) []string {
	rows, err := db.Query("show databases")
	checkErr(err)

	// Get schema list
	schemas := make([]string, 0)
	var database string
	for rows.Next() {
		err = rows.Scan(&database)
		checkErr(err)

		if database != "mysql" && database != "information_schema" && database != "performance_schema" {
			schemas = append(schemas, database)
		}
	}

	return schemas
}

// dumpSchema creates a file with the schema creation statement.
func dumpSchema(db *sql.DB, dumpdir string, schema string) {
	dir := path.Join(dumpdir, schema)
	var err error

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var ignore string
	var stmt string
	err = db.QueryRow("show create schema "+addQuotes(schema)).Scan(&ignore, &stmt)
	checkErr(err)

	file := path.Join(dir, schema+sqlExtension)
	err = ioutil.WriteFile(file, []byte(stmt+";\n"), filePerms)
	checkErr(err)
}

// dumpTables creates files containing table creation statements. It processes all tables for the schema passed to it. The /tables directory is hardcoded and expected by trite client code.
func dumpTables(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "tables")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select table_name from information_schema.tables where table_schema='" + schema + "' and table_type = 'BASE TABLE'")
	checkErr(err)

	var tableName string
	var ignore string
	var stmt string
	for rows.Next() {
		err = rows.Scan(&tableName)
		checkErr(err)

		err = db.QueryRow("show create table "+addQuotes(schema)+"."+addQuotes(tableName)).Scan(&ignore, &stmt)
		checkErr(err)

		file := path.Join(dir, tableName+sqlExtension)
		err = ioutil.WriteFile(file, []byte(stmt+";\n"), filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpProcs creates files containing procedure creation statements. It processes all procedures for the schema passed to it. The /procedures directory is hardcoded and expected by trite client code.
func dumpProcs(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "procedures")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select routine_name from information_schema.routines where routine_schema='" + schema + "' and routine_type = 'PROCEDURE'")
	checkErr(err)

	var procName string
	for rows.Next() {
		err = rows.Scan(&procName)
		checkErr(err)

		var procInfo createInfoStruct
		err = db.QueryRow("show create procedure "+addQuotes(schema)+"."+addQuotes(procName)).Scan(&procInfo.Name, &procInfo.SqlMode, &procInfo.Create, &procInfo.CharsetClient, &procInfo.Collation, &procInfo.DbCollation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(procInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, procName+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpFuncs creates files containing function creation statements. It processes all functions for the schema passed to it. The /functions directory is hardcoded and expected by trite client code.
func dumpFuncs(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "functions")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select routine_name from information_schema.routines where routine_schema='" + schema + "' and routine_type = 'FUNCTION'")
	checkErr(err)

	var funcName string
	for rows.Next() {
		err = rows.Scan(&funcName)
		checkErr(err)

		var funcInfo createInfoStruct
		err = db.QueryRow("show create function "+addQuotes(schema)+"."+addQuotes(funcName)).Scan(&funcInfo.Name, &funcInfo.SqlMode, &funcInfo.Create, &funcInfo.CharsetClient, &funcInfo.Collation, &funcInfo.DbCollation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(funcInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, funcName+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpTriggers creates files containing trigger creation statements. It processes all triggers for the schema passed to it. The /triggers directory is hardcoded and expected by trite client code.
func dumpTriggers(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "triggers")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select trigger_name from information_schema.triggers where trigger_schema='" + schema + "'")
	checkErr(err)

	var trigName string
	for rows.Next() {
		err = rows.Scan(&trigName)
		checkErr(err)

		var trigInfo createInfoStruct
		err = db.QueryRow("show create trigger "+addQuotes(schema)+"."+addQuotes(trigName)).Scan(&trigInfo.Name, &trigInfo.SqlMode, &trigInfo.Create, &trigInfo.CharsetClient, &trigInfo.Collation, &trigInfo.DbCollation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(trigInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, trigName+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}

// dumpViews creates files containing view creation statements. It processes all views for the schema passed to it. The /views directory is hardcoded and expected by trite client code.
func dumpViews(db *sql.DB, dumpdir string, schema string) int {
	dir := path.Join(dumpdir, schema, "views")
	var err error
	count := 0

	err = os.Mkdir(dir, dirPerms)
	checkErr(err)

	var rows *sql.Rows
	rows, err = db.Query("select table_name from information_schema.tables where table_schema='" + schema + "' and table_type = 'VIEW'")
	checkErr(err)

	var view string
	for rows.Next() {
		err = rows.Scan(&view)
		checkErr(err)

		var viewInfo createInfoStruct
		err = db.QueryRow("show create view "+addQuotes(schema)+"."+addQuotes(view)).Scan(&viewInfo.Name, &viewInfo.Create, &viewInfo.CharsetClient, &viewInfo.Collation)
		checkErr(err)

		var jbyte []byte
		jbyte, err = json.MarshalIndent(viewInfo, "", "  ")
		checkErr(err)

		file := path.Join(dir, view+sqlExtension)
		err = ioutil.WriteFile(file, jbyte, filePerms)
		checkErr(err)

		count++
	}

	return count
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/bmizerany/pat"
	"github.com/moovweb/gokogiri"
	"github.com/moovweb/gokogiri/css"
	"github.com/moovweb/gokogiri/xml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"syscall"
	"time"
)

type Item struct {
	Id   string
	Href string
	Type string
}

type Epub struct {
	Zip *zip.Writer
	// TODO identifier / rights
	Title        string
	Subject      string
	Date         string
	Creator      string
	Contributors []string
	Items        []Item
}

const (
	Host        = "linuxfr.org"
	ContentType = "application/epub+zip"
	Container   = `<?xml version="1.0" encoding="utf-8"?>
<container xmlns="urn:oasis:names:tc:opendocument:xmlns:container" version="1.0">
  <rootfiles>
    <rootfile full-path="EPUB/package.opf" media-type="application/oebps-package+xml"/>
  </rootfiles>
</container>`
)

// FIXME why html/template escapes the first '<' sign?!??
// TODO embed CSS & images
// TODO cover
var PackageTemplate = template.Must(template.New("package").Parse(
	`<?xml version="1.0" encoding="utf-8"?>
<package xmlns="http://www.idpf.org/2007/opf" unique-identifier="bookid" xml:lang="fr" version="3.0">
	<metadata xmlns:dc="http://purl.org/dc/elements/1.1/">
		<dc:language id="pub-language">fr</dc:language>
		<dc:identifier id="pub-identifier">xxx</dc:identifier>
		{{if .Title}}<dc:title id="pub-title">{{.Title}}</dc:title>{{end}}
		{{if .Date}}<dc:date>{{.Date}}</dc:date>{{end}}
		{{if .Creator}}<dc:creator id="pub-creator">{{.Creator}}</dc:creator>{{end}}
		{{range .Contributors}}<dc:contributor>{{.}}</dc:contributor>
		{{end}}
		<dc:rights>xxx</dc:rights>
	</metadata>
	<manifest>
		{{range .Items}}<item id="{{.Id}}" href="{{.Href}}" media-type="{{.Type}}"/>
		{{end}}
	</manifest>
	<spine>
		{{range .Items}}<itemref idref="{{.Id}}"/>
		{{end}}
	</spine>
</package>`))

func NewEpub(w io.Writer) (epub *Epub) {
	z := zip.NewWriter(w)
	epub = &Epub{Zip: z, Items: []Item{}}
	epub.AddMimetype()
	epub.AddFile("META-INF/container.xml", Container)
	return
}

func (epub *Epub) AddContent(article xml.Node) {
	xpath := css.Convert(".content", css.LOCAL)
	nodes, err := article.Search(xpath)
	if err != nil || len(nodes) == 0 {
		return
	}
	html := nodes[0].InnerHtml() // FIXME should be a complete HTML document
	filename := "EPUB/content.html"
	epub.Items = append(epub.Items, Item{"item-content", filename, "application/xhtml+xml"})
	epub.AddFile(filename, html)
}

func (epub *Epub) AddComments(article xml.Node) {
	list := article.NextSibling()
	xpath := css.Convert(".threads>li", css.LOCAL)
	threads, err := list.Search(xpath)
	if err != nil {
		return
	}
	for _, thread := range threads {
		html := thread.InnerHtml() // FIXME should be a complete HTML document
		id := thread.Attr("id")
		filename := fmt.Sprintf("EPUB/%s.html", id)
		epub.Items = append(epub.Items, Item{id, filename, "application/xhtml+xml"})
		epub.AddFile(filename, html)
	}
}

func (epub *Epub) FindMeta(article xml.Node, selector string) string {
	xpath := css.Convert(selector, css.LOCAL)
	nodes, err := article.Search(xpath)
	if err != nil || len(nodes) == 0 {
		return ""
	}
	return nodes[0].Content()
}

func (epub *Epub) FindMetas(article xml.Node, selector string) []string {
	xpath := css.Convert(selector, css.LOCAL)
	nodes, err := article.Search(xpath)
	if err != nil {
		return nil
	}
	metas := make([]string, len(nodes))
	for i, node := range nodes {
		metas[i] = node.Content()
	}
	return metas
}

func (epub *Epub) FillMeta(article xml.Node) {
	epub.Title = epub.FindMeta(article, "header h1 a:last-child")
	epub.Subject = epub.FindMeta(article, "header h1 a.topic")
	meta := epub.FindMeta(article, "header time.updated")
	// FIXME ParseInLocation
	date, err := time.Parse("le 02/01/06 à 15:04", meta)
	if err == nil {
		epub.Date = date.String()
	}
	epub.Creator = epub.FindMeta(article, "header .meta a[rel=\"author\"]")
	epub.Contributors = epub.FindMetas(article, "header .meta .edited_by a")
}

func (epub *Epub) AddMimetype() (err error) {
	header := &zip.FileHeader{Name: "mimetype", Method: zip.Store}
	f, err := epub.Zip.CreateHeader(header)
	if err != nil {
		log.Println(err)
		return
	}

	_, err = f.Write([]byte(ContentType))
	if err != nil {
		log.Println(err)
		return
	}

	return
}

func (epub *Epub) AddFile(filename, content string) (err error) {
	f, err := epub.Zip.Create(filename)
	if err != nil {
		log.Println(err)
		return
	}

	_, err = f.Write([]byte(content))
	if err != nil {
		log.Println(err)
		return
	}

	return
}

func (epub *Epub) Close() {
	var opf bytes.Buffer
	err := PackageTemplate.Execute(&opf, epub)
	if err != nil {
		log.Println(err)
		return
	}

	epub.AddFile("EPUB/package.opf", opf.String())
	err = epub.Zip.Close()
	if err != nil {
		log.Println(err)
	}
}

func fetchArticle(uri string) (article xml.Node, err error) {
	resp, err := http.Get(uri)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Error on ioutil.ReadAll for %s: %s\n", uri, err)
		return
	}

	doc, err := gokogiri.ParseHtml(body)
	if err != nil {
		log.Printf("Gokogiri error: %s\n", err)
		return
	}

	xpath := css.Convert("#contents article", css.LOCAL)
	articles, err := doc.Root().Search(xpath)
	if err != nil {
		log.Printf("Gokogiri error: %s\n", err)
		return
	}

	if len(articles) == 0 {
		err = errors.New("No article found in the page")
		return
	}

	article = articles[0]
	return
}

// Create an epub for a news
func News(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Content-Type", ContentType)

	slug := r.URL.Query().Get(":slug")
	uri := fmt.Sprintf("http://%s/news/%s", Host, slug)
	article, err := fetchArticle(uri)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	// 	buffer := make([]byte, 4096)
	// 	buffer, _ = doc.Root().ToHtml(xml.DefaultEncodingBytes, buffer)

	epub := NewEpub(w)
	epub.FillMeta(article)
	epub.AddContent(article)
	epub.AddComments(article)
	epub.Close()
}

// Returns 200 OK if the server is running (for monitoring)
func Status(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func main() {
	// Parse the command-line
	var addr string
	var logs string
	flag.StringVar(&addr, "a", "127.0.0.1:8000", "Bind to this address:port")
	flag.StringVar(&logs, "l", "-", "Use this file for logs")
	flag.Parse()

	// Logging
	if logs != "-" {
		f, err := os.OpenFile(logs, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			log.Fatal("OpenFile: ", err)
		}
		syscall.Dup2(int(f.Fd()), int(os.Stdout.Fd()))
		syscall.Dup2(int(f.Fd()), int(os.Stderr.Fd()))
	}

	// Routing
	m := pat.New()
	m.Get("/status", http.HandlerFunc(Status))
	m.Get("/news/:slug.epub", http.HandlerFunc(News))
	// TODO accept other content types
	http.Handle("/", m)

	// Start the HTTP server
	log.Printf("Listening on http://%s/\n", addr)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<commit_msg>Use multiple cores<commit_after>package main

import (
	"archive/zip"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/bmizerany/pat"
	"github.com/moovweb/gokogiri"
	"github.com/moovweb/gokogiri/css"
	"github.com/moovweb/gokogiri/xml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"runtime"
	"syscall"
	"time"
)

type Item struct {
	Id   string
	Href string
	Type string
}

type Epub struct {
	Zip *zip.Writer
	// TODO identifier / rights
	Title        string
	Subject      string
	Date         string
	Creator      string
	Contributors []string
	Items        []Item
}

const (
	Host        = "linuxfr.org"
	ContentType = "application/epub+zip"
	Container   = `<?xml version="1.0" encoding="utf-8"?>
<container xmlns="urn:oasis:names:tc:opendocument:xmlns:container" version="1.0">
  <rootfiles>
    <rootfile full-path="EPUB/package.opf" media-type="application/oebps-package+xml"/>
  </rootfiles>
</container>`
)

// FIXME why html/template escapes the first '<' sign?!??
// TODO embed CSS & images
// TODO cover
var PackageTemplate = template.Must(template.New("package").Parse(
	`<?xml version="1.0" encoding="utf-8"?>
<package xmlns="http://www.idpf.org/2007/opf" unique-identifier="bookid" xml:lang="fr" version="3.0">
	<metadata xmlns:dc="http://purl.org/dc/elements/1.1/">
		<dc:language id="pub-language">fr</dc:language>
		<dc:identifier id="pub-identifier">xxx</dc:identifier>
		{{if .Title}}<dc:title id="pub-title">{{.Title}}</dc:title>{{end}}
		{{if .Date}}<dc:date>{{.Date}}</dc:date>{{end}}
		{{if .Creator}}<dc:creator id="pub-creator">{{.Creator}}</dc:creator>{{end}}
		{{range .Contributors}}<dc:contributor>{{.}}</dc:contributor>
		{{end}}
		<dc:rights>xxx</dc:rights>
	</metadata>
	<manifest>
		{{range .Items}}<item id="{{.Id}}" href="{{.Href}}" media-type="{{.Type}}"/>
		{{end}}
	</manifest>
	<spine>
		{{range .Items}}<itemref idref="{{.Id}}"/>
		{{end}}
	</spine>
</package>`))

func NewEpub(w io.Writer) (epub *Epub) {
	z := zip.NewWriter(w)
	epub = &Epub{Zip: z, Items: []Item{}}
	epub.AddMimetype()
	epub.AddFile("META-INF/container.xml", Container)
	return
}

func (epub *Epub) AddContent(article xml.Node) {
	xpath := css.Convert(".content", css.LOCAL)
	nodes, err := article.Search(xpath)
	if err != nil || len(nodes) == 0 {
		return
	}
	html := nodes[0].InnerHtml() // FIXME should be a complete HTML document
	filename := "EPUB/content.html"
	epub.Items = append(epub.Items, Item{"item-content", filename, "application/xhtml+xml"})
	epub.AddFile(filename, html)
}

func (epub *Epub) AddComments(article xml.Node) {
	list := article.NextSibling()
	xpath := css.Convert(".threads>li", css.LOCAL)
	threads, err := list.Search(xpath)
	if err != nil {
		return
	}
	for _, thread := range threads {
		html := thread.InnerHtml() // FIXME should be a complete HTML document
		id := thread.Attr("id")
		filename := fmt.Sprintf("EPUB/%s.html", id)
		epub.Items = append(epub.Items, Item{id, filename, "application/xhtml+xml"})
		epub.AddFile(filename, html)
	}
}

func (epub *Epub) FindMeta(article xml.Node, selector string) string {
	xpath := css.Convert(selector, css.LOCAL)
	nodes, err := article.Search(xpath)
	if err != nil || len(nodes) == 0 {
		return ""
	}
	return nodes[0].Content()
}

func (epub *Epub) FindMetas(article xml.Node, selector string) []string {
	xpath := css.Convert(selector, css.LOCAL)
	nodes, err := article.Search(xpath)
	if err != nil {
		return nil
	}
	metas := make([]string, len(nodes))
	for i, node := range nodes {
		metas[i] = node.Content()
	}
	return metas
}

func (epub *Epub) FillMeta(article xml.Node) {
	epub.Title = epub.FindMeta(article, "header h1 a:last-child")
	epub.Subject = epub.FindMeta(article, "header h1 a.topic")
	meta := epub.FindMeta(article, "header time.updated")
	// FIXME ParseInLocation
	date, err := time.Parse("le 02/01/06 à 15:04", meta)
	if err == nil {
		epub.Date = date.String()
	}
	epub.Creator = epub.FindMeta(article, "header .meta a[rel=\"author\"]")
	epub.Contributors = epub.FindMetas(article, "header .meta .edited_by a")
}

func (epub *Epub) AddMimetype() (err error) {
	header := &zip.FileHeader{Name: "mimetype", Method: zip.Store}
	f, err := epub.Zip.CreateHeader(header)
	if err != nil {
		log.Println(err)
		return
	}

	_, err = f.Write([]byte(ContentType))
	if err != nil {
		log.Println(err)
		return
	}

	return
}

func (epub *Epub) AddFile(filename, content string) (err error) {
	f, err := epub.Zip.Create(filename)
	if err != nil {
		log.Println(err)
		return
	}

	_, err = f.Write([]byte(content))
	if err != nil {
		log.Println(err)
		return
	}

	return
}

func (epub *Epub) Close() {
	var opf bytes.Buffer
	err := PackageTemplate.Execute(&opf, epub)
	if err != nil {
		log.Println(err)
		return
	}

	epub.AddFile("EPUB/package.opf", opf.String())
	err = epub.Zip.Close()
	if err != nil {
		log.Println(err)
	}
}

func fetchArticle(uri string) (article xml.Node, err error) {
	resp, err := http.Get(uri)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Error on ioutil.ReadAll for %s: %s\n", uri, err)
		return
	}

	doc, err := gokogiri.ParseHtml(body)
	if err != nil {
		log.Printf("Gokogiri error: %s\n", err)
		return
	}

	xpath := css.Convert("#contents article", css.LOCAL)
	articles, err := doc.Root().Search(xpath)
	if err != nil {
		log.Printf("Gokogiri error: %s\n", err)
		return
	}

	if len(articles) == 0 {
		err = errors.New("No article found in the page")
		return
	}

	article = articles[0]
	return
}

// Create an epub for a news
func News(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Content-Type", ContentType)

	slug := r.URL.Query().Get(":slug")
	uri := fmt.Sprintf("http://%s/news/%s", Host, slug)
	article, err := fetchArticle(uri)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	// 	buffer := make([]byte, 4096)
	// 	buffer, _ = doc.Root().ToHtml(xml.DefaultEncodingBytes, buffer)

	epub := NewEpub(w)
	epub.FillMeta(article)
	epub.AddContent(article)
	epub.AddComments(article)
	epub.Close()
}

// Returns 200 OK if the server is running (for monitoring)
func Status(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Parse the command-line
	var addr string
	var logs string
	flag.StringVar(&addr, "a", "127.0.0.1:8000", "Bind to this address:port")
	flag.StringVar(&logs, "l", "-", "Use this file for logs")
	flag.Parse()

	// Logging
	if logs != "-" {
		f, err := os.OpenFile(logs, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
		if err != nil {
			log.Fatal("OpenFile: ", err)
		}
		syscall.Dup2(int(f.Fd()), int(os.Stdout.Fd()))
		syscall.Dup2(int(f.Fd()), int(os.Stderr.Fd()))
	}

	// Routing
	m := pat.New()
	m.Get("/status", http.HandlerFunc(Status))
	m.Get("/news/:slug.epub", http.HandlerFunc(News))
	// TODO accept other content types
	http.Handle("/", m)

	// Start the HTTP server
	log.Printf("Listening on http://%s/\n", addr)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}
<|endoftext|>