<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package json

import (
	"encoding/json"
	"errors"
	"net/http"

	"github.com/gorilla/rpc"
)

var null = json.RawMessage([]byte("null"))

// ----------------------------------------------------------------------------
// Request and Response
// ----------------------------------------------------------------------------

// serverRequest represents a JSON-RPC request received by the server.
type serverRequest struct {
	// A String containing the name of the method to be invoked.
	Method string `json:"method"`
	// An Array of objects to pass as arguments to the method.
	Params *json.RawMessage `json:"params"`
	// The request id. This can be of any type. It is used to match the
	// response with the request that it is replying to.
	Id *json.RawMessage `json:"id"`
}

// serverResponse represents a JSON-RPC response returned by the server.
type serverResponse struct {
	// The Object that was returned by the invoked method. This must be null
	// in case there was an error invoking the method.
	Result interface{} `json:"result"`
	// An Error object if there was an error invoking the method. It must be
	// null if there was no error.
	Error interface{} `json:"error"`
	// This must be the same id as the request it is responding to.
	Id *json.RawMessage `json:"id"`
}

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// NewCodec returns a new JSON Codec.
func NewCodec() *Codec {
	return &Codec{}
}

// Codec creates a CodecRequest to process each request.
type Codec struct {
}

// NewRequest returns a CodecRequest.
func (c *Codec) NewRequest(r *http.Request) rpc.CodecRequest {
	return newCodecRequest(r)
}

// ----------------------------------------------------------------------------
// CodecRequest
// ----------------------------------------------------------------------------

// newCodecRequest returns a new CodecRequest.
func newCodecRequest(r *http.Request) rpc.CodecRequest {
	// Decode the request body and check if RPC method is valid.
	req := new(serverRequest)
	err := json.NewDecoder(r.Body).Decode(req)
	r.Body.Close()
	return &CodecRequest{request: req, err: err}
}

// CodecRequest decodes and encodes a single request.
type CodecRequest struct {
	request *serverRequest
	err     error
}

// Method returns the RPC method for the current request.
//
// The method uses a dotted notation as in "Service.Method".
func (c *CodecRequest) Method() (string, error) {
	if c.err == nil {
		return c.request.Method, nil
	}
	return "", c.err
}

// ReadRequest fills the request object for the RPC method.
func (c *CodecRequest) ReadRequest(args interface{}) error {
	if c.err == nil {
		if c.request.Params != nil {
			// JSON params is array value. RPC params is struct.
			// Unmarshal into array containing the request struct.
			params := [1]interface{}{args}
			c.err = json.Unmarshal(*c.request.Params, &params)
		} else {
			c.err = errors.New("rpc: method request ill-formed: missing params field")
		}
	}
	return c.err
}

// WriteResponse encodes the response and writes it to the ResponseWriter.
//
// The err parameter is the error resulted from calling the RPC method,
// or nil if there was no error.
func (c *CodecRequest) WriteResponse(w http.ResponseWriter, reply interface{}, methodErr error) error {
	if c.err != nil {
		return c.err
	}
	res := &serverResponse{
		Result: reply,
		Error:  &null,
		Id:     c.request.Id,
	}
	if methodErr != nil {
		// Propagate error message as string.
		res.Error = methodErr.Error()
		// Result must be null if there was an error invoking the method.
		// http://json-rpc.org/wiki/specification#a1.2Response
		res.Result = &null
	}
	if c.request.Id == nil {
		// Id is null for notifications and they don't have a response.
		res.Id = &null
	} else {
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		encoder := json.NewEncoder(w)
		encoder.Encode(res)
	}
	return nil
}
<commit_msg>reveal encode errors #57<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package json

import (
	"encoding/json"
	"errors"
	"net/http"

	"github.com/gorilla/rpc"
)

var null = json.RawMessage([]byte("null"))

// ----------------------------------------------------------------------------
// Request and Response
// ----------------------------------------------------------------------------

// serverRequest represents a JSON-RPC request received by the server.
type serverRequest struct {
	// A String containing the name of the method to be invoked.
	Method string `json:"method"`
	// An Array of objects to pass as arguments to the method.
	Params *json.RawMessage `json:"params"`
	// The request id. This can be of any type. It is used to match the
	// response with the request that it is replying to.
	Id *json.RawMessage `json:"id"`
}

// serverResponse represents a JSON-RPC response returned by the server.
type serverResponse struct {
	// The Object that was returned by the invoked method. This must be null
	// in case there was an error invoking the method.
	Result interface{} `json:"result"`
	// An Error object if there was an error invoking the method. It must be
	// null if there was no error.
	Error interface{} `json:"error"`
	// This must be the same id as the request it is responding to.
	Id *json.RawMessage `json:"id"`
}

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// NewCodec returns a new JSON Codec.
func NewCodec() *Codec {
	return &Codec{}
}

// Codec creates a CodecRequest to process each request.
type Codec struct {
}

// NewRequest returns a CodecRequest.
func (c *Codec) NewRequest(r *http.Request) rpc.CodecRequest {
	return newCodecRequest(r)
}

// ----------------------------------------------------------------------------
// CodecRequest
// ----------------------------------------------------------------------------

// newCodecRequest returns a new CodecRequest.
func newCodecRequest(r *http.Request) rpc.CodecRequest {
	// Decode the request body and check if RPC method is valid.
	req := new(serverRequest)
	err := json.NewDecoder(r.Body).Decode(req)
	r.Body.Close()
	return &CodecRequest{request: req, err: err}
}

// CodecRequest decodes and encodes a single request.
type CodecRequest struct {
	request *serverRequest
	err     error
}

// Method returns the RPC method for the current request.
//
// The method uses a dotted notation as in "Service.Method".
func (c *CodecRequest) Method() (string, error) {
	if c.err == nil {
		return c.request.Method, nil
	}
	return "", c.err
}

// ReadRequest fills the request object for the RPC method.
func (c *CodecRequest) ReadRequest(args interface{}) error {
	if c.err == nil {
		if c.request.Params != nil {
			// JSON params is array value. RPC params is struct.
			// Unmarshal into array containing the request struct.
			params := [1]interface{}{args}
			c.err = json.Unmarshal(*c.request.Params, &params)
		} else {
			c.err = errors.New("rpc: method request ill-formed: missing params field")
		}
	}
	return c.err
}

// WriteResponse encodes the response and writes it to the ResponseWriter.
//
// The err parameter is the error resulted from calling the RPC method,
// or nil if there was no error.
func (c *CodecRequest) WriteResponse(w http.ResponseWriter, reply interface{}, methodErr error) error {
	if c.err != nil {
		return c.err
	}
	res := &serverResponse{
		Result: reply,
		Error:  &null,
		Id:     c.request.Id,
	}
	if methodErr != nil {
		// Propagate error message as string.
		res.Error = methodErr.Error()
		// Result must be null if there was an error invoking the method.
		// http://json-rpc.org/wiki/specification#a1.2Response
		res.Result = &null
	}
	if c.request.Id == nil {
		// Id is null for notifications and they don't have a response.
		res.Id = &null
	} else {
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		encoder := json.NewEncoder(w)
		c.err = encoder.Encode(res)
	}
	return c.err
}
<|endoftext|>


<commit_before>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 scheduler

import (
	"errors"
	"fmt"
	"os"
	"path"
	"testing"
	"time"

	"github.com/intelsdi-x/snap/control"
	"github.com/intelsdi-x/snap/core"
	"github.com/intelsdi-x/snap/core/cdata"
	"github.com/intelsdi-x/snap/pkg/promise"
	"github.com/intelsdi-x/snap/pkg/schedule"
	"github.com/intelsdi-x/snap/scheduler/wmap"

	log "github.com/Sirupsen/logrus"
	. "github.com/smartystreets/goconvey/convey"
)

var (
	SnapPath                     = os.Getenv("SNAP_PATH")
	snap_collector_mock2_path    = path.Join(SnapPath, "plugin", "snap-collector-mock2")
	snap_processor_passthru_path = path.Join(SnapPath, "plugin", "snap-processor-passthru")
	snap_publisher_file_path     = path.Join(SnapPath, "plugin", "snap-publisher-file")
)

type MockMetricType struct {
	namespace []string
}

func (m MockMetricType) Namespace() []string {
	return m.namespace
}

func (m MockMetricType) LastAdvertisedTime() time.Time {
	return time.Now()
}

func (m MockMetricType) Version() int {
	return 1
}

func (m MockMetricType) Config() *cdata.ConfigDataNode {
	return nil
}

func (m MockMetricType) Data() interface{} {
	return nil
}

func TestCollectPublishWorkflow(t *testing.T) {
	log.SetLevel(log.FatalLevel)
	Convey("Given a started plugin control", t, func() {

		c := control.New()
		c.Start()
		s := New()
		s.SetMetricManager(c)
		Convey("create a workflow", func() {
			rp, err := core.NewRequestedPlugin(snap_collector_mock2_path)
			So(err, ShouldBeNil)
			_, err = c.Load(rp)
			So(err, ShouldBeNil)
			rp2, err := core.NewRequestedPlugin(snap_publisher_file_path)
			So(err, ShouldBeNil)
			_, err = c.Load(rp2)
			So(err, ShouldBeNil)
			rp3, err := core.NewRequestedPlugin(snap_processor_passthru_path)
			So(err, ShouldBeNil)
			_, err = c.Load(rp3)
			So(err, ShouldBeNil)
			time.Sleep(100 * time.Millisecond)

			metrics, err2 := c.MetricCatalog()
			So(err2, ShouldBeNil)
			So(metrics, ShouldNotBeEmpty)

			w := wmap.NewWorkflowMap()
			w.CollectNode.AddMetric("/intel/mock/foo", 2)
			w.CollectNode.AddConfigItem("/intel/mock/foo", "password", "secret")

			pu := wmap.NewPublishNode("file", 3)
			pu.AddConfigItem("file", "/tmp/snap-TestCollectPublishWorkflow.out")

			pr := wmap.NewProcessNode("passthru", 1)
			time.Sleep(100 * time.Millisecond)

			pr.Add(pu)
			w.CollectNode.Add(pr)

			Convey("Start scheduler", func() {
				err := s.Start()
				So(err, ShouldBeNil)
				Convey("Create task", func() {
					t, err := s.CreateTask(schedule.NewSimpleSchedule(time.Millisecond*500), w, false)
					So(err.Errors(), ShouldBeEmpty)
					So(t, ShouldNotBeNil)
					t.(*task).Spin()
					time.Sleep(3 * time.Second)

				})
			})
		})
	})
}

// The mocks below are here for testing work submission
type Mock1 struct {
	count      int
	errorIndex int
	delay      time.Duration
	queue      map[string]int
}

func (m *Mock1) CollectMetrics([]core.Metric, time.Time, string) ([]core.Metric, []error) {
	return nil, nil
}

func (m *Mock1) Work(j job) queuedJob {
	m.queue[j.TypeString()]++
	return m
}

func (m *Mock1) Promise() promise.Promise {
	return m
}

func (m *Mock1) Await() []error {
	time.Sleep(m.delay)
	m.count++
	if m.count == m.errorIndex {
		return []error{errors.New("I am an error")}
	}
	return nil
}

func (m *Mock1) AwaitUntil(time.Duration) []error {
	return nil
}

func (m *Mock1) Complete([]error) {

}

func (m *Mock1) IsComplete() bool {
	return false
}

func (m *Mock1) Job() job {
	return nil
}

func (m *Mock1) AndThen(_ func([]error)) {
}

func (m *Mock1) AndThenUntil(_ time.Duration, _ func([]error)) {
}

func TestWorkJobs(t *testing.T) {
	Convey("Test speed and concurrency of TestWorkJobs\n", t, func() {
		Convey("submit multiple jobs\n", func() {
			m := &Mock1{queue: make(map[string]int)}
			m.delay = time.Millisecond * 100
			pj := newCollectorJob(nil, time.Second*1, m, nil, "")
			prs := make([]*processNode, 0)
			pus := make([]*publishNode, 0)
			counter := 0
			t := &task{manager: m, id: "1", name: "mock"}
			for x := 0; x < 3; x++ {
				n := cdata.NewNode()
				pr := &processNode{config: n, name: fmt.Sprintf("prjob%d", counter)}
				pu := &publishNode{config: n, name: fmt.Sprintf("pujob%d", counter)}
				counter++
				prs = append(prs, pr)
				pus = append(pus, pu)
			}
			workJobs(prs, pus, t, pj)
			So(m.queue["processor"], ShouldEqual, 3)
			So(m.queue["publisher"], ShouldEqual, 3)
			So(t.failedRuns, ShouldEqual, 0)
		})
		Convey("submit multiple jobs with nesting", func() {
			m := &Mock1{queue: make(map[string]int)}
			m.delay = time.Millisecond * 100
			pj := newCollectorJob(nil, time.Second*1, m, nil, "")
			prs := make([]*processNode, 0)
			pus := make([]*publishNode, 0)
			counter := 0
			t := &task{manager: m, id: "1", name: "mock"}
			// 3 proc + 3 pub
			for x := 0; x < 3; x++ {
				n := cdata.NewNode()
				pr := &processNode{config: n, name: fmt.Sprintf("prjob%d", counter)}
				pu := &publishNode{config: n, name: fmt.Sprintf("pujob%d", counter)}
				counter++
				prs = append(prs, pr)
				pus = append(pus, pu)
			}
			// 3 proc => 3 proc + 3 pub
			for _, pr := range prs {
				cprs := make([]*processNode, 0)
				cpus := make([]*publishNode, 0)
				for x := 0; x < 3; x++ {
					n := cdata.NewNode()
					cpr := &processNode{config: n, name: fmt.Sprintf("prjobchild%d", counter)}
					cpu := &publishNode{config: n, name: fmt.Sprintf("pujobchild%d", counter)}
					counter++
					cprs = append(cprs, cpr)
					cpus = append(cpus, cpu)
				}
				pr.ProcessNodes = cprs
				pr.PublishNodes = cpus
			}
			workJobs(prs, pus, t, pj)
			// (3*3)+3
			So(m.queue["processor"], ShouldEqual, 12)
			// (3*3)
			So(m.queue["publisher"], ShouldEqual, 12)
			So(t.failedRuns, ShouldEqual, 0)
		})
		Convey("submit multiple jobs where one has an error", func() {
			m := &Mock1{queue: make(map[string]int)}
			// make the 13th job fail
			m.errorIndex = 13
			m.delay = time.Millisecond * 100
			pj := newCollectorJob(nil, time.Second*1, m, nil, "")
			prs := make([]*processNode, 0)
			pus := make([]*publishNode, 0)
			counter := 0
			t := &task{manager: m, id: "1", name: "mock"}
			// 3 proc + 3 pub
			for x := 0; x < 3; x++ {
				n := cdata.NewNode()
				pr := &processNode{config: n, name: fmt.Sprintf("prjob%d", counter)}
				pu := &publishNode{config: n, name: fmt.Sprintf("pujob%d", counter)}
				counter++
				prs = append(prs, pr)
				pus = append(pus, pu)
			}
			// 3 proc => 3 proc + 3 pub
			for _, pr := range prs {
				cprs := make([]*processNode, 0)
				cpus := make([]*publishNode, 0)
				for x := 0; x < 3; x++ {
					n := cdata.NewNode()
					cpr := &processNode{config: n, name: fmt.Sprintf("prjobchild%d", counter)}
					cpu := &publishNode{config: n, name: fmt.Sprintf("pujobchild%d", counter)}
					counter++
					cprs = append(cprs, cpr)
					cpus = append(cpus, cpu)
				}
				pr.ProcessNodes = cprs
				pr.PublishNodes = cpus
			}
			workJobs(prs, pus, t, pj)
			// (3*3)+3
			So(m.queue["processor"], ShouldEqual, 12)
			// (3*3)
			So(m.queue["publisher"], ShouldEqual, 12)
			So(t.failedRuns, ShouldEqual, 1)
			So(t.lastFailureMessage, ShouldEqual, "I am an error")
		})

	})
}
<commit_msg>Added mutex to workflow job submission mock struc to make it thread safe. Removed the delay.<commit_after>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 scheduler

import (
	"errors"
	"fmt"
	"os"
	"path"
	"sync"
	"testing"
	"time"

	"github.com/intelsdi-x/snap/control"
	"github.com/intelsdi-x/snap/core"
	"github.com/intelsdi-x/snap/core/cdata"
	"github.com/intelsdi-x/snap/pkg/promise"
	"github.com/intelsdi-x/snap/pkg/schedule"
	"github.com/intelsdi-x/snap/scheduler/wmap"

	log "github.com/Sirupsen/logrus"
	. "github.com/smartystreets/goconvey/convey"
)

var (
	SnapPath                     = os.Getenv("SNAP_PATH")
	snap_collector_mock2_path    = path.Join(SnapPath, "plugin", "snap-collector-mock2")
	snap_processor_passthru_path = path.Join(SnapPath, "plugin", "snap-processor-passthru")
	snap_publisher_file_path     = path.Join(SnapPath, "plugin", "snap-publisher-file")
)

type MockMetricType struct {
	namespace []string
}

func (m MockMetricType) Namespace() []string {
	return m.namespace
}

func (m MockMetricType) LastAdvertisedTime() time.Time {
	return time.Now()
}

func (m MockMetricType) Version() int {
	return 1
}

func (m MockMetricType) Config() *cdata.ConfigDataNode {
	return nil
}

func (m MockMetricType) Data() interface{} {
	return nil
}

func TestCollectPublishWorkflow(t *testing.T) {
	log.SetLevel(log.FatalLevel)
	Convey("Given a started plugin control", t, func() {

		c := control.New()
		c.Start()
		s := New()
		s.SetMetricManager(c)
		Convey("create a workflow", func() {
			rp, err := core.NewRequestedPlugin(snap_collector_mock2_path)
			So(err, ShouldBeNil)
			_, err = c.Load(rp)
			So(err, ShouldBeNil)
			rp2, err := core.NewRequestedPlugin(snap_publisher_file_path)
			So(err, ShouldBeNil)
			_, err = c.Load(rp2)
			So(err, ShouldBeNil)
			rp3, err := core.NewRequestedPlugin(snap_processor_passthru_path)
			So(err, ShouldBeNil)
			_, err = c.Load(rp3)
			So(err, ShouldBeNil)
			time.Sleep(100 * time.Millisecond)

			metrics, err2 := c.MetricCatalog()
			So(err2, ShouldBeNil)
			So(metrics, ShouldNotBeEmpty)

			w := wmap.NewWorkflowMap()
			w.CollectNode.AddMetric("/intel/mock/foo", 2)
			w.CollectNode.AddConfigItem("/intel/mock/foo", "password", "secret")

			pu := wmap.NewPublishNode("file", 3)
			pu.AddConfigItem("file", "/tmp/snap-TestCollectPublishWorkflow.out")

			pr := wmap.NewProcessNode("passthru", 1)
			time.Sleep(100 * time.Millisecond)

			pr.Add(pu)
			w.CollectNode.Add(pr)

			Convey("Start scheduler", func() {
				err := s.Start()
				So(err, ShouldBeNil)
				Convey("Create task", func() {
					t, err := s.CreateTask(schedule.NewSimpleSchedule(time.Millisecond*500), w, false)
					So(err.Errors(), ShouldBeEmpty)
					So(t, ShouldNotBeNil)
					t.(*task).Spin()
					time.Sleep(3 * time.Second)

				})
			})
		})
	})
}

// The mocks below are here for testing work submission
type Mock1 struct {
	sync.Mutex
	count      int
	errorIndex int
	delay      time.Duration
	queue      map[string]int
	qstring    string
}

func (m *Mock1) CollectMetrics([]core.Metric, time.Time, string) ([]core.Metric, []error) {
	return nil, nil
}

func (m *Mock1) Work(j job) queuedJob {
	m.Lock()
	defer m.Unlock()
	time.Sleep(m.delay)
	m.queue[j.TypeString()]++
	m.qstring += "," + j.Name()
	return m
}

func (m *Mock1) Promise() promise.Promise {
	return m
}

func (m *Mock1) Await() []error {
	m.Lock()
	defer m.Unlock()
	m.count++
	if m.count == m.errorIndex {
		return []error{errors.New("I am an error")}
	}
	return nil
}

func (m *Mock1) AwaitUntil(time.Duration) []error {
	return nil
}

func (m *Mock1) Complete([]error) {

}

func (m *Mock1) IsComplete() bool {
	return false
}

func (m *Mock1) Job() job {
	return nil
}

func (m *Mock1) AndThen(_ func([]error)) {
}

func (m *Mock1) AndThenUntil(_ time.Duration, _ func([]error)) {
}

func TestWorkJobs(t *testing.T) {
	// log.SetLevel(log.DebugLevel)
	Convey("Test speed and concurrency of TestWorkJobs\n", t, func() {
		Convey("submit multiple jobs\n", func() {
			m1 := &Mock1{queue: make(map[string]int)}
			pj := newCollectorJob(nil, time.Second*1, m1, nil, "")
			prs := make([]*processNode, 0)
			pus := make([]*publishNode, 0)
			counter := 0
			t := &task{manager: m1, id: "1", name: "mock"}
			for x := 0; x < 3; x++ {
				n := cdata.NewNode()
				pr := &processNode{config: n, name: fmt.Sprintf("prjob%d", counter)}
				pu := &publishNode{config: n, name: fmt.Sprintf("pujob%d", counter)}
				counter++
				prs = append(prs, pr)
				pus = append(pus, pu)
			}
			workJobs(prs, pus, t, pj)
			So(t.failedRuns, ShouldEqual, 0)
			So(m1.queue["processor"], ShouldEqual, 3)
			So(m1.queue["publisher"], ShouldEqual, 3)
		})
		Convey("submit multiple jobs with nesting", func() {
			m2 := &Mock1{queue: make(map[string]int)}
			pj := newCollectorJob(nil, time.Second*1, m2, nil, "")
			prs := make([]*processNode, 0)
			pus := make([]*publishNode, 0)
			counter := 0
			t := &task{manager: m2, id: "1", name: "mock"}
			// 3 proc + 3 pub
			for x := 0; x < 3; x++ {
				n := cdata.NewNode()
				pr := &processNode{config: n, name: fmt.Sprintf("prjob%d", counter)}
				pu := &publishNode{config: n, name: fmt.Sprintf("pujob%d", counter)}
				counter++
				prs = append(prs, pr)
				pus = append(pus, pu)
			}
			// 3 proc => 3 proc + 3 pub
			for _, pr := range prs {
				cprs := make([]*processNode, 0)
				cpus := make([]*publishNode, 0)
				for x := 0; x < 3; x++ {
					n := cdata.NewNode()
					cpr := &processNode{config: n, name: fmt.Sprintf("prjobchild%d", counter)}
					cpu := &publishNode{config: n, name: fmt.Sprintf("pujobchild%d", counter)}
					counter++
					cprs = append(cprs, cpr)
					cpus = append(cpus, cpu)
				}
				pr.ProcessNodes = cprs
				pr.PublishNodes = cpus
			}
			workJobs(prs, pus, t, pj)
			So(t.failedRuns, ShouldEqual, 0)
			// (3*3)+3
			So(m2.queue["processor"], ShouldEqual, 12)
			// (3*3)
			So(m2.queue["publisher"], ShouldEqual, 12)

		})
		Convey("submit multiple jobs where one has an error", func() {
			m3 := &Mock1{queue: make(map[string]int)}
			// make the 13th job fail
			m3.errorIndex = 13
			m3.delay = time.Millisecond * 100
			pj := newCollectorJob(nil, time.Second*1, m3, nil, "")
			prs := make([]*processNode, 0)
			pus := make([]*publishNode, 0)
			counter := 0
			t := &task{manager: m3, id: "1", name: "mock"}
			// 3 proc + 3 pub
			for x := 0; x < 3; x++ {
				n := cdata.NewNode()
				pr := &processNode{config: n, name: fmt.Sprintf("prjob%d", counter)}
				pu := &publishNode{config: n, name: fmt.Sprintf("pujob%d", counter)}
				counter++
				prs = append(prs, pr)
				pus = append(pus, pu)
			}
			// 3 proc => 3 proc + 3 pub
			for _, pr := range prs {
				cprs := make([]*processNode, 0)
				cpus := make([]*publishNode, 0)
				for x := 0; x < 3; x++ {
					n := cdata.NewNode()
					cpr := &processNode{config: n, name: fmt.Sprintf("prjobchild%d", counter)}
					cpu := &publishNode{config: n, name: fmt.Sprintf("pujobchild%d", counter)}
					counter++
					cprs = append(cprs, cpr)
					cpus = append(cpus, cpu)
				}
				pr.ProcessNodes = cprs
				pr.PublishNodes = cpus
			}
			workJobs(prs, pus, t, pj)
			So(t.failedRuns, ShouldEqual, 1)
			So(t.lastFailureMessage, ShouldEqual, "I am an error")
			// (3*3)+3
			So(m3.queue["processor"], ShouldEqual, 12)
			// (3*3)
			So(m3.queue["publisher"], ShouldEqual, 12)

		})

	})
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package instrument

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"sync/atomic"
	"time"

	"github.com/m3db/m3x/process"

	"github.com/uber-go/tally"
)

// ExtendedMetricsType is a type of extended metrics to report.
type ExtendedMetricsType int

const (
	// NoExtendedMetrics describes no extended metrics.
	NoExtendedMetrics ExtendedMetricsType = iota

	// SimpleExtendedMetrics describes just a simple level of extended metrics:
	// - number of active goroutines
	// - number of configured gomaxprocs
	// - number of file descriptors
	SimpleExtendedMetrics

	// DetailedExtendedMetrics describes a detailed level of extended metrics:
	// - number of active goroutines
	// - number of configured gomaxprocs
	// - number of file descriptors
	// - memory allocated running count
	// - memory used by heap
	// - memory used by heap that is idle
	// - memory used by heap that is in use
	// - memory used by stack
	// - number of garbage collections
	// - GC pause times
	DetailedExtendedMetrics

	// DefaultExtendedMetricsType is the default extended metrics level.
	DefaultExtendedMetricsType = DetailedExtendedMetrics
)

var (
	validExtendedMetricsTypes = []ExtendedMetricsType{
		NoExtendedMetrics,
		SimpleExtendedMetrics,
		DetailedExtendedMetrics,
	}
)

func (t ExtendedMetricsType) String() string {
	switch t {
	case NoExtendedMetrics:
		return "none"
	case SimpleExtendedMetrics:
		return "simple"
	case DetailedExtendedMetrics:
		return "detailed"
	}
	return "unknown"
}

// UnmarshalYAML unmarshals an ExtendedMetricsType into a valid type from string.
func (t *ExtendedMetricsType) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var str string
	if err := unmarshal(&str); err != nil {
		return err
	}
	if str == "" {
		*t = DefaultExtendedMetricsType
		return nil
	}
	strs := make([]string, len(validExtendedMetricsTypes))
	for _, valid := range validExtendedMetricsTypes {
		if str == valid.String() {
			*t = valid
			return nil
		}
		strs = append(strs, "'"+valid.String()+"'")
	}
	return fmt.Errorf("invalid ExtendedMetricsType '%s' valid types are: %s",
		str, strings.Join(strs, ", "))
}

// StartReportingExtendedMetrics starts reporting extended metrics.
func StartReportingExtendedMetrics(
	scope tally.Scope,
	reportInterval time.Duration,
	metricsType ExtendedMetricsType,
) *ExtendedMetricsReporter {
	reporter := NewExtendedMetricsReporter(scope, reportInterval, metricsType)
	reporter.Start()
	return reporter
}

type runtimeMetrics struct {
	NumGoRoutines   tally.Gauge
	GoMaxProcs      tally.Gauge
	MemoryAllocated tally.Gauge
	MemoryHeap      tally.Gauge
	MemoryHeapIdle  tally.Gauge
	MemoryHeapInuse tally.Gauge
	MemoryStack     tally.Gauge
	NumGC           tally.Counter
	GcPauseMs       tally.Timer
	lastNumGC       uint32
}

type processMetrics struct {
	NumFDs      tally.Gauge
	NumFDErrors tally.Counter
	pid         int
}

// ExtendedMetricsReporter reports an extended set of metrics.
// The reporter is not thread-safe.
type ExtendedMetricsReporter struct {
	reportInterval time.Duration
	metricsType    ExtendedMetricsType
	runtime        runtimeMetrics
	process        processMetrics
	started        bool
	quit           chan struct{}
}

// NewExtendedMetricsReporter creates a new extended metrics reporter.
func NewExtendedMetricsReporter(
	scope tally.Scope,
	reportInterval time.Duration,
	metricsType ExtendedMetricsType,
) *ExtendedMetricsReporter {
	r := &ExtendedMetricsReporter{
		reportInterval: reportInterval,
		metricsType:    metricsType,
		started:        false,
		quit:           make(chan struct{}),
	}
	if r.metricsType == NoExtendedMetrics {
		return r
	}

	runtimeScope := scope.SubScope("runtime")
	processScope := scope.SubScope("process")
	r.runtime.NumGoRoutines = runtimeScope.Gauge("num-goroutines")
	r.runtime.GoMaxProcs = runtimeScope.Gauge("gomaxprocs")
	r.process.NumFDs = processScope.Gauge("num-fds")
	r.process.NumFDErrors = processScope.Counter("num-fd-errors")
	r.process.pid = os.Getpid()
	if r.metricsType == SimpleExtendedMetrics {
		return r
	}

	var memstats runtime.MemStats
	runtime.ReadMemStats(&memstats)
	memoryScope := runtimeScope.SubScope("memory")
	r.runtime.MemoryAllocated = memoryScope.Gauge("allocated")
	r.runtime.MemoryHeap = memoryScope.Gauge("heap")
	r.runtime.MemoryHeapIdle = memoryScope.Gauge("heapidle")
	r.runtime.MemoryHeapInuse = memoryScope.Gauge("heapinuse")
	r.runtime.MemoryStack = memoryScope.Gauge("stack")
	r.runtime.NumGC = memoryScope.Counter("num-gc")
	r.runtime.GcPauseMs = memoryScope.Timer("gc-pause-ms")
	r.runtime.lastNumGC = memstats.NumGC
	return r
}

// Start starts the reporter thread that periodically emits metrics.
func (r *ExtendedMetricsReporter) Start() {
	if r.started {
		return
	}
	go func() {
		ticker := time.NewTicker(r.reportInterval)
		for {
			select {
			case <-ticker.C:
				r.report()
			case <-r.quit:
				ticker.Stop()
				return
			}
		}
	}()
	r.started = true
}

// Stop stops reporting of runtime metrics.
// The reporter cannot be started again after it's been stopped.
func (r *ExtendedMetricsReporter) Stop() {
	close(r.quit)
}

func (r *ExtendedMetricsReporter) report() {
	r.reportRuntimeMetrics()
	r.reportProcessMetrics()
}

func (r *ExtendedMetricsReporter) reportRuntimeMetrics() {
	if r.metricsType == NoExtendedMetrics {
		return
	}

	r.runtime.NumGoRoutines.Update(float64(runtime.NumGoroutine()))
	r.runtime.GoMaxProcs.Update(float64(runtime.GOMAXPROCS(0)))
	if r.metricsType == SimpleExtendedMetrics {
		return
	}

	var memStats runtime.MemStats
	runtime.ReadMemStats(&memStats)
	r.runtime.MemoryAllocated.Update(float64(memStats.Alloc))
	r.runtime.MemoryHeap.Update(float64(memStats.HeapAlloc))
	r.runtime.MemoryHeapIdle.Update(float64(memStats.HeapIdle))
	r.runtime.MemoryHeapInuse.Update(float64(memStats.HeapInuse))
	r.runtime.MemoryStack.Update(float64(memStats.StackInuse))

	// memStats.NumGC is a perpetually incrementing counter (unless it wraps at 2^32).
	num := memStats.NumGC
	lastNum := atomic.SwapUint32(&r.runtime.lastNumGC, num)
	if delta := num - lastNum; delta > 0 {
		r.runtime.NumGC.Inc(int64(delta))
		if delta > 255 {
			// too many GCs happened, the timestamps buffer got wrapped around. Report only the last 256.
			lastNum = num - 256
		}
		for i := lastNum; i != num; i++ {
			pause := memStats.PauseNs[i%256]
			r.runtime.GcPauseMs.Record(time.Duration(pause))
		}
	}
}

func (r *ExtendedMetricsReporter) reportProcessMetrics() {
	if r.metricsType == NoExtendedMetrics {
		return
	}
	numFDs, err := process.NumFDs(r.process.pid)
	if err == nil {
		r.process.NumFDs.Update(float64(numFDs))
	} else {
		r.process.NumFDErrors.Inc(1)
	}
}
<commit_msg>Change default level to simple and move FD metrics to moderate level (#100)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package instrument

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"sync/atomic"
	"time"

	"github.com/m3db/m3x/process"

	"github.com/uber-go/tally"
)

// ExtendedMetricsType is a type of extended metrics to report.
type ExtendedMetricsType int

const (
	// NoExtendedMetrics describes no extended metrics.
	NoExtendedMetrics ExtendedMetricsType = iota

	// SimpleExtendedMetrics describes just a simple level of extended metrics:
	// - number of active goroutines
	// - number of configured gomaxprocs
	SimpleExtendedMetrics

	// ModerateExtendedMetrics describes a moderately verbose level of extended metrics:
	// - number of active goroutines
	// - number of configured gomaxprocs
	// - number of file descriptors
	ModerateExtendedMetrics

	// DetailedExtendedMetrics describes a detailed level of extended metrics:
	// - number of active goroutines
	// - number of configured gomaxprocs
	// - number of file descriptors
	// - memory allocated running count
	// - memory used by heap
	// - memory used by heap that is idle
	// - memory used by heap that is in use
	// - memory used by stack
	// - number of garbage collections
	// - GC pause times
	DetailedExtendedMetrics

	// DefaultExtendedMetricsType is the default extended metrics level.
	DefaultExtendedMetricsType = SimpleExtendedMetrics
)

var (
	validExtendedMetricsTypes = []ExtendedMetricsType{
		NoExtendedMetrics,
		SimpleExtendedMetrics,
		ModerateExtendedMetrics,
		DetailedExtendedMetrics,
	}
)

func (t ExtendedMetricsType) String() string {
	switch t {
	case NoExtendedMetrics:
		return "none"
	case SimpleExtendedMetrics:
		return "simple"
	case ModerateExtendedMetrics:
		return "moderate"
	case DetailedExtendedMetrics:
		return "detailed"
	}
	return "unknown"
}

// UnmarshalYAML unmarshals an ExtendedMetricsType into a valid type from string.
func (t *ExtendedMetricsType) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var str string
	if err := unmarshal(&str); err != nil {
		return err
	}
	if str == "" {
		*t = DefaultExtendedMetricsType
		return nil
	}
	strs := make([]string, len(validExtendedMetricsTypes))
	for _, valid := range validExtendedMetricsTypes {
		if str == valid.String() {
			*t = valid
			return nil
		}
		strs = append(strs, "'"+valid.String()+"'")
	}
	return fmt.Errorf("invalid ExtendedMetricsType '%s' valid types are: %s",
		str, strings.Join(strs, ", "))
}

// StartReportingExtendedMetrics starts reporting extended metrics.
func StartReportingExtendedMetrics(
	scope tally.Scope,
	reportInterval time.Duration,
	metricsType ExtendedMetricsType,
) *ExtendedMetricsReporter {
	reporter := NewExtendedMetricsReporter(scope, reportInterval, metricsType)
	reporter.Start()
	return reporter
}

type runtimeMetrics struct {
	NumGoRoutines   tally.Gauge
	GoMaxProcs      tally.Gauge
	MemoryAllocated tally.Gauge
	MemoryHeap      tally.Gauge
	MemoryHeapIdle  tally.Gauge
	MemoryHeapInuse tally.Gauge
	MemoryStack     tally.Gauge
	NumGC           tally.Counter
	GcPauseMs       tally.Timer
	lastNumGC       uint32
}

type processMetrics struct {
	NumFDs      tally.Gauge
	NumFDErrors tally.Counter
	pid         int
}

// ExtendedMetricsReporter reports an extended set of metrics.
// The reporter is not thread-safe.
type ExtendedMetricsReporter struct {
	reportInterval time.Duration
	metricsType    ExtendedMetricsType
	runtime        runtimeMetrics
	process        processMetrics
	started        bool
	quit           chan struct{}
}

// NewExtendedMetricsReporter creates a new extended metrics reporter.
func NewExtendedMetricsReporter(
	scope tally.Scope,
	reportInterval time.Duration,
	metricsType ExtendedMetricsType,
) *ExtendedMetricsReporter {
	r := &ExtendedMetricsReporter{
		reportInterval: reportInterval,
		metricsType:    metricsType,
		started:        false,
		quit:           make(chan struct{}),
	}
	if r.metricsType == NoExtendedMetrics {
		return r
	}

	runtimeScope := scope.SubScope("runtime")
	processScope := scope.SubScope("process")
	r.runtime.NumGoRoutines = runtimeScope.Gauge("num-goroutines")
	r.runtime.GoMaxProcs = runtimeScope.Gauge("gomaxprocs")
	r.process.NumFDs = processScope.Gauge("num-fds")
	r.process.NumFDErrors = processScope.Counter("num-fd-errors")
	r.process.pid = os.Getpid()
	if r.metricsType == SimpleExtendedMetrics {
		return r
	}

	var memstats runtime.MemStats
	runtime.ReadMemStats(&memstats)
	memoryScope := runtimeScope.SubScope("memory")
	r.runtime.MemoryAllocated = memoryScope.Gauge("allocated")
	r.runtime.MemoryHeap = memoryScope.Gauge("heap")
	r.runtime.MemoryHeapIdle = memoryScope.Gauge("heapidle")
	r.runtime.MemoryHeapInuse = memoryScope.Gauge("heapinuse")
	r.runtime.MemoryStack = memoryScope.Gauge("stack")
	r.runtime.NumGC = memoryScope.Counter("num-gc")
	r.runtime.GcPauseMs = memoryScope.Timer("gc-pause-ms")
	r.runtime.lastNumGC = memstats.NumGC
	return r
}

// Start starts the reporter thread that periodically emits metrics.
func (r *ExtendedMetricsReporter) Start() {
	if r.started {
		return
	}
	go func() {
		ticker := time.NewTicker(r.reportInterval)
		for {
			select {
			case <-ticker.C:
				r.report()
			case <-r.quit:
				ticker.Stop()
				return
			}
		}
	}()
	r.started = true
}

// Stop stops reporting of runtime metrics.
// The reporter cannot be started again after it's been stopped.
func (r *ExtendedMetricsReporter) Stop() {
	close(r.quit)
}

func (r *ExtendedMetricsReporter) report() {
	r.reportRuntimeMetrics()
	r.reportProcessMetrics()
}

func (r *ExtendedMetricsReporter) reportRuntimeMetrics() {
	if r.metricsType == NoExtendedMetrics {
		return
	}

	r.runtime.NumGoRoutines.Update(float64(runtime.NumGoroutine()))
	r.runtime.GoMaxProcs.Update(float64(runtime.GOMAXPROCS(0)))
	if r.metricsType < DetailedExtendedMetrics {
		return
	}

	var memStats runtime.MemStats
	runtime.ReadMemStats(&memStats)
	r.runtime.MemoryAllocated.Update(float64(memStats.Alloc))
	r.runtime.MemoryHeap.Update(float64(memStats.HeapAlloc))
	r.runtime.MemoryHeapIdle.Update(float64(memStats.HeapIdle))
	r.runtime.MemoryHeapInuse.Update(float64(memStats.HeapInuse))
	r.runtime.MemoryStack.Update(float64(memStats.StackInuse))

	// memStats.NumGC is a perpetually incrementing counter (unless it wraps at 2^32).
	num := memStats.NumGC
	lastNum := atomic.SwapUint32(&r.runtime.lastNumGC, num)
	if delta := num - lastNum; delta > 0 {
		r.runtime.NumGC.Inc(int64(delta))
		if delta > 255 {
			// too many GCs happened, the timestamps buffer got wrapped around. Report only the last 256.
			lastNum = num - 256
		}
		for i := lastNum; i != num; i++ {
			pause := memStats.PauseNs[i%256]
			r.runtime.GcPauseMs.Record(time.Duration(pause))
		}
	}
}

func (r *ExtendedMetricsReporter) reportProcessMetrics() {
	if r.metricsType < ModerateExtendedMetrics {
		return
	}
	numFDs, err := process.NumFDs(r.process.pid)
	if err == nil {
		r.process.NumFDs.Update(float64(numFDs))
	} else {
		r.process.NumFDErrors.Inc(1)
	}
}
<|endoftext|>


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

package main

import (
	"fmt"
	"log"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"github.com/minio-io/mc/pkg/s3"
)

const (
	recvFormat  = "2006-01-02T15:04:05.000Z"
	printFormat = "2006-01-02 15:04:05"
)

const (
	delimiter = "/"
)

func parseTime(t string) string {
	ti, _ := time.Parse(recvFormat, t)
	return ti.Format(printFormat)
}

func parseLastModified(t string) string {
	ti, _ := time.Parse(time.RFC1123, t)
	return ti.Format(printFormat)
}

func printBuckets(v []*s3.Bucket) {
	for _, b := range v {
		fmt.Printf("%s %s\n", parseTime(b.CreationDate), b.Name)
	}
}

func printObjects(v []*s3.Item) {
	sort.Sort(s3.BySize(v))
	for _, b := range v {
		fmt.Printf("%s   %d %s\n", parseTime(b.LastModified), b.Size, b.Key)
	}
}

func printPrefixes(v []*s3.Prefix) {
	for _, b := range v {
		fmt.Printf("                      PRE %s\n", b.Prefix)
	}
}

func printObject(v int64, date, key string) {
	fmt.Printf("%s   %d %s\n", parseLastModified(date), v, key)
}

func getBucketAndObject(p string) (bucket, object string) {
	i := strings.Index(p, delimiter)
	if i == -1 {
		bucket = p
	} else if i != -1 {
		bucket = p[:i]
		object = path.Base(p)
		if bucket == object {
			object = ""
		}
	}
	return
}

func doFsList(c *cli.Context) {
	var err error
	var auth *s3.Auth
	var items []*s3.Item
	var prefixes []*s3.Prefix

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

	var buckets []*s3.Bucket
	s3c := s3.NewS3Client(auth)
	switch len(c.Args()) {
	case 1:
		input := c.Args().Get(0)
		if path.IsAbs(input) {
			log.Fatal("Invalid bucket style")
		}
		bucket, object := getBucketAndObject(input)
		if object == "" {
			items, prefixes, err = s3c.GetBucket(bucket, "", "", "/", s3.MAX_OBJECT_LIST)
			if err != nil {
				log.Fatal(err)
			}
			printPrefixes(prefixes)
			printObjects(items)
		} else {
			var date string
			var size int64
			size, date, err = s3c.Stat(object, bucket)
			if err != nil {
				log.Fatal(err)
			}
			printObject(size, date, object)
		}
	default:
		buckets, err = s3c.Buckets()
		if err != nil {
			log.Fatal(err)
		}
		printBuckets(buckets)
	}
}
<commit_msg>if s3.Stat() fails, send a GetBucket with prefix set<commit_after>/*
 * Mini Object Storage, (C) 2014,2015 Minio, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"fmt"
	"log"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"github.com/minio-io/mc/pkg/s3"
)

const (
	recvFormat  = "2006-01-02T15:04:05.000Z"
	printFormat = "2006-01-02 15:04:05"
)

const (
	delimiter = "/"
)

func parseTime(t string) string {
	ti, _ := time.Parse(recvFormat, t)
	return ti.Format(printFormat)
}

func parseLastModified(t string) string {
	ti, _ := time.Parse(time.RFC1123, t)
	return ti.Format(printFormat)
}

func printBuckets(v []*s3.Bucket) {
	for _, b := range v {
		fmt.Printf("%s   %s\n", parseTime(b.CreationDate), b.Name)
	}
}

func printObjects(v []*s3.Item) {
	if len(v) > 0 {
		sort.Sort(s3.BySize(v))
		for _, b := range v {
			fmt.Printf("%s   %d %s\n", parseTime(b.LastModified), b.Size, b.Key)
		}
	}
}

func printPrefixes(v []*s3.Prefix) {
	if len(v) > 0 {
		for _, b := range v {
			fmt.Printf("                      PRE %s\n", b.Prefix)
		}
	}
}

func printObject(v int64, date, key string) {
	fmt.Printf("%s   %d %s\n", parseLastModified(date), v, key)
}

func getBucketAndObject(p string) (bucket, object string) {
	i := strings.Index(p, delimiter)
	if i == -1 {
		bucket = p
	} else if i != -1 {
		bucket = p[:i]
		object = path.Base(p)
		if bucket == object {
			object = ""
		}
	}
	return
}

func doFsList(c *cli.Context) {
	var err error
	var auth *s3.Auth
	var items []*s3.Item
	var prefixes []*s3.Prefix

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

	var buckets []*s3.Bucket
	s3c := s3.NewS3Client(auth)
	switch len(c.Args()) {
	case 1:
		input := c.Args().Get(0)
		if path.IsAbs(input) {
			log.Fatal("Invalid bucket style")
		}
		bucket, object := getBucketAndObject(input)
		if object == "" {
			items, prefixes, err = s3c.GetBucket(bucket, "", "", "/", s3.MAX_OBJECT_LIST)
			if err != nil {
				log.Fatal(err)
			}
			printPrefixes(prefixes)
			printObjects(items)
		} else {
			var date string
			var size int64
			size, date, err = s3c.Stat(object, bucket)
			if err != nil {
				items, prefixes, err = s3c.GetBucket(bucket, "", object, "/", s3.MAX_OBJECT_LIST)
				if err != nil {
					log.Fatal(err)
				}
				printPrefixes(prefixes)
				printObjects(items)
			} else {
				printObject(size, date, object)
			}
		}
	default:
		buckets, err = s3c.Buckets()
		if err != nil {
			log.Fatal(err)
		}
		printBuckets(buckets)
	}
}
<|endoftext|>


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

package diff

import (
	"fmt"
	"io"

	"cuelang.org/go/cue/errors"
	"cuelang.org/go/cue/format"
	"cuelang.org/go/internal/legacy/cue"
)

// Print the differences between two structs represented by an edit script.
func Print(w io.Writer, es *EditScript) error {
	p := printer{
		w:       w,
		margin:  2,
		context: 2,
	}
	p.script(es)
	return p.errs
}

type printer struct {
	w         io.Writer
	context   int
	margin    int
	indent    int
	prefix    string
	hasPrefix bool
	hasPrint  bool
	errs      errors.Error
}

func (p *printer) writeRaw(b []byte) {
	if len(b) == 0 {
		return
	}
	if !p.hasPrefix {
		io.WriteString(p.w, p.prefix)
		p.hasPrefix = true
	}
	if !p.hasPrint {
		fmt.Fprintf(p.w, "% [1]*s", p.indent+p.margin-len(p.prefix), "")
		p.hasPrint = true
	}
	p.w.Write(b)
}

func (p *printer) Write(b []byte) (n int, err error) {
	i, last := 0, 0
	for ; i < len(b); i++ {
		if b[i] != '\n' {
			continue
		}
		p.writeRaw(b[last:i])
		last = i + 1
		io.WriteString(p.w, "\n")
		p.hasPrefix = false
		p.hasPrint = false
	}
	p.writeRaw(b[last:])
	return len(b), nil
}

func (p *printer) write(b []byte) {
	_, _ = p.Write(b)
}

func (p *printer) printLen(align int, str string) {
	fmt.Fprintf(p, "% -[1]*s", align, str)
}

func (p *printer) println(s string) {
	fmt.Fprintln(p, s)
}

func (p *printer) printf(format string, args ...interface{}) {
	fmt.Fprintf(p, format, args...)
}

func (p *printer) script(e *EditScript) {
	switch e.x.Kind() {
	case cue.StructKind:
		p.printStruct(e)
	case cue.ListKind:
		p.printList(e)
	default:
		p.printf("BadExpr")
	}
}

func (p *printer) findRun(es *EditScript, i int) (start, end int) {
	lastEnd := i

	for ; i < es.Len() && es.edits[i].kind == Identity; i++ {
	}
	start = i

	// Find end of run
	include := p.context
	for ; i < es.Len(); i++ {
		e := es.edits[i]
		if e.kind != Identity {
			include = p.context + 1
			continue
		}
		if include--; include == 0 {
			break
		}
	}

	if i-start > 0 {
		// Adjust start of run
		if s := start - p.context; s > lastEnd {
			start = s
		} else {
			start = lastEnd
		}
	}
	return start, i
}

func (p *printer) printStruct(es *EditScript) {
	// TODO: consider not printing outer curlies, or make it an option.
	// if p.indent > 0 {
	p.println("{")
	defer p.println("}")
	// }
	p.indent += 4
	defer func() {
		p.indent -= 4
	}()

	var start, i int
	for i < es.Len() {
		lastEnd := i
		// Find provisional start of run.
		start, i = p.findRun(es, i)

		p.printSkipped(start - lastEnd)
		p.printFieldRun(es, start, i)
	}
	p.printSkipped(es.Len() - i)
}

func (p *printer) printList(es *EditScript) {
	p.println("[")
	p.indent += 4
	defer func() {
		p.indent -= 4
		p.println("]")
	}()

	x := getElems(es.x)
	y := getElems(es.y)

	var start, i int
	for i < es.Len() {
		lastEnd := i
		// Find provisional start of run.
		start, i = p.findRun(es, i)

		p.printSkipped(start - lastEnd)
		p.printElemRun(es, x, y, start, i)
	}
	p.printSkipped(es.Len() - i)
}

func getElems(x cue.Value) (a []cue.Value) {
	for i, _ := x.List(); i.Next(); {
		a = append(a, i.Value())
	}
	return a
}

func (p *printer) printSkipped(n int) {
	if n > 0 {
		p.printf("... // %d identical elements\n", n)
	}
}

func (p *printer) printValue(v cue.Value) {
	// TODO: have indent option.
	b, _ := format.Node(v.Syntax())
	p.write(b)
}

func (p *printer) printFieldRun(es *EditScript, start, end int) {
	// Determine max field len.
	for i := start; i < end; i++ {
		e := es.edits[i]

		switch e.kind {
		case UniqueX:
			p.printField("-", es, es.LabelX(i), es.ValueX(i))

		case UniqueY:
			p.printField("+", es, es.LabelY(i), es.ValueY(i))

		case Modified:
			if e.sub != nil {
				io.WriteString(p, es.LabelX(i))
				io.WriteString(p, " ")
				p.script(e.sub)
				break
			}
			// TODO: show per-line differences for multiline strings.
			p.printField("-", es, es.LabelX(i), es.ValueX(i))
			p.printField("+", es, es.LabelY(i), es.ValueY(i))

		case Identity:
			// TODO: write on one line
			p.printField("", es, es.LabelX(i), es.ValueX(i))
		}
	}
}

func (p *printer) printField(prefix string, es *EditScript, label string, v cue.Value) {
	p.prefix = prefix
	io.WriteString(p, label)
	io.WriteString(p, " ")
	p.printValue(v)
	io.WriteString(p, "\n")
	p.prefix = ""
}

func (p *printer) printElemRun(es *EditScript, x, y []cue.Value, start, end int) {
	for _, e := range es.edits[start:end] {
		switch e.kind {
		case UniqueX:
			p.printElem("-", x[e.XPos()])

		case UniqueY:
			p.printElem("+", y[e.YPos()])

		case Modified:
			if e.sub != nil {
				p.script(e.sub)
				break
			}
			// TODO: show per-line differences for multiline strings.
			p.printElem("-", x[e.XPos()])
			p.printElem("+", y[e.YPos()])

		case Identity:
			// TODO: write on one line
			p.printElem("", x[e.XPos()])
		}
	}
}

func (p *printer) printElem(prefix string, v cue.Value) {
	p.prefix = prefix
	p.printValue(v)
	io.WriteString(p, ",\n")
	p.prefix = ""
}
<commit_msg>internal/diff: use compact printing mode<commit_after>// Copyright 2019 CUE Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package diff

import (
	"fmt"
	"io"

	"cuelang.org/go/cue/errors"
	"cuelang.org/go/internal/legacy/cue"
)

// Print the differences between two structs represented by an edit script.
func Print(w io.Writer, es *EditScript) error {
	p := printer{
		w:       w,
		margin:  2,
		context: 2,
	}
	p.script(es)
	return p.errs
}

type printer struct {
	w         io.Writer
	context   int
	margin    int
	indent    int
	prefix    string
	hasPrefix bool
	hasPrint  bool
	errs      errors.Error
}

func (p *printer) writeRaw(b []byte) {
	if len(b) == 0 {
		return
	}
	if !p.hasPrefix {
		io.WriteString(p.w, p.prefix)
		p.hasPrefix = true
	}
	if !p.hasPrint {
		fmt.Fprintf(p.w, "% [1]*s", p.indent+p.margin-len(p.prefix), "")
		p.hasPrint = true
	}
	p.w.Write(b)
}

func (p *printer) Write(b []byte) (n int, err error) {
	i, last := 0, 0
	for ; i < len(b); i++ {
		if b[i] != '\n' {
			continue
		}
		p.writeRaw(b[last:i])
		last = i + 1
		io.WriteString(p.w, "\n")
		p.hasPrefix = false
		p.hasPrint = false
	}
	p.writeRaw(b[last:])
	return len(b), nil
}

func (p *printer) write(b []byte) {
	_, _ = p.Write(b)
}

func (p *printer) printLen(align int, str string) {
	fmt.Fprintf(p, "% -[1]*s", align, str)
}

func (p *printer) println(s string) {
	fmt.Fprintln(p, s)
}

func (p *printer) printf(format string, args ...interface{}) {
	fmt.Fprintf(p, format, args...)
}

func (p *printer) script(e *EditScript) {
	switch e.x.Kind() {
	case cue.StructKind:
		p.printStruct(e)
	case cue.ListKind:
		p.printList(e)
	default:
		p.printf("BadExpr")
	}
}

func (p *printer) findRun(es *EditScript, i int) (start, end int) {
	lastEnd := i

	for ; i < es.Len() && es.edits[i].kind == Identity; i++ {
	}
	start = i

	// Find end of run
	include := p.context
	for ; i < es.Len(); i++ {
		e := es.edits[i]
		if e.kind != Identity {
			include = p.context + 1
			continue
		}
		if include--; include == 0 {
			break
		}
	}

	if i-start > 0 {
		// Adjust start of run
		if s := start - p.context; s > lastEnd {
			start = s
		} else {
			start = lastEnd
		}
	}
	return start, i
}

func (p *printer) printStruct(es *EditScript) {
	// TODO: consider not printing outer curlies, or make it an option.
	// if p.indent > 0 {
	p.println("{")
	defer p.println("}")
	// }
	p.indent += 4
	defer func() {
		p.indent -= 4
	}()

	var start, i int
	for i < es.Len() {
		lastEnd := i
		// Find provisional start of run.
		start, i = p.findRun(es, i)

		p.printSkipped(start - lastEnd)
		p.printFieldRun(es, start, i)
	}
	p.printSkipped(es.Len() - i)
}

func (p *printer) printList(es *EditScript) {
	p.println("[")
	p.indent += 4
	defer func() {
		p.indent -= 4
		p.println("]")
	}()

	x := getElems(es.x)
	y := getElems(es.y)

	var start, i int
	for i < es.Len() {
		lastEnd := i
		// Find provisional start of run.
		start, i = p.findRun(es, i)

		p.printSkipped(start - lastEnd)
		p.printElemRun(es, x, y, start, i)
	}
	p.printSkipped(es.Len() - i)
}

func getElems(x cue.Value) (a []cue.Value) {
	for i, _ := x.List(); i.Next(); {
		a = append(a, i.Value())
	}
	return a
}

func (p *printer) printSkipped(n int) {
	if n > 0 {
		p.printf("... // %d identical elements\n", n)
	}
}

func (p *printer) printValue(v cue.Value) {
	// TODO: have indent option.
	s := fmt.Sprintf("%-v", v)
	io.WriteString(p, s)
}

func (p *printer) printFieldRun(es *EditScript, start, end int) {
	// Determine max field len.
	for i := start; i < end; i++ {
		e := es.edits[i]

		switch e.kind {
		case UniqueX:
			p.printField("-", es, es.LabelX(i), es.ValueX(i))

		case UniqueY:
			p.printField("+", es, es.LabelY(i), es.ValueY(i))

		case Modified:
			if e.sub != nil {
				io.WriteString(p, es.LabelX(i))
				io.WriteString(p, " ")
				p.script(e.sub)
				break
			}
			// TODO: show per-line differences for multiline strings.
			p.printField("-", es, es.LabelX(i), es.ValueX(i))
			p.printField("+", es, es.LabelY(i), es.ValueY(i))

		case Identity:
			// TODO: write on one line
			p.printField("", es, es.LabelX(i), es.ValueX(i))
		}
	}
}

func (p *printer) printField(prefix string, es *EditScript, label string, v cue.Value) {
	p.prefix = prefix
	io.WriteString(p, label)
	io.WriteString(p, " ")
	p.printValue(v)
	io.WriteString(p, "\n")
	p.prefix = ""
}

func (p *printer) printElemRun(es *EditScript, x, y []cue.Value, start, end int) {
	for _, e := range es.edits[start:end] {
		switch e.kind {
		case UniqueX:
			p.printElem("-", x[e.XPos()])

		case UniqueY:
			p.printElem("+", y[e.YPos()])

		case Modified:
			if e.sub != nil {
				p.script(e.sub)
				break
			}
			// TODO: show per-line differences for multiline strings.
			p.printElem("-", x[e.XPos()])
			p.printElem("+", y[e.YPos()])

		case Identity:
			// TODO: write on one line
			p.printElem("", x[e.XPos()])
		}
	}
}

func (p *printer) printElem(prefix string, v cue.Value) {
	p.prefix = prefix
	p.printValue(v)
	io.WriteString(p, ",\n")
	p.prefix = ""
}
<|endoftext|>


<commit_before>// Copyright 2013 Dmitry Chestnykh. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package safefile implements safe "atomic" saving of files.
//
// Instead of truncating and overwriting the destination file, it creates a
// temporary file in the same directory, writes to it, and then renames the
// temporary file to the original name when calling Commit.
//
// Example:
//
//  f, err := safefile.Create("/home/ken/report.txt")
//  if err != nil {
//  	// ...
//  }
//  // Created temporary file /home/ken/133a7876287381fa-0.tmp
//
//  defer f.Close()
//
//  _, err = io.WriteString(f, "Hello world")
//  if err != nil {
//  	// ...
//  }
//  // Wrote "Hello world" to /home/ken/133a7876287381fa-0.tmp
//
//  err = f.Commit()
//  if err != nil {
//      // ...
//  }
//  // Renamed /home/ken/133a7876287381fa-0.tmp to /home/ken/report.txt
//
package safefile

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"
)

// ErrAlreadyCommitted error is returned by when calling Commit on a file that
// has been already successfully committed.
var ErrAlreadyCommitted = errors.New("file already committed")

type File struct {
	*os.File
	origName    string
	closeFunc   func(*File) error
	isCommitted bool
}

func makeTempName(origname string, counter int) (tempname string, err error) {
	origname = filepath.Clean(origname)
	if len(origname) == 0 || origname[len(origname)-1] == filepath.Separator {
		return "", os.ErrInvalid
	}
	return filepath.Join(filepath.Dir(origname), fmt.Sprintf("%x-%d.tmp", time.Now().UnixNano(), counter)), nil
}

// Create creates a temporary file in the same directory as filename,
// which will be renamed to the given filename when calling Commit.
func Create(filename string, perm os.FileMode) (*File, error) {
	counter := 0
	for {
		tempname, err := makeTempName(filename, counter)
		if err != nil {
			return nil, err
		}
		f, err := os.OpenFile(tempname, os.O_RDWR|os.O_CREATE, perm)
		if err != nil {
			if os.IsExist(err) {
				counter++
				continue
			}
			return nil, err
		}
		return &File{
			File:      f,
			origName:  filename,
			closeFunc: closeUncommitted,
		}, nil
	}
}

// OrigName returns the original filename given to Create.
func (f *File) OrigName() string {
	return f.origName
}

// Close closes temporary file and removes it.
// If the file has been committed, Close is no-op.
func (f *File) Close() error {
	return f.closeFunc(f)
}

func closeUncommitted(f *File) error {
	err0 := f.File.Close()
	err1 := os.Remove(f.Name())
	f.closeFunc = closeAgainError
	if err0 != nil {
		return err0
	}
	return err1
}

func closeAfterFailedRename(f *File) error {
	// just remove temporary file.
	f.closeFunc = closeAgainError
	return os.Remove(f.Name())
}

func closeCommitted(f *File) error {
	// noop
	return nil
}

func closeAgainError(f *File) error {
	return os.ErrInvalid
}

// Commit safely commits data into the original file by syncing temporary
// file to disk, closing it and renaming to the original file name.
//
// In case of success, the temporary file is closed and no longer exists
// on disk. It is safe to call Close after Commit: the operation will do
// nothing.
//
// In case of error, the temporary file is still opened and exists on disk;
// it must be closed by callers by calling Close or by trying to commit again.
func (f *File) Commit() error {
	if f.isCommitted {
		return ErrAlreadyCommitted
	}
	// Sync to disk.
	err := f.Sync()
	if err != nil {
		return err
	}
	// Close underlying os.File.
	err = f.File.Close()
	if err != nil {
		return err
	}
	// Rename.
	err = os.Rename(f.Name(), f.origName)
	if err != nil {
		f.closeFunc = closeAfterFailedRename
		return err
	}
	f.closeFunc = closeCommitted
	f.isCommitted = true
	return nil
}

// WriteFile is a safe analog of ioutil.WriteFile.
func WriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := Create(filename, perm)
	if err != nil {
		return err
	}
	defer f.Close()
	n, err := f.Write(data)
	if err != nil {
		return err
	}
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
		return err
	}
	return f.Commit()
}
<commit_msg>Fix typo.<commit_after>// Copyright 2013 Dmitry Chestnykh. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package safefile implements safe "atomic" saving of files.
//
// Instead of truncating and overwriting the destination file, it creates a
// temporary file in the same directory, writes to it, and then renames the
// temporary file to the original name when calling Commit.
//
// Example:
//
//  f, err := safefile.Create("/home/ken/report.txt")
//  if err != nil {
//  	// ...
//  }
//  // Created temporary file /home/ken/133a7876287381fa-0.tmp
//
//  defer f.Close()
//
//  _, err = io.WriteString(f, "Hello world")
//  if err != nil {
//  	// ...
//  }
//  // Wrote "Hello world" to /home/ken/133a7876287381fa-0.tmp
//
//  err = f.Commit()
//  if err != nil {
//      // ...
//  }
//  // Renamed /home/ken/133a7876287381fa-0.tmp to /home/ken/report.txt
//
package safefile

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"
)

// ErrAlreadyCommitted error is returned when calling Commit on a file that
// has been already successfully committed.
var ErrAlreadyCommitted = errors.New("file already committed")

type File struct {
	*os.File
	origName    string
	closeFunc   func(*File) error
	isCommitted bool
}

func makeTempName(origname string, counter int) (tempname string, err error) {
	origname = filepath.Clean(origname)
	if len(origname) == 0 || origname[len(origname)-1] == filepath.Separator {
		return "", os.ErrInvalid
	}
	return filepath.Join(filepath.Dir(origname), fmt.Sprintf("%x-%d.tmp", time.Now().UnixNano(), counter)), nil
}

// Create creates a temporary file in the same directory as filename,
// which will be renamed to the given filename when calling Commit.
func Create(filename string, perm os.FileMode) (*File, error) {
	counter := 0
	for {
		tempname, err := makeTempName(filename, counter)
		if err != nil {
			return nil, err
		}
		f, err := os.OpenFile(tempname, os.O_RDWR|os.O_CREATE, perm)
		if err != nil {
			if os.IsExist(err) {
				counter++
				continue
			}
			return nil, err
		}
		return &File{
			File:      f,
			origName:  filename,
			closeFunc: closeUncommitted,
		}, nil
	}
}

// OrigName returns the original filename given to Create.
func (f *File) OrigName() string {
	return f.origName
}

// Close closes temporary file and removes it.
// If the file has been committed, Close is no-op.
func (f *File) Close() error {
	return f.closeFunc(f)
}

func closeUncommitted(f *File) error {
	err0 := f.File.Close()
	err1 := os.Remove(f.Name())
	f.closeFunc = closeAgainError
	if err0 != nil {
		return err0
	}
	return err1
}

func closeAfterFailedRename(f *File) error {
	// just remove temporary file.
	f.closeFunc = closeAgainError
	return os.Remove(f.Name())
}

func closeCommitted(f *File) error {
	// noop
	return nil
}

func closeAgainError(f *File) error {
	return os.ErrInvalid
}

// Commit safely commits data into the original file by syncing temporary
// file to disk, closing it and renaming to the original file name.
//
// In case of success, the temporary file is closed and no longer exists
// on disk. It is safe to call Close after Commit: the operation will do
// nothing.
//
// In case of error, the temporary file is still opened and exists on disk;
// it must be closed by callers by calling Close or by trying to commit again.
func (f *File) Commit() error {
	if f.isCommitted {
		return ErrAlreadyCommitted
	}
	// Sync to disk.
	err := f.Sync()
	if err != nil {
		return err
	}
	// Close underlying os.File.
	err = f.File.Close()
	if err != nil {
		return err
	}
	// Rename.
	err = os.Rename(f.Name(), f.origName)
	if err != nil {
		f.closeFunc = closeAfterFailedRename
		return err
	}
	f.closeFunc = closeCommitted
	f.isCommitted = true
	return nil
}

// WriteFile is a safe analog of ioutil.WriteFile.
func WriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := Create(filename, perm)
	if err != nil {
		return err
	}
	defer f.Close()
	n, err := f.Write(data)
	if err != nil {
		return err
	}
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
		return err
	}
	return f.Commit()
}
<|endoftext|>


<commit_before>/*
Copyright (C) 2013 Lucy

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.package main
*/

package main

import (
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"math/cmplx"
	"os"

	"github.com/jackvalmadre/go-fftw"
	flag "github.com/neeee/pflag"
	"github.com/neeee/termbox-go"
)

var (
	color = flag.StringP("color", "c", "blue", "Color to use")
	dim   = flag.BoolP("dim", "d", false, "Turn off bold")

	step   = flag.Int("step", 2, "Samples to average in each column (wave)")
	scale  = flag.Float64("scale", 2, "Scale divisor (spectrum)")
	icolor = flag.BoolP("intensitycolor", "i", false,
		"color bars based on intensity (spectrum)")
	imode = flag.String("imode", "dumb",
		"intensitycolor mode (dumb, 256 or grayscale)")

	filename = flag.StringP("file", "f", "/tmp/mpd.fifo",
		"Where to read fifo output from")
	vis = flag.StringP("viz", "v", "wave",
		"Visualization (spectrum or wave)")
)

var colors = map[string]termbox.Attribute{
	"default": termbox.ColorDefault,
	"black":   termbox.ColorBlack,
	"red":     termbox.ColorRed,
	"green":   termbox.ColorGreen,
	"yellow":  termbox.ColorYellow,
	"blue":    termbox.ColorBlue,
	"magenta": termbox.ColorMagenta,
	"cyan":    termbox.ColorCyan,
	"white":   termbox.ColorWhite,
}

var iColors []termbox.Attribute

var (
	on  = termbox.ColorDefault
	off = termbox.ColorDefault
)

func warn(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args...)
}

func main() {
	flag.Parse()
	var ok bool
	on, ok = colors[*color]
	if !ok {
		warn("Unknown color \"%s\"\n", *color)
		return
	}
	if !*dim {
		on = on | termbox.AttrBold
	}
	switch *imode {
	case "dumb":
		iColors = []termbox.Attribute{
			8 + termbox.ColorBlue,
			8 + termbox.ColorCyan,
			8 + termbox.ColorGreen,
			8 + termbox.ColorYellow,
			8 + termbox.ColorRed,
		}
	case "256":
		iColors = []termbox.Attribute{
			21, 27, 39, 45, 51, 86, 85, 84, 82,
			154, 192, 220, 214, 208, 202, 196,
		}
	case "grayscale":
		const num = 19
		iColors = make([]termbox.Attribute, num)
		for i := termbox.Attribute(0); i < num; i++ {
			iColors[i] = i + 255 - num
		}
	default:
		warn("Unsupported mode: \"%s\"\n", *imode)
		return
	}

	var draw func(chan int16)
	switch *vis {
	case "spectrum":
		draw = drawSpectrum
	case "wave":
		draw = drawWave
	default:
		warn("Unknown visualisation \"%s\"\n"+
			"Supported: spectrum, wave\n", *vis)
		return
	}

	file, err := os.Open(*filename)
	if err != nil {
		warn("%s\n", err)
		return
	}
	defer file.Close()

	err = termbox.Init()
	if err != nil {
		warn("%s\b", err)
		return
	}
	defer termbox.Close()

	ch := make(chan int16, 128)
	end := make(chan string)

	// drawer
	go draw(ch)

	// input handler
	go func() {
		for {
			ev := termbox.PollEvent()
			if ev.Ch == 0 && ev.Key == termbox.KeyCtrlC {
				close(end)
				return
			}
		}
	}()

	// file reader
	go func() {
		var i int16
		for binary.Read(file, binary.LittleEndian, &i) != io.EOF {
			ch <- i
		}
		close(end)
	}()

	<-end
}

func size() (int, int) {
	w, h := termbox.Size()
	return w, h * 2
}
func drawWave(c chan int16) {
	w, h := size()
	for pos := 0; ; pos++ {
		if pos >= w {
			pos = 0
			w, h = size()
			termbox.Flush()
			termbox.Clear(off, off)
		}

		var v float64
		for i := 0; i < *step; i++ {
			v += float64(<-c)
		}

		half_h := float64(h / 2)
		vi := int(v/float64(*step)/(math.MaxInt16/half_h) + half_h)
		if vi%2 == 0 {
			termbox.SetCell(pos, vi/2, '▀', on, off)
		} else {
			termbox.SetCell(pos, vi/2, '▄', on, off)
		}
	}
}

func drawSpectrum(c chan int16) {
	w, h := size()
	var (
		samples = (w - 1) * 2
		resn    = w
		in      = make([]float64, samples)
		out     = fftw.Alloc1d(resn)
		plan    = fftw.PlanDftR2C1d(in, out, fftw.Measure)
	)

	for {
		if resn != w && w != 1 {
			fftw.Free1d(out)
			resn = w
			samples = (w - 1) * 2
			in = make([]float64, samples)
			out = fftw.Alloc1d(resn)
			plan = fftw.PlanDftR2C1d(in, out, fftw.Measure)
		}

		for i := 0; i < samples; i++ {
			in[i] = float64(<-c)
		}

		plan.Execute()
		hf := float64(h)
		for i := 0; i < w; i++ {
			v := cmplx.Abs(out[i]) / 1e5 * hf / *scale
			vi := int(v)
			if *icolor {
				on = iColors[int(math.Min(float64(len(iColors)-1),
					(v/hf)*float64(len(iColors)-1)))]
			}
			for j := h - 1; j > h-vi; j-- {
				termbox.SetCell(i, j/2, '┃', on, off)
			}
			if vi%2 != 0 {
				termbox.SetCell(i, (h-vi)/2, '╻', on, off)
			}
		}

		termbox.Flush()
		termbox.Clear(off, off)
		w, h = size()
	}
}
<commit_msg>more refactoring<commit_after>/*
Copyright (C) 2013 Lucy

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.package main
*/

package main

import (
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"math/cmplx"
	"os"

	"github.com/jackvalmadre/go-fftw"
	flag "github.com/neeee/pflag"
	"github.com/neeee/termbox-go"
)

var (
	color = flag.StringP("color", "c", "default", "Color to use")
	dim   = flag.BoolP("dim", "d", false,
		"Turn off bright colors where possible")

	step  = flag.Int("step", 2, "Samples to average in each column (wave)")
	scale = flag.Float64("scale", 2, "Scale divisor (spectrum)")

	icolor = flag.BoolP("icolor", "i", false,
		"Color bars according to intensity (spectrum)")
	imode = flag.String("imode", "dumb",
		"Mode for intensity colorisation (dumb, 256 or grayscale)")

	filename = flag.StringP("file", "f", "/tmp/mpd.fifo",
		"Where to read pcm date from")
	vis = flag.StringP("viz", "v", "wave",
		"Visualisation (spectrum or wave)")
)

var colors = map[string]termbox.Attribute{
	"default": termbox.ColorDefault,
	"black":   termbox.ColorBlack,
	"red":     termbox.ColorRed,
	"green":   termbox.ColorGreen,
	"yellow":  termbox.ColorYellow,
	"blue":    termbox.ColorBlue,
	"magenta": termbox.ColorMagenta,
	"cyan":    termbox.ColorCyan,
	"white":   termbox.ColorWhite,
}

var iColors []termbox.Attribute

var (
	on  = termbox.ColorDefault
	off = termbox.ColorDefault
)

func warn(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format, args...)
}

func main() {
	flag.Parse()

	if cl, ok := colors[*color]; !ok {
		warn("Unknown color \"%s\"\n", *color)
		return
	} else {
		on = cl
	}

	if !*dim {
		on = on | termbox.AttrBold
	}

	switch *imode {
	case "dumb":
		iColors = []termbox.Attribute{
			termbox.ColorBlue,
			termbox.ColorCyan,
			termbox.ColorGreen,
			termbox.ColorYellow,
			termbox.ColorRed,
		}
		if !*dim {
			for i := range iColors {
				iColors[i] = iColors[i] + 8
			}
		}
	case "256":
		iColors = []termbox.Attribute{
			21, 27, 39, 45, 51, 86, 85, 84, 82,
			154, 192, 220, 214, 208, 202, 196,
		}
	case "grayscale":
		const num = 19
		iColors = make([]termbox.Attribute, num)
		for i := termbox.Attribute(0); i < num; i++ {
			iColors[i] = i + 255 - num
		}
	default:
		warn("Unsupported mode: \"%s\"\n", *imode)
		return
	}

	var draw func(chan int16)
	switch *vis {
	case "spectrum":
		draw = drawSpectrum
	case "wave":
		draw = drawWave
	default:
		warn("Unknown visualisation \"%s\"\n"+
			"Supported: spectrum, wave\n", *vis)
		return
	}

	file, err := os.Open(*filename)
	if err != nil {
		warn("%s\n", err)
		return
	}
	defer file.Close()

	err = termbox.Init()
	if err != nil {
		warn("%s\b", err)
		return
	}
	defer termbox.Close()

	var (
		ch  = make(chan int16, 128)
		end = make(chan string)
	)

	// drawer
	go draw(ch)

	// input handler
	go func() {
		for {
			ev := termbox.PollEvent()
			if ev.Ch == 0 && ev.Key == termbox.KeyCtrlC {
				close(end)
				return
			}
		}
	}()

	// file reader
	go func() {
		var i int16
		for binary.Read(file, binary.LittleEndian, &i) != io.EOF {
			ch <- i
		}
		close(end)
	}()

	<-end
}

func size() (int, int) {
	w, h := termbox.Size()
	return w, h * 2
}
func drawWave(c chan int16) {
	w, h := size()
	for pos := 0; ; pos++ {
		if pos >= w {
			pos = 0
			w, h = size()
			termbox.Flush()
			termbox.Clear(off, off)
		}

		var v float64
		for i := 0; i < *step; i++ {
			v += float64(<-c)
		}

		half_h := float64(h / 2)
		vi := int(v/float64(*step)/(math.MaxInt16/half_h) + half_h)
		if vi%2 == 0 {
			termbox.SetCell(pos, vi/2, '▀', on, off)
		} else {
			termbox.SetCell(pos, vi/2, '▄', on, off)
		}
	}
}

func drawSpectrum(c chan int16) {
	w, h := size()
	var (
		samples = (w - 1) * 2
		resn    = w
		in      = make([]float64, samples)
		out     = fftw.Alloc1d(resn)
		plan    = fftw.PlanDftR2C1d(in, out, fftw.Measure)
	)

	for {
		if resn != w && w != 1 {
			fftw.Free1d(out)
			resn = w
			samples = (w - 1) * 2
			in = make([]float64, samples)
			out = fftw.Alloc1d(resn)
			plan = fftw.PlanDftR2C1d(in, out, fftw.Measure)
		}

		for i := 0; i < samples; i++ {
			in[i] = float64(<-c)
		}

		plan.Execute()
		hf := float64(h)
		for i := 0; i < w; i++ {
			v := cmplx.Abs(out[i]) / 1e5 * hf / *scale
			vi := int(v)
			if *icolor {
				on = iColors[int(math.Min(float64(len(iColors)-1),
					(v/hf)*float64(len(iColors)-1)))]
			}
			for j := h - 1; j > h-vi; j-- {
				termbox.SetCell(i, j/2, '┃', on, off)
			}
			if vi%2 != 0 {
				termbox.SetCell(i, (h-vi)/2, '╻', on, off)
			}
		}

		termbox.Flush()
		termbox.Clear(off, off)
		w, h = size()
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// spoof contains logic to make polling HTTP requests against an endpoint with optional host spoofing.

package spoof

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"time"

	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"knative.dev/pkg/test/ingress"
	"knative.dev/pkg/test/logging"
	"knative.dev/pkg/test/zipkin"
	"knative.dev/pkg/tracing/propagation/tracecontextb3"

	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/trace"
)

// Response is a stripped down subset of http.Response. The is primarily useful
// for ResponseCheckers to inspect the response body without consuming it.
// Notably, Body is a byte slice instead of an io.ReadCloser.
type Response struct {
	Status     string
	StatusCode int
	Header     http.Header
	Body       []byte
}

func (r *Response) String() string {
	return fmt.Sprintf("status: %d, body: %s, headers: %v", r.StatusCode, string(r.Body), r.Header)
}

// https://medium.com/stupid-gopher-tricks/ensuring-go-interface-satisfaction-at-compile-time-1ed158e8fa17
var dialContext = (&net.Dialer{}).DialContext

// ResponseChecker is used to determine when SpoofingClient.Poll is done polling.
// This allows you to predicate wait.PollImmediate on the request's http.Response.
//
// See the apimachinery wait package:
// https://github.com/kubernetes/apimachinery/blob/cf7ae2f57dabc02a3d215f15ca61ae1446f3be8f/pkg/util/wait/wait.go#L172
type ResponseChecker func(resp *Response) (done bool, err error)

// ErrorRetryChecker is used to determine if an error should be retried or not.
// If an error should be retried, it should return true and the wrapped error to explain why to retry.
type ErrorRetryChecker func(e error) (retry bool, err error)

// SpoofingClient is a minimal HTTP client wrapper that spoofs the domain of requests
// for non-resolvable domains.
type SpoofingClient struct {
	Client          *http.Client
	RequestInterval time.Duration
	RequestTimeout  time.Duration
	Logf            logging.FormatLogger
}

// TransportOption allows callers to customize the http.Transport used by a SpoofingClient
type TransportOption func(transport *http.Transport) *http.Transport

// New returns a SpoofingClient that rewrites requests if the target domain is not `resolvable`.
// It does this by looking up the ingress at construction time, so reusing a client will not
// follow the ingress if it moves (or if there are multiple ingresses).
//
// If that's a problem, see test/request.go#WaitForEndpointState for oneshot spoofing.
func New(
	ctx context.Context,
	kubeClientset kubernetes.Interface,
	logf logging.FormatLogger,
	domain string,
	resolvable bool,
	endpointOverride string,
	requestInterval, requestTimeout time.Duration,
	opts ...TransportOption) (*SpoofingClient, error) {
	endpoint, mapper, err := ResolveEndpoint(ctx, kubeClientset, domain, resolvable, endpointOverride)
	if err != nil {
		return nil, fmt.Errorf("failed to get the cluster endpoint: %w", err)
	}

	// Spoof the hostname at the resolver level
	logf("Spoofing %s -> %s", domain, endpoint)
	transport := &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (conn net.Conn, e error) {
			_, port, err := net.SplitHostPort(addr)
			if err != nil {
				return nil, err
			}
			// The original hostname:port is spoofed by replacing the hostname by the value
			// returned by ResolveEndpoint.
			return dialContext(ctx, network, net.JoinHostPort(endpoint, mapper(port)))
		},
	}

	for _, opt := range opts {
		transport = opt(transport)
	}

	// Enable Zipkin tracing
	roundTripper := &ochttp.Transport{
		Base:        transport,
		Propagation: tracecontextb3.TraceContextB3Egress,
	}

	sc := &SpoofingClient{
		Client:          &http.Client{Transport: roundTripper},
		RequestInterval: requestInterval,
		RequestTimeout:  requestTimeout,
		Logf:            logf,
	}
	return sc, nil
}

// ResolveEndpoint resolves the endpoint address considering whether the domain is resolvable and taking into
// account whether the user overrode the endpoint address externally
func ResolveEndpoint(ctx context.Context, kubeClientset kubernetes.Interface, domain string, resolvable bool, endpointOverride string) (string, func(string) string, error) {
	id := func(in string) string { return in }
	// If the domain is resolvable, it can be used directly
	if resolvable {
		return domain, id, nil
	}
	// Otherwise, use the actual cluster endpoint
	return ingress.GetIngressEndpoint(ctx, kubeClientset, endpointOverride)
}

// Do dispatches to the underlying http.Client.Do, spoofing domains as needed
// and transforming the http.Response into a spoof.Response.
// Each response is augmented with "ZipkinTraceID" header that identifies the zipkin trace corresponding to the request.
func (sc *SpoofingClient) Do(req *http.Request, errorRetryCheckers ...ErrorRetryChecker) (*Response, error) {
	return sc.Poll(req, func(*Response) (bool, error) { return true, nil }, errorRetryCheckers...)
}

// Poll executes an http request until it satisfies the inState condition or, if there's an error,
// none of the error retry checkers permit a retry.
// If no retry checkers are specified `DefaultErrorRetryChecker` will be used.
func (sc *SpoofingClient) Poll(req *http.Request, inState ResponseChecker, errorRetryCheckers ...ErrorRetryChecker) (*Response, error) {
	if len(errorRetryCheckers) == 0 {
		errorRetryCheckers = []ErrorRetryChecker{DefaultErrorRetryChecker}
	}

	var resp *Response
	err := wait.PollImmediate(sc.RequestInterval, sc.RequestTimeout, func() (bool, error) {
		// Starting span to capture zipkin trace.
		traceContext, span := trace.StartSpan(req.Context(), "SpoofingClient-Trace")
		defer span.End()
		rawResp, err := sc.Client.Do(req.WithContext(traceContext))
		if err != nil {
			for _, checker := range errorRetryCheckers {
				retry, newErr := checker(err)
				if retry {
					sc.Logf("Retrying %s: %v", req.URL.String(), newErr)
					return false, nil
				}
			}
			sc.Logf("NOT Retrying %s: %v", req.URL.String(), err)
			return true, err
		}
		defer rawResp.Body.Close()

		body, err := ioutil.ReadAll(rawResp.Body)
		if err != nil {
			return true, err
		}
		rawResp.Header.Add(zipkin.ZipkinTraceIDHeader, span.SpanContext().TraceID.String())

		resp = &Response{
			Status:     rawResp.Status,
			StatusCode: rawResp.StatusCode,
			Header:     rawResp.Header,
			Body:       body,
		}
		return inState(resp)
	})

	if resp != nil {
		sc.logZipkinTrace(resp)
	}

	if err != nil {
		return resp, fmt.Errorf("response: %s did not pass checks: %w", resp, err)
	}
	return resp, nil
}

// DefaultErrorRetryChecker implements the defaults for retrying on error.
func DefaultErrorRetryChecker(err error) (bool, error) {
	if isTCPTimeout(err) {
		return true, fmt.Errorf("retrying for TCP timeout: %w", err)
	}
	// Retrying on DNS error, since we may be using sslip.io or nip.io in tests.
	if isDNSError(err) {
		return true, fmt.Errorf("retrying for DNS error: %w", err)
	}
	// Repeat the poll on `connection refused` errors, which are usually transient Istio errors.
	if isConnectionRefused(err) {
		return true, fmt.Errorf("retrying for connection refused: %w", err)
	}
	if isConnectionReset(err) {
		return true, fmt.Errorf("retrying for connection reset: %w", err)
	}
	// Retry on connection/network errors.
	if errors.Is(err, io.EOF) {
		return true, fmt.Errorf("retrying for: %w", err)
	}
	// No route to host errors are in the same category as connection refused errors and
	// are usually transient.
	if isNoRouteToHostError(err) {
		return true, fmt.Errorf("retrying for 'no route to host' error: %w", err)
	}
	return false, err
}

// logZipkinTrace provides support to log Zipkin Trace for param: spoofResponse
// We only log Zipkin trace for HTTP server errors i.e for HTTP status codes between 500 to 600
func (sc *SpoofingClient) logZipkinTrace(spoofResp *Response) {
	if !zipkin.ZipkinTracingEnabled || spoofResp.StatusCode < http.StatusInternalServerError || spoofResp.StatusCode >= 600 {
		return
	}

	traceID := spoofResp.Header.Get(zipkin.ZipkinTraceIDHeader)
	sc.Logf("Logging Zipkin Trace for: %s", traceID)

	json, err := zipkin.JSONTrace(traceID /* We don't know the expected number of spans */, -1, 5*time.Second)
	if err != nil {
		var errTimeout *zipkin.TimeoutError
		if !errors.As(err, &errTimeout) {
			sc.Logf("Error getting zipkin trace: %v", err)
		}
	}

	sc.Logf("%s", json)
}

func (sc *SpoofingClient) WaitForEndpointState(
	ctx context.Context,
	url *url.URL,
	inState ResponseChecker,
	desc string,
	opts ...RequestOption) (*Response, error) {

	return sc.endpointState(
		ctx,
		url,
		inState,
		desc,
		func(req *http.Request, check ResponseChecker) (*Response, error) { return sc.Poll(req, check) },
		"WaitForEndpointState",
		opts...)
}

func (sc *SpoofingClient) endpointState(
	ctx context.Context,
	url *url.URL,
	inState ResponseChecker,
	desc string,
	f func(*http.Request, ResponseChecker) (*Response, error),
	logName string,
	opts ...RequestOption) (*Response, error) {
	defer logging.GetEmitableSpan(ctx, logName+"/"+desc).End()

	if url.Scheme == "" || url.Host == "" {
		return nil, fmt.Errorf("invalid URL: %q", url.String())
	}

	req, err := http.NewRequest(http.MethodGet, url.String(), nil)
	if err != nil {
		return nil, err
	}

	for _, opt := range opts {
		opt(req)
	}

	return f(req, inState)
}

func (sc *SpoofingClient) Check(req *http.Request, inState ResponseChecker, errorRetryCheckers ...ErrorRetryChecker) (*Response, error) {
	if len(errorRetryCheckers) == 0 {
		errorRetryCheckers = []ErrorRetryChecker{DefaultErrorRetryChecker}
	}
	traceContext, span := trace.StartSpan(req.Context(), "SpoofingClient-Trace")
	defer span.End()
	var rawResp *http.Response
	err := wait.PollImmediate(sc.RequestInterval, sc.RequestTimeout, func() (bool, error) {
		var err error
		rawResp, err = sc.Client.Do(req.WithContext(traceContext))
		if err != nil {
			for _, checker := range errorRetryCheckers {
				retry, newErr := checker(err)
				if retry {
					sc.Logf("Retrying %s: %v", req.URL.String(), newErr)
					return false, nil
				}
			}
			sc.Logf("NOT Retrying %s: %v", req.URL.String(), err)
		}
		return true, err
	})

	if err != nil {
		return nil, err
	}
	defer rawResp.Body.Close()

	body, err := ioutil.ReadAll(rawResp.Body)
	if err != nil {
		return nil, err
	}
	rawResp.Header.Add(zipkin.ZipkinTraceIDHeader, span.SpanContext().TraceID.String())

	resp := &Response{
		Status:     rawResp.Status,
		StatusCode: rawResp.StatusCode,
		Header:     rawResp.Header,
		Body:       body,
	}
	ok, err := inState(resp)

	sc.logZipkinTrace(resp)

	if err != nil {
		return resp, fmt.Errorf("response: %s did not pass checks: %w", resp, err)
	}
	if ok {
		return resp, nil
	}

	return nil, err
}

func (sc *SpoofingClient) CheckEndpointState(
	ctx context.Context,
	url *url.URL,
	inState ResponseChecker,
	desc string,
	opts ...RequestOption) (*Response, error) {
	return sc.endpointState(
		ctx,
		url,
		inState,
		desc,
		func(req *http.Request, check ResponseChecker) (*Response, error) { return sc.Check(req, check) },
		"CheckEndpointState",
		opts...)
}
<commit_msg>Deduplicate Check and Poll functions on the spoof client (#2264)<commit_after>/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// spoof contains logic to make polling HTTP requests against an endpoint with optional host spoofing.

package spoof

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"time"

	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"knative.dev/pkg/test/ingress"
	"knative.dev/pkg/test/logging"
	"knative.dev/pkg/test/zipkin"
	"knative.dev/pkg/tracing/propagation/tracecontextb3"

	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/trace"
)

// Response is a stripped down subset of http.Response. The is primarily useful
// for ResponseCheckers to inspect the response body without consuming it.
// Notably, Body is a byte slice instead of an io.ReadCloser.
type Response struct {
	Status     string
	StatusCode int
	Header     http.Header
	Body       []byte
}

func (r *Response) String() string {
	return fmt.Sprintf("status: %d, body: %s, headers: %v", r.StatusCode, string(r.Body), r.Header)
}

// https://medium.com/stupid-gopher-tricks/ensuring-go-interface-satisfaction-at-compile-time-1ed158e8fa17
var dialContext = (&net.Dialer{}).DialContext

// ResponseChecker is used to determine when SpoofingClient.Poll is done polling.
// This allows you to predicate wait.PollImmediate on the request's http.Response.
//
// See the apimachinery wait package:
// https://github.com/kubernetes/apimachinery/blob/cf7ae2f57dabc02a3d215f15ca61ae1446f3be8f/pkg/util/wait/wait.go#L172
type ResponseChecker func(resp *Response) (done bool, err error)

// ErrorRetryChecker is used to determine if an error should be retried or not.
// If an error should be retried, it should return true and the wrapped error to explain why to retry.
type ErrorRetryChecker func(e error) (retry bool, err error)

// SpoofingClient is a minimal HTTP client wrapper that spoofs the domain of requests
// for non-resolvable domains.
type SpoofingClient struct {
	Client          *http.Client
	RequestInterval time.Duration
	RequestTimeout  time.Duration
	Logf            logging.FormatLogger
}

// TransportOption allows callers to customize the http.Transport used by a SpoofingClient
type TransportOption func(transport *http.Transport) *http.Transport

// New returns a SpoofingClient that rewrites requests if the target domain is not `resolvable`.
// It does this by looking up the ingress at construction time, so reusing a client will not
// follow the ingress if it moves (or if there are multiple ingresses).
//
// If that's a problem, see test/request.go#WaitForEndpointState for oneshot spoofing.
func New(
	ctx context.Context,
	kubeClientset kubernetes.Interface,
	logf logging.FormatLogger,
	domain string,
	resolvable bool,
	endpointOverride string,
	requestInterval, requestTimeout time.Duration,
	opts ...TransportOption) (*SpoofingClient, error) {
	endpoint, mapper, err := ResolveEndpoint(ctx, kubeClientset, domain, resolvable, endpointOverride)
	if err != nil {
		return nil, fmt.Errorf("failed to get the cluster endpoint: %w", err)
	}

	// Spoof the hostname at the resolver level
	logf("Spoofing %s -> %s", domain, endpoint)
	transport := &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (conn net.Conn, e error) {
			_, port, err := net.SplitHostPort(addr)
			if err != nil {
				return nil, err
			}
			// The original hostname:port is spoofed by replacing the hostname by the value
			// returned by ResolveEndpoint.
			return dialContext(ctx, network, net.JoinHostPort(endpoint, mapper(port)))
		},
	}

	for _, opt := range opts {
		transport = opt(transport)
	}

	// Enable Zipkin tracing
	roundTripper := &ochttp.Transport{
		Base:        transport,
		Propagation: tracecontextb3.TraceContextB3Egress,
	}

	sc := &SpoofingClient{
		Client:          &http.Client{Transport: roundTripper},
		RequestInterval: requestInterval,
		RequestTimeout:  requestTimeout,
		Logf:            logf,
	}
	return sc, nil
}

// ResolveEndpoint resolves the endpoint address considering whether the domain is resolvable and taking into
// account whether the user overrode the endpoint address externally
func ResolveEndpoint(ctx context.Context, kubeClientset kubernetes.Interface, domain string, resolvable bool, endpointOverride string) (string, func(string) string, error) {
	id := func(in string) string { return in }
	// If the domain is resolvable, it can be used directly
	if resolvable {
		return domain, id, nil
	}
	// Otherwise, use the actual cluster endpoint
	return ingress.GetIngressEndpoint(ctx, kubeClientset, endpointOverride)
}

// Do dispatches to the underlying http.Client.Do, spoofing domains as needed
// and transforming the http.Response into a spoof.Response.
// Each response is augmented with "ZipkinTraceID" header that identifies the zipkin trace corresponding to the request.
func (sc *SpoofingClient) Do(req *http.Request, errorRetryCheckers ...ErrorRetryChecker) (*Response, error) {
	return sc.Poll(req, func(*Response) (bool, error) { return true, nil }, errorRetryCheckers...)
}

// Poll executes an http request until it satisfies the inState condition or, if there's an error,
// none of the error retry checkers permit a retry.
// If no retry checkers are specified `DefaultErrorRetryChecker` will be used.
func (sc *SpoofingClient) Poll(req *http.Request, inState ResponseChecker, errorRetryCheckers ...ErrorRetryChecker) (*Response, error) {
	if len(errorRetryCheckers) == 0 {
		errorRetryCheckers = []ErrorRetryChecker{DefaultErrorRetryChecker}
	}

	var resp *Response
	err := wait.PollImmediate(sc.RequestInterval, sc.RequestTimeout, func() (bool, error) {
		// Starting span to capture zipkin trace.
		traceContext, span := trace.StartSpan(req.Context(), "SpoofingClient-Trace")
		defer span.End()
		rawResp, err := sc.Client.Do(req.WithContext(traceContext))
		if err != nil {
			for _, checker := range errorRetryCheckers {
				retry, newErr := checker(err)
				if retry {
					sc.Logf("Retrying %s: %v", req.URL.String(), newErr)
					return false, nil
				}
			}
			sc.Logf("NOT Retrying %s: %v", req.URL.String(), err)
			return true, err
		}
		defer rawResp.Body.Close()

		body, err := ioutil.ReadAll(rawResp.Body)
		if err != nil {
			return true, err
		}
		rawResp.Header.Add(zipkin.ZipkinTraceIDHeader, span.SpanContext().TraceID.String())

		resp = &Response{
			Status:     rawResp.Status,
			StatusCode: rawResp.StatusCode,
			Header:     rawResp.Header,
			Body:       body,
		}
		return inState(resp)
	})

	if resp != nil {
		sc.logZipkinTrace(resp)
	}

	if err != nil {
		return resp, fmt.Errorf("response: %s did not pass checks: %w", resp, err)
	}
	return resp, nil
}

// DefaultErrorRetryChecker implements the defaults for retrying on error.
func DefaultErrorRetryChecker(err error) (bool, error) {
	if isTCPTimeout(err) {
		return true, fmt.Errorf("retrying for TCP timeout: %w", err)
	}
	// Retrying on DNS error, since we may be using sslip.io or nip.io in tests.
	if isDNSError(err) {
		return true, fmt.Errorf("retrying for DNS error: %w", err)
	}
	// Repeat the poll on `connection refused` errors, which are usually transient Istio errors.
	if isConnectionRefused(err) {
		return true, fmt.Errorf("retrying for connection refused: %w", err)
	}
	if isConnectionReset(err) {
		return true, fmt.Errorf("retrying for connection reset: %w", err)
	}
	// Retry on connection/network errors.
	if errors.Is(err, io.EOF) {
		return true, fmt.Errorf("retrying for: %w", err)
	}
	// No route to host errors are in the same category as connection refused errors and
	// are usually transient.
	if isNoRouteToHostError(err) {
		return true, fmt.Errorf("retrying for 'no route to host' error: %w", err)
	}
	return false, err
}

// logZipkinTrace provides support to log Zipkin Trace for param: spoofResponse
// We only log Zipkin trace for HTTP server errors i.e for HTTP status codes between 500 to 600
func (sc *SpoofingClient) logZipkinTrace(spoofResp *Response) {
	if !zipkin.ZipkinTracingEnabled || spoofResp.StatusCode < http.StatusInternalServerError || spoofResp.StatusCode >= 600 {
		return
	}

	traceID := spoofResp.Header.Get(zipkin.ZipkinTraceIDHeader)
	sc.Logf("Logging Zipkin Trace for: %s", traceID)

	json, err := zipkin.JSONTrace(traceID /* We don't know the expected number of spans */, -1, 5*time.Second)
	if err != nil {
		var errTimeout *zipkin.TimeoutError
		if !errors.As(err, &errTimeout) {
			sc.Logf("Error getting zipkin trace: %v", err)
		}
	}

	sc.Logf("%s", json)
}

func (sc *SpoofingClient) WaitForEndpointState(
	ctx context.Context,
	url *url.URL,
	inState ResponseChecker,
	desc string,
	opts ...RequestOption) (*Response, error) {

	return sc.endpointState(
		ctx,
		url,
		inState,
		desc,
		func(req *http.Request, check ResponseChecker) (*Response, error) { return sc.Poll(req, check) },
		"WaitForEndpointState",
		opts...)
}

func (sc *SpoofingClient) endpointState(
	ctx context.Context,
	url *url.URL,
	inState ResponseChecker,
	desc string,
	f func(*http.Request, ResponseChecker) (*Response, error),
	logName string,
	opts ...RequestOption) (*Response, error) {
	defer logging.GetEmitableSpan(ctx, logName+"/"+desc).End()

	if url.Scheme == "" || url.Host == "" {
		return nil, fmt.Errorf("invalid URL: %q", url.String())
	}

	req, err := http.NewRequest(http.MethodGet, url.String(), nil)
	if err != nil {
		return nil, err
	}

	for _, opt := range opts {
		opt(req)
	}

	return f(req, inState)
}

func (sc *SpoofingClient) Check(req *http.Request, inState ResponseChecker, errorRetryCheckers ...ErrorRetryChecker) (*Response, error) {
	resp, err := sc.Do(req, errorRetryCheckers...)
	if err != nil {
		return nil, err
	}

	ok, err := inState(resp)
	if err != nil {
		return resp, fmt.Errorf("response: %s did not pass checks: %w", resp, err)
	}
	if ok {
		return resp, nil
	}

	return nil, err
}

func (sc *SpoofingClient) CheckEndpointState(
	ctx context.Context,
	url *url.URL,
	inState ResponseChecker,
	desc string,
	opts ...RequestOption) (*Response, error) {
	return sc.endpointState(
		ctx,
		url,
		inState,
		desc,
		func(req *http.Request, check ResponseChecker) (*Response, error) { return sc.Check(req, check) },
		"CheckEndpointState",
		opts...)
}
<|endoftext|>