<commit_before>package api

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/bytom/blockchain/pseudohsm"
	"github.com/bytom/blockchain/txbuilder"
	"github.com/bytom/consensus"
	"github.com/bytom/errors"
	"github.com/bytom/math/checked"
	"github.com/bytom/net/http/reqid"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

var defaultTxTTL = 5 * time.Minute

func (a *API) actionDecoder(action string) (func([]byte) (txbuilder.Action, error), bool) {
	var decoder func([]byte) (txbuilder.Action, error)
	switch action {
	case "control_address":
		decoder = txbuilder.DecodeControlAddressAction
	case "control_program":
		decoder = txbuilder.DecodeControlProgramAction
	case "control_receiver":
		decoder = txbuilder.DecodeControlReceiverAction
	case "issue":
		decoder = a.wallet.AssetReg.DecodeIssueAction
	case "retire":
		decoder = txbuilder.DecodeRetireAction
	case "spend_account":
		decoder = a.wallet.AccountMgr.DecodeSpendAction
	case "spend_account_unspent_output":
		decoder = a.wallet.AccountMgr.DecodeSpendUTXOAction
	default:
		return nil, false
	}
	return decoder, true
}

func mergeActions(req *BuildRequest) []map[string]interface{} {
	var actions []map[string]interface{}
	actionMap := make(map[string]map[string]interface{})

	for _, m := range req.Actions {
		if actionType := m["type"].(string); actionType != "spend_account" {
			actions = append(actions, m)
			continue
		}

		actionKey := m["asset_id"].(string) + m["account_id"].(string)
		amountNumber := m["amount"].(json.Number)
		amount, _ := amountNumber.Int64()

		if tmpM, ok := actionMap[actionKey]; ok {
			tmpNumber, _ := tmpM["amount"].(json.Number)
			tmpAmount, _ := tmpNumber.Int64()
			tmpM["amount"] = json.Number(fmt.Sprintf("%v", tmpAmount+amount))
		} else {
			actionMap[actionKey] = m
			actions = append(actions, m)
		}
	}

	return actions
}

func onlyHaveSpendActions(req *BuildRequest) bool {
	count := 0
	for _, m := range req.Actions {
		if actionType := m["type"].(string); strings.HasPrefix(actionType, "spend") {
			count++
		}
	}

	return count == len(req.Actions)
}

func (a *API) buildSingle(ctx context.Context, req *BuildRequest) (*txbuilder.Template, error) {
	err := a.filterAliases(ctx, req)
	if err != nil {
		return nil, err
	}

	if onlyHaveSpendActions(req) {
		return nil, errors.New("transaction only contain spend actions, didn't have output actions")
	}

	reqActions := mergeActions(req)
	actions := make([]txbuilder.Action, 0, len(reqActions))
	for i, act := range reqActions {
		typ, ok := act["type"].(string)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "no action type provided on action %d", i)
		}
		decoder, ok := a.actionDecoder(typ)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "unknown action type %q on action %d", typ, i)
		}

		// Remarshal to JSON, the action may have been modified when we
		// filtered aliases.
		b, err := json.Marshal(act)
		if err != nil {
			return nil, err
		}
		action, err := decoder(b)
		if err != nil {
			return nil, errors.WithDetailf(errBadAction, "%s on action %d", err.Error(), i)
		}
		actions = append(actions, action)
	}

	ttl := req.TTL.Duration
	if ttl == 0 {
		ttl = defaultTxTTL
	}
	maxTime := time.Now().Add(ttl)

	tpl, err := txbuilder.Build(ctx, req.Tx, actions, maxTime, req.TimeRange)
	if errors.Root(err) == txbuilder.ErrAction {
		// append each of the inner errors contained in the data.
		var Errs string
		for _, innerErr := range errors.Data(err)["actions"].([]error) {
			Errs = Errs + "<" + innerErr.Error() + ">"
		}
		err = errors.New(err.Error() + "-" + Errs)
	}
	if err != nil {
		return nil, err
	}

	// ensure null is never returned for signing instructions
	if tpl.SigningInstructions == nil {
		tpl.SigningInstructions = []*txbuilder.SigningInstruction{}
	}
	return tpl, nil
}

// POST /build-transaction
func (a *API) build(ctx context.Context, buildReqs *BuildRequest) Response {
	subctx := reqid.NewSubContext(ctx, reqid.New())

	tmpl, err := a.buildSingle(subctx, buildReqs)
	if err != nil {
		return NewErrorResponse(err)
	}

	return NewSuccessResponse(tmpl)
}

func (a *API) submitSingle(ctx context.Context, tpl *txbuilder.Template) (map[string]string, error) {
	if tpl.Transaction == nil {
		return nil, errors.Wrap(txbuilder.ErrMissingRawTx)
	}

	if err := txbuilder.FinalizeTx(ctx, a.chain, tpl.Transaction); err != nil {
		return nil, errors.Wrapf(err, "tx %s", tpl.Transaction.ID.String())
	}

	return map[string]string{"tx_id": tpl.Transaction.ID.String()}, nil
}

type submitTxResp struct {
	TxID *bc.Hash `json:"tx_id"`
}

// POST /submit-transaction
func (a *API) submit(ctx context.Context, ins struct {
	Tx types.Tx `json:"raw_transaction"`
}) Response {
	if err := txbuilder.FinalizeTx(ctx, a.chain, &ins.Tx); err != nil {
		return NewErrorResponse(err)
	}

	log.WithField("tx_id", ins.Tx.ID).Info("submit single tx")
	return NewSuccessResponse(&submitTxResp{TxID: &ins.Tx.ID})
}

// POST /sign-submit-transaction
func (a *API) signSubmit(ctx context.Context, x struct {
	Password string             `json:"password"`
	Txs      txbuilder.Template `json:"transaction"`
}) Response {
	if err := txbuilder.Sign(ctx, &x.Txs, nil, x.Password, a.pseudohsmSignTemplate); err != nil {
		log.WithField("build err", err).Error("fail on sign transaction.")
		return NewErrorResponse(err)
	}

	if signCount, complete := txbuilder.SignInfo(&x.Txs); !complete && signCount == 0 {
		return NewErrorResponse(pseudohsm.ErrLoadKey)
	}
	log.Info("Sign Transaction complete.")

	txID, err := a.submitSingle(nil, &x.Txs)
	if err != nil {
		log.WithField("err", err).Error("submit single tx")
		return NewErrorResponse(err)
	}

	log.WithField("tx_id", txID["tx_id"]).Info("submit single tx")
	return NewSuccessResponse(txID)
}

type calculateTxGasResp struct {
	LeftBTM     int64 `json:"left_btm"`
	ConsumedBTM int64 `json:"consumed_btm"`
	LeftGas     int64 `json:"left_gas"`
	ConsumedGas int64 `json:"consumed_gas"`
	StorageGas  int64 `json:"storage_gas"`
	VMGas       int64 `json:"vm_gas"`
}

// POST /estimate-transaction-gas
func (a *API) estimateTxGas(ctx context.Context, ins struct {
	Tx types.Tx `json:"raw_transaction"`
}) Response {
	gasState, err := txbuilder.EstimateTxGas(a.chain, &ins.Tx)
	if err != nil {
		return NewErrorResponse(err)
	}

	btmLeft, ok := checked.MulInt64(gasState.GasLeft, consensus.VMGasRate)
	if !ok {
		return NewErrorResponse(errors.New("calculate btmleft got a math error"))
	}

	txGasResp := &calculateTxGasResp{
		LeftBTM:     btmLeft,
		ConsumedBTM: int64(gasState.BTMValue) - btmLeft,
		LeftGas:     gasState.GasLeft,
		ConsumedGas: gasState.GasUsed,
		StorageGas:  gasState.StorageGas,
		VMGas:       gasState.GasUsed - gasState.StorageGas,
	}

	return NewSuccessResponse(txGasResp)
}
<commit_msg>optimise<commit_after>package api

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/bytom/blockchain/pseudohsm"
	"github.com/bytom/blockchain/txbuilder"
	"github.com/bytom/consensus"
	"github.com/bytom/errors"
	"github.com/bytom/math/checked"
	"github.com/bytom/net/http/reqid"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

var defaultTxTTL = 5 * time.Minute

func (a *API) actionDecoder(action string) (func([]byte) (txbuilder.Action, error), bool) {
	var decoder func([]byte) (txbuilder.Action, error)
	switch action {
	case "control_address":
		decoder = txbuilder.DecodeControlAddressAction
	case "control_program":
		decoder = txbuilder.DecodeControlProgramAction
	case "control_receiver":
		decoder = txbuilder.DecodeControlReceiverAction
	case "issue":
		decoder = a.wallet.AssetReg.DecodeIssueAction
	case "retire":
		decoder = txbuilder.DecodeRetireAction
	case "spend_account":
		decoder = a.wallet.AccountMgr.DecodeSpendAction
	case "spend_account_unspent_output":
		decoder = a.wallet.AccountMgr.DecodeSpendUTXOAction
	default:
		return nil, false
	}
	return decoder, true
}

func mergeActions(req *BuildRequest) []map[string]interface{} {
	var actions []map[string]interface{}
	actionMap := make(map[string]map[string]interface{})

	for _, m := range req.Actions {
		if actionType := m["type"].(string); actionType != "spend_account" {
			actions = append(actions, m)
			continue
		}

		actionKey := m["asset_id"].(string) + m["account_id"].(string)
		amountNumber := m["amount"].(json.Number)
		amount, _ := amountNumber.Int64()

		if tmpM, ok := actionMap[actionKey]; ok {
			tmpNumber, _ := tmpM["amount"].(json.Number)
			tmpAmount, _ := tmpNumber.Int64()
			tmpM["amount"] = json.Number(fmt.Sprintf("%v", tmpAmount+amount))
		} else {
			actionMap[actionKey] = m
			actions = append(actions, m)
		}
	}

	return actions
}

func onlyHaveSpendActions(req *BuildRequest) bool {
	count := 0
	for _, m := range req.Actions {
		if actionType := m["type"].(string); strings.HasPrefix(actionType, "spend") {
			count++
		}
	}

	return count == len(req.Actions)
}

func (a *API) buildSingle(ctx context.Context, req *BuildRequest) (*txbuilder.Template, error) {
	err := a.filterAliases(ctx, req)
	if err != nil {
		return nil, err
	}

	if onlyHaveSpendActions(req) {
		return nil, errors.New("transaction only contain spend actions, didn't have output actions")
	}

	reqActions := mergeActions(req)
	actions := make([]txbuilder.Action, 0, len(reqActions))
	for i, act := range reqActions {
		typ, ok := act["type"].(string)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "no action type provided on action %d", i)
		}
		decoder, ok := a.actionDecoder(typ)
		if !ok {
			return nil, errors.WithDetailf(errBadActionType, "unknown action type %q on action %d", typ, i)
		}

		// Remarshal to JSON, the action may have been modified when we
		// filtered aliases.
		b, err := json.Marshal(act)
		if err != nil {
			return nil, err
		}
		action, err := decoder(b)
		if err != nil {
			return nil, errors.WithDetailf(errBadAction, "%s on action %d", err.Error(), i)
		}
		actions = append(actions, action)
	}

	ttl := req.TTL.Duration
	if ttl == 0 {
		ttl = defaultTxTTL
	}
	maxTime := time.Now().Add(ttl)

	tpl, err := txbuilder.Build(ctx, req.Tx, actions, maxTime, req.TimeRange)
	if errors.Root(err) == txbuilder.ErrAction {
		// append each of the inner errors contained in the data.
		var Errs string
		for _, innerErr := range errors.Data(err)["actions"].([]error) {
			Errs = Errs + "<" + innerErr.Error() + ">"
		}
		err = errors.New(err.Error() + "-" + Errs)
	}
	if err != nil {
		return nil, err
	}

	// ensure null is never returned for signing instructions
	if tpl.SigningInstructions == nil {
		tpl.SigningInstructions = []*txbuilder.SigningInstruction{}
	}
	return tpl, nil
}

// POST /build-transaction
func (a *API) build(ctx context.Context, buildReqs *BuildRequest) Response {
	subctx := reqid.NewSubContext(ctx, reqid.New())

	tmpl, err := a.buildSingle(subctx, buildReqs)
	if err != nil {
		return NewErrorResponse(err)
	}

	return NewSuccessResponse(tmpl)
}

func (a *API) submitSingle(ctx context.Context, tpl *txbuilder.Template) (map[string]string, error) {
	if tpl.Transaction == nil {
		return nil, errors.Wrap(txbuilder.ErrMissingRawTx)
	}

	if err := txbuilder.FinalizeTx(ctx, a.chain, tpl.Transaction); err != nil {
		return nil, errors.Wrapf(err, "tx %s", tpl.Transaction.ID.String())
	}

	return map[string]string{"tx_id": tpl.Transaction.ID.String()}, nil
}

type submitTxResp struct {
	TxID *bc.Hash `json:"tx_id"`
}

// POST /submit-transaction
func (a *API) submit(ctx context.Context, ins struct {
	Tx types.Tx `json:"raw_transaction"`
}) Response {
	if err := txbuilder.FinalizeTx(ctx, a.chain, &ins.Tx); err != nil {
		return NewErrorResponse(err)
	}

	log.WithField("tx_id", ins.Tx.ID).Info("submit single tx")
	return NewSuccessResponse(&submitTxResp{TxID: &ins.Tx.ID})
}

// POST /sign-submit-transaction
func (a *API) signSubmit(ctx context.Context, x struct {
	Password string             `json:"password"`
	Txs      txbuilder.Template `json:"transaction"`
}) Response {
	if err := txbuilder.Sign(ctx, &x.Txs, nil, x.Password, a.pseudohsmSignTemplate); err != nil {
		log.WithField("build err", err).Error("fail on sign transaction.")
		return NewErrorResponse(err)
	}

	if signCount, complete := txbuilder.SignInfo(&x.Txs); !complete && signCount == 0 {
		return NewErrorResponse(pseudohsm.ErrLoadKey)
	}
	log.Info("Sign Transaction complete.")

	txID, err := a.submitSingle(nil, &x.Txs)
	if err != nil {
		log.WithField("err", err).Error("submit single tx")
		return NewErrorResponse(err)
	}

	log.WithField("tx_id", txID["tx_id"]).Info("submit single tx")
	return NewSuccessResponse(txID)
}

type EstimateTxGasResp struct {
	LeftBTM     int64 `json:"left_btm"`
	ConsumedBTM int64 `json:"consumed_btm"`
	LeftGas     int64 `json:"left_gas"`
	ConsumedGas int64 `json:"consumed_gas"`
	StorageGas  int64 `json:"storage_gas"`
	VMGas       int64 `json:"vm_gas"`
}

// POST /estimate-transaction-gas
func (a *API) estimateTxGas(ctx context.Context, ins struct {
	Tx types.Tx `json:"raw_transaction"`
}) Response {
	gasState, err := txbuilder.EstimateTxGas(a.chain, &ins.Tx)
	if err != nil {
		return NewErrorResponse(err)
	}

	btmLeft, ok := checked.MulInt64(gasState.GasLeft, consensus.VMGasRate)
	if !ok {
		return NewErrorResponse(errors.New("calculate btmleft got a math error"))
	}

	txGasResp := &EstimateTxGasResp{
		LeftBTM:     btmLeft,
		ConsumedBTM: int64(gasState.BTMValue) - btmLeft,
		LeftGas:     gasState.GasLeft,
		ConsumedGas: gasState.GasUsed,
		StorageGas:  gasState.StorageGas,
		VMGas:       gasState.GasUsed - gasState.StorageGas,
	}

	return NewSuccessResponse(txGasResp)
}
<|endoftext|>


<commit_before>package apiv1

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"encoding/base64"
	"bytes"
)

type SMIS struct {
	host string
	port string
	insecure bool
	username string
	password string
	httpClient *http.Client
}

func New(host string, port string, insecure bool, username string, password string) (*SMIS, error) {
	if host == "" || port == "" || username == "" || password == "" {
		return nill errors New("Missing host (SMIS Host IP), port (SMIS Host Port), username, or password \n Check Environment Variables..")
	}

	var client *http.Client
	if insecure {
		client = &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					InsecureSkipVerify: insecure,
				},
			},
		}
	}
	else {
		client= &http.Client{}
	}

	return &SMIS{host, port, insecure, username, password, client}, nil
}

func (smis *SMIS) query(httpType, objectPath string, body, resp interface{}) error {

	///////////////////////////////////////////
	//      Setup http/JSON header request   //
	///////////////////////////////////////////

	if smis.insecure {
		URL := strings.Join([]string{"http://", smis.host, ":", smis.port, objectPath}, "")
	}
	else {
		URL := strings.Join([]string{"https://", smis.host, ":", smis.port, objectPath}, "")
	}

	// Parse out body interface into JSON
	bodyBytes, _ := json.Marshal(body)

	// Create http resquest & add auth
	req, err := http.NewRequest(httpType, URL, bytes.NewBuffer(bodyBytes))
	req.SetBasicAuth(smis.username, smis.password)

	// Create Authorization token which is then added to header of request. Auth token is in the format of username:password then encoded in base64.
	encodedAuth := base64.StdEncoding.EncodeToString([]byte(smis.username + ":" + smis.password))

	// Add header specific items
	req.Header.Add("Authorization", "Basic" + encodedAuth)
	req.Header.Add("Accept","application/json")
	req.Header.Add("Content-Type","application/json")
	req.Header.Add("Connection","Keep-Alive")
	req.Header.Add("Host",smis.host + ":" + smis.port)

	// Perform request
	httpResp, err := smis.client.Do(req)
	if err != nill {
		return err
	}

	// Cleanup Response
	resp.Body.Close()

	// Deal with errors
	switch {
	case httpResp.StatusCode == 400:
		return errors.New("JSON Build Error")
	case httpResp.StatusCode == 401:
		return errors.New("JSON Auth Error")
	case httpResp.StatusCode == 404:
		return errors.New("Object Could not be found")
	case httpResp == nil:
		return nil
	// Decode JSON of response into our interface defined for the specific request sent
	case httpResp.StatusCode == 200 || httpResp.StatusCode == 201:
		err = json.NewDecoder(httpResp.Body).Decode(resp)
	default:
		return errors.New("JSON Build Error")
	}
}
<commit_msg>Changed SMIS query to handle interfaces<commit_after>package apiv1

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"encoding/base64"
	"bytes"
)

type SMIS struct {
	host string
	port string
	insecure bool
	username string
	password string
	httpClient *http.Client
}

func New(host string, port string, insecure bool, username string, password string) (*SMIS, error) {
	if host == "" || port == "" || username == "" || password == "" {
		return nill errors New("Missing host (SMIS Host IP), port (SMIS Host Port), username, or password \n Check Environment Variables..")
	}

	var client *http.Client
	if insecure {
		client = &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					InsecureSkipVerify: insecure,
				},
			},
		}
	}
	else {
		client= &http.Client{}
	}

	return &SMIS{host, port, insecure, username, password, client}, nil
}

func (smis *SMIS) query(httpType, objectPath string, body, resp interface{}) error {

	///////////////////////////////////////////
	//      Setup http/JSON header request   //
	///////////////////////////////////////////

	if smis.insecure {
		URL := strings.Join([]string{"http://", smis.host, ":", smis.port, objectPath}, "")
	}
	else {
		URL := strings.Join([]string{"https://", smis.host, ":", smis.port, objectPath}, "")
	}

	// Parse out body interface into JSON
	bodyBytes, _ := json.Marshal(body)

	// Create http resquest & add auth
	req, err := http.NewRequest(httpType, URL, bytes.NewBuffer(bodyBytes))
	req.SetBasicAuth(smis.username, smis.password)

	// Create Authorization token which is then added to header of request. Auth token is in the format of username:password then encoded in base64.
	encodedAuth := base64.StdEncoding.EncodeToString([]byte(smis.username + ":" + smis.password))

	// Add header specific items
	req.Header.Add("Authorization", "Basic" + encodedAuth)
	req.Header.Add("Accept","application/json")
	req.Header.Add("Content-Type","application/json")
	req.Header.Add("Connection","Keep-Alive")
	req.Header.Add("Host",smis.host + ":" + smis.port)

	// Perform request
	httpResp, err := smis.client.Do(req)
	if err != nill {
		return err
	}

	// Cleanup Response
	resp.Body.Close()

	// Deal with errors
	switch {
	case httpResp.StatusCode == 400:
		return errors.New("JSON Build Error")
	case httpResp.StatusCode == 401:
		return errors.New("JSON Auth Error")
	case httpResp.StatusCode == 404:
		return errors.New("Object Could not be found")
	case httpResp == nil:
		return nil
	// Decode JSON of response into our interface defined for the specific request sent
	case httpResp.StatusCode == 200 || httpResp.StatusCode == 201:
		err = json.NewDecoder(httpResp.Body).Decode(resp)
		return err
	default:
		return errors.New("JSON Build Error")
	}
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 harness implements the SDK side of the Beam FnAPI.
package harness

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/exec"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/hooks"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	fnpb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

// TODO(herohde) 2/8/2017: for now, assume we stage a full binary (not a plugin).

// Main is the main entrypoint for the Go harness. It runs at "runtime" -- not
// "pipeline-construction time" -- on each worker. It is a FnAPI client and
// ultimately responsible for correctly executing user code.
func Main(ctx context.Context, loggingEndpoint, controlEndpoint string) error {
	hooks.DeserializeHooksFromOptions(ctx)

	hooks.RunInitHooks(ctx)
	setupRemoteLogging(ctx, loggingEndpoint)
	recordHeader()

	// Connect to FnAPI control server. Receive and execute work.
	// TODO: setup data manager, DoFn register

	conn, err := dial(ctx, controlEndpoint, 60*time.Second)
	if err != nil {
		return errors.Wrap(err, "failed to connect")
	}
	defer conn.Close()

	client, err := fnpb.NewBeamFnControlClient(conn).Control(ctx)
	if err != nil {
		return errors.Wrapf(err, "failed to connect to control service")
	}

	log.Debugf(ctx, "Successfully connected to control @ %v", controlEndpoint)

	// Each ProcessBundle is a sub-graph of the original one.

	var wg sync.WaitGroup
	respc := make(chan *fnpb.InstructionResponse, 100)

	wg.Add(1)

	// gRPC requires all writers to a stream be the same goroutine, so this is the
	// goroutine for managing responses back to the control service.
	go func() {
		defer wg.Done()
		for resp := range respc {
			log.Debugf(ctx, "RESP: %v", proto.MarshalTextString(resp))

			if err := client.Send(resp); err != nil {
				log.Errorf(ctx, "control.Send: Failed to respond: %v", err)
			}
		}
	}()

	ctrl := &control{
		plans:  make(map[string]*exec.Plan),
		active: make(map[string]*exec.Plan),
		data:   &DataChannelManager{},
		state:  &StateChannelManager{},
	}

	// gRPC requires all readers of a stream be the same goroutine, so this goroutine
	// is responsible for managing the network data. All it does is pull data from
	// the stream, and hand off the message to a goroutine to actually be handled,
	// so as to avoid blocking the underlying network channel.
	for {
		req, err := client.Recv()
		if err != nil {
			close(respc)
			wg.Wait()

			if err == io.EOF {
				recordFooter()
				return nil
			}
			return errors.Wrapf(err, "control.Recv failed")
		}

		// Launch a goroutine to handle the control message.
		// TODO(wcn): implement a rate limiter for 'heavy' messages?
		fn := func(ctx context.Context, req *fnpb.InstructionRequest) {
			log.Debugf(ctx, "RECV: %v", proto.MarshalTextString(req))
			recordInstructionRequest(req)

			ctx = hooks.RunRequestHooks(ctx, req)
			resp := ctrl.handleInstruction(ctx, req)

			hooks.RunResponseHooks(ctx, req, resp)

			recordInstructionResponse(resp)
			if resp != nil {
				respc <- resp
			}
		}

		if req.GetProcessBundle() != nil {
			// Only process bundles in a goroutine. We at least need to process instructions for
			// each plan serially. Perhaps just invoke plan.Execute async?
			go fn(ctx, req)
		} else {
			fn(ctx, req)
		}
	}
}

type control struct {
	// plans that are candidates for execution.
	plans map[string]*exec.Plan // protected by mu
	// plans that are actively being executed.
	// a plan can only be in one of these maps at any time.
	active map[string]*exec.Plan // protected by mu
	mu     sync.Mutex

	data  *DataChannelManager
	state *StateChannelManager
}

func (c *control) handleInstruction(ctx context.Context, req *fnpb.InstructionRequest) *fnpb.InstructionResponse {
	id := req.GetInstructionId()
	ctx = setInstID(ctx, id)

	switch {
	case req.GetRegister() != nil:
		msg := req.GetRegister()

		for _, desc := range msg.GetProcessBundleDescriptor() {
			p, err := exec.UnmarshalPlan(desc)
			if err != nil {
				return fail(id, "Invalid bundle desc: %v", err)
			}

			pid := desc.GetId()
			log.Debugf(ctx, "Plan %v: %v", pid, p)

			c.mu.Lock()
			c.plans[pid] = p
			c.mu.Unlock()
		}

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_Register{
				Register: &fnpb.RegisterResponse{},
			},
		}

	case req.GetProcessBundle() != nil:
		msg := req.GetProcessBundle()

		// NOTE: the harness sends a 0-length process bundle request to sources (changed?)

		log.Debugf(ctx, "PB: %v", msg)

		ref := msg.GetProcessBundleDescriptorReference()
		c.mu.Lock()
		plan, ok := c.plans[ref]
		// Make the plan active, and remove it from candidates
		// since a plan can't be run concurrently.
		c.active[id] = plan
		delete(c.plans, ref)
		c.mu.Unlock()

		if !ok {
			return fail(id, "execution plan for %v not found", ref)
		}

		data := NewScopedDataManager(c.data, id)
		state := NewScopedStateReader(c.state, id)
		err := plan.Execute(ctx, id, exec.DataContext{Data: data, State: state})
		data.Close()
		state.Close()

		m := plan.Metrics()
		// Move the plan back to the candidate state
		c.mu.Lock()
		c.plans[plan.ID()] = plan
		delete(c.active, id)
		c.mu.Unlock()

		if err != nil {
			return fail(id, "execute failed: %v", err)
		}

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_ProcessBundle{
				ProcessBundle: &fnpb.ProcessBundleResponse{
					Metrics: m,
				},
			},
		}

	case req.GetProcessBundleProgress() != nil:
		msg := req.GetProcessBundleProgress()

		// log.Debugf(ctx, "PB Progress: %v", msg)

		ref := msg.GetInstructionReference()
		c.mu.Lock()
		plan, ok := c.active[ref]
		c.mu.Unlock()
		if !ok {
			return fail(id, "execution plan for %v not found", ref)
		}

		m := plan.Metrics()

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_ProcessBundleProgress{
				ProcessBundleProgress: &fnpb.ProcessBundleProgressResponse{
					Metrics: m,
				},
			},
		}

	case req.GetProcessBundleSplit() != nil:
		msg := req.GetProcessBundleSplit()

		log.Debugf(ctx, "PB Split: %v", msg)
		ref := msg.GetInstructionReference()
		c.mu.Lock()
		plan, ok := c.active[ref]
		c.mu.Unlock()
		if !ok {
			return fail(id, "execution plan for %v not found", ref)
		}

		// Get the desired splits for the root FnAPI read operation.
		ds := msg.GetDesiredSplits()[plan.SourcePTransformID()]
		if ds == nil {
			return fail(id, "failed to split: desired splits for root was empty.")
		}
		split, err := plan.Split(exec.SplitPoints{ds.GetAllowedSplitPoints(), ds.GetFractionOfRemainder()})

		if err != nil {
			return fail(id, "unable to split: %v", err)
		}

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_ProcessBundleSplit{
				ProcessBundleSplit: &fnpb.ProcessBundleSplitResponse{
					ChannelSplits: []*fnpb.ProcessBundleSplitResponse_ChannelSplit{
						&fnpb.ProcessBundleSplitResponse_ChannelSplit{
							LastPrimaryElement:   int32(split),
							FirstResidualElement: int32(split + 1),
						},
					},
				},
			},
		}

	default:
		return fail(id, "Unexpected request: %v", req)
	}
}

func fail(id, format string, args ...interface{}) *fnpb.InstructionResponse {
	dummy := &fnpb.InstructionResponse_Register{Register: &fnpb.RegisterResponse{}}

	return &fnpb.InstructionResponse{
		InstructionId: id,
		Error:         fmt.Sprintf(format, args...),
		Response:      dummy,
	}
}

// dial to the specified endpoint. if timeout <=0, call blocks until
// grpc.Dial succeeds.
func dial(ctx context.Context, endpoint string, timeout time.Duration) (*grpc.ClientConn, error) {
	log.Infof(ctx, "Connecting via grpc @ %s ...", endpoint)
	return grpcx.Dial(ctx, endpoint, timeout)
}
<commit_msg>Revert "[Go SDK] Off By One error on splitting"<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 harness implements the SDK side of the Beam FnAPI.
package harness

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/exec"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/hooks"
	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	fnpb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

// TODO(herohde) 2/8/2017: for now, assume we stage a full binary (not a plugin).

// Main is the main entrypoint for the Go harness. It runs at "runtime" -- not
// "pipeline-construction time" -- on each worker. It is a FnAPI client and
// ultimately responsible for correctly executing user code.
func Main(ctx context.Context, loggingEndpoint, controlEndpoint string) error {
	hooks.DeserializeHooksFromOptions(ctx)

	hooks.RunInitHooks(ctx)
	setupRemoteLogging(ctx, loggingEndpoint)
	recordHeader()

	// Connect to FnAPI control server. Receive and execute work.
	// TODO: setup data manager, DoFn register

	conn, err := dial(ctx, controlEndpoint, 60*time.Second)
	if err != nil {
		return errors.Wrap(err, "failed to connect")
	}
	defer conn.Close()

	client, err := fnpb.NewBeamFnControlClient(conn).Control(ctx)
	if err != nil {
		return errors.Wrapf(err, "failed to connect to control service")
	}

	log.Debugf(ctx, "Successfully connected to control @ %v", controlEndpoint)

	// Each ProcessBundle is a sub-graph of the original one.

	var wg sync.WaitGroup
	respc := make(chan *fnpb.InstructionResponse, 100)

	wg.Add(1)

	// gRPC requires all writers to a stream be the same goroutine, so this is the
	// goroutine for managing responses back to the control service.
	go func() {
		defer wg.Done()
		for resp := range respc {
			log.Debugf(ctx, "RESP: %v", proto.MarshalTextString(resp))

			if err := client.Send(resp); err != nil {
				log.Errorf(ctx, "control.Send: Failed to respond: %v", err)
			}
		}
	}()

	ctrl := &control{
		plans:  make(map[string]*exec.Plan),
		active: make(map[string]*exec.Plan),
		data:   &DataChannelManager{},
		state:  &StateChannelManager{},
	}

	// gRPC requires all readers of a stream be the same goroutine, so this goroutine
	// is responsible for managing the network data. All it does is pull data from
	// the stream, and hand off the message to a goroutine to actually be handled,
	// so as to avoid blocking the underlying network channel.
	for {
		req, err := client.Recv()
		if err != nil {
			close(respc)
			wg.Wait()

			if err == io.EOF {
				recordFooter()
				return nil
			}
			return errors.Wrapf(err, "control.Recv failed")
		}

		// Launch a goroutine to handle the control message.
		// TODO(wcn): implement a rate limiter for 'heavy' messages?
		fn := func(ctx context.Context, req *fnpb.InstructionRequest) {
			log.Debugf(ctx, "RECV: %v", proto.MarshalTextString(req))
			recordInstructionRequest(req)

			ctx = hooks.RunRequestHooks(ctx, req)
			resp := ctrl.handleInstruction(ctx, req)

			hooks.RunResponseHooks(ctx, req, resp)

			recordInstructionResponse(resp)
			if resp != nil {
				respc <- resp
			}
		}

		if req.GetProcessBundle() != nil {
			// Only process bundles in a goroutine. We at least need to process instructions for
			// each plan serially. Perhaps just invoke plan.Execute async?
			go fn(ctx, req)
		} else {
			fn(ctx, req)
		}
	}
}

type control struct {
	// plans that are candidates for execution.
	plans map[string]*exec.Plan // protected by mu
	// plans that are actively being executed.
	// a plan can only be in one of these maps at any time.
	active map[string]*exec.Plan // protected by mu
	mu     sync.Mutex

	data  *DataChannelManager
	state *StateChannelManager
}

func (c *control) handleInstruction(ctx context.Context, req *fnpb.InstructionRequest) *fnpb.InstructionResponse {
	id := req.GetInstructionId()
	ctx = setInstID(ctx, id)

	switch {
	case req.GetRegister() != nil:
		msg := req.GetRegister()

		for _, desc := range msg.GetProcessBundleDescriptor() {
			p, err := exec.UnmarshalPlan(desc)
			if err != nil {
				return fail(id, "Invalid bundle desc: %v", err)
			}

			pid := desc.GetId()
			log.Debugf(ctx, "Plan %v: %v", pid, p)

			c.mu.Lock()
			c.plans[pid] = p
			c.mu.Unlock()
		}

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_Register{
				Register: &fnpb.RegisterResponse{},
			},
		}

	case req.GetProcessBundle() != nil:
		msg := req.GetProcessBundle()

		// NOTE: the harness sends a 0-length process bundle request to sources (changed?)

		log.Debugf(ctx, "PB: %v", msg)

		ref := msg.GetProcessBundleDescriptorReference()
		c.mu.Lock()
		plan, ok := c.plans[ref]
		// Make the plan active, and remove it from candidates
		// since a plan can't be run concurrently.
		c.active[id] = plan
		delete(c.plans, ref)
		c.mu.Unlock()

		if !ok {
			return fail(id, "execution plan for %v not found", ref)
		}

		data := NewScopedDataManager(c.data, id)
		state := NewScopedStateReader(c.state, id)
		err := plan.Execute(ctx, id, exec.DataContext{Data: data, State: state})
		data.Close()
		state.Close()

		m := plan.Metrics()
		// Move the plan back to the candidate state
		c.mu.Lock()
		c.plans[plan.ID()] = plan
		delete(c.active, id)
		c.mu.Unlock()

		if err != nil {
			return fail(id, "execute failed: %v", err)
		}

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_ProcessBundle{
				ProcessBundle: &fnpb.ProcessBundleResponse{
					Metrics: m,
				},
			},
		}

	case req.GetProcessBundleProgress() != nil:
		msg := req.GetProcessBundleProgress()

		// log.Debugf(ctx, "PB Progress: %v", msg)

		ref := msg.GetInstructionReference()
		c.mu.Lock()
		plan, ok := c.active[ref]
		c.mu.Unlock()
		if !ok {
			return fail(id, "execution plan for %v not found", ref)
		}

		m := plan.Metrics()

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_ProcessBundleProgress{
				ProcessBundleProgress: &fnpb.ProcessBundleProgressResponse{
					Metrics: m,
				},
			},
		}

	case req.GetProcessBundleSplit() != nil:
		msg := req.GetProcessBundleSplit()

		log.Debugf(ctx, "PB Split: %v", msg)
		ref := msg.GetInstructionReference()
		c.mu.Lock()
		plan, ok := c.active[ref]
		c.mu.Unlock()
		if !ok {
			return fail(id, "execution plan for %v not found", ref)
		}

		// Get the desired splits for the root FnAPI read operation.
		ds := msg.GetDesiredSplits()[plan.SourcePTransformID()]
		if ds == nil {
			return fail(id, "failed to split: desired splits for root was empty.")
		}
		split, err := plan.Split(exec.SplitPoints{ds.GetAllowedSplitPoints(), ds.GetFractionOfRemainder()})

		if err != nil {
			return fail(id, "unable to split: %v", err)
		}

		return &fnpb.InstructionResponse{
			InstructionId: id,
			Response: &fnpb.InstructionResponse_ProcessBundleSplit{
				ProcessBundleSplit: &fnpb.ProcessBundleSplitResponse{
					ChannelSplits: []*fnpb.ProcessBundleSplitResponse_ChannelSplit{
						&fnpb.ProcessBundleSplitResponse_ChannelSplit{
							LastPrimaryElement:   int32(split - 1),
							FirstResidualElement: int32(split),
						},
					},
				},
			},
		}

	default:
		return fail(id, "Unexpected request: %v", req)
	}
}

func fail(id, format string, args ...interface{}) *fnpb.InstructionResponse {
	dummy := &fnpb.InstructionResponse_Register{Register: &fnpb.RegisterResponse{}}

	return &fnpb.InstructionResponse{
		InstructionId: id,
		Error:         fmt.Sprintf(format, args...),
		Response:      dummy,
	}
}

// dial to the specified endpoint. if timeout <=0, call blocks until
// grpc.Dial succeeds.
func dial(ctx context.Context, endpoint string, timeout time.Duration) (*grpc.ClientConn, error) {
	log.Infof(ctx, "Connecting via grpc @ %s ...", endpoint)
	return grpcx.Dial(ctx, endpoint, timeout)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/samurang87/playground_go/brackets/balanced_brackets"
)

func main() {

	var a int
	var b string

	fmt.Scanf("%d\n", &a)

	for i := 0; i < a; i++ {

		fmt.Scanf("%s\n", &b)

		response := balanced_brackets.Balance(b)

		if response == true {

			fmt.Println("YES")

		} else {

			fmt.Println("NO")
		}

	}

}
<commit_msg>balanced-brackets added reader from stdin<commit_after>package main

import (
	"fmt"
	"log"
	"github.com/samurang87/playground_go/brackets/balanced_brackets"
)

func Read() ([]string, error) {
	var length int

	_, err := fmt.Scanf("%d", &length)
	if err != nil {
		return nil, err
	}

	list := make([]string, length)

	for i := 0; i < len(list); i++ {
		_, err := fmt.Scanf("%v", &list[i])
		if err != nil {
			return nil, err
		}
	}

	return list, nil
}

func main() {

	list, err := Read()

	if err != nil {
		log.Fatal(err)
	}

	for _, item := range list {

		response := balanced_brackets.Balance(item)

		if response == true {
			fmt.Println("YES")

		} else {

			fmt.Println("NO")
		}

	}

}
<|endoftext|>


<commit_before>// Copyright 2015 Google 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 lease

import (
	"fmt"
	"io"
	"log"
	"sync"
)

// Create a ReadLease that never expires, unless voluntarily revoked or
// upgraded.
//
// The supplied function will be used to obtain the read lease contents, the
// first time and whenever the supplied file leaser decides to expire the
// temporary copy thus obtained. It must return the same contents every time,
// and the contents must be of the given size.
//
// This magic is not preserved after the lease is upgraded.
func NewAutoRefreshingReadLease(
	fl FileLeaser,
	size int64,
	f func() (io.ReadCloser, error)) (rl ReadLease) {
	rl = &autoRefreshingReadLease{
		leaser: fl,
		size:   size,
		f:      f,
	}

	return
}

type autoRefreshingReadLease struct {
	mu sync.Mutex

	/////////////////////////
	// Constant data
	/////////////////////////

	size int64

	/////////////////////////
	// Dependencies
	/////////////////////////

	leaser FileLeaser
	f      func() (io.ReadCloser, error)

	/////////////////////////
	// Mutable state
	/////////////////////////

	// Set to true when we've been revoked for good.
	//
	// GUARDED_BY(mu)
	revoked bool

	// The current wrapped lease, or nil if one has never been issued.
	//
	// GUARDED_BY(mu)
	wrapped ReadLease
}

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

// Attempt to clean up after the supplied read/write lease.
func destroyReadWriteLease(rwl ReadWriteLease) {
	var err error
	defer func() {
		if err != nil {
			log.Printf("Error destroying read/write lease: %v", err)
		}
	}()

	// Downgrade to a read lease.
	rl, err := rwl.Downgrade()
	if err != nil {
		err = fmt.Errorf("Downgrade: %v", err)
		return
	}

	// Revoke the read lease.
	rl.Revoke()
}

func isRevokedErr(err error) bool {
	_, ok := err.(*RevokedError)
	return ok
}

// Set up a read/write lease and fill in our contents.
//
// REQUIRES: The caller has observed that rl.lease has expired.
//
// LOCKS_REQUIRED(rl.mu)
func (rl *autoRefreshingReadLease) getContents() (
	rwl ReadWriteLease, err error) {
	// Obtain some space to write the contents.
	rwl, err = rl.leaser.NewFile()
	if err != nil {
		err = fmt.Errorf("NewFile: %v", err)
		return
	}

	// Attempt to clean up if we exit early.
	defer func() {
		if err != nil {
			destroyReadWriteLease(rwl)
		}
	}()

	// Obtain the reader for our contents.
	rc, err := rl.f()
	if err != nil {
		err = fmt.Errorf("User function: %v", err)
		return
	}

	defer func() {
		closeErr := rc.Close()
		if closeErr != nil && err == nil {
			err = fmt.Errorf("Close: %v", closeErr)
		}
	}()

	// Copy into the read/write lease.
	copied, err := io.Copy(rwl, rc)
	if err != nil {
		err = fmt.Errorf("Copy: %v", err)
		return
	}

	// Did the user lie about the size?
	if copied != rl.Size() {
		err = fmt.Errorf("Copied %v bytes; expected %v", copied, rl.Size())
		return
	}

	return
}

// Downgrade and save the supplied read/write lease obtained with getContents
// for later use.
//
// LOCKS_REQUIRED(rl.mu)
func (rl *autoRefreshingReadLease) saveContents(rwl ReadWriteLease) {
	downgraded, err := rwl.Downgrade()
	if err != nil {
		log.Printf("Failed to downgrade write lease (%q); abandoning.", err.Error())
		return
	}

	rl.wrapped = downgraded
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rl *autoRefreshingReadLease) Read(p []byte) (n int, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		n, err = rl.wrapped.Read(p)
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Get hold of a read/write lease containing our contents.
	rwl, err := rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	defer rl.saveContents(rwl)

	// Serve from the read/write lease.
	n, err = rwl.Read(p)

	return
}

func (rl *autoRefreshingReadLease) Seek(
	offset int64,
	whence int) (off int64, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		off, err = rl.wrapped.Seek(offset, whence)
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Get hold of a read/write lease containing our contents.
	rwl, err := rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	defer rl.saveContents(rwl)

	// Serve from the read/write lease.
	off, err = rwl.Seek(offset, whence)

	return
}

func (rl *autoRefreshingReadLease) ReadAt(
	p []byte,
	off int64) (n int, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		n, err = rl.wrapped.ReadAt(p, off)
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Get hold of a read/write lease containing our contents.
	rwl, err := rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	defer rl.saveContents(rwl)

	// Serve from the read/write lease.
	n, err = rwl.ReadAt(p, off)

	return
}

func (rl *autoRefreshingReadLease) Size() (size int64) {
	size = rl.size
	return
}

func (rl *autoRefreshingReadLease) Revoked() (revoked bool) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	revoked = rl.revoked
	return
}

func (rl *autoRefreshingReadLease) Upgrade() (rwl ReadWriteLease, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// If we succeed, we are now revoked.
	defer func() {
		if err == nil {
			rl.revoked = true
		}
	}()

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		rwl, err = rl.wrapped.Upgrade()
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Build the read/write lease anew.
	rwl, err = rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	return
}

func (rl *autoRefreshingReadLease) Revoke() {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	rl.revoked = true
	if rl.wrapped != nil {
		rl.wrapped.Revoke()
	}
}
<commit_msg>Pass a context to the refresh function.<commit_after>// Copyright 2015 Google 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 lease

import (
	"fmt"
	"io"
	"log"
	"sync"

	"golang.org/x/net/context"
)

// A function used by leases returned from NewAutoRefreshingReadLease to
// refresh the contents of the read lease. See the notes on that function.
type RefreshContentsFunc func(context.Context) (io.ReadCloser, error)

// Create a ReadLease that never expires, unless voluntarily revoked or
// upgraded.
//
// The supplied function will be used to obtain the read lease contents, the
// first time and whenever the supplied file leaser decides to expire the
// temporary copy thus obtained. It must return the same contents every time,
// and the contents must be of the given size.
//
// This magic is not preserved after the lease is upgraded.
func NewAutoRefreshingReadLease(
	fl FileLeaser,
	size int64,
	refresh RefreshContentsFunc) (rl ReadLease) {
	rl = &autoRefreshingReadLease{
		leaser:  fl,
		size:    size,
		refresh: refresh,
	}

	return
}

type autoRefreshingReadLease struct {
	mu sync.Mutex

	/////////////////////////
	// Constant data
	/////////////////////////

	size int64

	/////////////////////////
	// Dependencies
	/////////////////////////

	leaser FileLeaser
	f      func() (io.ReadCloser, error)

	/////////////////////////
	// Mutable state
	/////////////////////////

	// Set to true when we've been revoked for good.
	//
	// GUARDED_BY(mu)
	revoked bool

	// The current wrapped lease, or nil if one has never been issued.
	//
	// GUARDED_BY(mu)
	wrapped ReadLease
}

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

// Attempt to clean up after the supplied read/write lease.
func destroyReadWriteLease(rwl ReadWriteLease) {
	var err error
	defer func() {
		if err != nil {
			log.Printf("Error destroying read/write lease: %v", err)
		}
	}()

	// Downgrade to a read lease.
	rl, err := rwl.Downgrade()
	if err != nil {
		err = fmt.Errorf("Downgrade: %v", err)
		return
	}

	// Revoke the read lease.
	rl.Revoke()
}

func isRevokedErr(err error) bool {
	_, ok := err.(*RevokedError)
	return ok
}

// Set up a read/write lease and fill in our contents.
//
// REQUIRES: The caller has observed that rl.lease has expired.
//
// LOCKS_REQUIRED(rl.mu)
func (rl *autoRefreshingReadLease) getContents() (
	rwl ReadWriteLease, err error) {
	// Obtain some space to write the contents.
	rwl, err = rl.leaser.NewFile()
	if err != nil {
		err = fmt.Errorf("NewFile: %v", err)
		return
	}

	// Attempt to clean up if we exit early.
	defer func() {
		if err != nil {
			destroyReadWriteLease(rwl)
		}
	}()

	// Obtain the reader for our contents.
	rc, err := rl.f()
	if err != nil {
		err = fmt.Errorf("User function: %v", err)
		return
	}

	defer func() {
		closeErr := rc.Close()
		if closeErr != nil && err == nil {
			err = fmt.Errorf("Close: %v", closeErr)
		}
	}()

	// Copy into the read/write lease.
	copied, err := io.Copy(rwl, rc)
	if err != nil {
		err = fmt.Errorf("Copy: %v", err)
		return
	}

	// Did the user lie about the size?
	if copied != rl.Size() {
		err = fmt.Errorf("Copied %v bytes; expected %v", copied, rl.Size())
		return
	}

	return
}

// Downgrade and save the supplied read/write lease obtained with getContents
// for later use.
//
// LOCKS_REQUIRED(rl.mu)
func (rl *autoRefreshingReadLease) saveContents(rwl ReadWriteLease) {
	downgraded, err := rwl.Downgrade()
	if err != nil {
		log.Printf("Failed to downgrade write lease (%q); abandoning.", err.Error())
		return
	}

	rl.wrapped = downgraded
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rl *autoRefreshingReadLease) Read(p []byte) (n int, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		n, err = rl.wrapped.Read(p)
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Get hold of a read/write lease containing our contents.
	rwl, err := rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	defer rl.saveContents(rwl)

	// Serve from the read/write lease.
	n, err = rwl.Read(p)

	return
}

func (rl *autoRefreshingReadLease) Seek(
	offset int64,
	whence int) (off int64, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		off, err = rl.wrapped.Seek(offset, whence)
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Get hold of a read/write lease containing our contents.
	rwl, err := rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	defer rl.saveContents(rwl)

	// Serve from the read/write lease.
	off, err = rwl.Seek(offset, whence)

	return
}

func (rl *autoRefreshingReadLease) ReadAt(
	p []byte,
	off int64) (n int, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		n, err = rl.wrapped.ReadAt(p, off)
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Get hold of a read/write lease containing our contents.
	rwl, err := rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	defer rl.saveContents(rwl)

	// Serve from the read/write lease.
	n, err = rwl.ReadAt(p, off)

	return
}

func (rl *autoRefreshingReadLease) Size() (size int64) {
	size = rl.size
	return
}

func (rl *autoRefreshingReadLease) Revoked() (revoked bool) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	revoked = rl.revoked
	return
}

func (rl *autoRefreshingReadLease) Upgrade() (rwl ReadWriteLease, err error) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Special case: have we been permanently revoked?
	if rl.revoked {
		err = &RevokedError{}
		return
	}

	// If we succeed, we are now revoked.
	defer func() {
		if err == nil {
			rl.revoked = true
		}
	}()

	// Common case: is the existing lease still valid?
	if rl.wrapped != nil {
		rwl, err = rl.wrapped.Upgrade()
		if !isRevokedErr(err) {
			return
		}

		// Clear the revoked error.
		err = nil
	}

	// Build the read/write lease anew.
	rwl, err = rl.getContents()
	if err != nil {
		err = fmt.Errorf("getContents: %v", err)
		return
	}

	return
}

func (rl *autoRefreshingReadLease) Revoke() {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	rl.revoked = true
	if rl.wrapped != nil {
		rl.wrapped.Revoke()
	}
}
<|endoftext|>


<commit_before>package service

import (
	"reflect"
	"time"

	log "github.com/Cepave/open-falcon-backend/common/logruslog"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/common/utils"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/model"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/rdb"
)

var logger = log.NewDefaultLogger("INFO")

type mode byte

const (
	_DRAIN mode = 1
	_FLUSH mode = 2
)

var NqmQueue *nqmAgentUpdateService

func InitNqmHeartbeat(c *commonQueue.Config) {
	NqmQueue = newNqmAgentUpdateService(c)
	NqmQueue.updateToDatabase = updateNqmAgentHeartbeatImpl
	NqmQueue.Start()
}

func CloseNqmHeartbeat() {
	logger.Info("Closing NQM heartbeat queue service...")
	NqmQueue.Stop()
	logger.Info("Finish.")
}

var typeOfNqmAgentHeartbeat = reflect.TypeOf(new(model.NqmAgentHeartbeatRequest))
type nqmAgentUpdateService struct {
	q       *commonQueue.Queue
	c       *commonQueue.Config
	cnt     uint64 // counter for the dequeued elements
	running bool
	flush   chan struct{}
	done    chan struct{}
	updateToDatabase func([]*model.NqmAgentHeartbeatRequest)
}

func newNqmAgentUpdateService(c *commonQueue.Config) *nqmAgentUpdateService {
	return &nqmAgentUpdateService{
		q:     commonQueue.New(),
		c:     c,
		done:  make(chan struct{}),
		flush: make(chan struct{}),
	}
}

// Gets the number of consumed updating requests(not guarantee on database)
func (q *nqmAgentUpdateService) ConsumedCount() uint64 {
	return q.cnt
}

// Gets the number of pending request of heartbeats
func (q *nqmAgentUpdateService) PendingLen() int {
	return q.q.Len()
}

func (q *nqmAgentUpdateService) Start() {
	if q.running {
		return
	}
	q.running = true
	go q.draining()
}

func (q *nqmAgentUpdateService) Stop() {
	if !q.running {
		return
	}
	q.running = false

	time.Sleep(q.c.Dur) // for all `q.q.Enqueue()`s to be done

	close(q.flush)
	<-q.done
}

func (q *nqmAgentUpdateService) Put(req *model.NqmAgentHeartbeatRequest) {
	if !q.running {
		return
	}
	q.q.Enqueue(req)
}

func (q *nqmAgentUpdateService) draining() {
	for {
		select {
		default:
			q.syncToDatabase(_DRAIN)
		case <-q.flush:
			q.syncToDatabase(_FLUSH)
			close(q.done)
			return
		}
	}
}

func (q *nqmAgentUpdateService) syncToDatabase(m mode) {
	var config commonQueue.Config = *q.c

	var reqs []*model.NqmAgentHeartbeatRequest

	switch m {
	case _FLUSH:
		config.Dur = 0
		reqs = q.drainFromQueue(&config)
		if len(reqs) > 0 {
			logger.Infof("Flushing [%d] heartbeats of NQM agent from queue", len(reqs))
		}
	default:
		reqs = q.drainFromQueue(&config)
	}

	q.cnt += uint64(len(reqs))
	q.updateToDatabase(reqs)

	if len(reqs) > 0 {
		logger.Debugf("[%d] heartbeats of NQM agent from queue", len(reqs))

		if m == _FLUSH {
			q.syncToDatabase(m)
		}
	}
}

func (q *nqmAgentUpdateService) drainFromQueue(config *commonQueue.Config) []*model.NqmAgentHeartbeatRequest {
	return q.q.DrainNWithDurationByType(
		config, typeOfNqmAgentHeartbeat,
	).([]*model.NqmAgentHeartbeatRequest)
}

func updateNqmAgentHeartbeatImpl(reqs []*model.NqmAgentHeartbeatRequest) {
	utils.BuildPanicCapture(
		func() {
			rdb.UpdateNqmAgentHeartbeat(reqs)
		},
		func(p interface{}) {
			logger.Errorf("[PANIC] Update heartbeats of NQM agent[#%d]: %v", len(reqs), p)
		},
	)()
}
<commit_msg>[OWL-1730][nqm-mng] `cnt` counts after db queries<commit_after>package service

import (
	"reflect"
	"time"

	log "github.com/Cepave/open-falcon-backend/common/logruslog"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/common/utils"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/model"
	"github.com/Cepave/open-falcon-backend/modules/nqm-mng/rdb"
)

var logger = log.NewDefaultLogger("INFO")

type mode byte

const (
	_DRAIN mode = 1
	_FLUSH mode = 2
)

var NqmQueue *nqmAgentUpdateService

func InitNqmHeartbeat(c *commonQueue.Config) {
	NqmQueue = newNqmAgentUpdateService(c)
	NqmQueue.updateToDatabase = updateNqmAgentHeartbeatImpl
	NqmQueue.Start()
}

func CloseNqmHeartbeat() {
	logger.Info("Closing NQM heartbeat queue service...")
	NqmQueue.Stop()
	logger.Info("Finish.")
}

var typeOfNqmAgentHeartbeat = reflect.TypeOf(new(model.NqmAgentHeartbeatRequest))

type nqmAgentUpdateService struct {
	q                *commonQueue.Queue
	c                *commonQueue.Config
	cnt              uint64 // counter for the dequeued elements
	running          bool
	flush            chan struct{}
	done             chan struct{}
	updateToDatabase func([]*model.NqmAgentHeartbeatRequest)
}

func newNqmAgentUpdateService(c *commonQueue.Config) *nqmAgentUpdateService {
	return &nqmAgentUpdateService{
		q:     commonQueue.New(),
		c:     c,
		done:  make(chan struct{}),
		flush: make(chan struct{}),
	}
}

// Gets the number of consumed updating requests(not guarantee on database)
func (q *nqmAgentUpdateService) ConsumedCount() uint64 {
	return q.cnt
}

// Gets the number of pending request of heartbeats
func (q *nqmAgentUpdateService) PendingLen() int {
	return q.q.Len()
}

func (q *nqmAgentUpdateService) Start() {
	if q.running {
		return
	}
	q.running = true
	go q.draining()
}

func (q *nqmAgentUpdateService) Stop() {
	if !q.running {
		return
	}
	q.running = false

	time.Sleep(q.c.Dur) // for all `q.q.Enqueue()`s to be done

	close(q.flush)
	<-q.done
}

func (q *nqmAgentUpdateService) Put(req *model.NqmAgentHeartbeatRequest) {
	if !q.running {
		return
	}
	q.q.Enqueue(req)
}

func (q *nqmAgentUpdateService) draining() {
	for {
		select {
		default:
			q.syncToDatabase(_DRAIN)
		case <-q.flush:
			q.syncToDatabase(_FLUSH)
			close(q.done)
			return
		}
	}
}

func (q *nqmAgentUpdateService) syncToDatabase(m mode) {
	var config commonQueue.Config = *q.c

	var reqs []*model.NqmAgentHeartbeatRequest

	switch m {
	case _FLUSH:
		config.Dur = 0
		reqs = q.drainFromQueue(&config)
		if len(reqs) > 0 {
			logger.Infof("Flushing [%d] heartbeats of NQM agent from queue", len(reqs))
		}
	default:
		reqs = q.drainFromQueue(&config)
	}

	q.updateToDatabase(reqs)
	q.cnt += uint64(len(reqs))

	if len(reqs) > 0 {
		logger.Debugf("[%d] heartbeats of NQM agent from queue", len(reqs))

		if m == _FLUSH {
			q.syncToDatabase(m)
		}
	}
}

func (q *nqmAgentUpdateService) drainFromQueue(config *commonQueue.Config) []*model.NqmAgentHeartbeatRequest {
	return q.q.DrainNWithDurationByType(
		config, typeOfNqmAgentHeartbeat,
	).([]*model.NqmAgentHeartbeatRequest)
}

func updateNqmAgentHeartbeatImpl(reqs []*model.NqmAgentHeartbeatRequest) {
	utils.BuildPanicCapture(
		func() {
			rdb.UpdateNqmAgentHeartbeat(reqs)
		},
		func(p interface{}) {
			logger.Errorf("[PANIC] Update heartbeats of NQM agent[#%d]: %v", len(reqs), p)
		},
	)()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

// TestGofmt tests that all files are formatted.
func TestGofmt(t *testing.T) {
	root, fileMap := copyTree(t)

	gofmt := exec.Command("gofmt", "-w", ".")
	gofmt.Dir = root
	gofmt.Stdout, gofmt.Stderr = os.Stdout, os.Stderr
	if err := gofmt.Run(); err != nil {
		t.Fatalf("gofmt failed: %v", err)
	}

	// Diff the trees.
	if diffFiles(t, fileMap) {
		t.Errorf("Files are not gofmt clean. Please run gofmt.")
	}
}

// TestGenerated tests that all generated files are up-to-date.
func TestGenerated(t *testing.T) {
	root, fileMap := copyTree(t)

	// Build bitstringer
	build := exec.Command("go", "build")
	build.Dir = filepath.Join(root, "cmd/bitstringer")
	build.Stdout, build.Stderr = os.Stdout, os.Stderr
	if err := build.Run(); err != nil {
		t.Fatalf("go build failed: %v", err)
	}

	gen := exec.Command("go", "generate", "./...")
	gen.Dir = root
	gen.Stdout, gen.Stderr = os.Stdout, os.Stderr
	if err := gen.Run(); err != nil {
		t.Fatalf("go generate failed: %v", err)
	}

	// Diff the trees.
	if diffFiles(t, fileMap) {
		t.Errorf("Please run go generate.")
	}
}

func copyTree(t *testing.T) (string, map[string]string) {
	src, err := os.Getwd()
	if err != nil {
		t.Fatalf("getting working directory: %v", err)
	}
	dst := t.TempDir()

	// Ensure src ends with "/"
	src = fmt.Sprintf("%s%c", filepath.Clean(src), filepath.Separator)

	fileMap := make(map[string]string)
	err = filepath.WalkDir(src, func(path string, d fs.DirEntry, err error) error {
		if !strings.HasPrefix(path, src) {
			panic(fmt.Sprintf("WalkDir path %q does not start with source path %q", path, src))
		}

		rel := path[len(src):]
		if d.Name() == ".git" {
			return filepath.SkipDir
		}
		if d.IsDir() {
			if rel == "" {
				// This is the root of the tree, so
				// the destination already exists.
				return nil
			}
			return os.Mkdir(filepath.Join(dst, rel), 0777)
		}
		// Only copy .go and related files.
		if n := d.Name(); !(filepath.Ext(n) == ".go" || n == "go.mod" || n == "go.sum") {
			return nil
		}

		// Copy file.
		fileMap[path] = filepath.Join(dst, rel)
		return copyFile(path, filepath.Join(dst, rel))
	})
	if err != nil {
		t.Fatalf("error copying source tree: %v", err)
	}

	t.Logf("copied source tree to %s", dst)

	return dst, fileMap
}

func copyFile(src, dst string) error {
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	return out.Close()
}

func diffFiles(t *testing.T, fileMap map[string]string) bool {
	diffs := 0
	for orig, new := range fileMap {
		diff := exec.Command("diff", "-u", orig, new)
		diff.Stdout = os.Stdout
		diff.Stderr = os.Stderr
		if err := diff.Run(); err != nil {
			switch err := err.(type) {
			case *exec.ExitError:
				if err.ExitCode() == 1 {
					diffs++
					continue
				}
			}
			t.Errorf("diff failed: %v", err)
		}
	}
	return diffs != 0
}
<commit_msg>all: another shot at fixing Windows testing<commit_after>package main

import (
	"fmt"
	"io"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

// TestGofmt tests that all files are formatted.
func TestGofmt(t *testing.T) {
	root, fileMap := copyTree(t)

	gofmt := exec.Command("gofmt", "-w", ".")
	gofmt.Dir = root
	gofmt.Stdout, gofmt.Stderr = os.Stdout, os.Stderr
	if err := gofmt.Run(); err != nil {
		t.Fatalf("gofmt failed: %v", err)
	}

	// Diff the trees.
	if diffFiles(t, fileMap) {
		t.Errorf("Files are not gofmt clean. Please run gofmt.")
	}
}

// TestGenerated tests that all generated files are up-to-date.
func TestGenerated(t *testing.T) {
	root, fileMap := copyTree(t)

	// Build bitstringer
	build := exec.Command("go", "build")
	build.Dir = filepath.Join(root, "cmd/bitstringer")
	build.Stdout, build.Stderr = os.Stdout, os.Stderr
	if err := build.Run(); err != nil {
		t.Fatalf("go build failed: %v", err)
	}

	gen := exec.Command("go", "generate", "./...")
	gen.Dir = root
	gen.Stdout, gen.Stderr = os.Stdout, os.Stderr
	if err := gen.Run(); err != nil {
		t.Fatalf("go generate failed: %v", err)
	}

	// Diff the trees.
	if diffFiles(t, fileMap) {
		t.Errorf("Please run go generate.")
	}
}

func copyTree(t *testing.T) (string, map[string]string) {
	src, err := os.Getwd()
	if err != nil {
		t.Fatalf("getting working directory: %v", err)
	}
	dst := t.TempDir()

	// Ensure src ends with "/"
	src = fmt.Sprintf("%s%c", filepath.Clean(src), filepath.Separator)

	fileMap := make(map[string]string)
	err = filepath.WalkDir(src, func(path string, d fs.DirEntry, err error) error {
		if !strings.HasPrefix(path, src) {
			panic(fmt.Sprintf("WalkDir path %q does not start with source path %q", path, src))
		}

		rel := path[len(src):]
		if d.Name() == ".git" {
			return filepath.SkipDir
		}
		if d.IsDir() {
			if rel == "" {
				// This is the root of the tree, so
				// the destination already exists.
				return nil
			}
			return os.Mkdir(filepath.Join(dst, rel), 0777)
		}
		// Only copy .go and related files.
		if n := d.Name(); !(filepath.Ext(n) == ".go" || n == "go.mod" || n == "go.sum") {
			return nil
		}

		// Copy file.
		fileMap[path] = filepath.Join(dst, rel)
		return copyFile(path, filepath.Join(dst, rel))
	})
	if err != nil {
		t.Fatalf("error copying source tree: %v", err)
	}

	t.Logf("copied source tree to %s", dst)

	return dst, fileMap
}

func copyFile(src, dst string) error {
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()

	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		return err
	}

	return out.Close()
}

func diffFiles(t *testing.T, fileMap map[string]string) bool {
	diffs := 0
	for orig, new := range fileMap {
		// --strip-trailing-cr is important on Windows where
		// git and gofmt may disagree on newline conventions.
		diff := exec.Command("diff", "-u", "--strip-trailing-cr", orig, new)
		diff.Stdout = os.Stdout
		diff.Stderr = os.Stderr
		if err := diff.Run(); err != nil {
			switch err := err.(type) {
			case *exec.ExitError:
				if err.ExitCode() == 1 {
					diffs++
					continue
				}
			}
			t.Errorf("diff failed: %v", err)
		}
	}
	return diffs != 0
}
<|endoftext|>


<commit_before>// Copyright 2015 Hajime Hoshi
//
// 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 main

import (
	"bytes"
	"fmt"
	"image/color"
	"log"
	"math"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	"github.com/hajimehoshi/ebiten/examples/common"
)

const (
	screenWidth  = 320
	screenHeight = 240
	sampleRate   = 44100
)

var audioContext *audio.Context

func init() {
	var err error
	audioContext, err = audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}
}

var pcm = make([]float64, 4*sampleRate)

const baseFreq = 220

func init() {
	s := float64(sampleRate)
	amp := []float64{1.0, 0.8, 0.6, 0.4, 0.2}
	x := []float64{4.0, 2.0, 1.0, 0.5, 0.25}
	for i := 0; i < len(pcm); i++ {
		v := 0.0
		twoPiF := 2.0 * math.Pi * baseFreq
		for j := 0; j < len(amp); j++ {
			a := amp[j] * math.Exp(-5*float64(i)/(x[j]*s))
			v += a * math.Sin(float64(i)*twoPiF*float64(j+1)/s)
		}
		pcm[i] = v / 5.0
	}
}

var (
	noteCache = map[int][]byte{}
)

func toBytes(l, r []int16) []byte {
	if len(l) != len(r) {
		panic("len(l) must equal to len(r)")
	}
	b := make([]byte, len(l)*4)
	for i, _ := range l {
		b[4*i] = byte(l[i])
		b[4*i+1] = byte(l[i] >> 8)
		b[4*i+2] = byte(r[i])
		b[4*i+3] = byte(r[i] >> 8)
	}
	return b
}

type stream struct {
	*bytes.Reader
}

func (s *stream) Close() error {
	s.Reader = nil
	return nil
}

func addNote(freq float64, vol float64) error {
	f := int(freq)
	if n, ok := noteCache[f]; ok {
		p, err := audio.NewPlayer(audioContext, &stream{bytes.NewReader(n)})
		if err != nil {
			return err
		}
		p.Play()
		return nil
	}
	length := len(pcm) * baseFreq / f
	l := make([]int16, length)
	r := make([]int16, length)
	j := 0
	jj := 0
	for i := 0; i < len(l); i++ {
		p := pcm[j]
		l[i] = int16(p * vol * math.MaxInt16)
		r[i] = l[i]
		jj += f
		j = jj / baseFreq
	}
	n := toBytes(l, r)
	noteCache[f] = n
	p, err := audio.NewPlayer(audioContext, &stream{bytes.NewReader(n)})
	if err != nil {
		return err
	}
	p.Play()
	return nil
}

var keys = []ebiten.Key{
	ebiten.KeyQ,
	ebiten.KeyA,
	ebiten.KeyW,
	ebiten.KeyS,
	ebiten.KeyD,
	ebiten.KeyR,
	ebiten.KeyF,
	ebiten.KeyT,
	ebiten.KeyG,
	ebiten.KeyH,
	ebiten.KeyU,
	ebiten.KeyJ,
	ebiten.KeyI,
	ebiten.KeyK,
	ebiten.KeyO,
	ebiten.KeyL,
}

var keyStates = map[ebiten.Key]int{}

func init() {
	for _, key := range keys {
		keyStates[key] = 0
	}
}

func updateInput() {
	for _, key := range keys {
		if !ebiten.IsKeyPressed(key) {
			keyStates[key] = 0
			continue
		}
		keyStates[key]++
	}
}

var (
	imagePiano *ebiten.Image
)

func init() {
	var err error
	imageEmpty, err := ebiten.NewImage(16, 16, ebiten.FilterNearest)
	if err != nil {
		panic(err)
	}
	imageEmpty.Fill(color.White)
	imagePiano, err = ebiten.NewImage(screenWidth, screenHeight, ebiten.FilterNearest)
	if err != nil {
		panic(err)
	}
	whiteKeys := []string{"A", "S", "D", "F", "G", "H", "J", "K", "L"}
	width := 24
	y := 48
	for i, k := range whiteKeys {
		x := i*width + 36
		height := 112
		op := &ebiten.DrawImageOptions{}
		w, h := imageEmpty.Size()
		op.GeoM.Scale(float64(width-1)/float64(w), float64(height)/float64(h))
		op.GeoM.Translate(float64(x), float64(y))
		op.ColorM.Scale(1, 1, 1, 1)
		imagePiano.DrawImage(imageEmpty, op)
		common.ArcadeFont.DrawText(imagePiano, k, x+8, y+height-16, 1, color.Black)
	}

	blackKeys := []string{"Q", "W", "", "R", "T", "", "U", "I", "O"}
	for i, k := range blackKeys {
		if k == "" {
			continue
		}
		x := i*width + 24
		height := 64
		op := &ebiten.DrawImageOptions{}
		w, h := imageEmpty.Size()
		op.GeoM.Scale(float64(width-1)/float64(w), float64(height)/float64(h))
		op.GeoM.Translate(float64(x), float64(y))
		op.ColorM.Scale(0, 0, 0, 1)
		imagePiano.DrawImage(imageEmpty, op)
		common.ArcadeFont.DrawText(imagePiano, k, x+8, y+height-16, 1, color.White)
	}
}

func update(screen *ebiten.Image) error {
	updateInput()
	for i, key := range keys {
		if keyStates[key] != 1 {
			continue
		}
		if err := addNote(220*math.Exp2(float64(i-1)/12.0), 1.0); err != nil {
			return err
		}
	}

	screen.Fill(color.RGBA{0x80, 0x80, 0xc0, 0xff})
	screen.DrawImage(imagePiano, nil)

	ebitenutil.DebugPrint(screen, fmt.Sprintf("FPS: %0.2f", ebiten.CurrentFPS()))

	if err := audioContext.Update(); err != nil {
		return err
	}
	return nil
}

func main() {
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Piano (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>examples/piano: Add comments<commit_after>// Copyright 2015 Hajime Hoshi
//
// 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 main

import (
	"bytes"
	"fmt"
	"image/color"
	"log"
	"math"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/audio"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	"github.com/hajimehoshi/ebiten/examples/common"
)

const (
	screenWidth  = 320
	screenHeight = 240
	sampleRate   = 44100
)

var audioContext *audio.Context

func init() {
	var err error
	audioContext, err = audio.NewContext(sampleRate)
	if err != nil {
		log.Fatal(err)
	}
}

var pcm = make([]float64, 4*sampleRate)

const baseFreq = 220

func init() {
	s := float64(sampleRate)
	amp := []float64{1.0, 0.8, 0.6, 0.4, 0.2}
	x := []float64{4.0, 2.0, 1.0, 0.5, 0.25}
	for i := 0; i < len(pcm); i++ {
		v := 0.0
		twoPiF := 2.0 * math.Pi * baseFreq
		for j := 0; j < len(amp); j++ {
			a := amp[j] * math.Exp(-5*float64(i)/(x[j]*s))
			v += a * math.Sin(float64(i)*twoPiF*float64(j+1)/s)
		}
		pcm[i] = v / 5.0
	}
}

var (
	noteCache = map[int][]byte{}
)

func toBytes(l, r []int16) []byte {
	if len(l) != len(r) {
		panic("len(l) must equal to len(r)")
	}
	b := make([]byte, len(l)*4)
	for i, _ := range l {
		b[4*i] = byte(l[i])
		b[4*i+1] = byte(l[i] >> 8)
		b[4*i+2] = byte(r[i])
		b[4*i+3] = byte(r[i] >> 8)
	}
	return b
}

type stream struct {
	*bytes.Reader
}

func (s *stream) Close() error {
	s.Reader = nil
	return nil
}

func addNote(freq float64, vol float64) error {
	// TODO: Call Close method of *audio.Player.
	// However, this works without Close because Close is automatically called when GC
	// collects a *audio.Player object.
	f := int(freq)
	if n, ok := noteCache[f]; ok {
		p, err := audio.NewPlayer(audioContext, &stream{bytes.NewReader(n)})
		if err != nil {
			return err
		}
		p.Play()
		return nil
	}
	length := len(pcm) * baseFreq / f
	l := make([]int16, length)
	r := make([]int16, length)
	j := 0
	jj := 0
	for i := 0; i < len(l); i++ {
		p := pcm[j]
		l[i] = int16(p * vol * math.MaxInt16)
		r[i] = l[i]
		jj += f
		j = jj / baseFreq
	}
	n := toBytes(l, r)
	noteCache[f] = n
	p, err := audio.NewPlayer(audioContext, &stream{bytes.NewReader(n)})
	if err != nil {
		return err
	}
	p.Play()
	return nil
}

var keys = []ebiten.Key{
	ebiten.KeyQ,
	ebiten.KeyA,
	ebiten.KeyW,
	ebiten.KeyS,
	ebiten.KeyD,
	ebiten.KeyR,
	ebiten.KeyF,
	ebiten.KeyT,
	ebiten.KeyG,
	ebiten.KeyH,
	ebiten.KeyU,
	ebiten.KeyJ,
	ebiten.KeyI,
	ebiten.KeyK,
	ebiten.KeyO,
	ebiten.KeyL,
}

var keyStates = map[ebiten.Key]int{}

func init() {
	for _, key := range keys {
		keyStates[key] = 0
	}
}

func updateInput() {
	for _, key := range keys {
		if !ebiten.IsKeyPressed(key) {
			keyStates[key] = 0
			continue
		}
		keyStates[key]++
	}
}

var (
	imagePiano *ebiten.Image
)

func init() {
	var err error
	imageEmpty, err := ebiten.NewImage(16, 16, ebiten.FilterNearest)
	if err != nil {
		panic(err)
	}
	imageEmpty.Fill(color.White)
	imagePiano, err = ebiten.NewImage(screenWidth, screenHeight, ebiten.FilterNearest)
	if err != nil {
		panic(err)
	}
	whiteKeys := []string{"A", "S", "D", "F", "G", "H", "J", "K", "L"}
	width := 24
	y := 48
	for i, k := range whiteKeys {
		x := i*width + 36
		height := 112
		op := &ebiten.DrawImageOptions{}
		w, h := imageEmpty.Size()
		op.GeoM.Scale(float64(width-1)/float64(w), float64(height)/float64(h))
		op.GeoM.Translate(float64(x), float64(y))
		op.ColorM.Scale(1, 1, 1, 1)
		imagePiano.DrawImage(imageEmpty, op)
		common.ArcadeFont.DrawText(imagePiano, k, x+8, y+height-16, 1, color.Black)
	}

	blackKeys := []string{"Q", "W", "", "R", "T", "", "U", "I", "O"}
	for i, k := range blackKeys {
		if k == "" {
			continue
		}
		x := i*width + 24
		height := 64
		op := &ebiten.DrawImageOptions{}
		w, h := imageEmpty.Size()
		op.GeoM.Scale(float64(width-1)/float64(w), float64(height)/float64(h))
		op.GeoM.Translate(float64(x), float64(y))
		op.ColorM.Scale(0, 0, 0, 1)
		imagePiano.DrawImage(imageEmpty, op)
		common.ArcadeFont.DrawText(imagePiano, k, x+8, y+height-16, 1, color.White)
	}
}

func update(screen *ebiten.Image) error {
	updateInput()
	for i, key := range keys {
		if keyStates[key] != 1 {
			continue
		}
		if err := addNote(220*math.Exp2(float64(i-1)/12.0), 1.0); err != nil {
			return err
		}
	}

	screen.Fill(color.RGBA{0x80, 0x80, 0xc0, 0xff})
	screen.DrawImage(imagePiano, nil)

	ebitenutil.DebugPrint(screen, fmt.Sprintf("FPS: %0.2f", ebiten.CurrentFPS()))

	if err := audioContext.Update(); err != nil {
		return err
	}
	return nil
}

func main() {
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Piano (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package requests

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cactus/go-statsd-client/statsd"
	"github.com/twitchscience/scoop_protocol/spade"
	"github.com/twitchscience/spade_edge/loggers"
)

var (
	xDomainContents = func() []byte {
		filename := os.Getenv("CROSS_DOMAIN_LOCATION")
		if filename == "" {
			filename = "../build/config/crossdomain.xml"
		}
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			log.Fatalln("Cross domain file not found: ", err)
		}
		return b
	}()
	xmlApplicationType = mime.TypeByExtension(".xml")
	xarth              = []byte("XARTH")
	dataFlag           = []byte("data=")
)

const corsMaxAge = "86400" // One day

// EdgeLoggers represent the different kind of loggers for Spade events
type EdgeLoggers struct {
	sync.WaitGroup
	closed             chan struct{}
	S3AuditLogger      loggers.SpadeEdgeLogger
	S3EventLogger      loggers.SpadeEdgeLogger
	KinesisEventLogger loggers.SpadeEdgeLogger
}

// NewEdgeLoggers returns a new instance of an EdgeLoggers struct pre-filled
// wuth UndefinedLogger logger instances
func NewEdgeLoggers() *EdgeLoggers {
	return &EdgeLoggers{
		closed:             make(chan struct{}),
		S3AuditLogger:      loggers.UndefinedLogger{},
		S3EventLogger:      loggers.UndefinedLogger{},
		KinesisEventLogger: loggers.UndefinedLogger{},
	}
}

func (e *EdgeLoggers) log(event *spade.Event, context *requestContext) error {
	e.Add(1)
	defer e.Done()

	// If reading from the `closed` channel succeeds, the logger is closed.
	select {
	case <-e.closed:
		return errors.New("Loggers are shutting down")
	default: // Make this a non-blocking select
	}

	auditErr := e.S3AuditLogger.Log(event)
	eventErr := e.S3EventLogger.Log(event)
	kinesisErr := e.KinesisEventLogger.Log(event)

	context.recordLoggerAttempt(auditErr, "audit")
	context.recordLoggerAttempt(eventErr, "event")
	context.recordLoggerAttempt(kinesisErr, "kinesis")

	if eventErr != nil &&
		kinesisErr != nil {
		return errors.New("Failed to store the event in any of the loggers")
	}

	return nil
}

// Close closes the loggers
func (e *EdgeLoggers) Close() {
	close(e.closed)
	e.Wait()

	e.KinesisEventLogger.Close()
	e.S3AuditLogger.Close()
	e.S3EventLogger.Close()
}

// SpadeHandler handles http requests and forwards them to the EdgeLoggers
type SpadeHandler struct {
	StatLogger  statsd.Statter
	EdgeLoggers *EdgeLoggers
	Time        func() time.Time // Defaults to time.Now
	instanceID  string
	corsOrigins map[string]bool

	// eventCount counts the number of event requests handled. It is used in
	// uuid generation. eventCount is read and written from multiple go routines
	// so any access to it should go through sync/atomic
	eventCount uint64
}

// NewSpadeHandler returns a new instance of SpadeHandler
func NewSpadeHandler(stats statsd.Statter, loggers *EdgeLoggers, instanceID string, CORSOrigins []string) *SpadeHandler {
	h := &SpadeHandler{
		StatLogger:  stats,
		EdgeLoggers: loggers,
		Time:        time.Now,
		instanceID:  instanceID,
		corsOrigins: make(map[string]bool),
	}

	for _, origin := range CORSOrigins {
		trimmedOrigin := strings.TrimSpace(origin)
		if trimmedOrigin != "" {
			h.corsOrigins[trimmedOrigin] = true
		}
	}
	return h
}

func parseLastForwarder(header string) net.IP {
	var clientIP string
	comma := strings.LastIndex(header, ",")
	if comma > -1 && comma < len(header)+1 {
		clientIP = header[comma+1:]
	} else {
		clientIP = header
	}

	return net.ParseIP(strings.TrimSpace(clientIP))
}

func (s *SpadeHandler) handleSpadeRequests(r *http.Request, context *requestContext) int {
	statTimer := newTimerInstance()

	xForwardedFor := r.Header.Get(context.IPHeader)
	clientIP := parseLastForwarder(xForwardedFor)

	context.Timers["ip"] = statTimer.stopTiming()

	err := r.ParseForm()
	if err != nil {
		if err.Error() == "http: request body too large" {
			_ = s.StatLogger.Inc("large_request", 1, 1.0)
			log.Printf("Request larger than %d bytes, rejecting.", maxBytesPerRequest)
			return http.StatusRequestEntityTooLarge
		}
		return http.StatusBadRequest
	}

	data := r.Form.Get("data")
	if data == "" && r.Method == "POST" {
		// if we're here then our clients have POSTed us something weird,
		// for example, something that maybe
		// application/x-www-form-urlencoded but with the Content-Type
		// header set incorrectly... best effort here on out

		var b []byte
		b, err = ioutil.ReadAll(r.Body)
		if err != nil {
			if err.Error() == "http: request body too large" {
				_ = s.StatLogger.Inc("large_request", 1, 1.0)
				log.Printf("Request larger than %d bytes, rejecting.", maxBytesPerRequest)
				return http.StatusRequestEntityTooLarge
			}
			return http.StatusBadRequest
		}
		if bytes.Equal(b[:5], dataFlag) {
			context.BadClient = true
			b = b[5:]
		}
		data = string(b)

	}
	if data == "" {
		return http.StatusBadRequest
	}

	context.Timers["data"] = statTimer.stopTiming()

	count := atomic.AddUint64(&s.eventCount, 1)
	uuid := fmt.Sprintf("%s-%08x-%08x", s.instanceID, context.Now.Unix(), count)
	context.Timers["uuid"] = statTimer.stopTiming()

	event := spade.NewEvent(
		context.Now,
		clientIP,
		xForwardedFor,
		uuid,
		data,
	)

	defer func() {
		context.Timers["write"] = statTimer.stopTiming()
	}()

	err = s.EdgeLoggers.log(event, context)

	if err != nil {
		return http.StatusBadRequest
	}
	return http.StatusNoContent
}

const (
	ipForwardHeader    = "X-Forwarded-For"
	badEndpoint        = "FourOhFour"
	nTimers            = 5
	maxBytesPerRequest = 500 * 1024
)

var allowedMethods = map[string]bool{
	"GET":     true,
	"POST":    true,
	"OPTIONS": true,
}
var allowedMethodsHeader string // Comma-separated version of allowedMethods

func (s *SpadeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxBytesPerRequest)
	if !allowedMethods[r.Method] {
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	w.Header().Set("Vary", "Origin")

	origin := r.Header.Get("Origin")
	if s.corsOrigins[origin] {
		w.Header().Set("Access-Control-Allow-Origin", origin)
		w.Header().Set("Access-Control-Allow-Methods", allowedMethodsHeader)
	}

	if r.Method == "OPTIONS" {
		w.Header().Set("Access-Control-Max-Age", corsMaxAge)
		w.WriteHeader(http.StatusOK)
		return
	}

	context := &requestContext{
		Now:       s.Time(),
		Method:    r.Method,
		Endpoint:  r.URL.Path,
		IPHeader:  ipForwardHeader,
		Timers:    make(map[string]time.Duration, nTimers),
		BadClient: false,
	}
	timer := newTimerInstance()
	context.setStatus(s.serve(w, r, context))
	context.Timers["http"] = timer.stopTiming()

	context.recordStats(s.StatLogger)
}

func (s *SpadeHandler) serve(w http.ResponseWriter, r *http.Request, context *requestContext) int {
	var status int
	switch r.URL.Path {
	case "/crossdomain.xml":
		w.Header().Add("Content-Type", xmlApplicationType)
		_, err := w.Write(xDomainContents)
		if err != nil {
			log.Printf("Unable to write crossdomain.xml contents: %v", err)
			return http.StatusInternalServerError
		}
		return http.StatusOK
	case "/healthcheck":
		status = http.StatusOK
	case "/xarth":
		_, err := w.Write(xarth)
		if err != nil {
			log.Printf("Error writing XARTH response: %v", err)
			return http.StatusInternalServerError
		}
		return http.StatusOK
	// Accepted tracking endpoints.
	case "/":
		status = s.handleSpadeRequests(r, context)
	case "/track":
		status = s.handleSpadeRequests(r, context)
	case "/track/":
		status = s.handleSpadeRequests(r, context)
	// dont track everything else
	default:
		context.Endpoint = badEndpoint
		status = http.StatusNotFound
	}
	w.WriteHeader(status)
	return status
}

func init() {
	var allowedMethodsList []string
	for k := range allowedMethods {
		allowedMethodsList = append(allowedMethodsList, k)
	}
	allowedMethodsHeader = strings.Join(allowedMethodsList, ", ")
}
<commit_msg>Fix up request_handler.go<commit_after>package requests

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"mime"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cactus/go-statsd-client/statsd"
	"github.com/twitchscience/aws_utils/logger"
	"github.com/twitchscience/scoop_protocol/spade"
	"github.com/twitchscience/spade_edge/loggers"
)

var (
	xDomainContents = func() (b []byte) {
		filename := os.Getenv("CROSS_DOMAIN_LOCATION")
		if filename == "" {
			filename = "../build/config/crossdomain.xml"
		}
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			logger.WithError(err).Fatal("Cross domain file not found")
		}
		return
	}()
	xmlApplicationType = mime.TypeByExtension(".xml")
	xarth              = []byte("XARTH")
	dataFlag           = []byte("data=")
)

const corsMaxAge = "86400" // One day

// EdgeLoggers represent the different kind of loggers for Spade events
type EdgeLoggers struct {
	sync.WaitGroup
	closed             chan struct{}
	S3AuditLogger      loggers.SpadeEdgeLogger
	S3EventLogger      loggers.SpadeEdgeLogger
	KinesisEventLogger loggers.SpadeEdgeLogger
}

// NewEdgeLoggers returns a new instance of an EdgeLoggers struct pre-filled
// wuth UndefinedLogger logger instances
func NewEdgeLoggers() *EdgeLoggers {
	return &EdgeLoggers{
		closed:             make(chan struct{}),
		S3AuditLogger:      loggers.UndefinedLogger{},
		S3EventLogger:      loggers.UndefinedLogger{},
		KinesisEventLogger: loggers.UndefinedLogger{},
	}
}

func (e *EdgeLoggers) log(event *spade.Event, context *requestContext) error {
	e.Add(1)
	defer e.Done()

	// If reading from the `closed` channel succeeds, the logger is closed.
	select {
	case <-e.closed:
		return errors.New("Loggers are shutting down")
	default: // Make this a non-blocking select
	}

	auditErr := e.S3AuditLogger.Log(event)
	eventErr := e.S3EventLogger.Log(event)
	kinesisErr := e.KinesisEventLogger.Log(event)

	context.recordLoggerAttempt(auditErr, "audit")
	context.recordLoggerAttempt(eventErr, "event")
	context.recordLoggerAttempt(kinesisErr, "kinesis")

	if eventErr != nil &&
		kinesisErr != nil {
		return errors.New("Failed to store the event in any of the loggers")
	}

	return nil
}

// Close closes the loggers
func (e *EdgeLoggers) Close() {
	close(e.closed)
	e.Wait()

	e.KinesisEventLogger.Close()
	e.S3AuditLogger.Close()
	e.S3EventLogger.Close()
}

// SpadeHandler handles http requests and forwards them to the EdgeLoggers
type SpadeHandler struct {
	StatLogger  statsd.Statter
	EdgeLoggers *EdgeLoggers
	Time        func() time.Time // Defaults to time.Now
	instanceID  string
	corsOrigins map[string]bool

	// eventCount counts the number of event requests handled. It is used in
	// uuid generation. eventCount is read and written from multiple go routines
	// so any access to it should go through sync/atomic
	eventCount uint64
}

// NewSpadeHandler returns a new instance of SpadeHandler
func NewSpadeHandler(stats statsd.Statter, loggers *EdgeLoggers, instanceID string, CORSOrigins []string) *SpadeHandler {
	h := &SpadeHandler{
		StatLogger:  stats,
		EdgeLoggers: loggers,
		Time:        time.Now,
		instanceID:  instanceID,
		corsOrigins: make(map[string]bool),
	}

	for _, origin := range CORSOrigins {
		trimmedOrigin := strings.TrimSpace(origin)
		if trimmedOrigin != "" {
			h.corsOrigins[trimmedOrigin] = true
		}
	}
	return h
}

func parseLastForwarder(header string) net.IP {
	var clientIP string
	comma := strings.LastIndex(header, ",")
	if comma > -1 && comma < len(header)+1 {
		clientIP = header[comma+1:]
	} else {
		clientIP = header
	}

	return net.ParseIP(strings.TrimSpace(clientIP))
}

func (s *SpadeHandler) handleSpadeRequests(r *http.Request, context *requestContext) int {
	statTimer := newTimerInstance()

	xForwardedFor := r.Header.Get(context.IPHeader)
	clientIP := parseLastForwarder(xForwardedFor)

	context.Timers["ip"] = statTimer.stopTiming()

	err := r.ParseForm()
	if err != nil {
		if err.Error() == "http: request body too large" {
			_ = s.StatLogger.Inc("large_request", 1, 1.0)
			logger.Warn(maxBytesErrorString)
			return http.StatusRequestEntityTooLarge
		}
		return http.StatusBadRequest
	}

	data := r.Form.Get("data")
	if data == "" && r.Method == "POST" {
		// if we're here then our clients have POSTed us something weird,
		// for example, something that maybe
		// application/x-www-form-urlencoded but with the Content-Type
		// header set incorrectly... best effort here on out

		var b []byte
		b, err = ioutil.ReadAll(r.Body)
		if err != nil {
			if err.Error() == "http: request body too large" {
				_ = s.StatLogger.Inc("large_request", 1, 1.0)
				logger.Warn(maxBytesErrorString)
				return http.StatusRequestEntityTooLarge
			}
			return http.StatusBadRequest
		}
		if bytes.Equal(b[:5], dataFlag) {
			context.BadClient = true
			b = b[5:]
		}
		data = string(b)

	}
	if data == "" {
		return http.StatusBadRequest
	}

	context.Timers["data"] = statTimer.stopTiming()

	count := atomic.AddUint64(&s.eventCount, 1)
	uuid := fmt.Sprintf("%s-%08x-%08x", s.instanceID, context.Now.Unix(), count)
	context.Timers["uuid"] = statTimer.stopTiming()

	event := spade.NewEvent(
		context.Now,
		clientIP,
		xForwardedFor,
		uuid,
		data,
	)

	defer func() {
		context.Timers["write"] = statTimer.stopTiming()
	}()

	err = s.EdgeLoggers.log(event, context)

	if err != nil {
		return http.StatusBadRequest
	}
	return http.StatusNoContent
}

const (
	ipForwardHeader     = "X-Forwarded-For"
	badEndpoint         = "FourOhFour"
	nTimers             = 5
	maxBytesPerRequest  = 500 * 1024
	maxBytesErrorString = "Request larger than 500 KB, rejecting."
)

var allowedMethods = map[string]bool{
	"GET":     true,
	"POST":    true,
	"OPTIONS": true,
}
var allowedMethodsHeader string // Comma-separated version of allowedMethods

func (s *SpadeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, maxBytesPerRequest)
	if !allowedMethods[r.Method] {
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	w.Header().Set("Vary", "Origin")

	origin := r.Header.Get("Origin")
	if s.corsOrigins[origin] {
		w.Header().Set("Access-Control-Allow-Origin", origin)
		w.Header().Set("Access-Control-Allow-Methods", allowedMethodsHeader)
	}

	if r.Method == "OPTIONS" {
		w.Header().Set("Access-Control-Max-Age", corsMaxAge)
		w.WriteHeader(http.StatusOK)
		return
	}

	context := &requestContext{
		Now:       s.Time(),
		Method:    r.Method,
		Endpoint:  r.URL.Path,
		IPHeader:  ipForwardHeader,
		Timers:    make(map[string]time.Duration, nTimers),
		BadClient: false,
	}
	timer := newTimerInstance()
	context.setStatus(s.serve(w, r, context))
	context.Timers["http"] = timer.stopTiming()

	context.recordStats(s.StatLogger)
}

func (s *SpadeHandler) serve(w http.ResponseWriter, r *http.Request, context *requestContext) int {
	var status int
	switch r.URL.Path {
	case "/crossdomain.xml":
		w.Header().Add("Content-Type", xmlApplicationType)
		_, err := w.Write(xDomainContents)
		if err != nil {
			logger.WithError(err).Error("Unable to write crossdomain.xml contents")
			return http.StatusInternalServerError
		}
		return http.StatusOK
	case "/healthcheck":
		status = http.StatusOK
	case "/xarth":
		_, err := w.Write(xarth)
		if err != nil {
			logger.WithError(err).Error("Error writing XARTH response")
			return http.StatusInternalServerError
		}
		return http.StatusOK
	// Accepted tracking endpoints.
	case "/", "/track", "/track/":
		status = s.handleSpadeRequests(r, context)
	// dont track everything else
	default:
		context.Endpoint = badEndpoint
		status = http.StatusNotFound
	}
	w.WriteHeader(status)
	return status
}

func init() {
	var allowedMethodsList []string
	for k := range allowedMethods {
		allowedMethodsList = append(allowedMethodsList, k)
	}
	allowedMethodsHeader = strings.Join(allowedMethodsList, ", ")
}
<|endoftext|>


<commit_before>package combine

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"path/filepath"
	"strings"

	"github.com/minodisk/go-jsonschema/tools/encoding"
	"github.com/minodisk/go-watcher"

	"gopkg.in/yaml.v2"
)

type Options struct {
	Input    string
	Meta     string
	Output   string
	Encoding encoding.Encoding
	IsWatch  bool
}

func Run(o Options) (err error) {
	err = run(o)
	if err != nil {
		return err
	}
	if o.IsWatch {
		done := make(chan int)
		w := watcher.New()
		go func() {
			for {
				select {
				case <-w.Events:
					err = run(o)
					if err != nil {
						log.Println("Running error:", err)
					}
				case err := <-w.Errors:
					log.Println("Error:", err)
				}
			}
		}()
		err = w.Watch([]string{o.Input, o.Meta})
		if err != nil {
			return err
		}
		<-done
	}
	return nil
}

func run(o Options) (err error) {
	b, err := Combine(o.Input, o.Meta, o.Encoding)
	if err != nil {
		return err
	}

	if o.Output == "" {
		fmt.Println(string(b))
		return nil
	}

	return ioutil.WriteFile(o.Output, b, 0644)
}

func Combine(input string, meta string, enc encoding.Encoding) (combined []byte, err error) {
	d, err := filepath.Abs(input)
	if err != nil {
		return nil, err
	}

	meta, err = filepath.Abs(meta)
	if err != nil {
		return nil, err
	}

	files := map[string][]byte{}
	if err := readFiles(d, &files); err != nil {
		return nil, err
	}
	b, ok := files[meta]
	if !ok {
		return nil, fmt.Errorf("meta file is required: %s", meta)
	}
	delete(files, meta)

	var schema struct {
		Schema      string `yaml:"$schema",json:"$schema"`
		Type        interface{}
		Title       string
		Description string
		Links       []map[string]string
		Definitions map[string]interface{}
		Properties  map[string]map[string]interface{}
	}

	metaEnc, err := encoding.NewWithFilename(meta)
	if err != nil {
		return nil, err
	}
	switch metaEnc {
	default:
		return nil, fmt.Errorf("unsupported encoding '%s'", metaEnc)
	case encoding.JSON:
		if err := json.Unmarshal(b, &schema); err != nil {
			return nil, err
		}
	case encoding.YAML:
		if err := yaml.Unmarshal(b, &schema); err != nil {
			return nil, err
		}
	}

	schema.Definitions = make(map[string]interface{})
	schema.Properties = make(map[string]map[string]interface{})

	for filename, b := range files {
		var s interface{}
		if err := yaml.Unmarshal(b, &s); err != nil {
			return nil, err
		}
		subSchema, err := encoding.KeyValueMap(s)
		if err != nil {
			return nil, err
		}
		i, ok := subSchema.(map[string]interface{})["id"]
		var id string
		if ok {
			id = i.(string)
		} else {
			id = strings.TrimSuffix(filepath.Base(filename), filepath.Ext(filename))
		}
		schema.Definitions[id] = subSchema
		schema.Properties[id] = make(map[string]interface{})
		schema.Properties[id]["$ref"] = fmt.Sprintf("#/definitions/%s", id)
	}

	switch enc {
	default:
		return nil, fmt.Errorf("unsupported encoding '%s'", enc)
	case encoding.JSON:
		return json.MarshalIndent(schema, "", "  ")
	case encoding.YAML:
		return yaml.Marshal(schema)
	}
}

func readFiles(dirname string, files *map[string][]byte) error {
	fs, err := ioutil.ReadDir(dirname)
	if err != nil {
		return err
	}
	for _, f := range fs {
		name := filepath.Join(dirname, f.Name())
		if f.IsDir() {
			readFiles(name, files)
			continue
		}
		b, err := ioutil.ReadFile(name)
		if err != nil {
			return err
		}
		(*files)[name] = b
	}
	return nil
}
<commit_msg>Add json tags<commit_after>package combine

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"path/filepath"
	"strings"

	"github.com/minodisk/go-jsonschema/tools/encoding"
	"github.com/minodisk/go-watcher"

	"gopkg.in/yaml.v2"
)

type Options struct {
	Input    string
	Meta     string
	Output   string
	Encoding encoding.Encoding
	IsWatch  bool
}

func Run(o Options) (err error) {
	err = run(o)
	if err != nil {
		return err
	}
	if o.IsWatch {
		done := make(chan int)
		w := watcher.New()
		go func() {
			for {
				select {
				case <-w.Events:
					err = run(o)
					if err != nil {
						log.Println("Running error:", err)
					}
				case err := <-w.Errors:
					log.Println("Error:", err)
				}
			}
		}()
		err = w.Watch([]string{o.Input, o.Meta})
		if err != nil {
			return err
		}
		<-done
	}
	return nil
}

func run(o Options) (err error) {
	b, err := Combine(o.Input, o.Meta, o.Encoding)
	if err != nil {
		return err
	}

	if o.Output == "" {
		fmt.Println(string(b))
		return nil
	}

	return ioutil.WriteFile(o.Output, b, 0644)
}

func Combine(input string, meta string, enc encoding.Encoding) (combined []byte, err error) {
	d, err := filepath.Abs(input)
	if err != nil {
		return nil, err
	}

	meta, err = filepath.Abs(meta)
	if err != nil {
		return nil, err
	}

	files := map[string][]byte{}
	if err := readFiles(d, &files); err != nil {
		return nil, err
	}
	b, ok := files[meta]
	if !ok {
		return nil, fmt.Errorf("meta file is required: %s", meta)
	}
	delete(files, meta)

	var schema struct {
		Schema      string                            `yaml:"$schema" json:"$schema"`
		Type        interface{}                       `json:"type"`
		Title       string                            `json:"title"`
		Description string                            `json:"description"`
		Links       []map[string]string               `json:"links"`
		Definitions map[string]interface{}            `json:"definitions"`
		Properties  map[string]map[string]interface{} `json:"properties"`
	}

	metaEnc, err := encoding.NewWithFilename(meta)
	if err != nil {
		return nil, err
	}
	switch metaEnc {
	default:
		return nil, fmt.Errorf("unsupported encoding '%s'", metaEnc)
	case encoding.JSON:
		if err := json.Unmarshal(b, &schema); err != nil {
			return nil, err
		}
	case encoding.YAML:
		if err := yaml.Unmarshal(b, &schema); err != nil {
			return nil, err
		}
	}

	schema.Definitions = make(map[string]interface{})
	schema.Properties = make(map[string]map[string]interface{})

	for filename, b := range files {
		var s interface{}
		if err := yaml.Unmarshal(b, &s); err != nil {
			return nil, err
		}
		subSchema, err := encoding.KeyValueMap(s)
		if err != nil {
			return nil, err
		}
		i, ok := subSchema.(map[string]interface{})["id"]
		var id string
		if ok {
			id = i.(string)
		} else {
			id = strings.TrimSuffix(filepath.Base(filename), filepath.Ext(filename))
		}
		schema.Definitions[id] = subSchema
		schema.Properties[id] = make(map[string]interface{})
		schema.Properties[id]["$ref"] = fmt.Sprintf("#/definitions/%s", id)
	}

	switch enc {
	default:
		return nil, fmt.Errorf("unsupported encoding '%s'", enc)
	case encoding.JSON:
		return json.MarshalIndent(schema, "", "  ")
	case encoding.YAML:
		return yaml.Marshal(schema)
	}
}

func readFiles(dirname string, files *map[string][]byte) error {
	fs, err := ioutil.ReadDir(dirname)
	if err != nil {
		return err
	}
	for _, f := range fs {
		name := filepath.Join(dirname, f.Name())
		if f.IsDir() {
			readFiles(name, files)
			continue
		}
		b, err := ioutil.ReadFile(name)
		if err != nil {
			return err
		}
		(*files)[name] = b
	}
	return nil
}
<|endoftext|>