<commit_before>// Copyright (c) 2015-2016 Jeevanandam M (jeeva@myjeeva.com), All rights reserved.
// resty source code and usage is governed by a MIT style
// license that can be found in the LICENSE file.

package resty

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"strings"
)

//
// Request Middleware(s)
//

func parseRequestURL(c *Client, r *Request) error {
	// Parsing request URL
	reqURL, err := url.Parse(r.URL)
	if err != nil {
		return err
	}

	// If Request.Url is relative path then added c.HostUrl into
	// the request URL otherwise Request.Url will be used as-is
	if !reqURL.IsAbs() {
		if !strings.HasPrefix(r.URL, "/") {
			r.URL = "/" + r.URL
		}

		reqURL, err = url.Parse(c.HostURL + r.URL)
		if err != nil {
			return err
		}
	}

	// Adding Query Param
	query := reqURL.Query()
	for k, v := range c.QueryParam {
		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	for k, v := range r.QueryParam {
		// remove query param from client level by key
		// since overrides happens for that key in the request
		query.Del(k)

		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	reqURL.RawQuery = query.Encode()
	r.URL = reqURL.String()

	return nil
}

func parseRequestHeader(c *Client, r *Request) error {
	hdr := http.Header{}
	for k := range c.Header {
		hdr.Set(k, c.Header.Get(k))
	}
	for k := range r.Header {
		hdr.Set(k, r.Header.Get(k))
	}

	if IsStringEmpty(hdr.Get(hdrUserAgentKey)) {
		hdr.Set(hdrUserAgentKey, fmt.Sprintf(hdrUserAgentValue, Version))
	} else {
		hdr.Set("X-"+hdrUserAgentKey, fmt.Sprintf(hdrUserAgentValue, Version))
	}

	if IsStringEmpty(hdr.Get(hdrAcceptKey)) && !IsStringEmpty(hdr.Get(hdrContentTypeKey)) {
		hdr.Set(hdrAcceptKey, hdr.Get(hdrContentTypeKey))
	}

	r.Header = hdr

	return nil
}

func parseRequestBody(c *Client, r *Request) (err error) {
	if isPayloadSupported(r.Method) {

		// Handling Multipart
		if r.isMultiPart && !(r.Method == PATCH) {
			if err = handleMultipart(c, r); err != nil {
				return
			}

			goto CL
		}

		// Handling Form Data
		if len(c.FormData) > 0 || len(r.FormData) > 0 {
			handleFormData(c, r)

			goto CL
		}

		// Handling Request body
		if r.Body != nil {
			handleContentType(c, r)

			if err = handleRequestBody(c, r); err != nil {
				return
			}
		}
	} else {
		r.Header.Del(hdrContentTypeKey)
	}

CL:
	// by default resty won't set content length, you can if you want to :)
	if c.setContentLength || r.setContentLength {
		r.Header.Set(hdrContentLengthKey, fmt.Sprintf("%d", r.bodyBuf.Len()))
	}

	return
}

func createHTTPRequest(c *Client, r *Request) (err error) {
	if r.bodyBuf == nil {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, nil)
	} else {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, r.bodyBuf)
	}

	if err == nil {
		r.RawRequest.Close = c.closeConnection
	}

	// Add headers into http request
	r.RawRequest.Header = r.Header

	// Add cookies into http request
	for _, cookie := range c.Cookies {
		r.RawRequest.AddCookie(cookie)
	}

	if r.RawRequest.URL != nil && r.RawRequest.URL.Scheme == "" {
		r.RawRequest.URL.Scheme = c.scheme
		r.RawRequest.URL.Host = r.URL
	}

	return
}

func addCredentials(c *Client, r *Request) error {
	var isBasicAuth bool
	// Basic Auth
	if r.UserInfo != nil { // takes precedence
		r.RawRequest.SetBasicAuth(r.UserInfo.Username, r.UserInfo.Password)
		isBasicAuth = true
	} else if c.UserInfo != nil {
		r.RawRequest.SetBasicAuth(c.UserInfo.Username, c.UserInfo.Password)
		isBasicAuth = true
	}

	if !c.DisableWarn {
		if isBasicAuth && !strings.HasPrefix(r.URL, "https") {
			c.Log.Println("WARNING - Using Basic Auth in HTTP mode is not secure.")
		}
	}

	// Token Auth
	if !IsStringEmpty(r.Token) { // takes precedence
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+r.Token)
	} else if !IsStringEmpty(c.Token) {
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+c.Token)
	}

	return nil
}

func requestLogger(c *Client, r *Request) error {
	if c.Debug {
		rr := r.RawRequest
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- REQUEST LOG -----------------------")
		c.Log.Printf("%s  %s  %s\n", r.Method, rr.URL.RequestURI(), rr.Proto)
		c.Log.Printf("HOST   : %s", rr.URL.Host)
		c.Log.Println("HEADERS:")
		for h, v := range rr.Header {
			c.Log.Printf("%25s: %v", h, strings.Join(v, ", "))
		}
		c.Log.Printf("BODY   :\n%v", r.fmtBodyString())
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

//
// Response Middleware(s)
//

func responseLogger(c *Client, res *Response) error {
	if c.Debug {
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- RESPONSE LOG -----------------------")
		c.Log.Printf("STATUS 		: %s", res.Status())
		c.Log.Printf("RECEIVED AT	: %v", res.ReceivedAt())
		c.Log.Printf("RESPONSE TIME	: %v", res.Time())
		c.Log.Println("HEADERS:")
		for h, v := range res.Header() {
			c.Log.Printf("%30s: %v", h, strings.Join(v, ", "))
		}
		if res.Request.isSaveResponse {
			c.Log.Printf("BODY   :\n***** RESPONSE WRITTEN INTO FILE *****")
		} else {
			c.Log.Printf("BODY   :\n%v", res.fmtBodyString())
		}
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

func parseResponseBody(c *Client, res *Response) (err error) {
	// Handles only JSON or XML content type
	ct := res.Header().Get(hdrContentTypeKey)
	if IsJSONType(ct) || IsXMLType(ct) {
		// Considered as Result
		if res.StatusCode() > 199 && res.StatusCode() < 300 {
			if res.Request.Result != nil {
				err = Unmarshal(ct, res.body, res.Request.Result)
			}
		}

		// Considered as Error
		if res.StatusCode() > 399 {
			// global error interface
			if res.Request.Error == nil && c.Error != nil {
				res.Request.Error = reflect.New(c.Error).Interface()
			}

			if res.Request.Error != nil {
				err = Unmarshal(ct, res.body, res.Request.Error)
			}
		}
	}

	return
}

func handleMultipart(c *Client, r *Request) (err error) {
	r.bodyBuf = &bytes.Buffer{}
	w := multipart.NewWriter(r.bodyBuf)

	for p := range c.FormData {
		w.WriteField(p, c.FormData.Get(p))
	}

	for p := range r.FormData {
		if strings.HasPrefix(p, "@") { // file
			err = addFile(w, p[1:], r.FormData.Get(p))
			if err != nil {
				return
			}
		} else { // form value
			w.WriteField(p, r.FormData.Get(p))
		}
	}

	// #21 - adding io.Reader support
	if len(r.multipartFiles) > 0 {
		for _, f := range r.multipartFiles {
			err = addFileReader(w, f)
			if err != nil {
				return
			}
		}
	}

	r.Header.Set(hdrContentTypeKey, w.FormDataContentType())
	err = w.Close()

	return
}

func handleFormData(c *Client, r *Request) {
	formData := url.Values{}

	for p := range c.FormData {
		formData.Set(p, c.FormData.Get(p))
	}

	for p := range r.FormData { // takes precedence
		formData.Set(p, r.FormData.Get(p))
	}

	r.bodyBuf = bytes.NewBuffer([]byte(formData.Encode()))
	r.Header.Set(hdrContentTypeKey, formContentType)
	r.isFormData = true
}

func handleContentType(c *Client, r *Request) {
	contentType := r.Header.Get(hdrContentTypeKey)
	if IsStringEmpty(contentType) {
		contentType = DetectContentType(r.Body)
		r.Header.Set(hdrContentTypeKey, contentType)
	}
}

func handleRequestBody(c *Client, r *Request) (err error) {
	var bodyBytes []byte
	contentType := r.Header.Get(hdrContentTypeKey)
	kind := kindOf(r.Body)

	if reader, ok := r.Body.(io.Reader); ok {
		r.bodyBuf = &bytes.Buffer{}
		r.bodyBuf.ReadFrom(reader)
	} else if b, ok := r.Body.([]byte); ok {
		bodyBytes = b
	} else if s, ok := r.Body.(string); ok {
		bodyBytes = []byte(s)
	} else if IsJSONType(contentType) &&
		(kind == reflect.Struct || kind == reflect.Map || kind == reflect.Slice) {
		bodyBytes, err = json.Marshal(r.Body)
	} else if IsXMLType(contentType) && (kind == reflect.Struct) {
		bodyBytes, err = xml.Marshal(r.Body)
	}

	if bodyBytes == nil && r.bodyBuf == nil {
		err = errors.New("Unsupported 'Body' type/value")
	}

	// if any errors during body bytes handling, return it
	if err != nil {
		return
	}

	// []byte into Buffer
	if bodyBytes != nil && r.bodyBuf == nil {
		r.bodyBuf = bytes.NewBuffer(bodyBytes)
	}

	return
}

func saveResponseIntoFile(c *Client, res *Response) error {
	if res.Request.isSaveResponse {
		file := ""

		if len(c.outputDirectory) > 0 && !filepath.IsAbs(res.Request.outputFile) {
			file += c.outputDirectory + string(filepath.Separator)
		}

		file = filepath.Clean(file + res.Request.outputFile)
		err := createDirectory(filepath.Dir(file))
		if err != nil {
			return err
		}

		outFile, err := os.Create(file)
		if err != nil {
			return err
		}
		defer outFile.Close()

		// io.Copy reads maximum 32kb size, it is perfect for large file download too
		defer res.RawResponse.Body.Close()
		written, err := io.Copy(outFile, res.RawResponse.Body)
		if err != nil {
			return err
		}

		res.size = written
	}

	return nil
}
<commit_msg>code improvement at middleware#createHTTPRequest<commit_after>// Copyright (c) 2015-2016 Jeevanandam M (jeeva@myjeeva.com), All rights reserved.
// resty source code and usage is governed by a MIT style
// license that can be found in the LICENSE file.

package resty

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"reflect"
	"strings"
)

//
// Request Middleware(s)
//

func parseRequestURL(c *Client, r *Request) error {
	// Parsing request URL
	reqURL, err := url.Parse(r.URL)
	if err != nil {
		return err
	}

	// If Request.Url is relative path then added c.HostUrl into
	// the request URL otherwise Request.Url will be used as-is
	if !reqURL.IsAbs() {
		if !strings.HasPrefix(r.URL, "/") {
			r.URL = "/" + r.URL
		}

		reqURL, err = url.Parse(c.HostURL + r.URL)
		if err != nil {
			return err
		}
	}

	// Adding Query Param
	query := reqURL.Query()
	for k, v := range c.QueryParam {
		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	for k, v := range r.QueryParam {
		// remove query param from client level by key
		// since overrides happens for that key in the request
		query.Del(k)

		for _, iv := range v {
			query.Add(k, iv)
		}
	}

	reqURL.RawQuery = query.Encode()
	r.URL = reqURL.String()

	return nil
}

func parseRequestHeader(c *Client, r *Request) error {
	hdr := http.Header{}
	for k := range c.Header {
		hdr.Set(k, c.Header.Get(k))
	}
	for k := range r.Header {
		hdr.Set(k, r.Header.Get(k))
	}

	if IsStringEmpty(hdr.Get(hdrUserAgentKey)) {
		hdr.Set(hdrUserAgentKey, fmt.Sprintf(hdrUserAgentValue, Version))
	} else {
		hdr.Set("X-"+hdrUserAgentKey, fmt.Sprintf(hdrUserAgentValue, Version))
	}

	if IsStringEmpty(hdr.Get(hdrAcceptKey)) && !IsStringEmpty(hdr.Get(hdrContentTypeKey)) {
		hdr.Set(hdrAcceptKey, hdr.Get(hdrContentTypeKey))
	}

	r.Header = hdr

	return nil
}

func parseRequestBody(c *Client, r *Request) (err error) {
	if isPayloadSupported(r.Method) {

		// Handling Multipart
		if r.isMultiPart && !(r.Method == PATCH) {
			if err = handleMultipart(c, r); err != nil {
				return
			}

			goto CL
		}

		// Handling Form Data
		if len(c.FormData) > 0 || len(r.FormData) > 0 {
			handleFormData(c, r)

			goto CL
		}

		// Handling Request body
		if r.Body != nil {
			handleContentType(c, r)

			if err = handleRequestBody(c, r); err != nil {
				return
			}
		}
	} else {
		r.Header.Del(hdrContentTypeKey)
	}

CL:
	// by default resty won't set content length, you can if you want to :)
	if c.setContentLength || r.setContentLength {
		r.Header.Set(hdrContentLengthKey, fmt.Sprintf("%d", r.bodyBuf.Len()))
	}

	return
}

func createHTTPRequest(c *Client, r *Request) (err error) {
	if r.bodyBuf == nil {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, nil)
	} else {
		r.RawRequest, err = http.NewRequest(r.Method, r.URL, r.bodyBuf)
	}

	if err != nil {
		return
	}

	// Assign close connection option
	r.RawRequest.Close = c.closeConnection

	// Add headers into http request
	r.RawRequest.Header = r.Header

	// Add cookies into http request
	for _, cookie := range c.Cookies {
		r.RawRequest.AddCookie(cookie)
	}

	// it's for non-http scheme option
	if r.RawRequest.URL != nil && r.RawRequest.URL.Scheme == "" {
		r.RawRequest.URL.Scheme = c.scheme
		r.RawRequest.URL.Host = r.URL
	}

	return
}

func addCredentials(c *Client, r *Request) error {
	var isBasicAuth bool
	// Basic Auth
	if r.UserInfo != nil { // takes precedence
		r.RawRequest.SetBasicAuth(r.UserInfo.Username, r.UserInfo.Password)
		isBasicAuth = true
	} else if c.UserInfo != nil {
		r.RawRequest.SetBasicAuth(c.UserInfo.Username, c.UserInfo.Password)
		isBasicAuth = true
	}

	if !c.DisableWarn {
		if isBasicAuth && !strings.HasPrefix(r.URL, "https") {
			c.Log.Println("WARNING - Using Basic Auth in HTTP mode is not secure.")
		}
	}

	// Token Auth
	if !IsStringEmpty(r.Token) { // takes precedence
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+r.Token)
	} else if !IsStringEmpty(c.Token) {
		r.RawRequest.Header.Set(hdrAuthorizationKey, "Bearer "+c.Token)
	}

	return nil
}

func requestLogger(c *Client, r *Request) error {
	if c.Debug {
		rr := r.RawRequest
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- REQUEST LOG -----------------------")
		c.Log.Printf("%s  %s  %s\n", r.Method, rr.URL.RequestURI(), rr.Proto)
		c.Log.Printf("HOST   : %s", rr.URL.Host)
		c.Log.Println("HEADERS:")
		for h, v := range rr.Header {
			c.Log.Printf("%25s: %v", h, strings.Join(v, ", "))
		}
		c.Log.Printf("BODY   :\n%v", r.fmtBodyString())
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

//
// Response Middleware(s)
//

func responseLogger(c *Client, res *Response) error {
	if c.Debug {
		c.Log.Println()
		c.disableLogPrefix()
		c.Log.Println("---------------------- RESPONSE LOG -----------------------")
		c.Log.Printf("STATUS 		: %s", res.Status())
		c.Log.Printf("RECEIVED AT	: %v", res.ReceivedAt())
		c.Log.Printf("RESPONSE TIME	: %v", res.Time())
		c.Log.Println("HEADERS:")
		for h, v := range res.Header() {
			c.Log.Printf("%30s: %v", h, strings.Join(v, ", "))
		}
		if res.Request.isSaveResponse {
			c.Log.Printf("BODY   :\n***** RESPONSE WRITTEN INTO FILE *****")
		} else {
			c.Log.Printf("BODY   :\n%v", res.fmtBodyString())
		}
		c.Log.Println("----------------------------------------------------------")
		c.enableLogPrefix()
	}

	return nil
}

func parseResponseBody(c *Client, res *Response) (err error) {
	// Handles only JSON or XML content type
	ct := res.Header().Get(hdrContentTypeKey)
	if IsJSONType(ct) || IsXMLType(ct) {
		// Considered as Result
		if res.StatusCode() > 199 && res.StatusCode() < 300 {
			if res.Request.Result != nil {
				err = Unmarshal(ct, res.body, res.Request.Result)
			}
		}

		// Considered as Error
		if res.StatusCode() > 399 {
			// global error interface
			if res.Request.Error == nil && c.Error != nil {
				res.Request.Error = reflect.New(c.Error).Interface()
			}

			if res.Request.Error != nil {
				err = Unmarshal(ct, res.body, res.Request.Error)
			}
		}
	}

	return
}

func handleMultipart(c *Client, r *Request) (err error) {
	r.bodyBuf = &bytes.Buffer{}
	w := multipart.NewWriter(r.bodyBuf)

	for p := range c.FormData {
		w.WriteField(p, c.FormData.Get(p))
	}

	for p := range r.FormData {
		if strings.HasPrefix(p, "@") { // file
			err = addFile(w, p[1:], r.FormData.Get(p))
			if err != nil {
				return
			}
		} else { // form value
			w.WriteField(p, r.FormData.Get(p))
		}
	}

	// #21 - adding io.Reader support
	if len(r.multipartFiles) > 0 {
		for _, f := range r.multipartFiles {
			err = addFileReader(w, f)
			if err != nil {
				return
			}
		}
	}

	r.Header.Set(hdrContentTypeKey, w.FormDataContentType())
	err = w.Close()

	return
}

func handleFormData(c *Client, r *Request) {
	formData := url.Values{}

	for p := range c.FormData {
		formData.Set(p, c.FormData.Get(p))
	}

	for p := range r.FormData { // takes precedence
		formData.Set(p, r.FormData.Get(p))
	}

	r.bodyBuf = bytes.NewBuffer([]byte(formData.Encode()))
	r.Header.Set(hdrContentTypeKey, formContentType)
	r.isFormData = true
}

func handleContentType(c *Client, r *Request) {
	contentType := r.Header.Get(hdrContentTypeKey)
	if IsStringEmpty(contentType) {
		contentType = DetectContentType(r.Body)
		r.Header.Set(hdrContentTypeKey, contentType)
	}
}

func handleRequestBody(c *Client, r *Request) (err error) {
	var bodyBytes []byte
	contentType := r.Header.Get(hdrContentTypeKey)
	kind := kindOf(r.Body)

	if reader, ok := r.Body.(io.Reader); ok {
		r.bodyBuf = &bytes.Buffer{}
		r.bodyBuf.ReadFrom(reader)
	} else if b, ok := r.Body.([]byte); ok {
		bodyBytes = b
	} else if s, ok := r.Body.(string); ok {
		bodyBytes = []byte(s)
	} else if IsJSONType(contentType) &&
		(kind == reflect.Struct || kind == reflect.Map || kind == reflect.Slice) {
		bodyBytes, err = json.Marshal(r.Body)
	} else if IsXMLType(contentType) && (kind == reflect.Struct) {
		bodyBytes, err = xml.Marshal(r.Body)
	}

	if bodyBytes == nil && r.bodyBuf == nil {
		err = errors.New("Unsupported 'Body' type/value")
	}

	// if any errors during body bytes handling, return it
	if err != nil {
		return
	}

	// []byte into Buffer
	if bodyBytes != nil && r.bodyBuf == nil {
		r.bodyBuf = bytes.NewBuffer(bodyBytes)
	}

	return
}

func saveResponseIntoFile(c *Client, res *Response) error {
	if res.Request.isSaveResponse {
		file := ""

		if len(c.outputDirectory) > 0 && !filepath.IsAbs(res.Request.outputFile) {
			file += c.outputDirectory + string(filepath.Separator)
		}

		file = filepath.Clean(file + res.Request.outputFile)
		err := createDirectory(filepath.Dir(file))
		if err != nil {
			return err
		}

		outFile, err := os.Create(file)
		if err != nil {
			return err
		}
		defer outFile.Close()

		// io.Copy reads maximum 32kb size, it is perfect for large file download too
		defer res.RawResponse.Body.Close()
		written, err := io.Copy(outFile, res.RawResponse.Body)
		if err != nil {
			return err
		}

		res.size = written
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 The Go Cloud Development Kit 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
//
//     https://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 pubsub_test

import (
	"context"
	"errors"
	"math/rand"
	"sync"
	"testing"
	"time"

	"gocloud.dev/gcerrors"
	"gocloud.dev/internal/gcerr"
	"gocloud.dev/pubsub"
	"gocloud.dev/pubsub/driver"
)

type ackingDriverSub struct {
	driver.Subscription
	q        []*driver.Message
	sendAcks func(context.Context, []driver.AckID) error
}

func (s *ackingDriverSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	if len(s.q) <= maxMessages {
		ms := s.q
		s.q = nil
		return ms, nil
	}
	ms := s.q[:maxMessages]
	s.q = s.q[maxMessages:]
	return ms, nil
}

func (s *ackingDriverSub) SendAcks(ctx context.Context, ackIDs []driver.AckID) error {
	return s.sendAcks(ctx, ackIDs)
}

func (s *ackingDriverSub) IsRetryable(error) bool { return false }

func (s *ackingDriverSub) ErrorCode(error) gcerrors.ErrorCode { return gcerrors.Internal }

func TestAckTriggersDriverSendAcksForOneMessage(t *testing.T) {
	ctx := context.Background()
	var mu sync.Mutex
	var sentAcks []driver.AckID
	id := rand.Int()
	m := &driver.Message{AckID: id}
	ackChan := make(chan struct{})
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			mu.Lock()
			defer mu.Unlock()
			sentAcks = ackIDs
			ackChan <- struct{}{}
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	m2, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	m2.Ack()
	<-ackChan
	if len(sentAcks) != 1 {
		t.Fatalf("len(sentAcks) = %d, want exactly 1", len(sentAcks))
	}
	if sentAcks[0] != id {
		t.Errorf("sentAcks[0] = %d, want %d", sentAcks[0], id)
	}
}

func TestMultipleAcksCanGoIntoASingleBatch(t *testing.T) {
	ctx := context.Background()
	var wg sync.WaitGroup
	var mu sync.Mutex
	sentAcks := make(map[driver.AckID]int)
	ids := []int{1, 2}
	ds := &ackingDriverSub{
		q: []*driver.Message{{AckID: ids[0]}, {AckID: ids[1]}},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			mu.Lock()
			defer mu.Unlock()
			for _, id := range ackIDs {
				sentAcks[id]++
				wg.Done()
			}
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)

	// Receive and ack the messages concurrently.
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			mr, err := sub.Receive(ctx)
			if err != nil {
				t.Error(err)
				return
			}
			mr.Ack()
		}()
	}
	wg.Wait()

	if len(sentAcks) != 2 {
		t.Errorf("len(sentAcks) = %d, want exactly 2", len(sentAcks))
	}
	for _, id := range ids {
		if sentAcks[id] != 1 {
			t.Errorf("sentAcks[%v] = %d, want 1", id, sentAcks[id])
		}
	}
}

func TestTooManyAcksForASingleBatchGoIntoMultipleBatches(t *testing.T) {
	ctx := context.Background()
	var mu sync.Mutex
	var wg sync.WaitGroup
	var sentAckBatches [][]driver.AckID
	// This value of n is chosen large enough that it should create more
	// than one batch. Admittedly, there is currently no explicit guarantee
	// of this.
	n := 1000
	var ms []*driver.Message
	for i := 0; i < n; i++ {
		ms = append(ms, &driver.Message{AckID: i})
	}
	ds := &ackingDriverSub{
		q: ms,
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			mu.Lock()
			defer mu.Unlock()
			sentAckBatches = append(sentAckBatches, ackIDs)
			for i := 0; i < len(ackIDs); i++ {
				wg.Done()
			}
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)

	// Receive and ack the messages concurrently.
	recv := func() {
		mr, err := sub.Receive(ctx)
		if err != nil {
			t.Fatal(err)
		}
		mr.Ack()
	}
	wg.Add(n)
	for i := 0; i < n; i++ {
		go recv()
	}
	wg.Wait()

	if len(sentAckBatches) < 2 {
		t.Errorf("got %d batches, want at least 2", len(sentAckBatches))
	}
}

func TestAckDoesNotBlock(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			<-ctx.Done()
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	defer cancel()
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	// If Ack blocks here, waiting for sendAcks to finish, then the
	// deferred cancel() will never run, so sendAcks can never finish. That
	// would cause the test to hang. Thus hanging is how this test signals
	// failure.
	mr.Ack()
}

func TestDoubleAckCausesPanic(t *testing.T) {
	ctx := context.Background()
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	mr.Ack()
	defer func() {
		if r := recover(); r != nil {
			// ok, panic was expected.
			return
		}
		t.Errorf("second ack failed to panic")
	}()
	mr.Ack()
}

// For best results, run this test with -race.
func TestConcurrentDoubleAckCausesPanic(t *testing.T) {
	ctx := context.Background()
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}

	// Spin up some goroutines to ack the message.
	var mu sync.Mutex
	panics := 0
	var wg sync.WaitGroup
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			defer func() {
				if r := recover(); r != nil {
					mu.Lock()
					defer mu.Unlock()
					panics++
				}
			}()
			mr.Ack()
		}()
	}
	wg.Wait()

	// Check that one of the goroutines panicked.
	if panics != 1 {
		t.Errorf("panics = %d, want %d", panics, 1)
	}
}

func TestSubShutdownCanBeCanceledEvenWithHangingSendAcks(t *testing.T) {
	ctx := context.Background()
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(ctx context.Context, ackIDs []driver.AckID) error {
			<-ctx.Done()
			return ctx.Err()
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	mr.Ack()

	done := make(chan struct{})
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
		defer cancel()
		sub.Shutdown(ctx)
		close(done)
	}()
	tooLong := 5 * time.Second
	select {
	case <-done:
	case <-time.After(tooLong):
		t.Fatalf("waited too long (%v) for Shutdown to run", tooLong)
	}
}

func TestReceiveReturnsErrorFromSendAcks(t *testing.T) {
	// If SendAcks fails, the error is returned via receive.
	ctx := context.Background()
	serr := errors.New("SendAcks failed")
	ackChan := make(chan struct{})
	ds := &ackingDriverSub{
		q: []*driver.Message{
			&driver.Message{AckID: 0},
			&driver.Message{AckID: 1},
			&driver.Message{AckID: 2},
			&driver.Message{AckID: 3},
		},
		sendAcks: func(context.Context, []driver.AckID) error {
			close(ackChan)
			return serr
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	m, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	m.Ack()
	// Wait for the ack to be sent.
	<-ackChan
	// It might take a bit longer for the logic after SendAcks returns to happen, so
	// keep calling Receive.
	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	for {
		_, err = sub.Receive(ctx)
		if gcerrors.Code(err) == gcerrors.Internal && err.(*gcerr.Error).Unwrap() == serr {
			break // success
		}
		if err != nil {
			t.Fatalf("got %v, want %v", err, serr)
		}
	}
}
<commit_msg>pubsub: add TestReceiveReturnsAckErrorOnNoMoreMessages (#1283)<commit_after>// Copyright 2018 The Go Cloud Development Kit 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
//
//     https://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 pubsub_test

import (
	"context"
	"errors"
	"math/rand"
	"sync"
	"testing"
	"time"

	"golang.org/x/xerrors"

	"gocloud.dev/gcerrors"
	"gocloud.dev/internal/gcerr"
	"gocloud.dev/pubsub"
	"gocloud.dev/pubsub/driver"
)

type ackingDriverSub struct {
	driver.Subscription
	q        []*driver.Message
	sendAcks func(context.Context, []driver.AckID) error
}

func (s *ackingDriverSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	if len(s.q) <= maxMessages {
		ms := s.q
		s.q = nil
		return ms, nil
	}
	ms := s.q[:maxMessages]
	s.q = s.q[maxMessages:]
	return ms, nil
}

func (s *ackingDriverSub) SendAcks(ctx context.Context, ackIDs []driver.AckID) error {
	return s.sendAcks(ctx, ackIDs)
}

func (s *ackingDriverSub) IsRetryable(error) bool { return false }

func (s *ackingDriverSub) ErrorCode(error) gcerrors.ErrorCode { return gcerrors.Internal }

func TestAckTriggersDriverSendAcksForOneMessage(t *testing.T) {
	ctx := context.Background()
	var mu sync.Mutex
	var sentAcks []driver.AckID
	id := rand.Int()
	m := &driver.Message{AckID: id}
	ackChan := make(chan struct{})
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			mu.Lock()
			defer mu.Unlock()
			sentAcks = ackIDs
			ackChan <- struct{}{}
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	m2, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	m2.Ack()
	<-ackChan
	if len(sentAcks) != 1 {
		t.Fatalf("len(sentAcks) = %d, want exactly 1", len(sentAcks))
	}
	if sentAcks[0] != id {
		t.Errorf("sentAcks[0] = %d, want %d", sentAcks[0], id)
	}
}

func TestMultipleAcksCanGoIntoASingleBatch(t *testing.T) {
	ctx := context.Background()
	var wg sync.WaitGroup
	var mu sync.Mutex
	sentAcks := make(map[driver.AckID]int)
	ids := []int{1, 2}
	ds := &ackingDriverSub{
		q: []*driver.Message{{AckID: ids[0]}, {AckID: ids[1]}},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			mu.Lock()
			defer mu.Unlock()
			for _, id := range ackIDs {
				sentAcks[id]++
				wg.Done()
			}
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)

	// Receive and ack the messages concurrently.
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			mr, err := sub.Receive(ctx)
			if err != nil {
				t.Error(err)
				return
			}
			mr.Ack()
		}()
	}
	wg.Wait()

	if len(sentAcks) != 2 {
		t.Errorf("len(sentAcks) = %d, want exactly 2", len(sentAcks))
	}
	for _, id := range ids {
		if sentAcks[id] != 1 {
			t.Errorf("sentAcks[%v] = %d, want 1", id, sentAcks[id])
		}
	}
}

func TestTooManyAcksForASingleBatchGoIntoMultipleBatches(t *testing.T) {
	ctx := context.Background()
	var mu sync.Mutex
	var wg sync.WaitGroup
	var sentAckBatches [][]driver.AckID
	// This value of n is chosen large enough that it should create more
	// than one batch. Admittedly, there is currently no explicit guarantee
	// of this.
	n := 1000
	var ms []*driver.Message
	for i := 0; i < n; i++ {
		ms = append(ms, &driver.Message{AckID: i})
	}
	ds := &ackingDriverSub{
		q: ms,
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			mu.Lock()
			defer mu.Unlock()
			sentAckBatches = append(sentAckBatches, ackIDs)
			for i := 0; i < len(ackIDs); i++ {
				wg.Done()
			}
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)

	// Receive and ack the messages concurrently.
	recv := func() {
		mr, err := sub.Receive(ctx)
		if err != nil {
			t.Fatal(err)
		}
		mr.Ack()
	}
	wg.Add(n)
	for i := 0; i < n; i++ {
		go recv()
	}
	wg.Wait()

	if len(sentAckBatches) < 2 {
		t.Errorf("got %d batches, want at least 2", len(sentAckBatches))
	}
}

func TestAckDoesNotBlock(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			<-ctx.Done()
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	defer cancel()
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	// If Ack blocks here, waiting for sendAcks to finish, then the
	// deferred cancel() will never run, so sendAcks can never finish. That
	// would cause the test to hang. Thus hanging is how this test signals
	// failure.
	mr.Ack()
}

func TestDoubleAckCausesPanic(t *testing.T) {
	ctx := context.Background()
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	mr.Ack()
	defer func() {
		if r := recover(); r != nil {
			// ok, panic was expected.
			return
		}
		t.Errorf("second ack failed to panic")
	}()
	mr.Ack()
}

// For best results, run this test with -race.
func TestConcurrentDoubleAckCausesPanic(t *testing.T) {
	ctx := context.Background()
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(_ context.Context, ackIDs []driver.AckID) error {
			return nil
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}

	// Spin up some goroutines to ack the message.
	var mu sync.Mutex
	panics := 0
	var wg sync.WaitGroup
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			defer func() {
				if r := recover(); r != nil {
					mu.Lock()
					defer mu.Unlock()
					panics++
				}
			}()
			mr.Ack()
		}()
	}
	wg.Wait()

	// Check that one of the goroutines panicked.
	if panics != 1 {
		t.Errorf("panics = %d, want %d", panics, 1)
	}
}

func TestSubShutdownCanBeCanceledEvenWithHangingSendAcks(t *testing.T) {
	ctx := context.Background()
	m := &driver.Message{AckID: 0} // the batcher doesn't like nil interfaces
	ds := &ackingDriverSub{
		q: []*driver.Message{m},
		sendAcks: func(ctx context.Context, ackIDs []driver.AckID) error {
			<-ctx.Done()
			return ctx.Err()
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	mr, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	mr.Ack()

	done := make(chan struct{})
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
		defer cancel()
		sub.Shutdown(ctx)
		close(done)
	}()
	tooLong := 5 * time.Second
	select {
	case <-done:
	case <-time.After(tooLong):
		t.Fatalf("waited too long (%v) for Shutdown to run", tooLong)
	}
}

func TestReceiveReturnsErrorFromSendAcks(t *testing.T) {
	// If SendAcks fails, the error is returned via receive.
	ctx := context.Background()
	serr := errors.New("SendAcks failed")
	ackChan := make(chan struct{})
	ds := &ackingDriverSub{
		q: []*driver.Message{
			&driver.Message{AckID: 0},
			&driver.Message{AckID: 1},
			&driver.Message{AckID: 2},
			&driver.Message{AckID: 3},
		},
		sendAcks: func(context.Context, []driver.AckID) error {
			close(ackChan)
			return serr
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	m, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	m.Ack()
	// Wait for the ack to be sent.
	<-ackChan
	// It might take a bit longer for the logic after SendAcks returns to happen, so
	// keep calling Receive.
	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	for {
		_, err = sub.Receive(ctx)
		if gcerrors.Code(err) == gcerrors.Internal && err.(*gcerr.Error).Unwrap() == serr {
			break // success
		}
		if err != nil {
			t.Fatalf("got %v, want %v", err, serr)
		}
	}
}

// callbackDriverSub implements driver.Subscription and allows something like
// monkey patching of both its ReceiveBatch and SendAcks methods.
type callbackDriverSub struct {
	driver.Subscription
	mu           sync.Mutex
	receiveBatch func(context.Context) ([]*driver.Message, error)
	sendAcks     func(context.Context, []driver.AckID) error
}

func (s *callbackDriverSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	return s.receiveBatch(ctx)
}

func (s *callbackDriverSub) SendAcks(ctx context.Context, acks []driver.AckID) error {
	return s.sendAcks(ctx, acks)
}

func (s *callbackDriverSub) IsRetryable(error) bool { return false }

func (s *callbackDriverSub) ErrorCode(error) gcerrors.ErrorCode { return gcerrors.Internal }

// This test detects the root cause of
// https://github.com/google/go-cloud/issues/1238.
// If the issue is present, this test times out. The problem was that when
// there were no messages available from the provider,
// pubsub.Subscription.Receive would spin trying to get more messages without
// checking to see if an unrecoverable error had occurred while sending a batch
// of acks to the provider.
func TestReceiveReturnsAckErrorOnNoMoreMessages(t *testing.T) {
	// If SendAcks fails, the error is returned via receive.
	ctx := context.Background()
	serr := errors.New("unrecoverable error")
	receiveHappened := make(chan struct{})
	ackHappened := make(chan struct{})
	var ds = &callbackDriverSub{
		// First call to receiveBatch will return a single message.
		receiveBatch: func(context.Context) ([]*driver.Message, error) {
			ms := []*driver.Message{{AckID: 1}}
			return ms, nil
		},
		sendAcks: func(context.Context, []driver.AckID) error {
			ackHappened <- struct{}{}
			return serr
		},
	}
	sub := pubsub.NewSubscription(ds, nil)
	defer sub.Shutdown(ctx)
	m, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	m.Ack()

	// Second call to receiveBatch will wait for the pull from the
	// receiveHappened channel below, and return a nil slice of messages.
	ds.mu.Lock()
	ds.receiveBatch = func(context.Context) ([]*driver.Message, error) {
		ds.mu.Lock()
		// Subsequent calls to receiveBatch won't wait on receiveHappened,
		// and will also return nil slices of messages.
		ds.receiveBatch = func(context.Context) ([]*driver.Message, error) {
			return nil, nil
		}
		ds.mu.Unlock()
		receiveHappened <- struct{}{}
		return nil, nil
	}
	ds.mu.Unlock()

	errc := make(chan error)
	go func() {
		_, err := sub.Receive(ctx)
		errc <- err
	}()

	// sub.Receive has to start running first and then we need to trigger the unrecoverable error.
	<-receiveHappened

	// Trigger the unrecoverable error.
	<-ackHappened

	// Wait for sub.Receive to return so we can check the error it returns against serr.
	err = <-errc

	// Check the error returned from sub.Receive.
	if got := gcerrors.Code(err); got != gcerrors.Internal {
		t.Fatalf("error code = %v; want %v", got, gcerrors.Internal)
	}
	if got := xerrors.Unwrap(err); got != serr {
		t.Errorf("error = %v; want %v", got, serr)
	}
}
<|endoftext|>


<commit_before>package namespaces

import (
	"fmt"
	"testing"
)

func TestSendErrorFromChild(t *testing.T) {
	pipe, err := NewSyncPipe()
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		if err := pipe.Close(); err != nil {
			t.Fatal(err)
		}
	}()

	expected := "something bad happened"

	pipe.ReportChildError(fmt.Errorf(expected))

	childError := pipe.ReadFromChild()
	if childError == nil {
		t.Fatal("expected an error to be returned but did not receive anything")
	}

	if childError.Error() != expected {
		t.Fatalf("expected %q but received error message %q", expected, childError.Error())
	}
}

func TestSendPayloadToChild(t *testing.T) {
	pipe, err := NewSyncPipe()
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		if err := pipe.Close(); err != nil {
			t.Fatal(err)
		}
	}()

	expected := "libcontainer"

	if err := pipe.SendToChild(map[string]string{"name": expected}); err != nil {
		t.Fatal(err)
	}

	payload, err := pipe.ReadFromParent()
	if err != nil {
		t.Fatal(err)
	}

	if len(payload) != 1 {
		t.Fatalf("expected to only have one value in the payload but received %d", len(payload))
	}

	if name := payload["name"]; name != expected {
		t.Fatalf("expected name %q but received %q", expected, name)
	}
}
<commit_msg>Update tests for network state Docker-DCO-1.1-Signed-off-by: Michael Crosby <michael@docker.com> (github: crosbymichael)<commit_after>package namespaces

import (
	"fmt"
	"testing"

	"github.com/docker/libcontainer/network"
)

func TestSendErrorFromChild(t *testing.T) {
	pipe, err := NewSyncPipe()
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		if err := pipe.Close(); err != nil {
			t.Fatal(err)
		}
	}()

	expected := "something bad happened"

	pipe.ReportChildError(fmt.Errorf(expected))

	childError := pipe.ReadFromChild()
	if childError == nil {
		t.Fatal("expected an error to be returned but did not receive anything")
	}

	if childError.Error() != expected {
		t.Fatalf("expected %q but received error message %q", expected, childError.Error())
	}
}

func TestSendPayloadToChild(t *testing.T) {
	pipe, err := NewSyncPipe()
	if err != nil {
		t.Fatal(err)
	}

	defer func() {
		if err := pipe.Close(); err != nil {
			t.Fatal(err)
		}
	}()

	expected := "libcontainer"

	if err := pipe.SendToChild(&network.NetworkState{VethHost: expected}); err != nil {
		t.Fatal(err)
	}

	payload, err := pipe.ReadFromParent()
	if err != nil {
		t.Fatal(err)
	}

	if payload.VethHost != expected {
		t.Fatalf("expected veth host %q but received %q", expected, payload.VethHost)
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package proposervm

import (
	"errors"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/choices"
	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
	"github.com/ava-labs/avalanchego/vms/proposervm/block"
	"github.com/ava-labs/avalanchego/vms/proposervm/proposer"
)

const (
	// allowable block issuance in the future
	maxSkew = 10 * time.Second
)

var (
	errUnsignedChild            = errors.New("expected child to be signed")
	errUnexpectedBlockType      = errors.New("unexpected proposer block type")
	errInnerParentMismatch      = errors.New("inner parentID didn't match expected parent")
	errTimeNotMonotonic         = errors.New("time must monotonically increase")
	errPChainHeightNotMonotonic = errors.New("non monotonically increasing P-chain height")
	errPChainHeightNotReached   = errors.New("block P-chain height larger than current P-chain height")
	errTimeTooAdvanced          = errors.New("time is too far advanced")
	errProposerWindowNotStarted = errors.New("proposer window hasn't started")
	errProposersNotActivated    = errors.New("proposers haven't been activated yet")
	errPChainHeightTooLow       = errors.New("block P-chain height is too low")
)

type Block interface {
	snowman.Block

	getInnerBlk() snowman.Block

	verifyPreForkChild(child *preForkBlock) error
	verifyPostForkChild(child *postForkBlock) error
	verifyPostForkOption(child *postForkOption) error

	buildChild() (Block, error)

	pChainHeight() (uint64, error)
}

type PostForkBlock interface {
	Block

	setStatus(choices.Status)
	getStatelessBlk() block.Block
	setInnerBlk(snowman.Block)
}

// field of postForkBlock and postForkOption
type postForkCommonComponents struct {
	vm       *VM
	innerBlk snowman.Block
	status   choices.Status
}

// Return the inner block's height
func (p *postForkCommonComponents) Height() uint64 {
	return p.innerBlk.Height()
}

// Verify returns nil if:
// 1) [p]'s inner block is not an oracle block
// 2) [child]'s P-Chain height >= [parentPChainHeight]
// 3) [p]'s inner block is the parent of [c]'s inner block
// 4) [child]'s timestamp isn't before [p]'s timestamp
// 5) [child]'s timestamp is within the skew bound
// 6) [childPChainHeight] <= the current P-Chain height
// 7) [child]'s timestamp is within its proposer's window
// 8) [child] has a valid signature from its proposer
// 9) [child]'s inner block is valid
func (p *postForkCommonComponents) Verify(parentTimestamp time.Time, parentPChainHeight uint64, child *postForkBlock) error {
	if err := verifyIsNotOracleBlock(p.innerBlk); err != nil {
		return err
	}

	childPChainHeight := child.PChainHeight()
	if childPChainHeight < parentPChainHeight {
		return errPChainHeightNotMonotonic
	}

	expectedInnerParentID := p.innerBlk.ID()
	innerParentID := child.innerBlk.Parent()
	if innerParentID != expectedInnerParentID {
		return errInnerParentMismatch
	}

	childTimestamp := child.Timestamp()
	if childTimestamp.Before(parentTimestamp) {
		return errTimeNotMonotonic
	}

	maxTimestamp := p.vm.Time().Add(maxSkew)
	if childTimestamp.After(maxTimestamp) {
		return errTimeTooAdvanced
	}

	// If the node is currently bootstrapping - we don't assume that the P-chain
	// has been synced up to this point yet.
	if p.vm.bootstrapped {
		childID := child.ID()
		currentPChainHeight, err := p.vm.ctx.ValidatorState.GetCurrentHeight()
		if err != nil {
			p.vm.ctx.Log.Error("failed to get current P-Chain height while processing %s: %s",
				childID, err)
			return err
		}
		if childPChainHeight > currentPChainHeight {
			return errPChainHeightNotReached
		}

		childHeight := child.Height()
		proposerID := child.Proposer()
		minDelay, err := p.vm.Windower.Delay(childHeight, parentPChainHeight, proposerID)
		if err != nil {
			return err
		}

		delay := childTimestamp.Sub(parentTimestamp)
		if delay < minDelay {
			return errProposerWindowNotStarted
		}

		// Verify the signature of the node
		shouldHaveProposer := delay < proposer.MaxDelay
		if err := child.SignedBlock.Verify(shouldHaveProposer, p.vm.ctx.ChainID); err != nil {
			return err
		}

		p.vm.ctx.Log.Debug("verified post-fork block %s - parent timestamp %v, expected delay %v, block timestamp %v",
			childID, parentTimestamp, minDelay, childTimestamp)
	}

	return p.vm.verifyAndRecordInnerBlk(child)
}

// Return the child (a *postForkBlock) of this block
func (p *postForkCommonComponents) buildChild(
	parentID ids.ID,
	parentTimestamp time.Time,
	parentPChainHeight uint64,
) (Block, error) {
	// Child's timestamp is the later of now and this block's timestamp
	newTimestamp := p.vm.Time().Truncate(time.Second)
	if newTimestamp.Before(parentTimestamp) {
		newTimestamp = parentTimestamp
	}

	// The child's P-Chain height is proposed as the optimal P-Chain height that
	// is at least the parent's P-Chain height
	pChainHeight, err := p.vm.optimalPChainHeight(parentPChainHeight)
	if err != nil {
		return nil, err
	}

	delay := newTimestamp.Sub(parentTimestamp)
	if delay < proposer.MaxDelay {
		parentHeight := p.innerBlk.Height()
		proposerID := p.vm.ctx.NodeID
		minDelay, err := p.vm.Windower.Delay(parentHeight+1, parentPChainHeight, proposerID)
		if err != nil {
			return nil, err
		}

		if delay < minDelay {
			// It's not our turn to propose a block yet. This is likely caused
			// by having previously notified the consensus engine to attempt to
			// build a block on top of a block that is no longer the preferred
			// block.
			p.vm.ctx.Log.Debug("build block dropped; parent timestamp %s, expected delay %s, block timestamp %s",
				parentTimestamp, minDelay, newTimestamp)

			// In case the inner VM only issued one pendingTxs message, we
			// should attempt to re-handle that once it is our turn to build the
			// block.
			p.vm.notifyInnerBlockReady()
			return nil, errProposerWindowNotStarted
		}
	}

	innerBlock, err := p.vm.ChainVM.BuildBlock()
	if err != nil {
		return nil, err
	}

	// Build the child
	var statelessChild block.SignedBlock
	if delay >= proposer.MaxDelay {
		statelessChild, err = block.BuildUnsigned(
			parentID,
			newTimestamp,
			pChainHeight,
			innerBlock.Bytes(),
		)
		if err != nil {
			return nil, err
		}
	} else {
		statelessChild, err = block.Build(
			parentID,
			newTimestamp,
			pChainHeight,
			p.vm.ctx.StakingCertLeaf,
			innerBlock.Bytes(),
			p.vm.ctx.ChainID,
			p.vm.ctx.StakingLeafSigner,
		)
		if err != nil {
			return nil, err
		}
	}

	child := &postForkBlock{
		SignedBlock: statelessChild,
		postForkCommonComponents: postForkCommonComponents{
			vm:       p.vm,
			innerBlk: innerBlock,
			status:   choices.Processing,
		},
	}

	p.vm.ctx.Log.Info("built block %s - parent timestamp %v, block timestamp %v",
		child.ID(), parentTimestamp, newTimestamp)
	// Persist the child
	return child, nil
}

func (p *postForkCommonComponents) getInnerBlk() snowman.Block {
	return p.innerBlk
}

func (p *postForkCommonComponents) setInnerBlk(innerBlk snowman.Block) {
	p.innerBlk = innerBlk
}

func verifyIsOracleBlock(b snowman.Block) error {
	oracle, ok := b.(snowman.OracleBlock)
	if !ok {
		return errUnexpectedBlockType
	}
	_, err := oracle.Options()
	return err
}

func verifyIsNotOracleBlock(b snowman.Block) error {
	oracle, ok := b.(snowman.OracleBlock)
	if !ok {
		return nil
	}
	_, err := oracle.Options()
	switch err {
	case nil:
		return errUnexpectedBlockType
	case snowman.ErrNotOracle:
		return nil
	default:
		return err
	}
}
<commit_msg>Remove Incorrect Comment now (#1014)<commit_after>// Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package proposervm

import (
	"errors"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/choices"
	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
	"github.com/ava-labs/avalanchego/vms/proposervm/block"
	"github.com/ava-labs/avalanchego/vms/proposervm/proposer"
)

const (
	// allowable block issuance in the future
	maxSkew = 10 * time.Second
)

var (
	errUnsignedChild            = errors.New("expected child to be signed")
	errUnexpectedBlockType      = errors.New("unexpected proposer block type")
	errInnerParentMismatch      = errors.New("inner parentID didn't match expected parent")
	errTimeNotMonotonic         = errors.New("time must monotonically increase")
	errPChainHeightNotMonotonic = errors.New("non monotonically increasing P-chain height")
	errPChainHeightNotReached   = errors.New("block P-chain height larger than current P-chain height")
	errTimeTooAdvanced          = errors.New("time is too far advanced")
	errProposerWindowNotStarted = errors.New("proposer window hasn't started")
	errProposersNotActivated    = errors.New("proposers haven't been activated yet")
	errPChainHeightTooLow       = errors.New("block P-chain height is too low")
)

type Block interface {
	snowman.Block

	getInnerBlk() snowman.Block

	verifyPreForkChild(child *preForkBlock) error
	verifyPostForkChild(child *postForkBlock) error
	verifyPostForkOption(child *postForkOption) error

	buildChild() (Block, error)

	pChainHeight() (uint64, error)
}

type PostForkBlock interface {
	Block

	setStatus(choices.Status)
	getStatelessBlk() block.Block
	setInnerBlk(snowman.Block)
}

// field of postForkBlock and postForkOption
type postForkCommonComponents struct {
	vm       *VM
	innerBlk snowman.Block
	status   choices.Status
}

// Return the inner block's height
func (p *postForkCommonComponents) Height() uint64 {
	return p.innerBlk.Height()
}

// Verify returns nil if:
// 1) [p]'s inner block is not an oracle block
// 2) [child]'s P-Chain height >= [parentPChainHeight]
// 3) [p]'s inner block is the parent of [c]'s inner block
// 4) [child]'s timestamp isn't before [p]'s timestamp
// 5) [child]'s timestamp is within the skew bound
// 6) [childPChainHeight] <= the current P-Chain height
// 7) [child]'s timestamp is within its proposer's window
// 8) [child] has a valid signature from its proposer
// 9) [child]'s inner block is valid
func (p *postForkCommonComponents) Verify(parentTimestamp time.Time, parentPChainHeight uint64, child *postForkBlock) error {
	if err := verifyIsNotOracleBlock(p.innerBlk); err != nil {
		return err
	}

	childPChainHeight := child.PChainHeight()
	if childPChainHeight < parentPChainHeight {
		return errPChainHeightNotMonotonic
	}

	expectedInnerParentID := p.innerBlk.ID()
	innerParentID := child.innerBlk.Parent()
	if innerParentID != expectedInnerParentID {
		return errInnerParentMismatch
	}

	childTimestamp := child.Timestamp()
	if childTimestamp.Before(parentTimestamp) {
		return errTimeNotMonotonic
	}

	maxTimestamp := p.vm.Time().Add(maxSkew)
	if childTimestamp.After(maxTimestamp) {
		return errTimeTooAdvanced
	}

	// If the node is currently bootstrapping - we don't assume that the P-chain
	// has been synced up to this point yet.
	if p.vm.bootstrapped {
		childID := child.ID()
		currentPChainHeight, err := p.vm.ctx.ValidatorState.GetCurrentHeight()
		if err != nil {
			p.vm.ctx.Log.Error("failed to get current P-Chain height while processing %s: %s",
				childID, err)
			return err
		}
		if childPChainHeight > currentPChainHeight {
			return errPChainHeightNotReached
		}

		childHeight := child.Height()
		proposerID := child.Proposer()
		minDelay, err := p.vm.Windower.Delay(childHeight, parentPChainHeight, proposerID)
		if err != nil {
			return err
		}

		delay := childTimestamp.Sub(parentTimestamp)
		if delay < minDelay {
			return errProposerWindowNotStarted
		}

		// Verify the signature of the node
		shouldHaveProposer := delay < proposer.MaxDelay
		if err := child.SignedBlock.Verify(shouldHaveProposer, p.vm.ctx.ChainID); err != nil {
			return err
		}

		p.vm.ctx.Log.Debug("verified post-fork block %s - parent timestamp %v, expected delay %v, block timestamp %v",
			childID, parentTimestamp, minDelay, childTimestamp)
	}

	return p.vm.verifyAndRecordInnerBlk(child)
}

// Return the child (a *postForkBlock) of this block
func (p *postForkCommonComponents) buildChild(
	parentID ids.ID,
	parentTimestamp time.Time,
	parentPChainHeight uint64,
) (Block, error) {
	// Child's timestamp is the later of now and this block's timestamp
	newTimestamp := p.vm.Time().Truncate(time.Second)
	if newTimestamp.Before(parentTimestamp) {
		newTimestamp = parentTimestamp
	}

	// The child's P-Chain height is proposed as the optimal P-Chain height that
	// is at least the parent's P-Chain height
	pChainHeight, err := p.vm.optimalPChainHeight(parentPChainHeight)
	if err != nil {
		return nil, err
	}

	delay := newTimestamp.Sub(parentTimestamp)
	if delay < proposer.MaxDelay {
		parentHeight := p.innerBlk.Height()
		proposerID := p.vm.ctx.NodeID
		minDelay, err := p.vm.Windower.Delay(parentHeight+1, parentPChainHeight, proposerID)
		if err != nil {
			return nil, err
		}

		if delay < minDelay {
			// It's not our turn to propose a block yet. This is likely caused
			// by having previously notified the consensus engine to attempt to
			// build a block on top of a block that is no longer the preferred
			// block.
			p.vm.ctx.Log.Debug("build block dropped; parent timestamp %s, expected delay %s, block timestamp %s",
				parentTimestamp, minDelay, newTimestamp)

			// In case the inner VM only issued one pendingTxs message, we
			// should attempt to re-handle that once it is our turn to build the
			// block.
			p.vm.notifyInnerBlockReady()
			return nil, errProposerWindowNotStarted
		}
	}

	innerBlock, err := p.vm.ChainVM.BuildBlock()
	if err != nil {
		return nil, err
	}

	// Build the child
	var statelessChild block.SignedBlock
	if delay >= proposer.MaxDelay {
		statelessChild, err = block.BuildUnsigned(
			parentID,
			newTimestamp,
			pChainHeight,
			innerBlock.Bytes(),
		)
		if err != nil {
			return nil, err
		}
	} else {
		statelessChild, err = block.Build(
			parentID,
			newTimestamp,
			pChainHeight,
			p.vm.ctx.StakingCertLeaf,
			innerBlock.Bytes(),
			p.vm.ctx.ChainID,
			p.vm.ctx.StakingLeafSigner,
		)
		if err != nil {
			return nil, err
		}
	}

	child := &postForkBlock{
		SignedBlock: statelessChild,
		postForkCommonComponents: postForkCommonComponents{
			vm:       p.vm,
			innerBlk: innerBlock,
			status:   choices.Processing,
		},
	}

	p.vm.ctx.Log.Info("built block %s - parent timestamp %v, block timestamp %v",
		child.ID(), parentTimestamp, newTimestamp)
	return child, nil
}

func (p *postForkCommonComponents) getInnerBlk() snowman.Block {
	return p.innerBlk
}

func (p *postForkCommonComponents) setInnerBlk(innerBlk snowman.Block) {
	p.innerBlk = innerBlk
}

func verifyIsOracleBlock(b snowman.Block) error {
	oracle, ok := b.(snowman.OracleBlock)
	if !ok {
		return errUnexpectedBlockType
	}
	_, err := oracle.Options()
	return err
}

func verifyIsNotOracleBlock(b snowman.Block) error {
	oracle, ok := b.(snowman.OracleBlock)
	if !ok {
		return nil
	}
	_, err := oracle.Options()
	switch err {
	case nil:
		return errUnexpectedBlockType
	case snowman.ErrNotOracle:
		return nil
	default:
		return err
	}
}
<|endoftext|>


<commit_before>package minimalirc

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/textproto"
	"strings"
	"time"
	"unicode"
)

/*
 * minimalirc.go
 * small library to connect to an IRC server
 * by J. Stuart McMurray
 * created 20141130
 * last modified 20141202
 *
 * The MIT License (MIT)
 *
 * Copyright (c) 2014 J. Stuart McMurray
 *
 * 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.
 */

/* IRC represents a connection to an IRC server via OpenSSL */
type IRC struct {
	r       *textproto.Reader /* Reads messages from server */
	w       *textproto.Writer /* Writes messages to server */
	C       <-chan string     /* Messages from the server are sent here */
	E       <-chan error      /* Receives an error before close(C) */
	S       net.Conn          /* Represents the connection to the server */
	c       chan string       /* Sendable, closable C */
	e       chan error        /* Sendable E */
	Msglen  int               /* Size of an IRC message */
	Default string            /* Default target for privmsgs */
	rng     *rand.Rand        /* Random number generator */
	snick   string            /* The server's idea of our nick */

	/* Configs and defauls.  These may be changed at any time. */
	Host          string /* Host to which to connect */
	Port          uint16 /* Port to which to connect */
	Ssl           bool   /* True to use SSL/TLS */
	Hostname      string /* Hostname to verify on server's certificate */
	Nick          string /* For NICK */
	Username      string /* For USER */
	Realname      string /* For USER */
	IdNick        string /* To auth to NickServ */
	IdPass        string /* To auth to NickServ */
	Channel       string /* For JOIN */
	Chanpass      string /* For JOIN */
	Txp           string /* Prefix for logging sent messages */
	Rxp           string /* Prefix for logging received messages */
	Pongs         bool   /* Automatic ping responses */
	RandomNumbers bool   /* Append random numbers to the nick */
	QuitMessage   string /* Message to send when the client QUITs */
}

// New allocates, initializes, and returns a pointer to a new IRC struct.  hostname will be ignored if ssl is false, or assumed to be the same as host if it is the empty string and ssl is true.
func New(host string, port uint16, ssl bool, hostname string,
	nick, username, realname string) *IRC {
	/* Struct to return */
	i := &IRC{}
	/* Random number generator */
	i.rng = rand.New(rand.NewSource(time.Now().Unix()))
	/* Default max message length */
	i.Msglen = 467
	/* I/O channels */
	i.c = make(chan string)
	i.C = i.c
	i.e = make(chan error, 1)
	i.E = i.e
	i.Host = host
	i.Port = port
	i.Ssl = ssl
	if i.Ssl && "" == hostname {
		hostname = host
	}
	i.Hostname = hostname
	i.Nick = nick
	i.Username = username
	i.Realname = realname

	return i
}

// Connect connects to the server, and calls Handshake().  After connect returns, messages sent by the IRC server will be available on i.C.  If i.Rxp is set, received messages from the server will be logged via log.Printf prefixed by i.Rxp, separated by a space.  If an error is encountered reading messages from the IRC server, i.C will be closed and the error will be sent on i.E.  i.S represents the connection to the server.
func (i *IRC) Connect() error {
	/* Dial the server */
	h := net.JoinHostPort(i.Host, fmt.Sprintf("%v", i.Port))
	if i.Ssl { /* SSL requested */
		var err error
		i.S, err = tls.Dial("tcp", h,
			&tls.Config{ServerName: i.Hostname})
		if nil != err {
			return errors.New(fmt.Sprintf("unable to make ssl "+
				"connection to %v: %v", h, err))
		}
	} else { /* Plaintext connection */
		var err error
		i.S, err = net.Dial("tcp", h)
		if nil != err {
			return errors.New(fmt.Sprintf("unable to make "+
				"plaintext connection to %v: %v", h, err))
		}
	}

	/* Make a reader and a writer */
	i.r = textproto.NewReader(bufio.NewReader(i.S))
	i.w = textproto.NewWriter(bufio.NewWriter(i.S))

	/* Send nick and user */
	if err := i.Handshake(); nil != err {
		return errors.New(fmt.Sprintf("unable to handshake: %v", err))
	}

	/* Start reads from server into channel */
	go func() {
		for {
			/* Get a line from the reader */
			line, err := i.r.ReadLine()
			/* Close the channel on error */
			if nil != err {
				i.e <- err
				close(i.c)
				return
			}
			/* Log the line if needed */
			if "" != i.Rxp {
				log.Printf("%v %v", i.Rxp, line)
			}
			/* Handle pings if desired */
			if i.Pongs && strings.HasPrefix(strings.ToLower(line),
				"ping ") {
				/* Try to send pong */
				err := i.PrintfLine("PONG %v",
					strings.SplitN(line, " ", 2)[1])
				/* A send error is as bad as a read error */
				if nil != err {
					i.e <- err
					close(i.c)
				}
				return
			}
			/* Maybe get a nick */
			parts := strings.SplitN(line, " ", 4)
			/* If the 2nd bit is a 3-digit number, the 3rd bit is
			our nick */
			if 4 == len(parts) {
				n := []rune(parts[1])
				if 3 == len(n) &&
					unicode.IsNumber(n[0]) &&
					unicode.IsDigit(n[1]) &&
					unicode.IsDigit(n[2]) {
					i.snick = parts[2]
				}
			}

			/* Send out the line */
			i.c <- line
		}
	}()
	return nil
}

// ID sets the nick and user from the values in i, and sends a NICK command without any parameters (to get an easy-to-parse response with the nick as the server knows it).  If i.Nick, i.Username or i.Realname are the empty string, this is a no-op.
func (i *IRC) ID() error {
	if "" == i.Nick || "" == i.Username || "" == i.Realname {
		return nil
	}
	/* Add some numbers to the nick */
	nick := i.Nick
	if i.RandomNumbers {
		nick = fmt.Sprintf("%v-%v", nick, i.rng.Int63())
	}
	/* Iterate over the commands to send */
	for _, line := range []string{
		fmt.Sprintf("NICK :%v", nick),
		fmt.Sprintf("USER %v x x :%v", i.Username, i.Realname),
		"NICK",
	} {
		/* Try to send the line */
		if err := i.PrintfLine(line); nil != err {
			return errors.New(fmt.Sprintf("error sending ID "+
				"line %v: %v", line, err))
		}
	}
	return nil
}

// Auth authenticates to NickServ with the values in i.  If either i.IdNick or i.IdPass are the empty string, this is a no-op.
func (i *IRC) Auth() error {
	/* Don't auth with blank creds */
	if "" == i.IdNick || "" == i.IdPass {
		return nil
	}
	l := fmt.Sprintf("PRIVMSG NickServ :identify %v %v", i.IdNick,
		i.IdPass)
	if err := i.PrintfLine(l); nil != err {
		return errors.New(fmt.Sprintf("error authenticating to "+
			"services: %v", err))
	}
	return nil
}

// Join joins the channel with the optional password (which may be the empty string).  If the channel is the empty string, the value from i.Channel and i.Chanpass will be used.  If channel and i.Channel are both the empty string, this is a no-op.
func (i *IRC) Join(channel, pass string) error {
	/* If not specified, try the channel in i */
	if "" == channel {
		channel = i.Channel
		pass = i.Chanpass
	}
	/* If still no channel, no-op */
	if "" == channel {
		return nil
	}
	l := fmt.Sprintf("JOIN %v %v", channel, pass)
	if err := i.PrintfLine(l); nil != err {
		return errors.New(fmt.Sprintf("error joining %v: %v",
			channel, err))
	}
	return nil
}

// Handshake is a shorthand for ID, Auth, and Join, in that order, using the values in i.
func (i *IRC) Handshake() error {
	/* Set nick and user */
	if err := i.ID(); nil != err {
		return errors.New(fmt.Sprintf("handshake error (ID): %v", err))
	}
	/* Auth to services */
	if err := i.Auth(); err != nil {
		return errors.New(fmt.Sprintf("handshake error (Auth): %v",
			err))
	}
	/* Join the channel */
	if err := i.Join("", ""); err != nil {
		return errors.New(fmt.Sprintf("handshake error (Join): %v",
			err))
	}
	return nil
}

// PrintfLine sends the formatted string to the IRC server.  The message should be a raw IRC protocol message (like WHOIS or CAP).  It is not wrapped in PRIVMSG or anything else.  For PRIVMSGs, see Privmsg  .If i.Txp is not the empty string, successfully sent lines will be logged via log.Printf() prefixed by i.Txp, separated by a space.  Note that all the functions used to send protocol messages use PrintfLine.
func (i *IRC) PrintfLine(f string, args ...interface{}) error {
	/* Form the line into a string */
	line := fmt.Sprintf(f, args...)
	/* Try to send the line */
	if err := i.w.PrintfLine(line); err != nil {
		return err
	}
	/* Log if desired */
	if "" != i.Txp {
		log.Printf("%v %v", i.Txp, line)
	}
	return nil
}

// Target returns a target suitable for use in Privmsg, or "" if there is none.
func (i *IRC) target(target string) string {
	/* Use the default target if none was given */
	if "" == target {
		target = i.Default
	}
	/* If no default, use channel */
	if "" == target {
		target = i.Channel
	}
	/* Nop if no default target */
	if "" == target {
		return ""
	}
	return target
}

// Privmsg sends a PRIVMSG to the target, which may be a nick or a channel.  If the target is an empty string, the message will be sent to i.Target, unless that is also an empty string, in which case nothing is sent.
func (i *IRC) Privmsg(msg, target string) error {
	/* Get the target */
	t := i.target(target)
	if "" == t {
		return nil
	}
	/* Send the message */
	return i.PrintfLine("PRIVMSG %v :%v", t, msg)
}

// PrivmsgSize returns the length of the message that can be shoved into a PRIVMSG to the target.  i.Msglen may be changed to override the default size of an IRC message (467 bytes, determined experimentally on freenode, 510 should be it, though).  See Privmsg for the meaning of target.
func (i *IRC) PrivmsgSize(target string) int {
	/* Get the target */
	t := i.target(target)
	if "" == t {
		return -1
	}
	return i.Msglen - len([]byte(fmt.Sprintf("PRIVMSG %v :", t)))
}

// Nick returns a guess as to what the server thinks the nick is.  This is handy for servers that truncate nicks when RandomNumbers is true.  This is, however, only a guess (albiet a good one).  It should be called after setting the nick with Nick() or Handshake().  Note this is based on passive inspection of received messagess, which requires reading due to the read channel being unbuffered. */
func (i *IRC) SNick() string {
	return i.snick
}

// Quit sends a QUIT command to the IRC server, with the optional msg as the quit message and closes the connection if the send succeeds.  If msg is the empty string, i.QuitMessage will be used, unless it's also the empty string, in which case no message is sent with the QUIT command.
func (i *IRC) Quit(msg string) error {
	/* Use the stored message if msg is empty */
	if "" == msg && "" != i.QuitMessage {
		msg = i.QuitMessage
	}
	/* Make the message protocolish */
	if "" != msg {
		msg = " :" + msg
	}
	/* Send the quit message */
	if err := i.PrintfLine("QUIT%v", msg); nil != err {
		return err
	}
	/* Close the connection */
	if err := i.S.Close(); nil != err {
		return err
	}

	return nil
}
<commit_msg>bugfix<commit_after>package minimalirc

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/textproto"
	"strings"
	"time"
	"unicode"
)

/*
 * minimalirc.go
 * small library to connect to an IRC server
 * by J. Stuart McMurray
 * created 20141130
 * last modified 20141202
 *
 * The MIT License (MIT)
 *
 * Copyright (c) 2014 J. Stuart McMurray
 *
 * 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.
 */

/* IRC represents a connection to an IRC server via OpenSSL */
type IRC struct {
	r       *textproto.Reader /* Reads messages from server */
	w       *textproto.Writer /* Writes messages to server */
	C       <-chan string     /* Messages from the server are sent here */
	E       <-chan error      /* Receives an error before close(C) */
	S       net.Conn          /* Represents the connection to the server */
	c       chan string       /* Sendable, closable C */
	e       chan error        /* Sendable E */
	Msglen  int               /* Size of an IRC message */
	Default string            /* Default target for privmsgs */
	rng     *rand.Rand        /* Random number generator */
	snick   string            /* The server's idea of our nick */

	/* Configs and defauls.  These may be changed at any time. */
	Host          string /* Host to which to connect */
	Port          uint16 /* Port to which to connect */
	Ssl           bool   /* True to use SSL/TLS */
	Hostname      string /* Hostname to verify on server's certificate */
	Nick          string /* For NICK */
	Username      string /* For USER */
	Realname      string /* For USER */
	IdNick        string /* To auth to NickServ */
	IdPass        string /* To auth to NickServ */
	Channel       string /* For JOIN */
	Chanpass      string /* For JOIN */
	Txp           string /* Prefix for logging sent messages */
	Rxp           string /* Prefix for logging received messages */
	Pongs         bool   /* Automatic ping responses */
	RandomNumbers bool   /* Append random numbers to the nick */
	QuitMessage   string /* Message to send when the client QUITs */
}

// New allocates, initializes, and returns a pointer to a new IRC struct.  hostname will be ignored if ssl is false, or assumed to be the same as host if it is the empty string and ssl is true.
func New(host string, port uint16, ssl bool, hostname string,
	nick, username, realname string) *IRC {
	/* Struct to return */
	i := &IRC{}
	/* Random number generator */
	i.rng = rand.New(rand.NewSource(time.Now().Unix()))
	/* Default max message length */
	i.Msglen = 467
	/* I/O channels */
	i.c = make(chan string)
	i.C = i.c
	i.e = make(chan error, 1)
	i.E = i.e
	i.Host = host
	i.Port = port
	i.Ssl = ssl
	if i.Ssl && "" == hostname {
		hostname = host
	}
	i.Hostname = hostname
	i.Nick = nick
	i.Username = username
	i.Realname = realname

	return i
}

// Connect connects to the server, and calls Handshake().  After connect returns, messages sent by the IRC server will be available on i.C.  If i.Rxp is set, received messages from the server will be logged via log.Printf prefixed by i.Rxp, separated by a space.  If an error is encountered reading messages from the IRC server, i.C will be closed and the error will be sent on i.E.  i.S represents the connection to the server.
func (i *IRC) Connect() error {
	/* Dial the server */
	h := net.JoinHostPort(i.Host, fmt.Sprintf("%v", i.Port))
	if i.Ssl { /* SSL requested */
		var err error
		i.S, err = tls.Dial("tcp", h,
			&tls.Config{ServerName: i.Hostname})
		if nil != err {
			return errors.New(fmt.Sprintf("unable to make ssl "+
				"connection to %v: %v", h, err))
		}
	} else { /* Plaintext connection */
		var err error
		i.S, err = net.Dial("tcp", h)
		if nil != err {
			return errors.New(fmt.Sprintf("unable to make "+
				"plaintext connection to %v: %v", h, err))
		}
	}

	/* Make a reader and a writer */
	i.r = textproto.NewReader(bufio.NewReader(i.S))
	i.w = textproto.NewWriter(bufio.NewWriter(i.S))

	/* Send nick and user */
	if err := i.Handshake(); nil != err {
		return errors.New(fmt.Sprintf("unable to handshake: %v", err))
	}

	/* Start reads from server into channel */
	go func() {
		/* Close the connection when we're done with it */
		defer i.S.Close()
		for {
			/* Get a line from the reader */
			line, err := i.r.ReadLine()
			/* Close the channel on error */
			if nil != err {
				i.e <- err
				close(i.c)
				return
			}
			/* Log the line if needed */
			if "" != i.Rxp {
				log.Printf("%v %v", i.Rxp, line)
			}
			/* Handle pings if desired */
			if i.Pongs && strings.HasPrefix(strings.ToLower(line),
				"ping ") {
				/* Try to send pong */
				err := i.PrintfLine("PONG %v",
					strings.SplitN(line, " ", 2)[1])
				/* A send error is as bad as a read error */
				if nil != err {
					i.e <- err
					close(i.c)
					return
				}
			}
			/* Maybe get a nick */
			parts := strings.SplitN(line, " ", 4)
			/* If the 2nd bit is a 3-digit number, the 3rd bit is
			our nick */
			if 4 == len(parts) {
				n := []rune(parts[1])
				if 3 == len(n) &&
					unicode.IsDigit(n[0]) &&
					unicode.IsDigit(n[1]) &&
					unicode.IsDigit(n[2]) {
					i.snick = parts[2]
				}
			}

			/* Send out the line */
			i.c <- line
		}
	}()
	return nil
}

// ID sets the nick and user from the values in i, and sends a NICK command without any parameters (to get an easy-to-parse response with the nick as the server knows it).  If i.Nick, i.Username or i.Realname are the empty string, this is a no-op.
func (i *IRC) ID() error {
	if "" == i.Nick || "" == i.Username || "" == i.Realname {
		return nil
	}
	/* Add some numbers to the nick */
	nick := i.Nick
	if i.RandomNumbers {
		nick = fmt.Sprintf("%v-%v", nick, i.rng.Int63())
	}
	/* Iterate over the commands to send */
	for _, line := range []string{
		fmt.Sprintf("NICK :%v", nick),
		fmt.Sprintf("USER %v x x :%v", i.Username, i.Realname),
		"NICK",
	} {
		/* Try to send the line */
		if err := i.PrintfLine(line); nil != err {
			return errors.New(fmt.Sprintf("error sending ID "+
				"line %v: %v", line, err))
		}
	}
	return nil
}

// Auth authenticates to NickServ with the values in i.  If either i.IdNick or i.IdPass are the empty string, this is a no-op.
func (i *IRC) Auth() error {
	/* Don't auth with blank creds */
	if "" == i.IdNick || "" == i.IdPass {
		return nil
	}
	l := fmt.Sprintf("PRIVMSG NickServ :identify %v %v", i.IdNick,
		i.IdPass)
	if err := i.PrintfLine(l); nil != err {
		return errors.New(fmt.Sprintf("error authenticating to "+
			"services: %v", err))
	}
	return nil
}

// Join joins the channel with the optional password (which may be the empty string).  If the channel is the empty string, the value from i.Channel and i.Chanpass will be used.  If channel and i.Channel are both the empty string, this is a no-op.
func (i *IRC) Join(channel, pass string) error {
	/* If not specified, try the channel in i */
	if "" == channel {
		channel = i.Channel
		pass = i.Chanpass
	}
	/* If still no channel, no-op */
	if "" == channel {
		return nil
	}
	l := fmt.Sprintf("JOIN %v %v", channel, pass)
	if err := i.PrintfLine(l); nil != err {
		return errors.New(fmt.Sprintf("error joining %v: %v",
			channel, err))
	}
	return nil
}

// Handshake is a shorthand for ID, Auth, and Join, in that order, using the values in i.
func (i *IRC) Handshake() error {
	/* Set nick and user */
	if err := i.ID(); nil != err {
		return errors.New(fmt.Sprintf("handshake error (ID): %v", err))
	}
	/* Auth to services */
	if err := i.Auth(); err != nil {
		return errors.New(fmt.Sprintf("handshake error (Auth): %v",
			err))
	}
	/* Join the channel */
	if err := i.Join("", ""); err != nil {
		return errors.New(fmt.Sprintf("handshake error (Join): %v",
			err))
	}
	return nil
}

// PrintfLine sends the formatted string to the IRC server.  The message should be a raw IRC protocol message (like WHOIS or CAP).  It is not wrapped in PRIVMSG or anything else.  For PRIVMSGs, see Privmsg  .If i.Txp is not the empty string, successfully sent lines will be logged via log.Printf() prefixed by i.Txp, separated by a space.  Note that all the functions used to send protocol messages use PrintfLine.
func (i *IRC) PrintfLine(f string, args ...interface{}) error {
	/* Form the line into a string */
	line := fmt.Sprintf(f, args...)
	/* Try to send the line */
	if err := i.w.PrintfLine(line); err != nil {
		return err
	}
	/* Log if desired */
	if "" != i.Txp {
		log.Printf("%v %v", i.Txp, line)
	}
	return nil
}

// Target returns a target suitable for use in Privmsg, or "" if there is none.
func (i *IRC) target(target string) string {
	/* Use the default target if none was given */
	if "" == target {
		target = i.Default
	}
	/* If no default, use channel */
	if "" == target {
		target = i.Channel
	}
	/* Nop if no default target */
	if "" == target {
		return ""
	}
	return target
}

// Privmsg sends a PRIVMSG to the target, which may be a nick or a channel.  If the target is an empty string, the message will be sent to i.Target, unless that is also an empty string, in which case nothing is sent.
func (i *IRC) Privmsg(msg, target string) error {
	/* Get the target */
	t := i.target(target)
	if "" == t {
		return nil
	}
	/* Send the message */
	return i.PrintfLine("PRIVMSG %v :%v", t, msg)
}

// PrivmsgSize returns the length of the message that can be shoved into a PRIVMSG to the target.  i.Msglen may be changed to override the default size of an IRC message (467 bytes, determined experimentally on freenode, 510 should be it, though).  See Privmsg for the meaning of target.
func (i *IRC) PrivmsgSize(target string) int {
	/* Get the target */
	t := i.target(target)
	if "" == t {
		return -1
	}
	return i.Msglen - len([]byte(fmt.Sprintf("PRIVMSG %v :", t)))
}

// Nick returns a guess as to what the server thinks the nick is.  This is handy for servers that truncate nicks when RandomNumbers is true.  This is, however, only a guess (albiet a good one).  It should be called after setting the nick with Nick() or Handshake().  Note this is based on passive inspection of received messagess, which requires reading due to the read channel being unbuffered. */
func (i *IRC) SNick() string {
	return i.snick
}

// Quit sends a QUIT command to the IRC server, with the optional msg as the quit message and closes the connection if the send succeeds.  If msg is the empty string, i.QuitMessage will be used, unless it's also the empty string, in which case no message is sent with the QUIT command.
func (i *IRC) Quit(msg string) error {
	/* Use the stored message if msg is empty */
	if "" == msg && "" != i.QuitMessage {
		msg = i.QuitMessage
	}
	/* Make the message protocolish */
	if "" != msg {
		msg = " :" + msg
	}
	/* Send the quit message */
	if err := i.PrintfLine("QUIT%v", msg); nil != err {
		return err
	}
	/* Close the connection */
	if err := i.S.Close(); nil != err {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 api

import (
	"time"
)

// Config is a structure that represents configuration
// for a single test scenario.
type Config struct {
	// AutomanagedNamespaces is a number of automanaged namespaces.
	AutomanagedNamespaces int32 `json: automanagedNamespaces`
	// Steps is a sequence of test steps executed in serial.
	Steps []Step `json: steps`
	// TuningSets is a collection of tuning sets that can be used by steps.
	TuningSets []TuningSet `json: tuningSets`
}

// Step represents encapsulation of some actions. These actions could be
// object declarations or measurement usages.
// Exactly one field (Phases or Measurements) should be non-empty.
type Step struct {
	// Phases is a collection of declarative definitions of objects.
	// Phases will be executed in parallel.
	Phases []Phase `json: phases`
	// Measurements is a collection of parallel measurement calls.
	Measurements []Measurement `json: measurements`
}

// Phase is a structure that declaratively defines state of objects.
// In a given namespace range (or cluster scope if no range is specified)
// it defines the number and the configuration of managed objects.
type Phase struct {
	// NamespaceRange defines the set of namespaces in which objects
	// should be reconciled.
	// If null, objects are assumed to be cluster scoped.
	NamespaceRange *NamespaceRange `json: namespaceRange`
	// ReplicasPerNamespace is a number of instances of a given object
	// to exist in each of referenced namespaces.
	ReplicasPerNamespace int32 `json: replicasPerNamespace`
	// TuningSet is the name of TuningSet to be used.
	TuningSet string `json: tuningSet`
	// ObjectBundle declaratively defines a set of objects.
	// For every specified namespace and for every required replica,
	// these objects will be reconciled in serial.
	ObjectBundle []Object `json: objectBundle`
}

// Object is a structure that defines the object managed be the tests.
type Object struct {
	// Basename is a string from which names of objects will be created.
	Basename string `json: basename`
	// ObjectTemplatePath specifies the path to object definition.
	ObjectTemplatePath string `json: objectTemplatePath`
}

// NamespaceRange specifies the range of namespaces [Min, Max].
type NamespaceRange struct {
	// Min is the lower index of namespace range.
	Min int32 `json: min`
	// Min is the upper index of namespace range.
	Max int32 `json: max`
	// Basename defines the group of selected namespaces.
	// All of the namespaces, with name "<Basename>-<i>"
	// where <i> in [Min, Max], will be selected.
	// If no Basename is specified, automanaged namespace is assumed.
	Basename *string
}

// TuningSet defines the specific parameterization for the simulated load limit.
// It is required to have exactly one of the load structure provided.
type TuningSet struct {
	// Name by which the TuningSet will be referenced.
	Name string `json: name`
	// InitialDelay specifies the waiting time before starting phase execution.
	InitialDelay time.Duration `json: initialDelay`
	// QpsLoad is a definition for QpsLoad tuning set.
	QpsLoad *QpsLoad `json: qpsLoad`
	// RandomizedLoad is a definition for RandomizedLoad tuning set.
	RandomizedLoad *RandomizedLoad `json: randomizedLoad`
	// SteppedLoad is a definition for SteppedLoad tuning set.
	SteppedLoad *SteppedLoad `json: steppedLoad`
}

// Measurement is a structure that defines the measurement method call.
// This method call will either start or stop process of collecting specific data samples.
type Measurement struct {
	// Method is a name of a method registered in the ClusterLoader factory.
	Method string
}

// QpsLoad defines a uniform load with a given QPS.
type QpsLoad struct {
	// Qps specifies requested qps.
	Qps float64 `json: qps`
}

// RandomizedLoad defines a load that is spread randomly
// across a given total time.
type RandomizedLoad struct {
	// AverageQps specifies the expected average qps.
	AverageQps float64 `json: averageQps`
}

// SteppedLoad defines a load that generates a burst of
// a given size every X seconds.
type SteppedLoad struct {
	// BurstSize specifies the qps peek.
	BurstSize int32 `json: burstSize`
	// StepDelay specifies the interval between peeks.
	StepDelay time.Duration `json: stepDelay`
}
<commit_msg>adding mapping to api<commit_after>/*
Copyright 2018 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 api

import (
	"time"
)

// Config is a structure that represents configuration
// for a single test scenario.
type Config struct {
	// AutomanagedNamespaces is a number of automanaged namespaces.
	AutomanagedNamespaces int32 `json: automanagedNamespaces`
	// Steps is a sequence of test steps executed in serial.
	Steps []Step `json: steps`
	// TuningSets is a collection of tuning sets that can be used by steps.
	TuningSets []TuningSet `json: tuningSets`
}

// Step represents encapsulation of some actions. These actions could be
// object declarations or measurement usages.
// Exactly one field (Phases or Measurements) should be non-empty.
type Step struct {
	// Phases is a collection of declarative definitions of objects.
	// Phases will be executed in parallel.
	Phases []Phase `json: phases`
	// Measurements is a collection of parallel measurement calls.
	Measurements []Measurement `json: measurements`
}

// Phase is a structure that declaratively defines state of objects.
// In a given namespace range (or cluster scope if no range is specified)
// it defines the number and the configuration of managed objects.
type Phase struct {
	// NamespaceRange defines the set of namespaces in which objects
	// should be reconciled.
	// If null, objects are assumed to be cluster scoped.
	NamespaceRange *NamespaceRange `json: namespaceRange`
	// ReplicasPerNamespace is a number of instances of a given object
	// to exist in each of referenced namespaces.
	ReplicasPerNamespace int32 `json: replicasPerNamespace`
	// TuningSet is the name of TuningSet to be used.
	TuningSet string `json: tuningSet`
	// ObjectBundle declaratively defines a set of objects.
	// For every specified namespace and for every required replica,
	// these objects will be reconciled in serial.
	ObjectBundle []Object `json: objectBundle`
}

// Object is a structure that defines the object managed be the tests.
type Object struct {
	// Basename is a string from which names of objects will be created.
	Basename string `json: basename`
	// ObjectTemplatePath specifies the path to object definition.
	ObjectTemplatePath string `json: objectTemplatePath`
	// TemplateFillMap specifies for each placeholder what value should it be replaced with.
	TemplateFillMap map[string]interface{} `json: templateFillMap`
}

// NamespaceRange specifies the range of namespaces [Min, Max].
type NamespaceRange struct {
	// Min is the lower index of namespace range.
	Min int32 `json: min`
	// Min is the upper index of namespace range.
	Max int32 `json: max`
	// Basename defines the group of selected namespaces.
	// All of the namespaces, with name "<Basename>-<i>"
	// where <i> in [Min, Max], will be selected.
	// If no Basename is specified, automanaged namespace is assumed.
	Basename *string
}

// TuningSet defines the specific parameterization for the simulated load limit.
// It is required to have exactly one of the load structure provided.
type TuningSet struct {
	// Name by which the TuningSet will be referenced.
	Name string `json: name`
	// InitialDelay specifies the waiting time before starting phase execution.
	InitialDelay time.Duration `json: initialDelay`
	// QpsLoad is a definition for QpsLoad tuning set.
	QpsLoad *QpsLoad `json: qpsLoad`
	// RandomizedLoad is a definition for RandomizedLoad tuning set.
	RandomizedLoad *RandomizedLoad `json: randomizedLoad`
	// SteppedLoad is a definition for SteppedLoad tuning set.
	SteppedLoad *SteppedLoad `json: steppedLoad`
}

// Measurement is a structure that defines the measurement method call.
// This method call will either start or stop process of collecting specific data samples.
type Measurement struct {
	// Method is a name of a method registered in the ClusterLoader factory.
	Method string
}

// QpsLoad defines a uniform load with a given QPS.
type QpsLoad struct {
	// Qps specifies requested qps.
	Qps float64 `json: qps`
}

// RandomizedLoad defines a load that is spread randomly
// across a given total time.
type RandomizedLoad struct {
	// AverageQps specifies the expected average qps.
	AverageQps float64 `json: averageQps`
}

// SteppedLoad defines a load that generates a burst of
// a given size every X seconds.
type SteppedLoad struct {
	// BurstSize specifies the qps peek.
	BurstSize int32 `json: burstSize`
	// StepDelay specifies the interval between peeks.
	StepDelay time.Duration `json: stepDelay`
}
<|endoftext|>


<commit_before>package ambex

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

type harness struct {
	T               *testing.T
	C               chan time.Time
	version         int         // for generating versions
	updates         chan Update // we push versions here
	pushed          chan int    // versions that are pushed end up here
	usage           int         // simulated memory usage
	clock           time.Time   // current simulated time
	expectedVersion int
}

var drainTime = 10 * time.Minute

func newHarness(t *testing.T) *harness {
	C := make(chan time.Time)
	h := &harness{t, C, 0, make(chan Update), make(chan int, 10000), 0, time.Now(), 1}
	go updaterWithTicker(context.Background(), h.updates, func() int { return h.usage }, drainTime,
		&time.Ticker{C: C}, func() time.Time { return h.clock })
	return h
}

// Simulate a timer tick after the given duration.
func (h *harness) tick(d time.Duration) {
	h.clock = h.clock.Add(d)
	h.C <- h.clock
}

// Simulate an update with the specified version after the specified duration.
func (h *harness) update(d time.Duration) int {
	h.version++
	version := h.version
	h.clock = h.clock.Add(d)
	h.updates <- Update{fmt.Sprintf("%d", version), func() error {
		h.pushed <- version
		return nil
	}}
	return version
}

// Assert that a contiguous sequences of updates were pushed up to and including the specified
// version.
func (h *harness) expectUntil(version int) {
	for {
		select {
		case v := <-h.pushed:
			assert.Equal(h.T, h.expectedVersion, v)
			if v < version {
				h.expectedVersion++
				continue
			}

			if v != version {
				assert.Fail(h.T, fmt.Sprintf("expected version %d, got %d", version, v))
			}
			return
		case <-time.After(1 * time.Second):
			assert.Fail(h.T, fmt.Sprintf("expected version %d, but timed out", h.expectedVersion))
			return
		}
	}

}

// Assert that the next version is exactly the supplied version.
func (h *harness) expectExact(version int) {
	select {
	case v := <-h.pushed:
		assert.Equal(h.T, version, v)
		h.expectedVersion = version + 1
	case <-time.After(1 * time.Second):
		assert.Fail(h.T, fmt.Sprintf("expected version %d, but timed out", version))
		return
	}
}

// Assert that no version was pushed.
func (h *harness) expectNone() {
	select {
	case v := <-h.pushed:
		assert.Fail(h.T, fmt.Sprintf("expected no pushes, but go version %d", v))
	case <-time.After(1 * time.Second):
		return
	}
}

// Check to see when memory usage is zero we don't throttle at all.
func TestHappyPath(t *testing.T) {
	h := newHarness(t)
	var version int
	for i := 0; i < 1000; i++ {
		version = h.update(0)
	}
	h.expectUntil(version)
}

// Progress through the various levels of constraint and check that the correct updates are dropped.
func TestConstrained(t *testing.T) {
	h := newHarness(t)

	h.usage = 50
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 50% memory usage we only allow for 120 stale configs in memory at a time.
	h.expectUntil(120)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(1000)

	h.usage = 60
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 60% memory usage we only allow for 60 stale configs in memory at a time.
	h.expectUntil(1059)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(2000)

	h.usage = 70
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 70% memory usage we only allow for 30 stale configs in memory at a time.
	h.expectUntil(2029)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(3000)

	h.usage = 80
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 80% memory usage we only allow for 15 stale configs in memory at a time.
	h.expectUntil(3014)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(4000)

	h.usage = 90
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 90% memory usage we only allow for 1 stale config in memory at a time.
	h.expectNone()
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(5000)

	// Check that we go back to passing through everything when usage drops again.
	h.usage = 25
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	h.expectUntil(6000)
}
<commit_msg>(from AES) fix data races<commit_after>package ambex

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

type harness struct {
	T               *testing.T
	C               chan time.Time
	version         int         // for generating versions
	updates         chan Update // we push versions here
	pushed          chan int    // versions that are pushed end up here
	expectedVersion int

	mutex sync.Mutex // to proect the clock and usage
	usage int        // simulated memory usage
	clock time.Time  // current simulated time
}

var drainTime = 10 * time.Minute

func newHarness(t *testing.T) *harness {
	C := make(chan time.Time)
	h := &harness{t, C, 0, make(chan Update), make(chan int, 10000), 1, sync.Mutex{}, 0, time.Now()}
	go updaterWithTicker(context.Background(), h.updates, h.getUsage, drainTime, &time.Ticker{C: C}, h.time)
	return h
}

func (h *harness) advance(d time.Duration) time.Time {
	h.mutex.Lock()
	result := h.clock.Add(d)
	h.clock = result
	h.mutex.Unlock()
	return result
}

func (h *harness) time() time.Time {
	return h.advance(0)
}

func (h *harness) getUsage() int {
	h.mutex.Lock()
	defer h.mutex.Unlock()
	return h.usage
}

func (h *harness) setUsage(u int) {
	h.mutex.Lock()
	defer h.mutex.Unlock()
	h.usage = u
}

// Simulate a timer tick after the given duration.
func (h *harness) tick(d time.Duration) {
	h.C <- h.advance(d)
}

// Simulate an update with the specified version after the specified duration.
func (h *harness) update(d time.Duration) int {
	h.version++
	version := h.version
	h.advance(d)
	h.updates <- Update{fmt.Sprintf("%d", version), func() error {
		h.pushed <- version
		return nil
	}}
	return version
}

// Assert that a contiguous sequences of updates were pushed up to and including the specified
// version.
func (h *harness) expectUntil(version int) {
	for {
		select {
		case v := <-h.pushed:
			assert.Equal(h.T, h.expectedVersion, v)
			if v < version {
				h.expectedVersion++
				continue
			}

			if v != version {
				assert.Fail(h.T, fmt.Sprintf("expected version %d, got %d", version, v))
			}
			return
		case <-time.After(1 * time.Second):
			assert.Fail(h.T, fmt.Sprintf("expected version %d, but timed out", h.expectedVersion))
			return
		}
	}

}

// Assert that the next version is exactly the supplied version.
func (h *harness) expectExact(version int) {
	select {
	case v := <-h.pushed:
		assert.Equal(h.T, version, v)
		h.expectedVersion = version + 1
	case <-time.After(1 * time.Second):
		assert.Fail(h.T, fmt.Sprintf("expected version %d, but timed out", version))
		return
	}
}

// Assert that no version was pushed.
func (h *harness) expectNone() {
	select {
	case v := <-h.pushed:
		assert.Fail(h.T, fmt.Sprintf("expected no pushes, but go version %d", v))
	case <-time.After(1 * time.Second):
		return
	}
}

// Check to see when memory usage is zero we don't throttle at all.
func TestHappyPath(t *testing.T) {
	h := newHarness(t)
	var version int
	for i := 0; i < 1000; i++ {
		version = h.update(0)
	}
	h.expectUntil(version)
}

// Progress through the various levels of constraint and check that the correct updates are dropped.
func TestConstrained(t *testing.T) {
	h := newHarness(t)

	h.setUsage(50)
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 50% memory usage we only allow for 120 stale configs in memory at a time.
	h.expectUntil(120)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(1000)

	h.setUsage(60)
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 60% memory usage we only allow for 60 stale configs in memory at a time.
	h.expectUntil(1059)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(2000)

	h.setUsage(70)
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 70% memory usage we only allow for 30 stale configs in memory at a time.
	h.expectUntil(2029)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(3000)

	h.setUsage(80)
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 80% memory usage we only allow for 15 stale configs in memory at a time.
	h.expectUntil(3014)
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(4000)

	h.setUsage(90)
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	// Above 90% memory usage we only allow for 1 stale config in memory at a time.
	h.expectNone()
	// Fast forward by drainTime and check that the most recent update made it eventually.
	h.tick(drainTime)
	h.expectExact(5000)

	// Check that we go back to passing through everything when usage drops again.
	h.setUsage(25)
	for i := 0; i < 1000; i++ {
		h.update(0)
	}
	h.expectUntil(6000)
}
<|endoftext|>