<commit_before>package main

import (
	"database/sql"
	"fmt"
	"time"

	"github.com/mattn/go-sqlite3"

	"github.com/lxc/lxd/shared"
)

var (
	// DbErrAlreadyDefined hapens when the given entry already exists,
	// for example a container.
	DbErrAlreadyDefined = fmt.Errorf("The container/snapshot already exists")

	/* NoSuchObjectError is in the case of joins (and probably other) queries,
	 * we don't get back sql.ErrNoRows when no rows are returned, even though we do
	 * on selects without joins. Instead, you can use this error to
	 * propagate up and generate proper 404s to the client when something
	 * isn't found so we don't abuse sql.ErrNoRows any more than we
	 * already do.
	 */
	NoSuchObjectError = fmt.Errorf("No such object")
)

// CURRENT_SCHEMA contains the current SQLite SQL Schema.
const CURRENT_SCHEMA string = `
CREATE TABLE IF NOT EXISTS certificates (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    fingerprint VARCHAR(255) NOT NULL,
    type INTEGER NOT NULL,
    name VARCHAR(255) NOT NULL,
    certificate TEXT NOT NULL,
    UNIQUE (fingerprint)
);
CREATE TABLE IF NOT EXISTS config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    UNIQUE (key)
);
CREATE TABLE IF NOT EXISTS containers (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    architecture INTEGER NOT NULL,
    type INTEGER NOT NULL,
    ephemeral INTEGER NOT NULL DEFAULT 0,
    stateful INTEGER NOT NULL DEFAULT 0,
    creation_date DATETIME,
    last_use_date DATETIME,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS containers_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    container_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    FOREIGN KEY (container_id) REFERENCES containers (id) ON DELETE CASCADE,
    UNIQUE (container_id, key)
);
CREATE TABLE IF NOT EXISTS containers_devices (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    container_id INTEGER NOT NULL,
    name VARCHAR(255) NOT NULL,
    type INTEGER NOT NULL default 0,
    FOREIGN KEY (container_id) REFERENCES containers (id) ON DELETE CASCADE,
    UNIQUE (container_id, name)
);
CREATE TABLE IF NOT EXISTS containers_devices_config (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    container_device_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    FOREIGN KEY (container_device_id) REFERENCES containers_devices (id) ON DELETE CASCADE,
    UNIQUE (container_device_id, key)
);
CREATE TABLE IF NOT EXISTS containers_profiles (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    container_id INTEGER NOT NULL,
    profile_id INTEGER NOT NULL,
    apply_order INTEGER NOT NULL default 0,
    UNIQUE (container_id, profile_id),
    FOREIGN KEY (container_id) REFERENCES containers(id) ON DELETE CASCADE,
    FOREIGN KEY (profile_id) REFERENCES profiles(id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS images (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    cached INTEGER NOT NULL DEFAULT 0,
    fingerprint VARCHAR(255) NOT NULL,
    filename VARCHAR(255) NOT NULL,
    size INTEGER NOT NULL,
    public INTEGER NOT NULL DEFAULT 0,
    auto_update INTEGER NOT NULL DEFAULT 0,
    architecture INTEGER NOT NULL,
    creation_date DATETIME,
    expiry_date DATETIME,
    upload_date DATETIME NOT NULL,
    last_use_date DATETIME,
    UNIQUE (fingerprint)
);
CREATE TABLE IF NOT EXISTS images_aliases (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    image_id INTEGER NOT NULL,
    description VARCHAR(255),
    FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS images_properties (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    image_id INTEGER NOT NULL,
    type INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS images_source (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    image_id INTEGER NOT NULL,
    server TEXT NOT NULL,
    protocol INTEGER NOT NULL,
    certificate TEXT NOT NULL,
    alias VARCHAR(255) NOT NULL,
    FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS networks (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS networks_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    network_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    UNIQUE (network_id, key),
    FOREIGN KEY (network_id) REFERENCES networks (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS patches (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    applied_at DATETIME NOT NULL,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS profiles (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    description TEXT,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS profiles_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    profile_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value VARCHAR(255),
    UNIQUE (profile_id, key),
    FOREIGN KEY (profile_id) REFERENCES profiles(id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS profiles_devices (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    profile_id INTEGER NOT NULL,
    name VARCHAR(255) NOT NULL,
    type INTEGER NOT NULL default 0,
    UNIQUE (profile_id, name),
    FOREIGN KEY (profile_id) REFERENCES profiles (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS profiles_devices_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    profile_device_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    UNIQUE (profile_device_id, key),
    FOREIGN KEY (profile_device_id) REFERENCES profiles_devices (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS schema (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    version INTEGER NOT NULL,
    updated_at DATETIME NOT NULL,
    UNIQUE (version)
);`

// Create the initial (current) schema for a given SQLite DB connection.
func createDb(db *sql.DB) (err error) {
	latestVersion := dbGetSchema(db)

	if latestVersion != 0 {
		return nil
	}

	_, err = db.Exec(CURRENT_SCHEMA)
	if err != nil {
		return err
	}

	// There isn't an entry for schema version, let's put it in.
	insertStmt := `INSERT INTO schema (version, updated_at) values (?, strftime("%s"));`
	_, err = db.Exec(insertStmt, dbGetLatestSchema())
	if err != nil {
		return err
	}

	// Mark all existing patches as applied
	for _, p := range patches {
		dbPatchesMarkApplied(db, p.name)
	}

	err = dbProfileCreateDefault(db)
	if err != nil {
		return err
	}

	return dbProfileCreateDocker(db)
}

func dbGetSchema(db *sql.DB) (v int) {
	arg1 := []interface{}{}
	arg2 := []interface{}{&v}
	q := "SELECT max(version) FROM schema"
	err := dbQueryRowScan(db, q, arg1, arg2)
	if err != nil {
		return 0
	}
	return v
}

func dbGetLatestSchema() int {
	return dbUpdates[len(dbUpdates)-1].version
}

// Create a database connection object and return it.
func initializeDbObject(d *Daemon, path string) (err error) {
	var openPath string

	timeout := 5 // TODO - make this command-line configurable?

	// These are used to tune the transaction BEGIN behavior instead of using the
	// similar "locking_mode" pragma (locking for the whole database connection).
	openPath = fmt.Sprintf("%s?_busy_timeout=%d&_txlock=exclusive", path, timeout*1000)

	// Open the database. If the file doesn't exist it is created.
	d.db, err = sql.Open("sqlite3", openPath)
	if err != nil {
		return err
	}

	// Create the DB if it doesn't exist.
	err = createDb(d.db)
	if err != nil {
		return fmt.Errorf("Error creating database: %s", err)
	}

	// Run PRAGMA statements now since they are *per-connection*.
	d.db.Exec("PRAGMA foreign_keys=ON;") // This allows us to use ON DELETE CASCADE

	// Apply any update
	err = dbUpdatesApplyAll(d)
	if err != nil {
		return err
	}

	return nil
}

func isDbLockedError(err error) bool {
	if err == nil {
		return false
	}
	if err == sqlite3.ErrLocked || err == sqlite3.ErrBusy {
		return true
	}
	if err.Error() == "database is locked" {
		return true
	}
	return false
}

func isNoMatchError(err error) bool {
	if err == nil {
		return false
	}
	if err.Error() == "sql: no rows in result set" {
		return true
	}
	return false
}

func dbBegin(db *sql.DB) (*sql.Tx, error) {
	for i := 0; i < 100; i++ {
		tx, err := db.Begin()
		if err == nil {
			return tx, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbBegin: error %q", err)
			return nil, err
		}
		time.Sleep(100 * time.Millisecond)
	}

	shared.LogDebugf("DbBegin: DB still locked")
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}

func txCommit(tx *sql.Tx) error {
	for i := 0; i < 100; i++ {
		err := tx.Commit()
		if err == nil {
			return nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("Txcommit: error %q", err)
			return err
		}
		time.Sleep(100 * time.Millisecond)
	}

	shared.LogDebugf("Txcommit: db still locked")
	shared.PrintStack()
	return fmt.Errorf("DB is locked")
}

func dbQueryRowScan(db *sql.DB, q string, args []interface{}, outargs []interface{}) error {
	for i := 0; i < 100; i++ {
		err := db.QueryRow(q, args...).Scan(outargs...)
		if err == nil {
			return nil
		}
		if isNoMatchError(err) {
			return err
		}
		if !isDbLockedError(err) {
			return err
		}
		time.Sleep(100 * time.Millisecond)
	}

	shared.LogDebugf("DbQueryRowScan: query %q args %q, DB still locked", q, args)
	shared.PrintStack()
	return fmt.Errorf("DB is locked")
}

func dbQuery(db *sql.DB, q string, args ...interface{}) (*sql.Rows, error) {
	for i := 0; i < 100; i++ {
		result, err := db.Query(q, args...)
		if err == nil {
			return result, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbQuery: query %q error %q", q, err)
			return nil, err
		}
		time.Sleep(100 * time.Millisecond)
	}

	shared.LogDebugf("DbQuery: query %q args %q, DB still locked", q, args)
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}

func doDbQueryScan(db *sql.DB, q string, args []interface{}, outargs []interface{}) ([][]interface{}, error) {
	rows, err := db.Query(q, args...)
	if err != nil {
		return [][]interface{}{}, err
	}
	defer rows.Close()
	result := [][]interface{}{}
	for rows.Next() {
		ptrargs := make([]interface{}, len(outargs))
		for i := range outargs {
			switch t := outargs[i].(type) {
			case string:
				str := ""
				ptrargs[i] = &str
			case int:
				integer := 0
				ptrargs[i] = &integer
			case int64:
				integer := int64(0)
				ptrargs[i] = &integer
			default:
				return [][]interface{}{}, fmt.Errorf("Bad interface type: %s", t)
			}
		}
		err = rows.Scan(ptrargs...)
		if err != nil {
			return [][]interface{}{}, err
		}
		newargs := make([]interface{}, len(outargs))
		for i := range ptrargs {
			switch t := outargs[i].(type) {
			case string:
				newargs[i] = *ptrargs[i].(*string)
			case int:
				newargs[i] = *ptrargs[i].(*int)
			case int64:
				newargs[i] = *ptrargs[i].(*int64)
			default:
				return [][]interface{}{}, fmt.Errorf("Bad interface type: %s", t)
			}
		}
		result = append(result, newargs)
	}
	err = rows.Err()
	if err != nil {
		return [][]interface{}{}, err
	}
	return result, nil
}

/*
 * . q is the database query
 * . inargs is an array of interfaces containing the query arguments
 * . outfmt is an array of interfaces containing the right types of output
 *   arguments, i.e.
 *      var arg1 string
 *      var arg2 int
 *      outfmt := {}interface{}{arg1, arg2}
 *
 * The result will be an array (one per output row) of arrays (one per output argument)
 * of interfaces, containing pointers to the actual output arguments.
 */
func dbQueryScan(db *sql.DB, q string, inargs []interface{}, outfmt []interface{}) ([][]interface{}, error) {
	for i := 0; i < 100; i++ {
		result, err := doDbQueryScan(db, q, inargs, outfmt)
		if err == nil {
			return result, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbQuery: query %q error %q", q, err)
			return nil, err
		}
		time.Sleep(100 * time.Millisecond)
	}

	shared.LogDebugf("DbQueryscan: query %q inargs %q, DB still locked", q, inargs)
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}

func dbExec(db *sql.DB, q string, args ...interface{}) (sql.Result, error) {
	for i := 0; i < 100; i++ {
		result, err := db.Exec(q, args...)
		if err == nil {
			return result, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbExec: query %q error %q", q, err)
			return nil, err
		}
		time.Sleep(100 * time.Millisecond)
	}

	shared.LogDebugf("DbExec: query %q args %q, DB still locked", q, args)
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}
<commit_msg>Raise DB lock timeout to 30s, retry every 30ms<commit_after>package main

import (
	"database/sql"
	"fmt"
	"time"

	"github.com/mattn/go-sqlite3"

	"github.com/lxc/lxd/shared"
)

var (
	// DbErrAlreadyDefined hapens when the given entry already exists,
	// for example a container.
	DbErrAlreadyDefined = fmt.Errorf("The container/snapshot already exists")

	/* NoSuchObjectError is in the case of joins (and probably other) queries,
	 * we don't get back sql.ErrNoRows when no rows are returned, even though we do
	 * on selects without joins. Instead, you can use this error to
	 * propagate up and generate proper 404s to the client when something
	 * isn't found so we don't abuse sql.ErrNoRows any more than we
	 * already do.
	 */
	NoSuchObjectError = fmt.Errorf("No such object")
)

// CURRENT_SCHEMA contains the current SQLite SQL Schema.
const CURRENT_SCHEMA string = `
CREATE TABLE IF NOT EXISTS certificates (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    fingerprint VARCHAR(255) NOT NULL,
    type INTEGER NOT NULL,
    name VARCHAR(255) NOT NULL,
    certificate TEXT NOT NULL,
    UNIQUE (fingerprint)
);
CREATE TABLE IF NOT EXISTS config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    UNIQUE (key)
);
CREATE TABLE IF NOT EXISTS containers (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    architecture INTEGER NOT NULL,
    type INTEGER NOT NULL,
    ephemeral INTEGER NOT NULL DEFAULT 0,
    stateful INTEGER NOT NULL DEFAULT 0,
    creation_date DATETIME,
    last_use_date DATETIME,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS containers_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    container_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    FOREIGN KEY (container_id) REFERENCES containers (id) ON DELETE CASCADE,
    UNIQUE (container_id, key)
);
CREATE TABLE IF NOT EXISTS containers_devices (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    container_id INTEGER NOT NULL,
    name VARCHAR(255) NOT NULL,
    type INTEGER NOT NULL default 0,
    FOREIGN KEY (container_id) REFERENCES containers (id) ON DELETE CASCADE,
    UNIQUE (container_id, name)
);
CREATE TABLE IF NOT EXISTS containers_devices_config (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    container_device_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    FOREIGN KEY (container_device_id) REFERENCES containers_devices (id) ON DELETE CASCADE,
    UNIQUE (container_device_id, key)
);
CREATE TABLE IF NOT EXISTS containers_profiles (
    id INTEGER primary key AUTOINCREMENT NOT NULL,
    container_id INTEGER NOT NULL,
    profile_id INTEGER NOT NULL,
    apply_order INTEGER NOT NULL default 0,
    UNIQUE (container_id, profile_id),
    FOREIGN KEY (container_id) REFERENCES containers(id) ON DELETE CASCADE,
    FOREIGN KEY (profile_id) REFERENCES profiles(id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS images (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    cached INTEGER NOT NULL DEFAULT 0,
    fingerprint VARCHAR(255) NOT NULL,
    filename VARCHAR(255) NOT NULL,
    size INTEGER NOT NULL,
    public INTEGER NOT NULL DEFAULT 0,
    auto_update INTEGER NOT NULL DEFAULT 0,
    architecture INTEGER NOT NULL,
    creation_date DATETIME,
    expiry_date DATETIME,
    upload_date DATETIME NOT NULL,
    last_use_date DATETIME,
    UNIQUE (fingerprint)
);
CREATE TABLE IF NOT EXISTS images_aliases (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    image_id INTEGER NOT NULL,
    description VARCHAR(255),
    FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS images_properties (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    image_id INTEGER NOT NULL,
    type INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS images_source (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    image_id INTEGER NOT NULL,
    server TEXT NOT NULL,
    protocol INTEGER NOT NULL,
    certificate TEXT NOT NULL,
    alias VARCHAR(255) NOT NULL,
    FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS networks (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS networks_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    network_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    UNIQUE (network_id, key),
    FOREIGN KEY (network_id) REFERENCES networks (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS patches (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    applied_at DATETIME NOT NULL,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS profiles (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    name VARCHAR(255) NOT NULL,
    description TEXT,
    UNIQUE (name)
);
CREATE TABLE IF NOT EXISTS profiles_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    profile_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value VARCHAR(255),
    UNIQUE (profile_id, key),
    FOREIGN KEY (profile_id) REFERENCES profiles(id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS profiles_devices (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    profile_id INTEGER NOT NULL,
    name VARCHAR(255) NOT NULL,
    type INTEGER NOT NULL default 0,
    UNIQUE (profile_id, name),
    FOREIGN KEY (profile_id) REFERENCES profiles (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS profiles_devices_config (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    profile_device_id INTEGER NOT NULL,
    key VARCHAR(255) NOT NULL,
    value TEXT,
    UNIQUE (profile_device_id, key),
    FOREIGN KEY (profile_device_id) REFERENCES profiles_devices (id) ON DELETE CASCADE
);
CREATE TABLE IF NOT EXISTS schema (
    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
    version INTEGER NOT NULL,
    updated_at DATETIME NOT NULL,
    UNIQUE (version)
);`

// Create the initial (current) schema for a given SQLite DB connection.
func createDb(db *sql.DB) (err error) {
	latestVersion := dbGetSchema(db)

	if latestVersion != 0 {
		return nil
	}

	_, err = db.Exec(CURRENT_SCHEMA)
	if err != nil {
		return err
	}

	// There isn't an entry for schema version, let's put it in.
	insertStmt := `INSERT INTO schema (version, updated_at) values (?, strftime("%s"));`
	_, err = db.Exec(insertStmt, dbGetLatestSchema())
	if err != nil {
		return err
	}

	// Mark all existing patches as applied
	for _, p := range patches {
		dbPatchesMarkApplied(db, p.name)
	}

	err = dbProfileCreateDefault(db)
	if err != nil {
		return err
	}

	return dbProfileCreateDocker(db)
}

func dbGetSchema(db *sql.DB) (v int) {
	arg1 := []interface{}{}
	arg2 := []interface{}{&v}
	q := "SELECT max(version) FROM schema"
	err := dbQueryRowScan(db, q, arg1, arg2)
	if err != nil {
		return 0
	}
	return v
}

func dbGetLatestSchema() int {
	return dbUpdates[len(dbUpdates)-1].version
}

// Create a database connection object and return it.
func initializeDbObject(d *Daemon, path string) (err error) {
	var openPath string

	timeout := 5 // TODO - make this command-line configurable?

	// These are used to tune the transaction BEGIN behavior instead of using the
	// similar "locking_mode" pragma (locking for the whole database connection).
	openPath = fmt.Sprintf("%s?_busy_timeout=%d&_txlock=exclusive", path, timeout*1000)

	// Open the database. If the file doesn't exist it is created.
	d.db, err = sql.Open("sqlite3", openPath)
	if err != nil {
		return err
	}

	// Create the DB if it doesn't exist.
	err = createDb(d.db)
	if err != nil {
		return fmt.Errorf("Error creating database: %s", err)
	}

	// Run PRAGMA statements now since they are *per-connection*.
	d.db.Exec("PRAGMA foreign_keys=ON;") // This allows us to use ON DELETE CASCADE

	// Apply any update
	err = dbUpdatesApplyAll(d)
	if err != nil {
		return err
	}

	return nil
}

func isDbLockedError(err error) bool {
	if err == nil {
		return false
	}
	if err == sqlite3.ErrLocked || err == sqlite3.ErrBusy {
		return true
	}
	if err.Error() == "database is locked" {
		return true
	}
	return false
}

func isNoMatchError(err error) bool {
	if err == nil {
		return false
	}
	if err.Error() == "sql: no rows in result set" {
		return true
	}
	return false
}

func dbBegin(db *sql.DB) (*sql.Tx, error) {
	for i := 0; i < 1000; i++ {
		tx, err := db.Begin()
		if err == nil {
			return tx, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbBegin: error %q", err)
			return nil, err
		}
		time.Sleep(30 * time.Millisecond)
	}

	shared.LogDebugf("DbBegin: DB still locked")
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}

func txCommit(tx *sql.Tx) error {
	for i := 0; i < 1000; i++ {
		err := tx.Commit()
		if err == nil {
			return nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("Txcommit: error %q", err)
			return err
		}
		time.Sleep(30 * time.Millisecond)
	}

	shared.LogDebugf("Txcommit: db still locked")
	shared.PrintStack()
	return fmt.Errorf("DB is locked")
}

func dbQueryRowScan(db *sql.DB, q string, args []interface{}, outargs []interface{}) error {
	for i := 0; i < 1000; i++ {
		err := db.QueryRow(q, args...).Scan(outargs...)
		if err == nil {
			return nil
		}
		if isNoMatchError(err) {
			return err
		}
		if !isDbLockedError(err) {
			return err
		}
		time.Sleep(30 * time.Millisecond)
	}

	shared.LogDebugf("DbQueryRowScan: query %q args %q, DB still locked", q, args)
	shared.PrintStack()
	return fmt.Errorf("DB is locked")
}

func dbQuery(db *sql.DB, q string, args ...interface{}) (*sql.Rows, error) {
	for i := 0; i < 1000; i++ {
		result, err := db.Query(q, args...)
		if err == nil {
			return result, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbQuery: query %q error %q", q, err)
			return nil, err
		}
		time.Sleep(30 * time.Millisecond)
	}

	shared.LogDebugf("DbQuery: query %q args %q, DB still locked", q, args)
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}

func doDbQueryScan(db *sql.DB, q string, args []interface{}, outargs []interface{}) ([][]interface{}, error) {
	rows, err := db.Query(q, args...)
	if err != nil {
		return [][]interface{}{}, err
	}
	defer rows.Close()
	result := [][]interface{}{}
	for rows.Next() {
		ptrargs := make([]interface{}, len(outargs))
		for i := range outargs {
			switch t := outargs[i].(type) {
			case string:
				str := ""
				ptrargs[i] = &str
			case int:
				integer := 0
				ptrargs[i] = &integer
			case int64:
				integer := int64(0)
				ptrargs[i] = &integer
			default:
				return [][]interface{}{}, fmt.Errorf("Bad interface type: %s", t)
			}
		}
		err = rows.Scan(ptrargs...)
		if err != nil {
			return [][]interface{}{}, err
		}
		newargs := make([]interface{}, len(outargs))
		for i := range ptrargs {
			switch t := outargs[i].(type) {
			case string:
				newargs[i] = *ptrargs[i].(*string)
			case int:
				newargs[i] = *ptrargs[i].(*int)
			case int64:
				newargs[i] = *ptrargs[i].(*int64)
			default:
				return [][]interface{}{}, fmt.Errorf("Bad interface type: %s", t)
			}
		}
		result = append(result, newargs)
	}
	err = rows.Err()
	if err != nil {
		return [][]interface{}{}, err
	}
	return result, nil
}

/*
 * . q is the database query
 * . inargs is an array of interfaces containing the query arguments
 * . outfmt is an array of interfaces containing the right types of output
 *   arguments, i.e.
 *      var arg1 string
 *      var arg2 int
 *      outfmt := {}interface{}{arg1, arg2}
 *
 * The result will be an array (one per output row) of arrays (one per output argument)
 * of interfaces, containing pointers to the actual output arguments.
 */
func dbQueryScan(db *sql.DB, q string, inargs []interface{}, outfmt []interface{}) ([][]interface{}, error) {
	for i := 0; i < 1000; i++ {
		result, err := doDbQueryScan(db, q, inargs, outfmt)
		if err == nil {
			return result, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbQuery: query %q error %q", q, err)
			return nil, err
		}
		time.Sleep(30 * time.Millisecond)
	}

	shared.LogDebugf("DbQueryscan: query %q inargs %q, DB still locked", q, inargs)
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}

func dbExec(db *sql.DB, q string, args ...interface{}) (sql.Result, error) {
	for i := 0; i < 1000; i++ {
		result, err := db.Exec(q, args...)
		if err == nil {
			return result, nil
		}
		if !isDbLockedError(err) {
			shared.LogDebugf("DbExec: query %q error %q", q, err)
			return nil, err
		}
		time.Sleep(30 * time.Millisecond)
	}

	shared.LogDebugf("DbExec: query %q args %q, DB still locked", q, args)
	shared.PrintStack()
	return nil, fmt.Errorf("DB is locked")
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"os"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/route53"
)

const (
	// StateAvailable is the AWS string returned when machine is available
	StateAvailable = ec2.StateAvailable
	// Ubuntu1604LTSEast is the AMI for Ubuntu 16.04 LTS
	Ubuntu1604LTSEast = AMI("ami-29f96d3e")
	// CentOS7East is the AMI for CentOS 7
	CentOS7East = AMI("ami-6d1c2007")
	// RedHat7East is the AMI for RedHat 7
	RedHat7East = AMI("ami-b63769a1")
	// T2Micro is the T2 Micro instance type
	T2Micro = InstanceType(ec2.InstanceTypeT2Micro)
	// T2Medium is the T2 Medium instance type
	T2Medium = InstanceType(ec2.InstanceTypeT2Medium)
	// exponentialBackoffMaxAttempts is the number of times will try before failing
	// Exponential backoff for AWS eventual consistency
	exponentialBackoffMaxEC2Attempts     = 5
	exponentialBackoffMaxRoute53Attempts = 10
)

// A Node on AWS
type Node struct {
	PrivateDNSName string
	PrivateIP      string
	PublicIP       string
	SSHUser        string
	State          string
}

// DNSRecord in Router53 on AWS
type DNSRecord struct {
	Name   string
	Values []string
}

// AMI is the Amazon Machine Image
type AMI string

// InstanceType is the type of the Amazon machine
type InstanceType string

// ClientConfig of the AWS client
type ClientConfig struct {
	Region          string
	SubnetID        string
	Keyname         string
	SecurityGroupID string
	HostedZoneID    string
}

// Credentials to be used for accessing the AI
type Credentials struct {
	ID     string
	Secret string
}

// Client for provisioning machines on AWS
type Client struct {
	Config      ClientConfig
	Credentials Credentials
	session     *session.Session
}

func (c *Client) getEC2APIClient() (*ec2.EC2, error) {
	if err := c.prepareSession(); err != nil {
		return nil, err
	}
	return ec2.New(c.session), nil
}

func (c *Client) getRoute53APIClient() (*route53.Route53, error) {
	if err := c.prepareSession(); err != nil {
		return nil, err
	}
	return route53.New(c.session), nil
}

func (c *Client) prepareSession() error {
	if c.session == nil {
		creds := credentials.NewStaticCredentials(c.Credentials.ID, c.Credentials.Secret, "")
		_, err := creds.Get()
		if err != nil {
			return fmt.Errorf("Error with credentials provided: %v", err)
		}
		config := aws.NewConfig().WithRegion(c.Config.Region).WithCredentials(creds).WithMaxRetries(10)
		c.session = session.New(config)
	}
	return nil
}

// CreateNode is for creating a machine on AWS using the given AMI and InstanceType.
// Returns the ID of the newly created machine.
func (c Client) CreateNode(ami AMI, instanceType InstanceType) (string, error) {
	api, err := c.getEC2APIClient()
	if err != nil {
		return "", err
	}
	req := &ec2.RunInstancesInput{
		ImageId: aws.String(string(ami)),
		BlockDeviceMappings: []*ec2.BlockDeviceMapping{
			{
				DeviceName: aws.String("/dev/sda1"),
				Ebs: &ec2.EbsBlockDevice{
					DeleteOnTermination: aws.Bool(true),
					VolumeSize:          aws.Int64(10),
				},
			},
		},
		InstanceType:     aws.String(string(instanceType)),
		MinCount:         aws.Int64(1),
		MaxCount:         aws.Int64(1),
		SubnetId:         aws.String(c.Config.SubnetID),
		KeyName:          aws.String(c.Config.Keyname),
		SecurityGroupIds: []*string{aws.String(c.Config.SecurityGroupID)},
	}
	res, err := api.RunInstances(req)
	if err != nil {
		return "", err
	}
	instanceID := res.Instances[0].InstanceId
	// Modify the node
	modifyReq := &ec2.ModifyInstanceAttributeInput{
		InstanceId: instanceID,
		SourceDestCheck: &ec2.AttributeBooleanValue{
			Value: aws.Bool(false),
		},
	}
	err = RetryWithBackoff(func() error {
		var err2 error
		_, err2 = api.ModifyInstanceAttribute(modifyReq)
		return err2
	}, exponentialBackoffMaxEC2Attempts)
	if err != nil {
		fmt.Println("Failed to modify instance attributes")
		if err = c.DestroyNodes([]string{*instanceID}); err != nil {
			fmt.Printf("AWS NODE %q MUST BE CLEANED UP MANUALLY\n", *instanceID)
		}
		return "", err
	}
	// Tag the nodes
	thisHost, _ := os.Hostname()
	tagReq := &ec2.CreateTagsInput{
		Resources: []*string{instanceID},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("ApprendaTeam"),
				Value: aws.String("Kismatic"),
			},
			{
				Key:   aws.String("CreatedBy"),
				Value: aws.String(thisHost),
			},
		},
	}
	err = RetryWithBackoff(func() error {
		var err2 error
		_, err2 = api.CreateTags(tagReq)
		return err2
	}, exponentialBackoffMaxEC2Attempts)
	if err != nil {
		fmt.Println("Failed to tag instance")
		if err = c.DestroyNodes([]string{*instanceID}); err != nil {
			fmt.Printf("AWS NODE %q MUST BE CLEANED UP MANUALLY\n", *instanceID)
		}
		return "", err
	}
	return *res.Instances[0].InstanceId, nil
}

// GetNode returns information about a specific node. The consumer of this method
// is responsible for checking that the information it needs has been returned
// in the Node. (i.e. it's possible for the hostname, public IP to be empty)
func (c Client) GetNode(id string) (*Node, error) {
	api, err := c.getEC2APIClient()
	if err != nil {
		return nil, err
	}
	req := &ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(id)},
	}
	var resp *ec2.DescribeInstancesOutput
	err = RetryWithBackoff(func() error {
		var err2 error
		resp, err2 = api.DescribeInstances(req)
		return err2
	}, exponentialBackoffMaxEC2Attempts)
	if err != nil {
		fmt.Println("Failed to get node information")
		return nil, err
	}
	if len(resp.Reservations) != 1 {
		return nil, fmt.Errorf("Attempted to get a single node, but API returned %d reservations", len(resp.Reservations))
	}
	if len(resp.Reservations[0].Instances) != 1 {
		return nil, fmt.Errorf("Attempted to get a single node, but API returned %d instances", len(resp.Reservations[0].Instances))
	}
	instance := resp.Reservations[0].Instances[0]

	var privateDNSName string
	if instance.PrivateDnsName != nil {
		privateDNSName = *instance.PrivateDnsName
	}
	var privateIP string
	if instance.PrivateIpAddress != nil {
		privateIP = *instance.PrivateIpAddress
	}
	var publicIP string
	if instance.PublicIpAddress != nil {
		publicIP = *instance.PublicIpAddress
	}

	return &Node{
		PrivateDNSName: privateDNSName,
		PrivateIP:      privateIP,
		PublicIP:       publicIP,
		SSHUser:        defaultSSHUserForAMI(AMI(*instance.ImageId)),
		State:          instance.State.GoString(),
	}, nil
}

// DestroyNodes destroys the nodes identified by the ID.
func (c Client) DestroyNodes(nodeIDs []string) error {
	api, err := c.getEC2APIClient()
	if err != nil {
		return err
	}
	req := &ec2.TerminateInstancesInput{
		InstanceIds: aws.StringSlice(nodeIDs),
	}
	_, err = api.TerminateInstances(req)
	if err != nil {
		return err
	}
	return nil
}

// CreateDNSRecords generates an Route53 configured with the master nodes
func (c Client) CreateDNSRecords(nodeIPs []string) (*DNSRecord, error) {
	api, err := c.getRoute53APIClient()
	if err != nil {
		return nil, err
	}
	// Setup variables to modify hosted zone
	name := strconv.FormatInt(time.Now().Unix(), 10) + ".kismatic.integration."
	dnsRecord := &DNSRecord{Name: name, Values: nodeIPs}
	err = modifyHostedZone(dnsRecord, route53.ChangeActionUpsert, c.Config.HostedZoneID, api)
	if err != nil {
		return nil, err
	}

	return dnsRecord, nil
}

// GetDNSRecords returns all Record Sets for the Hosted Zone
func (c Client) GetDNSRecords() ([]*route53.ResourceRecordSet, error) {
	api, err := c.getRoute53APIClient()
	if err != nil {
		return nil, err
	}
	req := &route53.ListResourceRecordSetsInput{
		HostedZoneId: aws.String(c.Config.HostedZoneID), // Required
	}

	resp, err := api.ListResourceRecordSets(req)
	if err != nil {
		return nil, err
	}
	return resp.ResourceRecordSets, nil
}

// DeleteDNSRecords deletes the specified Record Set
func (c Client) DeleteDNSRecords(dnsRecord *DNSRecord) error {
	api, err := c.getRoute53APIClient()
	if err != nil {
		return err
	}
	err = modifyHostedZone(dnsRecord, route53.ChangeActionDelete, c.Config.HostedZoneID, api)
	if err != nil {
		return err
	}

	return nil
}

func modifyHostedZone(record *DNSRecord, action string, hostedZoneID string, api *route53.Route53) error {
	if len(record.Values) == 0 && action != route53.ChangeActionDelete {
		return fmt.Errorf("Values cannot be empty when setting up DNS records")
	}
	var records []*route53.ResourceRecord
	for _, value := range record.Values {
		records = append(records, &route53.ResourceRecord{Value: aws.String(value)})
	}

	req := &route53.ChangeResourceRecordSetsInput{
		ChangeBatch: &route53.ChangeBatch{
			Changes: []*route53.Change{
				{
					Action: aws.String(action),
					ResourceRecordSet: &route53.ResourceRecordSet{
						Name:            aws.String(record.Name),
						Type:            aws.String(route53.RRTypeA),
						ResourceRecords: records,
						TTL:             aws.Int64(300),
					},
				},
			},
		},
		HostedZoneId: aws.String(hostedZoneID),
	}

	resp, err := api.ChangeResourceRecordSets(req)
	if err != nil {
		return err
	}
	changeID := resp.ChangeInfo.Id
	if changeID == nil || *changeID == "" {
		return fmt.Errorf("Something went wrong, DNS change ID is nil")
	}

	changeReq := &route53.GetChangeInput{
		Id: aws.String(*changeID),
	}

	err = RetryWithBackoff(func() error {
		changeResp, err2 := api.GetChange(changeReq)
		if err2 != nil {
			return err2
		}
		if *changeResp.ChangeInfo.Status != route53.ChangeStatusInsync {
			return fmt.Errorf("DNS change status is still %s, took too long", *changeResp.ChangeInfo.Status)
		}
		return nil
	}, exponentialBackoffMaxRoute53Attempts)
	if err != nil {
		return err
	}

	return nil
}

// RetryWithBackoff will retry a function specified number of times
func RetryWithBackoff(fn func() error, retries uint) error {
	var attempts uint
	var err error
	for {
		err = fn()
		if err == nil {
			break
		}
		if attempts == retries {
			break
		}
		time.Sleep((1 << attempts) * time.Second)
		attempts++
	}
	return err
}

func defaultSSHUserForAMI(ami AMI) string {
	switch ami {
	case Ubuntu1604LTSEast:
		return "ubuntu"
	case CentOS7East:
		return "centos"
	case RedHat7East:
		return "ec2-user"
	default:
		panic(fmt.Sprintf("unsupported AMI: %q", ami))
	}
}
<commit_msg>Reduce the number of Route53 retries to 5<commit_after>package aws

import (
	"fmt"
	"os"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/route53"
)

const (
	// StateAvailable is the AWS string returned when machine is available
	StateAvailable = ec2.StateAvailable
	// Ubuntu1604LTSEast is the AMI for Ubuntu 16.04 LTS
	Ubuntu1604LTSEast = AMI("ami-29f96d3e")
	// CentOS7East is the AMI for CentOS 7
	CentOS7East = AMI("ami-6d1c2007")
	// RedHat7East is the AMI for RedHat 7
	RedHat7East = AMI("ami-b63769a1")
	// T2Micro is the T2 Micro instance type
	T2Micro = InstanceType(ec2.InstanceTypeT2Micro)
	// T2Medium is the T2 Medium instance type
	T2Medium = InstanceType(ec2.InstanceTypeT2Medium)
	// exponentialBackoffMaxAttempts is the number of times will try before failing
	// Exponential backoff for AWS eventual consistency
	exponentialBackoffMaxEC2Attempts     = 5
	exponentialBackoffMaxRoute53Attempts = 5
)

// A Node on AWS
type Node struct {
	PrivateDNSName string
	PrivateIP      string
	PublicIP       string
	SSHUser        string
	State          string
}

// DNSRecord in Router53 on AWS
type DNSRecord struct {
	Name   string
	Values []string
}

// AMI is the Amazon Machine Image
type AMI string

// InstanceType is the type of the Amazon machine
type InstanceType string

// ClientConfig of the AWS client
type ClientConfig struct {
	Region          string
	SubnetID        string
	Keyname         string
	SecurityGroupID string
	HostedZoneID    string
}

// Credentials to be used for accessing the AI
type Credentials struct {
	ID     string
	Secret string
}

// Client for provisioning machines on AWS
type Client struct {
	Config      ClientConfig
	Credentials Credentials
	session     *session.Session
}

func (c *Client) getEC2APIClient() (*ec2.EC2, error) {
	if err := c.prepareSession(); err != nil {
		return nil, err
	}
	return ec2.New(c.session), nil
}

func (c *Client) getRoute53APIClient() (*route53.Route53, error) {
	if err := c.prepareSession(); err != nil {
		return nil, err
	}
	return route53.New(c.session), nil
}

func (c *Client) prepareSession() error {
	if c.session == nil {
		creds := credentials.NewStaticCredentials(c.Credentials.ID, c.Credentials.Secret, "")
		_, err := creds.Get()
		if err != nil {
			return fmt.Errorf("Error with credentials provided: %v", err)
		}
		config := aws.NewConfig().WithRegion(c.Config.Region).WithCredentials(creds).WithMaxRetries(10)
		c.session = session.New(config)
	}
	return nil
}

// CreateNode is for creating a machine on AWS using the given AMI and InstanceType.
// Returns the ID of the newly created machine.
func (c Client) CreateNode(ami AMI, instanceType InstanceType) (string, error) {
	api, err := c.getEC2APIClient()
	if err != nil {
		return "", err
	}
	req := &ec2.RunInstancesInput{
		ImageId: aws.String(string(ami)),
		BlockDeviceMappings: []*ec2.BlockDeviceMapping{
			{
				DeviceName: aws.String("/dev/sda1"),
				Ebs: &ec2.EbsBlockDevice{
					DeleteOnTermination: aws.Bool(true),
					VolumeSize:          aws.Int64(10),
				},
			},
		},
		InstanceType:     aws.String(string(instanceType)),
		MinCount:         aws.Int64(1),
		MaxCount:         aws.Int64(1),
		SubnetId:         aws.String(c.Config.SubnetID),
		KeyName:          aws.String(c.Config.Keyname),
		SecurityGroupIds: []*string{aws.String(c.Config.SecurityGroupID)},
	}
	res, err := api.RunInstances(req)
	if err != nil {
		return "", err
	}
	instanceID := res.Instances[0].InstanceId
	// Modify the node
	modifyReq := &ec2.ModifyInstanceAttributeInput{
		InstanceId: instanceID,
		SourceDestCheck: &ec2.AttributeBooleanValue{
			Value: aws.Bool(false),
		},
	}
	err = RetryWithBackoff(func() error {
		var err2 error
		_, err2 = api.ModifyInstanceAttribute(modifyReq)
		return err2
	}, exponentialBackoffMaxEC2Attempts)
	if err != nil {
		fmt.Println("Failed to modify instance attributes")
		if err = c.DestroyNodes([]string{*instanceID}); err != nil {
			fmt.Printf("AWS NODE %q MUST BE CLEANED UP MANUALLY\n", *instanceID)
		}
		return "", err
	}
	// Tag the nodes
	thisHost, _ := os.Hostname()
	tagReq := &ec2.CreateTagsInput{
		Resources: []*string{instanceID},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("ApprendaTeam"),
				Value: aws.String("Kismatic"),
			},
			{
				Key:   aws.String("CreatedBy"),
				Value: aws.String(thisHost),
			},
		},
	}
	err = RetryWithBackoff(func() error {
		var err2 error
		_, err2 = api.CreateTags(tagReq)
		return err2
	}, exponentialBackoffMaxEC2Attempts)
	if err != nil {
		fmt.Println("Failed to tag instance")
		if err = c.DestroyNodes([]string{*instanceID}); err != nil {
			fmt.Printf("AWS NODE %q MUST BE CLEANED UP MANUALLY\n", *instanceID)
		}
		return "", err
	}
	return *res.Instances[0].InstanceId, nil
}

// GetNode returns information about a specific node. The consumer of this method
// is responsible for checking that the information it needs has been returned
// in the Node. (i.e. it's possible for the hostname, public IP to be empty)
func (c Client) GetNode(id string) (*Node, error) {
	api, err := c.getEC2APIClient()
	if err != nil {
		return nil, err
	}
	req := &ec2.DescribeInstancesInput{
		InstanceIds: []*string{aws.String(id)},
	}
	var resp *ec2.DescribeInstancesOutput
	err = RetryWithBackoff(func() error {
		var err2 error
		resp, err2 = api.DescribeInstances(req)
		return err2
	}, exponentialBackoffMaxEC2Attempts)
	if err != nil {
		fmt.Println("Failed to get node information")
		return nil, err
	}
	if len(resp.Reservations) != 1 {
		return nil, fmt.Errorf("Attempted to get a single node, but API returned %d reservations", len(resp.Reservations))
	}
	if len(resp.Reservations[0].Instances) != 1 {
		return nil, fmt.Errorf("Attempted to get a single node, but API returned %d instances", len(resp.Reservations[0].Instances))
	}
	instance := resp.Reservations[0].Instances[0]

	var privateDNSName string
	if instance.PrivateDnsName != nil {
		privateDNSName = *instance.PrivateDnsName
	}
	var privateIP string
	if instance.PrivateIpAddress != nil {
		privateIP = *instance.PrivateIpAddress
	}
	var publicIP string
	if instance.PublicIpAddress != nil {
		publicIP = *instance.PublicIpAddress
	}

	return &Node{
		PrivateDNSName: privateDNSName,
		PrivateIP:      privateIP,
		PublicIP:       publicIP,
		SSHUser:        defaultSSHUserForAMI(AMI(*instance.ImageId)),
		State:          instance.State.GoString(),
	}, nil
}

// DestroyNodes destroys the nodes identified by the ID.
func (c Client) DestroyNodes(nodeIDs []string) error {
	api, err := c.getEC2APIClient()
	if err != nil {
		return err
	}
	req := &ec2.TerminateInstancesInput{
		InstanceIds: aws.StringSlice(nodeIDs),
	}
	_, err = api.TerminateInstances(req)
	if err != nil {
		return err
	}
	return nil
}

// CreateDNSRecords generates an Route53 configured with the master nodes
func (c Client) CreateDNSRecords(nodeIPs []string) (*DNSRecord, error) {
	api, err := c.getRoute53APIClient()
	if err != nil {
		return nil, err
	}
	// Setup variables to modify hosted zone
	name := strconv.FormatInt(time.Now().Unix(), 10) + ".kismatic.integration."
	dnsRecord := &DNSRecord{Name: name, Values: nodeIPs}
	err = modifyHostedZone(dnsRecord, route53.ChangeActionUpsert, c.Config.HostedZoneID, api)
	if err != nil {
		return nil, err
	}

	return dnsRecord, nil
}

// GetDNSRecords returns all Record Sets for the Hosted Zone
func (c Client) GetDNSRecords() ([]*route53.ResourceRecordSet, error) {
	api, err := c.getRoute53APIClient()
	if err != nil {
		return nil, err
	}
	req := &route53.ListResourceRecordSetsInput{
		HostedZoneId: aws.String(c.Config.HostedZoneID), // Required
	}

	resp, err := api.ListResourceRecordSets(req)
	if err != nil {
		return nil, err
	}
	return resp.ResourceRecordSets, nil
}

// DeleteDNSRecords deletes the specified Record Set
func (c Client) DeleteDNSRecords(dnsRecord *DNSRecord) error {
	api, err := c.getRoute53APIClient()
	if err != nil {
		return err
	}
	err = modifyHostedZone(dnsRecord, route53.ChangeActionDelete, c.Config.HostedZoneID, api)
	if err != nil {
		return err
	}

	return nil
}

func modifyHostedZone(record *DNSRecord, action string, hostedZoneID string, api *route53.Route53) error {
	if len(record.Values) == 0 && action != route53.ChangeActionDelete {
		return fmt.Errorf("Values cannot be empty when setting up DNS records")
	}
	var records []*route53.ResourceRecord
	for _, value := range record.Values {
		records = append(records, &route53.ResourceRecord{Value: aws.String(value)})
	}

	req := &route53.ChangeResourceRecordSetsInput{
		ChangeBatch: &route53.ChangeBatch{
			Changes: []*route53.Change{
				{
					Action: aws.String(action),
					ResourceRecordSet: &route53.ResourceRecordSet{
						Name:            aws.String(record.Name),
						Type:            aws.String(route53.RRTypeA),
						ResourceRecords: records,
						TTL:             aws.Int64(300),
					},
				},
			},
		},
		HostedZoneId: aws.String(hostedZoneID),
	}

	resp, err := api.ChangeResourceRecordSets(req)
	if err != nil {
		return err
	}
	changeID := resp.ChangeInfo.Id
	if changeID == nil || *changeID == "" {
		return fmt.Errorf("Something went wrong, DNS change ID is nil")
	}

	changeReq := &route53.GetChangeInput{
		Id: aws.String(*changeID),
	}

	err = RetryWithBackoff(func() error {
		changeResp, err2 := api.GetChange(changeReq)
		if err2 != nil {
			return err2
		}
		if *changeResp.ChangeInfo.Status != route53.ChangeStatusInsync {
			return fmt.Errorf("DNS change status is still %s, took too long", *changeResp.ChangeInfo.Status)
		}
		return nil
	}, exponentialBackoffMaxRoute53Attempts)
	if err != nil {
		return err
	}

	return nil
}

// RetryWithBackoff will retry a function specified number of times
func RetryWithBackoff(fn func() error, retries uint) error {
	var attempts uint
	var err error
	for {
		err = fn()
		if err == nil {
			break
		}
		if attempts == retries {
			break
		}
		time.Sleep((1 << attempts) * time.Second)
		attempts++
	}
	return err
}

func defaultSSHUserForAMI(ami AMI) string {
	switch ami {
	case Ubuntu1604LTSEast:
		return "ubuntu"
	case CentOS7East:
		return "centos"
	case RedHat7East:
		return "ec2-user"
	default:
		panic(fmt.Sprintf("unsupported AMI: %q", ami))
	}
}
<|endoftext|>


<commit_before>// Package padbuf implements a simple buffer that knows how to pad/unpad
// itself so that the buffer size aligns with an arbitrary block size.
package padbuf

import (
	"bytes"
	"errors"
)

type PadBuffer []byte

func (pb PadBuffer) Len() int {
	return len(pb)
}

func (pb PadBuffer) Pad(n int) PadBuffer {
	rem := n - pb.Len() % n
	if rem == 0 {
		return pb
	}

	newpb := pb.Resize(pb.Len() + rem)
	copy(newpb[pb.Len():], bytes.Repeat([]byte{byte(rem)}, rem))
	return newpb
}

func (pb PadBuffer) Resize(newlen int) PadBuffer {
	if pb.Len() == newlen {
		return pb
	}

	buf := make([]byte, newlen)
	copy(buf, pb)
	return PadBuffer(buf)
}

func (pb PadBuffer) Unpad(n int) (PadBuffer, error) {
	rem := pb.Len() % n
	if rem != 0 {
		return pb, errors.New("buffer should be multiple block size")
	}

	last := pb[pb.Len()-1]
	pad := bytes.Repeat([]byte{last}, int(last))
	if !bytes.HasSuffix(pb, pad) {
		return pb, errors.New("invalid padding")
	}

	return PadBuffer(pb[:pb.Len()-int(last)]), nil
}<commit_msg>Remove bytes<commit_after>// Package padbuf implements a simple buffer that knows how to pad/unpad
// itself so that the buffer size aligns with an arbitrary block size.
package padbuf

import "errors"

type PadBuffer []byte

func (pb PadBuffer) Len() int {
	return len(pb)
}

func (pb PadBuffer) Pad(n int) PadBuffer {
	rem := n - pb.Len()%n
	if rem == 0 {
		return pb
	}

	newpb := pb.Resize(pb.Len() + rem)

	pad := make([]byte, rem)
	for i := 0; i < rem; i++ {
		pad[i] = byte(rem)
	}
	copy(newpb[pb.Len():], pad)

	return newpb
}

func (pb PadBuffer) Resize(newlen int) PadBuffer {
	if pb.Len() == newlen {
		return pb
	}

	buf := make([]byte, newlen)
	copy(buf, pb)
	return PadBuffer(buf)
}

func (pb PadBuffer) Unpad(n int) (PadBuffer, error) {
	rem := pb.Len() % n
	if rem != 0 {
		return pb, errors.New("buffer should be multiple block size")
	}

	last := pb[pb.Len()-1]

	count := 0
	for i := pb.Len() - 1; i >= 0; i-- {
		if pb[i] != last {
			break
		}
		count++
	}

	if count != int(last) {
		return pb, errors.New("invalid padding")
	}

	return PadBuffer(pb[:pb.Len()-int(last)]), nil
}<|endoftext|>


<commit_before>package mpflume

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin"
)

const (
	// Channel is Flume Channel Type
	Channel = "CHANNEL"
	// Sink is Flume Sink Type
	Sink    = "SINK"
	// Source is Flume Source Type
	Source  = "SOURCE"
)

// FlumePlugin mackerel plugin
type FlumePlugin struct {
	URI    string
	Prefix string
}

// MetricKeyPrefix interface for PluginWithPrefix
func (p *FlumePlugin) MetricKeyPrefix() string {
	if p.Prefix == "" {
		p.Prefix = "flume"
	}
	return p.Prefix
}

// GraphDefinition interface for mackerelplugin
func (p *FlumePlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(p.Prefix)
	return map[string]mp.Graphs{
		"channel.capacity.#": {
			Label: labelPrefix + " Channel Capacity",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ChannelCapacity", Label: "Channel Capacity"},
				{Name: "ChannelSize", Label: "Channel Size"},
			},
		},
		"channel.use_rate.#": {
			Label: labelPrefix + " Channel Use Rate",
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				{Name: "ChannelFillPercentage", Label: "Channel Fill Percentage"},
			},
		},
		"channel.event_put_num.#": {
			Label: labelPrefix + " Channel Event Put Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventPutAttemptCount", Label: "Attempt Count", Diff: true},
				{Name: "EventPutSuccessCount", Label: "Success Count", Diff: true},
			},
		},
		"channel.event_take_num.#": {
			Label: labelPrefix + " Channel Event Take Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventTakeAttemptCount", Label: "Attempt Count", Diff: true},
				{Name: "EventTakeSuccessCount", Label: "Success Count", Diff: true},
			},
		},
		"sink.batch_num.#": {
			Label: labelPrefix + " Sink Batch Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "BatchCompleteCount", Label: "Complete Count", Diff: true},
				{Name: "BatchEmptyCount", Label: "Empty Count", Diff: true},
				{Name: "BatchUnderflowCount", Label: "Underflow Count", Diff: true},
			},
		},
		"sink.connection.#": {
			Label: labelPrefix + " Sink Connection",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ConnectionCreatedCount", Label: "Created Count", Diff: true},
				{Name: "ConnectionClosedCount", Label: "Closed Count", Diff: true},
				{Name: "ConnectionFailedCount", Label: "Failed Count", Diff: true},
			},
		},
		"sink.event_drain_num.#": {
			Label: labelPrefix + " Sink Event Drain Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventDrainAttemptCount", Label: "Attempt Count", Diff: true},
				{Name: "EventDrainSuccessCount", Label: "Success Count", Diff: true},
			},
		},
		"source.append_num.#": {
			Label: labelPrefix + " Source Append Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "AppendAcceptedCount", Label: "Accepted Count", Diff: true},
				{Name: "AppendReceivedCount", Label: "Received Count", Diff: true},
			},
		},
		"source.append_batch_num.#": {
			Label: labelPrefix + " Source Append Batch Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "AppendBatchAcceptedCount", Label: "Accepted Count", Diff: true},
				{Name: "AppendBatchReceivedCount", Label: "Received Count", Diff: true},
			},
		},
		"source.event_num.#": {
			Label: labelPrefix + " Source Event Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventAcceptedCount", Label: "Accepted Count", Diff: true},
				{Name: "EventReceivedCount", Label: "Received Count", Diff: true},
			},
		},
		"source.connection.#": {
			Label: labelPrefix + " Source Connection",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "OpenConnectionCount", Label: "Open Count"},
			},
		},
	}
}

// FetchMetrics interface for mackerelplugin
func (p *FlumePlugin) FetchMetrics() (map[string]float64, error) {
	m, err := p.getMetrics()
	if err != nil {
		return nil, err
	}
	return p.parseMetrics(m), nil
}

func (p *FlumePlugin) getMetrics() (map[string]interface{}, error) {
	res, err := http.Get(p.URI)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	var m map[string]interface{}
	dec := json.NewDecoder(res.Body)
	if err := dec.Decode(&m); err != nil {
		return nil, err
	}
	return m, nil
}

func (p *FlumePlugin) parseMetrics(metrics map[string]interface{}) map[string]float64 {
	ret := make(map[string]float64)
	for k, v := range metrics {
		array := strings.Split(k, ".")
		typeName := array[0]
		componentName := array[1]
		switch typeName {
		case Channel:
			p.parseChannel(ret, componentName, v.(map[string]interface{}))
		case Sink:
			p.parseSink(ret, componentName, v.(map[string]interface{}))
		case Source:
			p.parseSource(ret, componentName, v.(map[string]interface{}))
		}
	}

	return ret
}

func (p *FlumePlugin) convertFloat64(value string) float64 {
	f, _ := strconv.ParseFloat(value, 64)
	return f
}

func (p *FlumePlugin) parseChannel(ret map[string]float64, componentName string, value map[string]interface{}) {
	ret["channel.capacity." + componentName + ".ChannelCapacity"]             = p.convertFloat64(value["ChannelCapacity"].(string))
	ret["channel.capacity." + componentName + ".ChannelSize"]                 = p.convertFloat64(value["ChannelSize"].(string))
	ret["channel.use_rate." + componentName + ".ChannelFillPercentage"]       = p.convertFloat64(value["ChannelFillPercentage"].(string))
	ret["channel.event_put_num." + componentName + ".EventPutAttemptCount"]   = p.convertFloat64(value["EventPutAttemptCount"].(string))
	ret["channel.event_put_num." + componentName + ".EventPutSuccessCount"]   = p.convertFloat64(value["EventPutSuccessCount"].(string))
	ret["channel.event_take_num." + componentName + ".EventTakeAttemptCount"] = p.convertFloat64(value["EventTakeAttemptCount"].(string))
	ret["channel.event_take_num." + componentName + ".EventTakeSuccessCount"] = p.convertFloat64(value["EventTakeSuccessCount"].(string))
}

func (p *FlumePlugin) parseSink(ret map[string]float64, componentName string, value map[string]interface{}) {
	ret["sink.batch_num." + componentName + ".BatchCompleteCount"]           = p.convertFloat64(value["BatchCompleteCount"].(string))
	ret["sink.batch_num." + componentName + ".BatchEmptyCount"]              = p.convertFloat64(value["BatchEmptyCount"].(string))
	ret["sink.batch_num." + componentName + ".BatchUnderflowCount"]          = p.convertFloat64(value["BatchUnderflowCount"].(string))
	ret["sink.connection." + componentName + ".ConnectionCreatedCount"]      = p.convertFloat64(value["ConnectionCreatedCount"].(string))
	ret["sink.connection." + componentName + ".ConnectionClosedCount"]       = p.convertFloat64(value["ConnectionClosedCount"].(string))
	ret["sink.connection." + componentName + ".ConnectionFailedCount"]       = p.convertFloat64(value["ConnectionFailedCount"].(string))
	ret["sink.event_drain_num." + componentName + ".EventDrainAttemptCount"] = p.convertFloat64(value["EventDrainAttemptCount"].(string))
	ret["sink.event_drain_num." + componentName + ".EventDrainSuccessCount"] = p.convertFloat64(value["EventDrainSuccessCount"].(string))
}

func (p *FlumePlugin) parseSource(ret map[string]float64, componentName string, value map[string]interface{}) {
	ret["source.append_num." + componentName + ".AppendAcceptedCount"]            = p.convertFloat64(value["AppendAcceptedCount"].(string))
	ret["source.append_num." + componentName + ".AppendReceivedCount"]            = p.convertFloat64(value["AppendReceivedCount"].(string))
	ret["source.append_batch_num." + componentName + ".AppendBatchAcceptedCount"] = p.convertFloat64(value["AppendBatchAcceptedCount"].(string))
	ret["source.append_batch_num." + componentName + ".AppendBatchReceivedCount"] = p.convertFloat64(value["AppendBatchReceivedCount"].(string))
	ret["source.event_num." + componentName + ".EventAcceptedCount"]              = p.convertFloat64(value["EventAcceptedCount"].(string))
	ret["source.event_num." + componentName + ".EventReceivedCount"]              = p.convertFloat64(value["EventReceivedCount"].(string))
	ret["source.connection." + componentName + ".OpenConnectionCount"]            = p.convertFloat64(value["OpenConnectionCount"].(string))
}

// Do the plugin
func Do() {
	optHost   := flag.String("host", "localhost", "Host Name")
	optPort   := flag.String("port", "41414", "Port")
	optPrefix := flag.String("metric-key-prefix", "", "Metric key prefix")
	flag.Parse()
	mp.NewMackerelPlugin(&FlumePlugin{
		URI:    fmt.Sprintf("http://%s:%s/metrics", *optHost, *optPort),
		Prefix: *optPrefix,
	}).Run()
}
<commit_msg>modify metrics label<commit_after>package mpflume

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin"
)

const (
	// Channel is Flume Channel Type
	Channel = "CHANNEL"
	// Sink is Flume Sink Type
	Sink    = "SINK"
	// Source is Flume Source Type
	Source  = "SOURCE"
)

// FlumePlugin mackerel plugin
type FlumePlugin struct {
	URI    string
	Prefix string
}

// MetricKeyPrefix interface for PluginWithPrefix
func (p *FlumePlugin) MetricKeyPrefix() string {
	if p.Prefix == "" {
		p.Prefix = "flume"
	}
	return p.Prefix
}

// GraphDefinition interface for mackerelplugin
func (p *FlumePlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(p.Prefix)
	return map[string]mp.Graphs{
		"channel.capacity.#": {
			Label: labelPrefix + " Channel Capacity",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ChannelCapacity", Label: "capacity"},
				{Name: "ChannelSize", Label: "size"},
			},
		},
		"channel.use_rate.#": {
			Label: labelPrefix + " Channel Use Rate",
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				{Name: "ChannelFillPercentage", Label: "fill percentage"},
			},
		},
		"channel.event_put_num.#": {
			Label: labelPrefix + " Channel Event Put Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventPutAttemptCount", Label: "attempt", Diff: true},
				{Name: "EventPutSuccessCount", Label: "success", Diff: true},
			},
		},
		"channel.event_take_num.#": {
			Label: labelPrefix + " Channel Event Take Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventTakeAttemptCount", Label: "attempt", Diff: true},
				{Name: "EventTakeSuccessCount", Label: "success", Diff: true},
			},
		},
		"sink.batch_num.#": {
			Label: labelPrefix + " Sink Batch Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "BatchCompleteCount", Label: "complete", Diff: true},
				{Name: "BatchEmptyCount", Label: "empty", Diff: true},
				{Name: "BatchUnderflowCount", Label: "underflow", Diff: true},
			},
		},
		"sink.connection.#": {
			Label: labelPrefix + " Sink Connection",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ConnectionCreatedCount", Label: "created", Diff: true},
				{Name: "ConnectionClosedCount", Label: "closed", Diff: true},
				{Name: "ConnectionFailedCount", Label: "failed", Diff: true},
			},
		},
		"sink.event_drain_num.#": {
			Label: labelPrefix + " Sink Event Drain Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventDrainAttemptCount", Label: "attempt", Diff: true},
				{Name: "EventDrainSuccessCount", Label: "success", Diff: true},
			},
		},
		"source.append_num.#": {
			Label: labelPrefix + " Source Append Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "AppendAcceptedCount", Label: "accepted", Diff: true},
				{Name: "AppendReceivedCount", Label: "received", Diff: true},
			},
		},
		"source.append_batch_num.#": {
			Label: labelPrefix + " Source Append Batch Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "AppendBatchAcceptedCount", Label: "accepted", Diff: true},
				{Name: "AppendBatchReceivedCount", Label: "received", Diff: true},
			},
		},
		"source.event_num.#": {
			Label: labelPrefix + " Source Event Num",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "EventAcceptedCount", Label: "accepted", Diff: true},
				{Name: "EventReceivedCount", Label: "received", Diff: true},
			},
		},
		"source.connection.#": {
			Label: labelPrefix + " Source Connection",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "OpenConnectionCount", Label: "open"},
			},
		},
	}
}

// FetchMetrics interface for mackerelplugin
func (p *FlumePlugin) FetchMetrics() (map[string]float64, error) {
	m, err := p.getMetrics()
	if err != nil {
		return nil, err
	}
	return p.parseMetrics(m), nil
}

func (p *FlumePlugin) getMetrics() (map[string]interface{}, error) {
	res, err := http.Get(p.URI)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	var m map[string]interface{}
	dec := json.NewDecoder(res.Body)
	if err := dec.Decode(&m); err != nil {
		return nil, err
	}
	return m, nil
}

func (p *FlumePlugin) parseMetrics(metrics map[string]interface{}) map[string]float64 {
	ret := make(map[string]float64)
	for k, v := range metrics {
		array := strings.Split(k, ".")
		typeName := array[0]
		componentName := array[1]
		switch typeName {
		case Channel:
			p.parseChannel(ret, componentName, v.(map[string]interface{}))
		case Sink:
			p.parseSink(ret, componentName, v.(map[string]interface{}))
		case Source:
			p.parseSource(ret, componentName, v.(map[string]interface{}))
		}
	}

	return ret
}

func (p *FlumePlugin) convertFloat64(value string) float64 {
	f, _ := strconv.ParseFloat(value, 64)
	return f
}

func (p *FlumePlugin) parseChannel(ret map[string]float64, componentName string, value map[string]interface{}) {
	ret["channel.capacity." + componentName + ".ChannelCapacity"]             = p.convertFloat64(value["ChannelCapacity"].(string))
	ret["channel.capacity." + componentName + ".ChannelSize"]                 = p.convertFloat64(value["ChannelSize"].(string))
	ret["channel.use_rate." + componentName + ".ChannelFillPercentage"]       = p.convertFloat64(value["ChannelFillPercentage"].(string))
	ret["channel.event_put_num." + componentName + ".EventPutAttemptCount"]   = p.convertFloat64(value["EventPutAttemptCount"].(string))
	ret["channel.event_put_num." + componentName + ".EventPutSuccessCount"]   = p.convertFloat64(value["EventPutSuccessCount"].(string))
	ret["channel.event_take_num." + componentName + ".EventTakeAttemptCount"] = p.convertFloat64(value["EventTakeAttemptCount"].(string))
	ret["channel.event_take_num." + componentName + ".EventTakeSuccessCount"] = p.convertFloat64(value["EventTakeSuccessCount"].(string))
}

func (p *FlumePlugin) parseSink(ret map[string]float64, componentName string, value map[string]interface{}) {
	ret["sink.batch_num." + componentName + ".BatchCompleteCount"]           = p.convertFloat64(value["BatchCompleteCount"].(string))
	ret["sink.batch_num." + componentName + ".BatchEmptyCount"]              = p.convertFloat64(value["BatchEmptyCount"].(string))
	ret["sink.batch_num." + componentName + ".BatchUnderflowCount"]          = p.convertFloat64(value["BatchUnderflowCount"].(string))
	ret["sink.connection." + componentName + ".ConnectionCreatedCount"]      = p.convertFloat64(value["ConnectionCreatedCount"].(string))
	ret["sink.connection." + componentName + ".ConnectionClosedCount"]       = p.convertFloat64(value["ConnectionClosedCount"].(string))
	ret["sink.connection." + componentName + ".ConnectionFailedCount"]       = p.convertFloat64(value["ConnectionFailedCount"].(string))
	ret["sink.event_drain_num." + componentName + ".EventDrainAttemptCount"] = p.convertFloat64(value["EventDrainAttemptCount"].(string))
	ret["sink.event_drain_num." + componentName + ".EventDrainSuccessCount"] = p.convertFloat64(value["EventDrainSuccessCount"].(string))
}

func (p *FlumePlugin) parseSource(ret map[string]float64, componentName string, value map[string]interface{}) {
	ret["source.append_num." + componentName + ".AppendAcceptedCount"]            = p.convertFloat64(value["AppendAcceptedCount"].(string))
	ret["source.append_num." + componentName + ".AppendReceivedCount"]            = p.convertFloat64(value["AppendReceivedCount"].(string))
	ret["source.append_batch_num." + componentName + ".AppendBatchAcceptedCount"] = p.convertFloat64(value["AppendBatchAcceptedCount"].(string))
	ret["source.append_batch_num." + componentName + ".AppendBatchReceivedCount"] = p.convertFloat64(value["AppendBatchReceivedCount"].(string))
	ret["source.event_num." + componentName + ".EventAcceptedCount"]              = p.convertFloat64(value["EventAcceptedCount"].(string))
	ret["source.event_num." + componentName + ".EventReceivedCount"]              = p.convertFloat64(value["EventReceivedCount"].(string))
	ret["source.connection." + componentName + ".OpenConnectionCount"]            = p.convertFloat64(value["OpenConnectionCount"].(string))
}

// Do the plugin
func Do() {
	optHost   := flag.String("host", "localhost", "Host Name")
	optPort   := flag.String("port", "41414", "Port")
	optPrefix := flag.String("metric-key-prefix", "", "Metric key prefix")
	flag.Parse()
	mp.NewMackerelPlugin(&FlumePlugin{
		URI:    fmt.Sprintf("http://%s:%s/metrics", *optHost, *optPort),
		Prefix: *optPrefix,
	}).Run()
}
<|endoftext|>


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

package sgdisk

import (
	"fmt"
	"os/exec"

	"github.com/coreos/ignition/internal/log"
)

const sgdiskPath = "/sbin/sgdisk"

type Operation struct {
	logger *log.Logger
	dev    string
	wipe   bool
	parts  []Partition
}

type Partition struct {
	Number   int
	Offset   uint64 // 512-byte sectors
	Length   uint64 // 512-byte sectors
	Label    string
	TypeGUID string
}

// Begin begins an sgdisk operation
func Begin(logger *log.Logger, dev string) *Operation {
	return &Operation{logger: logger, dev: dev}
}

// CreatePartition adds the supplied partition to the list of partitions to be created as part of an operation.
func (op *Operation) CreatePartition(p Partition) {
	// XXX(vc): no checking is performed here, since we perform checking at json parsing, Commit() will just fail on badness.
	op.parts = append(op.parts, p)
}

// WipeTable toggles if the table is to be wiped first when commiting this operation.
func (op *Operation) WipeTable(wipe bool) {
	op.wipe = wipe
}

// Commit commits an partitioning operation.
func (op *Operation) Commit() error {
	if op.wipe {
		cmd := exec.Command(sgdiskPath, "--zap-all", op.dev)
		if err := op.logger.LogCmd(cmd, "wiping table on %q", op.dev); err != nil {
			op.logger.Info("potential error encountered while wiping table... retrying")
			cmd = exec.Command(sgdiskPath, "--zap-all", op.dev)
			if err := op.logger.LogCmd(cmd, "wiping table on %q", op.dev); err != nil {
				return fmt.Errorf("wipe failed: %v", err)
			}
		}
	}

	if len(op.parts) != 0 {
		opts := []string{}
		for _, p := range op.parts {
			opts = append(opts, fmt.Sprintf("--new=%d:%d:+%d", p.Number, p.Offset, p.Length))
			opts = append(opts, fmt.Sprintf("--change-name=%d:%s", p.Number, p.Label))
			if p.TypeGUID != "" {
				opts = append(opts, fmt.Sprintf("--partition-guid=%d:%s", p.Number, p.TypeGUID))
			}
		}
		opts = append(opts, op.dev)
		cmd := exec.Command(sgdiskPath, opts...)
		if err := op.logger.LogCmd(cmd, "creating %d partitions on %q", len(op.parts), op.dev); err != nil {
			return fmt.Errorf("create partitions failed: %v", err)
		}
	}

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

package sgdisk

import (
	"fmt"
	"os/exec"

	"github.com/coreos/ignition/internal/log"
)

const sgdiskPath = "/sbin/sgdisk"

type Operation struct {
	logger *log.Logger
	dev    string
	wipe   bool
	parts  []Partition
}

type Partition struct {
	Number   int
	Offset   uint64 // 512-byte sectors
	Length   uint64 // 512-byte sectors
	Label    string
	TypeGUID string
}

// Begin begins an sgdisk operation
func Begin(logger *log.Logger, dev string) *Operation {
	return &Operation{logger: logger, dev: dev}
}

// CreatePartition adds the supplied partition to the list of partitions to be created as part of an operation.
func (op *Operation) CreatePartition(p Partition) {
	// XXX(vc): no checking is performed here, since we perform checking at json parsing, Commit() will just fail on badness.
	op.parts = append(op.parts, p)
}

// WipeTable toggles if the table is to be wiped first when commiting this operation.
func (op *Operation) WipeTable(wipe bool) {
	op.wipe = wipe
}

// Commit commits an partitioning operation.
func (op *Operation) Commit() error {
	if op.wipe {
		cmd := exec.Command(sgdiskPath, "--zap-all", op.dev)
		if err := op.logger.LogCmd(cmd, "wiping table on %q", op.dev); err != nil {
			op.logger.Info("potential error encountered while wiping table... retrying")
			cmd = exec.Command(sgdiskPath, "--zap-all", op.dev)
			if err := op.logger.LogCmd(cmd, "wiping table on %q", op.dev); err != nil {
				return fmt.Errorf("wipe failed: %v", err)
			}
		}
	}

	if len(op.parts) != 0 {
		opts := []string{}
		for _, p := range op.parts {
			opts = append(opts, fmt.Sprintf("--new=%d:%d:+%d", p.Number, p.Offset, p.Length))
			opts = append(opts, fmt.Sprintf("--change-name=%d:%s", p.Number, p.Label))
			if p.TypeGUID != "" {
				opts = append(opts, fmt.Sprintf("--typecode=%d:%s", p.Number, p.TypeGUID))
			}
		}
		opts = append(opts, op.dev)
		cmd := exec.Command(sgdiskPath, opts...)
		if err := op.logger.LogCmd(cmd, "creating %d partitions on %q", len(op.parts), op.dev); err != nil {
			return fmt.Errorf("create partitions failed: %v", err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"os"
	"reflect"
	"strconv"
	"strings"
)

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"mongodb.background_flushing": mp.Graphs{
		Label: "MongoDB Command",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "duration_ms", Label: "Duration in ms", Diff: true},
		},
	},
	"mongodb.connections": mp.Graphs{
		Label: "MongoDB Connections",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "connections_current", Label: "current"},
		},
	},
	"mongodb.index_counters.btree": mp.Graphs{
		Label: "MongoDB Index Counters Btree",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "btree_hits", Label: "hits", Diff: true},
		},
	},
	"mongodb.opcounters": mp.Graphs{
		Label: "MongoDB opcounters",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "opcounters_insert", Label: "Insert", Diff: true},
			mp.Metrics{Name: "opcounters_query", Label: "Query", Diff: true},
			mp.Metrics{Name: "opcounters_update", Label: "Update", Diff: true},
			mp.Metrics{Name: "opcounters_delete", Label: "Delete", Diff: true},
			mp.Metrics{Name: "opcounters_getmore", Label: "Getmore", Diff: true},
			mp.Metrics{Name: "opcounters_command", Label: "Command", Diff: true},
		},
	},
}

var graphdef30 map[string](mp.Graphs) = map[string](mp.Graphs){
	"mongodb.background_flushing": mp.Graphs{
		Label: "MongoDB Command",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "duration_ms", Label: "Duration in ms", Diff: true},
		},
	},
	"mongodb.connections": mp.Graphs{
		Label: "MongoDB Connections",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "connections_current", Label: "current"},
		},
	},
	"mongodb.index_counters.btree.miss_ratio": mp.Graphs{
		Label: "MongoDB Index Counters Btree Miss-ratio",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "miss_ratio", Label: "Miss ratio"},
		},
	},
	"mongodb.opcounters": mp.Graphs{
		Label: "MongoDB opcounters",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "opcounters_insert", Label: "Insert", Diff: true},
			mp.Metrics{Name: "opcounters_query", Label: "Query", Diff: true},
			mp.Metrics{Name: "opcounters_update", Label: "Update", Diff: true},
			mp.Metrics{Name: "opcounters_delete", Label: "Delete", Diff: true},
			mp.Metrics{Name: "opcounters_getmore", Label: "Getmore", Diff: true},
			mp.Metrics{Name: "opcounters_command", Label: "Command", Diff: true},
		},
	},
}

var metricPlace22 map[string][]string = map[string][]string{
	"duration_ms":         []string{"backgroundFlushing", "total_ms"},
	"connections_current": []string{"connections", "current"},
	"btree_hits":          []string{"indexCounters", "btree", "hits"},
	"opcounters_insert":   []string{"opcounters", "insert"},
	"opcounters_query":    []string{"opcounters", "query"},
	"opcounters_update":   []string{"opcounters", "update"},
	"opcounters_delete":   []string{"opcounters", "delete"},
	"opcounters_getmore":  []string{"opcounters", "getmore"},
	"opcounters_command":  []string{"opcounters", "command"},
}

var metricPlace24 map[string][]string = map[string][]string{
	"duration_ms":         []string{"backgroundFlushing", "total_ms"},
	"connections_current": []string{"connections", "current"},
	"btree_hits":          []string{"indexCounters", "hits"},
	"opcounters_insert":   []string{"opcounters", "insert"},
	"opcounters_query":    []string{"opcounters", "query"},
	"opcounters_update":   []string{"opcounters", "update"},
	"opcounters_delete":   []string{"opcounters", "delete"},
	"opcounters_getmore":  []string{"opcounters", "getmore"},
	"opcounters_command":  []string{"opcounters", "command"},
}

// indexCounters is removed from mongodb 3.0.
// ref. http://stackoverflow.com/questions/29428793/where-is-the-indexcounter-in-db-serverstatus-on-mongodb-3-0
var metricPlace30 map[string][]string = map[string][]string{
	"duration_ms":         []string{"backgroundFlushing", "total_ms"},
	"connections_current": []string{"connections", "current"},
	"opcounters_insert":   []string{"opcounters", "insert"},
	"opcounters_query":    []string{"opcounters", "query"},
	"opcounters_update":   []string{"opcounters", "update"},
	"opcounters_delete":   []string{"opcounters", "delete"},
	"opcounters_getmore":  []string{"opcounters", "getmore"},
	"opcounters_command":  []string{"opcounters", "command"},
}

func GetFloatValue(s map[string]interface{}, keys []string) (float64, error) {
	var val float64
	sm := s
	var err error
	for i, k := range keys {
		if i+1 < len(keys) {
			switch sm[k].(type) {
			case bson.M:
				sm = sm[k].(bson.M)
			default:
				return 0, fmt.Errorf("Cannot handle as a hash for %s", k)
			}
		} else {
			val, err = strconv.ParseFloat(fmt.Sprint(sm[k]), 64)
			if err != nil {
				return 0, err
			}
		}
	}

	return val, nil
}

type MongoDBPlugin struct {
	Url     string
	Verbose bool
}

func (m MongoDBPlugin) FetchStatus() (bson.M, error) {
	session, err := mgo.Dial(m.Url)
	if err != nil {
		return nil, err
	}

	defer session.Close()
	serverStatus := bson.M{}
	if err := session.Run("serverStatus", &serverStatus); err != nil {
		return nil, err
	}
	if m.Verbose {
		str, err := json.Marshal(serverStatus)
		if err != nil {
			fmt.Println(fmt.Errorf("Marshaling error: %s", err.Error()))
		}
		fmt.Println(string(str))
	}
	return serverStatus, nil
}

func (m MongoDBPlugin) FetchMetrics() (map[string]float64, error) {
	serverStatus, err := m.FetchStatus()
	if err != nil {
		return nil, err
	}
	return m.ParseStatus(serverStatus)
}

func (m MongoDBPlugin) getVersion(serverStatus bson.M) string {
	if reflect.TypeOf(serverStatus["version"]).String() == "string" {
		version := serverStatus["version"].(string)
		return version
	}
	return ""
}

func (m MongoDBPlugin) ParseStatus(serverStatus bson.M) (map[string]float64, error) {
	stat := make(map[string]float64)
	metricPlace := &metricPlace22
	version := m.getVersion(serverStatus)
	if strings.HasPrefix(version, "2.4") {
		metricPlace = &metricPlace24
	} else if strings.HasPrefix(version, "2.6") {
		metricPlace = &metricPlace24
	} else if strings.HasPrefix(version, "3.0") {
		metricPlace = &metricPlace30
	}

	for k, v := range *metricPlace {
		val, err := GetFloatValue(serverStatus, v)
		if err != nil {
			return nil, err
		}

		stat[k] = val
	}

	return stat, nil
}

func (m MongoDBPlugin) GraphDefinition() map[string](mp.Graphs) {
	serverStatus, err := m.FetchStatus()
	if err != nil {
		return graphdef
	}
	version := m.getVersion(serverStatus)
	if strings.HasPrefix(version, "3.0") {
		return graphdef30
	}
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "27017", "Port")
	optUser := flag.String("username", "", "Username")
	optPass := flag.String("password", "", "Password")
	optVerbose := flag.Bool("v", false, "Verbose mode")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var mongodb MongoDBPlugin
	mongodb.Verbose = *optVerbose
	if *optUser == "" && *optPass == "" {
		mongodb.Url = fmt.Sprintf("mongodb://%s:%s", *optHost, *optPort)
	} else {
		mongodb.Url = fmt.Sprintf("mongodb://%s:%s@%s:%s", *optUser, *optPass, *optHost, *optPort)
	}

	helper := mp.NewMackerelPlugin(mongodb)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-mongodb-%s-%s", *optHost, *optPort)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>use new helper<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"reflect"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"mongodb.background_flushing": mp.Graphs{
		Label: "MongoDB Command",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "duration_ms", Label: "Duration in ms", Diff: true, Type: "uint64"},
		},
	},
	"mongodb.connections": mp.Graphs{
		Label: "MongoDB Connections",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "connections_current", Label: "current"},
		},
	},
	"mongodb.index_counters.btree": mp.Graphs{
		Label: "MongoDB Index Counters Btree",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "btree_hits", Label: "hits", Diff: true, Type: "uint64"},
		},
	},
	"mongodb.opcounters": mp.Graphs{
		Label: "MongoDB opcounters",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "opcounters_insert", Label: "Insert", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_query", Label: "Query", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_update", Label: "Update", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_delete", Label: "Delete", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_getmore", Label: "Getmore", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_command", Label: "Command", Diff: true, Type: "uint64"},
		},
	},
}

var graphdef30 map[string](mp.Graphs) = map[string](mp.Graphs){
	"mongodb.background_flushing": mp.Graphs{
		Label: "MongoDB Command",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "duration_ms", Label: "Duration in ms", Diff: true, Type: "uint64"},
		},
	},
	"mongodb.connections": mp.Graphs{
		Label: "MongoDB Connections",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "connections_current", Label: "current"},
		},
	},
	"mongodb.index_counters.btree.miss_ratio": mp.Graphs{
		Label: "MongoDB Index Counters Btree Miss-ratio",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "miss_ratio", Label: "Miss ratio"},
		},
	},
	"mongodb.opcounters": mp.Graphs{
		Label: "MongoDB opcounters",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "opcounters_insert", Label: "Insert", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_query", Label: "Query", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_update", Label: "Update", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_delete", Label: "Delete", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_getmore", Label: "Getmore", Diff: true, Type: "uint64"},
			mp.Metrics{Name: "opcounters_command", Label: "Command", Diff: true, Type: "uint64"},
		},
	},
}

var metricPlace22 map[string][]string = map[string][]string{
	"duration_ms":         []string{"backgroundFlushing", "total_ms"},
	"connections_current": []string{"connections", "current"},
	"btree_hits":          []string{"indexCounters", "btree", "hits"},
	"opcounters_insert":   []string{"opcounters", "insert"},
	"opcounters_query":    []string{"opcounters", "query"},
	"opcounters_update":   []string{"opcounters", "update"},
	"opcounters_delete":   []string{"opcounters", "delete"},
	"opcounters_getmore":  []string{"opcounters", "getmore"},
	"opcounters_command":  []string{"opcounters", "command"},
}

var metricPlace24 map[string][]string = map[string][]string{
	"duration_ms":         []string{"backgroundFlushing", "total_ms"},
	"connections_current": []string{"connections", "current"},
	"btree_hits":          []string{"indexCounters", "hits"},
	"opcounters_insert":   []string{"opcounters", "insert"},
	"opcounters_query":    []string{"opcounters", "query"},
	"opcounters_update":   []string{"opcounters", "update"},
	"opcounters_delete":   []string{"opcounters", "delete"},
	"opcounters_getmore":  []string{"opcounters", "getmore"},
	"opcounters_command":  []string{"opcounters", "command"},
}

// indexCounters is removed from mongodb 3.0.
// ref. http://stackoverflow.com/questions/29428793/where-is-the-indexcounter-in-db-serverstatus-on-mongodb-3-0
var metricPlace30 map[string][]string = map[string][]string{
	"duration_ms":         []string{"backgroundFlushing", "total_ms"},
	"connections_current": []string{"connections", "current"},
	"opcounters_insert":   []string{"opcounters", "insert"},
	"opcounters_query":    []string{"opcounters", "query"},
	"opcounters_update":   []string{"opcounters", "update"},
	"opcounters_delete":   []string{"opcounters", "delete"},
	"opcounters_getmore":  []string{"opcounters", "getmore"},
	"opcounters_command":  []string{"opcounters", "command"},
}

func GetFloatValue(s map[string]interface{}, keys []string) (float64, error) {
	var val float64
	sm := s
	var err error
	for i, k := range keys {
		if i+1 < len(keys) {
			switch sm[k].(type) {
			case bson.M:
				sm = sm[k].(bson.M)
			default:
				return 0, fmt.Errorf("Cannot handle as a hash for %s", k)
			}
		} else {
			val, err = strconv.ParseFloat(fmt.Sprint(sm[k]), 64)
			if err != nil {
				return 0, err
			}
		}
	}

	return val, nil
}

type MongoDBPlugin struct {
	Url     string
	Verbose bool
}

func (m MongoDBPlugin) FetchStatus() (bson.M, error) {
	session, err := mgo.Dial(m.Url)
	if err != nil {
		return nil, err
	}

	defer session.Close()
	serverStatus := bson.M{}
	if err := session.Run("serverStatus", &serverStatus); err != nil {
		return nil, err
	}
	if m.Verbose {
		str, err := json.Marshal(serverStatus)
		if err != nil {
			fmt.Println(fmt.Errorf("Marshaling error: %s", err.Error()))
		}
		fmt.Println(string(str))
	}
	return serverStatus, nil
}

func (m MongoDBPlugin) FetchMetrics() (map[string]interface{}, error) {
	serverStatus, err := m.FetchStatus()
	if err != nil {
		return nil, err
	}
	return m.ParseStatus(serverStatus)
}

func (m MongoDBPlugin) getVersion(serverStatus bson.M) string {
	if reflect.TypeOf(serverStatus["version"]).String() == "string" {
		version := serverStatus["version"].(string)
		return version
	}
	return ""
}

func (m MongoDBPlugin) ParseStatus(serverStatus bson.M) (map[string]interface{}, error) {
	stat := make(map[string]interface{})
	metricPlace := &metricPlace22
	version := m.getVersion(serverStatus)
	if strings.HasPrefix(version, "2.4") {
		metricPlace = &metricPlace24
	} else if strings.HasPrefix(version, "2.6") {
		metricPlace = &metricPlace24
	} else if strings.HasPrefix(version, "3.0") {
		metricPlace = &metricPlace30
	}

	for k, v := range *metricPlace {
		val, err := GetFloatValue(serverStatus, v)
		if err != nil {
			return nil, err
		}

		stat[k] = val
	}

	return stat, nil
}

func (m MongoDBPlugin) GraphDefinition() map[string](mp.Graphs) {
	serverStatus, err := m.FetchStatus()
	if err != nil {
		return graphdef
	}
	version := m.getVersion(serverStatus)
	if strings.HasPrefix(version, "3.0") {
		return graphdef30
	}
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "27017", "Port")
	optUser := flag.String("username", "", "Username")
	optPass := flag.String("password", "", "Password")
	optVerbose := flag.Bool("v", false, "Verbose mode")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var mongodb MongoDBPlugin
	mongodb.Verbose = *optVerbose
	if *optUser == "" && *optPass == "" {
		mongodb.Url = fmt.Sprintf("mongodb://%s:%s", *optHost, *optPort)
	} else {
		mongodb.Url = fmt.Sprintf("mongodb://%s:%s@%s:%s", *optUser, *optPass, *optHost, *optPort)
	}

	helper := mp.NewMackerelPlugin(mongodb)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-mongodb-%s-%s", *optHost, *optPort)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>