<commit_before>package checkntpoffset

import (
	"bufio"
	"fmt"
	"io"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/beevik/ntp"
	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
)

var opts struct {
	Crit       float64 `short:"c" long:"critical" default:"100" description:"Critical threshold of ntp offset(ms)"`
	Warn       float64 `short:"w" long:"warning" default:"50" description:"Warning threshold of ntp offset(ms)"`
	NTPServers string  `short:"s" long:"ntp-servers" default:"" description:"Use specified NTP Servers(plural servers can be set separated by ,). When set plural servers, use first response. If not set, use local command just like ntpd/chronyd."`
	NTPTimeout int     `short:"t" long:"ntp-timeout" default:"15" description:"Timeout of NTP Server Querying(in seconds)."`
}

var ntpTimeout int

func init() {
	ntpTimeout = 15
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "NTP"
	ckr.Exit()
}

func run(args []string) *checkers.Checker {
	_, err := flags.ParseArgs(&opts, args)
	if err != nil {
		os.Exit(1)
	}
	ntpTimeout = opts.NTPTimeout

	offset, err := getNTPOffset(opts.NTPServers)
	if err != nil {
		return checkers.Unknown(err.Error())
	}

	var chkSt checkers.Status
	var msg string
	stratum, err := getNTPStratum()
	if err != nil {
		msg = fmt.Sprintf("ntp stratum not found offset is %f(actual), %f(warning threshold), %f(critical threshold)", math.Abs(offset), opts.Warn, opts.Crit)
		return checkers.NewChecker(checkers.WARNING, msg)
	}

	if opts.Crit < math.Abs(offset) {
		msg = fmt.Sprintf("ntp offset is over %f(actual) > %f(threshold)", math.Abs(offset), opts.Crit)
		chkSt = checkers.CRITICAL
	} else if opts.Warn < math.Abs(offset) {
		msg = fmt.Sprintf("ntp offset is over %f(actual) > %f(threshold)", math.Abs(offset), opts.Warn)
		chkSt = checkers.WARNING
	} else {
		msg = fmt.Sprintf("ntp stratum=%s, offset is %f(actual) < %f(warning threshold), %f(critial threshold)", stratum, math.Abs(offset), opts.Warn, opts.Crit)
		chkSt = checkers.OK
	}

	return checkers.NewChecker(chkSt, msg)
}

func withCmd(cmd *exec.Cmd, fn func(io.Reader) error) error {
	out, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return err
	}
	if err := fn(out); err != nil {
		return err
	}
	return cmd.Wait()
}

const (
	ntpNTPD    = "ntpd"
	ntpChronyd = "chronyd"

	cmdNTPq    = "ntpq"
	cmdChronyc = "chronyc"
)

func hasCommand(cmd string) bool {
	_, err := exec.LookPath(cmd)
	return err == nil
}

func detectNTPDname() (ntpdName string, err error) {
	err = withCmd(exec.Command("ps", "-eo", "comm"), func(out io.Reader) error {
		scr := bufio.NewScanner(out)
		for scr.Scan() {
			switch filepath.Base(scr.Text()) {
			case ntpChronyd:
				if hasCommand(cmdChronyc) {
					ntpdName = ntpChronyd
					return nil
				}
			case ntpNTPD:
				if hasCommand(cmdNTPq) {
					ntpdName = ntpNTPD
					return nil
				}
			}
		}
		return fmt.Errorf("no ntp daemons detected")
	})
	return ntpdName, err
}

func getNTPOffset(ntpServers string) (float64, error) {
	if ntpServers != "" {
		return getNTPOffsetFromNTPServers(ntpServers)
	}
	ntpdName, err := detectNTPDname()
	if err != nil {
		return 0.0, err
	}
	switch ntpdName {
	case ntpNTPD:
		return getNTPOffsetFromNTPD()
	case ntpChronyd:
		return getNTPOffsetFromChrony()
	}
	return 0.0, fmt.Errorf("unsupported ntp daemon %q", ntpdName)
}

// getNTPOffsetFromNTPServers ask time to ntp servers and return NTP Offset.
// Use first response, ignore others
//
// FIXME need fluent cancel mechanism
func getNTPOffsetFromNTPServers(ntpServers string) (offset float64, err error) {
	resultChan := make(chan float64)
	for _, ntpServer := range strings.Split(ntpServers, ",") {
		go func(ntpServer string) error {
			ntpServer = strings.Trim(ntpServer, " ")
			options := ntp.QueryOptions{Timeout: time.Duration(ntpTimeout) * time.Second}
			response, err := ntp.QueryWithOptions(ntpServer, options)
			if err != nil {
				return err
			}
			resultChan <- float64(response.ClockOffset / time.Millisecond)
			return nil
		}(ntpServer)
	}

	select {
	case <-time.After(time.Duration(ntpTimeout) * time.Second):
		// return error only when all NTPServers are failed
		return 0.0, fmt.Errorf("NTP offset cannot get from %q", ntpServers)
	case offset = <-resultChan:
		return offset, nil
	}
}

func getNTPOffsetFromNTPD() (offset float64, err error) {
	err = withCmd(exec.Command(cmdNTPq, "-c", "rv 0 offset"), func(out io.Reader) error {
		offset, err = parseNTPOffsetFromNTPD(out)
		return err
	})
	return offset, err
}

func parseNTPOffsetFromNTPD(out io.Reader) (float64, error) {
	scr := bufio.NewScanner(out)
	const offsetPrefix = "offset="
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, offsetPrefix) {
			return strconv.ParseFloat(strings.TrimPrefix(line, offsetPrefix), 64)
		}
	}
	return 0.0, fmt.Errorf("couldn't get ntp offset. ntpd process may be down")
}

func getNTPOffsetFromChrony() (offset float64, err error) {
	err = withCmd(exec.Command(cmdChronyc, "tracking"), func(out io.Reader) error {
		offset, err = parseNTPOffsetFromChrony(out)
		return err
	})
	return offset, err
}

func parseNTPOffsetFromChrony(out io.Reader) (offset float64, err error) {
	scr := bufio.NewScanner(out)
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, "Last offset") {
			flds := strings.Fields(line)
			if len(flds) != 5 {
				return 0.0, fmt.Errorf("failed to get ntp offset")
			}
			offset, err = strconv.ParseFloat(flds[3], 64)
			if err != nil {
				return 0.0, err
			}
			return offset * 1000, nil
		}
	}
	return 0.0, fmt.Errorf("failed to get ntp offset")
}

func getNTPStratum() (string, error) {
	ntpdName, err := detectNTPDname()
	if err != nil {
		return "unknown", err
	}
	switch ntpdName {
	case ntpNTPD:
		return getNTPStratumFromNTPD()
	case ntpChronyd:
		return getNTPStratumFromChrony()
	}
	return "unknown", fmt.Errorf("unsupported ntp daemon %q", ntpdName)
}

func getNTPStratumFromNTPD() (stratum string, err error) {
	err = withCmd(exec.Command(cmdNTPq, "-np"), func(out io.Reader) error {
		stratum, err = parseNTPStratumFromNTPD(out)
		return err
	})
	return stratum, err
}

func parseNTPStratumFromNTPD(out io.Reader) (string, error) {
	scr := bufio.NewScanner(out)
        const stratumPrefix = "*"
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, stratumPrefix) {
                        flds := strings.Fields(line)
			if len(flds) != 10 {
				return "unknown", fmt.Errorf("failed to get ntp stratum")
			}
                        return strings.TrimPrefix(flds[0], stratumPrefix), nil
		}
	}
	return "unknown", fmt.Errorf("couldn't get ntp stratum.")
}

func getNTPStratumFromChrony() (stratum string, err error) {
	err = withCmd(exec.Command(cmdChronyc, "-n", "sources"), func(out io.Reader) error {
		stratum, err = parseNTPStratumFromChrony(out)
		return err
	})
	return stratum, err
}

func parseNTPStratumFromChrony(out io.Reader) (string, error) {
	scr := bufio.NewScanner(out)
	const stratumPrefix = "^*"
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, stratumPrefix) {
                        flds := strings.Fields(line)
			if len(flds) != 10 {
				return "unknown", fmt.Errorf("failed to get ntp stratum")
			}
                        return flds[1], nil
		}
	}
	return "unknown", fmt.Errorf("couldn't get ntp stratum.")
}
<commit_msg>fix for lint<commit_after>package checkntpoffset

import (
	"bufio"
	"fmt"
	"io"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/beevik/ntp"
	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
)

var opts struct {
	Crit       float64 `short:"c" long:"critical" default:"100" description:"Critical threshold of ntp offset(ms)"`
	Warn       float64 `short:"w" long:"warning" default:"50" description:"Warning threshold of ntp offset(ms)"`
	NTPServers string  `short:"s" long:"ntp-servers" default:"" description:"Use specified NTP Servers(plural servers can be set separated by ,). When set plural servers, use first response. If not set, use local command just like ntpd/chronyd."`
	NTPTimeout int     `short:"t" long:"ntp-timeout" default:"15" description:"Timeout of NTP Server Querying(in seconds)."`
}

var ntpTimeout int

func init() {
	ntpTimeout = 15
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "NTP"
	ckr.Exit()
}

func run(args []string) *checkers.Checker {
	_, err := flags.ParseArgs(&opts, args)
	if err != nil {
		os.Exit(1)
	}
	ntpTimeout = opts.NTPTimeout

	offset, err := getNTPOffset(opts.NTPServers)
	if err != nil {
		return checkers.Unknown(err.Error())
	}

	var chkSt checkers.Status
	var msg string
	stratum, err := getNTPStratum()
	if err != nil {
		msg = fmt.Sprintf("ntp stratum not found offset is %f(actual), %f(warning threshold), %f(critical threshold)", math.Abs(offset), opts.Warn, opts.Crit)
		return checkers.NewChecker(checkers.WARNING, msg)
	}

	if opts.Crit < math.Abs(offset) {
		msg = fmt.Sprintf("ntp offset is over %f(actual) > %f(threshold)", math.Abs(offset), opts.Crit)
		chkSt = checkers.CRITICAL
	} else if opts.Warn < math.Abs(offset) {
		msg = fmt.Sprintf("ntp offset is over %f(actual) > %f(threshold)", math.Abs(offset), opts.Warn)
		chkSt = checkers.WARNING
	} else {
		msg = fmt.Sprintf("ntp stratum=%s, offset is %f(actual) < %f(warning threshold), %f(critial threshold)", stratum, math.Abs(offset), opts.Warn, opts.Crit)
		chkSt = checkers.OK
	}

	return checkers.NewChecker(chkSt, msg)
}

func withCmd(cmd *exec.Cmd, fn func(io.Reader) error) error {
	out, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return err
	}
	if err := fn(out); err != nil {
		return err
	}
	return cmd.Wait()
}

const (
	ntpNTPD    = "ntpd"
	ntpChronyd = "chronyd"

	cmdNTPq    = "ntpq"
	cmdChronyc = "chronyc"
)

func hasCommand(cmd string) bool {
	_, err := exec.LookPath(cmd)
	return err == nil
}

func detectNTPDname() (ntpdName string, err error) {
	err = withCmd(exec.Command("ps", "-eo", "comm"), func(out io.Reader) error {
		scr := bufio.NewScanner(out)
		for scr.Scan() {
			switch filepath.Base(scr.Text()) {
			case ntpChronyd:
				if hasCommand(cmdChronyc) {
					ntpdName = ntpChronyd
					return nil
				}
			case ntpNTPD:
				if hasCommand(cmdNTPq) {
					ntpdName = ntpNTPD
					return nil
				}
			}
		}
		return fmt.Errorf("no ntp daemons detected")
	})
	return ntpdName, err
}

func getNTPOffset(ntpServers string) (float64, error) {
	if ntpServers != "" {
		return getNTPOffsetFromNTPServers(ntpServers)
	}
	ntpdName, err := detectNTPDname()
	if err != nil {
		return 0.0, err
	}
	switch ntpdName {
	case ntpNTPD:
		return getNTPOffsetFromNTPD()
	case ntpChronyd:
		return getNTPOffsetFromChrony()
	}
	return 0.0, fmt.Errorf("unsupported ntp daemon %q", ntpdName)
}

// getNTPOffsetFromNTPServers ask time to ntp servers and return NTP Offset.
// Use first response, ignore others
//
// FIXME need fluent cancel mechanism
func getNTPOffsetFromNTPServers(ntpServers string) (offset float64, err error) {
	resultChan := make(chan float64)
	for _, ntpServer := range strings.Split(ntpServers, ",") {
		go func(ntpServer string) error {
			ntpServer = strings.Trim(ntpServer, " ")
			options := ntp.QueryOptions{Timeout: time.Duration(ntpTimeout) * time.Second}
			response, err := ntp.QueryWithOptions(ntpServer, options)
			if err != nil {
				return err
			}
			resultChan <- float64(response.ClockOffset / time.Millisecond)
			return nil
		}(ntpServer)
	}

	select {
	case <-time.After(time.Duration(ntpTimeout) * time.Second):
		// return error only when all NTPServers are failed
		return 0.0, fmt.Errorf("NTP offset cannot get from %q", ntpServers)
	case offset = <-resultChan:
		return offset, nil
	}
}

func getNTPOffsetFromNTPD() (offset float64, err error) {
	err = withCmd(exec.Command(cmdNTPq, "-c", "rv 0 offset"), func(out io.Reader) error {
		offset, err = parseNTPOffsetFromNTPD(out)
		return err
	})
	return offset, err
}

func parseNTPOffsetFromNTPD(out io.Reader) (float64, error) {
	scr := bufio.NewScanner(out)
	const offsetPrefix = "offset="
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, offsetPrefix) {
			return strconv.ParseFloat(strings.TrimPrefix(line, offsetPrefix), 64)
		}
	}
	return 0.0, fmt.Errorf("couldn't get ntp offset. ntpd process may be down")
}

func getNTPOffsetFromChrony() (offset float64, err error) {
	err = withCmd(exec.Command(cmdChronyc, "tracking"), func(out io.Reader) error {
		offset, err = parseNTPOffsetFromChrony(out)
		return err
	})
	return offset, err
}

func parseNTPOffsetFromChrony(out io.Reader) (offset float64, err error) {
	scr := bufio.NewScanner(out)
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, "Last offset") {
			flds := strings.Fields(line)
			if len(flds) != 5 {
				return 0.0, fmt.Errorf("failed to get ntp offset")
			}
			offset, err = strconv.ParseFloat(flds[3], 64)
			if err != nil {
				return 0.0, err
			}
			return offset * 1000, nil
		}
	}
	return 0.0, fmt.Errorf("failed to get ntp offset")
}

func getNTPStratum() (string, error) {
	ntpdName, err := detectNTPDname()
	if err != nil {
		return "unknown", err
	}
	switch ntpdName {
	case ntpNTPD:
		return getNTPStratumFromNTPD()
	case ntpChronyd:
		return getNTPStratumFromChrony()
	}
	return "unknown", fmt.Errorf("unsupported ntp daemon %q", ntpdName)
}

func getNTPStratumFromNTPD() (stratum string, err error) {
	err = withCmd(exec.Command(cmdNTPq, "-np"), func(out io.Reader) error {
		stratum, err = parseNTPStratumFromNTPD(out)
		return err
	})
	return stratum, err
}

func parseNTPStratumFromNTPD(out io.Reader) (string, error) {
	scr := bufio.NewScanner(out)
        const stratumPrefix = "*"
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, stratumPrefix) {
                        flds := strings.Fields(line)
			if len(flds) != 10 {
				return "unknown", fmt.Errorf("failed to get ntp stratum")
			}
                        return strings.TrimPrefix(flds[0], stratumPrefix), nil
		}
	}
	return "unknown", fmt.Errorf("failed to get ntp stratum")
}

func getNTPStratumFromChrony() (stratum string, err error) {
	err = withCmd(exec.Command(cmdChronyc, "-n", "sources"), func(out io.Reader) error {
		stratum, err = parseNTPStratumFromChrony(out)
		return err
	})
	return stratum, err
}

func parseNTPStratumFromChrony(out io.Reader) (string, error) {
	scr := bufio.NewScanner(out)
	const stratumPrefix = "^*"
	for scr.Scan() {
		line := scr.Text()
		if strings.HasPrefix(line, stratumPrefix) {
                        flds := strings.Fields(line)
			if len(flds) != 10 {
				return "unknown", fmt.Errorf("failed to get ntp stratum")
			}
                        return flds[1], nil
		}
	}
	return "unknown", fmt.Errorf("failed to get ntp stratum")
}
<|endoftext|>


<commit_before>package apply

import (
	"database/sql"
	"errors"
	"github.com/elwinar/rambler/lib"
)

var (
	ErrNilMigration = errors.New("nil migration")
	ErrNilTransaction = errors.New("nil transaction")
	ErrNilDatabase = errors.New("nil database")
)

func Apply(migration *lib.Migration, tx *sql.Tx) (error, error) {
	return apply(migration, tx)
}

type migration interface {
	Scan(string) []string
}

type txer interface {
	Exec(string, ...interface{}) (sql.Result, error)
	Commit() error
	Rollback() error
}

func apply(migration migration, tx txer) (err error, sqlerr error) {
	if migration == nil {
		return ErrNilMigration, nil
	}

	if tx == nil {
		return ErrNilTransaction, nil
	}

	for _, statement := range migration.Scan("up") {
		_, sqlerr := tx.Exec(statement)
		if sqlerr != nil {
			err := tx.Rollback()
			return err, sqlerr
		}
	}

	err = tx.Commit()
	return err, nil
}
<commit_msg>Removed unused error<commit_after>package apply

import (
	"database/sql"
	"errors"
	"github.com/elwinar/rambler/lib"
)

var (
	ErrNilMigration = errors.New("nil migration")
	ErrNilTransaction = errors.New("nil transaction")
)

func Apply(migration *lib.Migration, tx *sql.Tx) (error, error) {
	return apply(migration, tx)
}

type migration interface {
	Scan(string) []string
}

type txer interface {
	Exec(string, ...interface{}) (sql.Result, error)
	Commit() error
	Rollback() error
}

func apply(migration migration, tx txer) (err error, sqlerr error) {
	if migration == nil {
		return ErrNilMigration, nil
	}

	if tx == nil {
		return ErrNilTransaction, nil
	}

	for _, statement := range migration.Scan("up") {
		_, sqlerr := tx.Exec(statement)
		if sqlerr != nil {
			err := tx.Rollback()
			return err, sqlerr
		}
	}

	err = tx.Commit()
	return err, nil
}
<|endoftext|>


<commit_before>package notify

import (
	"errors"
	"fmt"
	"log"
	"sync"

	"github.com/godbus/dbus/v5"
)

const (
	dbusRemoveMatch            = "org.freedesktop.DBus.RemoveMatch"
	dbusAddMatch               = "org.freedesktop.DBus.AddMatch"
	dbusObjectPath             = "/org/freedesktop/Notifications" // the DBUS object path
	dbusNotificationsInterface = "org.freedesktop.Notifications"  // DBUS Interface
	signalNotificationClosed   = "org.freedesktop.Notifications.NotificationClosed"
	signalActionInvoked        = "org.freedesktop.Notifications.ActionInvoked"
	callGetCapabilities        = "org.freedesktop.Notifications.GetCapabilities"
	callCloseNotification      = "org.freedesktop.Notifications.CloseNotification"
	callNotify                 = "org.freedesktop.Notifications.Notify"
	callGetServerInformation   = "org.freedesktop.Notifications.GetServerInformation"

	channelBufferSize = 2
)

// Notification holds all information needed for creating a notification
type Notification struct {
	AppName string
	// Setting ReplacesID atomically replaces the notification with this ID.
	// Optional.
	ReplacesID uint32
	// See predefined icons here: http://standards.freedesktop.org/icon-naming-spec/icon-naming-spec-latest.html
	// Optional.
	AppIcon string
	Summary string
	Body    string
	// Actions are tuples of (action_key, label), e.g.: []string{"cancel", "Cancel", "open", "Open"}
	Actions []string
	Hints   map[string]dbus.Variant
	// ExpireTimeout: milliseconds to show notification
	ExpireTimeout int32
}

// SendNotification is provided for convenience.
// Use if you only want to deliver a notification and dont care about events.
func SendNotification(conn *dbus.Conn, note Notification) (uint32, error) {
	actions := len(note.Actions)
	if (actions % 2) != 0 {
		return 0, errors.New("actions must be pairs of (key, label)")
	}

	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callNotify, 0,
		note.AppName,
		note.ReplacesID,
		note.AppIcon,
		note.Summary,
		note.Body,
		note.Actions,
		note.Hints,
		note.ExpireTimeout)
	if call.Err != nil {
		return 0, fmt.Errorf("error sending notification: %w", call.Err)
	}
	var ret uint32
	err := call.Store(&ret)
	if err != nil {
		return ret, fmt.Errorf("error getting uint32 ret value: %w", err)
	}
	return ret, nil
}

// ServerInformation is a holder for information returned by
// GetServerInformation call.
type ServerInformation struct {
	Name        string
	Vendor      string
	Version     string
	SpecVersion string
}

// GetServerInformation returns the information on the server.
//
// org.freedesktop.Notifications.GetServerInformation
//
//  GetServerInformation Return Values
//
//		Name		 Type	  Description
//		name		 STRING	  The product name of the server.
//		vendor		 STRING	  The vendor name. For example, "KDE," "GNOME," "freedesktop.org," or "Microsoft."
//		version		 STRING	  The server's version number.
//		spec_version STRING	  The specification version the server is compliant with.
//
func GetServerInformation(conn *dbus.Conn) (ServerInformation, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	if obj == nil {
		return ServerInformation{}, errors.New("error creating dbus call object")
	}
	call := obj.Call(callGetServerInformation, 0)
	if call.Err != nil {
		return ServerInformation{}, fmt.Errorf("error calling %v: %v", callGetServerInformation, call.Err)
	}

	ret := ServerInformation{}
	err := call.Store(&ret.Name, &ret.Vendor, &ret.Version, &ret.SpecVersion)
	if err != nil {
		return ret, fmt.Errorf("error reading %v return values: %v", callGetServerInformation, err)
	}
	return ret, nil
}

// GetCapabilities gets the capabilities of the notification server.
// This call takes no parameters.
// It returns an array of strings. Each string describes an optional capability implemented by the server.
//
// See also: https://developer.gnome.org/notification-spec/
// GetCapabilities provide an exported method for this operation
func GetCapabilities(conn *dbus.Conn) ([]string, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callGetCapabilities, 0)
	if call.Err != nil {
		return []string{}, call.Err
	}
	var ret []string
	err := call.Store(&ret)
	if err != nil {
		return ret, err
	}
	return ret, nil
}

// Notifier is an interface for implementing the operations supported by the
// freedesktop DBus Notifications object.
//
// New() sets up a Notifier that listens on dbus' signals regarding
// Notifications: NotificationClosed and ActionInvoked.
//
// Note this also means the caller MUST consume output from these channels,
// given in methods NotificationClosed() and ActionInvoked().
// Users that only want to send a simple notification, but don't care about
// interactions, see exported method: SendNotification(conn, Notification)
//
// Caller is also responsible to call Close() before exiting,
// to shut down event loop and cleanup.
type Notifier interface {
	SendNotification(n Notification) (uint32, error)
	GetCapabilities() ([]string, error)
	GetServerInformation() (ServerInformation, error)
	CloseNotification(id uint32) (bool, error)
	Close() error
}

// NotificationClosedHandler is called when we receive a NotificationClosed signal
type NotificationClosedHandler func(*NotificationClosedSignal)

// ActionInvokedHandler is called when we receive a signal that one of the actions passed was invoked.
//
// Note that invoking an action often also produces a NotificationClosedSignal,
// so you might receive both a Closed signal and a ActionInvoked signal.
// I suspect this detail is implementation specific for the UI interaction,
// and does at least happen on XFCE4.
type ActionInvokedHandler func(*ActionInvokedSignal)

// ActionInvokedSignal holds data from any signal received regarding Actions invoked
type ActionInvokedSignal struct {
	// ID of the Notification the action was invoked for
	ID uint32
	// Key from the tuple (action_key, label)
	ActionKey string
}

// notifier implements Notifier interface
type notifier struct {
	conn     *dbus.Conn
	signal   chan *dbus.Signal
	done     chan bool
	onClosed NotificationClosedHandler
	onAction ActionInvokedHandler
	wg       *sync.WaitGroup
	log      logger
}

type logger interface {
	Printf(format string, v ...interface{})
}

// option overrides certain parts of a Notifier
type option func(*notifier)

// WithLogger sets a new logger func
func WithLogger(logz logger) option {
	return func(n *notifier) {
		n.log = logz
	}
}

// WithOnAction sets ActionInvokedHandler handler
func WithOnAction(h ActionInvokedHandler) option {
	return func(n *notifier) {
		n.onAction = h
	}
}

// WithOnClosed sets NotificationClosed handler
func WithOnClosed(h NotificationClosedHandler) option {
	return func(n *notifier) {
		n.onClosed = h
	}
}

// New creates a new Notifier using conn.
// See also: Notifier
func New(conn *dbus.Conn, opts ...option) (Notifier, error) {
	n := &notifier{
		conn:     conn,
		signal:   make(chan *dbus.Signal, channelBufferSize),
		done:     make(chan bool),
		wg:       &sync.WaitGroup{},
		onClosed: func(s *NotificationClosedSignal) {},
		onAction: func(s *ActionInvokedSignal) {},
		log:      &loggerWrapper{"notify: "},
	}

	for _, val := range opts {
		val(n)
	}

	// add a listener (matcher) in dbus for signals to Notification interface.
	err := n.conn.AddMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)
	if err != nil {
		return nil, fmt.Errorf("error registering for signals in dbus: %w", err)
	}
	// register in dbus for signal delivery
	n.conn.Signal(n.signal)

	// start eventloop
	go n.eventLoop()

	return n, nil
}

func (n notifier) eventLoop() {
	n.wg.Add(1)
	defer n.wg.Done()
	for {
		select {
		case signal := <-n.signal:
			n.handleSignal(signal)
		case <-n.done:
			n.log.Printf("Got Close() signal, shutting down...")
			return
		}
	}
}

// signal handler that translates and sends notifications to channels
func (n notifier) handleSignal(signal *dbus.Signal) {
	switch signal.Name {
	case signalNotificationClosed:
		nc := &NotificationClosedSignal{
			ID:     signal.Body[0].(uint32),
			Reason: Reason(signal.Body[1].(uint32)),
		}
		n.onClosed(nc)
	case signalActionInvoked:
		is := &ActionInvokedSignal{
			ID:        signal.Body[0].(uint32),
			ActionKey: signal.Body[1].(string),
		}
		n.onAction(is)
	default:
		n.log.Printf("Received unknown signal: %+v", signal)
	}
}

func (n *notifier) GetCapabilities() ([]string, error) {
	return GetCapabilities(n.conn)
}
func (n *notifier) GetServerInformation() (ServerInformation, error) {
	return GetServerInformation(n.conn)
}

// SendNotification sends a notification to the notification server and returns the ID or an error.
//
// Implements dbus call:
//
//     UINT32 org.freedesktop.Notifications.Notify (
//	       STRING app_name,
//	       UINT32 replaces_id,
//	       STRING app_icon,
//	       STRING summary,
//	       STRING body,
//	       ARRAY  actions,
//	       DICT   hints,
//	       INT32  expire_timeout
//     );
//
//		Name	    	Type	Description
//		app_name		STRING	The optional name of the application sending the notification. Can be blank.
//		replaces_id	    UINT32	The optional notification ID that this notification replaces. The server must atomically (ie with no flicker or other visual cues) replace the given notification with this one. This allows clients to effectively modify the notification while it's active. A value of value of 0 means that this notification won't replace any existing notifications.
//		app_icon		STRING	The optional program icon of the calling application. Can be an empty string, indicating no icon.
//		summary		    STRING	The summary text briefly describing the notification.
//		body			STRING	The optional detailed body text. Can be empty.
//		actions		    ARRAY	Actions are sent over as a list of pairs. Each even element in the list (starting at index 0) represents the identifier for the action. Each odd element in the list is the localized string that will be displayed to the user.
//		hints	        DICT	Optional hints that can be passed to the server from the client program. Although clients and servers should never assume each other supports any specific hints, they can be used to pass along information, such as the process PID or window ID, that the server may be able to make use of. See Hints. Can be empty.
//      expire_timeout  INT32   The timeout time in milliseconds since the display of the notification at which the notification should automatically close.
//								If -1, the notification's expiration time is dependent on the notification server's settings, and may vary for the type of notification. If 0, never expire.
//
// If replaces_id is 0, the return value is a UINT32 that represent the notification.
// It is unique, and will not be reused unless a MAXINT number of notifications have been generated.
// An acceptable implementation may just use an incrementing counter for the ID.
// The returned ID is always greater than zero. Servers must make sure not to return zero as an ID.
//
// If replaces_id is not 0, the returned value is the same value as replaces_id.
func (n *notifier) SendNotification(note Notification) (uint32, error) {
	return SendNotification(n.conn, note)
}

// CloseNotification causes a notification to be forcefully closed and removed from the user's view.
// It can be used, for example, in the event that what the notification pertains to is no longer relevant,
// or to cancel a notification with no expiration time.
//
// The NotificationClosed (dbus) signal is emitted by this method.
// If the notification no longer exists, an empty D-BUS Error message is sent back.
func (n *notifier) CloseNotification(id uint32) (bool, error) {
	obj := n.conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callCloseNotification, 0, id)
	if call.Err != nil {
		return false, call.Err
	}
	return true, nil
}

// NotificationClosedSignal holds data for *Closed callbacks from Notifications Interface.
type NotificationClosedSignal struct {
	// ID of the Notification the signal was invoked for
	ID uint32
	// A reason given if known
	Reason Reason
}

// Reason for the closed notification
type Reason uint32

const (
	// ReasonExpired when a notification expired
	ReasonExpired Reason = 1

	// ReasonDismissedByUser when a notification has been dismissed by a user
	ReasonDismissedByUser Reason = 2

	// ReasonClosedByCall when a notification has been closed by a call to CloseNotification
	ReasonClosedByCall Reason = 3

	// ReasonUnknown when as notification has been closed for an unknown reason
	ReasonUnknown Reason = 4
)

func (r Reason) String() string {
	switch r {
	case ReasonExpired:
		return "Expired"
	case ReasonDismissedByUser:
		return "DismissedByUser"
	case ReasonClosedByCall:
		return "ClosedByCall"
	case ReasonUnknown:
		return "Unknown"
	default:
		return "Other"
	}
}

// Close cleans up and shuts down signal delivery loop
func (n *notifier) Close() error {
	n.done <- true

	// remove signal reception
	n.conn.RemoveSignal(n.signal)

	// unregister in dbus:
	errRemoveMatch := n.conn.RemoveMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)

	close(n.done)

	// wait for eventloop to shut down...
	n.wg.Wait()

	return errRemoveMatch
}

type loggerWrapper struct {
	prefix string
}

func (l *loggerWrapper) Printf(format string, v ...interface{}) {
	log.Printf(l.prefix+format, v...)
}
<commit_msg>Update notification.go<commit_after>package notify

import (
	"errors"
	"fmt"
	"log"
	"sync"

	"github.com/godbus/dbus/v5"
)

const (
	dbusRemoveMatch            = "org.freedesktop.DBus.RemoveMatch"
	dbusAddMatch               = "org.freedesktop.DBus.AddMatch"
	dbusObjectPath             = "/org/freedesktop/Notifications" // the DBUS object path
	dbusNotificationsInterface = "org.freedesktop.Notifications"  // DBUS Interface
	signalNotificationClosed   = "org.freedesktop.Notifications.NotificationClosed"
	signalActionInvoked        = "org.freedesktop.Notifications.ActionInvoked"
	callGetCapabilities        = "org.freedesktop.Notifications.GetCapabilities"
	callCloseNotification      = "org.freedesktop.Notifications.CloseNotification"
	callNotify                 = "org.freedesktop.Notifications.Notify"
	callGetServerInformation   = "org.freedesktop.Notifications.GetServerInformation"

	channelBufferSize = 2
)

// Notification holds all information needed for creating a notification
type Notification struct {
	AppName string
	// Setting ReplacesID atomically replaces the notification with this ID.
	// Optional.
	ReplacesID uint32
	// See predefined icons here: http://standards.freedesktop.org/icon-naming-spec/icon-naming-spec-latest.html
	// Optional.
	AppIcon string
	Summary string
	Body    string
	// Actions are tuples of (action_key, label), e.g.: []string{"cancel", "Cancel", "open", "Open"}
	Actions []string
	Hints   map[string]dbus.Variant
	// ExpireTimeout: milliseconds to show notification
	ExpireTimeout int32
}

// SendNotification is provided for convenience.
// Use if you only want to deliver a notification and dont care about events.
func SendNotification(conn *dbus.Conn, note Notification) (uint32, error) {
	actions := len(note.Actions)
	if (actions % 2) != 0 {
		return 0, errors.New("actions must be pairs of (key, label)")
	}

	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callNotify, 0,
		note.AppName,
		note.ReplacesID,
		note.AppIcon,
		note.Summary,
		note.Body,
		note.Actions,
		note.Hints,
		note.ExpireTimeout)
	if call.Err != nil {
		return 0, fmt.Errorf("error sending notification: %w", call.Err)
	}
	var ret uint32
	err := call.Store(&ret)
	if err != nil {
		return ret, fmt.Errorf("error getting uint32 ret value: %w", err)
	}
	return ret, nil
}

// ServerInformation is a holder for information returned by
// GetServerInformation call.
type ServerInformation struct {
	Name        string
	Vendor      string
	Version     string
	SpecVersion string
}

// GetServerInformation returns the information on the server.
//
// org.freedesktop.Notifications.GetServerInformation
//
//  GetServerInformation Return Values
//
//		Name		 Type	  Description
//		name		 STRING	  The product name of the server.
//		vendor		 STRING	  The vendor name. For example, "KDE," "GNOME," "freedesktop.org," or "Microsoft."
//		version		 STRING	  The server's version number.
//		spec_version STRING	  The specification version the server is compliant with.
//
func GetServerInformation(conn *dbus.Conn) (ServerInformation, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	if obj == nil {
		return ServerInformation{}, errors.New("error creating dbus call object")
	}
	call := obj.Call(callGetServerInformation, 0)
	if call.Err != nil {
		return ServerInformation{}, fmt.Errorf("error calling %v: %v", callGetServerInformation, call.Err)
	}

	ret := ServerInformation{}
	err := call.Store(&ret.Name, &ret.Vendor, &ret.Version, &ret.SpecVersion)
	if err != nil {
		return ret, fmt.Errorf("error reading %v return values: %v", callGetServerInformation, err)
	}
	return ret, nil
}

// GetCapabilities gets the capabilities of the notification server.
// This call takes no parameters.
// It returns an array of strings. Each string describes an optional capability implemented by the server.
//
// See also: https://developer.gnome.org/notification-spec/
// GetCapabilities provide an exported method for this operation
func GetCapabilities(conn *dbus.Conn) ([]string, error) {
	obj := conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callGetCapabilities, 0)
	if call.Err != nil {
		return []string{}, call.Err
	}
	var ret []string
	err := call.Store(&ret)
	if err != nil {
		return ret, err
	}
	return ret, nil
}

// Notifier is an interface implementing the operations supported by the
// Freedesktop DBus Notifications object.
//
// New() sets up a Notifier that listens on dbus' signals regarding
// Notifications: NotificationClosed and ActionInvoked.
//
// Signal delivery works by subscribing to all signals regarding Notifications,
// which means you will see signals for Notifications also from other sources,
// not just the latest you sent
//
// Users that only want to send a simple notification, but don't care about
// interacting with signals, can use exported method: SendNotification(conn, Notification)
//
// Caller is responsible for calling Close() before exiting,
// to shut down event loop and cleanup dbus registration.
type Notifier interface {
	SendNotification(n Notification) (uint32, error)
	GetCapabilities() ([]string, error)
	GetServerInformation() (ServerInformation, error)
	CloseNotification(id uint32) (bool, error)
	Close() error
}

// NotificationClosedHandler is called when we receive a NotificationClosed signal
type NotificationClosedHandler func(*NotificationClosedSignal)

// ActionInvokedHandler is called when we receive a signal that one of the action_keys was invoked.
//
// Note that invoking an action often also produces a NotificationClosedSignal,
// so you might receive both a Closed signal and a ActionInvoked signal.
//
// I suspect this detail is implementation specific for the UI interaction,
// and does at least happen on XFCE4.
type ActionInvokedHandler func(*ActionInvokedSignal)

// ActionInvokedSignal holds data from any signal received regarding Actions invoked
type ActionInvokedSignal struct {
	// ID of the Notification the action was invoked for
	ID uint32
	// Key from the tuple (action_key, label)
	ActionKey string
}

// notifier implements Notifier interface
type notifier struct {
	conn     *dbus.Conn
	signal   chan *dbus.Signal
	done     chan bool
	onClosed NotificationClosedHandler
	onAction ActionInvokedHandler
	wg       *sync.WaitGroup
	log      logger
}

type logger interface {
	Printf(format string, v ...interface{})
}

// option overrides certain parts of a Notifier
type option func(*notifier)

// WithLogger sets a new logger func
func WithLogger(logz logger) option {
	return func(n *notifier) {
		n.log = logz
	}
}

// WithOnAction sets ActionInvokedHandler handler
func WithOnAction(h ActionInvokedHandler) option {
	return func(n *notifier) {
		n.onAction = h
	}
}

// WithOnClosed sets NotificationClosed handler
func WithOnClosed(h NotificationClosedHandler) option {
	return func(n *notifier) {
		n.onClosed = h
	}
}

// New creates a new Notifier using conn.
// See also: Notifier
func New(conn *dbus.Conn, opts ...option) (Notifier, error) {
	n := &notifier{
		conn:     conn,
		signal:   make(chan *dbus.Signal, channelBufferSize),
		done:     make(chan bool),
		wg:       &sync.WaitGroup{},
		onClosed: func(s *NotificationClosedSignal) {},
		onAction: func(s *ActionInvokedSignal) {},
		log:      &loggerWrapper{"notify: "},
	}

	for _, val := range opts {
		val(n)
	}

	// add a listener (matcher) in dbus for signals to Notification interface.
	err := n.conn.AddMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)
	if err != nil {
		return nil, fmt.Errorf("error registering for signals in dbus: %w", err)
	}
	// register in dbus for signal delivery
	n.conn.Signal(n.signal)

	// start eventloop
	go n.eventLoop()

	return n, nil
}

func (n notifier) eventLoop() {
	n.wg.Add(1)
	defer n.wg.Done()
	for {
		select {
		case signal := <-n.signal:
			n.handleSignal(signal)
		case <-n.done:
			n.log.Printf("Got Close() signal, shutting down...")
			return
		}
	}
}

// signal handler that translates and sends notifications to channels
func (n notifier) handleSignal(signal *dbus.Signal) {
	switch signal.Name {
	case signalNotificationClosed:
		nc := &NotificationClosedSignal{
			ID:     signal.Body[0].(uint32),
			Reason: Reason(signal.Body[1].(uint32)),
		}
		n.onClosed(nc)
	case signalActionInvoked:
		is := &ActionInvokedSignal{
			ID:        signal.Body[0].(uint32),
			ActionKey: signal.Body[1].(string),
		}
		n.onAction(is)
	default:
		n.log.Printf("Received unknown signal: %+v", signal)
	}
}

func (n *notifier) GetCapabilities() ([]string, error) {
	return GetCapabilities(n.conn)
}
func (n *notifier) GetServerInformation() (ServerInformation, error) {
	return GetServerInformation(n.conn)
}

// SendNotification sends a notification to the notification server and returns the ID or an error.
//
// Implements dbus call:
//
//     UINT32 org.freedesktop.Notifications.Notify (
//	       STRING app_name,
//	       UINT32 replaces_id,
//	       STRING app_icon,
//	       STRING summary,
//	       STRING body,
//	       ARRAY  actions,
//	       DICT   hints,
//	       INT32  expire_timeout
//     );
//
//		Name	    	Type	Description
//		app_name		STRING	The optional name of the application sending the notification. Can be blank.
//		replaces_id	    UINT32	The optional notification ID that this notification replaces. The server must atomically (ie with no flicker or other visual cues) replace the given notification with this one. This allows clients to effectively modify the notification while it's active. A value of value of 0 means that this notification won't replace any existing notifications.
//		app_icon		STRING	The optional program icon of the calling application. Can be an empty string, indicating no icon.
//		summary		    STRING	The summary text briefly describing the notification.
//		body			STRING	The optional detailed body text. Can be empty.
//		actions		    ARRAY	Actions are sent over as a list of pairs. Each even element in the list (starting at index 0) represents the identifier for the action. Each odd element in the list is the localized string that will be displayed to the user.
//		hints	        DICT	Optional hints that can be passed to the server from the client program. Although clients and servers should never assume each other supports any specific hints, they can be used to pass along information, such as the process PID or window ID, that the server may be able to make use of. See Hints. Can be empty.
//      expire_timeout  INT32   The timeout time in milliseconds since the display of the notification at which the notification should automatically close.
//								If -1, the notification's expiration time is dependent on the notification server's settings, and may vary for the type of notification. If 0, never expire.
//
// If replaces_id is 0, the return value is a UINT32 that represent the notification.
// It is unique, and will not be reused unless a MAXINT number of notifications have been generated.
// An acceptable implementation may just use an incrementing counter for the ID.
// The returned ID is always greater than zero. Servers must make sure not to return zero as an ID.
//
// If replaces_id is not 0, the returned value is the same value as replaces_id.
func (n *notifier) SendNotification(note Notification) (uint32, error) {
	return SendNotification(n.conn, note)
}

// CloseNotification causes a notification to be forcefully closed and removed from the user's view.
// It can be used, for example, in the event that what the notification pertains to is no longer relevant,
// or to cancel a notification with no expiration time.
//
// The NotificationClosed (dbus) signal is emitted by this method.
// If the notification no longer exists, an empty D-BUS Error message is sent back.
func (n *notifier) CloseNotification(id uint32) (bool, error) {
	obj := n.conn.Object(dbusNotificationsInterface, dbusObjectPath)
	call := obj.Call(callCloseNotification, 0, id)
	if call.Err != nil {
		return false, call.Err
	}
	return true, nil
}

// NotificationClosedSignal holds data for *Closed callbacks from Notifications Interface.
type NotificationClosedSignal struct {
	// ID of the Notification the signal was invoked for
	ID uint32
	// A reason given if known
	Reason Reason
}

// Reason for the closed notification
type Reason uint32

const (
	// ReasonExpired when a notification expired
	ReasonExpired Reason = 1

	// ReasonDismissedByUser when a notification has been dismissed by a user
	ReasonDismissedByUser Reason = 2

	// ReasonClosedByCall when a notification has been closed by a call to CloseNotification
	ReasonClosedByCall Reason = 3

	// ReasonUnknown when as notification has been closed for an unknown reason
	ReasonUnknown Reason = 4
)

func (r Reason) String() string {
	switch r {
	case ReasonExpired:
		return "Expired"
	case ReasonDismissedByUser:
		return "DismissedByUser"
	case ReasonClosedByCall:
		return "ClosedByCall"
	case ReasonUnknown:
		return "Unknown"
	default:
		return "Other"
	}
}

// Close cleans up and shuts down signal delivery loop
func (n *notifier) Close() error {
	n.done <- true

	// remove signal reception
	n.conn.RemoveSignal(n.signal)

	// unregister in dbus:
	errRemoveMatch := n.conn.RemoveMatchSignal(
		dbus.WithMatchObjectPath(dbusObjectPath),
		dbus.WithMatchInterface(dbusNotificationsInterface),
	)

	close(n.done)

	// wait for eventloop to shut down...
	n.wg.Wait()

	return errRemoveMatch
}

type loggerWrapper struct {
	prefix string
}

func (l *loggerWrapper) Printf(format string, v ...interface{}) {
	log.Printf(l.prefix+format, v...)
}
<|endoftext|>


<commit_before>package logs

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/state"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

//
// EKSS3LogsClient corresponds with the aws logs driver
// for ECS and returns logs for runs
//
type EKSS3LogsClient struct {
	logRetentionInDays int64
	logNamespace       string
	s3Client           *s3.S3
	s3Bucket           string
	s3BucketRootDir    string
	logger             *log.Logger
	emrS3LogsBucket    string
	emrS3LogsBasePath  string
}

type s3Log struct {
	Log    string    `json:"log"`
	Stream string    `json:"stream"`
	Time   time.Time `json:"time"`
}

//
// Name returns the name of the logs client
//
func (lc *EKSS3LogsClient) Name() string {
	return "eks-s3"
}

//
// Initialize sets up the EKSS3LogsClient
//
func (lc *EKSS3LogsClient) Initialize(conf config.Config) error {
	confLogOptions := conf.GetStringMapString("eks.log.driver.options")

	awsRegion := confLogOptions["awslogs-region"]
	if len(awsRegion) == 0 {
		awsRegion = conf.GetString("aws_default_region")
	}

	if len(awsRegion) == 0 {
		return errors.Errorf(
			"EKSS3LogsClient needs one of [eks.log.driver.options.awslogs-region] or [aws_default_region] set in config")
	}

	flotillaMode := conf.GetString("flotilla_mode")
	if flotillaMode != "test" {
		sess := session.Must(session.NewSession(&aws.Config{
			Region: aws.String(awsRegion)}))

		lc.s3Client = s3.New(sess, aws.NewConfig().WithRegion(awsRegion))
	}
	lc.emrS3LogsBucket = conf.GetString("emr.log.bucket")
	lc.emrS3LogsBasePath = conf.GetString("emr.log.base_path")
	s3BucketName := confLogOptions["s3_bucket_name"]

	if len(s3BucketName) == 0 {
		return errors.Errorf(
			"EKSS3LogsClient needs [eks.log.driver.options.s3_bucket_name] set in config")
	}
	lc.s3Bucket = s3BucketName

	s3BucketRootDir := confLogOptions["s3_bucket_root_dir"]

	if len(s3BucketRootDir) == 0 {
		return errors.Errorf(
			"EKSS3LogsClient needs [eks.log.driver.options.s3_bucket_root_dir] set in config")
	}
	lc.s3BucketRootDir = s3BucketRootDir

	lc.logger = log.New(os.Stderr, "[s3logs] ",
		log.Ldate|log.Ltime|log.Lshortfile)
	return nil
}

func (lc *EKSS3LogsClient) emrLogsToMessageString(run state.Run, lastSeen *string, role *string, facility *string) (string, *string, error) {
	s3DirName, err := lc.emrDriverLogsPath(run)
	if err != nil {
		return "", aws.String(""), errors.Errorf("No logs")
	}

	params := &s3.ListObjectsV2Input{
		Bucket:  aws.String(lc.emrS3LogsBucket),
		Prefix:  aws.String(s3DirName),
		MaxKeys: aws.Int64(1000),
	}

	pageNum := 0
	lastModified := &time.Time{}
	var key *string

	err = lc.s3Client.ListObjectsV2Pages(params,
		func(result *s3.ListObjectsV2Output, lastPage bool) bool {
			pageNum++
			if result != nil {
				for _, content := range result.Contents {
					if strings.Contains(*content.Key, *role) && strings.Contains(*content.Key, *facility) && lastModified.Before(*content.LastModified) {
						if content != nil && *content.Size < int64(10000000) {
							key = content.Key
							lastModified = content.LastModified
						}
					}
				}
			}
			if lastPage {
				return false
			}
			return pageNum <= 10
		})

	if key == nil {
		lc.logger.Println(fmt.Sprintf("run=%s emr logging key not found for role=%s facility=%s", run.RunID, *role, *facility))
		return "", aws.String(""), errors.Errorf("No driver logs found")
	}

	startPosition := int64(0)
	if lastSeen != nil {
		parsed, err := strconv.ParseInt(*lastSeen, 10, 64)
		if err == nil {
			startPosition = parsed
		}
	}

	s3Obj, err := lc.s3Client.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(lc.emrS3LogsBucket),
		Key:    aws.String(*key),
	})

	if s3Obj != nil && err == nil && *s3Obj.ContentLength < int64(10000000) {
		defer s3Obj.Body.Close()
		gr, err := gzip.NewReader(s3Obj.Body)
		if err != nil {
			return "", aws.String(""), errors.Errorf("No driver logs found")
		}
		defer gr.Close()
		reader := bufio.NewReader(gr)
		var b0 bytes.Buffer
		counter := int64(0)
		for {
			line, err := reader.ReadBytes('\n')
			if err != nil {
				if err == io.EOF {
					err = nil
					return b0.String(), aws.String(fmt.Sprintf("%d", counter)), nil
				}

			} else {
				if counter >= startPosition {
					b0.Write(line)
				}
				counter = counter + 1
			}
		}
	}
	return "", aws.String(""), errors.Errorf("No driver logs found")
}

func (lc *EKSS3LogsClient) emrDriverLogsPath(run state.Run) (string, error) {
	if run.SparkExtension.EMRJobId != nil &&
		run.SparkExtension.VirtualClusterId != nil {
		return fmt.Sprintf("%s/%s/jobs/%s/",
			lc.emrS3LogsBasePath,
			*run.SparkExtension.VirtualClusterId,
			*run.SparkExtension.EMRJobId,
		), nil
	}
	return "", errors.New("couldn't construct s3 path.")
}

func (lc *EKSS3LogsClient) Logs(executable state.Executable, run state.Run, lastSeen *string, role *string, facility *string) (string, *string, error) {
	if *run.Engine == state.EKSSparkEngine {
		return lc.emrLogsToMessageString(run, lastSeen, role, facility)
	}

	result, err := lc.getS3Object(run)
	startPosition := int64(0)
	if lastSeen != nil {
		parsed, err := strconv.ParseInt(*lastSeen, 10, 64)
		if err == nil {
			startPosition = parsed
		}
	}

	if result != nil && err == nil {
		acc, position, err := lc.logsToMessageString(result, startPosition)
		newLastSeen := fmt.Sprintf("%d", position)
		return acc, &newLastSeen, err
	}

	return "", aws.String(""), errors.Errorf("No logs.")
}

//
// Logs returns all logs from the log stream identified by handle since lastSeen
//
func (lc *EKSS3LogsClient) LogsText(executable state.Executable, run state.Run, w http.ResponseWriter) error {
	if run.Engine == nil || *run.Engine == state.EKSEngine {
		result, err := lc.getS3Object(run)

		if result != nil && err == nil {
			return lc.logsToMessage(result, w)
		}
	}
	if *run.Engine == state.EKSSparkEngine {
		return lc.logsEMR(w)
	}
	return nil
}

//
// Fetch S3Object associated with the pod's log.
//
func (lc *EKSS3LogsClient) getS3Object(run state.Run) (*s3.GetObjectOutput, error) {
	//Pod isn't there yet - dont return a 404
	//if run.PodName == nil {
	//	return nil, errors.New("no pod associated with the run.")
	//}
	s3DirName := lc.toS3DirName(run)

	// Get list of S3 objects in the run_id folder.
	result, err := lc.s3Client.ListObjects(&s3.ListObjectsInput{
		Bucket: aws.String(lc.s3Bucket),
		Prefix: aws.String(s3DirName),
	})

	if err != nil {
		return nil, errors.Wrap(err, "problem getting logs")
	}

	if result == nil || result.Contents == nil || len(result.Contents) == 0 {
		return nil, errors.New("no s3 files associated with the run.")
	}
	var key *string
	lastModified := &time.Time{}

	//Find latest log file (could have multiple log files per pod - due to pod retries)
	for _, content := range result.Contents {
		if strings.Contains(*content.Key, run.RunID) && lastModified.Before(*content.LastModified) {
			key = content.Key
			lastModified = content.LastModified
		}
	}
	if key != nil {
		return lc.getS3Key(key)
	} else {
		return nil, errors.New("no s3 files associated with the run.")
	}
}

func (lc *EKSS3LogsClient) getS3Key(s3Key *string) (*s3.GetObjectOutput, error) {
	result, err := lc.s3Client.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(lc.s3Bucket),
		Key:    aws.String(*s3Key),
	})
	if err != nil {
		return nil, err
	}
	return result, nil
}

//
// Formulate dir name on S3.
//
func (lc *EKSS3LogsClient) toS3DirName(run state.Run) string {
	return fmt.Sprintf("%s/%s", lc.s3BucketRootDir, run.RunID)
}

//
// Converts log messages from S3 to strings - returns the contents of the entire file.
//
func (lc *EKSS3LogsClient) logsToMessage(result *s3.GetObjectOutput, w http.ResponseWriter) error {
	reader := bufio.NewReader(result.Body)
	for {
		line, err := reader.ReadBytes('\n')
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		} else {
			var parsedLine s3Log
			err := json.Unmarshal(line, &parsedLine)
			if err != nil {
				return err
			}
			_, err = io.WriteString(w, parsedLine.Log)
			if err != nil {
				return err
			}
		}
	}

}

func (lc *EKSS3LogsClient) logsEMR(w http.ResponseWriter) error {
	_, _ = io.WriteString(w, "todo!!!")
	return nil
}

//
// Converts log messages from S3 to strings, takes a starting offset.
//
func (lc *EKSS3LogsClient) logsToMessageString(result *s3.GetObjectOutput, startingPosition int64) (string, int64, error) {
	acc := ""
	currentPosition := int64(0)
	// if less than/equal to 0, read entire log.
	if startingPosition <= 0 {
		startingPosition = currentPosition
	}

	// No S3 file or object, return "", 0, err
	if result == nil {
		return acc, startingPosition, errors.New("s3 object not present.")
	}

	reader := bufio.NewReader(result.Body)

	// Reading until startingPosition and discard unneeded lines.
	for currentPosition < startingPosition {
		currentPosition = currentPosition + 1
		_, err := reader.ReadBytes('\n')
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return acc, startingPosition, err
		}
	}

	// Read upto MaxLogLines
	for currentPosition <= startingPosition+state.MaxLogLines {
		currentPosition = currentPosition + 1
		line, err := reader.ReadBytes('\n')
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return acc, currentPosition, err
		} else {
			var parsedLine s3Log
			err := json.Unmarshal(line, &parsedLine)
			if err == nil {
				acc = fmt.Sprintf("%s%s", acc, parsedLine.Log)
			}
		}
	}

	_ = result.Body.Close()
	return acc, currentPosition, nil
}
<commit_msg>adding more size limiters<commit_after>package logs

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/state"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

//
// EKSS3LogsClient corresponds with the aws logs driver
// for ECS and returns logs for runs
//
type EKSS3LogsClient struct {
	logRetentionInDays int64
	logNamespace       string
	s3Client           *s3.S3
	s3Bucket           string
	s3BucketRootDir    string
	logger             *log.Logger
	emrS3LogsBucket    string
	emrS3LogsBasePath  string
}

type s3Log struct {
	Log    string    `json:"log"`
	Stream string    `json:"stream"`
	Time   time.Time `json:"time"`
}

//
// Name returns the name of the logs client
//
func (lc *EKSS3LogsClient) Name() string {
	return "eks-s3"
}

//
// Initialize sets up the EKSS3LogsClient
//
func (lc *EKSS3LogsClient) Initialize(conf config.Config) error {
	confLogOptions := conf.GetStringMapString("eks.log.driver.options")

	awsRegion := confLogOptions["awslogs-region"]
	if len(awsRegion) == 0 {
		awsRegion = conf.GetString("aws_default_region")
	}

	if len(awsRegion) == 0 {
		return errors.Errorf(
			"EKSS3LogsClient needs one of [eks.log.driver.options.awslogs-region] or [aws_default_region] set in config")
	}

	flotillaMode := conf.GetString("flotilla_mode")
	if flotillaMode != "test" {
		sess := session.Must(session.NewSession(&aws.Config{
			Region: aws.String(awsRegion)}))

		lc.s3Client = s3.New(sess, aws.NewConfig().WithRegion(awsRegion))
	}
	lc.emrS3LogsBucket = conf.GetString("emr.log.bucket")
	lc.emrS3LogsBasePath = conf.GetString("emr.log.base_path")
	s3BucketName := confLogOptions["s3_bucket_name"]

	if len(s3BucketName) == 0 {
		return errors.Errorf(
			"EKSS3LogsClient needs [eks.log.driver.options.s3_bucket_name] set in config")
	}
	lc.s3Bucket = s3BucketName

	s3BucketRootDir := confLogOptions["s3_bucket_root_dir"]

	if len(s3BucketRootDir) == 0 {
		return errors.Errorf(
			"EKSS3LogsClient needs [eks.log.driver.options.s3_bucket_root_dir] set in config")
	}
	lc.s3BucketRootDir = s3BucketRootDir

	lc.logger = log.New(os.Stderr, "[s3logs] ",
		log.Ldate|log.Ltime|log.Lshortfile)
	return nil
}

func (lc *EKSS3LogsClient) emrLogsToMessageString(run state.Run, lastSeen *string, role *string, facility *string) (string, *string, error) {
	s3DirName, err := lc.emrDriverLogsPath(run)
	if err != nil {
		return "", aws.String(""), errors.Errorf("No logs")
	}

	params := &s3.ListObjectsV2Input{
		Bucket:  aws.String(lc.emrS3LogsBucket),
		Prefix:  aws.String(s3DirName),
		MaxKeys: aws.Int64(1000),
	}

	pageNum := 0
	lastModified := &time.Time{}
	var key *string

	err = lc.s3Client.ListObjectsV2Pages(params,
		func(result *s3.ListObjectsV2Output, lastPage bool) bool {
			pageNum++
			if result != nil {
				for _, content := range result.Contents {
					if strings.Contains(*content.Key, *role) && strings.Contains(*content.Key, *facility) && lastModified.Before(*content.LastModified) {
						if content != nil && *content.Size < int64(10000000) {
							key = content.Key
							lastModified = content.LastModified
						}
					}
				}
			}
			if lastPage {
				return false
			}
			return pageNum <= 10
		})

	if key == nil {
		lc.logger.Println(fmt.Sprintf("run=%s emr logging key not found for role=%s facility=%s", run.RunID, *role, *facility))
		return "", aws.String(""), errors.Errorf("No driver logs found")
	}

	startPosition := int64(0)
	if lastSeen != nil {
		parsed, err := strconv.ParseInt(*lastSeen, 10, 64)
		if err == nil {
			startPosition = parsed
		}
	}

	s3Obj, err := lc.s3Client.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(lc.emrS3LogsBucket),
		Key:    aws.String(*key),
	})

	if s3Obj != nil && err == nil && *s3Obj.ContentLength < int64(10000000) {
		defer s3Obj.Body.Close()
		gr, err := gzip.NewReader(s3Obj.Body)
		if err != nil {
			return "", aws.String(""), errors.Errorf("No driver logs found")
		}
		defer gr.Close()
		reader := bufio.NewReader(gr)
		var b0 bytes.Buffer
		counter := int64(0)
		for {
			line, err := reader.ReadBytes('\n')
			if err != nil {
				if err == io.EOF {
					err = nil
					return b0.String(), aws.String(fmt.Sprintf("%d", counter)), nil
				}

			} else {
				if counter >= startPosition {
					b0.Write(line)
				}
				counter = counter + 1
			}
		}
	}
	return "", aws.String(""), errors.Errorf("No driver logs found")
}

func (lc *EKSS3LogsClient) emrDriverLogsPath(run state.Run) (string, error) {
	if run.SparkExtension.EMRJobId != nil &&
		run.SparkExtension.VirtualClusterId != nil {
		return fmt.Sprintf("%s/%s/jobs/%s/",
			lc.emrS3LogsBasePath,
			*run.SparkExtension.VirtualClusterId,
			*run.SparkExtension.EMRJobId,
		), nil
	}
	return "", errors.New("couldn't construct s3 path.")
}

func (lc *EKSS3LogsClient) Logs(executable state.Executable, run state.Run, lastSeen *string, role *string, facility *string) (string, *string, error) {
	if *run.Engine == state.EKSSparkEngine {
		return lc.emrLogsToMessageString(run, lastSeen, role, facility)
	}

	result, err := lc.getS3Object(run)
	startPosition := int64(0)
	if lastSeen != nil {
		parsed, err := strconv.ParseInt(*lastSeen, 10, 64)
		if err == nil {
			startPosition = parsed
		}
	}

	if result != nil && err == nil {
		acc, position, err := lc.logsToMessageString(result, startPosition)
		newLastSeen := fmt.Sprintf("%d", position)
		return acc, &newLastSeen, err
	}

	return "", aws.String(""), errors.Errorf("No logs.")
}

//
// Logs returns all logs from the log stream identified by handle since lastSeen
//
func (lc *EKSS3LogsClient) LogsText(executable state.Executable, run state.Run, w http.ResponseWriter) error {
	if run.Engine == nil || *run.Engine == state.EKSEngine {
		result, err := lc.getS3Object(run)

		if result != nil && err == nil {
			return lc.logsToMessage(result, w)
		}
	}
	if *run.Engine == state.EKSSparkEngine {
		return lc.logsEMR(w)
	}
	return nil
}

//
// Fetch S3Object associated with the pod's log.
//
func (lc *EKSS3LogsClient) getS3Object(run state.Run) (*s3.GetObjectOutput, error) {
	//Pod isn't there yet - dont return a 404
	//if run.PodName == nil {
	//	return nil, errors.New("no pod associated with the run.")
	//}
	s3DirName := lc.toS3DirName(run)

	// Get list of S3 objects in the run_id folder.
	result, err := lc.s3Client.ListObjects(&s3.ListObjectsInput{
		Bucket: aws.String(lc.s3Bucket),
		Prefix: aws.String(s3DirName),
	})

	if err != nil {
		return nil, errors.Wrap(err, "problem getting logs")
	}

	if result == nil || result.Contents == nil || len(result.Contents) == 0 {
		return nil, errors.New("no s3 files associated with the run.")
	}
	var key *string
	lastModified := &time.Time{}

	//Find latest log file (could have multiple log files per pod - due to pod retries)
	for _, content := range result.Contents {
		if strings.Contains(*content.Key, run.RunID) && lastModified.Before(*content.LastModified) {
			if content != nil && *content.Size < int64(10000000) {
				key = content.Key
				lastModified = content.LastModified
			}
		}
	}
	if key != nil {
		return lc.getS3Key(key)
	} else {
		return nil, errors.New("no s3 files associated with the run.")
	}
}

func (lc *EKSS3LogsClient) getS3Key(s3Key *string) (*s3.GetObjectOutput, error) {
	result, err := lc.s3Client.GetObject(&s3.GetObjectInput{
		Bucket: aws.String(lc.s3Bucket),
		Key:    aws.String(*s3Key),
	})
	if err != nil {
		return nil, err
	}
	return result, nil
}

//
// Formulate dir name on S3.
//
func (lc *EKSS3LogsClient) toS3DirName(run state.Run) string {
	return fmt.Sprintf("%s/%s", lc.s3BucketRootDir, run.RunID)
}

//
// Converts log messages from S3 to strings - returns the contents of the entire file.
//
func (lc *EKSS3LogsClient) logsToMessage(result *s3.GetObjectOutput, w http.ResponseWriter) error {
	reader := bufio.NewReader(result.Body)
	for {
		line, err := reader.ReadBytes('\n')
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return err
		} else {
			var parsedLine s3Log
			err := json.Unmarshal(line, &parsedLine)
			if err != nil {
				return err
			}
			_, err = io.WriteString(w, parsedLine.Log)
			if err != nil {
				return err
			}
		}
	}

}

func (lc *EKSS3LogsClient) logsEMR(w http.ResponseWriter) error {
	_, _ = io.WriteString(w, "todo!!!")
	return nil
}

//
// Converts log messages from S3 to strings, takes a starting offset.
//
func (lc *EKSS3LogsClient) logsToMessageString(result *s3.GetObjectOutput, startingPosition int64) (string, int64, error) {
	acc := ""
	currentPosition := int64(0)
	// if less than/equal to 0, read entire log.
	if startingPosition <= 0 {
		startingPosition = currentPosition
	}

	// No S3 file or object, return "", 0, err
	if result == nil {
		return acc, startingPosition, errors.New("s3 object not present.")
	}

	reader := bufio.NewReader(result.Body)

	// Reading until startingPosition and discard unneeded lines.
	for currentPosition < startingPosition {
		currentPosition = currentPosition + 1
		_, err := reader.ReadBytes('\n')
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return acc, startingPosition, err
		}
	}

	// Read upto MaxLogLines
	for currentPosition <= startingPosition+state.MaxLogLines {
		currentPosition = currentPosition + 1
		line, err := reader.ReadBytes('\n')
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return acc, currentPosition, err
		} else {
			var parsedLine s3Log
			err := json.Unmarshal(line, &parsedLine)
			if err == nil {
				acc = fmt.Sprintf("%s%s", acc, parsedLine.Log)
			}
		}
	}

	_ = result.Body.Close()
	return acc, currentPosition, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/renstrom/dedent"
	"github.com/spf13/cobra"

	"k8s.io/client-go/tools/clientcmd"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset"
	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/kubernetes/pkg/util/wait"
	_ "k8s.io/kubernetes/plugin/pkg/scheduler/algorithmprovider"

	"github.com/kubernetes-incubator/cluster-capacity/cmd/cluster-capacity/app/options"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/apiserver"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/apiserver/cache"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/framework"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/framework/store"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/utils"
)

var (
	clusterCapacityLong = dedent.Dedent(`
		Cluster-capacity simulates an API server with initial state copied from the Kubernetes enviroment
		with its configuration specified in KUBECONFIG. The simulated API server tries to schedule the number of
		pods specified by --max-limits flag. If the --max-limits flag is not specified, pods are scheduled until
		the simulated API server runs out of resources.
	`)

	MAXREPORTS = 100
)

func NewClusterCapacityCommand() *cobra.Command {
	opt := options.NewClusterCapacityOptions()
	cmd := &cobra.Command{
		Use:   "cluster-capacity --kubeconfig KUBECONFIG --podspec PODSPEC",
		Short: "Cluster-capacity is used for simulating scheduling of one or multiple pods",
		Long:  clusterCapacityLong,
		Run: func(cmd *cobra.Command, args []string) {
			err := Validate(opt)
			if err != nil {
				fmt.Println(err)
				cmd.Help()
				return
			}
			err = Run(opt)
			if err != nil {
				fmt.Println(err)
			}
		},
	}
	opt.AddFlags(cmd.Flags())
	return cmd
}

func Validate(opt *options.ClusterCapacityOptions) error {

	if len(opt.AdmissionControl) > 0 {
		admissionsNames := strings.Split(opt.AdmissionControl, ",")
		admissionNamesSets := sets.NewString(admissionsNames...)
		if !options.SupportedAdmissionControllers.IsSuperset(admissionNamesSets) {
			return fmt.Errorf("Requested not supported admission control plugin. Supported admission control plugins are: %v",
				strings.Join(options.SupportedAdmissionControllers.List(), ", "))
		}
	}

	if len(opt.PodSpecFile) == 0 {
		return fmt.Errorf("Pod spec file is missing")
	}

	if len(opt.Kubeconfig) == 0 {
		return fmt.Errorf("kubeconfig is missing")
	}

	if opt.ResourceSpaceMode != "" && opt.ResourceSpaceMode != "ResourceSpaceFull" && opt.ResourceSpaceMode != "ResourceSpacePartial" {
		return fmt.Errorf("Resource space mode not recognized. Valid values are: ResourceSpaceFull, ResourceSpacePartial")
	}

	return nil
}

func Run(opt *options.ClusterCapacityOptions) error {
	conf := options.NewClusterCapacityConfig(opt)
	err := conf.ParseAPISpec()
	if err != nil {
		return fmt.Errorf("Failed to parse pod spec file: %v ", err)
	}

	err = conf.SetDefaultScheduler()
	if err != nil {
		return fmt.Errorf("Failed to set default scheduler config: %v ", err)
	}
	err = conf.ParseAdditionalSchedulerConfigs()
	if err != nil {
		return fmt.Errorf("Failed to parse config file: %v ", err)
	}

	master, err := utils.GetMasterFromKubeConfig(conf.Options.Kubeconfig)
	if err != nil {
		return fmt.Errorf("Failed to parse kubeconfig file: %v ", err)
	}
	conf.KubeClient, err = getKubeClient(master, conf.Options.Kubeconfig)

	if err != nil {
		return err
	}

	if opt.Period == 0 {
		report, err := runSimulator(conf, true)
		if err != nil {
			return err
		}
		if err := framework.ClusterCapacityReviewPrint(report, conf.Options.Verbose, conf.Options.OutputFormat); err != nil {
			return fmt.Errorf("Error while printing: %v", err)
		}
		return nil
	}

	conf.Reports = cache.NewCache(MAXREPORTS)

	r := apiserver.NewResource(conf)

	go func() {
		log.Fatal(apiserver.ListenAndServe(r))
	}()

	// sync and watch apiserver
	conf.ResourceStore = store.NewResourceReflectors(conf.KubeClient, wait.NeverStop)

	runSimulation := func(syncWithClient bool) {
		report, err := runSimulator(conf, syncWithClient)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}

		conf.Reports.Add(report)
		r.PutStatus(report)

		if conf.Options.Verbose {
			framework.ClusterCapacityReviewPrint(report, conf.Options.Verbose, conf.Options.OutputFormat)
		}
	}

	runSimulation(true)
	time.Sleep(time.Duration(opt.Period) * time.Second)

	for {
		runSimulation(false)
		time.Sleep(time.Duration(opt.Period) * time.Second)
	}
}

func getKubeClient(master string, config string) (clientset.Interface, error) {
	cfg, err := clientcmd.BuildConfigFromFlags(master, config)
	if err != nil {
		return nil, fmt.Errorf("Unable to build config: %v", err)
	}
	kubeClient, err := clientset.NewForConfig(cfg)
	if err != nil {
		return nil, fmt.Errorf("Invalid API configuration: %v", err)
	}

	if _, err = kubeClient.Discovery().ServerVersion(); err != nil {
		return nil, fmt.Errorf("Unable to get server version: %v\n", err)
	}

	return kubeClient, nil
}

func runSimulator(s *options.ClusterCapacityConfig, syncWithClient bool) (*framework.ClusterCapacityReview, error) {
	mode, err := framework.StringToResourceSpaceMode(s.Options.ResourceSpaceMode)
	if err != nil {
		return nil, err
	}

	cc, err := framework.New(s.DefaultScheduler, s.Pod, s.Options.MaxLimit, mode, s.Options.AdmissionControl)
	if err != nil {
		return nil, err
	}

	for i := 0; i < len(s.Schedulers); i++ {
		if err = cc.AddScheduler(s.Schedulers[i]); err != nil {
			return nil, err
		}
	}

	if syncWithClient {
		err = cc.SyncWithClient(s.KubeClient)
	} else {
		err = cc.SyncWithStore(s.ResourceStore)
	}
	if err != nil {
		return nil, err
	}

	err = cc.Run()
	if err != nil {
		return nil, err
	}

	report := cc.Report()
	return report, nil
}
<commit_msg>Update cmd/cluster-capacity/app/server.go for 1.6 rebase.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package app

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/renstrom/dedent"
	"github.com/spf13/cobra"

	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/tools/clientcmd"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset"
	_ "k8s.io/kubernetes/plugin/pkg/scheduler/algorithmprovider"

	"github.com/kubernetes-incubator/cluster-capacity/cmd/cluster-capacity/app/options"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/apiserver"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/apiserver/cache"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/framework"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/framework/store"
	"github.com/kubernetes-incubator/cluster-capacity/pkg/utils"
)

var (
	clusterCapacityLong = dedent.Dedent(`
		Cluster-capacity simulates an API server with initial state copied from the Kubernetes enviroment
		with its configuration specified in KUBECONFIG. The simulated API server tries to schedule the number of
		pods specified by --max-limits flag. If the --max-limits flag is not specified, pods are scheduled until
		the simulated API server runs out of resources.
	`)

	MAXREPORTS = 100
)

func NewClusterCapacityCommand() *cobra.Command {
	opt := options.NewClusterCapacityOptions()
	cmd := &cobra.Command{
		Use:   "cluster-capacity --kubeconfig KUBECONFIG --podspec PODSPEC",
		Short: "Cluster-capacity is used for simulating scheduling of one or multiple pods",
		Long:  clusterCapacityLong,
		Run: func(cmd *cobra.Command, args []string) {
			err := Validate(opt)
			if err != nil {
				fmt.Println(err)
				cmd.Help()
				return
			}
			err = Run(opt)
			if err != nil {
				fmt.Println(err)
			}
		},
	}
	opt.AddFlags(cmd.Flags())
	return cmd
}

func Validate(opt *options.ClusterCapacityOptions) error {

	if len(opt.AdmissionControl) > 0 {
		admissionsNames := strings.Split(opt.AdmissionControl, ",")
		admissionNamesSets := sets.NewString(admissionsNames...)
		if !options.SupportedAdmissionControllers.IsSuperset(admissionNamesSets) {
			return fmt.Errorf("Requested not supported admission control plugin. Supported admission control plugins are: %v",
				strings.Join(options.SupportedAdmissionControllers.List(), ", "))
		}
	}

	if len(opt.PodSpecFile) == 0 {
		return fmt.Errorf("Pod spec file is missing")
	}

	if len(opt.Kubeconfig) == 0 {
		return fmt.Errorf("kubeconfig is missing")
	}

	if opt.ResourceSpaceMode != "" && opt.ResourceSpaceMode != "ResourceSpaceFull" && opt.ResourceSpaceMode != "ResourceSpacePartial" {
		return fmt.Errorf("Resource space mode not recognized. Valid values are: ResourceSpaceFull, ResourceSpacePartial")
	}

	return nil
}

func Run(opt *options.ClusterCapacityOptions) error {
	conf := options.NewClusterCapacityConfig(opt)
	err := conf.ParseAPISpec()
	if err != nil {
		return fmt.Errorf("Failed to parse pod spec file: %v ", err)
	}

	err = conf.SetDefaultScheduler()
	if err != nil {
		return fmt.Errorf("Failed to set default scheduler config: %v ", err)
	}
	err = conf.ParseAdditionalSchedulerConfigs()
	if err != nil {
		return fmt.Errorf("Failed to parse config file: %v ", err)
	}

	master, err := utils.GetMasterFromKubeConfig(conf.Options.Kubeconfig)
	if err != nil {
		return fmt.Errorf("Failed to parse kubeconfig file: %v ", err)
	}
	conf.KubeClient, err = getKubeClient(master, conf.Options.Kubeconfig)

	if err != nil {
		return err
	}

	if opt.Period == 0 {
		report, err := runSimulator(conf, true)
		if err != nil {
			return err
		}
		if err := framework.ClusterCapacityReviewPrint(report, conf.Options.Verbose, conf.Options.OutputFormat); err != nil {
			return fmt.Errorf("Error while printing: %v", err)
		}
		return nil
	}

	conf.Reports = cache.NewCache(MAXREPORTS)

	r := apiserver.NewResource(conf)

	go func() {
		log.Fatal(apiserver.ListenAndServe(r))
	}()

	// sync and watch apiserver
	conf.ResourceStore = store.NewResourceReflectors(conf.KubeClient, wait.NeverStop)

	runSimulation := func(syncWithClient bool) {
		report, err := runSimulator(conf, syncWithClient)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}

		conf.Reports.Add(report)
		r.PutStatus(report)

		if conf.Options.Verbose {
			framework.ClusterCapacityReviewPrint(report, conf.Options.Verbose, conf.Options.OutputFormat)
		}
	}

	runSimulation(true)
	time.Sleep(time.Duration(opt.Period) * time.Second)

	for {
		runSimulation(false)
		time.Sleep(time.Duration(opt.Period) * time.Second)
	}
}

func getKubeClient(master string, config string) (clientset.Interface, error) {
	cfg, err := clientcmd.BuildConfigFromFlags(master, config)
	if err != nil {
		return nil, fmt.Errorf("Unable to build config: %v", err)
	}
	kubeClient, err := clientset.NewForConfig(cfg)
	if err != nil {
		return nil, fmt.Errorf("Invalid API configuration: %v", err)
	}

	if _, err = kubeClient.Discovery().ServerVersion(); err != nil {
		return nil, fmt.Errorf("Unable to get server version: %v\n", err)
	}

	return kubeClient, nil
}

func runSimulator(s *options.ClusterCapacityConfig, syncWithClient bool) (*framework.ClusterCapacityReview, error) {
	mode, err := framework.StringToResourceSpaceMode(s.Options.ResourceSpaceMode)
	if err != nil {
		return nil, err
	}

	cc, err := framework.New(s.DefaultScheduler, s.Pod, s.Options.MaxLimit, mode, s.Options.AdmissionControl)
	if err != nil {
		return nil, err
	}

	for i := 0; i < len(s.Schedulers); i++ {
		if err = cc.AddScheduler(s.Schedulers[i]); err != nil {
			return nil, err
		}
	}

	if syncWithClient {
		err = cc.SyncWithClient(s.KubeClient)
	} else {
		err = cc.SyncWithStore(s.ResourceStore)
	}
	if err != nil {
		return nil, err
	}

	err = cc.Run()
	if err != nil {
		return nil, err
	}

	report := cc.Report()
	return report, nil
}
<|endoftext|>