<commit_before>package web

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/eyecuelab/kit/maputil"
	"github.com/google/jsonapi"
	"github.com/labstack/echo"
	"github.com/eyecuelab/kit/flect"
	"errors"
)

var reNotJsonApi = regexp.MustCompile("not a jsonapi|EOF")

func notJsonApi(err error) bool {
	return reNotJsonApi.MatchString(err.Error())
}

type (
	ApiContext interface {
		echo.Context

		Payload() *jsonapi.OnePayload
		Attrs(permitted ...string) map[string]interface{}
		AttrKeys() []string
		RequireAttrs(...string) error
		BindAndValidate(interface{}) error
		BindMulti(interface{}) ([]interface{}, error)
		BindIdParam(*int, ...string) error
		JsonApi(interface{}, int) error
		JsonApiOK(interface{}, ...interface{}) error
		ApiError(string, ...int) *echo.HTTPError
		JsonAPIError(string, int, string) *jsonapi.ErrorObject
		QueryParamTrue(string) (bool, bool)

		RequiredQueryParams(...string) (map[string]string, error)
		OptionalQueryParams(...string) map[string]string
		QParams(...string) (map[string]string, error)
	}

	apiContext struct {
		echo.Context

		payload     *jsonapi.OnePayload
		manyPayload *jsonapi.ManyPayload
	}

	CommonExtendable interface {
		CommonExtend(interface{}) error
	}

	Extendable interface {
		Extend(interface{}) error
	}

	CommonMetable interface {
		CommonMeta() error
	}

	Metable interface {
		Meta() error
	}

	CommonLinkable interface {
		CommonLinks() error
	}

	Linkable interface {
		Links() error
	}
)

func (c *apiContext) Payload() *jsonapi.OnePayload {
	return c.payload
}

func (c *apiContext) Attrs(permitted ...string) map[string]interface{} {
	//TODO: remove this once all refactoring is complete
	if len(permitted) == 0 {
		return c.payload.Data.Attributes
	}

	permittedAttrs := make(map[string]interface{})
	for _, p := range permitted {
		if val, ok := c.payload.Data.Attributes[p]; ok {
			permittedAttrs[p] = val
		}
	}
	return permittedAttrs
}

func (c *apiContext) AttrKeys() []string {
	return maputil.Keys(c.Attrs())
}

func (c *apiContext) RequireAttrs(required ...string) error {
	missing := make([]string, 0, len(required))

	for _, key := range required {
		if c.payload.Data.Attributes[key] == nil {
			missing = append(missing, key)
			continue
		}
	}

	if len(missing) > 0 {
		return fmt.Errorf("missing required attributes: %v", missing)
	}

	return nil
}

//Before binding we make a copy of the req body and restore it after binding.
//This allows the body to be used again later
func (c *apiContext) Bind(i interface{}) error {
	body, err := c.readRestoreBody()
	if err != nil {
		return err
	}

	ctype := c.Request().Header.Get(echo.HeaderContentType)
	if isJSONAPI(ctype) {
		err = jsonAPIBind(c, i)
	} else {
		err = c.defaultBind(i)
	}

	c.restoreBody(body)

	return err
}

func (c *apiContext) BindMulti(containedType interface{}) ([]interface{}, error) {
	body, err := c.readRestoreBody()
	if err != nil {
		return nil, err
	}

	ctype := c.Request().Header.Get(echo.HeaderContentType)

	if !isJSONAPI(ctype) {
		return nil, errors.New("BindMulti only supports JSONApi, use Bind")
	}

	i, err := jsonAPIBindMulti(c, containedType)

	c.restoreBody(body)

	return i, err
}

func (c *apiContext) readRestoreBody() ([]byte, error) {
	b, err := ioutil.ReadAll(c.Request().Body)
	c.restoreBody(b)
	return b, err
}

func (c *apiContext) restoreBody(b []byte) {
	c.Request().Body = ioutil.NopCloser(bytes.NewBuffer(b))
}

func (c *apiContext) defaultBind(i interface{}) error {
	db := new(echo.DefaultBinder)
	return db.Bind(i, c)
}

func isJSONAPI(s string) bool {
	const MIMEJsonAPI = "application/vnd.api+json"
	return strings.HasPrefix(s, MIMEJsonAPI)
}

func (c *apiContext) BindAndValidate(i interface{}) error {
	if err := c.Bind(i); err != nil {
		return err
	}
	if err := c.Validate(i); err != nil {
		return err
	}
	return nil
}

func (c *apiContext) JsonApi(i interface{}, status int) error {
	c.Response().Header().Set(echo.HeaderContentType, jsonapi.MediaType)
	c.Response().WriteHeader(status)

	return jsonapi.MarshalPayload(c.Response().Writer, i)
}

func applyCommon(i interface{}, extendData interface{}) error {
	if casted, ok := i.(CommonExtendable); ok {
		if err := casted.CommonExtend(extendData); err != nil {
			return err
		}
	}

	if casted, ok := i.(CommonMetable); ok {
		if err := casted.CommonMeta(); err != nil {
			return err
		}
	}

	if casted, ok := i.(CommonLinkable); ok {
		if err := casted.CommonLinks(); err != nil {
			return err
		}
	}
	return nil
}

func apply(i interface{}, extendData interface{}) error {
	if casted, ok := i.(Extendable); ok {
		if err := casted.Extend(extendData); err != nil {
			return err
		}
	}

	if casted, ok := i.(Metable); ok {
		if err := casted.Meta(); err != nil {
			return err
		}
	}

	if casted, ok := i.(Linkable); ok {
		if err := casted.Links(); err != nil {
			return err
		}
	}
	return nil
}

func extendAndExtract(i interface{}, extendData interface{}) (data interface{}, err error) {
	if flect.IsSlice(i) {
		slice := reflect.ValueOf(i)
		for idx := 0; idx < slice.Len(); idx++ {
			elementInterface := slice.Index(idx).Interface()
			if err := applyCommon(elementInterface, extendData); err != nil {
				return nil, err
			}
		}
		return i, nil
	}

	if err := applyCommon(i, extendData); err != nil {
		return nil, err
	}

	if err := apply(i, extendData); err != nil {
		return nil, err
	}
	return i,nil
}

func (c *apiContext) JsonApiOK(i interface{}, extendData ...interface{}) error {
	var ed interface{}
	if len(extendData) > 0 {
		ed = extendData[0]
	}
	data, err := extendAndExtract(i, ed)
	if err != nil {
		return err
	}
	return c.JsonApi(data, http.StatusOK)
}

func (c *apiContext) BindIdParam(idValue *int, named ...string) (err error) {
	paramName := "id"
	if len(named) > 0 {
		paramName = named[0]
	}
	*idValue, err = strconv.Atoi(c.Param(paramName))
	return err
}

func (c *apiContext) QueryParamTrue(name string) (val, ok bool) {
	switch strings.ToLower(c.QueryParam(name)) {
	case "true", "1":
		return true, true
	case "false", "0":
		return false, true
	default:
		return false, false
	}
}

func jsonAPIBindMulti(c *apiContext, elementType interface{}) ([]interface{}, error) {
	buf := new(bytes.Buffer)
	tee := io.TeeReader(c.Request().Body, buf)

	unmarshaled, err := jsonapi.UnmarshalManyPayload(tee, reflect.TypeOf(elementType))
	if err != nil {
		return nil, err
	}

	c.manyPayload = new(jsonapi.ManyPayload)
	return unmarshaled, json.Unmarshal(buf.Bytes(), c.manyPayload)
}

func jsonAPIBind(c *apiContext, i interface{}) error {
	buf := new(bytes.Buffer)
	tee := io.TeeReader(c.Request().Body, buf)

	rType := reflect.TypeOf(i)

	if rType.Kind() == reflect.Slice {
		value := reflect.TypeOf(rType.Elem())

		unmarshaled, err := jsonapi.UnmarshalManyPayload(tee, value)
		if err != nil {
			return err
		}
		i = unmarshaled
	} else {
		if err := jsonapi.UnmarshalPayload(tee, i); err != nil {
			if notJsonApi(err) {
				return c.ApiError("Request Body is not valid JsonAPI")
			}
			return err
		}
	}

	c.payload = new(jsonapi.OnePayload)
	return json.Unmarshal(buf.Bytes(), c.payload)
}

func (c *apiContext) ApiError(msg string, codes ...int) *echo.HTTPError {
	if len(codes) > 0 {
		return echo.NewHTTPError(codes[0], msg)
	}
	// TODO: return jsonapi error instead
	return echo.NewHTTPError(http.StatusBadRequest, msg)
}

func (c *apiContext) JsonAPIError(msg string, code int, param string) *jsonapi.ErrorObject {
	return &jsonapi.ErrorObject{
		Status: fmt.Sprintf("%d", code),
		Title:  http.StatusText(code),
		Detail: msg,
		Meta: &map[string]interface{}{
			"parameter": param,
		},
	}
}

func (c *apiContext) RequiredQueryParams(required ...string) (map[string]string, error) {
	missing := make([]string, 0, len(required))
	params := make(map[string]string)

	for _, key := range required {
		val := c.QueryParam(key)
		if val == "" {
			missing = append(missing, key)
			continue
		}
		params[key] = val
	}

	if len(missing) > 0 {
		return nil, fmt.Errorf("missing required params: %v", missing)
	}

	return params, nil
}

func (c *apiContext) QParams(required ...string) (map[string]string, error) {
	return QParams(c, required...)
}

func QParams(c echo.Context, required ...string) (map[string]string, error) {
	missing := make([]string, 0, len(required))
	params := make(map[string]string)

	for k := range c.QueryParams() {
		params[k] = c.QueryParam(k)
	}

	for _, k := range required {
		if _, ok := params[k]; !ok {
			missing = append(missing, k)
		}
	}

	if len(missing) > 0 {
		return nil, fmt.Errorf("missing required params: %v", missing)
	}

	return params, nil
}

func (c *apiContext) OptionalQueryParams(optional ...string) map[string]string {
	params := make(map[string]string)
	for _, key := range optional {
		val := c.QueryParam(key)
		params[key] = val
	}
	return params
}

func ApiContextMiddleWare() func(echo.HandlerFunc) echo.HandlerFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			return next(&apiContext{c, nil, nil})
		}
	}
}

func restrictedValue(value string, allowed []string, errorText string) (string, error) {
	if contains(allowed, value) {
		return value, nil
	}
	return "", fmt.Errorf(errorText, value)
}

func contains(set []string, s string) bool {
	for _, v := range set {
		if s == v {
			return true
		}
	}
	return false
}
<commit_msg>Allow passing pagination details into link data<commit_after>package web

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/eyecuelab/kit/maputil"
	"github.com/google/jsonapi"
	"github.com/labstack/echo"
	"github.com/eyecuelab/kit/flect"
	"errors"
	"github.com/eyecuelab/kit/web/meta"
)

var reNotJsonApi = regexp.MustCompile("not a jsonapi|EOF")

func notJsonApi(err error) bool {
	return reNotJsonApi.MatchString(err.Error())
}

type (
	ApiContext interface {
		echo.Context

		Payload() *jsonapi.OnePayload
		Attrs(permitted ...string) map[string]interface{}
		AttrKeys() []string
		RequireAttrs(...string) error
		BindAndValidate(interface{}) error
		BindMulti(interface{}) ([]interface{}, error)
		BindIdParam(*int, ...string) error
		JsonApi(interface{}, int) error
		JsonApiOK(interface{}, ...interface{}) error
		JsonApiOKPaged(interface{}, meta.Pagination, ...interface{}) error
		ApiError(string, ...int) *echo.HTTPError
		JsonAPIError(string, int, string) *jsonapi.ErrorObject
		QueryParamTrue(string) (bool, bool)

		RequiredQueryParams(...string) (map[string]string, error)
		OptionalQueryParams(...string) map[string]string
		QParams(...string) (map[string]string, error)
	}

	apiContext struct {
		echo.Context

		payload     *jsonapi.OnePayload
		manyPayload *jsonapi.ManyPayload
	}

	CommonExtendable interface {
		CommonExtend(interface{}) error
	}

	Extendable interface {
		Extend(interface{}) error
	}

	CommonMetable interface {
		CommonMeta() error
	}

	Metable interface {
		Meta() error
	}

	CommonLinkable interface {
		CommonLinks(*meta.Pagination) error
	}

	Linkable interface {
		Links(*meta.Pagination) error
	}
)

func (c *apiContext) Payload() *jsonapi.OnePayload {
	return c.payload
}

func (c *apiContext) Attrs(permitted ...string) map[string]interface{} {
	//TODO: remove this once all refactoring is complete
	if len(permitted) == 0 {
		return c.payload.Data.Attributes
	}

	permittedAttrs := make(map[string]interface{})
	for _, p := range permitted {
		if val, ok := c.payload.Data.Attributes[p]; ok {
			permittedAttrs[p] = val
		}
	}
	return permittedAttrs
}

func (c *apiContext) AttrKeys() []string {
	return maputil.Keys(c.Attrs())
}

func (c *apiContext) RequireAttrs(required ...string) error {
	missing := make([]string, 0, len(required))

	for _, key := range required {
		if c.payload.Data.Attributes[key] == nil {
			missing = append(missing, key)
			continue
		}
	}

	if len(missing) > 0 {
		return fmt.Errorf("missing required attributes: %v", missing)
	}

	return nil
}

//Before binding we make a copy of the req body and restore it after binding.
//This allows the body to be used again later
func (c *apiContext) Bind(i interface{}) error {
	body, err := c.readRestoreBody()
	if err != nil {
		return err
	}

	ctype := c.Request().Header.Get(echo.HeaderContentType)
	if isJSONAPI(ctype) {
		err = jsonAPIBind(c, i)
	} else {
		err = c.defaultBind(i)
	}

	c.restoreBody(body)

	return err
}

func (c *apiContext) BindMulti(containedType interface{}) ([]interface{}, error) {
	body, err := c.readRestoreBody()
	if err != nil {
		return nil, err
	}

	ctype := c.Request().Header.Get(echo.HeaderContentType)

	if !isJSONAPI(ctype) {
		return nil, errors.New("BindMulti only supports JSONApi, use Bind")
	}

	i, err := jsonAPIBindMulti(c, containedType)

	c.restoreBody(body)

	return i, err
}

func (c *apiContext) readRestoreBody() ([]byte, error) {
	b, err := ioutil.ReadAll(c.Request().Body)
	c.restoreBody(b)
	return b, err
}

func (c *apiContext) restoreBody(b []byte) {
	c.Request().Body = ioutil.NopCloser(bytes.NewBuffer(b))
}

func (c *apiContext) defaultBind(i interface{}) error {
	db := new(echo.DefaultBinder)
	return db.Bind(i, c)
}

func isJSONAPI(s string) bool {
	const MIMEJsonAPI = "application/vnd.api+json"
	return strings.HasPrefix(s, MIMEJsonAPI)
}

func (c *apiContext) BindAndValidate(i interface{}) error {
	if err := c.Bind(i); err != nil {
		return err
	}
	if err := c.Validate(i); err != nil {
		return err
	}
	return nil
}

func (c *apiContext) JsonApi(i interface{}, status int) error {
	c.Response().Header().Set(echo.HeaderContentType, jsonapi.MediaType)
	c.Response().WriteHeader(status)

	return jsonapi.MarshalPayload(c.Response().Writer, i)
}

func applyCommon(i interface{}, page *meta.Pagination,  extendData ...interface{}) error {
	if casted, ok := i.(CommonExtendable); ok {
		if err := casted.CommonExtend(extendData); err != nil {
			return err
		}
	}

	if casted, ok := i.(CommonMetable); ok {
		if err := casted.CommonMeta(); err != nil {
			return err
		}
	}

	if casted, ok := i.(CommonLinkable); ok {
		if err := casted.CommonLinks(page); err != nil {
			return err
		}
	}
	return nil
}

func apply(i interface{}, page *meta.Pagination, extendData interface{}) error {
	if casted, ok := i.(Extendable); ok {
		if err := casted.Extend(extendData); err != nil {
			return err
		}
	}

	if casted, ok := i.(Metable); ok {
		if err := casted.Meta(); err != nil {
			return err
		}
	}

	if casted, ok := i.(Linkable); ok {
		if err := casted.Links(page); err != nil {
			return err
		}
	}
	return nil
}

func extendAndExtract(i interface{}, page *meta.Pagination, extendData interface{}) (data interface{}, err error) {
	if flect.IsSlice(i) {
		slice := reflect.ValueOf(i)
		for idx := 0; idx < slice.Len(); idx++ {
			elementInterface := slice.Index(idx).Interface()
			if err := applyCommon(elementInterface, page, extendData); err != nil {
				return nil, err
			}
		}
		return i, nil
	}

	if err := applyCommon(i, page, extendData); err != nil {
		return nil, err
	}

	if err := apply(i, page, extendData); err != nil {
		return nil, err
	}
	return i,nil
}

func (c *apiContext) JsonApiOK(i interface{}, extendData ...interface{}) error {
	var ed interface{}
	if len(extendData) > 0 {
		ed = extendData[0]
	}
	data, err := extendAndExtract(i, nil, ed)
	if err != nil {
		return err
	}
	return c.JsonApi(data, http.StatusOK)
}

func (c *apiContext) JsonApiOKPaged(i interface{}, page *meta.Pagination, extendData ...interface{}) error {
	var ed interface{}
	if len(extendData) > 0 {
		ed = extendData[0]
	}
	data, err := extendAndExtract(i, page, ed)
	if err != nil {
		return err
	}
	return c.JsonApi(data, http.StatusOK)
}

func (c *apiContext) BindIdParam(idValue *int, named ...string) (err error) {
	paramName := "id"
	if len(named) > 0 {
		paramName = named[0]
	}
	*idValue, err = strconv.Atoi(c.Param(paramName))
	return err
}

func (c *apiContext) QueryParamTrue(name string) (val, ok bool) {
	switch strings.ToLower(c.QueryParam(name)) {
	case "true", "1":
		return true, true
	case "false", "0":
		return false, true
	default:
		return false, false
	}
}

func jsonAPIBindMulti(c *apiContext, elementType interface{}) ([]interface{}, error) {
	buf := new(bytes.Buffer)
	tee := io.TeeReader(c.Request().Body, buf)

	unmarshaled, err := jsonapi.UnmarshalManyPayload(tee, reflect.TypeOf(elementType))
	if err != nil {
		return nil, err
	}

	c.manyPayload = new(jsonapi.ManyPayload)
	return unmarshaled, json.Unmarshal(buf.Bytes(), c.manyPayload)
}

func jsonAPIBind(c *apiContext, i interface{}) error {
	buf := new(bytes.Buffer)
	tee := io.TeeReader(c.Request().Body, buf)

	rType := reflect.TypeOf(i)

	if rType.Kind() == reflect.Slice {
		value := reflect.TypeOf(rType.Elem())

		unmarshaled, err := jsonapi.UnmarshalManyPayload(tee, value)
		if err != nil {
			return err
		}
		i = unmarshaled
	} else {
		if err := jsonapi.UnmarshalPayload(tee, i); err != nil {
			if notJsonApi(err) {
				return c.ApiError("Request Body is not valid JsonAPI")
			}
			return err
		}
	}

	c.payload = new(jsonapi.OnePayload)
	return json.Unmarshal(buf.Bytes(), c.payload)
}

func (c *apiContext) ApiError(msg string, codes ...int) *echo.HTTPError {
	if len(codes) > 0 {
		return echo.NewHTTPError(codes[0], msg)
	}
	// TODO: return jsonapi error instead
	return echo.NewHTTPError(http.StatusBadRequest, msg)
}

func (c *apiContext) JsonAPIError(msg string, code int, param string) *jsonapi.ErrorObject {
	return &jsonapi.ErrorObject{
		Status: fmt.Sprintf("%d", code),
		Title:  http.StatusText(code),
		Detail: msg,
		Meta: &map[string]interface{}{
			"parameter": param,
		},
	}
}

func (c *apiContext) RequiredQueryParams(required ...string) (map[string]string, error) {
	missing := make([]string, 0, len(required))
	params := make(map[string]string)

	for _, key := range required {
		val := c.QueryParam(key)
		if val == "" {
			missing = append(missing, key)
			continue
		}
		params[key] = val
	}

	if len(missing) > 0 {
		return nil, fmt.Errorf("missing required params: %v", missing)
	}

	return params, nil
}

func (c *apiContext) QParams(required ...string) (map[string]string, error) {
	return QParams(c, required...)
}

func QParams(c echo.Context, required ...string) (map[string]string, error) {
	missing := make([]string, 0, len(required))
	params := make(map[string]string)

	for k := range c.QueryParams() {
		params[k] = c.QueryParam(k)
	}

	for _, k := range required {
		if _, ok := params[k]; !ok {
			missing = append(missing, k)
		}
	}

	if len(missing) > 0 {
		return nil, fmt.Errorf("missing required params: %v", missing)
	}

	return params, nil
}

func (c *apiContext) OptionalQueryParams(optional ...string) map[string]string {
	params := make(map[string]string)
	for _, key := range optional {
		val := c.QueryParam(key)
		params[key] = val
	}
	return params
}

func ApiContextMiddleWare() func(echo.HandlerFunc) echo.HandlerFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			return next(&apiContext{c, nil, nil})
		}
	}
}

func restrictedValue(value string, allowed []string, errorText string) (string, error) {
	if contains(allowed, value) {
		return value, nil
	}
	return "", fmt.Errorf(errorText, value)
}

func contains(set []string, s string) bool {
	for _, v := range set {
		if s == v {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package web

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"github.com/evepraisal/go-evepraisal"

	"golang.org/x/oauth2"
)

func (ctx *Context) GetCurrentUser(r *http.Request) *evepraisal.User {
	user, ok := ctx.getSessionValue(r, "user").(evepraisal.User)
	if !ok {
		return nil
	}
	return &user
}

func (ctx *Context) HandleLogin(w http.ResponseWriter, r *http.Request) {
	if ctx.OauthConfig == nil {
		ctx.renderErrorPage(r, w, http.StatusBadRequest, "Feature Unavailable", "SSO is not configured")
		return
	}
	url := ctx.OauthConfig.AuthCodeURL("state", oauth2.AccessTypeOnline)
	http.Redirect(w, r, url, http.StatusTemporaryRedirect)
}

func (ctx *Context) HandleLogout(w http.ResponseWriter, r *http.Request) {
	if ctx.OauthConfig == nil {
		ctx.renderErrorPage(r, w, http.StatusBadRequest, "Feature Unavailable", "SSO is not configured")
		return
	}
	ctx.setSessionValue(r, w, "user", nil)
	http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
}

func (ctx *Context) HandleAuthCallback(w http.ResponseWriter, r *http.Request) {
	if ctx.OauthConfig == nil {
		ctx.renderErrorPage(r, w, http.StatusBadRequest, "Feature Unavailable", "SSO is not configured")
		return
	}
	code := r.FormValue("code")
	tok, err := ctx.OauthConfig.Exchange(r.Context(), code)
	if err != nil {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", err.Error())
	}

	client := ctx.OauthConfig.Client(r.Context(), tok)
	resp, err := client.Get(ctx.OauthVerifyURL)
	if err != nil {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", err.Error())
	}

	if resp.StatusCode != 200 {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", fmt.Sprint("Invalid response when verifying identity (%d)", resp.StatusCode))
	}

	user := &evepraisal.User{}
	err = json.NewDecoder(resp.Body).Decode(user)
	defer resp.Body.Close()
	if err != nil {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", err.Error())
	}

	ctx.setSessionValue(r, w, "user", &user)
	log.Printf("User logged in: %s", user.CharacterName)

	http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
}
<commit_msg>Fix small lint problem<commit_after>package web

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"github.com/evepraisal/go-evepraisal"

	"golang.org/x/oauth2"
)

func (ctx *Context) GetCurrentUser(r *http.Request) *evepraisal.User {
	user, ok := ctx.getSessionValue(r, "user").(evepraisal.User)
	if !ok {
		return nil
	}
	return &user
}

func (ctx *Context) HandleLogin(w http.ResponseWriter, r *http.Request) {
	if ctx.OauthConfig == nil {
		ctx.renderErrorPage(r, w, http.StatusBadRequest, "Feature Unavailable", "SSO is not configured")
		return
	}
	url := ctx.OauthConfig.AuthCodeURL("state", oauth2.AccessTypeOnline)
	http.Redirect(w, r, url, http.StatusTemporaryRedirect)
}

func (ctx *Context) HandleLogout(w http.ResponseWriter, r *http.Request) {
	if ctx.OauthConfig == nil {
		ctx.renderErrorPage(r, w, http.StatusBadRequest, "Feature Unavailable", "SSO is not configured")
		return
	}
	ctx.setSessionValue(r, w, "user", nil)
	http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
}

func (ctx *Context) HandleAuthCallback(w http.ResponseWriter, r *http.Request) {
	if ctx.OauthConfig == nil {
		ctx.renderErrorPage(r, w, http.StatusBadRequest, "Feature Unavailable", "SSO is not configured")
		return
	}
	code := r.FormValue("code")
	tok, err := ctx.OauthConfig.Exchange(r.Context(), code)
	if err != nil {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", err.Error())
	}

	client := ctx.OauthConfig.Client(r.Context(), tok)
	resp, err := client.Get(ctx.OauthVerifyURL)
	if err != nil {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", err.Error())
	}

	if resp.StatusCode != 200 {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", fmt.Sprintf("Invalid response when verifying identity (%d)", resp.StatusCode))
	}

	user := &evepraisal.User{}
	err = json.NewDecoder(resp.Body).Decode(user)
	defer resp.Body.Close()
	if err != nil {
		ctx.renderErrorPage(r, w, http.StatusInternalServerError, "Something bad happened", err.Error())
	}

	ctx.setSessionValue(r, w, "user", &user)
	log.Printf("User logged in: %s", user.CharacterName)

	http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
}
<|endoftext|>


<commit_before>// bulk_load_mongo loads a Mongo daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"gopkg.in/mgo.v2"

	"github.com/influxdata/influxdb-comparisons/mongo_serialization"
)

// Program option vars:
var (
	daemonUrl    string
	workers      int
	batchSize    int
	doLoad       bool
	writeTimeout time.Duration
)

// Global vars
var (
	batchChan    chan *Batch
	inputDone    chan struct{}
	workersGroup sync.WaitGroup
)

// bufPool holds []byte instances to reduce heap churn.
var bufPool = &sync.Pool{
	New: func() interface{} {
		return make([]byte, 0, 256)
	},
}

// Batch holds byte slices that will become mongo_serialization.Item instances.
type Batch [][]byte

func (b *Batch) ClearReferences() {
	*b = (*b)[:0]
}

// bufPool holds *Batch instances to reduce heap churn.
var batchPool = &sync.Pool{
	New: func() interface{} {
		return &Batch{}
	},
}

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:27017", "Mongo URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.DurationVar(&writeTimeout, "write-timeout", 10*time.Second, "Write timeout.")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.Parse()
}

func main() {
	if doLoad {
		//createKeyspace(daemonUrl)
	}

	var session *mgo.Session

	if doLoad {
		var err error
		session, err = mgo.Dial(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}
		defer session.Close()
	}

	batchChan = make(chan *Batch, workers*100)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processItems(session)
	}

	start := time.Now()
	itemsRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)
	rate := float64(itemsRead) / float64(took.Seconds())

	fmt.Printf("loaded %d items in %fsec with %d workers (mean rate %f/sec)\n", itemsRead, took.Seconds(), workers, rate)
}

// scan reads length-delimited flatbuffers items from stdin.
func scan(session *mgo.Session, itemsPerBatch int) int64 {
	//var batch *gocql.Batch
	if doLoad {
		//batch = session.NewBatch(gocql.LoggedBatch)
	}

	var n int
	var itemsRead int64
	r := bufio.NewReaderSize(os.Stdin, 32<<20)

	batch := batchPool.Get().(*Batch)
	lenBuf := make([]byte, 8)

	for {
		// get the serialized item length (this is the framing format)
		_, err := r.Read(lenBuf)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err.Error())
		}

		// ensure correct len of receiving buffer
		l := int(binary.LittleEndian.Uint64(lenBuf))
		itemBuf := bufPool.Get().([]byte)
		if cap(itemBuf) < l {
			itemBuf = make([]byte, l)
		}
		itemBuf = itemBuf[:l]

		// read the bytes and init the flatbuffer
		totRead := 0
		for totRead < l {
			m, err := r.Read(itemBuf[totRead:])
			// (EOF is also fatal)
			if err != nil {
				log.Fatal(err.Error())
			}
			totRead += m
		}
		if totRead != len(itemBuf) {
			panic(fmt.Sprintf("reader/writer logic error, %d != %d", n, len(itemBuf)))
		}

		*batch = append(*batch, itemBuf)

		itemsRead++
		n++

		if n >= batchSize {
			batchChan <- batch
			n = 0
			batch = batchPool.Get().(*Batch)
		}
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead
}

// processItems reads byte buffers from itemChan, interprets them and writes
// them to the target server. Note that mgo incurs a lot of overhead.
func processItems(session *mgo.Session) {
	db := session.DB("benchmark_db")
	item := &mongo_serialization.Item{}

	type pointValue struct {
		timestamp int64
		value     interface{}
	}

	pv := &pointValue{}
	for batch := range batchChan {
		for _, itemBuf := range *batch {
			// this ui could be improved on the library side:
			n := flatbuffers.GetUOffsetT(itemBuf)
			item.Init(itemBuf, n)

			switch item.ValueType() {
			case mongo_serialization.ValueTypeLong:
				pv.value = item.LongValue()
			case mongo_serialization.ValueTypeDouble:
				pv.value = item.DoubleValue()
			default:
				panic("logic error")
			}

			if doLoad {
				collection := db.C(unsafeBytesToString(item.SeriesIdBytes()))
				err := collection.Insert(pv)
				if err != nil {
					log.Fatalf("Error writing: %s\n", err.Error())
				}
			}
		}

		for _, itemBuf := range *batch {
			bufPool.Put(itemBuf)
		}
		batch.ClearReferences()
		batchPool.Put(batch)
	}
	workersGroup.Done()
}
<commit_msg>submit timestamp to mongo<commit_after>// bulk_load_mongo loads a Mongo daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"gopkg.in/mgo.v2"

	"github.com/influxdata/influxdb-comparisons/mongo_serialization"
)

// Program option vars:
var (
	daemonUrl    string
	workers      int
	batchSize    int
	doLoad       bool
	writeTimeout time.Duration
)

// Global vars
var (
	batchChan    chan *Batch
	inputDone    chan struct{}
	workersGroup sync.WaitGroup
)

// bufPool holds []byte instances to reduce heap churn.
var bufPool = &sync.Pool{
	New: func() interface{} {
		return make([]byte, 0, 256)
	},
}

// Batch holds byte slices that will become mongo_serialization.Item instances.
type Batch [][]byte

func (b *Batch) ClearReferences() {
	*b = (*b)[:0]
}

// bufPool holds *Batch instances to reduce heap churn.
var batchPool = &sync.Pool{
	New: func() interface{} {
		return &Batch{}
	},
}

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:27017", "Mongo URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.DurationVar(&writeTimeout, "write-timeout", 10*time.Second, "Write timeout.")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.Parse()
}

func main() {
	if doLoad {
		//createKeyspace(daemonUrl)
	}

	var session *mgo.Session

	if doLoad {
		var err error
		session, err = mgo.Dial(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}
		defer session.Close()
	}

	batchChan = make(chan *Batch, workers*100)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processItems(session)
	}

	start := time.Now()
	itemsRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)
	rate := float64(itemsRead) / float64(took.Seconds())

	fmt.Printf("loaded %d items in %fsec with %d workers (mean rate %f/sec)\n", itemsRead, took.Seconds(), workers, rate)
}

// scan reads length-delimited flatbuffers items from stdin.
func scan(session *mgo.Session, itemsPerBatch int) int64 {
	//var batch *gocql.Batch
	if doLoad {
		//batch = session.NewBatch(gocql.LoggedBatch)
	}

	var n int
	var itemsRead int64
	r := bufio.NewReaderSize(os.Stdin, 32<<20)

	batch := batchPool.Get().(*Batch)
	lenBuf := make([]byte, 8)

	for {
		// get the serialized item length (this is the framing format)
		_, err := r.Read(lenBuf)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err.Error())
		}

		// ensure correct len of receiving buffer
		l := int(binary.LittleEndian.Uint64(lenBuf))
		itemBuf := bufPool.Get().([]byte)
		if cap(itemBuf) < l {
			itemBuf = make([]byte, l)
		}
		itemBuf = itemBuf[:l]

		// read the bytes and init the flatbuffer
		totRead := 0
		for totRead < l {
			m, err := r.Read(itemBuf[totRead:])
			// (EOF is also fatal)
			if err != nil {
				log.Fatal(err.Error())
			}
			totRead += m
		}
		if totRead != len(itemBuf) {
			panic(fmt.Sprintf("reader/writer logic error, %d != %d", n, len(itemBuf)))
		}

		*batch = append(*batch, itemBuf)

		itemsRead++
		n++

		if n >= batchSize {
			batchChan <- batch
			n = 0
			batch = batchPool.Get().(*Batch)
		}
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead
}

// processItems reads byte buffers from itemChan, interprets them and writes
// them to the target server. Note that mgo incurs a lot of overhead.
func processItems(session *mgo.Session) {
	db := session.DB("benchmark_db")
	item := &mongo_serialization.Item{}

	type pointValue struct {
		timestamp int64
		value     interface{}
	}

	pv := &pointValue{}
	for batch := range batchChan {
		for _, itemBuf := range *batch {
			// this ui could be improved on the library side:
			n := flatbuffers.GetUOffsetT(itemBuf)
			item.Init(itemBuf, n)

			pv.timestamp = item.TimestampNanos()

			switch item.ValueType() {
			case mongo_serialization.ValueTypeLong:
				pv.value = item.LongValue()
			case mongo_serialization.ValueTypeDouble:
				pv.value = item.DoubleValue()
			default:
				panic("logic error")
			}

			if doLoad {
				collection := db.C(unsafeBytesToString(item.SeriesIdBytes()))
				err := collection.Insert(pv)
				if err != nil {
					log.Fatalf("Error writing: %s\n", err.Error())
				}
			}
		}

		for _, itemBuf := range *batch {
			bufPool.Put(itemBuf)
		}
		batch.ClearReferences()
		batchPool.Put(batch)
	}
	workersGroup.Done()
}
<|endoftext|>


<commit_before>package workers

import (
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"time"
)

type APTRestoreSpotTest struct {
	Context          *context.Context
	CreatedBefore    time.Time
	NotRestoredSince time.Time
	MaxSize          int64
}

// NewAPTRestoreSpotTest creates a new restore spot test worker.
// This is meant to run as a cron job.
//
// Param maxSize tells the worker to choose bags no larger than
// maxSize for restoration. Param createdBefore means choose bags
// created before this date. Param notRestoredSince means choose
// bags that have not been restored since this date (which helps
// prevent us restoring the same bag again and again).
func NewAPTRestoreSpotTest(_context *context.Context, maxSize int64, createdBefore, notRestoredSince time.Time) *APTRestoreSpotTest {
	return &APTRestoreSpotTest{
		Context:          _context,
		CreatedBefore:    createdBefore,
		NotRestoredSince: notRestoredSince,
		MaxSize:          maxSize,
	}
}

// Run runs the spot test by choosing ONE bag from each institution
// that matches the specified criteria (smaller than maxSize, createdBefore
// the specified date and notRestoredSince the specified date).
// It creates a Restore WorkItem for each bag, and returns the WorkItems
// it created. The caller can get the WorkItem.Id and object identifier from there.
func (restoreTest *APTRestoreSpotTest) Run() []*models.WorkItem {

	return nil
}
<commit_msg>Working on restore spot test<commit_after>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/models"
	"net/url"
	"strconv"
	"time"
)

type APTRestoreSpotTest struct {
	Context          *context.Context
	CreatedBefore    time.Time
	NotRestoredSince time.Time
	MaxSize          int64
}

// NewAPTRestoreSpotTest creates a new restore spot test worker.
// This is meant to run as a cron job.
//
// Param maxSize tells the worker to choose bags no larger than
// maxSize for restoration. Param createdBefore means choose bags
// created before this date. Param notRestoredSince means choose
// bags that have not been restored since this date (which helps
// prevent us restoring the same bag again and again).
func NewAPTRestoreSpotTest(_context *context.Context, maxSize int64, createdBefore, notRestoredSince time.Time) *APTRestoreSpotTest {
	return &APTRestoreSpotTest{
		Context:          _context,
		CreatedBefore:    createdBefore,
		NotRestoredSince: notRestoredSince,
		MaxSize:          maxSize,
	}
}

// Run runs the spot test by choosing ONE bag from each institution
// that matches the specified criteria (smaller than maxSize, createdBefore
// the specified date and notRestoredSince the specified date).
// It creates a Restore WorkItem for each bag, and returns the WorkItems
// it created. The caller can get the WorkItem.Id and object identifier from there.
func (restoreTest *APTRestoreSpotTest) Run() ([]*models.WorkItem, error) {
	workItems := make([]*models.WorkItem, 0)
	institutions, err := restoreTest.GetInstitutions()
	if err != nil {
		return workItems, fmt.Errorf("Error getting list of institutions from Pharos: %v", err)
	}
	for _, inst := range institutions {
		obj, err := restoreTest.GetObjectFor(inst.Identifier)
		if err != nil {
			restoreTest.Context.MessageLog.Error(err.Error())
			continue
		}
		if obj == nil {
			restoreTest.Context.MessageLog.Info("No suitable objects for %s", inst.Identifier)
			continue
		}
		workItem, err := restoreTest.CreateWorkItem(obj)
		if err != nil {
			restoreTest.Context.MessageLog.Error("Error creating new Restore WorkItem for %s: %v",
				obj.Identifier, err)
			continue
		}
		workItems = append(workItems, workItem)
	}

	return workItems, nil
}

// GetInstitutions returns a list of all depositing institutions from Pharos.
func (restoreTest *APTRestoreSpotTest) GetInstitutions() ([]*models.Institution, error) {
	resp := restoreTest.Context.PharosClient.InstitutionList(url.Values{})
	if resp.Error != nil {
		return nil, resp.Error
	}
	institutions := resp.Institutions()
	restoreTest.Context.MessageLog.Info("Got %d institutions from Pharos", len(institutions))
	return institutions, nil
}

// GetObjectFor returns the IntellectualObject we should restore for the
// specified institution. This object will be the first one we find in
// Pharos that matches all of the criteria, which include:
//
// institution - it belongs the specified institution
//
// createdBefore - it was created before the specified date
//
// notRestoredSince - it has not been restored since the specified date
//
// maxSize - it's total size is less than or equal to this many bytes
//
// state - is "A" for active
//
// access - is not "restricted"
func (restoreTest *APTRestoreSpotTest) GetObjectFor(institution string) (*models.IntellectualObject, error) {
	var obj *models.IntellectualObject
	var err error
	hasMoreResults := true
	pageNumber := 1
	for hasMoreResults {
		obj, hasMoreResults, err = restoreTest.findOne(institution, pageNumber)
		if obj != nil || err != nil {
			break
		}
		pageNumber++
	}
	return obj, err
}

// findOne finds one IntellectualObject from the specified page of results
// that meets our criteria for restoration spot tests.
//
// Returns an IntellectualObject (or nil if none match our criteria), a boolean
// indicating whether Pharos has more results to fetch, and an error
// if there is one.
func (restoreTest *APTRestoreSpotTest) findOne(institution string, pageNumber int) (*models.IntellectualObject, bool, error) {
	var selectedObject *models.IntellectualObject
	params := url.Values{}
	params.Set("institution", institution)
	params.Set("state", "A")
	params.Set("createdBefore", restoreTest.CreatedBefore.Format(time.RFC3339))
	params.Set("page", strconv.Itoa(pageNumber))
	params.Set("per_page", "100")
	resp := restoreTest.Context.PharosClient.IntellectualObjectList(params)
	if resp.Error != nil {
		return nil, false, resp.Error
	}
	objects := resp.IntellectualObjects()
	restoreTest.Context.MessageLog.Info("Found %d object candidates for %s", len(objects), institution)
	for _, obj := range objects {
		if obj.Access == "restricted" {
			restoreTest.Context.MessageLog.Info("Skipping %s: restricted", obj.Identifier)
			continue
		}
		if obj.FileSize > restoreTest.MaxSize {
			restoreTest.Context.MessageLog.Info("Skipping %s: size %d is greater than max %d",
				obj.Identifier, obj.FileSize, restoreTest.MaxSize)
			continue
		}
		hasCompletedRestore, err := restoreTest.HasCompletedRestore(obj.Identifier)
		if err != nil {
			restoreTest.Context.MessageLog.Warning("Error checking restore WorkItem for %s: %v",
				obj.Identifier, err)
			continue
		}
		if !hasCompletedRestore {
			selectedObject = obj
			break
		}
	}
	return selectedObject, resp.HasNextPage(), nil
}

// HasCompletedRestore returns true if the object with the specified identifier
// has been successfully restored since NotRestoredSince.
func (restoreTest *APTRestoreSpotTest) HasCompletedRestore(objIdentifier string) (bool, error) {
	params := url.Values{}
	params.Set("object_identifier", objIdentifier)
	params.Set("action", constants.ActionRestore)
	params.Set("status", constants.StatusSuccess)
	params.Set("updated_after", restoreTest.NotRestoredSince.Format(time.RFC3339))
	params.Set("page", "1")
	params.Set("per_page", "1")
	resp := restoreTest.Context.PharosClient.WorkItemList(params)
	if resp.Error != nil {
		return false, resp.Error
	}
	hasRestore := resp.WorkItem() != nil
	return hasRestore, nil
}

func (restoreTest *APTRestoreSpotTest) GetLastIngestWorkItem(objIdentifier string) (*models.WorkItem, error) {
	params := url.Values{}
	params.Set("object_identifier", objIdentifier)
	params.Set("action", constants.ActionIngest)
	params.Set("status", constants.StatusSuccess)
	params.Set("sort", "date") // Sorts by date_processed desc
	params.Set("page", "1")
	params.Set("per_page", "1")
	resp := restoreTest.Context.PharosClient.WorkItemList(params)
	if resp.Error != nil {
		return nil, resp.Error
	}
	return resp.WorkItem(), nil
}

// CreateWorkItem creates the Restore WorkItem for the specified object identifier.
func (restoreTest *APTRestoreSpotTest) CreateWorkItem(obj *models.IntellectualObject) (*models.WorkItem, error) {
	lastIngestItem, err := restoreTest.GetLastIngestWorkItem(obj.Identifier)
	if err != nil {
		return nil, fmt.Errorf("Cannot find last ingest WorkItem for %s: %v", obj.Identifier, err)
	}
	if lastIngestItem == nil {
		return nil, fmt.Errorf("Last ingest WorkItem is missing for %s", obj.Identifier)
	}
	workItem := &models.WorkItem{
		ObjectIdentifier: obj.Identifier,
		Name:             lastIngestItem.Name,
		Bucket:           lastIngestItem.Bucket,
		ETag:             lastIngestItem.ETag,
		Size:             obj.FileSize, // may have changed with file deletions or overwrites
		BagDate:          lastIngestItem.BagDate,
		InstitutionId:    lastIngestItem.InstitutionId,
		User:             "system@aptrust.org", // because this is an automated spot test
		Action:           constants.ActionRestore,
		Stage:            constants.StageRequested,
		Status:           constants.StatusPending,
		Outcome:          "Pending",
		Retry:            true,
		Note:             "Automated object restoration spot test created by system",
	}
	resp := restoreTest.Context.PharosClient.WorkItemSave(workItem)
	if resp.Error != nil {
		return nil, resp.Error
	}
	return resp.WorkItem(), nil
}
<|endoftext|>


<commit_before><commit_msg>Sqrt(x)<commit_after><|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/pachyderm/pachyderm/src/client"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pfs/fuse"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmd"
	pfssync "github.com/pachyderm/pachyderm/src/server/pkg/sync"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"

	"github.com/spf13/cobra"
	"go.pedge.io/env"
	"go.pedge.io/lion"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

const (
	// PFSInputPrefix is where the input repos reside
	PFSInputPrefix = "/pfs"
	// PFSOutputPrefix is where the output data resides
	PFSOutputPrefix = "/pfs/out"
	// FUSEMountPoint is where we mount FUSE
	FUSEMountPoint = "/pfs/fuse"
)

type appEnv struct {
	PachydermAddress string `env:"PACHD_PORT_650_TCP_ADDR,required"`
	PodName          string `env:"PPS_POD_NAME,required"`
}

func main() {
	env.Main(do, &appEnv{})
}

func downloadInput(c *client.APIClient, commitMounts []*fuse.CommitMount) error {
	var g errgroup.Group
	for _, commitMount := range commitMounts {
		commitMount := commitMount
		if commitMount.Alias == "prev" || commitMount.Alias == "out" {
			continue
		}
		g.Go(func() error {
			return pfssync.Pull(context.Background(), c.PfsAPIClient, filepath.Join(PFSInputPrefix, commitMount.Commit.Repo.Name), commitMount.Commit, commitMount.DiffMethod, commitMount.Shard)
		})
	}
	return g.Wait()
}

func uploadOutput(c *client.APIClient, out *fuse.CommitMount, overwrite bool) error {
	return pfssync.Push(context.Background(), c.PfsAPIClient, PFSOutputPrefix, out.Commit, overwrite)
}

func do(appEnvObj interface{}) error {
	appEnv := appEnvObj.(*appEnv)
	rootCmd := &cobra.Command{
		Use:   os.Args[0] + " job-id",
		Short: `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Long:  `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Run: cmd.RunFixedArgs(1, func(args []string) (retErr error) {
			defer func() {
				// Kubernetes' jobs have the semantic that a pod is always restarted
				// until it exits successfully.  Therefore, we always return a non-zero
				// exit code to ensure that there's always N pods running, where N is
				// the user-specified parallelism.
				//
				// A pod only ever truly terminates when the job is finished.  It learns
				// that through ContinuePod.
				if retErr == nil {
					os.Exit(1)
				}
			}()
			ppsClient, err := ppsserver.NewInternalPodAPIClientFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}
			job := &ppsclient.Job{
				ID: args[0],
			}

			var chunkID string
			var chunkIDLock sync.Mutex

			// Start sending ContinuePod so PPS knows that we are alive.
			// We want to start sending ContinuePod before we send StartPod,
			// because StartPod can block if the job has been finished.  Therefore
			// we need to send ContinuePod to find out if the job has been finished.
			go func() {
				tick := time.Tick(10 * time.Second)
				for {
					<-tick
					chunkIDLock.Lock()
					res, err := ppsClient.ContinuePod(
						context.Background(),
						&ppsserver.ContinuePodRequest{
							Job:     job,
							ChunkID: chunkID,
							PodName: appEnv.PodName,
						},
					)
					chunkIDLock.Unlock()
					if err != nil {
						lion.Errorf("error from ContinuePod: %s", err.Error())
					}
					fmt.Printf("ContinuePod response: %v\n", res)
					switch res.Exit {
					case ppsserver.ExitMode_EM_TERMINATE:
						os.Exit(0)
					case ppsserver.ExitMode_EM_RESTART:
						os.Exit(1)
					case ppsserver.ExitMode_EM_NONE:
						continue
					}
				}
			}()

			response, err := ppsClient.StartPod(
				context.Background(),
				&ppsserver.StartPodRequest{
					Job:     job,
					PodName: appEnv.PodName,
				})
			if err != nil {
				lion.Errorf("error from StartPod: %s", err.Error())
				return err
			}
			chunkIDLock.Lock()
			chunkID = response.ChunkID
			chunkIDLock.Unlock()

			if response.Transform.Debug {
				lion.SetLevel(lion.LevelDebug)
			}
			// We want to make sure that we only send FinishPod once.
			// The most bulletproof way would be to check that on server side,
			// but this is easier.
			var finished bool
			// Make sure that we call FinishPod even if something caused a panic
			defer func() {
				if r := recover(); r != nil && !finished {
					fmt.Println("job shim crashed; this is like a bug in pachyderm")
					if _, err := ppsClient.FinishPod(
						context.Background(),
						&ppsserver.FinishPodRequest{
							ChunkID: response.ChunkID,
							PodName: appEnv.PodName,
							Success: false,
						},
					); err != nil && retErr == nil {
						retErr = err
					}
				}
			}()

			c, err := client.NewFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}

			if err := downloadInput(c, response.CommitMounts); err != nil {
				return err
			}

			// Create /pfs/out
			if err := os.MkdirAll(PFSOutputPrefix, 0777); err != nil {
				return err
			}

			// mount /pfs/fuse
			// right now we only have /pfs/fuse/prev
			var mounts []*fuse.CommitMount
			for _, m := range response.CommitMounts {
				if m.Alias == "prev" {
					mounts = append(mounts, m)
				}
			}
			mounter := fuse.NewMounter(appEnv.PachydermAddress, c)
			ready := make(chan bool)
			errCh := make(chan error)
			go func() {
				if err := mounter.MountAndCreate(
					FUSEMountPoint,
					nil,
					mounts,
					ready,
					response.Transform.Debug,
					false,
				); err != nil {
					errCh <- err
				}
			}()
			select {
			case <-ready:
			case err := <-errCh:
				return err
			}
			defer func() {
				if err := mounter.Unmount(FUSEMountPoint); err != nil && retErr == nil {
					retErr = err
				}
			}()

			var readers []io.Reader
			for _, line := range response.Transform.Stdin {
				readers = append(readers, strings.NewReader(line+"\n"))
			}
			if len(response.Transform.Cmd) == 0 {
				fmt.Println("unable to run; a cmd needs to be provided")
				if _, err := ppsClient.FinishPod(
					context.Background(),
					&ppsserver.FinishPodRequest{
						ChunkID: response.ChunkID,
						PodName: appEnv.PodName,
						Success: false,
					},
				); err != nil {
					return err
				}
				finished = true
				return
			}

			cmdCh := make(chan bool)
			go func() {
				cmd := exec.Command(response.Transform.Cmd[0], response.Transform.Cmd[1:]...)
				cmd.Stdin = io.MultiReader(readers...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				success := true
				if err := cmd.Run(); err != nil {
					success = false
					if exiterr, ok := err.(*exec.ExitError); ok {
						if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
							for _, returnCode := range response.Transform.AcceptReturnCode {
								if int(returnCode) == status.ExitStatus() {
									success = true
								}
							}
						}
					}
					if !success {
						fmt.Fprintf(os.Stderr, "Error from exec: %s\n", err.Error())
					}
				}
				cmdCh <- success
			}()

			success := <-cmdCh
			var outputMount *fuse.CommitMount
			for _, c := range response.CommitMounts {
				if c.Alias == "out" {
					outputMount = c
					break
				}
			}
			if err := uploadOutput(c, outputMount, response.Transform.Overwrite); err != nil {
				fmt.Printf("err from uploading output: %s\n", err)
				success = false
			}

			res, err := ppsClient.FinishPod(
				context.Background(),
				&ppsserver.FinishPodRequest{
					ChunkID: response.ChunkID,
					PodName: appEnv.PodName,
					Success: success,
				},
			)
			if err != nil {
				return err
			}
			finished = true
			if res.Fail {
				return errors.New("restarting")
			}
			return nil
		}),
	}

	return rootCmd.Execute()
}
<commit_msg>Revert "Launch ContinueJob before StartJob"<commit_after>package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/pachyderm/pachyderm/src/client"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pfs/fuse"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmd"
	"github.com/pachyderm/pachyderm/src/server/pkg/sync"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"

	"github.com/spf13/cobra"
	"go.pedge.io/env"
	"go.pedge.io/lion"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

const (
	// PFSInputPrefix is where the input repos reside
	PFSInputPrefix = "/pfs"
	// PFSOutputPrefix is where the output data resides
	PFSOutputPrefix = "/pfs/out"
	// FUSEMountPoint is where we mount FUSE
	FUSEMountPoint = "/pfs/fuse"
)

type appEnv struct {
	PachydermAddress string `env:"PACHD_PORT_650_TCP_ADDR,required"`
	PodName          string `env:"PPS_POD_NAME,required"`
}

func main() {
	env.Main(do, &appEnv{})
}

func downloadInput(c *client.APIClient, commitMounts []*fuse.CommitMount) error {
	var g errgroup.Group
	for _, commitMount := range commitMounts {
		commitMount := commitMount
		if commitMount.Alias == "prev" || commitMount.Alias == "out" {
			continue
		}
		g.Go(func() error {
			return sync.Pull(context.Background(), c.PfsAPIClient, filepath.Join(PFSInputPrefix, commitMount.Commit.Repo.Name), commitMount.Commit, commitMount.DiffMethod, commitMount.Shard)
		})
	}
	return g.Wait()
}

func uploadOutput(c *client.APIClient, out *fuse.CommitMount, overwrite bool) error {
	return sync.Push(context.Background(), c.PfsAPIClient, PFSOutputPrefix, out.Commit, overwrite)
}

func do(appEnvObj interface{}) error {
	appEnv := appEnvObj.(*appEnv)
	rootCmd := &cobra.Command{
		Use:   os.Args[0] + " job-id",
		Short: `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Long:  `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Run: cmd.RunFixedArgs(1, func(args []string) (retErr error) {
			defer func() {
				// Kubernetes' jobs have the semantic that a pod is always restarted
				// until it exits successfully.  Therefore, we always return a non-zero
				// exit code to ensure that there's always N pods running, where N is
				// the user-specified parallelism.
				//
				// A pod only ever truly terminates when the job is finished.  It learns
				// that through ContinuePod.
				if retErr == nil {
					os.Exit(1)
				}
			}()
			ppsClient, err := ppsserver.NewInternalPodAPIClientFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}
			job := &ppsclient.Job{
				ID: args[0],
			}
			response, err := ppsClient.StartPod(
				context.Background(),
				&ppsserver.StartPodRequest{
					Job:     job,
					PodName: appEnv.PodName,
				})
			if err != nil {
				lion.Errorf("error from StartPod: %s", err.Error())
				return err
			}

			// Start sending ContinuePod so PPS knows that we are alive
			go func() {
				tick := time.Tick(10 * time.Second)
				for {
					<-tick
					res, err := ppsClient.ContinuePod(
						context.Background(),
						&ppsserver.ContinuePodRequest{
							Job:     job,
							ChunkID: response.ChunkID,
							PodName: appEnv.PodName,
						},
					)
					if err != nil {
						lion.Errorf("error from ContinuePod: %s", err.Error())
					}
					switch res.Exit {
					case ppsserver.ExitMode_EM_TERMINATE:
						os.Exit(0)
					case ppsserver.ExitMode_EM_RESTART:
						os.Exit(1)
					case ppsserver.ExitMode_EM_NONE:
						continue
					}
				}
			}()

			if response.Transform.Debug {
				lion.SetLevel(lion.LevelDebug)
			}
			// We want to make sure that we only send FinishPod once.
			// The most bulletproof way would be to check that on server side,
			// but this is easier.
			var finished bool
			// Make sure that we call FinishPod even if something caused a panic
			defer func() {
				if r := recover(); r != nil && !finished {
					fmt.Println("job shim crashed; this is like a bug in pachyderm")
					if _, err := ppsClient.FinishPod(
						context.Background(),
						&ppsserver.FinishPodRequest{
							ChunkID: response.ChunkID,
							PodName: appEnv.PodName,
							Success: false,
						},
					); err != nil && retErr == nil {
						retErr = err
					}
				}
			}()

			c, err := client.NewFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}

			if err := downloadInput(c, response.CommitMounts); err != nil {
				return err
			}

			// Create /pfs/out
			if err := os.MkdirAll(PFSOutputPrefix, 0777); err != nil {
				return err
			}

			// mount /pfs/fuse
			// right now we only have /pfs/fuse/prev
			var mounts []*fuse.CommitMount
			for _, m := range response.CommitMounts {
				if m.Alias == "prev" {
					mounts = append(mounts, m)
				}
			}
			mounter := fuse.NewMounter(appEnv.PachydermAddress, c)
			ready := make(chan bool)
			errCh := make(chan error)
			go func() {
				if err := mounter.MountAndCreate(
					FUSEMountPoint,
					nil,
					mounts,
					ready,
					response.Transform.Debug,
					false,
				); err != nil {
					errCh <- err
				}
			}()
			select {
			case <-ready:
			case err := <-errCh:
				return err
			}
			defer func() {
				if err := mounter.Unmount(FUSEMountPoint); err != nil && retErr == nil {
					retErr = err
				}
			}()

			var readers []io.Reader
			for _, line := range response.Transform.Stdin {
				readers = append(readers, strings.NewReader(line+"\n"))
			}
			if len(response.Transform.Cmd) == 0 {
				fmt.Println("unable to run; a cmd needs to be provided")
				if _, err := ppsClient.FinishPod(
					context.Background(),
					&ppsserver.FinishPodRequest{
						ChunkID: response.ChunkID,
						PodName: appEnv.PodName,
						Success: false,
					},
				); err != nil {
					return err
				}
				finished = true
				return
			}

			cmdCh := make(chan bool)
			go func() {
				cmd := exec.Command(response.Transform.Cmd[0], response.Transform.Cmd[1:]...)
				cmd.Stdin = io.MultiReader(readers...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				success := true
				if err := cmd.Run(); err != nil {
					success = false
					if exiterr, ok := err.(*exec.ExitError); ok {
						if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
							for _, returnCode := range response.Transform.AcceptReturnCode {
								if int(returnCode) == status.ExitStatus() {
									success = true
								}
							}
						}
					}
					if !success {
						fmt.Fprintf(os.Stderr, "Error from exec: %s\n", err.Error())
					}
				}
				cmdCh <- success
			}()

			success := <-cmdCh
			var outputMount *fuse.CommitMount
			for _, c := range response.CommitMounts {
				if c.Alias == "out" {
					outputMount = c
					break
				}
			}
			if err := uploadOutput(c, outputMount, response.Transform.Overwrite); err != nil {
				fmt.Printf("err from uploading output: %s\n", err)
				success = false
			}

			res, err := ppsClient.FinishPod(
				context.Background(),
				&ppsserver.FinishPodRequest{
					ChunkID: response.ChunkID,
					PodName: appEnv.PodName,
					Success: success,
				},
			)
			if err != nil {
				return err
			}
			finished = true
			if res.Fail {
				return errors.New("restarting")
			}
			return nil
		}),
	}

	return rootCmd.Execute()
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/pachyderm/pachyderm/src/client"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pfs/fuse"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmd"
	"github.com/pachyderm/pachyderm/src/server/pkg/sync"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"

	"github.com/spf13/cobra"
	"go.pedge.io/env"
	"go.pedge.io/lion"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

const (
	// PFSInputPrefix is where the input repos reside
	PFSInputPrefix = "/pfs"
	// PFSOutputPrefix is where the output data resides
	PFSOutputPrefix = "/pfs/out"
	// FUSEMountPoint is where we mount FUSE
	FUSEMountPoint = "/pfs/prev"
)

type appEnv struct {
	PachydermAddress string `env:"PACHD_PORT_650_TCP_ADDR,required"`
	PodName          string `env:"PPS_POD_NAME,required"`
}

func main() {
	env.Main(do, &appEnv{})
}

func downloadInput(c *client.APIClient, commitMounts []*fuse.CommitMount) error {
	var g errgroup.Group
	for _, commitMount := range commitMounts {
		commitMount := commitMount
		if commitMount.Alias == "prev" || commitMount.Alias == "out" {
			continue
		}
		g.Go(func() error {
			return sync.Pull(context.Background(), c.PfsAPIClient, filepath.Join(PFSInputPrefix, commitMount.Commit.Repo.Name),
				commitMount.Commit, commitMount.DiffMethod, commitMount.Shard, commitMount.Lazy)
		})
	}
	return g.Wait()
}

func uploadOutput(c *client.APIClient, out *fuse.CommitMount, overwrite bool) error {
	return sync.Push(context.Background(), c.PfsAPIClient, PFSOutputPrefix, out.Commit, overwrite)
}

func do(appEnvObj interface{}) error {
	appEnv := appEnvObj.(*appEnv)
	rootCmd := &cobra.Command{
		Use:   os.Args[0] + " job-id",
		Short: `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Long:  `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Run: cmd.RunFixedArgs(1, func(args []string) (retErr error) {
			ppsClient, err := ppsserver.NewInternalPodAPIClientFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}
			response, err := ppsClient.StartPod(
				context.Background(),
				&ppsserver.StartPodRequest{
					Job: &ppsclient.Job{
						ID: args[0],
					},
					PodName: appEnv.PodName,
				})
			if err != nil {
				lion.Errorf("error from StartPod: %s", err.Error())
				return err
			}

			// Start sending ContinuePod to PPS to signal that we are alive
			exitCh := make(chan struct{})
			go func() {
				tick := time.Tick(10 * time.Second)
				for {
					<-tick
					res, err := ppsClient.ContinuePod(
						context.Background(),
						&ppsserver.ContinuePodRequest{
							ChunkID: response.ChunkID,
							PodName: appEnv.PodName,
						},
					)
					if err != nil {
						lion.Errorf("error from ContinuePod: %s; restarting...", err.Error())
					}
					if res != nil && res.Restart {
						lion.Errorf("chunk was revoked. restarting...")
					}
					if err != nil || res != nil && res.Restart {
						select {
						case exitCh <- struct{}{}:
							// If someone received this signal, then they are
							// responsible to exiting the program and release
							// all resources.
							lion.Errorf("releasing resources...")
							return
						default:
							// Otherwise, we just terminate the program.
							// We use a non-zero exit code so k8s knows to create
							// a new pod.
							lion.Errorf("terminating...")
							os.Exit(1)
						}
					}
				}
			}()

			if response.Transform.Debug {
				lion.SetLevel(lion.LevelDebug)
			}
			// We want to make sure that we only send FinishPod once.
			// The most bulletproof way would be to check that on server side,
			// but this is easier.
			var finished bool
			// Make sure that we call FinishPod even if something caused a panic
			defer func() {
				if r := recover(); r != nil && !finished {
					lion.Errorf("job shim crashed; this is like a bug in pachyderm")
					if _, err := ppsClient.FinishPod(
						context.Background(),
						&ppsserver.FinishPodRequest{
							ChunkID: response.ChunkID,
							PodName: appEnv.PodName,
							Success: false,
						},
					); err != nil && retErr == nil {
						retErr = err
					}
				}
			}()

			c, err := client.NewFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}

			if err := downloadInput(c, response.CommitMounts); err != nil {
				return err
			}

			// Create /pfs/out
			if err := os.MkdirAll(PFSOutputPrefix, 0777); err != nil {
				return err
			}

			// mount /pfs/prev
			var mounts []*fuse.CommitMount
			for _, m := range response.CommitMounts {
				if m.Alias == "prev" {
					mounts = append(mounts, m)
				}
			}
			if len(mounts) > 0 {
				mounter := fuse.NewMounter(appEnv.PachydermAddress, c)
				ready := make(chan bool)
				errCh := make(chan error)
				go func() {
					if err := mounter.MountAndCreate(
						FUSEMountPoint,
						nil,
						mounts,
						ready,
						response.Transform.Debug,
						false,
						true,
					); err != nil {
						errCh <- err
					}
				}()
				select {
				case <-ready:
				case err := <-errCh:
					return err
				}
				defer func() {
					errCh := make(chan error)
					go func() {
						if err := mounter.Unmount(FUSEMountPoint); err != nil {
							errCh <- err
						}
					}()
					select {
					case err := <-errCh:
						if err != nil && retErr == nil {
							retErr = err
						}
					case <-time.After(time.Duration(10 * time.Second)):
						lion.Errorf("unable to unmount FUSE")
					}
				}()

			}
			var readers []io.Reader
			for _, line := range response.Transform.Stdin {
				readers = append(readers, strings.NewReader(line+"\n"))
			}
			if len(response.Transform.Cmd) == 0 {
				lion.Errorf("unable to run; a cmd needs to be provided")
				if _, err := ppsClient.FinishPod(
					context.Background(),
					&ppsserver.FinishPodRequest{
						ChunkID: response.ChunkID,
						PodName: appEnv.PodName,
						Success: false,
					},
				); err != nil {
					return err
				}
				finished = true
				return
			}

			cmdCh := make(chan bool)
			go func() {
				cmd := exec.Command(response.Transform.Cmd[0], response.Transform.Cmd[1:]...)
				cmd.Stdin = io.MultiReader(readers...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				success := true
				if err := cmd.Run(); err != nil {
					success = false
					if exiterr, ok := err.(*exec.ExitError); ok {
						if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
							for _, returnCode := range response.Transform.AcceptReturnCode {
								if int(returnCode) == status.ExitStatus() {
									success = true
								}
							}
						}
					}
					if !success {
						fmt.Fprintf(os.Stderr, "Error from exec: %s\n", err.Error())
					}
				}
				cmdCh <- success
			}()

			var success bool
			select {
			case <-exitCh:
				// Returning an error to ensure that this pod will be restarted
				return errors.New("")
			case success = <-cmdCh:
			}
			var outputMount *fuse.CommitMount
			for _, c := range response.CommitMounts {
				if c.Alias == "out" {
					outputMount = c
					break
				}
			}
			if err := uploadOutput(c, outputMount, response.Transform.Overwrite); err != nil {
				lion.Errorf("err from uploading output: %s\n", err)
				success = false
			}

			res, err := ppsClient.FinishPod(
				context.Background(),
				&ppsserver.FinishPodRequest{
					ChunkID: response.ChunkID,
					PodName: appEnv.PodName,
					Success: success,
				},
			)
			if err != nil {
				return err
			}
			finished = true
			if res.Restart {
				return errors.New("restarting")
			}
			return nil
		}),
	}

	return rootCmd.Execute()
}
<commit_msg>Fix a bug where an error message is printed even when there's no error<commit_after>package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/pachyderm/pachyderm/src/client"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pfs/fuse"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmd"
	"github.com/pachyderm/pachyderm/src/server/pkg/sync"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"

	"github.com/spf13/cobra"
	"go.pedge.io/env"
	"go.pedge.io/lion"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

const (
	// PFSInputPrefix is where the input repos reside
	PFSInputPrefix = "/pfs"
	// PFSOutputPrefix is where the output data resides
	PFSOutputPrefix = "/pfs/out"
	// FUSEMountPoint is where we mount FUSE
	FUSEMountPoint = "/pfs/prev"
)

type appEnv struct {
	PachydermAddress string `env:"PACHD_PORT_650_TCP_ADDR,required"`
	PodName          string `env:"PPS_POD_NAME,required"`
}

func main() {
	env.Main(do, &appEnv{})
}

func downloadInput(c *client.APIClient, commitMounts []*fuse.CommitMount) error {
	var g errgroup.Group
	for _, commitMount := range commitMounts {
		commitMount := commitMount
		if commitMount.Alias == "prev" || commitMount.Alias == "out" {
			continue
		}
		g.Go(func() error {
			return sync.Pull(context.Background(), c.PfsAPIClient, filepath.Join(PFSInputPrefix, commitMount.Commit.Repo.Name),
				commitMount.Commit, commitMount.DiffMethod, commitMount.Shard, commitMount.Lazy)
		})
	}
	return g.Wait()
}

func uploadOutput(c *client.APIClient, out *fuse.CommitMount, overwrite bool) error {
	return sync.Push(context.Background(), c.PfsAPIClient, PFSOutputPrefix, out.Commit, overwrite)
}

func do(appEnvObj interface{}) error {
	appEnv := appEnvObj.(*appEnv)
	rootCmd := &cobra.Command{
		Use:   os.Args[0] + " job-id",
		Short: `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Long:  `Pachyderm job-shim, coordinates with ppsd to create an output commit and run user work.`,
		Run: cmd.RunFixedArgs(1, func(args []string) (retErr error) {
			ppsClient, err := ppsserver.NewInternalPodAPIClientFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}
			response, err := ppsClient.StartPod(
				context.Background(),
				&ppsserver.StartPodRequest{
					Job: &ppsclient.Job{
						ID: args[0],
					},
					PodName: appEnv.PodName,
				})
			if err != nil {
				lion.Errorf("error from StartPod: %s", err.Error())
				return err
			}

			// Start sending ContinuePod to PPS to signal that we are alive
			exitCh := make(chan struct{})
			go func() {
				tick := time.Tick(10 * time.Second)
				for {
					<-tick
					res, err := ppsClient.ContinuePod(
						context.Background(),
						&ppsserver.ContinuePodRequest{
							ChunkID: response.ChunkID,
							PodName: appEnv.PodName,
						},
					)
					if err != nil {
						lion.Errorf("error from ContinuePod: %s; restarting...", err.Error())
					}
					if res != nil && res.Restart {
						lion.Errorf("chunk was revoked. restarting...")
					}
					if err != nil || res != nil && res.Restart {
						select {
						case exitCh <- struct{}{}:
							// If someone received this signal, then they are
							// responsible to exiting the program and release
							// all resources.
							lion.Errorf("releasing resources...")
							return
						default:
							// Otherwise, we just terminate the program.
							// We use a non-zero exit code so k8s knows to create
							// a new pod.
							lion.Errorf("terminating...")
							os.Exit(1)
						}
					}
				}
			}()

			if response.Transform.Debug {
				lion.SetLevel(lion.LevelDebug)
			}
			// We want to make sure that we only send FinishPod once.
			// The most bulletproof way would be to check that on server side,
			// but this is easier.
			var finished bool
			// Make sure that we call FinishPod even if something caused a panic
			defer func() {
				if r := recover(); r != nil && !finished {
					lion.Errorf("job shim crashed; this is like a bug in pachyderm")
					if _, err := ppsClient.FinishPod(
						context.Background(),
						&ppsserver.FinishPodRequest{
							ChunkID: response.ChunkID,
							PodName: appEnv.PodName,
							Success: false,
						},
					); err != nil && retErr == nil {
						retErr = err
					}
				}
			}()

			c, err := client.NewFromAddress(fmt.Sprintf("%v:650", appEnv.PachydermAddress))
			if err != nil {
				return err
			}

			if err := downloadInput(c, response.CommitMounts); err != nil {
				return err
			}

			// Create /pfs/out
			if err := os.MkdirAll(PFSOutputPrefix, 0777); err != nil {
				return err
			}

			// mount /pfs/prev
			var mounts []*fuse.CommitMount
			for _, m := range response.CommitMounts {
				if m.Alias == "prev" {
					mounts = append(mounts, m)
				}
			}
			if len(mounts) > 0 {
				mounter := fuse.NewMounter(appEnv.PachydermAddress, c)
				ready := make(chan bool)
				errCh := make(chan error)
				go func() {
					if err := mounter.MountAndCreate(
						FUSEMountPoint,
						nil,
						mounts,
						ready,
						response.Transform.Debug,
						false,
						true,
					); err != nil {
						errCh <- err
					}
				}()
				select {
				case <-ready:
				case err := <-errCh:
					return err
				}
				defer func() {
					errCh := make(chan error)
					go func() {
						errCh <- mounter.Unmount(FUSEMountPoint)
					}()
					select {
					case err := <-errCh:
						if err != nil && retErr == nil {
							retErr = err
						}
					case <-time.After(time.Duration(10 * time.Second)):
						lion.Errorf("unable to unmount FUSE")
					}
				}()

			}
			var readers []io.Reader
			for _, line := range response.Transform.Stdin {
				readers = append(readers, strings.NewReader(line+"\n"))
			}
			if len(response.Transform.Cmd) == 0 {
				lion.Errorf("unable to run; a cmd needs to be provided")
				if _, err := ppsClient.FinishPod(
					context.Background(),
					&ppsserver.FinishPodRequest{
						ChunkID: response.ChunkID,
						PodName: appEnv.PodName,
						Success: false,
					},
				); err != nil {
					return err
				}
				finished = true
				return
			}

			cmdCh := make(chan bool)
			go func() {
				cmd := exec.Command(response.Transform.Cmd[0], response.Transform.Cmd[1:]...)
				cmd.Stdin = io.MultiReader(readers...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				success := true
				if err := cmd.Run(); err != nil {
					success = false
					if exiterr, ok := err.(*exec.ExitError); ok {
						if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
							for _, returnCode := range response.Transform.AcceptReturnCode {
								if int(returnCode) == status.ExitStatus() {
									success = true
								}
							}
						}
					}
					if !success {
						fmt.Fprintf(os.Stderr, "Error from exec: %s\n", err.Error())
					}
				}
				cmdCh <- success
			}()

			var success bool
			select {
			case <-exitCh:
				// Returning an error to ensure that this pod will be restarted
				return errors.New("")
			case success = <-cmdCh:
			}
			var outputMount *fuse.CommitMount
			for _, c := range response.CommitMounts {
				if c.Alias == "out" {
					outputMount = c
					break
				}
			}
			if err := uploadOutput(c, outputMount, response.Transform.Overwrite); err != nil {
				lion.Errorf("err from uploading output: %s\n", err)
				success = false
			}

			res, err := ppsClient.FinishPod(
				context.Background(),
				&ppsserver.FinishPodRequest{
					ChunkID: response.ChunkID,
					PodName: appEnv.PodName,
					Success: success,
				},
			)
			if err != nil {
				return err
			}
			finished = true
			if res.Restart {
				return errors.New("restarting")
			}
			return nil
		}),
	}

	return rootCmd.Execute()
}
<|endoftext|>


<commit_before>package pretty

import (
	"fmt"
	"html/template"
	"io"
	"os"

	"github.com/docker/go-units"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/pretty"
)

// PrintRepoHeader prints a repo header.
func PrintRepoHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tCREATED\tSIZE\t\n")
}

// PrintRepoInfo pretty-prints repo info.
func PrintRepoInfo(w io.Writer, repoInfo *pfs.RepoInfo) {
	fmt.Fprintf(w, "%s\t", repoInfo.Repo.Name)
	fmt.Fprintf(
		w,
		"%s\t",
		pretty.Ago(repoInfo.Created),
	)
	fmt.Fprintf(w, "%s\t\n", units.BytesSize(float64(repoInfo.SizeBytes)))
}

// PrintDetailedRepoInfo pretty-prints detailed repo info.
func PrintDetailedRepoInfo(repoInfo *pfs.RepoInfo) error {
	template, err := template.New("RepoInfo").Funcs(funcMap).Parse(
		`Name: {{.Repo.Name}}
Created: {{prettyAgo .Created}}
Size: {{prettySize .SizeBytes}}{{if .Provenance}}
Provenance: {{range .Provenance}} {{.Name}} {{end}} {{end}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, repoInfo)
	if err != nil {
		return err
	}
	return nil
}

// PrintBranchHeader prints a branch header.
func PrintBranchHeader(w io.Writer) {
	fmt.Fprint(w, "BRANCH\tHEAD\t\n")
}

// PrintBranch pretty-prints a Branch.
func PrintBranch(w io.Writer, branch *pfs.Branch) {
	fmt.Fprintf(w, "%s\t", branch.Name)
	fmt.Fprintf(w, "%s\t\n", branch.Head.ID)
}

// PrintCommitInfoHeader prints a commit info header.
func PrintCommitInfoHeader(w io.Writer) {
	fmt.Fprint(w, "REPO\tID\tPARENT\tSTARTED\tDURATION\tSIZE\t\n")
}

// PrintCommitInfo pretty-prints commit info.
func PrintCommitInfo(w io.Writer, commitInfo *pfs.CommitInfo) {
	fmt.Fprintf(w, "%s\t", commitInfo.Commit.Repo.Name)
	fmt.Fprintf(w, "%s\t", commitInfo.Commit.ID)
	if commitInfo.ParentCommit != nil {
		fmt.Fprintf(w, "%s\t", commitInfo.ParentCommit.ID)
	} else {
		fmt.Fprint(w, "<none>\t")
	}
	fmt.Fprintf(
		w,
		"%s\t",
		pretty.Ago(commitInfo.Started),
	)
	duration := "-\t"
	if commitInfo.Finished != nil {
		duration = fmt.Sprintf("%s\t", pretty.Duration(commitInfo.Started, commitInfo.Finished))
	}
	fmt.Fprintf(w, duration)
	fmt.Fprintf(w, "%s\t\n", units.BytesSize(float64(commitInfo.SizeBytes)))
}

// PrintDetailedCommitInfo pretty-prints detailed commit info.
func PrintDetailedCommitInfo(commitInfo *pfs.CommitInfo) error {
	template, err := template.New("CommitInfo").Funcs(funcMap).Parse(
		`Commit: {{.Commit.Repo.Name}}/{{.Commit.ID}}{{if .ParentCommit}}
Parent: {{.ParentCommit.ID}} {{end}}
Started: {{prettyAgo .Started}}{{if .Finished}}
Finished: {{prettyAgo .Finished}} {{end}}
Size: {{prettySize .SizeBytes}}{{if .Provenance}}
Provenance: {{range .Provenance}} {{.Repo.Name}}/{{.ID}} {{end}} {{end}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, commitInfo)
	if err != nil {
		return err
	}
	return nil
}

// PrintFileInfoHeader prints a file info header.
func PrintFileInfoHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tTYPE\tSIZE\t\n")
}

// PrintFileInfo pretty-prints file info.
// If recurse is false and directory size is 0, display "-" instead
// If fast is true and file size is 0, display "-" instead
func PrintFileInfo(w io.Writer, fileInfo *pfs.FileInfo) {
	fmt.Fprintf(w, "%s\t", fileInfo.File.Path)
	if fileInfo.FileType == pfs.FileType_FILE {
		fmt.Fprint(w, "file\t")
	} else {
		fmt.Fprint(w, "dir\t")
	}
	fmt.Fprintf(w, "%s\t\n", units.BytesSize(float64(fileInfo.SizeBytes)))
}

// PrintDetailedFileInfo pretty-prints detailed file info.
func PrintDetailedFileInfo(fileInfo *pfs.FileInfo) error {
	template, err := template.New("FileInfo").Funcs(funcMap).Parse(
		`Path: {{.File.Path}}
Type: {{fileType .FileType}}
Size: {{prettySize .SizeBytes}}
Children: {{range .Children}} {{.Path}} {{end}}
`)
	if err != nil {
		return err
	}
	if err := template.Execute(os.Stdout, fileInfo); err != nil {
		return err
	}
	return nil
}

// PrintBlockInfoHeader prints a block info header.
func PrintBlockInfoHeader(w io.Writer) {
	fmt.Fprintf(w, "HASH\tCREATED\tSIZE\t\n")
}

// PrintBlockInfo pretty-prints block info.
func PrintBlockInfo(w io.Writer, blockInfo *pfs.BlockInfo) {
	fmt.Fprintf(w, "%s\t", blockInfo.Block.Hash)
	fmt.Fprintf(
		w,
		"%s\t",
		pretty.Ago(blockInfo.Created),
	)
	fmt.Fprintf(w, "%s\t\n", units.BytesSize(float64(blockInfo.SizeBytes)))
}

type uint64Slice []uint64

func (s uint64Slice) Len() int           { return len(s) }
func (s uint64Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s uint64Slice) Less(i, j int) bool { return s[i] < s[j] }

func fileType(fileType pfs.FileType) string {
	if fileType == pfs.FileType_FILE {
		return "file"
	}
	return "dir"
}

var funcMap = template.FuncMap{
	"prettyAgo":  pretty.Ago,
	"prettySize": pretty.Size,
	"fileType":   fileType,
}
<commit_msg>Get pretty printing compiling.<commit_after>package pretty

import (
	"fmt"
	"html/template"
	"io"
	"os"

	"github.com/docker/go-units"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/pretty"
)

// PrintRepoHeader prints a repo header.
func PrintRepoHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tCREATED\tSIZE\t\n")
}

// PrintRepoInfo pretty-prints repo info.
func PrintRepoInfo(w io.Writer, repoInfo *pfs.RepoInfo) {
	fmt.Fprintf(w, "%s\t", repoInfo.Repo.Name)
	fmt.Fprintf(
		w,
		"%s\t",
		pretty.Ago(repoInfo.Created),
	)
	fmt.Fprintf(w, "%s\t\n", units.BytesSize(float64(repoInfo.SizeBytes)))
}

// PrintDetailedRepoInfo pretty-prints detailed repo info.
func PrintDetailedRepoInfo(repoInfo *pfs.RepoInfo) error {
	template, err := template.New("RepoInfo").Funcs(funcMap).Parse(
		`Name: {{.Repo.Name}}
Created: {{prettyAgo .Created}}
Size: {{prettySize .SizeBytes}}{{if .Provenance}}
Provenance: {{range .Provenance}} {{.Name}} {{end}} {{end}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, repoInfo)
	if err != nil {
		return err
	}
	return nil
}

// PrintBranchHeader prints a branch header.
func PrintBranchHeader(w io.Writer) {
	fmt.Fprint(w, "BRANCH\tHEAD\t\n")
}

// PrintBranch pretty-prints a Branch.
func PrintBranch(w io.Writer, branch *pfs.Branch) {
	fmt.Fprintf(w, "%s\t", branch.Name)
	fmt.Fprintf(w, "%s\t\n", branch.Head.ID)
}

// PrintCommitInfoHeader prints a commit info header.
func PrintCommitInfoHeader(w io.Writer) {
	fmt.Fprint(w, "REPO\tID\tPARENT\tSTARTED\tDURATION\tSIZE\t\n")
}

// PrintCommitInfo pretty-prints commit info.
func PrintCommitInfo(w io.Writer, commitInfo *pfs.CommitInfo) {
	fmt.Fprintf(w, "%s\t", commitInfo.Commit.Repo.Name)
	fmt.Fprintf(w, "%s\t", commitInfo.Commit.ID)
	if commitInfo.ParentCommit != nil {
		fmt.Fprintf(w, "%s\t", commitInfo.ParentCommit.ID)
	} else {
		fmt.Fprint(w, "<none>\t")
	}
	fmt.Fprintf(
		w,
		"%s\t",
		pretty.Ago(commitInfo.Started),
	)
	duration := "-\t"
	if commitInfo.Finished != nil {
		duration = fmt.Sprintf("%s\t", pretty.Duration(commitInfo.Started, commitInfo.Finished))
	}
	fmt.Fprintf(w, duration)
	fmt.Fprintf(w, "%s\t\n", units.BytesSize(float64(commitInfo.SizeBytes)))
}

// PrintDetailedCommitInfo pretty-prints detailed commit info.
func PrintDetailedCommitInfo(commitInfo *pfs.CommitInfo) error {
	template, err := template.New("CommitInfo").Funcs(funcMap).Parse(
		`Commit: {{.Commit.Repo.Name}}/{{.Commit.ID}}{{if .ParentCommit}}
Parent: {{.ParentCommit.ID}} {{end}}
Started: {{prettyAgo .Started}}{{if .Finished}}
Finished: {{prettyAgo .Finished}} {{end}}
Size: {{prettySize .SizeBytes}}{{if .Provenance}}
Provenance: {{range .Provenance}} {{.Repo.Name}}/{{.ID}} {{end}} {{end}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, commitInfo)
	if err != nil {
		return err
	}
	return nil
}

// PrintFileInfoHeader prints a file info header.
func PrintFileInfoHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tTYPE\tSIZE\t\n")
}

// PrintFileInfo pretty-prints file info.
// If recurse is false and directory size is 0, display "-" instead
// If fast is true and file size is 0, display "-" instead
func PrintFileInfo(w io.Writer, fileInfo *pfs.FileInfo) {
	fmt.Fprintf(w, "%s\t", fileInfo.File.Path)
	if fileInfo.FileType == pfs.FileType_FILE {
		fmt.Fprint(w, "file\t")
	} else {
		fmt.Fprint(w, "dir\t")
	}
	fmt.Fprintf(w, "%s\t\n", units.BytesSize(float64(fileInfo.SizeBytes)))
}

// PrintDetailedFileInfo pretty-prints detailed file info.
func PrintDetailedFileInfo(fileInfo *pfs.FileInfo) error {
	template, err := template.New("FileInfo").Funcs(funcMap).Parse(
		`Path: {{.File.Path}}
Type: {{fileType .FileType}}
Size: {{prettySize .SizeBytes}}
Children: {{range .Children}} {{.Path}} {{end}}
`)
	if err != nil {
		return err
	}
	if err := template.Execute(os.Stdout, fileInfo); err != nil {
		return err
	}
	return nil
}

type uint64Slice []uint64

func (s uint64Slice) Len() int           { return len(s) }
func (s uint64Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s uint64Slice) Less(i, j int) bool { return s[i] < s[j] }

func fileType(fileType pfs.FileType) string {
	if fileType == pfs.FileType_FILE {
		return "file"
	}
	return "dir"
}

var funcMap = template.FuncMap{
	"prettyAgo":  pretty.Ago,
	"prettySize": pretty.Size,
	"fileType":   fileType,
}
<|endoftext|>


<commit_before><commit_msg>fix links<commit_after>package tensor3

// component type
type BaseType float32

const scale = 1

func baseScale(v BaseType) BaseType {
	return v
}

func baseUnscale(v BaseType) BaseType {
	return v
}

func vectorScale(v *Vector) {
	return
}

func vectorUnscale(v *Vector) {
	return
}

func New(x,y,z float32) *Vector{
	return &Vector{BaseType(x),BaseType(y),BaseType(z)}
}

<|endoftext|>


<commit_before>package server

import (
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/shard"
	pfsserver "github.com/pachyderm/pachyderm/src/server/pfs"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"go.pedge.io/lion/proto"
)

var (
	blockSize = 8 * 1024 * 1024 // 8 Megabytes
)

type APIServer interface {
	pfsclient.APIServer // SJ: WOW this is bad naming
	shard.Frontend
}

type InternalAPIServer interface {
	pfsclient.InternalAPIServer // SJ: also bad naming
	shard.Server
}

func NewAPIServer(hasher *pfsserver.Hasher, router shard.Router) APIServer {
	return newAPIServer(hasher, router)
}

func NewInternalAPIServer(hasher *pfsserver.Hasher, router shard.Router, driver drive.Driver) InternalAPIServer {
	return newInternalAPIServer(hasher, router, driver)
}

func NewLocalBlockAPIServer(dir string) (pfsclient.BlockAPIServer, error) { // SJ: also bad naming
	return newLocalBlockAPIServer(dir)
}

func NewObjBlockAPIServer(dir string, objClient obj.Client) (pfsclient.BlockAPIServer, error) { // SJ: Also bad naming
	return newObjBlockAPIServer(dir, objClient)
}

// NewBlockAPIServer creates a BlockAPIServer using the credentials it finds in
// the environment
func NewBlockAPIServer(dir string) (pfsclient.BlockAPIServer, error) {
	if blockAPIServer, err := newAmazonBlockAPIServer(dir); err == nil {
		return blockAPIServer, nil
	} else {
		protolion.Errorf("error create Amazon block backend: %s", err.Error())
	}
	if blockAPIServer, err := newGoogleBlockAPIServer(dir); err == nil {
		return blockAPIServer, nil
	} else {
		protolion.Errorf("error create Google block backend: %s", err.Error())
	}
	protolion.Errorf("failed to create obj backend, falling back to local")
	return NewLocalBlockAPIServer(dir)
}
<commit_msg>Happify the linter.<commit_after>package server

import (
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/shard"
	pfsserver "github.com/pachyderm/pachyderm/src/server/pfs"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"
	"go.pedge.io/lion/proto"
)

var (
	blockSize = 8 * 1024 * 1024 // 8 Megabytes
)

type APIServer interface {
	pfsclient.APIServer // SJ: WOW this is bad naming
	shard.Frontend
}

type InternalAPIServer interface {
	pfsclient.InternalAPIServer // SJ: also bad naming
	shard.Server
}

func NewAPIServer(hasher *pfsserver.Hasher, router shard.Router) APIServer {
	return newAPIServer(hasher, router)
}

func NewInternalAPIServer(hasher *pfsserver.Hasher, router shard.Router, driver drive.Driver) InternalAPIServer {
	return newInternalAPIServer(hasher, router, driver)
}

func NewLocalBlockAPIServer(dir string) (pfsclient.BlockAPIServer, error) { // SJ: also bad naming
	return newLocalBlockAPIServer(dir)
}

func NewObjBlockAPIServer(dir string, objClient obj.Client) (pfsclient.BlockAPIServer, error) { // SJ: Also bad naming
	return newObjBlockAPIServer(dir, objClient)
}

// NewBlockAPIServer creates a BlockAPIServer using the credentials it finds in
// the environment
func NewBlockAPIServer(dir string) (pfsclient.BlockAPIServer, error) {
	var blockAPIServer pfsclient.BlockAPIServer
	var err error
	if blockAPIServer, err = newAmazonBlockAPIServer(dir); err == nil {
		return blockAPIServer, nil
	}
	protolion.Errorf("error create Amazon block backend: %s", err.Error())
	if blockAPIServer, err = newGoogleBlockAPIServer(dir); err == nil {
		return blockAPIServer, nil
	}
	protolion.Errorf("error create Google block backend: %s", err.Error())
	protolion.Errorf("failed to create obj backend, falling back to local")
	return NewLocalBlockAPIServer(dir)
}
<|endoftext|>