<commit_before>package gorm

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"reflect"
	"strings"
	"time"

	"github.com/lib/pq/hstore"
)

type postgres struct {
	commonDialect
}

func (postgres) BinVar(i int) string {
	return fmt.Sprintf("$%v", i)
}

func (postgres) SupportLastInsertId() bool {
	return false
}

func (postgres) SqlTag(value reflect.Value, size int, autoIncrease bool) string {
	switch value.Kind() {
	case reflect.Bool:
		return "boolean"
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uintptr:
		if autoIncrease {
			return "serial"
		}
		return "integer"
	case reflect.Int64, reflect.Uint64:
		if autoIncrease {
			return "bigserial"
		}
		return "bigint"
	case reflect.Float32, reflect.Float64:
		return "numeric"
	case reflect.String:
		if size > 0 && size < 65532 {
			return fmt.Sprintf("varchar(%d)", size)
		}
		return "text"
	case reflect.Struct:
		if _, ok := value.Interface().(time.Time); ok {
			return "timestamp with time zone"
		}
	case reflect.Map:
		if value.Type() == hstoreType {
			return "hstore"
		}
	default:
		if isByteArrayOrSlice(value) {
			if isUUID(value) {
				return "uuid"
			}
			return "bytea"
		}
	}
	panic(fmt.Sprintf("invalid sql type %s (%s) for postgres", value.Type().Name(), value.Kind().String()))
}

var byteType = reflect.TypeOf(uint8(0))

func isByteArrayOrSlice(value reflect.Value) bool {
	return (value.Kind() == reflect.Array || value.Kind() == reflect.Slice) && value.Type().Elem() == byteType
}

func isUUID(value reflect.Value) bool {
	if value.Type().Len() != 16 {
		return false
	}
	typename := value.Type().Name()
	lower := strings.ToLower(typename)
	return "uuid" == lower || "guid" == lower
}

func (s postgres) ReturningStr(tableName, key string) string {
	return fmt.Sprintf("RETURNING %v.%v", tableName, key)
}

func (s postgres) HasTable(scope *Scope, tableName string) bool {
	var count int
	s.RawScanInt(scope, &count, "SELECT count(*) FROM INFORMATION_SCHEMA.tables WHERE table_name = ? AND table_type = 'BASE TABLE'", tableName)
	return count > 0
}

func (s postgres) HasColumn(scope *Scope, tableName string, columnName string) bool {
	var count int
	s.RawScanInt(scope, &count, "SELECT count(*) FROM INFORMATION_SCHEMA.columns WHERE table_name = ? AND column_name = ?", tableName, columnName)
	return count > 0
}

func (postgres) RemoveIndex(scope *Scope, indexName string) {
	scope.Err(scope.NewDB().Exec(fmt.Sprintf("DROP INDEX %v", indexName)).Error)
}

func (s postgres) HasIndex(scope *Scope, tableName string, indexName string) bool {
	var count int
	s.RawScanInt(scope, &count, "SELECT count(*) FROM pg_indexes WHERE tablename = ? AND indexname = ?", tableName, indexName)
	return count > 0
}

func (s postgres) CurrentDatabase(scope *Scope) (name string) {
	s.RawScanString(scope, &name, "SELECT CURRENT_DATABASE()")
	return
}

var hstoreType = reflect.TypeOf(Hstore{})

type Hstore map[string]*string

func (h Hstore) Value() (driver.Value, error) {
	hstore := hstore.Hstore{Map: map[string]sql.NullString{}}
	if len(h) == 0 {
		return nil, nil
	}

	for key, value := range h {
		var s sql.NullString
		if value != nil {
			s.String = *value
			s.Valid = true
		}
		hstore.Map[key] = s
	}
	return hstore.Value()
}

func (h *Hstore) Scan(value interface{}) error {
	hstore := hstore.Hstore{}

	if err := hstore.Scan(value); err != nil {
		return err
	}

	if len(hstore.Map) == 0 {
		return nil
	}

	*h = Hstore{}
	for k := range hstore.Map {
		if hstore.Map[k].Valid {
			s := hstore.Map[k].String
			(*h)[k] = &s
		} else {
			(*h)[k] = nil
		}
	}

	return nil
}
<commit_msg>Fix panic() in isUUID()<commit_after>package gorm

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"reflect"
	"strings"
	"time"

	"github.com/lib/pq/hstore"
)

type postgres struct {
	commonDialect
}

func (postgres) BinVar(i int) string {
	return fmt.Sprintf("$%v", i)
}

func (postgres) SupportLastInsertId() bool {
	return false
}

func (postgres) SqlTag(value reflect.Value, size int, autoIncrease bool) string {
	switch value.Kind() {
	case reflect.Bool:
		return "boolean"
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uintptr:
		if autoIncrease {
			return "serial"
		}
		return "integer"
	case reflect.Int64, reflect.Uint64:
		if autoIncrease {
			return "bigserial"
		}
		return "bigint"
	case reflect.Float32, reflect.Float64:
		return "numeric"
	case reflect.String:
		if size > 0 && size < 65532 {
			return fmt.Sprintf("varchar(%d)", size)
		}
		return "text"
	case reflect.Struct:
		if _, ok := value.Interface().(time.Time); ok {
			return "timestamp with time zone"
		}
	case reflect.Map:
		if value.Type() == hstoreType {
			return "hstore"
		}
	default:
		if isByteArrayOrSlice(value) {
			if isUUID(value) {
				return "uuid"
			}
			return "bytea"
		}
	}
	panic(fmt.Sprintf("invalid sql type %s (%s) for postgres", value.Type().Name(), value.Kind().String()))
}

var byteType = reflect.TypeOf(uint8(0))

func isByteArrayOrSlice(value reflect.Value) bool {
	return (value.Kind() == reflect.Array || value.Kind() == reflect.Slice) && value.Type().Elem() == byteType
}

func isUUID(value reflect.Value) bool {
	if value.Kind() != reflect.Array || value.Type().Len() != 16 {
		return false
	}
	typename := value.Type().Name()
	lower := strings.ToLower(typename)
	return "uuid" == lower || "guid" == lower
}

func (s postgres) ReturningStr(tableName, key string) string {
	return fmt.Sprintf("RETURNING %v.%v", tableName, key)
}

func (s postgres) HasTable(scope *Scope, tableName string) bool {
	var count int
	s.RawScanInt(scope, &count, "SELECT count(*) FROM INFORMATION_SCHEMA.tables WHERE table_name = ? AND table_type = 'BASE TABLE'", tableName)
	return count > 0
}

func (s postgres) HasColumn(scope *Scope, tableName string, columnName string) bool {
	var count int
	s.RawScanInt(scope, &count, "SELECT count(*) FROM INFORMATION_SCHEMA.columns WHERE table_name = ? AND column_name = ?", tableName, columnName)
	return count > 0
}

func (postgres) RemoveIndex(scope *Scope, indexName string) {
	scope.Err(scope.NewDB().Exec(fmt.Sprintf("DROP INDEX %v", indexName)).Error)
}

func (s postgres) HasIndex(scope *Scope, tableName string, indexName string) bool {
	var count int
	s.RawScanInt(scope, &count, "SELECT count(*) FROM pg_indexes WHERE tablename = ? AND indexname = ?", tableName, indexName)
	return count > 0
}

func (s postgres) CurrentDatabase(scope *Scope) (name string) {
	s.RawScanString(scope, &name, "SELECT CURRENT_DATABASE()")
	return
}

var hstoreType = reflect.TypeOf(Hstore{})

type Hstore map[string]*string

func (h Hstore) Value() (driver.Value, error) {
	hstore := hstore.Hstore{Map: map[string]sql.NullString{}}
	if len(h) == 0 {
		return nil, nil
	}

	for key, value := range h {
		var s sql.NullString
		if value != nil {
			s.String = *value
			s.Valid = true
		}
		hstore.Map[key] = s
	}
	return hstore.Value()
}

func (h *Hstore) Scan(value interface{}) error {
	hstore := hstore.Hstore{}

	if err := hstore.Scan(value); err != nil {
		return err
	}

	if len(hstore.Map) == 0 {
		return nil
	}

	*h = Hstore{}
	for k := range hstore.Map {
		if hstore.Map[k].Valid {
			s := hstore.Map[k].String
			(*h)[k] = &s
		} else {
			(*h)[k] = nil
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package cli

import (
	"testing"

	"github.com/gruntwork-io/terragrunt/errors"
	"github.com/hashicorp/go-version"
	"github.com/stretchr/testify/assert"
)

// Terraform Version Checking
func TestCheckTerraformVersionMeetsConstraintEqual(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.9.3", ">= v0.9.3", true)
}

func TestCheckTerraformVersionMeetsConstraintGreaterPatch(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.9.4", ">= v0.9.3", true)
}

func TestCheckTerraformVersionMeetsConstraintGreaterMajor(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v1.0.0", ">= v0.9.3", true)
}

func TestCheckTerraformVersionMeetsConstraintLessPatch(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.9.2", ">= v0.9.3", false)
}

func TestCheckTerraformVersionMeetsConstraintLessMajor(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.8.8", ">= v0.9.3", false)
}

func TestParseTerraformVersionNormal(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.3", "v0.9.3", nil)
}

func TestParseTerraformVersionWithoutV(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform 0.9.3", "0.9.3", nil)
}

func TestParseTerraformVersionWithDebug(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.4 cad024a5fe131a546936674ef85445215bbc4226", "v0.9.4", nil)
}

func TestParseTerraformVersionWithChanges(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.4-dev (cad024a5fe131a546936674ef85445215bbc4226+CHANGES)", "v0.9.4", nil)
}

func TestParseTerraformVersionWithDev(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.4-dev", "v0.9.4", nil)
}

func TestParseTerraformVersionInvalidSyntax(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "invalid-syntax", "", InvalidTerraformVersionSyntax("invalid-syntax"))
}

func testCheckTerraformVersionMeetsConstraint(t *testing.T, currentVersion string, versionConstraint string, versionMeetsConstraint bool) {
	current, err := version.NewVersion(currentVersion)
	if err != nil {
		t.Fatalf("Invalid current version specified in test: %v", err)
	}

	err = checkTerraformVersionMeetsConstraint(current, versionConstraint)
	if versionMeetsConstraint && err != nil {
		assert.Nil(t, err, "Expected Terraform version %s to meet constraint %s, but got error: %v", currentVersion, versionConstraint, err)
	} else if !versionMeetsConstraint && err == nil {
		assert.NotNil(t, err, "Expected Terraform version %s to NOT meet constraint %s, but got back a nil error", currentVersion, versionConstraint)
	}
}

func testParseTerraformVersion(t *testing.T, versionString string, expectedVersion string, expectedErr error) {
	actualVersion, actualErr := parseTerraformVersion(versionString)
	if expectedErr == nil {
		expected, err := version.NewVersion(expectedVersion)
		if err != nil {
			t.Fatalf("Invalid expected version specified in test: %v", err)
		}

		assert.Nil(t, actualErr)
		assert.Equal(t, expected, actualVersion)
	} else {
		assert.True(t, errors.IsError(actualErr, expectedErr))
	}
}

// Terragrunt Version Checking
func TestCheckTerragruntVersionMeetsConstraintEqual(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.23.18", ">= v0.23.18", true)
}

func TestCheckTerragruntVersionMeetsConstraintGreaterPatch(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.23.18", ">= v0.23.9", true)
}

func TestCheckTerragruntVersionMeetsConstraintGreaterMajor(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v1.0.0", ">= v0.23.18", true)
}

func TestCheckTerragruntVersionMeetsConstraintLessPatch(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.23.17", ">= v0.23.18", false)
}

func TestCheckTerragruntVersionMeetsConstraintLessMajor(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.22.15", ">= v0.23.18", false)
}

func testCheckTerragruntVersionMeetsConstraint(t *testing.T, currentVersion string, versionConstraint string, versionMeetsConstraint bool) {
	current, err := version.NewVersion(currentVersion)
	if err != nil {
		t.Fatalf("Invalid current version specified in test: %v", err)
	}

	err = checkTerragruntVersionMeetsConstraint(current, versionConstraint)
	if versionMeetsConstraint && err != nil {
		assert.Nil(t, err, "Expected Terragrunt version %s to meet constraint %s, but got error: %v", currentVersion, versionConstraint, err)
	} else if !versionMeetsConstraint && err == nil {
		assert.NotNil(t, err, "Expected Terragrunt version %s to NOT meet constraint %s, but got back a nil error", currentVersion, versionConstraint)
	}
}
<commit_msg>Updated to use `t.Fatalf`<commit_after>package cli

import (
	"testing"

	"github.com/gruntwork-io/terragrunt/errors"
	"github.com/hashicorp/go-version"
	"github.com/stretchr/testify/assert"
)

// Terraform Version Checking
func TestCheckTerraformVersionMeetsConstraintEqual(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.9.3", ">= v0.9.3", true)
}

func TestCheckTerraformVersionMeetsConstraintGreaterPatch(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.9.4", ">= v0.9.3", true)
}

func TestCheckTerraformVersionMeetsConstraintGreaterMajor(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v1.0.0", ">= v0.9.3", true)
}

func TestCheckTerraformVersionMeetsConstraintLessPatch(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.9.2", ">= v0.9.3", false)
}

func TestCheckTerraformVersionMeetsConstraintLessMajor(t *testing.T) {
	t.Parallel()
	testCheckTerraformVersionMeetsConstraint(t, "v0.8.8", ">= v0.9.3", false)
}

func TestParseTerraformVersionNormal(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.3", "v0.9.3", nil)
}

func TestParseTerraformVersionWithoutV(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform 0.9.3", "0.9.3", nil)
}

func TestParseTerraformVersionWithDebug(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.4 cad024a5fe131a546936674ef85445215bbc4226", "v0.9.4", nil)
}

func TestParseTerraformVersionWithChanges(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.4-dev (cad024a5fe131a546936674ef85445215bbc4226+CHANGES)", "v0.9.4", nil)
}

func TestParseTerraformVersionWithDev(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "Terraform v0.9.4-dev", "v0.9.4", nil)
}

func TestParseTerraformVersionInvalidSyntax(t *testing.T) {
	t.Parallel()
	testParseTerraformVersion(t, "invalid-syntax", "", InvalidTerraformVersionSyntax("invalid-syntax"))
}

func testCheckTerraformVersionMeetsConstraint(t *testing.T, currentVersion string, versionConstraint string, versionMeetsConstraint bool) {
	current, err := version.NewVersion(currentVersion)
	if err != nil {
		t.Fatalf("Invalid current version specified in test: %v", err)
	}

	err = checkTerraformVersionMeetsConstraint(current, versionConstraint)
	if versionMeetsConstraint && err != nil {
		assert.Nil(t, err, "Expected Terraform version %s to meet constraint %s, but got error: %v", currentVersion, versionConstraint, err)
	} else if !versionMeetsConstraint && err == nil {
		assert.NotNil(t, err, "Expected Terraform version %s to NOT meet constraint %s, but got back a nil error", currentVersion, versionConstraint)
	}
}

func testParseTerraformVersion(t *testing.T, versionString string, expectedVersion string, expectedErr error) {
	actualVersion, actualErr := parseTerraformVersion(versionString)
	if expectedErr == nil {
		expected, err := version.NewVersion(expectedVersion)
		if err != nil {
			t.Fatalf("Invalid expected version specified in test: %v", err)
		}

		assert.Nil(t, actualErr)
		assert.Equal(t, expected, actualVersion)
	} else {
		assert.True(t, errors.IsError(actualErr, expectedErr))
	}
}

// Terragrunt Version Checking
func TestCheckTerragruntVersionMeetsConstraintEqual(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.23.18", ">= v0.23.18", true)
}

func TestCheckTerragruntVersionMeetsConstraintGreaterPatch(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.23.18", ">= v0.23.9", true)
}

func TestCheckTerragruntVersionMeetsConstraintGreaterMajor(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v1.0.0", ">= v0.23.18", true)
}

func TestCheckTerragruntVersionMeetsConstraintLessPatch(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.23.17", ">= v0.23.18", false)
}

func TestCheckTerragruntVersionMeetsConstraintLessMajor(t *testing.T) {
	t.Parallel()
	testCheckTerragruntVersionMeetsConstraint(t, "v0.22.15", ">= v0.23.18", false)
}

func testCheckTerragruntVersionMeetsConstraint(t *testing.T, currentVersion string, versionConstraint string, versionMeetsConstraint bool) {
	current, err := version.NewVersion(currentVersion)
	if err != nil {
		t.Fatalf("Invalid current version specified in test: %v", err)
	}

	err = checkTerragruntVersionMeetsConstraint(current, versionConstraint)
	if versionMeetsConstraint && err != nil {
		t.Fatalf("Expected Terragrunt version %s to meet constraint %s, but got error: %v", currentVersion, versionConstraint, err)
	} else if !versionMeetsConstraint && err == nil {
		t.Fatalf("Expected Terragrunt version %s to NOT meet constraint %s, but got back a nil error", currentVersion, versionConstraint)
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013-2014, Jeremy Bingham (<jbingham@gmail.com>)
 *
 * 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 shovey

/* MySQL funcs for shovey */

import (
	"github.com/ctdk/goiardi/datastore"
	"github.com/go-sql-driver/mysql"
	"github.com/ctdk/goiardi/util"
	"net/http"
	"time"
)

func (s *Shovey) fillShoveyFromMySQL(row datastore.ResRow) error {

}

func (s *ShoveyRun) fillShoveyRunFromMySQL(row datastore.ResRow) error {

}

func (s *ShoveyRunStream) fillShoveyRunStreamFromMySQL(row datastore.ResRow) error {

}

func (s *Shovey) saveMySQL() util.Gerror {

}

func (sr *ShoveyRun) saveMySQL() util.Gerror {

}
<commit_msg>Time to make the mysql shovey bits now<commit_after>/*
 * Copyright (c) 2013-2014, Jeremy Bingham (<jbingham@gmail.com>)
 *
 * 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 shovey

/* MySQL funcs for shovey */

import (
	"github.com/ctdk/goiardi/datastore"
	"github.com/go-sql-driver/mysql"
	"github.com/ctdk/goiardi/util"
	"net/http"
	"time"
)

func (s *Shovey) fillShoveyFromMySQL(row datastore.ResRow) error {
	var ca, ua mysql.NullTime
	var nn util.StringSlice
	var tm int64
	err := row.Scan(&s.RunID, &nn, &s.Command, &ca, &ua, &s.Status, &tm, &s.Quorum)
	if err != nil {
		return err
	}
	if ca.Valid {
		s.CreatedAt = ca.Time
	}
	if ua.Valid {
		s.UpdatedAt = ua.Time
	}
	s.Timeout = time.Duration(tm)

	s.NodeNames = nn

	return nil
}

func (s *ShoveyRun) fillShoveyRunFromMySQL(row datastore.ResRow) error {
	var at, et mysql.NullTime
	err := row.Scan(&s.ID, &s.ShoveyUUID, &s.NodeName, &s.Status, &at, &et, &s.Output, &s.Error, &s.Stderr, &s.ExitStatus)
	if err != nil {
		return err
	}
	if at.Valid {
		s.AckTime = at.Time
	}
	if et.Valid {
		s.EndTime = et.Time
	}
	return nil
}

func (s *ShoveyRunStream) fillShoveyRunStreamFromMySQL(row datastore.ResRow) error {
	var ca mysql.NullTime
	err := row.Scan(&s.ShoveyUUID, &s.NodeName, &s.Seq, &s.OutputType, &s.Output, &s.IsLast, &ca)
	if err != nil {
		return err
	}
	if ca.Valid {
		s.CreatedAt = ca.Time
	}
	return nil
}

func (s *Shovey) saveMySQL() util.Gerror {
	tx, err := datastore.Dbh.Begin()
	if err != nil {
		gerr := util.CastErr(err)
		gerr.SetStatus(http.StatusInternalServerError)
		return gerr
	}
}

func (sr *ShoveyRun) saveMySQL() util.Gerror {
	tx, err := datastore.Dbh.Begin()
	if err != nil {
		gerr := util.CastErr(err)
		gerr.SetStatus(http.StatusInternalServerError)
		return gerr
	}
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd 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 compression

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"sync"

	"github.com/containerd/containerd/log"
)

type (
	// Compression is the state represents if compressed or not.
	Compression int
)

const (
	// Uncompressed represents the uncompressed.
	Uncompressed Compression = iota
	// Gzip is gzip compression algorithm.
	Gzip
)

const disablePigzEnv = "CONTAINERD_DISABLE_PIGZ"

var (
	initPigz   sync.Once
	unpigzPath string
)

var (
	bufioReader32KPool = &sync.Pool{
		New: func() interface{} { return bufio.NewReaderSize(nil, 32*1024) },
	}
)

// DecompressReadCloser include the stream after decompress and the compress method detected.
type DecompressReadCloser interface {
	io.ReadCloser
	// GetCompression returns the compress method which is used before decompressing
	GetCompression() Compression
}

type readCloserWrapper struct {
	io.Reader
	compression Compression
	closer      func() error
}

func (r *readCloserWrapper) Close() error {
	if r.closer != nil {
		return r.closer()
	}
	return nil
}

func (r *readCloserWrapper) GetCompression() Compression {
	return r.compression
}

type writeCloserWrapper struct {
	io.Writer
	closer func() error
}

func (w *writeCloserWrapper) Close() error {
	if w.closer != nil {
		w.closer()
	}
	return nil
}

type bufferedReader struct {
	buf *bufio.Reader
}

func newBufferedReader(r io.Reader) *bufferedReader {
	buf := bufioReader32KPool.Get().(*bufio.Reader)
	buf.Reset(r)
	return &bufferedReader{buf}
}

func (r *bufferedReader) Read(p []byte) (n int, err error) {
	if r.buf == nil {
		return 0, io.EOF
	}
	n, err = r.buf.Read(p)
	if err == io.EOF {
		r.buf.Reset(nil)
		bufioReader32KPool.Put(r.buf)
		r.buf = nil
	}
	return
}

func (r *bufferedReader) Peek(n int) ([]byte, error) {
	if r.buf == nil {
		return nil, io.EOF
	}
	return r.buf.Peek(n)
}

// DetectCompression detects the compression algorithm of the source.
func DetectCompression(source []byte) Compression {
	for compression, m := range map[Compression][]byte{
		Gzip: {0x1F, 0x8B, 0x08},
	} {
		if len(source) < len(m) {
			// Len too short
			continue
		}
		if bytes.Equal(m, source[:len(m)]) {
			return compression
		}
	}
	return Uncompressed
}

// DecompressStream decompresses the archive and returns a ReaderCloser with the decompressed archive.
func DecompressStream(archive io.Reader) (DecompressReadCloser, error) {
	buf := newBufferedReader(archive)
	bs, err := buf.Peek(10)
	if err != nil && err != io.EOF {
		// Note: we'll ignore any io.EOF error because there are some odd
		// cases where the layer.tar file will be empty (zero bytes) and
		// that results in an io.EOF from the Peek() call. So, in those
		// cases we'll just treat it as a non-compressed stream and
		// that means just create an empty layer.
		// See Issue docker/docker#18170
		return nil, err
	}

	switch compression := DetectCompression(bs); compression {
	case Uncompressed:
		return &readCloserWrapper{
			Reader:      buf,
			compression: compression,
		}, nil
	case Gzip:
		ctx, cancel := context.WithCancel(context.Background())
		gzReader, err := gzipDecompress(ctx, buf)
		if err != nil {
			cancel()
			return nil, err
		}

		return &readCloserWrapper{
			Reader:      gzReader,
			compression: compression,
			closer: func() error {
				cancel()
				return gzReader.Close()
			},
		}, nil

	default:
		return nil, fmt.Errorf("unsupported compression format %s", (&compression).Extension())
	}
}

// CompressStream compresseses the dest with specified compression algorithm.
func CompressStream(dest io.Writer, compression Compression) (io.WriteCloser, error) {
	switch compression {
	case Uncompressed:
		return &writeCloserWrapper{dest, nil}, nil
	case Gzip:
		return gzip.NewWriter(dest), nil
	default:
		return nil, fmt.Errorf("unsupported compression format %s", (&compression).Extension())
	}
}

// Extension returns the extension of a file that uses the specified compression algorithm.
func (compression *Compression) Extension() string {
	switch *compression {
	case Gzip:
		return "gz"
	}
	return ""
}

func gzipDecompress(ctx context.Context, buf io.Reader) (io.ReadCloser, error) {
	initPigz.Do(func() {
		if unpigzPath = detectPigz(); unpigzPath != "" {
			log.L.Debug("using pigz for decompression")
		}
	})

	if unpigzPath == "" {
		return gzip.NewReader(buf)
	}

	return cmdStream(exec.CommandContext(ctx, unpigzPath, "-d", "-c"), buf)
}

func cmdStream(cmd *exec.Cmd, in io.Reader) (io.ReadCloser, error) {
	reader, writer := io.Pipe()

	cmd.Stdin = in
	cmd.Stdout = writer

	var errBuf bytes.Buffer
	cmd.Stderr = &errBuf

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	go func() {
		if err := cmd.Wait(); err != nil {
			writer.CloseWithError(fmt.Errorf("%s: %s", err, errBuf.String()))
		} else {
			writer.Close()
		}
	}()

	return reader, nil
}

func detectPigz() string {
	path, err := exec.LookPath("unpigz")
	if err != nil {
		log.L.WithError(err).Debug("unpigz not found, falling back to go gzip")
		return ""
	}

	// Check if pigz disabled via CONTAINERD_DISABLE_PIGZ env variable
	value := os.Getenv(disablePigzEnv)
	if value == "" {
		return path
	}

	disable, err := strconv.ParseBool(value)
	if err != nil {
		log.L.WithError(err).Warnf("could not parse %s: %s", disablePigzEnv, value)
		return path
	}

	if disable {
		return ""
	}

	return path
}
<commit_msg>fix wrong spells in compression.go<commit_after>/*
   Copyright The containerd 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 compression

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"sync"

	"github.com/containerd/containerd/log"
)

type (
	// Compression is the state represents if compressed or not.
	Compression int
)

const (
	// Uncompressed represents the uncompressed.
	Uncompressed Compression = iota
	// Gzip is gzip compression algorithm.
	Gzip
)

const disablePigzEnv = "CONTAINERD_DISABLE_PIGZ"

var (
	initPigz   sync.Once
	unpigzPath string
)

var (
	bufioReader32KPool = &sync.Pool{
		New: func() interface{} { return bufio.NewReaderSize(nil, 32*1024) },
	}
)

// DecompressReadCloser include the stream after decompress and the compress method detected.
type DecompressReadCloser interface {
	io.ReadCloser
	// GetCompression returns the compress method which is used before decompressing
	GetCompression() Compression
}

type readCloserWrapper struct {
	io.Reader
	compression Compression
	closer      func() error
}

func (r *readCloserWrapper) Close() error {
	if r.closer != nil {
		return r.closer()
	}
	return nil
}

func (r *readCloserWrapper) GetCompression() Compression {
	return r.compression
}

type writeCloserWrapper struct {
	io.Writer
	closer func() error
}

func (w *writeCloserWrapper) Close() error {
	if w.closer != nil {
		w.closer()
	}
	return nil
}

type bufferedReader struct {
	buf *bufio.Reader
}

func newBufferedReader(r io.Reader) *bufferedReader {
	buf := bufioReader32KPool.Get().(*bufio.Reader)
	buf.Reset(r)
	return &bufferedReader{buf}
}

func (r *bufferedReader) Read(p []byte) (n int, err error) {
	if r.buf == nil {
		return 0, io.EOF
	}
	n, err = r.buf.Read(p)
	if err == io.EOF {
		r.buf.Reset(nil)
		bufioReader32KPool.Put(r.buf)
		r.buf = nil
	}
	return
}

func (r *bufferedReader) Peek(n int) ([]byte, error) {
	if r.buf == nil {
		return nil, io.EOF
	}
	return r.buf.Peek(n)
}

// DetectCompression detects the compression algorithm of the source.
func DetectCompression(source []byte) Compression {
	for compression, m := range map[Compression][]byte{
		Gzip: {0x1F, 0x8B, 0x08},
	} {
		if len(source) < len(m) {
			// Len too short
			continue
		}
		if bytes.Equal(m, source[:len(m)]) {
			return compression
		}
	}
	return Uncompressed
}

// DecompressStream decompresses the archive and returns a ReaderCloser with the decompressed archive.
func DecompressStream(archive io.Reader) (DecompressReadCloser, error) {
	buf := newBufferedReader(archive)
	bs, err := buf.Peek(10)
	if err != nil && err != io.EOF {
		// Note: we'll ignore any io.EOF error because there are some odd
		// cases where the layer.tar file will be empty (zero bytes) and
		// that results in an io.EOF from the Peek() call. So, in those
		// cases we'll just treat it as a non-compressed stream and
		// that means just create an empty layer.
		// See Issue docker/docker#18170
		return nil, err
	}

	switch compression := DetectCompression(bs); compression {
	case Uncompressed:
		return &readCloserWrapper{
			Reader:      buf,
			compression: compression,
		}, nil
	case Gzip:
		ctx, cancel := context.WithCancel(context.Background())
		gzReader, err := gzipDecompress(ctx, buf)
		if err != nil {
			cancel()
			return nil, err
		}

		return &readCloserWrapper{
			Reader:      gzReader,
			compression: compression,
			closer: func() error {
				cancel()
				return gzReader.Close()
			},
		}, nil

	default:
		return nil, fmt.Errorf("unsupported compression format %s", (&compression).Extension())
	}
}

// CompressStream compresses the dest with specified compression algorithm.
func CompressStream(dest io.Writer, compression Compression) (io.WriteCloser, error) {
	switch compression {
	case Uncompressed:
		return &writeCloserWrapper{dest, nil}, nil
	case Gzip:
		return gzip.NewWriter(dest), nil
	default:
		return nil, fmt.Errorf("unsupported compression format %s", (&compression).Extension())
	}
}

// Extension returns the extension of a file that uses the specified compression algorithm.
func (compression *Compression) Extension() string {
	switch *compression {
	case Gzip:
		return "gz"
	}
	return ""
}

func gzipDecompress(ctx context.Context, buf io.Reader) (io.ReadCloser, error) {
	initPigz.Do(func() {
		if unpigzPath = detectPigz(); unpigzPath != "" {
			log.L.Debug("using pigz for decompression")
		}
	})

	if unpigzPath == "" {
		return gzip.NewReader(buf)
	}

	return cmdStream(exec.CommandContext(ctx, unpigzPath, "-d", "-c"), buf)
}

func cmdStream(cmd *exec.Cmd, in io.Reader) (io.ReadCloser, error) {
	reader, writer := io.Pipe()

	cmd.Stdin = in
	cmd.Stdout = writer

	var errBuf bytes.Buffer
	cmd.Stderr = &errBuf

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	go func() {
		if err := cmd.Wait(); err != nil {
			writer.CloseWithError(fmt.Errorf("%s: %s", err, errBuf.String()))
		} else {
			writer.Close()
		}
	}()

	return reader, nil
}

func detectPigz() string {
	path, err := exec.LookPath("unpigz")
	if err != nil {
		log.L.WithError(err).Debug("unpigz not found, falling back to go gzip")
		return ""
	}

	// Check if pigz disabled via CONTAINERD_DISABLE_PIGZ env variable
	value := os.Getenv(disablePigzEnv)
	if value == "" {
		return path
	}

	disable, err := strconv.ParseBool(value)
	if err != nil {
		log.L.WithError(err).Warnf("could not parse %s: %s", disablePigzEnv, value)
		return path
	}

	if disable {
		return ""
	}

	return path
}
<|endoftext|>


<commit_before>package clicommand

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/buildkite/agent/agent"
	"github.com/buildkite/agent/cliconfig"
	"github.com/buildkite/agent/logger"
	"github.com/urfave/cli"
)

var StartDescription = `Usage:

   buildkite-agent start [arguments...]

Description:

   When a job is ready to run it will call the "bootstrap-script"
   and pass it all the environment variables required for the job to run.
   This script is responsible for checking out the code, and running the
   actual build script defined in the pipeline.

   The agent will run any jobs within a PTY (pseudo terminal) if available.

Example:

   $ buildkite-agent start --token xxx`

type AgentStartConfig struct {
	Config                    string   `cli:"config"`
	Token                     string   `cli:"token" validate:"required"`
	Name                      string   `cli:"name"`
	Priority                  string   `cli:"priority"`
	DisconnectAfterJob        bool     `cli:"disconnect-after-job"`
	DisconnectAfterJobTimeout int      `cli:"disconnect-after-job-timeout"`
	BootstrapScript           string   `cli:"bootstrap-script" normalize:"filepath"`
	BuildPath                 string   `cli:"build-path" normalize:"filepath" validate:"required"`
	HooksPath                 string   `cli:"hooks-path" normalize:"filepath"`
	PluginsPath               string   `cli:"plugins-path" normalize:"filepath"`
	Shell                     string   `cli:"shell"`
	Tags                      []string `cli:"tags"`
	TagsFromEC2               bool     `cli:"tags-from-ec2"`
	TagsFromEC2Tags           bool     `cli:"tags-from-ec2-tags"`
	TagsFromGCP               bool     `cli:"tags-from-gcp"`
	WaitForEC2TagsTimeout     string   `cli:"wait-for-ec2-tags-timeout"`
	GitCloneFlags             string   `cli:"git-clone-flags"`
	GitCleanFlags             string   `cli:"git-clean-flags"`
	NoColor                   bool     `cli:"no-color"`
	NoSSHKeyscan              bool     `cli:"no-ssh-keyscan"`
	NoCommandEval             bool     `cli:"no-command-eval"`
	NoPlugins                 bool     `cli:"no-plugins"`
	NoPTY                     bool     `cli:"no-pty"`
	TimestampLines            bool     `cli:"timestamp-lines"`
	Endpoint                  string   `cli:"endpoint" validate:"required"`
	Debug                     bool     `cli:"debug"`
	DebugHTTP                 bool     `cli:"debug-http"`
	Experiments               []string `cli:"experiment"`

	/* Deprecated */
	NoSSHFingerprintVerification bool     `cli:"no-automatic-ssh-fingerprint-verification" deprecated-and-renamed-to:"NoSSHKeyscan"`
	MetaData                     []string `cli:"meta-data" deprecated-and-renamed-to:"Tags"`
	MetaDataEC2                  bool     `cli:"meta-data-ec2" deprecated-and-renamed-to:"TagsFromEC2"`
	MetaDataEC2Tags              bool     `cli:"meta-data-ec2-tags" deprecated-and-renamed-to:"TagsFromEC2Tags"`
	MetaDataGCP                  bool     `cli:"meta-data-gcp" deprecated-and-renamed-to:"TagsFromGCP"`
}

func DefaultShell() string {
	// https://github.com/golang/go/blob/master/src/go/build/syslist.go#L7
	switch runtime.GOOS {
	case "windows":
		return `C:\Windows\System32\CMD.exe /S /C`
	case "freebsd", "openbsd", "netbsd":
		return `/usr/local/bin/bash -e -c`
	default:
		return `/bin/bash -e -c`
	}
}

func DefaultConfigFilePaths() (paths []string) {
	// Toggle beetwen windows an *nix paths
	if runtime.GOOS == "windows" {
		paths = []string{
			"$USERPROFILE\\AppData\\Local\\BuildkiteAgent\\buildkite-agent.cfg",
		}
	} else {
		paths = []string{
			"$HOME/.buildkite-agent/buildkite-agent.cfg",
			"/usr/local/etc/buildkite-agent/buildkite-agent.cfg",
			"/etc/buildkite-agent/buildkite-agent.cfg",
		}
	}

	// Also check to see if there's a buildkite-agent.cfg in the folder
	// that the binary is running in.
	pathToBinary, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err == nil {
		pathToRelativeConfig := filepath.Join(pathToBinary, "buildkite-agent.cfg")
		paths = append([]string{pathToRelativeConfig}, paths...)
	}

	return
}

var AgentStartCommand = cli.Command{
	Name:        "start",
	Usage:       "Starts a Buildkite agent",
	Description: StartDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "config",
			Value:  "",
			Usage:  "Path to a configuration file",
			EnvVar: "BUILDKITE_AGENT_CONFIG",
		},
		cli.StringFlag{
			Name:   "token",
			Value:  "",
			Usage:  "Your account agent token",
			EnvVar: "BUILDKITE_AGENT_TOKEN",
		},
		cli.StringFlag{
			Name:   "name",
			Value:  "",
			Usage:  "The name of the agent",
			EnvVar: "BUILDKITE_AGENT_NAME",
		},
		cli.StringFlag{
			Name:   "priority",
			Value:  "",
			Usage:  "The priority of the agent (higher priorities are assigned work first)",
			EnvVar: "BUILDKITE_AGENT_PRIORITY",
		},
		cli.BoolFlag{
			Name:   "disconnect-after-job",
			Usage:  "Disconnect the agent after running a job",
			EnvVar: "BUILDKITE_AGENT_DISCONNECT_AFTER_JOB",
		},
		cli.IntFlag{
			Name:   "disconnect-after-job-timeout",
			Value:  120,
			Usage:  "When --disconnect-after-job is specified, the number of seconds to wait for a job before shutting down",
			EnvVar: "BUILDKITE_AGENT_DISCONNECT_AFTER_JOB_TIMEOUT",
		},
		cli.StringFlag{
			Name:   "shell",
			Value:  DefaultShell(),
			Usage:  "The shell to use to interpret build commands",
			EnvVar: "BUILDKITE_SHELL",
		},
		cli.StringSliceFlag{
			Name:   "tags",
			Value:  &cli.StringSlice{},
			Usage:  "A comma-separated list of tags for the agent (e.g. \"linux\" or \"mac,xcode=8\")",
			EnvVar: "BUILDKITE_AGENT_TAGS",
		},
		cli.BoolFlag{
			Name:   "tags-from-ec2",
			Usage:  "Include the host's EC2 meta-data as tags (instance-id, instance-type, and ami-id)",
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_EC2",
		},
		cli.BoolFlag{
			Name:   "tags-from-ec2-tags",
			Usage:  "Include the host's EC2 tags as tags",
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_EC2_TAGS",
		},
		cli.BoolFlag{
			Name:   "tags-from-gcp",
			Usage:  "Include the host's Google Cloud meta-data as tags (instance-id, machine-type, preemptible, project-id, region, and zone)",
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_GCP",
		},
		cli.DurationFlag{
			Name:   "wait-for-ec2-tags-timeout",
			Usage:  "The amount of time to wait for tags from EC2 before proceeding",
			EnvVar: "BUILDKITE_AGENT_WAIT_FOR_EC2_TAGS_TIMEOUT",
			Value:  time.Second * 10,
		},
		cli.StringFlag{
			Name:   "git-clone-flags",
			Value:  "-v",
			Usage:  "Flags to pass to the \"git clone\" command",
			EnvVar: "BUILDKITE_GIT_CLONE_FLAGS",
		},
		cli.StringFlag{
			Name:   "git-clean-flags",
			Value:  "-fxdq",
			Usage:  "Flags to pass to \"git clean\" command",
			EnvVar: "BUILDKITE_GIT_CLEAN_FLAGS",
		},
		cli.StringFlag{
			Name:   "bootstrap-script",
			Value:  "",
			Usage:  "The command that is executed for bootstrapping a job, defaults to the bootstrap sub-command of this binary",
			EnvVar: "BUILDKITE_BOOTSTRAP_SCRIPT_PATH",
		},
		cli.StringFlag{
			Name:   "build-path",
			Value:  "",
			Usage:  "Path to where the builds will run from",
			EnvVar: "BUILDKITE_BUILD_PATH",
		},
		cli.StringFlag{
			Name:   "hooks-path",
			Value:  "",
			Usage:  "Directory where the hook scripts are found",
			EnvVar: "BUILDKITE_HOOKS_PATH",
		},
		cli.StringFlag{
			Name:   "plugins-path",
			Value:  "",
			Usage:  "Directory where the plugins are saved to",
			EnvVar: "BUILDKITE_PLUGINS_PATH",
		},
		cli.BoolFlag{
			Name:   "timestamp-lines",
			Usage:  "Prepend timestamps on each line of output.",
			EnvVar: "BUILDKITE_TIMESTAMP_LINES",
		},
		cli.BoolFlag{
			Name:   "no-pty",
			Usage:  "Do not run jobs within a pseudo terminal",
			EnvVar: "BUILDKITE_NO_PTY",
		},
		cli.BoolFlag{
			Name:   "no-ssh-keyscan",
			Usage:  "Don't automatically run ssh-keyscan before checkout",
			EnvVar: "BUILDKITE_NO_SSH_KEYSCAN",
		},
		cli.BoolFlag{
			Name:   "no-command-eval",
			Usage:  "Don't allow this agent to run arbitrary console commands",
			EnvVar: "BUILDKITE_NO_COMMAND_EVAL",
		},
		cli.BoolFlag{
			Name:   "no-plugins",
			Usage:  "Don't allow this agent to load plugins",
			EnvVar: "BUILDKITE_NO_PLUGINS",
		},
		ExperimentsFlag,
		EndpointFlag,
		NoColorFlag,
		DebugFlag,
		DebugHTTPFlag,
		/* Deprecated flags which will be removed in v4 */
		cli.StringSliceFlag{
			Name:   "meta-data",
			Value:  &cli.StringSlice{},
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_META_DATA",
		},
		cli.BoolFlag{
			Name:   "meta-data-ec2",
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_META_DATA_EC2",
		},
		cli.BoolFlag{
			Name:   "meta-data-ec2-tags",
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_EC2_TAGS",
		},
		cli.BoolFlag{
			Name:   "meta-data-gcp",
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_META_DATA_GCP",
		},
		cli.BoolFlag{
			Name:   "no-automatic-ssh-fingerprint-verification",
			Hidden: true,
			EnvVar: "BUILDKITE_NO_AUTOMATIC_SSH_FINGERPRINT_VERIFICATION",
		},
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := AgentStartConfig{}

		// Setup the config loader. You'll see that we also path paths to
		// potential config files. The loader will use the first one it finds.
		loader := cliconfig.Loader{
			CLI:                    c,
			Config:                 &cfg,
			DefaultConfigFilePaths: DefaultConfigFilePaths(),
		}

		// Load the configuration
		if err := loader.Load(); err != nil {
			logger.Fatal("%s", err)
		}

		// Setup the any global configuration options
		HandleGlobalFlags(cfg)

		// Force some settings if on Windows (these aren't supported
		// yet)
		if runtime.GOOS == "windows" {
			cfg.NoPTY = true
		}

		// Set a useful default for the bootstrap script
		if cfg.BootstrapScript == "" {
			cfg.BootstrapScript = fmt.Sprintf("%q bootstrap", os.Args[0])
		}

		// Guess the shell if none is provided
		if cfg.Shell == "" {
			cfg.Shell = DefaultShell()
		}

		logger.Debug("Using shell %q", cfg.Shell)

		// Make sure the DisconnectAfterJobTimeout value is correct
		if cfg.DisconnectAfterJob && cfg.DisconnectAfterJobTimeout < 120 {
			logger.Fatal("The timeout for `disconnect-after-job` must be at least 120 seconds")
		}

		var ec2TagTimeout time.Duration
		if t := cfg.WaitForEC2TagsTimeout; t != "" {
			var err error
			ec2TagTimeout, err = time.ParseDuration(t)
			if err != nil {
				logger.Fatal("Failed to parse ec2 tag timeout: %v", err)
			}
		}

		// Setup the agent
		pool := agent.AgentPool{
			Token:                 cfg.Token,
			Name:                  cfg.Name,
			Priority:              cfg.Priority,
			Tags:                  cfg.Tags,
			TagsFromEC2:           cfg.TagsFromEC2,
			TagsFromEC2Tags:       cfg.TagsFromEC2Tags,
			TagsFromGCP:           cfg.TagsFromGCP,
			WaitForEC2TagsTimeout: ec2TagTimeout,
			Endpoint:              cfg.Endpoint,
			AgentConfiguration: &agent.AgentConfiguration{
				BootstrapScript:           cfg.BootstrapScript,
				BuildPath:                 cfg.BuildPath,
				HooksPath:                 cfg.HooksPath,
				PluginsPath:               cfg.PluginsPath,
				GitCloneFlags:             cfg.GitCloneFlags,
				GitCleanFlags:             cfg.GitCleanFlags,
				SSHKeyscan:                !cfg.NoSSHKeyscan,
				CommandEval:               !cfg.NoCommandEval,
				PluginsEnabled:            !cfg.NoPlugins,
				RunInPty:                  !cfg.NoPTY,
				TimestampLines:            cfg.TimestampLines,
				DisconnectAfterJob:        cfg.DisconnectAfterJob,
				DisconnectAfterJobTimeout: cfg.DisconnectAfterJobTimeout,
				Shell: cfg.Shell,
			},
		}

		// Store the loaded config file path on the pool so we can
		// show it when the agent starts
		if loader.File != nil {
			pool.ConfigFilePath = loader.File.Path
		}

		// Start the agent pool
		if err := pool.Start(); err != nil {
			logger.Fatal("%s", err)
		}
	},
}
<commit_msg>Fixes default bootstrap script path on Windows<commit_after>package clicommand

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/buildkite/agent/agent"
	"github.com/buildkite/agent/cliconfig"
	"github.com/buildkite/agent/logger"
	"github.com/urfave/cli"
)

var StartDescription = `Usage:

   buildkite-agent start [arguments...]

Description:

   When a job is ready to run it will call the "bootstrap-script"
   and pass it all the environment variables required for the job to run.
   This script is responsible for checking out the code, and running the
   actual build script defined in the pipeline.

   The agent will run any jobs within a PTY (pseudo terminal) if available.

Example:

   $ buildkite-agent start --token xxx`

type AgentStartConfig struct {
	Config                    string   `cli:"config"`
	Token                     string   `cli:"token" validate:"required"`
	Name                      string   `cli:"name"`
	Priority                  string   `cli:"priority"`
	DisconnectAfterJob        bool     `cli:"disconnect-after-job"`
	DisconnectAfterJobTimeout int      `cli:"disconnect-after-job-timeout"`
	BootstrapScript           string   `cli:"bootstrap-script" normalize:"filepath"`
	BuildPath                 string   `cli:"build-path" normalize:"filepath" validate:"required"`
	HooksPath                 string   `cli:"hooks-path" normalize:"filepath"`
	PluginsPath               string   `cli:"plugins-path" normalize:"filepath"`
	Shell                     string   `cli:"shell"`
	Tags                      []string `cli:"tags"`
	TagsFromEC2               bool     `cli:"tags-from-ec2"`
	TagsFromEC2Tags           bool     `cli:"tags-from-ec2-tags"`
	TagsFromGCP               bool     `cli:"tags-from-gcp"`
	WaitForEC2TagsTimeout     string   `cli:"wait-for-ec2-tags-timeout"`
	GitCloneFlags             string   `cli:"git-clone-flags"`
	GitCleanFlags             string   `cli:"git-clean-flags"`
	NoColor                   bool     `cli:"no-color"`
	NoSSHKeyscan              bool     `cli:"no-ssh-keyscan"`
	NoCommandEval             bool     `cli:"no-command-eval"`
	NoPlugins                 bool     `cli:"no-plugins"`
	NoPTY                     bool     `cli:"no-pty"`
	TimestampLines            bool     `cli:"timestamp-lines"`
	Endpoint                  string   `cli:"endpoint" validate:"required"`
	Debug                     bool     `cli:"debug"`
	DebugHTTP                 bool     `cli:"debug-http"`
	Experiments               []string `cli:"experiment"`

	/* Deprecated */
	NoSSHFingerprintVerification bool     `cli:"no-automatic-ssh-fingerprint-verification" deprecated-and-renamed-to:"NoSSHKeyscan"`
	MetaData                     []string `cli:"meta-data" deprecated-and-renamed-to:"Tags"`
	MetaDataEC2                  bool     `cli:"meta-data-ec2" deprecated-and-renamed-to:"TagsFromEC2"`
	MetaDataEC2Tags              bool     `cli:"meta-data-ec2-tags" deprecated-and-renamed-to:"TagsFromEC2Tags"`
	MetaDataGCP                  bool     `cli:"meta-data-gcp" deprecated-and-renamed-to:"TagsFromGCP"`
}

func DefaultShell() string {
	// https://github.com/golang/go/blob/master/src/go/build/syslist.go#L7
	switch runtime.GOOS {
	case "windows":
		return `C:\Windows\System32\CMD.exe /S /C`
	case "freebsd", "openbsd", "netbsd":
		return `/usr/local/bin/bash -e -c`
	default:
		return `/bin/bash -e -c`
	}
}

func DefaultConfigFilePaths() (paths []string) {
	// Toggle beetwen windows an *nix paths
	if runtime.GOOS == "windows" {
		paths = []string{
			"$USERPROFILE\\AppData\\Local\\BuildkiteAgent\\buildkite-agent.cfg",
		}
	} else {
		paths = []string{
			"$HOME/.buildkite-agent/buildkite-agent.cfg",
			"/usr/local/etc/buildkite-agent/buildkite-agent.cfg",
			"/etc/buildkite-agent/buildkite-agent.cfg",
		}
	}

	// Also check to see if there's a buildkite-agent.cfg in the folder
	// that the binary is running in.
	pathToBinary, err := filepath.Abs(filepath.Dir(os.Args[0]))
	if err == nil {
		pathToRelativeConfig := filepath.Join(pathToBinary, "buildkite-agent.cfg")
		paths = append([]string{pathToRelativeConfig}, paths...)
	}

	return
}

var AgentStartCommand = cli.Command{
	Name:        "start",
	Usage:       "Starts a Buildkite agent",
	Description: StartDescription,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:   "config",
			Value:  "",
			Usage:  "Path to a configuration file",
			EnvVar: "BUILDKITE_AGENT_CONFIG",
		},
		cli.StringFlag{
			Name:   "token",
			Value:  "",
			Usage:  "Your account agent token",
			EnvVar: "BUILDKITE_AGENT_TOKEN",
		},
		cli.StringFlag{
			Name:   "name",
			Value:  "",
			Usage:  "The name of the agent",
			EnvVar: "BUILDKITE_AGENT_NAME",
		},
		cli.StringFlag{
			Name:   "priority",
			Value:  "",
			Usage:  "The priority of the agent (higher priorities are assigned work first)",
			EnvVar: "BUILDKITE_AGENT_PRIORITY",
		},
		cli.BoolFlag{
			Name:   "disconnect-after-job",
			Usage:  "Disconnect the agent after running a job",
			EnvVar: "BUILDKITE_AGENT_DISCONNECT_AFTER_JOB",
		},
		cli.IntFlag{
			Name:   "disconnect-after-job-timeout",
			Value:  120,
			Usage:  "When --disconnect-after-job is specified, the number of seconds to wait for a job before shutting down",
			EnvVar: "BUILDKITE_AGENT_DISCONNECT_AFTER_JOB_TIMEOUT",
		},
		cli.StringFlag{
			Name:   "shell",
			Value:  DefaultShell(),
			Usage:  "The shell to use to interpret build commands",
			EnvVar: "BUILDKITE_SHELL",
		},
		cli.StringSliceFlag{
			Name:   "tags",
			Value:  &cli.StringSlice{},
			Usage:  "A comma-separated list of tags for the agent (e.g. \"linux\" or \"mac,xcode=8\")",
			EnvVar: "BUILDKITE_AGENT_TAGS",
		},
		cli.BoolFlag{
			Name:   "tags-from-ec2",
			Usage:  "Include the host's EC2 meta-data as tags (instance-id, instance-type, and ami-id)",
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_EC2",
		},
		cli.BoolFlag{
			Name:   "tags-from-ec2-tags",
			Usage:  "Include the host's EC2 tags as tags",
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_EC2_TAGS",
		},
		cli.BoolFlag{
			Name:   "tags-from-gcp",
			Usage:  "Include the host's Google Cloud meta-data as tags (instance-id, machine-type, preemptible, project-id, region, and zone)",
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_GCP",
		},
		cli.DurationFlag{
			Name:   "wait-for-ec2-tags-timeout",
			Usage:  "The amount of time to wait for tags from EC2 before proceeding",
			EnvVar: "BUILDKITE_AGENT_WAIT_FOR_EC2_TAGS_TIMEOUT",
			Value:  time.Second * 10,
		},
		cli.StringFlag{
			Name:   "git-clone-flags",
			Value:  "-v",
			Usage:  "Flags to pass to the \"git clone\" command",
			EnvVar: "BUILDKITE_GIT_CLONE_FLAGS",
		},
		cli.StringFlag{
			Name:   "git-clean-flags",
			Value:  "-fxdq",
			Usage:  "Flags to pass to \"git clean\" command",
			EnvVar: "BUILDKITE_GIT_CLEAN_FLAGS",
		},
		cli.StringFlag{
			Name:   "bootstrap-script",
			Value:  "",
			Usage:  "The command that is executed for bootstrapping a job, defaults to the bootstrap sub-command of this binary",
			EnvVar: "BUILDKITE_BOOTSTRAP_SCRIPT_PATH",
		},
		cli.StringFlag{
			Name:   "build-path",
			Value:  "",
			Usage:  "Path to where the builds will run from",
			EnvVar: "BUILDKITE_BUILD_PATH",
		},
		cli.StringFlag{
			Name:   "hooks-path",
			Value:  "",
			Usage:  "Directory where the hook scripts are found",
			EnvVar: "BUILDKITE_HOOKS_PATH",
		},
		cli.StringFlag{
			Name:   "plugins-path",
			Value:  "",
			Usage:  "Directory where the plugins are saved to",
			EnvVar: "BUILDKITE_PLUGINS_PATH",
		},
		cli.BoolFlag{
			Name:   "timestamp-lines",
			Usage:  "Prepend timestamps on each line of output.",
			EnvVar: "BUILDKITE_TIMESTAMP_LINES",
		},
		cli.BoolFlag{
			Name:   "no-pty",
			Usage:  "Do not run jobs within a pseudo terminal",
			EnvVar: "BUILDKITE_NO_PTY",
		},
		cli.BoolFlag{
			Name:   "no-ssh-keyscan",
			Usage:  "Don't automatically run ssh-keyscan before checkout",
			EnvVar: "BUILDKITE_NO_SSH_KEYSCAN",
		},
		cli.BoolFlag{
			Name:   "no-command-eval",
			Usage:  "Don't allow this agent to run arbitrary console commands",
			EnvVar: "BUILDKITE_NO_COMMAND_EVAL",
		},
		cli.BoolFlag{
			Name:   "no-plugins",
			Usage:  "Don't allow this agent to load plugins",
			EnvVar: "BUILDKITE_NO_PLUGINS",
		},
		ExperimentsFlag,
		EndpointFlag,
		NoColorFlag,
		DebugFlag,
		DebugHTTPFlag,
		/* Deprecated flags which will be removed in v4 */
		cli.StringSliceFlag{
			Name:   "meta-data",
			Value:  &cli.StringSlice{},
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_META_DATA",
		},
		cli.BoolFlag{
			Name:   "meta-data-ec2",
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_META_DATA_EC2",
		},
		cli.BoolFlag{
			Name:   "meta-data-ec2-tags",
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_TAGS_FROM_EC2_TAGS",
		},
		cli.BoolFlag{
			Name:   "meta-data-gcp",
			Hidden: true,
			EnvVar: "BUILDKITE_AGENT_META_DATA_GCP",
		},
		cli.BoolFlag{
			Name:   "no-automatic-ssh-fingerprint-verification",
			Hidden: true,
			EnvVar: "BUILDKITE_NO_AUTOMATIC_SSH_FINGERPRINT_VERIFICATION",
		},
	},
	Action: func(c *cli.Context) {
		// The configuration will be loaded into this struct
		cfg := AgentStartConfig{}

		// Setup the config loader. You'll see that we also path paths to
		// potential config files. The loader will use the first one it finds.
		loader := cliconfig.Loader{
			CLI:                    c,
			Config:                 &cfg,
			DefaultConfigFilePaths: DefaultConfigFilePaths(),
		}

		// Load the configuration
		if err := loader.Load(); err != nil {
			logger.Fatal("%s", err)
		}

		// Setup the any global configuration options
		HandleGlobalFlags(cfg)

		// Force some settings if on Windows (these aren't supported
		// yet)
		if runtime.GOOS == "windows" {
			cfg.NoPTY = true
		}

		// Set a useful default for the bootstrap script
		if cfg.BootstrapScript == "" {
			// "%q" on Windows doesn't work quite well for escaping
			// paths, since it'll escape `\` which is the Windows
			// path seperator (and we don't want that escaped).
			if runtime.GOOS == "windows" {
				cfg.BootstrapScript = fmt.Sprintf("\"%s\" bootstrap", os.Args[0])
			} else {
				cfg.BootstrapScript = fmt.Sprintf("%q bootstrap", os.Args[0])
			}
		}

		// Guess the shell if none is provided
		if cfg.Shell == "" {
			cfg.Shell = DefaultShell()
		}

		logger.Debug("Using shell %q", cfg.Shell)

		// Make sure the DisconnectAfterJobTimeout value is correct
		if cfg.DisconnectAfterJob && cfg.DisconnectAfterJobTimeout < 120 {
			logger.Fatal("The timeout for `disconnect-after-job` must be at least 120 seconds")
		}

		var ec2TagTimeout time.Duration
		if t := cfg.WaitForEC2TagsTimeout; t != "" {
			var err error
			ec2TagTimeout, err = time.ParseDuration(t)
			if err != nil {
				logger.Fatal("Failed to parse ec2 tag timeout: %v", err)
			}
		}

		// Setup the agent
		pool := agent.AgentPool{
			Token:                 cfg.Token,
			Name:                  cfg.Name,
			Priority:              cfg.Priority,
			Tags:                  cfg.Tags,
			TagsFromEC2:           cfg.TagsFromEC2,
			TagsFromEC2Tags:       cfg.TagsFromEC2Tags,
			TagsFromGCP:           cfg.TagsFromGCP,
			WaitForEC2TagsTimeout: ec2TagTimeout,
			Endpoint:              cfg.Endpoint,
			AgentConfiguration: &agent.AgentConfiguration{
				BootstrapScript:           cfg.BootstrapScript,
				BuildPath:                 cfg.BuildPath,
				HooksPath:                 cfg.HooksPath,
				PluginsPath:               cfg.PluginsPath,
				GitCloneFlags:             cfg.GitCloneFlags,
				GitCleanFlags:             cfg.GitCleanFlags,
				SSHKeyscan:                !cfg.NoSSHKeyscan,
				CommandEval:               !cfg.NoCommandEval,
				PluginsEnabled:            !cfg.NoPlugins,
				RunInPty:                  !cfg.NoPTY,
				TimestampLines:            cfg.TimestampLines,
				DisconnectAfterJob:        cfg.DisconnectAfterJob,
				DisconnectAfterJobTimeout: cfg.DisconnectAfterJobTimeout,
				Shell: cfg.Shell,
			},
		}

		// Store the loaded config file path on the pool so we can
		// show it when the agent starts
		if loader.File != nil {
			pool.ConfigFilePath = loader.File.Path
		}

		// Start the agent pool
		if err := pool.Start(); err != nil {
			logger.Fatal("%s", err)
		}
	},
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"

	"github.com/google/go-github/github"
)

type MirroredPR struct {
	DownstreamID int
	UpstreamID   int
}

type PRMirror struct {
	GitHubClient  *github.Client
	Context       *context.Context
	Configuration *Config
}

func (p PRMirror) Run() {
	events, _, err := p.GitHubClient.Activity.ListRepositoryEvents(*p.Context, p.Configuration.UpstreamOwner, p.Configuration.UpstreamRepo, nil)
	if _, ok := err.(*github.RateLimitError); ok {
		// TODO: Maybe add some context here
		log.Error("The github.com rate limit has been hit")
	}

	for _, event := range events {

		if *event.Type == "PullRequestEvent" {
			prEvent := github.PullRequestEvent{}

			err = json.Unmarshal(event.GetRawPayload(), &prEvent)
			if err != nil {
				panic(err)
			}

			prAction := prEvent.GetAction()

			log.Debugf("%s\n", prEvent.PullRequest.GetURL())

			if prAction == "opened" {
				//TODO: Check if we already have an open PR for this and add a comment saying upstream reopened it and remove the upsteam closed tag
				p.MirrorPR(&prEvent)
			} else if prAction == "closed" {

				//AddLabel("Upstream Closed")
			}
		}
	}
}

func (p PRMirror) MirrorPR(PREvent *github.PullRequestEvent) {
	log.Infof("Mirroring PR [%d]: %s from ", PREvent.PullRequest.GetNumber(), PREvent.PullRequest.GetTitle(), PREvent.PullRequest.User.GetLogin())
}

func (p PRMirror) AddLabels(id int, tags []string) bool {
	_, _, err := p.GitHubClient.Issues.AddLabelsToIssue(*p.Context, p.Configuration.UpstreamOwner, p.Configuration.UpstreamRepo, id, tags)
	if err != nil {
		log.Errorf("Error while adding a label to issue#:%d - %s", id, err.Error())
		return false
	}

	return true
}
<commit_msg>Refactor the PR event<commit_after>package main

import (
	"context"
	"encoding/json"

	"github.com/google/go-github/github"
)

type MirroredPR struct {
	DownstreamID int
	UpstreamID   int
}

type PRMirror struct {
	GitHubClient  *github.Client
	Context       *context.Context
	Configuration *Config
}

func (p PRMirror) HandlePREvent(prEvent *github.PullRequestEvent) {
	prAction := prEvent.GetAction()

	log.Debugf("%s\n", prEvent.PullRequest.GetURL())

	if prAction == "opened" {
		//TODO: Check if we already have an open PR for this and add a comment saying upstream reopened it and remove the upsteam closed tag
		p.MirrorPR(prEvent)
	} else if prAction == "closed" {

		//AddLabel("Upstream Closed")
	}
}

func (p PRMirror) Run() {
	events, _, err := p.GitHubClient.Activity.ListRepositoryEvents(*p.Context, p.Configuration.UpstreamOwner, p.Configuration.UpstreamRepo, nil)
	if _, ok := err.(*github.RateLimitError); ok {
		// TODO: Maybe add some context here
		log.Error("The github.com rate limit has been hit")
	}

	for _, event := range events {
		eventType := event.GetType()

		if eventType == "PullRequestEvent" {
			prEvent := github.PullRequestEvent{}
			err = json.Unmarshal(event.GetRawPayload(), &prEvent)
			if err != nil {
				panic(err)
			}

			p.HandlePREvent(&prEvent)
		}
	}
}

func (p PRMirror) MirrorPR(PREvent *github.PullRequestEvent) {
	log.Infof("Mirroring PR [%d]: %s from ", PREvent.PullRequest.GetNumber(), PREvent.PullRequest.GetTitle(), PREvent.PullRequest.User.GetLogin())
}

func (p PRMirror) AddLabels(id int, tags []string) bool {
	_, _, err := p.GitHubClient.Issues.AddLabelsToIssue(*p.Context, p.Configuration.UpstreamOwner, p.Configuration.UpstreamRepo, id, tags)
	if err != nil {
		log.Errorf("Error while adding a label to issue#:%d - %s", id, err.Error())
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package nsq

import (
	"errors"
	"io"
	"log"
	"sync"
	"time"
)

// ProducerConfig carries the different variables to tune a newly started
// producer.
type ProducerConfig struct {
	Address         string
	MaxConcurrency  int
	DialTimeout     time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration
	MaxRetryTimeout time.Duration
	MinRetryTimeout time.Duration
}

// Producer provide an abstraction around using direct connections to nsqd
// nodes to send messages.
type Producer struct {
	// Communication channels of the producer.
	reqs chan ProducerRequest
	done chan struct{}
	once sync.Once
	join sync.WaitGroup

	// Immutable state of the producer.
	address         string
	dialTimeout     time.Duration
	readTimeout     time.Duration
	writeTimeout    time.Duration
	maxRetryTimeout time.Duration
	minRetryTimeout time.Duration
}

// ProducerRequest are used to represent operations that are submitted to
// producers.
type ProducerRequest struct {
	Topic    string
	Message  []byte
	Response chan<- error
}

// StartProducer starts and returns a new producer p, configured with the
// variables from the config parameter, or returning an non-nil error if
// some of the configuration variables were invalid.
func StartProducer(config ProducerConfig) (p *Producer, err error) {
	if len(config.Address) == 0 {
		config.Address = "localhost:4151"
	}

	if config.MaxConcurrency == 0 {
		config.MaxConcurrency = DefaultMaxConcurrency
	}

	if config.DialTimeout == 0 {
		config.DialTimeout = DefaultDialTimeout
	}

	if config.ReadTimeout == 0 {
		config.ReadTimeout = DefaultReadTimeout
	}

	if config.WriteTimeout == 0 {
		config.WriteTimeout = DefaultWriteTimeout
	}

	if config.MaxRetryTimeout == 0 {
		config.MaxRetryTimeout = DefaultMaxRetryTimeout
	}

	if config.MinRetryTimeout == 0 {
		config.MinRetryTimeout = DefaultMinRetryTimeout
	}

	p = &Producer{
		reqs:            make(chan ProducerRequest, config.MaxConcurrency),
		done:            make(chan struct{}),
		address:         config.Address,
		dialTimeout:     config.DialTimeout,
		readTimeout:     config.ReadTimeout,
		writeTimeout:    config.WriteTimeout,
		maxRetryTimeout: config.MaxRetryTimeout,
		minRetryTimeout: config.MinRetryTimeout,
	}
	p.join.Add(config.MaxConcurrency)

	for i := 0; i != config.MaxConcurrency; i++ {
		go p.run()
	}

	return
}

// Stop gracefully shutsdown the producer, cancelling all inflight requests and
// waiting for all backend connections to be closed.
//
// It is safe to call the method multiple times and from multiple goroutines,
// they will all block until the producer has been completely shutdown.
func (p *Producer) Stop() {
	p.once.Do(p.stop)
	err := errors.New("publishing to a producer that was already stopped")

	for req := range p.reqs {
		req.complete(err)
	}

	p.join.Wait()
}

// Publish sends a message using the producer p, returning an error if it was
// already closed or if an error occurred while publishing the message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) Publish(topic string, message []byte) (err error) {
	defer func() {
		if recover() != nil {
			err = errors.New("publishing to a producer that was already stopped")
		}
	}()

	res := make(chan error, 1)

	p.reqs <- ProducerRequest{
		Topic:    topic,
		Message:  message,
		Response: res,
	}

	err = <-res
	return
}

// Requests returns a write-only channel that can be used to submit requests to p.
//
// This method is useful when the publish operation needs to be associated with
// other operations on channels in a select statement for example, or to publish
// in a non-blocking fashion.
func (p *Producer) Requests() chan<- ProducerRequest {
	return p.reqs
}

func (p *Producer) stop() {
	close(p.done)
	close(p.reqs)
}

func (p *Producer) run() {
	var conn *Conn
	var pipe chan ProducerRequest
	var ping chan struct{}
	var retry time.Duration

	defer p.join.Done()
	defer func() {
		if conn != nil {
			conn.Close()
		}
	}()

	for {
		select {
		case <-p.done:
			return

		case <-ping:
			if err := p.ping(conn); err != nil {
				conn.Close()
				conn = nil
				ping = nil
				pipe = nil
				continue
			}

		case req, ok := <-p.reqs:
			if !ok {
				return
			}

			if conn == nil {
				var err error

				if conn, err = DialTimeout(p.address, p.dialTimeout); err != nil {
					req.complete(err)

					log.Printf("failed to connect to %s, retrying after %s: %s", p.address, retry, err)
					retry = p.sleep(retry)
					continue
				}

				retry = 0
				pipe = make(chan ProducerRequest)
				ping = make(chan struct{})
				go p.flush(conn, pipe, ping)
			}

			if err := p.publish(conn, req.Topic, req.Message); err != nil {
				conn.Close()
				conn = nil
				ping = nil
				pipe = nil
				req.complete(err)
				continue
			}

			pipe <- req
		}
	}
}

func (p *Producer) flush(conn *Conn, pipe <-chan ProducerRequest, ping chan<- struct{}) {
	defer conn.Close()

	for {
		var frame Frame
		var err error

		if frame, err = conn.ReadFrame(); err != nil {
			if err != io.EOF && err != io.ErrUnexpectedEOF {
				log.Print(err)
			}
			return
		}

		switch f := frame.(type) {
		case Response:
			switch f {
			case OK:
				req := <-pipe
				req.complete(nil)
				continue

			case Heartbeat:
				ping <- struct{}{}
				continue

			case CloseWait:
				continue
			}

		case Error:
			req := <-pipe
			req.complete(f)
			log.Printf("closing connection after receiving an error from %s: %s", conn.RemoteAddr(), f)
			return

		case Message:
			log.Printf("closing connection after receiving an unexpected message from %s: %s", conn.RemoteAddr(), f.FrameType())
			return

		default:
			log.Printf("closing connection after receiving an unsupported frame from %s: %s", conn.RemoteAddr(), f.FrameType())
			return
		}
	}
}

func (p *Producer) publish(conn *Conn, topic string, message []byte) error {
	return p.write(conn, Pub{Topic: topic, Message: message})
}

func (p *Producer) ping(conn *Conn) error {
	return p.write(conn, Nop{})
}

func (p *Producer) write(conn *Conn, cmd Command) (err error) {
	if err = conn.SetDeadline(time.Now().Add(p.writeTimeout)); err == nil {
		err = conn.WriteCommand(cmd)
	}
	return
}

func (p *Producer) sleep(d time.Duration) time.Duration {
	if d < p.minRetryTimeout {
		d = p.minRetryTimeout
	}

	t := time.NewTimer(d)
	defer t.Stop()

	select {
	case <-t.C:
	case <-p.done:
	}

	if d *= 2; d > p.maxRetryTimeout {
		d = p.maxRetryTimeout
	}

	return d
}

func (r ProducerRequest) complete(err error) {
	if r.Response != nil {
		r.Response <- err
	}
}
<commit_msg>fix producer shutdown logic<commit_after>package nsq

import (
	"io"
	"log"
	"sync"
	"time"

	"github.com/pkg/errors"
)

// ProducerConfig carries the different variables to tune a newly started
// producer.
type ProducerConfig struct {
	Address         string
	MaxConcurrency  int
	DialTimeout     time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration
	MaxRetryTimeout time.Duration
	MinRetryTimeout time.Duration
}

// Producer provide an abstraction around using direct connections to nsqd
// nodes to send messages.
type Producer struct {
	// Communication channels of the producer.
	reqs chan ProducerRequest
	done chan struct{}
	once sync.Once
	join sync.WaitGroup

	// Immutable state of the producer.
	address         string
	dialTimeout     time.Duration
	readTimeout     time.Duration
	writeTimeout    time.Duration
	maxRetryTimeout time.Duration
	minRetryTimeout time.Duration
}

// ProducerRequest are used to represent operations that are submitted to
// producers.
type ProducerRequest struct {
	Topic    string
	Message  []byte
	Response chan<- error
}

// StartProducer starts and returns a new producer p, configured with the
// variables from the config parameter, or returning an non-nil error if
// some of the configuration variables were invalid.
func StartProducer(config ProducerConfig) (p *Producer, err error) {
	if len(config.Address) == 0 {
		config.Address = "localhost:4151"
	}

	if config.MaxConcurrency == 0 {
		config.MaxConcurrency = DefaultMaxConcurrency
	}

	if config.DialTimeout == 0 {
		config.DialTimeout = DefaultDialTimeout
	}

	if config.ReadTimeout == 0 {
		config.ReadTimeout = DefaultReadTimeout
	}

	if config.WriteTimeout == 0 {
		config.WriteTimeout = DefaultWriteTimeout
	}

	if config.MaxRetryTimeout == 0 {
		config.MaxRetryTimeout = DefaultMaxRetryTimeout
	}

	if config.MinRetryTimeout == 0 {
		config.MinRetryTimeout = DefaultMinRetryTimeout
	}

	p = &Producer{
		reqs:            make(chan ProducerRequest, config.MaxConcurrency),
		done:            make(chan struct{}),
		address:         config.Address,
		dialTimeout:     config.DialTimeout,
		readTimeout:     config.ReadTimeout,
		writeTimeout:    config.WriteTimeout,
		maxRetryTimeout: config.MaxRetryTimeout,
		minRetryTimeout: config.MinRetryTimeout,
	}
	p.join.Add(config.MaxConcurrency)

	for i := 0; i != config.MaxConcurrency; i++ {
		go p.run()
	}

	return
}

// Stop gracefully shutsdown the producer, cancelling all inflight requests and
// waiting for all backend connections to be closed.
//
// It is safe to call the method multiple times and from multiple goroutines,
// they will all block until the producer has been completely shutdown.
func (p *Producer) Stop() {
	p.once.Do(p.stop)
	err := errors.New("publishing to a producer that was already stopped")

	for req := range p.reqs {
		req.complete(err)
	}

	p.join.Wait()
}

// Publish sends a message using the producer p, returning an error if it was
// already closed or if an error occurred while publishing the message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) Publish(topic string, message []byte) (err error) {
	defer func() {
		if recover() != nil {
			err = errors.New("publishing to a producer that was already stopped")
		}
	}()

	res := make(chan error, 1)

	p.reqs <- ProducerRequest{
		Topic:    topic,
		Message:  message,
		Response: res,
	}

	err = <-res
	return
}

// Requests returns a write-only channel that can be used to submit requests to p.
//
// This method is useful when the publish operation needs to be associated with
// other operations on channels in a select statement for example, or to publish
// in a non-blocking fashion.
func (p *Producer) Requests() chan<- ProducerRequest {
	return p.reqs
}

func (p *Producer) stop() {
	close(p.done)
	close(p.reqs)
}

func (p *Producer) run() {
	var conn *Conn
	var pipe chan ProducerRequest
	var ping chan struct{}
	var retry time.Duration

	defer p.join.Done()
	defer func() {
		if conn != nil {
			conn.Close()
		}
		if pipe != nil {
			close(pipe)
		}
	}()

	for {
		select {
		case <-p.done:
			return

		case <-ping:
			if err := p.ping(conn); err != nil {
				close(pipe)
				close(ping)
				conn.Close()
				conn = nil
				ping = nil
				pipe = nil
				continue
			}

		case req, ok := <-p.reqs:
			if !ok {
				return
			}

			if conn == nil {
				var err error

				if conn, err = DialTimeout(p.address, p.dialTimeout); err != nil {
					req.complete(err)

					log.Printf("failed to connect to %s, retrying after %s: %s", p.address, retry, err)
					retry = p.sleep(retry)
					continue
				}

				retry = 0
				pipe = make(chan ProducerRequest)
				ping = make(chan struct{})
				go p.flush(conn, pipe, ping)
			}

			if err := p.publish(conn, req.Topic, req.Message); err != nil {
				close(pipe)
				close(ping)
				conn.Close()
				conn = nil
				ping = nil
				pipe = nil
				req.complete(err)
				continue
			}

			pipe <- req
		}
	}
}

func (p *Producer) flush(conn *Conn, pipe <-chan ProducerRequest, ping chan<- struct{}) {
	var err error

	defer conn.Close()
	defer func() { recover() }() // may happen when the ping channel is closed
	defer func() {
		if err == nil {
			err = io.ErrUnexpectedEOF
		}
		for req := range pipe {
			req.Response <- err
		}
	}()

	for {
		var frame Frame

		if frame, err = conn.ReadFrame(); err != nil {
			return
		}

		switch f := frame.(type) {
		case Response:
			switch f {
			case OK:
				req := <-pipe
				req.complete(nil)
				continue

			case Heartbeat:
				ping <- struct{}{}
				continue

			case CloseWait:
				continue
			}

		case Error:
			err = errors.Errorf("closing connection after receiving an error from %s: %s", conn.RemoteAddr(), f)
			return

		case Message:
			err = errors.Errorf("closing connection after receiving an unexpected message from %s: %s", conn.RemoteAddr(), f.FrameType())
			return

		default:
			err = errors.Errorf("closing connection after receiving an unsupported frame from %s: %s", conn.RemoteAddr(), f.FrameType())
			return
		}
	}
}

func (p *Producer) publish(conn *Conn, topic string, message []byte) error {
	return p.write(conn, Pub{Topic: topic, Message: message})
}

func (p *Producer) ping(conn *Conn) error {
	return p.write(conn, Nop{})
}

func (p *Producer) write(conn *Conn, cmd Command) (err error) {
	if err = conn.SetDeadline(time.Now().Add(p.writeTimeout)); err == nil {
		err = conn.WriteCommand(cmd)
	}
	return
}

func (p *Producer) sleep(d time.Duration) time.Duration {
	if d < p.minRetryTimeout {
		d = p.minRetryTimeout
	}

	t := time.NewTimer(d)
	defer t.Stop()

	select {
	case <-t.C:
	case <-p.done:
	}

	if d *= 2; d > p.maxRetryTimeout {
		d = p.maxRetryTimeout
	}

	return d
}

func (r ProducerRequest) complete(err error) {
	if r.Response != nil {
		r.Response <- err
	}
}
<|endoftext|>


<commit_before>package nsq

import (
	"log"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

// ProducerConfig carries the different variables to tune a newly started
// producer.
type ProducerConfig struct {
	Address         string
	Topic           string
	MaxConcurrency  int
	ForceConnect    bool
	DialTimeout     time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration
	MaxRetryTimeout time.Duration
	MinRetryTimeout time.Duration
}

// Producer provide an abstraction around using direct connections to nsqd
// nodes to send messages.
type Producer struct {
	// Communication channels of the producer.
	reqs chan ProducerRequest
	done chan struct{}
	once sync.Once
	join sync.WaitGroup
	ok   uint32

	// Immutable state of the producer.
	address         string
	topic           string
	forceConnect    bool
	dialTimeout     time.Duration
	readTimeout     time.Duration
	writeTimeout    time.Duration
	maxRetryTimeout time.Duration
	minRetryTimeout time.Duration
}

// ProducerRequest are used to represent operations that are submitted to
// producers.
type ProducerRequest struct {
	Topic    string
	Message  []byte
	Response chan<- error
	Deadline time.Time
}

// StartProducer starts and returns a new producer p, configured with the
// variables from the config parameter, or returning an non-nil error if
// some of the configuration variables were invalid.
func StartProducer(config ProducerConfig) (p *Producer, err error) {
	if len(config.Address) == 0 {
		config.Address = "localhost:4151"
	}

	if config.MaxConcurrency == 0 {
		config.MaxConcurrency = DefaultMaxConcurrency
	}

	if config.DialTimeout == 0 {
		config.DialTimeout = DefaultDialTimeout
	}

	if config.ReadTimeout == 0 {
		config.ReadTimeout = DefaultReadTimeout
	}

	if config.WriteTimeout == 0 {
		config.WriteTimeout = DefaultWriteTimeout
	}

	if config.MaxRetryTimeout == 0 {
		config.MaxRetryTimeout = DefaultMaxRetryTimeout
	}

	if config.MinRetryTimeout == 0 {
		config.MinRetryTimeout = DefaultMinRetryTimeout
	}

	p = &Producer{
		reqs:            make(chan ProducerRequest, config.MaxConcurrency),
		done:            make(chan struct{}),
		address:         config.Address,
		topic:           config.Topic,
		forceConnect:    config.ForceConnect,
		dialTimeout:     config.DialTimeout,
		readTimeout:     config.ReadTimeout,
		writeTimeout:    config.WriteTimeout,
		maxRetryTimeout: config.MaxRetryTimeout,
		minRetryTimeout: config.MinRetryTimeout,
	}
	p.join.Add(config.MaxConcurrency)

	for i := 0; i != config.MaxConcurrency; i++ {
		go p.run()
	}

	return
}

// Stop gracefully shutsdown the producer, cancelling all inflight requests and
// waiting for all backend connections to be closed.
//
// It is safe to call the method multiple times and from multiple goroutines,
// they will all block until the producer has been completely shutdown.
func (p *Producer) Stop() {
	p.once.Do(p.stop)
	err := errors.New("publishing to a producer that was already stopped")

	for req := range p.reqs {
		req.complete(err)
	}

	p.join.Wait()
}

// Publish sends a message using the producer p, returning an error if it was
// already closed or if an error occurred while publishing the message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) Publish(message []byte) (err error) {
	return p.PublishTo(p.topic, message)
}

// PublishTo sends a message to a specific topic using the producer p, returning
// an error if it was already closed or if an error occurred while publishing the
// message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) PublishTo(topic string, message []byte) (err error) {
	defer func() {
		if recover() != nil {
			err = errors.New("publishing to a producer that was already stopped")
		}
	}()

	if len(topic) == 0 {
		return errors.New("no topic set for publishing message")
	}

	response := make(chan error, 1)
	deadline := time.Now().Add(p.dialTimeout + p.readTimeout + p.writeTimeout)

	// Attempts to queue the request so one of the active connections can pick
	// it up.
	p.reqs <- ProducerRequest{
		Topic:    topic,
		Message:  message,
		Response: response,
		Deadline: deadline,
	}

	// This will always trigger, either if the connection was lost or if a
	// response was successfully sent.
	err = <-response
	return
}

// Requests returns a write-only channel that can be used to submit requests to p.
//
// This method is useful when the publish operation needs to be associated with
// other operations on channels in a select statement for example, or to publish
// in a non-blocking fashion.
func (p *Producer) Requests() chan<- ProducerRequest {
	return p.reqs
}

// Connected returns true if the producer has successfully established a
// connection to nsqd, false otherwise.
func (p *Producer) Connected() bool {
	return atomic.LoadUint32(&p.ok) != 0
}

func (p *Producer) stop() {
	close(p.done)
	close(p.reqs)
}

func (p *Producer) run() {
	var conn *Conn
	var resChan chan Frame
	var pending []ProducerRequest
	var retry time.Duration

	shutdown := func(err error) {
		atomic.StoreUint32(&p.ok, 0)

		if conn != nil {
			close(resChan)
			conn.Close()
			conn = nil
			resChan = nil
			pending = completeAllProducerRequests(pending, err)
		}
	}

	connect := func() (err error) {
		if conn, err = DialTimeout(p.address, p.dialTimeout); err != nil {
			log.Printf("failed to connect to %s, retrying after %s: %s", p.address, retry, err)
			retry = p.sleep(retry)
			return
		}

		retry = 0
		resChan = make(chan Frame, 16)
		go p.flush(conn, resChan)

		atomic.StoreUint32(&p.ok, 1)
		return
	}

	defer p.join.Done()
	defer shutdown(nil)

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-p.done:
			return

		case now := <-ticker.C:
			if p.forceConnect && conn == nil {
				connect()
			}

			if producerRequestsTimedOut(pending, now) {
				shutdown(errors.New("timeout"))
				continue
			}

		case req, ok := <-p.reqs:
			if !ok {
				return
			}

			if conn == nil {
				if err := connect(); err != nil {
					req.complete(err)
					continue
				}
			}

			if err := p.publish(conn, req.Topic, req.Message); err != nil {
				req.complete(err)
				shutdown(err)
				continue
			}

			pending = append(pending, req)

		case frame := <-resChan:
			switch f := frame.(type) {
			case Response:
				switch f {
				case OK:
					pending = completeNextProducerRequest(pending, nil)

				case Heartbeat:
					if err := p.ping(conn); err != nil {
						shutdown(err)
						continue
					}

				case CloseWait:
					return

				default:
					shutdown(errors.Errorf("closing connection after receiving an unexpected response from %s: %s", conn.RemoteAddr(), f))
					continue
				}

			case Error:
				shutdown(errors.Errorf("closing connection after receiving an error from %s: %s", conn.RemoteAddr(), f))
				continue

			case Message:
				shutdown(errors.Errorf("closing connection after receiving an unexpected message from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue

			default:
				shutdown(errors.Errorf("closing connection after receiving an unsupported frame from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue
			}
		}
	}
}

func (p *Producer) flush(conn *Conn, resChan chan<- Frame) {
	defer func() { recover() }() // may happen if the resChan is closed
	defer conn.Close()

	for {
		frame, err := conn.ReadFrame()

		if err != nil {
			resChan <- Error(err.Error())
			return
		}

		resChan <- frame
	}
}

func (p *Producer) write(conn *Conn, cmd Command) (err error) {
	if err = conn.SetDeadline(time.Now().Add(p.writeTimeout)); err == nil {
		err = conn.WriteCommand(cmd)
	}
	return
}

func (p *Producer) publish(conn *Conn, topic string, message []byte) error {
	return p.write(conn, Pub{
		Topic:   topic,
		Message: message,
	})
}

func (p *Producer) ping(conn *Conn) error {
	return p.write(conn, Nop{})
}

func (p *Producer) sleep(d time.Duration) time.Duration {
	if d < p.minRetryTimeout {
		d = p.minRetryTimeout
	}

	t := time.NewTimer(d)
	defer t.Stop()

	select {
	case <-t.C:
	case <-p.done:
	}

	if d *= 2; d > p.maxRetryTimeout {
		d = p.maxRetryTimeout
	}

	return d
}

func (r ProducerRequest) complete(err error) {
	if r.Response != nil {
		r.Response <- err
	}
}

func completeNextProducerRequest(reqs []ProducerRequest, err error) []ProducerRequest {
	reqs[0].complete(err)
	return reqs[1:]
}

func completeAllProducerRequests(reqs []ProducerRequest, err error) []ProducerRequest {
	for _, req := range reqs {
		req.complete(err)
	}
	return nil
}

func producerRequestsTimedOut(reqs []ProducerRequest, now time.Time) bool {
	return len(reqs) != 0 && now.After(reqs[0].Deadline)
}
<commit_msg>always try to connect when ForceConnect is set<commit_after>package nsq

import (
	"log"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

// ProducerConfig carries the different variables to tune a newly started
// producer.
type ProducerConfig struct {
	Address         string
	Topic           string
	MaxConcurrency  int
	ForceConnect    bool
	DialTimeout     time.Duration
	ReadTimeout     time.Duration
	WriteTimeout    time.Duration
	MaxRetryTimeout time.Duration
	MinRetryTimeout time.Duration
}

// Producer provide an abstraction around using direct connections to nsqd
// nodes to send messages.
type Producer struct {
	// Communication channels of the producer.
	reqs chan ProducerRequest
	done chan struct{}
	once sync.Once
	join sync.WaitGroup
	ok   uint32

	// Immutable state of the producer.
	address         string
	topic           string
	forceConnect    bool
	dialTimeout     time.Duration
	readTimeout     time.Duration
	writeTimeout    time.Duration
	maxRetryTimeout time.Duration
	minRetryTimeout time.Duration
}

// ProducerRequest are used to represent operations that are submitted to
// producers.
type ProducerRequest struct {
	Topic    string
	Message  []byte
	Response chan<- error
	Deadline time.Time
}

// StartProducer starts and returns a new producer p, configured with the
// variables from the config parameter, or returning an non-nil error if
// some of the configuration variables were invalid.
func StartProducer(config ProducerConfig) (p *Producer, err error) {
	if len(config.Address) == 0 {
		config.Address = "localhost:4151"
	}

	if config.MaxConcurrency == 0 {
		config.MaxConcurrency = DefaultMaxConcurrency
	}

	if config.DialTimeout == 0 {
		config.DialTimeout = DefaultDialTimeout
	}

	if config.ReadTimeout == 0 {
		config.ReadTimeout = DefaultReadTimeout
	}

	if config.WriteTimeout == 0 {
		config.WriteTimeout = DefaultWriteTimeout
	}

	if config.MaxRetryTimeout == 0 {
		config.MaxRetryTimeout = DefaultMaxRetryTimeout
	}

	if config.MinRetryTimeout == 0 {
		config.MinRetryTimeout = DefaultMinRetryTimeout
	}

	p = &Producer{
		reqs:            make(chan ProducerRequest, config.MaxConcurrency),
		done:            make(chan struct{}),
		address:         config.Address,
		topic:           config.Topic,
		forceConnect:    config.ForceConnect,
		dialTimeout:     config.DialTimeout,
		readTimeout:     config.ReadTimeout,
		writeTimeout:    config.WriteTimeout,
		maxRetryTimeout: config.MaxRetryTimeout,
		minRetryTimeout: config.MinRetryTimeout,
	}
	p.join.Add(config.MaxConcurrency)

	for i := 0; i != config.MaxConcurrency; i++ {
		go p.run()
	}

	return
}

// Stop gracefully shutsdown the producer, cancelling all inflight requests and
// waiting for all backend connections to be closed.
//
// It is safe to call the method multiple times and from multiple goroutines,
// they will all block until the producer has been completely shutdown.
func (p *Producer) Stop() {
	p.once.Do(p.stop)
	err := errors.New("publishing to a producer that was already stopped")

	for req := range p.reqs {
		req.complete(err)
	}

	p.join.Wait()
}

// Publish sends a message using the producer p, returning an error if it was
// already closed or if an error occurred while publishing the message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) Publish(message []byte) (err error) {
	return p.PublishTo(p.topic, message)
}

// PublishTo sends a message to a specific topic using the producer p, returning
// an error if it was already closed or if an error occurred while publishing the
// message.
//
// Note that no retry is done internally, the producer will fail after the
// first unsuccessful attempt to publish the message. It is the responsibility
// of the caller to retry if necessary.
func (p *Producer) PublishTo(topic string, message []byte) (err error) {
	defer func() {
		if recover() != nil {
			err = errors.New("publishing to a producer that was already stopped")
		}
	}()

	if len(topic) == 0 {
		return errors.New("no topic set for publishing message")
	}

	response := make(chan error, 1)
	deadline := time.Now().Add(p.dialTimeout + p.readTimeout + p.writeTimeout)

	// Attempts to queue the request so one of the active connections can pick
	// it up.
	p.reqs <- ProducerRequest{
		Topic:    topic,
		Message:  message,
		Response: response,
		Deadline: deadline,
	}

	// This will always trigger, either if the connection was lost or if a
	// response was successfully sent.
	err = <-response
	return
}

// Requests returns a write-only channel that can be used to submit requests to p.
//
// This method is useful when the publish operation needs to be associated with
// other operations on channels in a select statement for example, or to publish
// in a non-blocking fashion.
func (p *Producer) Requests() chan<- ProducerRequest {
	return p.reqs
}

// Connected returns true if the producer has successfully established a
// connection to nsqd, false otherwise.
func (p *Producer) Connected() bool {
	return atomic.LoadUint32(&p.ok) != 0
}

func (p *Producer) stop() {
	close(p.done)
	close(p.reqs)
}

func (p *Producer) run() {
	var conn *Conn
	var resChan chan Frame
	var pending []ProducerRequest
	var retry time.Duration

	shutdown := func(err error) {
		atomic.StoreUint32(&p.ok, 0)

		if conn != nil {
			close(resChan)
			conn.Close()
			conn = nil
			resChan = nil
			pending = completeAllProducerRequests(pending, err)
		}
	}

	connect := func() (err error) {
		if conn, err = DialTimeout(p.address, p.dialTimeout); err != nil {
			log.Printf("failed to connect to %s, retrying after %s: %s", p.address, retry, err)
			retry = p.sleep(retry)
			return
		}

		retry = 0
		resChan = make(chan Frame, 16)
		go p.flush(conn, resChan)

		atomic.StoreUint32(&p.ok, 1)
		return
	}

	defer p.join.Done()
	defer shutdown(nil)

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	if p.forceConnect {
		connect()
	}

	for {
		select {
		case <-p.done:
			return

		case now := <-ticker.C:
			if p.forceConnect && conn == nil {
				connect()
			}

			if producerRequestsTimedOut(pending, now) {
				shutdown(errors.New("timeout"))
				continue
			}

		case req, ok := <-p.reqs:
			if !ok {
				return
			}

			if conn == nil {
				if err := connect(); err != nil {
					req.complete(err)
					continue
				}
			}

			if err := p.publish(conn, req.Topic, req.Message); err != nil {
				req.complete(err)
				shutdown(err)
				continue
			}

			pending = append(pending, req)

		case frame := <-resChan:
			switch f := frame.(type) {
			case Response:
				switch f {
				case OK:
					pending = completeNextProducerRequest(pending, nil)

				case Heartbeat:
					if err := p.ping(conn); err != nil {
						shutdown(err)
						continue
					}

				case CloseWait:
					return

				default:
					shutdown(errors.Errorf("closing connection after receiving an unexpected response from %s: %s", conn.RemoteAddr(), f))
					continue
				}

			case Error:
				shutdown(errors.Errorf("closing connection after receiving an error from %s: %s", conn.RemoteAddr(), f))
				continue

			case Message:
				shutdown(errors.Errorf("closing connection after receiving an unexpected message from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue

			default:
				shutdown(errors.Errorf("closing connection after receiving an unsupported frame from %s: %s", conn.RemoteAddr(), f.FrameType()))
				continue
			}
		}
	}
}

func (p *Producer) flush(conn *Conn, resChan chan<- Frame) {
	defer func() { recover() }() // may happen if the resChan is closed
	defer conn.Close()

	for {
		frame, err := conn.ReadFrame()

		if err != nil {
			resChan <- Error(err.Error())
			return
		}

		resChan <- frame
	}
}

func (p *Producer) write(conn *Conn, cmd Command) (err error) {
	if err = conn.SetDeadline(time.Now().Add(p.writeTimeout)); err == nil {
		err = conn.WriteCommand(cmd)
	}
	return
}

func (p *Producer) publish(conn *Conn, topic string, message []byte) error {
	return p.write(conn, Pub{
		Topic:   topic,
		Message: message,
	})
}

func (p *Producer) ping(conn *Conn) error {
	return p.write(conn, Nop{})
}

func (p *Producer) sleep(d time.Duration) time.Duration {
	if d < p.minRetryTimeout {
		d = p.minRetryTimeout
	}

	t := time.NewTimer(d)
	defer t.Stop()

	select {
	case <-t.C:
	case <-p.done:
	}

	if d *= 2; d > p.maxRetryTimeout {
		d = p.maxRetryTimeout
	}

	return d
}

func (r ProducerRequest) complete(err error) {
	if r.Response != nil {
		r.Response <- err
	}
}

func completeNextProducerRequest(reqs []ProducerRequest, err error) []ProducerRequest {
	reqs[0].complete(err)
	return reqs[1:]
}

func completeAllProducerRequests(reqs []ProducerRequest, err error) []ProducerRequest {
	for _, req := range reqs {
		req.complete(err)
	}
	return nil
}

func producerRequestsTimedOut(reqs []ProducerRequest, now time.Time) bool {
	return len(reqs) != 0 && now.After(reqs[0].Deadline)
}
<|endoftext|>