<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-sql"
	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/pq/hstore"
	"github.com/flynn/flynn/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/flynn/flynn/controller/name"
	"github.com/flynn/flynn/controller/schema"
	ct "github.com/flynn/flynn/controller/types"
	logaggc "github.com/flynn/flynn/logaggregator/client"
	"github.com/flynn/flynn/pkg/ctxhelper"
	"github.com/flynn/flynn/pkg/httphelper"
	"github.com/flynn/flynn/pkg/postgres"
	"github.com/flynn/flynn/pkg/random"
	"github.com/flynn/flynn/pkg/sse"
	routerc "github.com/flynn/flynn/router/client"
	"github.com/flynn/flynn/router/types"
)

type AppRepo struct {
	router        routerc.Client
	defaultDomain string

	db *postgres.DB
}

type appUpdate map[string]interface{}

func NewAppRepo(db *postgres.DB, defaultDomain string, router routerc.Client) *AppRepo {
	return &AppRepo{db: db, defaultDomain: defaultDomain, router: router}
}

var appNamePattern = regexp.MustCompile(`^[a-z\d]+(-[a-z\d]+)*$`)

func (r *AppRepo) Add(data interface{}) error {
	app := data.(*ct.App)
	if app.Name == "" {
		var nameID uint32
		if err := r.db.QueryRow("SELECT nextval('name_ids')").Scan(&nameID); err != nil {
			return err
		}
		app.Name = name.Get(nameID)
	}
	if len(app.Name) > 100 || !appNamePattern.MatchString(app.Name) {
		return ct.ValidationError{Field: "name", Message: "is invalid"}
	}
	if app.ID == "" {
		app.ID = random.UUID()
	}
	if app.Strategy == "" {
		app.Strategy = "all-at-once"
	}
	meta := metaToHstore(app.Meta)
	if err := r.db.QueryRow("INSERT INTO apps (app_id, name, meta, strategy) VALUES ($1, $2, $3, $4) RETURNING created_at, updated_at", app.ID, app.Name, meta, app.Strategy).Scan(&app.CreatedAt, &app.UpdatedAt); err != nil {
		if postgres.IsUniquenessError(err, "apps_name_idx") {
			return httphelper.ObjectExistsErr(fmt.Sprintf("application %q already exists", app.Name))
		}
		return err
	}
	app.ID = postgres.CleanUUID(app.ID)
	if !app.System() && r.defaultDomain != "" {
		route := (&router.HTTPRoute{
			Domain:  fmt.Sprintf("%s.%s", app.Name, r.defaultDomain),
			Service: app.Name + "-web",
		}).ToRoute()
		route.ParentRef = routeParentRef(app.ID)
		if err := r.router.CreateRoute(route); err != nil {
			log.Printf("Error creating default route for %s: %s", app.Name, err)
		}
	}
	return nil
}

func scanApp(s postgres.Scanner) (*ct.App, error) {
	app := &ct.App{}
	var meta hstore.Hstore
	err := s.Scan(&app.ID, &app.Name, &meta, &app.Strategy, &app.CreatedAt, &app.UpdatedAt)
	if err == sql.ErrNoRows {
		err = ErrNotFound
	}
	if len(meta.Map) > 0 {
		app.Meta = make(map[string]string, len(meta.Map))
		for k, v := range meta.Map {
			app.Meta[k] = v.String
		}
	}
	app.ID = postgres.CleanUUID(app.ID)
	return app, err
}

var idPattern = regexp.MustCompile(`^[a-f0-9]{8}-?([a-f0-9]{4}-?){3}[a-f0-9]{12}$`)

type rowQueryer interface {
	QueryRow(query string, args ...interface{}) postgres.Scanner
}

func selectApp(db rowQueryer, id string, update bool) (*ct.App, error) {
	var row postgres.Scanner
	query := "SELECT app_id, name, meta, strategy, created_at, updated_at FROM apps WHERE deleted_at IS NULL AND "
	var suffix string
	if update {
		suffix = " FOR UPDATE"
	}
	if idPattern.MatchString(id) {
		row = db.QueryRow(query+"(app_id = $1 OR name = $2) LIMIT 1"+suffix, id, id)
	} else {
		row = db.QueryRow(query+"name = $1"+suffix, id)
	}
	return scanApp(row)
}

func (r *AppRepo) Get(id string) (interface{}, error) {
	return selectApp(r.db, id, false)
}

func (r *AppRepo) Update(id string, data map[string]interface{}) (interface{}, error) {
	tx, err := r.db.Begin()
	if err != nil {
		return nil, err
	}
	app, err := selectApp(tx, id, true)
	if err != nil {
		tx.Rollback()
		return nil, err
	}

	for k, v := range data {
		switch k {
		case "strategy":
			strategy, ok := v.(string)
			if !ok {
				tx.Rollback()
				return nil, fmt.Errorf("controller: expected string, got %T", v)
			}
			if _, err := tx.Exec("UPDATE apps SET strategy = $2, updated_at = now() WHERE app_id = $1", app.ID, strategy); err != nil {
				tx.Rollback()
				return nil, err
			}
		case "meta":
			data, ok := v.(map[string]interface{})
			if !ok {
				tx.Rollback()
				return nil, fmt.Errorf("controller: expected map[string]interface{}, got %T", v)
			}
			var meta hstore.Hstore
			meta.Map = make(map[string]sql.NullString, len(data))
			app.Meta = make(map[string]string, len(data))
			for k, v := range data {
				s, ok := v.(string)
				if !ok {
					tx.Rollback()
					return nil, fmt.Errorf("controller: expected string, got %T", v)
				}
				meta.Map[k] = sql.NullString{String: s, Valid: true}
				app.Meta[k] = s
			}
			if _, err := tx.Exec("UPDATE apps SET meta = $2, updated_at = now() WHERE app_id = $1", app.ID, meta); err != nil {
				tx.Rollback()
				return nil, err
			}
		}
	}

	return app, tx.Commit()
}

func (r *AppRepo) Remove(id string) error {
	tx, err := r.db.Begin()
	if err != nil {
		return err
	}

	if !idPattern.MatchString(id) {
		app, err := selectApp(r.db, id, false)
		if err != nil {
			return err
		}
		id = app.ID
	}
	_, err = tx.Exec("UPDATE apps SET deleted_at = now() WHERE app_id = $1 AND deleted_at IS NULL", id)
	if err != nil {
		tx.Rollback()
		return err
	}

	_, err = tx.Exec("UPDATE formations SET deleted_at = now(), processes = NULL, updated_at = now() WHERE app_id = $1 AND deleted_at IS NULL", id)
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("UPDATE app_resources SET deleted_at = now() WHERE app_id = $1 AND deleted_at IS NULL", id)
	if err != nil {
		tx.Rollback()
		return err
	}
	return tx.Commit()
}

func (r *AppRepo) List() (interface{}, error) {
	rows, err := r.db.Query("SELECT app_id, name, meta, strategy, created_at, updated_at FROM apps WHERE deleted_at IS NULL ORDER BY created_at DESC")
	if err != nil {
		return nil, err
	}
	apps := []*ct.App{}
	for rows.Next() {
		app, err := scanApp(rows)
		if err != nil {
			rows.Close()
			return nil, err
		}
		apps = append(apps, app)
	}
	return apps, rows.Err()
}

func (r *AppRepo) SetRelease(appID string, releaseID string) error {
	return r.db.Exec("UPDATE apps SET release_id = $2, updated_at = now() WHERE app_id = $1", appID, releaseID)
}

func (r *AppRepo) GetRelease(id string) (*ct.Release, error) {
	row := r.db.QueryRow("SELECT r.release_id, r.artifact_id, r.data, r.created_at FROM apps a JOIN releases r USING (release_id) WHERE a.app_id = $1", id)
	return scanRelease(row)
}

func (c *controllerAPI) UpdateApp(ctx context.Context, rw http.ResponseWriter, req *http.Request) {
	params, _ := ctxhelper.ParamsFromContext(ctx)

	var data appUpdate
	if err := httphelper.DecodeJSON(req, &data); err != nil {
		respondWithError(rw, err)
		return
	}

	if err := schema.Validate(data); err != nil {
		respondWithError(rw, err)
		return
	}

	app, err := c.appRepo.Update(params.ByName("apps_id"), data)
	if err != nil {
		respondWithError(rw, err)
		return
	}
	httphelper.JSON(rw, 200, app)
}

func (c *controllerAPI) AppLog(ctx context.Context, w http.ResponseWriter, req *http.Request) {
	ctx, cancel := context.WithCancel(ctx)

	opts := logaggc.LogOpts{
		Follow: req.FormValue("follow") == "true",
		JobID:  req.FormValue("job_id"),
	}
	if vals, ok := req.Form["process_type"]; ok && len(vals) > 0 {
		opts.ProcessType = &vals[len(vals)-1]
	}
	if strLines := req.FormValue("lines"); strLines != "" {
		lines, err := strconv.Atoi(req.FormValue("lines"))
		if err != nil {
			respondWithError(w, err)
			return
		}
		opts.Lines = &lines
	}
	rc, err := c.logaggc.GetLog(c.getApp(ctx).ID, &opts)
	if err != nil {
		respondWithError(w, err)
		return
	}

	if cn, ok := w.(http.CloseNotifier); ok {
		go func() {
			select {
			case <-cn.CloseNotify():
				rc.Close()
			case <-ctx.Done():
			}
		}()
	}
	defer cancel()
	defer rc.Close()

	if !strings.Contains(req.Header.Get("Accept"), "text/event-stream") {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(200)
		// Send headers right away if following
		if wf, ok := w.(http.Flusher); ok && opts.Follow {
			wf.Flush()
		}

		fw := httphelper.FlushWriter{Writer: w, Enabled: opts.Follow}
		io.Copy(fw, rc)
		return
	}

	ch := make(chan *sseLogChunk)
	l, _ := ctxhelper.LoggerFromContext(ctx)
	s := sse.NewStream(w, ch, l)
	defer s.Close()
	s.Serve()

	msgc := make(chan *json.RawMessage)
	go func() {
		defer close(msgc)
		dec := json.NewDecoder(rc)
		for {
			var m json.RawMessage
			if err := dec.Decode(&m); err != nil {
				if err != io.EOF {
					l.Error("decoding logagg stream", err)
				}
				return
			}
			msgc <- &m
		}
	}()

	for {
		select {
		case m := <-msgc:
			if m == nil {
				ch <- &sseLogChunk{Event: "eof"}
				return
			}
			// write to sse
			select {
			case ch <- &sseLogChunk{Event: "message", Data: *m}:
			case <-s.Done:
				return
			case <-ctx.Done():
				return
			}
		case <-s.Done:
			return
		case <-ctx.Done():
			return
		}
	}
}
<commit_msg>controller: special handling of host-prefixed JobID for app logs<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-sql"
	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/pq/hstore"
	"github.com/flynn/flynn/Godeps/_workspace/src/golang.org/x/net/context"
	"github.com/flynn/flynn/controller/name"
	"github.com/flynn/flynn/controller/schema"
	ct "github.com/flynn/flynn/controller/types"
	logaggc "github.com/flynn/flynn/logaggregator/client"
	"github.com/flynn/flynn/pkg/ctxhelper"
	"github.com/flynn/flynn/pkg/httphelper"
	"github.com/flynn/flynn/pkg/postgres"
	"github.com/flynn/flynn/pkg/random"
	"github.com/flynn/flynn/pkg/sse"
	routerc "github.com/flynn/flynn/router/client"
	"github.com/flynn/flynn/router/types"
)

type AppRepo struct {
	router        routerc.Client
	defaultDomain string

	db *postgres.DB
}

type appUpdate map[string]interface{}

func NewAppRepo(db *postgres.DB, defaultDomain string, router routerc.Client) *AppRepo {
	return &AppRepo{db: db, defaultDomain: defaultDomain, router: router}
}

var appNamePattern = regexp.MustCompile(`^[a-z\d]+(-[a-z\d]+)*$`)

func (r *AppRepo) Add(data interface{}) error {
	app := data.(*ct.App)
	if app.Name == "" {
		var nameID uint32
		if err := r.db.QueryRow("SELECT nextval('name_ids')").Scan(&nameID); err != nil {
			return err
		}
		app.Name = name.Get(nameID)
	}
	if len(app.Name) > 100 || !appNamePattern.MatchString(app.Name) {
		return ct.ValidationError{Field: "name", Message: "is invalid"}
	}
	if app.ID == "" {
		app.ID = random.UUID()
	}
	if app.Strategy == "" {
		app.Strategy = "all-at-once"
	}
	meta := metaToHstore(app.Meta)
	if err := r.db.QueryRow("INSERT INTO apps (app_id, name, meta, strategy) VALUES ($1, $2, $3, $4) RETURNING created_at, updated_at", app.ID, app.Name, meta, app.Strategy).Scan(&app.CreatedAt, &app.UpdatedAt); err != nil {
		if postgres.IsUniquenessError(err, "apps_name_idx") {
			return httphelper.ObjectExistsErr(fmt.Sprintf("application %q already exists", app.Name))
		}
		return err
	}
	app.ID = postgres.CleanUUID(app.ID)
	if !app.System() && r.defaultDomain != "" {
		route := (&router.HTTPRoute{
			Domain:  fmt.Sprintf("%s.%s", app.Name, r.defaultDomain),
			Service: app.Name + "-web",
		}).ToRoute()
		route.ParentRef = routeParentRef(app.ID)
		if err := r.router.CreateRoute(route); err != nil {
			log.Printf("Error creating default route for %s: %s", app.Name, err)
		}
	}
	return nil
}

func scanApp(s postgres.Scanner) (*ct.App, error) {
	app := &ct.App{}
	var meta hstore.Hstore
	err := s.Scan(&app.ID, &app.Name, &meta, &app.Strategy, &app.CreatedAt, &app.UpdatedAt)
	if err == sql.ErrNoRows {
		err = ErrNotFound
	}
	if len(meta.Map) > 0 {
		app.Meta = make(map[string]string, len(meta.Map))
		for k, v := range meta.Map {
			app.Meta[k] = v.String
		}
	}
	app.ID = postgres.CleanUUID(app.ID)
	return app, err
}

var idPattern = regexp.MustCompile(`^[a-f0-9]{8}-?([a-f0-9]{4}-?){3}[a-f0-9]{12}$`)

type rowQueryer interface {
	QueryRow(query string, args ...interface{}) postgres.Scanner
}

func selectApp(db rowQueryer, id string, update bool) (*ct.App, error) {
	var row postgres.Scanner
	query := "SELECT app_id, name, meta, strategy, created_at, updated_at FROM apps WHERE deleted_at IS NULL AND "
	var suffix string
	if update {
		suffix = " FOR UPDATE"
	}
	if idPattern.MatchString(id) {
		row = db.QueryRow(query+"(app_id = $1 OR name = $2) LIMIT 1"+suffix, id, id)
	} else {
		row = db.QueryRow(query+"name = $1"+suffix, id)
	}
	return scanApp(row)
}

func (r *AppRepo) Get(id string) (interface{}, error) {
	return selectApp(r.db, id, false)
}

func (r *AppRepo) Update(id string, data map[string]interface{}) (interface{}, error) {
	tx, err := r.db.Begin()
	if err != nil {
		return nil, err
	}
	app, err := selectApp(tx, id, true)
	if err != nil {
		tx.Rollback()
		return nil, err
	}

	for k, v := range data {
		switch k {
		case "strategy":
			strategy, ok := v.(string)
			if !ok {
				tx.Rollback()
				return nil, fmt.Errorf("controller: expected string, got %T", v)
			}
			if _, err := tx.Exec("UPDATE apps SET strategy = $2, updated_at = now() WHERE app_id = $1", app.ID, strategy); err != nil {
				tx.Rollback()
				return nil, err
			}
		case "meta":
			data, ok := v.(map[string]interface{})
			if !ok {
				tx.Rollback()
				return nil, fmt.Errorf("controller: expected map[string]interface{}, got %T", v)
			}
			var meta hstore.Hstore
			meta.Map = make(map[string]sql.NullString, len(data))
			app.Meta = make(map[string]string, len(data))
			for k, v := range data {
				s, ok := v.(string)
				if !ok {
					tx.Rollback()
					return nil, fmt.Errorf("controller: expected string, got %T", v)
				}
				meta.Map[k] = sql.NullString{String: s, Valid: true}
				app.Meta[k] = s
			}
			if _, err := tx.Exec("UPDATE apps SET meta = $2, updated_at = now() WHERE app_id = $1", app.ID, meta); err != nil {
				tx.Rollback()
				return nil, err
			}
		}
	}

	return app, tx.Commit()
}

func (r *AppRepo) Remove(id string) error {
	tx, err := r.db.Begin()
	if err != nil {
		return err
	}

	if !idPattern.MatchString(id) {
		app, err := selectApp(r.db, id, false)
		if err != nil {
			return err
		}
		id = app.ID
	}
	_, err = tx.Exec("UPDATE apps SET deleted_at = now() WHERE app_id = $1 AND deleted_at IS NULL", id)
	if err != nil {
		tx.Rollback()
		return err
	}

	_, err = tx.Exec("UPDATE formations SET deleted_at = now(), processes = NULL, updated_at = now() WHERE app_id = $1 AND deleted_at IS NULL", id)
	if err != nil {
		tx.Rollback()
		return err
	}
	_, err = tx.Exec("UPDATE app_resources SET deleted_at = now() WHERE app_id = $1 AND deleted_at IS NULL", id)
	if err != nil {
		tx.Rollback()
		return err
	}
	return tx.Commit()
}

func (r *AppRepo) List() (interface{}, error) {
	rows, err := r.db.Query("SELECT app_id, name, meta, strategy, created_at, updated_at FROM apps WHERE deleted_at IS NULL ORDER BY created_at DESC")
	if err != nil {
		return nil, err
	}
	apps := []*ct.App{}
	for rows.Next() {
		app, err := scanApp(rows)
		if err != nil {
			rows.Close()
			return nil, err
		}
		apps = append(apps, app)
	}
	return apps, rows.Err()
}

func (r *AppRepo) SetRelease(appID string, releaseID string) error {
	return r.db.Exec("UPDATE apps SET release_id = $2, updated_at = now() WHERE app_id = $1", appID, releaseID)
}

func (r *AppRepo) GetRelease(id string) (*ct.Release, error) {
	row := r.db.QueryRow("SELECT r.release_id, r.artifact_id, r.data, r.created_at FROM apps a JOIN releases r USING (release_id) WHERE a.app_id = $1", id)
	return scanRelease(row)
}

func (c *controllerAPI) UpdateApp(ctx context.Context, rw http.ResponseWriter, req *http.Request) {
	params, _ := ctxhelper.ParamsFromContext(ctx)

	var data appUpdate
	if err := httphelper.DecodeJSON(req, &data); err != nil {
		respondWithError(rw, err)
		return
	}

	if err := schema.Validate(data); err != nil {
		respondWithError(rw, err)
		return
	}

	app, err := c.appRepo.Update(params.ByName("apps_id"), data)
	if err != nil {
		respondWithError(rw, err)
		return
	}
	httphelper.JSON(rw, 200, app)
}

func (c *controllerAPI) AppLog(ctx context.Context, w http.ResponseWriter, req *http.Request) {
	ctx, cancel := context.WithCancel(ctx)

	opts := logaggc.LogOpts{
		Follow: req.FormValue("follow") == "true",
		JobID:  req.FormValue("job_id"),
	}
	if opts.JobID != "" {
		// Temporary handling of combined JobID format. Logs are sent to aggregator
		// without host- prefix. https://github.com/flynn/flynn/issues/1238
		index := strings.LastIndex(opts.JobID, "-")
		if len(opts.JobID) > index {
			opts.JobID = opts.JobID[index+1:]
		}
	}
	if vals, ok := req.Form["process_type"]; ok && len(vals) > 0 {
		opts.ProcessType = &vals[len(vals)-1]
	}
	if strLines := req.FormValue("lines"); strLines != "" {
		lines, err := strconv.Atoi(req.FormValue("lines"))
		if err != nil {
			respondWithError(w, err)
			return
		}
		opts.Lines = &lines
	}
	rc, err := c.logaggc.GetLog(c.getApp(ctx).ID, &opts)
	if err != nil {
		respondWithError(w, err)
		return
	}

	if cn, ok := w.(http.CloseNotifier); ok {
		go func() {
			select {
			case <-cn.CloseNotify():
				rc.Close()
			case <-ctx.Done():
			}
		}()
	}
	defer cancel()
	defer rc.Close()

	if !strings.Contains(req.Header.Get("Accept"), "text/event-stream") {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(200)
		// Send headers right away if following
		if wf, ok := w.(http.Flusher); ok && opts.Follow {
			wf.Flush()
		}

		fw := httphelper.FlushWriter{Writer: w, Enabled: opts.Follow}
		io.Copy(fw, rc)
		return
	}

	ch := make(chan *sseLogChunk)
	l, _ := ctxhelper.LoggerFromContext(ctx)
	s := sse.NewStream(w, ch, l)
	defer s.Close()
	s.Serve()

	msgc := make(chan *json.RawMessage)
	go func() {
		defer close(msgc)
		dec := json.NewDecoder(rc)
		for {
			var m json.RawMessage
			if err := dec.Decode(&m); err != nil {
				if err != io.EOF {
					l.Error("decoding logagg stream", err)
				}
				return
			}
			msgc <- &m
		}
	}()

	for {
		select {
		case m := <-msgc:
			if m == nil {
				ch <- &sseLogChunk{Event: "eof"}
				return
			}
			// write to sse
			select {
			case ch <- &sseLogChunk{Event: "message", Data: *m}:
			case <-s.Done:
				return
			case <-ctx.Done():
				return
			}
		case <-s.Done:
			return
		case <-ctx.Done():
			return
		}
	}
}
<|endoftext|>


<commit_before>package mppostgres

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"regexp"
	"strconv"

	"github.com/jmoiron/sqlx"
	// PostgreSQL Driver
	_ "github.com/lib/pq"
	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/mackerel-agent/logging"
)

var logger = logging.GetLogger("metrics.plugin.postgres")

var graphdef = map[string]mp.Graphs{
	"postgres.connections": {
		Label: "Postgres Connections",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "active", Label: "Active", Diff: false, Stacked: true},
			{Name: "active_waiting", Label: "Active waiting", Diff: false, Stacked: true},
			{Name: "idle", Label: "Idle", Diff: false, Stacked: true},
			{Name: "idle_in_transaction", Label: "Idle in transaction", Diff: false, Stacked: true},
			{Name: "idle_in_transaction_aborted_", Label: "Idle in transaction (aborted)", Diff: false, Stacked: true},
			{Name: "fastpath_function_call", Label: "fast-path function call", Diff: false, Stacked: true},
			{Name: "disabled", Label: "Disabled", Diff: false, Stacked: true},
		},
	},
	"postgres.commits": {
		Label: "Postgres Commits",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "xact_commit", Label: "Xact Commit", Diff: true, Stacked: false},
			{Name: "xact_rollback", Label: "Xact Rollback", Diff: true, Stacked: false},
		},
	},
	"postgres.blocks": {
		Label: "Postgres Blocks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "blks_read", Label: "Blocks Read", Diff: true, Stacked: false},
			{Name: "blks_hit", Label: "Blocks Hit", Diff: true, Stacked: false},
		},
	},
	"postgres.rows": {
		Label: "Postgres Rows",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "tup_returned", Label: "Returned Rows", Diff: true, Stacked: false},
			{Name: "tup_fetched", Label: "Fetched Rows", Diff: true, Stacked: true},
			{Name: "tup_inserted", Label: "Inserted Rows", Diff: true, Stacked: true},
			{Name: "tup_updated", Label: "Updated Rows", Diff: true, Stacked: true},
			{Name: "tup_deleted", Label: "Deleted Rows", Diff: true, Stacked: true},
		},
	},
	"postgres.size": {
		Label: "Postgres Data Size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "total_size", Label: "Total Size", Diff: false, Stacked: false},
		},
	},
	"postgres.deadlocks": {
		Label: "Postgres Dead Locks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "deadlocks", Label: "Deadlocks", Diff: true, Stacked: false},
		},
	},
	"postgres.iotime": {
		Label: "Postgres Block I/O time",
		Unit:  "float",
		Metrics: []mp.Metrics{
			{Name: "blk_read_time", Label: "Block Read Time (ms)", Diff: true, Stacked: false},
			{Name: "blk_write_time", Label: "Block Write Time (ms)", Diff: true, Stacked: false},
		},
	},
	"postgres.tempfile": {
		Label: "Postgres Temporary file",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "temp_bytes", Label: "Temporary file size (byte)", Diff: true, Stacked: false},
		},
	},
}

// PostgresPlugin mackerel plugin for PostgreSQL
type PostgresPlugin struct {
	Host     string
	Port     string
	Username string
	Password string
	SSLmode  string
	Timeout  int
	Tempfile string
	Option   string
}

func fetchStatDatabase(db *sqlx.DB) (map[string]interface{}, error) {
	db = db.Unsafe()
	rows, err := db.Queryx(`SELECT * FROM pg_stat_database`)
	if err != nil {
		logger.Errorf("Failed to select pg_stat_database. %s", err)
		return nil, err
	}

	type pgStat struct {
		XactCommit   uint64   `db:"xact_commit"`
		XactRollback uint64   `db:"xact_rollback"`
		BlksRead     uint64   `db:"blks_read"`
		BlksHit      uint64   `db:"blks_hit"`
		BlkReadTime  *float64 `db:"blk_read_time"`
		BlkWriteTime *float64 `db:"blk_write_time"`
		TupReturned  uint64   `db:"tup_returned"`
		TupFetched   uint64   `db:"tup_fetched"`
		TupInserted  uint64   `db:"tup_inserted"`
		TupUpdated   uint64   `db:"tup_updated"`
		TupDeleted   uint64   `db:"tup_deleted"`
		Deadlocks    *uint64  `db:"deadlocks"`
		TempBytes    *uint64  `db:"temp_bytes"`
	}

	totalStat := pgStat{}
	for rows.Next() {
		p := pgStat{}
		if err := rows.StructScan(&p); err != nil {
			logger.Warningf("Failed to scan. %s", err)
			continue
		}
		totalStat.XactCommit += p.XactCommit
		totalStat.XactRollback += p.XactRollback
		totalStat.BlksRead += p.BlksRead
		totalStat.BlksHit += p.BlksHit
		if p.BlkReadTime != nil {
			if totalStat.BlkReadTime == nil {
				totalStat.BlkReadTime = p.BlkReadTime
			} else {
				*totalStat.BlkReadTime += *p.BlkReadTime
			}
		}
		if p.BlkWriteTime != nil {
			if totalStat.BlkWriteTime == nil {
				totalStat.BlkWriteTime = p.BlkWriteTime
			} else {
				*totalStat.BlkWriteTime += *p.BlkWriteTime
			}
		}
		totalStat.TupReturned += p.TupReturned
		totalStat.TupFetched += p.TupFetched
		totalStat.TupInserted += p.TupInserted
		totalStat.TupUpdated += p.TupUpdated
		totalStat.TupDeleted += p.TupDeleted
		if p.Deadlocks != nil {
			if totalStat.Deadlocks == nil {
				totalStat.Deadlocks = p.Deadlocks
			} else {
				*totalStat.Deadlocks += *p.Deadlocks
			}
		}
		if p.TempBytes != nil {
			if totalStat.TempBytes == nil {
				totalStat.TempBytes = p.TempBytes
			} else {
				*totalStat.TempBytes += *p.TempBytes
			}
		}
	}
	stat := make(map[string]interface{})
	stat["xact_commit"] = totalStat.XactCommit
	stat["xact_rollback"] = totalStat.XactRollback
	stat["blks_read"] = totalStat.BlksRead
	stat["blks_hit"] = totalStat.BlksHit
	if totalStat.BlkReadTime != nil {
		stat["blk_read_time"] = *totalStat.BlkReadTime
	}
	if totalStat.BlkWriteTime != nil {
		stat["blk_write_time"] = *totalStat.BlkWriteTime
	}
	stat["tup_returned"] = totalStat.TupReturned
	stat["tup_fetched"] = totalStat.TupFetched
	stat["tup_inserted"] = totalStat.TupInserted
	stat["tup_updated"] = totalStat.TupUpdated
	stat["tup_deleted"] = totalStat.TupDeleted
	if totalStat.Deadlocks != nil {
		stat["deadlocks"] = *totalStat.Deadlocks
	}
	if totalStat.TempBytes != nil {
		stat["temp_bytes"] = *totalStat.TempBytes
	}
	return stat, nil
}

func fetchConnections(db *sqlx.DB, version version) (map[string]interface{}, error) {
	var query string

	if version.first > 9 || version.first == 9 && version.second >= 6 {
		query = `select count(*), state, wait_event is not null from pg_stat_activity group by state, wait_event is not null`
	} else {
		query = `select count(*), state, waiting from pg_stat_activity group by state, waiting`
	}
	rows, err := db.Query(query)
	if err != nil {
		logger.Errorf("Failed to select pg_stat_activity. %s", err)
		return nil, err
	}

	stat := make(map[string]interface{})

	normalizeRe := regexp.MustCompile("[^a-zA-Z0-9_-]+")

	for rows.Next() {
		var count float64
		var waiting bool
		var state string
		if err := rows.Scan(&count, &state, &waiting); err != nil {
			logger.Warningf("Failed to scan %s", err)
			continue
		}
		state = normalizeRe.ReplaceAllString(state, "_")
		if waiting {
			state += "_waiting"
		}
		stat[state] = float64(count)
	}

	return stat, nil
}

func fetchDatabaseSize(db *sqlx.DB) (map[string]interface{}, error) {
	rows, err := db.Query("select sum(pg_database_size(datname)) as dbsize from pg_database")
	if err != nil {
		logger.Errorf("Failed to select pg_database_size. %s", err)
		return nil, err
	}

	var totalSize float64
	for rows.Next() {
		var dbsize float64
		if err := rows.Scan(&dbsize); err != nil {
			logger.Warningf("Failed to scan %s", err)
			continue
		}
		totalSize += dbsize
	}

	return map[string]interface{}{
		"total_size": totalSize,
	}, nil
}

var versionRe = regexp.MustCompile("PostgreSQL (\\d+)\\.(\\d+)(\\.(\\d+))? ")

type version struct {
	first  uint
	second uint
	thrird uint
}

func fetchVersion(db *sqlx.DB) (version, error) {

	res := version{}

	rows, err := db.Query("select version()")
	if err != nil {
		logger.Errorf("Failed to select version(). %s", err)
		return res, err
	}

	for rows.Next() {
		var versionStr string
		var first, second, third uint64
		if err := rows.Scan(&versionStr); err != nil {
			return res, err
		}

		// ref. https://www.postgresql.org/support/versioning/

		submatch := versionRe.FindStringSubmatch(versionStr)
		if len(submatch) >= 4 {
			first, err = strconv.ParseUint(submatch[1], 10, 0)
			if err != nil {
				return res, err
			}
			second, err = strconv.ParseUint(submatch[2], 10, 0)
			if err != nil {
				return res, err
			}
			if len(submatch) == 5 {
				third, err = strconv.ParseUint(submatch[4], 10, 0)
				if err != nil {
					return res, err
				}
			}
			res = version{uint(first), uint(second), uint(third)}
			return res, err
		}
	}
	return res, errors.New("failed to select version()")
}

func mergeStat(dst, src map[string]interface{}) {
	for k, v := range src {
		dst[k] = v
	}
}

// FetchMetrics interface for mackerelplugin
func (p PostgresPlugin) FetchMetrics() (map[string]interface{}, error) {

	db, err := sqlx.Connect("postgres", fmt.Sprintf("user=%s password=%s host=%s port=%s sslmode=%s connect_timeout=%d %s", p.Username, p.Password, p.Host, p.Port, p.SSLmode, p.Timeout, p.Option))
	if err != nil {
		logger.Errorf("FetchMetrics: %s", err)
		return nil, err
	}
	defer db.Close()

	version, err := fetchVersion(db)
	if err != nil {
		logger.Warningf("FetchMetrics: %s", err)
		return nil, err
	}

	statStatDatabase, err := fetchStatDatabase(db)
	if err != nil {
		return nil, err
	}
	statConnections, err := fetchConnections(db, version)
	if err != nil {
		return nil, err
	}
	statDatabaseSize, err := fetchDatabaseSize(db)
	if err != nil {
		return nil, err
	}

	stat := make(map[string]interface{})
	mergeStat(stat, statStatDatabase)
	mergeStat(stat, statConnections)
	mergeStat(stat, statDatabaseSize)

	return stat, err
}

// GraphDefinition interface for mackerelplugin
func (p PostgresPlugin) GraphDefinition() map[string]mp.Graphs {
	return graphdef
}

// Do the plugin
func Do() {
	optHost := flag.String("hostname", "localhost", "Hostname to login to")
	optPort := flag.String("port", "5432", "Database port")
	optUser := flag.String("user", "", "Postgres User")
	optDatabase := flag.String("database", "", "Database name")
	optPass := flag.String("password", "", "Postgres Password")
	optSSLmode := flag.String("sslmode", "disable", "Whether or not to use SSL")
	optConnectTimeout := flag.Int("connect_timeout", 5, "Maximum wait for connection, in seconds.")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	if *optUser == "" {
		logger.Warningf("user is required")
		flag.PrintDefaults()
		os.Exit(1)
	}
	if *optPass == "" {
		logger.Warningf("password is required")
		flag.PrintDefaults()
		os.Exit(1)
	}
	option := ""
	if *optDatabase != "" {
		option = fmt.Sprintf("dbname=%s", *optDatabase)
	}

	var postgres PostgresPlugin
	postgres.Host = *optHost
	postgres.Port = *optPort
	postgres.Username = *optUser
	postgres.Password = *optPass
	postgres.SSLmode = *optSSLmode
	postgres.Timeout = *optConnectTimeout
	postgres.Option = option

	helper := mp.NewMackerelPlugin(postgres)

	helper.Tempfile = *optTempfile
	helper.Run()
}
<commit_msg>postgres: zero fill<commit_after>package mppostgres

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"regexp"
	"strconv"

	"github.com/jmoiron/sqlx"
	// PostgreSQL Driver
	_ "github.com/lib/pq"
	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/mackerel-agent/logging"
)

var logger = logging.GetLogger("metrics.plugin.postgres")

var graphdef = map[string]mp.Graphs{
	"postgres.connections": {
		Label: "Postgres Connections",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "active", Label: "Active", Diff: false, Stacked: true},
			{Name: "active_waiting", Label: "Active waiting", Diff: false, Stacked: true},
			{Name: "idle", Label: "Idle", Diff: false, Stacked: true},
			{Name: "idle_in_transaction", Label: "Idle in transaction", Diff: false, Stacked: true},
			{Name: "idle_in_transaction_aborted_", Label: "Idle in transaction (aborted)", Diff: false, Stacked: true},
			{Name: "fastpath_function_call", Label: "fast-path function call", Diff: false, Stacked: true},
			{Name: "disabled", Label: "Disabled", Diff: false, Stacked: true},
		},
	},
	"postgres.commits": {
		Label: "Postgres Commits",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "xact_commit", Label: "Xact Commit", Diff: true, Stacked: false},
			{Name: "xact_rollback", Label: "Xact Rollback", Diff: true, Stacked: false},
		},
	},
	"postgres.blocks": {
		Label: "Postgres Blocks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "blks_read", Label: "Blocks Read", Diff: true, Stacked: false},
			{Name: "blks_hit", Label: "Blocks Hit", Diff: true, Stacked: false},
		},
	},
	"postgres.rows": {
		Label: "Postgres Rows",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "tup_returned", Label: "Returned Rows", Diff: true, Stacked: false},
			{Name: "tup_fetched", Label: "Fetched Rows", Diff: true, Stacked: true},
			{Name: "tup_inserted", Label: "Inserted Rows", Diff: true, Stacked: true},
			{Name: "tup_updated", Label: "Updated Rows", Diff: true, Stacked: true},
			{Name: "tup_deleted", Label: "Deleted Rows", Diff: true, Stacked: true},
		},
	},
	"postgres.size": {
		Label: "Postgres Data Size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "total_size", Label: "Total Size", Diff: false, Stacked: false},
		},
	},
	"postgres.deadlocks": {
		Label: "Postgres Dead Locks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "deadlocks", Label: "Deadlocks", Diff: true, Stacked: false},
		},
	},
	"postgres.iotime": {
		Label: "Postgres Block I/O time",
		Unit:  "float",
		Metrics: []mp.Metrics{
			{Name: "blk_read_time", Label: "Block Read Time (ms)", Diff: true, Stacked: false},
			{Name: "blk_write_time", Label: "Block Write Time (ms)", Diff: true, Stacked: false},
		},
	},
	"postgres.tempfile": {
		Label: "Postgres Temporary file",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "temp_bytes", Label: "Temporary file size (byte)", Diff: true, Stacked: false},
		},
	},
}

// PostgresPlugin mackerel plugin for PostgreSQL
type PostgresPlugin struct {
	Host     string
	Port     string
	Username string
	Password string
	SSLmode  string
	Timeout  int
	Tempfile string
	Option   string
}

func fetchStatDatabase(db *sqlx.DB) (map[string]interface{}, error) {
	db = db.Unsafe()
	rows, err := db.Queryx(`SELECT * FROM pg_stat_database`)
	if err != nil {
		logger.Errorf("Failed to select pg_stat_database. %s", err)
		return nil, err
	}

	type pgStat struct {
		XactCommit   uint64   `db:"xact_commit"`
		XactRollback uint64   `db:"xact_rollback"`
		BlksRead     uint64   `db:"blks_read"`
		BlksHit      uint64   `db:"blks_hit"`
		BlkReadTime  *float64 `db:"blk_read_time"`
		BlkWriteTime *float64 `db:"blk_write_time"`
		TupReturned  uint64   `db:"tup_returned"`
		TupFetched   uint64   `db:"tup_fetched"`
		TupInserted  uint64   `db:"tup_inserted"`
		TupUpdated   uint64   `db:"tup_updated"`
		TupDeleted   uint64   `db:"tup_deleted"`
		Deadlocks    *uint64  `db:"deadlocks"`
		TempBytes    *uint64  `db:"temp_bytes"`
	}

	totalStat := pgStat{}
	for rows.Next() {
		p := pgStat{}
		if err := rows.StructScan(&p); err != nil {
			logger.Warningf("Failed to scan. %s", err)
			continue
		}
		totalStat.XactCommit += p.XactCommit
		totalStat.XactRollback += p.XactRollback
		totalStat.BlksRead += p.BlksRead
		totalStat.BlksHit += p.BlksHit
		if p.BlkReadTime != nil {
			if totalStat.BlkReadTime == nil {
				totalStat.BlkReadTime = p.BlkReadTime
			} else {
				*totalStat.BlkReadTime += *p.BlkReadTime
			}
		}
		if p.BlkWriteTime != nil {
			if totalStat.BlkWriteTime == nil {
				totalStat.BlkWriteTime = p.BlkWriteTime
			} else {
				*totalStat.BlkWriteTime += *p.BlkWriteTime
			}
		}
		totalStat.TupReturned += p.TupReturned
		totalStat.TupFetched += p.TupFetched
		totalStat.TupInserted += p.TupInserted
		totalStat.TupUpdated += p.TupUpdated
		totalStat.TupDeleted += p.TupDeleted
		if p.Deadlocks != nil {
			if totalStat.Deadlocks == nil {
				totalStat.Deadlocks = p.Deadlocks
			} else {
				*totalStat.Deadlocks += *p.Deadlocks
			}
		}
		if p.TempBytes != nil {
			if totalStat.TempBytes == nil {
				totalStat.TempBytes = p.TempBytes
			} else {
				*totalStat.TempBytes += *p.TempBytes
			}
		}
	}
	stat := make(map[string]interface{})
	stat["xact_commit"] = totalStat.XactCommit
	stat["xact_rollback"] = totalStat.XactRollback
	stat["blks_read"] = totalStat.BlksRead
	stat["blks_hit"] = totalStat.BlksHit
	if totalStat.BlkReadTime != nil {
		stat["blk_read_time"] = *totalStat.BlkReadTime
	}
	if totalStat.BlkWriteTime != nil {
		stat["blk_write_time"] = *totalStat.BlkWriteTime
	}
	stat["tup_returned"] = totalStat.TupReturned
	stat["tup_fetched"] = totalStat.TupFetched
	stat["tup_inserted"] = totalStat.TupInserted
	stat["tup_updated"] = totalStat.TupUpdated
	stat["tup_deleted"] = totalStat.TupDeleted
	if totalStat.Deadlocks != nil {
		stat["deadlocks"] = *totalStat.Deadlocks
	}
	if totalStat.TempBytes != nil {
		stat["temp_bytes"] = *totalStat.TempBytes
	}
	return stat, nil
}

func fetchConnections(db *sqlx.DB, version version) (map[string]interface{}, error) {
	var query string

	if version.first > 9 || version.first == 9 && version.second >= 6 {
		query = `select count(*), state, wait_event is not null from pg_stat_activity group by state, wait_event is not null`
	} else {
		query = `select count(*), state, waiting from pg_stat_activity group by state, waiting`
	}
	rows, err := db.Query(query)
	if err != nil {
		logger.Errorf("Failed to select pg_stat_activity. %s", err)
		return nil, err
	}

	stat := map[string]interface{}{
		"active":                       0.0,
		"active_waiting":               0.0,
		"idle":                         0.0,
		"idle_in_transaction":          0.0,
		"idle_in_transaction_aborted_": 0.0,
		"fastpath_function_call":       0.0,
		"disabled":                     0.0,
	}

	normalizeRe := regexp.MustCompile("[^a-zA-Z0-9_-]+")

	for rows.Next() {
		var count float64
		var waiting bool
		var state string
		if err := rows.Scan(&count, &state, &waiting); err != nil {
			logger.Warningf("Failed to scan %s", err)
			continue
		}
		state = normalizeRe.ReplaceAllString(state, "_")
		if waiting {
			state += "_waiting"
		}
		stat[state] = float64(count)
	}

	return stat, nil
}

func fetchDatabaseSize(db *sqlx.DB) (map[string]interface{}, error) {
	rows, err := db.Query("select sum(pg_database_size(datname)) as dbsize from pg_database")
	if err != nil {
		logger.Errorf("Failed to select pg_database_size. %s", err)
		return nil, err
	}

	var totalSize float64
	for rows.Next() {
		var dbsize float64
		if err := rows.Scan(&dbsize); err != nil {
			logger.Warningf("Failed to scan %s", err)
			continue
		}
		totalSize += dbsize
	}

	return map[string]interface{}{
		"total_size": totalSize,
	}, nil
}

var versionRe = regexp.MustCompile("PostgreSQL (\\d+)\\.(\\d+)(\\.(\\d+))? ")

type version struct {
	first  uint
	second uint
	thrird uint
}

func fetchVersion(db *sqlx.DB) (version, error) {

	res := version{}

	rows, err := db.Query("select version()")
	if err != nil {
		logger.Errorf("Failed to select version(). %s", err)
		return res, err
	}

	for rows.Next() {
		var versionStr string
		var first, second, third uint64
		if err := rows.Scan(&versionStr); err != nil {
			return res, err
		}

		// ref. https://www.postgresql.org/support/versioning/

		submatch := versionRe.FindStringSubmatch(versionStr)
		if len(submatch) >= 4 {
			first, err = strconv.ParseUint(submatch[1], 10, 0)
			if err != nil {
				return res, err
			}
			second, err = strconv.ParseUint(submatch[2], 10, 0)
			if err != nil {
				return res, err
			}
			if len(submatch) == 5 {
				third, err = strconv.ParseUint(submatch[4], 10, 0)
				if err != nil {
					return res, err
				}
			}
			res = version{uint(first), uint(second), uint(third)}
			return res, err
		}
	}
	return res, errors.New("failed to select version()")
}

func mergeStat(dst, src map[string]interface{}) {
	for k, v := range src {
		dst[k] = v
	}
}

// FetchMetrics interface for mackerelplugin
func (p PostgresPlugin) FetchMetrics() (map[string]interface{}, error) {

	db, err := sqlx.Connect("postgres", fmt.Sprintf("user=%s password=%s host=%s port=%s sslmode=%s connect_timeout=%d %s", p.Username, p.Password, p.Host, p.Port, p.SSLmode, p.Timeout, p.Option))
	if err != nil {
		logger.Errorf("FetchMetrics: %s", err)
		return nil, err
	}
	defer db.Close()

	version, err := fetchVersion(db)
	if err != nil {
		logger.Warningf("FetchMetrics: %s", err)
		return nil, err
	}

	statStatDatabase, err := fetchStatDatabase(db)
	if err != nil {
		return nil, err
	}
	statConnections, err := fetchConnections(db, version)
	if err != nil {
		return nil, err
	}
	statDatabaseSize, err := fetchDatabaseSize(db)
	if err != nil {
		return nil, err
	}

	stat := make(map[string]interface{})
	mergeStat(stat, statStatDatabase)
	mergeStat(stat, statConnections)
	mergeStat(stat, statDatabaseSize)

	return stat, err
}

// GraphDefinition interface for mackerelplugin
func (p PostgresPlugin) GraphDefinition() map[string]mp.Graphs {
	return graphdef
}

// Do the plugin
func Do() {
	optHost := flag.String("hostname", "localhost", "Hostname to login to")
	optPort := flag.String("port", "5432", "Database port")
	optUser := flag.String("user", "", "Postgres User")
	optDatabase := flag.String("database", "", "Database name")
	optPass := flag.String("password", "", "Postgres Password")
	optSSLmode := flag.String("sslmode", "disable", "Whether or not to use SSL")
	optConnectTimeout := flag.Int("connect_timeout", 5, "Maximum wait for connection, in seconds.")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	if *optUser == "" {
		logger.Warningf("user is required")
		flag.PrintDefaults()
		os.Exit(1)
	}
	if *optPass == "" {
		logger.Warningf("password is required")
		flag.PrintDefaults()
		os.Exit(1)
	}
	option := ""
	if *optDatabase != "" {
		option = fmt.Sprintf("dbname=%s", *optDatabase)
	}

	var postgres PostgresPlugin
	postgres.Host = *optHost
	postgres.Port = *optPort
	postgres.Username = *optUser
	postgres.Password = *optPass
	postgres.SSLmode = *optSSLmode
	postgres.Timeout = *optConnectTimeout
	postgres.Option = option

	helper := mp.NewMackerelPlugin(postgres)

	helper.Tempfile = *optTempfile
	helper.Run()
}
<|endoftext|>


<commit_before>//
// Copyright © 2016-2018 Guy M. Allard
//
// 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 sngecomm provides common functionality used in the stompngo_examples
project.
*/
package sngecomm

import (
	"bytes"
	"log"
	"os"
	"strconv"
	"sync"
	//
	// "github.com/gmallard/stompngo"
)

var (
	//
	nqs     = 1               // Default number of queues for multi-queue demo(s)
	nqsLock sync.Mutex        // nqs variable lock
	mdml    = 1024 * 32       // Message data max length of variable message, 32K
	md      = make([]byte, 1) // Additional message data, primed during init()
	pbc     = 64              // Number of bytes to print (used in some examples that receive).

	ngors    = 1  // Number of go routines to use (publish)
	gorsleep = "" // If non-empty, go routines will sleep (publish)

	//
	sendFact = 1.0 // Send sleep time factor
	recvFact = 1.0 // Receive sleep time factor
	//
	ackMode = "auto" // The default ack mode
	//
	pprof = false // Do not do profiling

	// TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256 = 0xC0,0x2F
	// TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256 = 0xC0,0x2B
	// TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384 = 0xC0,0x30
	cipherSuites = []uint16{
		0xc02f,
		0xc02b,
		0xc030,
	}

	useCustomCiphers = false // Set Custom Cipher Suite list

	memprof = "" // memory profile file
	cpuprof = "" // cpu profile file
)

const (
	// EOFMsg is the EOF message body
	EOFMsg = "STOMP_EOF"
)

// Initialization
func init() {
	p := "_123456789ABCDEF"
	c := mdml / len(p)
	b := []byte(p)
	md = bytes.Repeat(b, c) // A long string
	//
	memprof = os.Getenv("STOMP_MEMPROF")
	cpuprof = os.Getenv("STOMP_CPUPROF")
}

// Ngors sets the number of go routines
func Ngors() int {
	//
	if s := os.Getenv("STOMP_NGORS"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "NGORS conversion error", e)
		} else {
			ngors = int(i)
		}
	}
	return ngors
}

// Nqs sets the number of queues
func Nqs() int {
	//
	nqsLock.Lock()
	defer nqsLock.Unlock()
	if s := os.Getenv("STOMP_NQS"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "NQS conversion error", e)
		} else {
			nqs = int(i)
		}
	}
	return nqs
}

// Mdml sets the Max Data Message Length
func Mdml() int {
	if s := os.Getenv("STOMP_MDML"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "MDML conversion error", e)
		} else {
			mdml = int(i)
			p := "_123456789ABCDEF"
			c := mdml / len(p)
			b := []byte(p)
			md = bytes.Repeat(b, c) // A long string
		}
	}
	return mdml
}

// Pprof indicates whether to use profiling or not
func Pprof() bool {
	if am := os.Getenv("STOMP_PPROF"); am != "" {
		pprof = true
	}
	return pprof
}

// Memprof returns the memory profile file name
func Memprof() string {
	return memprof
}

// Cpuprof returns the CPU profile file name
func Cpuprof() string {
	return cpuprof
}

// AckMode returns an ACK mode value for those examples that use it.
func AckMode() string {
	if am := os.Getenv("STOMP_ACKMODE"); am != "" {
		if am == "auto" || am == "client" || am == "client-individual" {
			ackMode = am
		} else {
			log.Printf("v1:%v v2:%v\n", "ACKMODE error", am)
		}
	}
	return ackMode
}

// SendFactor returns the send sleep factor
func SendFactor() float64 {
	if s := os.Getenv("STOMP_SENDFACT"); s != "" {
		f, e := strconv.ParseFloat(s, 64)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "SENDFACT conversion error", e)
		} else {
			sendFact = float64(f)
		}
	}
	return sendFact
}

// RecvFactor returns the recv sleep factor
func RecvFactor() float64 {
	if s := os.Getenv("STOMP_RECVFACT"); s != "" {
		f, e := strconv.ParseFloat(s, 64)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "RECVFACT conversion error", e)
		} else {
			recvFact = float64(f)
		}
	}
	return recvFact
}

// Partial returns the partial byte slice for logging, random length
func Partial() []byte {
	r := int(ValueBetween(1, int64(mdml-1), 1.0))
	return md[0:r]
}

// PartialSubstr returns the partial string for logging, fixed length
func PartialSubstr(l int) []byte {
	return md[0:l]
}

// Pbc returns the byte count to log
func Pbc() int {
	if s := os.Getenv("STOMP_PBC"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "PBC conversion error", e)
		} else {
			pbc = int(i)
		}
	}
	return pbc
}

// Gorsleep returns an indication of whether go routines will sleep or not
func Gorsleep() string {
	gorsleep = os.Getenv("STOMP_GORSLEEP")
	return gorsleep
}

// RecvWait indicates whether to wait in receives to simulate message processing
func RecvWait() bool {
	f := os.Getenv("STOMP_RECVWAIT")
	if f == "" {
		return true
	}
	return false
}

// SendWait indicates whether to wait in sends to simulate message processing
func SendWait() bool {
	f := os.Getenv("STOMP_SENDWAIT")
	if f == "" {
		return true
	}
	return false
}

// SetMAXPROCS returns true if max procs are to be set
func SetMAXPROCS() bool {
	f := os.Getenv("STOMP_SETMAXPROCS")
	if f == "" {
		return false
	}
	return true
}

// UseCustomCiphers returns true if custom ciphers are to be used
func UseCustomCiphers() bool {
	f := os.Getenv("STOMP_USECUSTOMCIPHERS")
	if f == "" {
		return useCustomCiphers
	}
	useCustomCiphers = true
	return useCustomCiphers
}

// CustomCiphers returns a slice of custom ciphers
func CustomCiphers() []uint16 {
	if UseCustomCiphers() {
		return cipherSuites
	}
	return []uint16{}
}

// Logger returns an indication of whether to do logging
func Logger() string {
	return os.Getenv("STOMP_LOGGER")
}

// LogFile returns a log fle name or the empty string
func LogFile() string {
	return os.Getenv("STOMP_LOGFILE")
}

// UseEOF returns true if en EOF message is to be used.
func UseEOF() bool {
	if os.Getenv("STOMP_USEEOF") != "" {
		return true
	}
	return false
}
<commit_msg>Correct Max data lenght bug.<commit_after>//
// Copyright © 2016-2018 Guy M. Allard
//
// 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 sngecomm provides common functionality used in the stompngo_examples
project.
*/
package sngecomm

import (
	"bytes"
	"log"
	"os"
	"strconv"
	"sync"
	//
	// "github.com/gmallard/stompngo"
)

var (
	//
	nqs     = 1               // Default number of queues for multi-queue demo(s)
	nqsLock sync.Mutex        // nqs variable lock
	mdml    = 1024 * 32       // Message data max length of variable message, 32K
	md      = make([]byte, 1) // Additional message data, primed during init()
	pbc     = 64              // Number of bytes to print (used in some examples that receive).

	ngors    = 1  // Number of go routines to use (publish)
	gorsleep = "" // If non-empty, go routines will sleep (publish)

	//
	sendFact = 1.0 // Send sleep time factor
	recvFact = 1.0 // Receive sleep time factor
	//
	ackMode = "auto" // The default ack mode
	//
	pprof = false // Do not do profiling

	// TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256 = 0xC0,0x2F
	// TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256 = 0xC0,0x2B
	// TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384 = 0xC0,0x30
	cipherSuites = []uint16{
		0xc02f,
		0xc02b,
		0xc030,
	}

	useCustomCiphers = false // Set Custom Cipher Suite list

	memprof = "" // memory profile file
	cpuprof = "" // cpu profile file
)

const (
	// EOFMsg is the EOF message body
	EOFMsg = "STOMP_EOF"
)

// Initialization
func init() {
	p := "_123456789ABCDEF"
	// c := mdml / len(p)
	c := Mdml() / len(p)
	b := []byte(p)
	md = bytes.Repeat(b, c) // A long string
	//
	memprof = os.Getenv("STOMP_MEMPROF")
	cpuprof = os.Getenv("STOMP_CPUPROF")
}

// Ngors sets the number of go routines
func Ngors() int {
	//
	if s := os.Getenv("STOMP_NGORS"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "NGORS conversion error", e)
		} else {
			ngors = int(i)
		}
	}
	return ngors
}

// Nqs sets the number of queues
func Nqs() int {
	//
	nqsLock.Lock()
	defer nqsLock.Unlock()
	if s := os.Getenv("STOMP_NQS"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "NQS conversion error", e)
		} else {
			nqs = int(i)
		}
	}
	return nqs
}

// Mdml sets the Max Data Message Length
func Mdml() int {
	if s := os.Getenv("STOMP_MDML"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "MDML conversion error", e)
		} else {
			mdml = int(i)
			p := "_123456789ABCDEF"
			c := mdml / len(p)
			b := []byte(p)
			md = bytes.Repeat(b, c) // A long string
		}
	}
	return mdml
}

// Mdml sets the Max Data Message Length
func GetMdml() int {
	return mdml
}

// Pprof indicates whether to use profiling or not
func Pprof() bool {
	if am := os.Getenv("STOMP_PPROF"); am != "" {
		pprof = true
	}
	return pprof
}

// Memprof returns the memory profile file name
func Memprof() string {
	return memprof
}

// Cpuprof returns the CPU profile file name
func Cpuprof() string {
	return cpuprof
}

// AckMode returns an ACK mode value for those examples that use it.
func AckMode() string {
	if am := os.Getenv("STOMP_ACKMODE"); am != "" {
		if am == "auto" || am == "client" || am == "client-individual" {
			ackMode = am
		} else {
			log.Printf("v1:%v v2:%v\n", "ACKMODE error", am)
		}
	}
	return ackMode
}

// SendFactor returns the send sleep factor
func SendFactor() float64 {
	if s := os.Getenv("STOMP_SENDFACT"); s != "" {
		f, e := strconv.ParseFloat(s, 64)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "SENDFACT conversion error", e)
		} else {
			sendFact = float64(f)
		}
	}
	return sendFact
}

// RecvFactor returns the recv sleep factor
func RecvFactor() float64 {
	if s := os.Getenv("STOMP_RECVFACT"); s != "" {
		f, e := strconv.ParseFloat(s, 64)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "RECVFACT conversion error", e)
		} else {
			recvFact = float64(f)
		}
	}
	return recvFact
}

// Partial returns the partial byte slice for logging, random length
func Partial() []byte {
	r := int(ValueBetween(1, int64(mdml-1), 1.0))
	return md[0:r]
}

// PartialSubstr returns the partial string for logging, fixed length
func PartialSubstr(l int) []byte {
	return md[0:l]
}

// Pbc returns the byte count to log
func Pbc() int {
	if s := os.Getenv("STOMP_PBC"); s != "" {
		i, e := strconv.ParseInt(s, 10, 32)
		if nil != e {
			log.Printf("v1:%v v2:%v\n", "PBC conversion error", e)
		} else {
			pbc = int(i)
		}
	}
	return pbc
}

// Gorsleep returns an indication of whether go routines will sleep or not
func Gorsleep() string {
	gorsleep = os.Getenv("STOMP_GORSLEEP")
	return gorsleep
}

// RecvWait indicates whether to wait in receives to simulate message processing
func RecvWait() bool {
	f := os.Getenv("STOMP_RECVWAIT")
	if f == "" {
		return true
	}
	return false
}

// SendWait indicates whether to wait in sends to simulate message processing
func SendWait() bool {
	f := os.Getenv("STOMP_SENDWAIT")
	if f == "" {
		return true
	}
	return false
}

// SetMAXPROCS returns true if max procs are to be set
func SetMAXPROCS() bool {
	f := os.Getenv("STOMP_SETMAXPROCS")
	if f == "" {
		return false
	}
	return true
}

// UseCustomCiphers returns true if custom ciphers are to be used
func UseCustomCiphers() bool {
	f := os.Getenv("STOMP_USECUSTOMCIPHERS")
	if f == "" {
		return useCustomCiphers
	}
	useCustomCiphers = true
	return useCustomCiphers
}

// CustomCiphers returns a slice of custom ciphers
func CustomCiphers() []uint16 {
	if UseCustomCiphers() {
		return cipherSuites
	}
	return []uint16{}
}

// Logger returns an indication of whether to do logging
func Logger() string {
	return os.Getenv("STOMP_LOGGER")
}

// LogFile returns a log fle name or the empty string
func LogFile() string {
	return os.Getenv("STOMP_LOGFILE")
}

// UseEOF returns true if en EOF message is to be used.
func UseEOF() bool {
	if os.Getenv("STOMP_USEEOF") != "" {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package clrs

import (
	"math/rand"
	"reflect"
	"testing"
)

func TestYoungTableau(t *testing.T) {
	yt := &youngTableau{}
	yt.n = 4
	yt.m = 4
	yt.a = [][]cell{
		[]cell{{2, false}, {3, false}, {12, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {16, false}, {0, true}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	want := []cell{
		{2, false}, {3, false}, {4, false}, {5, false},
		{8, false}, {9, false}, {12, false}, {14, false},
		{16, false}, {0, true}, {0, true}, {0, true},
		{0, true}, {0, true}, {0, true}, {0, true},
	}
	index := 0
	for i := 0; i < 16; i++ {
		c, err := yt.extractMin()
		if err != nil || !reflect.DeepEqual(c, want[index]) {
			t.Errorf(" err %v", err)
			t.Errorf(" got %v", c)
			t.Errorf("want %v", want[index])
		}
		index = index + 1
	}
	// reset
	yt.a = [][]cell{
		[]cell{{2, false}, {3, false}, {12, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {16, false}, {0, true}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	want1 := [][]cell{
		[]cell{{2, false}, {3, false}, {6, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {12, false}, {16, false}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	yt.insert(cell{6, false})
	if !reflect.DeepEqual(yt.a, want1) {
		t.Errorf(" got %v", yt.a)
		t.Errorf("want %v", want1)
	}
	want2 := [][]cell{
		[]cell{{1, false}, {3, false}, {6, false}, {14, false}},
		[]cell{{2, false}, {4, false}, {8, false}, {16, false}},
		[]cell{{5, false}, {9, false}, {12, false}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	yt.insert(cell{1, false})
	if !reflect.DeepEqual(yt.a, want2) {
		t.Errorf(" got %v", yt.a)
		t.Errorf("want %v", want1)
	}
}
func TestYoungTableauSort(t *testing.T) {
	want := []int{0, 1, 2, 3, 4, 5, 6, 7, 8}
	a := rand.Perm(9)
	yt := &youngTableau{}
	yt.n = 3
	yt.m = 3
	yt.a = make([][]cell, yt.m)
	for i := 0; i < yt.m; i++ {
		yt.a[i] = make([]cell, yt.n)
		for j := 0; j < yt.n; j++ {
			yt.a[i][j] = cell{0, true}
		}
	}
	for i := 0; i < 9; i++ {
		yt.insert(cell{a[i], false})
	}
	got := make([]int, 0)
	for i := 0; i < 9; i++ {
		c, err := yt.extractMin()
		if err == nil {
			got = append(got, c.value)
		} else {
			t.Errorf("extract min error %v", err)
		}
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf(" got %v", got)
		t.Errorf("want %v", want)
	}
}
func TestYoungTableauFind(t *testing.T) {
	yt := &youngTableau{}
	yt.n = 4
	yt.m = 4
	yt.a = [][]cell{
		[]cell{{2, false}, {3, false}, {12, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {16, false}, {0, true}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}

	for _, v := range []int{2, 3, 4, 5, 8, 9, 12, 14, 16} {
		if !yt.find(v) {
			t.Errorf("find %v", v)
		}
	}
	for _, v := range []int{0, 100} {
		if yt.find(v) {
			t.Errorf("find %v", v)
		}
	}
}
<commit_msg>fix golint<commit_after>package clrs

import (
	"math/rand"
	"reflect"
	"testing"
)

func TestYoungTableau(t *testing.T) {
	yt := &youngTableau{}
	yt.n = 4
	yt.m = 4
	yt.a = [][]cell{
		[]cell{{2, false}, {3, false}, {12, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {16, false}, {0, true}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	want := []cell{
		{2, false}, {3, false}, {4, false}, {5, false},
		{8, false}, {9, false}, {12, false}, {14, false},
		{16, false}, {0, true}, {0, true}, {0, true},
		{0, true}, {0, true}, {0, true}, {0, true},
	}
	index := 0
	for i := 0; i < 16; i++ {
		c, err := yt.extractMin()
		if err != nil || !reflect.DeepEqual(c, want[index]) {
			t.Errorf(" err %v", err)
			t.Errorf(" got %v", c)
			t.Errorf("want %v", want[index])
		}
		index = index + 1
	}
	// reset
	yt.a = [][]cell{
		[]cell{{2, false}, {3, false}, {12, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {16, false}, {0, true}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	want1 := [][]cell{
		[]cell{{2, false}, {3, false}, {6, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {12, false}, {16, false}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	yt.insert(cell{6, false})
	if !reflect.DeepEqual(yt.a, want1) {
		t.Errorf(" got %v", yt.a)
		t.Errorf("want %v", want1)
	}
	want2 := [][]cell{
		[]cell{{1, false}, {3, false}, {6, false}, {14, false}},
		[]cell{{2, false}, {4, false}, {8, false}, {16, false}},
		[]cell{{5, false}, {9, false}, {12, false}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}
	yt.insert(cell{1, false})
	if !reflect.DeepEqual(yt.a, want2) {
		t.Errorf(" got %v", yt.a)
		t.Errorf("want %v", want1)
	}
}
func TestYoungTableauSort(t *testing.T) {
	want := []int{0, 1, 2, 3, 4, 5, 6, 7, 8}
	a := rand.Perm(9)
	yt := &youngTableau{}
	yt.n = 3
	yt.m = 3
	yt.a = make([][]cell, yt.m)
	for i := 0; i < yt.m; i++ {
		yt.a[i] = make([]cell, yt.n)
		for j := 0; j < yt.n; j++ {
			yt.a[i][j] = cell{0, true}
		}
	}
	for i := 0; i < 9; i++ {
		yt.insert(cell{a[i], false})
	}
	var got []int
	for i := 0; i < 9; i++ {
		c, err := yt.extractMin()
		if err == nil {
			got = append(got, c.value)
		} else {
			t.Errorf("extract min error %v", err)
		}
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf(" got %v", got)
		t.Errorf("want %v", want)
	}
}
func TestYoungTableauFind(t *testing.T) {
	yt := &youngTableau{}
	yt.n = 4
	yt.m = 4
	yt.a = [][]cell{
		[]cell{{2, false}, {3, false}, {12, false}, {14, false}},
		[]cell{{4, false}, {8, false}, {16, false}, {0, true}},
		[]cell{{5, false}, {9, false}, {0, true}, {0, true}},
		[]cell{{0, true}, {0, true}, {0, true}, {0, true}},
	}

	for _, v := range []int{2, 3, 4, 5, 8, 9, 12, 14, 16} {
		if !yt.find(v) {
			t.Errorf("find %v", v)
		}
	}
	for _, v := range []int{0, 100} {
		if yt.find(v) {
			t.Errorf("find %v", v)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

type Validate struct {
	root    string // work dir
	assets  string // submission dir
	Train   string // path of training executable
	Pred    string // path of prediction executable
	Trained string // path to trained data

	Readme string // name of the README file

	DoTraining bool

	trainfile string // path to training.csv file
	testfile  string // path to test.csv file
}

func NewValidate(dir string, train bool) (Validate, error) {
	var err error

	v := Validate{
		root:       dir,
		Trained:    "trained.dat",
		DoTraining: train,
	}

	// FIXME: handle multiple-submissions zipfiles
	//        presumably: 1 directory per submission.

	exes := make([]string, 0)
	// find executables
	err = filepath.Walk(dir, func(path string, fi os.FileInfo, err error) error {
		if fi.IsDir() {
			return nil
		}
		if strings.Contains(strings.ToLower(path), "readme") {
			v.Readme = path
		}

		if strings.Contains(strings.ToLower(path), "trained.dat") {
			v.Trained = path
		}

		// FIXME: better way ?
		if !strings.Contains(fi.Mode().String(), "x") {
			return nil
		}
		exes = append(exes, path)
		// printf(">>> %s\n", path)
		if strings.Contains(strings.ToLower(path), "higgsml-train") {
			v.Train = path
		}
		if strings.Contains(strings.ToLower(path), "higgsml-pred") {
			v.Pred = path
		}
		return err
	})

	if len(exes) <= 0 {
		return v, fmt.Errorf("hml: could not find any suitable executable in zip-file")
	}

	if v.Train == "" && v.Pred == "" {
		// take first one
		v.Train = exes[0]
		v.Pred = exes[0]
	}

	if v.Train == "" && v.Pred != "" {
		v.Train = v.Pred
	}

	if v.Train != "" && v.Pred == "" {
		v.Pred = v.Train
	}

	v.assets = filepath.Dir(v.Pred)

	return v, err
}

func (v Validate) Run() error {
	var err error

	// printf("root:   [%s]\n", v.root)
	// printf("assets: [%s]\n", v.assets)

	dir := filepath.Join(v.assets, ".higgsml-work")
	err = os.MkdirAll(dir, 0755)
	if err != nil {
		return err
	}

	if v.DoTraining {
		printf("\n")
		err = v.run_training(dir)
		if err != nil {
			return err
		}
	}

	printf("\n")
	err = v.run_pred(dir)
	if err != nil {
		return err
	}

	return err
}

func (v Validate) run_training(dir string) error {
	var err error
	errch := make(chan error)

	printf("::: run training...\n")

	cmd := exec.Command(v.Train, v.wdir("training.csv"), pdir(dir, "trained.dat"))
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Dir = v.assets

	start := time.Now()
	go func() {
		err = cmd.Start()
		if err != nil {
			errch <- err
			return
		}
		errch <- cmd.Wait()
	}()

	duration := 1 * time.Hour
	select {
	case <-time.After(duration):
		cmd.Process.Kill()
		return fmt.Errorf("hml: training timed out (%v)\n", duration)
	case err = <-errch:
		break
	}

	if err != nil {
		printf("::: run training... [ERR] (delta=%v)\n", time.Since(start))
		return err
	}

	printf("::: run training... [ok] (delta=%v)\n", time.Since(start))
	return err
}

func (v Validate) run_pred(dir string) error {
	var err error
	errch := make(chan error)

	printf("::: run prediction...\n")

	cmd := exec.Command(v.Pred, v.wdir("test.csv"), v.Trained, pdir(dir, "scores_test.csv"))
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Dir = v.assets

	start := time.Now()
	go func() {
		err = cmd.Start()
		if err != nil {
			errch <- err
			return
		}
		errch <- cmd.Wait()
	}()

	duration := 1 * time.Hour
	select {
	case <-time.After(duration):
		cmd.Process.Kill()
		return fmt.Errorf("hml: prediction timed out (%v)\n", duration)
	case err = <-errch:
		break
	}

	if err != nil {
		printf("::: run prediction... [ERR] (delta=%v)\n", time.Since(start))
		return err
	}

	printf("::: run prediction... [ok] (delta=%v)\n", time.Since(start))
	return err
}

func (v Validate) wdir(fname string) string {
	pwd, err := os.Getwd()
	if err != nil {
		return fname
	}
	return filepath.Join(pwd, fname)
}

func pdir(dirs ...string) string {
	return filepath.Join(dirs...)
}
<commit_msg>main: define higgsml-run and higgsml-train<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

const (
	runscript   = "higgsml-run"
	trainscript = "higgsml-train"
)

type Validate struct {
	root    string // work dir
	assets  string // submission dir
	Train   string // path of training executable
	Pred    string // path of prediction executable
	Trained string // path to trained data

	Readme string // name of the README file

	DoTraining bool

	trainfile string // path to training.csv file
	testfile  string // path to test.csv file
}

func NewValidate(dir string, train bool) (Validate, error) {
	var err error

	v := Validate{
		root:       dir,
		Trained:    "trained.dat",
		DoTraining: train,
	}

	// FIXME: handle multiple-submissions zipfiles
	//        presumably: 1 directory per submission.

	exes := make([]string, 0)
	// find executables
	err = filepath.Walk(dir, func(path string, fi os.FileInfo, err error) error {
		if fi.IsDir() {
			return nil
		}
		if strings.Contains(strings.ToLower(path), "readme") {
			v.Readme = path
		}

		if strings.Contains(strings.ToLower(path), "trained.dat") {
			v.Trained = path
		}

		// FIXME: better way ?
		if !strings.Contains(fi.Mode().String(), "x") {
			return nil
		}
		exes = append(exes, path)
		// printf(">>> %s\n", path)
		if strings.Contains(strings.ToLower(path), "higgsml-train") {
			v.Train = path
		}
		if strings.Contains(strings.ToLower(path), "higgsml-pred") {
			v.Pred = path
		}
		return err
	})

	if len(exes) <= 0 {
		return v, fmt.Errorf("hml: could not find any suitable executable in zip-file")
	}

	if v.Train == "" && v.Pred == "" {
		// take first one
		v.Train = exes[0]
		v.Pred = exes[0]
	}

	if v.Train == "" && v.Pred != "" {
		v.Train = v.Pred
	}

	if v.Train != "" && v.Pred == "" {
		v.Pred = v.Train
	}

	v.assets = filepath.Dir(v.Pred)

	return v, err
}

func (v Validate) Run() error {
	var err error

	// printf("root:   [%s]\n", v.root)
	// printf("assets: [%s]\n", v.assets)

	dir := filepath.Join(v.assets, ".higgsml-work")
	err = os.MkdirAll(dir, 0755)
	if err != nil {
		return err
	}

	if v.DoTraining {
		printf("\n")
		err = v.run_training(dir)
		if err != nil {
			return err
		}
	}

	printf("\n")
	err = v.run_pred(dir)
	if err != nil {
		return err
	}

	return err
}

func (v Validate) run_training(dir string) error {
	var err error
	errch := make(chan error)

	printf("::: run training...\n")

	cmd := exec.Command(v.Train, v.wdir("training.csv"), pdir(dir, "trained.dat"))
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Dir = v.assets

	start := time.Now()
	go func() {
		err = cmd.Start()
		if err != nil {
			errch <- err
			return
		}
		errch <- cmd.Wait()
	}()

	duration := 1 * time.Hour
	select {
	case <-time.After(duration):
		cmd.Process.Kill()
		return fmt.Errorf("hml: training timed out (%v)\n", duration)
	case err = <-errch:
		break
	}

	if err != nil {
		printf("::: run training... [ERR] (delta=%v)\n", time.Since(start))
		return err
	}

	printf("::: run training... [ok] (delta=%v)\n", time.Since(start))
	return err
}

func (v Validate) run_pred(dir string) error {
	var err error
	errch := make(chan error)

	printf("::: run prediction...\n")

	cmd := exec.Command(v.Pred, v.wdir("test.csv"), v.Trained, pdir(dir, "scores_test.csv"))
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Dir = v.assets

	start := time.Now()
	go func() {
		err = cmd.Start()
		if err != nil {
			errch <- err
			return
		}
		errch <- cmd.Wait()
	}()

	duration := 1 * time.Hour
	select {
	case <-time.After(duration):
		cmd.Process.Kill()
		return fmt.Errorf("hml: prediction timed out (%v)\n", duration)
	case err = <-errch:
		break
	}

	if err != nil {
		printf("::: run prediction... [ERR] (delta=%v)\n", time.Since(start))
		return err
	}

	printf("::: run prediction... [ok] (delta=%v)\n", time.Since(start))
	return err
}

func (v Validate) wdir(fname string) string {
	pwd, err := os.Getwd()
	if err != nil {
		return fname
	}
	return filepath.Join(pwd, fname)
}

func pdir(dirs ...string) string {
	return filepath.Join(dirs...)
}
<|endoftext|>


<commit_before>package sudoku

import (
	"testing"
)

//TODO: test a few more puzzles to make sure I'm exercising it correctly.

func TestSwordfishCol(t *testing.T) {

	grid := NewGrid()

	puzzleName := "swordfish_example.sdk"

	if !grid.LoadFromFile(puzzlePath(puzzleName)) {
		t.Fatal("Couldn't load puzzle ", puzzleName)
	}

	//Set up the grid correctly for the Swordfish technique to work. The
	//example we use is a grid that has other work done to exclude
	//possibilities from certain cells.

	//TODO: it's a smell that there's no way to serialize and load up a grid
	//with extra excludes set.
	excludedConfig := map[cellRef]IntSlice{
		cellRef{0, 0}: IntSlice{1, 8},
		cellRef{1, 3}: IntSlice{1},
		cellRef{1, 4}: IntSlice{1, 8},
		cellRef{2, 3}: IntSlice{1},
		cellRef{2, 5}: IntSlice{1, 8},
		cellRef{3, 0}: IntSlice{2, 8},
		cellRef{4, 0}: IntSlice{7},
		cellRef{4, 1}: IntSlice{7},
		cellRef{7, 3}: IntSlice{1, 6},
		cellRef{7, 5}: IntSlice{1},
	}

	for ref, ints := range excludedConfig {
		cell := ref.Cell(grid)
		for _, exclude := range ints {
			cell.SetExcluded(exclude, true)
		}
	}

	options := solveTechniqueTestHelperOptions{}
	options.stepsToCheck.grid = grid

	_, _, _ = humanSolveTechniqueTestHelperStepGenerator(t, "NOOP", "Swordfish", options)
}

//TODO: TestSwordfishRow (and implement Row!)
<commit_msg>Added the variants test for Swordfish.<commit_after>package sudoku

import (
	"testing"
)

//TODO: test a few more puzzles to make sure I'm exercising it correctly.

func TestSwordfishCol(t *testing.T) {

	techniqueVariantsTestHelper(t, "Swordfish")

	grid := NewGrid()

	puzzleName := "swordfish_example.sdk"

	if !grid.LoadFromFile(puzzlePath(puzzleName)) {
		t.Fatal("Couldn't load puzzle ", puzzleName)
	}

	//Set up the grid correctly for the Swordfish technique to work. The
	//example we use is a grid that has other work done to exclude
	//possibilities from certain cells.

	//TODO: it's a smell that there's no way to serialize and load up a grid
	//with extra excludes set.
	excludedConfig := map[cellRef]IntSlice{
		cellRef{0, 0}: IntSlice{1, 8},
		cellRef{1, 3}: IntSlice{1},
		cellRef{1, 4}: IntSlice{1, 8},
		cellRef{2, 3}: IntSlice{1},
		cellRef{2, 5}: IntSlice{1, 8},
		cellRef{3, 0}: IntSlice{2, 8},
		cellRef{4, 0}: IntSlice{7},
		cellRef{4, 1}: IntSlice{7},
		cellRef{7, 3}: IntSlice{1, 6},
		cellRef{7, 5}: IntSlice{1},
	}

	for ref, ints := range excludedConfig {
		cell := ref.Cell(grid)
		for _, exclude := range ints {
			cell.SetExcluded(exclude, true)
		}
	}

	options := solveTechniqueTestHelperOptions{}
	options.stepsToCheck.grid = grid

	_, _, _ = humanSolveTechniqueTestHelperStepGenerator(t, "NOOP", "Swordfish", options)
}

//TODO: TestSwordfishRow (and implement Row!)
<|endoftext|>


<commit_before>package ehttprouter

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/creack/ehttp"
	"github.com/julienschmidt/httprouter"
)

func assertInt(t *testing.T, expect, got int) {
	if expect != got {
		t.Errorf("Unexpected result.\nExpect:\t%d\nGot:\t%d\n", expect, got)
	}
}

func assertString(t *testing.T, expect, got string) {
	expect, got = strings.TrimSpace(expect), strings.TrimSpace(got)
	if expect != got {
		t.Errorf("Unexpected result.\nExpect:\t%s\nGot:\t%s\n", expect, got)
	}
}

func TestMWErrorPanicCommon(t *testing.T) {
	hdlr := func(w http.ResponseWriter, req *http.Request, p httprouter.Params) error {
		panic(fmt.Errorf("fail"))
	}
	router := httprouter.New()
	router.GET("/", MWErrorPanic(hdlr))

	ts := httptest.NewServer(router)
	defer ts.Close()

	resp, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	assertInt(t, http.StatusInternalServerError, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}
	assertString(t, "fail", string(body))
}

func TestMWErrorPanicEHTTP(t *testing.T) {
	hdlr := func(w http.ResponseWriter, req *http.Request, p httprouter.Params) error {
		panic(ehttp.NewErrorf(418, "fail"))
	}
	router := httprouter.New()
	router.GET("/", MWErrorPanic(hdlr))

	ts := httptest.NewServer(router)
	defer ts.Close()

	resp, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	assertInt(t, 418, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}
	assertString(t, "fail", string(body))
}

func TestMWErrorPanicInt(t *testing.T) {
	hdlr := func(w http.ResponseWriter, req *http.Request, p httprouter.Params) error {
		panic(418)
	}
	router := httprouter.New()
	router.GET("/", MWErrorPanic(hdlr))

	ts := httptest.NewServer(router)
	defer ts.Close()

	resp, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	assertInt(t, http.StatusInternalServerError, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}
	assertString(t, "(int) 418", string(body))
}
<commit_msg>Fix ehttprouter<commit_after>package ehttprouter

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"path"
	"runtime"
	"strings"
	"testing"

	"github.com/creack/ehttp"
	"github.com/julienschmidt/httprouter"
)

func assertInt(t *testing.T, expect, got int) {
	_, file, line := getCallstack(1)
	if expect != got {
		t.Errorf("[%s:%d] Unexpected result.\nExpect:\t%d\nGot:\t%d\n", file, line, expect, got)
	}
}

func assertString(t *testing.T, expect, got string) {
	_, file, line := getCallstack(1)
	expect, got = strings.TrimSpace(expect), strings.TrimSpace(got)
	if expect != got {
		t.Errorf("[%s:%d] Unexpected result.\nExpect:\t%s\nGot:\t%s\n", file, line, expect, got)
	}
}

func assertJSONError(t *testing.T, expect, got string) {
	_, file, line := getCallstack(1)
	expect, got = strings.TrimSpace(expect), strings.TrimSpace(got)

	jErr := ehttp.JSONError{}
	if err := json.Unmarshal([]byte(got), &jErr); err != nil {
		t.Errorf("[%s:%d] Error parsing json error: %s\n", file, line, expect)
	}
	for _, errStr := range jErr.Errors {
		if errStr == expect {
			return
		}
	}
	t.Errorf("[%s:%d] Unexpected error.\nExpect:\t%s\nGot:\t%s\n", file, line, expect, got)
}

func TestMWErrorPanicCommon(t *testing.T) {
	hdlr := func(w http.ResponseWriter, req *http.Request, p httprouter.Params) error {
		panic(fmt.Errorf("fail"))
	}
	router := httprouter.New()
	router.GET("/", MWErrorPanic(hdlr))

	ts := httptest.NewServer(router)
	defer ts.Close()

	resp, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	assertInt(t, http.StatusInternalServerError, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}
	assertJSONError(t, "fail", string(body))
}

func TestMWErrorPanicEHTTP(t *testing.T) {
	hdlr := func(w http.ResponseWriter, req *http.Request, p httprouter.Params) error {
		panic(ehttp.NewErrorf(418, "fail"))
	}
	router := httprouter.New()
	router.GET("/", MWErrorPanic(hdlr))

	ts := httptest.NewServer(router)
	defer ts.Close()

	resp, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	assertInt(t, 418, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}
	assertJSONError(t, "fail", string(body))
}

func TestMWErrorPanicInt(t *testing.T) {
	hdlr := func(w http.ResponseWriter, req *http.Request, p httprouter.Params) error {
		panic(418)
	}
	router := httprouter.New()
	router.GET("/", MWErrorPanic(hdlr))

	ts := httptest.NewServer(router)
	defer ts.Close()

	resp, err := http.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	assertInt(t, http.StatusInternalServerError, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	_ = resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}
	assertJSONError(t, "(int) 418", string(body))
}

func getCallstack(skip int) (string, string, int) {
	var name string
	pc, file, line, ok := runtime.Caller(1 + skip)
	if !ok {
		name, file, line = "<unkown>", "<unknown>", -1
	} else {
		name = runtime.FuncForPC(pc).Name()
		name = path.Base(name)
		file = path.Base(file)
	}
	return name, file, line
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	caf "github.com/cascades-fbp/cascades-caf"
	httputils "github.com/cascades-fbp/cascades-http/utils"
	"github.com/cascades-fbp/cascades/components/utils"
	"github.com/cascades-fbp/cascades/runtime"
	zmq "github.com/pebbe/zmq4"
)

var (
	intervalEndpoint = flag.String("port.int", "10s", "Component's input port endpoint")
	requestEndpoint  = flag.String("port.req", "", "Component's input port endpoint")
	templateEndpoint = flag.String("port.tmpl", "", "Component's input port endpoint")
	propertyEndpoint = flag.String("port.prop", "", "Component's output port endpoint")
	responseEndpoint = flag.String("port.resp", "", "Component's output port endpoint")
	bodyEndpoint     = flag.String("port.body", "", "Component's output port endpoint")
	errorEndpoint    = flag.String("port.err", "", "Component's error port endpoint")
	jsonFlag         = flag.Bool("json", false, "Print component documentation in JSON")
	debug            = flag.Bool("debug", false, "Enable debug mode")
)

type requestIP struct {
	URL         string              `json:"url"`
	Method      string              `json:"method"`
	ContentType string              `json:"content-type"`
	Headers     map[string][]string `json:"headers"`
}

func assertError(err error) {
	if err != nil {
		fmt.Println("ERROR:", err.Error())
		os.Exit(1)
	}
}

func main() {
	flag.Parse()

	if *jsonFlag {
		doc, _ := registryEntry.JSON()
		fmt.Println(string(doc))
		os.Exit(0)
	}

	if *requestEndpoint == "" || *templateEndpoint == "" {
		flag.Usage()
		os.Exit(1)
	}
	if *propertyEndpoint == "" && *responseEndpoint == "" && *bodyEndpoint == "" {
		flag.Usage()
		os.Exit(1)
	}

	log.SetFlags(0)
	if *debug {
		log.SetOutput(os.Stdout)
	} else {
		log.SetOutput(ioutil.Discard)
	}

	var err error

	defer zmq.Term()

	// Input sockets
	// Interval socket
	intSock, err := utils.CreateInputPort(*intervalEndpoint)
	assertError(err)
	defer intSock.Close()

	// Request socket
	reqSock, err := utils.CreateInputPort(*requestEndpoint)
	assertError(err)
	defer reqSock.Close()

	// Property template socket
	tmplSock, err := utils.CreateInputPort(*templateEndpoint)
	assertError(err)
	defer tmplSock.Close()

	// Output sockets
	// Property socket
	var propSock *zmq.Socket
	if *propertyEndpoint != "" {
		propSock, err = utils.CreateOutputPort(*propertyEndpoint)
		assertError(err)
		defer propSock.Close()
	}

	// Response socket
	var respSock *zmq.Socket
	if *responseEndpoint != "" {
		respSock, err = utils.CreateOutputPort(*responseEndpoint)
		assertError(err)
		defer respSock.Close()
	}
	// Response body socket
	var bodySock *zmq.Socket
	if *bodyEndpoint != "" {
		bodySock, err = utils.CreateOutputPort(*bodyEndpoint)
		assertError(err)
		defer bodySock.Close()
	}
	// Error socket
	var errSock *zmq.Socket
	if *errorEndpoint != "" {
		errSock, err = utils.CreateOutputPort(*errorEndpoint)
		assertError(err)
		defer errSock.Close()
	}

	// Ctrl+C handling
	utils.HandleInterruption()

	//TODO: setup input ports monitoring to close sockets when upstreams are disconnected

	// Setup socket poll items
	poller := zmq.NewPoller()
	poller.Add(intSock, zmq.POLLIN)
	poller.Add(reqSock, zmq.POLLIN)
	poller.Add(tmplSock, zmq.POLLIN)

	// This is obviously dangerous but we need it to deal with our custom CA's
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: tr}
	client.Timeout = 30 * time.Second

	var (
		interval     time.Duration
		ip           [][]byte
		request      *requestIP
		propTemplate *caf.PropertyTemplate
		httpRequest  *http.Request
	)

	for {
		sockets, err := poller.Poll(-1)
		if err != nil {
			log.Println("Error polling ports:", err.Error())
			continue
		}
		for _, socket := range sockets {
			if socket.Socket == nil {
				log.Println("ERROR: could not find socket in polling items array")
				continue
			}
			ip, err = socket.Socket.RecvMessageBytes(0)
			if err != nil {
				log.Println("Error receiving message:", err.Error())
				continue
			}
			if !runtime.IsValidIP(ip) {
				log.Println("Invalid IP:", ip)
				continue
			}
			switch socket.Socket {
			case intSock:
				interval, err = time.ParseDuration(string(ip[1]))
				log.Println("Interval specified:", interval)
			case reqSock:
				err = json.Unmarshal(ip[1], &request)
				if err != nil {
					log.Println("ERROR: failed to unmarshal request:", err.Error())
					continue
				}
				log.Println("Request specified:", request)
			case tmplSock:
				err = json.Unmarshal(ip[1], &propTemplate)
				if err != nil {
					log.Println("ERROR: failed to unmarshal template:", err.Error())
					continue
				}
				log.Println("Template specified:", propTemplate)

			default:
				log.Println("ERROR: IP from unhandled socket received!")
				continue
			}
		}
		if interval > 0 && request != nil && propTemplate != nil {
			log.Println("Component configured. Moving on...")
			break
		}
	}

	log.Println("Started...")
	ticker := time.NewTicker(interval)
	for _ = range ticker.C {
		httpRequest, err = http.NewRequest(request.Method, request.URL, nil)
		assertError(err)

		// Set the accepted Content-Type
		if request.ContentType != "" {
			httpRequest.Header.Add("Content-Type", request.ContentType)
		}

		// Set any additional headers if provided
		for k, v := range request.Headers {
			httpRequest.Header.Add(k, v[0])
		}

		response, err := client.Do(httpRequest)
		if err != nil {
			log.Printf("ERROR performing HTTP %s %s: %s\n", request.Method, request.URL, err.Error())
			if errSock != nil {
				errSock.SendMessageDontwait(runtime.NewPacket([]byte(err.Error())))
			}
			continue
		}

		resp, err := httputils.Response2Response(response)
		if err != nil {
			log.Println("ERROR converting response to reply:", err.Error())
			if errSock != nil {
				errSock.SendMessageDontwait(runtime.NewPacket([]byte(err.Error())))
			}
			continue
		}

		// Property output socket
		if propSock != nil {
			var (
				data interface{}
				buf  bytes.Buffer
				out  []byte
			)
			ts := time.Now().Unix()
			prop := &caf.Property{
				ID:        propTemplate.ID,
				Name:      propTemplate.Name,
				Group:     propTemplate.Group,
				Timestamp: &ts,
			}

			tmpl, err := template.New("value").Parse(propTemplate.Template)
			if err != nil {
				log.Println("ERROR parsing the template:", err.Error())
				continue
			}
			if strings.HasSuffix(request.ContentType, "json") {
				err = json.Unmarshal(resp.Body, &data)
				if err != nil {
					log.Println("ERROR unmarshaling the JSON response:", err.Error())
					continue
				}
			} else {
				// TODO: support other content-types
				log.Printf("WARNING processing of %s is not supported", request.ContentType)
				continue
			}

			err = tmpl.Execute(&buf, data)
			if err != nil {
				log.Println("ERROR executing the template:", err.Error())
				continue
			}

			switch propTemplate.Type {
			case caf.PropTypeString:
				v := buf.String()
				prop.StringValue = &v

			case caf.PropTypeFloat:
				v, err := strconv.ParseFloat(buf.String(), 64)
				prop.Value = &v
				if err != nil {
					log.Println("ERROR parsing float:", err.Error())
					continue
				}

			case caf.PropTypeBool:
				v, err := strconv.ParseBool(buf.String())
				prop.BoolValue = &v
				if err != nil {
					log.Println("ERROR parsing bool:", err.Error())
					continue
				}

			case caf.PropTypeJSON:
				err = json.Unmarshal(buf.Bytes(), prop.JSONValue)

				if err != nil {
					log.Println("ERROR marshaling the result in JSON:", err.Error())
					continue
				}

			default:
				log.Printf("WARNING marshaling to %s is not supported", propTemplate.Type)
				continue
			}
			out, _ = json.Marshal(prop)
			propSock.SendMessage(runtime.NewPacket(out))
		}

		// Extra output sockets (e.g., for debugging)
		if respSock != nil {
			ip, err = httputils.Response2IP(resp)
			if err != nil {
				log.Println("ERROR converting reply to IP:", err.Error())
				if errSock != nil {
					errSock.SendMessageDontwait(runtime.NewPacket([]byte(err.Error())))
				}
			} else {
				respSock.SendMessage(ip)
			}
		}
		if bodySock != nil {
			bodySock.SendMessage(runtime.NewPacket(resp.Body))
		}
	}
}
<commit_msg>http-property: error handling & text/template<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"

	caf "github.com/cascades-fbp/cascades-caf"
	httputils "github.com/cascades-fbp/cascades-http/utils"
	"github.com/cascades-fbp/cascades/components/utils"
	"github.com/cascades-fbp/cascades/runtime"
	zmq "github.com/pebbe/zmq4"
)

var (
	intervalEndpoint = flag.String("port.int", "10s", "Component's input port endpoint")
	requestEndpoint  = flag.String("port.req", "", "Component's input port endpoint")
	templateEndpoint = flag.String("port.tmpl", "", "Component's input port endpoint")
	propertyEndpoint = flag.String("port.prop", "", "Component's output port endpoint")
	responseEndpoint = flag.String("port.resp", "", "Component's output port endpoint")
	bodyEndpoint     = flag.String("port.body", "", "Component's output port endpoint")
	errorEndpoint    = flag.String("port.err", "", "Component's error port endpoint")
	jsonFlag         = flag.Bool("json", false, "Print component documentation in JSON")
	debug            = flag.Bool("debug", false, "Enable debug mode")
)

type requestIP struct {
	URL         string              `json:"url"`
	Method      string              `json:"method"`
	ContentType string              `json:"content-type"`
	Headers     map[string][]string `json:"headers"`
}

func assertError(err error) {
	if err != nil {
		fmt.Println("ERROR:", err.Error())
		os.Exit(1)
	}
}

func main() {
	flag.Parse()

	if *jsonFlag {
		doc, _ := registryEntry.JSON()
		fmt.Println(string(doc))
		os.Exit(0)
	}

	if *requestEndpoint == "" || *templateEndpoint == "" {
		flag.Usage()
		os.Exit(1)
	}
	if *propertyEndpoint == "" && *responseEndpoint == "" && *bodyEndpoint == "" {
		flag.Usage()
		os.Exit(1)
	}

	log.SetFlags(0)
	if *debug {
		log.SetOutput(os.Stdout)
	} else {
		log.SetOutput(ioutil.Discard)
	}

	var err error

	defer zmq.Term()

	// Input sockets
	// Interval socket
	intSock, err := utils.CreateInputPort(*intervalEndpoint)
	assertError(err)
	defer intSock.Close()

	// Request socket
	reqSock, err := utils.CreateInputPort(*requestEndpoint)
	assertError(err)
	defer reqSock.Close()

	// Property template socket
	tmplSock, err := utils.CreateInputPort(*templateEndpoint)
	assertError(err)
	defer tmplSock.Close()

	// Output sockets
	// Property socket
	var propSock *zmq.Socket
	if *propertyEndpoint != "" {
		propSock, err = utils.CreateOutputPort(*propertyEndpoint)
		assertError(err)
		defer propSock.Close()
	}

	// Response socket
	var respSock *zmq.Socket
	if *responseEndpoint != "" {
		respSock, err = utils.CreateOutputPort(*responseEndpoint)
		assertError(err)
		defer respSock.Close()
	}
	// Response body socket
	var bodySock *zmq.Socket
	if *bodyEndpoint != "" {
		bodySock, err = utils.CreateOutputPort(*bodyEndpoint)
		assertError(err)
		defer bodySock.Close()
	}
	// Error socket
	var errSock *zmq.Socket
	if *errorEndpoint != "" {
		errSock, err = utils.CreateOutputPort(*errorEndpoint)
		assertError(err)
		defer errSock.Close()
	}

	// Ctrl+C handling
	utils.HandleInterruption()

	//TODO: setup input ports monitoring to close sockets when upstreams are disconnected

	// Setup socket poll items
	poller := zmq.NewPoller()
	poller.Add(intSock, zmq.POLLIN)
	poller.Add(reqSock, zmq.POLLIN)
	poller.Add(tmplSock, zmq.POLLIN)

	// This is obviously dangerous but we need it to deal with our custom CA's
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: tr}
	client.Timeout = 30 * time.Second

	var (
		interval     time.Duration
		ip           [][]byte
		request      *requestIP
		propTemplate *caf.PropertyTemplate
		httpRequest  *http.Request
	)

	for {
		sockets, err := poller.Poll(-1)
		if err != nil {
			log.Println("Error polling ports:", err.Error())
			continue
		}
		for _, socket := range sockets {
			if socket.Socket == nil {
				log.Println("ERROR: could not find socket in polling items array")
				continue
			}
			ip, err = socket.Socket.RecvMessageBytes(0)
			if err != nil {
				log.Println("Error receiving message:", err.Error())
				continue
			}
			if !runtime.IsValidIP(ip) {
				log.Println("Invalid IP:", ip)
				continue
			}
			switch socket.Socket {
			case intSock:
				interval, err = time.ParseDuration(string(ip[1]))
				log.Println("Interval specified:", interval)
			case reqSock:
				err = json.Unmarshal(ip[1], &request)
				if err != nil {
					log.Println("ERROR: failed to unmarshal request:", err.Error())
					continue
				}
				log.Println("Request specified:", request)
			case tmplSock:
				err = json.Unmarshal(ip[1], &propTemplate)
				if err != nil {
					log.Println("ERROR: failed to unmarshal template:", err.Error())
					continue
				}
				log.Println("Template specified:", propTemplate)

			default:
				log.Println("ERROR: IP from unhandled socket received!")
				continue
			}
		}
		if interval > 0 && request != nil && propTemplate != nil {
			log.Println("Component configured. Moving on...")
			break
		}
	}

	log.Println("Started...")
	ticker := time.NewTicker(interval)
	for _ = range ticker.C {
		httpRequest, err = http.NewRequest(request.Method, request.URL, nil)
		assertError(err)

		// Set the accepted Content-Type
		if request.ContentType != "" {
			httpRequest.Header.Add("Content-Type", request.ContentType)
		}

		// Set any additional headers if provided
		for k, v := range request.Headers {
			httpRequest.Header.Add(k, v[0])
		}

		response, err := client.Do(httpRequest)
		if err != nil {
			log.Printf("ERROR performing HTTP %s %s: %s\n", request.Method, request.URL, err.Error())
			if errSock != nil {
				errSock.SendMessageDontwait(runtime.NewPacket([]byte(err.Error())))
			}
			continue
		}

		resp, err := httputils.Response2Response(response)
		if err != nil {
			log.Println("ERROR converting response to reply:", err.Error())
			if errSock != nil {
				errSock.SendMessageDontwait(runtime.NewPacket([]byte(err.Error())))
			}
			continue
		}

		// Property output socket
		if propSock != nil {
			var (
				data interface{}
				buf  bytes.Buffer
				out  []byte
			)
			ts := time.Now().Unix()
			prop := &caf.Property{
				ID:        propTemplate.ID,
				Name:      propTemplate.Name,
				Group:     propTemplate.Group,
				Timestamp: &ts,
			}

			tmpl, err := template.New("value").Parse(propTemplate.Template)
			if err != nil {
				log.Println("ERROR parsing the template:", err.Error())
				continue
			}
			if strings.HasSuffix(request.ContentType, "json") {
				err = json.Unmarshal(resp.Body, &data)
				if err != nil {
					log.Println("ERROR unmarshaling the JSON response:", err.Error())
					continue
				}
			} else {
				// TODO: support other content-types
				log.Printf("WARNING processing of %s is not supported", request.ContentType)
				continue
			}

			err = tmpl.Execute(&buf, data)
			if err != nil {
				log.Println("ERROR executing the template:", err.Error())
				continue
			}

			switch propTemplate.Type {
			case caf.PropTypeString:
				v := buf.String()
				prop.StringValue = &v

			case caf.PropTypeFloat:
				v, err := strconv.ParseFloat(buf.String(), 64)
				if err != nil {
					log.Println("ERROR parsing float:", err.Error())
					continue
				}
				prop.Value = &v

			case caf.PropTypeBool:
				v, err := strconv.ParseBool(buf.String())
				if err != nil {
					log.Println("ERROR parsing bool:", err.Error())
					continue
				}
				prop.BoolValue = &v

			case caf.PropTypeJSON:
				err = json.Unmarshal(buf.Bytes(), prop.JSONValue)

				if err != nil {
					log.Println("ERROR marshaling the result in JSON:", err.Error())
					continue
				}

			default:
				log.Printf("WARNING marshaling to %s is not supported", propTemplate.Type)
				continue
			}
			out, _ = json.Marshal(prop)
			propSock.SendMessage(runtime.NewPacket(out))
		}

		// Extra output sockets (e.g., for debugging)
		if respSock != nil {
			ip, err = httputils.Response2IP(resp)
			if err != nil {
				log.Println("ERROR converting reply to IP:", err.Error())
				if errSock != nil {
					errSock.SendMessageDontwait(runtime.NewPacket([]byte(err.Error())))
				}
			} else {
				respSock.SendMessage(ip)
			}
		}
		if bodySock != nil {
			bodySock.SendMessage(runtime.NewPacket(resp.Body))
		}
	}
}
<|endoftext|>


<commit_before>/*
   Copyright 2013 Juliano Martinez <juliano@martinez.io>

   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.

      Based on http://github.com/nf/webfront

   @author: Juliano Martinez
*/

package http_server

import (
	"github.com/fiorix/go-redis/redis"
	hpr_utils "github.com/ncode/hot-potato-router/utils"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"
	"sync"
	"time"
)

var (
	cfg = hpr_utils.NewConfig()
	rc  = redis.New(cfg.Options["redis"]["server_list"])
)

func xff(req *http.Request) string {
	remote_addr := strings.Split(req.RemoteAddr, ":")
	if len(remote_addr) == 0 {
		return ""
	}
	return remote_addr[0]
}

type Server struct {
	mu      sync.RWMutex
	last    time.Time
	proxy   map[string][]Proxy
	backend map[string]int
}

type Proxy struct {
	Connections int64
	Backend     string
	handler     http.Handler
}

func Listen(fd int, addr string) net.Listener {
	var l net.Listener
	var err error
	if fd >= 3 {
		l, err = net.FileListener(os.NewFile(uintptr(fd), "http"))
	} else {
		l, err = net.Listen("tcp", addr)
	}
	if err != nil {
		log.Fatal(err)
	}
	return l
}

func NewServer(probe time.Duration) (*Server, error) {
	s := new(Server)
	s.proxy = make(map[string][]Proxy)
	go s.probe_backends(probe)
	return s, nil
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if h := s.handler(r); h != nil {
		r.Header.Add("X-Forwarded-For", xff(r))
		r.Header.Add("X-Real-IP", xff(r))
		h.ServeHTTP(w, r)
		return
	}
	http.Error(w, "Not found.", http.StatusNotFound)
}

func (s *Server) handler(req *http.Request) http.Handler {
	s.mu.RLock()
	defer s.mu.RUnlock()
	h := req.Host
	if i := strings.Index(h, ":"); i >= 0 {
		h = h[:i]
	}

	_, ok := s.proxy[h]
	if !ok {
		f, _ := rc.Get(h)
		if f == "" {
			return nil
		}
		s.proxy[h] = append(s.proxy[h], Proxy{0, f, makeHandler(f)})
	}
	return s.Next(h)
}

/* TODO: Implement at least one decent loadbalance algorithm
Using n+1 to see it working*/
func (s *Server) Next(h string) http.Handler {
	s.mu.Lock()
	defer s.mu.Unlock()
	total := len(s.proxy[h])
	if s.backend[h] == total {
		s.backend[h] = 0
	}
	s.backend[h]++
	return s.proxy[h][s.backend[h]].handler
}

func (s *Server) probe_backends(probe time.Duration) {
	for {
		s.mu.Lock()
		s.mu.Unlock()
		time.Sleep(probe)
	}
}

func makeHandler(f string) http.Handler {
	if f != "" {
		return &httputil.ReverseProxy{
			Director: func(req *http.Request) {
				req.URL.Scheme = "http"
				req.URL.Host = f
			},
		}
	}
	return nil
}
<commit_msg>update comment<commit_after>/*
   Copyright 2013 Juliano Martinez <juliano@martinez.io>

   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.

      Based on http://github.com/nf/webfront

   @author: Juliano Martinez
*/

package http_server

import (
	"github.com/fiorix/go-redis/redis"
	hpr_utils "github.com/ncode/hot-potato-router/utils"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"
	"sync"
	"time"
)

var (
	cfg = hpr_utils.NewConfig()
	rc  = redis.New(cfg.Options["redis"]["server_list"])
)

func xff(req *http.Request) string {
	remote_addr := strings.Split(req.RemoteAddr, ":")
	if len(remote_addr) == 0 {
		return ""
	}
	return remote_addr[0]
}

type Server struct {
	mu      sync.RWMutex
	last    time.Time
	proxy   map[string][]Proxy
	backend map[string]int
}

type Proxy struct {
	Connections int64
	Backend     string
	handler     http.Handler
}

func Listen(fd int, addr string) net.Listener {
	var l net.Listener
	var err error
	if fd >= 3 {
		l, err = net.FileListener(os.NewFile(uintptr(fd), "http"))
	} else {
		l, err = net.Listen("tcp", addr)
	}
	if err != nil {
		log.Fatal(err)
	}
	return l
}

func NewServer(probe time.Duration) (*Server, error) {
	s := new(Server)
	s.proxy = make(map[string][]Proxy)
	go s.probe_backends(probe)
	return s, nil
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if h := s.handler(r); h != nil {
		r.Header.Add("X-Forwarded-For", xff(r))
		r.Header.Add("X-Real-IP", xff(r))
		h.ServeHTTP(w, r)
		return
	}
	http.Error(w, "Not found.", http.StatusNotFound)
}

func (s *Server) handler(req *http.Request) http.Handler {
	s.mu.RLock()
	defer s.mu.RUnlock()
	h := req.Host
	if i := strings.Index(h, ":"); i >= 0 {
		h = h[:i]
	}

	_, ok := s.proxy[h]
	if !ok {
		f, _ := rc.Get(h)
		if f == "" {
			return nil
		}
		s.proxy[h] = append(s.proxy[h], Proxy{0, f, makeHandler(f)})
	}
	return s.Next(h)
}

/* TODO: Implement more balance algorithms */
func (s *Server) Next(h string) http.Handler {
	s.mu.Lock()
	defer s.mu.Unlock()
	total := len(s.proxy[h])
	if s.backend[h] == total {
		s.backend[h] = 0
	}
	s.backend[h]++
	return s.proxy[h][s.backend[h]].handler
}

func (s *Server) probe_backends(probe time.Duration) {
	for {
		s.mu.Lock()
		s.mu.Unlock()
		time.Sleep(probe)
	}
}

func makeHandler(f string) http.Handler {
	if f != "" {
		return &httputil.ReverseProxy{
			Director: func(req *http.Request) {
				req.URL.Scheme = "http"
				req.URL.Host = f
			},
		}
	}
	return nil
}
<|endoftext|>