<commit_before>/*
Package portping provides simple functions to ping TCP ports.
It also includes a simple command line interface.
 */
package portping

import (
	"net"
	"time"
)

// Ping connects to the specified host and port
// using net.DialTimeout and network "tcp".
func Ping(host, port string) error {
	addr := net.JoinHostPort(host, port)
	a, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	conn, err := net.DialTimeout(a.Network(), a.String(), 10*time.Second)
	if conn != nil {
		defer conn.Close()
	}
	if err != nil {
		return err
	}
	return nil
}

// PingN calls Ping the specified number of times,
// and sends the results to the given channel.
func PingN(host, port string, count int, c chan error) {
	for i := 0; i < count; i++ {
		c <- Ping(host, port)
	}
}

// FormatResult converts the result returned by Ping to string.
func FormatResult(err error) string {
	if err == nil {
		return "success"
	}
	switch err := err.(type) {
	case *net.OpError:
		return err.Err.Error()
	default:
		return err.Error()
	}
}
<commit_msg>dropped unnecessary resolve tcp addr<commit_after>/*
Package portping provides simple functions to ping TCP ports.
It also includes a simple command line interface.
 */
package portping

import (
	"net"
	"time"
)

// Ping connects to the specified host and port
// using net.DialTimeout and network "tcp".
func Ping(host, port string) error {
	addr := net.JoinHostPort(host, port)
	conn, err := net.DialTimeout("tcp", addr, 10*time.Second)
	if conn != nil {
		defer conn.Close()
	}
	if err != nil {
		return err
	}
	return nil
}

// PingN calls Ping the specified number of times,
// and sends the results to the given channel.
func PingN(host, port string, count int, c chan error) {
	for i := 0; i < count; i++ {
		c <- Ping(host, port)
	}
}

// FormatResult converts the result returned by Ping to string.
func FormatResult(err error) string {
	if err == nil {
		return "success"
	}
	switch err := err.(type) {
	case *net.OpError:
		return err.Err.Error()
	default:
		return err.Error()
	}
}
<|endoftext|>


<commit_before>package mastodon

import (
	"context"
	"fmt"

	"github.com/jmoiron/sqlx"
)

// InitializePostgres initializes the connections to the Postgres host.
func InitializePostgres(ctx context.Context, config *Config) (*sqlx.DB, error) {
	dataSourceName := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=%s connect_timeout=%d",
		config.Host,
		config.Port,
		config.User,
		config.Password,
		config.Database,
		config.SSLMode,
		config.ConnectTimeout,
	)

	// Make the connection.
	connection, err := sqlx.ConnectContext(ctx, "postgres", dataSourceName)
	if err != nil {
		return nil, err
	}

	// Set the connection limits.
	connection.SetMaxIdleConns(config.MaxIdleConnections)
	connection.SetMaxOpenConns(config.MaxConnections)

	return connection, nil
}
<commit_msg>Added an initialization function to allow for connections, with context, to Postgres.<commit_after>package mastodon

import (
	"context"
	"fmt"

	"github.com/jmoiron/sqlx"
)

// InitializePostgres initializes the connections to the Postgres host.
func InitializePostgres(config *Config) (*sqlx.DB, error) {
	dataSourceName := generateDataSourceName(config)

	// Make the connection.
	connection, err := sqlx.Connect("postgres", dataSourceName)
	if err != nil {
		return nil, err
	}

	// Set the connection limits.
	connection.SetMaxIdleConns(config.MaxIdleConnections)
	connection.SetMaxOpenConns(config.MaxConnections)

	return connection, nil
}

// InitializePostgresWithContext initializes the connections to the Postgres host -- with context.
func InitializePostgresWithContext(ctx context.Context, config *Config) (*sqlx.DB, error) {
	dataSourceName := generateDataSourceName(config)

	// Make the connection.
	connection, err := sqlx.ConnectContext(ctx, "postgres", dataSourceName)
	if err != nil {
		return nil, err
	}

	// Set the connection limits.
	connection.SetMaxIdleConns(config.MaxIdleConnections)
	connection.SetMaxOpenConns(config.MaxConnections)

	return connection, nil
}

func generateDataSourceName(config *Config) string {
	return fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=%s connect_timeout=%d",
		config.Host,
		config.Port,
		config.User,
		config.Password,
		config.Database,
		config.SSLMode,
		config.ConnectTimeout,
	)
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"regexp"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sns"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

// The preferences are account-wide, so the tests must be serialized
func TestAccAWSSNSSMSPreferences(t *testing.T) {
	testCases := map[string]func(t *testing.T){
		"almostAll":      testAccAWSSNSSMSPreferences_almostAll,
		"defaultSMSType": testAccAWSSNSSMSPreferences_defaultSMSType,
		"deliveryRole":   testAccAWSSNSSMSPreferences_deliveryRole,
		"empty":          testAccAWSSNSSMSPreferences_empty,
	}

	for name, tc := range testCases {
		tc := tc
		t.Run(name, func(t *testing.T) {
			tc(t)
		})
	}
}

func testAccAWSSNSSMSPreferences_empty(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_empty,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "monthly_spend_limit"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_iam_role_arn"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_success_sampling_rate"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "default_sender_id"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "default_sms_type"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "usage_report_s3_bucket"),
				),
			},
		},
	})
}

func testAccAWSSNSSMSPreferences_defaultSMSType(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_defSMSType,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "monthly_spend_limit"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_iam_role_arn"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_success_sampling_rate"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "default_sender_id"),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "default_sms_type", "Transactional"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "usage_report_s3_bucket"),
				),
			},
		},
	})
}

func testAccAWSSNSSMSPreferences_almostAll(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_almostAll,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "monthly_spend_limit", "1"),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "default_sms_type", "Transactional"),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "usage_report_s3_bucket", "some-bucket"),
				),
			},
		},
	})
}

func testAccAWSSNSSMSPreferences_deliveryRole(t *testing.T) {
	arnRole := regexp.MustCompile(`^arn:aws:iam::\d+:role/test_smsdelivery_role$`)
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_deliveryRole,
				Check: resource.ComposeTestCheckFunc(
					resource.TestMatchResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_iam_role_arn", arnRole),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_success_sampling_rate", "75"),
				),
			},
		},
	})
}

func testAccCheckAWSSNSSMSPrefsDestroy(s *terraform.State) error {
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_sns_sms_preferences" {
			continue
		}

		conn := testAccProvider.Meta().(*AWSClient).snsconn
		attrs, err := conn.GetSMSAttributes(&sns.GetSMSAttributesInput{})
		if err != nil {
			return fmt.Errorf("error getting SMS attributes: %s", err)
		}
		if attrs == nil || len(attrs.Attributes) == 0 {
			return nil
		}

		for attrName, attrValue := range attrs.Attributes {
			if aws.StringValue(attrValue) != "" {
				return fmt.Errorf("expected SMS attribute %q to be empty, but received: %s", attrName, aws.StringValue(attrValue))
			}
		}

		return nil
	}

	return nil
}

const testAccAWSSNSSMSPreferencesConfig_empty = `
resource "aws_sns_sms_preferences" "test_pref" {}
`
const testAccAWSSNSSMSPreferencesConfig_defSMSType = `
resource "aws_sns_sms_preferences" "test_pref" {
	default_sms_type = "Transactional"
}
`
const testAccAWSSNSSMSPreferencesConfig_almostAll = `
resource "aws_sns_sms_preferences" "test_pref" {
	monthly_spend_limit = "1",
	default_sms_type = "Transactional",
	usage_report_s3_bucket = "some-bucket",
}
`
const testAccAWSSNSSMSPreferencesConfig_deliveryRole = `
resource "aws_iam_role" "test_smsdelivery_role" {
    name = "test_smsdelivery_role"
    path = "/"
    assume_role_policy = <<POLICY
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "sns.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
POLICY
}

resource "aws_iam_role_policy" "test_smsdelivery_role_policy" {
  name   = "test_smsdelivery_role_policy"
  role   = "${aws_iam_role.test_smsdelivery_role.id}"
  policy = <<POLICY
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": ["logs:CreateLogGroup","logs:CreateLogStream","logs:PutLogEvents","logs:PutMetricFilter","logs:PutRetentionPolicy"],
      "Resource": "*",
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
POLICY
}

resource "aws_sns_sms_preferences" "test_pref" {
	delivery_status_iam_role_arn = "${aws_iam_role.test_smsdelivery_role.arn}",
	delivery_status_success_sampling_rate = "75",
}
`
<commit_msg>tests/resource/aws_sns_sms_preferences: Fix Terraform 0.12 syntax<commit_after>package aws

import (
	"fmt"
	"regexp"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/sns"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

// The preferences are account-wide, so the tests must be serialized
func TestAccAWSSNSSMSPreferences(t *testing.T) {
	testCases := map[string]func(t *testing.T){
		"almostAll":      testAccAWSSNSSMSPreferences_almostAll,
		"defaultSMSType": testAccAWSSNSSMSPreferences_defaultSMSType,
		"deliveryRole":   testAccAWSSNSSMSPreferences_deliveryRole,
		"empty":          testAccAWSSNSSMSPreferences_empty,
	}

	for name, tc := range testCases {
		tc := tc
		t.Run(name, func(t *testing.T) {
			tc(t)
		})
	}
}

func testAccAWSSNSSMSPreferences_empty(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_empty,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "monthly_spend_limit"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_iam_role_arn"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_success_sampling_rate"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "default_sender_id"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "default_sms_type"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "usage_report_s3_bucket"),
				),
			},
		},
	})
}

func testAccAWSSNSSMSPreferences_defaultSMSType(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_defSMSType,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "monthly_spend_limit"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_iam_role_arn"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_success_sampling_rate"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "default_sender_id"),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "default_sms_type", "Transactional"),
					resource.TestCheckNoResourceAttr("aws_sns_sms_preferences.test_pref", "usage_report_s3_bucket"),
				),
			},
		},
	})
}

func testAccAWSSNSSMSPreferences_almostAll(t *testing.T) {
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_almostAll,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "monthly_spend_limit", "1"),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "default_sms_type", "Transactional"),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "usage_report_s3_bucket", "some-bucket"),
				),
			},
		},
	})
}

func testAccAWSSNSSMSPreferences_deliveryRole(t *testing.T) {
	arnRole := regexp.MustCompile(`^arn:aws:iam::\d+:role/test_smsdelivery_role$`)
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSSNSSMSPrefsDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSSNSSMSPreferencesConfig_deliveryRole,
				Check: resource.ComposeTestCheckFunc(
					resource.TestMatchResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_iam_role_arn", arnRole),
					resource.TestCheckResourceAttr("aws_sns_sms_preferences.test_pref", "delivery_status_success_sampling_rate", "75"),
				),
			},
		},
	})
}

func testAccCheckAWSSNSSMSPrefsDestroy(s *terraform.State) error {
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_sns_sms_preferences" {
			continue
		}

		conn := testAccProvider.Meta().(*AWSClient).snsconn
		attrs, err := conn.GetSMSAttributes(&sns.GetSMSAttributesInput{})
		if err != nil {
			return fmt.Errorf("error getting SMS attributes: %s", err)
		}
		if attrs == nil || len(attrs.Attributes) == 0 {
			return nil
		}

		for attrName, attrValue := range attrs.Attributes {
			if aws.StringValue(attrValue) != "" {
				return fmt.Errorf("expected SMS attribute %q to be empty, but received: %s", attrName, aws.StringValue(attrValue))
			}
		}

		return nil
	}

	return nil
}

const testAccAWSSNSSMSPreferencesConfig_empty = `
resource "aws_sns_sms_preferences" "test_pref" {}
`
const testAccAWSSNSSMSPreferencesConfig_defSMSType = `
resource "aws_sns_sms_preferences" "test_pref" {
	default_sms_type = "Transactional"
}
`
const testAccAWSSNSSMSPreferencesConfig_almostAll = `
resource "aws_sns_sms_preferences" "test_pref" {
	monthly_spend_limit = "1"
	default_sms_type = "Transactional"
	usage_report_s3_bucket = "some-bucket"
}
`
const testAccAWSSNSSMSPreferencesConfig_deliveryRole = `
resource "aws_iam_role" "test_smsdelivery_role" {
    name = "test_smsdelivery_role"
    path = "/"
    assume_role_policy = <<POLICY
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": "sns.amazonaws.com"
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
POLICY
}

resource "aws_iam_role_policy" "test_smsdelivery_role_policy" {
  name   = "test_smsdelivery_role_policy"
  role   = "${aws_iam_role.test_smsdelivery_role.id}"
  policy = <<POLICY
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": ["logs:CreateLogGroup","logs:CreateLogStream","logs:PutLogEvents","logs:PutMetricFilter","logs:PutRetentionPolicy"],
      "Resource": "*",
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
POLICY
}

resource "aws_sns_sms_preferences" "test_pref" {
	delivery_status_iam_role_arn = "${aws_iam_role.test_smsdelivery_role.arn}"
	delivery_status_success_sampling_rate = "75"
}
`
<|endoftext|>


<commit_before>package blockchain

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"sync"

	wire "github.com/tendermint/go-wire"
	"github.com/tendermint/tendermint/types"
	. "github.com/tendermint/tmlibs/common"
	dbm "github.com/tendermint/tmlibs/db"
)

/*
Simple low level store for blocks.

There are three types of information stored:
 - BlockMeta:   Meta information about each block
 - Block part:  Parts of each block, aggregated w/ PartSet
 - Commit:      The commit part of each block, for gossiping precommit votes

Currently the precommit signatures are duplicated in the Block parts as
well as the Commit.  In the future this may change, perhaps by moving
the Commit data outside the Block.

Panics indicate probable corruption in the data
*/
type BlockStore struct {
	db dbm.DB

	mtx    sync.RWMutex
	height int
}

func NewBlockStore(db dbm.DB) *BlockStore {
	bsjson := LoadBlockStoreStateJSON(db)
	return &BlockStore{
		height: bsjson.Height,
		db:     db,
	}
}

// Height() returns the last known contiguous block height.
func (bs *BlockStore) Height() int {
	bs.mtx.RLock()
	defer bs.mtx.RUnlock()
	return bs.height
}

func (bs *BlockStore) GetReader(key []byte) io.Reader {
	bytez := bs.db.Get(key)
	if bytez == nil {
		return nil
	}
	return bytes.NewReader(bytez)
}

func (bs *BlockStore) LoadBlock(height int) *types.Block {
	var n int
	var err error
	r := bs.GetReader(calcBlockMetaKey(height))
	if r == nil {
		return nil
	}
	blockMeta := wire.ReadBinary(&types.BlockMeta{}, r, 0, &n, &err).(*types.BlockMeta)
	if err != nil {
		PanicCrisis(Fmt("Error reading block meta: %v", err))
	}
	bytez := []byte{}
	for i := 0; i < blockMeta.BlockID.PartsHeader.Total; i++ {
		part := bs.LoadBlockPart(height, i)
		bytez = append(bytez, part.Bytes...)
	}
	block := wire.ReadBinary(&types.Block{}, bytes.NewReader(bytez), 0, &n, &err).(*types.Block)
	if err != nil {
		PanicCrisis(Fmt("Error reading block: %v", err))
	}
	return block
}

func (bs *BlockStore) LoadBlockPart(height int, index int) *types.Part {
	var n int
	var err error
	r := bs.GetReader(calcBlockPartKey(height, index))
	if r == nil {
		return nil
	}
	part := wire.ReadBinary(&types.Part{}, r, 0, &n, &err).(*types.Part)
	if err != nil {
		PanicCrisis(Fmt("Error reading block part: %v", err))
	}
	return part
}

func (bs *BlockStore) LoadBlockMeta(height int) *types.BlockMeta {
	var n int
	var err error
	r := bs.GetReader(calcBlockMetaKey(height))
	if r == nil {
		return nil
	}
	blockMeta := wire.ReadBinary(&types.BlockMeta{}, r, 0, &n, &err).(*types.BlockMeta)
	if err != nil {
		PanicCrisis(Fmt("Error reading block meta: %v", err))
	}
	return blockMeta
}

// The +2/3 and other Precommit-votes for block at `height`.
// This Commit comes from block.LastCommit for `height+1`.
func (bs *BlockStore) LoadBlockCommit(height int) *types.Commit {
	var n int
	var err error
	r := bs.GetReader(calcBlockCommitKey(height))
	if r == nil {
		return nil
	}
	commit := wire.ReadBinary(&types.Commit{}, r, 0, &n, &err).(*types.Commit)
	if err != nil {
		PanicCrisis(Fmt("Error reading commit: %v", err))
	}
	return commit
}

// NOTE: the Precommit-vote heights are for the block at `height`
func (bs *BlockStore) LoadSeenCommit(height int) *types.Commit {
	var n int
	var err error
	r := bs.GetReader(calcSeenCommitKey(height))
	if r == nil {
		return nil
	}
	commit := wire.ReadBinary(&types.Commit{}, r, 0, &n, &err).(*types.Commit)
	if err != nil {
		PanicCrisis(Fmt("Error reading commit: %v", err))
	}
	return commit
}

// blockParts: Must be parts of the block
// seenCommit: The +2/3 precommits that were seen which committed at height.
//             If all the nodes restart after committing a block,
//             we need this to reload the precommits to catch-up nodes to the
//             most recent height.  Otherwise they'd stall at H-1.
func (bs *BlockStore) SaveBlock(block *types.Block, blockParts *types.PartSet, seenCommit *types.Commit) {
	height := block.Height
	if height != bs.Height()+1 {
		PanicSanity(Fmt("BlockStore can only save contiguous blocks. Wanted %v, got %v", bs.Height()+1, height))
	}
	if !blockParts.IsComplete() {
		PanicSanity(Fmt("BlockStore can only save complete block part sets"))
	}

	// Save block meta
	blockMeta := types.NewBlockMeta(block, blockParts)
	metaBytes := wire.BinaryBytes(blockMeta)
	bs.db.Set(calcBlockMetaKey(height), metaBytes)

	// Save block parts
	for i := 0; i < blockParts.Total(); i++ {
		bs.saveBlockPart(height, i, blockParts.GetPart(i))
	}

	// Save block commit (duplicate and separate from the Block)
	blockCommitBytes := wire.BinaryBytes(block.LastCommit)
	bs.db.Set(calcBlockCommitKey(height-1), blockCommitBytes)

	// Save seen commit (seen +2/3 precommits for block)
	// NOTE: we can delete this at a later height
	seenCommitBytes := wire.BinaryBytes(seenCommit)
	bs.db.Set(calcSeenCommitKey(height), seenCommitBytes)

	// Save new BlockStoreStateJSON descriptor
	BlockStoreStateJSON{Height: height}.Save(bs.db)

	// Done!
	bs.mtx.Lock()
	bs.height = height
	bs.mtx.Unlock()

	// Flush
	bs.db.SetSync(nil, nil)
}

func (bs *BlockStore) saveBlockPart(height int, index int, part *types.Part) {
	if height != bs.Height()+1 {
		PanicSanity(Fmt("BlockStore can only save contiguous blocks. Wanted %v, got %v", bs.Height()+1, height))
	}
	partBytes := wire.BinaryBytes(part)
	bs.db.Set(calcBlockPartKey(height, index), partBytes)
}

//-----------------------------------------------------------------------------

func calcBlockMetaKey(height int) []byte {
	return []byte(fmt.Sprintf("H:%v", height))
}

func calcBlockPartKey(height int, partIndex int) []byte {
	return []byte(fmt.Sprintf("P:%v:%v", height, partIndex))
}

func calcBlockCommitKey(height int) []byte {
	return []byte(fmt.Sprintf("C:%v", height))
}

func calcSeenCommitKey(height int) []byte {
	return []byte(fmt.Sprintf("SC:%v", height))
}

//-----------------------------------------------------------------------------

var blockStoreKey = []byte("blockStore")

type BlockStoreStateJSON struct {
	Height int
}

func (bsj BlockStoreStateJSON) Save(db dbm.DB) {
	bytes, err := json.Marshal(bsj)
	if err != nil {
		PanicSanity(Fmt("Could not marshal state bytes: %v", err))
	}
	db.SetSync(blockStoreKey, bytes)
}

func LoadBlockStoreStateJSON(db dbm.DB) BlockStoreStateJSON {
	bytes := db.Get(blockStoreKey)
	if bytes == nil {
		return BlockStoreStateJSON{
			Height: 0,
		}
	}
	bsj := BlockStoreStateJSON{}
	err := json.Unmarshal(bytes, &bsj)
	if err != nil {
		PanicCrisis(Fmt("Could not unmarshal bytes: %X", bytes))
	}
	return bsj
}
<commit_msg>blockchain/store: comment about panics<commit_after>package blockchain

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"sync"

	wire "github.com/tendermint/go-wire"
	"github.com/tendermint/tendermint/types"
	. "github.com/tendermint/tmlibs/common"
	dbm "github.com/tendermint/tmlibs/db"
)

/*
Simple low level store for blocks.

There are three types of information stored:
 - BlockMeta:   Meta information about each block
 - Block part:  Parts of each block, aggregated w/ PartSet
 - Commit:      The commit part of each block, for gossiping precommit votes

Currently the precommit signatures are duplicated in the Block parts as
well as the Commit.  In the future this may change, perhaps by moving
the Commit data outside the Block.

// NOTE: BlockStore methods will panic if they encounter errors
// deserializing loaded data, indicating probable corruption on disk.
*/
type BlockStore struct {
	db dbm.DB

	mtx    sync.RWMutex
	height int
}

func NewBlockStore(db dbm.DB) *BlockStore {
	bsjson := LoadBlockStoreStateJSON(db)
	return &BlockStore{
		height: bsjson.Height,
		db:     db,
	}
}

// Height() returns the last known contiguous block height.
func (bs *BlockStore) Height() int {
	bs.mtx.RLock()
	defer bs.mtx.RUnlock()
	return bs.height
}

func (bs *BlockStore) GetReader(key []byte) io.Reader {
	bytez := bs.db.Get(key)
	if bytez == nil {
		return nil
	}
	return bytes.NewReader(bytez)
}

func (bs *BlockStore) LoadBlock(height int) *types.Block {
	var n int
	var err error
	r := bs.GetReader(calcBlockMetaKey(height))
	if r == nil {
		return nil
	}
	blockMeta := wire.ReadBinary(&types.BlockMeta{}, r, 0, &n, &err).(*types.BlockMeta)
	if err != nil {
		PanicCrisis(Fmt("Error reading block meta: %v", err))
	}
	bytez := []byte{}
	for i := 0; i < blockMeta.BlockID.PartsHeader.Total; i++ {
		part := bs.LoadBlockPart(height, i)
		bytez = append(bytez, part.Bytes...)
	}
	block := wire.ReadBinary(&types.Block{}, bytes.NewReader(bytez), 0, &n, &err).(*types.Block)
	if err != nil {
		PanicCrisis(Fmt("Error reading block: %v", err))
	}
	return block
}

func (bs *BlockStore) LoadBlockPart(height int, index int) *types.Part {
	var n int
	var err error
	r := bs.GetReader(calcBlockPartKey(height, index))
	if r == nil {
		return nil
	}
	part := wire.ReadBinary(&types.Part{}, r, 0, &n, &err).(*types.Part)
	if err != nil {
		PanicCrisis(Fmt("Error reading block part: %v", err))
	}
	return part
}

func (bs *BlockStore) LoadBlockMeta(height int) *types.BlockMeta {
	var n int
	var err error
	r := bs.GetReader(calcBlockMetaKey(height))
	if r == nil {
		return nil
	}
	blockMeta := wire.ReadBinary(&types.BlockMeta{}, r, 0, &n, &err).(*types.BlockMeta)
	if err != nil {
		PanicCrisis(Fmt("Error reading block meta: %v", err))
	}
	return blockMeta
}

// The +2/3 and other Precommit-votes for block at `height`.
// This Commit comes from block.LastCommit for `height+1`.
func (bs *BlockStore) LoadBlockCommit(height int) *types.Commit {
	var n int
	var err error
	r := bs.GetReader(calcBlockCommitKey(height))
	if r == nil {
		return nil
	}
	commit := wire.ReadBinary(&types.Commit{}, r, 0, &n, &err).(*types.Commit)
	if err != nil {
		PanicCrisis(Fmt("Error reading commit: %v", err))
	}
	return commit
}

// NOTE: the Precommit-vote heights are for the block at `height`
func (bs *BlockStore) LoadSeenCommit(height int) *types.Commit {
	var n int
	var err error
	r := bs.GetReader(calcSeenCommitKey(height))
	if r == nil {
		return nil
	}
	commit := wire.ReadBinary(&types.Commit{}, r, 0, &n, &err).(*types.Commit)
	if err != nil {
		PanicCrisis(Fmt("Error reading commit: %v", err))
	}
	return commit
}

// blockParts: Must be parts of the block
// seenCommit: The +2/3 precommits that were seen which committed at height.
//             If all the nodes restart after committing a block,
//             we need this to reload the precommits to catch-up nodes to the
//             most recent height.  Otherwise they'd stall at H-1.
func (bs *BlockStore) SaveBlock(block *types.Block, blockParts *types.PartSet, seenCommit *types.Commit) {
	height := block.Height
	if height != bs.Height()+1 {
		PanicSanity(Fmt("BlockStore can only save contiguous blocks. Wanted %v, got %v", bs.Height()+1, height))
	}
	if !blockParts.IsComplete() {
		PanicSanity(Fmt("BlockStore can only save complete block part sets"))
	}

	// Save block meta
	blockMeta := types.NewBlockMeta(block, blockParts)
	metaBytes := wire.BinaryBytes(blockMeta)
	bs.db.Set(calcBlockMetaKey(height), metaBytes)

	// Save block parts
	for i := 0; i < blockParts.Total(); i++ {
		bs.saveBlockPart(height, i, blockParts.GetPart(i))
	}

	// Save block commit (duplicate and separate from the Block)
	blockCommitBytes := wire.BinaryBytes(block.LastCommit)
	bs.db.Set(calcBlockCommitKey(height-1), blockCommitBytes)

	// Save seen commit (seen +2/3 precommits for block)
	// NOTE: we can delete this at a later height
	seenCommitBytes := wire.BinaryBytes(seenCommit)
	bs.db.Set(calcSeenCommitKey(height), seenCommitBytes)

	// Save new BlockStoreStateJSON descriptor
	BlockStoreStateJSON{Height: height}.Save(bs.db)

	// Done!
	bs.mtx.Lock()
	bs.height = height
	bs.mtx.Unlock()

	// Flush
	bs.db.SetSync(nil, nil)
}

func (bs *BlockStore) saveBlockPart(height int, index int, part *types.Part) {
	if height != bs.Height()+1 {
		PanicSanity(Fmt("BlockStore can only save contiguous blocks. Wanted %v, got %v", bs.Height()+1, height))
	}
	partBytes := wire.BinaryBytes(part)
	bs.db.Set(calcBlockPartKey(height, index), partBytes)
}

//-----------------------------------------------------------------------------

func calcBlockMetaKey(height int) []byte {
	return []byte(fmt.Sprintf("H:%v", height))
}

func calcBlockPartKey(height int, partIndex int) []byte {
	return []byte(fmt.Sprintf("P:%v:%v", height, partIndex))
}

func calcBlockCommitKey(height int) []byte {
	return []byte(fmt.Sprintf("C:%v", height))
}

func calcSeenCommitKey(height int) []byte {
	return []byte(fmt.Sprintf("SC:%v", height))
}

//-----------------------------------------------------------------------------

var blockStoreKey = []byte("blockStore")

type BlockStoreStateJSON struct {
	Height int
}

func (bsj BlockStoreStateJSON) Save(db dbm.DB) {
	bytes, err := json.Marshal(bsj)
	if err != nil {
		PanicSanity(Fmt("Could not marshal state bytes: %v", err))
	}
	db.SetSync(blockStoreKey, bytes)
}

func LoadBlockStoreStateJSON(db dbm.DB) BlockStoreStateJSON {
	bytes := db.Get(blockStoreKey)
	if bytes == nil {
		return BlockStoreStateJSON{
			Height: 0,
		}
	}
	bsj := BlockStoreStateJSON{}
	err := json.Unmarshal(bytes, &bsj)
	if err != nil {
		PanicCrisis(Fmt("Could not unmarshal bytes: %X", bytes))
	}
	return bsj
}
<|endoftext|>


<commit_before>package bitswap

import (
	"sync"
	"time"

	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
	engine "github.com/ipfs/go-ipfs/exchange/bitswap/decision"
	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	wantlist "github.com/ipfs/go-ipfs/exchange/bitswap/wantlist"
	peer "github.com/ipfs/go-ipfs/p2p/peer"
	u "github.com/ipfs/go-ipfs/util"
)

type WantManager struct {
	// sync channels for Run loop
	incoming   chan []*bsmsg.Entry
	connect    chan peer.ID // notification channel for new peers connecting
	disconnect chan peer.ID // notification channel for peers disconnecting

	// synchronized by Run loop, only touch inside there
	peers map[peer.ID]*msgQueue
	wl    *wantlist.Wantlist

	network bsnet.BitSwapNetwork
	ctx     context.Context
}

func NewWantManager(ctx context.Context, network bsnet.BitSwapNetwork) *WantManager {
	return &WantManager{
		incoming:   make(chan []*bsmsg.Entry, 10),
		connect:    make(chan peer.ID, 10),
		disconnect: make(chan peer.ID, 10),
		peers:      make(map[peer.ID]*msgQueue),
		wl:         wantlist.New(),
		network:    network,
		ctx:        ctx,
	}
}

type msgPair struct {
	to  peer.ID
	msg bsmsg.BitSwapMessage
}

type cancellation struct {
	who peer.ID
	blk u.Key
}

type msgQueue struct {
	p peer.ID

	outlk   sync.Mutex
	out     bsmsg.BitSwapMessage
	network bsnet.BitSwapNetwork

	work chan struct{}
	done chan struct{}
}

func (pm *WantManager) WantBlocks(ks []u.Key) {
	log.Infof("want blocks: %s", ks)
	pm.addEntries(ks, false)
}

func (pm *WantManager) CancelWants(ks []u.Key) {
	pm.addEntries(ks, true)
}

func (pm *WantManager) addEntries(ks []u.Key, cancel bool) {
	var entries []*bsmsg.Entry
	for i, k := range ks {
		entries = append(entries, &bsmsg.Entry{
			Cancel: cancel,
			Entry: wantlist.Entry{
				Key:      k,
				Priority: kMaxPriority - i,
			},
		})
	}
	select {
	case pm.incoming <- entries:
	case <-pm.ctx.Done():
	}
}

func (pm *WantManager) SendBlock(ctx context.Context, env *engine.Envelope) {
	// Blocks need to be sent synchronously to maintain proper backpressure
	// throughout the network stack
	defer env.Sent()

	msg := bsmsg.New(false)
	msg.AddBlock(env.Block)
	log.Infof("Sending block %s to %s", env.Peer, env.Block)
	err := pm.network.SendMessage(ctx, env.Peer, msg)
	if err != nil {
		log.Noticef("sendblock error: %s", err)
	}
}

func (pm *WantManager) startPeerHandler(p peer.ID) *msgQueue {
	_, ok := pm.peers[p]
	if ok {
		// TODO: log an error?
		return nil
	}

	mq := pm.newMsgQueue(p)

	// new peer, we will want to give them our full wantlist
	fullwantlist := bsmsg.New(true)
	for _, e := range pm.wl.Entries() {
		fullwantlist.AddEntry(e.Key, e.Priority)
	}
	mq.out = fullwantlist
	mq.work <- struct{}{}

	pm.peers[p] = mq
	go mq.runQueue(pm.ctx)
	return mq
}

func (pm *WantManager) stopPeerHandler(p peer.ID) {
	pq, ok := pm.peers[p]
	if !ok {
		// TODO: log error?
		return
	}

	close(pq.done)
	delete(pm.peers, p)
}

func (mq *msgQueue) runQueue(ctx context.Context) {
	for {
		select {
		case <-mq.work: // there is work to be done

			err := mq.network.ConnectTo(ctx, mq.p)
			if err != nil {
				log.Errorf("cant connect to peer %s: %s", mq.p, err)
				// TODO: cant connect, what now?
				continue
			}

			// grab outgoing message
			mq.outlk.Lock()
			wlm := mq.out
			if wlm == nil || wlm.Empty() {
				mq.outlk.Unlock()
				continue
			}
			mq.out = nil
			mq.outlk.Unlock()

			// send wantlist updates
			err = mq.network.SendMessage(ctx, mq.p, wlm)
			if err != nil {
				log.Noticef("bitswap send error: %s", err)
				// TODO: what do we do if this fails?
			}
		case <-mq.done:
			return
		}
	}
}

func (pm *WantManager) Connected(p peer.ID) {
	pm.connect <- p
}

func (pm *WantManager) Disconnected(p peer.ID) {
	pm.disconnect <- p
}

// TODO: use goprocess here once i trust it
func (pm *WantManager) Run() {
	tock := time.NewTicker(rebroadcastDelay.Get())
	for {
		select {
		case entries := <-pm.incoming:

			// add changes to our wantlist
			for _, e := range entries {
				if e.Cancel {
					pm.wl.Remove(e.Key)
				} else {
					pm.wl.Add(e.Key, e.Priority)
				}
			}

			// broadcast those wantlist changes
			for _, p := range pm.peers {
				p.addMessage(entries)
			}

		case <-tock.C:
			// resend entire wantlist every so often (REALLY SHOULDNT BE NECESSARY)
			var es []*bsmsg.Entry
			for _, e := range pm.wl.Entries() {
				es = append(es, &bsmsg.Entry{Entry: e})
			}
			for _, p := range pm.peers {
				p.outlk.Lock()
				p.out = bsmsg.New(true)
				p.outlk.Unlock()

				p.addMessage(es)
			}
		case p := <-pm.connect:
			pm.startPeerHandler(p)
		case p := <-pm.disconnect:
			pm.stopPeerHandler(p)
		case <-pm.ctx.Done():
			return
		}
	}
}

func (wm *WantManager) newMsgQueue(p peer.ID) *msgQueue {
	mq := new(msgQueue)
	mq.done = make(chan struct{})
	mq.work = make(chan struct{}, 1)
	mq.network = wm.network
	mq.p = p

	return mq
}

func (mq *msgQueue) addMessage(entries []*bsmsg.Entry) {
	mq.outlk.Lock()
	defer func() {
		mq.outlk.Unlock()
		select {
		case mq.work <- struct{}{}:
		default:
		}
	}()

	// if we have no message held, or the one we are given is full
	// overwrite the one we are holding
	if mq.out == nil {
		mq.out = bsmsg.New(false)
	}

	// TODO: add a msg.Combine(...) method
	// otherwise, combine the one we are holding with the
	// one passed in
	for _, e := range entries {
		if e.Cancel {
			mq.out.Cancel(e.Key)
		} else {
			mq.out.AddEntry(e.Key, e.Priority)
		}
	}
}
<commit_msg>defer tock.Stop()<commit_after>package bitswap

import (
	"sync"
	"time"

	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
	engine "github.com/ipfs/go-ipfs/exchange/bitswap/decision"
	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	wantlist "github.com/ipfs/go-ipfs/exchange/bitswap/wantlist"
	peer "github.com/ipfs/go-ipfs/p2p/peer"
	u "github.com/ipfs/go-ipfs/util"
)

type WantManager struct {
	// sync channels for Run loop
	incoming   chan []*bsmsg.Entry
	connect    chan peer.ID // notification channel for new peers connecting
	disconnect chan peer.ID // notification channel for peers disconnecting

	// synchronized by Run loop, only touch inside there
	peers map[peer.ID]*msgQueue
	wl    *wantlist.Wantlist

	network bsnet.BitSwapNetwork
	ctx     context.Context
}

func NewWantManager(ctx context.Context, network bsnet.BitSwapNetwork) *WantManager {
	return &WantManager{
		incoming:   make(chan []*bsmsg.Entry, 10),
		connect:    make(chan peer.ID, 10),
		disconnect: make(chan peer.ID, 10),
		peers:      make(map[peer.ID]*msgQueue),
		wl:         wantlist.New(),
		network:    network,
		ctx:        ctx,
	}
}

type msgPair struct {
	to  peer.ID
	msg bsmsg.BitSwapMessage
}

type cancellation struct {
	who peer.ID
	blk u.Key
}

type msgQueue struct {
	p peer.ID

	outlk   sync.Mutex
	out     bsmsg.BitSwapMessage
	network bsnet.BitSwapNetwork

	work chan struct{}
	done chan struct{}
}

func (pm *WantManager) WantBlocks(ks []u.Key) {
	log.Infof("want blocks: %s", ks)
	pm.addEntries(ks, false)
}

func (pm *WantManager) CancelWants(ks []u.Key) {
	pm.addEntries(ks, true)
}

func (pm *WantManager) addEntries(ks []u.Key, cancel bool) {
	var entries []*bsmsg.Entry
	for i, k := range ks {
		entries = append(entries, &bsmsg.Entry{
			Cancel: cancel,
			Entry: wantlist.Entry{
				Key:      k,
				Priority: kMaxPriority - i,
			},
		})
	}
	select {
	case pm.incoming <- entries:
	case <-pm.ctx.Done():
	}
}

func (pm *WantManager) SendBlock(ctx context.Context, env *engine.Envelope) {
	// Blocks need to be sent synchronously to maintain proper backpressure
	// throughout the network stack
	defer env.Sent()

	msg := bsmsg.New(false)
	msg.AddBlock(env.Block)
	log.Infof("Sending block %s to %s", env.Peer, env.Block)
	err := pm.network.SendMessage(ctx, env.Peer, msg)
	if err != nil {
		log.Noticef("sendblock error: %s", err)
	}
}

func (pm *WantManager) startPeerHandler(p peer.ID) *msgQueue {
	_, ok := pm.peers[p]
	if ok {
		// TODO: log an error?
		return nil
	}

	mq := pm.newMsgQueue(p)

	// new peer, we will want to give them our full wantlist
	fullwantlist := bsmsg.New(true)
	for _, e := range pm.wl.Entries() {
		fullwantlist.AddEntry(e.Key, e.Priority)
	}
	mq.out = fullwantlist
	mq.work <- struct{}{}

	pm.peers[p] = mq
	go mq.runQueue(pm.ctx)
	return mq
}

func (pm *WantManager) stopPeerHandler(p peer.ID) {
	pq, ok := pm.peers[p]
	if !ok {
		// TODO: log error?
		return
	}

	close(pq.done)
	delete(pm.peers, p)
}

func (mq *msgQueue) runQueue(ctx context.Context) {
	for {
		select {
		case <-mq.work: // there is work to be done

			err := mq.network.ConnectTo(ctx, mq.p)
			if err != nil {
				log.Errorf("cant connect to peer %s: %s", mq.p, err)
				// TODO: cant connect, what now?
				continue
			}

			// grab outgoing message
			mq.outlk.Lock()
			wlm := mq.out
			if wlm == nil || wlm.Empty() {
				mq.outlk.Unlock()
				continue
			}
			mq.out = nil
			mq.outlk.Unlock()

			// send wantlist updates
			err = mq.network.SendMessage(ctx, mq.p, wlm)
			if err != nil {
				log.Noticef("bitswap send error: %s", err)
				// TODO: what do we do if this fails?
			}
		case <-mq.done:
			return
		}
	}
}

func (pm *WantManager) Connected(p peer.ID) {
	pm.connect <- p
}

func (pm *WantManager) Disconnected(p peer.ID) {
	pm.disconnect <- p
}

// TODO: use goprocess here once i trust it
func (pm *WantManager) Run() {
	tock := time.NewTicker(rebroadcastDelay.Get())
	defer tock.Stop()
	for {
		select {
		case entries := <-pm.incoming:

			// add changes to our wantlist
			for _, e := range entries {
				if e.Cancel {
					pm.wl.Remove(e.Key)
				} else {
					pm.wl.Add(e.Key, e.Priority)
				}
			}

			// broadcast those wantlist changes
			for _, p := range pm.peers {
				p.addMessage(entries)
			}

		case <-tock.C:
			// resend entire wantlist every so often (REALLY SHOULDNT BE NECESSARY)
			var es []*bsmsg.Entry
			for _, e := range pm.wl.Entries() {
				es = append(es, &bsmsg.Entry{Entry: e})
			}
			for _, p := range pm.peers {
				p.outlk.Lock()
				p.out = bsmsg.New(true)
				p.outlk.Unlock()

				p.addMessage(es)
			}
		case p := <-pm.connect:
			pm.startPeerHandler(p)
		case p := <-pm.disconnect:
			pm.stopPeerHandler(p)
		case <-pm.ctx.Done():
			return
		}
	}
}

func (wm *WantManager) newMsgQueue(p peer.ID) *msgQueue {
	mq := new(msgQueue)
	mq.done = make(chan struct{})
	mq.work = make(chan struct{}, 1)
	mq.network = wm.network
	mq.p = p

	return mq
}

func (mq *msgQueue) addMessage(entries []*bsmsg.Entry) {
	mq.outlk.Lock()
	defer func() {
		mq.outlk.Unlock()
		select {
		case mq.work <- struct{}{}:
		default:
		}
	}()

	// if we have no message held, or the one we are given is full
	// overwrite the one we are holding
	if mq.out == nil {
		mq.out = bsmsg.New(false)
	}

	// TODO: add a msg.Combine(...) method
	// otherwise, combine the one we are holding with the
	// one passed in
	for _, e := range entries {
		if e.Cancel {
			mq.out.Cancel(e.Key)
		} else {
			mq.out.AddEntry(e.Key, e.Priority)
		}
	}
}
<|endoftext|>


<commit_before>package testlog

import (
	"os"
	"sync"
	"testing"

	"github.com/getlantern/golog"
)

// Capture captures logs to the given testing.T's Log function.
// Returns a function that stops capturing logs.
//
// Typical usage:
//
//    func MyTest(t *testing.T) {
//        stopCapture := testlog.Capture(t)
//        defer stopCapture()
//        // do stuff
//    }
//
func Capture(t *testing.T) func() {
	w := &testLogWriter{T: t}
	reset := golog.SetOutputs(w, w)
	return func() {
		reset()
		w.stop()
	}
}

type testLogWriter struct {
	*testing.T
	mu      sync.RWMutex
	stopped bool
}

func (w *testLogWriter) Write(p []byte) (n int, err error) {
	w.mu.RLock()
	defer w.mu.RUnlock()
	if w.stopped {
		// After writer stopped, just log to console
		p = append([]byte("(logged after test capture stopped) "), p...)
		_, err := os.Stderr.Write(p)
		return len(p), err
	}
	return len(p), nil
}

func (w *testLogWriter) stop() {
	w.mu.Lock()
	w.stopped = true
	w.mu.Unlock()
}
<commit_msg>actually call T.Log<commit_after>package testlog

import (
	"os"
	"sync"
	"testing"

	"github.com/getlantern/golog"
)

// Capture captures logs to the given testing.T's Log function.
// Returns a function that stops capturing logs.
//
// Typical usage:
//
//    func MyTest(t *testing.T) {
//        stopCapture := testlog.Capture(t)
//        defer stopCapture()
//        // do stuff
//    }
//
func Capture(t *testing.T) func() {
	w := &testLogWriter{T: t}
	reset := golog.SetOutputs(w, w)
	return func() {
		reset()
		w.stop()
	}
}

type testLogWriter struct {
	*testing.T
	mu      sync.RWMutex
	stopped bool
}

func (w *testLogWriter) Write(p []byte) (n int, err error) {
	w.mu.RLock()
	defer w.mu.RUnlock()
	if w.stopped {
		// After writer stopped, just log to console
		p = append([]byte("(logged after test capture stopped) "), p...)
		_, err := os.Stderr.Write(p)
		return len(p), err
	}
	w.Log(string(p))
	return len(p), nil
}

func (w *testLogWriter) stop() {
	w.mu.Lock()
	w.stopped = true
	w.mu.Unlock()
}
<|endoftext|>


<commit_before>package vagrantkite

import (
	"errors"
	"fmt"
	"net/url"
	"strconv"
	"testing"

	"github.com/koding/kite"

	"github.com/hashicorp/terraform/helper/resource"
)

const (
	testVagrantKiteBuildConfig = `
resource "vagrantkite_build" "myfirstvm" {
    kiteURL = "////////test"
}
`
	vagrantFile = `# -*- mode: ruby -*-
# vi: set ft=ruby :

VAGRANTFILE_API_VERSION = "2"

Vagrant.configure(VAGRANTFILE_API_VERSION) do |config|
  config.vm.box = "ubuntu/trusty64"
  config.vm.hostname = "vagrant"

  config.vm.provider "virtualbox" do |vb|
    # Use VBoxManage to customize the VM. For example to change memory:
    vb.customize ["modifyvm", :id, "--memory", "2048", "--cpus", "2"]
  end
end
`
	vagrantFilePath = "/home/etc"
)

func TestAccGithubAddUser_Basic(t *testing.T) {
	resource.Test(t, resource.TestCase{
		// PreCheck:  func() { testAccPreCheck(t) },
		Providers: testVagrantResourceProviders,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testVagrantKiteBuildConfig,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"vagrantkite_build.myfirstvm",
						"kiteURL",
						"////////test",
					),
				),
			},
		},
	})
}

func withClient(t *testing.T, f func(c *Client) error) {
	client, err := NewClient()
	if err != nil {
		t.Errorf(err.Error())
	}

	client.Kite.Config.DisableAuthentication = true
	client.Kite.Config.Port = 5000
	client.Kite.HandleFunc(klientFuncName, mockHandler)

	go client.Kite.Run()
	<-client.Kite.ServerReadyNotify()

	kiteURL := &url.URL{
		Scheme: "http",
		Host:   "localhost:" + strconv.Itoa(client.Kite.Port()),
		Path:   "/kite",
	}

	if _, err := client.Kite.Register(kiteURL); err != nil {
		t.Errorf("couldnt register to kontrol %s", err.Error())
	}

	err = f(client)
	client.Kite.Close()
	if err != nil {
		t.Errorf("failed with %s", err.Error())
	}
}

func TestSendingCommandSuccess(t *testing.T) {
	withClient(t, func(c *Client) error {

		args := &vagrantKiteReq{
			VagrantFile: vagrantFile,
			FilePath:    vagrantFilePath,
		}

		queryString := c.Kite.Kite().String()

		if err := sendCommand(klientFuncName, queryString, args); err != nil {
			return err
		}

		return nil
	})
}

func TestSendingCommandFailure(t *testing.T) {
	withClient(t, func(c *Client) error {

		args := &vagrantKiteReq{
			VagrantFile: vagrantFile + "1",
			FilePath:    vagrantFilePath,
		}

		queryString := c.Kite.Kite().String()

		if err := sendCommand(klientFuncName, queryString, args); err != nil {
			return nil
		}

		return errors.New("failure should happen")
	})
}

var mockHandler = func(r *kite.Request) (interface{}, error) {
	var res []vagrantKiteReq // another slice??
	if err := r.Args.Unmarshal(&res); err != nil {
		return nil, fmt.Errorf("err while unmarshalling: %s", err.Error())
	}

	if res[0].FilePath != vagrantFilePath {
		return nil, fmt.Errorf("filePath is %+v, expected %+v", res[0].FilePath, vagrantFilePath)
	}

	if res[0].VagrantFile != vagrantFile {
		return nil, fmt.Errorf("vagrantFile is %+v, expected %+v", res[0].VagrantFile, vagrantFile)
	}

	return res, nil
}
<commit_msg>Kite: added more test cases<commit_after>package vagrantkite

import (
	"errors"
	"fmt"
	"net/url"
	"strconv"
	"testing"

	"github.com/koding/kite"

	"github.com/hashicorp/terraform/helper/resource"
)

const (
	testVagrantKiteBuildConfig = `
resource "vagrantkite_build" "myfirstvm" {
	filePath = ` + vagrantFilePath + `
    queryString = ` + queryString + `
    vagrantFile = ` + vagrantFile + `
}
`
	queryString = "///////8c396fd6-c91c-4454-45c2-5c461ad32645"

	vagrantFile = `# -*- mode: ruby -*-
# vi: set ft=ruby :

VAGRANTFILE_API_VERSION = "2"

Vagrant.configure(VAGRANTFILE_API_VERSION) do |config|
  config.vm.box = "ubuntu/trusty64"
  config.vm.hostname = "vagrant"

  config.vm.provider "virtualbox" do |vb|
    # Use VBoxManage to customize the VM. For example to change memory:
    vb.customize ["modifyvm", :id, "--memory", "2048", "--cpus", "2"]
  end
end
`
	vagrantFilePath = "/home/etc/Vagrantfile"
)

func TestAccGithubAddUser_Basic(t *testing.T) {
	resource.Test(t, resource.TestCase{
		// PreCheck:  func() { testAccPreCheck(t) },
		Providers: testVagrantResourceProviders,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testVagrantKiteBuildConfig,
				Check: resource.ComposeTestCheckFunc(
					resource.TestCheckResourceAttr(
						"vagrantkite_build.myfirstvm",
						"filePath",
						vagrantFilePath,
					),
					resource.TestCheckResourceAttr(
						"vagrantkite_build.myfirstvm",
						"queryString",
						queryString,
					),
					resource.TestCheckResourceAttr(
						"vagrantkite_build.myfirstvm",
						"vagrantFile",
						vagrantFile,
					),
				),
			},
		},
	})
}

func withClient(t *testing.T, f func(c *Client) error) {
	client, err := NewClient()
	if err != nil {
		t.Errorf(err.Error())
	}

	client.Kite.Config.DisableAuthentication = true
	client.Kite.Config.Port = 5000
	client.Kite.HandleFunc(klientFuncName, mockHandler)

	go client.Kite.Run()
	<-client.Kite.ServerReadyNotify()

	queryString := &url.URL{
		Scheme: "http",
		Host:   "localhost:" + strconv.Itoa(client.Kite.Port()),
		Path:   "/kite",
	}

	if _, err := client.Kite.Register(queryString); err != nil {
		t.Errorf("couldnt register to kontrol %s", err.Error())
	}

	err = f(client)
	client.Kite.Close()
	if err != nil {
		t.Errorf("failed with %s", err.Error())
	}
}

func TestSendingCommandSuccess(t *testing.T) {
	withClient(t, func(c *Client) error {

		args := &vagrantKiteReq{
			VagrantFile: vagrantFile,
			FilePath:    vagrantFilePath,
		}

		queryString := c.Kite.Kite().String()

		if err := sendCommand(klientFuncName, queryString, args); err != nil {
			return err
		}

		return nil
	})
}

func TestSendingCommandFailure(t *testing.T) {
	withClient(t, func(c *Client) error {

		args := &vagrantKiteReq{
			VagrantFile: vagrantFile + "1",
			FilePath:    vagrantFilePath,
		}

		queryString := c.Kite.Kite().String()

		if err := sendCommand(klientFuncName, queryString, args); err != nil {
			return nil
		}

		return errors.New("failure should happen")
	})
}

var mockHandler = func(r *kite.Request) (interface{}, error) {
	var res []vagrantKiteReq // another slice??
	if err := r.Args.Unmarshal(&res); err != nil {
		return nil, fmt.Errorf("err while unmarshalling: %s", err.Error())
	}

	if res[0].FilePath != vagrantFilePath {
		return nil, fmt.Errorf("filePath is %+v, expected %+v", res[0].FilePath, vagrantFilePath)
	}

	if res[0].VagrantFile != vagrantFile {
		return nil, fmt.Errorf("vagrantFile is %+v, expected %+v", res[0].VagrantFile, vagrantFile)
	}

	return res, nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013, 2014 Canonical Ltd.
// Copyright 2014 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.

package cloudinit

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/names"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/cloudinit"
	"github.com/juju/juju/version"
)

type UserdataConfig interface {
	// Configure is a convenience function that updates the cloudinit.Config
	// with appropriate configuration. It will run ConfigureBasic() and
	// ConfigureJuju()
	Configure() error
	// ConfigureBasic updates the provided cloudinit.Config with
	// basic configuration to initialise an OS image.
	ConfigureBasic() error
	// ConfigureJuju updates the provided cloudinit.Config with configuration
	// to initialise a Juju machine agent.
	ConfigureJuju() error
	// Render renders the cloudinit/cloudbase-init userdata needed to initialize
	// the juju agent
	Render() ([]byte, error)
}

// addAgentInfo adds agent-required information to the agent's directory
// and returns the agent directory name.
func addAgentInfo(
	cfg *MachineConfig,
	c *cloudinit.Config,
	tag names.Tag,
	toolsVersion version.Number,
) (agent.Config, error) {
	acfg, err := cfg.agentConfig(tag, toolsVersion)
	if err != nil {
		return nil, err
	}
	acfg.SetValue(agent.AgentServiceName, cfg.MachineAgentServiceName)
	cmds, err := acfg.WriteCommands(cfg.Series)
	if err != nil {
		return nil, errors.Annotate(err, "failed to write commands")
	}
	c.AddScripts(cmds...)
	return acfg, nil
}

func NewUserdataConfig(mcfg *MachineConfig, conf *cloudinit.Config) (UserdataConfig, error) {
	// TODO(ericsnow) bug #1426217
	// Protect mcfg and conf better.
	operatingSystem, err := version.GetOSFromSeries(mcfg.Series)
	if err != nil {
		return nil, err
	}

	base := baseConfigure{
		mcfg: mcfg,
		conf: conf,
		os:   operatingSystem,
	}
	if err := base.init(); err != nil {
		return nil, errors.Trace(err)
	}

	switch operatingSystem {
	case version.Ubuntu:
		return &ubuntuConfigure{base}, nil
	case version.Windows:
		return &windowsConfigure{base}, nil
	default:
		return nil, errors.Errorf("Unsupported OS %s", mcfg.Series)
	}
}

type baseConfigure struct {
	mcfg     *MachineConfig
	conf     *cloudinit.Config
	renderer cloudinit.Renderer
	os       version.OSType
}

func (c *baseConfigure) init() error {
	renderer, err := cloudinit.NewRenderer(c.mcfg.Series)
	if err != nil {
		return err
	}
	c.renderer = renderer
	return nil
}

func (c *baseConfigure) Render() ([]byte, error) {
	return c.renderer.Render(c.conf)
}

func (c *baseConfigure) addMachineAgentToBoot(tag string) error {
	svc, toolsDir, err := c.mcfg.initService()
	if err != nil {
		return errors.Trace(err)
	}

	// Make the agent run via a symbolic link to the actual tools
	// directory, so it can upgrade itself without needing to change
	// the init script.
	c.conf.AddScripts(c.toolsSymlinkCommand(toolsDir))

	cmds, err := svc.InstallCommands()
	if err != nil {
		return errors.Annotatef(err, "cannot make cloud-init init script for the %s agent", tag)
	}

	name := c.mcfg.MachineAgentServiceName
	c.conf.AddRunCmd(cloudinit.LogProgressCmd("Starting Juju machine agent (%s)", name))
	c.conf.AddScripts(cmds...)
	return nil
}

func (c *baseConfigure) toolsSymlinkCommand(toolsDir string) string {
	switch c.os {
	case version.Windows:
		return fmt.Sprintf(
			`cmd.exe /C mklink /D %s %v`,
			c.renderer.FromSlash(toolsDir),
			c.mcfg.Tools.Version,
		)
	default:
		// TODO(dfc) ln -nfs, so it doesn't fail if for some reason that
		// the target already exists.
		return fmt.Sprintf(
			"ln -s %v %s",
			c.mcfg.Tools.Version,
			shquote(toolsDir),
		)
	}
}
<commit_msg>Fix an error message.<commit_after>// Copyright 2012, 2013, 2014 Canonical Ltd.
// Copyright 2014 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.

package cloudinit

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/names"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/cloudinit"
	"github.com/juju/juju/version"
)

type UserdataConfig interface {
	// Configure is a convenience function that updates the cloudinit.Config
	// with appropriate configuration. It will run ConfigureBasic() and
	// ConfigureJuju()
	Configure() error
	// ConfigureBasic updates the provided cloudinit.Config with
	// basic configuration to initialise an OS image.
	ConfigureBasic() error
	// ConfigureJuju updates the provided cloudinit.Config with configuration
	// to initialise a Juju machine agent.
	ConfigureJuju() error
	// Render renders the cloudinit/cloudbase-init userdata needed to initialize
	// the juju agent
	Render() ([]byte, error)
}

// addAgentInfo adds agent-required information to the agent's directory
// and returns the agent directory name.
func addAgentInfo(
	cfg *MachineConfig,
	c *cloudinit.Config,
	tag names.Tag,
	toolsVersion version.Number,
) (agent.Config, error) {
	acfg, err := cfg.agentConfig(tag, toolsVersion)
	if err != nil {
		return nil, err
	}
	acfg.SetValue(agent.AgentServiceName, cfg.MachineAgentServiceName)
	cmds, err := acfg.WriteCommands(cfg.Series)
	if err != nil {
		return nil, errors.Annotate(err, "failed to write commands")
	}
	c.AddScripts(cmds...)
	return acfg, nil
}

func NewUserdataConfig(mcfg *MachineConfig, conf *cloudinit.Config) (UserdataConfig, error) {
	// TODO(ericsnow) bug #1426217
	// Protect mcfg and conf better.
	operatingSystem, err := version.GetOSFromSeries(mcfg.Series)
	if err != nil {
		return nil, err
	}

	base := baseConfigure{
		mcfg: mcfg,
		conf: conf,
		os:   operatingSystem,
	}
	if err := base.init(); err != nil {
		return nil, errors.Trace(err)
	}

	switch operatingSystem {
	case version.Ubuntu:
		return &ubuntuConfigure{base}, nil
	case version.Windows:
		return &windowsConfigure{base}, nil
	default:
		return nil, errors.NotSupportedf("OS %s", mcfg.Series)
	}
}

type baseConfigure struct {
	mcfg     *MachineConfig
	conf     *cloudinit.Config
	renderer cloudinit.Renderer
	os       version.OSType
}

func (c *baseConfigure) init() error {
	renderer, err := cloudinit.NewRenderer(c.mcfg.Series)
	if err != nil {
		return err
	}
	c.renderer = renderer
	return nil
}

func (c *baseConfigure) Render() ([]byte, error) {
	return c.renderer.Render(c.conf)
}

func (c *baseConfigure) addMachineAgentToBoot(tag string) error {
	svc, toolsDir, err := c.mcfg.initService()
	if err != nil {
		return errors.Trace(err)
	}

	// Make the agent run via a symbolic link to the actual tools
	// directory, so it can upgrade itself without needing to change
	// the init script.
	c.conf.AddScripts(c.toolsSymlinkCommand(toolsDir))

	cmds, err := svc.InstallCommands()
	if err != nil {
		return errors.Annotatef(err, "cannot make cloud-init init script for the %s agent", tag)
	}

	name := c.mcfg.MachineAgentServiceName
	c.conf.AddRunCmd(cloudinit.LogProgressCmd("Starting Juju machine agent (%s)", name))
	c.conf.AddScripts(cmds...)
	return nil
}

func (c *baseConfigure) toolsSymlinkCommand(toolsDir string) string {
	switch c.os {
	case version.Windows:
		return fmt.Sprintf(
			`cmd.exe /C mklink /D %s %v`,
			c.renderer.FromSlash(toolsDir),
			c.mcfg.Tools.Version,
		)
	default:
		// TODO(dfc) ln -nfs, so it doesn't fail if for some reason that
		// the target already exists.
		return fmt.Sprintf(
			"ln -s %v %s",
			c.mcfg.Tools.Version,
			shquote(toolsDir),
		)
	}
}
<|endoftext|>


<commit_before>package oled96x96

func ExampleHelloWorld() {
	bus := &i2c.Devfs{
		// change the following value if you use another bus
		Dev:  "/dev/i2c-1",
		Addr: Address,
	}

	display, err := New(bus)
	if err != nil {
		panic(err)
	}
	defer display.Close()

	display.Write("Hello World!")
}
<commit_msg>bringing back the test example<commit_after>package oled96x96_test

import "github.com/goiot/drivers/i2c/grove/oled96x96"

func Example() {
	bus := &i2c.Devfs{
		// change the following value if you use another bus
		Dev:  "/dev/i2c-1",
		Addr: oled96x96.Address,
	}

	display, err := oled96x96.New(bus)
	if err != nil {
		panic(err)
	}
	defer display.Close()

	display.Write("Hello World!")
}
<|endoftext|>


<commit_before>package gitbot

import "fmt"

type PayloadWatch struct {
	Action string      `json:"action"`     // The action that was performed. Currently, can only be started.
	Repo   *Repository `json:"repository"` // Repo
	Sender *User       `json:"sender"`     // Sender
}

func (s PayloadWatch) String() string {
	return fmt.Sprintf("[%s] was starred by %s", s.Repo, s.Sender)
}
<commit_msg>Watch: Check if action was started in case more actions are added in future<commit_after>package gitbot

import "fmt"

type PayloadWatch struct {
	Action string      `json:"action"`     // The action that was performed. Currently, can only be started.
	Repo   *Repository `json:"repository"` // Repo
	Sender *User       `json:"sender"`     // Sender
}

func (s PayloadWatch) String() string {
	if s.Action == "started" {
		return fmt.Sprintf("[%s] was starred by %s", s.Repo, s.Sender)
	}
}
<|endoftext|>


<commit_before>package mux_test

import (
	"log"
	"net/http"

	"github.com/gorilla/mux"
)

// Define our struct
type authenticationMiddleware struct {
	tokenUsers map[string]string
}

// Initialize it somewhere
func (amw *authenticationMiddleware) Populate() {
	amw.tokenUsers["00000000"] = "user0"
	amw.tokenUsers["aaaaaaaa"] = "userA"
	amw.tokenUsers["05f717e5"] = "randomUser"
	amw.tokenUsers["deadbeef"] = "user0"
}

// Middleware function, which will be called for each request
func (amw *authenticationMiddleware) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		token := r.Header.Get("X-Session-Token")

		if user, found := amw.tokenUsers[token]; found {
			// We found the token in our map
			log.Printf("Authenticated user %s\n", user)
			next.ServeHTTP(w, r)
		} else {
			http.Error(w, "Forbidden", http.StatusForbidden)
		}
	})
}

func Example_authenticationMiddleware() {
	r := mux.NewRouter()
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Do something here
	})
	amw := authenticationMiddleware{}
	amw.Populate()
	r.Use(amw.Middleware)
}
<commit_msg>Initialize user map (#371)<commit_after>package mux_test

import (
	"log"
	"net/http"

	"github.com/gorilla/mux"
)

// Define our struct
type authenticationMiddleware struct {
	tokenUsers map[string]string
}

// Initialize it somewhere
func (amw *authenticationMiddleware) Populate() {
	amw.tokenUsers["00000000"] = "user0"
	amw.tokenUsers["aaaaaaaa"] = "userA"
	amw.tokenUsers["05f717e5"] = "randomUser"
	amw.tokenUsers["deadbeef"] = "user0"
}

// Middleware function, which will be called for each request
func (amw *authenticationMiddleware) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		token := r.Header.Get("X-Session-Token")

		if user, found := amw.tokenUsers[token]; found {
			// We found the token in our map
			log.Printf("Authenticated user %s\n", user)
			next.ServeHTTP(w, r)
		} else {
			http.Error(w, "Forbidden", http.StatusForbidden)
		}
	})
}

func Example_authenticationMiddleware() {
	r := mux.NewRouter()
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Do something here
	})
	amw := authenticationMiddleware{make(map[string]string)}
	amw.Populate()
	r.Use(amw.Middleware)
}
<|endoftext|>


<commit_before>package rabbitgo

import (
  "time"
  //"errors"
  "github.com/streadway/amqp"
  "github.com/entropyx/rabbitgo/utils"
)

type Producer struct {
  conn  *Connection
  ch    *amqp.Channel
  e     *Exchange
  q     *Queue
  pc    *ProducerConfig
}

type ProducerConfig struct {
  // The key that when publishing a message to a exchange/queue will be only delivered to
	// given routing key listeners
	RoutingKey string
	// Publishing tagpackage
	Tag string
  // Maximum waiting time in miliseconds
  Timeout int
	// Queue should be on the server/broker
	Mandatory bool
	// Consumer should be bound to server
	Immediate bool
}

// TODO: Should we use this instead of amqp.Publishing?
type Publishing struct {
   // Application or exchange specific fields,
   // the headers exchange will inspect this field.
   // TODO: convert to amqp.Table
   Headers map[string]interface{}
   // Properties
   ContentType      string    // MIME content type
   ContentEncoding  string    // MIME content encoding
   DeliveryMode     uint8     // Transient (0 or 1) or Persistent (2)
   Priority         uint8     // 0 to 9
   CorrelationId    string    // correlation identifier
   ReplyTo          string    // address to to reply to (ex: RPC)
   Expiration       string    // message expiration spec
   MessageId        string    // message identifier
   Timestamp        time.Time // message timestamp
   Type             string    // message type name
   UserId           string    // creating user id - ex: "guest"
   AppId            string    // creating application id
   // The application specific payload of the message
   Body []byte
}

func (c *Connection) NewProducer(e *Exchange, q *Queue, pc *ProducerConfig) (*Producer, error) {
  ch, err := c.conn.Channel()
  if err != nil {
    return nil, err
  }
  return &Producer{
    conn:  c,
    ch:    ch,
    e:     e,
    q:     q,
    pc:    pc,
  }, nil
}

func (p *Producer) Publish(publishing *amqp.Publishing) error {
  routingKey := p.pc.RoutingKey
	// if exchange name is empty, this means we are gonna publish
	// this mesage to a queue, every queue has a binding to default exchange
	if p.e.Name == "" {
		routingKey = p.q.Name
	}
  err := p.ch.Publish(
		p.e.Name,       // publish to an exchange(it can be default exchange)
		routingKey,   // routing to 0 or more queues
		p.pc.Mandatory, // mandatory, if no queue than err
		p.pc.Immediate, // immediate, if no consumer than err
		*publishing,
	)
  return err
}

// PublishRPC accepts a handler function for every message streamed from RabbitMq
// as a reply after publishing a message.
func (p *Producer) PublishRPC(publishing *amqp.Publishing, handler func(delivery amqp.Delivery)) error {
  randString := utils.RandomString(35)
	queue := &Queue{
    Name: "queue_" + randString,
		AutoDelete: true,
    Exclusive: true,
	}
	consumerConfig := &ConsumerConfig{
		Tag: "consumer_" + randString,
    Timeout: p.pc.Timeout,
	}
	consumer, err := p.conn.NewConsumer(nil, queue, nil, consumerConfig)
	if err != nil {
		return err
	}
	//defer consumer.Shutdown()
  routingKey := p.pc.RoutingKey
	// if exchange name is empty, this means we are gonna publish
	// this mesage to a queue, every queue has a binding to default exchange
	if p.e.Name == "" {
		routingKey = p.q.Name
	}
  publishing.CorrelationId = randString
  publishing.ReplyTo = queue.Name
  err = p.ch.Publish(
		p.e.Name,       // publish to an exchange(it can be default exchange)
		routingKey,   // routing to 0 or more queues
		p.pc.Mandatory, // mandatory, if no queue than err
		p.pc.Immediate, // immediate, if no consumer than err
		*publishing,
	)
  err = consumer.Consume(func(d amqp.Delivery) {
    if randString == d.CorrelationId {
      handler(d)
      d.Ack(true)
      consumer.Shutdown()
    }
  })
  return err
}

func (p *Producer) Shutdown() {
  p.ch.Close()
}

// NotifyReturn captures a message when a Publishing is unable to be
// delivered either due to the `mandatory` flag set
// and no route found, or `immediate` flag set and no free consumer.
func (p *Producer) NotifyReturn(notifier func(message amqp.Return)) {
	go func() {
		for res := range p.ch.NotifyReturn(make(chan amqp.Return)) {
			notifier(res)
		}
	}()
}

func (p *Producer) NotifyPublish(confirmer func(message amqp.Confirmation))  {
  go func() {
    for res := range p.ch.NotifyPublish(make(chan amqp.Confirmation)) {
      confirmer(res)
    }
  }()
}
<commit_msg>Shutdown now returns a possible error.<commit_after>package rabbitgo

import (
  "time"
  //"errors"
  "github.com/streadway/amqp"
  "github.com/entropyx/rabbitgo/utils"
)

type Producer struct {
  conn  *Connection
  ch    *amqp.Channel
  e     *Exchange
  q     *Queue
  pc    *ProducerConfig
}

type ProducerConfig struct {
  // The key that when publishing a message to a exchange/queue will be only delivered to
	// given routing key listeners
	RoutingKey string
	// Publishing tagpackage
	Tag string
  // Maximum waiting time in miliseconds
  Timeout int
	// Queue should be on the server/broker
	Mandatory bool
	// Consumer should be bound to server
	Immediate bool
}

// TODO: Should we use this instead of amqp.Publishing?
type Publishing struct {
   // Application or exchange specific fields,
   // the headers exchange will inspect this field.
   // TODO: convert to amqp.Table
   Headers map[string]interface{}
   // Properties
   ContentType      string    // MIME content type
   ContentEncoding  string    // MIME content encoding
   DeliveryMode     uint8     // Transient (0 or 1) or Persistent (2)
   Priority         uint8     // 0 to 9
   CorrelationId    string    // correlation identifier
   ReplyTo          string    // address to to reply to (ex: RPC)
   Expiration       string    // message expiration spec
   MessageId        string    // message identifier
   Timestamp        time.Time // message timestamp
   Type             string    // message type name
   UserId           string    // creating user id - ex: "guest"
   AppId            string    // creating application id
   // The application specific payload of the message
   Body []byte
}

func (c *Connection) NewProducer(e *Exchange, q *Queue, pc *ProducerConfig) (*Producer, error) {
  ch, err := c.conn.Channel()
  if err != nil {
    return nil, err
  }
  return &Producer{
    conn:  c,
    ch:    ch,
    e:     e,
    q:     q,
    pc:    pc,
  }, nil
}

func (p *Producer) Publish(publishing *amqp.Publishing) error {
  routingKey := p.pc.RoutingKey
	// if exchange name is empty, this means we are gonna publish
	// this mesage to a queue, every queue has a binding to default exchange
	if p.e.Name == "" {
		routingKey = p.q.Name
	}
  err := p.ch.Publish(
		p.e.Name,       // publish to an exchange(it can be default exchange)
		routingKey,   // routing to 0 or more queues
		p.pc.Mandatory, // mandatory, if no queue than err
		p.pc.Immediate, // immediate, if no consumer than err
		*publishing,
	)
  return err
}

// PublishRPC accepts a handler function for every message streamed from RabbitMq
// as a reply after publishing a message.
func (p *Producer) PublishRPC(publishing *amqp.Publishing, handler func(delivery amqp.Delivery)) error {
  randString := utils.RandomString(35)
	queue := &Queue{
    Name: "queue_" + randString,
		AutoDelete: true,
    Exclusive: true,
	}
	consumerConfig := &ConsumerConfig{
		Tag: "consumer_" + randString,
    Timeout: p.pc.Timeout,
	}
	consumer, err := p.conn.NewConsumer(nil, queue, nil, consumerConfig)
	if err != nil {
		return err
	}
	//defer consumer.Shutdown()
  routingKey := p.pc.RoutingKey
	// if exchange name is empty, this means we are gonna publish
	// this mesage to a queue, every queue has a binding to default exchange
	if p.e.Name == "" {
		routingKey = p.q.Name
	}
  publishing.CorrelationId = randString
  publishing.ReplyTo = queue.Name
  err = p.ch.Publish(
		p.e.Name,       // publish to an exchange(it can be default exchange)
		routingKey,   // routing to 0 or more queues
		p.pc.Mandatory, // mandatory, if no queue than err
		p.pc.Immediate, // immediate, if no consumer than err
		*publishing,
	)
  err = consumer.Consume(func(d amqp.Delivery) {
    if randString == d.CorrelationId {
      handler(d)
      d.Ack(true)
      consumer.Shutdown()
    }
  })
  return err
}

func (p *Producer) Shutdown() error {
  if err := p.ch.Close(); err != nil {
    return err
  }
  return nil
}

// NotifyReturn captures a message when a Publishing is unable to be
// delivered either due to the `mandatory` flag set
// and no route found, or `immediate` flag set and no free consumer.
func (p *Producer) NotifyReturn(notifier func(message amqp.Return)) {
	go func() {
		for res := range p.ch.NotifyReturn(make(chan amqp.Return)) {
			notifier(res)
		}
	}()
}

func (p *Producer) NotifyPublish(confirmer func(message amqp.Confirmation))  {
  go func() {
    for res := range p.ch.NotifyPublish(make(chan amqp.Confirmation)) {
      confirmer(res)
    }
  }()
}
<|endoftext|>


<commit_before>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	prr = 150 * time.Millisecond // default RefreshRate
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	refreshCh    chan time.Time
	once         sync.Once
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// following are provided/overrided by user
	idCount          int
	reqWidth         int
	popPriority      int
	popCompleted     bool
	outputDiscarded  bool
	rr               time.Duration
	uwg              *sync.WaitGroup
	externalRefresh  <-chan interface{}
	renderDelay      <-chan struct{}
	shutdownNotifier chan struct{}
	queueBars        map[*Bar]*Bar
	output           io.Writer
	debugOut         io.Writer
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		bHeap:       priorityQueue{},
		rr:          prr,
		queueBars:   make(map[*Bar]*Bar),
		output:      os.Stdout,
		popPriority: math.MinInt32,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
	}

	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar with provided BarFillerBuilder.
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	return p.Add(total, builder.Build(), options...)
}

// Add creates a bar which renders itself by provided filler.
// If `total <= 0` triggering complete event by increment methods is disabled.
// Panics if *Progress instance is done, i.e. called after (*Progress).Wait().
func (p *Progress) Add(total int64, filler BarFiller, options ...BarOption) *Bar {
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediatly, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container.
// After this method has been called, there is no way to reuse *Progress
// instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()

	p.once.Do(p.shutdown)

	// wait for container to quit
	p.cwg.Wait()
}

func (p *Progress) shutdown() {
	close(p.done)
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	render := func(debugOut io.Writer) {
		err := s.render(cw)
		for err != nil {
			if debugOut != nil {
				_, err = fmt.Fprintln(debugOut, err)
			} else {
				panic(err)
			}
			debugOut = nil
		}
	}

	p.refreshCh = s.newTicker(p.done)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-p.refreshCh:
			render(s.debugOut)
		case <-s.shutdownNotifier:
			for s.heapUpdated {
				render(s.debugOut)
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	return s.flush(cw, height)
}

func (s *pState) flush(cw *cwriter.Writer, height int) error {
	var wg sync.WaitGroup
	var popCount int
	rows := make([]io.Reader, 0, height)
	pool := make([]*Bar, 0, s.bHeap.Len())
	for s.bHeap.Len() > 0 {
		var frameRowsUsed int
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if len(rows) < height {
				rows = append(rows, frame.rows[i])
				frameRowsUsed++
			} else {
				wg.Add(1)
				go func(discardRow io.Reader) {
					_, _ = io.Copy(io.Discard, discardRow)
					wg.Done()
				}(frame.rows[i])
			}
		}
		if frame.shutdown != 0 {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			drop := b.bs.dropOnComplete
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				pool = append(pool, qb)
				drop = true
			} else if s.popCompleted && !b.bs.noPop {
				if frame.shutdown > 1 {
					popCount += frameRowsUsed
					drop = true
				} else {
					s.popPriority++
					b.priority = s.popPriority
				}
			}
			if drop {
				s.heapUpdated = true
				continue
			}
		}
		pool = append(pool, b)
	}

	for _, b := range pool {
		heap.Push(&s.bHeap, b)
	}

	for i := len(rows) - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(rows[i])
		if err != nil {
			return err
		}
	}

	wg.Wait()
	return cw.Flush(len(rows) - popCount)
}

func (s *pState) newTicker(done <-chan struct{}) chan time.Time {
	ch := make(chan time.Time)
	if s.shutdownNotifier == nil {
		s.shutdownNotifier = make(chan struct{})
	}
	go func() {
		if s.renderDelay != nil {
			<-s.renderDelay
		}
		var internalRefresh <-chan time.Time
		if !s.outputDiscarded {
			if s.externalRefresh == nil {
				ticker := time.NewTicker(s.rr)
				defer ticker.Stop()
				internalRefresh = ticker.C
			}
		} else {
			s.externalRefresh = nil
		}
		for {
			select {
			case t := <-internalRefresh:
				ch <- t
			case x := <-s.externalRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-done:
				close(s.shutdownNotifier)
				return
			}
		}
	}()
	return ch
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:       s.idCount,
		priority: s.idCount,
		reqWidth: s.reqWidth,
		total:    total,
		filler:   filler,
		debugOut: s.debugOut,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		go maxWidthDistributor(column)
	}
}

func maxWidthDistributor(column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
}
<commit_msg>minor: var refactoring<commit_after>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	prr = 150 * time.Millisecond // default RefreshRate
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	refreshCh    chan time.Time
	once         sync.Once
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// following are provided/overrided by user
	idCount          int
	reqWidth         int
	popPriority      int
	popCompleted     bool
	outputDiscarded  bool
	rr               time.Duration
	uwg              *sync.WaitGroup
	externalRefresh  <-chan interface{}
	renderDelay      <-chan struct{}
	shutdownNotifier chan struct{}
	queueBars        map[*Bar]*Bar
	output           io.Writer
	debugOut         io.Writer
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		bHeap:       priorityQueue{},
		rr:          prr,
		queueBars:   make(map[*Bar]*Bar),
		output:      os.Stdout,
		popPriority: math.MinInt32,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
	}

	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar with provided BarFillerBuilder.
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	return p.Add(total, builder.Build(), options...)
}

// Add creates a bar which renders itself by provided filler.
// If `total <= 0` triggering complete event by increment methods is disabled.
// Panics if *Progress instance is done, i.e. called after (*Progress).Wait().
func (p *Progress) Add(total int64, filler BarFiller, options ...BarOption) *Bar {
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediatly, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container.
// After this method has been called, there is no way to reuse *Progress
// instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()

	p.once.Do(p.shutdown)

	// wait for container to quit
	p.cwg.Wait()
}

func (p *Progress) shutdown() {
	close(p.done)
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	render := func(debugOut io.Writer) {
		err := s.render(cw)
		for err != nil {
			if debugOut != nil {
				_, err = fmt.Fprintln(debugOut, err)
			} else {
				panic(err)
			}
			debugOut = nil
		}
	}

	p.refreshCh = s.newTicker(p.done)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-p.refreshCh:
			render(s.debugOut)
		case <-s.shutdownNotifier:
			for s.heapUpdated {
				render(s.debugOut)
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	return s.flush(cw, height)
}

func (s *pState) flush(cw *cwriter.Writer, height int) error {
	var wg sync.WaitGroup
	var popCount int
	rows := make([]io.Reader, 0, height)
	pool := make([]*Bar, 0, s.bHeap.Len())
	for s.bHeap.Len() > 0 {
		var usedRows int
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if len(rows) < height {
				rows = append(rows, frame.rows[i])
				usedRows++
			} else {
				wg.Add(1)
				go func(discardRow io.Reader) {
					_, _ = io.Copy(io.Discard, discardRow)
					wg.Done()
				}(frame.rows[i])
			}
		}
		if frame.shutdown != 0 {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			drop := b.bs.dropOnComplete
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				pool = append(pool, qb)
				drop = true
			} else if s.popCompleted && !b.bs.noPop {
				if frame.shutdown > 1 {
					popCount += usedRows
					drop = true
				} else {
					s.popPriority++
					b.priority = s.popPriority
				}
			}
			if drop {
				s.heapUpdated = true
				continue
			}
		}
		pool = append(pool, b)
	}

	for _, b := range pool {
		heap.Push(&s.bHeap, b)
	}

	for i := len(rows) - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(rows[i])
		if err != nil {
			return err
		}
	}

	wg.Wait()
	return cw.Flush(len(rows) - popCount)
}

func (s *pState) newTicker(done <-chan struct{}) chan time.Time {
	ch := make(chan time.Time)
	if s.shutdownNotifier == nil {
		s.shutdownNotifier = make(chan struct{})
	}
	go func() {
		if s.renderDelay != nil {
			<-s.renderDelay
		}
		var internalRefresh <-chan time.Time
		if !s.outputDiscarded {
			if s.externalRefresh == nil {
				ticker := time.NewTicker(s.rr)
				defer ticker.Stop()
				internalRefresh = ticker.C
			}
		} else {
			s.externalRefresh = nil
		}
		for {
			select {
			case t := <-internalRefresh:
				ch <- t
			case x := <-s.externalRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-done:
				close(s.shutdownNotifier)
				return
			}
		}
	}()
	return ch
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:       s.idCount,
		priority: s.idCount,
		reqWidth: s.reqWidth,
		total:    total,
		filler:   filler,
		debugOut: s.debugOut,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		go maxWidthDistributor(column)
	}
}

func maxWidthDistributor(column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
}
<|endoftext|>