<commit_before>package scheduler

import (
	"fmt"
	"log"

	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// maxServiceScheduleAttempts is used to limit the number of times
	// we will attempt to schedule if we continue to hit conflicts for services.
	maxServiceScheduleAttempts = 5

	// maxBatchScheduleAttempts is used to limit the number of times
	// we will attempt to schedule if we continue to hit conflicts for batch.
	maxBatchScheduleAttempts = 2

	// allocNotNeeded is the status used when a job no longer requires an allocation
	allocNotNeeded = "alloc not needed due to job update"

	// allocMigrating is the status used when we must migrate an allocation
	allocMigrating = "alloc is being migrated"

	// allocUpdating is the status used when a job requires an update
	allocUpdating = "alloc is being updated due to job update"
)

// SetStatusError is used to set the status of the evaluation to the given error
type SetStatusError struct {
	Err        error
	EvalStatus string
}

func (s *SetStatusError) Error() string {
	return s.Err.Error()
}

// GenericScheduler is used for 'service' and 'batch' type jobs. This scheduler is
// designed for long-lived services, and as such spends more time attemping
// to make a high quality placement. This is the primary scheduler for
// most workloads. It also supports a 'batch' mode to optimize for fast decision
// making at the cost of quality.
type GenericScheduler struct {
	logger  *log.Logger
	state   State
	planner Planner
	batch   bool

	eval *structs.Evaluation
	plan *structs.Plan
}

// NewServiceScheduler is a factory function to instantiate a new service scheduler
func NewServiceScheduler(logger *log.Logger, state State, planner Planner) Scheduler {
	s := &GenericScheduler{
		logger:  logger,
		state:   state,
		planner: planner,
		batch:   false,
	}
	return s
}

// NewBatchScheduler is a factory function to instantiate a new batch scheduler
func NewBatchScheduler(logger *log.Logger, state State, planner Planner) Scheduler {
	s := &GenericScheduler{
		logger:  logger,
		state:   state,
		planner: planner,
		batch:   true,
	}
	return s
}

// setStatus is used to update the status of the evaluation
func (s *GenericScheduler) setStatus(status, desc string) error {
	s.logger.Printf("[DEBUG] sched: %#v: setting status to %s (%s)", s.eval, status, desc)
	newEval := s.eval.Copy()
	newEval.Status = status
	newEval.StatusDescription = desc
	return s.planner.UpdateEval(newEval)
}

// Process is used to handle a single evaluation
func (s *GenericScheduler) Process(eval *structs.Evaluation) error {
	// Verify the evaluation trigger reason is understood
	switch eval.TriggeredBy {
	case structs.EvalTriggerJobRegister, structs.EvalTriggerNodeUpdate,
		structs.EvalTriggerJobDeregister:
	default:
		desc := fmt.Sprintf("scheduler cannot handle '%s' evaluation reason",
			eval.TriggeredBy)
		return s.setStatus(structs.EvalStatusFailed, desc)
	}

	// Store the evaluation
	s.eval = eval

	// Retry up to the maxScheduleAttempts
	limit := maxServiceScheduleAttempts
	if s.batch {
		limit = maxBatchScheduleAttempts
	}
	if err := retryMax(limit, s.process); err != nil {
		if statusErr, ok := err.(*SetStatusError); ok {
			return s.setStatus(statusErr.EvalStatus, err.Error())
		}
		return err
	}

	// Update the status to complete
	return s.setStatus(structs.EvalStatusComplete, "")
}

// process is wrapped in retryMax to iteratively run the handler until we have no
// further work or we've made the maximum number of attempts.
func (s *GenericScheduler) process() (bool, error) {
	// Lookup the Job by ID
	job, err := s.state.JobByID(s.eval.JobID)
	if err != nil {
		return false, fmt.Errorf("failed to get job '%s': %v",
			s.eval.JobID, err)
	}

	// Create a plan
	s.plan = s.eval.MakePlan(job)

	// Compute the target job allocations
	if err := s.computeJobAllocs(job); err != nil {
		s.logger.Printf("[ERR] sched: %#v: %v", s.eval, err)
		return false, err
	}

	// If the plan is a no-op, we can bail
	if s.plan.IsNoOp() {
		return true, nil
	}

	// Submit the plan
	result, newState, err := s.planner.SubmitPlan(s.plan)
	if err != nil {
		return false, err
	}

	// If we got a state refresh, try again since we have stale data
	if newState != nil {
		s.logger.Printf("[DEBUG] sched: %#v: refresh forced", s.eval)
		s.state = newState
		return false, nil
	}

	// Try again if the plan was not fully committed, potential conflict
	fullCommit, expected, actual := result.FullCommit(s.plan)
	if !fullCommit {
		s.logger.Printf("[DEBUG] sched: %#v: attempted %d placements, %d placed",
			s.eval, expected, actual)
		return false, nil
	}

	// Success!
	return true, nil
}

// computeJobAllocs is used to reconcile differences between the job,
// existing allocations and node status to update the allocations.
func (s *GenericScheduler) computeJobAllocs(job *structs.Job) error {
	// Materialize all the task groups, job could be missing if deregistered
	var groups map[string]*structs.TaskGroup
	if job != nil {
		groups = materializeTaskGroups(job)
	}

	// Lookup the allocations by JobID
	allocs, err := s.state.AllocsByJob(s.eval.JobID)
	if err != nil {
		return fmt.Errorf("failed to get allocs for job '%s': %v",
			s.eval.JobID, err)
	}

	// Determine the tainted nodes containing job allocs
	tainted, err := taintedNodes(s.state, allocs)
	if err != nil {
		return fmt.Errorf("failed to get tainted nodes for job '%s': %v",
			s.eval.JobID, err)
	}

	// Diff the required and existing allocations
	diff := diffAllocs(job, tainted, groups, allocs)
	s.logger.Printf("[DEBUG] sched: %#v: %#v", s.eval, diff)

	// Add all the allocs to stop
	for _, e := range diff.stop {
		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocNotNeeded)
	}

	// For simplicity, we treat all migrates as an evict + place.
	// XXX: This could probably be done more intelligently?
	for _, e := range diff.migrate {
		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocMigrating)
	}
	diff.place = append(diff.place, diff.migrate...)

	// For simplicity, we treat all updates as an evict + place.
	// XXX: This should be done with rolling in-place updates instead.
	for _, e := range diff.update {
		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocUpdating)
	}
	diff.place = append(diff.place, diff.update...)

	// Nothing remaining to do if placement is not required
	if len(diff.place) == 0 {
		return nil
	}

	// Compute the placements
	return s.computePlacements(job, diff.place)
}

func (s *GenericScheduler) computePlacements(job *structs.Job, place []allocTuple) error {
	// Create an evaluation context
	ctx := NewEvalContext(s.state, s.plan, s.logger)

	// Get the base nodes
	nodes, err := readyNodesInDCs(s.state, job.Datacenters)
	if err != nil {
		return err
	}

	// Construct the placement stack
	stack := NewGenericStack(s.batch, ctx, nodes)

	// Track the failed task groups so that we can coalesce
	// the failures together to avoid creating many failed allocs.
	failedTG := make(map[*structs.TaskGroup]*structs.Allocation)

	for _, missing := range place {
		// Check if this task group has already failed
		if alloc, ok := failedTG[missing.TaskGroup]; ok {
			alloc.Metrics.CoalescedFailures += 1
			continue
		}

		// Attempt to match the task group
		option, size := stack.Select(missing.TaskGroup)

		// Handle a placement failure
		var nodeID, status, desc, clientStatus string
		if option == nil {
			status = structs.AllocDesiredStatusFailed
			desc = "failed to find a node for placement"
			clientStatus = structs.AllocClientStatusFailed
		} else {
			nodeID = option.Node.ID
			status = structs.AllocDesiredStatusRun
			clientStatus = structs.AllocClientStatusPending
		}

		// Create an allocation for this
		alloc := &structs.Allocation{
			ID:                 mock.GenerateUUID(),
			EvalID:             s.eval.ID,
			Name:               missing.Name,
			NodeID:             nodeID,
			JobID:              job.ID,
			Job:                job,
			TaskGroup:          missing.TaskGroup.Name,
			Resources:          size,
			Metrics:            ctx.Metrics(),
			DesiredStatus:      status,
			DesiredDescription: desc,
			ClientStatus:       clientStatus,
		}
		if nodeID != "" {
			s.plan.AppendAlloc(alloc)
		} else {
			s.plan.AppendFailed(alloc)
			failedTG[missing.TaskGroup] = alloc
		}
	}
	return nil
}
<commit_msg>scheduler: share context and stack<commit_after>package scheduler

import (
	"fmt"
	"log"

	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// maxServiceScheduleAttempts is used to limit the number of times
	// we will attempt to schedule if we continue to hit conflicts for services.
	maxServiceScheduleAttempts = 5

	// maxBatchScheduleAttempts is used to limit the number of times
	// we will attempt to schedule if we continue to hit conflicts for batch.
	maxBatchScheduleAttempts = 2

	// allocNotNeeded is the status used when a job no longer requires an allocation
	allocNotNeeded = "alloc not needed due to job update"

	// allocMigrating is the status used when we must migrate an allocation
	allocMigrating = "alloc is being migrated"

	// allocUpdating is the status used when a job requires an update
	allocUpdating = "alloc is being updated due to job update"
)

// SetStatusError is used to set the status of the evaluation to the given error
type SetStatusError struct {
	Err        error
	EvalStatus string
}

func (s *SetStatusError) Error() string {
	return s.Err.Error()
}

// GenericScheduler is used for 'service' and 'batch' type jobs. This scheduler is
// designed for long-lived services, and as such spends more time attemping
// to make a high quality placement. This is the primary scheduler for
// most workloads. It also supports a 'batch' mode to optimize for fast decision
// making at the cost of quality.
type GenericScheduler struct {
	logger  *log.Logger
	state   State
	planner Planner
	batch   bool

	eval  *structs.Evaluation
	plan  *structs.Plan
	ctx   *EvalContext
	stack *GenericStack
}

// NewServiceScheduler is a factory function to instantiate a new service scheduler
func NewServiceScheduler(logger *log.Logger, state State, planner Planner) Scheduler {
	s := &GenericScheduler{
		logger:  logger,
		state:   state,
		planner: planner,
		batch:   false,
	}
	return s
}

// NewBatchScheduler is a factory function to instantiate a new batch scheduler
func NewBatchScheduler(logger *log.Logger, state State, planner Planner) Scheduler {
	s := &GenericScheduler{
		logger:  logger,
		state:   state,
		planner: planner,
		batch:   true,
	}
	return s
}

// setStatus is used to update the status of the evaluation
func (s *GenericScheduler) setStatus(status, desc string) error {
	s.logger.Printf("[DEBUG] sched: %#v: setting status to %s (%s)", s.eval, status, desc)
	newEval := s.eval.Copy()
	newEval.Status = status
	newEval.StatusDescription = desc
	return s.planner.UpdateEval(newEval)
}

// Process is used to handle a single evaluation
func (s *GenericScheduler) Process(eval *structs.Evaluation) error {
	// Verify the evaluation trigger reason is understood
	switch eval.TriggeredBy {
	case structs.EvalTriggerJobRegister, structs.EvalTriggerNodeUpdate,
		structs.EvalTriggerJobDeregister:
	default:
		desc := fmt.Sprintf("scheduler cannot handle '%s' evaluation reason",
			eval.TriggeredBy)
		return s.setStatus(structs.EvalStatusFailed, desc)
	}

	// Store the evaluation
	s.eval = eval

	// Retry up to the maxScheduleAttempts
	limit := maxServiceScheduleAttempts
	if s.batch {
		limit = maxBatchScheduleAttempts
	}
	if err := retryMax(limit, s.process); err != nil {
		if statusErr, ok := err.(*SetStatusError); ok {
			return s.setStatus(statusErr.EvalStatus, err.Error())
		}
		return err
	}

	// Update the status to complete
	return s.setStatus(structs.EvalStatusComplete, "")
}

// process is wrapped in retryMax to iteratively run the handler until we have no
// further work or we've made the maximum number of attempts.
func (s *GenericScheduler) process() (bool, error) {
	// Lookup the Job by ID
	job, err := s.state.JobByID(s.eval.JobID)
	if err != nil {
		return false, fmt.Errorf("failed to get job '%s': %v",
			s.eval.JobID, err)
	}

	// Create a plan
	s.plan = s.eval.MakePlan(job)

	// Create an evaluation context
	s.ctx = NewEvalContext(s.state, s.plan, s.logger)

	// Construct the placement stack
	s.stack = NewGenericStack(s.batch, s.ctx, nil)

	// Compute the target job allocations
	if err := s.computeJobAllocs(job); err != nil {
		s.logger.Printf("[ERR] sched: %#v: %v", s.eval, err)
		return false, err
	}

	// If the plan is a no-op, we can bail
	if s.plan.IsNoOp() {
		return true, nil
	}

	// Submit the plan
	result, newState, err := s.planner.SubmitPlan(s.plan)
	if err != nil {
		return false, err
	}

	// If we got a state refresh, try again since we have stale data
	if newState != nil {
		s.logger.Printf("[DEBUG] sched: %#v: refresh forced", s.eval)
		s.state = newState
		return false, nil
	}

	// Try again if the plan was not fully committed, potential conflict
	fullCommit, expected, actual := result.FullCommit(s.plan)
	if !fullCommit {
		s.logger.Printf("[DEBUG] sched: %#v: attempted %d placements, %d placed",
			s.eval, expected, actual)
		return false, nil
	}

	// Success!
	return true, nil
}

// computeJobAllocs is used to reconcile differences between the job,
// existing allocations and node status to update the allocations.
func (s *GenericScheduler) computeJobAllocs(job *structs.Job) error {
	// Materialize all the task groups, job could be missing if deregistered
	var groups map[string]*structs.TaskGroup
	if job != nil {
		groups = materializeTaskGroups(job)
	}

	// Lookup the allocations by JobID
	allocs, err := s.state.AllocsByJob(s.eval.JobID)
	if err != nil {
		return fmt.Errorf("failed to get allocs for job '%s': %v",
			s.eval.JobID, err)
	}

	// Determine the tainted nodes containing job allocs
	tainted, err := taintedNodes(s.state, allocs)
	if err != nil {
		return fmt.Errorf("failed to get tainted nodes for job '%s': %v",
			s.eval.JobID, err)
	}

	// Diff the required and existing allocations
	diff := diffAllocs(job, tainted, groups, allocs)
	s.logger.Printf("[DEBUG] sched: %#v: %#v", s.eval, diff)

	// Add all the allocs to stop
	for _, e := range diff.stop {
		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocNotNeeded)
	}

	// For simplicity, we treat all migrates as an evict + place.
	// XXX: This could probably be done more intelligently?
	for _, e := range diff.migrate {
		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocMigrating)
	}
	diff.place = append(diff.place, diff.migrate...)

	// For simplicity, we treat all updates as an evict + place.
	// XXX: This should be done with rolling in-place updates instead.
	for _, e := range diff.update {
		s.plan.AppendUpdate(e.Alloc, structs.AllocDesiredStatusStop, allocUpdating)
	}
	diff.place = append(diff.place, diff.update...)

	// Nothing remaining to do if placement is not required
	if len(diff.place) == 0 {
		return nil
	}

	// Compute the placements
	return s.computePlacements(job, diff.place)
}

func (s *GenericScheduler) computePlacements(job *structs.Job, place []allocTuple) error {
	// Create an evaluation context
	ctx := NewEvalContext(s.state, s.plan, s.logger)

	// Get the base nodes
	nodes, err := readyNodesInDCs(s.state, job.Datacenters)
	if err != nil {
		return err
	}

	// Update the set of placement ndoes
	s.stack.SetNodes(nodes)

	// Track the failed task groups so that we can coalesce
	// the failures together to avoid creating many failed allocs.
	failedTG := make(map[*structs.TaskGroup]*structs.Allocation)

	for _, missing := range place {
		// Check if this task group has already failed
		if alloc, ok := failedTG[missing.TaskGroup]; ok {
			alloc.Metrics.CoalescedFailures += 1
			continue
		}

		// Attempt to match the task group
		option, size := s.stack.Select(missing.TaskGroup)

		// Handle a placement failure
		var nodeID, status, desc, clientStatus string
		if option == nil {
			status = structs.AllocDesiredStatusFailed
			desc = "failed to find a node for placement"
			clientStatus = structs.AllocClientStatusFailed
		} else {
			nodeID = option.Node.ID
			status = structs.AllocDesiredStatusRun
			clientStatus = structs.AllocClientStatusPending
		}

		// Create an allocation for this
		alloc := &structs.Allocation{
			ID:                 mock.GenerateUUID(),
			EvalID:             s.eval.ID,
			Name:               missing.Name,
			NodeID:             nodeID,
			JobID:              job.ID,
			Job:                job,
			TaskGroup:          missing.TaskGroup.Name,
			Resources:          size,
			Metrics:            s.ctx.Metrics(),
			DesiredStatus:      status,
			DesiredDescription: desc,
			ClientStatus:       clientStatus,
		}
		if nodeID != "" {
			s.plan.AppendAlloc(alloc)
		} else {
			s.plan.AppendFailed(alloc)
			failedTG[missing.TaskGroup] = alloc
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016-2020 The Libsacloud Authors
//
// 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 sacloud

import (
	"encoding/json"

	"github.com/sacloud/libsacloud/v2/sacloud/naked"
	"github.com/sacloud/libsacloud/v2/sacloud/search"
)

// UnmarshalJSON APIからの戻り値でレスポンスボディ直下にデータを持つことへの対応
func (a *authStatusReadResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias authStatusReadResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedAuthStatus naked.AuthStatus
	if err := json.Unmarshal(data, &nakedAuthStatus); err != nil {
		return err
	}
	tmp.AuthStatus = &nakedAuthStatus

	*a = authStatusReadResponseEnvelope(tmp)
	return nil
}

func (b *billDetailsCSVResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias billDetailsCSVResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedBillDetailCSV naked.BillDetailCSV
	if err := json.Unmarshal(data, &nakedBillDetailCSV); err != nil {
		return err
	}
	tmp.CSV = &nakedBillDetailCSV

	*b = billDetailsCSVResponseEnvelope(tmp)
	return nil
}

func (m *mobileGatewaySetSIMRoutesRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias struct {
		SIMRoutes []*naked.MobileGatewaySIMRoute `json:"sim_routes"`
	}
	tmp := &alias{
		SIMRoutes: m.SIMRoutes,
	}
	if len(tmp.SIMRoutes) == 0 {
		tmp.SIMRoutes = make([]*naked.MobileGatewaySIMRoute, 0)
	}
	return json.Marshal(tmp)
}

// UnmarshalJSON APIからの戻り値でレスポンスボディ直下にデータを持つことへの対応
func (s *serverGetVNCProxyResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias serverGetVNCProxyResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedVNCProxy naked.VNCProxyInfo
	if err := json.Unmarshal(data, &nakedVNCProxy); err != nil {
		return err
	}
	tmp.VNCProxyInfo = &nakedVNCProxy

	*s = serverGetVNCProxyResponseEnvelope(tmp)
	return nil
}

/*
 * 検索時に固定パラメータを設定するための実装
 */

func (s autoBackupFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias autoBackupFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "autobackup"
	return json.Marshal(tmp)
}

func (s containerRegistryFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias containerRegistryFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "containerregistry"
	return json.Marshal(tmp)
}

func (s dNSFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias dNSFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "dns"
	return json.Marshal(tmp)
}

func (s simpleMonitorFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias simpleMonitorFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "simplemon"
	return json.Marshal(tmp)
}

func (s gSLBFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias gSLBFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "gslb"
	return json.Marshal(tmp)
}

func (s proxyLBFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias proxyLBFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "proxylb"
	return json.Marshal(tmp)
}

func (s sIMFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias sIMFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "sim"
	return json.Marshal(tmp)
}

func (s localRouterFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias localRouterFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "localrouter"
	return json.Marshal(tmp)
}

func (s databaseFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias databaseFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "database"
	return json.Marshal(tmp)
}

func (s loadBalancerFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias loadBalancerFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "loadbalancer"
	return json.Marshal(tmp)
}

func (s vPCRouterFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias vPCRouterFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "vpcrouter"
	return json.Marshal(tmp)
}

func (s nFSFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias nFSFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "nfs"
	return json.Marshal(tmp)
}

func (s mobileGatewayFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias mobileGatewayFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "mobilegateway"
	return json.Marshal(tmp)
}

/*
 * for Shared Archive
 */

func (s archiveShareRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias archiveShareRequestEnvelope
	tmp := alias(s)
	tmp.Shared = true
	return json.Marshal(tmp)
}

// UnmarshalJSON APIからの戻り値でレスポンスボディ直下にデータを持つことへの対応
func (s *archiveShareResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias archiveShareResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedData naked.ArchiveShareInfo
	if err := json.Unmarshal(data, &nakedData); err != nil {
		return err
	}
	tmp.ArchiveShareInfo = &nakedData

	*s = archiveShareResponseEnvelope(tmp)
	return nil
}
<commit_msg>Add const parameter for finding ESME<commit_after>// Copyright 2016-2020 The Libsacloud Authors
//
// 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 sacloud

import (
	"encoding/json"

	"github.com/sacloud/libsacloud/v2/sacloud/naked"
	"github.com/sacloud/libsacloud/v2/sacloud/search"
)

// UnmarshalJSON APIからの戻り値でレスポンスボディ直下にデータを持つことへの対応
func (a *authStatusReadResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias authStatusReadResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedAuthStatus naked.AuthStatus
	if err := json.Unmarshal(data, &nakedAuthStatus); err != nil {
		return err
	}
	tmp.AuthStatus = &nakedAuthStatus

	*a = authStatusReadResponseEnvelope(tmp)
	return nil
}

func (b *billDetailsCSVResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias billDetailsCSVResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedBillDetailCSV naked.BillDetailCSV
	if err := json.Unmarshal(data, &nakedBillDetailCSV); err != nil {
		return err
	}
	tmp.CSV = &nakedBillDetailCSV

	*b = billDetailsCSVResponseEnvelope(tmp)
	return nil
}

func (m *mobileGatewaySetSIMRoutesRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias struct {
		SIMRoutes []*naked.MobileGatewaySIMRoute `json:"sim_routes"`
	}
	tmp := &alias{
		SIMRoutes: m.SIMRoutes,
	}
	if len(tmp.SIMRoutes) == 0 {
		tmp.SIMRoutes = make([]*naked.MobileGatewaySIMRoute, 0)
	}
	return json.Marshal(tmp)
}

// UnmarshalJSON APIからの戻り値でレスポンスボディ直下にデータを持つことへの対応
func (s *serverGetVNCProxyResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias serverGetVNCProxyResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedVNCProxy naked.VNCProxyInfo
	if err := json.Unmarshal(data, &nakedVNCProxy); err != nil {
		return err
	}
	tmp.VNCProxyInfo = &nakedVNCProxy

	*s = serverGetVNCProxyResponseEnvelope(tmp)
	return nil
}

/*
 * 検索時に固定パラメータを設定するための実装
 */

func (s autoBackupFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias autoBackupFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "autobackup"
	return json.Marshal(tmp)
}

func (s containerRegistryFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias containerRegistryFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "containerregistry"
	return json.Marshal(tmp)
}

func (s eSMEFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias eSMEFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "esme"
	return json.Marshal(tmp)
}

func (s dNSFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias dNSFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "dns"
	return json.Marshal(tmp)
}

func (s simpleMonitorFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias simpleMonitorFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "simplemon"
	return json.Marshal(tmp)
}

func (s gSLBFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias gSLBFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "gslb"
	return json.Marshal(tmp)
}

func (s proxyLBFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias proxyLBFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "proxylb"
	return json.Marshal(tmp)
}

func (s sIMFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias sIMFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "sim"
	return json.Marshal(tmp)
}

func (s localRouterFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias localRouterFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Provider.Class")] = "localrouter"
	return json.Marshal(tmp)
}

func (s databaseFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias databaseFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "database"
	return json.Marshal(tmp)
}

func (s loadBalancerFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias loadBalancerFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "loadbalancer"
	return json.Marshal(tmp)
}

func (s vPCRouterFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias vPCRouterFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "vpcrouter"
	return json.Marshal(tmp)
}

func (s nFSFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias nFSFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "nfs"
	return json.Marshal(tmp)
}

func (s mobileGatewayFindRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias mobileGatewayFindRequestEnvelope
	tmp := alias(s)
	if tmp.Filter == nil {
		tmp.Filter = search.Filter{}
	}
	tmp.Filter[search.Key("Class")] = "mobilegateway"
	return json.Marshal(tmp)
}

/*
 * for Shared Archive
 */

func (s archiveShareRequestEnvelope) MarshalJSON() ([]byte, error) {
	type alias archiveShareRequestEnvelope
	tmp := alias(s)
	tmp.Shared = true
	return json.Marshal(tmp)
}

// UnmarshalJSON APIからの戻り値でレスポンスボディ直下にデータを持つことへの対応
func (s *archiveShareResponseEnvelope) UnmarshalJSON(data []byte) error {
	type alias archiveShareResponseEnvelope

	var tmp alias
	if err := json.Unmarshal(data, &tmp); err != nil {
		return err
	}

	var nakedData naked.ArchiveShareInfo
	if err := json.Unmarshal(data, &nakedData); err != nil {
		return err
	}
	tmp.ArchiveShareInfo = &nakedData

	*s = archiveShareResponseEnvelope(tmp)
	return nil
}
<|endoftext|>


<commit_before>package controller

import (
	"fmt"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"strings"

	"github.com/hkohlsaat/vtr/model"
	"github.com/julienschmidt/httprouter"
)

func PostPlan(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	r.ParseMultipartForm(65536)
	if r.Form.Get("passwort") != os.Args[2] {
		http.Error(w, "falsches Passwort", http.StatusUnauthorized)
		return
	}
	file, _, err := r.FormFile("vertretungsplan")
	if err != nil {
		log.Printf("error recieving file from upload: %v\n", err)
		return
	}
	go processPlan(file)
}

func processPlan(file multipart.File) {
	defer file.Close()

	plan, err := model.ToPlan(file)
	if err != nil {
		log.Printf("can't make an object of the plan: %v\n", err)
		return
	}

	file.Seek(0, 0)
	bytes, _ := ioutil.ReadAll(file)
	plan.Create(bytes)

	for _, part := range plan.Parts {
		for _, s := range part.Substitutions {
			if !s.SubstTeacher.Exists() {
				unknownTeacher := model.UnknownTeacher{Short: s.SubstTeacher.Short}
				unknownTeacher.Create()
			}
			if !s.InstdTeacher.Exists() {
				unknownTeacher := model.UnknownTeacher{Short: s.InstdTeacher.Short}
				unknownTeacher.Create()
			}
			if !s.InstdSubject.Exists() {
				unknownSubject := model.UnknownSubject{Short: s.InstdSubject.Short}
				unknownSubject.Create()
			}
		}
	}

	if len(os.Args) > 3 {
		//	sentToFirebase()
	}
}

func sentToFirebase() {
	json := `{
		"to":"/topics/newplan",
		"collapse_key":"newplan",
		"delay_while_idle":true,
		"data":{"event":"newplan"}
		}`
	client := &http.Client{}
	req, _ := http.NewRequest("POST", "https://fcm.googleapis.com/fcm/send", strings.NewReader(json))
	req.Header.Set("content-type", "application/json")
	req.Header.Add("Authorization", "key="+os.Args[3])
	resp, err := client.Do(req)
	//	if resp.StatusCode != 200 {
	log.Printf("calling firebase: response code: %d", resp.StatusCode)
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)
	log.Print(string(body))
	//	}
	if err != nil {
		log.Printf("error calling firebase: %v", err)
	}
}

func GetPlan(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	lastPlan := model.LastPlanJSON()
	if lastPlan == "" {
		http.NotFound(w, r)
		return
	} else {
		w.Header().Set("content-type", "application/json; charset=utf-8")
		fmt.Fprintln(w, lastPlan)
	}
}
<commit_msg>enabling push notifications<commit_after>package controller

import (
	"fmt"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"strings"

	"github.com/hkohlsaat/vtr/model"
	"github.com/julienschmidt/httprouter"
)

func PostPlan(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	r.ParseMultipartForm(65536)
	if r.Form.Get("passwort") != os.Args[2] {
		http.Error(w, "falsches Passwort", http.StatusUnauthorized)
		return
	}
	file, _, err := r.FormFile("vertretungsplan")
	if err != nil {
		log.Printf("error recieving file from upload: %v\n", err)
		return
	}
	go processPlan(file)
}

func processPlan(file multipart.File) {
	defer file.Close()

	plan, err := model.ToPlan(file)
	if err != nil {
		log.Printf("can't make an object of the plan: %v\n", err)
		return
	}

	file.Seek(0, 0)
	bytes, _ := ioutil.ReadAll(file)
	plan.Create(bytes)

	for _, part := range plan.Parts {
		for _, s := range part.Substitutions {
			if !s.SubstTeacher.Exists() {
				unknownTeacher := model.UnknownTeacher{Short: s.SubstTeacher.Short}
				unknownTeacher.Create()
			}
			if !s.InstdTeacher.Exists() {
				unknownTeacher := model.UnknownTeacher{Short: s.InstdTeacher.Short}
				unknownTeacher.Create()
			}
			if !s.InstdSubject.Exists() {
				unknownSubject := model.UnknownSubject{Short: s.InstdSubject.Short}
				unknownSubject.Create()
			}
		}
	}

	if len(os.Args) > 3 {
		sentToFirebase()
	}
}

func sentToFirebase() {
	json := `{
		"to":"/topics/newplan",
		"collapse_key":"newplan",
		"delay_while_idle":true,
		"data":{"event":"newplan"}
		}`
	client := &http.Client{}
	req, _ := http.NewRequest("POST", "https://fcm.googleapis.com/fcm/send", strings.NewReader(json))
	req.Header.Set("content-type", "application/json")
	req.Header.Add("Authorization", "key="+os.Args[3])
	resp, err := client.Do(req)
	if resp.StatusCode != 200 {
		log.Printf("calling firebase: response code: %d", resp.StatusCode)
		defer resp.Body.Close()
		body, _ := ioutil.ReadAll(resp.Body)
		log.Print(string(body))
	}
	if err != nil {
		log.Printf("error calling firebase: %v", err)
	}
}

func GetPlan(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	lastPlan := model.LastPlanJSON()
	if lastPlan == "" {
		http.NotFound(w, r)
		return
	} else {
		w.Header().Set("content-type", "application/json; charset=utf-8")
		fmt.Fprintln(w, lastPlan)
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 oglemock_test

import (
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type errorReport struct {
	fileName string
	lineNumber int
	err error
}

type fakeErrorReporter struct {
	errorsReported []errorReport
}

func (r *fakeErrorReporter) ReportError(fileName string, lineNumber int, err error) {
	report := errorReport{fileName, lineNumber, err}
	r.errorsReported = append(r.errorsReported, report)
}

type trivialMockObject struct {
	id uintptr
	desc string
}

func (o *trivialMockObject) Oglemock_Id() uintptr {
	return o.id
}

func (o *trivialMockObject) Oglemock_Description() string {
	return o.desc
}

type ControllerTest struct {
	reporter fakeErrorReporter
	controller Controller

	mock1 MockObject
	mock2 MockObject
}

func (t *ControllerTest) SetUp() {
	t.reporter.errorsReported = make([]errorReport, 0)
	t.controller = NewController(&t.reporter)

	t.mock1 = &trivialMockObject{17, "taco"}
	t.mock2 = &trivialMockObject{19, "burrito"}
}

func init() { RegisterTestSuite(&ControllerTest{}) }

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func (t *ControllerTest) FinishWithoutAnyEvents() {
	t.controller.Finish()
	ExpectThat(len(t.reporter.errorsReported), Equals(0))
}

func (t *ControllerTest) HandleCallForUnknownObject() {
	p := []byte{255}
	t.controller.HandleMethodCall(
		t.mock1,
		"Read",
		"taco.go",
		112,
		[]interface{}{p})

	// The error should be reported immediately.
	ExpectThat(len(t.reporter.errorsReported), Equals(1))
	ExpectThat(t.reporter.errorsReported[0].fileName, Equals("taco.go"))
	ExpectThat(t.reporter.errorsReported[0].lineNumber, Equals(112))
	ExpectThat(t.reporter.errorsReported[0].err, Error(HasSubstr("Unexpected")))
	ExpectThat(t.reporter.errorsReported[0].err, Error(HasSubstr("Read")))
	ExpectThat(t.reporter.errorsReported[0].err, Error(HasSubstr("[255]")))
}

func (t *ControllerTest) ExpectCallForUnknownMethod() {
}

func (t *ControllerTest) PartialExpectationGivenWrongNumberOfArgs() {
}

func (t *ControllerTest) PartialExpectationCalledTwice() {
}

func (t *ControllerTest) ExpectThenNonMatchingCall() {
}

func (t *ControllerTest) ExplicitCardinalityNotSatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionCountNotSatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionLowerBoundNotSatisfied() {
}

func (t *ControllerTest) ImplicitCardinalityOfOneNotSatisfied() {
}

func (t *ControllerTest) ExplicitCardinalitySatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionCountSatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionLowerBoundSatisfied() {
}

func (t *ControllerTest) FallbackActionConfiguredWithZeroCalls() {
}

func (t *ControllerTest) FallbackActionConfiguredWithMultipleCalls() {
}

func (t *ControllerTest) ImplicitCardinalityOfOneSatisfied() {
}

func (t *ControllerTest) InvokesOneTimeActions() {
}

func (t *ControllerTest) InvokesFallbackActions() {
}

func (t *ControllerTest) InvokesImplicitActions() {
}

func (t *ControllerTest) ExpectationsAreMatchedLastToFirst() {
}

func (t *ControllerTest) ExpectationsAreSegregatedByMockObject() {
}

func (t *ControllerTest) ExpectationsAreSegregatedByMethodName() {
}
<commit_msg>ControllerTest.ExpectCallForUnknownMethod<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 oglemock_test

import (
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type errorReport struct {
	fileName string
	lineNumber int
	err error
}

type fakeErrorReporter struct {
	errorsReported []errorReport
}

func (r *fakeErrorReporter) ReportError(fileName string, lineNumber int, err error) {
	report := errorReport{fileName, lineNumber, err}
	r.errorsReported = append(r.errorsReported, report)
}

type trivialMockObject struct {
	id uintptr
	desc string
}

func (o *trivialMockObject) Oglemock_Id() uintptr {
	return o.id
}

func (o *trivialMockObject) Oglemock_Description() string {
	return o.desc
}

type ControllerTest struct {
	reporter fakeErrorReporter
	controller Controller

	mock1 MockObject
	mock2 MockObject
}

func (t *ControllerTest) SetUp() {
	t.reporter.errorsReported = make([]errorReport, 0)
	t.controller = NewController(&t.reporter)

	t.mock1 = &trivialMockObject{17, "taco"}
	t.mock2 = &trivialMockObject{19, "burrito"}
}

func init() { RegisterTestSuite(&ControllerTest{}) }

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func (t *ControllerTest) FinishWithoutAnyEvents() {
	t.controller.Finish()
	ExpectThat(len(t.reporter.errorsReported), Equals(0))
}

func (t *ControllerTest) HandleCallForUnknownObject() {
	p := []byte{255}
	t.controller.HandleMethodCall(
		t.mock1,
		"Read",
		"taco.go",
		112,
		[]interface{}{p})

	// The error should be reported immediately.
	ExpectThat(len(t.reporter.errorsReported), Equals(1))
	ExpectThat(t.reporter.errorsReported[0].fileName, Equals("taco.go"))
	ExpectThat(t.reporter.errorsReported[0].lineNumber, Equals(112))
	ExpectThat(t.reporter.errorsReported[0].err, Error(HasSubstr("Unexpected")))
	ExpectThat(t.reporter.errorsReported[0].err, Error(HasSubstr("Read")))
	ExpectThat(t.reporter.errorsReported[0].err, Error(HasSubstr("[255]")))
}

func (t *ControllerTest) ExpectCallForUnknownMethod() {
	ExpectThat(
		func() {
			t.controller.ExpectCall(t.mock1, "Frobnicate")
		},
		Panics(HasSubstr("Unknown method: Frobnicate")))
}

func (t *ControllerTest) PartialExpectationGivenWrongNumberOfArgs() {
}

func (t *ControllerTest) PartialExpectationCalledTwice() {
}

func (t *ControllerTest) ExpectThenNonMatchingCall() {
}

func (t *ControllerTest) ExplicitCardinalityNotSatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionCountNotSatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionLowerBoundNotSatisfied() {
}

func (t *ControllerTest) ImplicitCardinalityOfOneNotSatisfied() {
}

func (t *ControllerTest) ExplicitCardinalitySatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionCountSatisfied() {
}

func (t *ControllerTest) ImplicitOneTimeActionLowerBoundSatisfied() {
}

func (t *ControllerTest) FallbackActionConfiguredWithZeroCalls() {
}

func (t *ControllerTest) FallbackActionConfiguredWithMultipleCalls() {
}

func (t *ControllerTest) ImplicitCardinalityOfOneSatisfied() {
}

func (t *ControllerTest) InvokesOneTimeActions() {
}

func (t *ControllerTest) InvokesFallbackActions() {
}

func (t *ControllerTest) InvokesImplicitActions() {
}

func (t *ControllerTest) ExpectationsAreMatchedLastToFirst() {
}

func (t *ControllerTest) ExpectationsAreSegregatedByMockObject() {
}

func (t *ControllerTest) ExpectationsAreSegregatedByMethodName() {
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 core

import (
	"errors"
	"fmt"
	"time"

	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logroot"
	"github.com/ligato/cn-infra/utils/safeclose"
	"github.com/namsral/flag"
)

// variables set by the Makefile using ldflags
var (
	BuildVersion string
	BuildDate    string
)

// Agent implements startup & shutdown procedures.
type Agent struct {
	// plugin list
	plugins []*NamedPlugin
	logging.Logger
	// agent startup details
	startup
}

type startup struct {
	// The startup/initialization must take no longer that maxStartup.
	MaxStartupTime time.Duration
	// successfully initialized plugins
	initDuration time.Duration
	// successfully after-initialized plugins
	afterInitDuration time.Duration
	// the field is set before initialization of every plugin with its name
	currentlyProcessing string
}

const (
	logErrorFmt        = "plugin %s: Init error '%s', took %v"
	logSuccessFmt      = "plugin %s: Init took %v"
	logSkippedFmt      = "plugin %s: Init skipped due to previous error"
	logAfterSkippedFmt = "plugin %s: AfterInit skipped due to previous error"
	logAfterErrorFmt   = "plugin %s: AfterInit error '%s', took %v"
	logAfterSuccessFmt = "plugin %s: AfterInit took %v"
	logNoAfterInitFmt  = "plugin %s: not implement AfterInit"
	logTimeoutFmt      = "plugin %s not completed before timeout"
	// The default value serves as an indicator for timer still running even after MaxStartupTime. Used in case
	// some plugin lasts long time to load or is stuck
	defaultTimerValue = -1
)

// NewAgent returns a new instance of the Agent with plugins. Use options if needed:
// <WithLogger() option> will be used to log messages related to the agent life-cycle,
// but not for the plugins themselves.
// <WithTimeout() option> puts a time limit on initialization of all provided plugins.
// Agent.Start() returns ErrPluginsInitTimeout error if one or more plugins fail
// to initialize inside the specified time limit.
// <WithPlugins() option> is a variable list of plugins to load. ListPluginsInFlavor() helper
// method can be used to obtain the list from a given flavor.
//
// Example 1 (existing flavor - or use alias rpc.NewAgent()):
//
//    core.NewAgent(&FlavorRPC{}, core.WithTimeout(5 * time.Second), rpc.WithPlugins(func(flavor *FlavorRPC) []*core.NamedPlugins {
//		return []*core.NamedPlugins{{"customization": &CustomPlugin{DependencyXY: &flavor.GRPC}}}
//    })
//
// Example 2 (custom flavor):
//
//    core.NewAgent(&MyFlavor{}, core.WithTimeout(5 * time.Second), my.WithPlugins(func(flavor *MyFlavor) []*core.NamedPlugins {
//		return []*core.NamedPlugins{{"customization": &CustomPlugin{DependencyXY: &flavor.XY}}}
//    })
func NewAgent(flavor Flavor, opts ...Option) *Agent {
	plugins := flavor.Plugins()

	var agentCoreLogger logging.Logger
	maxStartup := 15 * time.Second

	var flavors []Flavor
	if fs, ok := flavor.(flavors); ok {
		flavors = fs
	} else {
		flavors = []Flavor{flavor}
	}

	flavor.Inject()

	for _, opt := range opts {
		switch opt.(type) {
		case WithPluginsOpt:
			plugins = append(plugins, opt.(WithPluginsOpt).Plugins(flavors...)...)
		case *WithTimeoutOpt:
			ms := opt.(*WithTimeoutOpt).Timeout
			if ms > 0 {
				maxStartup = ms
			}
		case *WithLoggerOpt:
			agentCoreLogger = opt.(*WithLoggerOpt).Logger
		}
	}

	if logRegGet, ok := flavor.(logRegistryGetter); ok && logRegGet != nil {
		logReg := logRegGet.LogRegistry()

		if logReg != nil {
			agentCoreLogger = logReg.NewLogger("agent_core")
		} else {
			agentCoreLogger = logroot.StandardLogger()
		}
	} else {
		agentCoreLogger = logroot.StandardLogger()
	}

	a := Agent{
		plugins,
		agentCoreLogger,
		startup{MaxStartupTime: maxStartup},
	}
	return &a
}

// NewAgentDeprecated older & deprecated version of a constructor
// <logger> will be used to log messages related to the agent life-cycle,
// but not for the plugins themselves.
// <maxStartup> puts a time limit on initialization of all provided plugins.
// Agent.Start() returns ErrPluginsInitTimeout error if one or more plugins fail
// to initialize inside the specified time limit.
// <plugins> is a variable list of plugins to load. ListPluginsInFlavor() helper
// method can be used to obtain the list from a given flavor.
func NewAgentDeprecated(logger logging.Logger, maxStartup time.Duration, plugins ...*NamedPlugin) *Agent {
	a := Agent{
		plugins,
		logger,
		startup{MaxStartupTime: maxStartup},
	}
	return &a
}

type logRegistryGetter interface {
	// LogRegistry is a getter for log registry instance
	LogRegistry() logging.Registry
}

// Start starts/initializes all selected plugins.
// The first iteration tries to run Init() method on every plugin from the list.
// If any of the plugins fails to initialize (Init() return non-nil error),
// initialization is cancelled by calling Close() method for already initialized
// plugins in the reverse order. The encountered error is returned by this
// function as-is.
// The second iteration does the same for the AfterInit() method. The difference
// is that AfterInit() is an optional method (not required by the Plugin
// interface, only suggested by PostInit interface) and therefore not necessarily
// called on every plugin.
// The startup/initialization must take no longer than maxStartup time limit,
// otherwise ErrPluginsInitTimeout error is returned.
func (agent *Agent) Start() error {
	agent.WithFields(logging.Fields{"BuildVersion": BuildVersion, "BuildDate": BuildDate}).Info("Starting the agent...")

	doneChannel := make(chan struct{}, 0)
	errChannel := make(chan error, 0)

	if !flag.Parsed() {
		flag.Parse()
	}

	go func() {
		err := agent.initPlugins()
		if err != nil {
			errChannel <- err
			return
		}
		err = agent.handleAfterInit()
		if err != nil {
			errChannel <- err
			return
		}
		close(doneChannel)
	}()

	//block until all Plugins are initialized or timeout expires
	select {
	case err := <-errChannel:
		agent.WithField("durationInNs", agent.initDuration.Nanoseconds()).Infof("Agent Init took %v", agent.initDuration)
		agent.WithField("durationInNs", agent.afterInitDuration.Nanoseconds()).Infof("Agent AfterInit took %v", agent.afterInitDuration)
		return err
	case <-doneChannel:
		agent.WithField("durationInNs", agent.initDuration.Nanoseconds()).Infof("Agent Init took %v", agent.initDuration)
		agent.WithField("durationInNs", agent.afterInitDuration.Nanoseconds()).Infof("Agent AfterInit took %v", agent.afterInitDuration)
		duration := agent.initDuration + agent.afterInitDuration
		agent.WithField("durationInNs", duration.Nanoseconds()).Info(fmt.Sprintf("All plugins initialized successfully, took %v", duration))
		return nil
	case <-time.After(agent.MaxStartupTime):
		if agent.initDuration == defaultTimerValue {
			agent.Infof("Agent Init took > %v", agent.MaxStartupTime)
			agent.WithField("durationInNs", agent.afterInitDuration.Nanoseconds()).Infof("Agent AfterInit took %v", agent.afterInitDuration)
		} else if agent.afterInitDuration == defaultTimerValue {
			agent.WithField("durationInNs", agent.initDuration.Nanoseconds()).Infof("Agent Init took %v", agent.initDuration)
			agent.Infof("Agent AfterInit took > %v", agent.MaxStartupTime)
		}

		return fmt.Errorf(logTimeoutFmt, agent.currentlyProcessing)
	}
}

// Stop gracefully shuts down the Agent. It is called usually when the user
// interrupts the Agent from the EventLoopWithInterrupt().
//
// This implementation tries to call Close() method on every plugin on the list
// in the reverse order. It continues even if some error occurred.
func (agent *Agent) Stop() error {
	agent.Info("Stopping agent...")
	errMsg := ""
	for i := len(agent.plugins) - 1; i >= 0; i-- {
		agent.WithField("pluginName", agent.plugins[i].PluginName).Debug("Stopping plugin begin")
		err := safeclose.Close(agent.plugins[i].Plugin)
		if err != nil {
			if len(errMsg) > 0 {
				errMsg += "; "
			}
			errMsg += string(agent.plugins[i].PluginName)
			errMsg += ": " + err.Error()
		}
		agent.WithField("pluginName", agent.plugins[i].PluginName).Debug("Stopping plugin end ", err)
	}

	agent.Debug("Agent stopped")

	if len(errMsg) > 0 {
		return errors.New(errMsg)
	}
	return nil
}

// initPlugins calls Init() an all plugins on the list
func (agent *Agent) initPlugins() error {
	// Flag indicates that some of the plugins failed to initialize
	var initPluginCounter int
	var pluginFailed bool
	var wasError error

	agent.initDuration = defaultTimerValue
	initStartTime := time.Now()
	for index, plugin := range agent.plugins {
		initPluginCounter = index

		// set currently initialized plugin name
		agent.currentlyProcessing = string(plugin.PluginName)

		// skip all other plugins if some of them failed
		if pluginFailed {
			agent.Info(fmt.Sprintf(logSkippedFmt, plugin.PluginName))
			continue
		}

		pluginStartTime := time.Now()
		err := plugin.Init()
		if err != nil {
			pluginErrTime := time.Since(pluginStartTime)
			agent.WithField("durationInNs", pluginErrTime.Nanoseconds()).Errorf(logErrorFmt, plugin.PluginName, err, pluginErrTime)

			pluginFailed = true
			wasError = fmt.Errorf(logErrorFmt, plugin.PluginName, err, pluginErrTime)
		} else {
			pluginSuccTime := time.Since(pluginStartTime)
			agent.WithField("durationInNs", pluginSuccTime.Nanoseconds()).Infof(logSuccessFmt, plugin.PluginName, pluginSuccTime)
		}
	}
	agent.initDuration = time.Since(initStartTime)

	if wasError != nil {
		//Stop the plugins that are initialized
		for i := initPluginCounter; i >= 0; i-- {
			agent.Debugf("Closing %v", agent.plugins[i])
			err := safeclose.Close(agent.plugins[i])
			if err != nil {
				wasError = err
			}
		}
		return wasError
	}
	return nil
}

// handleAfterInit calls the AfterInit handlers for plugins that can only
// finish their initialization after  all other plugins have been initialized.
func (agent *Agent) handleAfterInit() error {
	// Flag indicates that some of the plugins failed to after-initialize
	var pluginFailed bool
	var wasError error

	agent.afterInitDuration = defaultTimerValue
	afterInitStartTime := time.Now()
	for _, plug := range agent.plugins {
		// set currently after-initialized plugin name
		agent.currentlyProcessing = string(plug.PluginName)

		// skip all other plugins if some of them failed
		if pluginFailed {
			agent.Info(fmt.Sprintf(logAfterSkippedFmt, plug.PluginName))
			continue
		}

		// Check if plugin implements AfterInit()
		if plugin, ok := plug.Plugin.(PostInit); ok {
			pluginStartTime := time.Now()
			err := plugin.AfterInit()
			if err != nil {
				pluginErrTime := time.Since(pluginStartTime)
				agent.WithField("durationInNs", pluginErrTime.Nanoseconds()).Errorf(logAfterErrorFmt, plug.PluginName, err, pluginErrTime)

				pluginFailed = true
				wasError = fmt.Errorf(logAfterErrorFmt, plug.PluginName, err, pluginErrTime)
			} else {
				pluginSuccTime := time.Since(pluginStartTime)
				agent.WithField("durationInNs", pluginSuccTime.Nanoseconds()).Infof(logAfterSuccessFmt, plug.PluginName, pluginSuccTime)
			}
		} else {
			agent.Info(fmt.Sprintf(logNoAfterInitFmt, plug.PluginName))
		}
	}
	agent.afterInitDuration = time.Since(afterInitStartTime)

	if wasError != nil {
		agent.Stop()
		return wasError
	}
	return nil
}
<commit_msg>ODPM-641 core.Inject() for multiple flavors - fix<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 core

import (
	"errors"
	"fmt"
	"time"

	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logroot"
	"github.com/ligato/cn-infra/utils/safeclose"
	"github.com/namsral/flag"
)

// variables set by the Makefile using ldflags
var (
	BuildVersion string
	BuildDate    string
)

// Agent implements startup & shutdown procedures.
type Agent struct {
	// plugin list
	plugins []*NamedPlugin
	logging.Logger
	// agent startup details
	startup
}

type startup struct {
	// The startup/initialization must take no longer that maxStartup.
	MaxStartupTime time.Duration
	// successfully initialized plugins
	initDuration time.Duration
	// successfully after-initialized plugins
	afterInitDuration time.Duration
	// the field is set before initialization of every plugin with its name
	currentlyProcessing string
}

const (
	logErrorFmt        = "plugin %s: Init error '%s', took %v"
	logSuccessFmt      = "plugin %s: Init took %v"
	logSkippedFmt      = "plugin %s: Init skipped due to previous error"
	logAfterSkippedFmt = "plugin %s: AfterInit skipped due to previous error"
	logAfterErrorFmt   = "plugin %s: AfterInit error '%s', took %v"
	logAfterSuccessFmt = "plugin %s: AfterInit took %v"
	logNoAfterInitFmt  = "plugin %s: not implement AfterInit"
	logTimeoutFmt      = "plugin %s not completed before timeout"
	// The default value serves as an indicator for timer still running even after MaxStartupTime. Used in case
	// some plugin lasts long time to load or is stuck
	defaultTimerValue = -1
)

// NewAgent returns a new instance of the Agent with plugins. Use options if needed:
// <WithLogger() option> will be used to log messages related to the agent life-cycle,
// but not for the plugins themselves.
// <WithTimeout() option> puts a time limit on initialization of all provided plugins.
// Agent.Start() returns ErrPluginsInitTimeout error if one or more plugins fail
// to initialize inside the specified time limit.
// <WithPlugins() option> is a variable list of plugins to load. ListPluginsInFlavor() helper
// method can be used to obtain the list from a given flavor.
//
// Example 1 (existing flavor - or use alias rpc.NewAgent()):
//
//    core.NewAgent(&FlavorRPC{}, core.WithTimeout(5 * time.Second), rpc.WithPlugins(func(flavor *FlavorRPC) []*core.NamedPlugins {
//		return []*core.NamedPlugins{{"customization": &CustomPlugin{DependencyXY: &flavor.GRPC}}}
//    })
//
// Example 2 (custom flavor):
//
//    core.NewAgent(&MyFlavor{}, core.WithTimeout(5 * time.Second), my.WithPlugins(func(flavor *MyFlavor) []*core.NamedPlugins {
//		return []*core.NamedPlugins{{"customization": &CustomPlugin{DependencyXY: &flavor.XY}}}
//    })
func NewAgent(flavor Flavor, opts ...Option) *Agent {
	plugins := flavor.Plugins()

	var agentCoreLogger logging.Logger
	maxStartup := 15 * time.Second

	var flavors []Flavor
	if fs, ok := flavor.(flavors); ok {
		flavors = fs.fs
	} else {
		flavors = []Flavor{flavor}
	}

	flavor.Inject()

	for _, opt := range opts {
		switch opt.(type) {
		case WithPluginsOpt:
			plugins = append(plugins, opt.(WithPluginsOpt).Plugins(flavors...)...)
		case *WithTimeoutOpt:
			ms := opt.(*WithTimeoutOpt).Timeout
			if ms > 0 {
				maxStartup = ms
			}
		case *WithLoggerOpt:
			agentCoreLogger = opt.(*WithLoggerOpt).Logger
		}
	}

	if logRegGet, ok := flavor.(logRegistryGetter); ok && logRegGet != nil {
		logReg := logRegGet.LogRegistry()

		if logReg != nil {
			agentCoreLogger = logReg.NewLogger("agent_core")
		} else {
			agentCoreLogger = logroot.StandardLogger()
		}
	} else {
		agentCoreLogger = logroot.StandardLogger()
	}

	a := Agent{
		plugins,
		agentCoreLogger,
		startup{MaxStartupTime: maxStartup},
	}
	return &a
}

// NewAgentDeprecated older & deprecated version of a constructor
// <logger> will be used to log messages related to the agent life-cycle,
// but not for the plugins themselves.
// <maxStartup> puts a time limit on initialization of all provided plugins.
// Agent.Start() returns ErrPluginsInitTimeout error if one or more plugins fail
// to initialize inside the specified time limit.
// <plugins> is a variable list of plugins to load. ListPluginsInFlavor() helper
// method can be used to obtain the list from a given flavor.
func NewAgentDeprecated(logger logging.Logger, maxStartup time.Duration, plugins ...*NamedPlugin) *Agent {
	a := Agent{
		plugins,
		logger,
		startup{MaxStartupTime: maxStartup},
	}
	return &a
}

type logRegistryGetter interface {
	// LogRegistry is a getter for log registry instance
	LogRegistry() logging.Registry
}

// Start starts/initializes all selected plugins.
// The first iteration tries to run Init() method on every plugin from the list.
// If any of the plugins fails to initialize (Init() return non-nil error),
// initialization is cancelled by calling Close() method for already initialized
// plugins in the reverse order. The encountered error is returned by this
// function as-is.
// The second iteration does the same for the AfterInit() method. The difference
// is that AfterInit() is an optional method (not required by the Plugin
// interface, only suggested by PostInit interface) and therefore not necessarily
// called on every plugin.
// The startup/initialization must take no longer than maxStartup time limit,
// otherwise ErrPluginsInitTimeout error is returned.
func (agent *Agent) Start() error {
	agent.WithFields(logging.Fields{"BuildVersion": BuildVersion, "BuildDate": BuildDate}).Info("Starting the agent...")

	doneChannel := make(chan struct{}, 0)
	errChannel := make(chan error, 0)

	if !flag.Parsed() {
		flag.Parse()
	}

	go func() {
		err := agent.initPlugins()
		if err != nil {
			errChannel <- err
			return
		}
		err = agent.handleAfterInit()
		if err != nil {
			errChannel <- err
			return
		}
		close(doneChannel)
	}()

	//block until all Plugins are initialized or timeout expires
	select {
	case err := <-errChannel:
		agent.WithField("durationInNs", agent.initDuration.Nanoseconds()).Infof("Agent Init took %v", agent.initDuration)
		agent.WithField("durationInNs", agent.afterInitDuration.Nanoseconds()).Infof("Agent AfterInit took %v", agent.afterInitDuration)
		return err
	case <-doneChannel:
		agent.WithField("durationInNs", agent.initDuration.Nanoseconds()).Infof("Agent Init took %v", agent.initDuration)
		agent.WithField("durationInNs", agent.afterInitDuration.Nanoseconds()).Infof("Agent AfterInit took %v", agent.afterInitDuration)
		duration := agent.initDuration + agent.afterInitDuration
		agent.WithField("durationInNs", duration.Nanoseconds()).Info(fmt.Sprintf("All plugins initialized successfully, took %v", duration))
		return nil
	case <-time.After(agent.MaxStartupTime):
		if agent.initDuration == defaultTimerValue {
			agent.Infof("Agent Init took > %v", agent.MaxStartupTime)
			agent.WithField("durationInNs", agent.afterInitDuration.Nanoseconds()).Infof("Agent AfterInit took %v", agent.afterInitDuration)
		} else if agent.afterInitDuration == defaultTimerValue {
			agent.WithField("durationInNs", agent.initDuration.Nanoseconds()).Infof("Agent Init took %v", agent.initDuration)
			agent.Infof("Agent AfterInit took > %v", agent.MaxStartupTime)
		}

		return fmt.Errorf(logTimeoutFmt, agent.currentlyProcessing)
	}
}

// Stop gracefully shuts down the Agent. It is called usually when the user
// interrupts the Agent from the EventLoopWithInterrupt().
//
// This implementation tries to call Close() method on every plugin on the list
// in the reverse order. It continues even if some error occurred.
func (agent *Agent) Stop() error {
	agent.Info("Stopping agent...")
	errMsg := ""
	for i := len(agent.plugins) - 1; i >= 0; i-- {
		agent.WithField("pluginName", agent.plugins[i].PluginName).Debug("Stopping plugin begin")
		err := safeclose.Close(agent.plugins[i].Plugin)
		if err != nil {
			if len(errMsg) > 0 {
				errMsg += "; "
			}
			errMsg += string(agent.plugins[i].PluginName)
			errMsg += ": " + err.Error()
		}
		agent.WithField("pluginName", agent.plugins[i].PluginName).Debug("Stopping plugin end ", err)
	}

	agent.Debug("Agent stopped")

	if len(errMsg) > 0 {
		return errors.New(errMsg)
	}
	return nil
}

// initPlugins calls Init() an all plugins on the list
func (agent *Agent) initPlugins() error {
	// Flag indicates that some of the plugins failed to initialize
	var initPluginCounter int
	var pluginFailed bool
	var wasError error

	agent.initDuration = defaultTimerValue
	initStartTime := time.Now()
	for index, plugin := range agent.plugins {
		initPluginCounter = index

		// set currently initialized plugin name
		agent.currentlyProcessing = string(plugin.PluginName)

		// skip all other plugins if some of them failed
		if pluginFailed {
			agent.Info(fmt.Sprintf(logSkippedFmt, plugin.PluginName))
			continue
		}

		pluginStartTime := time.Now()
		err := plugin.Init()
		if err != nil {
			pluginErrTime := time.Since(pluginStartTime)
			agent.WithField("durationInNs", pluginErrTime.Nanoseconds()).Errorf(logErrorFmt, plugin.PluginName, err, pluginErrTime)

			pluginFailed = true
			wasError = fmt.Errorf(logErrorFmt, plugin.PluginName, err, pluginErrTime)
		} else {
			pluginSuccTime := time.Since(pluginStartTime)
			agent.WithField("durationInNs", pluginSuccTime.Nanoseconds()).Infof(logSuccessFmt, plugin.PluginName, pluginSuccTime)
		}
	}
	agent.initDuration = time.Since(initStartTime)

	if wasError != nil {
		//Stop the plugins that are initialized
		for i := initPluginCounter; i >= 0; i-- {
			agent.Debugf("Closing %v", agent.plugins[i])
			err := safeclose.Close(agent.plugins[i])
			if err != nil {
				wasError = err
			}
		}
		return wasError
	}
	return nil
}

// handleAfterInit calls the AfterInit handlers for plugins that can only
// finish their initialization after  all other plugins have been initialized.
func (agent *Agent) handleAfterInit() error {
	// Flag indicates that some of the plugins failed to after-initialize
	var pluginFailed bool
	var wasError error

	agent.afterInitDuration = defaultTimerValue
	afterInitStartTime := time.Now()
	for _, plug := range agent.plugins {
		// set currently after-initialized plugin name
		agent.currentlyProcessing = string(plug.PluginName)

		// skip all other plugins if some of them failed
		if pluginFailed {
			agent.Info(fmt.Sprintf(logAfterSkippedFmt, plug.PluginName))
			continue
		}

		// Check if plugin implements AfterInit()
		if plugin, ok := plug.Plugin.(PostInit); ok {
			pluginStartTime := time.Now()
			err := plugin.AfterInit()
			if err != nil {
				pluginErrTime := time.Since(pluginStartTime)
				agent.WithField("durationInNs", pluginErrTime.Nanoseconds()).Errorf(logAfterErrorFmt, plug.PluginName, err, pluginErrTime)

				pluginFailed = true
				wasError = fmt.Errorf(logAfterErrorFmt, plug.PluginName, err, pluginErrTime)
			} else {
				pluginSuccTime := time.Since(pluginStartTime)
				agent.WithField("durationInNs", pluginSuccTime.Nanoseconds()).Infof(logAfterSuccessFmt, plug.PluginName, pluginSuccTime)
			}
		} else {
			agent.Info(fmt.Sprintf(logNoAfterInitFmt, plug.PluginName))
		}
	}
	agent.afterInitDuration = time.Since(afterInitStartTime)

	if wasError != nil {
		agent.Stop()
		return wasError
	}
	return nil
}
<|endoftext|>


<commit_before>package mitake

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

var (
	// mux is the HTTP request multiplexer used with the test server.
	mux *http.ServeMux

	// client is the eztalk client being tested.
	client *Client

	// server is a test HTTP server used to provide mock API responses.
	server *httptest.Server
)

func setup() {
	// test server
	mux = http.NewServeMux()
	server = httptest.NewServer(mux)

	// mitake client configured to use test server
	baseURL, _ := url.Parse(server.URL)
	client = &Client{
		client:  http.DefaultClient,
		BaseURL: baseURL,
	}
}

// teardown closes the test HTTP server.
func teardown() {
	server.Close()
}

func testMethod(t *testing.T, r *http.Request, want string) {
	if got := r.Method; got != want {
		t.Errorf("Request method is %v, want %v", got, want)
	}
}

func testINI(t *testing.T, r *http.Request, want string) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		t.Errorf("Request parameters error: %v", err)
	}
	defer r.Body.Close()

	if got := string(body); got != want {
		t.Errorf("Request parameters is %v, want %v", got, want)
	}
}

func TestNewClient(t *testing.T) {
	c := NewClient("username", "password", http.DefaultClient)

	if got, want := c.BaseURL.String(), defaultBaseURL; got != want {
		t.Errorf("NewClient BaseURL is %v, want %v", got, want)
	}
}

func TestClient_NewRequest(t *testing.T) {
	c := NewClient("username", "password", http.DefaultClient)

	inURL, outURL := "/foo", defaultBaseURL+"foo"
	inBody, outBody := "Hello, 世界", "Hello, 世界"
	req, _ := c.NewRequest("GET", inURL, strings.NewReader(inBody))

	if got, want := req.URL.String(), outURL; got != want {
		t.Errorf("NewRequest(%q) URL is %v, want %v", inURL, got, want)
	}

	body, _ := ioutil.ReadAll(req.Body)
	if got, want := string(body), outBody; got != want {
		t.Errorf("NewRequest(%q) Body is %v, want %v", inBody, got, want)
	}

	if got, want := req.Header.Get("User-Agent"), c.UserAgent; got != want {
		t.Errorf("NewRequest() User-Agent is %v, want %v", got, want)
	}
}

func TestClient_Do(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if m := "GET"; m != r.Method {
			t.Errorf("Request method is %v, want %v", r.Method, m)
		}
		fmt.Fprint(w, "Hello, 世界")
	})

	req, _ := client.NewRequest("GET", "/", nil)
	resp, err := client.Do(req)
	if err != nil {

	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)

	want := "Hello, 世界"
	if !reflect.DeepEqual(string(body), want) {
		t.Errorf("Response body is %s, want %s", body, want)
	}
}

func TestClient_Do_httpError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "Bad Request", 400)
	})

	req, _ := client.NewRequest("GET", "/", nil)
	_, err := client.Do(req)

	if err == nil {
		t.Error("Expected HTTP 400 error.")
	}
}

func TestClient_Do_noContent(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "")
	})

	req, _ := client.NewRequest("GET", "/", nil)
	_, err := client.Do(req)

	if err == nil {
		t.Error("Expected empty body error.")
	}
}
<commit_msg>Add missing error in TestClient_Do<commit_after>package mitake

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

var (
	// mux is the HTTP request multiplexer used with the test server.
	mux *http.ServeMux

	// client is the eztalk client being tested.
	client *Client

	// server is a test HTTP server used to provide mock API responses.
	server *httptest.Server
)

func setup() {
	// test server
	mux = http.NewServeMux()
	server = httptest.NewServer(mux)

	// mitake client configured to use test server
	baseURL, _ := url.Parse(server.URL)
	client = &Client{
		client:  http.DefaultClient,
		BaseURL: baseURL,
	}
}

// teardown closes the test HTTP server.
func teardown() {
	server.Close()
}

func testMethod(t *testing.T, r *http.Request, want string) {
	if got := r.Method; got != want {
		t.Errorf("Request method is %v, want %v", got, want)
	}
}

func testINI(t *testing.T, r *http.Request, want string) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		t.Errorf("Request parameters error: %v", err)
	}
	defer r.Body.Close()

	if got := string(body); got != want {
		t.Errorf("Request parameters is %v, want %v", got, want)
	}
}

func TestNewClient(t *testing.T) {
	c := NewClient("username", "password", http.DefaultClient)

	if got, want := c.BaseURL.String(), defaultBaseURL; got != want {
		t.Errorf("NewClient BaseURL is %v, want %v", got, want)
	}
}

func TestClient_NewRequest(t *testing.T) {
	c := NewClient("username", "password", http.DefaultClient)

	inURL, outURL := "/foo", defaultBaseURL+"foo"
	inBody, outBody := "Hello, 世界", "Hello, 世界"
	req, _ := c.NewRequest("GET", inURL, strings.NewReader(inBody))

	if got, want := req.URL.String(), outURL; got != want {
		t.Errorf("NewRequest(%q) URL is %v, want %v", inURL, got, want)
	}

	body, _ := ioutil.ReadAll(req.Body)
	if got, want := string(body), outBody; got != want {
		t.Errorf("NewRequest(%q) Body is %v, want %v", inBody, got, want)
	}

	if got, want := req.Header.Get("User-Agent"), c.UserAgent; got != want {
		t.Errorf("NewRequest() User-Agent is %v, want %v", got, want)
	}
}

func TestClient_Do(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if m := "GET"; m != r.Method {
			t.Errorf("Request method is %v, want %v", r.Method, m)
		}
		fmt.Fprint(w, "Hello, 世界")
	})

	req, _ := client.NewRequest("GET", "/", nil)
	resp, err := client.Do(req)
	if err != nil {
		t.Errorf("Do returned unexpected error: %v", err)
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)

	want := "Hello, 世界"
	if !reflect.DeepEqual(string(body), want) {
		t.Errorf("Response body is %s, want %s", body, want)
	}
}

func TestClient_Do_httpError(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "Bad Request", 400)
	})

	req, _ := client.NewRequest("GET", "/", nil)
	_, err := client.Do(req)

	if err == nil {
		t.Error("Expected HTTP 400 error.")
	}
}

func TestClient_Do_noContent(t *testing.T) {
	setup()
	defer teardown()

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "")
	})

	req, _ := client.NewRequest("GET", "/", nil)
	_, err := client.Do(req)

	if err == nil {
		t.Error("Expected empty body error.")
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 PingCAP, 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 executor_test

import (
	"time"

	. "github.com/pingcap/check"
	"github.com/pingcap/tidb/util/testkit"
)

func (s *testSuiteP2) TestQueryTime(c *C) {
	tk := testkit.NewTestKit(c, s.store)
	tk.MustExec("use test")

	costTime := time.Since(tk.Se.GetSessionVars().StartTime)
	c.Assert(costTime < 1*time.Second, IsTrue)

	tk.MustExec("drop table if exists t")
	tk.MustExec("create table t(a int)")
	tk.MustExec("insert into t values(1), (1), (1), (1), (1)")
	tk.MustExec("select * from t t1 join t t2 on t1.a = t2.a")

	costTime = time.Since(tk.Se.GetSessionVars().StartTime)
	c.Assert(costTime < 1*time.Second, IsTrue)
}
<commit_msg>executor: migrate test infra to testify for adpater_test.go (#28539)<commit_after>// Copyright 2019 PingCAP, 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 executor_test

import (
	"testing"
	"time"

	"github.com/pingcap/tidb/testkit"
	"github.com/stretchr/testify/require"
)

func TestQueryTime(t *testing.T) {
	t.Parallel()
	store, clean := testkit.CreateMockStore(t)
	defer clean()

	tk := testkit.NewTestKit(t, store)
	tk.MustExec("use test")

	costTime := time.Since(tk.Session().GetSessionVars().StartTime)
	require.Less(t, costTime, time.Second)

	tk.MustExec("drop table if exists t")
	tk.MustExec("create table t(a int)")
	tk.MustExec("insert into t values(1), (1), (1), (1), (1)")
	tk.MustExec("select * from t t1 join t t2 on t1.a = t2.a")

	costTime = time.Since(tk.Session().GetSessionVars().StartTime)
	require.Less(t, costTime, time.Second)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/HewlettPackard/oneview-golang/ov"
	"github.com/HewlettPackard/oneview-golang/utils"
	"os"
	"strconv"
)

func main() {
	var (
		ClientOV *ov.OVClient
	)
	apiversion, _ := strconv.Atoi(os.Getenv("ONEVIEW_APIVERSION"))
	ovc := ClientOV.NewOVClient(
		os.Getenv("ONEVIEW_OV_USER"),
		os.Getenv("ONEVIEW_OV_PASSWORD"),
		os.Getenv("ONEVIEW_OV_DOMAIN"),
		os.Getenv("ONEVIEW_OV_ENDPOINT"),
		false,
		apiversion,
		"*")

	// Gets all labels
	responseAllLabels, err := ovc.GetAllLabels("", "", "", "", "")
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Get all labels----------------#")
		jsonResponse, _ := json.MarshalIndent(responseAllLabels, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	labels := []ov.Label{}
	labels = append(labels, ov.Label{
		Name: "NewestLabel",
	})

	// Creates new labels
	label := ov.AssignedLabel{
		ResourceUri: utils.Nstring("/rest/server-profile-templates/b6777c57-34f1-4491-93c4-8bec773f286c"),
		Labels:	     labels,
	}
	responseLabel, err := ovc.CreateLabel(label)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Labels created----------------#")
		jsonResponse, _ := json.MarshalIndent(responseLabel, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Gets Assigned Labels
	responseLabel, err = ovc.GetAssignedLabels(responseLabel.ResourceUri)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Assigned labels----------------#")
		jsonResponse, _ := json.MarshalIndent(responseLabel, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Updates Assigned Labels
	responseLabel.Labels[0].Name = "UpdatedLabel"
	responseLabel, err = ovc.UpdateAssignedLabels(responseLabel)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Updated assigned labels----------------#")
		jsonResponse, _ := json.MarshalIndent(responseLabel, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Gets a Label By Uri
	response, err := ovc.GetLabelByURI(responseLabel.Labels[0].Uri)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Get label by uri----------------#")
		jsonResponse, _ := json.MarshalIndent(response, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Deletes all labels from the resource
	err = ovc.DeleteAssignedLabel(string(responseLabel.ResourceUri))
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------All labels to assigned recource are deleted----------------#")
	}
}
<commit_msg>fixed fmt issues<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/HewlettPackard/oneview-golang/ov"
	"github.com/HewlettPackard/oneview-golang/utils"
	"os"
	"strconv"
)

func main() {
	var (
		ClientOV *ov.OVClient
	)
	apiversion, _ := strconv.Atoi(os.Getenv("ONEVIEW_APIVERSION"))
	ovc := ClientOV.NewOVClient(
		os.Getenv("ONEVIEW_OV_USER"),
		os.Getenv("ONEVIEW_OV_PASSWORD"),
		os.Getenv("ONEVIEW_OV_DOMAIN"),
		os.Getenv("ONEVIEW_OV_ENDPOINT"),
		false,
		apiversion,
		"*")

	// Gets all labels
	responseAllLabels, err := ovc.GetAllLabels("", "", "", "", "")
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Get all labels----------------#")
		jsonResponse, _ := json.MarshalIndent(responseAllLabels, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	labels := []ov.Label{}
	labels = append(labels, ov.Label{
		Name: "NewestLabel",
	})

	// Creates new labels
	label := ov.AssignedLabel{
		ResourceUri: utils.Nstring("/rest/server-profile-templates/b6777c57-34f1-4491-93c4-8bec773f286c"),
		Labels:      labels,
	}
	responseLabel, err := ovc.CreateLabel(label)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Labels created----------------#")
		jsonResponse, _ := json.MarshalIndent(responseLabel, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Gets Assigned Labels
	responseLabel, err = ovc.GetAssignedLabels(responseLabel.ResourceUri)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Assigned labels----------------#")
		jsonResponse, _ := json.MarshalIndent(responseLabel, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Updates Assigned Labels
	responseLabel.Labels[0].Name = "UpdatedLabel"
	responseLabel, err = ovc.UpdateAssignedLabels(responseLabel)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Updated assigned labels----------------#")
		jsonResponse, _ := json.MarshalIndent(responseLabel, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Gets a Label By Uri
	response, err := ovc.GetLabelByURI(responseLabel.Labels[0].Uri)
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------Get label by uri----------------#")
		jsonResponse, _ := json.MarshalIndent(response, "", "  ")
		fmt.Print(string(jsonResponse), "\n\n")
	}

	// Deletes all labels from the resource
	err = ovc.DeleteAssignedLabel(string(responseLabel.ResourceUri))
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("#-------------All labels to assigned recource are deleted----------------#")
	}
}
<|endoftext|>


<commit_before>// +build linux,!noaio

// The above build tag specifies this file is only to be built on linux (because
// goaio only supports linux). If you are having difficulty building on linux
// or want to build without librados present, then use
//    go build -tags 'noaio'

package nbd

import (
	"github.com/traetox/goaio"
	"golang.org/x/net/context"
	"os"
)

// AioFileBackend implements Backend
type AioFileBackend struct {
	aio  *goaio.AIO
	size uint64
}

// WriteAt implements Backend.WriteAt
func (afb *AioFileBackend) WriteAt(ctx context.Context, b []byte, offset int64, fua bool) (int, error) {
	if err := afb.aio.Wait(); err != nil {
		return 0, err
	}
	requestId, err := afb.aio.WriteAt(b, offset)
	if err != nil {
		return 0, err
	}
	if err := afb.aio.WaitFor(requestId); err != nil {
		return 0, err
	}
	if fua {
		if err := afb.aio.Flush(); err != nil {
			return 0, err
		}
	}
	return len(b), err
}

// ReadAt implements Backend.ReadAt
func (afb *AioFileBackend) ReadAt(ctx context.Context, b []byte, offset int64) (int, error) {
	if err := afb.aio.Wait(); err != nil {
		return 0, err
	}
	requestId, err := afb.aio.ReadAt(b, offset)
	if err != nil {
		return 0, err
	}
	if err := afb.aio.WaitFor(requestId); err != nil {
		return 0, err
	}
	return len(b), err
}

// TrimAt implements Backend.TrimAt
func (afb *AioFileBackend) TrimAt(ctx context.Context, length int, offset int64) (int, error) {
	return length, nil
}

// Flush implements Backend.Flush
func (afb *AioFileBackend) Flush(ctx context.Context) error {
	return afb.aio.Flush()
}

// Close implements Backend.Close
func (afb *AioFileBackend) Close(ctx context.Context) error {
	return afb.aio.Close()
}

// Size implements Backend.Size
func (afb *AioFileBackend) Geometry(ctx context.Context) (uint64, uint64, uint64, uint64, error) {
	return afb.size, 1, 32 * 1024, 128 * 1024 * 1024, nil
}

// Size implements Backend.HasFua
func (afb *AioFileBackend) HasFua(ctx context.Context) bool {
	return false
}

// Size implements Backend.HasFua
func (afb *AioFileBackend) HasFlush(ctx context.Context) bool {
	return false
}

// Generate a new aio backend
func NewAioFileBackend(ctx context.Context, ec *ExportConfig) (Backend, error) {
	perms := os.O_RDWR
	if ec.ReadOnly {
		perms = os.O_RDONLY
	}
	if s, err := isTrue(ec.DriverParameters["sync"]); err != nil {
		return nil, err
	} else if s {
		perms |= os.O_SYNC
	}
	aio, err := goaio.NewAIO(ec.DriverParameters["path"], perms, 0666)
	if err != nil {
		return nil, err
	}
	stat, err := aio.FD().Stat()
	if err != nil {
		aio.Close()
		return nil, err
	}
	return &AioFileBackend{
		aio:  aio,
		size: uint64(stat.Size()),
	}, nil
}

// Register our backend
func init() {
	RegisterBackend("aiofile", NewAioFileBackend)
}
<commit_msg>aiofile plugin to use updated goaio API<commit_after>// +build linux,!noaio

// The above build tag specifies this file is only to be built on linux (because
// goaio only supports linux). If you are having difficulty building on linux
// or want to build without librados present, then use
//    go build -tags 'noaio'

package nbd

import (
	"github.com/traetox/goaio"
	"golang.org/x/net/context"
	"os"
)

// AioFileBackend implements Backend
type AioFileBackend struct {
	aio  *goaio.AIO
	size uint64
}

// WriteAt implements Backend.WriteAt
func (afb *AioFileBackend) WriteAt(ctx context.Context, b []byte, offset int64, fua bool) (int, error) {
	if err := afb.aio.Wait(); err != nil {
		return 0, err
	}
	requestId, err := afb.aio.WriteAt(b, offset)
	if err != nil {
		return 0, err
	}
	n, err := afb.aio.WaitFor(requestId)
	if err != nil {
		return 0, err
	}
	if fua {
		if err := afb.aio.Flush(); err != nil {
			return 0, err
		}
	}
	return n, err
}

// ReadAt implements Backend.ReadAt
func (afb *AioFileBackend) ReadAt(ctx context.Context, b []byte, offset int64) (int, error) {
	if err := afb.aio.Wait(); err != nil {
		return 0, err
	}
	requestId, err := afb.aio.ReadAt(b, offset)
	if err != nil {
		return 0, err
	}
	n, err := afb.aio.WaitFor(requestId)
	if err != nil {
		return 0, err
	}
	return n, err
}

// TrimAt implements Backend.TrimAt
func (afb *AioFileBackend) TrimAt(ctx context.Context, length int, offset int64) (int, error) {
	return length, nil
}

// Flush implements Backend.Flush
func (afb *AioFileBackend) Flush(ctx context.Context) error {
	return afb.aio.Flush()
}

// Close implements Backend.Close
func (afb *AioFileBackend) Close(ctx context.Context) error {
	return afb.aio.Close()
}

// Size implements Backend.Size
func (afb *AioFileBackend) Geometry(ctx context.Context) (uint64, uint64, uint64, uint64, error) {
	return afb.size, 1, 32 * 1024, 128 * 1024 * 1024, nil
}

// Size implements Backend.HasFua
func (afb *AioFileBackend) HasFua(ctx context.Context) bool {
	return false
}

// Size implements Backend.HasFua
func (afb *AioFileBackend) HasFlush(ctx context.Context) bool {
	return false
}

// Generate a new aio backend
func NewAioFileBackend(ctx context.Context, ec *ExportConfig) (Backend, error) {
	perms := os.O_RDWR
	if ec.ReadOnly {
		perms = os.O_RDONLY
	}
	if s, err := isTrue(ec.DriverParameters["sync"]); err != nil {
		return nil, err
	} else if s {
		perms |= os.O_SYNC
	}
	aio, err := goaio.NewAIO(ec.DriverParameters["path"], perms, 0666)
	if err != nil {
		return nil, err
	}
	stat, err := aio.FD().Stat()
	if err != nil {
		aio.Close()
		return nil, err
	}
	return &AioFileBackend{
		aio:  aio,
		size: uint64(stat.Size()),
	}, nil
}

// Register our backend
func init() {
	RegisterBackend("aiofile", NewAioFileBackend)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Jan Broer. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package hosts

import (
	"bufio"
	"io"
	"net"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/shawn1m/overture/core/common"
)

type hostsLine struct {
	domain string
	ip     net.IP
	ipv6   bool
}

type hostsLines struct {
	data []*hostsLine
	hash map[string]struct{}
}

func newHostsLineList(r io.Reader) *hostsLines {
	resultLines := new(hostsLines)
	resultLines.hash = make(map[string]struct{})

	defer log.Debugf("%s took %s", "Load hosts", time.Since(time.Now()))

	reader := bufio.NewReader(r)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				log.Errorf("Error reading hosts file: %s", err)
			} else {
				log.Debug("Reading hosts file reached EOF")
			}
			break
		}

		if host := parseLine(line); host != nil {
			if err := resultLines.add(host); err != nil {
				log.Warnf("Bad formatted hosts file line: %s", err)
			}
		}
	}

	return resultLines
}

func (hl *hostsLines) FindHosts(name string) (ipv4List []net.IP, ipv6List []net.IP) {
	for _, h := range hl.data {
		if common.IsDomainMatchRule(h.domain, name) {
			log.WithFields(log.Fields{
				"question": name,
				"domain":   h.domain,
				"ip":       h.ip,
			}).Debug("Matched")
			if h.ip.To4() != nil {
				ipv4List = append(ipv4List, h.ip)
			} else {
				ipv6List = append(ipv6List, h.ip)
			}
		}
	}
	return
}

func (hl *hostsLines) add(h *hostsLine) error {
	if _, ok := hl.hash[h.domain]; !ok {
		hl.data = append(hl.data, h)
		hl.hash[h.domain] = struct{}{}
	} else {
		log.Warnf("Duplicate entry for host %s in hosts file, ignored value: %s", h.domain, h.ip.String())
	}
	return nil
}

func parseLine(line string) *hostsLine {
	if len(line) == 0 {
		return nil
	}

	// Parse leading # for disabled lines
	if line[0:1] == "#" {
		return nil
	}

	// Parse other #s for actual comments
	line = strings.Split(line, "#")[0]

	// Replace tabs and multispaces with single spaces throughout
	line = strings.Replace(line, "\t", " ", -1)
	for strings.Contains(line, "  ") {
		line = strings.Replace(line, "  ", " ", -1)
	}

	line = strings.TrimSpace(line)

	// Break line into words
	words := strings.Split(line, " ")

	if len(words) < 2 {
		return nil
	}
	for i, word := range words {
		words[i] = strings.TrimSpace(word)
	}
	// Separate the first bit (the ip) from the other bits (the domains)
	a, h := words[0], words[1]

	ip := net.ParseIP(a)

	var isIPv6 bool

	switch {
	case ip.To4() != nil:
		isIPv6 = false
	case ip.To16() != nil:
		isIPv6 = true
	default:
		log.Warnf("Invalid IP address found in hosts file: %s", a)
		return nil
	}

	return &hostsLine{h, ip, isIPv6}
}
<commit_msg>Host file support one domain have ipv4 and ipv6 same time (#185) (#189)<commit_after>// Copyright (c) 2015 Jan Broer. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package hosts

import (
	"bufio"
	"io"
	"net"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/shawn1m/overture/core/common"
)

type hostsLine struct {
	domain string
	ip     net.IP
	ipv6   bool
}

type hostsLines struct {
	data []*hostsLine
	hash map[string]struct{}
}

func newHostsLineList(r io.Reader) *hostsLines {
	resultLines := new(hostsLines)
	resultLines.hash = make(map[string]struct{})

	defer log.Debugf("%s took %s", "Load hosts", time.Since(time.Now()))

	reader := bufio.NewReader(r)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				log.Errorf("Error reading hosts file: %s", err)
			} else {
				log.Debug("Reading hosts file reached EOF")
			}
			break
		}

		if host := parseLine(line); host != nil {
			if err := resultLines.add(host); err != nil {
				log.Warnf("Bad formatted hosts file line: %s", err)
			}
		}
	}

	return resultLines
}

func (hl *hostsLines)  FindHosts(name string) (ipv4List []net.IP, ipv6List []net.IP) {
	for _, h := range hl.data {
		if common.IsDomainMatchRule(h.domain, name) {
			log.WithFields(log.Fields{
				"question": name,
				"domain":   h.domain,
				"ip":       h.ip,
			}).Debug("Matched")
			if h.ip.To4() != nil {
				ipv4List = append(ipv4List, h.ip)
			} else {
				ipv6List = append(ipv6List, h.ip)
			}
		}
	}
	return
}

func (hl *hostsLines) add(h *hostsLine) error {
	if _, ok := hl.hash[h.domain]; !ok {
		hl.data = append(hl.data, h)
		hl.hash[h.domain] = struct{}{}
	} else {
		if h.ipv6 == hl.data[0].ipv6 {
			log.Warnf("Duplicate entry for host %s in hosts file, ignored value: %s", h.domain, h.ip.String())
			return nil
		}
		hl.data = append(hl.data, h)
	}
	return nil
}

func parseLine(line string) *hostsLine {
	if len(line) == 0 {
		return nil
	}

	// Parse leading # for disabled lines
	if line[0:1] == "#" {
		return nil
	}

	// Parse other #s for actual comments
	line = strings.Split(line, "#")[0]

	// Replace tabs and multispaces with single spaces throughout
	line = strings.Replace(line, "\t", " ", -1)
	for strings.Contains(line, "  ") {
		line = strings.Replace(line, "  ", " ", -1)
	}

	line = strings.TrimSpace(line)

	// Break line into words
	words := strings.Split(line, " ")

	if len(words) < 2 {
		return nil
	}
	for i, word := range words {
		words[i] = strings.TrimSpace(word)
	}
	// Separate the first bit (the ip) from the other bits (the domains)
	a, h := words[0], words[1]

	ip := net.ParseIP(a)

	var isIPv6 bool

	switch {
	case ip.To4() != nil:
		isIPv6 = false
	case ip.To16() != nil:
		isIPv6 = true
	default:
		log.Warnf("Invalid IP address found in hosts file: %s", a)
		return nil
	}

	return &hostsLine{h, ip, isIPv6}
}
<|endoftext|>