<commit_before>// Package bq includes all code related to BigQuery.
//
// NB: NOTES ON MEMORY USE AND HTTP SIZE
// The bigquery library uses JSON encoding of data, which appears to be the only
// option at this time.  Furthermore, it uses intermediate data representations,
// eventually creating a map[string]Value (unless you pass that in to begin with).
// In general, when we start pumping large volumes of data, both the map and the
// JSON will cause some memory pressure, and likely pretty severe limits on the size
// of the insert we can send, likely on the order of a couple MB of actual row footprint
// in BQ.
//
// Passing in slice of structs makes memory pressure a bit worse, but probably isn't
// worth worrying about.

package bq

import (
	"log"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"

	"cloud.google.com/go/bigquery"
	"golang.org/x/net/context"

	"github.com/m-lab/etl/etl"
	"github.com/m-lab/etl/metrics"
)

// TODO(dev) Use a more thoughtful setting for buffer size.
// TODO(P3) Include the project name in the parameters.
// For now, 10K per row times 100 results is 1MB, which is an order of
// magnitude below our 10MB max, so 100 might not be such a bad
// default.
func NewInserter(dataset string, dt etl.DataType, partition time.Time) (etl.Inserter, error) {
	suffix := ""
	table := etl.DataTypeToTable[dt]
	if time.Since(partition) < 30*24*time.Hour {
		// If within past 30 days, we can stream directly to partition.
		suffix = "$" + partition.Format("20060102")
	} else {
		// Otherwise, we use a templated table, and must merge it later.
		suffix = "_" + partition.Format("20060102")
	}

	return NewBQInserter(
		etl.InserterParams{Dataset: dataset, Table: table, Suffix: suffix,
			Timeout: 15 * time.Minute, BufferSize: etl.DataTypeToBQBufferSize[dt], RetryDelay: 30 * time.Second},
		nil)
}

// TODO - improve the naming between here and NewInserter.
// Pass in nil uploader for normal use, custom uploader for custom behavior
func NewBQInserter(params etl.InserterParams, uploader etl.Uploader) (etl.Inserter, error) {
	if uploader == nil {
		client := MustGetClient(params.Timeout)
		table := params.Table
		if params.Suffix[0] == '$' {
			// Suffix starting with $ is just a partition spec.
			table += params.Suffix
		}
		u := client.Dataset(params.Dataset).Table(table).Uploader()
		if params.Suffix[0] == '_' {
			// Suffix starting with _ is a template suffix.
			u.TableTemplateSuffix = params.Suffix
		}
		// This avoids problems when a single row of the insert has invalid
		// data.  We then have to carefully parse the returned error object.
		u.SkipInvalidRows = true
		uploader = u
	}
	in := BQInserter{params: params, uploader: uploader, timeout: params.Timeout}
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return &in, nil
}

//===============================================================================
var (
	clientOnce sync.Once // This avoids a race on setting bqClient.
	bqClient   *bigquery.Client
)

// Returns the Singleton bigquery client for this process.
// TODO - is there any advantage to using more than one client?
func MustGetClient(timeout time.Duration) *bigquery.Client {
	// We do this here, instead of in init(), because we only want to do it
	// when we actually want to access the bigquery backend.
	clientOnce.Do(func() {
		ctx, _ := context.WithTimeout(context.Background(), timeout)
		project, ok := os.LookupEnv("BIGQUERY_PROJECT")
		if !ok {
			project = os.Getenv("GCLOUD_PROJECT")
		}

		log.Printf("Using project: %s\n", project)
		// Heavyweight!
		var err error
		bqClient, err = bigquery.NewClient(ctx, project)
		if err != nil {
			panic(err.Error())
		}
	})
	return bqClient
}

//===============================================================================

// Generic implementation of bq.ValueSaver, based on map.  This avoids extra
// conversion steps in the bigquery library (except for the JSON conversion).
// IMPLEMENTS: bigquery.ValueSaver
type MapSaver struct {
	Values map[string]bigquery.Value
}

func (s *MapSaver) Save() (row map[string]bigquery.Value, insertID string, err error) {
	return s.Values, "", nil
}

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

type BQInserter struct {
	etl.Inserter
	params   etl.InserterParams
	uploader etl.Uploader // May be a BQ Uploader, or a test Uploader
	timeout  time.Duration
	rows     []interface{}
	inserted int // Number of rows successfully inserted.
	badRows  int // Number of row failures, including rows in full failures.
	failures int // Number of complete insert failures.
}

// Caller should check error, and take appropriate action before calling again.
func (in *BQInserter) InsertRow(data interface{}) error {
	return in.InsertRows([]interface{}{data})
}

// Caller should check error, and take appropriate action before calling again.
// TODO - should this return a specific error to indicate that a flush is needed
// instead of flushing internally?  The "handle errors in the middle" would
// be easier, though other complications would ensue.
func (in *BQInserter) InsertRows(data []interface{}) error {
	metrics.WorkerState.WithLabelValues("insert").Inc()
	defer metrics.WorkerState.WithLabelValues("insert").Dec()

	for len(data)+len(in.rows) >= in.params.BufferSize {
		// space >= len(data)
		space := cap(in.rows) - len(in.rows)
		var add []interface{}
		add, data = data[:space], data[space:] // does this break?
		in.rows = append(in.rows, add...)
		err := in.Flush()
		if err != nil {
			// TODO - handle errors in middle better?
			return err
		}
	}
	in.rows = append(in.rows, data...)
	return nil
}

func (in *BQInserter) HandleInsertErrors(err error) error {
	switch typedErr := err.(type) {
	case bigquery.PutMultiError:
		if len(typedErr) == len(in.rows) {
			log.Printf("%v\n", err)
			metrics.BackendFailureCount.WithLabelValues(
				in.TableBase(), "failed insert").Inc()
			in.failures += 1
		}
		// If ALL rows failed, and number of rows is large, just report single failure.
		if len(typedErr) > 10 && len(typedErr) == len(in.rows) {
			log.Printf("Insert error: %v\n", err)
			metrics.ErrorCount.WithLabelValues(
				in.TableBase(), "unknown", "insert row error").
				Add(float64(len(typedErr)))
		} else {
			// Handle each error individually.
			// TODO Should we try to handle large numbers of row errors?
			for _, rowError := range typedErr {
				// These are rowInsertionErrors
				log.Printf("%v\n", rowError)
				// rowError.Errors is a MultiError
				for _, oneErr := range rowError.Errors {
					log.Printf("Insert error: %v\n", oneErr)
					metrics.ErrorCount.WithLabelValues(
						in.TableBase(), "unknown", "insert row error").Inc()
				}
			}
		}
		in.inserted += len(in.rows) - len(typedErr)
		in.badRows += len(typedErr)
		err = nil
	default:
		log.Printf("Unhandled insert error %v\n", typedErr)
		metrics.BackendFailureCount.WithLabelValues(
			in.TableBase(), "failed insert").Inc()
		metrics.ErrorCount.WithLabelValues(
			in.TableBase(), "unknown", "UNHANDLED insert error").Inc()
		// TODO - Conservative, but possibly not correct.
		// This at least preserves the count invariance.
		in.badRows += len(in.rows)
		err = nil
	}
	return err
}

// TODO(dev) Should have a recovery mechanism for failed inserts.
func (in *BQInserter) Flush() error {
	metrics.WorkerState.WithLabelValues("flush").Inc()
	defer metrics.WorkerState.WithLabelValues("flush").Dec()

	if len(in.rows) == 0 {
		return nil
	}

	// If we exceed the quota, this basically backs off and tries again.  When
	// operating near quota, this will fire enough times to slow down each task
	// enough to stay within quota.  It may result in AppEngine reducing the
	// number of workers, but that is fine - it will also result in staying
	// under the quota.
	// Analysis:
	//  We can handle a minimum of 10 inserts per second, because
	//   the default quota is 100MB/sec, and the limit on requests
	//   is 10MB per request.  Since generally inserts are smaller,
	//   the typical number is more like 20 inserts/sec.
	//  The net effect we need to see is that, if the pipeline capacity
	//   exceeds the quota by 10%, then the pipeline needs to slow down
	//   by roughly 10% to fit within the quota.  The incoming request
	//   rate is dictated by the task queue, and ultimately the handler
	//   must reject 10% of the incoming requests.  This only happens
	//   when 10% of the instances have hit MAX_WORKERS.
	//  If the capacity of the pipeline is, e.g., 2X the task queue rate,
	//   then each task will need to be slowed down to the point that it
	//   takes roughly 2.2X longer than it would without any Quota exceeded
	//   errors.  For NDT, the 100MB tasks require about 35 concurrent tasks
	//   to process 60 tasks/min, indicating that they require about 35
	//   seconds per task.  There are about 70 tests/task, so this is about
	//   7 buffer flushes per second (of 10 tests each), or on average, about
	//   one buffer flush every 5 seconds for each task.
	//  The batch job might have 50 instances, and process 900 tasks
	//   concurrently.  If this had to be scaled back to 50%, the tasks
	//   would have to spend 50% of their time sleeping between Put requests.
	//   Since each task typically takes about 35 seconds, each task would
	//   on average experience just over one 'Quota exceeded' error in order
	//   to slow the pipeline down by 50%.
	//  Note that a secondary effect is to reduce the CPU utilization, which
	//   will likely trigger a reduction in the number of instances running.
	//   Under these conditions, AppEngine would reduce the number of instances
	//   until the target utilization is reaches, reducing the number of
	//   concurrent tasks, and thus the frequency at which the tasks would
	//   experience 'Quota error' events.

	var err error
	for i := 0; i < 10; i++ {
		// This is heavyweight, and may run forever without a context deadline.
		ctx, _ := context.WithTimeout(context.Background(), in.timeout)
		err = in.uploader.Put(ctx, in.rows)
		log.Println(err)
		if err == nil || !strings.Contains(err.Error(), "Quota exceeded:") {
			break
		}
		metrics.WarningCount.WithLabelValues(in.TableBase(), "", "Quota Exceeded").Inc()
		// Use some randomness to reduce risk of synchronization across tasks.
		time.Sleep(time.Seconds((0.5 + rand.Float64()) * in.params.RetryDelay.Seconds()))
	}

	// If there is still an error, then handle it.
	if err == nil {
		in.inserted += len(in.rows)
	} else {
		// This adjusts the inserted count, failure count, and updates in.rows.
		err = in.HandleInsertErrors(err)
	}
	// Allocate new slice of rows.  Any failed rows are lost.
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return err
}

func (in *BQInserter) FullTableName() string {
	return in.TableBase() + in.TableSuffix()
}
func (in *BQInserter) TableBase() string {
	return in.params.Table
}

// The $ or _ suffix.
func (in *BQInserter) TableSuffix() string {
	return in.params.Suffix
}
func (in *BQInserter) Dataset() string {
	return in.params.Dataset
}
func (in *BQInserter) RowsInBuffer() int {
	return len(in.rows)
}
func (in *BQInserter) Accepted() int {
	return in.inserted + in.badRows + len(in.rows)
}
func (in *BQInserter) Committed() int {
	return in.inserted
}
func (in *BQInserter) Failed() int {
	return in.badRows
}

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

// Inserter wrapper that handles flush metrics.
// TODO - add prometheus counters for attempts, number of rows.
type DurationWrapper struct {
	etl.Inserter
}

func (dw DurationWrapper) Flush() error {
	t := time.Now()
	status := "succeed"
	err := dw.Inserter.Flush()
	if err != nil {
		status = "fail"
	}
	metrics.InsertionHistogram.WithLabelValues(
		dw.TableBase(), status).Observe(time.Since(t).Seconds())
	return err
}
<commit_msg>Fix time arithmetic<commit_after>// Package bq includes all code related to BigQuery.
//
// NB: NOTES ON MEMORY USE AND HTTP SIZE
// The bigquery library uses JSON encoding of data, which appears to be the only
// option at this time.  Furthermore, it uses intermediate data representations,
// eventually creating a map[string]Value (unless you pass that in to begin with).
// In general, when we start pumping large volumes of data, both the map and the
// JSON will cause some memory pressure, and likely pretty severe limits on the size
// of the insert we can send, likely on the order of a couple MB of actual row footprint
// in BQ.
//
// Passing in slice of structs makes memory pressure a bit worse, but probably isn't
// worth worrying about.

package bq

import (
	"log"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"

	"cloud.google.com/go/bigquery"
	"golang.org/x/net/context"

	"github.com/m-lab/etl/etl"
	"github.com/m-lab/etl/metrics"
)

// TODO(dev) Use a more thoughtful setting for buffer size.
// TODO(P3) Include the project name in the parameters.
// For now, 10K per row times 100 results is 1MB, which is an order of
// magnitude below our 10MB max, so 100 might not be such a bad
// default.
func NewInserter(dataset string, dt etl.DataType, partition time.Time) (etl.Inserter, error) {
	suffix := ""
	table := etl.DataTypeToTable[dt]
	if time.Since(partition) < 30*24*time.Hour {
		// If within past 30 days, we can stream directly to partition.
		suffix = "$" + partition.Format("20060102")
	} else {
		// Otherwise, we use a templated table, and must merge it later.
		suffix = "_" + partition.Format("20060102")
	}

	return NewBQInserter(
		etl.InserterParams{Dataset: dataset, Table: table, Suffix: suffix,
			Timeout: 15 * time.Minute, BufferSize: etl.DataTypeToBQBufferSize[dt], RetryDelay: 30 * time.Second},
		nil)
}

// TODO - improve the naming between here and NewInserter.
// Pass in nil uploader for normal use, custom uploader for custom behavior
func NewBQInserter(params etl.InserterParams, uploader etl.Uploader) (etl.Inserter, error) {
	if uploader == nil {
		client := MustGetClient(params.Timeout)
		table := params.Table
		if params.Suffix[0] == '$' {
			// Suffix starting with $ is just a partition spec.
			table += params.Suffix
		}
		u := client.Dataset(params.Dataset).Table(table).Uploader()
		if params.Suffix[0] == '_' {
			// Suffix starting with _ is a template suffix.
			u.TableTemplateSuffix = params.Suffix
		}
		// This avoids problems when a single row of the insert has invalid
		// data.  We then have to carefully parse the returned error object.
		u.SkipInvalidRows = true
		uploader = u
	}
	in := BQInserter{params: params, uploader: uploader, timeout: params.Timeout}
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return &in, nil
}

//===============================================================================
var (
	clientOnce sync.Once // This avoids a race on setting bqClient.
	bqClient   *bigquery.Client
)

// Returns the Singleton bigquery client for this process.
// TODO - is there any advantage to using more than one client?
func MustGetClient(timeout time.Duration) *bigquery.Client {
	// We do this here, instead of in init(), because we only want to do it
	// when we actually want to access the bigquery backend.
	clientOnce.Do(func() {
		ctx, _ := context.WithTimeout(context.Background(), timeout)
		project, ok := os.LookupEnv("BIGQUERY_PROJECT")
		if !ok {
			project = os.Getenv("GCLOUD_PROJECT")
		}

		log.Printf("Using project: %s\n", project)
		// Heavyweight!
		var err error
		bqClient, err = bigquery.NewClient(ctx, project)
		if err != nil {
			panic(err.Error())
		}
	})
	return bqClient
}

//===============================================================================

// Generic implementation of bq.ValueSaver, based on map.  This avoids extra
// conversion steps in the bigquery library (except for the JSON conversion).
// IMPLEMENTS: bigquery.ValueSaver
type MapSaver struct {
	Values map[string]bigquery.Value
}

func (s *MapSaver) Save() (row map[string]bigquery.Value, insertID string, err error) {
	return s.Values, "", nil
}

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

type BQInserter struct {
	etl.Inserter
	params   etl.InserterParams
	uploader etl.Uploader // May be a BQ Uploader, or a test Uploader
	timeout  time.Duration
	rows     []interface{}
	inserted int // Number of rows successfully inserted.
	badRows  int // Number of row failures, including rows in full failures.
	failures int // Number of complete insert failures.
}

// Caller should check error, and take appropriate action before calling again.
func (in *BQInserter) InsertRow(data interface{}) error {
	return in.InsertRows([]interface{}{data})
}

// Caller should check error, and take appropriate action before calling again.
// TODO - should this return a specific error to indicate that a flush is needed
// instead of flushing internally?  The "handle errors in the middle" would
// be easier, though other complications would ensue.
func (in *BQInserter) InsertRows(data []interface{}) error {
	metrics.WorkerState.WithLabelValues("insert").Inc()
	defer metrics.WorkerState.WithLabelValues("insert").Dec()

	for len(data)+len(in.rows) >= in.params.BufferSize {
		// space >= len(data)
		space := cap(in.rows) - len(in.rows)
		var add []interface{}
		add, data = data[:space], data[space:] // does this break?
		in.rows = append(in.rows, add...)
		err := in.Flush()
		if err != nil {
			// TODO - handle errors in middle better?
			return err
		}
	}
	in.rows = append(in.rows, data...)
	return nil
}

func (in *BQInserter) HandleInsertErrors(err error) error {
	switch typedErr := err.(type) {
	case bigquery.PutMultiError:
		if len(typedErr) == len(in.rows) {
			log.Printf("%v\n", err)
			metrics.BackendFailureCount.WithLabelValues(
				in.TableBase(), "failed insert").Inc()
			in.failures += 1
		}
		// If ALL rows failed, and number of rows is large, just report single failure.
		if len(typedErr) > 10 && len(typedErr) == len(in.rows) {
			log.Printf("Insert error: %v\n", err)
			metrics.ErrorCount.WithLabelValues(
				in.TableBase(), "unknown", "insert row error").
				Add(float64(len(typedErr)))
		} else {
			// Handle each error individually.
			// TODO Should we try to handle large numbers of row errors?
			for _, rowError := range typedErr {
				// These are rowInsertionErrors
				log.Printf("%v\n", rowError)
				// rowError.Errors is a MultiError
				for _, oneErr := range rowError.Errors {
					log.Printf("Insert error: %v\n", oneErr)
					metrics.ErrorCount.WithLabelValues(
						in.TableBase(), "unknown", "insert row error").Inc()
				}
			}
		}
		in.inserted += len(in.rows) - len(typedErr)
		in.badRows += len(typedErr)
		err = nil
	default:
		log.Printf("Unhandled insert error %v\n", typedErr)
		metrics.BackendFailureCount.WithLabelValues(
			in.TableBase(), "failed insert").Inc()
		metrics.ErrorCount.WithLabelValues(
			in.TableBase(), "unknown", "UNHANDLED insert error").Inc()
		// TODO - Conservative, but possibly not correct.
		// This at least preserves the count invariance.
		in.badRows += len(in.rows)
		err = nil
	}
	return err
}

// TODO(dev) Should have a recovery mechanism for failed inserts.
func (in *BQInserter) Flush() error {
	metrics.WorkerState.WithLabelValues("flush").Inc()
	defer metrics.WorkerState.WithLabelValues("flush").Dec()

	if len(in.rows) == 0 {
		return nil
	}

	// If we exceed the quota, this basically backs off and tries again.  When
	// operating near quota, this will fire enough times to slow down each task
	// enough to stay within quota.  It may result in AppEngine reducing the
	// number of workers, but that is fine - it will also result in staying
	// under the quota.
	// Analysis:
	//  We can handle a minimum of 10 inserts per second, because
	//   the default quota is 100MB/sec, and the limit on requests
	//   is 10MB per request.  Since generally inserts are smaller,
	//   the typical number is more like 20 inserts/sec.
	//  The net effect we need to see is that, if the pipeline capacity
	//   exceeds the quota by 10%, then the pipeline needs to slow down
	//   by roughly 10% to fit within the quota.  The incoming request
	//   rate is dictated by the task queue, and ultimately the handler
	//   must reject 10% of the incoming requests.  This only happens
	//   when 10% of the instances have hit MAX_WORKERS.
	//  If the capacity of the pipeline is, e.g., 2X the task queue rate,
	//   then each task will need to be slowed down to the point that it
	//   takes roughly 2.2X longer than it would without any Quota exceeded
	//   errors.  For NDT, the 100MB tasks require about 35 concurrent tasks
	//   to process 60 tasks/min, indicating that they require about 35
	//   seconds per task.  There are about 70 tests/task, so this is about
	//   7 buffer flushes per second (of 10 tests each), or on average, about
	//   one buffer flush every 5 seconds for each task.
	//  The batch job might have 50 instances, and process 900 tasks
	//   concurrently.  If this had to be scaled back to 50%, the tasks
	//   would have to spend 50% of their time sleeping between Put requests.
	//   Since each task typically takes about 35 seconds, each task would
	//   on average experience just over one 'Quota exceeded' error in order
	//   to slow the pipeline down by 50%.
	//  Note that a secondary effect is to reduce the CPU utilization, which
	//   will likely trigger a reduction in the number of instances running.
	//   Under these conditions, AppEngine would reduce the number of instances
	//   until the target utilization is reaches, reducing the number of
	//   concurrent tasks, and thus the frequency at which the tasks would
	//   experience 'Quota error' events.

	var err error
	for i := 0; i < 10; i++ {
		// This is heavyweight, and may run forever without a context deadline.
		ctx, _ := context.WithTimeout(context.Background(), in.timeout)
		err = in.uploader.Put(ctx, in.rows)
		log.Println(err)
		if err == nil || !strings.Contains(err.Error(), "Quota exceeded:") {
			break
		}
		metrics.WarningCount.WithLabelValues(in.TableBase(), "", "Quota Exceeded").Inc()
		// Use some randomness to reduce risk of synchronization across tasks.
		t := in.params.RetryDelay.Seconds() * (0.5 + rand.Float64()) // between 0.5 and 1.5 * RetryDelay
		time.Sleep(time.Duration(1000000*t) * time.Microsecond)
	}

	// If there is still an error, then handle it.
	if err == nil {
		in.inserted += len(in.rows)
	} else {
		// This adjusts the inserted count, failure count, and updates in.rows.
		err = in.HandleInsertErrors(err)
	}
	// Allocate new slice of rows.  Any failed rows are lost.
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return err
}

func (in *BQInserter) FullTableName() string {
	return in.TableBase() + in.TableSuffix()
}
func (in *BQInserter) TableBase() string {
	return in.params.Table
}

// The $ or _ suffix.
func (in *BQInserter) TableSuffix() string {
	return in.params.Suffix
}
func (in *BQInserter) Dataset() string {
	return in.params.Dataset
}
func (in *BQInserter) RowsInBuffer() int {
	return len(in.rows)
}
func (in *BQInserter) Accepted() int {
	return in.inserted + in.badRows + len(in.rows)
}
func (in *BQInserter) Committed() int {
	return in.inserted
}
func (in *BQInserter) Failed() int {
	return in.badRows
}

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

// Inserter wrapper that handles flush metrics.
// TODO - add prometheus counters for attempts, number of rows.
type DurationWrapper struct {
	etl.Inserter
}

func (dw DurationWrapper) Flush() error {
	t := time.Now()
	status := "succeed"
	err := dw.Inserter.Flush()
	if err != nil {
		status = "fail"
	}
	metrics.InsertionHistogram.WithLabelValues(
		dw.TableBase(), status).Observe(time.Since(t).Seconds())
	return err
}
<|endoftext|>


<commit_before>package connmgr

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	ifconnmgr "github.com/libp2p/go-libp2p-interface-connmgr"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

var SilencePeriod = 10 * time.Second

var log = logging.Logger("connmgr")

// BasicConnMgr is a ConnManager that trims connections whenever the count exceeds the
// high watermark. New connections are given a grace period before they're subject
// to trimming. Trims are automatically run on demand, only if the time from the
// previous trim is higher than 10 seconds. Furthermore, trims can be explicitly
// requested through the public interface of this struct (see TrimOpenConns).
//
// See configuration parameters in NewConnManager.
type BasicConnMgr struct {
	lk          sync.Mutex
	highWater   int
	lowWater    int
	connCount   int
	gracePeriod time.Duration
	peers       map[peer.ID]*peerInfo

	plk       sync.RWMutex
	protected map[peer.ID]map[string]struct{}

	// channel-based semaphore that enforces only a single trim is in progress
	trimRunningCh chan struct{}
	lastTrim      time.Time
	silencePeriod time.Duration
}

var _ ifconnmgr.ConnManager = (*BasicConnMgr)(nil)

// NewConnManager creates a new BasicConnMgr with the provided params:
// * lo and hi are watermarks governing the number of connections that'll be maintained.
//   When the peer count exceeds the 'high watermark', as many peers will be pruned (and
//   their connections terminated) until 'low watermark' peers remain.
// * grace is the amount of time a newly opened connection is given before it becomes
//   subject to pruning.
func NewConnManager(low, hi int, grace time.Duration) *BasicConnMgr {
	return &BasicConnMgr{
		highWater:     hi,
		lowWater:      low,
		gracePeriod:   grace,
		peers:         make(map[peer.ID]*peerInfo),
		trimRunningCh: make(chan struct{}, 1),
		protected:     make(map[peer.ID]map[string]struct{}, 16),
		silencePeriod: SilencePeriod,
	}
}

func (cm *BasicConnMgr) Protect(id peer.ID, tag string) {
	cm.plk.Lock()
	defer cm.plk.Unlock()

	tags, ok := cm.protected[id]
	if !ok {
		tags = make(map[string]struct{}, 2)
		cm.protected[id] = tags
	}
	tags[tag] = struct{}{}
}

func (cm *BasicConnMgr) Unprotect(id peer.ID, tag string) (protected bool) {
	cm.plk.Lock()
	defer cm.plk.Unlock()

	tags, ok := cm.protected[id]
	if !ok {
		return false
	}
	if delete(tags, tag); len(tags) == 0 {
		delete(cm.protected, id)
		return false
	}
	return true
}

// peerInfo stores metadata for a given peer.
type peerInfo struct {
	id    peer.ID
	tags  map[string]int // value for each tag
	value int            // cached sum of all tag values

	conns map[inet.Conn]time.Time // start time of each connection

	firstSeen time.Time // timestamp when we began tracking this peer.
}

// TrimOpenConns closes the connections of as many peers as needed to make the peer count
// equal the low watermark. Peers are sorted in ascending order based on their total value,
// pruning those peers with the lowest scores first, as long as they are not within their
// grace period.
//
// TODO: error return value so we can cleanly signal we are aborting because:
// (a) there's another trim in progress, or (b) the silence period is in effect.
func (cm *BasicConnMgr) TrimOpenConns(ctx context.Context) {
	select {
	case cm.trimRunningCh <- struct{}{}:
	default:
		return
	}
	defer func() { <-cm.trimRunningCh }()
	if time.Since(cm.lastTrim) < cm.silencePeriod {
		// skip this attempt to trim as the last one just took place.
		return
	}
	defer log.EventBegin(ctx, "connCleanup").Done()
	for _, c := range cm.getConnsToClose(ctx) {
		log.Info("closing conn: ", c.RemotePeer())
		log.Event(ctx, "closeConn", c.RemotePeer())
		c.Close()
	}
	cm.lastTrim = time.Now()
}

// getConnsToClose runs the heuristics described in TrimOpenConns and returns the
// connections to close.
func (cm *BasicConnMgr) getConnsToClose(ctx context.Context) []inet.Conn {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	if cm.lowWater == 0 || cm.highWater == 0 {
		// disabled
		return nil
	}
	now := time.Now()
	if len(cm.peers) < cm.lowWater {
		log.Info("open connection count below limit")
		return nil
	}

	var candidates []*peerInfo
	cm.plk.RLock()
	for id, inf := range cm.peers {
		if _, ok := cm.protected[id]; ok {
			// skip over protected peer.
			continue
		}
		candidates = append(candidates, inf)
	}
	cm.plk.RUnlock()

	// Sort peers according to their value.
	sort.Slice(candidates, func(i, j int) bool {
		return candidates[i].value < candidates[j].value
	})

	target := len(cm.peers) - cm.lowWater

	// 2x number of peers we're disconnecting from because we may have more
	// than one connection per peer. Slightly over allocating isn't an issue
	// as this is a very short-lived array.
	selected := make([]inet.Conn, 0, target*2)

	for _, inf := range candidates {
		// TODO: should we be using firstSeen or the time associated with the connection itself?
		if inf.firstSeen.Add(cm.gracePeriod).After(now) {
			continue
		}

		for c := range inf.conns {
			selected = append(selected, c)
		}

		target--
		if target == 0 {
			break
		}
	}

	return selected
}

// GetTagInfo is called to fetch the tag information associated with a given
// peer, nil is returned if p refers to an unknown peer.
func (cm *BasicConnMgr) GetTagInfo(p peer.ID) *ifconnmgr.TagInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		return nil
	}

	out := &ifconnmgr.TagInfo{
		FirstSeen: pi.firstSeen,
		Value:     pi.value,
		Tags:      make(map[string]int),
		Conns:     make(map[string]time.Time),
	}

	for t, v := range pi.tags {
		out.Tags[t] = v
	}
	for c, t := range pi.conns {
		out.Conns[c.RemoteMultiaddr().String()] = t
	}

	return out
}

// TagPeer is called to associate a string and integer with a given peer.
func (cm *BasicConnMgr) TagPeer(p peer.ID, tag string, val int) {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		log.Info("tried to tag conn from untracked peer: ", p)
		return
	}

	// Update the total value of the peer.
	pi.value += (val - pi.tags[tag])
	pi.tags[tag] = val
}

// UntagPeer is called to disassociate a string and integer from a given peer.
func (cm *BasicConnMgr) UntagPeer(p peer.ID, tag string) {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		log.Info("tried to remove tag from untracked peer: ", p)
		return
	}

	// Update the total value of the peer.
	pi.value -= pi.tags[tag]
	delete(pi.tags, tag)
}

// UpsertTag is called to insert/update a peer tag
func (cm *BasicConnMgr) UpsertTag(p peer.ID, tag string, upsert func(int) int) {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		log.Info("tried to upsert tag from untracked peer: ", p)
		return
	}

	oldval := pi.tags[tag]
	newval := upsert(oldval)
	pi.value += (newval - oldval)
	pi.tags[tag] = newval
}

// CMInfo holds the configuration for BasicConnMgr, as well as status data.
type CMInfo struct {
	// The low watermark, as described in NewConnManager.
	LowWater int

	// The high watermark, as described in NewConnManager.
	HighWater int

	// The timestamp when the last trim was triggered.
	LastTrim time.Time

	// The configured grace period, as described in NewConnManager.
	GracePeriod time.Duration

	// The current connection count.
	ConnCount int
}

// GetInfo returns the configuration and status data for this connection manager.
func (cm *BasicConnMgr) GetInfo() CMInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	return CMInfo{
		HighWater:   cm.highWater,
		LowWater:    cm.lowWater,
		LastTrim:    cm.lastTrim,
		GracePeriod: cm.gracePeriod,
		ConnCount:   cm.connCount,
	}
}

// Notifee returns a sink through which Notifiers can inform the BasicConnMgr when
// events occur. Currently, the notifee only reacts upon connection events
// {Connected, Disconnected}.
func (cm *BasicConnMgr) Notifee() inet.Notifiee {
	return (*cmNotifee)(cm)
}

type cmNotifee BasicConnMgr

func (nn *cmNotifee) cm() *BasicConnMgr {
	return (*BasicConnMgr)(nn)
}

// Connected is called by notifiers to inform that a new connection has been established.
// The notifee updates the BasicConnMgr to start tracking the connection. If the new connection
// count exceeds the high watermark, a trim may be triggered.
func (nn *cmNotifee) Connected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	cm.lk.Lock()
	defer cm.lk.Unlock()

	pinfo, ok := cm.peers[c.RemotePeer()]
	if !ok {
		pinfo = &peerInfo{
			firstSeen: time.Now(),
			tags:      make(map[string]int),
			conns:     make(map[inet.Conn]time.Time),
		}
		cm.peers[c.RemotePeer()] = pinfo
	}

	_, ok = pinfo.conns[c]
	if ok {
		log.Error("received connected notification for conn we are already tracking: ", c.RemotePeer())
		return
	}

	pinfo.conns[c] = time.Now()
	cm.connCount++

	if cm.connCount > nn.highWater {
		go cm.TrimOpenConns(context.Background())
	}
}

// Disconnected is called by notifiers to inform that an existing connection has been closed or terminated.
// The notifee updates the BasicConnMgr accordingly to stop tracking the connection, and performs housekeeping.
func (nn *cmNotifee) Disconnected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	cm.lk.Lock()
	defer cm.lk.Unlock()

	cinf, ok := cm.peers[c.RemotePeer()]
	if !ok {
		log.Error("received disconnected notification for peer we are not tracking: ", c.RemotePeer())
		return
	}

	_, ok = cinf.conns[c]
	if !ok {
		log.Error("received disconnected notification for conn we are not tracking: ", c.RemotePeer())
		return
	}

	delete(cinf.conns, c)
	cm.connCount--
	if len(cinf.conns) == 0 {
		delete(cm.peers, c.RemotePeer())
	}
}

// Listen is no-op in this implementation.
func (nn *cmNotifee) Listen(n inet.Network, addr ma.Multiaddr) {}

// ListenClose is no-op in this implementation.
func (nn *cmNotifee) ListenClose(n inet.Network, addr ma.Multiaddr) {}

// OpenedStream is no-op in this implementation.
func (nn *cmNotifee) OpenedStream(inet.Network, inet.Stream) {}

// ClosedStream is no-op in this implementation.
func (nn *cmNotifee) ClosedStream(inet.Network, inet.Stream) {}
<commit_msg>implement segmented lock to eliminate big lock contention<commit_after>package connmgr

import (
	"context"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	logging "github.com/ipfs/go-log"
	ifconnmgr "github.com/libp2p/go-libp2p-interface-connmgr"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

var SilencePeriod = 10 * time.Second

var log = logging.Logger("connmgr")

// BasicConnMgr is a ConnManager that trims connections whenever the count exceeds the
// high watermark. New connections are given a grace period before they're subject
// to trimming. Trims are automatically run on demand, only if the time from the
// previous trim is higher than 10 seconds. Furthermore, trims can be explicitly
// requested through the public interface of this struct (see TrimOpenConns).
//
// See configuration parameters in NewConnManager.
type BasicConnMgr struct {
	highWater   int
	lowWater    int
	connCount   int32
	gracePeriod time.Duration
	segments    segments

	plk       sync.RWMutex
	protected map[peer.ID]map[string]struct{}

	// channel-based semaphore that enforces only a single trim is in progress
	trimRunningCh chan struct{}
	lastTrim      time.Time
	silencePeriod time.Duration
}

var _ ifconnmgr.ConnManager = (*BasicConnMgr)(nil)

type segment struct {
	sync.Mutex
	peers map[peer.ID]*peerInfo
}

type segments [256]*segment

func (s *segments) get(id peer.ID) *segment {
	b := []byte(id)
	return s[b[len(b)-1]]
}

func (s *segments) countPeers() (count int) {
	for _, seg := range s {
		seg.Lock()
		count += len(seg.peers)
		seg.Unlock()
	}
	return count
}

// NewConnManager creates a new BasicConnMgr with the provided params:
// * lo and hi are watermarks governing the number of connections that'll be maintained.
//   When the peer count exceeds the 'high watermark', as many peers will be pruned (and
//   their connections terminated) until 'low watermark' peers remain.
// * grace is the amount of time a newly opened connection is given before it becomes
//   subject to pruning.
func NewConnManager(low, hi int, grace time.Duration) *BasicConnMgr {
	return &BasicConnMgr{
		highWater:     hi,
		lowWater:      low,
		gracePeriod:   grace,
		trimRunningCh: make(chan struct{}, 1),
		protected:     make(map[peer.ID]map[string]struct{}, 16),
		silencePeriod: SilencePeriod,
		segments: func() (ret segments) {
			for i := range ret {
				ret[i] = &segment{
					peers: make(map[peer.ID]*peerInfo),
				}
			}
			return ret
		}(),
	}
}

func (cm *BasicConnMgr) Protect(id peer.ID, tag string) {
	cm.plk.Lock()
	defer cm.plk.Unlock()

	tags, ok := cm.protected[id]
	if !ok {
		tags = make(map[string]struct{}, 2)
		cm.protected[id] = tags
	}
	tags[tag] = struct{}{}
}

func (cm *BasicConnMgr) Unprotect(id peer.ID, tag string) (protected bool) {
	cm.plk.Lock()
	defer cm.plk.Unlock()

	tags, ok := cm.protected[id]
	if !ok {
		return false
	}
	if delete(tags, tag); len(tags) == 0 {
		delete(cm.protected, id)
		return false
	}
	return true
}

// peerInfo stores metadata for a given peer.
type peerInfo struct {
	id    peer.ID
	tags  map[string]int // value for each tag
	value int            // cached sum of all tag values

	conns map[inet.Conn]time.Time // start time of each connection

	firstSeen time.Time // timestamp when we began tracking this peer.
}

// TrimOpenConns closes the connections of as many peers as needed to make the peer count
// equal the low watermark. Peers are sorted in ascending order based on their total value,
// pruning those peers with the lowest scores first, as long as they are not within their
// grace period.
//
// TODO: error return value so we can cleanly signal we are aborting because:
// (a) there's another trim in progress, or (b) the silence period is in effect.
func (cm *BasicConnMgr) TrimOpenConns(ctx context.Context) {
	select {
	case cm.trimRunningCh <- struct{}{}:
	default:
		return
	}
	defer func() { <-cm.trimRunningCh }()
	if time.Since(cm.lastTrim) < cm.silencePeriod {
		// skip this attempt to trim as the last one just took place.
		return
	}
	defer log.EventBegin(ctx, "connCleanup").Done()
	for _, c := range cm.getConnsToClose(ctx) {
		log.Info("closing conn: ", c.RemotePeer())
		log.Event(ctx, "closeConn", c.RemotePeer())
		c.Close()
	}
	cm.lastTrim = time.Now()
}

// getConnsToClose runs the heuristics described in TrimOpenConns and returns the
// connections to close.
func (cm *BasicConnMgr) getConnsToClose(ctx context.Context) []inet.Conn {
	if cm.lowWater == 0 || cm.highWater == 0 {
		// disabled
		return nil
	}
	now := time.Now()
	npeers := cm.segments.countPeers()
	if npeers < cm.lowWater {
		log.Info("open connection count below limit")
		return nil
	}

	var candidates []*peerInfo
	cm.plk.RLock()
	for _, s := range cm.segments {
		s.Lock()
		for id, inf := range s.peers {
			if _, ok := cm.protected[id]; ok {
				// skip over protected peer.
				continue
			}
			candidates = append(candidates, inf)
		}
		s.Unlock()
	}
	cm.plk.RUnlock()

	// Sort peers according to their value.
	sort.Slice(candidates, func(i, j int) bool {
		return candidates[i].value < candidates[j].value
	})

	target := npeers - cm.lowWater

	// 2x number of peers we're disconnecting from because we may have more
	// than one connection per peer. Slightly over allocating isn't an issue
	// as this is a very short-lived array.
	selected := make([]inet.Conn, 0, target*2)

	for _, inf := range candidates {
		// TODO: should we be using firstSeen or the time associated with the connection itself?
		if inf.firstSeen.Add(cm.gracePeriod).After(now) {
			continue
		}

		for c := range inf.conns {
			selected = append(selected, c)
		}

		target--
		if target == 0 {
			break
		}
	}

	return selected
}

// GetTagInfo is called to fetch the tag information associated with a given
// peer, nil is returned if p refers to an unknown peer.
func (cm *BasicConnMgr) GetTagInfo(p peer.ID) *ifconnmgr.TagInfo {
	s := cm.segments.get(p)
	s.Lock()
	defer s.Unlock()

	pi, ok := s.peers[p]
	if !ok {
		return nil
	}

	out := &ifconnmgr.TagInfo{
		FirstSeen: pi.firstSeen,
		Value:     pi.value,
		Tags:      make(map[string]int),
		Conns:     make(map[string]time.Time),
	}

	for t, v := range pi.tags {
		out.Tags[t] = v
	}
	for c, t := range pi.conns {
		out.Conns[c.RemoteMultiaddr().String()] = t
	}

	return out
}

// TagPeer is called to associate a string and integer with a given peer.
func (cm *BasicConnMgr) TagPeer(p peer.ID, tag string, val int) {
	s := cm.segments.get(p)
	s.Lock()
	defer s.Unlock()

	pi, ok := s.peers[p]
	if !ok {
		log.Info("tried to tag conn from untracked peer: ", p)
		return
	}

	// Update the total value of the peer.
	pi.value += (val - pi.tags[tag])
	pi.tags[tag] = val
}

// UntagPeer is called to disassociate a string and integer from a given peer.
func (cm *BasicConnMgr) UntagPeer(p peer.ID, tag string) {
	s := cm.segments.get(p)
	s.Lock()
	defer s.Unlock()

	pi, ok := s.peers[p]
	if !ok {
		log.Info("tried to remove tag from untracked peer: ", p)
		return
	}

	// Update the total value of the peer.
	pi.value -= pi.tags[tag]
	delete(pi.tags, tag)
}

// UpsertTag is called to insert/update a peer tag
func (cm *BasicConnMgr) UpsertTag(p peer.ID, tag string, upsert func(int) int) {
	s := cm.segments.get(p)
	s.Lock()
	defer s.Unlock()

	pi, ok := s.peers[p]
	if !ok {
		log.Info("tried to upsert tag from untracked peer: ", p)
		return
	}

	oldval := pi.tags[tag]
	newval := upsert(oldval)
	pi.value += (newval - oldval)
	pi.tags[tag] = newval
}

// CMInfo holds the configuration for BasicConnMgr, as well as status data.
type CMInfo struct {
	// The low watermark, as described in NewConnManager.
	LowWater int

	// The high watermark, as described in NewConnManager.
	HighWater int

	// The timestamp when the last trim was triggered.
	LastTrim time.Time

	// The configured grace period, as described in NewConnManager.
	GracePeriod time.Duration

	// The current connection count.
	ConnCount int
}

// GetInfo returns the configuration and status data for this connection manager.
func (cm *BasicConnMgr) GetInfo() CMInfo {
	return CMInfo{
		HighWater:   cm.highWater,
		LowWater:    cm.lowWater,
		LastTrim:    cm.lastTrim,
		GracePeriod: cm.gracePeriod,
		ConnCount:   int(atomic.LoadInt32(&cm.connCount)),
	}
}

// Notifee returns a sink through which Notifiers can inform the BasicConnMgr when
// events occur. Currently, the notifee only reacts upon connection events
// {Connected, Disconnected}.
func (cm *BasicConnMgr) Notifee() inet.Notifiee {
	return (*cmNotifee)(cm)
}

type cmNotifee BasicConnMgr

func (nn *cmNotifee) cm() *BasicConnMgr {
	return (*BasicConnMgr)(nn)
}

// Connected is called by notifiers to inform that a new connection has been established.
// The notifee updates the BasicConnMgr to start tracking the connection. If the new connection
// count exceeds the high watermark, a trim may be triggered.
func (nn *cmNotifee) Connected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	p := c.RemotePeer()
	s := cm.segments.get(p)
	s.Lock()
	defer s.Unlock()

	pinfo, ok := s.peers[p]
	if !ok {
		pinfo = &peerInfo{
			firstSeen: time.Now(),
			tags:      make(map[string]int),
			conns:     make(map[inet.Conn]time.Time),
		}
		s.peers[p] = pinfo
	}

	_, ok = pinfo.conns[c]
	if ok {
		log.Error("received connected notification for conn we are already tracking: ", p)
		return
	}

	pinfo.conns[c] = time.Now()
	connCount := atomic.AddInt32(&cm.connCount, 1)

	if int(connCount) > nn.highWater {
		go cm.TrimOpenConns(context.Background())
	}
}

// Disconnected is called by notifiers to inform that an existing connection has been closed or terminated.
// The notifee updates the BasicConnMgr accordingly to stop tracking the connection, and performs housekeeping.
func (nn *cmNotifee) Disconnected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	p := c.RemotePeer()
	s := cm.segments.get(p)
	s.Lock()
	defer s.Unlock()

	cinf, ok := s.peers[p]
	if !ok {
		log.Error("received disconnected notification for peer we are not tracking: ", p)
		return
	}

	_, ok = cinf.conns[c]
	if !ok {
		log.Error("received disconnected notification for conn we are not tracking: ", p)
		return
	}

	delete(cinf.conns, c)
	if len(cinf.conns) == 0 {
		delete(s.peers, p)
	}
	atomic.AddInt32(&cm.connCount, -1)
}

// Listen is no-op in this implementation.
func (nn *cmNotifee) Listen(n inet.Network, addr ma.Multiaddr) {}

// ListenClose is no-op in this implementation.
func (nn *cmNotifee) ListenClose(n inet.Network, addr ma.Multiaddr) {}

// OpenedStream is no-op in this implementation.
func (nn *cmNotifee) OpenedStream(inet.Network, inet.Stream) {}

// ClosedStream is no-op in this implementation.
func (nn *cmNotifee) ClosedStream(inet.Network, inet.Stream) {}
<|endoftext|>


<commit_before>package connmgr

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	ifconnmgr "github.com/libp2p/go-libp2p-interface-connmgr"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

var log = logging.Logger("connmgr")

type BasicConnMgr struct {
	highWater int
	lowWater  int

	gracePeriod time.Duration

	peers     map[peer.ID]*peerInfo
	connCount int

	lk sync.Mutex

	lastTrim time.Time
}

var _ ifconnmgr.ConnManager = (*BasicConnMgr)(nil)

func NewConnManager(low, hi int, grace time.Duration) *BasicConnMgr {
	return &BasicConnMgr{
		highWater:   hi,
		lowWater:    low,
		gracePeriod: grace,
		peers:       make(map[peer.ID]*peerInfo),
	}
}

type peerInfo struct {
	tags  map[string]int
	value int

	conns map[inet.Conn]time.Time

	firstSeen time.Time
}

func (cm *BasicConnMgr) TrimOpenConns(ctx context.Context) {
	defer log.EventBegin(ctx, "connCleanup").Done()
	for _, c := range cm.getConnsToClose(ctx) {
		log.Info("closing conn: ", c.RemotePeer())
		log.Event(ctx, "closeConn", c.RemotePeer())
		c.Close()
	}
}

func (cm *BasicConnMgr) getConnsToClose(ctx context.Context) []inet.Conn {
	cm.lk.Lock()
	defer cm.lk.Unlock()
	if cm.lowWater == 0 || cm.highWater == 0 {
		// disabled
		return nil
	}
	now := time.Now()
	cm.lastTrim = now

	if len(cm.peers) < cm.lowWater {
		log.Info("open connection count below limit")
		return nil
	}

	var infos []*peerInfo

	for _, inf := range cm.peers {
		infos = append(infos, inf)
	}

	sort.Slice(infos, func(i, j int) bool {
		return infos[i].value < infos[j].value
	})

	closeCount := len(infos) - cm.lowWater
	toclose := infos[:closeCount]

	// 2x number of peers we're disconnecting from because we may have more
	// than one connection per peer. Slightly over allocating isn't an issue
	// as this is a very short-lived array.
	closed := make([]inet.Conn, 0, len(toclose)*2)

	for _, inf := range toclose {
		// TODO: should we be using firstSeen or the time associated with the connection itself?
		if inf.firstSeen.Add(cm.gracePeriod).After(now) {
			continue
		}

		// TODO: if a peer has more than one connection, maybe only close one?
		for c := range inf.conns {
			// TODO: probably don't want to always do this in a goroutine
			closed = append(closed, c)
		}
	}

	return closed
}

func (cm *BasicConnMgr) GetTagInfo(p peer.ID) *ifconnmgr.TagInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		return nil
	}

	out := &ifconnmgr.TagInfo{
		FirstSeen: pi.firstSeen,
		Value:     pi.value,
		Tags:      make(map[string]int),
		Conns:     make(map[string]time.Time),
	}

	for t, v := range pi.tags {
		out.Tags[t] = v
	}
	for c, t := range pi.conns {
		out.Conns[c.RemoteMultiaddr().String()] = t
	}

	return out
}

func (cm *BasicConnMgr) TagPeer(p peer.ID, tag string, val int) {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		log.Error("tried to tag conn from untracked peer: ", p)
		return
	}

	pi.value += (val - pi.tags[tag])
	pi.tags[tag] = val
}

func (cm *BasicConnMgr) UntagPeer(p peer.ID, tag string) {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		log.Error("tried to remove tag from untracked peer: ", p)
		return
	}

	pi.value -= pi.tags[tag]
	delete(pi.tags, tag)
}

type CMInfo struct {
	LowWater    int
	HighWater   int
	LastTrim    time.Time
	GracePeriod time.Duration
	ConnCount   int
}

func (cm *BasicConnMgr) GetInfo() CMInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	return CMInfo{
		HighWater:   cm.highWater,
		LowWater:    cm.lowWater,
		LastTrim:    cm.lastTrim,
		GracePeriod: cm.gracePeriod,
		ConnCount:   cm.connCount,
	}
}

func (cm *BasicConnMgr) Notifee() inet.Notifiee {
	return (*cmNotifee)(cm)
}

type cmNotifee BasicConnMgr

func (nn *cmNotifee) cm() *BasicConnMgr {
	return (*BasicConnMgr)(nn)
}

func (nn *cmNotifee) Connected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	cm.lk.Lock()
	defer cm.lk.Unlock()

	pinfo, ok := cm.peers[c.RemotePeer()]
	if !ok {
		pinfo = &peerInfo{
			firstSeen: time.Now(),
			tags:      make(map[string]int),
			conns:     make(map[inet.Conn]time.Time),
		}
		cm.peers[c.RemotePeer()] = pinfo
	}

	_, ok = pinfo.conns[c]
	if ok {
		log.Error("received connected notification for conn we are already tracking: ", c.RemotePeer())
		return
	}

	pinfo.conns[c] = time.Now()
	cm.connCount++

	if cm.connCount > nn.highWater {
		if cm.lastTrim.IsZero() || time.Since(cm.lastTrim) > time.Second*10 {
			go cm.TrimOpenConns(context.Background())
		}
	}
}

func (nn *cmNotifee) Disconnected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	cm.lk.Lock()
	defer cm.lk.Unlock()

	cinf, ok := cm.peers[c.RemotePeer()]
	if !ok {
		log.Error("received disconnected notification for peer we are not tracking: ", c.RemotePeer())
		return
	}

	_, ok = cinf.conns[c]
	if !ok {
		log.Error("received disconnected notification for conn we are not tracking: ", c.RemotePeer())
		return
	}

	delete(cinf.conns, c)
	cm.connCount--
	if len(cinf.conns) == 0 {
		delete(cm.peers, c.RemotePeer())
	}
}

func (nn *cmNotifee) Listen(n inet.Network, addr ma.Multiaddr)      {}
func (nn *cmNotifee) ListenClose(n inet.Network, addr ma.Multiaddr) {}
func (nn *cmNotifee) OpenedStream(inet.Network, inet.Stream)        {}
func (nn *cmNotifee) ClosedStream(inet.Network, inet.Stream)        {}
<commit_msg>demote error message to info<commit_after>package connmgr

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	ifconnmgr "github.com/libp2p/go-libp2p-interface-connmgr"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
)

var log = logging.Logger("connmgr")

type BasicConnMgr struct {
	highWater int
	lowWater  int

	gracePeriod time.Duration

	peers     map[peer.ID]*peerInfo
	connCount int

	lk sync.Mutex

	lastTrim time.Time
}

var _ ifconnmgr.ConnManager = (*BasicConnMgr)(nil)

func NewConnManager(low, hi int, grace time.Duration) *BasicConnMgr {
	return &BasicConnMgr{
		highWater:   hi,
		lowWater:    low,
		gracePeriod: grace,
		peers:       make(map[peer.ID]*peerInfo),
	}
}

type peerInfo struct {
	tags  map[string]int
	value int

	conns map[inet.Conn]time.Time

	firstSeen time.Time
}

func (cm *BasicConnMgr) TrimOpenConns(ctx context.Context) {
	defer log.EventBegin(ctx, "connCleanup").Done()
	for _, c := range cm.getConnsToClose(ctx) {
		log.Info("closing conn: ", c.RemotePeer())
		log.Event(ctx, "closeConn", c.RemotePeer())
		c.Close()
	}
}

func (cm *BasicConnMgr) getConnsToClose(ctx context.Context) []inet.Conn {
	cm.lk.Lock()
	defer cm.lk.Unlock()
	if cm.lowWater == 0 || cm.highWater == 0 {
		// disabled
		return nil
	}
	now := time.Now()
	cm.lastTrim = now

	if len(cm.peers) < cm.lowWater {
		log.Info("open connection count below limit")
		return nil
	}

	var infos []*peerInfo

	for _, inf := range cm.peers {
		infos = append(infos, inf)
	}

	sort.Slice(infos, func(i, j int) bool {
		return infos[i].value < infos[j].value
	})

	closeCount := len(infos) - cm.lowWater
	toclose := infos[:closeCount]

	// 2x number of peers we're disconnecting from because we may have more
	// than one connection per peer. Slightly over allocating isn't an issue
	// as this is a very short-lived array.
	closed := make([]inet.Conn, 0, len(toclose)*2)

	for _, inf := range toclose {
		// TODO: should we be using firstSeen or the time associated with the connection itself?
		if inf.firstSeen.Add(cm.gracePeriod).After(now) {
			continue
		}

		// TODO: if a peer has more than one connection, maybe only close one?
		for c := range inf.conns {
			// TODO: probably don't want to always do this in a goroutine
			closed = append(closed, c)
		}
	}

	return closed
}

func (cm *BasicConnMgr) GetTagInfo(p peer.ID) *ifconnmgr.TagInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		return nil
	}

	out := &ifconnmgr.TagInfo{
		FirstSeen: pi.firstSeen,
		Value:     pi.value,
		Tags:      make(map[string]int),
		Conns:     make(map[string]time.Time),
	}

	for t, v := range pi.tags {
		out.Tags[t] = v
	}
	for c, t := range pi.conns {
		out.Conns[c.RemoteMultiaddr().String()] = t
	}

	return out
}

func (cm *BasicConnMgr) TagPeer(p peer.ID, tag string, val int) {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		log.Info("tried to tag conn from untracked peer: ", p)
		return
	}

	pi.value += (val - pi.tags[tag])
	pi.tags[tag] = val
}

func (cm *BasicConnMgr) UntagPeer(p peer.ID, tag string) {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	pi, ok := cm.peers[p]
	if !ok {
		log.Info("tried to remove tag from untracked peer: ", p)
		return
	}

	pi.value -= pi.tags[tag]
	delete(pi.tags, tag)
}

type CMInfo struct {
	LowWater    int
	HighWater   int
	LastTrim    time.Time
	GracePeriod time.Duration
	ConnCount   int
}

func (cm *BasicConnMgr) GetInfo() CMInfo {
	cm.lk.Lock()
	defer cm.lk.Unlock()

	return CMInfo{
		HighWater:   cm.highWater,
		LowWater:    cm.lowWater,
		LastTrim:    cm.lastTrim,
		GracePeriod: cm.gracePeriod,
		ConnCount:   cm.connCount,
	}
}

func (cm *BasicConnMgr) Notifee() inet.Notifiee {
	return (*cmNotifee)(cm)
}

type cmNotifee BasicConnMgr

func (nn *cmNotifee) cm() *BasicConnMgr {
	return (*BasicConnMgr)(nn)
}

func (nn *cmNotifee) Connected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	cm.lk.Lock()
	defer cm.lk.Unlock()

	pinfo, ok := cm.peers[c.RemotePeer()]
	if !ok {
		pinfo = &peerInfo{
			firstSeen: time.Now(),
			tags:      make(map[string]int),
			conns:     make(map[inet.Conn]time.Time),
		}
		cm.peers[c.RemotePeer()] = pinfo
	}

	_, ok = pinfo.conns[c]
	if ok {
		log.Error("received connected notification for conn we are already tracking: ", c.RemotePeer())
		return
	}

	pinfo.conns[c] = time.Now()
	cm.connCount++

	if cm.connCount > nn.highWater {
		if cm.lastTrim.IsZero() || time.Since(cm.lastTrim) > time.Second*10 {
			go cm.TrimOpenConns(context.Background())
		}
	}
}

func (nn *cmNotifee) Disconnected(n inet.Network, c inet.Conn) {
	cm := nn.cm()

	cm.lk.Lock()
	defer cm.lk.Unlock()

	cinf, ok := cm.peers[c.RemotePeer()]
	if !ok {
		log.Error("received disconnected notification for peer we are not tracking: ", c.RemotePeer())
		return
	}

	_, ok = cinf.conns[c]
	if !ok {
		log.Error("received disconnected notification for conn we are not tracking: ", c.RemotePeer())
		return
	}

	delete(cinf.conns, c)
	cm.connCount--
	if len(cinf.conns) == 0 {
		delete(cm.peers, c.RemotePeer())
	}
}

func (nn *cmNotifee) Listen(n inet.Network, addr ma.Multiaddr)      {}
func (nn *cmNotifee) ListenClose(n inet.Network, addr ma.Multiaddr) {}
func (nn *cmNotifee) OpenedStream(inet.Network, inet.Stream)        {}
func (nn *cmNotifee) ClosedStream(inet.Network, inet.Stream)        {}
<|endoftext|>


<commit_before>// Copyright 2016 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package ast

import "fmt"

// Transformer defines the interface for transforming AST elements. If the
// transformer returns nil and does not indicate an error, the AST element will
// be set to nil and no transformations will be applied to children of the
// element.
type Transformer interface {
	Transform(v interface{}) (interface{}, error)
}

// Transform iterates the AST and calls the Transform function on the
// Transformer t for x before recursing.
func Transform(t Transformer, x interface{}) (interface{}, error) {

	if term, ok := x.(*Term); ok {
		return Transform(t, term.Value)
	}

	y, err := t.Transform(x)
	if err != nil {
		return x, err
	}

	if y == nil {
		return nil, nil
	}

	var ok bool
	switch y := y.(type) {
	case *Module:
		p, err := Transform(t, y.Package)
		if err != nil {
			return nil, err
		}
		if y.Package, ok = p.(*Package); !ok {
			return nil, fmt.Errorf("illegal transform: %T != %T", y.Package, p)
		}
		for i := range y.Imports {
			imp, err := Transform(t, y.Imports[i])
			if err != nil {
				return nil, err
			}
			if y.Imports[i], ok = imp.(*Import); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Imports[i], imp)
			}
		}
		for i := range y.Rules {
			rule, err := Transform(t, y.Rules[i])
			if err != nil {
				return nil, err
			}
			if y.Rules[i], ok = rule.(*Rule); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Rules[i], rule)
			}
		}
		for i := range y.Comments {
			comment, err := Transform(t, y.Comments[i])
			if err != nil {
				return nil, err
			}
			if y.Comments[i], ok = comment.(*Comment); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Comments[i], comment)
			}
		}
		return y, nil
	case *Package:
		ref, err := Transform(t, y.Path)
		if err != nil {
			return nil, err
		}
		if y.Path, ok = ref.(Ref); !ok {
			return nil, fmt.Errorf("illegal transform: %T != %T", y.Path, ref)
		}
		return y, nil
	case *Import:
		y.Path, err = transformTerm(t, y.Path)
		if err != nil {
			return nil, err
		}
		if y.Alias, err = transformVar(t, y.Alias); err != nil {
			return nil, err
		}
		return y, nil
	case *Rule:
		if y.Head, err = transformHead(t, y.Head); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		if y.Else != nil {
			rule, err := Transform(t, y.Else)
			if err != nil {
				return nil, err
			}
			if y.Else, ok = rule.(*Rule); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Else, rule)
			}
		}
		return y, nil
	case *Head:
		if y.Name, err = transformVar(t, y.Name); err != nil {
			return nil, err
		}
		if y.Args, err = transformArgs(t, y.Args); err != nil {
			return nil, err
		}
		if y.Key != nil {
			if y.Key, err = transformTerm(t, y.Key); err != nil {
				return nil, err
			}
		}
		if y.Value != nil {
			if y.Value, err = transformTerm(t, y.Value); err != nil {
				return nil, err
			}
		}
		return y, nil
	case Args:
		for i := range y {
			if y[i], err = transformTerm(t, y[i]); err != nil {
				return nil, err
			}
		}
		return y, nil
	case Body:
		for i, e := range y {
			e, err := Transform(t, e)
			if err != nil {
				return nil, err
			}
			if y[i], ok = e.(*Expr); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y[i], e)
			}
		}
		return y, nil
	case *Expr:
		switch ts := y.Terms.(type) {
		case []*Term:
			for i := range ts {
				if ts[i], err = transformTerm(t, ts[i]); err != nil {
					return nil, err
				}
			}
		case *Term:
			if y.Terms, err = transformTerm(t, ts); err != nil {
				return nil, err
			}
		}
		for i, w := range y.With {
			w, err := Transform(t, w)
			if err != nil {
				return nil, err
			}
			if y.With[i], ok = w.(*With); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.With[i], w)
			}
		}
		return y, nil
	case *With:
		if y.Target, err = transformTerm(t, y.Target); err != nil {
			return nil, err
		}
		if y.Value, err = transformTerm(t, y.Value); err != nil {
			return nil, err
		}
		return y, nil
	case Ref:
		for i, term := range y {
			if y[i], err = transformTerm(t, term); err != nil {
				return nil, err
			}
		}
		return y, nil
	case Object:
		for i, elem := range y {
			k, err := transformTerm(t, elem[0])
			if err != nil {
				return nil, err
			}
			v, err := transformTerm(t, elem[1])
			if err != nil {
				return nil, err
			}
			y[i] = Item(k, v)
		}
		return y, nil
	case Array:
		for i := range y {
			if y[i], err = transformTerm(t, y[i]); err != nil {
				return nil, err
			}
		}
		return y, nil
	case *Set:
		y, err = y.Map(func(term *Term) (*Term, error) {
			return transformTerm(t, term)
		})
		if err != nil {
			return nil, err
		}
		return y, nil
	case *ArrayComprehension:
		if y.Term, err = transformTerm(t, y.Term); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		return y, nil
	case *ObjectComprehension:
		if y.Key, err = transformTerm(t, y.Key); err != nil {
			return nil, err
		}
		if y.Value, err = transformTerm(t, y.Value); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		return y, nil
	case *SetComprehension:
		if y.Term, err = transformTerm(t, y.Term); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		return y, nil
	default:
		return y, nil
	}
}

// TransformRefs calls the function f on all references under x.
func TransformRefs(x interface{}, f func(Ref) (Value, error)) (interface{}, error) {
	t := &GenericTransformer{func(x interface{}) (interface{}, error) {
		if r, ok := x.(Ref); ok {
			return f(r)
		}
		return x, nil
	}}
	return Transform(t, x)
}

// TransformComprehensions calls the functio nf on all comprehensions under x.
func TransformComprehensions(x interface{}, f func(interface{}) (Value, error)) (interface{}, error) {
	t := &GenericTransformer{func(x interface{}) (interface{}, error) {
		switch x := x.(type) {
		case *ArrayComprehension:
			return f(x)
		case *SetComprehension:
			return f(x)
		case *ObjectComprehension:
			return f(x)
		}
		return x, nil
	}}
	return Transform(t, x)
}

// GenericTransformer implements the Transformer interface to provide a utility
// to transform AST nodes using a closure.
type GenericTransformer struct {
	f func(x interface{}) (interface{}, error)
}

// Transform calls the function f on the GenericTransformer.
func (t *GenericTransformer) Transform(x interface{}) (interface{}, error) {
	return t.f(x)
}

func transformHead(t Transformer, head *Head) (*Head, error) {
	y, err := Transform(t, head)
	if err != nil {
		return nil, err
	}
	h, ok := y.(*Head)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", head, y)
	}
	return h, nil
}
func transformArgs(t Transformer, args Args) (Args, error) {
	y, err := Transform(t, args)
	if err != nil {
		return nil, err
	}
	a, ok := y.(Args)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", args, y)
	}
	return a, nil
}

func transformBody(t Transformer, body Body) (Body, error) {
	y, err := Transform(t, body)
	if err != nil {
		return nil, err
	}
	r, ok := y.(Body)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", body, y)
	}
	return r, nil
}

func transformTerm(t Transformer, term *Term) (*Term, error) {
	v, err := transformValue(t, term.Value)
	if err != nil {
		return nil, err
	}
	r := &Term{
		Value:    v,
		Location: term.Location,
	}
	return r, nil
}

func transformValue(t Transformer, v Value) (Value, error) {
	v1, err := Transform(t, v)
	if err != nil {
		return nil, err
	}
	r, ok := v1.(Value)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", v, v1)
	}
	return r, nil
}

func transformVar(t Transformer, v Var) (Var, error) {
	v1, err := Transform(t, v)
	if err != nil {
		return "", err
	}
	r, ok := v1.(Var)
	if !ok {
		return "", fmt.Errorf("illegal transform: %T != %T", v, v1)
	}
	return r, nil
}
<commit_msg>Add generic transformer constructor<commit_after>// Copyright 2016 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package ast

import "fmt"

// Transformer defines the interface for transforming AST elements. If the
// transformer returns nil and does not indicate an error, the AST element will
// be set to nil and no transformations will be applied to children of the
// element.
type Transformer interface {
	Transform(v interface{}) (interface{}, error)
}

// Transform iterates the AST and calls the Transform function on the
// Transformer t for x before recursing.
func Transform(t Transformer, x interface{}) (interface{}, error) {

	if term, ok := x.(*Term); ok {
		return Transform(t, term.Value)
	}

	y, err := t.Transform(x)
	if err != nil {
		return x, err
	}

	if y == nil {
		return nil, nil
	}

	var ok bool
	switch y := y.(type) {
	case *Module:
		p, err := Transform(t, y.Package)
		if err != nil {
			return nil, err
		}
		if y.Package, ok = p.(*Package); !ok {
			return nil, fmt.Errorf("illegal transform: %T != %T", y.Package, p)
		}
		for i := range y.Imports {
			imp, err := Transform(t, y.Imports[i])
			if err != nil {
				return nil, err
			}
			if y.Imports[i], ok = imp.(*Import); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Imports[i], imp)
			}
		}
		for i := range y.Rules {
			rule, err := Transform(t, y.Rules[i])
			if err != nil {
				return nil, err
			}
			if y.Rules[i], ok = rule.(*Rule); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Rules[i], rule)
			}
		}
		for i := range y.Comments {
			comment, err := Transform(t, y.Comments[i])
			if err != nil {
				return nil, err
			}
			if y.Comments[i], ok = comment.(*Comment); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Comments[i], comment)
			}
		}
		return y, nil
	case *Package:
		ref, err := Transform(t, y.Path)
		if err != nil {
			return nil, err
		}
		if y.Path, ok = ref.(Ref); !ok {
			return nil, fmt.Errorf("illegal transform: %T != %T", y.Path, ref)
		}
		return y, nil
	case *Import:
		y.Path, err = transformTerm(t, y.Path)
		if err != nil {
			return nil, err
		}
		if y.Alias, err = transformVar(t, y.Alias); err != nil {
			return nil, err
		}
		return y, nil
	case *Rule:
		if y.Head, err = transformHead(t, y.Head); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		if y.Else != nil {
			rule, err := Transform(t, y.Else)
			if err != nil {
				return nil, err
			}
			if y.Else, ok = rule.(*Rule); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.Else, rule)
			}
		}
		return y, nil
	case *Head:
		if y.Name, err = transformVar(t, y.Name); err != nil {
			return nil, err
		}
		if y.Args, err = transformArgs(t, y.Args); err != nil {
			return nil, err
		}
		if y.Key != nil {
			if y.Key, err = transformTerm(t, y.Key); err != nil {
				return nil, err
			}
		}
		if y.Value != nil {
			if y.Value, err = transformTerm(t, y.Value); err != nil {
				return nil, err
			}
		}
		return y, nil
	case Args:
		for i := range y {
			if y[i], err = transformTerm(t, y[i]); err != nil {
				return nil, err
			}
		}
		return y, nil
	case Body:
		for i, e := range y {
			e, err := Transform(t, e)
			if err != nil {
				return nil, err
			}
			if y[i], ok = e.(*Expr); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y[i], e)
			}
		}
		return y, nil
	case *Expr:
		switch ts := y.Terms.(type) {
		case []*Term:
			for i := range ts {
				if ts[i], err = transformTerm(t, ts[i]); err != nil {
					return nil, err
				}
			}
		case *Term:
			if y.Terms, err = transformTerm(t, ts); err != nil {
				return nil, err
			}
		}
		for i, w := range y.With {
			w, err := Transform(t, w)
			if err != nil {
				return nil, err
			}
			if y.With[i], ok = w.(*With); !ok {
				return nil, fmt.Errorf("illegal transform: %T != %T", y.With[i], w)
			}
		}
		return y, nil
	case *With:
		if y.Target, err = transformTerm(t, y.Target); err != nil {
			return nil, err
		}
		if y.Value, err = transformTerm(t, y.Value); err != nil {
			return nil, err
		}
		return y, nil
	case Ref:
		for i, term := range y {
			if y[i], err = transformTerm(t, term); err != nil {
				return nil, err
			}
		}
		return y, nil
	case Object:
		for i, elem := range y {
			k, err := transformTerm(t, elem[0])
			if err != nil {
				return nil, err
			}
			v, err := transformTerm(t, elem[1])
			if err != nil {
				return nil, err
			}
			y[i] = Item(k, v)
		}
		return y, nil
	case Array:
		for i := range y {
			if y[i], err = transformTerm(t, y[i]); err != nil {
				return nil, err
			}
		}
		return y, nil
	case *Set:
		y, err = y.Map(func(term *Term) (*Term, error) {
			return transformTerm(t, term)
		})
		if err != nil {
			return nil, err
		}
		return y, nil
	case *ArrayComprehension:
		if y.Term, err = transformTerm(t, y.Term); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		return y, nil
	case *ObjectComprehension:
		if y.Key, err = transformTerm(t, y.Key); err != nil {
			return nil, err
		}
		if y.Value, err = transformTerm(t, y.Value); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		return y, nil
	case *SetComprehension:
		if y.Term, err = transformTerm(t, y.Term); err != nil {
			return nil, err
		}
		if y.Body, err = transformBody(t, y.Body); err != nil {
			return nil, err
		}
		return y, nil
	default:
		return y, nil
	}
}

// TransformRefs calls the function f on all references under x.
func TransformRefs(x interface{}, f func(Ref) (Value, error)) (interface{}, error) {
	t := &GenericTransformer{func(x interface{}) (interface{}, error) {
		if r, ok := x.(Ref); ok {
			return f(r)
		}
		return x, nil
	}}
	return Transform(t, x)
}

// TransformComprehensions calls the functio nf on all comprehensions under x.
func TransformComprehensions(x interface{}, f func(interface{}) (Value, error)) (interface{}, error) {
	t := &GenericTransformer{func(x interface{}) (interface{}, error) {
		switch x := x.(type) {
		case *ArrayComprehension:
			return f(x)
		case *SetComprehension:
			return f(x)
		case *ObjectComprehension:
			return f(x)
		}
		return x, nil
	}}
	return Transform(t, x)
}

// GenericTransformer implements the Transformer interface to provide a utility
// to transform AST nodes using a closure.
type GenericTransformer struct {
	f func(x interface{}) (interface{}, error)
}

// NewGenericTransformer returns a new GenericTransformer that will transform
// AST nodes using the function f.
func NewGenericTransformer(f func(x interface{}) (interface{}, error)) *GenericTransformer {
	return &GenericTransformer{
		f: f,
	}
}

// Transform calls the function f on the GenericTransformer.
func (t *GenericTransformer) Transform(x interface{}) (interface{}, error) {
	return t.f(x)
}

func transformHead(t Transformer, head *Head) (*Head, error) {
	y, err := Transform(t, head)
	if err != nil {
		return nil, err
	}
	h, ok := y.(*Head)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", head, y)
	}
	return h, nil
}
func transformArgs(t Transformer, args Args) (Args, error) {
	y, err := Transform(t, args)
	if err != nil {
		return nil, err
	}
	a, ok := y.(Args)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", args, y)
	}
	return a, nil
}

func transformBody(t Transformer, body Body) (Body, error) {
	y, err := Transform(t, body)
	if err != nil {
		return nil, err
	}
	r, ok := y.(Body)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", body, y)
	}
	return r, nil
}

func transformTerm(t Transformer, term *Term) (*Term, error) {
	v, err := transformValue(t, term.Value)
	if err != nil {
		return nil, err
	}
	r := &Term{
		Value:    v,
		Location: term.Location,
	}
	return r, nil
}

func transformValue(t Transformer, v Value) (Value, error) {
	v1, err := Transform(t, v)
	if err != nil {
		return nil, err
	}
	r, ok := v1.(Value)
	if !ok {
		return nil, fmt.Errorf("illegal transform: %T != %T", v, v1)
	}
	return r, nil
}

func transformVar(t Transformer, v Var) (Var, error) {
	v1, err := Transform(t, v)
	if err != nil {
		return "", err
	}
	r, ok := v1.(Var)
	if !ok {
		return "", fmt.Errorf("illegal transform: %T != %T", v, v1)
	}
	return r, nil
}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/dotcloud/docker/utils"
	"io"
	"io/ioutil"
	"os"
	"path"
	"reflect"
	"strings"
)

type BuildFile interface {
	Build(io.Reader, io.Reader) (string, error)
	CmdFrom(string) error
	CmdRun(string) error
}

type buildFile struct {
	runtime *Runtime
	builder *Builder
	srv     *Server

	image      string
	maintainer string
	config     *Config
	context    string

	tmpContainers map[string]struct{}
	tmpImages     map[string]struct{}

	out io.Writer
}

func (b *buildFile) clearTmp(containers, images map[string]struct{}) {
	for c := range containers {
		tmp := b.runtime.Get(c)
		b.runtime.Destroy(tmp)
		utils.Debugf("Removing container %s", c)
	}
	for i := range images {
		b.runtime.graph.Delete(i)
		utils.Debugf("Removing image %s", i)
	}
}

func (b *buildFile) CmdFrom(name string) error {
	image, err := b.runtime.repositories.LookupImage(name)
	if err != nil {
		if b.runtime.graph.IsNotExist(err) {

			var tag, remote string
			if strings.Contains(name, ":") {
				remoteParts := strings.Split(name, ":")
				tag = remoteParts[1]
				remote = remoteParts[0]
			} else {
				remote = name
			}

			if err := b.srv.ImagePull(remote, tag, "", b.out, utils.NewStreamFormatter(false)); err != nil {
				return err
			}

			image, err = b.runtime.repositories.LookupImage(name)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}
	b.image = image.ID
	b.config = &Config{}
	return nil
}

func (b *buildFile) CmdMaintainer(name string) error {
	b.maintainer = name
	return b.commit("", b.config.Cmd, fmt.Sprintf("MAINTAINER %s", name))
}

func (b *buildFile) CmdRun(args string) error {
	if b.image == "" {
		return fmt.Errorf("Please provide a source image with `from` prior to run")
	}
	config, _, err := ParseRun([]string{b.image, "/bin/sh", "-c", args}, nil)
	if err != nil {
		return err
	}

	cmd := b.config.Cmd
	b.config.Cmd = nil
	MergeConfig(b.config, config)

	utils.Debugf("Command to be executed: %v", b.config.Cmd)

	if cache, err := b.srv.ImageGetCached(b.image, b.config); err != nil {
		return err
	} else if cache != nil {
		utils.Debugf("[BUILDER] Use cached version")
		b.image = cache.ID
		return nil
	} else {
		utils.Debugf("[BUILDER] Cache miss")
	}

	cid, err := b.run()
	if err != nil {
		return err
	}
	if err := b.commit(cid, cmd, "run"); err != nil {
		return err
	}
	b.config.Cmd = cmd
	return nil
}

func (b *buildFile) CmdEnv(args string) error {
	tmp := strings.SplitN(args, " ", 2)
	if len(tmp) != 2 {
		return fmt.Errorf("Invalid ENV format")
	}
	key := strings.Trim(tmp[0], " ")
	value := strings.Trim(tmp[1], " ")

	for i, elem := range b.config.Env {
		if strings.HasPrefix(elem, key+"=") {
			b.config.Env[i] = key + "=" + value
			return nil
		}
	}
	b.config.Env = append(b.config.Env, key+"="+value)
	return b.commit("", b.config.Cmd, fmt.Sprintf("ENV %s=%s", key, value))
}

func (b *buildFile) CmdCmd(args string) error {
	var cmd []string
	if err := json.Unmarshal([]byte(args), &cmd); err != nil {
		utils.Debugf("Error unmarshalling: %s, using /bin/sh -c", err)
		cmd = []string{"/bin/sh", "-c", args}
	}
	if err := b.commit("", cmd, fmt.Sprintf("CMD %v", cmd)); err != nil {
		return err
	}
	b.config.Cmd = cmd
	return nil
}

func (b *buildFile) CmdExpose(args string) error {
	ports := strings.Split(args, " ")
	b.config.PortSpecs = append(ports, b.config.PortSpecs...)
	return b.commit("", b.config.Cmd, fmt.Sprintf("EXPOSE %v", ports))
}

func (b *buildFile) CmdInsert(args string) error {
	return fmt.Errorf("INSERT has been deprecated. Please use ADD instead")
}

func (b *buildFile) CmdCopy(args string) error {
	return fmt.Errorf("COPY has been deprecated. Please use ADD instead")
}

func (b *buildFile) isUrl(str string) bool {
	return strings.HasPrefix(str, "http://") || strings.HasPrefix(str, "https://")
}

func (b *buildFile) addRemote(container *Container, orig, dest string) error {
	file, err := utils.Download(orig, ioutil.Discard)
	if err != nil {
		return err
	}
	defer file.Body.Close()

	return container.Inject(file.Body, dest)
}

func (b *buildFile) addContext(container *Container, orig, dest string) error {
	origPath := path.Join(b.context, orig)
	destPath := path.Join(container.RootfsPath(), dest)

	fi, err := os.Stat(origPath)
	if err != nil {
		return err
	}
	if fi.IsDir() {
		if err := os.MkdirAll(destPath, 0700); err != nil {
			return err
		}

		files, err := ioutil.ReadDir(path.Join(b.context, orig))
		if err != nil {
			return err
		}
		for _, fi := range files {
			if err := utils.CopyDirectory(path.Join(origPath, fi.Name()), path.Join(destPath, fi.Name())); err != nil {
				return err
			}
		}
	} else {
		if err := os.MkdirAll(path.Dir(destPath), 0700); err != nil {
			return err
		}
		if err := utils.CopyDirectory(origPath, destPath); err != nil {
			return err
		}
	}
	return nil
}

func (b *buildFile) CmdAdd(args string) error {
	if b.context == "" {
		return fmt.Errorf("No context given. Impossible to use ADD")
	}
	tmp := strings.SplitN(args, " ", 2)
	if len(tmp) != 2 {
		return fmt.Errorf("Invalid ADD format")
	}
	orig := strings.Trim(tmp[0], " ")
	dest := strings.Trim(tmp[1], " ")

	cmd := b.config.Cmd
	b.config.Cmd = []string{"/bin/sh", "-c", fmt.Sprintf("#(nop) ADD %s in %s", orig, dest)}
	cid, err := b.run()
	if err != nil {
		return err
	}

	container := b.runtime.Get(cid)
	if container == nil {
		return fmt.Errorf("Error while creating the container (CmdAdd)")
	}
	if err := container.EnsureMounted(); err != nil {
		return err
	}
	defer container.Unmount()

	if b.isUrl(orig) {
		if err := b.addRemote(container, orig, dest); err != nil {
			return err
		}
	} else {
		if err := b.addContext(container, orig, dest); err != nil {
			return err
		}
	}

	if err := b.commit(cid, cmd, fmt.Sprintf("ADD %s in %s", orig, dest)); err != nil {
		return err
	}
	b.config.Cmd = cmd
	return nil
}

func (b *buildFile) run() (string, error) {
	if b.image == "" {
		return "", fmt.Errorf("Please provide a source image with `from` prior to run")
	}
	b.config.Image = b.image

	// Create the container and start it
	c, err := b.builder.Create(b.config)
	if err != nil {
		return "", err
	}
	b.tmpContainers[c.ID] = struct{}{}

	//start the container
	if err := c.Start(); err != nil {
		return "", err
	}

	// Wait for it to finish
	if ret := c.Wait(); ret != 0 {
		return "", fmt.Errorf("The command %v returned a non-zero code: %d", b.config.Cmd, ret)
	}

	return c.ID, nil
}

// Commit the container <id> with the autorun command <autoCmd>
func (b *buildFile) commit(id string, autoCmd []string, comment string) error {
	if b.image == "" {
		return fmt.Errorf("Please provide a source image with `from` prior to commit")
	}
	b.config.Image = b.image
	if id == "" {
		b.config.Cmd = []string{"/bin/sh", "-c", "#(nop) " + comment}

		if cache, err := b.srv.ImageGetCached(b.image, b.config); err != nil {
			return err
		} else if cache != nil {
			utils.Debugf("[BUILDER] Use cached version")
			b.image = cache.ID
			return nil
		} else {
			utils.Debugf("[BUILDER] Cache miss")
		}

		cid, err := b.run()
		if err != nil {
			return err
		}
		id = cid
	}

	container := b.runtime.Get(id)
	if container == nil {
		return fmt.Errorf("An error occured while creating the container")
	}

	// Note: Actually copy the struct
	autoConfig := *b.config
	autoConfig.Cmd = autoCmd
	// Commit the container
	image, err := b.builder.Commit(container, "", "", "", b.maintainer, &autoConfig)
	if err != nil {
		return err
	}
	b.tmpImages[image.ID] = struct{}{}
	b.image = image.ID
	return nil
}

func (b *buildFile) Build(dockerfile, context io.Reader) (string, error) {
	if context != nil {
		name, err := ioutil.TempDir("/tmp", "docker-build")
		if err != nil {
			return "", err
		}
		if err := Untar(context, name); err != nil {
			return "", err
		}
		defer os.RemoveAll(name)
		b.context = name
	}
	file := bufio.NewReader(dockerfile)
	for {
		line, err := file.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return "", err
		}
		line = strings.TrimSpace(strings.Replace(line, "	", " ", -1))
		// Skip comments and empty line
		if len(line) == 0 || line[0] == '#' {
			continue
		}
		tmp := strings.SplitN(line, " ", 2)
		if len(tmp) != 2 {
			return "", fmt.Errorf("Invalid Dockerfile format")
		}
		instruction := strings.ToLower(strings.Trim(tmp[0], " "))
		arguments := strings.Trim(tmp[1], " ")

		fmt.Fprintf(b.out, "%s %s (%s)\n", strings.ToUpper(instruction), arguments, b.image)

		method, exists := reflect.TypeOf(b).MethodByName("Cmd" + strings.ToUpper(instruction[:1]) + strings.ToLower(instruction[1:]))
		if !exists {
			fmt.Fprintf(b.out, "Skipping unknown instruction %s\n", strings.ToUpper(instruction))
			continue
		}
		ret := method.Func.Call([]reflect.Value{reflect.ValueOf(b), reflect.ValueOf(arguments)})[0].Interface()
		if ret != nil {
			return "", ret.(error)
		}

		fmt.Fprintf(b.out, "===> %v\n", b.image)
	}
	if b.image != "" {
		fmt.Fprintf(b.out, "Build successful.\n===> %s\n", b.image)
		return b.image, nil
	}
	return "", fmt.Errorf("An error occured during the build\n")
}

func NewBuildFile(srv *Server, out io.Writer) BuildFile {
	return &buildFile{
		builder:       NewBuilder(srv.runtime),
		runtime:       srv.runtime,
		srv:           srv,
		config:        &Config{},
		out:           out,
		tmpContainers: make(map[string]struct{}),
		tmpImages:     make(map[string]struct{}),
	}
}
<commit_msg>Fix issue EXPOSE override CMD within builder<commit_after>package docker

import (
	"bufio"
	"encoding/json"
	"fmt"
	"github.com/dotcloud/docker/utils"
	"io"
	"io/ioutil"
	"os"
	"path"
	"reflect"
	"strings"
)

type BuildFile interface {
	Build(io.Reader, io.Reader) (string, error)
	CmdFrom(string) error
	CmdRun(string) error
}

type buildFile struct {
	runtime *Runtime
	builder *Builder
	srv     *Server

	image      string
	maintainer string
	config     *Config
	context    string

	tmpContainers map[string]struct{}
	tmpImages     map[string]struct{}

	out io.Writer
}

func (b *buildFile) clearTmp(containers, images map[string]struct{}) {
	for c := range containers {
		tmp := b.runtime.Get(c)
		b.runtime.Destroy(tmp)
		utils.Debugf("Removing container %s", c)
	}
	for i := range images {
		b.runtime.graph.Delete(i)
		utils.Debugf("Removing image %s", i)
	}
}

func (b *buildFile) CmdFrom(name string) error {
	image, err := b.runtime.repositories.LookupImage(name)
	if err != nil {
		if b.runtime.graph.IsNotExist(err) {

			var tag, remote string
			if strings.Contains(name, ":") {
				remoteParts := strings.Split(name, ":")
				tag = remoteParts[1]
				remote = remoteParts[0]
			} else {
				remote = name
			}

			if err := b.srv.ImagePull(remote, tag, "", b.out, utils.NewStreamFormatter(false)); err != nil {
				return err
			}

			image, err = b.runtime.repositories.LookupImage(name)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}
	b.image = image.ID
	b.config = &Config{}
	return nil
}

func (b *buildFile) CmdMaintainer(name string) error {
	b.maintainer = name
	return b.commit("", b.config.Cmd, fmt.Sprintf("MAINTAINER %s", name))
}

func (b *buildFile) CmdRun(args string) error {
	if b.image == "" {
		return fmt.Errorf("Please provide a source image with `from` prior to run")
	}
	config, _, err := ParseRun([]string{b.image, "/bin/sh", "-c", args}, nil)
	if err != nil {
		return err
	}

	cmd := b.config.Cmd
	b.config.Cmd = nil
	MergeConfig(b.config, config)

	utils.Debugf("Command to be executed: %v", b.config.Cmd)

	if cache, err := b.srv.ImageGetCached(b.image, b.config); err != nil {
		return err
	} else if cache != nil {
		utils.Debugf("[BUILDER] Use cached version")
		b.image = cache.ID
		return nil
	} else {
		utils.Debugf("[BUILDER] Cache miss")
	}

	cid, err := b.run()
	if err != nil {
		return err
	}
	if err := b.commit(cid, cmd, "run"); err != nil {
		return err
	}
	b.config.Cmd = cmd
	return nil
}

func (b *buildFile) CmdEnv(args string) error {
	tmp := strings.SplitN(args, " ", 2)
	if len(tmp) != 2 {
		return fmt.Errorf("Invalid ENV format")
	}
	key := strings.Trim(tmp[0], " ")
	value := strings.Trim(tmp[1], " ")

	for i, elem := range b.config.Env {
		if strings.HasPrefix(elem, key+"=") {
			b.config.Env[i] = key + "=" + value
			return nil
		}
	}
	b.config.Env = append(b.config.Env, key+"="+value)
	return b.commit("", b.config.Cmd, fmt.Sprintf("ENV %s=%s", key, value))
}

func (b *buildFile) CmdCmd(args string) error {
	var cmd []string
	if err := json.Unmarshal([]byte(args), &cmd); err != nil {
		utils.Debugf("Error unmarshalling: %s, using /bin/sh -c", err)
		cmd = []string{"/bin/sh", "-c", args}
	}
	if err := b.commit("", cmd, fmt.Sprintf("CMD %v", cmd)); err != nil {
		return err
	}
	b.config.Cmd = cmd
	return nil
}

func (b *buildFile) CmdExpose(args string) error {
	ports := strings.Split(args, " ")
	b.config.PortSpecs = append(ports, b.config.PortSpecs...)
	return b.commit("", b.config.Cmd, fmt.Sprintf("EXPOSE %v", ports))
}

func (b *buildFile) CmdInsert(args string) error {
	return fmt.Errorf("INSERT has been deprecated. Please use ADD instead")
}

func (b *buildFile) CmdCopy(args string) error {
	return fmt.Errorf("COPY has been deprecated. Please use ADD instead")
}

func (b *buildFile) isUrl(str string) bool {
	return strings.HasPrefix(str, "http://") || strings.HasPrefix(str, "https://")
}

func (b *buildFile) addRemote(container *Container, orig, dest string) error {
	file, err := utils.Download(orig, ioutil.Discard)
	if err != nil {
		return err
	}
	defer file.Body.Close()

	return container.Inject(file.Body, dest)
}

func (b *buildFile) addContext(container *Container, orig, dest string) error {
	origPath := path.Join(b.context, orig)
	destPath := path.Join(container.RootfsPath(), dest)

	fi, err := os.Stat(origPath)
	if err != nil {
		return err
	}
	if fi.IsDir() {
		if err := os.MkdirAll(destPath, 0700); err != nil {
			return err
		}

		files, err := ioutil.ReadDir(path.Join(b.context, orig))
		if err != nil {
			return err
		}
		for _, fi := range files {
			if err := utils.CopyDirectory(path.Join(origPath, fi.Name()), path.Join(destPath, fi.Name())); err != nil {
				return err
			}
		}
	} else {
		if err := os.MkdirAll(path.Dir(destPath), 0700); err != nil {
			return err
		}
		if err := utils.CopyDirectory(origPath, destPath); err != nil {
			return err
		}
	}
	return nil
}

func (b *buildFile) CmdAdd(args string) error {
	if b.context == "" {
		return fmt.Errorf("No context given. Impossible to use ADD")
	}
	tmp := strings.SplitN(args, " ", 2)
	if len(tmp) != 2 {
		return fmt.Errorf("Invalid ADD format")
	}
	orig := strings.Trim(tmp[0], " ")
	dest := strings.Trim(tmp[1], " ")

	cmd := b.config.Cmd
	b.config.Cmd = []string{"/bin/sh", "-c", fmt.Sprintf("#(nop) ADD %s in %s", orig, dest)}
	cid, err := b.run()
	if err != nil {
		return err
	}

	container := b.runtime.Get(cid)
	if container == nil {
		return fmt.Errorf("Error while creating the container (CmdAdd)")
	}
	if err := container.EnsureMounted(); err != nil {
		return err
	}
	defer container.Unmount()

	if b.isUrl(orig) {
		if err := b.addRemote(container, orig, dest); err != nil {
			return err
		}
	} else {
		if err := b.addContext(container, orig, dest); err != nil {
			return err
		}
	}

	if err := b.commit(cid, cmd, fmt.Sprintf("ADD %s in %s", orig, dest)); err != nil {
		return err
	}
	b.config.Cmd = cmd
	return nil
}

func (b *buildFile) run() (string, error) {
	if b.image == "" {
		return "", fmt.Errorf("Please provide a source image with `from` prior to run")
	}
	b.config.Image = b.image

	// Create the container and start it
	c, err := b.builder.Create(b.config)
	if err != nil {
		return "", err
	}
	b.tmpContainers[c.ID] = struct{}{}

	//start the container
	if err := c.Start(); err != nil {
		return "", err
	}

	// Wait for it to finish
	if ret := c.Wait(); ret != 0 {
		return "", fmt.Errorf("The command %v returned a non-zero code: %d", b.config.Cmd, ret)
	}

	return c.ID, nil
}

// Commit the container <id> with the autorun command <autoCmd>
func (b *buildFile) commit(id string, autoCmd []string, comment string) error {
	if b.image == "" {
		return fmt.Errorf("Please provide a source image with `from` prior to commit")
	}
	b.config.Image = b.image
	if id == "" {
		cmd := b.config.Cmd
		b.config.Cmd = []string{"/bin/sh", "-c", "#(nop) " + comment}
		defer func(cmd []string) { b.config.Cmd = cmd }(cmd)

		if cache, err := b.srv.ImageGetCached(b.image, b.config); err != nil {
			return err
		} else if cache != nil {
			utils.Debugf("[BUILDER] Use cached version")
			b.image = cache.ID
			return nil
		} else {
			utils.Debugf("[BUILDER] Cache miss")
		}

		cid, err := b.run()
		if err != nil {
			return err
		}
		id = cid
	}

	container := b.runtime.Get(id)
	if container == nil {
		return fmt.Errorf("An error occured while creating the container")
	}

	// Note: Actually copy the struct
	autoConfig := *b.config
	autoConfig.Cmd = autoCmd
	// Commit the container
	image, err := b.builder.Commit(container, "", "", "", b.maintainer, &autoConfig)
	if err != nil {
		return err
	}
	b.tmpImages[image.ID] = struct{}{}
	b.image = image.ID
	return nil
}

func (b *buildFile) Build(dockerfile, context io.Reader) (string, error) {
	if context != nil {
		name, err := ioutil.TempDir("/tmp", "docker-build")
		if err != nil {
			return "", err
		}
		if err := Untar(context, name); err != nil {
			return "", err
		}
		defer os.RemoveAll(name)
		b.context = name
	}
	file := bufio.NewReader(dockerfile)
	for {
		line, err := file.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return "", err
		}
		line = strings.TrimSpace(strings.Replace(line, "	", " ", -1))
		// Skip comments and empty line
		if len(line) == 0 || line[0] == '#' {
			continue
		}
		tmp := strings.SplitN(line, " ", 2)
		if len(tmp) != 2 {
			return "", fmt.Errorf("Invalid Dockerfile format")
		}
		instruction := strings.ToLower(strings.Trim(tmp[0], " "))
		arguments := strings.Trim(tmp[1], " ")

		fmt.Fprintf(b.out, "%s %s (%s)\n", strings.ToUpper(instruction), arguments, b.image)

		method, exists := reflect.TypeOf(b).MethodByName("Cmd" + strings.ToUpper(instruction[:1]) + strings.ToLower(instruction[1:]))
		if !exists {
			fmt.Fprintf(b.out, "Skipping unknown instruction %s\n", strings.ToUpper(instruction))
			continue
		}
		ret := method.Func.Call([]reflect.Value{reflect.ValueOf(b), reflect.ValueOf(arguments)})[0].Interface()
		if ret != nil {
			return "", ret.(error)
		}

		fmt.Fprintf(b.out, "===> %v\n", b.image)
	}
	if b.image != "" {
		fmt.Fprintf(b.out, "Build successful.\n===> %s\n", b.image)
		return b.image, nil
	}
	return "", fmt.Errorf("An error occured during the build\n")
}

func NewBuildFile(srv *Server, out io.Writer) BuildFile {
	return &buildFile{
		builder:       NewBuilder(srv.runtime),
		runtime:       srv.runtime,
		srv:           srv,
		config:        &Config{},
		out:           out,
		tmpContainers: make(map[string]struct{}),
		tmpImages:     make(map[string]struct{}),
	}
}
<|endoftext|>


<commit_before>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2021-11-25 04:12:42.562465 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "581cd5b5d4939cab6f3abc2a492420a79153d9c5"

// SpartaGitShortHash is the short version of SpartaGitHash
const SpartaGitShortHash = "581cd5b"
<commit_msg>"Autogenerated build info"<commit_after>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2021-12-06 04:39:32.070052 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "ebfe5f0b981aa536a7f07d0d73a70f2ba874c1e2"

// SpartaGitShortHash is the short version of SpartaGitHash
const SpartaGitShortHash = "ebfe5f0"
<|endoftext|>


<commit_before>// Copyright 2019 xgfone
//
// 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 execution

import (
	"context"
	"fmt"
)

func ExampleExecute() {
	if err := Execute(context.TODO(), "ls", "."); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleExecutes() {
	if err := Executes(context.TODO(), []string{"ls"}); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleOutput() {
	if _, err := Output(context.TODO(), "ls", "."); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleOutputs() {
	if _, err := Outputs(context.TODO(), []string{"ls"}); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleHook() {
	filterCmd := func(name string, args ...string) bool {
		// Disable run: rm -rf /
		if name == "rm" && len(args) >= 2 && args[0] == "-rf" && args[1] == "/" {
			return false
		}
		return true
	}
	printCmd := func(name string, args ...string) bool {
		// Print the cmd
		fmt.Printf("Run: %v\n", append([]string{name}, args...))
		return true
	}
	AppendHooks(filterCmd, printCmd)

	RunCmd(context.TODO(), "ls")
	if _, _, err := RunCmd(context.TODO(), "rm", "-rf", "/"); err == ErrDeny {
		fmt.Println(`deny to run "rm -rf /"`)
	}

	// Output:
	// Run: [ls]
	// deny to run "rm -rf /"
}
<commit_msg>fix the test<commit_after>// Copyright 2019 xgfone
//
// 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 execution

import (
	"context"
	"fmt"
)

func ExampleExecute() {
	if err := Execute(context.TODO(), "ls", "."); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleExecutes() {
	if err := Executes(context.TODO(), []string{"ls"}); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleOutput() {
	if _, err := Output(context.TODO(), "ls", "."); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleOutputs() {
	if _, err := Outputs(context.TODO(), []string{"ls"}); err != nil {
		fmt.Println("ERROR")
	} else {
		fmt.Println("OK")
	}

	// Output:
	// OK
}

func ExampleHook() {
	filterCmd := func(name string, args ...string) bool {
		// Disable run: rm -rf /
		if name == "rm" && len(args) >= 2 && args[0] == "-rf" && args[1] == "/" {
			return false
		}
		return true
	}
	printCmd := func(name string, args ...string) bool {
		// Print the cmd
		fmt.Printf("Run: %v\n", append([]string{name}, args...))
		return true
	}
	AppendHooks(filterCmd, printCmd)

	RunCmd(context.TODO(), "ls")
	if _, _, err := RunCmd(context.TODO(), "rm", "-rf", "/"); err.(CmdError).Err == ErrDeny {
		fmt.Println(`deny to run "rm -rf /"`)
	}

	// Output:
	// Run: [ls]
	// deny to run "rm -rf /"
}
<|endoftext|>