<commit_before>package retrie

import (
	"fmt"
	"regexp"
	"regexp/syntax"

	"github.com/philharnish/forge/src/data/graph/bloom/node"
	"github.com/philharnish/forge/src/data/graph/bloom/query"
	"github.com/philharnish/forge/src/data/graph/bloom/weight"
)

type reTrie struct {
	rootTrieNode *reTrieNode
	original     *regexp.Regexp // Used in Metadata().
	captureNames []string
	directory    *reTrieDirectory
}

func NewReTrie(regularExpression string, matchWeight weight.Weight) *reTrie {
	re, err := syntax.Parse(regularExpression, syntax.Perl)
	if err != nil {
		panic(err)
	}
	captureNames := processCaptureNames(re.CapNames())

	re = re.Simplify()
	directory := newDfaDirectory()
	matchNode := directory.addNode(node.NewNode(matchWeight))
	rootTrieNode := directory.linker(nil, matchNode, re, false)

	return &reTrie{
		rootTrieNode: rootTrieNode,
		original:     regexp.MustCompile(regularExpression),
		captureNames: captureNames,
		directory:    directory,
	}
}

func (root *reTrie) Items(acceptor node.NodeAcceptor) node.NodeItems {
	return root.rootTrieNode.Items(acceptor)
}

func (root *reTrie) Root() *node.Node {
	return root.rootTrieNode.Root()
}

func (root *reTrie) Header() query.QueryRowHeader {
	return root
}

func (root *reTrie) Labels() []string {
	return root.captureNames
}

func (root *reTrie) Metadata(path string) []weight.WeightedString {
	if len(root.captureNames) == 0 {
		return nil
	}
	submatches := root.original.FindStringSubmatch(path)
	result := make([]weight.WeightedString, len(submatches)-1)
	for i, submatch := range submatches[1:] {
		result[i].String = submatch
		result[i].Weight = 1
	}
	return result
}

func (root *reTrie) String() string {
	return root.rootTrieNode.String()
}

func processCaptureNames(captureNames []string) []string {
	captureNames = captureNames[1:]
	for i, name := range captureNames {
		if name == "" {
			captureNames[i] = fmt.Sprintf("%d", i+1)
		} else {
			captureNames[i] = name
		}
	}
	return captureNames
}
<commit_msg>Complete clean for retrie.go.<commit_after>package retrie

import (
	"fmt"
	"regexp"
	"regexp/syntax"

	"github.com/philharnish/forge/src/data/graph/bloom/node"
	"github.com/philharnish/forge/src/data/graph/bloom/query"
	"github.com/philharnish/forge/src/data/graph/bloom/weight"
)

type reTrie struct {
	root         *reTrieNode
	original     *regexp.Regexp // Used in Metadata().
	captureNames []string
	directory    *reTrieDirectory
}

func NewReTrie(regularExpression string, matchWeight weight.Weight) *reTrie {
	re, err := syntax.Parse(regularExpression, syntax.Perl)
	if err != nil {
		panic(err)
	}
	captureNames := processCaptureNames(re.CapNames())

	re = re.Simplify()
	directory := newDfaDirectory()
	matchNode := directory.addNode(node.NewNode(matchWeight))
	rootTrieNode := directory.linker(nil, matchNode, re, false)

	return &reTrie{
		root:         rootTrieNode,
		original:     regexp.MustCompile(regularExpression),
		captureNames: captureNames,
		directory:    directory,
	}
}

func (root *reTrie) Items(acceptor node.NodeAcceptor) node.NodeItems {
	return root.root.Items(acceptor)
}

func (root *reTrie) Root() *node.Node {
	return root.root.Root()
}

func (root *reTrie) Header() query.QueryRowHeader {
	return root
}

func (root *reTrie) Labels() []string {
	return root.captureNames
}

func (root *reTrie) Metadata(path string) []weight.WeightedString {
	if len(root.captureNames) == 0 {
		return nil
	}
	submatches := root.original.FindStringSubmatch(path)
	result := make([]weight.WeightedString, len(submatches)-1)
	for i, submatch := range submatches[1:] {
		result[i].String = submatch
		result[i].Weight = 1
	}
	return result
}

func (root *reTrie) String() string {
	return root.root.String()
}

func processCaptureNames(captureNames []string) []string {
	captureNames = captureNames[1:]
	for i, name := range captureNames {
		if name == "" {
			captureNames[i] = fmt.Sprintf("%d", i+1)
		} else {
			captureNames[i] = name
		}
	}
	return captureNames
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 VMware, Inc. All Rights Reserved.
//
// 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 job

import (
	"fmt"
	"sync"

	"github.com/vmware/harbor/src/common/models"
	"github.com/vmware/harbor/src/common/utils/log"
	"github.com/vmware/harbor/src/jobservice/config"
	"github.com/vmware/harbor/src/jobservice/replication"
	"github.com/vmware/harbor/src/jobservice/scan"
)

// SM is the state machine to handle job, it handles one job at a time.
type SM struct {
	CurrentJob    Job
	CurrentState  string
	PreviousState string
	//The states that don't have to exist in transition map, such as "Error", "Canceled"
	ForcedStates map[string]struct{}
	Transitions  map[string]map[string]struct{}
	Handlers     map[string]StateHandler
	desiredState string
	Logger       *log.Logger
	lock         *sync.Mutex
}

// EnterState transit the statemachine from the current state to the state in parameter.
// It returns the next state the statemachine should tranit to.
func (sm *SM) EnterState(s string) (string, error) {
	log.Debugf("Job: %v, transiting from State: %s, to State: %s", sm.CurrentJob, sm.CurrentState, s)
	targets, ok := sm.Transitions[sm.CurrentState]
	_, exist := targets[s]
	_, isForced := sm.ForcedStates[s]
	if !exist && !isForced {
		return "", fmt.Errorf("job: %v, transition from %s to %s does not exist", sm.CurrentJob, sm.CurrentState, s)
	}
	exitHandler, ok := sm.Handlers[sm.CurrentState]
	if ok {
		if err := exitHandler.Exit(); err != nil {
			return "", err
		}
	} else {
		log.Debugf("Job: %v, no exit handler found for state:%s, skip", sm.CurrentJob, sm.CurrentState)
	}
	enterHandler, ok := sm.Handlers[s]
	var next = models.JobContinue
	var err error
	if ok {
		if next, err = enterHandler.Enter(); err != nil {
			return "", err
		}
	} else {
		log.Debugf("Job: %v, no handler found for state:%s, skip", sm.CurrentJob, s)
	}
	sm.PreviousState = sm.CurrentState
	sm.CurrentState = s
	log.Debugf("Job: %v, transition succeeded, current state: %s", sm.CurrentJob, s)
	return next, nil
}

// Start kicks off the statemachine to transit from current state to s, and moves on
// It will search the transit map if the next state is "_continue", and
// will enter error state if there's more than one possible path when next state is "_continue"
func (sm *SM) Start(s string) {
	n, err := sm.EnterState(s)
	log.Debugf("Job: %v, next state from handler: %s", sm.CurrentJob, n)
	for len(n) > 0 && err == nil {
		if d := sm.getDesiredState(); len(d) > 0 {
			log.Debugf("Job: %v, Desired state: %s, will ignore the next state from handler", sm.CurrentJob, d)
			n = d
			sm.setDesiredState("")
			continue
		}
		if n == models.JobContinue && len(sm.Transitions[sm.CurrentState]) == 1 {
			for n = range sm.Transitions[sm.CurrentState] {
				break
			}
			log.Debugf("Job: %v, Continue to state: %s", sm.CurrentJob, n)
			continue
		}
		if n == models.JobContinue && len(sm.Transitions[sm.CurrentState]) != 1 {
			log.Errorf("Job: %v, next state is continue but there are %d possible next states in transition table", sm.CurrentJob, len(sm.Transitions[sm.CurrentState]))
			err = fmt.Errorf("Unable to continue")
			break
		}
		n, err = sm.EnterState(n)
		log.Debugf("Job: %v, next state from handler: %s", sm.CurrentJob, n)
	}
	if err != nil {
		log.Warningf("Job: %v, the statemachine will enter error state due to error: %v", sm.CurrentJob, err)
		sm.EnterState(models.JobError)
	}
}

// AddTransition add a transition to the transition table of state machine, the handler is the handler of target state "to"
func (sm *SM) AddTransition(from string, to string, h StateHandler) {
	_, ok := sm.Transitions[from]
	if !ok {
		sm.Transitions[from] = make(map[string]struct{})
	}
	sm.Transitions[from][to] = struct{}{}
	sm.Handlers[to] = h
}

// RemoveTransition removes a transition from transition table of the state machine
func (sm *SM) RemoveTransition(from string, to string) {
	_, ok := sm.Transitions[from]
	if !ok {
		return
	}
	delete(sm.Transitions[from], to)
}

// Stop will set the desired state as "stopped" such that when next tranisition happen the state machine will stop handling the current job
// and the worker can release itself to the workerpool.
func (sm *SM) Stop(job Job) {
	log.Debugf("Trying to stop the job: %v", job)
	sm.lock.Lock()
	defer sm.lock.Unlock()
	//need to check if the sm switched to other job
	if job.ID() == sm.CurrentJob.ID() && job.Type() == sm.CurrentJob.Type() {
		sm.desiredState = models.JobStopped
		log.Debugf("Desired state of job %v is set to stopped", job)
	} else {
		log.Debugf("State machine has switched to job %v, so the action to stop job %v will be ignored", sm.CurrentJob, job)
	}
}

func (sm *SM) getDesiredState() string {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	return sm.desiredState
}

func (sm *SM) setDesiredState(s string) {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	sm.desiredState = s
}

// Init initialzie the state machine, it will be called once in the lifecycle of state machine.
func (sm *SM) Init() {
	sm.lock = &sync.Mutex{}
	sm.Handlers = make(map[string]StateHandler)
	sm.Transitions = make(map[string]map[string]struct{})
	sm.ForcedStates = map[string]struct{}{
		models.JobError:    struct{}{},
		models.JobStopped:  struct{}{},
		models.JobCanceled: struct{}{},
		models.JobRetrying: struct{}{},
	}
}

// Reset resets the state machine and after prereq checking, it will start handling the job.
func (sm *SM) Reset(j Job) error {
	//To ensure the Job visible to the thread to stop the SM
	sm.lock.Lock()
	sm.CurrentJob = j
	sm.desiredState = ""
	sm.lock.Unlock()

	var err error
	sm.Logger, err = NewLogger(j)
	if err != nil {
		return err
	}
	//init states handlers
	sm.Handlers = make(map[string]StateHandler)
	sm.Transitions = make(map[string]map[string]struct{})
	sm.CurrentState = models.JobPending

	sm.AddTransition(models.JobPending, models.JobRunning, StatusUpdater{sm.CurrentJob, models.JobRunning})
	sm.AddTransition(models.JobRetrying, models.JobRunning, StatusUpdater{sm.CurrentJob, models.JobRunning})
	sm.Handlers[models.JobError] = StatusUpdater{sm.CurrentJob, models.JobError}
	sm.Handlers[models.JobStopped] = StatusUpdater{sm.CurrentJob, models.JobStopped}
	sm.Handlers[models.JobCanceled] = StatusUpdater{sm.CurrentJob, models.JobCanceled}
	sm.Handlers[models.JobRetrying] = Retry{sm.CurrentJob}
	if err := sm.CurrentJob.Init(); err != nil {
		return err
	}
	if err := sm.initTransitions(); err != nil {
		return err
	}
	return sm.kickOff()
}

func (sm *SM) kickOff() error {
	if repJob, ok := sm.CurrentJob.(*RepJob); ok {
		if repJob.parm.Enabled == 0 {
			log.Debugf("The policy of job:%v is disabled, will cancel the job", repJob)
			_, err := sm.EnterState(models.JobCanceled)
			if err != nil {
				log.Warningf("For job: %v, failed to update state to 'canceled', error: %v", repJob, err)
			}
			return err
		}
	}
	log.Debugf("In kickOff: will start job: %v", sm.CurrentJob)
	sm.Start(models.JobRunning)
	return nil
}

func (sm *SM) initTransitions() error {
	switch sm.CurrentJob.Type() {
	case ReplicationType:
		repJob, ok := sm.CurrentJob.(*RepJob)
		if !ok {
			//Shouldn't be here.
			return fmt.Errorf("The job: %v is not a type of RepJob", sm.CurrentJob)
		}
		jobParm := repJob.parm
		if jobParm.Operation == models.RepOpTransfer {
			addImgTransferTransition(sm, jobParm)
		} else if jobParm.Operation == models.RepOpDelete {
			addImgDeleteTransition(sm, jobParm)
		} else {
			return fmt.Errorf("unsupported operation: %s", jobParm.Operation)
		}
	case ScanType:
		scanJob, ok := sm.CurrentJob.(*ScanJob)
		if !ok {
			//Shouldn't be here.
			return fmt.Errorf("The job: %v is not a type of ScanJob", sm.CurrentJob)
		}
		addImgScanTransition(sm, scanJob.parm)
		return nil
	default:
		return fmt.Errorf("Unsupported job type: %v", sm.CurrentJob.Type())
	}
	return nil
}

//for testing onlly
/*
func addTestTransition(sm *SM) error {
	sm.AddTransition(models.JobRunning, "pull-img", ImgPuller{img: sm.Parms.Repository, logger: sm.Logger})
	return nil
}
*/

func addImgScanTransition(sm *SM, parm *ScanJobParm) {
	ctx := &scan.JobContext{
		Repository: parm.Repository,
		Tag:        parm.Tag,
		Digest:     parm.Digest,
		JobID:      sm.CurrentJob.ID(),
		Logger:     sm.Logger,
	}

	layerScanHandler := &scan.LayerScanHandler{Context: ctx}
	sm.AddTransition(models.JobRunning, scan.StateInitialize, &scan.Initializer{Context: ctx})
	sm.AddTransition(scan.StateInitialize, scan.StateScanLayer, layerScanHandler)
	sm.AddTransition(scan.StateScanLayer, scan.StateScanLayer, layerScanHandler)
	sm.AddTransition(scan.StateScanLayer, scan.StateSummarize, &scan.SummarizeHandler{Context: ctx})
	sm.AddTransition(scan.StateSummarize, models.JobFinished, &StatusUpdater{sm.CurrentJob, models.JobFinished})
}

func addImgTransferTransition(sm *SM, parm *RepJobParm) {
	base := replication.InitBaseHandler(parm.Repository, parm.LocalRegURL, config.JobserviceSecret(),
		parm.TargetURL, parm.TargetUsername, parm.TargetPassword,
		parm.Insecure, parm.Tags, sm.Logger)

	sm.AddTransition(models.JobRunning, replication.StateInitialize, &replication.Initializer{BaseHandler: base})
	sm.AddTransition(replication.StateInitialize, replication.StateCheck, &replication.Checker{BaseHandler: base})
	sm.AddTransition(replication.StateCheck, replication.StatePullManifest, &replication.ManifestPuller{BaseHandler: base})
	sm.AddTransition(replication.StatePullManifest, replication.StateTransferBlob, &replication.BlobTransfer{BaseHandler: base})
	sm.AddTransition(replication.StatePullManifest, models.JobFinished, &StatusUpdater{sm.CurrentJob, models.JobFinished})
	sm.AddTransition(replication.StateTransferBlob, replication.StatePushManifest, &replication.ManifestPusher{BaseHandler: base})
	sm.AddTransition(replication.StatePushManifest, replication.StatePullManifest, &replication.ManifestPuller{BaseHandler: base})
}

func addImgDeleteTransition(sm *SM, parm *RepJobParm) {
	deleter := replication.NewDeleter(parm.Repository, parm.Tags, parm.TargetURL,
		parm.TargetUsername, parm.TargetPassword, parm.Insecure, sm.Logger)

	sm.AddTransition(models.JobRunning, replication.StateDelete, deleter)
	sm.AddTransition(replication.StateDelete, models.JobFinished, &StatusUpdater{sm.CurrentJob, models.JobFinished})
}
<commit_msg>State machine recover from panic and set job state to error<commit_after>// Copyright (c) 2017 VMware, Inc. All Rights Reserved.
//
// 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 job

import (
	"fmt"
	"sync"

	"github.com/vmware/harbor/src/common/models"
	"github.com/vmware/harbor/src/common/utils/log"
	"github.com/vmware/harbor/src/jobservice/config"
	"github.com/vmware/harbor/src/jobservice/replication"
	"github.com/vmware/harbor/src/jobservice/scan"
)

// SM is the state machine to handle job, it handles one job at a time.
type SM struct {
	CurrentJob    Job
	CurrentState  string
	PreviousState string
	//The states that don't have to exist in transition map, such as "Error", "Canceled"
	ForcedStates map[string]struct{}
	Transitions  map[string]map[string]struct{}
	Handlers     map[string]StateHandler
	desiredState string
	Logger       *log.Logger
	lock         *sync.Mutex
}

// EnterState transit the statemachine from the current state to the state in parameter.
// It returns the next state the statemachine should tranit to.
func (sm *SM) EnterState(s string) (string, error) {
	log.Debugf("Job: %v, transiting from State: %s, to State: %s", sm.CurrentJob, sm.CurrentState, s)
	targets, ok := sm.Transitions[sm.CurrentState]
	_, exist := targets[s]
	_, isForced := sm.ForcedStates[s]
	if !exist && !isForced {
		return "", fmt.Errorf("job: %v, transition from %s to %s does not exist", sm.CurrentJob, sm.CurrentState, s)
	}
	exitHandler, ok := sm.Handlers[sm.CurrentState]
	if ok {
		if err := exitHandler.Exit(); err != nil {
			return "", err
		}
	} else {
		log.Debugf("Job: %v, no exit handler found for state:%s, skip", sm.CurrentJob, sm.CurrentState)
	}
	enterHandler, ok := sm.Handlers[s]
	var next = models.JobContinue
	var err error
	if ok {
		if next, err = enterHandler.Enter(); err != nil {
			return "", err
		}
	} else {
		log.Debugf("Job: %v, no handler found for state:%s, skip", sm.CurrentJob, s)
	}
	sm.PreviousState = sm.CurrentState
	sm.CurrentState = s
	log.Debugf("Job: %v, transition succeeded, current state: %s", sm.CurrentJob, s)
	return next, nil
}

// Start kicks off the statemachine to transit from current state to s, and moves on
// It will search the transit map if the next state is "_continue", and
// will enter error state if there's more than one possible path when next state is "_continue"
func (sm *SM) Start(s string) {
	defer func() {
		if r := recover(); r != nil {
			sm.Logger.Errorf("Panic: %v, entering error state", r)
			log.Warningf("Panic when handling job: %v, panic: %v, entering error state", sm.CurrentJob, r)
			sm.EnterState(models.JobError)
		}
	}()
	n, err := sm.EnterState(s)
	log.Debugf("Job: %v, next state from handler: %s", sm.CurrentJob, n)
	for len(n) > 0 && err == nil {
		if d := sm.getDesiredState(); len(d) > 0 {
			log.Debugf("Job: %v, Desired state: %s, will ignore the next state from handler", sm.CurrentJob, d)
			n = d
			sm.setDesiredState("")
			continue
		}
		if n == models.JobContinue && len(sm.Transitions[sm.CurrentState]) == 1 {
			for n = range sm.Transitions[sm.CurrentState] {
				break
			}
			log.Debugf("Job: %v, Continue to state: %s", sm.CurrentJob, n)
			continue
		}
		if n == models.JobContinue && len(sm.Transitions[sm.CurrentState]) != 1 {
			log.Errorf("Job: %v, next state is continue but there are %d possible next states in transition table", sm.CurrentJob, len(sm.Transitions[sm.CurrentState]))
			err = fmt.Errorf("Unable to continue")
			break
		}
		n, err = sm.EnterState(n)
		log.Debugf("Job: %v, next state from handler: %s", sm.CurrentJob, n)
	}
	if err != nil {
		log.Warningf("Job: %v, the statemachine will enter error state due to error: %v", sm.CurrentJob, err)
		sm.EnterState(models.JobError)
	}
}

// AddTransition add a transition to the transition table of state machine, the handler is the handler of target state "to"
func (sm *SM) AddTransition(from string, to string, h StateHandler) {
	_, ok := sm.Transitions[from]
	if !ok {
		sm.Transitions[from] = make(map[string]struct{})
	}
	sm.Transitions[from][to] = struct{}{}
	sm.Handlers[to] = h
}

// RemoveTransition removes a transition from transition table of the state machine
func (sm *SM) RemoveTransition(from string, to string) {
	_, ok := sm.Transitions[from]
	if !ok {
		return
	}
	delete(sm.Transitions[from], to)
}

// Stop will set the desired state as "stopped" such that when next tranisition happen the state machine will stop handling the current job
// and the worker can release itself to the workerpool.
func (sm *SM) Stop(job Job) {
	log.Debugf("Trying to stop the job: %v", job)
	sm.lock.Lock()
	defer sm.lock.Unlock()
	//need to check if the sm switched to other job
	if job.ID() == sm.CurrentJob.ID() && job.Type() == sm.CurrentJob.Type() {
		sm.desiredState = models.JobStopped
		log.Debugf("Desired state of job %v is set to stopped", job)
	} else {
		log.Debugf("State machine has switched to job %v, so the action to stop job %v will be ignored", sm.CurrentJob, job)
	}
}

func (sm *SM) getDesiredState() string {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	return sm.desiredState
}

func (sm *SM) setDesiredState(s string) {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	sm.desiredState = s
}

// Init initialzie the state machine, it will be called once in the lifecycle of state machine.
func (sm *SM) Init() {
	sm.lock = &sync.Mutex{}
	sm.Handlers = make(map[string]StateHandler)
	sm.Transitions = make(map[string]map[string]struct{})
	sm.ForcedStates = map[string]struct{}{
		models.JobError:    struct{}{},
		models.JobStopped:  struct{}{},
		models.JobCanceled: struct{}{},
		models.JobRetrying: struct{}{},
	}
}

// Reset resets the state machine and after prereq checking, it will start handling the job.
func (sm *SM) Reset(j Job) error {
	//To ensure the Job visible to the thread to stop the SM
	sm.lock.Lock()
	sm.CurrentJob = j
	sm.desiredState = ""
	sm.lock.Unlock()

	var err error
	sm.Logger, err = NewLogger(j)
	if err != nil {
		return err
	}
	//init states handlers
	sm.Handlers = make(map[string]StateHandler)
	sm.Transitions = make(map[string]map[string]struct{})
	sm.CurrentState = models.JobPending

	sm.AddTransition(models.JobPending, models.JobRunning, StatusUpdater{sm.CurrentJob, models.JobRunning})
	sm.AddTransition(models.JobRetrying, models.JobRunning, StatusUpdater{sm.CurrentJob, models.JobRunning})
	sm.Handlers[models.JobError] = StatusUpdater{sm.CurrentJob, models.JobError}
	sm.Handlers[models.JobStopped] = StatusUpdater{sm.CurrentJob, models.JobStopped}
	sm.Handlers[models.JobCanceled] = StatusUpdater{sm.CurrentJob, models.JobCanceled}
	sm.Handlers[models.JobRetrying] = Retry{sm.CurrentJob}
	if err := sm.CurrentJob.Init(); err != nil {
		return err
	}
	if err := sm.initTransitions(); err != nil {
		return err
	}
	return sm.kickOff()
}

func (sm *SM) kickOff() error {
	if repJob, ok := sm.CurrentJob.(*RepJob); ok {
		if repJob.parm.Enabled == 0 {
			log.Debugf("The policy of job:%v is disabled, will cancel the job", repJob)
			_, err := sm.EnterState(models.JobCanceled)
			if err != nil {
				log.Warningf("For job: %v, failed to update state to 'canceled', error: %v", repJob, err)
			}
			return err
		}
	}
	log.Debugf("In kickOff: will start job: %v", sm.CurrentJob)
	sm.Start(models.JobRunning)
	return nil
}

func (sm *SM) initTransitions() error {
	switch sm.CurrentJob.Type() {
	case ReplicationType:
		repJob, ok := sm.CurrentJob.(*RepJob)
		if !ok {
			//Shouldn't be here.
			return fmt.Errorf("The job: %v is not a type of RepJob", sm.CurrentJob)
		}
		jobParm := repJob.parm
		if jobParm.Operation == models.RepOpTransfer {
			addImgTransferTransition(sm, jobParm)
		} else if jobParm.Operation == models.RepOpDelete {
			addImgDeleteTransition(sm, jobParm)
		} else {
			return fmt.Errorf("unsupported operation: %s", jobParm.Operation)
		}
	case ScanType:
		scanJob, ok := sm.CurrentJob.(*ScanJob)
		if !ok {
			//Shouldn't be here.
			return fmt.Errorf("The job: %v is not a type of ScanJob", sm.CurrentJob)
		}
		addImgScanTransition(sm, scanJob.parm)
		return nil
	default:
		return fmt.Errorf("Unsupported job type: %v", sm.CurrentJob.Type())
	}
	return nil
}

//for testing onlly
/*
func addTestTransition(sm *SM) error {
	sm.AddTransition(models.JobRunning, "pull-img", ImgPuller{img: sm.Parms.Repository, logger: sm.Logger})
	return nil
}
*/

func addImgScanTransition(sm *SM, parm *ScanJobParm) {
	ctx := &scan.JobContext{
		Repository: parm.Repository,
		Tag:        parm.Tag,
		Digest:     parm.Digest,
		JobID:      sm.CurrentJob.ID(),
		Logger:     sm.Logger,
	}

	layerScanHandler := &scan.LayerScanHandler{Context: ctx}
	sm.AddTransition(models.JobRunning, scan.StateInitialize, &scan.Initializer{Context: ctx})
	sm.AddTransition(scan.StateInitialize, scan.StateScanLayer, layerScanHandler)
	sm.AddTransition(scan.StateScanLayer, scan.StateScanLayer, layerScanHandler)
	sm.AddTransition(scan.StateScanLayer, scan.StateSummarize, &scan.SummarizeHandler{Context: ctx})
	sm.AddTransition(scan.StateSummarize, models.JobFinished, &StatusUpdater{sm.CurrentJob, models.JobFinished})
}

func addImgTransferTransition(sm *SM, parm *RepJobParm) {
	base := replication.InitBaseHandler(parm.Repository, parm.LocalRegURL, config.JobserviceSecret(),
		parm.TargetURL, parm.TargetUsername, parm.TargetPassword,
		parm.Insecure, parm.Tags, sm.Logger)

	sm.AddTransition(models.JobRunning, replication.StateInitialize, &replication.Initializer{BaseHandler: base})
	sm.AddTransition(replication.StateInitialize, replication.StateCheck, &replication.Checker{BaseHandler: base})
	sm.AddTransition(replication.StateCheck, replication.StatePullManifest, &replication.ManifestPuller{BaseHandler: base})
	sm.AddTransition(replication.StatePullManifest, replication.StateTransferBlob, &replication.BlobTransfer{BaseHandler: base})
	sm.AddTransition(replication.StatePullManifest, models.JobFinished, &StatusUpdater{sm.CurrentJob, models.JobFinished})
	sm.AddTransition(replication.StateTransferBlob, replication.StatePushManifest, &replication.ManifestPusher{BaseHandler: base})
	sm.AddTransition(replication.StatePushManifest, replication.StatePullManifest, &replication.ManifestPuller{BaseHandler: base})
}

func addImgDeleteTransition(sm *SM, parm *RepJobParm) {
	deleter := replication.NewDeleter(parm.Repository, parm.Tags, parm.TargetURL,
		parm.TargetUsername, parm.TargetPassword, parm.Insecure, sm.Logger)

	sm.AddTransition(models.JobRunning, replication.StateDelete, deleter)
	sm.AddTransition(replication.StateDelete, models.JobFinished, &StatusUpdater{sm.CurrentJob, models.JobFinished})
}
<|endoftext|>


<commit_before>package gcache

import (
	"bytes"
	"errors"
	"math/rand"
	"net/http"
	"time"

	"google.golang.org/api/drive/v3"
	"google.golang.org/api/googleapi"
)

var (
	random *rand.Rand
)

func init() {
	random = rand.New(rand.NewSource(time.Now().UnixNano()))
}

// StoreGDrive stores a file to Google Drive.
func StoreGDrive(
	r *http.Request,
	file *drive.File,
	payload *[]byte,
) (
	*drive.File,
	error,
) {

	if file.Name == "" {
		return nil, errors.New("`file.Name` must be enough")
	}

	existFile, service, err := getGDriveFile(r, file.Name, "")
	switch err.(type) {
	case nil:
	case *DriveFileDoesNotExistError:
	default:
		return nil, err
	}

	if existFile == nil {
		file.MimeType = MimeGSuiteDoc
		if file.Parents == nil {
			folderID, err := getRootFolderID(r)
			if err != nil {
				return nil, err
			}
			file.Parents = []string{folderID}
		}
	}

	contentType := googleapi.ContentType(mimeTxt)
	var (
		newFile    *drive.File
		clearToken bool
	)

	n := 1
retry:
	payloadReader := bytes.NewReader(*payload)
	if existFile == nil {
		<-tokenBucketGDriveAPI
		newFile, err = service.Files.Create(file).Media(payloadReader, contentType).Do()
	} else {
		<-tokenBucketGDriveAPI
		newFile, err = service.Files.Update(existFile.Id, existFile).Media(payloadReader, contentType).Do()
	}

	if err == nil {
		return newFile, nil
	}
	clearToken, n, err = Triable(n, err)
	if err != nil {
		return nil, err
	}
	if clearToken {
		service, err = GetGDriveService(r)
		if err != nil {
			return nil, err
		}
	}
	goto retry
}
<commit_msg>refactor<commit_after>package gcache

import (
	"bytes"
	"errors"
	"math/rand"
	"net/http"
	"time"

	"google.golang.org/api/drive/v3"
	"google.golang.org/api/googleapi"
)

var (
	random *rand.Rand
)

func init() {
	random = rand.New(rand.NewSource(time.Now().UnixNano()))
}

// StoreGDrive stores a file to Google Drive.
func StoreGDrive(
	r *http.Request,
	file *drive.File,
	payload *[]byte,
) (
	*drive.File,
	error,
) {

	if file.Name == "" {
		return nil, errors.New("`file.Name` must be enough")
	}

	existFile, service, err := getGDriveFile(r, file.Name, "")
	switch err.(type) {
	case nil:
	case *DriveFileDoesNotExistError:
		file.MimeType = MimeGSuiteDoc
		if file.Parents == nil {
			folderID, err := getRootFolderID(r)
			if err != nil {
				return nil, err
			}
			file.Parents = []string{folderID}
		}
	default:
		return nil, err
	}

	contentType := googleapi.ContentType(mimeTxt)
	var (
		newFile    *drive.File
		clearToken bool
	)

	n := 1
retry:
	payloadReader := bytes.NewReader(*payload)
	if existFile == nil {
		<-tokenBucketGDriveAPI
		newFile, err = service.Files.Create(file).Media(payloadReader, contentType).Do()
	} else {
		<-tokenBucketGDriveAPI
		newFile, err = service.Files.Update(existFile.Id, existFile).Media(payloadReader, contentType).Do()
	}

	if err == nil {
		return newFile, nil
	}
	clearToken, n, err = Triable(n, err)
	if err != nil {
		return nil, err
	}
	if clearToken {
		service, err = GetGDriveService(r)
		if err != nil {
			return nil, err
		}
	}
	goto retry
}
<|endoftext|>


<commit_before>package s3api

import (
	"crypto/md5"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	weed_server "github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
)

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConnsPerHost: 1024,
	}}
}

func (s3a *S3ApiServer) PutObjectHandler(w http.ResponseWriter, r *http.Request) {

	// http://docs.aws.amazon.com/AmazonS3/latest/dev/UploadingObjects.html

	bucket, object := getBucketAndObject(r)

	_, err := validateContentMd5(r.Header)
	if err != nil {
		writeErrorResponse(w, ErrInvalidDigest, r.URL)
		return
	}

	rAuthType := getRequestAuthType(r)
	dataReader := r.Body
	var s3ErrCode ErrorCode
	switch rAuthType {
	case authTypeStreamingSigned:
		dataReader, s3ErrCode = s3a.iam.newSignV4ChunkedReader(r)
	case authTypeSignedV2, authTypePresignedV2:
		_, s3ErrCode = s3a.iam.isReqAuthenticatedV2(r)
	case authTypePresigned, authTypeSigned:
		_, s3ErrCode = s3a.iam.reqSignatureV4Verify(r)
	}
	if s3ErrCode != ErrNone {
		writeErrorResponse(w, s3ErrCode, r.URL)
		return
	}
	defer dataReader.Close()

	uploadUrl := fmt.Sprintf("http://%s%s/%s%s", s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	etag, errCode := s3a.putToFiler(r, uploadUrl, dataReader)

	if errCode != ErrNone {
		writeErrorResponse(w, errCode, r.URL)
		return
	}

	setEtag(w, etag)

	writeSuccessResponseEmpty(w)
}

func (s3a *S3ApiServer) GetObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	if strings.HasSuffix(r.URL.Path, "/") {
		writeErrorResponse(w, ErrNotImplemented, r.URL)
		return
	}

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) HeadObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) DeleteObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s?recursive=true",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, func(proxyResponse *http.Response, w http.ResponseWriter) {
		for k, v := range proxyResponse.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(http.StatusNoContent)
	})

}

// / ObjectIdentifier carries key name for the object to delete.
type ObjectIdentifier struct {
	ObjectName string `xml:"Key"`
}

// DeleteObjectsRequest - xml carrying the object key names which needs to be deleted.
type DeleteObjectsRequest struct {
	// Element to enable quiet mode for the request
	Quiet bool
	// List of objects to be deleted
	Objects []ObjectIdentifier `xml:"Object"`
}

// DeleteError structure.
type DeleteError struct {
	Code    string
	Message string
	Key     string
}

// DeleteObjectsResponse container for multiple object deletes.
type DeleteObjectsResponse struct {
	XMLName xml.Name `xml:"http://s3.amazonaws.com/doc/2006-03-01/ DeleteResult" json:"-"`

	// Collection of all deleted objects
	DeletedObjects []ObjectIdentifier `xml:"Deleted,omitempty"`

	// Collection of errors deleting certain objects.
	Errors []DeleteError `xml:"Error,omitempty"`
}

// DeleteMultipleObjectsHandler - Delete multiple objects
func (s3a *S3ApiServer) DeleteMultipleObjectsHandler(w http.ResponseWriter, r *http.Request) {

	bucket, _ := getBucketAndObject(r)

	deleteXMLBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeErrorResponse(w, ErrInternalError, r.URL)
		return
	}

	deleteObjects := &DeleteObjectsRequest{}
	if err := xml.Unmarshal(deleteXMLBytes, deleteObjects); err != nil {
		writeErrorResponse(w, ErrMalformedXML, r.URL)
		return
	}

	var deletedObjects []ObjectIdentifier
	var deleteErrors []DeleteError

	s3a.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		for _, object := range deleteObjects.Objects {
			lastSeparator := strings.LastIndex(object.ObjectName, "/")
			parentDirectoryPath, entryName, isDeleteData, isRecursive := "/", object.ObjectName, true, true
			if lastSeparator > 0 && lastSeparator+1 < len(object.ObjectName) {
				entryName = object.ObjectName[lastSeparator+1:]
				parentDirectoryPath = "/" + object.ObjectName[:lastSeparator]
			}
			parentDirectoryPath = fmt.Sprintf("%s/%s%s", s3a.option.BucketsPath, bucket, parentDirectoryPath)

			err := doDeleteEntry(client, parentDirectoryPath, entryName, isDeleteData, isRecursive)
			if err == nil {
				deletedObjects = append(deletedObjects, object)
			} else {
				deleteErrors = append(deleteErrors, DeleteError{
					Code:    "",
					Message: err.Error(),
					Key:     object.ObjectName,
				})
			}
		}
		return nil
	})

	deleteResp := DeleteObjectsResponse{}
	if !deleteObjects.Quiet {
		deleteResp.DeletedObjects = deletedObjects
	}
	deleteResp.Errors = deleteErrors

	writeSuccessResponseXML(w, encodeResponse(deleteResp))

}

func (s3a *S3ApiServer) proxyToFiler(w http.ResponseWriter, r *http.Request, destUrl string, responseFn func(proxyResponse *http.Response, w http.ResponseWriter)) {

	glog.V(2).Infof("s3 proxying %s to %s", r.Method, destUrl)

	proxyReq, err := http.NewRequest(r.Method, destUrl, r.Body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", destUrl, err)
		writeErrorResponse(w, ErrInternalError, r.URL)
		return
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		writeErrorResponse(w, ErrInternalError, r.URL)
		return
	}
	defer util.CloseResponse(resp)

	responseFn(resp, w)

}
func passThroughResponse(proxyResponse *http.Response, w http.ResponseWriter) {
	for k, v := range proxyResponse.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(proxyResponse.StatusCode)
	io.Copy(w, proxyResponse.Body)
}

func (s3a *S3ApiServer) putToFiler(r *http.Request, uploadUrl string, dataReader io.Reader) (etag string, code ErrorCode) {

	hash := md5.New()
	var body = io.TeeReader(dataReader, hash)

	proxyReq, err := http.NewRequest("PUT", uploadUrl, body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", uploadUrl, err)
		return "", ErrInternalError
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		return "", ErrInternalError
	}
	defer resp.Body.Close()

	etag = fmt.Sprintf("%x", hash.Sum(nil))

	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		glog.Errorf("upload to filer response read: %v", ra_err)
		return etag, ErrInternalError
	}
	var ret weed_server.FilerPostResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.Errorf("failing to read upload to %s : %v", uploadUrl, string(resp_body))
		return "", ErrInternalError
	}
	if ret.Error != "" {
		glog.Errorf("upload to filer error: %v", ret.Error)
		return "", ErrInternalError
	}

	return etag, ErrNone
}

func setEtag(w http.ResponseWriter, etag string) {
	if etag != "" {
		if strings.HasPrefix(etag, "\"") {
			w.Header().Set("ETag", etag)
		} else {
			w.Header().Set("ETag", "\""+etag+"\"")
		}
	}
}

func getBucketAndObject(r *http.Request) (bucket, object string) {
	vars := mux.Vars(r)
	bucket = vars["bucket"]
	object = vars["object"]
	if !strings.HasPrefix(object, "/") {
		object = "/" + object
	}

	return
}
<commit_msg>PutObject: fix create folder bug<commit_after>package s3api

import (
	"crypto/md5"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	weed_server "github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
)

var (
	client *http.Client
)

func init() {
	client = &http.Client{Transport: &http.Transport{
		MaxIdleConnsPerHost: 1024,
	}}
}

func (s3a *S3ApiServer) PutObjectHandler(w http.ResponseWriter, r *http.Request) {

	// http://docs.aws.amazon.com/AmazonS3/latest/dev/UploadingObjects.html

	bucket, object := getBucketAndObject(r)

	_, err := validateContentMd5(r.Header)
	if err != nil {
		writeErrorResponse(w, ErrInvalidDigest, r.URL)
		return
	}

	rAuthType := getRequestAuthType(r)
	dataReader := r.Body
	var s3ErrCode ErrorCode
	switch rAuthType {
	case authTypeStreamingSigned:
		dataReader, s3ErrCode = s3a.iam.newSignV4ChunkedReader(r)
	case authTypeSignedV2, authTypePresignedV2:
		_, s3ErrCode = s3a.iam.isReqAuthenticatedV2(r)
	case authTypePresigned, authTypeSigned:
		_, s3ErrCode = s3a.iam.reqSignatureV4Verify(r)
	}
	if s3ErrCode != ErrNone {
		writeErrorResponse(w, s3ErrCode, r.URL)
		return
	}
	defer dataReader.Close()

	if strings.HasSuffix(object, "/") {
		if err := s3a.mkdir(s3a.option.BucketsPath, bucket+object, nil); err != nil {
			writeErrorResponse(w, ErrInternalError, r.URL)
			return
		}
	} else {
		uploadUrl := fmt.Sprintf("http://%s%s/%s%s", s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

		etag, errCode := s3a.putToFiler(r, uploadUrl, dataReader)

		if errCode != ErrNone {
			writeErrorResponse(w, errCode, r.URL)
			return
		}

		setEtag(w, etag)
	}

	writeSuccessResponseEmpty(w)
}

func (s3a *S3ApiServer) GetObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	if strings.HasSuffix(r.URL.Path, "/") {
		writeErrorResponse(w, ErrNotImplemented, r.URL)
		return
	}

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) HeadObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, passThroughResponse)

}

func (s3a *S3ApiServer) DeleteObjectHandler(w http.ResponseWriter, r *http.Request) {

	bucket, object := getBucketAndObject(r)

	destUrl := fmt.Sprintf("http://%s%s/%s%s?recursive=true",
		s3a.option.Filer, s3a.option.BucketsPath, bucket, object)

	s3a.proxyToFiler(w, r, destUrl, func(proxyResponse *http.Response, w http.ResponseWriter) {
		for k, v := range proxyResponse.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(http.StatusNoContent)
	})

}

// / ObjectIdentifier carries key name for the object to delete.
type ObjectIdentifier struct {
	ObjectName string `xml:"Key"`
}

// DeleteObjectsRequest - xml carrying the object key names which needs to be deleted.
type DeleteObjectsRequest struct {
	// Element to enable quiet mode for the request
	Quiet bool
	// List of objects to be deleted
	Objects []ObjectIdentifier `xml:"Object"`
}

// DeleteError structure.
type DeleteError struct {
	Code    string
	Message string
	Key     string
}

// DeleteObjectsResponse container for multiple object deletes.
type DeleteObjectsResponse struct {
	XMLName xml.Name `xml:"http://s3.amazonaws.com/doc/2006-03-01/ DeleteResult" json:"-"`

	// Collection of all deleted objects
	DeletedObjects []ObjectIdentifier `xml:"Deleted,omitempty"`

	// Collection of errors deleting certain objects.
	Errors []DeleteError `xml:"Error,omitempty"`
}

// DeleteMultipleObjectsHandler - Delete multiple objects
func (s3a *S3ApiServer) DeleteMultipleObjectsHandler(w http.ResponseWriter, r *http.Request) {

	bucket, _ := getBucketAndObject(r)

	deleteXMLBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeErrorResponse(w, ErrInternalError, r.URL)
		return
	}

	deleteObjects := &DeleteObjectsRequest{}
	if err := xml.Unmarshal(deleteXMLBytes, deleteObjects); err != nil {
		writeErrorResponse(w, ErrMalformedXML, r.URL)
		return
	}

	var deletedObjects []ObjectIdentifier
	var deleteErrors []DeleteError

	s3a.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		for _, object := range deleteObjects.Objects {
			lastSeparator := strings.LastIndex(object.ObjectName, "/")
			parentDirectoryPath, entryName, isDeleteData, isRecursive := "/", object.ObjectName, true, true
			if lastSeparator > 0 && lastSeparator+1 < len(object.ObjectName) {
				entryName = object.ObjectName[lastSeparator+1:]
				parentDirectoryPath = "/" + object.ObjectName[:lastSeparator]
			}
			parentDirectoryPath = fmt.Sprintf("%s/%s%s", s3a.option.BucketsPath, bucket, parentDirectoryPath)

			err := doDeleteEntry(client, parentDirectoryPath, entryName, isDeleteData, isRecursive)
			if err == nil {
				deletedObjects = append(deletedObjects, object)
			} else {
				deleteErrors = append(deleteErrors, DeleteError{
					Code:    "",
					Message: err.Error(),
					Key:     object.ObjectName,
				})
			}
		}
		return nil
	})

	deleteResp := DeleteObjectsResponse{}
	if !deleteObjects.Quiet {
		deleteResp.DeletedObjects = deletedObjects
	}
	deleteResp.Errors = deleteErrors

	writeSuccessResponseXML(w, encodeResponse(deleteResp))

}

func (s3a *S3ApiServer) proxyToFiler(w http.ResponseWriter, r *http.Request, destUrl string, responseFn func(proxyResponse *http.Response, w http.ResponseWriter)) {

	glog.V(2).Infof("s3 proxying %s to %s", r.Method, destUrl)

	proxyReq, err := http.NewRequest(r.Method, destUrl, r.Body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", destUrl, err)
		writeErrorResponse(w, ErrInternalError, r.URL)
		return
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		writeErrorResponse(w, ErrInternalError, r.URL)
		return
	}
	defer util.CloseResponse(resp)

	responseFn(resp, w)

}
func passThroughResponse(proxyResponse *http.Response, w http.ResponseWriter) {
	for k, v := range proxyResponse.Header {
		w.Header()[k] = v
	}
	w.WriteHeader(proxyResponse.StatusCode)
	io.Copy(w, proxyResponse.Body)
}

func (s3a *S3ApiServer) putToFiler(r *http.Request, uploadUrl string, dataReader io.Reader) (etag string, code ErrorCode) {

	hash := md5.New()
	var body = io.TeeReader(dataReader, hash)

	proxyReq, err := http.NewRequest("PUT", uploadUrl, body)

	if err != nil {
		glog.Errorf("NewRequest %s: %v", uploadUrl, err)
		return "", ErrInternalError
	}

	proxyReq.Header.Set("Host", s3a.option.Filer)
	proxyReq.Header.Set("X-Forwarded-For", r.RemoteAddr)

	for header, values := range r.Header {
		for _, value := range values {
			proxyReq.Header.Add(header, value)
		}
	}

	resp, postErr := client.Do(proxyReq)

	if postErr != nil {
		glog.Errorf("post to filer: %v", postErr)
		return "", ErrInternalError
	}
	defer resp.Body.Close()

	etag = fmt.Sprintf("%x", hash.Sum(nil))

	resp_body, ra_err := ioutil.ReadAll(resp.Body)
	if ra_err != nil {
		glog.Errorf("upload to filer response read: %v", ra_err)
		return etag, ErrInternalError
	}
	var ret weed_server.FilerPostResult
	unmarshal_err := json.Unmarshal(resp_body, &ret)
	if unmarshal_err != nil {
		glog.Errorf("failing to read upload to %s : %v", uploadUrl, string(resp_body))
		return "", ErrInternalError
	}
	if ret.Error != "" {
		glog.Errorf("upload to filer error: %v", ret.Error)
		return "", ErrInternalError
	}

	return etag, ErrNone
}

func setEtag(w http.ResponseWriter, etag string) {
	if etag != "" {
		if strings.HasPrefix(etag, "\"") {
			w.Header().Set("ETag", etag)
		} else {
			w.Header().Set("ETag", "\""+etag+"\"")
		}
	}
}

func getBucketAndObject(r *http.Request) (bucket, object string) {
	vars := mux.Vars(r)
	bucket = vars["bucket"]
	object = vars["object"]
	if !strings.HasPrefix(object, "/") {
		object = "/" + object
	}

	return
}
<|endoftext|>


<commit_before>// Package local is a file system backed store
package file

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/micro/go-micro/v2/store"
	bolt "go.etcd.io/bbolt"
)

var (
	// DefaultDatabase is the namespace that the bbolt store
	// will use if no namespace is provided.
	DefaultDatabase = "micro"
	// DefaultTable when none is specified
	DefaultTable = "micro"
	// DefaultDir is the default directory for bbolt files
	DefaultDir = filepath.Join(os.TempDir(), "micro", "store")

	// bucket used for data storage
	dataBucket = "data"
)

// NewStore returns a memory store
func NewStore(opts ...store.Option) store.Store {
	s := &fileStore{
		handles: make(map[string]*fileHandle),
	}
	s.init(opts...)
	return s
}

type fileStore struct {
	options store.Options
	dir     string

	// the database handle
	sync.RWMutex
	handles map[string]*fileHandle
}

type fileHandle struct {
	key string
	db  *bolt.DB
}

// record stored by us
type record struct {
	Key       string
	Value     []byte
	Metadata  map[string]interface{}
	ExpiresAt time.Time
}

func key(database, table string) string {
	return database + ":" + table
}

func (m *fileStore) delete(fd *fileHandle, key string) error {
	return fd.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}
		return b.Delete([]byte(key))
	})
}

func (m *fileStore) init(opts ...store.Option) error {
	for _, o := range opts {
		o(&m.options)
	}

	if m.options.Database == "" {
		m.options.Database = DefaultDatabase
	}

	if m.options.Table == "" {
		// bbolt requires bucketname to not be empty
		m.options.Table = DefaultTable
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, m.options.Database)
	// Ignoring this as the folder might exist.
	// Reads/Writes updates will return with sensible error messages
	// about the dir not existing in case this cannot create the path anyway
	os.MkdirAll(dir, 0700)

	return nil
}

func (f *fileStore) getDB(database, table string) (*fileHandle, error) {
	if len(database) == 0 {
		database = f.options.Database
	}
	if len(table) == 0 {
		table = f.options.Table
	}

	k := key(database, table)
	f.RLock()
	fd, ok := f.handles[k]
	f.RUnlock()

	// return the file handle
	if ok {
		return fd, nil
	}

	// double check locking
	f.Lock()
	defer f.Unlock()
	if fd, ok := f.handles[k]; ok {
		return fd, nil
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, database)
	// create the database handle
	fname := table + ".db"
	// make the dir
	os.MkdirAll(dir, 0700)
	// database path
	dbPath := filepath.Join(dir, fname)

	// create new db handle
	// Bolt DB only allows one process to open the file R/W so make sure we're doing this under a lock
	db, err := bolt.Open(dbPath, 0700, &bolt.Options{Timeout: 5 * time.Second})
	if err != nil {
		return nil, err
	}
	fd = &fileHandle{
		key: k,
		db:  db,
	}
	f.handles[k] = fd

	return fd, nil
}

func (m *fileStore) list(fd *fileHandle, limit, offset uint) []string {
	var allItems []string

	fd.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		// nothing to read
		if b == nil {
			return nil
		}

		// @todo very inefficient
		if err := b.ForEach(func(k, v []byte) error {
			storedRecord := &record{}

			if err := json.Unmarshal(v, storedRecord); err != nil {
				return err
			}

			if !storedRecord.ExpiresAt.IsZero() {
				if storedRecord.ExpiresAt.Before(time.Now()) {
					return nil
				}
			}

			allItems = append(allItems, string(k))

			return nil
		}); err != nil {
			return err
		}

		return nil
	})

	allKeys := make([]string, len(allItems))

	for i, k := range allItems {
		allKeys[i] = k
	}

	if limit != 0 || offset != 0 {
		sort.Slice(allKeys, func(i, j int) bool { return allKeys[i] < allKeys[j] })
		min := func(i, j uint) uint {
			if i < j {
				return i
			}
			return j
		}
		return allKeys[offset:min(limit, uint(len(allKeys)))]
	}

	return allKeys
}

func (m *fileStore) get(fd *fileHandle, k string) (*store.Record, error) {
	var value []byte

	fd.db.View(func(tx *bolt.Tx) error {
		// @todo this is still very experimental...
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}

		value = b.Get([]byte(k))
		return nil
	})

	if value == nil {
		return nil, store.ErrNotFound
	}

	storedRecord := &record{}

	if err := json.Unmarshal(value, storedRecord); err != nil {
		return nil, err
	}

	newRecord := &store.Record{}
	newRecord.Key = storedRecord.Key
	newRecord.Value = storedRecord.Value
	newRecord.Metadata = make(map[string]interface{})

	for k, v := range storedRecord.Metadata {
		newRecord.Metadata[k] = v
	}

	if !storedRecord.ExpiresAt.IsZero() {
		if storedRecord.ExpiresAt.Before(time.Now()) {
			return nil, store.ErrNotFound
		}
		newRecord.Expiry = time.Until(storedRecord.ExpiresAt)
	}

	return newRecord, nil
}

func (m *fileStore) set(fd *fileHandle, r *store.Record) error {
	// copy the incoming record and then
	// convert the expiry in to a hard timestamp
	item := &record{}
	item.Key = r.Key
	item.Value = r.Value
	item.Metadata = make(map[string]interface{})

	if r.Expiry != 0 {
		item.ExpiresAt = time.Now().Add(r.Expiry)
	}

	for k, v := range r.Metadata {
		item.Metadata[k] = v
	}

	// marshal the data
	data, _ := json.Marshal(item)

	return fd.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			var err error
			b, err = tx.CreateBucketIfNotExists([]byte(dataBucket))
			if err != nil {
				return err
			}
		}
		return b.Put([]byte(r.Key), data)
	})
}

func (f *fileStore) Close() error {
	f.Lock()
	defer f.Unlock()
	for k, v := range f.handles {
		v.db.Close()
		delete(f.handles, k)
	}
	return nil
}

func (f *fileStore) Init(opts ...store.Option) error {
	return f.init(opts...)
}

func (m *fileStore) Delete(key string, opts ...store.DeleteOption) error {
	var deleteOptions store.DeleteOptions
	for _, o := range opts {
		o(&deleteOptions)
	}

	fd, err := m.getDB(deleteOptions.Database, deleteOptions.Table)
	if err != nil {
		return err
	}

	return m.delete(fd, key)
}

func (m *fileStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var readOpts store.ReadOptions
	for _, o := range opts {
		o(&readOpts)
	}

	fd, err := m.getDB(readOpts.Database, readOpts.Table)
	if err != nil {
		return nil, err
	}

	var keys []string

	// Handle Prefix / suffix
	// TODO: do range scan here rather than listing all keys
	if readOpts.Prefix || readOpts.Suffix {
		// list the keys
		k := m.list(fd, readOpts.Limit, readOpts.Offset)

		// check for prefix and suffix
		for _, v := range k {
			if readOpts.Prefix && !strings.HasPrefix(v, key) {
				continue
			}
			if readOpts.Suffix && !strings.HasSuffix(v, key) {
				continue
			}
			keys = append(keys, v)
		}
	} else {
		keys = []string{key}
	}

	var results []*store.Record

	for _, k := range keys {
		r, err := m.get(fd, k)
		if err != nil {
			return results, err
		}
		results = append(results, r)
	}

	return results, nil
}

func (m *fileStore) Write(r *store.Record, opts ...store.WriteOption) error {
	var writeOpts store.WriteOptions
	for _, o := range opts {
		o(&writeOpts)
	}

	fd, err := m.getDB(writeOpts.Database, writeOpts.Table)
	if err != nil {
		return err
	}

	if len(opts) > 0 {
		// Copy the record before applying options, or the incoming record will be mutated
		newRecord := store.Record{}
		newRecord.Key = r.Key
		newRecord.Value = r.Value
		newRecord.Metadata = make(map[string]interface{})
		newRecord.Expiry = r.Expiry

		if !writeOpts.Expiry.IsZero() {
			newRecord.Expiry = time.Until(writeOpts.Expiry)
		}
		if writeOpts.TTL != 0 {
			newRecord.Expiry = writeOpts.TTL
		}

		for k, v := range r.Metadata {
			newRecord.Metadata[k] = v
		}

		return m.set(fd, &newRecord)
	}

	return m.set(fd, r)
}

func (m *fileStore) Options() store.Options {
	return m.options
}

func (m *fileStore) List(opts ...store.ListOption) ([]string, error) {
	var listOptions store.ListOptions

	for _, o := range opts {
		o(&listOptions)
	}

	fd, err := m.getDB(listOptions.Database, listOptions.Table)
	if err != nil {
		return nil, err
	}

	// TODO apply prefix/suffix in range query
	allKeys := m.list(fd, listOptions.Limit, listOptions.Offset)

	if len(listOptions.Prefix) > 0 {
		var prefixKeys []string
		for _, k := range allKeys {
			if strings.HasPrefix(k, listOptions.Prefix) {
				prefixKeys = append(prefixKeys, k)
			}
		}
		allKeys = prefixKeys
	}

	if len(listOptions.Suffix) > 0 {
		var suffixKeys []string
		for _, k := range allKeys {
			if strings.HasSuffix(k, listOptions.Suffix) {
				suffixKeys = append(suffixKeys, k)
			}
		}
		allKeys = suffixKeys
	}

	return allKeys, nil
}

func (m *fileStore) String() string {
	return "file"
}
<commit_msg>Let bolt do locking per each List Get etc op, instead of managing fil… (#1831)<commit_after>// Package local is a file system backed store
package file

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/micro/go-micro/v2/store"
	bolt "go.etcd.io/bbolt"
)

var (
	// DefaultDatabase is the namespace that the bbolt store
	// will use if no namespace is provided.
	DefaultDatabase = "micro"
	// DefaultTable when none is specified
	DefaultTable = "micro"
	// DefaultDir is the default directory for bbolt files
	DefaultDir = filepath.Join(os.TempDir(), "micro", "store")

	// bucket used for data storage
	dataBucket = "data"
)

// NewStore returns a memory store
func NewStore(opts ...store.Option) store.Store {
	s := &fileStore{}
	s.init(opts...)
	return s
}

type fileStore struct {
	options store.Options
	dir     string
}

type fileHandle struct {
	key string
	db  *bolt.DB
}

// record stored by us
type record struct {
	Key       string
	Value     []byte
	Metadata  map[string]interface{}
	ExpiresAt time.Time
}

func key(database, table string) string {
	return database + ":" + table
}

func (m *fileStore) delete(db *bolt.DB, key string) error {
	return db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}
		return b.Delete([]byte(key))
	})
}

func (m *fileStore) init(opts ...store.Option) error {
	for _, o := range opts {
		o(&m.options)
	}

	if m.options.Database == "" {
		m.options.Database = DefaultDatabase
	}

	if m.options.Table == "" {
		// bbolt requires bucketname to not be empty
		m.options.Table = DefaultTable
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, m.options.Database)
	// Ignoring this as the folder might exist.
	// Reads/Writes updates will return with sensible error messages
	// about the dir not existing in case this cannot create the path anyway
	os.MkdirAll(dir, 0700)

	return nil
}

func (f *fileStore) getDB(database, table string) (*bolt.DB, error) {
	if len(database) == 0 {
		database = f.options.Database
	}
	if len(table) == 0 {
		table = f.options.Table
	}

	// create a directory /tmp/micro
	dir := filepath.Join(DefaultDir, database)
	// create the database handle
	fname := table + ".db"
	// make the dir
	os.MkdirAll(dir, 0700)
	// database path
	dbPath := filepath.Join(dir, fname)

	// create new db handle
	// Bolt DB only allows one process to open the file R/W so make sure we're doing this under a lock
	return bolt.Open(dbPath, 0700, &bolt.Options{Timeout: 5 * time.Second})
}

func (m *fileStore) list(db *bolt.DB, limit, offset uint) []string {
	var allItems []string

	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		// nothing to read
		if b == nil {
			return nil
		}

		// @todo very inefficient
		if err := b.ForEach(func(k, v []byte) error {
			storedRecord := &record{}

			if err := json.Unmarshal(v, storedRecord); err != nil {
				return err
			}

			if !storedRecord.ExpiresAt.IsZero() {
				if storedRecord.ExpiresAt.Before(time.Now()) {
					return nil
				}
			}

			allItems = append(allItems, string(k))

			return nil
		}); err != nil {
			return err
		}

		return nil
	})

	allKeys := make([]string, len(allItems))

	for i, k := range allItems {
		allKeys[i] = k
	}

	if limit != 0 || offset != 0 {
		sort.Slice(allKeys, func(i, j int) bool { return allKeys[i] < allKeys[j] })
		min := func(i, j uint) uint {
			if i < j {
				return i
			}
			return j
		}
		return allKeys[offset:min(limit, uint(len(allKeys)))]
	}

	return allKeys
}

func (m *fileStore) get(db *bolt.DB, k string) (*store.Record, error) {
	var value []byte

	db.View(func(tx *bolt.Tx) error {
		// @todo this is still very experimental...
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			return nil
		}

		value = b.Get([]byte(k))
		return nil
	})

	if value == nil {
		return nil, store.ErrNotFound
	}

	storedRecord := &record{}

	if err := json.Unmarshal(value, storedRecord); err != nil {
		return nil, err
	}

	newRecord := &store.Record{}
	newRecord.Key = storedRecord.Key
	newRecord.Value = storedRecord.Value
	newRecord.Metadata = make(map[string]interface{})

	for k, v := range storedRecord.Metadata {
		newRecord.Metadata[k] = v
	}

	if !storedRecord.ExpiresAt.IsZero() {
		if storedRecord.ExpiresAt.Before(time.Now()) {
			return nil, store.ErrNotFound
		}
		newRecord.Expiry = time.Until(storedRecord.ExpiresAt)
	}

	return newRecord, nil
}

func (m *fileStore) set(db *bolt.DB, r *store.Record) error {
	// copy the incoming record and then
	// convert the expiry in to a hard timestamp
	item := &record{}
	item.Key = r.Key
	item.Value = r.Value
	item.Metadata = make(map[string]interface{})

	if r.Expiry != 0 {
		item.ExpiresAt = time.Now().Add(r.Expiry)
	}

	for k, v := range r.Metadata {
		item.Metadata[k] = v
	}

	// marshal the data
	data, _ := json.Marshal(item)

	return db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(dataBucket))
		if b == nil {
			var err error
			b, err = tx.CreateBucketIfNotExists([]byte(dataBucket))
			if err != nil {
				return err
			}
		}
		return b.Put([]byte(r.Key), data)
	})
}

func (f *fileStore) Close() error {
	return nil
}

func (f *fileStore) Init(opts ...store.Option) error {
	return f.init(opts...)
}

func (m *fileStore) Delete(key string, opts ...store.DeleteOption) error {
	var deleteOptions store.DeleteOptions
	for _, o := range opts {
		o(&deleteOptions)
	}

	db, err := m.getDB(deleteOptions.Database, deleteOptions.Table)
	if err != nil {
		return err
	}
	defer db.Close()

	return m.delete(db, key)
}

func (m *fileStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var readOpts store.ReadOptions
	for _, o := range opts {
		o(&readOpts)
	}

	db, err := m.getDB(readOpts.Database, readOpts.Table)
	if err != nil {
		return nil, err
	}
	defer db.Close()

	var keys []string

	// Handle Prefix / suffix
	// TODO: do range scan here rather than listing all keys
	if readOpts.Prefix || readOpts.Suffix {
		// list the keys
		k := m.list(db, readOpts.Limit, readOpts.Offset)

		// check for prefix and suffix
		for _, v := range k {
			if readOpts.Prefix && !strings.HasPrefix(v, key) {
				continue
			}
			if readOpts.Suffix && !strings.HasSuffix(v, key) {
				continue
			}
			keys = append(keys, v)
		}
	} else {
		keys = []string{key}
	}

	var results []*store.Record

	for _, k := range keys {
		r, err := m.get(db, k)
		if err != nil {
			return results, err
		}
		results = append(results, r)
	}

	return results, nil
}

func (m *fileStore) Write(r *store.Record, opts ...store.WriteOption) error {
	var writeOpts store.WriteOptions
	for _, o := range opts {
		o(&writeOpts)
	}

	db, err := m.getDB(writeOpts.Database, writeOpts.Table)
	if err != nil {
		return err
	}
	defer db.Close()

	if len(opts) > 0 {
		// Copy the record before applying options, or the incoming record will be mutated
		newRecord := store.Record{}
		newRecord.Key = r.Key
		newRecord.Value = r.Value
		newRecord.Metadata = make(map[string]interface{})
		newRecord.Expiry = r.Expiry

		if !writeOpts.Expiry.IsZero() {
			newRecord.Expiry = time.Until(writeOpts.Expiry)
		}
		if writeOpts.TTL != 0 {
			newRecord.Expiry = writeOpts.TTL
		}

		for k, v := range r.Metadata {
			newRecord.Metadata[k] = v
		}

		return m.set(db, &newRecord)
	}

	return m.set(db, r)
}

func (m *fileStore) Options() store.Options {
	return m.options
}

func (m *fileStore) List(opts ...store.ListOption) ([]string, error) {
	var listOptions store.ListOptions

	for _, o := range opts {
		o(&listOptions)
	}

	db, err := m.getDB(listOptions.Database, listOptions.Table)
	if err != nil {
		return nil, err
	}
	defer db.Close()

	// TODO apply prefix/suffix in range query
	allKeys := m.list(db, listOptions.Limit, listOptions.Offset)

	if len(listOptions.Prefix) > 0 {
		var prefixKeys []string
		for _, k := range allKeys {
			if strings.HasPrefix(k, listOptions.Prefix) {
				prefixKeys = append(prefixKeys, k)
			}
		}
		allKeys = prefixKeys
	}

	if len(listOptions.Suffix) > 0 {
		var suffixKeys []string
		for _, k := range allKeys {
			if strings.HasSuffix(k, listOptions.Suffix) {
				suffixKeys = append(suffixKeys, k)
			}
		}
		allKeys = suffixKeys
	}

	return allKeys, nil
}

func (m *fileStore) String() string {
	return "file"
}
<|endoftext|>


<commit_before>package reform

import (
	"fmt"
	"strings"
)

func filteredColumnsAndValues(str Struct, columnsIn []string, isUpdate bool) (columns []string, values []interface{}, err error) {
	columnsSet := make(map[string]struct{}, len(columnsIn))
	for _, c := range columnsIn {
		columnsSet[c] = struct{}{}
	}

	// select columns from set and collect values
	view := str.View()
	allColumns := view.Columns()
	allValues := str.Values()
	columns = make([]string, 0, len(columnsSet))
	values = make([]interface{}, 0, len(columns))

	record, _ := str.(Record)
	var pk uint
	if record != nil {
		pk = view.(Table).PKColumnIndex()
	}

	for i, c := range allColumns {
		if _, ok := columnsSet[c]; ok {
			if isUpdate && record != nil && i == int(pk) {
				err = fmt.Errorf("reform: will not update PK column: %s", c)
				return
			}
			delete(columnsSet, c)
			columns = append(columns, c)
			values = append(values, allValues[i])
		}
	}

	// make error for extra columns
	if len(columnsSet) > 0 {
		columns = make([]string, 0, len(columnsSet))
		for c := range columnsSet {
			columns = append(columns, c)
		}
		// TODO make exported type for that error
		err = fmt.Errorf("reform: unexpected columns: %v", columns)
		return
	}

	return
}

func (q *Querier) insert(str Struct, columns []string, values []interface{}) error {
	for i, c := range columns {
		columns[i] = q.QuoteIdentifier(c)
	}
	placeholders := q.Placeholders(1, len(columns))

	view := str.View()
	record, _ := str.(Record)
	lastInsertIdMethod := q.LastInsertIdMethod()
	defaultValuesMethod := q.DefaultValuesMethod()

	var pk uint
	if record != nil {
		pk = view.(Table).PKColumnIndex()
	}

	// make query
	query := q.startQuery("INSERT") + " INTO " + q.QualifiedView(view)
	if len(columns) != 0 || defaultValuesMethod == EmptyLists {
		query += " (" + strings.Join(columns, ", ") + ")"
	}
	if record != nil && lastInsertIdMethod == OutputInserted {
		query += fmt.Sprintf(" OUTPUT INSERTED.%s", q.QuoteIdentifier(view.Columns()[pk]))
	}
	if len(placeholders) != 0 || defaultValuesMethod == EmptyLists {
		query += fmt.Sprintf(" VALUES (%s)", strings.Join(placeholders, ", "))
	} else {
		query += " DEFAULT VALUES"
	}
	if record != nil && lastInsertIdMethod == Returning {
		query += fmt.Sprintf(" RETURNING %s", q.QuoteIdentifier(view.Columns()[pk]))
	}

	switch lastInsertIdMethod {
	case LastInsertId:
		res, err := q.Exec(query, values...)
		if err != nil {
			return err
		}
		if record != nil && !record.HasPK() {
			id, err := res.LastInsertId()
			if err != nil {
				return err
			}
			record.SetPK(id)
		}
		return nil

	case Returning, OutputInserted:
		var err error
		if record != nil {
			err = q.QueryRow(query, values...).Scan(record.PKPointer())
		} else {
			_, err = q.Exec(query, values...)
		}
		return err

	default:
		panic("reform: Unhandled LastInsertIdMethod. Please report this bug.")
	}
}

func (q *Querier) beforeInsert(str Struct) error {
	return q.callStructMethod(str, "BeforeInsert")
}

func (q *Querier) afterInsert(str Struct) error {
	return q.callStructMethod(str, "AfterInsert")
}

// Insert inserts a struct into SQL database table.
// If str has valid method "BeforeInsert", it calls BeforeInsert() before doing so.
// If str has valid method "AfterInsert", it calls AfterInsert() after doing so.
//
// It fills record's primary key field.
func (q *Querier) Insert(str Struct) error {
	if err := q.beforeInsert(str); err != nil {
		return err
	}

	view := str.View()
	values := str.Values()
	columns := view.Columns()
	record, _ := str.(Record)

	if record != nil {
		pk := view.(Table).PKColumnIndex()

		// cut primary key
		if !record.HasPK() {
			values = append(values[:pk], values[pk+1:]...)
			columns = append(columns[:pk], columns[pk+1:]...)
		}
	}

	err := q.insert(str, columns, values)

	if err == nil {
		return q.afterInsert(str)
	}

	return nil
}

// InsertColumns inserts a struct into SQL database table with specified columns.
// Other columns are omitted from generated INSERT statement.
// If str has valid method "BeforeInsert", it calls BeforeInsert() before doing so.
// If str has valid method "AfterInsert", it calls AfterInsert() after doing so.
//
// It fills record's primary key field.
func (q *Querier) InsertColumns(str Struct, columns ...string) error {
	if err := q.beforeInsert(str); err != nil {
		return err
	}

	columns, values, err := filteredColumnsAndValues(str, columns, false)
	if err != nil {
		return err
	}

	err = q.insert(str, columns, values)

	if err == nil {
		return q.afterInsert(str)
	}

	return nil
}

// InsertMulti inserts several structs into SQL database table with single query.
// If they has valid method "BeforeInsert", it calls BeforeInsert() before doing so.
//
// All structs should belong to the same view/table.
// All records should either have or not have primary key set.
// It doesn't fill primary key fields.
// Given all these limitations, most users should use Querier.Insert in a loop, not this method.
func (q *Querier) InsertMulti(structs ...Struct) error {
	if len(structs) == 0 {
		return nil
	}

	// check that view is the same
	view := structs[0].View()
	for _, str := range structs {
		if str.View() != view {
			return fmt.Errorf("reform: different tables in InsertMulti: %s and %s", view.Name(), str.View().Name())
		}
	}

	for _, str := range structs {
		err := q.callStructMethod(str, "BeforeInsert")
		if err != nil {
			return err
		}
	}

	// check if all PK are present or all are absent
	record, _ := structs[0].(Record)
	if record != nil {
		for _, str := range structs {
			rec, _ := str.(Record)
			if record.HasPK() != rec.HasPK() {
				return fmt.Errorf("reform: PK in present in one struct and absent in other: first: %s, second: %s",
					record, rec)
			}
		}
	}

	columns := view.Columns()
	for i, c := range columns {
		columns[i] = q.QuoteIdentifier(c)
	}

	var pk uint
	if record != nil && !record.HasPK() {
		pk = view.(Table).PKColumnIndex()
		columns = append(columns[:pk], columns[pk+1:]...)
	}

	placeholders := q.Placeholders(1, len(columns)*len(structs))
	query := fmt.Sprintf("%s INTO %s (%s) VALUES ",
		q.startQuery("INSERT"),
		q.QualifiedView(view),
		strings.Join(columns, ", "),
	)
	for i := 0; i < len(structs); i++ {
		query += fmt.Sprintf("(%s), ", strings.Join(placeholders[len(columns)*i:len(columns)*(i+1)], ", "))
	}
	query = query[:len(query)-2] // cut last ", "

	values := make([]interface{}, 0, len(placeholders))
	for _, str := range structs {
		v := str.Values()
		if record != nil && !record.HasPK() {
			v = append(v[:pk], v[pk+1:]...)
		}
		values = append(values, v...)
	}

	_, err := q.Exec(query, values...)
	return err
}

func (q *Querier) update(record Record, columns []string, values []interface{}) error {
	for i, c := range columns {
		columns[i] = q.QuoteIdentifier(c)
	}
	placeholders := q.Placeholders(1, len(columns))

	p := make([]string, len(columns))
	for i, c := range columns {
		p[i] = c + " = " + placeholders[i]
	}
	table := record.Table()
	query := fmt.Sprintf("%s %s SET %s WHERE %s = %s",
		q.startQuery("UPDATE"),
		q.QualifiedView(table),
		strings.Join(p, ", "),
		q.QuoteIdentifier(table.Columns()[table.PKColumnIndex()]),
		q.Placeholder(len(columns)+1),
	)

	args := append(values, record.PKValue())
	res, err := q.Exec(query, args...)
	if err != nil {
		return err
	}
	ra, err := res.RowsAffected()
	if err != nil {
		return err
	}
	if ra == 0 {
		return ErrNoRows
	}
	if ra > 1 {
		panic(fmt.Sprintf("reform: %d rows by UPDATE by primary key. Please report this bug.", ra))
	}
	return nil
}

func (q *Querier) beforeUpdate(record Record) error {
	if !record.HasPK() {
		return ErrNoPK
	}

	return q.callStructMethod(record, "BeforeUpdate")
}

func (q *Querier) afterUpdate(record Record) error {
	return q.callStructMethod(record, "AfterUpdate")
}

// Update updates all columns of row specified by primary key in SQL database table with given record.
// If record has valid method "BeforeUpdate", it calls BeforeUpdate() before doing so.
// If record has valid method "AfterUpdate", it calls AfterUpdate() before doing so.
//
// Method returns ErrNoRows if no rows were updated.
// Method returns ErrNoPK if primary key is not set.
func (q *Querier) Update(record Record) error {
	if err := q.beforeUpdate(record); err != nil {
		return err
	}

	table := record.Table()
	values := record.Values()
	columns := table.Columns()

	// cut primary key
	pk := table.PKColumnIndex()
	values = append(values[:pk], values[pk+1:]...)
	columns = append(columns[:pk], columns[pk+1:]...)

	err := q.update(record, columns, values)

	if err == nil {
		return q.afterUpdate(record)
	}
	return nil
}

// UpdateColumns updates specified columns of row specified by primary key in SQL database table with given record.
// Other columns are omitted from generated UPDATE statement.
// If record has valid method "BeforeUpdate", it calls BeforeUpdate() before doing so.
// If record has valid method "AfterUpdate", it calls AfterUpdate() before doing so.
//
// Method returns ErrNoRows if no rows were updated.
// Method returns ErrNoPK if primary key is not set.
func (q *Querier) UpdateColumns(record Record, columns ...string) error {
	if err := q.beforeUpdate(record); err != nil {
		return err
	}

	columns, values, err := filteredColumnsAndValues(record, columns, true)
	if err != nil {
		return err
	}

	if len(values) == 0 {
		// TODO make exported type for that error
		return fmt.Errorf("reform: nothing to update")
	}

	err = q.update(record, columns, values)

	if err == nil {
		return q.afterUpdate(record)
	}
	return nil
}

// Save saves record in SQL database table.
// If primary key is set, it first calls Update and checks if row was updated.
// If primary key is absent or no row was updated, it calls Insert.
func (q *Querier) Save(record Record) error {
	if record.HasPK() {
		err := q.Update(record)
		if err != ErrNoRows {
			return err
		}
	}

	return q.Insert(record)
}

func (q *Querier) beforeDelete(record Record) error {
	if !record.HasPK() {
		return ErrNoPK
	}

	return q.callStructMethod(record, "BeforeDelete")
}

func (q *Querier) afterDelete(record Record) error {
	return q.callStructMethod(record, "AfterDelete")
}

// Delete deletes record from SQL database table by primary key.
// If record has valid method "BeforeDelete", it calls BeforeDelete() before doing so.
// If record has valid method "AfterDelete", it calls AfterDelete() before doing so.
//
// Method returns ErrNoRows if no rows were deleted.
// Method returns ErrNoPK if primary key is not set.
func (q *Querier) Delete(record Record) error {
	err := q.beforeDelete(record)
	if err != nil {
		return err
	}

	table := record.Table()
	pk := table.PKColumnIndex()
	query := fmt.Sprintf("%s FROM %s WHERE %s = %s",
		q.startQuery("DELETE"),
		q.QualifiedView(table),
		q.QuoteIdentifier(table.Columns()[pk]),
		q.Placeholder(1),
	)

	res, err := q.Exec(query, record.PKValue())
	if err != nil {
		return err
	}
	ra, err := res.RowsAffected()
	if err != nil {
		return err
	}
	if ra == 0 {
		return ErrNoRows
	}
	if ra > 1 {
		panic(fmt.Sprintf("reform: %d rows by DELETE by primary key. Please report this bug.", ra))
	}
	return q.afterDelete(record)
}

// DeleteFrom deletes rows from view with tail and args and returns a number of deleted rows.
//
// Method never returns ErrNoRows.
func (q *Querier) DeleteFrom(view View, tail string, args ...interface{}) (uint, error) {
	query := fmt.Sprintf("%s FROM %s %s",
		q.startQuery("DELETE"),
		q.QualifiedView(view),
		tail,
	)

	res, err := q.Exec(query, args...)
	if err != nil {
		return 0, err
	}
	ra, err := res.RowsAffected()
	if err != nil {
		return 0, err
	}
	return uint(ra), nil
}
<commit_msg>A bugfix: Insert*()/Update*() didn't return errors<commit_after>package reform

import (
	"fmt"
	"strings"
)

func filteredColumnsAndValues(str Struct, columnsIn []string, isUpdate bool) (columns []string, values []interface{}, err error) {
	columnsSet := make(map[string]struct{}, len(columnsIn))
	for _, c := range columnsIn {
		columnsSet[c] = struct{}{}
	}

	// select columns from set and collect values
	view := str.View()
	allColumns := view.Columns()
	allValues := str.Values()
	columns = make([]string, 0, len(columnsSet))
	values = make([]interface{}, 0, len(columns))

	record, _ := str.(Record)
	var pk uint
	if record != nil {
		pk = view.(Table).PKColumnIndex()
	}

	for i, c := range allColumns {
		if _, ok := columnsSet[c]; ok {
			if isUpdate && record != nil && i == int(pk) {
				err = fmt.Errorf("reform: will not update PK column: %s", c)
				return
			}
			delete(columnsSet, c)
			columns = append(columns, c)
			values = append(values, allValues[i])
		}
	}

	// make error for extra columns
	if len(columnsSet) > 0 {
		columns = make([]string, 0, len(columnsSet))
		for c := range columnsSet {
			columns = append(columns, c)
		}
		// TODO make exported type for that error
		err = fmt.Errorf("reform: unexpected columns: %v", columns)
		return
	}

	return
}

func (q *Querier) insert(str Struct, columns []string, values []interface{}) error {
	for i, c := range columns {
		columns[i] = q.QuoteIdentifier(c)
	}
	placeholders := q.Placeholders(1, len(columns))

	view := str.View()
	record, _ := str.(Record)
	lastInsertIdMethod := q.LastInsertIdMethod()
	defaultValuesMethod := q.DefaultValuesMethod()

	var pk uint
	if record != nil {
		pk = view.(Table).PKColumnIndex()
	}

	// make query
	query := q.startQuery("INSERT") + " INTO " + q.QualifiedView(view)
	if len(columns) != 0 || defaultValuesMethod == EmptyLists {
		query += " (" + strings.Join(columns, ", ") + ")"
	}
	if record != nil && lastInsertIdMethod == OutputInserted {
		query += fmt.Sprintf(" OUTPUT INSERTED.%s", q.QuoteIdentifier(view.Columns()[pk]))
	}
	if len(placeholders) != 0 || defaultValuesMethod == EmptyLists {
		query += fmt.Sprintf(" VALUES (%s)", strings.Join(placeholders, ", "))
	} else {
		query += " DEFAULT VALUES"
	}
	if record != nil && lastInsertIdMethod == Returning {
		query += fmt.Sprintf(" RETURNING %s", q.QuoteIdentifier(view.Columns()[pk]))
	}

	switch lastInsertIdMethod {
	case LastInsertId:
		res, err := q.Exec(query, values...)
		if err != nil {
			return err
		}
		if record != nil && !record.HasPK() {
			id, err := res.LastInsertId()
			if err != nil {
				return err
			}
			record.SetPK(id)
		}
		return nil

	case Returning, OutputInserted:
		var err error
		if record != nil {
			err = q.QueryRow(query, values...).Scan(record.PKPointer())
		} else {
			_, err = q.Exec(query, values...)
		}
		return err

	default:
		panic("reform: Unhandled LastInsertIdMethod. Please report this bug.")
	}
}

func (q *Querier) beforeInsert(str Struct) error {
	return q.callStructMethod(str, "BeforeInsert")
}

func (q *Querier) afterInsert(str Struct) error {
	return q.callStructMethod(str, "AfterInsert")
}

// Insert inserts a struct into SQL database table.
// If str has valid method "BeforeInsert", it calls BeforeInsert() before doing so.
// If str has valid method "AfterInsert", it calls AfterInsert() after doing so.
//
// It fills record's primary key field.
func (q *Querier) Insert(str Struct) error {
	if err := q.beforeInsert(str); err != nil {
		return err
	}

	view := str.View()
	values := str.Values()
	columns := view.Columns()
	record, _ := str.(Record)

	if record != nil {
		pk := view.(Table).PKColumnIndex()

		// cut primary key
		if !record.HasPK() {
			values = append(values[:pk], values[pk+1:]...)
			columns = append(columns[:pk], columns[pk+1:]...)
		}
	}

	err := q.insert(str, columns, values)

	if err == nil {
		return q.afterInsert(str)
	}

	return err
}

// InsertColumns inserts a struct into SQL database table with specified columns.
// Other columns are omitted from generated INSERT statement.
// If str has valid method "BeforeInsert", it calls BeforeInsert() before doing so.
// If str has valid method "AfterInsert", it calls AfterInsert() after doing so.
//
// It fills record's primary key field.
func (q *Querier) InsertColumns(str Struct, columns ...string) error {
	if err := q.beforeInsert(str); err != nil {
		return err
	}

	columns, values, err := filteredColumnsAndValues(str, columns, false)
	if err != nil {
		return err
	}

	err = q.insert(str, columns, values)

	if err == nil {
		return q.afterInsert(str)
	}

	return err
}

// InsertMulti inserts several structs into SQL database table with single query.
// If they has valid method "BeforeInsert", it calls BeforeInsert() before doing so.
//
// All structs should belong to the same view/table.
// All records should either have or not have primary key set.
// It doesn't fill primary key fields.
// Given all these limitations, most users should use Querier.Insert in a loop, not this method.
func (q *Querier) InsertMulti(structs ...Struct) error {
	if len(structs) == 0 {
		return nil
	}

	// check that view is the same
	view := structs[0].View()
	for _, str := range structs {
		if str.View() != view {
			return fmt.Errorf("reform: different tables in InsertMulti: %s and %s", view.Name(), str.View().Name())
		}
	}

	for _, str := range structs {
		err := q.callStructMethod(str, "BeforeInsert")
		if err != nil {
			return err
		}
	}

	// check if all PK are present or all are absent
	record, _ := structs[0].(Record)
	if record != nil {
		for _, str := range structs {
			rec, _ := str.(Record)
			if record.HasPK() != rec.HasPK() {
				return fmt.Errorf("reform: PK in present in one struct and absent in other: first: %s, second: %s",
					record, rec)
			}
		}
	}

	columns := view.Columns()
	for i, c := range columns {
		columns[i] = q.QuoteIdentifier(c)
	}

	var pk uint
	if record != nil && !record.HasPK() {
		pk = view.(Table).PKColumnIndex()
		columns = append(columns[:pk], columns[pk+1:]...)
	}

	placeholders := q.Placeholders(1, len(columns)*len(structs))
	query := fmt.Sprintf("%s INTO %s (%s) VALUES ",
		q.startQuery("INSERT"),
		q.QualifiedView(view),
		strings.Join(columns, ", "),
	)
	for i := 0; i < len(structs); i++ {
		query += fmt.Sprintf("(%s), ", strings.Join(placeholders[len(columns)*i:len(columns)*(i+1)], ", "))
	}
	query = query[:len(query)-2] // cut last ", "

	values := make([]interface{}, 0, len(placeholders))
	for _, str := range structs {
		v := str.Values()
		if record != nil && !record.HasPK() {
			v = append(v[:pk], v[pk+1:]...)
		}
		values = append(values, v...)
	}

	_, err := q.Exec(query, values...)
	return err
}

func (q *Querier) update(record Record, columns []string, values []interface{}) error {
	for i, c := range columns {
		columns[i] = q.QuoteIdentifier(c)
	}
	placeholders := q.Placeholders(1, len(columns))

	p := make([]string, len(columns))
	for i, c := range columns {
		p[i] = c + " = " + placeholders[i]
	}
	table := record.Table()
	query := fmt.Sprintf("%s %s SET %s WHERE %s = %s",
		q.startQuery("UPDATE"),
		q.QualifiedView(table),
		strings.Join(p, ", "),
		q.QuoteIdentifier(table.Columns()[table.PKColumnIndex()]),
		q.Placeholder(len(columns)+1),
	)

	args := append(values, record.PKValue())
	res, err := q.Exec(query, args...)
	if err != nil {
		return err
	}
	ra, err := res.RowsAffected()
	if err != nil {
		return err
	}
	if ra == 0 {
		return ErrNoRows
	}
	if ra > 1 {
		panic(fmt.Sprintf("reform: %d rows by UPDATE by primary key. Please report this bug.", ra))
	}
	return nil
}

func (q *Querier) beforeUpdate(record Record) error {
	if !record.HasPK() {
		return ErrNoPK
	}

	return q.callStructMethod(record, "BeforeUpdate")
}

func (q *Querier) afterUpdate(record Record) error {
	return q.callStructMethod(record, "AfterUpdate")
}

// Update updates all columns of row specified by primary key in SQL database table with given record.
// If record has valid method "BeforeUpdate", it calls BeforeUpdate() before doing so.
// If record has valid method "AfterUpdate", it calls AfterUpdate() before doing so.
//
// Method returns ErrNoRows if no rows were updated.
// Method returns ErrNoPK if primary key is not set.
func (q *Querier) Update(record Record) error {
	if err := q.beforeUpdate(record); err != nil {
		return err
	}

	table := record.Table()
	values := record.Values()
	columns := table.Columns()

	// cut primary key
	pk := table.PKColumnIndex()
	values = append(values[:pk], values[pk+1:]...)
	columns = append(columns[:pk], columns[pk+1:]...)

	err := q.update(record, columns, values)

	if err == nil {
		return q.afterUpdate(record)
	}
	return err
}

// UpdateColumns updates specified columns of row specified by primary key in SQL database table with given record.
// Other columns are omitted from generated UPDATE statement.
// If record has valid method "BeforeUpdate", it calls BeforeUpdate() before doing so.
// If record has valid method "AfterUpdate", it calls AfterUpdate() before doing so.
//
// Method returns ErrNoRows if no rows were updated.
// Method returns ErrNoPK if primary key is not set.
func (q *Querier) UpdateColumns(record Record, columns ...string) error {
	if err := q.beforeUpdate(record); err != nil {
		return err
	}

	columns, values, err := filteredColumnsAndValues(record, columns, true)
	if err != nil {
		return err
	}

	if len(values) == 0 {
		// TODO make exported type for that error
		return fmt.Errorf("reform: nothing to update")
	}

	err = q.update(record, columns, values)

	if err == nil {
		return q.afterUpdate(record)
	}
	return err
}

// Save saves record in SQL database table.
// If primary key is set, it first calls Update and checks if row was updated.
// If primary key is absent or no row was updated, it calls Insert.
func (q *Querier) Save(record Record) error {
	if record.HasPK() {
		err := q.Update(record)
		if err != ErrNoRows {
			return err
		}
	}

	return q.Insert(record)
}

func (q *Querier) beforeDelete(record Record) error {
	if !record.HasPK() {
		return ErrNoPK
	}

	return q.callStructMethod(record, "BeforeDelete")
}

func (q *Querier) afterDelete(record Record) error {
	return q.callStructMethod(record, "AfterDelete")
}

// Delete deletes record from SQL database table by primary key.
// If record has valid method "BeforeDelete", it calls BeforeDelete() before doing so.
// If record has valid method "AfterDelete", it calls AfterDelete() before doing so.
//
// Method returns ErrNoRows if no rows were deleted.
// Method returns ErrNoPK if primary key is not set.
func (q *Querier) Delete(record Record) error {
	err := q.beforeDelete(record)
	if err != nil {
		return err
	}

	table := record.Table()
	pk := table.PKColumnIndex()
	query := fmt.Sprintf("%s FROM %s WHERE %s = %s",
		q.startQuery("DELETE"),
		q.QualifiedView(table),
		q.QuoteIdentifier(table.Columns()[pk]),
		q.Placeholder(1),
	)

	res, err := q.Exec(query, record.PKValue())
	if err != nil {
		return err
	}
	ra, err := res.RowsAffected()
	if err != nil {
		return err
	}
	if ra == 0 {
		return ErrNoRows
	}
	if ra > 1 {
		panic(fmt.Sprintf("reform: %d rows by DELETE by primary key. Please report this bug.", ra))
	}
	return q.afterDelete(record)
}

// DeleteFrom deletes rows from view with tail and args and returns a number of deleted rows.
//
// Method never returns ErrNoRows.
func (q *Querier) DeleteFrom(view View, tail string, args ...interface{}) (uint, error) {
	query := fmt.Sprintf("%s FROM %s %s",
		q.startQuery("DELETE"),
		q.QualifiedView(view),
		tail,
	)

	res, err := q.Exec(query, args...)
	if err != nil {
		return 0, err
	}
	ra, err := res.RowsAffected()
	if err != nil {
		return 0, err
	}
	return uint(ra), nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"fmt"
	"bytes"
	"io"
)

const (
	FORMAT_DEFAULT = "[%D %T] [%L] (%S) %M"
	FORMAT_SHORT   = "[%t %d] [%L] %M"
	FORMAT_ABBREV  = "[%L] %M"
)

type formatCacheType struct {
	LastUpdateSeconds    int64
	shortTime, shortDate string
	longTime, longDate   string
}

var formatCache = &formatCacheType{}

// Known format codes:
// %T - Time (15:04:05 MST)
// %t - Time (15:04)
// %D - Date (2006/01/02)
// %d - Date (01/02/06)
// %L - Level (FNST, FINE, DEBG, TRAC, WARN, EROR, CRIT)
// %S - Source
// %M - Message
// Ignores unknown formats
// Recommended: "[%D %T] [%L] (%S) %M"
func FormatLogRecord(format string, rec *LogRecord) string {
	if rec == nil {
		return "<nil>"
	}
	if len(format) == 0 {
		return ""
	}

	out := bytes.NewBuffer(make([]byte, 0, 64))
	secs := rec.Created.UnixNano() / 1e9

	cache := *formatCache
	if cache.LastUpdateSeconds != secs {
		month, day, year := rec.Created.Month(), rec.Created.Day(), rec.Created.Year()
		hour, minute, second := rec.Created.Hour(), rec.Created.Minute(), rec.Created.Second()
		zone, _ := rec.Created.Zone()
		updated := &formatCacheType{
			LastUpdateSeconds: secs,
			shortTime:         fmt.Sprintf("%02d:%02d", hour, minute),
			shortDate:         fmt.Sprintf("%02d/%02d/%02d", month, day, year%100),
			longTime:          fmt.Sprintf("%02d:%02d:%02d %s", hour, minute, second, zone),
			longDate:          fmt.Sprintf("%04d/%02d/%02d", year, month, day),
		}
		cache = *updated
		formatCache = updated
	}

	// Split the string into pieces by % signs
	pieces := bytes.Split([]byte(format), []byte{'%'})

	// Iterate over the pieces, replacing known formats
	for i, piece := range pieces {
		if i > 0 && len(piece) > 0 {
			switch piece[0] {
			case 'T':
				out.WriteString(cache.longTime)
			case 't':
				out.WriteString(cache.shortTime)
			case 'D':
				out.WriteString(cache.longDate)
			case 'd':
				out.WriteString(cache.shortDate)
			case 'L':
				out.WriteString(levelStrings[rec.Level])
			case 'S':
				out.WriteString(rec.Source)
			case 'M':
				out.WriteString(rec.Message)
			}
			if len(piece) > 1 {
				out.Write(piece[1:])
			}
		} else if len(piece) > 0 {
			out.Write(piece)
		}
	}
	out.WriteByte('\n')

	return out.String()
}

// This is the standard writer that prints to standard output.
type FormatLogWriter chan *LogRecord

// This creates a new FormatLogWriter
func NewFormatLogWriter(out io.Writer, format string) FormatLogWriter {
	records := make(FormatLogWriter, LogBufferLength)
	go records.run(out, format)
	return records
}

func (w FormatLogWriter) run(out io.Writer, format string) {
	for rec := range w {
		fmt.Fprint(out, FormatLogRecord(format, rec))
	}
}

// This is the FormatLogWriter's output method.  This will block if the output
// buffer is full.
func (w FormatLogWriter) LogWrite(rec *LogRecord) {
	w <- rec
}

// Close stops the logger from sending messages to standard output.  Attempts to
// send log messages to this logger after a Close have undefined behavior.
func (w FormatLogWriter) Close() {
	close(w)
}
<commit_msg>修改 longtime 时间格式 添加 longTime 纳秒<commit_after>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"bytes"
	"fmt"
	"io"
	"strings"
)

const (
	FORMAT_DEFAULT = "[%D %T] [%L] (%S) %M"
	FORMAT_SHORT   = "[%t %d] [%L] %M"
	FORMAT_ABBREV  = "[%L] %M"
)

type formatCacheType struct {
	LastUpdateSeconds    int64
	shortTime, shortDate string
	longTime, longDate   string
}

var formatCache = &formatCacheType{}

// Known format codes:
// %T - Time (15:04:05 MST)
// %t - Time (15:04)
// %D - Date (2006/01/02)
// %d - Date (01/02/06)
// %L - Level (FNST, FINE, DEBG, TRAC, WARN, EROR, CRIT)
// %S - Source
// %M - Message
// Ignores unknown formats
// Recommended: "[%D %T] [%L] (%S) %M"
func FormatLogRecord(format string, rec *LogRecord) string {
	if rec == nil {
		return "<nil>"
	}
	if len(format) == 0 {
		return ""
	}

	out := bytes.NewBuffer(make([]byte, 0, 64))
	secs := rec.Created.UnixNano() / 1e9

	cache := *formatCache
	if cache.LastUpdateSeconds != secs {
		month, day, year := rec.Created.Month(), rec.Created.Day(), rec.Created.Year()
		hour, minute, second, nanosce := rec.Created.Hour(), rec.Created.Minute(), rec.Created.Second(), rec.Created.Nanosecond()
		zone, _ := rec.Created.Zone()
		updated := &formatCacheType{
			LastUpdateSeconds: secs,
			shortTime:         fmt.Sprintf("%02d:%02d", hour, minute),
			shortDate:         fmt.Sprintf("%02d/%02d/%02d", day, month, year%100),
			longTime:          fmt.Sprintf("%02d:%02d:%02d.%09d %s", hour, minute, second, nanosce, zone),
			longDate:          fmt.Sprintf("%04d/%02d/%02d", year, month, day),
		}
		cache = *updated
		formatCache = updated
	}

	// Split the string into pieces by % signs
	pieces := bytes.Split([]byte(format), []byte{'%'})

	// Iterate over the pieces, replacing known formats
	for i, piece := range pieces {
		if i > 0 && len(piece) > 0 {
			switch piece[0] {
			case 'T':
				out.WriteString(cache.longTime)
			case 't':
				out.WriteString(cache.shortTime)
			case 'D':
				out.WriteString(cache.longDate)
			case 'd':
				out.WriteString(cache.shortDate)
			case 'L':
				out.WriteString(levelStrings[rec.Level])
			case 'S':
				out.WriteString(rec.Source)
			case 's':
				slice := strings.Split(rec.Source, "/")
				out.WriteString(slice[len(slice)-1])
			case 'M':
				out.WriteString(rec.Message)
			}
			if len(piece) > 1 {
				out.Write(piece[1:])
			}
		} else if len(piece) > 0 {
			out.Write(piece)
		}
	}
	out.WriteByte('\n')

	return out.String()
}

// This is the standard writer that prints to standard output.
type FormatLogWriter chan *LogRecord

// This creates a new FormatLogWriter
func NewFormatLogWriter(out io.Writer, format string) FormatLogWriter {
	records := make(FormatLogWriter, LogBufferLength)
	go records.run(out, format)
	return records
}

func (w FormatLogWriter) run(out io.Writer, format string) {
	for rec := range w {
		fmt.Fprint(out, FormatLogRecord(format, rec))
	}
}

// This is the FormatLogWriter's output method.  This will block if the output
// buffer is full.
func (w FormatLogWriter) LogWrite(rec *LogRecord) {
	w <- rec
}

// Close stops the logger from sending messages to standard output.  Attempts to
// send log messages to this logger after a Close have undefined behavior.
func (w FormatLogWriter) Close() {
	close(w)
}
<|endoftext|>


<commit_before>package pb

import "fmt"

const (
	// AverageMetricAge average over a one-minute period, which means the average
	// age of the metrics is in the period of 30 seconds
	AverageMetricAge float64 = 30.0

	// Decay formula for computing the decay factor for average metric age
	Decay float64 = 2 / (float64(AverageMetricAge) + 1)
)

func init() {
	fmt.Println("ewma decay = ", Decay)
}

// EWMA represents the exponentially weighted moving average of a series of numbers.
type EWMA struct {
	value float64 // The current value of the average.
}

// Add a value to the series and update the moving average.
func (e *EWMA) Add(value float64) {
	if e.value == 0 { // perhaps first input, no decay factor needed
		e.value = value
		return
	}
	e.value = (value * Decay) + (e.value * (1 - Decay))
}

// Value returns the current value of the moving average.
func (e *EWMA) Value() float64 {
	return e.value
}
<commit_msg>Don't print ewma decay<commit_after>package pb

import "fmt"

const (
	// AverageMetricAge average over a one-minute period, which means the average
	// age of the metrics is in the period of 30 seconds
	AverageMetricAge float64 = 30.0

	// Decay formula for computing the decay factor for average metric age
	Decay float64 = 2 / (float64(AverageMetricAge) + 1)
)

// EWMA represents the exponentially weighted moving average of a series of numbers.
type EWMA struct {
	value float64 // The current value of the average.
}

// Add a value to the series and update the moving average.
func (e *EWMA) Add(value float64) {
	if e.value == 0 { // perhaps first input, no decay factor needed
		e.value = value
		return
	}
	e.value = (value * Decay) + (e.value * (1 - Decay))
}

// Value returns the current value of the moving average.
func (e *EWMA) Value() float64 {
	return e.value
}
<|endoftext|>