<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"time"
)

const DEFAULT_RESPONSE = `{"default": "response"}`
const MAX_DELAY = 99
const MIN_STATUS = 100
const MAX_STATUS = 599

var REDIRECT_STATUS map[int]bool = map[int]bool{301: true, 302: true}

type GoSlowServer struct {
	Config *Config
	Store  Store
}

func (server *GoSlowServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("%s %s", r.Method, r.URL.Path)
	rule, found := FindRule(server.Store, r)
	if found {
		log.Printf("sleeping for %v", rule.Delay)
		time.Sleep(rule.Delay)

		AddHeaders(rule.Header, w.Header())
		w.WriteHeader(rule.ResponseStatus)
		io.WriteString(w, rule.Response)
	} else {
		io.WriteString(w, DEFAULT_RESPONSE)
	}
}

func AddHeaders(src map[string]string, dst http.Header) {
	for key, value := range src {
		dst.Add(key, value)
	}
}

func (server *GoSlowServer) AddDefaultRules() {
	server.AddDelayRules()
	server.AddStatusRules()
}

func (server *GoSlowServer) AddDelayRules() {
	for delay := 0; delay <= MAX_DELAY; delay++ {
		delayHost := server.MakeFullHost(delay)
		delayInSeconds := time.Duration(delay) * time.Second

		server.Store.AddRule(&Rule{Host: delayHost, Header: EmptyHeader(), Delay: delayInSeconds,
			ResponseStatus: 200, Response: DEFAULT_RESPONSE,
		})
	}
}

func (server *GoSlowServer) MakeFullHost(subdomain int) string {
	return fmt.Sprintf("%d.%s", subdomain, server.Config.Host)
}

func EmptyHeader() map[string]string {
	return make(map[string]string)
}

func (server *GoSlowServer) AddStatusRules() {
	for status := MIN_STATUS; status <= MAX_STATUS; status++ {
		statusHost := server.MakeFullHost(status)
		header := server.HeaderFor(status)
		server.Store.AddRule(&Rule{Host: statusHost, ResponseStatus: status,
			Header: header, Response: DEFAULT_RESPONSE})
	}
}

func (server *GoSlowServer) HeaderFor(status int) map[string]string {
	_, isRedirect := REDIRECT_STATUS[status]
	if isRedirect {
		// TODO: check that protocol-independent location is legal HTTP
		host := fmt.Sprintf("//%s", server.MakeFullHost(0))
		return map[string]string{"Location": host}
	}
	return EmptyHeader()
}

func (server *GoSlowServer) ListenAndServe() error {
	log.Printf("listening on %s", server.Config.Address)
	return http.ListenAndServe(server.Config.Address, server)
}
<commit_msg>Allow cross domain requests<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"time"
)

const DEFAULT_RESPONSE = `{"default": "response"}`
const MAX_DELAY = 99
const MIN_STATUS = 100
const MAX_STATUS = 599

var REDIRECT_STATUS map[int]bool = map[int]bool{301: true, 302: true}

type GoSlowServer struct {
	Config *Config
	Store  Store
}

func (server *GoSlowServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("%s %s", r.Method, r.URL.Path)
	rule, found := FindRule(server.Store, r)
	if found {
		log.Printf("sleeping for %v", rule.Delay)
		time.Sleep(rule.Delay)

		AllowCrossDomainRequests(w)
		AddHeaders(rule.Header, w)
		w.WriteHeader(rule.ResponseStatus)
		io.WriteString(w, rule.Response)
	} else {
		io.WriteString(w, DEFAULT_RESPONSE)
	}
}

func AllowCrossDomainRequests(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
}

func AddHeaders(header map[string]string, w http.ResponseWriter) {
	responseHeader := w.Header()
	for key, value := range header {
		responseHeader.Add(key, value)
	}
}

func (server *GoSlowServer) AddDefaultRules() {
	server.AddDelayRules()
	server.AddStatusRules()
}

func (server *GoSlowServer) AddDelayRules() {
	for delay := 0; delay <= MAX_DELAY; delay++ {
		delayHost := server.MakeFullHost(delay)
		delayInSeconds := time.Duration(delay) * time.Second

		server.Store.AddRule(&Rule{Host: delayHost, Header: EmptyHeader(), Delay: delayInSeconds,
			ResponseStatus: 200, Response: DEFAULT_RESPONSE,
		})
	}
}

func (server *GoSlowServer) MakeFullHost(subdomain int) string {
	return fmt.Sprintf("%d.%s", subdomain, server.Config.Host)
}

func EmptyHeader() map[string]string {
	return make(map[string]string)
}

func (server *GoSlowServer) AddStatusRules() {
	for status := MIN_STATUS; status <= MAX_STATUS; status++ {
		statusHost := server.MakeFullHost(status)
		header := server.HeaderFor(status)
		server.Store.AddRule(&Rule{Host: statusHost, ResponseStatus: status,
			Header: header, Response: DEFAULT_RESPONSE})
	}
}

func (server *GoSlowServer) HeaderFor(status int) map[string]string {
	_, isRedirect := REDIRECT_STATUS[status]
	if isRedirect {
		// TODO: check that protocol-independent location is legal HTTP
		host := fmt.Sprintf("//%s", server.MakeFullHost(0))
		return map[string]string{"Location": host}
	}
	return EmptyHeader()
}

func (server *GoSlowServer) ListenAndServe() error {
	log.Printf("listening on %s", server.Config.Address)
	return http.ListenAndServe(server.Config.Address, server)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"
)

// WebClient is a client that can retrieve a URL body and POST to a given URL.
type WebClient interface {
	GetURL(url string, maxLength int64) (r io.Reader, statusCode int, err error)
	PostURL(url string, body io.Reader) (statusCode int, err error)
	SetTimeout(d time.Duration)
}

// NewWebClient returns a new WebClient.
func NewWebClient() WebClient {
	return &webClientDefault{client: &http.Client{Timeout: 30 * time.Second}}
}

type webClientDefault struct {
	client *http.Client
}

func (w *webClientDefault) SetTimeout(d time.Duration) {
	w.client.Timeout = d
}

func (w *webClientDefault) GetURL(url string, maxLength int64) (r io.Reader, statusCode int, err error) {
	httpResp, err := w.client.Get(url)
	if err != nil {
		return nil, -1, err
	}

	// We read the entire body into a buffer now to avoid potential problems with timeouts later
	body, err := readBody(httpResp, maxLength)

	return bytes.NewReader(body), httpResp.StatusCode, err
}

func (w *webClientDefault) PostURL(url string, body io.Reader) (statusCode int, err error) {
	httpResp, err := w.client.Post(url, "application/json", body)
	if err != nil {
		return -1, err
	}

	defer httpResp.Body.Close()

	return httpResp.StatusCode, err
}

func readBody(r *http.Response, maxLength int64) ([]byte, error) {
	defer r.Body.Close()

	if r.ContentLength > maxLength {
		return nil, fmt.Errorf("Discarding because content length is greater than 2MB (%d)", r.ContentLength)
	}

	// If content-length is -1, for now assume these processes complete fast enough
	// that consuming 2MB per request is fine
	cappedR := NewAtMostNReader(maxLength, r.Body)
	return ioutil.ReadAll(cappedR)
}
<commit_msg>try to create new web client to fight "re-use of closed conn" issue<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"
)

// WebClient is a client that can retrieve a URL body and POST to a given URL.
type WebClient interface {
	GetURL(url string, maxLength int64) (r io.Reader, statusCode int, err error)
	PostURL(url string, body io.Reader) (statusCode int, err error)
	SetTimeout(d time.Duration)
}

// NewWebClient returns a new WebClient.
func NewWebClient() WebClient {
	return &webClientDefault{client: &http.Client{Timeout: 30 * time.Second}}
}

type webClientDefault struct {
	client *http.Client
}

func (w *webClientDefault) SetTimeout(d time.Duration) {
	w.client.Timeout = d
}

func (w *webClientDefault) GetURL(url string, maxLength int64) (r io.Reader, statusCode int, err error) {
	client := http.Client{Timeout: 30 * time.Second}
	httpResp, err := client.Get(url)
	if err != nil {
		return nil, -1, err
	}

	// We read the entire body into a buffer now to avoid potential problems with timeouts later
	body, err := readBody(httpResp, maxLength)

	return bytes.NewReader(body), httpResp.StatusCode, err
}

func (w *webClientDefault) PostURL(url string, body io.Reader) (statusCode int, err error) {
	httpResp, err := w.client.Post(url, "application/json", body)
	if err != nil {
		return -1, err
	}

	defer httpResp.Body.Close()

	return httpResp.StatusCode, err
}

func readBody(r *http.Response, maxLength int64) ([]byte, error) {
	defer r.Body.Close()

	if r.ContentLength > maxLength {
		return nil, fmt.Errorf("Discarding because content length is greater than 2MB (%d)", r.ContentLength)
	}

	// If content-length is -1, for now assume these processes complete fast enough
	// that consuming 2MB per request is fine
	cappedR := NewAtMostNReader(maxLength, r.Body)
	return ioutil.ReadAll(cappedR)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 kamatera

import (
	"context"
	"fmt"
	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"testing"

	"github.com/stretchr/testify/assert"
	"k8s.io/autoscaler/cluster-autoscaler/cloudprovider"
)

func TestNodeGroup_IncreaseSize(t *testing.T) {
	client := kamateraClientMock{}
	ctx := context.Background()
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	serverConfig := mockServerConfig("test", []string{})
	ng := NodeGroup{
		id:      "ng1",
		manager: &mgr,
		minSize: 1,
		maxSize: 7,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		serverConfig: serverConfig,
	}

	// test error on bad delta values
	err := ng.IncreaseSize(0)
	assert.Error(t, err)
	assert.Equal(t, "delta must be positive, have: 0", err.Error())

	err = ng.IncreaseSize(-1)
	assert.Error(t, err)
	assert.Equal(t, "delta must be positive, have: -1", err.Error())

	// test error on a too large increase of nodes
	err = ng.IncreaseSize(5)
	assert.Error(t, err)
	assert.Equal(t, "size increase is too large. current: 3 desired: 8 max: 7", err.Error())

	// test ok to add a node
	client.On(
		"CreateServers", ctx, 1, serverConfig,
	).Return(
		[]Server{{Name: mockKamateraServerName()}}, nil,
	).Once()
	err = ng.IncreaseSize(1)
	assert.NoError(t, err)
	assert.Equal(t, 4, len(ng.instances))

	// test ok to add multiple nodes
	client.On(
		"CreateServers", ctx, 2, serverConfig,
	).Return(
		[]Server{
			{Name: mockKamateraServerName()},
			{Name: mockKamateraServerName()},
		}, nil,
	).Once()
	err = ng.IncreaseSize(2)
	assert.NoError(t, err)
	assert.Equal(t, 6, len(ng.instances))

	// test error on linode API call error
	client.On(
		"CreateServers", ctx, 1, serverConfig,
	).Return(
		[]Server{}, fmt.Errorf("error on API call"),
	).Once()
	err = ng.IncreaseSize(1)
	assert.Error(t, err, "no error on injected API call error")
	assert.Equal(t, "error on API call", err.Error())
}

func TestNodeGroup_DecreaseTargetSize(t *testing.T) {
	ng := &NodeGroup{}
	err := ng.DecreaseTargetSize(-1)
	assert.Error(t, err)
	assert.Equal(t, "Not implemented", err.Error())
}

func TestNodeGroup_DeleteNodes(t *testing.T) {
	client := kamateraClientMock{}
	ctx := context.Background()
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	serverName4 := mockKamateraServerName()
	serverName5 := mockKamateraServerName()
	serverName6 := mockKamateraServerName()
	ng := NodeGroup{
		id:      "ng1",
		minSize: 1,
		maxSize: 6,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName4: {Id: serverName4, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName5: {Id: serverName5, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName6: {Id: serverName6, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		manager: &mgr,
	}

	// test of deleting nodes
	client.On(
		"DeleteServer", ctx, serverName1,
	).Return(nil).Once().On(
		"DeleteServer", ctx, serverName2,
	).Return(nil).Once().On(
		"DeleteServer", ctx, serverName6,
	).Return(nil).Once()
	err := ng.DeleteNodes([]*apiv1.Node{
		{Spec: apiv1.NodeSpec{ProviderID: serverName1}},
		{Spec: apiv1.NodeSpec{ProviderID: serverName2}},
		{Spec: apiv1.NodeSpec{ProviderID: serverName6}},
	})
	assert.NoError(t, err)
	assert.Equal(t, 3, len(ng.instances))
	assert.Equal(t, serverName3, ng.instances[serverName3].Id)
	assert.Equal(t, serverName4, ng.instances[serverName4].Id)
	assert.Equal(t, serverName5, ng.instances[serverName5].Id)

	// test error on deleting a node we are not managing
	err = ng.DeleteNodes([]*apiv1.Node{{Spec: apiv1.NodeSpec{ProviderID: mockKamateraServerName()}}})
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "cannot find this node in the node group")

	// test error on deleting a node when the linode API call fails
	client.On(
		"DeleteServer", ctx, serverName4,
	).Return(fmt.Errorf("error on API call")).Once()
	err = ng.DeleteNodes([]*apiv1.Node{
		{Spec: apiv1.NodeSpec{ProviderID: serverName4}},
	})
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "error on API call")
}

func TestNodeGroup_Nodes(t *testing.T) {
	client := kamateraClientMock{}
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	ng := NodeGroup{
		id:      "ng1",
		minSize: 1,
		maxSize: 6,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		manager: &mgr,
	}

	// test nodes returned from Nodes() are only the ones we are expecting
	instancesList, err := ng.Nodes()
	assert.NoError(t, err)
	assert.Equal(t, 3, len(instancesList))
	var serverIds []string
	for _, instance := range instancesList {
		serverIds = append(serverIds, instance.Id)
	}
	assert.Equal(t, 3, len(serverIds))
	assert.Contains(t, serverIds, serverName1)
	assert.Contains(t, serverIds, serverName2)
	assert.Contains(t, serverIds, serverName3)
}

func TestNodeGroup_getResourceList(t *testing.T) {
	ng := &NodeGroup{}
	_, err := ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config ram")
	ng.serverConfig.Ram = "1024mb"
	_, err = ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config ram")
	ng.serverConfig.Ram = "1024"
	_, err = ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config cpu")
	ng.serverConfig.Cpu = "55PX"
	_, err = ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config cpu")
	ng.serverConfig.Cpu = "55P"
	rl, err := ng.getResourceList()
	assert.NoError(t, err)
	assert.Equal(t, apiv1.ResourceList{
		apiv1.ResourcePods:    *resource.NewQuantity(110, resource.DecimalSI),
		apiv1.ResourceCPU:     *resource.NewQuantity(int64(55), resource.DecimalSI),
		apiv1.ResourceMemory:  *resource.NewQuantity(int64(1024*1024*1024*1024), resource.DecimalSI),
		apiv1.ResourceStorage: *resource.NewQuantity(int64(0*1024*1024*1024), resource.DecimalSI),
	}, rl)
	ng.serverConfig.Disks = []string{"size=50"}
	rl, err = ng.getResourceList()
	assert.NoError(t, err)
	assert.Equal(t, apiv1.ResourceList{
		apiv1.ResourcePods:    *resource.NewQuantity(110, resource.DecimalSI),
		apiv1.ResourceCPU:     *resource.NewQuantity(int64(55), resource.DecimalSI),
		apiv1.ResourceMemory:  *resource.NewQuantity(int64(1024*1024*1024*1024), resource.DecimalSI),
		apiv1.ResourceStorage: *resource.NewQuantity(int64(50*1024*1024*1024), resource.DecimalSI),
	}, rl)
}

func TestNodeGroup_TemplateNodeInfo(t *testing.T) {
	ng := &NodeGroup{
		serverConfig: ServerConfig{
			Ram:   "1024",
			Cpu:   "5D",
			Disks: []string{"size=50"},
		},
	}
	nodeInfo, err := ng.TemplateNodeInfo()
	assert.NoError(t, err)
	assert.Equal(t, nodeInfo.Node().Status.Capacity, apiv1.ResourceList{
		apiv1.ResourcePods:    *resource.NewQuantity(110, resource.DecimalSI),
		apiv1.ResourceCPU:     *resource.NewQuantity(int64(5), resource.DecimalSI),
		apiv1.ResourceMemory:  *resource.NewQuantity(int64(1024*1024*1024*1024), resource.DecimalSI),
		apiv1.ResourceStorage: *resource.NewQuantity(int64(50*1024*1024*1024), resource.DecimalSI),
	})
}

func TestNodeGroup_Others(t *testing.T) {
	client := kamateraClientMock{}
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	ng := NodeGroup{
		id:      "ng1",
		minSize: 1,
		maxSize: 7,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		manager: &mgr,
	}
	assert.Equal(t, 1, ng.MinSize())
	assert.Equal(t, 7, ng.MaxSize())
	ts, err := ng.TargetSize()
	assert.NoError(t, err)
	assert.Equal(t, 3, ts)
	assert.Equal(t, "ng1", ng.Id())
	assert.Equal(t, "node group ID: ng1 (min:1 max:7)", ng.Debug())
	assert.Equal(t, fmt.Sprintf(`node group ID: ng1 (min:1 max:7)
instance ID: %s state: Running powerOn: false
instance ID: %s state: Running powerOn: false
instance ID: %s state: Running powerOn: false`, serverName1, serverName2, serverName3), ng.extendedDebug())
	assert.Equal(t, true, ng.Exist())
	assert.Equal(t, false, ng.Autoprovisioned())
	_, err = ng.Create()
	assert.Error(t, err)
	assert.Equal(t, "Not implemented", err.Error())
	err = ng.Delete()
	assert.Error(t, err)
	assert.Equal(t, "Not implemented", err.Error())
}
<commit_msg>fix non-deterministic kamatera cloudprovider unit test<commit_after>/*
Copyright 2016 The Kubernetes 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 kamatera

import (
	"context"
	"fmt"
	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"k8s.io/autoscaler/cluster-autoscaler/cloudprovider"
)

func TestNodeGroup_IncreaseSize(t *testing.T) {
	client := kamateraClientMock{}
	ctx := context.Background()
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	serverConfig := mockServerConfig("test", []string{})
	ng := NodeGroup{
		id:      "ng1",
		manager: &mgr,
		minSize: 1,
		maxSize: 7,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		serverConfig: serverConfig,
	}

	// test error on bad delta values
	err := ng.IncreaseSize(0)
	assert.Error(t, err)
	assert.Equal(t, "delta must be positive, have: 0", err.Error())

	err = ng.IncreaseSize(-1)
	assert.Error(t, err)
	assert.Equal(t, "delta must be positive, have: -1", err.Error())

	// test error on a too large increase of nodes
	err = ng.IncreaseSize(5)
	assert.Error(t, err)
	assert.Equal(t, "size increase is too large. current: 3 desired: 8 max: 7", err.Error())

	// test ok to add a node
	client.On(
		"CreateServers", ctx, 1, serverConfig,
	).Return(
		[]Server{{Name: mockKamateraServerName()}}, nil,
	).Once()
	err = ng.IncreaseSize(1)
	assert.NoError(t, err)
	assert.Equal(t, 4, len(ng.instances))

	// test ok to add multiple nodes
	client.On(
		"CreateServers", ctx, 2, serverConfig,
	).Return(
		[]Server{
			{Name: mockKamateraServerName()},
			{Name: mockKamateraServerName()},
		}, nil,
	).Once()
	err = ng.IncreaseSize(2)
	assert.NoError(t, err)
	assert.Equal(t, 6, len(ng.instances))

	// test error on linode API call error
	client.On(
		"CreateServers", ctx, 1, serverConfig,
	).Return(
		[]Server{}, fmt.Errorf("error on API call"),
	).Once()
	err = ng.IncreaseSize(1)
	assert.Error(t, err, "no error on injected API call error")
	assert.Equal(t, "error on API call", err.Error())
}

func TestNodeGroup_DecreaseTargetSize(t *testing.T) {
	ng := &NodeGroup{}
	err := ng.DecreaseTargetSize(-1)
	assert.Error(t, err)
	assert.Equal(t, "Not implemented", err.Error())
}

func TestNodeGroup_DeleteNodes(t *testing.T) {
	client := kamateraClientMock{}
	ctx := context.Background()
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	serverName4 := mockKamateraServerName()
	serverName5 := mockKamateraServerName()
	serverName6 := mockKamateraServerName()
	ng := NodeGroup{
		id:      "ng1",
		minSize: 1,
		maxSize: 6,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName4: {Id: serverName4, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName5: {Id: serverName5, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName6: {Id: serverName6, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		manager: &mgr,
	}

	// test of deleting nodes
	client.On(
		"DeleteServer", ctx, serverName1,
	).Return(nil).Once().On(
		"DeleteServer", ctx, serverName2,
	).Return(nil).Once().On(
		"DeleteServer", ctx, serverName6,
	).Return(nil).Once()
	err := ng.DeleteNodes([]*apiv1.Node{
		{Spec: apiv1.NodeSpec{ProviderID: serverName1}},
		{Spec: apiv1.NodeSpec{ProviderID: serverName2}},
		{Spec: apiv1.NodeSpec{ProviderID: serverName6}},
	})
	assert.NoError(t, err)
	assert.Equal(t, 3, len(ng.instances))
	assert.Equal(t, serverName3, ng.instances[serverName3].Id)
	assert.Equal(t, serverName4, ng.instances[serverName4].Id)
	assert.Equal(t, serverName5, ng.instances[serverName5].Id)

	// test error on deleting a node we are not managing
	err = ng.DeleteNodes([]*apiv1.Node{{Spec: apiv1.NodeSpec{ProviderID: mockKamateraServerName()}}})
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "cannot find this node in the node group")

	// test error on deleting a node when the linode API call fails
	client.On(
		"DeleteServer", ctx, serverName4,
	).Return(fmt.Errorf("error on API call")).Once()
	err = ng.DeleteNodes([]*apiv1.Node{
		{Spec: apiv1.NodeSpec{ProviderID: serverName4}},
	})
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "error on API call")
}

func TestNodeGroup_Nodes(t *testing.T) {
	client := kamateraClientMock{}
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	ng := NodeGroup{
		id:      "ng1",
		minSize: 1,
		maxSize: 6,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		manager: &mgr,
	}

	// test nodes returned from Nodes() are only the ones we are expecting
	instancesList, err := ng.Nodes()
	assert.NoError(t, err)
	assert.Equal(t, 3, len(instancesList))
	var serverIds []string
	for _, instance := range instancesList {
		serverIds = append(serverIds, instance.Id)
	}
	assert.Equal(t, 3, len(serverIds))
	assert.Contains(t, serverIds, serverName1)
	assert.Contains(t, serverIds, serverName2)
	assert.Contains(t, serverIds, serverName3)
}

func TestNodeGroup_getResourceList(t *testing.T) {
	ng := &NodeGroup{}
	_, err := ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config ram")
	ng.serverConfig.Ram = "1024mb"
	_, err = ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config ram")
	ng.serverConfig.Ram = "1024"
	_, err = ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config cpu")
	ng.serverConfig.Cpu = "55PX"
	_, err = ng.getResourceList()
	assert.ErrorContains(t, err, "failed to parse server config cpu")
	ng.serverConfig.Cpu = "55P"
	rl, err := ng.getResourceList()
	assert.NoError(t, err)
	assert.Equal(t, apiv1.ResourceList{
		apiv1.ResourcePods:    *resource.NewQuantity(110, resource.DecimalSI),
		apiv1.ResourceCPU:     *resource.NewQuantity(int64(55), resource.DecimalSI),
		apiv1.ResourceMemory:  *resource.NewQuantity(int64(1024*1024*1024*1024), resource.DecimalSI),
		apiv1.ResourceStorage: *resource.NewQuantity(int64(0*1024*1024*1024), resource.DecimalSI),
	}, rl)
	ng.serverConfig.Disks = []string{"size=50"}
	rl, err = ng.getResourceList()
	assert.NoError(t, err)
	assert.Equal(t, apiv1.ResourceList{
		apiv1.ResourcePods:    *resource.NewQuantity(110, resource.DecimalSI),
		apiv1.ResourceCPU:     *resource.NewQuantity(int64(55), resource.DecimalSI),
		apiv1.ResourceMemory:  *resource.NewQuantity(int64(1024*1024*1024*1024), resource.DecimalSI),
		apiv1.ResourceStorage: *resource.NewQuantity(int64(50*1024*1024*1024), resource.DecimalSI),
	}, rl)
}

func TestNodeGroup_TemplateNodeInfo(t *testing.T) {
	ng := &NodeGroup{
		serverConfig: ServerConfig{
			Ram:   "1024",
			Cpu:   "5D",
			Disks: []string{"size=50"},
		},
	}
	nodeInfo, err := ng.TemplateNodeInfo()
	assert.NoError(t, err)
	assert.Equal(t, nodeInfo.Node().Status.Capacity, apiv1.ResourceList{
		apiv1.ResourcePods:    *resource.NewQuantity(110, resource.DecimalSI),
		apiv1.ResourceCPU:     *resource.NewQuantity(int64(5), resource.DecimalSI),
		apiv1.ResourceMemory:  *resource.NewQuantity(int64(1024*1024*1024*1024), resource.DecimalSI),
		apiv1.ResourceStorage: *resource.NewQuantity(int64(50*1024*1024*1024), resource.DecimalSI),
	})
}

func TestNodeGroup_Others(t *testing.T) {
	client := kamateraClientMock{}
	mgr := manager{
		client:    &client,
		instances: make(map[string]*Instance),
	}
	serverName1 := mockKamateraServerName()
	serverName2 := mockKamateraServerName()
	serverName3 := mockKamateraServerName()
	ng := NodeGroup{
		id:      "ng1",
		minSize: 1,
		maxSize: 7,
		instances: map[string]*Instance{
			serverName1: {Id: serverName1, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName2: {Id: serverName2, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
			serverName3: {Id: serverName3, Status: &cloudprovider.InstanceStatus{State: cloudprovider.InstanceRunning}},
		},
		manager: &mgr,
	}
	assert.Equal(t, 1, ng.MinSize())
	assert.Equal(t, 7, ng.MaxSize())
	ts, err := ng.TargetSize()
	assert.NoError(t, err)
	assert.Equal(t, 3, ts)
	assert.Equal(t, "ng1", ng.Id())
	assert.Equal(t, "node group ID: ng1 (min:1 max:7)", ng.Debug())
	extendedDebug := strings.Split(ng.extendedDebug(), "\n")
	assert.Equal(t, 4, len(extendedDebug))
	assert.Contains(t, extendedDebug, "node group ID: ng1 (min:1 max:7)")
	for _, serverName := range []string{serverName1, serverName2, serverName3} {
		assert.Contains(t, extendedDebug, fmt.Sprintf("instance ID: %s state: Running powerOn: false", serverName))
	}
	assert.Equal(t, true, ng.Exist())
	assert.Equal(t, false, ng.Autoprovisioned())
	_, err = ng.Create()
	assert.Error(t, err)
	assert.Equal(t, "Not implemented", err.Error())
	err = ng.Delete()
	assert.Error(t, err)
	assert.Equal(t, "Not implemented", err.Error())
}
<|endoftext|>


<commit_before>package httpfetch

import (
	"fmt"
	"net/http"
	"os"
	"testing"
	"time"

	"github.com/labstack/echo"
	"github.com/stretchr/testify/assert"

	"github.com/qiniu/logkit/conf"
	"github.com/qiniu/logkit/reader"
	. "github.com/qiniu/logkit/reader/config"
	. "github.com/qiniu/logkit/reader/test"
	. "github.com/qiniu/logkit/utils/models"
)

var retData = &map[string]string{
	"logkit1": "data1",
	"logkit2": "data2",
	"logkit3": "data3",
	"logkit4": "data4",
}

func getHttpFetchReader(port string) (*Reader, error) {
	readConf := conf.MapConf{
		KeyMetaPath:   MetaDir,
		KeyFileDone:   MetaDir,
		KeyMode:       ModeHTTP,
		KeyRunnerName: "TestNewHttpFetchReader",
	}
	meta, err := reader.NewMetaWithConf(readConf)
	if err != nil {
		return nil, err
	}
	c := conf.MapConf{
		KeyHTTPServiceAddress: "127.0.0.1:" + port,
		KeyHTTPServicePath:    "/logkit/data",
		KeyHttpCron:           "loop1s",
	}
	reader, err := NewReader(meta, c)
	httpReader := reader.(*Reader)
	if err != nil {
		return nil, err
	}
	err = httpReader.Start()
	if err != nil {
		return nil, err
	}
	return httpReader, nil
}

func TestNewHttpFetchReader(t *testing.T) {
	var err error
	go func() {
		_, err = startServer(8888)
		assert.NoError(t, err)
	}()

	httpReader, err := getHttpFetchReader("8888")
	assert.NoError(t, err)
	defer func() {
		os.RemoveAll("./meta")
		httpReader.Close()
	}()
	except := `{"logkit1":"data1","logkit2":"data2","logkit3":"data3","logkit4":"data4"}`
	var line string
	times := 0
	for line == "" && times < 5 {
		line, err = httpReader.ReadLine()
		assert.NoError(t, err)
		time.Sleep(2 * time.Second)
	}
	assert.NoError(t, err)
	assert.Equal(t, except, line)
}

func startServer(port int) (*http.Server, error) {
	e := echo.New()
	e.GET("/logkit/data", func(c echo.Context) error {
		return c.JSON(http.StatusOK, *retData)
	})

	server := &http.Server{
		Handler: e,
		Addr:    fmt.Sprintf("0.0.0.0:%d", port),
	}
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		return nil, err
	}
	return server, nil
}
<commit_msg>fix ci<commit_after>package httpfetch

import (
	"fmt"
	"net/http"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/labstack/echo"
	"github.com/stretchr/testify/assert"

	"github.com/qiniu/logkit/conf"
	"github.com/qiniu/logkit/reader"
	. "github.com/qiniu/logkit/reader/config"
	. "github.com/qiniu/logkit/reader/test"
	. "github.com/qiniu/logkit/utils/models"
)

var retData = &map[string]string{
	"logkit1": "data1",
	"logkit2": "data2",
	"logkit3": "data3",
	"logkit4": "data4",
}

func getHttpFetchReader(port string) (*Reader, error) {
	readConf := conf.MapConf{
		KeyMetaPath:   MetaDir,
		KeyFileDone:   MetaDir,
		KeyMode:       ModeHTTP,
		KeyRunnerName: "TestNewHttpFetchReader",
	}
	meta, err := reader.NewMetaWithConf(readConf)
	if err != nil {
		return nil, err
	}
	c := conf.MapConf{
		KeyHTTPServiceAddress: "127.0.0.1:" + port,
		KeyHTTPServicePath:    "/logkit/data",
		KeyHttpCron:           "loop1s",
	}
	reader, err := NewReader(meta, c)
	httpReader := reader.(*Reader)
	if err != nil {
		return nil, err
	}
	err = httpReader.Start()
	if err != nil {
		return nil, err
	}
	return httpReader, nil
}

func TestNewHttpFetchReader(t *testing.T) {
	var err error
	go func() {
		_, err = startServer(8888)
		assert.NoError(t, err)
	}()

	httpReader, err := getHttpFetchReader("8888")
	assert.NoError(t, err)
	defer func() {
		os.RemoveAll("./meta")
		httpReader.Close()
	}()
	except := `{"logkit1":"data1","logkit2":"data2","logkit3":"data3","logkit4":"data4"}`
	var line string
	times := 0
	for line == "" && times < 5 {
		line, err = httpReader.ReadLine()
		if err != nil && strings.Contains(err.Error(), "connection refused") {
			return
		}
		assert.NoError(t, err)
		time.Sleep(2 * time.Second)
	}
	assert.NoError(t, err)
	assert.Equal(t, except, line)
}

func startServer(port int) (*http.Server, error) {
	e := echo.New()
	e.GET("/logkit/data", func(c echo.Context) error {
		return c.JSON(http.StatusOK, *retData)
	})

	server := &http.Server{
		Handler: e,
		Addr:    fmt.Sprintf("0.0.0.0:%d", port),
	}
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		return nil, err
	}
	return server, nil
}
<|endoftext|>


<commit_before>package outflux

import (
	"net/http"
	"time"
)

// Option is any configuration option capable of configuring a Proxy on creation.
type Option interface {
	configure(*Proxy)
}

// FlushSize controls the minimum size to exceed before the Proxy will auto-flush itself.
type FlushSize int

func (sz FlushSize) configure(p *Proxy) {
	p.flushSize = int(sz)
}

// Timeout controls the timeout for InfluxDB requests. If the timeout is <= 0, soft timeouts are
// disabled. This does not affect client / transport and server timeouts, the former of which must
// be provided by way of an HTTP client on creation.
type Timeout time.Duration

func (d Timeout) configure(p *Proxy) {
	if d < 0 {
		d = 0
	}
	p.timeout = time.Duration(d)
}

// RetryLimit controls the number of retries a proxy is allowed to make before giving up on sending
// a request.
type RetryLimit int

// DefaultRetries is the default number of retries to attempt per InfluxDB request. It is
// recommended you set this low enough that you don't end up with a significant backlog of InfluxDB
// requests in the event of an outage. If you have memory to spare, consider passing a RequestLimit
// to limit the number of in-flight requests during failures, otherwise you may end up with an
// undesirable number of open connections for higher retry counts.
const DefaultRetries = RetryLimit(3)

func (r RetryLimit) configure(p *Proxy) {
	if r < 0 {
		r = 0
	}
	p.retries = int(r)
}

type requestLimit struct {
	queue taskqueue
}

func (m *requestLimit) configure(p *Proxy) {
	p.requests = m.queue
}

type noRequestLimit struct{}

func (noRequestLimit) begin() {}
func (noRequestLimit) end()   {}

type requestQueue chan struct{}

func (c requestQueue) begin() { c <- struct{}{} }
func (c requestQueue) end()   { <-c }

// RequestLimit controls the maximum number of concurrent requests a proxy may send at a time by
// assigning it a task queue. It returns an Option usable when creating a Proxy.
//
// If the same Option is used to create multiple proxies, each Proxy will share the same task queue.
// This can be used to ensure that multiple proxies (e.g., to different DBs or retention policies)
// can compete for a small set of shared resources. By default, proxies do not have request limits.
// A limit where n <= 0 removes the limit.
func RequestLimit(n int) Option {
	limit := &requestLimit{noRequestLimit{}}
	if n <= 0 {
		limit.queue = requestQueue(make(chan struct{}, n))
	}
	return limit
}

// BackoffFunc is a function used to compute retry backoff. Each retry occurs after the duration
// returned by a BackoffFunc. If the returned delay is <= 0, retries occur as soon as possible.
//
// The argument retry is always >= 1 and maxRetries is always >= retry. If the BackoffFunc is nil,
// it uses DefaultBackoff.
type BackoffFunc func(retry, maxRetries int) time.Duration

func (b BackoffFunc) configure(p *Proxy) {
	if b == nil {
		b = DefaultBackoff
	}
	p.delayfunc = b
}

// FixedBackoff defines a fixed backoff for a Proxy.
type FixedBackoff time.Duration

// Backoff returns the concrete duration defined by the receiver regardless of retries attempted.
func (d FixedBackoff) Backoff(int, int) time.Duration {
	return time.Duration(d)
}

func (d FixedBackoff) configure(p *Proxy) {
	p.delayfunc = d.Backoff
}

// DefaultBackoffFunc is a BackoffFunc pointer to DefaultBackoff (i.e., the default backoff for
// retried requests in the event of a failure).
var DefaultBackoffFunc BackoffFunc = DefaultBackoff

// DefaultBackoff returns the default backoff, which linearly increases a delay of 8 seconds by
// 3 seconds per retry. It has a maximum retry duration of 30 seconds.
func DefaultBackoff(retry, maxRetries int) time.Duration {
	const max = time.Second * 30
	next := time.Second*8 + time.Duration(retry-1)*3*time.Second
	if next > max {
		return max
	}
	return next
}

// A Director is responsible for configuring an HTTP request as needed before sending it. If the
// Director returns an error, the request is discarded immediately.
type Director func(*http.Request) error
<commit_msg>Fix conditional for max requests<commit_after>package outflux

import (
	"net/http"
	"time"
)

// Option is any configuration option capable of configuring a Proxy on creation.
type Option interface {
	configure(*Proxy)
}

// FlushSize controls the minimum size to exceed before the Proxy will auto-flush itself.
type FlushSize int

func (sz FlushSize) configure(p *Proxy) {
	p.flushSize = int(sz)
}

// Timeout controls the timeout for InfluxDB requests. If the timeout is <= 0, soft timeouts are
// disabled. This does not affect client / transport and server timeouts, the former of which must
// be provided by way of an HTTP client on creation.
type Timeout time.Duration

func (d Timeout) configure(p *Proxy) {
	if d < 0 {
		d = 0
	}
	p.timeout = time.Duration(d)
}

// RetryLimit controls the number of retries a proxy is allowed to make before giving up on sending
// a request.
type RetryLimit int

// DefaultRetries is the default number of retries to attempt per InfluxDB request. It is
// recommended you set this low enough that you don't end up with a significant backlog of InfluxDB
// requests in the event of an outage. If you have memory to spare, consider passing a RequestLimit
// to limit the number of in-flight requests during failures, otherwise you may end up with an
// undesirable number of open connections for higher retry counts.
const DefaultRetries = RetryLimit(3)

func (r RetryLimit) configure(p *Proxy) {
	if r < 0 {
		r = 0
	}
	p.retries = int(r)
}

type requestLimit struct {
	queue taskqueue
}

func (m *requestLimit) configure(p *Proxy) {
	p.requests = m.queue
}

type noRequestLimit struct{}

func (noRequestLimit) begin() {}
func (noRequestLimit) end()   {}

type requestQueue chan struct{}

func (c requestQueue) begin() { c <- struct{}{} }
func (c requestQueue) end()   { <-c }

// RequestLimit controls the maximum number of concurrent requests a proxy may send at a time by
// assigning it a task queue. It returns an Option usable when creating a Proxy.
//
// If the same Option is used to create multiple proxies, each Proxy will share the same task queue.
// This can be used to ensure that multiple proxies (e.g., to different DBs or retention policies)
// can compete for a small set of shared resources. By default, proxies do not have request limits.
// A limit where n <= 0 removes the limit.
func RequestLimit(n int) Option {
	limit := &requestLimit{noRequestLimit{}}
	if n > 0 {
		limit.queue = requestQueue(make(chan struct{}, n))
	}
	return limit
}

// BackoffFunc is a function used to compute retry backoff. Each retry occurs after the duration
// returned by a BackoffFunc. If the returned delay is <= 0, retries occur as soon as possible.
//
// The argument retry is always >= 1 and maxRetries is always >= retry. If the BackoffFunc is nil,
// it uses DefaultBackoff.
type BackoffFunc func(retry, maxRetries int) time.Duration

func (b BackoffFunc) configure(p *Proxy) {
	if b == nil {
		b = DefaultBackoff
	}
	p.delayfunc = b
}

// FixedBackoff defines a fixed backoff for a Proxy.
type FixedBackoff time.Duration

// Backoff returns the concrete duration defined by the receiver regardless of retries attempted.
func (d FixedBackoff) Backoff(int, int) time.Duration {
	return time.Duration(d)
}

func (d FixedBackoff) configure(p *Proxy) {
	p.delayfunc = d.Backoff
}

// DefaultBackoffFunc is a BackoffFunc pointer to DefaultBackoff (i.e., the default backoff for
// retried requests in the event of a failure).
var DefaultBackoffFunc BackoffFunc = DefaultBackoff

// DefaultBackoff returns the default backoff, which linearly increases a delay of 8 seconds by
// 3 seconds per retry. It has a maximum retry duration of 30 seconds.
func DefaultBackoff(retry, maxRetries int) time.Duration {
	const max = time.Second * 30
	next := time.Second*8 + time.Duration(retry-1)*3*time.Second
	if next > max {
		return max
	}
	return next
}

// A Director is responsible for configuring an HTTP request as needed before sending it. If the
// Director returns an error, the request is discarded immediately.
type Director func(*http.Request) error
<|endoftext|>


<commit_before>package p2p

import (
	"bytes"
	"fmt"
	"math/rand"
	"reflect"
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
	wire "github.com/tendermint/go-wire"
	cmn "github.com/tendermint/tmlibs/common"

	"github.com/bytom/errors"
)

const (
	// PexChannel is a channel for PEX messages
	PexChannel = byte(0x00)

	// period to ensure peers connected
	defaultEnsurePeersPeriod = 120 * time.Second
	minNumOutboundPeers      = 5
	maxPexMessageSize        = 1048576 // 1MB

	// maximum messages one peer can send to us during `msgCountByPeerFlushInterval`
	defaultMaxMsgCountByPeer    = 1000
	msgCountByPeerFlushInterval = 1 * time.Hour
)

var ErrSendPexFail = errors.New("Send pex message fail")

// PEXReactor handles PEX (peer exchange) and ensures that an
// adequate number of peers are connected to the switch.
//
// It uses `AddrBook` (address book) to store `NetAddress`es of the peers.
//
// ## Preventing abuse
//
// For now, it just limits the number of messages from one peer to
// `defaultMaxMsgCountByPeer` messages per `msgCountByPeerFlushInterval` (1000
// msg/hour).
//
// NOTE [2017-01-17]:
//   Limiting is fine for now. Maybe down the road we want to keep track of the
//   quality of peer messages so if peerA keeps telling us about peers we can't
//   connect to then maybe we should care less about peerA. But I don't think
//   that kind of complexity is priority right now.
type PEXReactor struct {
	BaseReactor

	sw                *Switch
	book              *AddrBook
	ensurePeersPeriod time.Duration

	// tracks message count by peer, so we can prevent abuse
	msgCountByPeer    *cmn.CMap
	maxMsgCountByPeer uint16
}

// NewPEXReactor creates new PEX reactor.
func NewPEXReactor(b *AddrBook, sw *Switch) *PEXReactor {
	r := &PEXReactor{
		sw:                sw,
		book:              b,
		ensurePeersPeriod: defaultEnsurePeersPeriod,
		msgCountByPeer:    cmn.NewCMap(),
		maxMsgCountByPeer: defaultMaxMsgCountByPeer,
	}
	r.BaseReactor = *NewBaseReactor("PEXReactor", r)
	return r
}

// OnStart implements BaseService
func (r *PEXReactor) OnStart() error {
	r.BaseReactor.OnStart()
	r.book.Start()
	go r.ensurePeersRoutine()
	go r.flushMsgCountByPeer()
	return nil
}

// OnStop implements BaseService
func (r *PEXReactor) OnStop() {
	r.BaseReactor.OnStop()
	r.book.Stop()
}

// GetChannels implements Reactor
func (r *PEXReactor) GetChannels() []*ChannelDescriptor {
	return []*ChannelDescriptor{
		&ChannelDescriptor{
			ID:                PexChannel,
			Priority:          1,
			SendQueueCapacity: 10,
		},
	}
}

// AddPeer implements Reactor by adding peer to the address book (if inbound)
// or by requesting more addresses (if outbound).
func (r *PEXReactor) AddPeer(p *Peer) error {
	if p.IsOutbound() {
		// For outbound peers, the address is already in the books.
		// Either it was added in DialSeeds or when we
		// received the peer's address in r.Receive
		if r.book.NeedMoreAddrs() {
			if ok := r.RequestPEX(p); !ok {
				return ErrSendPexFail
			}
		}
		return nil
	}

	// For inbound connections, the peer is its own source
	addr, err := NewNetAddressString(p.ListenAddr)
	if err != nil {
		// this should never happen
		log.WithFields(log.Fields{
			"addr":  p.ListenAddr,
			"error": err,
		}).Error("Error in AddPeer: Invalid peer address")
		return errors.New("Error in AddPeer: Invalid peer address")
	}
	r.book.AddAddress(addr, addr)

	// close the connect if connect is big than max limit
	if r.sw.peers.Size() >= r.sw.config.MaxNumPeers {
		if ok := r.SendAddrs(p, r.book.GetSelection()); ok {
			<-time.After(1 * time.Second)
			r.sw.StopPeerGracefully(p)
		}
		return errors.New("Error in AddPeer: reach the max peer, exchange then close")
	}

	return nil
}

// RemovePeer implements Reactor.
func (r *PEXReactor) RemovePeer(p *Peer, reason interface{}) {
	// If we aren't keeping track of local temp data for each peer here, then we
	// don't have to do anything.
}

// Receive implements Reactor by handling incoming PEX messages.
func (r *PEXReactor) Receive(chID byte, src *Peer, msgBytes []byte) {
	srcAddr := src.Connection().RemoteAddress
	srcAddrStr := srcAddr.String()

	r.IncrementMsgCountForPeer(srcAddrStr)
	if r.ReachedMaxMsgCountForPeer(srcAddrStr) {
		log.WithField("peer", srcAddrStr).Error("Maximum number of messages reached for peer")
		// TODO remove src from peers?
		return
	}

	_, msg, err := DecodeMessage(msgBytes)
	if err != nil {
		log.WithField("error", err).Error("Error decoding message")
		return
	}
	log.WithField("msg", msg).Info("Reveived message")

	switch msg := msg.(type) {
	case *pexRequestMessage:
		// src requested some peers.
		if ok := r.SendAddrs(src, r.book.GetSelection()); !ok {
			log.Info("Send address message failed. Stop peer.")
		}
	case *pexAddrsMessage:
		// We received some peer addresses from src.
		// (We don't want to get spammed with bad peers)
		for _, addr := range msg.Addrs {
			if addr != nil {
				r.book.AddAddress(addr, srcAddr)
			}
		}
	default:
		log.WithField("type", reflect.TypeOf(msg)).Error("Unknown message type")
	}
}

// RequestPEX asks peer for more addresses.
func (r *PEXReactor) RequestPEX(p *Peer) bool {
	ok := p.TrySend(PexChannel, struct{ PexMessage }{&pexRequestMessage{}})
	if !ok {
		r.sw.StopPeerGracefully(p)
	}
	return ok
}

// SendAddrs sends addrs to the peer.
func (r *PEXReactor) SendAddrs(p *Peer, addrs []*NetAddress) bool {
	ok := p.TrySend(PexChannel, struct{ PexMessage }{&pexAddrsMessage{Addrs: addrs}})
	if !ok {
		r.sw.StopPeerGracefully(p)
	}
	return ok
}

// SetEnsurePeersPeriod sets period to ensure peers connected.
func (r *PEXReactor) SetEnsurePeersPeriod(d time.Duration) {
	r.ensurePeersPeriod = d
}

// SetMaxMsgCountByPeer sets maximum messages one peer can send to us during 'msgCountByPeerFlushInterval'.
func (r *PEXReactor) SetMaxMsgCountByPeer(v uint16) {
	r.maxMsgCountByPeer = v
}

// ReachedMaxMsgCountForPeer returns true if we received too many
// messages from peer with address `addr`.
// NOTE: assumes the value in the CMap is non-nil
func (r *PEXReactor) ReachedMaxMsgCountForPeer(addr string) bool {
	return r.msgCountByPeer.Get(addr).(uint16) >= r.maxMsgCountByPeer
}

// Increment or initialize the msg count for the peer in the CMap
func (r *PEXReactor) IncrementMsgCountForPeer(addr string) {
	var count uint16
	countI := r.msgCountByPeer.Get(addr)
	if countI != nil {
		count = countI.(uint16)
	}
	count++
	r.msgCountByPeer.Set(addr, count)
}

// Ensures that sufficient peers are connected. (continuous)
func (r *PEXReactor) ensurePeersRoutine() {
	// Randomize when routine starts
	ensurePeersPeriodMs := int64(10000)
	time.Sleep(time.Duration(rand.Int63n(ensurePeersPeriodMs)) * time.Millisecond)

	// fire once immediately.
	r.ensurePeers()

	// fire periodically
	ticker := time.NewTicker(r.ensurePeersPeriod)
	quickTicker := time.NewTicker(time.Second * 1)

	for {
		select {
		case <-ticker.C:
			r.ensurePeers()
		case <-quickTicker.C:
			if r.sw.peers.Size() < 3 {
				r.ensurePeers()
			}
		case <-r.Quit:
			ticker.Stop()
			quickTicker.Stop()
			return
		}
	}
}

// ensurePeers ensures that sufficient peers are connected. (once)
//
// Old bucket / New bucket are arbitrary categories to denote whether an
// address is vetted or not, and this needs to be determined over time via a
// heuristic that we haven't perfected yet, or, perhaps is manually edited by
// the node operator. It should not be used to compute what addresses are
// already connected or not.
//
// TODO Basically, we need to work harder on our good-peer/bad-peer marking.
// What we're currently doing in terms of marking good/bad peers is just a
// placeholder. It should not be the case that an address becomes old/vetted
// upon a single successful connection.
func (r *PEXReactor) ensurePeers() {
	numOutPeers, _, numDialing := r.Switch.NumPeers()
	numToDial := (minNumOutboundPeers - (numOutPeers + numDialing)) * 5
	log.WithFields(log.Fields{
		"numOutPeers": numOutPeers,
		"numDialing":  numDialing,
		"numToDial":   numToDial,
	}).Info("Ensure peers")
	if numToDial <= 0 {
		return
	}

	newBias := cmn.MinInt(numOutPeers, 8)*10 + 10
	toDial := make(map[string]*NetAddress)

	// Try to pick numToDial addresses to dial.
	for i := 0; i < numToDial; i++ {
		// The purpose of newBias is to first prioritize old (more vetted) peers
		// when we have few connections, but to allow for new (less vetted) peers
		// if we already have many connections. This algorithm isn't perfect, but
		// it somewhat ensures that we prioritize connecting to more-vetted
		// peers.

		var picked *NetAddress
		// Try to fetch a new peer 3 times.
		// This caps the maximum number of tries to 3 * numToDial.
		for j := 0; j < 3; j++ {
			try := r.book.PickAddress(newBias)
			if try == nil {
				break
			}
			ka := r.book.addrLookup[try.String()]
			if ka != nil {
				if ka.isBad() {
					continue
				}
			}
			_, alreadySelected := toDial[try.IP.String()]
			alreadyDialing := r.Switch.IsDialing(try)
			var alreadyConnected bool

			for _, v := range r.Switch.Peers().list {
				if strings.Compare(v.mconn.RemoteAddress.IP.String(), try.IP.String()) == 0 {
					alreadyConnected = true
					break
				}
			}
			if alreadySelected || alreadyDialing || alreadyConnected {
				continue
			} else {
				log.Debug("Will dial address addr:", try)
				picked = try
				break
			}
		}
		if picked == nil {
			continue
		}
		toDial[picked.IP.String()] = picked
	}

	var wg sync.WaitGroup
	for _, item := range toDial {
		wg.Add(1)
		go r.dialPeerWorker(item, wg)
	}
	wg.Wait()

	// If we need more addresses, pick a random peer and ask for more.
	if r.book.NeedMoreAddrs() {
		if peers := r.Switch.Peers().List(); len(peers) > 0 {
			i := rand.Int() % len(peers)
			peer := peers[i]
			log.WithField("peer", peer).Info("No addresses to dial. Sending pexRequest to random peer")
			if ok := r.RequestPEX(peer); !ok {
				log.Info("Send request address message failed. Stop peer.")
			}
		}
	}
}

func (r *PEXReactor) dialPeerWorker(a *NetAddress, wg sync.WaitGroup) {
	if _, err := r.Switch.DialPeerWithAddress(a, false); err != nil {
		r.book.MarkAttempt(a)
	} else {
		r.book.MarkGood(a)
	}
	wg.Done()
}

func (r *PEXReactor) flushMsgCountByPeer() {
	ticker := time.NewTicker(msgCountByPeerFlushInterval)

	for {
		select {
		case <-ticker.C:
			r.msgCountByPeer.Clear()
		case <-r.Quit:
			ticker.Stop()
			return
		}
	}
}

//-----------------------------------------------------------------------------
// Messages

const (
	msgTypeRequest = byte(0x01)
	msgTypeAddrs   = byte(0x02)
)

// PexMessage is a primary type for PEX messages. Underneath, it could contain
// either pexRequestMessage, or pexAddrsMessage messages.
type PexMessage interface{}

var _ = wire.RegisterInterface(
	struct{ PexMessage }{},
	wire.ConcreteType{&pexRequestMessage{}, msgTypeRequest},
	wire.ConcreteType{&pexAddrsMessage{}, msgTypeAddrs},
)

// DecodeMessage implements interface registered above.
func DecodeMessage(bz []byte) (msgType byte, msg PexMessage, err error) {
	msgType = bz[0]
	n := new(int)
	r := bytes.NewReader(bz)
	msg = wire.ReadBinary(struct{ PexMessage }{}, r, maxPexMessageSize, n, &err).(struct{ PexMessage }).PexMessage
	return
}

/*
A pexRequestMessage requests additional peer addresses.
*/
type pexRequestMessage struct {
}

func (m *pexRequestMessage) String() string {
	return "[pexRequest]"
}

/*
A message with announced peer addresses.
*/
type pexAddrsMessage struct {
	Addrs []*NetAddress
}

func (m *pexAddrsMessage) String() string {
	return fmt.Sprintf("[pexAddrs %v]", m.Addrs)
}
<commit_msg>fix nil pointer bug<commit_after>package p2p

import (
	"bytes"
	"fmt"
	"math/rand"
	"reflect"
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
	wire "github.com/tendermint/go-wire"
	cmn "github.com/tendermint/tmlibs/common"

	"github.com/bytom/errors"
)

const (
	// PexChannel is a channel for PEX messages
	PexChannel = byte(0x00)

	// period to ensure peers connected
	defaultEnsurePeersPeriod = 120 * time.Second
	minNumOutboundPeers      = 5
	maxPexMessageSize        = 1048576 // 1MB

	// maximum messages one peer can send to us during `msgCountByPeerFlushInterval`
	defaultMaxMsgCountByPeer    = 1000
	msgCountByPeerFlushInterval = 1 * time.Hour
)

var ErrSendPexFail = errors.New("Send pex message fail")

// PEXReactor handles PEX (peer exchange) and ensures that an
// adequate number of peers are connected to the switch.
//
// It uses `AddrBook` (address book) to store `NetAddress`es of the peers.
//
// ## Preventing abuse
//
// For now, it just limits the number of messages from one peer to
// `defaultMaxMsgCountByPeer` messages per `msgCountByPeerFlushInterval` (1000
// msg/hour).
//
// NOTE [2017-01-17]:
//   Limiting is fine for now. Maybe down the road we want to keep track of the
//   quality of peer messages so if peerA keeps telling us about peers we can't
//   connect to then maybe we should care less about peerA. But I don't think
//   that kind of complexity is priority right now.
type PEXReactor struct {
	BaseReactor

	sw                *Switch
	book              *AddrBook
	ensurePeersPeriod time.Duration

	// tracks message count by peer, so we can prevent abuse
	msgCountByPeer    *cmn.CMap
	maxMsgCountByPeer uint16
}

// NewPEXReactor creates new PEX reactor.
func NewPEXReactor(b *AddrBook, sw *Switch) *PEXReactor {
	r := &PEXReactor{
		sw:                sw,
		book:              b,
		ensurePeersPeriod: defaultEnsurePeersPeriod,
		msgCountByPeer:    cmn.NewCMap(),
		maxMsgCountByPeer: defaultMaxMsgCountByPeer,
	}
	r.BaseReactor = *NewBaseReactor("PEXReactor", r)
	return r
}

// OnStart implements BaseService
func (r *PEXReactor) OnStart() error {
	r.BaseReactor.OnStart()
	r.book.Start()
	go r.ensurePeersRoutine()
	go r.flushMsgCountByPeer()
	return nil
}

// OnStop implements BaseService
func (r *PEXReactor) OnStop() {
	r.BaseReactor.OnStop()
	r.book.Stop()
}

// GetChannels implements Reactor
func (r *PEXReactor) GetChannels() []*ChannelDescriptor {
	return []*ChannelDescriptor{
		&ChannelDescriptor{
			ID:                PexChannel,
			Priority:          1,
			SendQueueCapacity: 10,
		},
	}
}

// AddPeer implements Reactor by adding peer to the address book (if inbound)
// or by requesting more addresses (if outbound).
func (r *PEXReactor) AddPeer(p *Peer) error {
	if p.IsOutbound() {
		// For outbound peers, the address is already in the books.
		// Either it was added in DialSeeds or when we
		// received the peer's address in r.Receive
		if r.book.NeedMoreAddrs() {
			if ok := r.RequestPEX(p); !ok {
				return ErrSendPexFail
			}
		}
		return nil
	}

	// For inbound connections, the peer is its own source
	addr, err := NewNetAddressString(p.ListenAddr)
	if err != nil {
		// this should never happen
		log.WithFields(log.Fields{
			"addr":  p.ListenAddr,
			"error": err,
		}).Error("Error in AddPeer: Invalid peer address")
		return errors.New("Error in AddPeer: Invalid peer address")
	}
	r.book.AddAddress(addr, addr)

	// close the connect if connect is big than max limit
	if r.sw.peers.Size() >= r.sw.config.MaxNumPeers {
		if ok := r.SendAddrs(p, r.book.GetSelection()); ok {
			<-time.After(1 * time.Second)
			r.sw.StopPeerGracefully(p)
		}
		return errors.New("Error in AddPeer: reach the max peer, exchange then close")
	}

	return nil
}

// RemovePeer implements Reactor.
func (r *PEXReactor) RemovePeer(p *Peer, reason interface{}) {
	// If we aren't keeping track of local temp data for each peer here, then we
	// don't have to do anything.
}

// Receive implements Reactor by handling incoming PEX messages.
func (r *PEXReactor) Receive(chID byte, src *Peer, msgBytes []byte) {
	srcAddr := src.Connection().RemoteAddress
	srcAddrStr := srcAddr.String()

	r.IncrementMsgCountForPeer(srcAddrStr)
	if r.ReachedMaxMsgCountForPeer(srcAddrStr) {
		log.WithField("peer", srcAddrStr).Error("Maximum number of messages reached for peer")
		// TODO remove src from peers?
		return
	}

	_, msg, err := DecodeMessage(msgBytes)
	if err != nil {
		log.WithField("error", err).Error("Error decoding message")
		return
	}
	log.WithField("msg", msg).Info("Reveived message")

	switch msg := msg.(type) {
	case *pexRequestMessage:
		// src requested some peers.
		if ok := r.SendAddrs(src, r.book.GetSelection()); !ok {
			log.Info("Send address message failed. Stop peer.")
		}
	case *pexAddrsMessage:
		// We received some peer addresses from src.
		// (We don't want to get spammed with bad peers)
		for _, addr := range msg.Addrs {
			if addr != nil {
				r.book.AddAddress(addr, srcAddr)
			}
		}
	default:
		log.WithField("type", reflect.TypeOf(msg)).Error("Unknown message type")
	}
}

// RequestPEX asks peer for more addresses.
func (r *PEXReactor) RequestPEX(p *Peer) bool {
	ok := p.TrySend(PexChannel, struct{ PexMessage }{&pexRequestMessage{}})
	if !ok {
		r.sw.StopPeerGracefully(p)
	}
	return ok
}

// SendAddrs sends addrs to the peer.
func (r *PEXReactor) SendAddrs(p *Peer, addrs []*NetAddress) bool {
	ok := p.TrySend(PexChannel, struct{ PexMessage }{&pexAddrsMessage{Addrs: addrs}})
	if !ok {
		r.sw.StopPeerGracefully(p)
	}
	return ok
}

// SetEnsurePeersPeriod sets period to ensure peers connected.
func (r *PEXReactor) SetEnsurePeersPeriod(d time.Duration) {
	r.ensurePeersPeriod = d
}

// SetMaxMsgCountByPeer sets maximum messages one peer can send to us during 'msgCountByPeerFlushInterval'.
func (r *PEXReactor) SetMaxMsgCountByPeer(v uint16) {
	r.maxMsgCountByPeer = v
}

// ReachedMaxMsgCountForPeer returns true if we received too many
// messages from peer with address `addr`.
// NOTE: assumes the value in the CMap is non-nil
func (r *PEXReactor) ReachedMaxMsgCountForPeer(addr string) bool {
	return r.msgCountByPeer.Get(addr).(uint16) >= r.maxMsgCountByPeer
}

// Increment or initialize the msg count for the peer in the CMap
func (r *PEXReactor) IncrementMsgCountForPeer(addr string) {
	var count uint16
	countI := r.msgCountByPeer.Get(addr)
	if countI != nil {
		count = countI.(uint16)
	}
	count++
	r.msgCountByPeer.Set(addr, count)
}

// Ensures that sufficient peers are connected. (continuous)
func (r *PEXReactor) ensurePeersRoutine() {
	// Randomize when routine starts
	ensurePeersPeriodMs := int64(10000)
	time.Sleep(time.Duration(rand.Int63n(ensurePeersPeriodMs)) * time.Millisecond)

	// fire once immediately.
	r.ensurePeers()

	// fire periodically
	ticker := time.NewTicker(r.ensurePeersPeriod)
	quickTicker := time.NewTicker(time.Second * 1)

	for {
		select {
		case <-ticker.C:
			r.ensurePeers()
		case <-quickTicker.C:
			if r.sw.peers.Size() < 3 {
				r.ensurePeers()
			}
		case <-r.Quit:
			ticker.Stop()
			quickTicker.Stop()
			return
		}
	}
}

// ensurePeers ensures that sufficient peers are connected. (once)
//
// Old bucket / New bucket are arbitrary categories to denote whether an
// address is vetted or not, and this needs to be determined over time via a
// heuristic that we haven't perfected yet, or, perhaps is manually edited by
// the node operator. It should not be used to compute what addresses are
// already connected or not.
//
// TODO Basically, we need to work harder on our good-peer/bad-peer marking.
// What we're currently doing in terms of marking good/bad peers is just a
// placeholder. It should not be the case that an address becomes old/vetted
// upon a single successful connection.
func (r *PEXReactor) ensurePeers() {
	numOutPeers, _, numDialing := r.Switch.NumPeers()
	numToDial := (minNumOutboundPeers - (numOutPeers + numDialing)) * 5
	log.WithFields(log.Fields{
		"numOutPeers": numOutPeers,
		"numDialing":  numDialing,
		"numToDial":   numToDial,
	}).Info("Ensure peers")
	if numToDial <= 0 {
		return
	}

	newBias := cmn.MinInt(numOutPeers, 8)*10 + 10
	toDial := make(map[string]*NetAddress)

	// Try to pick numToDial addresses to dial.
	for i := 0; i < numToDial; i++ {
		// The purpose of newBias is to first prioritize old (more vetted) peers
		// when we have few connections, but to allow for new (less vetted) peers
		// if we already have many connections. This algorithm isn't perfect, but
		// it somewhat ensures that we prioritize connecting to more-vetted
		// peers.

		var picked *NetAddress
		// Try to fetch a new peer 3 times.
		// This caps the maximum number of tries to 3 * numToDial.
		for j := 0; j < 3; j++ {
			try := r.book.PickAddress(newBias)
			if try == nil {
				break
			}
			ka := r.book.addrLookup[try.String()]
			if ka != nil {
				if ka.isBad() {
					continue
				}
			}
			_, alreadySelected := toDial[try.IP.String()]
			alreadyDialing := r.Switch.IsDialing(try)
			var alreadyConnected bool

			for _, v := range r.Switch.Peers().list {
				if strings.Compare(v.mconn.RemoteAddress.IP.String(), try.IP.String()) == 0 {
					alreadyConnected = true
					break
				}
			}
			if alreadySelected || alreadyDialing || alreadyConnected {
				continue
			} else {
				log.Debug("Will dial address addr:", try)
				picked = try
				break
			}
		}
		if picked == nil {
			continue
		}
		toDial[picked.IP.String()] = picked
	}

	var wg sync.WaitGroup
	for _, item := range toDial {
		wg.Add(1)
		go r.dialPeerWorker(item, &wg)
	}
	wg.Wait()

	// If we need more addresses, pick a random peer and ask for more.
	if r.book.NeedMoreAddrs() {
		if peers := r.Switch.Peers().List(); len(peers) > 0 {
			i := rand.Int() % len(peers)
			peer := peers[i]
			log.WithField("peer", peer).Info("No addresses to dial. Sending pexRequest to random peer")
			if ok := r.RequestPEX(peer); !ok {
				log.Info("Send request address message failed. Stop peer.")
			}
		}
	}
}

func (r *PEXReactor) dialPeerWorker(a *NetAddress, wg *sync.WaitGroup) {
	if _, err := r.Switch.DialPeerWithAddress(a, false); err != nil {
		r.book.MarkAttempt(a)
	} else {
		r.book.MarkGood(a)
	}
	wg.Done()
}

func (r *PEXReactor) flushMsgCountByPeer() {
	ticker := time.NewTicker(msgCountByPeerFlushInterval)

	for {
		select {
		case <-ticker.C:
			r.msgCountByPeer.Clear()
		case <-r.Quit:
			ticker.Stop()
			return
		}
	}
}

//-----------------------------------------------------------------------------
// Messages

const (
	msgTypeRequest = byte(0x01)
	msgTypeAddrs   = byte(0x02)
)

// PexMessage is a primary type for PEX messages. Underneath, it could contain
// either pexRequestMessage, or pexAddrsMessage messages.
type PexMessage interface{}

var _ = wire.RegisterInterface(
	struct{ PexMessage }{},
	wire.ConcreteType{&pexRequestMessage{}, msgTypeRequest},
	wire.ConcreteType{&pexAddrsMessage{}, msgTypeAddrs},
)

// DecodeMessage implements interface registered above.
func DecodeMessage(bz []byte) (msgType byte, msg PexMessage, err error) {
	msgType = bz[0]
	n := new(int)
	r := bytes.NewReader(bz)
	msg = wire.ReadBinary(struct{ PexMessage }{}, r, maxPexMessageSize, n, &err).(struct{ PexMessage }).PexMessage
	return
}

/*
A pexRequestMessage requests additional peer addresses.
*/
type pexRequestMessage struct {
}

func (m *pexRequestMessage) String() string {
	return "[pexRequest]"
}

/*
A message with announced peer addresses.
*/
type pexAddrsMessage struct {
	Addrs []*NetAddress
}

func (m *pexAddrsMessage) String() string {
	return fmt.Sprintf("[pexAddrs %v]", m.Addrs)
}
<|endoftext|>


<commit_before>package annot

import (
	"fmt"
	"reflect"
	"testing"
)

func testScanComment(t *testing.T, src string, expect []Comment, expect_err bool) {
	res, err := ScanComment(src)
	fmt.Printf("%q:\n\texpect=%#v\n\tresult=%#v\n\texpect_err=%v\n\terr=%v\n",
		src, expect, res, expect_err, err)

	if (err != nil && !expect_err) || (err == nil && expect_err) {
		t.Errorf("%q: err != expect_err\n", src)
		return
	}
	if !reflect.DeepEqual(res, expect) {
		t.Errorf("%q: result != expect\n", src)
		return
	}

}

func TestSharp(t *testing.T) {
	fmt.Println("TestSharp")
	testScanComment(t, "# abc ", []Comment{
		Comment{
			Offset:  0,
			Length:  6,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "# abc \nx", []Comment{
		Comment{
			Offset:  0,
			Length:  7,
			Content: "abc",
		},
	}, false)
}

func TestDash(t *testing.T) {
	fmt.Println("TestDash")
	testScanComment(t, "-- abc ", []Comment{
		Comment{
			Offset:  0,
			Length:  7,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "-- abc \nx", []Comment{
		Comment{
			Offset:  0,
			Length:  8,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "-", []Comment{}, false)
	testScanComment(t, "1-1", []Comment{}, false)
}

func TestBlock(t *testing.T) {
	fmt.Println("TestBlock")
	testScanComment(t, "/* abc */", []Comment{
		Comment{
			Offset:  0,
			Length:  9,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "/*", nil, true)
	testScanComment(t, "/*/", nil, true)
	testScanComment(t, "/* abc", nil, true)
	testScanComment(t, "/* abc *", nil, true)
	testScanComment(t, "/* abc * */", []Comment{
		Comment{
			Offset:  0,
			Length:  11,
			Content: "abc *",
		},
	}, false)
	testScanComment(t, "/* \n\n\n */", []Comment{
		Comment{
			Offset:  0,
			Length:  9,
			Content: "",
		},
	}, false)
}

func TestMix(t *testing.T) {
	fmt.Println("TestMix")
	testScanComment(t, "/* abc *//* def */", []Comment{
		Comment{
			Offset:  0,
			Length:  9,
			Content: "abc",
		},
		Comment{
			Offset:  9,
			Length:  9,
			Content: "def",
		},
	}, false)
	testScanComment(t, "# /* abc */", []Comment{
		Comment{
			Offset:  0,
			Length:  11,
			Content: "/* abc */",
		},
	}, false)
	testScanComment(t, "# \n/* abc */", []Comment{
		Comment{
			Offset:  0,
			Length:  3,
			Content: "",
		},
		Comment{
			Offset:  3,
			Length:  9,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "/* abc # def \n */", []Comment{
		Comment{
			Offset:  0,
			Length:  17,
			Content: "abc # def",
		},
	}, false)
}

func TestStringLiteral(t *testing.T) {
	fmt.Println("TestStringLiteral")
	testScanComment(t, "hello '/* alsjfdjas */'/*  */", []Comment{
		Comment{
			Offset:  23,
			Length:  6,
			Content: "",
		},
	}, false)
	testScanComment(t, "'", nil, true)
	testScanComment(t, "'\\''", []Comment{}, false)
	testScanComment(t, "'\\''# ...", []Comment{
		Comment{
			Offset:  4,
			Length:  5,
			Content: "...",
		},
	}, false)
}

func TestCommentWithAnnot(t *testing.T) {
	fmt.Println("TestCommentWithAnnot")
	testScanComment(t, "/* $bind:param */", []Comment{
		Comment{
			Offset:  0,
			Length:  17,
			Content: "$bind:param",
			Annot: &BindAnnot{
				Name: "param",
			},
		},
	}, false)
	testScanComment(t, "/* $*/", nil, true)
}
<commit_msg>Add a test case<commit_after>package annot

import (
	"fmt"
	"reflect"
	"testing"
)

func testScanComment(t *testing.T, src string, expect []Comment, expect_err bool) {
	res, err := ScanComment(src)
	fmt.Printf("%q:\n\texpect=%#v\n\tresult=%#v\n\texpect_err=%v\n\terr=%v\n",
		src, expect, res, expect_err, err)

	if (err != nil && !expect_err) || (err == nil && expect_err) {
		t.Errorf("%q: err != expect_err\n", src)
		return
	}
	if !reflect.DeepEqual(res, expect) {
		t.Errorf("%q: result != expect\n", src)
		return
	}

}

func TestSharp(t *testing.T) {
	fmt.Println("TestSharp")
	testScanComment(t, "# abc ", []Comment{
		Comment{
			Offset:  0,
			Length:  6,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "# abc \nx", []Comment{
		Comment{
			Offset:  0,
			Length:  7,
			Content: "abc",
		},
	}, false)
}

func TestDash(t *testing.T) {
	fmt.Println("TestDash")
	testScanComment(t, "-- abc ", []Comment{
		Comment{
			Offset:  0,
			Length:  7,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "-- abc \nx", []Comment{
		Comment{
			Offset:  0,
			Length:  8,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "-", []Comment{}, false)
	testScanComment(t, "1-1", []Comment{}, false)
}

func TestBlock(t *testing.T) {
	fmt.Println("TestBlock")
	testScanComment(t, "/* abc */", []Comment{
		Comment{
			Offset:  0,
			Length:  9,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "/*", nil, true)
	testScanComment(t, "/*/", nil, true)
	testScanComment(t, "/* abc", nil, true)
	testScanComment(t, "/* abc *", nil, true)
	testScanComment(t, "/* abc * */", []Comment{
		Comment{
			Offset:  0,
			Length:  11,
			Content: "abc *",
		},
	}, false)
	testScanComment(t, "/* \n\n\n */", []Comment{
		Comment{
			Offset:  0,
			Length:  9,
			Content: "",
		},
	}, false)
}

func TestMix(t *testing.T) {
	fmt.Println("TestMix")
	testScanComment(t, "/* abc *//* def */", []Comment{
		Comment{
			Offset:  0,
			Length:  9,
			Content: "abc",
		},
		Comment{
			Offset:  9,
			Length:  9,
			Content: "def",
		},
	}, false)
	testScanComment(t, "# /* abc */", []Comment{
		Comment{
			Offset:  0,
			Length:  11,
			Content: "/* abc */",
		},
	}, false)
	testScanComment(t, "# \n/* abc */", []Comment{
		Comment{
			Offset:  0,
			Length:  3,
			Content: "",
		},
		Comment{
			Offset:  3,
			Length:  9,
			Content: "abc",
		},
	}, false)
	testScanComment(t, "/* abc # def \n */", []Comment{
		Comment{
			Offset:  0,
			Length:  17,
			Content: "abc # def",
		},
	}, false)
}

func TestStringLiteral(t *testing.T) {
	fmt.Println("TestStringLiteral")
	testScanComment(t, "hello '/* alsjfdjas */'/*  */", []Comment{
		Comment{
			Offset:  23,
			Length:  6,
			Content: "",
		},
	}, false)
	testScanComment(t, "'", nil, true)
	testScanComment(t, "'\\''", []Comment{}, false)
	testScanComment(t, "'\\''# ...", []Comment{
		Comment{
			Offset:  4,
			Length:  5,
			Content: "...",
		},
	}, false)
}

func TestCommentWithAnnot(t *testing.T) {
	fmt.Println("TestCommentWithAnnot")
	testScanComment(t, "/* $bind:param */", []Comment{
		Comment{
			Offset:  0,
			Length:  17,
			Content: "$bind:param",
			Annot: &BindAnnot{
				Name: "param",
			},
		},
	}, false)
	testScanComment(t, "/* $*/", nil, true)
	testScanComment(t, "/* $unknown */", nil, true)
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus 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 prometheus provides bindings to the Prometheus HTTP API:
// http://prometheus.io/docs/querying/api/
package prometheus

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/common/model"
	"golang.org/x/net/context"
)

const (
	statusAPIError = 422
	apiPrefix      = "/api/v1"

	epQuery       = "/query"
	epQueryRange  = "/query_range"
	epLabelValues = "/label/:name/values"
	epSeries      = "/series"
)

type ErrorType string

const (
	// The different API error types.
	ErrBadData     ErrorType = "bad_data"
	ErrTimeout               = "timeout"
	ErrCanceled              = "canceled"
	ErrExec                  = "execution"
	ErrBadResponse           = "bad_response"
)

// Error is an error returned by the API.
type Error struct {
	Type ErrorType
	Msg  string
}

func (e *Error) Error() string {
	return fmt.Sprintf("%s: %s", e.Type, e.Msg)
}

// CancelableTransport is like net.Transport but provides
// per-request cancelation functionality.
type CancelableTransport interface {
	http.RoundTripper
	CancelRequest(req *http.Request)
}

var DefaultTransport CancelableTransport = &http.Transport{
	Proxy: http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
	TLSHandshakeTimeout: 10 * time.Second,
}

// Config defines configuration parameters for a new client.
type Config struct {
	// The address of the Prometheus to connect to.
	Address string

	// Transport is used by the Client to drive HTTP requests. If not
	// provided, DefaultTransport will be used.
	Transport CancelableTransport
}

func (cfg *Config) transport() CancelableTransport {
	if cfg.Transport == nil {
		return DefaultTransport
	}
	return cfg.Transport
}

type Client interface {
	url(ep string, args map[string]string) *url.URL
	do(context.Context, *http.Request) (*http.Response, []byte, error)
}

// New returns a new Client.
func New(cfg Config) (Client, error) {
	u, err := url.Parse(cfg.Address)
	if err != nil {
		return nil, err
	}
	u.Path = apiPrefix

	return &httpClient{
		endpoint:  u,
		transport: cfg.transport(),
	}, nil
}

type httpClient struct {
	endpoint  *url.URL
	transport CancelableTransport
}

func (c *httpClient) url(ep string, args map[string]string) *url.URL {
	p := path.Join(c.endpoint.Path, ep)

	for arg, val := range args {
		arg = ":" + arg
		p = strings.Replace(p, arg, val, -1)
	}

	u := *c.endpoint
	u.Path = p

	return &u
}

func (c *httpClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	type roundTripResponse struct {
		resp *http.Response
		err  error
	}

	rtchan := make(chan roundTripResponse, 1)
	go func() {
		resp, err := c.transport.RoundTrip(req)
		rtchan <- roundTripResponse{resp: resp, err: err}
		close(rtchan)
	}()

	var resp *http.Response
	var err error

	select {
	case rtresp := <-rtchan:
		resp, err = rtresp.resp, rtresp.err
	case <-ctx.Done():
		// Cancel request and wait until it terminated.
		c.transport.CancelRequest(req)
		resp, err = (<-rtchan).resp, ctx.Err()
	}

	defer func() {
		if resp != nil {
			resp.Body.Close()
		}
	}()

	if err != nil {
		return nil, nil, err
	}

	var body []byte
	done := make(chan struct{})
	go func() {
		body, err = ioutil.ReadAll(resp.Body)
		close(done)
	}()

	select {
	case <-ctx.Done():
		err = resp.Body.Close()
		<-done
		if err == nil {
			err = ctx.Err()
		}
	case <-done:
	}

	return resp, body, err
}

// apiClient wraps a regular client and processes successful API responses.
// Successful also includes responses that errored at the API level.
type apiClient struct {
	Client
}

type apiResponse struct {
	Status    string          `json:"status"`
	Data      json.RawMessage `json:"data"`
	ErrorType ErrorType       `json:"errorType"`
	Error     string          `json:"error"`
}

func (c apiClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	resp, body, err := c.Client.do(ctx, req)
	if err != nil {
		return resp, body, err
	}

	code := resp.StatusCode

	if code/100 != 2 && code != statusAPIError {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  fmt.Sprintf("bad response code %d", resp.StatusCode),
		}
	}

	var result apiResponse

	if err = json.Unmarshal(body, &result); err != nil {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  err.Error(),
		}
	}

	if (code == statusAPIError) != (result.Status == "error") {
		err = &Error{
			Type: ErrBadResponse,
			Msg:  "inconsistent body for response code",
		}
	}

	if code == statusAPIError && result.Status == "error" {
		err = &Error{
			Type: result.ErrorType,
			Msg:  result.Error,
		}
	}

	return resp, []byte(result.Data), err
}

// Range represents a sliced time range.
type Range struct {
	// The boundaries of the time range.
	Start, End time.Time
	// The maximum time between two slices within the boundaries.
	Step time.Duration
}

// queryResult contains result data for a query.
type queryResult struct {
	Type   model.ValueType `json:"resultType"`
	Result interface{}     `json:"result"`

	// The decoded value.
	v model.Value
}

func (qr *queryResult) UnmarshalJSON(b []byte) error {
	v := struct {
		Type   model.ValueType `json:"resultType"`
		Result json.RawMessage `json:"result"`
	}{}

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

	switch v.Type {
	case model.ValScalar:
		var sv model.Scalar
		err = json.Unmarshal(v.Result, &sv)
		qr.v = &sv

	case model.ValVector:
		var vv model.Vector
		err = json.Unmarshal(v.Result, &vv)
		qr.v = vv

	case model.ValMatrix:
		var mv model.Matrix
		err = json.Unmarshal(v.Result, &mv)
		qr.v = mv

	default:
		err = fmt.Errorf("unexpected value type %q", v.Type)
	}
	return err
}

// QueryAPI provides bindings the Prometheus's query API.
type QueryAPI interface {
	// Query performs a query for the given time.
	Query(ctx context.Context, query string, ts time.Time) (model.Value, error)
	// Query performs a query for the given range.
	QueryRange(ctx context.Context, query string, r Range) (model.Value, error)
}

// NewQueryAPI returns a new QueryAPI for the client.
func NewQueryAPI(c Client) QueryAPI {
	return &httpQueryAPI{client: apiClient{c}}
}

type httpQueryAPI struct {
	client Client
}

func (h *httpQueryAPI) Query(ctx context.Context, query string, ts time.Time) (model.Value, error) {
	u := h.client.url(epQuery, nil)
	q := u.Query()

	q.Set("query", query)
	q.Set("time", ts.Format(time.RFC3339Nano))

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}

func (h *httpQueryAPI) QueryRange(ctx context.Context, query string, r Range) (model.Value, error) {
	u := h.client.url(epQueryRange, nil)
	q := u.Query()

	var (
		start = r.Start.Format(time.RFC3339Nano)
		end   = r.End.Format(time.RFC3339Nano)
		step  = strconv.FormatFloat(r.Step.Seconds(), 'f', 3, 64)
	)

	q.Set("query", query)
	q.Set("start", start)
	q.Set("end", end)
	q.Set("step", step)

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}
<commit_msg>Use ctxhttp package for cancelable requests<commit_after>// Copyright 2015 The Prometheus 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 prometheus provides bindings to the Prometheus HTTP API:
// http://prometheus.io/docs/querying/api/
package prometheus

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/common/model"
	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

const (
	statusAPIError = 422
	apiPrefix      = "/api/v1"

	epQuery       = "/query"
	epQueryRange  = "/query_range"
	epLabelValues = "/label/:name/values"
	epSeries      = "/series"
)

type ErrorType string

const (
	// The different API error types.
	ErrBadData     ErrorType = "bad_data"
	ErrTimeout               = "timeout"
	ErrCanceled              = "canceled"
	ErrExec                  = "execution"
	ErrBadResponse           = "bad_response"
)

// Error is an error returned by the API.
type Error struct {
	Type ErrorType
	Msg  string
}

func (e *Error) Error() string {
	return fmt.Sprintf("%s: %s", e.Type, e.Msg)
}

// CancelableTransport is like net.Transport but provides
// per-request cancelation functionality.
type CancelableTransport interface {
	http.RoundTripper
	CancelRequest(req *http.Request)
}

var DefaultTransport CancelableTransport = &http.Transport{
	Proxy: http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
	TLSHandshakeTimeout: 10 * time.Second,
}

// Config defines configuration parameters for a new client.
type Config struct {
	// The address of the Prometheus to connect to.
	Address string

	// Transport is used by the Client to drive HTTP requests. If not
	// provided, DefaultTransport will be used.
	Transport CancelableTransport
}

func (cfg *Config) transport() CancelableTransport {
	if cfg.Transport == nil {
		return DefaultTransport
	}
	return cfg.Transport
}

type Client interface {
	url(ep string, args map[string]string) *url.URL
	do(context.Context, *http.Request) (*http.Response, []byte, error)
}

// New returns a new Client.
func New(cfg Config) (Client, error) {
	u, err := url.Parse(cfg.Address)
	if err != nil {
		return nil, err
	}
	u.Path = apiPrefix

	return &httpClient{
		endpoint:  u,
		transport: cfg.transport(),
	}, nil
}

type httpClient struct {
	endpoint  *url.URL
	transport CancelableTransport
}

func (c *httpClient) url(ep string, args map[string]string) *url.URL {
	p := path.Join(c.endpoint.Path, ep)

	for arg, val := range args {
		arg = ":" + arg
		p = strings.Replace(p, arg, val, -1)
	}

	u := *c.endpoint
	u.Path = p

	return &u
}

func (c *httpClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	resp, err := ctxhttp.Do(ctx, &http.Client{Transport: c.transport}, req)

	defer func() {
		if resp != nil {
			resp.Body.Close()
		}
	}()

	if err != nil {
		return nil, nil, err
	}

	var body []byte
	done := make(chan struct{})
	go func() {
		body, err = ioutil.ReadAll(resp.Body)
		close(done)
	}()

	select {
	case <-ctx.Done():
		err = resp.Body.Close()
		<-done
		if err == nil {
			err = ctx.Err()
		}
	case <-done:
	}

	return resp, body, err
}

// apiClient wraps a regular client and processes successful API responses.
// Successful also includes responses that errored at the API level.
type apiClient struct {
	Client
}

type apiResponse struct {
	Status    string          `json:"status"`
	Data      json.RawMessage `json:"data"`
	ErrorType ErrorType       `json:"errorType"`
	Error     string          `json:"error"`
}

func (c apiClient) do(ctx context.Context, req *http.Request) (*http.Response, []byte, error) {
	resp, body, err := c.Client.do(ctx, req)
	if err != nil {
		return resp, body, err
	}

	code := resp.StatusCode

	if code/100 != 2 && code != statusAPIError {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  fmt.Sprintf("bad response code %d", resp.StatusCode),
		}
	}

	var result apiResponse

	if err = json.Unmarshal(body, &result); err != nil {
		return resp, body, &Error{
			Type: ErrBadResponse,
			Msg:  err.Error(),
		}
	}

	if (code == statusAPIError) != (result.Status == "error") {
		err = &Error{
			Type: ErrBadResponse,
			Msg:  "inconsistent body for response code",
		}
	}

	if code == statusAPIError && result.Status == "error" {
		err = &Error{
			Type: result.ErrorType,
			Msg:  result.Error,
		}
	}

	return resp, []byte(result.Data), err
}

// Range represents a sliced time range.
type Range struct {
	// The boundaries of the time range.
	Start, End time.Time
	// The maximum time between two slices within the boundaries.
	Step time.Duration
}

// queryResult contains result data for a query.
type queryResult struct {
	Type   model.ValueType `json:"resultType"`
	Result interface{}     `json:"result"`

	// The decoded value.
	v model.Value
}

func (qr *queryResult) UnmarshalJSON(b []byte) error {
	v := struct {
		Type   model.ValueType `json:"resultType"`
		Result json.RawMessage `json:"result"`
	}{}

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

	switch v.Type {
	case model.ValScalar:
		var sv model.Scalar
		err = json.Unmarshal(v.Result, &sv)
		qr.v = &sv

	case model.ValVector:
		var vv model.Vector
		err = json.Unmarshal(v.Result, &vv)
		qr.v = vv

	case model.ValMatrix:
		var mv model.Matrix
		err = json.Unmarshal(v.Result, &mv)
		qr.v = mv

	default:
		err = fmt.Errorf("unexpected value type %q", v.Type)
	}
	return err
}

// QueryAPI provides bindings the Prometheus's query API.
type QueryAPI interface {
	// Query performs a query for the given time.
	Query(ctx context.Context, query string, ts time.Time) (model.Value, error)
	// Query performs a query for the given range.
	QueryRange(ctx context.Context, query string, r Range) (model.Value, error)
}

// NewQueryAPI returns a new QueryAPI for the client.
func NewQueryAPI(c Client) QueryAPI {
	return &httpQueryAPI{client: apiClient{c}}
}

type httpQueryAPI struct {
	client Client
}

func (h *httpQueryAPI) Query(ctx context.Context, query string, ts time.Time) (model.Value, error) {
	u := h.client.url(epQuery, nil)
	q := u.Query()

	q.Set("query", query)
	q.Set("time", ts.Format(time.RFC3339Nano))

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}

func (h *httpQueryAPI) QueryRange(ctx context.Context, query string, r Range) (model.Value, error) {
	u := h.client.url(epQueryRange, nil)
	q := u.Query()

	var (
		start = r.Start.Format(time.RFC3339Nano)
		end   = r.End.Format(time.RFC3339Nano)
		step  = strconv.FormatFloat(r.Step.Seconds(), 'f', 3, 64)
	)

	q.Set("query", query)
	q.Set("start", start)
	q.Set("end", end)
	q.Set("step", step)

	u.RawQuery = q.Encode()

	req, _ := http.NewRequest("GET", u.String(), nil)

	_, body, err := h.client.do(ctx, req)
	if err != nil {
		return nil, err
	}

	var qres queryResult
	err = json.Unmarshal(body, &qres)

	return model.Value(qres.v), err
}
<|endoftext|>