<commit_before>package testing

import (
	"fmt"

	"sync"

	"github.com/aws/aws-sdk-go/service/swf"
	"github.com/sclasen/swfsm/activity"
	"github.com/sclasen/swfsm/fsm"
	. "github.com/sclasen/swfsm/log"
	. "github.com/sclasen/swfsm/sugar"
)

const (
	StubWorkflow      = "stub"
	ShortStubWorkflow = "stub"
	StubVersion       = "1"
)

var (
	StubTaskList      = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", StubWorkflow, StubVersion))}
	ShortStubTaskList = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", ShortStubWorkflow, StubVersion))}
)

type DecisionOutcome struct {
	DecisionTask *swf.PollForDecisionTaskOutput
	State        *fsm.SerializedState
	Decisions    []*swf.Decision
}

type StateData struct {
	State string
	Data  interface{}
}

func StubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       StubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: StubState()})
	return f
}

func StubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h *swf.HistoryEvent, data interface{}) fsm.Outcome {
		Log.Printf("at=stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.Stay(data, ctx.EmptyDecisions())
	}
}

func ShortStubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       ShortStubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: ShortStubState()})
	return f
}

func ShortStubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h *swf.HistoryEvent, data interface{}) fsm.Outcome {
		Log.Printf("at=short-stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.CompleteWorkflow(data)
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
func TestDecisionInterceptor(testId string, stubbedWorkflows, stubbedShortWorkflows []string) fsm.DecisionInterceptor {
	stubbed := make(map[string]struct{})
	stubbedShort := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflows {
		stubbed[s] = v
	}
	for _, s := range stubbedShortWorkflows {
		stubbedShort[s] = v
	}
	return &fsm.FuncInterceptor{
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartChildWorkflowExecution:
					if _, ok := stubbed[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(StubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = StubTaskList
					}
					if _, ok := stubbedShort[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(ShortStubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = ShortStubTaskList
					}
				case swf.DecisionTypeScheduleActivityTask:
					d.ScheduleActivityTaskDecisionAttributes.TaskList = &swf.TaskList{Name: S(*d.ScheduleActivityTaskDecisionAttributes.TaskList.Name + testId)}
				case swf.DecisionTypeContinueAsNewWorkflowExecution:
					d.ContinueAsNewWorkflowExecutionDecisionAttributes.TaskList = &swf.TaskList{Name: S(testId)}
				}
			}
		},
	}
}

func NoOpActivityInterceptor() activity.ActivityInterceptor {
	return &activity.FuncInterceptor{}
}

// interceptor that fails the activity once per activity and returns to actual result subsequent time
// used to test error handling and retries of activities in fsms
func TestFailOnceActivityInterceptor() activity.ActivityInterceptor {
	mutex := sync.Mutex{}
	tried := map[string]bool{}
	return &activity.FuncInterceptor{
		AfterTaskFn: func(t *swf.PollForActivityTaskOutput, result interface{}, err error) (interface{}, error) {
			mutex.Lock()
			defer mutex.Unlock()

			if err != nil || tried[*t.ActivityId] {
				Log.Printf("interceptor.test.fail-once at=passthrough activity-id=%q", *t.ActivityId)
				return result, err
			}

			tried[*t.ActivityId] = true
			msg := fmt.Sprintf("interceptor.test.fail-once at=fail activity-id=%q", *t.ActivityId)
			Log.Println(msg)
			return nil, fmt.Errorf(msg)
		},
	}
}

func TestThrotteSignalsOnceInterceptor() fsm.DecisionInterceptor {
	throttledSignals := make(map[string]*fsm.SignalInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeSignalExternalWorkflowExecutionFailed {
					if info, found := throttledSignals[*h.SignalExternalWorkflowExecutionFailedEventAttributes.WorkflowId]; found {
						h.SignalExternalWorkflowExecutionFailedEventAttributes.WorkflowId = &info.WorkflowId
						//rewrite workflow not found to throttling.
						h.SignalExternalWorkflowExecutionFailedEventAttributes.Cause = S(swf.SignalExternalWorkflowExecutionFailedCauseSignalExternalWorkflowExecutionRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeSignalExternalWorkflowExecution:
					signalWF := *d.SignalExternalWorkflowExecutionDecisionAttributes.WorkflowId
					signalName := *d.SignalExternalWorkflowExecutionDecisionAttributes.SignalName
					throttle := fmt.Sprintf("fail-on-purpose-%s-%s", signalWF, signalName)
					if _, found := throttledSignals[throttle]; !found {
						//if we havent failed the signal yet, add it and swap the workflow id so it will fail
						throttledSignals[throttle] = &fsm.SignalInfo{WorkflowId: signalWF, SignalName: signalName}
						d.SignalExternalWorkflowExecutionDecisionAttributes.WorkflowId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledSignals, throttle)
					}
				}

			}
		},
	}
}

func TestThrotteCancelsOnceInterceptor() fsm.DecisionInterceptor {
	throttledCancels := make(map[string]*fsm.CancellationInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeRequestCancelExternalWorkflowExecutionFailed {
					if info, found := throttledCancels[*h.RequestCancelExternalWorkflowExecutionFailedEventAttributes.WorkflowId]; found {
						h.RequestCancelExternalWorkflowExecutionFailedEventAttributes.WorkflowId = &info.WorkflowId
						//rewrite workflow not found to throttling.
						h.RequestCancelExternalWorkflowExecutionFailedEventAttributes.Cause = S(swf.RequestCancelExternalWorkflowExecutionFailedCauseRequestCancelExternalWorkflowExecutionRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeRequestCancelExternalWorkflowExecution:
					cancelWF := *d.RequestCancelExternalWorkflowExecutionDecisionAttributes.WorkflowId
					throttle := fmt.Sprintf("fail-on-purpose-%s", cancelWF)
					if _, found := throttledCancels[throttle]; !found {
						//if we havent failed the signal yet, add it and swap the workflow id so it will fail
						throttledCancels[throttle] = &fsm.CancellationInfo{WorkflowId: cancelWF}
						d.RequestCancelExternalWorkflowExecutionDecisionAttributes.WorkflowId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledCancels, throttle)
					}
				}

			}
		},
	}
}

func TestThrotteChildrenOnceInterceptor() fsm.DecisionInterceptor {
	throttledChildren := make(map[string]*fsm.ChildInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeStartChildWorkflowExecutionFailed {
					if info, found := throttledChildren[*h.StartChildWorkflowExecutionFailedEventAttributes.WorkflowId]; found {
						h.StartChildWorkflowExecutionFailedEventAttributes.WorkflowId = &info.WorkflowId
						//rewrite workflow not found to throttling.
						h.StartChildWorkflowExecutionFailedEventAttributes.Cause = S(swf.StartChildWorkflowExecutionFailedCauseChildCreationRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartChildWorkflowExecution:
					childWF := *d.StartChildWorkflowExecutionDecisionAttributes.WorkflowId
					throttle := fmt.Sprintf("fail-on-purpose-%s", childWF)
					if _, found := throttledChildren[throttle]; !found {
						//if we havent failed the child yet, add it and swap the workflow id so it will fail
						throttledChildren[throttle] = &fsm.ChildInfo{WorkflowId: childWF}
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledChildren, throttle)
					}
				}

			}
		},
	}
}

//This one is a bit weird. OnStart, we need to schedule a bunch of timers that we can use to
//cause timer id in use failures, that we rewrite to throttles
func TestThrotteTimersOnceInterceptor(numTimers int) fsm.DecisionInterceptor {
	timerNames := make(chan string, numTimers*2)
	throttledTimers := make(map[string]*fsm.TimerInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeStartTimerFailed {
					if info, found := throttledTimers[*h.StartTimerFailedEventAttributes.TimerId]; found {
						h.StartTimerFailedEventAttributes.TimerId = &info.TimerId
						//rewrite timer in use to throttling.
						h.StartTimerFailedEventAttributes.Cause = S(swf.StartTimerFailedCauseTimerCreationRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				//when the workflow starts schedule numTimers timers to use for causing fake throttles
				if *h.EventType == swf.EventTypeWorkflowExecutionStarted {
					for i := 0; i < numTimers; i++ {
						timer := fmt.Sprintf("throttle-test-%d", i)
						timerNames <- timer
						decision := &swf.Decision{
							DecisionType: S(swf.DecisionTypeStartTimer),
							StartTimerDecisionAttributes: &swf.StartTimerDecisionAttributes{
								TimerId:            S(timer),
								StartToFireTimeout: S("3600"),
							},
						}
						outcome.Decisions = append(outcome.Decisions, decision)
					}
				}
			}
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartTimer:
					timerId := *d.StartTimerDecisionAttributes.TimerId
					throttle := <-timerNames
					if _, found := throttledTimers[throttle]; !found {
						//if we havent failed the signal yet, add it and swap the workflow id so it wont fail
						throttledTimers[throttle] = &fsm.TimerInfo{TimerId: timerId}
						d.StartTimerDecisionAttributes.TimerId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledTimers, throttle)
						timerNames <- throttle
					}
				}

			}
		},
	}
}

func TestReplicator(decisionOutcomes chan DecisionOutcome) fsm.ReplicationHandler {
	return func(ctx *fsm.FSMContext, task *swf.PollForDecisionTaskOutput, outcome *swf.RespondDecisionTaskCompletedInput, state *fsm.SerializedState) error {
		decisionOutcomes <- DecisionOutcome{State: state, DecisionTask: task, Decisions: outcome.Decisions}
		return nil
	}
}

func TestSWF(client fsm.ClientSWFOps, stubbedWorkflow ...string) fsm.ClientSWFOps {
	stubbed := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflow {
		stubbed[s] = v
	}
	return &StubSWFClient{
		ClientSWFOps:     client,
		stubbedWorkflows: stubbed,
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
type StubSWFClient struct {
	fsm.ClientSWFOps
	stubbedWorkflows map[string]struct{}
}

func (s *StubSWFClient) StartWorkflowExecution(req *swf.StartWorkflowExecutionInput) (resp *swf.StartWorkflowExecutionOutput, err error) {
	if _, ok := s.stubbedWorkflows[*req.WorkflowType.Name]; ok {
		req.WorkflowType.Name = S(StubWorkflow)
		req.WorkflowType.Version = S(StubVersion)
		req.ExecutionStartToCloseTimeout = S("360")
		req.TaskList = StubTaskList
	}
	return s.ClientSWFOps.StartWorkflowExecution(req)
}
<commit_msg>dont throttle continue timer creation<commit_after>package testing

import (
	"fmt"

	"sync"

	"github.com/aws/aws-sdk-go/service/swf"
	"github.com/sclasen/swfsm/activity"
	"github.com/sclasen/swfsm/fsm"
	. "github.com/sclasen/swfsm/log"
	. "github.com/sclasen/swfsm/sugar"
)

const (
	StubWorkflow      = "stub"
	ShortStubWorkflow = "stub"
	StubVersion       = "1"
)

var (
	StubTaskList      = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", StubWorkflow, StubVersion))}
	ShortStubTaskList = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", ShortStubWorkflow, StubVersion))}
)

type DecisionOutcome struct {
	DecisionTask *swf.PollForDecisionTaskOutput
	State        *fsm.SerializedState
	Decisions    []*swf.Decision
}

type StateData struct {
	State string
	Data  interface{}
}

func StubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       StubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: StubState()})
	return f
}

func StubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h *swf.HistoryEvent, data interface{}) fsm.Outcome {
		Log.Printf("at=stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.Stay(data, ctx.EmptyDecisions())
	}
}

func ShortStubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       ShortStubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: ShortStubState()})
	return f
}

func ShortStubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h *swf.HistoryEvent, data interface{}) fsm.Outcome {
		Log.Printf("at=short-stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.CompleteWorkflow(data)
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
func TestDecisionInterceptor(testId string, stubbedWorkflows, stubbedShortWorkflows []string) fsm.DecisionInterceptor {
	stubbed := make(map[string]struct{})
	stubbedShort := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflows {
		stubbed[s] = v
	}
	for _, s := range stubbedShortWorkflows {
		stubbedShort[s] = v
	}
	return &fsm.FuncInterceptor{
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartChildWorkflowExecution:
					if _, ok := stubbed[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(StubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = StubTaskList
					}
					if _, ok := stubbedShort[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(ShortStubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = ShortStubTaskList
					}
				case swf.DecisionTypeScheduleActivityTask:
					d.ScheduleActivityTaskDecisionAttributes.TaskList = &swf.TaskList{Name: S(*d.ScheduleActivityTaskDecisionAttributes.TaskList.Name + testId)}
				case swf.DecisionTypeContinueAsNewWorkflowExecution:
					d.ContinueAsNewWorkflowExecutionDecisionAttributes.TaskList = &swf.TaskList{Name: S(testId)}
				}
			}
		},
	}
}

func NoOpActivityInterceptor() activity.ActivityInterceptor {
	return &activity.FuncInterceptor{}
}

// interceptor that fails the activity once per activity and returns to actual result subsequent time
// used to test error handling and retries of activities in fsms
func TestFailOnceActivityInterceptor() activity.ActivityInterceptor {
	mutex := sync.Mutex{}
	tried := map[string]bool{}
	return &activity.FuncInterceptor{
		AfterTaskFn: func(t *swf.PollForActivityTaskOutput, result interface{}, err error) (interface{}, error) {
			mutex.Lock()
			defer mutex.Unlock()

			if err != nil || tried[*t.ActivityId] {
				Log.Printf("interceptor.test.fail-once at=passthrough activity-id=%q", *t.ActivityId)
				return result, err
			}

			tried[*t.ActivityId] = true
			msg := fmt.Sprintf("interceptor.test.fail-once at=fail activity-id=%q", *t.ActivityId)
			Log.Println(msg)
			return nil, fmt.Errorf(msg)
		},
	}
}

func TestThrotteSignalsOnceInterceptor() fsm.DecisionInterceptor {
	throttledSignals := make(map[string]*fsm.SignalInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeSignalExternalWorkflowExecutionFailed {
					if info, found := throttledSignals[*h.SignalExternalWorkflowExecutionFailedEventAttributes.WorkflowId]; found {
						h.SignalExternalWorkflowExecutionFailedEventAttributes.WorkflowId = &info.WorkflowId
						//rewrite workflow not found to throttling.
						h.SignalExternalWorkflowExecutionFailedEventAttributes.Cause = S(swf.SignalExternalWorkflowExecutionFailedCauseSignalExternalWorkflowExecutionRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeSignalExternalWorkflowExecution:
					signalWF := *d.SignalExternalWorkflowExecutionDecisionAttributes.WorkflowId
					signalName := *d.SignalExternalWorkflowExecutionDecisionAttributes.SignalName
					throttle := fmt.Sprintf("fail-on-purpose-%s-%s", signalWF, signalName)
					if _, found := throttledSignals[throttle]; !found {
						//if we havent failed the signal yet, add it and swap the workflow id so it will fail
						throttledSignals[throttle] = &fsm.SignalInfo{WorkflowId: signalWF, SignalName: signalName}
						d.SignalExternalWorkflowExecutionDecisionAttributes.WorkflowId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledSignals, throttle)
					}
				}

			}
		},
	}
}

func TestThrotteCancelsOnceInterceptor() fsm.DecisionInterceptor {
	throttledCancels := make(map[string]*fsm.CancellationInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeRequestCancelExternalWorkflowExecutionFailed {
					if info, found := throttledCancels[*h.RequestCancelExternalWorkflowExecutionFailedEventAttributes.WorkflowId]; found {
						h.RequestCancelExternalWorkflowExecutionFailedEventAttributes.WorkflowId = &info.WorkflowId
						//rewrite workflow not found to throttling.
						h.RequestCancelExternalWorkflowExecutionFailedEventAttributes.Cause = S(swf.RequestCancelExternalWorkflowExecutionFailedCauseRequestCancelExternalWorkflowExecutionRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeRequestCancelExternalWorkflowExecution:
					cancelWF := *d.RequestCancelExternalWorkflowExecutionDecisionAttributes.WorkflowId
					throttle := fmt.Sprintf("fail-on-purpose-%s", cancelWF)
					if _, found := throttledCancels[throttle]; !found {
						//if we havent failed the signal yet, add it and swap the workflow id so it will fail
						throttledCancels[throttle] = &fsm.CancellationInfo{WorkflowId: cancelWF}
						d.RequestCancelExternalWorkflowExecutionDecisionAttributes.WorkflowId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledCancels, throttle)
					}
				}

			}
		},
	}
}

func TestThrotteChildrenOnceInterceptor() fsm.DecisionInterceptor {
	throttledChildren := make(map[string]*fsm.ChildInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeStartChildWorkflowExecutionFailed {
					if info, found := throttledChildren[*h.StartChildWorkflowExecutionFailedEventAttributes.WorkflowId]; found {
						h.StartChildWorkflowExecutionFailedEventAttributes.WorkflowId = &info.WorkflowId
						//rewrite workflow not found to throttling.
						h.StartChildWorkflowExecutionFailedEventAttributes.Cause = S(swf.StartChildWorkflowExecutionFailedCauseChildCreationRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartChildWorkflowExecution:
					childWF := *d.StartChildWorkflowExecutionDecisionAttributes.WorkflowId
					throttle := fmt.Sprintf("fail-on-purpose-%s", childWF)
					if _, found := throttledChildren[throttle]; !found {
						//if we havent failed the child yet, add it and swap the workflow id so it will fail
						throttledChildren[throttle] = &fsm.ChildInfo{WorkflowId: childWF}
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledChildren, throttle)
					}
				}

			}
		},
	}
}

//This one is a bit weird. OnStart, we need to schedule a bunch of timers that we can use to
//cause timer id in use failures, that we rewrite to throttles
func TestThrotteTimersOnceInterceptor(numTimers int) fsm.DecisionInterceptor {
	timerNames := make(chan string, numTimers*2)
	throttledTimers := make(map[string]*fsm.TimerInfo)
	return &fsm.FuncInterceptor{
		BeforeDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				if *h.EventType == swf.EventTypeStartTimerFailed {
					if info, found := throttledTimers[*h.StartTimerFailedEventAttributes.TimerId]; found {
						h.StartTimerFailedEventAttributes.TimerId = &info.TimerId
						//rewrite timer in use to throttling.
						h.StartTimerFailedEventAttributes.Cause = S(swf.StartTimerFailedCauseTimerCreationRateExceeded)
					}
				}
			}

		},
		AfterDecisionFn: func(decision *swf.PollForDecisionTaskOutput, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, h := range decision.Events {
				if *h.EventId <= *decision.PreviousStartedEventId {
					break
				}
				//when the workflow starts schedule numTimers timers to use for causing fake throttles
				if *h.EventType == swf.EventTypeWorkflowExecutionStarted {
					for i := 0; i < numTimers; i++ {
						timer := fmt.Sprintf("throttle-test-%d", i)
						timerNames <- timer
						decision := &swf.Decision{
							DecisionType: S(swf.DecisionTypeStartTimer),
							StartTimerDecisionAttributes: &swf.StartTimerDecisionAttributes{
								TimerId:            S(timer),
								StartToFireTimeout: S("3600"),
							},
						}
						outcome.Decisions = append(outcome.Decisions, decision)
					}
				}
			}
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartTimer:
					timerId := *d.StartTimerDecisionAttributes.TimerId
					if timerId == fsm.ContinueTimer {
						continue
					}
					throttle := <-timerNames
					if _, found := throttledTimers[throttle]; !found {
						//if we havent failed the signal yet, add it and swap the workflow id so it wont fail
						throttledTimers[throttle] = &fsm.TimerInfo{TimerId: timerId}
						d.StartTimerDecisionAttributes.TimerId = S(throttle)
					} else {
						//we are retrying after the forced failure, remove it
						delete(throttledTimers, throttle)
						timerNames <- throttle
					}
				}

			}
		},
	}
}

func TestReplicator(decisionOutcomes chan DecisionOutcome) fsm.ReplicationHandler {
	return func(ctx *fsm.FSMContext, task *swf.PollForDecisionTaskOutput, outcome *swf.RespondDecisionTaskCompletedInput, state *fsm.SerializedState) error {
		decisionOutcomes <- DecisionOutcome{State: state, DecisionTask: task, Decisions: outcome.Decisions}
		return nil
	}
}

func TestSWF(client fsm.ClientSWFOps, stubbedWorkflow ...string) fsm.ClientSWFOps {
	stubbed := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflow {
		stubbed[s] = v
	}
	return &StubSWFClient{
		ClientSWFOps:     client,
		stubbedWorkflows: stubbed,
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
type StubSWFClient struct {
	fsm.ClientSWFOps
	stubbedWorkflows map[string]struct{}
}

func (s *StubSWFClient) StartWorkflowExecution(req *swf.StartWorkflowExecutionInput) (resp *swf.StartWorkflowExecutionOutput, err error) {
	if _, ok := s.stubbedWorkflows[*req.WorkflowType.Name]; ok {
		req.WorkflowType.Name = S(StubWorkflow)
		req.WorkflowType.Version = S(StubVersion)
		req.ExecutionStartToCloseTimeout = S("360")
		req.TaskList = StubTaskList
	}
	return s.ClientSWFOps.StartWorkflowExecution(req)
}
<|endoftext|>


<commit_before>// Copyright 2016 Michal Witkowski. All Rights Reserved.
// See LICENSE for licensing terms.

/*

The validator plugin generates a Validate method for each message.
By default, if none of the message's fields are annotated with the gogo validator annotation, it returns a nil.
In case some of the fields are annotated, the Validate function returns nil upon sucessful validation, or an error
describing why the validation failed.
The Validate method is called recursively for all submessage of the message.

TODO(michal): ADD COMMENTS.

Equal is enabled using the following extensions:

  - equal
  - equal_all

While VerboseEqual is enable dusing the following extensions:

  - verbose_equal
  - verbose_equal_all

The equal plugin also generates a test given it is enabled using one of the following extensions:

  - testgen
  - testgen_all

Let us look at:

  github.com/gogo/protobuf/test/example/example.proto

Btw all the output can be seen at:

  github.com/gogo/protobuf/test/example/*

The following message:

given to the equal plugin, will generate the following code:

and the following test code:

*/
package plugin

import (
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/gogo/protobuf/gogoproto"
	"github.com/gogo/protobuf/proto"
	descriptor "github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
	"github.com/gogo/protobuf/protoc-gen-gogo/generator"
	"github.com/gogo/protobuf/vanity"
	"github.com/mwitkow/go-proto-validators"
)

func init() {
	generator.RegisterPlugin(NewPlugin(true))
}

type plugin struct {
	*generator.Generator
	generator.PluginImports
	regexPkg      generator.Single
	fmtPkg        generator.Single
	protoPkg      generator.Single
	validatorPkg  generator.Single
	useGogoImport bool
}

func NewPlugin(useGogoImport bool) generator.Plugin {
	return &plugin{useGogoImport: useGogoImport}
}

func (p *plugin) Name() string {
	return "validator"
}

func (p *plugin) Init(g *generator.Generator) {
	p.Generator = g
}

func (p *plugin) Generate(file *generator.FileDescriptor) {
	if !p.useGogoImport {
		vanity.TurnOffGogoImport(file.FileDescriptorProto)
	}
	p.PluginImports = generator.NewPluginImports(p.Generator)
	p.regexPkg = p.NewImport("regexp")
	p.fmtPkg = p.NewImport("fmt")
	p.validatorPkg = p.NewImport("github.com/mwitkow/go-proto-validators")

	for _, msg := range file.Messages() {
		if msg.DescriptorProto.GetOptions().GetMapEntry() {
			continue
		}
		p.generateRegexVars(file, msg)
		if gogoproto.IsProto3(file.FileDescriptorProto) {
			p.generateProto3Message(file, msg)
		} else {
			p.generateProto2Message(file, msg)
		}

	}
}

func getFieldValidatorIfAny(field *descriptor.FieldDescriptorProto) *validator.FieldValidator {
	if field.Options != nil {
		v, err := proto.GetExtension(field.Options, validator.E_Field)
		if err == nil && v.(*validator.FieldValidator) != nil {
			return (v.(*validator.FieldValidator))
		}
	}
	return nil
}

func (p *plugin) isSupportedInt(field *descriptor.FieldDescriptorProto) bool {
	switch *(field.Type) {
	case descriptor.FieldDescriptorProto_TYPE_INT32, descriptor.FieldDescriptorProto_TYPE_INT64:
		return true
	case descriptor.FieldDescriptorProto_TYPE_UINT32, descriptor.FieldDescriptorProto_TYPE_UINT64:
		return true
	case descriptor.FieldDescriptorProto_TYPE_SINT32, descriptor.FieldDescriptorProto_TYPE_SINT64:
		return true
	}
	return false
}

func (p *plugin) generateRegexVars(file *generator.FileDescriptor, message *generator.Descriptor) {
	ccTypeName := generator.CamelCaseSlice(message.TypeName())
	for _, field := range message.Field {
		validator := getFieldValidatorIfAny(field)
		if validator != nil && validator.Regex != nil {
			fieldName := p.GetFieldName(message, field)
			p.P(`var `, p.regexName(ccTypeName, fieldName), ` = `, p.regexPkg.Use(), `.MustCompile(`, strconv.Quote(*validator.Regex), `)`)
		}
	}
}

func (p *plugin) generateProto2Message(file *generator.FileDescriptor, message *generator.Descriptor) {
	ccTypeName := generator.CamelCaseSlice(message.TypeName())

	p.P(`func (this *`, ccTypeName, `) Validate() error {`)
	p.In()
	for _, field := range message.Field {
		fieldName := p.GetFieldName(message, field)
		fieldValudator := getFieldValidatorIfAny(field)
		if fieldValudator == nil && !field.IsMessage() {
			continue
		}
		if p.validatorWithMessageExists(fieldValudator) {
			fmt.Fprintf(os.Stderr, "WARNING: field %v.%v is a proto2 message, validator.msg_exists has no effect\n", ccTypeName, fieldName)
		}
		variableName := "this." + fieldName
		repeated := field.IsRepeated()
		nullable := gogoproto.IsNullable(field)
		// For proto2 syntax, only Gogo generates non-pointer fields
		nonpointer := gogoproto.ImportsGoGoProto(file.FileDescriptorProto) && !gogoproto.IsNullable(field)
		if repeated {
			p.P(`for _, item := range `, variableName, `{`)
			p.In()
			variableName = "item"
		} else if nullable {
			p.P(`if `, variableName, ` != nil {`)
			p.In()
			variableName = "*(" + variableName + ")"
		} else if nonpointer {
			// can use the field directly
		} else if !field.IsMessage() {
			variableName = `this.Get` + fieldName + `()`
		}
		if field.IsString() {
			p.generateStringValidator(variableName, ccTypeName, fieldName, fieldValudator)
		} else if p.isSupportedInt(field) {
			p.generateIntValidator(variableName, ccTypeName, fieldName, fieldValudator)
		} else if field.IsMessage() {
			if repeated && nullable {
				variableName = "*(item)"
			}
			p.P(`if err := `, p.validatorPkg.Use(), `.CallValidatorIfExists(&(`, variableName, `)); err != nil {`)
			p.In()
			p.P(`return err`)
			p.Out()
			p.P(`}`)
		}
		if repeated {
			// end the repeated loop
			p.Out()
			p.P(`}`)
		} else if nullable {
			// end the if around nullable
			p.Out()
			p.P(`}`)
		}
	}
	p.P(`return nil`)
	p.Out()
	p.P(`}`)
}

func (p *plugin) generateProto3Message(file *generator.FileDescriptor, message *generator.Descriptor) {
	ccTypeName := generator.CamelCaseSlice(message.TypeName())
	p.P(`func (this *`, ccTypeName, `) Validate() error {`)
	p.In()
	for _, field := range message.Field {
		fieldValidator := getFieldValidatorIfAny(field)
		if fieldValidator == nil && !field.IsMessage() {
			continue
		}
		fieldName := p.GetFieldName(message, field)
		variableName := "this." + fieldName
		repeated := field.IsRepeated()
		// Golang's proto3 has no concept of unset primitive fields
		nullable := (gogoproto.IsNullable(field) || !gogoproto.ImportsGoGoProto(file.FileDescriptorProto)) && field.IsMessage()
		if p.fieldIsProto3Map(file, message, field) {
			p.P(`// Validation of proto3 map<> fields is unsupported.`)
			continue
		}
		if repeated {
			p.P(`for _, item := range `, variableName, `{`)
			p.In()
			variableName = "item"
		}
		if field.IsString() {
			p.generateStringValidator(variableName, ccTypeName, fieldName, fieldValidator)
		} else if p.isSupportedInt(field) {
			p.generateIntValidator(variableName, ccTypeName, fieldName, fieldValidator)
		} else if field.IsMessage() {
			if p.validatorWithMessageExists(fieldValidator) {
				if nullable && !repeated {
					p.P(`if nil == `, variableName, `{`)
					p.In()
					p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, ccTypeName+"."+fieldName, ` message must exist")`)
					p.Out()
					p.P(`}`)
				} else if repeated {
					fmt.Fprintf(os.Stderr, "WARNING: field %v.%v is a repeated, validator.msg_exists has no effect\n", ccTypeName, fieldName)
				} else if !nullable {
					fmt.Fprintf(os.Stderr, "WARNING: field %v.%v is a nullable=false, validator.msg_exists has no effect\n", ccTypeName, fieldName)
				}
			}
			if nullable {
				p.P(`if `, variableName, ` != nil {`)
				p.In()
			} else {
				// non-nullable fields in proto3 store actual structs, we need pointers to operate on interfaces
				variableName = "&(" + variableName + ")"
			}
			p.P(`if err := `, p.validatorPkg.Use(), `.CallValidatorIfExists(`, variableName, `); err != nil {`)
			p.In()
			p.P(`return err`)
			p.Out()
			p.P(`}`)
			if nullable {
				p.Out()
				p.P(`}`)
			}
		}
		if repeated {
			// end the repeated loop
			p.Out()
			p.P(`}`)
		}
	}
	p.P(`return nil`)
	p.Out()
	p.P(`}`)
}

func (p *plugin) generateIntValidator(variableName string, ccTypeName string, fieldName string, fv *validator.FieldValidator) {
	fieldIdentifier := ccTypeName + "." + fieldName
	if fv.IntGt != nil {
		p.P(`if !(`, variableName, ` > `, fv.IntGt, `){`)
		p.In()
		p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, fieldIdentifier, ` must be greater than '`, fv.IntGt, `'")`)
		p.Out()
		p.P(`}`)
	}
	if fv.IntLt != nil {
		p.P(`if !(`, variableName, ` < `, fv.IntLt, `){`)
		p.In()
		p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, fieldIdentifier, ` must be less than '`, fv.IntLt, `'")`)
		p.Out()
		p.P(`}`)
	}
}

func (p *plugin) generateStringValidator(variableName string, ccTypeName string, fieldName string, fv *validator.FieldValidator) {
	fieldIdentifier := ccTypeName + "." + fieldName
	if fv.Regex != nil {
		p.P(`if !`, p.regexName(ccTypeName, fieldName), `.MatchString(`, variableName, `) {`)
		p.In()
		p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, fieldIdentifier, ` must conform to regex " +`, strconv.Quote(*fv.Regex), `)`)
		p.Out()
		p.P(`}`)
	}
}

func (p *plugin) fieldIsProto3Map(file *generator.FileDescriptor, message *generator.Descriptor, field *descriptor.FieldDescriptorProto) bool {
	// Context from descriptor.proto
	// Whether the message is an automatically generated map entry type for the
	// maps field.
	//
	// For maps fields:
	//     map<KeyType, ValueType> map_field = 1;
	// The parsed descriptor looks like:
	//     message MapFieldEntry {
	//         option map_entry = true;
	//         optional KeyType key = 1;
	//         optional ValueType value = 2;
	//     }
	//     repeated MapFieldEntry map_field = 1;
	//
	// Implementations may choose not to generate the map_entry=true message, but
	// use a native map in the target language to hold the keys and values.
	// The reflection APIs in such implementions still need to work as
	// if the field is a repeated message field.
	//
	// NOTE: Do not set the option in .proto files. Always use the maps syntax
	// instead. The option should only be implicitly set by the proto compiler
	// parser.
	if field.GetType() != descriptor.FieldDescriptorProto_TYPE_MESSAGE || !field.IsRepeated() {
		return false
	}
	typeName := field.GetTypeName()
	var msg *descriptor.DescriptorProto
	if strings.HasPrefix(typeName, ".") {
		// Fully qualified case, look up in global map, must work or fail badly.
		msg = p.ObjectNamed(field.GetTypeName()).(*generator.Descriptor).DescriptorProto
	} else {
		// Nested, relative case.
		msg = file.GetNestedMessage(message.DescriptorProto, field.GetTypeName())
	}
	return msg.GetOptions().GetMapEntry()
}

func (p *plugin) validatorWithMessageExists(fv *validator.FieldValidator) bool {
	return fv != nil && fv.MsgExists != nil && *(fv.MsgExists)
}

func (p *plugin) regexName(ccTypeName string, fieldName string) string {
	return "_regex_" + ccTypeName + "_" + fieldName
}
<commit_msg>Remove unnecessary plugin registration.<commit_after>// Copyright 2016 Michal Witkowski. All Rights Reserved.
// See LICENSE for licensing terms.

/*

The validator plugin generates a Validate method for each message.
By default, if none of the message's fields are annotated with the gogo validator annotation, it returns a nil.
In case some of the fields are annotated, the Validate function returns nil upon sucessful validation, or an error
describing why the validation failed.
The Validate method is called recursively for all submessage of the message.

TODO(michal): ADD COMMENTS.

Equal is enabled using the following extensions:

  - equal
  - equal_all

While VerboseEqual is enable dusing the following extensions:

  - verbose_equal
  - verbose_equal_all

The equal plugin also generates a test given it is enabled using one of the following extensions:

  - testgen
  - testgen_all

Let us look at:

  github.com/gogo/protobuf/test/example/example.proto

Btw all the output can be seen at:

  github.com/gogo/protobuf/test/example/*

The following message:

given to the equal plugin, will generate the following code:

and the following test code:

*/
package plugin

import (
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/gogo/protobuf/gogoproto"
	"github.com/gogo/protobuf/proto"
	descriptor "github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
	"github.com/gogo/protobuf/protoc-gen-gogo/generator"
	"github.com/gogo/protobuf/vanity"
	"github.com/mwitkow/go-proto-validators"
)

type plugin struct {
	*generator.Generator
	generator.PluginImports
	regexPkg      generator.Single
	fmtPkg        generator.Single
	protoPkg      generator.Single
	validatorPkg  generator.Single
	useGogoImport bool
}

func NewPlugin(useGogoImport bool) generator.Plugin {
	return &plugin{useGogoImport: useGogoImport}
}

func (p *plugin) Name() string {
	return "validator"
}

func (p *plugin) Init(g *generator.Generator) {
	p.Generator = g
}

func (p *plugin) Generate(file *generator.FileDescriptor) {
	if !p.useGogoImport {
		vanity.TurnOffGogoImport(file.FileDescriptorProto)
	}
	p.PluginImports = generator.NewPluginImports(p.Generator)
	p.regexPkg = p.NewImport("regexp")
	p.fmtPkg = p.NewImport("fmt")
	p.validatorPkg = p.NewImport("github.com/mwitkow/go-proto-validators")

	for _, msg := range file.Messages() {
		if msg.DescriptorProto.GetOptions().GetMapEntry() {
			continue
		}
		p.generateRegexVars(file, msg)
		if gogoproto.IsProto3(file.FileDescriptorProto) {
			p.generateProto3Message(file, msg)
		} else {
			p.generateProto2Message(file, msg)
		}

	}
}

func getFieldValidatorIfAny(field *descriptor.FieldDescriptorProto) *validator.FieldValidator {
	if field.Options != nil {
		v, err := proto.GetExtension(field.Options, validator.E_Field)
		if err == nil && v.(*validator.FieldValidator) != nil {
			return (v.(*validator.FieldValidator))
		}
	}
	return nil
}

func (p *plugin) isSupportedInt(field *descriptor.FieldDescriptorProto) bool {
	switch *(field.Type) {
	case descriptor.FieldDescriptorProto_TYPE_INT32, descriptor.FieldDescriptorProto_TYPE_INT64:
		return true
	case descriptor.FieldDescriptorProto_TYPE_UINT32, descriptor.FieldDescriptorProto_TYPE_UINT64:
		return true
	case descriptor.FieldDescriptorProto_TYPE_SINT32, descriptor.FieldDescriptorProto_TYPE_SINT64:
		return true
	}
	return false
}

func (p *plugin) generateRegexVars(file *generator.FileDescriptor, message *generator.Descriptor) {
	ccTypeName := generator.CamelCaseSlice(message.TypeName())
	for _, field := range message.Field {
		validator := getFieldValidatorIfAny(field)
		if validator != nil && validator.Regex != nil {
			fieldName := p.GetFieldName(message, field)
			p.P(`var `, p.regexName(ccTypeName, fieldName), ` = `, p.regexPkg.Use(), `.MustCompile(`, strconv.Quote(*validator.Regex), `)`)
		}
	}
}

func (p *plugin) generateProto2Message(file *generator.FileDescriptor, message *generator.Descriptor) {
	ccTypeName := generator.CamelCaseSlice(message.TypeName())

	p.P(`func (this *`, ccTypeName, `) Validate() error {`)
	p.In()
	for _, field := range message.Field {
		fieldName := p.GetFieldName(message, field)
		fieldValudator := getFieldValidatorIfAny(field)
		if fieldValudator == nil && !field.IsMessage() {
			continue
		}
		if p.validatorWithMessageExists(fieldValudator) {
			fmt.Fprintf(os.Stderr, "WARNING: field %v.%v is a proto2 message, validator.msg_exists has no effect\n", ccTypeName, fieldName)
		}
		variableName := "this." + fieldName
		repeated := field.IsRepeated()
		nullable := gogoproto.IsNullable(field)
		// For proto2 syntax, only Gogo generates non-pointer fields
		nonpointer := gogoproto.ImportsGoGoProto(file.FileDescriptorProto) && !gogoproto.IsNullable(field)
		if repeated {
			p.P(`for _, item := range `, variableName, `{`)
			p.In()
			variableName = "item"
		} else if nullable {
			p.P(`if `, variableName, ` != nil {`)
			p.In()
			variableName = "*(" + variableName + ")"
		} else if nonpointer {
			// can use the field directly
		} else if !field.IsMessage() {
			variableName = `this.Get` + fieldName + `()`
		}
		if field.IsString() {
			p.generateStringValidator(variableName, ccTypeName, fieldName, fieldValudator)
		} else if p.isSupportedInt(field) {
			p.generateIntValidator(variableName, ccTypeName, fieldName, fieldValudator)
		} else if field.IsMessage() {
			if repeated && nullable {
				variableName = "*(item)"
			}
			p.P(`if err := `, p.validatorPkg.Use(), `.CallValidatorIfExists(&(`, variableName, `)); err != nil {`)
			p.In()
			p.P(`return err`)
			p.Out()
			p.P(`}`)
		}
		if repeated {
			// end the repeated loop
			p.Out()
			p.P(`}`)
		} else if nullable {
			// end the if around nullable
			p.Out()
			p.P(`}`)
		}
	}
	p.P(`return nil`)
	p.Out()
	p.P(`}`)
}

func (p *plugin) generateProto3Message(file *generator.FileDescriptor, message *generator.Descriptor) {
	ccTypeName := generator.CamelCaseSlice(message.TypeName())
	p.P(`func (this *`, ccTypeName, `) Validate() error {`)
	p.In()
	for _, field := range message.Field {
		fieldValidator := getFieldValidatorIfAny(field)
		if fieldValidator == nil && !field.IsMessage() {
			continue
		}
		fieldName := p.GetFieldName(message, field)
		variableName := "this." + fieldName
		repeated := field.IsRepeated()
		// Golang's proto3 has no concept of unset primitive fields
		nullable := (gogoproto.IsNullable(field) || !gogoproto.ImportsGoGoProto(file.FileDescriptorProto)) && field.IsMessage()
		if p.fieldIsProto3Map(file, message, field) {
			p.P(`// Validation of proto3 map<> fields is unsupported.`)
			continue
		}
		if repeated {
			p.P(`for _, item := range `, variableName, `{`)
			p.In()
			variableName = "item"
		}
		if field.IsString() {
			p.generateStringValidator(variableName, ccTypeName, fieldName, fieldValidator)
		} else if p.isSupportedInt(field) {
			p.generateIntValidator(variableName, ccTypeName, fieldName, fieldValidator)
		} else if field.IsMessage() {
			if p.validatorWithMessageExists(fieldValidator) {
				if nullable && !repeated {
					p.P(`if nil == `, variableName, `{`)
					p.In()
					p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, ccTypeName+"."+fieldName, ` message must exist")`)
					p.Out()
					p.P(`}`)
				} else if repeated {
					fmt.Fprintf(os.Stderr, "WARNING: field %v.%v is a repeated, validator.msg_exists has no effect\n", ccTypeName, fieldName)
				} else if !nullable {
					fmt.Fprintf(os.Stderr, "WARNING: field %v.%v is a nullable=false, validator.msg_exists has no effect\n", ccTypeName, fieldName)
				}
			}
			if nullable {
				p.P(`if `, variableName, ` != nil {`)
				p.In()
			} else {
				// non-nullable fields in proto3 store actual structs, we need pointers to operate on interfaces
				variableName = "&(" + variableName + ")"
			}
			p.P(`if err := `, p.validatorPkg.Use(), `.CallValidatorIfExists(`, variableName, `); err != nil {`)
			p.In()
			p.P(`return err`)
			p.Out()
			p.P(`}`)
			if nullable {
				p.Out()
				p.P(`}`)
			}
		}
		if repeated {
			// end the repeated loop
			p.Out()
			p.P(`}`)
		}
	}
	p.P(`return nil`)
	p.Out()
	p.P(`}`)
}

func (p *plugin) generateIntValidator(variableName string, ccTypeName string, fieldName string, fv *validator.FieldValidator) {
	fieldIdentifier := ccTypeName + "." + fieldName
	if fv.IntGt != nil {
		p.P(`if !(`, variableName, ` > `, fv.IntGt, `){`)
		p.In()
		p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, fieldIdentifier, ` must be greater than '`, fv.IntGt, `'")`)
		p.Out()
		p.P(`}`)
	}
	if fv.IntLt != nil {
		p.P(`if !(`, variableName, ` < `, fv.IntLt, `){`)
		p.In()
		p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, fieldIdentifier, ` must be less than '`, fv.IntLt, `'")`)
		p.Out()
		p.P(`}`)
	}
}

func (p *plugin) generateStringValidator(variableName string, ccTypeName string, fieldName string, fv *validator.FieldValidator) {
	fieldIdentifier := ccTypeName + "." + fieldName
	if fv.Regex != nil {
		p.P(`if !`, p.regexName(ccTypeName, fieldName), `.MatchString(`, variableName, `) {`)
		p.In()
		p.P(`return `, p.fmtPkg.Use(), `.Errorf("validation error: `, fieldIdentifier, ` must conform to regex " +`, strconv.Quote(*fv.Regex), `)`)
		p.Out()
		p.P(`}`)
	}
}

func (p *plugin) fieldIsProto3Map(file *generator.FileDescriptor, message *generator.Descriptor, field *descriptor.FieldDescriptorProto) bool {
	// Context from descriptor.proto
	// Whether the message is an automatically generated map entry type for the
	// maps field.
	//
	// For maps fields:
	//     map<KeyType, ValueType> map_field = 1;
	// The parsed descriptor looks like:
	//     message MapFieldEntry {
	//         option map_entry = true;
	//         optional KeyType key = 1;
	//         optional ValueType value = 2;
	//     }
	//     repeated MapFieldEntry map_field = 1;
	//
	// Implementations may choose not to generate the map_entry=true message, but
	// use a native map in the target language to hold the keys and values.
	// The reflection APIs in such implementions still need to work as
	// if the field is a repeated message field.
	//
	// NOTE: Do not set the option in .proto files. Always use the maps syntax
	// instead. The option should only be implicitly set by the proto compiler
	// parser.
	if field.GetType() != descriptor.FieldDescriptorProto_TYPE_MESSAGE || !field.IsRepeated() {
		return false
	}
	typeName := field.GetTypeName()
	var msg *descriptor.DescriptorProto
	if strings.HasPrefix(typeName, ".") {
		// Fully qualified case, look up in global map, must work or fail badly.
		msg = p.ObjectNamed(field.GetTypeName()).(*generator.Descriptor).DescriptorProto
	} else {
		// Nested, relative case.
		msg = file.GetNestedMessage(message.DescriptorProto, field.GetTypeName())
	}
	return msg.GetOptions().GetMapEntry()
}

func (p *plugin) validatorWithMessageExists(fv *validator.FieldValidator) bool {
	return fv != nil && fv.MsgExists != nil && *(fv.MsgExists)
}

func (p *plugin) regexName(ccTypeName string, fieldName string) string {
	return "_regex_" + ccTypeName + "_" + fieldName
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"fmt"
	"testing"

	"github.com/hailocab/go-hostpool"
)

// Tests of the round-robin host selection policy implementation
func TestRoundRobinHostPolicy(t *testing.T) {
	policy := RoundRobinHostPolicy()

	hosts := [...]*HostInfo{
		{hostId: "0"},
		{hostId: "1"},
	}

	for _, host := range hosts {
		policy.AddHost(host)
	}

	// interleaved iteration should always increment the host
	iterA := policy.Pick(nil)
	if actual := iterA(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	iterB := policy.Pick(nil)
	if actual := iterB(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterB(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterA(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}

	iterC := policy.Pick(nil)
	if actual := iterC(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterC(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
}

// Tests of the token-aware host selection policy implementation with a
// round-robin host selection policy fallback.
func TestTokenAwareHostPolicy(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	query := &Query{}

	iter := policy.Pick(nil)
	if iter == nil {
		t.Fatal("host iterator was nil")
	}
	actual := iter()
	if actual != nil {
		t.Fatalf("expected nil from iterator, but was %v", actual)
	}

	// set the hosts
	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}

	// the token ring is not setup without the partitioner, but the fallback
	// should work
	if actual := policy.Pick(nil)(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}

	query.RoutingKey([]byte("30"))
	if actual := policy.Pick(query)(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}

	policy.SetPartitioner("OrderedPartitioner")

	// now the token ring is configured
	query.RoutingKey([]byte("20"))
	iter = policy.Pick(query)
	if actual := iter(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}
	// rest are round robin
	if actual := iter(); actual.Info().Peer() != "2" {
		t.Errorf("Expected peer 2 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "3" {
		t.Errorf("Expected peer 3 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}
}

// Tests of the host pool host selection policy implementation
func TestHostPoolHostPolicy(t *testing.T) {
	policy := HostPoolHostPolicy(hostpool.New(nil))

	hosts := [...]*HostInfo{
		{hostId: "0", peer: "0"},
		{hostId: "1", peer: "1"},
	}

	for _, host := range hosts {
		policy.AddHost(host)
	}

	// the first host selected is actually at [1], but this is ok for RR
	// interleaved iteration should always increment the host
	iter := policy.Pick(nil)
	actualA := iter()
	if actualA.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualA.Info().HostID())
	}
	actualA.Mark(nil)

	actualB := iter()
	if actualB.Info().HostID() != "1" {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actualB.Info().HostID())
	}
	actualB.Mark(fmt.Errorf("error"))

	actualC := iter()
	if actualC.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualC.Info().HostID())
	}
	actualC.Mark(nil)

	actualD := iter()
	if actualD.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualD.Info().HostID())
	}
	actualD.Mark(nil)
}

func TestRoundRobinNilHostInfo(t *testing.T) {
	policy := RoundRobinHostPolicy()

	host := &HostInfo{hostId: "host-1"}
	policy.AddHost(host)

	iter := policy.Pick(nil)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.HostID() != host.HostID() {
		t.Fatalf("expected host %v got %v", host, *v)
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestTokenAwareNilHostInfo(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}
	policy.SetPartitioner("OrderedPartitioner")

	query := &Query{}
	query.RoutingKey([]byte("20"))

	iter := policy.Pick(query)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.Peer() != "1" {
		t.Fatalf("expected peer 1 got %v", v.Peer())
	}

	// Empty the hosts to trigger the panic when using the fallback.
	for _, host := range hosts {
		policy.RemoveHost(host.Peer())
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestCOWList_Add(t *testing.T) {
	var cow cowHostList

	toAdd := [...]string{"peer1", "peer2", "peer3"}

	for _, addr := range toAdd {
		if !cow.add(&HostInfo{peer: addr}) {
			t.Fatal("did not add peer which was not in the set")
		}
	}

	hosts := cow.get()
	if len(hosts) != len(toAdd) {
		t.Fatalf("expected to have %d hosts got %d", len(toAdd), len(hosts))
	}

	set := make(map[string]bool)
	for _, host := range hosts {
		set[host.Peer()] = true
	}

	for _, addr := range toAdd {
		if !set[addr] {
			t.Errorf("addr was not in the host list: %q", addr)
		}
	}
}

func TestSimpleRetryPolicy(t *testing.T) {
	q := &Query{}
	rt := &SimpleRetryPolicy{NumRetries: 2}
	if !rt.Attempt(q) {
		t.Fatal("should allow retry after 0 attempts")
	}
	q.attempts = 5
	if rt.Attempt(q) {
		t.Fatal("should not allow retry after passing threshold")
	}
}
<commit_msg>dont copy value with mutex inside<commit_after>// Copyright (c) 2015 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"fmt"
	"testing"

	"github.com/hailocab/go-hostpool"
)

// Tests of the round-robin host selection policy implementation
func TestRoundRobinHostPolicy(t *testing.T) {
	policy := RoundRobinHostPolicy()

	hosts := [...]*HostInfo{
		{hostId: "0"},
		{hostId: "1"},
	}

	for _, host := range hosts {
		policy.AddHost(host)
	}

	// interleaved iteration should always increment the host
	iterA := policy.Pick(nil)
	if actual := iterA(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	iterB := policy.Pick(nil)
	if actual := iterB(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterB(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterA(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}

	iterC := policy.Pick(nil)
	if actual := iterC(); actual.Info() != hosts[0] {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actual.Info().HostID())
	}
	if actual := iterC(); actual.Info() != hosts[1] {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actual.Info().HostID())
	}
}

// Tests of the token-aware host selection policy implementation with a
// round-robin host selection policy fallback.
func TestTokenAwareHostPolicy(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	query := &Query{}

	iter := policy.Pick(nil)
	if iter == nil {
		t.Fatal("host iterator was nil")
	}
	actual := iter()
	if actual != nil {
		t.Fatalf("expected nil from iterator, but was %v", actual)
	}

	// set the hosts
	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}

	// the token ring is not setup without the partitioner, but the fallback
	// should work
	if actual := policy.Pick(nil)(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}

	query.RoutingKey([]byte("30"))
	if actual := policy.Pick(query)(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}

	policy.SetPartitioner("OrderedPartitioner")

	// now the token ring is configured
	query.RoutingKey([]byte("20"))
	iter = policy.Pick(query)
	if actual := iter(); actual.Info().Peer() != "1" {
		t.Errorf("Expected peer 1 but was %s", actual.Info().Peer())
	}
	// rest are round robin
	if actual := iter(); actual.Info().Peer() != "2" {
		t.Errorf("Expected peer 2 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "3" {
		t.Errorf("Expected peer 3 but was %s", actual.Info().Peer())
	}
	if actual := iter(); actual.Info().Peer() != "0" {
		t.Errorf("Expected peer 0 but was %s", actual.Info().Peer())
	}
}

// Tests of the host pool host selection policy implementation
func TestHostPoolHostPolicy(t *testing.T) {
	policy := HostPoolHostPolicy(hostpool.New(nil))

	hosts := [...]*HostInfo{
		{hostId: "0", peer: "0"},
		{hostId: "1", peer: "1"},
	}

	for _, host := range hosts {
		policy.AddHost(host)
	}

	// the first host selected is actually at [1], but this is ok for RR
	// interleaved iteration should always increment the host
	iter := policy.Pick(nil)
	actualA := iter()
	if actualA.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualA.Info().HostID())
	}
	actualA.Mark(nil)

	actualB := iter()
	if actualB.Info().HostID() != "1" {
		t.Errorf("Expected hosts[1] but was hosts[%s]", actualB.Info().HostID())
	}
	actualB.Mark(fmt.Errorf("error"))

	actualC := iter()
	if actualC.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualC.Info().HostID())
	}
	actualC.Mark(nil)

	actualD := iter()
	if actualD.Info().HostID() != "0" {
		t.Errorf("Expected hosts[0] but was hosts[%s]", actualD.Info().HostID())
	}
	actualD.Mark(nil)
}

func TestRoundRobinNilHostInfo(t *testing.T) {
	policy := RoundRobinHostPolicy()

	host := &HostInfo{hostId: "host-1"}
	policy.AddHost(host)

	iter := policy.Pick(nil)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.HostID() != host.HostID() {
		t.Fatalf("expected host %v got %v", host, v)
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestTokenAwareNilHostInfo(t *testing.T) {
	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())

	hosts := [...]*HostInfo{
		{peer: "0", tokens: []string{"00"}},
		{peer: "1", tokens: []string{"25"}},
		{peer: "2", tokens: []string{"50"}},
		{peer: "3", tokens: []string{"75"}},
	}
	for _, host := range hosts {
		policy.AddHost(host)
	}
	policy.SetPartitioner("OrderedPartitioner")

	query := &Query{}
	query.RoutingKey([]byte("20"))

	iter := policy.Pick(query)
	next := iter()
	if next == nil {
		t.Fatal("got nil host")
	} else if v := next.Info(); v == nil {
		t.Fatal("got nil HostInfo")
	} else if v.Peer() != "1" {
		t.Fatalf("expected peer 1 got %v", v.Peer())
	}

	// Empty the hosts to trigger the panic when using the fallback.
	for _, host := range hosts {
		policy.RemoveHost(host.Peer())
	}

	next = iter()
	if next != nil {
		t.Errorf("expected to get nil host got %+v", next)
		if next.Info() == nil {
			t.Fatalf("HostInfo is nil")
		}
	}
}

func TestCOWList_Add(t *testing.T) {
	var cow cowHostList

	toAdd := [...]string{"peer1", "peer2", "peer3"}

	for _, addr := range toAdd {
		if !cow.add(&HostInfo{peer: addr}) {
			t.Fatal("did not add peer which was not in the set")
		}
	}

	hosts := cow.get()
	if len(hosts) != len(toAdd) {
		t.Fatalf("expected to have %d hosts got %d", len(toAdd), len(hosts))
	}

	set := make(map[string]bool)
	for _, host := range hosts {
		set[host.Peer()] = true
	}

	for _, addr := range toAdd {
		if !set[addr] {
			t.Errorf("addr was not in the host list: %q", addr)
		}
	}
}

func TestSimpleRetryPolicy(t *testing.T) {
	q := &Query{}
	rt := &SimpleRetryPolicy{NumRetries: 2}
	if !rt.Attempt(q) {
		t.Fatal("should allow retry after 0 attempts")
	}
	q.attempts = 5
	if rt.Attempt(q) {
		t.Fatal("should not allow retry after passing threshold")
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Google 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.

/*
Probes package provides an interface to initialize probes using prober config.
*/
package probes

import (
	"context"
	"fmt"
	"regexp"
	"sync"
	"time"

	"cloud.google.com/go/compute/metadata"
	"github.com/golang/glog"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	"github.com/google/cloudprober/probes/dns"
	"github.com/google/cloudprober/probes/external"
	httpprobe "github.com/google/cloudprober/probes/http"
	"github.com/google/cloudprober/probes/ping"
	"github.com/google/cloudprober/probes/udp"
	"github.com/google/cloudprober/targets"
	"github.com/google/cloudprober/targets/lameduck"
)

const (
	logsNamePrefix = "cloudprober"
)

var (
	userDefinedProbes   = make(map[string]Probe)
	userDefinedProbesMu sync.Mutex
)

func runOnThisHost(runOn string, hostname string) bool {
	if runOn == "" {
		return true
	}
	r, err := regexp.Compile(runOn)
	if err != nil {
		glog.Exit(err)
	}
	return r.MatchString(hostname)
}

// Probe interface represents a probe.
//
// A probe is initilized using the Init() method. Init takes the name of the
// probe, a targets.Targets object, interval (how often to run the probe),
// timeout, a logger.Logger handle and the probe type specific config.
//
// Start() method starts the probe. Start is not expected to return for the
// lifetime of the prober. It takes a data channel that it writes the probe
// results on. Actual publishing of these results is handled by cloudprober
// itself.
type Probe interface {
	Init(name string, tgts targets.Targets, interval, timeout time.Duration, l *logger.Logger, config interface{}) error
	Start(ctx context.Context, dataChan chan *metrics.EventMetrics)
}

func newLogger(probeName string) (*logger.Logger, error) {
	return logger.New(context.Background(), logsNamePrefix+"."+probeName)
}

// Init initializes the probes defined in the config.
func Init(probeProtobufs []*ProbeDef, globalTargetsOpts *targets.GlobalTargetsOptions, sysVars map[string]string) map[string]Probe {
	globalTargetsLogger, err := newLogger("globalTargets")
	if err != nil {
		glog.Exitf("Error in initializing logger for the global targets. Err: %v", err)
	}
	if globalTargetsOpts != nil {
		if globalTargetsOpts.GetLameDuckOptions() != nil && metadata.OnGCE() {
			if err := lameduck.Init(globalTargetsOpts.GetLameDuckOptions(), globalTargetsLogger); err != nil {
				glog.Exitf("Error in initializing lameduck module. Err: %v", err)
			}
		}
	}

	probes := make(map[string]Probe)
	for _, p := range probeProtobufs {
		if !runOnThisHost(p.GetRunOn(), sysVars["hostname"]) {
			continue
		}
		if probes[p.GetName()] != nil {
			glog.Exitf("Bad config: probe %s is already defined", p.GetName())
		}
		l, err := newLogger(p.GetName())
		if err != nil {
			glog.Exitf("Error in initializing logger for the probe %s. Err: %v", p.GetName(), err)
		}

		interval := time.Duration(p.GetIntervalMsec()) * time.Millisecond
		timeout := time.Duration(p.GetTimeoutMsec()) * time.Millisecond
		tgts, err := targets.New(p.GetTargets(), globalTargetsOpts, globalTargetsLogger, l)
		if err != nil {
			glog.Exit(err)
		}
		probes[p.GetName()] = initProbe(p, tgts, interval, timeout, l)
	}
	return probes
}

func initProbe(p *ProbeDef, tgts targets.Targets, interval, timeout time.Duration, l *logger.Logger) (probe Probe) {
	glog.Infof("Creating a %s probe: %s", p.GetType(), p.GetName())
	var err error

	switch p.GetType() {
	case ProbeDef_PING:
		probe = &ping.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetPingProbe())
	case ProbeDef_HTTP:
		probe = &httpprobe.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetHttpProbe())
	case ProbeDef_DNS:
		probe = &dns.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetDnsProbe())
	case ProbeDef_EXTERNAL:
		probe = &external.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetExternalProbe())
	case ProbeDef_UDP:
		probe = &udp.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetUdpProbe())
	case ProbeDef_USER_DEFINED:
		userDefinedProbesMu.Lock()
		defer userDefinedProbesMu.Unlock()
		probe := userDefinedProbes[p.GetName()]
		if probe == nil {
			err = fmt.Errorf("unregistered user defined probe: %s", p.GetName())
		} else {
			err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetUserDefinedProbe())
		}
	default:
		glog.Exitf("Unknown probe type: %s", p.GetType())
	}

	if err != nil {
		glog.Exitf("Error in initializing probe %s from the config. Err: %v", p.GetName(), err)
	}
	return
}

// Register allows you to register a user defined probe with cloudprober.
// Example usage:
//	import (
//		"cloudprober/cloudprober"
//		"cloudprober/probes/probes"
//	)
//
//	p := &FancyProbe{}
//	probes.Register("fancy_probe", p)
//	pr, err := cloudprober.InitFromConfig(*configFile)
//	if err != nil {
//		log.Exitf("Error initializing cloudprober. Err: %v", err)
//	}
func Register(name string, probe Probe) {
	userDefinedProbesMu.Lock()
	defer userDefinedProbesMu.Unlock()
	userDefinedProbes[name] = probe
}
<commit_msg>Fix the import path in probes.Register example.<commit_after>// Copyright 2017 Google 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.

/*
Probes package provides an interface to initialize probes using prober config.
*/
package probes

import (
	"context"
	"fmt"
	"regexp"
	"sync"
	"time"

	"cloud.google.com/go/compute/metadata"
	"github.com/golang/glog"
	"github.com/google/cloudprober/logger"
	"github.com/google/cloudprober/metrics"
	"github.com/google/cloudprober/probes/dns"
	"github.com/google/cloudprober/probes/external"
	httpprobe "github.com/google/cloudprober/probes/http"
	"github.com/google/cloudprober/probes/ping"
	"github.com/google/cloudprober/probes/udp"
	"github.com/google/cloudprober/targets"
	"github.com/google/cloudprober/targets/lameduck"
)

const (
	logsNamePrefix = "cloudprober"
)

var (
	userDefinedProbes   = make(map[string]Probe)
	userDefinedProbesMu sync.Mutex
)

func runOnThisHost(runOn string, hostname string) bool {
	if runOn == "" {
		return true
	}
	r, err := regexp.Compile(runOn)
	if err != nil {
		glog.Exit(err)
	}
	return r.MatchString(hostname)
}

// Probe interface represents a probe.
//
// A probe is initilized using the Init() method. Init takes the name of the
// probe, a targets.Targets object, interval (how often to run the probe),
// timeout, a logger.Logger handle and the probe type specific config.
//
// Start() method starts the probe. Start is not expected to return for the
// lifetime of the prober. It takes a data channel that it writes the probe
// results on. Actual publishing of these results is handled by cloudprober
// itself.
type Probe interface {
	Init(name string, tgts targets.Targets, interval, timeout time.Duration, l *logger.Logger, config interface{}) error
	Start(ctx context.Context, dataChan chan *metrics.EventMetrics)
}

func newLogger(probeName string) (*logger.Logger, error) {
	return logger.New(context.Background(), logsNamePrefix+"."+probeName)
}

// Init initializes the probes defined in the config.
func Init(probeProtobufs []*ProbeDef, globalTargetsOpts *targets.GlobalTargetsOptions, sysVars map[string]string) map[string]Probe {
	globalTargetsLogger, err := newLogger("globalTargets")
	if err != nil {
		glog.Exitf("Error in initializing logger for the global targets. Err: %v", err)
	}
	if globalTargetsOpts != nil {
		if globalTargetsOpts.GetLameDuckOptions() != nil && metadata.OnGCE() {
			if err := lameduck.Init(globalTargetsOpts.GetLameDuckOptions(), globalTargetsLogger); err != nil {
				glog.Exitf("Error in initializing lameduck module. Err: %v", err)
			}
		}
	}

	probes := make(map[string]Probe)
	for _, p := range probeProtobufs {
		if !runOnThisHost(p.GetRunOn(), sysVars["hostname"]) {
			continue
		}
		if probes[p.GetName()] != nil {
			glog.Exitf("Bad config: probe %s is already defined", p.GetName())
		}
		l, err := newLogger(p.GetName())
		if err != nil {
			glog.Exitf("Error in initializing logger for the probe %s. Err: %v", p.GetName(), err)
		}

		interval := time.Duration(p.GetIntervalMsec()) * time.Millisecond
		timeout := time.Duration(p.GetTimeoutMsec()) * time.Millisecond
		tgts, err := targets.New(p.GetTargets(), globalTargetsOpts, globalTargetsLogger, l)
		if err != nil {
			glog.Exit(err)
		}
		probes[p.GetName()] = initProbe(p, tgts, interval, timeout, l)
	}
	return probes
}

func initProbe(p *ProbeDef, tgts targets.Targets, interval, timeout time.Duration, l *logger.Logger) (probe Probe) {
	glog.Infof("Creating a %s probe: %s", p.GetType(), p.GetName())
	var err error

	switch p.GetType() {
	case ProbeDef_PING:
		probe = &ping.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetPingProbe())
	case ProbeDef_HTTP:
		probe = &httpprobe.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetHttpProbe())
	case ProbeDef_DNS:
		probe = &dns.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetDnsProbe())
	case ProbeDef_EXTERNAL:
		probe = &external.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetExternalProbe())
	case ProbeDef_UDP:
		probe = &udp.Probe{}
		err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetUdpProbe())
	case ProbeDef_USER_DEFINED:
		userDefinedProbesMu.Lock()
		defer userDefinedProbesMu.Unlock()
		probe := userDefinedProbes[p.GetName()]
		if probe == nil {
			err = fmt.Errorf("unregistered user defined probe: %s", p.GetName())
		} else {
			err = probe.Init(p.GetName(), tgts, interval, timeout, l, p.GetUserDefinedProbe())
		}
	default:
		glog.Exitf("Unknown probe type: %s", p.GetType())
	}

	if err != nil {
		glog.Exitf("Error in initializing probe %s from the config. Err: %v", p.GetName(), err)
	}
	return
}

// Register allows you to register a user defined probe with cloudprober.
// Example usage:
//	import (
//		"github.com/google/cloudprober"
//		"github.com/google/cloudprober/probes"
//	)
//
//	p := &FancyProbe{}
//	probes.Register("fancy_probe", p)
//	pr, err := cloudprober.InitFromConfig(*configFile)
//	if err != nil {
//		log.Exitf("Error initializing cloudprober. Err: %v", err)
//	}
func Register(name string, probe Probe) {
	userDefinedProbesMu.Lock()
	defer userDefinedProbesMu.Unlock()
	userDefinedProbes[name] = probe
}
<|endoftext|>


<commit_before>package prompt

import (
	// Stdlib
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	// Internal
	"github.com/salsita/salsaflow/errs"
	"github.com/salsita/salsaflow/modules/common"

	// Other
	"gopkg.in/salsita/go-pivotaltracker.v0/v5/pivotal"
)

var ErrCanceled = errors.New("operation canceled")

type InvalidInputError struct {
	input string
}

func (i *InvalidInputError) Error() string {
	return "Invalid input: " + i.input
}

type OutOfBoundsError struct {
	input string
}

func (i *OutOfBoundsError) Error() string {
	return "Index out of bounds: " + i.input
}

func Confirm(question string) (bool, error) {
	printQuestion := func() {
		fmt.Print(question)
		fmt.Print(" [y/N]: ")
	}
	printQuestion()

	var line string
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line = strings.ToLower(scanner.Text())
		switch line {
		case "":
			line = "n"
		case "y":
		case "n":
		default:
			printQuestion()
			continue
		}
		break
	}
	if err := scanner.Err(); err != nil {
		return false, err
	}

	return line == "y", nil
}

func Prompt(msg string) (string, error) {
	fmt.Print(msg)
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	if err := scanner.Err(); err != nil {
		return "", err
	}
	return scanner.Text(), nil
}

func PromptIndex(msg string, min, max int) (int, error) {
	line, err := Prompt(msg)
	if err != nil {
		return 0, err
	}
	if line == "" {
		return 0, ErrCanceled
	}

	index, err := strconv.Atoi(line)
	if err != nil {
		return 0, &InvalidInputError{line}
	}

	if index < min || index > max {
		return 0, &OutOfBoundsError{line}
	}

	return index, nil
}

func PromptStory(msg string, stories []common.Story) (common.Story, error) {
	// Print the into message.
	fmt.Println(msg)
	fmt.Println()

	// Present the stories to the user.
	tw := tabwriter.NewWriter(os.Stdout, 0, 8, 4, '\t', 0)
	io.WriteString(tw, "  Index\tStory ID\tStory Title\n")
	io.WriteString(tw, "  =====\t========\t===========\n")
	for i, story := range stories {
		fmt.Fprintf(tw, "  %v\t%v\t%v\n", i, story.ReadableId(), story.Title())
	}
	io.WriteString(tw, "\n")
	tw.Flush()

	// Prompt the user to select a story to assign the commit with.
	task := "Prompt the user to select a story"
	index, err := PromptIndex("Choose a story by inserting its index: ", 0, len(stories)-1)
	if err != nil {
		if err == ErrCanceled {
			return nil, ErrCanceled
		}
		return nil, errs.NewError(task, err, nil)
	}
	return stories[index], nil
}

func ConfirmStories(headerLine string, stories []*pivotal.Story) (bool, error) {
	printStoriesConfirmationDialog(headerLine, stories)

	var line string
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line = strings.ToLower(scanner.Text())
		switch line {
		case "":
			line = "n"
		case "y":
		case "n":
		default:
			printStoriesConfirmationDialog(headerLine, stories)
			continue
		}
		break
	}
	if err := scanner.Err(); err != nil {
		return false, err
	}

	return line == "y", nil
}

func printStoriesConfirmationDialog(headerLine string, stories []*pivotal.Story) {
	tw := tabwriter.NewWriter(os.Stdout, 0, 8, 2, '\t', 0)

	io.WriteString(tw, "\n")
	io.WriteString(tw, headerLine)
	io.WriteString(tw, "\n\n")
	io.WriteString(tw, "Story Name\tStory URL\n")
	io.WriteString(tw, "==========\t=========\n")

	for _, story := range stories {
		fmt.Fprintf(tw, "%v\t%v\n", story.Name, story.URL)
	}

	io.WriteString(tw, "\nDo you want to proceed? [y/N]:")
	tw.Flush()
}
<commit_msg>prompt: Truncate story title when it's too long<commit_after>package prompt

import (
	// Stdlib
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	// Internal
	"github.com/salsita/salsaflow/errs"
	"github.com/salsita/salsaflow/modules/common"

	// Other
	"gopkg.in/salsita/go-pivotaltracker.v0/v5/pivotal"
)

// maxStoryTitleColumnWidth specifies the width of the story title column for story listing.
// The story title is truncated to this width in case it is too long.
const maxStoryTitleColumnWidth = 80

var ErrCanceled = errors.New("operation canceled")

type InvalidInputError struct {
	input string
}

func (i *InvalidInputError) Error() string {
	return "Invalid input: " + i.input
}

type OutOfBoundsError struct {
	input string
}

func (i *OutOfBoundsError) Error() string {
	return "Index out of bounds: " + i.input
}

func Confirm(question string) (bool, error) {
	printQuestion := func() {
		fmt.Print(question)
		fmt.Print(" [y/N]: ")
	}
	printQuestion()

	var line string
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line = strings.ToLower(scanner.Text())
		switch line {
		case "":
			line = "n"
		case "y":
		case "n":
		default:
			printQuestion()
			continue
		}
		break
	}
	if err := scanner.Err(); err != nil {
		return false, err
	}

	return line == "y", nil
}

func Prompt(msg string) (string, error) {
	fmt.Print(msg)
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	if err := scanner.Err(); err != nil {
		return "", err
	}
	return scanner.Text(), nil
}

func PromptIndex(msg string, min, max int) (int, error) {
	line, err := Prompt(msg)
	if err != nil {
		return 0, err
	}
	if line == "" {
		return 0, ErrCanceled
	}

	index, err := strconv.Atoi(line)
	if err != nil {
		return 0, &InvalidInputError{line}
	}

	if index < min || index > max {
		return 0, &OutOfBoundsError{line}
	}

	return index, nil
}

func PromptStory(msg string, stories []common.Story) (common.Story, error) {
	// Print the into message.
	fmt.Println(msg)
	fmt.Println()

	// Present the stories to the user.
	tw := tabwriter.NewWriter(os.Stdout, 0, 8, 4, '\t', 0)
	io.WriteString(tw, "  Index\tStory ID\tStory Title\n")
	io.WriteString(tw, "  =====\t========\t===========\n")
	for i, story := range stories {
		fmt.Fprintf(tw, "  %v\t%v\t%v\n", i, story.ReadableId(), formatStoryTitle(story.Title()))
	}
	io.WriteString(tw, "\n")
	tw.Flush()

	// Prompt the user to select a story to assign the commit with.
	task := "Prompt the user to select a story"
	index, err := PromptIndex("Choose a story by inserting its index: ", 0, len(stories)-1)
	if err != nil {
		if err == ErrCanceled {
			return nil, ErrCanceled
		}
		return nil, errs.NewError(task, err, nil)
	}
	return stories[index], nil
}

func ConfirmStories(headerLine string, stories []*pivotal.Story) (bool, error) {
	printStoriesConfirmationDialog(headerLine, stories)

	var line string
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line = strings.ToLower(scanner.Text())
		switch line {
		case "":
			line = "n"
		case "y":
		case "n":
		default:
			printStoriesConfirmationDialog(headerLine, stories)
			continue
		}
		break
	}
	if err := scanner.Err(); err != nil {
		return false, err
	}

	return line == "y", nil
}

func printStoriesConfirmationDialog(headerLine string, stories []*pivotal.Story) {
	tw := tabwriter.NewWriter(os.Stdout, 0, 8, 2, '\t', 0)

	io.WriteString(tw, "\n")
	io.WriteString(tw, headerLine)
	io.WriteString(tw, "\n\n")
	io.WriteString(tw, "Story Name\tStory URL\n")
	io.WriteString(tw, "==========\t=========\n")

	for _, story := range stories {
		fmt.Fprintf(tw, "%v\t%v\n", story.Name, story.URL)
	}

	io.WriteString(tw, "\nDo you want to proceed? [y/N]:")
	tw.Flush()
}

func formatStoryTitle(title string) string {
	if len(title) < maxStoryTitleColumnWidth {
		return title
	}

	// maxStoryTitleColumnWidth incorporates the trailing " ...",
	// so that is why we subtract len(" ...") when truncating.
	truncatedTitle := title[:maxStoryTitleColumnWidth-4]
	if title[maxStoryTitleColumnWidth-4] != ' ' {
		if i := strings.LastIndex(truncatedTitle, " "); i != -1 {
			truncatedTitle = truncatedTitle[:i]
		}
	}
	return truncatedTitle + " ..."
}
<|endoftext|>