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

package query

import (
	"errors"
	"fmt"

	"github.com/square/metrics/api"
	"github.com/square/metrics/query/aggregate"
)

// EvaluationContext is the central piece of logic, providing
// helper funcions & varaibles to evaluate a given piece of
// metrics query.
// * Contains Backend object, which can be used to fetch data
// from the backend system.s
// * Contains current timerange being queried for - this can be
// changed by say, application of time shift function.
type EvaluationContext struct {
	// Backend to fetch data from
	Backend api.Backend
	// Timerange to fetch data from
	Timerange api.Timerange
	// SampleMethod to use when up/downsampling to match the requested resolution
	SampleMethod api.SampleMethod
	Predicate    api.Predicate
}

// A value is the result of evaluating an expression.
// They can be floating point values, strings, or series lists.
type value interface {
	toSeriesList(api.Timerange) (api.SeriesList, error)
	toString() (string, error)
	toScalar() (float64, error)
}

type conversionError struct {
	from string
	to   string
}

func (e conversionError) Error() string {
	return fmt.Sprintf("cannot convert from type %s to type %s", e.from, e.to)
}

// A seriesListValue is a value which holds a SeriesList
type seriesListValue api.SeriesList

func (value seriesListValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList(value), nil
}
func (value seriesListValue) toString() (string, error) {
	return "", conversionError{"SeriesList", "string"}
}
func (value seriesListValue) toScalar() (float64, error) {
	return 0, conversionError{"SeriesList", "scalar"}
}

// A stringValue holds a string
type stringValue string

func (value stringValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList{}, conversionError{"string", "SeriesList"}
}
func (value stringValue) toString() (string, error) {
	return string(value), nil
}
func (value stringValue) toScalar() (float64, error) {
	return 0, conversionError{"string", "scalar"}
}

// A scalarValue holds a float and can be converted to a serieslist
type scalarValue float64

func (value scalarValue) toSeriesList(timerange api.Timerange) (api.SeriesList, error) {
	if !timerange.IsValid() {
		return api.SeriesList{}, errors.New("Invalid context.Timerange")
	}

	series := make([]float64, timerange.Slots())
	for i := range series {
		series[i] = float64(value)
	}

	return api.SeriesList{
		Series:    []api.Timeseries{api.Timeseries{series, api.NewTagSet()}},
		Timerange: timerange,
	}, nil
}

func (value scalarValue) toString() (string, error) {
	return "", conversionError{"scalar", "string"}
}

func (value scalarValue) toScalar() (float64, error) {
	return float64(value), nil
}

// Expression is a piece of code, which can be evaluated in a given
// EvaluationContext. EvaluationContext must never be changed in an Evalute().
//
// The contract of Expressions is that leaf nodes must sample a resulting
// timeseries according to the resolution specified in its EvaluationContext's
// Timerange. Internal nodes may assume that results from evaluating child
// Expressions correspond to the timerange in the current EvaluationContext.
type Expression interface {
	// Evaluate the given expression.
	Evaluate(context EvaluationContext) (value, error)
}

func evaluateToSeriesList(e Expression, context EvaluationContext) (api.SeriesList, error) {
	value, err := e.Evaluate(context)
	if err != nil {
		return api.SeriesList{}, err
	}
	return value.toSeriesList(context.Timerange)
}

// Implementations
// ===============

// Generates a Timeseries from the encapsulated scalar.
func (expr scalarExpression) Evaluate(context EvaluationContext) (value, error) {
	return scalarValue(expr.value), nil
}

func (expr *metricFetchExpression) Evaluate(context EvaluationContext) (value, error) {
	// Merge predicates appropriately
	var predicate api.Predicate
	if context.Predicate == nil {
		predicate = expr.predicate
	} else if expr.predicate == nil {
		predicate = context.Predicate
	} else {
		predicate = &andPredicate{[]api.Predicate{expr.predicate, context.Predicate}}
	}

	series, err := context.Backend.FetchSeries(api.TaggedMetric{api.MetricKey(expr.metricName), nil}, predicate, context.SampleMethod, context.Timerange)
	if err != nil {
		return seriesListValue{}, err
	}
	return seriesListValue(*series), err
}

func (expr *functionExpression) Evaluate(context EvaluationContext) (value, error) {

	name := expr.functionName
	switch name {
	case "+":
		fallthrough
	case "-":
		fallthrough
	case "*":
		fallthrough
	case "/":
		arguments := make([]value, len(expr.arguments))
		var err error
		for i := range arguments {
			arguments[i], err = expr.arguments[i].Evaluate(context)
			if err != nil {
				return nil, err
			}
		}
		// Evaluation of a binary operator:
		if len(arguments) != 2 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects 2 operands but received %d (%+v)", name, len(arguments), arguments))
		}
		left := arguments[0]
		right := arguments[1]
		operatorMap := map[string]func(float64, float64) float64{
			"+": func(x, y float64) float64 { return x + y },
			"-": func(x, y float64) float64 { return x - y },
			"*": func(x, y float64) float64 { return x * y },
			"/": func(x, y float64) float64 { return x / y },
		}
		return evaluateBinaryOperation(context, name, left, right, operatorMap[name])
	case "aggregate.sum":
		fallthrough
	case "aggregate.mean":
		fallthrough
	case "aggregate.min":
		fallthrough
	case "aggregate.max":
		arguments := make([]value, len(expr.arguments))
		var err error
		for i := range arguments {
			arguments[i], err = expr.arguments[i].Evaluate(context)
			if err != nil {
				return nil, err
			}
		}
		if len(arguments) != 1 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expected 1 operand but received %d (%+v)", name, len(arguments), arguments))
		}
		list, err := arguments[0].toSeriesList(context.Timerange)
		if err != nil {
			return nil, err
		}
		series, err := aggregate.AggregateBy(list, name, expr.groupBy)
		if err != nil {
			return nil, err
		}
		return seriesListValue(series), nil
	default:
		return nil, errors.New(fmt.Sprintf("Invalid function: %s", expr.functionName))
	}
}

//
// Auxiliary functions
//

// evaluateExpression wraps expr.Evaluate() to provide common messaging
// for errors. This can get pretty messy if the Expression we evaluate
// isn't a leaf node, but a leaf fails.
func evaluateExpression(context EvaluationContext, expr Expression) (value, error) {
	result, err := expr.Evaluate(context)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Evaluation of expression %+v failed:\n%s\n", expr, err.Error()))
	}
	return result, err
}

// evaluateBinaryOperation applies an arbirary binary operation to two
// Expressions.
func evaluateBinaryOperation(
	context EvaluationContext,
	functionName string,
	leftValue value,
	rightValue value,
	evaluate func(float64, float64) float64,
) (value, error) {

	leftList, err := leftValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}
	rightList, err := rightValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}

	joined := join([]api.SeriesList{leftList, rightList})

	result := make([]api.Timeseries, len(joined.Rows))

	for i, row := range joined.Rows {
		left := row.Row[0]
		right := row.Row[1]
		array := make([]float64, len(left.Values))
		for j := 0; j < len(left.Values); j++ {
			array[j] = evaluate(left.Values[j], right.Values[j])
		}
		result[i] = api.Timeseries{array, row.TagSet}
	}

	return seriesListValue(api.SeriesList{
		Series:    result,
		Timerange: context.Timerange,
	}), nil
}

// evaluateExpressions evaluates all provided Expressions in the
// EvaluationContext. If any evaluations error, evaluateExpressions will
// propagate that error. The resulting SeriesLists will be in an order
// corresponding to the provided Expresesions.
func evaluateExpressions(context EvaluationContext, expressions []Expression) ([]value, error) {
	if len(expressions) == 0 {
		return []value{}, nil
	}

	results := make([]value, len(expressions))
	for i, expr := range expressions {
		result, err := expr.Evaluate(context)
		if err != nil {
			return nil, err
		}

		results[i] = result
	}

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

package query

import (
	"errors"
	"fmt"

	"github.com/square/metrics/api"
	"github.com/square/metrics/query/aggregate"
)

// EvaluationContext is the central piece of logic, providing
// helper funcions & varaibles to evaluate a given piece of
// metrics query.
// * Contains Backend object, which can be used to fetch data
// from the backend system.s
// * Contains current timerange being queried for - this can be
// changed by say, application of time shift function.
type EvaluationContext struct {
	// Backend to fetch data from
	Backend api.Backend
	// Timerange to fetch data from
	Timerange api.Timerange
	// SampleMethod to use when up/downsampling to match the requested resolution
	SampleMethod api.SampleMethod
	Predicate    api.Predicate
}

// A value is the result of evaluating an expression.
// They can be floating point values, strings, or series lists.
type value interface {
	toSeriesList(api.Timerange) (api.SeriesList, error)
	toString() (string, error)
	toScalar() (float64, error)
}

type conversionError struct {
	from string
	to   string
}

func (e conversionError) Error() string {
	return fmt.Sprintf("cannot convert from type %s to type %s", e.from, e.to)
}

// A seriesListValue is a value which holds a SeriesList
type seriesListValue api.SeriesList

func (value seriesListValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList(value), nil
}
func (value seriesListValue) toString() (string, error) {
	return "", conversionError{"SeriesList", "string"}
}
func (value seriesListValue) toScalar() (float64, error) {
	return 0, conversionError{"SeriesList", "scalar"}
}

// A stringValue holds a string
type stringValue string

func (value stringValue) toSeriesList(time api.Timerange) (api.SeriesList, error) {
	return api.SeriesList{}, conversionError{"string", "SeriesList"}
}
func (value stringValue) toString() (string, error) {
	return string(value), nil
}
func (value stringValue) toScalar() (float64, error) {
	return 0, conversionError{"string", "scalar"}
}

// A scalarValue holds a float and can be converted to a serieslist
type scalarValue float64

func (value scalarValue) toSeriesList(timerange api.Timerange) (api.SeriesList, error) {
	if !timerange.IsValid() {
		return api.SeriesList{}, errors.New("Invalid context.Timerange")
	}

	series := make([]float64, timerange.Slots())
	for i := range series {
		series[i] = float64(value)
	}

	return api.SeriesList{
		Series:    []api.Timeseries{api.Timeseries{series, api.NewTagSet()}},
		Timerange: timerange,
	}, nil
}

func (value scalarValue) toString() (string, error) {
	return "", conversionError{"scalar", "string"}
}

func (value scalarValue) toScalar() (float64, error) {
	return float64(value), nil
}

// Expression is a piece of code, which can be evaluated in a given
// EvaluationContext. EvaluationContext must never be changed in an Evalute().
//
// The contract of Expressions is that leaf nodes must sample a resulting
// timeseries according to the resolution specified in its EvaluationContext's
// Timerange. Internal nodes may assume that results from evaluating child
// Expressions correspond to the timerange in the current EvaluationContext.
type Expression interface {
	// Evaluate the given expression.
	Evaluate(context EvaluationContext) (value, error)
}

func evaluateToSeriesList(e Expression, context EvaluationContext) (api.SeriesList, error) {
	value, err := e.Evaluate(context)
	if err != nil {
		return api.SeriesList{}, err
	}
	return value.toSeriesList(context.Timerange)
}

// Implementations
// ===============

// Generates a Timeseries from the encapsulated scalar.
func (expr scalarExpression) Evaluate(context EvaluationContext) (value, error) {
	return scalarValue(expr.value), nil
}

func (expr *metricFetchExpression) Evaluate(context EvaluationContext) (value, error) {
	// Merge predicates appropriately
	var predicate api.Predicate
	if context.Predicate == nil {
		predicate = expr.predicate
	} else if expr.predicate == nil {
		predicate = context.Predicate
	} else {
		predicate = &andPredicate{[]api.Predicate{expr.predicate, context.Predicate}}
	}

	series, err := context.Backend.FetchSeries(api.TaggedMetric{api.MetricKey(expr.metricName), nil}, predicate, context.SampleMethod, context.Timerange)
	if err != nil {
		return seriesListValue{}, err
	}
	return seriesListValue(*series), err
}

func (expr *functionExpression) Evaluate(context EvaluationContext) (value, error) {

	name := expr.functionName
	switch name {
	case "+", "-", "*", "/":
		arguments := make([]value, len(expr.arguments))
		var err error
		for i := range arguments {
			arguments[i], err = expr.arguments[i].Evaluate(context)
			if err != nil {
				return nil, err
			}
		}
		// Evaluation of a binary operator:
		if len(arguments) != 2 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expects 2 operands but received %d (%+v)", name, len(arguments), arguments))
		}
		left := arguments[0]
		right := arguments[1]
		operatorMap := map[string]func(float64, float64) float64{
			"+": func(x, y float64) float64 { return x + y },
			"-": func(x, y float64) float64 { return x - y },
			"*": func(x, y float64) float64 { return x * y },
			"/": func(x, y float64) float64 { return x / y },
		}
		return evaluateBinaryOperation(context, name, left, right, operatorMap[name])
	case "aggregate.sum", "aggregate.mean", "aggregate.min", "aggregate.max":
		arguments := make([]value, len(expr.arguments))
		var err error
		for i := range arguments {
			arguments[i], err = expr.arguments[i].Evaluate(context)
			if err != nil {
				return nil, err
			}
		}
		if len(arguments) != 1 {
			return nil, errors.New(fmt.Sprintf("Function `%s` expected 1 operand but received %d (%+v)", name, len(arguments), arguments))
		}
		list, err := arguments[0].toSeriesList(context.Timerange)
		if err != nil {
			return nil, err
		}
		series, err := aggregate.AggregateBy(list, name, expr.groupBy)
		if err != nil {
			return nil, err
		}
		return seriesListValue(series), nil
	default:
		return nil, errors.New(fmt.Sprintf("Invalid function: %s", expr.functionName))
	}
}

//
// Auxiliary functions
//

// evaluateExpression wraps expr.Evaluate() to provide common messaging
// for errors. This can get pretty messy if the Expression we evaluate
// isn't a leaf node, but a leaf fails.
func evaluateExpression(context EvaluationContext, expr Expression) (value, error) {
	result, err := expr.Evaluate(context)
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Evaluation of expression %+v failed:\n%s\n", expr, err.Error()))
	}
	return result, err
}

// evaluateBinaryOperation applies an arbirary binary operation to two
// Expressions.
func evaluateBinaryOperation(
	context EvaluationContext,
	functionName string,
	leftValue value,
	rightValue value,
	evaluate func(float64, float64) float64,
) (value, error) {

	leftList, err := leftValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}
	rightList, err := rightValue.toSeriesList(context.Timerange)
	if err != nil {
		return nil, err
	}

	joined := join([]api.SeriesList{leftList, rightList})

	result := make([]api.Timeseries, len(joined.Rows))

	for i, row := range joined.Rows {
		left := row.Row[0]
		right := row.Row[1]
		array := make([]float64, len(left.Values))
		for j := 0; j < len(left.Values); j++ {
			array[j] = evaluate(left.Values[j], right.Values[j])
		}
		result[i] = api.Timeseries{array, row.TagSet}
	}

	return seriesListValue(api.SeriesList{
		Series:    result,
		Timerange: context.Timerange,
	}), nil
}

// evaluateExpressions evaluates all provided Expressions in the
// EvaluationContext. If any evaluations error, evaluateExpressions will
// propagate that error. The resulting SeriesLists will be in an order
// corresponding to the provided Expresesions.
func evaluateExpressions(context EvaluationContext, expressions []Expression) ([]value, error) {
	if len(expressions) == 0 {
		return []value{}, nil
	}

	results := make([]value, len(expressions))
	for i, expr := range expressions {
		result, err := expr.Evaluate(context)
		if err != nil {
			return nil, err
		}

		results[i] = result
	}

	return results, nil
}
<|endoftext|>


<commit_before>package flows

import (
	"github.com/skygeario/skygear-server/pkg/auth/dependency/interaction"
	"github.com/skygeario/skygear-server/pkg/auth/model"
	"github.com/skygeario/skygear-server/pkg/core/authn"
)

type AuthAPIFlow struct {
	Interactions   InteractionProvider
	UserController *UserController
}

func (f *AuthAPIFlow) LoginWithLoginIDPassword(
	clientID string, loginIDKey string, loginID string, password string,
) (*AuthResult, error) {
	i, err := f.Interactions.NewInteractionLogin(&interaction.IntentLogin{
		Identity: interaction.IdentitySpec{
			Type: authn.IdentityTypeLoginID,
			Claims: map[string]interface{}{
				interaction.IdentityClaimLoginIDKey:   loginIDKey,
				interaction.IdentityClaimLoginIDValue: loginID,
			},
		},
	}, clientID)
	if err != nil {
		return nil, err
	}

	s, err := f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	} else if s.CurrentStep().Step != interaction.StepAuthenticatePrimary {
		return nil, ErrUnsupportedConfiguration
	}

	err = f.Interactions.PerformAction(
		i,
		interaction.StepAuthenticatePrimary,
		&interaction.ActionAuthenticate{
			Authenticator: interaction.AuthenticatorSpec{Type: interaction.AuthenticatorTypePassword},
			Secret:        password,
		},
	)
	if err != nil {
		return nil, err
	}

	if i.Error != nil {
		return nil, i.Error
	}

	s, err = f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	}

	switch s.CurrentStep().Step {
	case interaction.StepAuthenticateSecondary, interaction.StepSetupSecondaryAuthenticator:
		panic("interaction_flow_auth_api: TODO: handle MFA")

	case interaction.StepCommit:
		attrs, err := f.Interactions.Commit(i)
		if err != nil {
			return nil, err
		}

		result, err := f.UserController.CreateSession(i, attrs, true)
		if err != nil {
			return nil, err
		}

		return result, nil

	default:
		return nil, ErrUnsupportedConfiguration
	}
}

func (f *AuthAPIFlow) SignupWithLoginIDPassword(
	clientID string,
	loginIDKey string,
	loginID string,
	password string,
	metadata map[string]interface{},
	onUserDuplicate model.OnUserDuplicate,
) (*AuthResult, error) {
	i, err := f.Interactions.NewInteractionSignup(&interaction.IntentSignup{
		Identity: interaction.IdentitySpec{
			Type: authn.IdentityTypeLoginID,
			Claims: map[string]interface{}{
				interaction.IdentityClaimLoginIDKey:   loginIDKey,
				interaction.IdentityClaimLoginIDValue: loginID,
			},
		},
		OnUserDuplicate: onUserDuplicate,
		UserMetadata:    metadata,
	}, clientID)
	if err != nil {
		return nil, err
	}

	s, err := f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	} else if s.CurrentStep().Step != interaction.StepSetupPrimaryAuthenticator {
		return nil, ErrUnsupportedConfiguration
	}

	err = f.Interactions.PerformAction(
		i,
		interaction.StepSetupPrimaryAuthenticator,
		&interaction.ActionSetupAuthenticator{
			Authenticator: interaction.AuthenticatorSpec{Type: interaction.AuthenticatorTypePassword},
			Secret:        password,
		},
	)
	if err != nil {
		return nil, err
	}

	if i.Error != nil {
		return nil, i.Error
	}

	_, err = f.Interactions.Commit(i)
	if err != nil {
		return nil, err
	}

	// Login with new user:

	i, err = f.Interactions.NewInteractionLogin(&interaction.IntentLogin{
		Identity: i.Identity.ToSpec(),
		AuthenticatedAs: &interaction.IntentLoginAuthenticatedAs{
			UserID:               i.UserID,
			PrimaryAuthenticator: i.PrimaryAuthenticator.ToSpec(),
		},
		OriginalIntentType: i.Intent.Type(),
	}, i.ClientID)
	if err != nil {
		return nil, err
	}

	s, err = f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	}

	switch s.CurrentStep().Step {
	case interaction.StepAuthenticateSecondary, interaction.StepSetupSecondaryAuthenticator:
		panic("interaction_flow_auth_api: TODO: handle MFA")

	case interaction.StepCommit:
		attrs, err := f.Interactions.Commit(i)
		if err != nil {
			return nil, err
		}

		result, err := f.UserController.CreateSession(i, attrs, true)
		if err != nil {
			return nil, err
		}

		return result, nil

	default:
		return nil, ErrUnsupportedConfiguration
	}
}
<commit_msg>Add TODO for signup auth API<commit_after>package flows

import (
	"github.com/skygeario/skygear-server/pkg/auth/dependency/interaction"
	"github.com/skygeario/skygear-server/pkg/auth/model"
	"github.com/skygeario/skygear-server/pkg/core/authn"
)

type AuthAPIFlow struct {
	Interactions   InteractionProvider
	UserController *UserController
}

func (f *AuthAPIFlow) LoginWithLoginIDPassword(
	clientID string, loginIDKey string, loginID string, password string,
) (*AuthResult, error) {
	i, err := f.Interactions.NewInteractionLogin(&interaction.IntentLogin{
		Identity: interaction.IdentitySpec{
			Type: authn.IdentityTypeLoginID,
			Claims: map[string]interface{}{
				interaction.IdentityClaimLoginIDKey:   loginIDKey,
				interaction.IdentityClaimLoginIDValue: loginID,
			},
		},
	}, clientID)
	if err != nil {
		return nil, err
	}

	s, err := f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	} else if s.CurrentStep().Step != interaction.StepAuthenticatePrimary {
		return nil, ErrUnsupportedConfiguration
	}

	err = f.Interactions.PerformAction(
		i,
		interaction.StepAuthenticatePrimary,
		&interaction.ActionAuthenticate{
			Authenticator: interaction.AuthenticatorSpec{Type: interaction.AuthenticatorTypePassword},
			Secret:        password,
		},
	)
	if err != nil {
		return nil, err
	}

	if i.Error != nil {
		return nil, i.Error
	}

	s, err = f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	}

	switch s.CurrentStep().Step {
	case interaction.StepAuthenticateSecondary, interaction.StepSetupSecondaryAuthenticator:
		panic("interaction_flow_auth_api: TODO: handle MFA")

	case interaction.StepCommit:
		attrs, err := f.Interactions.Commit(i)
		if err != nil {
			return nil, err
		}

		result, err := f.UserController.CreateSession(i, attrs, true)
		if err != nil {
			return nil, err
		}

		return result, nil

	default:
		return nil, ErrUnsupportedConfiguration
	}
}

// TODO(interaction): support multiple login IDs
func (f *AuthAPIFlow) SignupWithLoginIDPassword(
	clientID string,
	loginIDKey string,
	loginID string,
	password string,
	metadata map[string]interface{},
	onUserDuplicate model.OnUserDuplicate,
) (*AuthResult, error) {
	i, err := f.Interactions.NewInteractionSignup(&interaction.IntentSignup{
		Identity: interaction.IdentitySpec{
			Type: authn.IdentityTypeLoginID,
			Claims: map[string]interface{}{
				interaction.IdentityClaimLoginIDKey:   loginIDKey,
				interaction.IdentityClaimLoginIDValue: loginID,
			},
		},
		OnUserDuplicate: onUserDuplicate,
		UserMetadata:    metadata,
	}, clientID)
	if err != nil {
		return nil, err
	}

	s, err := f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	} else if s.CurrentStep().Step != interaction.StepSetupPrimaryAuthenticator {
		return nil, ErrUnsupportedConfiguration
	}

	err = f.Interactions.PerformAction(
		i,
		interaction.StepSetupPrimaryAuthenticator,
		&interaction.ActionSetupAuthenticator{
			Authenticator: interaction.AuthenticatorSpec{Type: interaction.AuthenticatorTypePassword},
			Secret:        password,
		},
	)
	if err != nil {
		return nil, err
	}

	if i.Error != nil {
		return nil, i.Error
	}

	_, err = f.Interactions.Commit(i)
	if err != nil {
		return nil, err
	}

	// Login with new user:

	i, err = f.Interactions.NewInteractionLogin(&interaction.IntentLogin{
		Identity: i.Identity.ToSpec(),
		AuthenticatedAs: &interaction.IntentLoginAuthenticatedAs{
			UserID:               i.UserID,
			PrimaryAuthenticator: i.PrimaryAuthenticator.ToSpec(),
		},
		OriginalIntentType: i.Intent.Type(),
	}, i.ClientID)
	if err != nil {
		return nil, err
	}

	s, err = f.Interactions.GetInteractionState(i)
	if err != nil {
		return nil, err
	}

	switch s.CurrentStep().Step {
	case interaction.StepAuthenticateSecondary, interaction.StepSetupSecondaryAuthenticator:
		panic("interaction_flow_auth_api: TODO: handle MFA")

	case interaction.StepCommit:
		attrs, err := f.Interactions.Commit(i)
		if err != nil {
			return nil, err
		}

		result, err := f.UserController.CreateSession(i, attrs, true)
		if err != nil {
			return nil, err
		}

		return result, nil

	default:
		return nil, ErrUnsupportedConfiguration
	}
}
<|endoftext|>


<commit_before>package node

import (
	"fmt"
	"net"
	"strconv"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kerrors "k8s.io/apimachinery/pkg/util/errors"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	kubeletoptions "k8s.io/kubernetes/cmd/kubelet/app/options"
	"k8s.io/kubernetes/pkg/features"
	"k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig"
	kubeletcni "k8s.io/kubernetes/pkg/kubelet/network/cni"
	kubelettypes "k8s.io/kubernetes/pkg/kubelet/types"

	configapi "github.com/openshift/origin/pkg/cmd/server/api"
	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	cmdflags "github.com/openshift/origin/pkg/cmd/util/flags"
	"github.com/openshift/origin/pkg/cmd/util/variable"
	"github.com/openshift/origin/pkg/network"
)

// Build creates the core Kubernetes component configs for a given NodeConfig, or returns
// an error
func Build(options configapi.NodeConfig) (*kubeletoptions.KubeletServer, error) {
	imageTemplate := variable.NewDefaultImageTemplate()
	imageTemplate.Format = options.ImageConfig.Format
	imageTemplate.Latest = options.ImageConfig.Latest

	var path string
	var fileCheckInterval int64
	if options.PodManifestConfig != nil {
		path = options.PodManifestConfig.Path
		fileCheckInterval = options.PodManifestConfig.FileCheckIntervalSeconds
	}

	kubeAddressStr, kubePortStr, err := net.SplitHostPort(options.ServingInfo.BindAddress)
	if err != nil {
		return nil, fmt.Errorf("cannot parse node address: %v", err)
	}
	kubePort, err := strconv.Atoi(kubePortStr)
	if err != nil {
		return nil, fmt.Errorf("cannot parse node port: %v", err)
	}

	// Defaults are tested in TestKubeletDefaults
	server, err := kubeletoptions.NewKubeletServer()
	if err != nil {
		return nil, fmt.Errorf("cannot create kubelet server: %v", err)
	}
	// Adjust defaults
	server.RequireKubeConfig = true
	server.KubeConfig.Default(options.MasterKubeConfig)
	server.PodManifestPath = path
	server.RootDirectory = options.VolumeDirectory
	server.NodeIP = options.NodeIP
	server.HostnameOverride = options.NodeName
	server.AllowPrivileged = true
	server.RegisterNode = true
	server.Address = kubeAddressStr
	server.Port = int32(kubePort)
	server.ReadOnlyPort = 0        // no read only access
	server.CAdvisorPort = 0        // no unsecured cadvisor access
	server.HealthzPort = 0         // no unsecured healthz access
	server.HealthzBindAddress = "" // no unsecured healthz access
	server.ClusterDNS = []string{options.DNSIP}
	server.ClusterDomain = options.DNSDomain
	server.NetworkPluginName = options.NetworkConfig.NetworkPluginName
	server.HostNetworkSources = []string{kubelettypes.ApiserverSource, kubelettypes.FileSource}
	server.HostPIDSources = []string{kubelettypes.ApiserverSource, kubelettypes.FileSource}
	server.HostIPCSources = []string{kubelettypes.ApiserverSource, kubelettypes.FileSource}
	server.HTTPCheckFrequency = metav1.Duration{Duration: time.Duration(0)} // no remote HTTP pod creation access
	server.FileCheckFrequency = metav1.Duration{Duration: time.Duration(fileCheckInterval) * time.Second}
	server.KubeletFlags.ContainerRuntimeOptions.PodSandboxImage = imageTemplate.ExpandOrDie("pod")
	server.MaxPods = 250
	server.PodsPerCore = 10
	server.CgroupDriver = "systemd"
	server.DockerExecHandlerName = string(options.DockerConfig.ExecHandlerName)
	server.RemoteRuntimeEndpoint = options.DockerConfig.DockerShimSocket
	server.RemoteImageEndpoint = options.DockerConfig.DockerShimSocket
	server.DockershimRootDirectory = options.DockerConfig.DockershimRootDirectory

	// prevents kube from generating certs
	server.TLSCertFile = options.ServingInfo.ServerCert.CertFile
	server.TLSPrivateKeyFile = options.ServingInfo.ServerCert.KeyFile

	containerized := cmdutil.Env("OPENSHIFT_CONTAINERIZED", "") == "true"
	server.Containerized = containerized

	// force the authentication and authorization
	// Setup auth
	authnTTL, err := time.ParseDuration(options.AuthConfig.AuthenticationCacheTTL)
	if err != nil {
		return nil, err
	}
	server.Authentication = kubeletconfig.KubeletAuthentication{
		X509: kubeletconfig.KubeletX509Authentication{
			ClientCAFile: options.ServingInfo.ClientCA,
		},
		Webhook: kubeletconfig.KubeletWebhookAuthentication{
			Enabled:  true,
			CacheTTL: metav1.Duration{Duration: authnTTL},
		},
		Anonymous: kubeletconfig.KubeletAnonymousAuthentication{
			Enabled: true,
		},
	}
	authzTTL, err := time.ParseDuration(options.AuthConfig.AuthorizationCacheTTL)
	if err != nil {
		return nil, err
	}
	server.Authorization = kubeletconfig.KubeletAuthorization{
		Mode: kubeletconfig.KubeletAuthorizationModeWebhook,
		Webhook: kubeletconfig.KubeletWebhookAuthorization{
			CacheAuthorizedTTL:   metav1.Duration{Duration: authzTTL},
			CacheUnauthorizedTTL: metav1.Duration{Duration: authzTTL},
		},
	}

	// resolve extended arguments
	// TODO: this should be done in config validation (along with the above) so we can provide
	// proper errors
	if err := cmdflags.Resolve(options.KubeletArguments, server.AddFlags); len(err) > 0 {
		return nil, kerrors.NewAggregate(err)
	}

	// terminate early if feature gate is incorrect on the node
	if len(server.FeatureGates) > 0 {
		if err := utilfeature.DefaultFeatureGate.Set(server.FeatureGates); err != nil {
			return nil, err
		}
	}
	if utilfeature.DefaultFeatureGate.Enabled(features.RotateKubeletServerCertificate) {
		// Server cert rotation is ineffective if a cert is hardcoded.
		if len(server.CertDirectory) > 0 {
			server.TLSCertFile = ""
			server.TLSPrivateKeyFile = ""
		}
	}

	if network.IsOpenShiftNetworkPlugin(options.NetworkConfig.NetworkPluginName) {
		// SDN plugin pod setup/teardown is implemented as a CNI plugin
		server.NetworkPluginName = kubeletcni.CNIPluginName
		server.NetworkPluginDir = kubeletcni.DefaultNetDir
		server.CNIConfDir = kubeletcni.DefaultNetDir
		server.CNIBinDir = kubeletcni.DefaultCNIDir
		server.HairpinMode = kubeletconfig.HairpinNone
	}

	return server, nil
}

func ToFlags(config *kubeletoptions.KubeletServer) []string {
	server, _ := kubeletoptions.NewKubeletServer()
	return cmdflags.AsArgs(config.AddFlags, server.AddFlags)
}
<commit_msg>interesting: restore ability to start with swap on by default<commit_after>package node

import (
	"fmt"
	"net"
	"strconv"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kerrors "k8s.io/apimachinery/pkg/util/errors"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	kubeletoptions "k8s.io/kubernetes/cmd/kubelet/app/options"
	"k8s.io/kubernetes/pkg/features"
	"k8s.io/kubernetes/pkg/kubelet/apis/kubeletconfig"
	kubeletcni "k8s.io/kubernetes/pkg/kubelet/network/cni"
	kubelettypes "k8s.io/kubernetes/pkg/kubelet/types"

	configapi "github.com/openshift/origin/pkg/cmd/server/api"
	cmdutil "github.com/openshift/origin/pkg/cmd/util"
	cmdflags "github.com/openshift/origin/pkg/cmd/util/flags"
	"github.com/openshift/origin/pkg/cmd/util/variable"
	"github.com/openshift/origin/pkg/network"
)

// Build creates the core Kubernetes component configs for a given NodeConfig, or returns
// an error
func Build(options configapi.NodeConfig) (*kubeletoptions.KubeletServer, error) {
	imageTemplate := variable.NewDefaultImageTemplate()
	imageTemplate.Format = options.ImageConfig.Format
	imageTemplate.Latest = options.ImageConfig.Latest

	var path string
	var fileCheckInterval int64
	if options.PodManifestConfig != nil {
		path = options.PodManifestConfig.Path
		fileCheckInterval = options.PodManifestConfig.FileCheckIntervalSeconds
	}

	kubeAddressStr, kubePortStr, err := net.SplitHostPort(options.ServingInfo.BindAddress)
	if err != nil {
		return nil, fmt.Errorf("cannot parse node address: %v", err)
	}
	kubePort, err := strconv.Atoi(kubePortStr)
	if err != nil {
		return nil, fmt.Errorf("cannot parse node port: %v", err)
	}

	// Defaults are tested in TestKubeletDefaults
	server, err := kubeletoptions.NewKubeletServer()
	if err != nil {
		return nil, fmt.Errorf("cannot create kubelet server: %v", err)
	}
	// Adjust defaults
	server.RequireKubeConfig = true
	server.KubeConfig.Default(options.MasterKubeConfig)
	server.PodManifestPath = path
	server.RootDirectory = options.VolumeDirectory
	server.NodeIP = options.NodeIP
	server.HostnameOverride = options.NodeName
	server.AllowPrivileged = true
	server.RegisterNode = true
	server.Address = kubeAddressStr
	server.Port = int32(kubePort)
	server.ReadOnlyPort = 0        // no read only access
	server.CAdvisorPort = 0        // no unsecured cadvisor access
	server.HealthzPort = 0         // no unsecured healthz access
	server.HealthzBindAddress = "" // no unsecured healthz access
	server.ClusterDNS = []string{options.DNSIP}
	server.ClusterDomain = options.DNSDomain
	server.NetworkPluginName = options.NetworkConfig.NetworkPluginName
	server.HostNetworkSources = []string{kubelettypes.ApiserverSource, kubelettypes.FileSource}
	server.HostPIDSources = []string{kubelettypes.ApiserverSource, kubelettypes.FileSource}
	server.HostIPCSources = []string{kubelettypes.ApiserverSource, kubelettypes.FileSource}
	server.HTTPCheckFrequency = metav1.Duration{Duration: time.Duration(0)} // no remote HTTP pod creation access
	server.FileCheckFrequency = metav1.Duration{Duration: time.Duration(fileCheckInterval) * time.Second}
	server.KubeletFlags.ContainerRuntimeOptions.PodSandboxImage = imageTemplate.ExpandOrDie("pod")
	server.MaxPods = 250
	server.PodsPerCore = 10
	server.CgroupDriver = "systemd"
	server.DockerExecHandlerName = string(options.DockerConfig.ExecHandlerName)
	server.RemoteRuntimeEndpoint = options.DockerConfig.DockerShimSocket
	server.RemoteImageEndpoint = options.DockerConfig.DockerShimSocket
	server.DockershimRootDirectory = options.DockerConfig.DockershimRootDirectory

	// TODO: check/warn/fail in setup instead?
	// allows kubelet to continue to start in swap environments
	server.FailSwapOn = false

	// prevents kube from generating certs
	server.TLSCertFile = options.ServingInfo.ServerCert.CertFile
	server.TLSPrivateKeyFile = options.ServingInfo.ServerCert.KeyFile

	containerized := cmdutil.Env("OPENSHIFT_CONTAINERIZED", "") == "true"
	server.Containerized = containerized

	// force the authentication and authorization
	// Setup auth
	authnTTL, err := time.ParseDuration(options.AuthConfig.AuthenticationCacheTTL)
	if err != nil {
		return nil, err
	}
	server.Authentication = kubeletconfig.KubeletAuthentication{
		X509: kubeletconfig.KubeletX509Authentication{
			ClientCAFile: options.ServingInfo.ClientCA,
		},
		Webhook: kubeletconfig.KubeletWebhookAuthentication{
			Enabled:  true,
			CacheTTL: metav1.Duration{Duration: authnTTL},
		},
		Anonymous: kubeletconfig.KubeletAnonymousAuthentication{
			Enabled: true,
		},
	}
	authzTTL, err := time.ParseDuration(options.AuthConfig.AuthorizationCacheTTL)
	if err != nil {
		return nil, err
	}
	server.Authorization = kubeletconfig.KubeletAuthorization{
		Mode: kubeletconfig.KubeletAuthorizationModeWebhook,
		Webhook: kubeletconfig.KubeletWebhookAuthorization{
			CacheAuthorizedTTL:   metav1.Duration{Duration: authzTTL},
			CacheUnauthorizedTTL: metav1.Duration{Duration: authzTTL},
		},
	}

	// resolve extended arguments
	// TODO: this should be done in config validation (along with the above) so we can provide
	// proper errors
	if err := cmdflags.Resolve(options.KubeletArguments, server.AddFlags); len(err) > 0 {
		return nil, kerrors.NewAggregate(err)
	}

	// terminate early if feature gate is incorrect on the node
	if len(server.FeatureGates) > 0 {
		if err := utilfeature.DefaultFeatureGate.Set(server.FeatureGates); err != nil {
			return nil, err
		}
	}
	if utilfeature.DefaultFeatureGate.Enabled(features.RotateKubeletServerCertificate) {
		// Server cert rotation is ineffective if a cert is hardcoded.
		if len(server.CertDirectory) > 0 {
			server.TLSCertFile = ""
			server.TLSPrivateKeyFile = ""
		}
	}

	if network.IsOpenShiftNetworkPlugin(options.NetworkConfig.NetworkPluginName) {
		// SDN plugin pod setup/teardown is implemented as a CNI plugin
		server.NetworkPluginName = kubeletcni.CNIPluginName
		server.NetworkPluginDir = kubeletcni.DefaultNetDir
		server.CNIConfDir = kubeletcni.DefaultNetDir
		server.CNIBinDir = kubeletcni.DefaultCNIDir
		server.HairpinMode = kubeletconfig.HairpinNone
	}

	return server, nil
}

func ToFlags(config *kubeletoptions.KubeletServer) []string {
	server, _ := kubeletoptions.NewKubeletServer()
	return cmdflags.AsArgs(config.AddFlags, server.AddFlags)
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package device_manager

import (
	"fmt"
	"net"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/grpc"
	pluginapi "k8s.io/kubernetes/pkg/kubelet/apis/deviceplugin/v1alpha"

	"github.com/fsnotify/fsnotify"

	"kubevirt.io/kubevirt/pkg/log"
)

const (
	DeviceNamespace   = "devices.kubevirt.io"
	connectionTimeout = 5 * time.Second
)

type GenericDevicePlugin struct {
	counter    int
	devs       []*pluginapi.Device
	server     *grpc.Server
	socketPath string
	stop       chan struct{}
	health     chan string
	devicePath string
	deviceName string
}

func NewGenericDevicePlugin(deviceName string, devicePath string, maxDevices int) *GenericDevicePlugin {
	serverSock := fmt.Sprintf(pluginapi.DevicePluginPath+"kubevirt-%s.sock", deviceName)
	dpi := &GenericDevicePlugin{
		counter:    0,
		devs:       []*pluginapi.Device{},
		socketPath: serverSock,
		health:     make(chan string),
		deviceName: deviceName,
		devicePath: devicePath,
	}
	for i := 0; i < maxDevices; i++ {
		dpi.addNewGenericDevice()
	}

	return dpi
}

func waitForGrpcServer(socketPath string, timeout time.Duration) error {
	conn, err := connect(socketPath, timeout)
	if err != nil {
		return err
	}
	conn.Close()
	return nil
}

// dial establishes the gRPC communication with the registered device plugin.
func connect(socketPath string, timeout time.Duration) (*grpc.ClientConn, error) {
	c, err := grpc.Dial(socketPath,
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(timeout),
		grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
			return net.DialTimeout("unix", addr, timeout)
		}),
	)

	if err != nil {
		return nil, err
	}

	return c, nil
}

// Start starts the gRPC server of the device plugin
func (dpi *GenericDevicePlugin) Start(stop chan struct{}) error {
	if dpi.server != nil {
		return fmt.Errorf("gRPC server already started")
	}

	dpi.stop = stop

	err := dpi.cleanup()
	if err != nil {
		return err
	}

	sock, err := net.Listen("unix", dpi.socketPath)
	if err != nil {
		return err
	}

	dpi.server = grpc.NewServer([]grpc.ServerOption{}...)
	pluginapi.RegisterDevicePluginServer(dpi.server, dpi)

	dpi.Register()

	go dpi.server.Serve(sock)

	err = waitForGrpcServer(dpi.socketPath, connectionTimeout)

	go dpi.healthCheck()

	return err
}

// Stop stops the gRPC server
func (dpi *GenericDevicePlugin) Stop() error {
	if dpi.server == nil {
		return nil
	}

	dpi.server.Stop()
	dpi.server = nil

	return dpi.cleanup()
}

// Register registers the device plugin for the given resourceName with Kubelet.
func (dpi *GenericDevicePlugin) Register() error {
	conn, err := connect(pluginapi.KubeletSocket, connectionTimeout)
	if err != nil {
		return err
	}
	defer conn.Close()

	client := pluginapi.NewRegistrationClient(conn)
	reqt := &pluginapi.RegisterRequest{
		Version:      pluginapi.Version,
		Endpoint:     path.Base(dpi.socketPath),
		ResourceName: fmt.Sprintf("%s/%s", DeviceNamespace, dpi.deviceName),
	}

	_, err = client.Register(context.Background(), reqt)
	if err != nil {
		return err
	}
	return nil
}

func (dpi *GenericDevicePlugin) addNewGenericDevice() {
	deviceId := dpi.deviceName + strconv.Itoa(dpi.counter)
	dpi.devs = append(dpi.devs, &pluginapi.Device{
		ID:     deviceId,
		Health: pluginapi.Healthy,
	})

	dpi.counter += 1
}

func (dpi *GenericDevicePlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin_ListAndWatchServer) error {
	s.Send(&pluginapi.ListAndWatchResponse{Devices: dpi.devs})

	for {
		select {
		case health := <-dpi.health:
			// There's only one shared generic device
			// so update each plugin device to reflect overall device health
			for _, dev := range dpi.devs {
				dev.Health = health
			}
			s.Send(&pluginapi.ListAndWatchResponse{Devices: dpi.devs})
		case <-dpi.stop:
			return nil
		}
	}
}

func (dpi *GenericDevicePlugin) Allocate(ctx context.Context, r *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	var response pluginapi.AllocateResponse
	dev := new(pluginapi.DeviceSpec)
	dev.HostPath = dpi.devicePath
	dev.ContainerPath = dpi.devicePath
	dev.Permissions = "rw"
	response.Devices = append(response.Devices, dev)

	return &response, nil
}

func (dpi *GenericDevicePlugin) cleanup() error {
	if err := os.Remove(dpi.socketPath); err != nil && !os.IsNotExist(err) {
		return err
	}

	return nil
}

func (dpi *GenericDevicePlugin) healthCheck() error {
	logger := log.DefaultLogger()
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		logger.Errorf("Unable to create fsnotify watcher: %v", err)
		return nil
	}
	defer watcher.Close()

	healthy := pluginapi.Healthy
	_, err = os.Stat(dpi.devicePath)
	if err != nil {
		if !os.IsNotExist(err) {
			logger.Errorf("Unable to stat device: %v", err)
			return err
		}
		healthy = pluginapi.Unhealthy
	}

	dirName := filepath.Dir(dpi.devicePath)

	err = watcher.Add(dirName)
	if err != nil {
		logger.Errorf("Unable to add path to fsnotify watcher: %v", err)
		return err
	}

	for {
		select {
		case <-dpi.stop:
			return nil
		case event := <-watcher.Events:
			logger.Infof("health Event: %v", event)
			logger.Infof("health Event Name: %s", event.Name)
			if event.Name == dpi.devicePath {
				// Health in this case is if the device path actually exists
				if event.Op == fsnotify.Create {
					healthy = pluginapi.Healthy
				} else if (event.Op == fsnotify.Remove) || (event.Op == fsnotify.Rename) {
					healthy = pluginapi.Unhealthy
				}
				dpi.health <- healthy
			}
		}
	}
}
<commit_msg>Adapt generic device plugin to be compatible with v1beta1 api<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package device_manager

import (
	"fmt"
	"net"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/grpc"
	pluginapi "k8s.io/kubernetes/pkg/kubelet/apis/deviceplugin/v1beta1"

	"github.com/fsnotify/fsnotify"

	"kubevirt.io/kubevirt/pkg/log"
)

const (
	DeviceNamespace   = "devices.kubevirt.io"
	connectionTimeout = 5 * time.Second
)

type GenericDevicePlugin struct {
	counter    int
	devs       []*pluginapi.Device
	server     *grpc.Server
	socketPath string
	stop       chan struct{}
	health     chan string
	devicePath string
	deviceName string
}

func NewGenericDevicePlugin(deviceName string, devicePath string, maxDevices int) *GenericDevicePlugin {
	serverSock := fmt.Sprintf(pluginapi.DevicePluginPath+"kubevirt-%s.sock", deviceName)
	dpi := &GenericDevicePlugin{
		counter:    0,
		devs:       []*pluginapi.Device{},
		socketPath: serverSock,
		health:     make(chan string),
		deviceName: deviceName,
		devicePath: devicePath,
	}
	for i := 0; i < maxDevices; i++ {
		dpi.addNewGenericDevice()
	}

	return dpi
}

func waitForGrpcServer(socketPath string, timeout time.Duration) error {
	conn, err := connect(socketPath, timeout)
	if err != nil {
		return err
	}
	conn.Close()
	return nil
}

// dial establishes the gRPC communication with the registered device plugin.
func connect(socketPath string, timeout time.Duration) (*grpc.ClientConn, error) {
	c, err := grpc.Dial(socketPath,
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(timeout),
		grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
			return net.DialTimeout("unix", addr, timeout)
		}),
	)

	if err != nil {
		return nil, err
	}

	return c, nil
}

// Start starts the gRPC server of the device plugin
func (dpi *GenericDevicePlugin) Start(stop chan struct{}) error {
	if dpi.server != nil {
		return fmt.Errorf("gRPC server already started")
	}

	dpi.stop = stop

	err := dpi.cleanup()
	if err != nil {
		return err
	}

	sock, err := net.Listen("unix", dpi.socketPath)
	if err != nil {
		return err
	}

	dpi.server = grpc.NewServer([]grpc.ServerOption{}...)
	pluginapi.RegisterDevicePluginServer(dpi.server, dpi)

	dpi.Register()

	go dpi.server.Serve(sock)

	err = waitForGrpcServer(dpi.socketPath, connectionTimeout)

	go dpi.healthCheck()

	return err
}

// Stop stops the gRPC server
func (dpi *GenericDevicePlugin) Stop() error {
	if dpi.server == nil {
		return nil
	}

	dpi.server.Stop()
	dpi.server = nil

	return dpi.cleanup()
}

// Register registers the device plugin for the given resourceName with Kubelet.
func (dpi *GenericDevicePlugin) Register() error {
	conn, err := connect(pluginapi.KubeletSocket, connectionTimeout)
	if err != nil {
		return err
	}
	defer conn.Close()

	client := pluginapi.NewRegistrationClient(conn)
	reqt := &pluginapi.RegisterRequest{
		Version:      pluginapi.Version,
		Endpoint:     path.Base(dpi.socketPath),
		ResourceName: fmt.Sprintf("%s/%s", DeviceNamespace, dpi.deviceName),
	}

	_, err = client.Register(context.Background(), reqt)
	if err != nil {
		return err
	}
	return nil
}

func (dpi *GenericDevicePlugin) addNewGenericDevice() {
	deviceId := dpi.deviceName + strconv.Itoa(dpi.counter)
	dpi.devs = append(dpi.devs, &pluginapi.Device{
		ID:     deviceId,
		Health: pluginapi.Healthy,
	})

	dpi.counter += 1
}

func (dpi *GenericDevicePlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin_ListAndWatchServer) error {
	s.Send(&pluginapi.ListAndWatchResponse{Devices: dpi.devs})

	for {
		select {
		case health := <-dpi.health:
			// There's only one shared generic device
			// so update each plugin device to reflect overall device health
			for _, dev := range dpi.devs {
				dev.Health = health
			}
			s.Send(&pluginapi.ListAndWatchResponse{Devices: dpi.devs})
		case <-dpi.stop:
			return nil
		}
	}
}

func (dpi *GenericDevicePlugin) Allocate(ctx context.Context, r *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
	response := pluginapi.AllocateResponse{}
	containerResponse := new(pluginapi.ContainerAllocateResponse)

	dev := new(pluginapi.DeviceSpec)
	dev.HostPath = dpi.devicePath
	dev.ContainerPath = dpi.devicePath
	dev.Permissions = "rw"
	containerResponse.Devices = []*pluginapi.DeviceSpec{dev}

	response.ContainerResponses = []*pluginapi.ContainerAllocateResponse{containerResponse}

	return &response, nil
}

func (dpi *GenericDevicePlugin) cleanup() error {
	if err := os.Remove(dpi.socketPath); err != nil && !os.IsNotExist(err) {
		return err
	}

	return nil
}

func (dpi *GenericDevicePlugin) GetDevicePluginOptions(ctx context.Context, e *pluginapi.Empty) (*pluginapi.DevicePluginOptions, error) {
	options := &pluginapi.DevicePluginOptions{
		PreStartRequired: false,
	}
	return options, nil
}

func (dpi *GenericDevicePlugin) PreStartContainer(ctx context.Context, in *pluginapi.PreStartContainerRequest) (*pluginapi.PreStartContainerResponse, error) {
	res := &pluginapi.PreStartContainerResponse{}
	return res, nil
}

func (dpi *GenericDevicePlugin) healthCheck() error {
	logger := log.DefaultLogger()
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		logger.Errorf("Unable to create fsnotify watcher: %v", err)
		return nil
	}
	defer watcher.Close()

	healthy := pluginapi.Healthy
	_, err = os.Stat(dpi.devicePath)
	if err != nil {
		if !os.IsNotExist(err) {
			logger.Errorf("Unable to stat device: %v", err)
			return err
		}
		healthy = pluginapi.Unhealthy
	}

	dirName := filepath.Dir(dpi.devicePath)

	err = watcher.Add(dirName)
	if err != nil {
		logger.Errorf("Unable to add path to fsnotify watcher: %v", err)
		return err
	}

	for {
		select {
		case <-dpi.stop:
			return nil
		case event := <-watcher.Events:
			logger.Infof("health Event: %v", event)
			logger.Infof("health Event Name: %s", event.Name)
			if event.Name == dpi.devicePath {
				// Health in this case is if the device path actually exists
				if event.Op == fsnotify.Create {
					healthy = pluginapi.Healthy
				} else if (event.Op == fsnotify.Remove) || (event.Op == fsnotify.Rename) {
					healthy = pluginapi.Unhealthy
				}
				dpi.health <- healthy
			}
		}
	}
}
<|endoftext|>


<commit_before>package seq

import (
	"encoding/json"
	"fmt"
	"math"
	"strconv"
)

// HMM states in the Plan7 architecture.
const (
	Match HMMState = iota
	Deletion
	Insertion
	Begin
	End
)

type HMMState int

type HMM struct {
	// An ordered list of HMM nodes.
	Nodes []HMMNode

	// The alphabet as defined by an ordering of residues.
	// Indices in this slice correspond to indices in match/insertion emissions.
	Alphabet Alphabet

	// NULL model. (Amino acid background frequencies.)
	// HMMER hmm files don't have this, but HHsuite hhm files do.
	// In the case of HHsuite, the NULL model is used for insertion emissions
	// in every node.
	Null EProbs
}

type DynamicTable struct {
	scores []Prob
	nodes  int
}

// allocTable returns a freshly allocated dynamic programming table for use
// in HMM alogirthms like Viterbi. It is indexed by HMM state, node index and
// sequence length, in that order. The total size of the table is equal to
// (#states * (#nodes + 1) * (seqLen + 1)).
//
// The only states used are Match, Deletion and Insertion.
//
// Each value is initialized to a minimum probability.
func AllocTable(numNodes int, seqLen int) *DynamicTable {
	nodes := numNodes + 1
	t := &DynamicTable{
		scores: make([]Prob, 3*nodes*(seqLen+1)),
		nodes:  nodes,
	}
	t.reset()
	return t
}

func (t *DynamicTable) index(state HMMState, node int, obs int) int {
	return int(state) + 3*(node+t.nodes*obs)
}

func (t *DynamicTable) set(state HMMState, node int, obs int, p Prob) {
	i := t.index(state, node, obs)
	if t.scores[i].Less(p) {
		t.scores[i] = p
	}
}

func (t *DynamicTable) reset() {
	for i := 0; i < len(t.scores); i++ {
		t.scores[i] = MinProb
	}
}

func (hmm *HMM) ViterbiScore(seq Sequence) Prob {
	table := AllocTable(len(hmm.Nodes), seq.Len())
	return hmm.ViterbiScoreMem(seq, table)
}

func (hmm *HMM) ViterbiScoreMem(seq Sequence, table *DynamicTable) Prob {
	table.reset()
	table.scores[table.index(Match, 0, 0)] = Prob(0.0) // The begin node.

	var trans TProbs
	var residue Residue
	var memit, iemit, here Prob
	for node := 0; node < len(hmm.Nodes); node++ {
		for obs := 0; obs < seq.Len(); obs++ {
			trans = hmm.Nodes[node].Transitions
			residue = seq.Residues[obs]
			iemit = hmm.Nodes[node].InsEmit.Lookup(residue)
			if node+1 < len(hmm.Nodes) {
				memit = hmm.Nodes[node+1].MatEmit.Lookup(residue)
			} else {
				memit = 0.0 // Force into match state for end node.
			}

			here = table.scores[table.index(Match, node, obs)]
			table.set(Insertion, node, obs+1, here+trans.MI+iemit)
			table.set(Match, node+1, obs+1, here+trans.MM+memit)
			table.set(Deletion, node+1, obs, here+trans.MD)

			here = table.scores[table.index(Insertion, node, obs)]
			table.set(Insertion, node, obs+1, here+trans.II+iemit)
			table.set(Match, node+1, obs+1, here+trans.IM+memit)

			here = table.scores[table.index(Deletion, node, obs)]
			table.set(Match, node+1, obs+1, here+trans.DM+memit)
			table.set(Deletion, node+1, obs, here+trans.DD)
		}
	}
	return table.scores[table.index(Match, len(hmm.Nodes), seq.Len())]
}

type HMMNode struct {
	Residue             Residue
	NodeNum             int
	InsEmit             EProbs
	MatEmit             EProbs
	Transitions         TProbs
	NeffM, NeffI, NeffD Prob
}

// EProbs represents emission probabilities, as log-odds scores.
// The representation of EProbs should not be used by clients; it is exported
// only for convenience with marshalling APIs.
type EProbs struct {
	Offset Residue
	Probs  []Prob
}

// NewEProbs creates a new EProbs map from the given alphabet. Keys of the map
// are residues defined in the alphabet, and values are defaulted to the
// minimum probability.
func NewEProbs(alphabet Alphabet) EProbs {
	offset, max := Residue(255), Residue(0)
	for _, r := range alphabet {
		if r < offset {
			offset = r
		}
		if r > max {
			max = r
		}
	}
	probs := make([]Prob, 1+max-offset)
	for i := range probs {
		probs[i] = MinProb
	}
	return EProbs{offset, probs}
}

// Returns the emission probability for a particular residue.
func (ep EProbs) Lookup(r Residue) Prob {
	i := r - ep.Offset
	if i < 0 || int(i) >= len(ep.Probs) {
		return MinProb
	}
	return ep.Probs[r-ep.Offset]
}

// Set sets the emission probability of the given residue.
func (ep *EProbs) Set(r Residue, p Prob) {
	ep.Probs[r-ep.Offset] = p
}

// func (ep *EProbs) MarshalJSON() ([]byte, error) {
// strmap := make(map[string]Prob, len(*ep))
// for k, v := range *ep {
// strmap[string(k)] = v
// }
// return json.Marshal(strmap)
// }
//
// func (ep *EProbs) UnmarshalJSON(bs []byte) error {
// var strmap map[string]Prob
// if err := json.Unmarshal(bs, &strmap); err != nil {
// return err
// }
// if *ep == nil {
// *ep = make(EProbs, len(strmap))
// }
// for k, v := range strmap {
// (*ep)[Residue(k[0])] = v
// }
// return nil
// }

// TProbs represents transition probabilities, as log-odds scores.
// Note that ID and DI are omitted (Plan7).
type TProbs struct {
	MM, MI, MD, IM, II, DM, DD Prob
}

// Prob represents a transition or emission probability.
type Prob float64

var invalidProb = Prob(math.NaN())

// The value representing a minimum emission/transition probability.
// Remember, max in negative log space is minimum probability.
var MinProb = Prob(math.MaxFloat64)

// NewProb creates a new probability value from a string (usually read from
// an hmm or hhm file). If the string is equivalent to the special value "*",
// then the probability returned is guaranteed to be minimal. Otherwise, the
// string is parsed as a float, and an error returned if parsing fails.
func NewProb(fstr string) (Prob, error) {
	if fstr == "*" {
		return MinProb, nil
	}

	f, err := strconv.ParseFloat(fstr, 64)
	if err != nil {
		return invalidProb,
			fmt.Errorf("Could not convert '%s' to a log probability: %s",
				fstr, err)
	}
	return Prob(f), nil
}

// Less returns true if `p1` represents a smaller probability than `p2`.
func (p1 Prob) Less(p2 Prob) bool {
	return p1 > p2
}

// IsMin returns true if the probability is minimal.
func (p Prob) IsMin() bool {
	return p == MinProb
}

// Ratio returns the log probability as a ratio in the range [0, 1].
// (This assumes that `p` is a log-odds score.)
func (p Prob) Ratio() float64 {
	if p.IsMin() {
		return 0.0
	}
	return math.Exp(-(float64(p)))
}

// Distance returns the distance between `p1` and `p2`.
func (p1 Prob) Distance(p2 Prob) float64 {
	return math.Abs(float64(p1) - float64(p2))
}

// String returns a string representation of the probability.
// When `p` is the minimum probability, then "*" is used.
// Otherwise, the full number is written.
func (p Prob) String() string {
	if p.IsMin() {
		return "*"
	}
	return fmt.Sprintf("%v", float64(p))
}

func (p Prob) MarshalJSON() ([]byte, error) {
	return json.Marshal(p.String())
}

func (p *Prob) UnmarshalJSON(bs []byte) error {
	var str string
	var err error
	if err := json.Unmarshal(bs, &str); err != nil {
		return err
	}
	if *p, err = NewProb(str); err != nil {
		return err
	}
	return nil
}

// NewHMM creates a new HMM from a list of nodes, an ordered alphabet and a
// set of null probabilities (which may be nil).
func NewHMM(nodes []HMMNode, alphabet []Residue, null EProbs) *HMM {
	return &HMM{
		Nodes:    nodes,
		Alphabet: alphabet,
		Null:     null,
	}
}

// Slice returns a slice of the HMM given. A slice of an HMM returns only the
// HMM nodes (i.e., columns or match/delete states) in the region specified
// by the slice. Also, the transition probabilities of the last state are
// specially set: M->M = 0, M->I = *, M->D = *, I->M = 0, I->I = *, D->M = 0,
// and D->D = *.
// No other modifications are made.
func (hmm *HMM) Slice(start, end int) *HMM {
	nodes := make([]HMMNode, end-start)
	copy(nodes, hmm.Nodes[start:end])
	last := len(nodes) - 1
	nodes[last].Transitions.MM = 0
	nodes[last].Transitions.MI = MinProb
	nodes[last].Transitions.MD = MinProb
	nodes[last].Transitions.IM = 0
	nodes[last].Transitions.II = MinProb
	nodes[last].Transitions.DM = 0
	nodes[last].Transitions.DD = MinProb

	return &HMM{
		Nodes:    nodes,
		Alphabet: hmm.Alphabet,
		Null:     hmm.Null,
	}
}
<commit_msg>Don't need this any more since we changed the representation of EProbs.<commit_after>package seq

import (
	"encoding/json"
	"fmt"
	"math"
	"strconv"
)

// HMM states in the Plan7 architecture.
const (
	Match HMMState = iota
	Deletion
	Insertion
	Begin
	End
)

type HMMState int

type HMM struct {
	// An ordered list of HMM nodes.
	Nodes []HMMNode

	// The alphabet as defined by an ordering of residues.
	// Indices in this slice correspond to indices in match/insertion emissions.
	Alphabet Alphabet

	// NULL model. (Amino acid background frequencies.)
	// HMMER hmm files don't have this, but HHsuite hhm files do.
	// In the case of HHsuite, the NULL model is used for insertion emissions
	// in every node.
	Null EProbs
}

type DynamicTable struct {
	scores []Prob
	nodes  int
}

// allocTable returns a freshly allocated dynamic programming table for use
// in HMM alogirthms like Viterbi. It is indexed by HMM state, node index and
// sequence length, in that order. The total size of the table is equal to
// (#states * (#nodes + 1) * (seqLen + 1)).
//
// The only states used are Match, Deletion and Insertion.
//
// Each value is initialized to a minimum probability.
func AllocTable(numNodes int, seqLen int) *DynamicTable {
	nodes := numNodes + 1
	t := &DynamicTable{
		scores: make([]Prob, 3*nodes*(seqLen+1)),
		nodes:  nodes,
	}
	t.reset()
	return t
}

func (t *DynamicTable) index(state HMMState, node int, obs int) int {
	return int(state) + 3*(node+t.nodes*obs)
}

func (t *DynamicTable) set(state HMMState, node int, obs int, p Prob) {
	i := t.index(state, node, obs)
	if t.scores[i].Less(p) {
		t.scores[i] = p
	}
}

func (t *DynamicTable) reset() {
	for i := 0; i < len(t.scores); i++ {
		t.scores[i] = MinProb
	}
}

func (hmm *HMM) ViterbiScore(seq Sequence) Prob {
	table := AllocTable(len(hmm.Nodes), seq.Len())
	return hmm.ViterbiScoreMem(seq, table)
}

func (hmm *HMM) ViterbiScoreMem(seq Sequence, table *DynamicTable) Prob {
	table.reset()
	table.scores[table.index(Match, 0, 0)] = Prob(0.0) // The begin node.

	var trans TProbs
	var residue Residue
	var memit, iemit, here Prob
	for node := 0; node < len(hmm.Nodes); node++ {
		for obs := 0; obs < seq.Len(); obs++ {
			trans = hmm.Nodes[node].Transitions
			residue = seq.Residues[obs]
			iemit = hmm.Nodes[node].InsEmit.Lookup(residue)
			if node+1 < len(hmm.Nodes) {
				memit = hmm.Nodes[node+1].MatEmit.Lookup(residue)
			} else {
				memit = 0.0 // Force into match state for end node.
			}

			here = table.scores[table.index(Match, node, obs)]
			table.set(Insertion, node, obs+1, here+trans.MI+iemit)
			table.set(Match, node+1, obs+1, here+trans.MM+memit)
			table.set(Deletion, node+1, obs, here+trans.MD)

			here = table.scores[table.index(Insertion, node, obs)]
			table.set(Insertion, node, obs+1, here+trans.II+iemit)
			table.set(Match, node+1, obs+1, here+trans.IM+memit)

			here = table.scores[table.index(Deletion, node, obs)]
			table.set(Match, node+1, obs+1, here+trans.DM+memit)
			table.set(Deletion, node+1, obs, here+trans.DD)
		}
	}
	return table.scores[table.index(Match, len(hmm.Nodes), seq.Len())]
}

type HMMNode struct {
	Residue             Residue
	NodeNum             int
	InsEmit             EProbs
	MatEmit             EProbs
	Transitions         TProbs
	NeffM, NeffI, NeffD Prob
}

// EProbs represents emission probabilities, as log-odds scores.
// The representation of EProbs should not be used by clients; it is exported
// only for convenience with marshalling APIs.
type EProbs struct {
	Offset Residue
	Probs  []Prob
}

// NewEProbs creates a new EProbs map from the given alphabet. Keys of the map
// are residues defined in the alphabet, and values are defaulted to the
// minimum probability.
func NewEProbs(alphabet Alphabet) EProbs {
	offset, max := Residue(255), Residue(0)
	for _, r := range alphabet {
		if r < offset {
			offset = r
		}
		if r > max {
			max = r
		}
	}
	probs := make([]Prob, 1+max-offset)
	for i := range probs {
		probs[i] = MinProb
	}
	return EProbs{offset, probs}
}

// Returns the emission probability for a particular residue.
func (ep EProbs) Lookup(r Residue) Prob {
	i := r - ep.Offset
	if i < 0 || int(i) >= len(ep.Probs) {
		return MinProb
	}
	return ep.Probs[r-ep.Offset]
}

// Set sets the emission probability of the given residue.
func (ep *EProbs) Set(r Residue, p Prob) {
	ep.Probs[r-ep.Offset] = p
}

// TProbs represents transition probabilities, as log-odds scores.
// Note that ID and DI are omitted (Plan7).
type TProbs struct {
	MM, MI, MD, IM, II, DM, DD Prob
}

// Prob represents a transition or emission probability.
type Prob float64

var invalidProb = Prob(math.NaN())

// The value representing a minimum emission/transition probability.
// Remember, max in negative log space is minimum probability.
var MinProb = Prob(math.MaxFloat64)

// NewProb creates a new probability value from a string (usually read from
// an hmm or hhm file). If the string is equivalent to the special value "*",
// then the probability returned is guaranteed to be minimal. Otherwise, the
// string is parsed as a float, and an error returned if parsing fails.
func NewProb(fstr string) (Prob, error) {
	if fstr == "*" {
		return MinProb, nil
	}

	f, err := strconv.ParseFloat(fstr, 64)
	if err != nil {
		return invalidProb,
			fmt.Errorf("Could not convert '%s' to a log probability: %s",
				fstr, err)
	}
	return Prob(f), nil
}

// Less returns true if `p1` represents a smaller probability than `p2`.
func (p1 Prob) Less(p2 Prob) bool {
	return p1 > p2
}

// IsMin returns true if the probability is minimal.
func (p Prob) IsMin() bool {
	return p == MinProb
}

// Ratio returns the log probability as a ratio in the range [0, 1].
// (This assumes that `p` is a log-odds score.)
func (p Prob) Ratio() float64 {
	if p.IsMin() {
		return 0.0
	}
	return math.Exp(-(float64(p)))
}

// Distance returns the distance between `p1` and `p2`.
func (p1 Prob) Distance(p2 Prob) float64 {
	return math.Abs(float64(p1) - float64(p2))
}

// String returns a string representation of the probability.
// When `p` is the minimum probability, then "*" is used.
// Otherwise, the full number is written.
func (p Prob) String() string {
	if p.IsMin() {
		return "*"
	}
	return fmt.Sprintf("%v", float64(p))
}

func (p Prob) MarshalJSON() ([]byte, error) {
	return json.Marshal(p.String())
}

func (p *Prob) UnmarshalJSON(bs []byte) error {
	var str string
	var err error
	if err := json.Unmarshal(bs, &str); err != nil {
		return err
	}
	if *p, err = NewProb(str); err != nil {
		return err
	}
	return nil
}

// NewHMM creates a new HMM from a list of nodes, an ordered alphabet and a
// set of null probabilities (which may be nil).
func NewHMM(nodes []HMMNode, alphabet []Residue, null EProbs) *HMM {
	return &HMM{
		Nodes:    nodes,
		Alphabet: alphabet,
		Null:     null,
	}
}

// Slice returns a slice of the HMM given. A slice of an HMM returns only the
// HMM nodes (i.e., columns or match/delete states) in the region specified
// by the slice. Also, the transition probabilities of the last state are
// specially set: M->M = 0, M->I = *, M->D = *, I->M = 0, I->I = *, D->M = 0,
// and D->D = *.
// No other modifications are made.
func (hmm *HMM) Slice(start, end int) *HMM {
	nodes := make([]HMMNode, end-start)
	copy(nodes, hmm.Nodes[start:end])
	last := len(nodes) - 1
	nodes[last].Transitions.MM = 0
	nodes[last].Transitions.MI = MinProb
	nodes[last].Transitions.MD = MinProb
	nodes[last].Transitions.IM = 0
	nodes[last].Transitions.II = MinProb
	nodes[last].Transitions.DM = 0
	nodes[last].Transitions.DD = MinProb

	return &HMM{
		Nodes:    nodes,
		Alphabet: hmm.Alphabet,
		Null:     hmm.Null,
	}
}
<|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.

   @author: Juliano Martinez
*/

package main

import (
	"crypto/tls"
	hpr_http_server "github.com/ncode/hot-potato-router/http_server"
	hpr_utils "github.com/ncode/hot-potato-router/utils"
	"net/http"
	"strconv"
	"time"
)

var (
	cfg = hpr_utils.NewConfig()
)

func main() {
	hpr_utils.Log("Starting Hot Potato Router...")
	probe_interval, _ := strconv.Atoi(cfg.Options["hpr"]["probe_interval"])
	if probe_interval == 0 {
		probe_interval = 10
	}
	s, err := hpr_http_server.NewServer(time.Duration(probe_interval) * time.Second)
	hpr_utils.CheckPanic(err, "Unable to spawn")

	http_fd, _ := strconv.Atoi(cfg.Options["hpr"]["http_fd"])
	https_fd, _ := strconv.Atoi(cfg.Options["hpr"]["https_fd"])
	if https_fd >= 3 || cfg.Options["hpr"]["https_addr"] != "" {
		cert, err := tls.LoadX509KeyPair(cfg.Options["hpr"]["cert_file"], cfg.Options["hpr"]["key_file"])
		hpr_utils.CheckPanic(err, "Unable to load certificate")

		c := &tls.Config{Certificates: []tls.Certificate{cert}}
		l := tls.NewListener(hpr_http_server.Listen(https_fd, cfg.Options["hpr"]["https_addr"]), c)
		go func() {
			hpr_utils.CheckPanic(http.Serve(l, s), "Problem with https server")
		}()
	}
	hpr_utils.CheckPanic(
		http.Serve(hpr_http_server.Listen(http_fd, cfg.Options["hpr"]["http_addr"]), s),
		"Problem with http server")
}
<commit_msg>removing the fd idea and add some comments<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.

   @author: Juliano Martinez
*/

package main

import (
	"crypto/tls"
	hpr_http_server "github.com/ncode/hot-potato-router/http_server"
	hpr_utils "github.com/ncode/hot-potato-router/utils"
	"net/http"
	"strconv"
	"time"
)

var (
	cfg = hpr_utils.NewConfig()
)

func main() {
	hpr_utils.Log("Starting Hot Potato Router...")
	probe_interval, _ := strconv.Atoi(cfg.Options["hpr"]["probe_interval"])
	if probe_interval == 0 {
		probe_interval = 10
	}
	s, err := hpr_http_server.NewServer(time.Duration(probe_interval) * time.Second)
	hpr_utils.CheckPanic(err, "Unable to spawn")

	if cfg.Options["hpr"]["https_addr"] != "" {
		cert, err := tls.LoadX509KeyPair(cfg.Options["hpr"]["cert_file"], cfg.Options["hpr"]["key_file"])
		hpr_utils.CheckPanic(err, "Unable to load certificate")

		c := &tls.Config{Certificates: []tls.Certificate{cert}}
		l := tls.NewListener(hpr_http_server.Listen(https_fd, cfg.Options["hpr"]["https_addr"]), c)
		go func() {
			hpr_utils.CheckPanic(http.Serve(l, s), "Problem with https server")
		}()
	}
	hpr_utils.CheckPanic(
		http.Serve(hpr_http_server.Listen(http_fd, cfg.Options["hpr"]["http_addr"]), s),
		"Problem with http server")
}
<|endoftext|>


<commit_before>package tok

import (
	"expvar"
	"log"
)

var (
	expOnline = expvar.NewInt("tokOnline")
	expUp     = expvar.NewInt("tokUp")
	expDown   = expvar.NewInt("tokDown")
	expEnq    = expvar.NewInt("tokEnq")
	expDeq    = expvar.NewInt("tokDeq")
)

type checkFrame struct {
	uid    interface{}
	chBool chan bool //channel to return online status
}

type downFrame struct {
	uid   interface{}
	ttl   uint32
	data  []byte
	chErr chan error //channel to read send result from
}

type upFrame struct {
	dv   *Device
	data []byte
}

//HubConfig config struct for creating new Hub
type HubConfig struct {
	Actor Actor //Actor implement dispatch logic
	Q     Queue //Message Queue, if nil, message to offline user will not be cached
	Sso   bool  //If it's true, new connection  with same uid will kick off old ones
}

//Hub core of tok, dispatch message between connections
type Hub struct {
	sso           bool
	actor         Actor
	q             Queue
	cons          map[interface{}][]*connection //connection list
	chUp          chan *upFrame
	chDown        chan *downFrame
	chConState    chan *conState
	chReadSignal  chan interface{}
	chKick        chan interface{}
	chQueryOnline chan chan []interface{}
	chCheck       chan *checkFrame
}

func createHub(actor Actor, q Queue, sso bool) *Hub {
	if ReadTimeout > 0 {
		log.Println("[tok] read timeout is enabled, make sure it's greater than your client ping interval. otherwise you'll get read timeout err")
	} else {
		if actor.Ping() == nil {
			log.Fatalln("[tok] both read timeout and server ping have been disabled, server socket resource leak might happen")
		}
	}

	hub := &Hub{
		sso:           sso,
		actor:         actor,
		q:             q,
		cons:          make(map[interface{}][]*connection),
		chUp:          make(chan *upFrame),
		chDown:        make(chan *downFrame),
		chConState:    make(chan *conState),
		chReadSignal:  make(chan interface{}),
		chKick:        make(chan interface{}),
		chQueryOnline: make(chan chan []interface{}),
		chCheck:       make(chan *checkFrame),
	}
	go hub.run()
	return hub
}

func (p *Hub) run() {
	for {

		select {
		case state := <-p.chConState:
			//			log.Printf("connection state change: %v, %v \n", state.online, &state.con)

			if state.online {
				p.goOnline(state.con)
			} else {
				p.goOffline(state.con)
			}
			count := int64(len(p.cons))
			expOnline.Set(count)
		case f := <-p.chUp:
			//			log.Println("up data")
			expUp.Add(1)
			go func() {
				b, err := p.actor.BeforeReceive(f.dv, f.data)
				if err != nil {
					return
				}
				if b == nil {
					b = f.data
				}
				p.actor.OnReceive(f.dv, b)
			}()
		case ff := <-p.chDown:
			if l := p.cons[ff.uid]; len(l) > 0 {
				//online
				go p.down(ff, l)
			} else {
				//offline
				if ff.ttl == 0 {
					ff.chErr <- ErrOffline
					close(ff.chErr)
				} else {
					go p.cache(ff)
				}
			}
		case cf := <-p.chCheck:
			_, ok := p.cons[cf.uid]
			cf.chBool <- ok
			close(cf.chBool)
		case uid := <-p.chReadSignal:
			//only pop msg for online user
			if len(p.cons[uid]) > 0 {
				go p.popMsg(uid)
			}
		case uid := <-p.chKick:
			p.innerKick(uid)
		case chOnline := <-p.chQueryOnline:
			result := make([]interface{}, 0, len(p.cons))
			for uid := range p.cons {
				result = append(result, uid)
			}
			chOnline <- result
			close(chOnline)
		}
	}
}

func (p *Hub) popMsg(uid interface{}) {
	if p.q == nil {
		return
	}

	for {
		b, err := p.q.Deq(uid)
		if err != nil {
			log.Println("deq error", err)
			return
		}
		if len(b) == 0 {
			//no more data in queue
			return
		}
		expDeq.Add(1)
		if err := p.Send(uid, b, 0); err != nil {
			if err := p.q.Enq(uid, b); err != nil {
				log.Println("re-cache err", err, uid)
			}
			return
		}
	}
}

//Send message to someone.
//ttl is expiry seconds. 0 means only send to online user
//If ttl = 0 and user is offline, ErrOffline will be returned.
//If ttl > 0 and user is offline or online but send fail, message will be cached for ttl seconds.
func (p *Hub) Send(to interface{}, b []byte, ttl uint32) error {

	ff := &downFrame{uid: to, data: b, ttl: ttl, chErr: make(chan error)}
	p.chDown <- ff
	err := <-ff.chErr
	if ttl > 0 && err != nil {
		//online send err
		ff.chErr = make(chan error) //create new channel
		go p.cache(ff)
		return <-ff.chErr
	}
	return err
}

//CheckOnline return whether user online or not
func (p *Hub) CheckOnline(uid interface{}) bool {
	cf := &checkFrame{uid: uid, chBool: make(chan bool)}
	p.chCheck <- cf
	return <-cf.chBool
}

//Online query online user list
func (p *Hub) Online() []interface{} {
	ch := make(chan []interface{})
	p.chQueryOnline <- ch
	return <-ch
}

func (p *Hub) cache(ff *downFrame) {
	defer close(ff.chErr)
	expEnq.Add(1)
	if p.q == nil {
		ff.chErr <- ErrQueueRequired
		return
	}

	if err := p.q.Enq(ff.uid, ff.data, ff.ttl); err != nil {
		ff.chErr <- err
	}
}

func (p *Hub) down(f *downFrame, conns []*connection) {
	expDown.Add(1)

	for _, con := range conns {
		b, err := p.actor.BeforeSend(con.dv, f.data)
		if err != nil {
			return
		}
		if b == nil {
			b = f.data
		}

		if con == nil {
			log.Println("hub.down, conn is nil")
			continue
		}
		if err := con.Write(b); err != nil {
			f.chErr <- err
			return
		}
		go p.actor.OnSent(con.dv, f.data)
	}

}

func (p *Hub) goOffline(conn *connection) {
	l := p.cons[conn.uid()]
	rest := connExclude(l, conn)

	//this connection has gotten offline, ignore
	if len(l) == len(rest) {
		return
	}

	if len(rest) == 0 {
		delete(p.cons, conn.uid())
	} else {
		p.cons[conn.uid()] = rest
	}

	go p.close(conn)
}

func (p *Hub) innerKick(uid interface{}) {
	for _, conn := range p.cons[uid] {
		go p.close(conn)
	}
	delete(p.cons, uid)
}

func (p *Hub) byeThenClose(conn *connection) {
	b := p.actor.Bye(conn.dv, "sso")
	if b != nil {
		data, err := p.actor.BeforeSend(conn.dv, b)
		if err == nil {
			if data != nil {
				b = data
			}
			conn.Write(b)
		}
	}
	p.close(conn)
}

func (p *Hub) close(conn *connection) {
	conn.close()
	p.actor.OnClose(conn.dv)
}

func (p *Hub) goOnline(conn *connection) {
	defer func() {
		go p.tryDeliver(conn.uid())
	}()

	l := p.cons[conn.uid()]
	if l == nil {
		p.cons[conn.uid()] = []*connection{conn}
		return
	}

	if p.sso {
		for _, c := range l {
			//notify before close connection
			go p.byeThenClose(c)
		}
		p.cons[conn.uid()] = []*connection{conn}
		return
	}

	//it's a new connection
	if len(connExclude(l, conn)) == len(l) {
		l = append(l, conn)
		p.cons[conn.uid()] = l
	}
}

//tryDeliver try to deliver all messages, if uid is online
func (p *Hub) tryDeliver(uid interface{}) {
	p.chReadSignal <- uid
}

//Kick kick all connections of uid
func (p *Hub) Kick(uid interface{}) {
	p.chKick <- uid
}

func (p *Hub) stateChange(conn *connection, online bool) {
	p.chConState <- &conState{conn, online}
}

//receive data from user
func (p *Hub) receive(dv *Device, b []byte) {
	p.chUp <- &upFrame{dv: dv, data: b}
}

func connExclude(l []*connection, ex *connection) []*connection {
	rest := make([]*connection, 0, len(l))
	for _, c := range l {
		if c != ex {
			rest = append(rest, c)
		}
	}
	return rest
}
<commit_msg>minor update<commit_after>package tok

import (
	"expvar"
	"log"
)

var (
	expOnline = expvar.NewInt("tokOnline")
	expUp     = expvar.NewInt("tokUp")
	expDown   = expvar.NewInt("tokDown")
	expEnq    = expvar.NewInt("tokEnq")
	expDeq    = expvar.NewInt("tokDeq")
)

type checkFrame struct {
	uid    interface{}
	chBool chan bool //channel to return online status
}

type downFrame struct {
	uid   interface{}
	ttl   uint32
	data  []byte
	chErr chan error //channel to read send result from
}

type upFrame struct {
	dv   *Device
	data []byte
}

//HubConfig config struct for creating new Hub
type HubConfig struct {
	Actor Actor //Actor implement dispatch logic
	Q     Queue //Message Queue, if nil, message to offline user will not be cached
	Sso   bool  //If it's true, new connection  with same uid will kick off old ones
}

//Hub core of tok, dispatch message between connections
type Hub struct {
	sso           bool
	actor         Actor
	q             Queue
	cons          map[interface{}][]*connection //connection list
	chUp          chan *upFrame
	chDown        chan *downFrame
	chConState    chan *conState
	chReadSignal  chan interface{}
	chKick        chan interface{}
	chQueryOnline chan chan []interface{}
	chCheck       chan *checkFrame
}

func createHub(actor Actor, q Queue, sso bool) *Hub {
	if ReadTimeout > 0 {
		log.Println("[tok] read timeout is enabled, make sure it's greater than your client ping interval. otherwise you'll get read timeout err")
	} else {
		if actor.Ping() == nil {
			log.Fatalln("[tok] both read timeout and server ping have been disabled, server socket resource leak might happen")
		}
	}

	hub := &Hub{
		sso:           sso,
		actor:         actor,
		q:             q,
		cons:          make(map[interface{}][]*connection),
		chUp:          make(chan *upFrame),
		chDown:        make(chan *downFrame),
		chConState:    make(chan *conState),
		chReadSignal:  make(chan interface{}),
		chKick:        make(chan interface{}),
		chQueryOnline: make(chan chan []interface{}),
		chCheck:       make(chan *checkFrame),
	}
	go hub.run()
	return hub
}

func (p *Hub) run() {
	for {

		select {
		case state := <-p.chConState:
			//			log.Printf("connection state change: %v, %v \n", state.online, &state.con)

			if state.online {
				p.goOnline(state.con)
			} else {
				p.goOffline(state.con)
			}
			count := int64(len(p.cons))
			expOnline.Set(count)
		case f := <-p.chUp:
			//			log.Println("up data")
			expUp.Add(1)
			go func() {
				b, err := p.actor.BeforeReceive(f.dv, f.data)
				if err != nil {
					return
				}
				if b == nil {
					b = f.data
				}
				p.actor.OnReceive(f.dv, b)
			}()
		case ff := <-p.chDown:
			if l := p.cons[ff.uid]; len(l) > 0 {
				//online
				go p.down(ff, l)
			} else {
				//offline
				if ff.ttl == 0 {
					ff.chErr <- ErrOffline
					close(ff.chErr)
				} else {
					go p.cache(ff)
				}
			}
		case cf := <-p.chCheck:
			_, ok := p.cons[cf.uid]
			cf.chBool <- ok
			close(cf.chBool)
		case uid := <-p.chReadSignal:
			//only pop msg for online user
			if len(p.cons[uid]) > 0 {
				go p.popMsg(uid)
			}
		case uid := <-p.chKick:
			p.innerKick(uid)
		case chOnline := <-p.chQueryOnline:
			result := make([]interface{}, 0, len(p.cons))
			for uid := range p.cons {
				result = append(result, uid)
			}
			chOnline <- result
			close(chOnline)
		}
	}
}

func (p *Hub) popMsg(uid interface{}) {
	if p.q == nil {
		return
	}

	for {
		b, err := p.q.Deq(uid)
		if err != nil {
			log.Println("deq error", err)
			return
		}
		if len(b) == 0 {
			//no more data in queue
			return
		}
		expDeq.Add(1)
		if err := p.Send(uid, b, 0); err != nil {
			if err := p.q.Enq(uid, b); err != nil {
				log.Println("re-cache err", err, uid)
			}
			return
		}
	}
}

//Send message to someone.
//ttl is expiry seconds. 0 means only send to online user
//If ttl = 0 and user is offline, ErrOffline will be returned.
//If ttl > 0 and user is offline or online but send fail, message will be cached for ttl seconds.
func (p *Hub) Send(to interface{}, b []byte, ttl uint32) error {

	ff := &downFrame{uid: to, data: b, ttl: ttl, chErr: make(chan error)}
	p.chDown <- ff
	err := <-ff.chErr
	if ttl > 0 && err != nil {
		//online send err
		ff.chErr = make(chan error) //create new channel
		go p.cache(ff)
		return <-ff.chErr
	}
	return err
}

//CheckOnline return whether user online or not
func (p *Hub) CheckOnline(uid interface{}) bool {
	cf := &checkFrame{uid: uid, chBool: make(chan bool)}
	p.chCheck <- cf
	return <-cf.chBool
}

//Online query online user list
func (p *Hub) Online() []interface{} {
	ch := make(chan []interface{})
	p.chQueryOnline <- ch
	return <-ch
}

func (p *Hub) cache(ff *downFrame) {
	defer close(ff.chErr)
	expEnq.Add(1)
	if p.q == nil {
		ff.chErr <- ErrQueueRequired
		return
	}

	if err := p.q.Enq(ff.uid, ff.data, ff.ttl); err != nil {
		ff.chErr <- err
	}
}

func (p *Hub) down(f *downFrame, conns []*connection) {
	expDown.Add(1)

	for _, con := range conns {
		b, err := p.actor.BeforeSend(con.dv, f.data)
		if err != nil {
			return
		}
		if b == nil {
			b = f.data
		}

		if err := con.Write(b); err != nil {
			f.chErr <- err
			return
		}
		go p.actor.OnSent(con.dv, f.data)
	}

}

func (p *Hub) goOffline(conn *connection) {
	l := p.cons[conn.uid()]
	rest := connExclude(l, conn)

	//this connection has gotten offline, ignore
	if len(l) == len(rest) {
		return
	}

	if len(rest) == 0 {
		delete(p.cons, conn.uid())
	} else {
		p.cons[conn.uid()] = rest
	}

	go p.close(conn)
}

func (p *Hub) innerKick(uid interface{}) {
	for _, conn := range p.cons[uid] {
		go p.close(conn)
	}
	delete(p.cons, uid)
}

func (p *Hub) byeThenClose(conn *connection) {
	b := p.actor.Bye(conn.dv, "sso")
	if b != nil {
		data, err := p.actor.BeforeSend(conn.dv, b)
		if err == nil {
			if data != nil {
				b = data
			}
			conn.Write(b)
		}
	}
	p.close(conn)
}

func (p *Hub) close(conn *connection) {
	conn.close()
	p.actor.OnClose(conn.dv)
}

func (p *Hub) goOnline(conn *connection) {
	defer func() {
		go p.tryDeliver(conn.uid())
	}()

	l := p.cons[conn.uid()]
	if l == nil {
		p.cons[conn.uid()] = []*connection{conn}
		return
	}

	if p.sso {
		for _, c := range l {
			//notify before close connection
			go p.byeThenClose(c)
		}
		p.cons[conn.uid()] = []*connection{conn}
		return
	}

	//it's a new connection
	if len(connExclude(l, conn)) == len(l) {
		l = append(l, conn)
		p.cons[conn.uid()] = l
	}
}

//tryDeliver try to deliver all messages, if uid is online
func (p *Hub) tryDeliver(uid interface{}) {
	p.chReadSignal <- uid
}

//Kick kick all connections of uid
func (p *Hub) Kick(uid interface{}) {
	p.chKick <- uid
}

func (p *Hub) stateChange(conn *connection, online bool) {
	p.chConState <- &conState{conn, online}
}

//receive data from user
func (p *Hub) receive(dv *Device, b []byte) {
	p.chUp <- &upFrame{dv: dv, data: b}
}

func connExclude(l []*connection, ex *connection) []*connection {
	rest := make([]*connection, 0, len(l))
	for _, c := range l {
		if c != ex {
			rest = append(rest, c)
		}
	}
	return rest
}
<|endoftext|>


<commit_before>/**
 * hub, core of chat application
 */

package tok

import (
	"expvar"
	"log"
)

var (
	expOnline = expvar.NewInt("tokOnline")
	expUp     = expvar.NewInt("tokUp")
	expDown   = expvar.NewInt("tokDown")
	expEnq    = expvar.NewInt("tokEnq")
	expDeq    = expvar.NewInt("tokDeq")
)

type fatFrame struct {
	frame *frame //frame to be sent
	ttl   uint32
	chErr chan error //channel to read send result from
}

type frame struct {
	uid  interface{}
	data []byte
}

//Config to create new Hub
type HubConfig struct {
	Actor Actor //Actor implement dispatch logic
	Q     Queue //Message Queue, if nil, message to offline user will not be cached
	Sso   bool  //If it's true, new connection  with same uid will kick off old ones
}

//Dispatch message between connections
type Hub struct {
	sso           bool
	actor         Actor
	q             Queue
	cons          map[interface{}][]*connection //connection list
	chUp          chan *frame
	chDown        chan *fatFrame //for online user
	chDown2       chan *fatFrame //for all user
	chConState    chan *conState
	chReadSignal  chan interface{}
	chQueryOnline chan chan []interface{}
}

func createHub(actor Actor, q Queue, sso bool) *Hub {
	if READ_TIMEOUT > 0 {
		log.Println("[tok] read timeout is enabled, make sure it's greater than your client ping interval. otherwise you'll get read timeout err")
	} else {
		if actor.Ping() == nil {
			log.Fatalln("[tok] both read timeout and server ping have been disabled, server socket resource leak might happen")
		}
	}

	hub := &Hub{
		sso:           sso,
		actor:         actor,
		q:             q,
		cons:          make(map[interface{}][]*connection),
		chUp:          make(chan *frame),
		chDown:        make(chan *fatFrame),
		chDown2:       make(chan *fatFrame),
		chConState:    make(chan *conState),
		chReadSignal:  make(chan interface{}),
		chQueryOnline: make(chan chan []interface{}),
	}
	go hub.run()
	return hub
}

func (p *Hub) run() {
	for {
		select {
		case state := <-p.chConState:
			//			log.Printf("connection state change: %v, %v \n", state.online, &state.con)

			if state.online {
				p.goOnline(state.con)
			} else {
				p.goOffline(state.con)
			}
			count := int64(len(p.cons))
			expOnline.Set(count)
		case f := <-p.chUp:
			//			log.Println("up data")
			expUp.Add(1)
			go p.actor.OnReceive(f.uid, f.data)
		case ff := <-p.chDown:
			if len(p.cons[ff.frame.uid]) > 0 {
				expDown.Add(1)
				go p.down(ff)
			} else {
				ff.chErr <- ErrOffline
				close(ff.chErr)
			}
		case ff := <-p.chDown2:
			if len(p.cons[ff.frame.uid]) > 0 {
				expDown.Add(1)
				go p.down(ff)
			} else {
				expEnq.Add(1)
				go p.cache(ff)
			}

		case uid := <-p.chReadSignal:
			//only pop msg for online user
			if len(p.cons[uid]) > 0 {
				go p.popMsg(uid)
			}
		case chOnline := <-p.chQueryOnline:
			result := make([]interface{}, 0, len(p.cons))
			for uid := range p.cons {
				result = append(result, uid)
			}
			chOnline <- result
			close(chOnline)
		}
	}
}

func (p *Hub) popMsg(uid interface{}) {
	if p.q == nil {
		return
	}

	for {
		b, err := p.q.Deq(uid)
		if err != nil {
			log.Println("deq error", err)
			return
		}
		if b == nil {
			//no more data in queue
			return
		}
		expDeq.Add(1)
		if err := p.Send(uid, b); err != nil {
			log.Println("send err after deq")
			return
		}
	}
}

//Send message to someone,
//ttl is expiry seconds. 0 means forever
//if ttl >= 0 and user is offline, message will be cached for ttl seconds
//if ttl < 0 and user is offline, ErrOffline will be returned
//if ttl >=0 and user is online, but error occurred during send, message will be cached
//if ttl < 0 and user is online, but error occurred during send, the error will be returned
func (p *Hub) Send(to interface{}, b []byte, ttl ...int) error {
	t := 0
	if len(ttl) > 0 {
		t = ttl[0]
	}

	ff := &fatFrame{frame: &frame{uid: to, data: b}, ttl: uint32(t), chErr: make(chan error)}
	if t > 0 {
		p.chDown2 <- ff
	} else {
		p.chDown <- ff
	}
	err := <-ff.chErr
	if err == nil {
		return nil
	}

	if t < 0 || p.q == nil {
		return err
	}

	//try cache it in case error
	ff.chErr = make(chan error)
	go p.cache(ff)
	return <-ff.chErr
}

//Query online user list
func (p *Hub) Online() []interface{} {
	ch := make(chan []interface{})
	p.chQueryOnline <- ch
	return <-ch
}

func (p *Hub) cache(ff *fatFrame) {
	defer close(ff.chErr)

	if p.q == nil {
		ff.chErr <- ErrQueueRequired
		return
	}

	f := ff.frame
	if err := p.q.Enq(f.uid, f.data, ff.ttl); err != nil {
		ff.chErr <- err
	}
}

func (p *Hub) down(ff *fatFrame) {
	defer close(ff.chErr)

	for _, con := range p.cons[ff.frame.uid] {
		err := con.Write(ff.frame.data)
		if err != nil {
			ff.chErr <- err
			return
		}
	}
}

func (p *Hub) goOffline(conn *connection) {
	l := p.cons[conn.uid]
	rest := connExclude(l, conn)

	//this connection has gotten offline, ignore
	if len(l) == len(rest) {
		return
	}

	if len(rest) == 0 {
		delete(p.cons, conn.uid)
	} else {
		p.cons[conn.uid] = rest
	}

	go func(active int) {
		conn.close()
		p.actor.OnClose(conn.uid, active)
	}(len(rest))
}

func (p *Hub) goOnline(conn *connection) {
	l := p.cons[conn.uid]
	if l == nil {
		l = []*connection{conn}
	} else {
		if p.sso {
			for _, old := range l {
				log.Printf("kick %v\n", old)
				//notify before close connection
				go func() {
					old.Write(p.actor.Bye("sso"))
					old.close()
					//after sso kick, only one connection keep active
					p.actor.OnClose(conn.uid, 1)
				}()
			}
			l = []*connection{conn}
		} else {
			exists := false
			for _, c := range l {
				if c == conn {
					log.Println("warning, repeat online: ", c)
					exists = true
					break
				}
			}
			if !exists {
				l = append(l, conn)
			}
		}
	}
	p.cons[conn.uid] = l
	go p.startRead(conn.uid)
}

func (p *Hub) startRead(uid interface{}) {
	p.chReadSignal <- uid
}

func (p *Hub) stateChange(conn *connection, online bool) {
	p.chConState <- &conState{conn, online}
}

//receive data from user
func (p *Hub) receive(uid interface{}, b []byte) {
	p.chUp <- &frame{uid: uid, data: b}
}

func connExclude(l []*connection, ex *connection) []*connection {
	rest := make([]*connection, 0, len(l))
	for _, c := range l {
		if c != ex {
			rest = append(rest, c)
		}
	}
	return rest
}
<commit_msg>minor update<commit_after>/**
 * hub, core of chat application
 */

package tok

import (
	"expvar"
	"log"
)

var (
	expOnline = expvar.NewInt("tokOnline")
	expUp     = expvar.NewInt("tokUp")
	expDown   = expvar.NewInt("tokDown")
	expEnq    = expvar.NewInt("tokEnq")
	expDeq    = expvar.NewInt("tokDeq")
)

type fatFrame struct {
	frame *frame //frame to be sent
	ttl   uint32
	chErr chan error //channel to read send result from
}

type frame struct {
	uid  interface{}
	data []byte
}

//Config to create new Hub
type HubConfig struct {
	Actor Actor //Actor implement dispatch logic
	Q     Queue //Message Queue, if nil, message to offline user will not be cached
	Sso   bool  //If it's true, new connection  with same uid will kick off old ones
}

//Dispatch message between connections
type Hub struct {
	sso           bool
	actor         Actor
	q             Queue
	cons          map[interface{}][]*connection //connection list
	chUp          chan *frame
	chDown        chan *fatFrame //for online user
	chDown2       chan *fatFrame //for all user
	chConState    chan *conState
	chReadSignal  chan interface{}
	chQueryOnline chan chan []interface{}
}

func createHub(actor Actor, q Queue, sso bool) *Hub {
	if READ_TIMEOUT > 0 {
		log.Println("[tok] read timeout is enabled, make sure it's greater than your client ping interval. otherwise you'll get read timeout err")
	} else {
		if actor.Ping() == nil {
			log.Fatalln("[tok] both read timeout and server ping have been disabled, server socket resource leak might happen")
		}
	}

	hub := &Hub{
		sso:           sso,
		actor:         actor,
		q:             q,
		cons:          make(map[interface{}][]*connection),
		chUp:          make(chan *frame),
		chDown:        make(chan *fatFrame),
		chDown2:       make(chan *fatFrame),
		chConState:    make(chan *conState),
		chReadSignal:  make(chan interface{}),
		chQueryOnline: make(chan chan []interface{}),
	}
	go hub.run()
	return hub
}

func (p *Hub) run() {
	for {
		select {
		case state := <-p.chConState:
			//			log.Printf("connection state change: %v, %v \n", state.online, &state.con)

			if state.online {
				p.goOnline(state.con)
			} else {
				p.goOffline(state.con)
			}
			count := int64(len(p.cons))
			expOnline.Set(count)
		case f := <-p.chUp:
			//			log.Println("up data")
			expUp.Add(1)
			go p.actor.OnReceive(f.uid, f.data)
		case ff := <-p.chDown:
			if len(p.cons[ff.frame.uid]) > 0 {
				expDown.Add(1)
				go p.down(ff)
			} else {
				ff.chErr <- ErrOffline
				close(ff.chErr)
			}
		case ff := <-p.chDown2:
			if len(p.cons[ff.frame.uid]) > 0 {
				expDown.Add(1)
				go p.down(ff)
			} else {
				expEnq.Add(1)
				go p.cache(ff)
			}

		case uid := <-p.chReadSignal:
			//only pop msg for online user
			if len(p.cons[uid]) > 0 {
				go p.popMsg(uid)
			}
		case chOnline := <-p.chQueryOnline:
			result := make([]interface{}, 0, len(p.cons))
			for uid := range p.cons {
				result = append(result, uid)
			}
			chOnline <- result
			close(chOnline)
		}
	}
}

func (p *Hub) popMsg(uid interface{}) {
	if p.q == nil {
		return
	}

	for {
		b, err := p.q.Deq(uid)
		if err != nil {
			log.Println("deq error", err)
			return
		}
		if b == nil {
			//no more data in queue
			return
		}
		expDeq.Add(1)
		if err := p.Send(uid, b); err != nil {
			log.Println("send err after deq")
			return
		}
	}
}

//Send message to someone,
//ttl is expiry seconds. 0 means forever
//if ttl >= 0 and user is offline, message will be cached for ttl seconds
//if ttl < 0 and user is offline, ErrOffline will be returned
//if ttl >=0 and user is online, but error occurred during send, message will be cached
//if ttl < 0 and user is online, but error occurred during send, the error will be returned
func (p *Hub) Send(to interface{}, b []byte, ttl ...int) error {
	t := -1
	if len(ttl) > 0 {
		t = ttl[0]
	}

	ff := &fatFrame{frame: &frame{uid: to, data: b}, ttl: uint32(t), chErr: make(chan error)}
	if t > 0 {
		p.chDown2 <- ff
	} else {
		p.chDown <- ff
	}
	err := <-ff.chErr
	if err == nil {
		return nil
	}

	if t < 0 || p.q == nil {
		return err
	}

	//try cache it in case error
	ff.chErr = make(chan error)
	go p.cache(ff)
	return <-ff.chErr
}

//Query online user list
func (p *Hub) Online() []interface{} {
	ch := make(chan []interface{})
	p.chQueryOnline <- ch
	return <-ch
}

func (p *Hub) cache(ff *fatFrame) {
	defer close(ff.chErr)

	if p.q == nil {
		ff.chErr <- ErrQueueRequired
		return
	}

	f := ff.frame
	if err := p.q.Enq(f.uid, f.data, ff.ttl); err != nil {
		ff.chErr <- err
	}
}

func (p *Hub) down(ff *fatFrame) {
	defer close(ff.chErr)

	for _, con := range p.cons[ff.frame.uid] {
		err := con.Write(ff.frame.data)
		if err != nil {
			ff.chErr <- err
			return
		}
	}
}

func (p *Hub) goOffline(conn *connection) {
	l := p.cons[conn.uid]
	rest := connExclude(l, conn)

	//this connection has gotten offline, ignore
	if len(l) == len(rest) {
		return
	}

	if len(rest) == 0 {
		delete(p.cons, conn.uid)
	} else {
		p.cons[conn.uid] = rest
	}

	go func(active int) {
		conn.close()
		p.actor.OnClose(conn.uid, active)
	}(len(rest))
}

func (p *Hub) goOnline(conn *connection) {
	l := p.cons[conn.uid]
	if l == nil {
		l = []*connection{conn}
	} else {
		if p.sso {
			for _, old := range l {
				log.Printf("kick %v\n", old)
				//notify before close connection
				go func() {
					old.Write(p.actor.Bye("sso"))
					old.close()
					//after sso kick, only one connection keep active
					p.actor.OnClose(conn.uid, 1)
				}()
			}
			l = []*connection{conn}
		} else {
			exists := false
			for _, c := range l {
				if c == conn {
					log.Println("warning, repeat online: ", c)
					exists = true
					break
				}
			}
			if !exists {
				l = append(l, conn)
			}
		}
	}
	p.cons[conn.uid] = l
	go p.startRead(conn.uid)
}

func (p *Hub) startRead(uid interface{}) {
	p.chReadSignal <- uid
}

func (p *Hub) stateChange(conn *connection, online bool) {
	p.chConState <- &conState{conn, online}
}

//receive data from user
func (p *Hub) receive(uid interface{}, b []byte) {
	p.chUp <- &frame{uid: uid, data: b}
}

func connExclude(l []*connection, ex *connection) []*connection {
	rest := make([]*connection, 0, len(l))
	for _, c := range l {
		if c != ex {
			rest = append(rest, c)
		}
	}
	return rest
}
<|endoftext|>