<commit_before>package graph

import (
	"fmt"
	"geo"
	"mm"
	"path"
	"sort"
)

type OverlayGraphFile struct {
	*GraphFile
	Cluster          []uint16      // cluster id -> vertex indices
	VertexIndices    []int         // vertex indices -> cluster id
	Matrices         [][][]float32 // transport mode -> metric -> (cluster id, i, j) -> weight
	ClusterEdgeCount int           // combined boundary edge count of the clusters
	EdgeCounts       []int         // cluster id -> id of first edge inside the cluster
}

// I/O

func computeVertexIndices(g *OverlayGraphFile) {
	n, c := g.VertexCount(), 0
	g.VertexIndices = make([]int, n)
	for i := 0; i < n; i++ {
		for c+1 < len(g.Cluster) && i >= int(g.Cluster[c+1]) {
			c++
		}
		g.VertexIndices[i] = c
	}
}

func computeEdgeCounts(g *OverlayGraphFile) {
	g.EdgeCounts = make([]int, g.ClusterCount()+1)
	g.EdgeCounts[0] = g.GraphFile.EdgeCount()
	for i := 0; i < g.ClusterCount(); i++ {
		g.EdgeCounts[i+1] = g.EdgeCounts[i] + g.ClusterSize(i)*g.ClusterSize(i)
	}
}

func loadAllMatrices(g *OverlayGraphFile, base string) error {
	g.Matrices = make([][][]float32, TransportMax)
	for t := 0; t < int(TransportMax); t++ {
		g.Matrices[t] = make([][]float32, MetricMax)
		for m := 0; m < int(MetricMax); m++ {
			var matrixFile []float32
			fileName := fmt.Sprintf("matrices.trans%d.metric%d.ftf", t+1, m+1)
			err := mm.Open(path.Join(base, fileName), &matrixFile)
			if err != nil {
				return err
			}
			g.Matrices[t][m] = matrixFile
		}
	}
	return nil
}

func OpenOverlay(base string, loadMatrices, ignoreErrors bool) (*OverlayGraphFile, error) {
	overlayBaseDir := path.Join(base, "/overlay")
	g, err := OpenGraphFile(overlayBaseDir, ignoreErrors)
	if err != nil && !ignoreErrors {
		return nil, err
	}

	overlay := &OverlayGraphFile{GraphFile: g}
	files := []struct {
		name string
		p    interface{}
	}{
		{"partitions.ftf", &overlay.Cluster},
	}

	for _, file := range files {
		err = mm.Open(path.Join(overlayBaseDir, file.name), file.p)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	computeVertexIndices(overlay)
	computeEdgeCounts(overlay)
	if loadMatrices {
		err = loadAllMatrices(overlay, base)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	for i := 0; i < overlay.ClusterCount(); i++ {
		overlay.ClusterEdgeCount += overlay.ClusterSize(i) * overlay.ClusterSize(i)
	}

	return overlay, nil
}

func CloseOverlay(overlay *OverlayGraphFile) error {
	err := CloseGraphFile(overlay.GraphFile)
	if err != nil {
		return err
	}

	files := []interface{}{
		&overlay.Cluster,
	}

	for _, p := range files {
		err = mm.Close(p)
		if err != nil {
			return err
		}
	}

	return nil
}

// Graph Interface

func (g *OverlayGraphFile) EdgeCount() int {
	// Count edges and matrices...
	return g.GraphFile.EdgeCount() + g.ClusterEdgeCount
}

func (g *OverlayGraphFile) VertexEdges(v Vertex, forward bool, t Transport, buf []Edge) []Edge {
	// Add the cut edges
	result := g.GraphFile.VertexEdges(v, forward, t, buf)
	// Add the precomputed edges.
	cluster, indexInCluster := g.VertexCluster(v)
	clusterStart := g.EdgeCounts[cluster]
	clusterSize := g.ClusterSize(cluster)
	edgesStart := clusterStart + int(indexInCluster)*clusterSize
	for i := 0; i < clusterSize; i++ {
		result = append(result, Edge(edgesStart+i))
	}
	return result
}

func (g *OverlayGraphFile) IsCutEdge(e Edge) bool {
	return int(e) < g.GraphFile.EdgeCount()
}

func (g *OverlayGraphFile) EdgeOpposite(e Edge, v Vertex) Vertex {
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeOpposite(e, v)
	}
	// binary search for cluster id
	cluster := sort.Search(g.ClusterCount(), func(i int) bool { return int(e) < g.EdgeCounts[i+1] })

	e = e - Edge(g.EdgeCounts[cluster])
	vCheck := int(e) / g.ClusterSize(cluster)
	if int(v) != vCheck {
		panic("index of v is not as expected")
	}
	u := int(e) % g.ClusterSize(cluster)
	return Vertex(u)
}

func (g *OverlayGraphFile) EdgeSteps(e Edge, from Vertex) []geo.Coordinate {
	// Return nil unless the edge is a cross partition edge.
	// In this case, defer to the normal Graph interface.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeSteps(e, from)
	}
	return nil
}

func (g *OverlayGraphFile) EdgeWeight(e Edge, t Transport, m Metric) float64 {
	// Return the normal weight if e is a cross partition edge,
	// otherwise return the precomputed weight for t and m.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeWeight(e, t, m)
	}
	edgeIndex := int(e) - g.GraphFile.EdgeCount()
	return float64(g.Matrices[t][m][edgeIndex])
}

// Overlay Interface

func (g *OverlayGraphFile) ClusterCount() int {
	return len(g.Cluster) - 1
}

// actually: boundary vertex count
func (g *OverlayGraphFile) ClusterSize(i int) int {
	// cluster id -> number of vertices
	return int(g.Cluster[i+1] - g.Cluster[i])
}

func (g *OverlayGraphFile) VertexCluster(v Vertex) (int, Vertex) {
	// overlay vertex id -> cluster id, cluster vertex id
	i := g.VertexIndices[v]
	return i, v - Vertex(g.Cluster[i])
}

func (g *OverlayGraphFile) ClusterVertex(i int, v Vertex) Vertex {
	// cluster id, cluster vertex id -> overlay vertex id
	return Vertex(g.Cluster[i]) + v
}
<commit_msg>fixed computeVertexIndices and VertexEdges<commit_after>package graph

import (
	"fmt"
	"geo"
	"mm"
	"path"
	"sort"
)

type OverlayGraphFile struct {
	*GraphFile
	Cluster          []uint16      // cluster id -> vertex indices
	VertexIndices    []int         // vertex indices -> cluster id
	Matrices         [][][]float32 // transport mode -> metric -> (cluster id, i, j) -> weight
	ClusterEdgeCount int           // combined boundary edge count of the clusters
	EdgeCounts       []int         // cluster id -> id of first edge inside the cluster
}

// I/O

func computeVertexIndices(g *OverlayGraphFile) {
	n, c := g.VertexCount()-1, 0
	g.VertexIndices = make([]int, n)
	g.VertexIndices[0] = 0
	for i := 0; i < n; i++ {
		for c+1 < len(g.Cluster) && i+1 >= int(g.Cluster[c+1]) {
			c++
		}
		g.VertexIndices[i+1] = c
	}
}

func computeEdgeCounts(g *OverlayGraphFile) {
	g.EdgeCounts = make([]int, g.ClusterCount()+1)
	g.EdgeCounts[0] = g.GraphFile.EdgeCount()
	for i := 0; i < g.ClusterCount(); i++ {
		g.EdgeCounts[i+1] = g.EdgeCounts[i] + g.ClusterSize(i)*g.ClusterSize(i)
	}
}

func loadAllMatrices(g *OverlayGraphFile, base string) error {
	g.Matrices = make([][][]float32, TransportMax)
	for t := 0; t < int(TransportMax); t++ {
		g.Matrices[t] = make([][]float32, MetricMax)
		for m := 0; m < int(MetricMax); m++ {
			var matrixFile []float32
			fileName := fmt.Sprintf("matrices.trans%d.metric%d.ftf", t+1, m+1)
			err := mm.Open(path.Join(base, fileName), &matrixFile)
			if err != nil {
				return err
			}
			g.Matrices[t][m] = matrixFile
		}
	}
	return nil
}

func OpenOverlay(base string, loadMatrices, ignoreErrors bool) (*OverlayGraphFile, error) {
	overlayBaseDir := path.Join(base, "/overlay")
	g, err := OpenGraphFile(overlayBaseDir, ignoreErrors)
	if err != nil && !ignoreErrors {
		return nil, err
	}

	overlay := &OverlayGraphFile{GraphFile: g}
	files := []struct {
		name string
		p    interface{}
	}{
		{"partitions.ftf", &overlay.Cluster},
	}

	for _, file := range files {
		err = mm.Open(path.Join(overlayBaseDir, file.name), file.p)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	computeVertexIndices(overlay)
	computeEdgeCounts(overlay)
	if loadMatrices {
		err = loadAllMatrices(overlay, base)
		if err != nil && !ignoreErrors {
			return nil, err
		}
	}

	for i := 0; i < overlay.ClusterCount(); i++ {
		overlay.ClusterEdgeCount += overlay.ClusterSize(i) * overlay.ClusterSize(i)
	}

	return overlay, nil
}

func CloseOverlay(overlay *OverlayGraphFile) error {
	err := CloseGraphFile(overlay.GraphFile)
	if err != nil {
		return err
	}

	files := []interface{}{
		&overlay.Cluster,
	}

	for _, p := range files {
		err = mm.Close(p)
		if err != nil {
			return err
		}
	}

	return nil
}

// Graph Interface

func (g *OverlayGraphFile) EdgeCount() int {
	// Count edges and matrices...
	return g.GraphFile.EdgeCount() + g.ClusterEdgeCount
}

func (g *OverlayGraphFile) VertexEdges(v Vertex, forward bool, t Transport, buf []Edge) []Edge {
	// Add the cut edges
	result := g.GraphFile.VertexEdges(v, forward, t, buf)
	// Add the precomputed edges.
	cluster, indexInCluster := g.VertexCluster(v)
	clusterStart := g.EdgeCounts[cluster]
	clusterSize := g.ClusterSize(cluster)
	if forward {
		// out edges
		outEdgesStart := clusterStart + int(indexInCluster)*clusterSize
		for i := 0; i < clusterSize; i++ {
			result = append(result, Edge(outEdgesStart+i))
		}
	} else {
		// in edges
		for i := 0; i < clusterSize; i++ {
			result = append(result, Edge(i*clusterSize+int(indexInCluster)))
		}
	}
	return result
}

func (g *OverlayGraphFile) IsCutEdge(e Edge) bool {
	return int(e) < g.GraphFile.EdgeCount()
}

func (g *OverlayGraphFile) EdgeOpposite(e Edge, v Vertex) Vertex {
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeOpposite(e, v)
	}
	// binary search for cluster id
	cluster := sort.Search(g.ClusterCount(), func(i int) bool { return int(e) < g.EdgeCounts[i+1] })

	e = e - Edge(g.EdgeCounts[cluster])
	vCheck := int(e) / g.ClusterSize(cluster)
	if int(v) != vCheck {
		panic("index of v is not as expected")
	}
	u := int(e) % g.ClusterSize(cluster)
	return Vertex(u)
}

func (g *OverlayGraphFile) EdgeSteps(e Edge, from Vertex) []geo.Coordinate {
	// Return nil unless the edge is a cross partition edge.
	// In this case, defer to the normal Graph interface.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeSteps(e, from)
	}
	return nil
}

func (g *OverlayGraphFile) EdgeWeight(e Edge, t Transport, m Metric) float64 {
	// Return the normal weight if e is a cross partition edge,
	// otherwise return the precomputed weight for t and m.
	if g.IsCutEdge(e) {
		return g.GraphFile.EdgeWeight(e, t, m)
	}
	edgeIndex := int(e) - g.GraphFile.EdgeCount()
	return float64(g.Matrices[t][m][edgeIndex])
}

// Overlay Interface

func (g *OverlayGraphFile) ClusterCount() int {
	return len(g.Cluster) - 1
}

// actually: boundary vertex count
func (g *OverlayGraphFile) ClusterSize(i int) int {
	// cluster id -> number of vertices
	return int(g.Cluster[i+1] - g.Cluster[i])
}

func (g *OverlayGraphFile) VertexCluster(v Vertex) (int, Vertex) {
	// overlay vertex id -> cluster id, cluster vertex id
	i := g.VertexIndices[v]
	return i, v - Vertex(g.Cluster[i])
}

func (g *OverlayGraphFile) ClusterVertex(i int, v Vertex) Vertex {
	// cluster id, cluster vertex id -> overlay vertex id
	return Vertex(g.Cluster[i]) + v
}
<|endoftext|>


<commit_before>package uploader

import (
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/pivotal-golang/lager"
)

type Uploader interface {
	Upload(fileLocation string, destinationUrl *url.URL, logger lager.Logger) (int64, error)
}

type URLUploader struct {
	timeout time.Duration
}

func New(timeout time.Duration) Uploader {
	return &URLUploader{
		timeout: timeout,
	}
}

func (uploader *URLUploader) Upload(fileLocation string, url *url.URL, logger lager.Logger) (int64, error) {
	httpTransport := &http.Transport{
		ResponseHeaderTimeout: uploader.timeout,
	}
	httpClient := &http.Client{
		Transport: httpTransport,
	}

	var resp *http.Response
	var err error
	var uploadedBytes int64

	for attempt := 0; attempt < 3; attempt++ {
		var request *http.Request
		var sourceFile *os.File
		var fileInfo os.FileInfo

		if logger != nil {
			logger.Info("uploader.attempt", lager.Data{
				"attempt": attempt,
			})
		}

		sourceFile, err = os.Open(fileLocation)
		if err != nil {
			return 0, err
		}

		fileInfo, err = sourceFile.Stat()
		if err != nil {
			return 0, err
		}

		contentHash := md5.New()
		_, err = io.Copy(contentHash, sourceFile)
		if err != nil {
			return 0, err
		}

		_, err = sourceFile.Seek(0, 0)
		if err != nil {
			return 0, err
		}

		request, err = http.NewRequest("POST", url.String(), sourceFile)

		uploadedBytes = fileInfo.Size()
		request.ContentLength = uploadedBytes
		request.Header.Set("Content-Type", "application/octet-stream")
		request.Header.Set("Content-MD5", base64.StdEncoding.EncodeToString(contentHash.Sum(nil)))

		resp, err = httpClient.Do(request)
		if err == nil {
			defer resp.Body.Close()
			break
		}
	}

	if err != nil {
		return 0, err
	}

	if resp.StatusCode >= 400 {
		return 0, fmt.Errorf("Upload failed: Status code %d", resp.StatusCode)
	}

	return int64(uploadedBytes), nil
}
<commit_msg>Fix goroutine leak associated with http client<commit_after>package uploader

import (
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/pivotal-golang/lager"
)

type Uploader interface {
	Upload(fileLocation string, destinationUrl *url.URL, logger lager.Logger) (int64, error)
}

type URLUploader struct {
	httpClient *http.Client
}

func New(timeout time.Duration) Uploader {
	httpTransport := &http.Transport{
		ResponseHeaderTimeout: timeout,
	}

	httpClient := &http.Client{
		Transport: httpTransport,
	}

	return &URLUploader{
		httpClient: httpClient,
	}
}

func (uploader *URLUploader) Upload(fileLocation string, url *url.URL, logger lager.Logger) (int64, error) {
	var resp *http.Response
	var err error
	var uploadedBytes int64

	for attempt := 0; attempt < 3; attempt++ {
		var request *http.Request
		var sourceFile *os.File
		var fileInfo os.FileInfo

		if logger != nil {
			logger.Info("uploader.attempt", lager.Data{
				"attempt": attempt,
			})
		}

		sourceFile, err = os.Open(fileLocation)
		if err != nil {
			return 0, err
		}

		fileInfo, err = sourceFile.Stat()
		if err != nil {
			return 0, err
		}

		contentHash := md5.New()
		_, err = io.Copy(contentHash, sourceFile)
		if err != nil {
			return 0, err
		}

		_, err = sourceFile.Seek(0, 0)
		if err != nil {
			return 0, err
		}

		request, err = http.NewRequest("POST", url.String(), sourceFile)

		uploadedBytes = fileInfo.Size()
		request.ContentLength = uploadedBytes
		request.Header.Set("Content-Type", "application/octet-stream")
		request.Header.Set("Content-MD5", base64.StdEncoding.EncodeToString(contentHash.Sum(nil)))

		resp, err = uploader.httpClient.Do(request)
		if err == nil {
			defer resp.Body.Close()
			break
		}
	}

	if err != nil {
		return 0, err
	}

	if resp.StatusCode >= 400 {
		return 0, fmt.Errorf("Upload failed: Status code %d", resp.StatusCode)
	}

	return int64(uploadedBytes), nil
}
<|endoftext|>


<commit_before>package grpc

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

	"github.com/gogo/protobuf/proto"
	grpc_logging "github.com/grpc-ecosystem/go-grpc-middleware/logging"
	grpc_logrus "github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus"
	ctx_logrus "github.com/grpc-ecosystem/go-grpc-middleware/tags/logrus"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

func logRequest(entry *logrus.Entry, info string, pbMsg interface{}, ctx context.Context, logClaims bool) {
	if logClaims {
		if data, err := json.Marshal(ctx.Value("claims")); err == nil {
			entry = entry.WithField("grpc.request.claims", string(data))
		}
	}
	if p, ok := pbMsg.(proto.Message); ok {
		entry = entry.WithField("grpc.request.content", &jsonpbMarshalleble{p})
	}
	entry.Info(info)
}

type jsonpbMarshalleble struct {
	proto.Message
}

func (j *jsonpbMarshalleble) MarshalJSON() ([]byte, error) {
	b := &bytes.Buffer{}
	if err := grpc_logrus.JsonPbMarshaller.Marshal(b, j.Message); err != nil {
		return nil, fmt.Errorf("jsonpb serializer failed: %v", err)
	}
	return b.Bytes(), nil
}

type loggingServerStream struct {
	grpc.ServerStream
	entry     *logrus.Entry
	logClaims bool
	info      string
}

func (l *loggingServerStream) SendMsg(m interface{}) error {
	return l.ServerStream.SendMsg(m)
}

func (l *loggingServerStream) RecvMsg(m interface{}) error {
	err := l.ServerStream.RecvMsg(m)
	if err == nil {
		logRequest(l.entry, l.info, m, l.ServerStream.Context(), l.logClaims)
	}
	return err
}

func PayloadStreamServerInterceptor(entry *logrus.Entry, logClaims bool, decider grpc_logging.ServerPayloadLoggingDecider) grpc.StreamServerInterceptor {
	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if !decider(stream.Context(), info.FullMethod, srv) {
			return handler(srv, stream)
		}
		logEntry := entry.WithFields(ctx_logrus.Extract(stream.Context()).Data)
		newStream := &loggingServerStream{ServerStream: stream, entry: logEntry, logClaims: logClaims, info: fmt.Sprintf("received streaming call %s", info.FullMethod)}
		return handler(srv, newStream)
	}
}

func PayloadUnaryServerInterceptor(entry *logrus.Entry, logClaims bool, decider grpc_logging.ServerPayloadLoggingDecider) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		if !decider(ctx, info.FullMethod, info.Server) {
			return handler(ctx, req)
		}
		logEntry := entry.WithFields(ctx_logrus.Extract(ctx).Data)
		logRequest(logEntry, fmt.Sprintf("received unary call %s", info.FullMethod), req, ctx, logClaims)
		resp, err := handler(ctx, req)
		return resp, err
	}
}
<commit_msg>fix(logging.go): changing marshaler for JSON logging to use gogo (#4319)<commit_after>package grpc

import (
	"encoding/json"
	"fmt"

	"github.com/gogo/protobuf/proto"
	grpc_logging "github.com/grpc-ecosystem/go-grpc-middleware/logging"
	ctx_logrus "github.com/grpc-ecosystem/go-grpc-middleware/tags/logrus"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

func logRequest(entry *logrus.Entry, info string, pbMsg interface{}, ctx context.Context, logClaims bool) {
	if logClaims {
		if data, err := json.Marshal(ctx.Value("claims")); err == nil {
			entry = entry.WithField("grpc.request.claims", string(data))
		}
	}
	if p, ok := pbMsg.(proto.Message); ok {
		entry = entry.WithField("grpc.request.content", &jsonpbMarshalleble{p})
	}
	entry.Info(info)
}

type jsonpbMarshalleble struct {
	proto.Message
}

func (j *jsonpbMarshalleble) MarshalJSON() ([]byte, error) {
	b, err := proto.Marshal(j.Message)
	if err != nil {
		return nil, fmt.Errorf("jsonpb serializer failed: %v", err)
	}
	return b, nil
}

type loggingServerStream struct {
	grpc.ServerStream
	entry     *logrus.Entry
	logClaims bool
	info      string
}

func (l *loggingServerStream) SendMsg(m interface{}) error {
	return l.ServerStream.SendMsg(m)
}

func (l *loggingServerStream) RecvMsg(m interface{}) error {
	err := l.ServerStream.RecvMsg(m)
	if err == nil {
		logRequest(l.entry, l.info, m, l.ServerStream.Context(), l.logClaims)
	}
	return err
}

func PayloadStreamServerInterceptor(entry *logrus.Entry, logClaims bool, decider grpc_logging.ServerPayloadLoggingDecider) grpc.StreamServerInterceptor {
	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if !decider(stream.Context(), info.FullMethod, srv) {
			return handler(srv, stream)
		}
		logEntry := entry.WithFields(ctx_logrus.Extract(stream.Context()).Data)
		newStream := &loggingServerStream{ServerStream: stream, entry: logEntry, logClaims: logClaims, info: fmt.Sprintf("received streaming call %s", info.FullMethod)}
		return handler(srv, newStream)
	}
}

func PayloadUnaryServerInterceptor(entry *logrus.Entry, logClaims bool, decider grpc_logging.ServerPayloadLoggingDecider) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		if !decider(ctx, info.FullMethod, info.Server) {
			return handler(ctx, req)
		}
		logEntry := entry.WithFields(ctx_logrus.Extract(ctx).Data)
		logRequest(logEntry, fmt.Sprintf("received unary call %s", info.FullMethod), req, ctx, logClaims)
		resp, err := handler(ctx, req)
		return resp, err
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package stop

import (
	"fmt"
	"log"
	"sort"
	"strings"
	"sync"

	"github.com/cockroachdb/cockroach/util/caller"
)

// Closer is an interface for objects to attach to the stopper to
// be closed once the stopper completes.
type Closer interface {
	Close()
}

// CloserFn is type that allows any function to be a Closer.
type CloserFn func()

// Close implements the Closer interface.
func (f CloserFn) Close() {
	f()
}

type taskKey struct {
	file string
	line int
}

func (k taskKey) String() string {
	return fmt.Sprintf("%s:%d", k.file, k.line)
}

// A Stopper provides a channel-based mechanism to stop an arbitrary
// array of workers. Each worker is registered with the stopper via
// the RunWorker() method. The system further allows execution of functions
// through RunTask() and RunAsyncTask().
//
// Stopping occurs in two phases: the first is the request to stop, which moves
// the stopper into a draining phase. While draining, calls to RunTask() &
// RunAsyncTask() don't execute the function passed in and return false.
// When all outstanding tasks have been completed, the stopper
// closes its stopper channel, which signals all live workers that it's safe to
// shut down. When all workers have shutdown, the stopper is complete.
//
// An arbitrary list of objects implementing the Closer interface may
// be added to the stopper via AddCloser(), to be closed after the
// stopper has stopped.
type Stopper struct {
	drainer  chan struct{}  // Closed when draining
	stopper  chan struct{}  // Closed when stopping
	stopped  chan struct{}  // Closed when stopped completely
	stop     sync.WaitGroup // Incremented for outstanding workers
	mu       sync.Mutex     // Protects the fields below
	drain    *sync.Cond     // Conditional variable to wait for outstanding tasks
	draining bool           // true when Stop() has been called
	numTasks int            // number of outstanding tasks
	tasks    map[taskKey]int
	closers  []Closer
}

// NewStopper returns an instance of Stopper.
func NewStopper() *Stopper {
	s := &Stopper{
		drainer: make(chan struct{}),
		stopper: make(chan struct{}),
		stopped: make(chan struct{}),
		tasks:   map[taskKey]int{},
	}
	s.drain = sync.NewCond(&s.mu)
	return s
}

// RunWorker runs the supplied function as a "worker" to be stopped
// by the stopper. The function <f> is run in a goroutine.
func (s *Stopper) RunWorker(f func()) {
	s.stop.Add(1)
	go func() {
		defer s.stop.Done()
		f()
	}()
}

// AddCloser adds an object to close after the stopper has been stopped.
func (s *Stopper) AddCloser(c Closer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.closers = append(s.closers, c)
}

// RunTask adds one to the count of tasks left to drain in the system. Any
// worker which is a "first mover" when starting tasks must call this method
// before starting work on a new task. First movers include
// goroutines launched to do periodic work and the kv/db.go gateway which
// accepts external client requests.
//
// Returns false to indicate that the system is currently draining and
// function f was not called.
func (s *Stopper) RunTask(f func()) bool {
	file, line, _ := caller.Lookup(1)
	key := taskKey{file, line}
	if !s.runPrelude(key) {
		return false
	}
	// Call f.
	defer s.runPostlude(key)
	f()
	return true
}

// RunAsyncTask runs function f in a goroutine. It returns false when the
// Stopper is draining and the function is not executed.
func (s *Stopper) RunAsyncTask(f func()) bool {
	file, line, _ := caller.Lookup(1)
	key := taskKey{file, line}
	if !s.runPrelude(key) {
		return false
	}
	// Call f.
	go func() {
		defer s.runPostlude(key)
		f()
	}()
	return true
}

func (s *Stopper) runPrelude(key taskKey) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.draining {
		return false
	}
	s.numTasks++
	s.tasks[key]++
	return true
}

func (s *Stopper) runPostlude(key taskKey) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.numTasks--
	s.tasks[key]--
	s.drain.Broadcast()
}

// NumTasks returns the number of active tasks.
func (s *Stopper) NumTasks() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.numTasks
}

// A TaskMap is returned by RunningTasks().
type TaskMap map[string]int

// String implements fmt.Stringer and returns a sorted multi-line listing of
// the TaskMap.
func (tm TaskMap) String() string {
	var lines []string
	for location, num := range tm {
		lines = append(lines, fmt.Sprintf("%-6d %s", num, location))
	}
	sort.Sort(sort.Reverse(sort.StringSlice(lines)))
	return strings.Join(lines, "\n")
}

// RunningTasks returns a map containing the count of running tasks keyed by
// call site.
func (s *Stopper) RunningTasks() TaskMap {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.runningTasksLocked()
}

func (s *Stopper) runningTasksLocked() TaskMap {
	m := map[string]int{}
	for k := range s.tasks {
		if s.tasks[k] == 0 {
			continue
		}
		m[k.String()] = s.tasks[k]
	}
	return m
}

// Stop signals all live workers to stop and then waits for each to
// confirm it has stopped.
func (s *Stopper) Stop() {
	// Don't bother doing stuff cleanly if we're panicking, that would likely
	// block. Instead, best effort only. This cleans up the stack traces,
	// avoids stalls and helps some tests in `./cli` finish cleanly (where
	// panics happen on purpose).
	if r := recover(); r != nil {
		go s.Quiesce()
		close(s.stopper)
		close(s.stopped)
		for _, c := range s.closers {
			go c.Close()
		}
		panic(r)
	}

	s.Quiesce()
	close(s.stopper)
	s.stop.Wait()
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, c := range s.closers {
		c.Close()
	}
	close(s.stopped)
}

// ShouldDrain returns a channel which will be closed when Stop() has been
// invoked and outstanding tasks should begin to drain.
func (s *Stopper) ShouldDrain() <-chan struct{} {
	if s == nil {
		// A nil stopper will never signal ShouldDrain, but will also never panic.
		return nil
	}
	return s.drainer
}

// ShouldStop returns a channel which will be closed when Stop() has been
// invoked and outstanding tasks have drained.
func (s *Stopper) ShouldStop() <-chan struct{} {
	if s == nil {
		// A nil stopper will never signal ShouldStop, but will also never panic.
		return nil
	}
	return s.stopper
}

// IsStopped returns a channel which will be closed after Stop() has
// been invoked to full completion, meaning all workers have completed
// and all closers have been closed.
func (s *Stopper) IsStopped() <-chan struct{} {
	if s == nil {
		return nil
	}
	return s.stopped
}

// Quiesce moves the stopper to state draining and waits until all
// tasks complete. This is used from Stop() and unittests.
func (s *Stopper) Quiesce() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.draining {
		s.draining = true
		close(s.drainer)
	}
	for s.numTasks > 0 {
		// Use stdlib "log" instead of "cockroach/util/log" due to import cycles.
		log.Print("draining; tasks left:\n", s.runningTasksLocked())
		// Unlock s.mu, wait for the signal, and lock s.mu.
		s.drain.Wait()
	}
}
<commit_msg>Add /debug/stopper endpoint<commit_after>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package stop

import (
	"fmt"
	"log"
	"net/http"
	"sort"
	"strings"
	"sync"

	"github.com/cockroachdb/cockroach/util/caller"
)

func register(s *Stopper) {
	trackedStoppers.Lock()
	trackedStoppers.stoppers = append(trackedStoppers.stoppers, s)
	trackedStoppers.Unlock()
}

func unregister(s *Stopper) {
	trackedStoppers.Lock()
	defer trackedStoppers.Unlock()
	sl := trackedStoppers.stoppers
	for i, tracked := range sl {
		if tracked == s {
			trackedStoppers.stoppers = sl[:i+copy(sl[i:], sl[i+1:])]
			return
		}
	}
	panic("attempt to unregister untracked stopper")
}

var trackedStoppers struct {
	sync.Mutex
	stoppers []*Stopper
}

func handleDebug(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	trackedStoppers.Lock()
	defer trackedStoppers.Unlock()
	for _, s := range trackedStoppers.stoppers {
		s.mu.Lock()
		fmt.Fprintf(w, "%p: %d tasks\n%s", s, s.numTasks, s.runningTasksLocked())
		s.mu.Unlock()
	}
}

func init() {
	http.Handle("/debug/stopper", http.HandlerFunc(handleDebug))
}

// Closer is an interface for objects to attach to the stopper to
// be closed once the stopper completes.
type Closer interface {
	Close()
}

// CloserFn is type that allows any function to be a Closer.
type CloserFn func()

// Close implements the Closer interface.
func (f CloserFn) Close() {
	f()
}

type taskKey struct {
	file string
	line int
}

func (k taskKey) String() string {
	return fmt.Sprintf("%s:%d", k.file, k.line)
}

// A Stopper provides a channel-based mechanism to stop an arbitrary
// array of workers. Each worker is registered with the stopper via
// the RunWorker() method. The system further allows execution of functions
// through RunTask() and RunAsyncTask().
//
// Stopping occurs in two phases: the first is the request to stop, which moves
// the stopper into a draining phase. While draining, calls to RunTask() &
// RunAsyncTask() don't execute the function passed in and return false.
// When all outstanding tasks have been completed, the stopper
// closes its stopper channel, which signals all live workers that it's safe to
// shut down. When all workers have shutdown, the stopper is complete.
//
// An arbitrary list of objects implementing the Closer interface may
// be added to the stopper via AddCloser(), to be closed after the
// stopper has stopped.
type Stopper struct {
	drainer  chan struct{}  // Closed when draining
	stopper  chan struct{}  // Closed when stopping
	stopped  chan struct{}  // Closed when stopped completely
	stop     sync.WaitGroup // Incremented for outstanding workers
	mu       sync.Mutex     // Protects the fields below
	drain    *sync.Cond     // Conditional variable to wait for outstanding tasks
	draining bool           // true when Stop() has been called
	numTasks int            // number of outstanding tasks
	tasks    map[taskKey]int
	closers  []Closer
}

// NewStopper returns an instance of Stopper.
func NewStopper() *Stopper {
	s := &Stopper{
		drainer: make(chan struct{}),
		stopper: make(chan struct{}),
		stopped: make(chan struct{}),
		tasks:   map[taskKey]int{},
	}
	s.drain = sync.NewCond(&s.mu)
	register(s)
	return s
}

// RunWorker runs the supplied function as a "worker" to be stopped
// by the stopper. The function <f> is run in a goroutine.
func (s *Stopper) RunWorker(f func()) {
	s.stop.Add(1)
	go func() {
		defer s.stop.Done()
		f()
	}()
}

// AddCloser adds an object to close after the stopper has been stopped.
func (s *Stopper) AddCloser(c Closer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.closers = append(s.closers, c)
}

// RunTask adds one to the count of tasks left to drain in the system. Any
// worker which is a "first mover" when starting tasks must call this method
// before starting work on a new task. First movers include
// goroutines launched to do periodic work and the kv/db.go gateway which
// accepts external client requests.
//
// Returns false to indicate that the system is currently draining and
// function f was not called.
func (s *Stopper) RunTask(f func()) bool {
	file, line, _ := caller.Lookup(1)
	key := taskKey{file, line}
	if !s.runPrelude(key) {
		return false
	}
	// Call f.
	defer s.runPostlude(key)
	f()
	return true
}

// RunAsyncTask runs function f in a goroutine. It returns false when the
// Stopper is draining and the function is not executed.
func (s *Stopper) RunAsyncTask(f func()) bool {
	file, line, _ := caller.Lookup(1)
	key := taskKey{file, line}
	if !s.runPrelude(key) {
		return false
	}
	// Call f.
	go func() {
		defer s.runPostlude(key)
		f()
	}()
	return true
}

func (s *Stopper) runPrelude(key taskKey) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.draining {
		return false
	}
	s.numTasks++
	s.tasks[key]++
	return true
}

func (s *Stopper) runPostlude(key taskKey) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.numTasks--
	s.tasks[key]--
	s.drain.Broadcast()
}

// NumTasks returns the number of active tasks.
func (s *Stopper) NumTasks() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.numTasks
}

// A TaskMap is returned by RunningTasks().
type TaskMap map[string]int

// String implements fmt.Stringer and returns a sorted multi-line listing of
// the TaskMap.
func (tm TaskMap) String() string {
	var lines []string
	for location, num := range tm {
		lines = append(lines, fmt.Sprintf("%-6d %s", num, location))
	}
	sort.Sort(sort.Reverse(sort.StringSlice(lines)))
	return strings.Join(lines, "\n")
}

// RunningTasks returns a map containing the count of running tasks keyed by
// call site.
func (s *Stopper) RunningTasks() TaskMap {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.runningTasksLocked()
}

func (s *Stopper) runningTasksLocked() TaskMap {
	m := map[string]int{}
	for k := range s.tasks {
		if s.tasks[k] == 0 {
			continue
		}
		m[k.String()] = s.tasks[k]
	}
	return m
}

// Stop signals all live workers to stop and then waits for each to
// confirm it has stopped.
func (s *Stopper) Stop() {
	defer unregister(s)
	// Don't bother doing stuff cleanly if we're panicking, that would likely
	// block. Instead, best effort only. This cleans up the stack traces,
	// avoids stalls and helps some tests in `./cli` finish cleanly (where
	// panics happen on purpose).
	if r := recover(); r != nil {
		go s.Quiesce()
		close(s.stopper)
		close(s.stopped)
		for _, c := range s.closers {
			go c.Close()
		}
		panic(r)
	}

	s.Quiesce()
	close(s.stopper)
	s.stop.Wait()
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, c := range s.closers {
		c.Close()
	}
	close(s.stopped)
}

// ShouldDrain returns a channel which will be closed when Stop() has been
// invoked and outstanding tasks should begin to drain.
func (s *Stopper) ShouldDrain() <-chan struct{} {
	if s == nil {
		// A nil stopper will never signal ShouldDrain, but will also never panic.
		return nil
	}
	return s.drainer
}

// ShouldStop returns a channel which will be closed when Stop() has been
// invoked and outstanding tasks have drained.
func (s *Stopper) ShouldStop() <-chan struct{} {
	if s == nil {
		// A nil stopper will never signal ShouldStop, but will also never panic.
		return nil
	}
	return s.stopper
}

// IsStopped returns a channel which will be closed after Stop() has
// been invoked to full completion, meaning all workers have completed
// and all closers have been closed.
func (s *Stopper) IsStopped() <-chan struct{} {
	if s == nil {
		return nil
	}
	return s.stopped
}

// Quiesce moves the stopper to state draining and waits until all
// tasks complete. This is used from Stop() and unittests.
func (s *Stopper) Quiesce() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.draining {
		s.draining = true
		close(s.drainer)
	}
	for s.numTasks > 0 {
		// Use stdlib "log" instead of "cockroach/util/log" due to import cycles.
		log.Print("draining; tasks left:\n", s.runningTasksLocked())
		// Unlock s.mu, wait for the signal, and lock s.mu.
		s.drain.Wait()
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package util

import (
	"net"
	"testing"
	"time"
)

// verifyAddr starts a server listener at the specified addr and
// then dials a client to verify a connection is established.
func verifyAddr(addr net.Addr, t *testing.T) {
	ln, err := net.Listen(addr.Network(), addr.String())
	if err != nil {
		t.Error(err)
		return
	}

	acceptChan := make(chan struct{})
	go func() {
		_, err := ln.Accept()
		if err != nil {
			t.Error(err)
		}
		close(acceptChan)
	}()

	addr = ln.Addr()
	conn, err := net.Dial(addr.Network(), addr.String())
	if err != nil {
		t.Errorf("could not connect to %s", addr)
		return
	}
	select {
	case <-acceptChan:
		// success.
	case <-time.After(500 * time.Millisecond):
		t.Error("timeout waiting for client connection after 500ms")
	}
	conn.Close()
}

func TestCreateTestAddr(t *testing.T) {
	verifyAddr(CreateTestAddr("unix"), t)
	verifyAddr(CreateTestAddr("tcp"), t)
}

func TestTrueWithin(t *testing.T) {
	// Start with a function that never returns true.
	if err := IsTrueWithin(func() bool { return false }, 1*time.Millisecond); err == nil {
		t.Error("expected true within to fail on method which always returns false")
	}
	// Try a method which always returns true.
	if err := IsTrueWithin(func() bool { return true }, 1*time.Millisecond); err != nil {
		t.Errorf("unexpected error on method which always returns true: %v", err)
	}
	// Try a method which returns true on 5th invocation.
	count := 0
	if err := IsTrueWithin(func() bool { count++; return count >= 5 }, 1*time.Millisecond); err != nil {
		t.Errorf("unexpected error on method which returns true after 5 invocations: %v", err)
	}
}

func TestSucceedsWithin(t *testing.T) {
	// Try a method which always succeeds.
	SucceedsWithin(t, 1*time.Millisecond, func() error { return nil })

	// Try a method which succeeds after a known duration.
	start := time.Now()
	duration := time.Millisecond
	SucceedsWithin(t, 100*duration, func() error {
		elapsed := time.Since(start)
		if elapsed > duration {
			return nil
		}
		return Errorf("%s elapsed, waiting until %s elapses", elapsed, duration)
	})
}

func TestNoZeroField(t *testing.T) {
	type foo struct {
		X, Y int
	}
	testFoo := foo{1, 2}
	if err := NoZeroField(&testFoo); err != nil {
		t.Fatal(err)
	}
	if err := NoZeroField(interface{}(testFoo)); err != nil {
		t.Fatal(err)
	}
	testFoo.Y = 0
	if err := NoZeroField(&testFoo); err == nil {
		t.Fatal("expected an error")
	}
}
<commit_msg>Increase duration for SucceedsWithin test.<commit_after>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package util

import (
	"net"
	"testing"
	"time"
)

// verifyAddr starts a server listener at the specified addr and
// then dials a client to verify a connection is established.
func verifyAddr(addr net.Addr, t *testing.T) {
	ln, err := net.Listen(addr.Network(), addr.String())
	if err != nil {
		t.Error(err)
		return
	}

	acceptChan := make(chan struct{})
	go func() {
		_, err := ln.Accept()
		if err != nil {
			t.Error(err)
		}
		close(acceptChan)
	}()

	addr = ln.Addr()
	conn, err := net.Dial(addr.Network(), addr.String())
	if err != nil {
		t.Errorf("could not connect to %s", addr)
		return
	}
	select {
	case <-acceptChan:
		// success.
	case <-time.After(500 * time.Millisecond):
		t.Error("timeout waiting for client connection after 500ms")
	}
	conn.Close()
}

func TestCreateTestAddr(t *testing.T) {
	verifyAddr(CreateTestAddr("unix"), t)
	verifyAddr(CreateTestAddr("tcp"), t)
}

func TestTrueWithin(t *testing.T) {
	// Start with a function that never returns true.
	if err := IsTrueWithin(func() bool { return false }, 1*time.Millisecond); err == nil {
		t.Error("expected true within to fail on method which always returns false")
	}
	// Try a method which always returns true.
	if err := IsTrueWithin(func() bool { return true }, 1*time.Millisecond); err != nil {
		t.Errorf("unexpected error on method which always returns true: %v", err)
	}
	// Try a method which returns true on 5th invocation.
	count := 0
	if err := IsTrueWithin(func() bool { count++; return count >= 5 }, 1*time.Millisecond); err != nil {
		t.Errorf("unexpected error on method which returns true after 5 invocations: %v", err)
	}
}

func TestSucceedsWithin(t *testing.T) {
	// Try a method which always succeeds.
	SucceedsWithin(t, 1*time.Millisecond, func() error { return nil })

	// Try a method which succeeds after a known duration.
	start := time.Now()
	duration := time.Millisecond * 10
	SucceedsWithin(t, 100*duration, func() error {
		elapsed := time.Since(start)
		if elapsed > duration {
			return nil
		}
		return Errorf("%s elapsed, waiting until %s elapses", elapsed, duration)
	})
}

func TestNoZeroField(t *testing.T) {
	type foo struct {
		X, Y int
	}
	testFoo := foo{1, 2}
	if err := NoZeroField(&testFoo); err != nil {
		t.Fatal(err)
	}
	if err := NoZeroField(interface{}(testFoo)); err != nil {
		t.Fatal(err)
	}
	testFoo.Y = 0
	if err := NoZeroField(&testFoo); err == nil {
		t.Fatal("expected an error")
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"log"
	"os"

	"github.com/docker/machine/libmachine/drivers/plugin/localbinary"
	"github.com/tsuru/tsuru-client/tsuru/admin"
	"github.com/tsuru/tsuru-client/tsuru/client"
	"github.com/tsuru/tsuru-client/tsuru/installer"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/iaas/dockermachine"
	_ "github.com/tsuru/tsuru/provision/docker/cmds"
)

const (
	version = "1.5.0-rc1"
	header  = "Supported-Tsuru"
)

func buildManager(name string) *cmd.Manager {
	lookup := func(context *cmd.Context) error {
		return client.RunPlugin(context)
	}
	m := cmd.BuildBaseManager(name, version, header, lookup)
	m.Register(&client.AppRun{})
	m.Register(&client.AppInfo{})
	m.Register(&client.AppCreate{})
	m.Register(&client.AppRemove{})
	m.Register(&client.AppUpdate{})
	m.Register(&client.UnitAdd{})
	m.Register(&client.UnitRemove{})
	m.Register(&client.AppList{})
	m.Register(&client.AppLog{})
	m.Register(&client.AppGrant{})
	m.Register(&client.AppRevoke{})
	m.Register(&client.AppRestart{})
	m.Register(&client.AppStart{})
	m.Register(&client.AppStop{})
	m.Register(&admin.AppLockDelete{})
	m.Register(&client.CertificateSet{})
	m.Register(&client.CertificateUnset{})
	m.Register(&client.CertificateList{})
	m.Register(&client.CnameAdd{})
	m.Register(&client.CnameRemove{})
	m.Register(&client.EnvGet{})
	m.Register(&client.EnvSet{})
	m.Register(&client.EnvUnset{})
	m.Register(&client.KeyAdd{})
	m.Register(&client.KeyRemove{})
	m.Register(&client.KeyList{})
	m.Register(client.ServiceList{})
	m.Register(&client.ServiceInstanceAdd{})
	m.Register(&client.ServiceInstanceUpdate{})
	m.Register(&client.ServiceInstanceRemove{})
	m.Register(client.ServiceInfo{})
	m.Register(client.ServiceInstanceInfo{})
	m.Register(client.ServiceInstanceStatus{})
	m.Register(&client.ServiceInstanceGrant{})
	m.Register(&client.ServiceInstanceRevoke{})
	m.Register(&client.ServiceInstanceBind{})
	m.Register(&client.ServiceInstanceUnbind{})
	m.Register(&admin.PlatformList{})
	m.Register(&admin.PlatformAdd{})
	m.Register(&admin.PlatformUpdate{})
	m.Register(&admin.PlatformRemove{})
	m.Register(&client.PluginInstall{})
	m.Register(&client.PluginRemove{})
	m.Register(&client.PluginList{})
	m.Register(&client.AppSwap{})
	m.Register(&client.AppDeploy{})
	m.Register(&client.AppBuild{})
	m.Register(&client.PlanList{})
	m.Register(&client.UserCreate{})
	m.Register(&client.ResetPassword{})
	m.Register(&client.UserRemove{})
	m.Register(&client.ListUsers{})
	m.Register(&client.TeamCreate{})
	m.Register(&client.TeamUpdate{})
	m.Register(&client.TeamRemove{})
	m.Register(&client.TeamList{})
	m.Register(&client.TeamInfo{})
	m.Register(&client.ChangePassword{})
	m.Register(&client.ShowAPIToken{})
	m.Register(&client.RegenerateAPIToken{})
	m.Register(&client.AppDeployList{})
	m.Register(&client.AppDeployRollback{})
	m.Register(&client.AppDeployRollbackUpdate{})
	m.Register(&client.AppDeployRebuild{})
	m.Register(&cmd.ShellToContainerCmd{})
	m.Register(&client.PoolList{})
	m.Register(&client.PermissionList{})
	m.Register(&client.RoleAdd{})
	m.Register(&client.RoleUpdate{})
	m.Register(&client.RoleRemove{})
	m.Register(&client.RoleList{})
	m.Register(&client.RoleInfo{})
	m.Register(&client.RolePermissionAdd{})
	m.Register(&client.RolePermissionRemove{})
	m.Register(&client.RoleAssign{})
	m.Register(&client.RoleDissociate{})
	m.Register(&client.RoleDefaultAdd{})
	m.Register(&client.RoleDefaultList{})
	m.Register(&client.RoleDefaultRemove{})
	m.Register(&installer.Install{})
	m.Register(&installer.Uninstall{})
	m.Register(&installer.InstallHostList{})
	m.Register(&installer.InstallSSH{})
	m.Register(&installer.InstallConfigInit{})
	m.Register(&admin.AddPoolToSchedulerCmd{})
	m.Register(&client.EventList{})
	m.Register(&client.EventInfo{})
	m.Register(&client.EventCancel{})
	m.Register(&client.RoutersList{})
	m.Register(&admin.TemplateList{})
	m.Register(&admin.TemplateAdd{})
	m.Register(&admin.TemplateRemove{})
	m.Register(&admin.MachineList{})
	m.Register(&admin.MachineDestroy{})
	m.Register(&admin.TemplateUpdate{})
	m.Register(&admin.PlanCreate{})
	m.Register(&admin.PlanRemove{})
	m.Register(&admin.UpdatePoolToSchedulerCmd{})
	m.Register(&admin.RemovePoolFromSchedulerCmd{})
	m.Register(&admin.ServiceCreate{})
	m.Register(&admin.ServiceDestroy{})
	m.Register(&admin.ServiceUpdate{})
	m.Register(&admin.ServiceDocGet{})
	m.Register(&admin.ServiceDocAdd{})
	m.Register(&admin.ServiceTemplate{})
	m.Register(&admin.UserQuotaView{})
	m.Register(&admin.UserChangeQuota{})
	m.Register(&admin.AppQuotaView{})
	m.Register(&admin.AppQuotaChange{})
	m.Register(&admin.AppRoutesRebuild{})
	m.Register(&admin.PoolConstraintList{})
	m.Register(&admin.PoolConstraintSet{})
	m.Register(&admin.EventBlockList{})
	m.Register(&admin.EventBlockAdd{})
	m.Register(&admin.EventBlockRemove{})
	m.Register(&client.TagList{})
	m.Register(&admin.NodeContainerList{})
	m.Register(&admin.NodeContainerAdd{})
	m.Register(&admin.NodeContainerInfo{})
	m.Register(&admin.NodeContainerUpdate{})
	m.Register(&admin.NodeContainerDelete{})
	m.Register(&admin.NodeContainerUpgrade{})
	m.Register(&admin.ClusterAdd{})
	m.Register(&admin.ClusterUpdate{})
	m.Register(&admin.ClusterRemove{})
	m.Register(&admin.ClusterList{})
	m.Register(&client.VolumeCreate{})
	m.Register(&client.VolumeUpdate{})
	m.Register(&client.VolumeList{})
	m.Register(&client.VolumePlansList{})
	m.Register(&client.VolumeDelete{})
	m.Register(&client.VolumeBind{})
	m.Register(&client.VolumeUnbind{})
	m.Register(&client.AppRoutersList{})
	m.Register(&client.AppRoutersAdd{})
	m.Register(&client.AppRoutersRemove{})
	m.Register(&client.AppRoutersUpdate{})
	m.Register(&admin.InfoNodeCmd{})
	m.RegisterRemoved("bs-env-set", "You should use `tsuru node-container-update big-sibling` instead.")
	m.RegisterRemoved("bs-info", "You should use `tsuru node-container-info big-sibling` instead.")
	m.RegisterRemoved("bs-upgrade", "You should use `tsuru node-container-upgrade big-sibling` instead.")
	m.RegisterDeprecated(&admin.AddTeamsToPoolCmd{}, "pool-teams-add")
	m.RegisterDeprecated(&admin.RemoveTeamsFromPoolCmd{}, "pool-teams-remove")
	m.RegisterDeprecated(&admin.AddNodeCmd{}, "docker-node-add")
	m.RegisterDeprecated(&admin.RemoveNodeCmd{}, "docker-node-remove")
	m.RegisterDeprecated(&admin.UpdateNodeCmd{}, "docker-node-update")
	m.RegisterDeprecated(&admin.ListNodesCmd{}, "docker-node-list")
	m.RegisterDeprecated(&admin.GetNodeHealingConfigCmd{}, "docker-healing-info")
	m.RegisterDeprecated(&admin.SetNodeHealingConfigCmd{}, "docker-healing-update")
	m.RegisterDeprecated(&admin.DeleteNodeHealingConfigCmd{}, "docker-healing-delete")
	m.RegisterDeprecated(&admin.RebalanceNodeCmd{}, "containers-rebalance")
	m.RegisterDeprecated(&admin.AutoScaleRunCmd{}, "docker-autoscale-run")
	m.RegisterDeprecated(&admin.ListAutoScaleHistoryCmd{}, "docker-autoscale-list")
	m.RegisterDeprecated(&admin.AutoScaleInfoCmd{}, "docker-autoscale-info")
	m.RegisterDeprecated(&admin.AutoScaleSetRuleCmd{}, "docker-autoscale-rule-set")
	m.RegisterDeprecated(&admin.AutoScaleDeleteRuleCmd{}, "docker-autoscale-rule-remove")
	m.RegisterDeprecated(&admin.ListHealingHistoryCmd{}, "docker-healing-list")
	registerExtraCommands(m)
	return m
}

func registerExtraCommands(m *cmd.Manager) {
	for _, c := range cmd.ExtraCmds() {
		m.Register(c)
	}
}

func inDockerMachineDriverMode() bool {
	return os.Getenv(localbinary.PluginEnvKey) == localbinary.PluginEnvVal
}

func main() {
	if inDockerMachineDriverMode() {
		err := dockermachine.RunDriver(os.Getenv(localbinary.PluginEnvDriverName))
		if err != nil {
			log.Fatalf("Error running driver: %s", err)
		}
	} else {
		localbinary.CurrentBinaryIsDockerMachine = true
		name := cmd.ExtractProgramName(os.Args[0])
		m := buildManager(name)
		m.Run(os.Args[1:])
	}
}
<commit_msg>bump version to 1.5.0-rc2<commit_after>// Copyright 2017 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"log"
	"os"

	"github.com/docker/machine/libmachine/drivers/plugin/localbinary"
	"github.com/tsuru/tsuru-client/tsuru/admin"
	"github.com/tsuru/tsuru-client/tsuru/client"
	"github.com/tsuru/tsuru-client/tsuru/installer"
	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/iaas/dockermachine"
	_ "github.com/tsuru/tsuru/provision/docker/cmds"
)

const (
	version = "1.5.0-rc2"
	header  = "Supported-Tsuru"
)

func buildManager(name string) *cmd.Manager {
	lookup := func(context *cmd.Context) error {
		return client.RunPlugin(context)
	}
	m := cmd.BuildBaseManager(name, version, header, lookup)
	m.Register(&client.AppRun{})
	m.Register(&client.AppInfo{})
	m.Register(&client.AppCreate{})
	m.Register(&client.AppRemove{})
	m.Register(&client.AppUpdate{})
	m.Register(&client.UnitAdd{})
	m.Register(&client.UnitRemove{})
	m.Register(&client.AppList{})
	m.Register(&client.AppLog{})
	m.Register(&client.AppGrant{})
	m.Register(&client.AppRevoke{})
	m.Register(&client.AppRestart{})
	m.Register(&client.AppStart{})
	m.Register(&client.AppStop{})
	m.Register(&admin.AppLockDelete{})
	m.Register(&client.CertificateSet{})
	m.Register(&client.CertificateUnset{})
	m.Register(&client.CertificateList{})
	m.Register(&client.CnameAdd{})
	m.Register(&client.CnameRemove{})
	m.Register(&client.EnvGet{})
	m.Register(&client.EnvSet{})
	m.Register(&client.EnvUnset{})
	m.Register(&client.KeyAdd{})
	m.Register(&client.KeyRemove{})
	m.Register(&client.KeyList{})
	m.Register(client.ServiceList{})
	m.Register(&client.ServiceInstanceAdd{})
	m.Register(&client.ServiceInstanceUpdate{})
	m.Register(&client.ServiceInstanceRemove{})
	m.Register(client.ServiceInfo{})
	m.Register(client.ServiceInstanceInfo{})
	m.Register(client.ServiceInstanceStatus{})
	m.Register(&client.ServiceInstanceGrant{})
	m.Register(&client.ServiceInstanceRevoke{})
	m.Register(&client.ServiceInstanceBind{})
	m.Register(&client.ServiceInstanceUnbind{})
	m.Register(&admin.PlatformList{})
	m.Register(&admin.PlatformAdd{})
	m.Register(&admin.PlatformUpdate{})
	m.Register(&admin.PlatformRemove{})
	m.Register(&client.PluginInstall{})
	m.Register(&client.PluginRemove{})
	m.Register(&client.PluginList{})
	m.Register(&client.AppSwap{})
	m.Register(&client.AppDeploy{})
	m.Register(&client.AppBuild{})
	m.Register(&client.PlanList{})
	m.Register(&client.UserCreate{})
	m.Register(&client.ResetPassword{})
	m.Register(&client.UserRemove{})
	m.Register(&client.ListUsers{})
	m.Register(&client.TeamCreate{})
	m.Register(&client.TeamUpdate{})
	m.Register(&client.TeamRemove{})
	m.Register(&client.TeamList{})
	m.Register(&client.TeamInfo{})
	m.Register(&client.ChangePassword{})
	m.Register(&client.ShowAPIToken{})
	m.Register(&client.RegenerateAPIToken{})
	m.Register(&client.AppDeployList{})
	m.Register(&client.AppDeployRollback{})
	m.Register(&client.AppDeployRollbackUpdate{})
	m.Register(&client.AppDeployRebuild{})
	m.Register(&cmd.ShellToContainerCmd{})
	m.Register(&client.PoolList{})
	m.Register(&client.PermissionList{})
	m.Register(&client.RoleAdd{})
	m.Register(&client.RoleUpdate{})
	m.Register(&client.RoleRemove{})
	m.Register(&client.RoleList{})
	m.Register(&client.RoleInfo{})
	m.Register(&client.RolePermissionAdd{})
	m.Register(&client.RolePermissionRemove{})
	m.Register(&client.RoleAssign{})
	m.Register(&client.RoleDissociate{})
	m.Register(&client.RoleDefaultAdd{})
	m.Register(&client.RoleDefaultList{})
	m.Register(&client.RoleDefaultRemove{})
	m.Register(&installer.Install{})
	m.Register(&installer.Uninstall{})
	m.Register(&installer.InstallHostList{})
	m.Register(&installer.InstallSSH{})
	m.Register(&installer.InstallConfigInit{})
	m.Register(&admin.AddPoolToSchedulerCmd{})
	m.Register(&client.EventList{})
	m.Register(&client.EventInfo{})
	m.Register(&client.EventCancel{})
	m.Register(&client.RoutersList{})
	m.Register(&admin.TemplateList{})
	m.Register(&admin.TemplateAdd{})
	m.Register(&admin.TemplateRemove{})
	m.Register(&admin.MachineList{})
	m.Register(&admin.MachineDestroy{})
	m.Register(&admin.TemplateUpdate{})
	m.Register(&admin.PlanCreate{})
	m.Register(&admin.PlanRemove{})
	m.Register(&admin.UpdatePoolToSchedulerCmd{})
	m.Register(&admin.RemovePoolFromSchedulerCmd{})
	m.Register(&admin.ServiceCreate{})
	m.Register(&admin.ServiceDestroy{})
	m.Register(&admin.ServiceUpdate{})
	m.Register(&admin.ServiceDocGet{})
	m.Register(&admin.ServiceDocAdd{})
	m.Register(&admin.ServiceTemplate{})
	m.Register(&admin.UserQuotaView{})
	m.Register(&admin.UserChangeQuota{})
	m.Register(&admin.AppQuotaView{})
	m.Register(&admin.AppQuotaChange{})
	m.Register(&admin.AppRoutesRebuild{})
	m.Register(&admin.PoolConstraintList{})
	m.Register(&admin.PoolConstraintSet{})
	m.Register(&admin.EventBlockList{})
	m.Register(&admin.EventBlockAdd{})
	m.Register(&admin.EventBlockRemove{})
	m.Register(&client.TagList{})
	m.Register(&admin.NodeContainerList{})
	m.Register(&admin.NodeContainerAdd{})
	m.Register(&admin.NodeContainerInfo{})
	m.Register(&admin.NodeContainerUpdate{})
	m.Register(&admin.NodeContainerDelete{})
	m.Register(&admin.NodeContainerUpgrade{})
	m.Register(&admin.ClusterAdd{})
	m.Register(&admin.ClusterUpdate{})
	m.Register(&admin.ClusterRemove{})
	m.Register(&admin.ClusterList{})
	m.Register(&client.VolumeCreate{})
	m.Register(&client.VolumeUpdate{})
	m.Register(&client.VolumeList{})
	m.Register(&client.VolumePlansList{})
	m.Register(&client.VolumeDelete{})
	m.Register(&client.VolumeBind{})
	m.Register(&client.VolumeUnbind{})
	m.Register(&client.AppRoutersList{})
	m.Register(&client.AppRoutersAdd{})
	m.Register(&client.AppRoutersRemove{})
	m.Register(&client.AppRoutersUpdate{})
	m.Register(&admin.InfoNodeCmd{})
	m.RegisterRemoved("bs-env-set", "You should use `tsuru node-container-update big-sibling` instead.")
	m.RegisterRemoved("bs-info", "You should use `tsuru node-container-info big-sibling` instead.")
	m.RegisterRemoved("bs-upgrade", "You should use `tsuru node-container-upgrade big-sibling` instead.")
	m.RegisterDeprecated(&admin.AddTeamsToPoolCmd{}, "pool-teams-add")
	m.RegisterDeprecated(&admin.RemoveTeamsFromPoolCmd{}, "pool-teams-remove")
	m.RegisterDeprecated(&admin.AddNodeCmd{}, "docker-node-add")
	m.RegisterDeprecated(&admin.RemoveNodeCmd{}, "docker-node-remove")
	m.RegisterDeprecated(&admin.UpdateNodeCmd{}, "docker-node-update")
	m.RegisterDeprecated(&admin.ListNodesCmd{}, "docker-node-list")
	m.RegisterDeprecated(&admin.GetNodeHealingConfigCmd{}, "docker-healing-info")
	m.RegisterDeprecated(&admin.SetNodeHealingConfigCmd{}, "docker-healing-update")
	m.RegisterDeprecated(&admin.DeleteNodeHealingConfigCmd{}, "docker-healing-delete")
	m.RegisterDeprecated(&admin.RebalanceNodeCmd{}, "containers-rebalance")
	m.RegisterDeprecated(&admin.AutoScaleRunCmd{}, "docker-autoscale-run")
	m.RegisterDeprecated(&admin.ListAutoScaleHistoryCmd{}, "docker-autoscale-list")
	m.RegisterDeprecated(&admin.AutoScaleInfoCmd{}, "docker-autoscale-info")
	m.RegisterDeprecated(&admin.AutoScaleSetRuleCmd{}, "docker-autoscale-rule-set")
	m.RegisterDeprecated(&admin.AutoScaleDeleteRuleCmd{}, "docker-autoscale-rule-remove")
	m.RegisterDeprecated(&admin.ListHealingHistoryCmd{}, "docker-healing-list")
	registerExtraCommands(m)
	return m
}

func registerExtraCommands(m *cmd.Manager) {
	for _, c := range cmd.ExtraCmds() {
		m.Register(c)
	}
}

func inDockerMachineDriverMode() bool {
	return os.Getenv(localbinary.PluginEnvKey) == localbinary.PluginEnvVal
}

func main() {
	if inDockerMachineDriverMode() {
		err := dockermachine.RunDriver(os.Getenv(localbinary.PluginEnvDriverName))
		if err != nil {
			log.Fatalf("Error running driver: %s", err)
		}
	} else {
		localbinary.CurrentBinaryIsDockerMachine = true
		name := cmd.ExtractProgramName(os.Args[0])
		m := buildManager(name)
		m.Run(os.Args[1:])
	}
}
<|endoftext|>


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

package logging

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/ava-labs/avalanchego/utils/constants"
)

var (
	filePrefix = fmt.Sprintf("%s/", constants.AppName)
)

// Log implements the Logger interface
type Log struct {
	config Config

	messages []string
	size     int

	wg                               sync.WaitGroup
	flushLock, writeLock, configLock sync.Mutex
	needsFlush                       *sync.Cond

	closed bool

	writer RotatingWriter
}

// New returns a new logger set up according to [config]
func New(config Config) (*Log, error) {
	if err := os.MkdirAll(config.Directory, os.ModePerm); err != nil {
		return nil, err
	}
	l := &Log{
		config: config,
		writer: &fileWriter{},
	}
	l.needsFlush = sync.NewCond(&l.flushLock)

	l.wg.Add(1)

	go l.RecoverAndPanic(l.run)

	return l, nil
}

func (l *Log) run() {
	defer l.wg.Done()

	l.writeLock.Lock()
	defer l.writeLock.Unlock()

	currentSize, err := l.writer.Initialize(l.config)
	if err != nil {
		panic(err)
	}

	closed := false
	nextRotation := time.Now().Add(l.config.RotationInterval)

	for !closed {
		l.writeLock.Unlock()
		l.flushLock.Lock()
		for l.size < l.config.FlushSize && !l.closed {
			l.needsFlush.Wait()
		}
		closed = l.closed
		prevMessages := l.messages
		l.messages = nil
		l.size = 0
		l.flushLock.Unlock()
		l.writeLock.Lock()

		for _, msg := range prevMessages {
			n, _ := l.writer.WriteString(msg)
			currentSize += n
		}

		if !l.config.DisableFlushOnWrite {
			// attempt to flush after the write
			_ = l.writer.Flush()
		}

		if now := time.Now(); nextRotation.Before(now) || currentSize > l.config.FileSize {
			nextRotation = now.Add(l.config.RotationInterval)
			currentSize = 0
			// attempt to flush before closing
			_ = l.writer.Flush()
			// attempt to close the file
			_ = l.writer.Close()

			if err := l.writer.Rotate(); err != nil {
				panic(err)
			}
		}
	}
	// attempt to flush when exiting
	_ = l.writer.Flush()
	// attempt to close the file when exiting
	_ = l.writer.Close()
}

func (l *Log) Write(p []byte) (int, error) {
	if l == nil {
		return 0, nil
	}

	l.configLock.Lock()
	defer l.configLock.Unlock()

	if !l.config.DisableLogging {
		l.flushLock.Lock()
		l.messages = append(l.messages, string(p))
		l.size += len(p)
		l.needsFlush.Signal()
		l.flushLock.Unlock()
	}

	return len(p), nil
}

// Stop implements the Logger interface
func (l *Log) Stop() {
	l.flushLock.Lock()
	l.closed = true
	l.needsFlush.Signal()
	l.flushLock.Unlock()

	l.wg.Wait()
}

// Should only be called from [Level] functions.
func (l *Log) log(level Level, format string, args ...interface{}) {
	if l == nil {
		return
	}

	l.configLock.Lock()
	defer l.configLock.Unlock()

	shouldLog := !l.config.DisableLogging && level <= l.config.LogLevel
	shouldDisplay := (!l.config.DisableDisplaying && level <= l.config.DisplayLevel) || level == Fatal

	if !shouldLog && !shouldDisplay {
		return
	}

	output := l.format(level, format, args...)

	if shouldLog {
		l.flushLock.Lock()
		l.messages = append(l.messages, output)
		l.size += len(output)
		l.needsFlush.Signal()
		l.flushLock.Unlock()
	}

	if shouldDisplay {
		switch {
		case l.config.DisableContextualDisplaying:
			fmt.Println(fmt.Sprintf(format, args...))
		case l.config.DisplayHighlight == Plain:
			fmt.Print(output)
		default:
			fmt.Print(level.Color().Wrap(output))
		}
	}
}

func (l *Log) format(level Level, format string, args ...interface{}) string {
	loc := "?"
	if _, file, no, ok := runtime.Caller(3); ok {
		loc = fmt.Sprintf("%s#%d", file, no)
	}
	if i := strings.Index(loc, filePrefix); i != -1 {
		loc = loc[i+len(filePrefix):]
	}
	text := fmt.Sprintf("%s: %s", loc, fmt.Sprintf(format, args...))

	prefix := ""
	if l.config.MsgPrefix != "" {
		prefix = fmt.Sprintf(" <%s>", l.config.MsgPrefix)
	}

	return fmt.Sprintf("%s[%s]%s %s\n",
		level,
		time.Now().Format("01-02|15:04:05"),
		prefix,
		text)
}

// Fatal implements the Logger interface
func (l *Log) Fatal(format string, args ...interface{}) { l.log(Fatal, format, args...) }

// Error implements the Logger interface
func (l *Log) Error(format string, args ...interface{}) { l.log(Error, format, args...) }

// Warn implements the Logger interface
func (l *Log) Warn(format string, args ...interface{}) { l.log(Warn, format, args...) }

// Info implements the Logger interface
func (l *Log) Info(format string, args ...interface{}) { l.log(Info, format, args...) }

// Debug implements the Logger interface
func (l *Log) Debug(format string, args ...interface{}) { l.log(Debug, format, args...) }

// Verbo implements the Logger interface
func (l *Log) Verbo(format string, args ...interface{}) { l.log(Verbo, format, args...) }

// AssertNoError implements the Logger interface
func (l *Log) AssertNoError(err error) {
	if err != nil {
		l.log(Fatal, "%s", err)
	}
	if l.config.Assertions && err != nil {
		l.Stop()
		panic(err)
	}
}

// AssertTrue implements the Logger interface
func (l *Log) AssertTrue(b bool, format string, args ...interface{}) {
	if !b {
		l.log(Fatal, format, args...)
	}
	if l.config.Assertions && !b {
		l.Stop()
		panic(fmt.Sprintf(format, args...))
	}
}

// AssertDeferredTrue implements the Logger interface
func (l *Log) AssertDeferredTrue(f func() bool, format string, args ...interface{}) {
	// Note, the logger will only be notified here if assertions are enabled
	if l.config.Assertions && !f() {
		err := fmt.Sprintf(format, args...)
		l.log(Fatal, err)
		l.Stop()
		panic(err)
	}
}

// AssertDeferredNoError implements the Logger interface
func (l *Log) AssertDeferredNoError(f func() error) {
	if l.config.Assertions {
		err := f()
		if err != nil {
			l.log(Fatal, "%s", err)
		}
		if l.config.Assertions && err != nil {
			l.Stop()
			panic(err)
		}
	}
}

// StopOnPanic implements the Logger interface
func (l *Log) StopOnPanic() {
	if r := recover(); r != nil {
		l.Fatal("Panicking due to:\n%s\nFrom:\n%s", r, Stacktrace{})
		l.Stop()
		panic(r)
	}
}

// RecoverAndPanic implements the Logger interface
func (l *Log) RecoverAndPanic(f func()) { defer l.StopOnPanic(); f() }

func (l *Log) stopAndExit(exit func()) {
	if r := recover(); r != nil {
		l.Fatal("Panicking due to:\n%s\nFrom:\n%s", r, Stacktrace{})
		l.Stop()
		exit()
	}
}

// RecoverAndExit implements the Logger interface
func (l *Log) RecoverAndExit(f, exit func()) { defer l.stopAndExit(exit); f() }

// SetLogLevel ...
func (l *Log) SetLogLevel(lvl Level) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.LogLevel = lvl
}

// SetDisplayLevel implements the Logger interface
func (l *Log) SetDisplayLevel(lvl Level) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisplayLevel = lvl
}

// SetPrefix implements the Logger interface
func (l *Log) SetPrefix(prefix string) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.MsgPrefix = prefix
}

// SetLoggingEnabled implements the Logger interface
func (l *Log) SetLoggingEnabled(enabled bool) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisableLogging = !enabled
}

// SetDisplayingEnabled implements the Logger interface
func (l *Log) SetDisplayingEnabled(enabled bool) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisableDisplaying = !enabled
}

// SetContextualDisplayingEnabled implements the Logger interface
func (l *Log) SetContextualDisplayingEnabled(enabled bool) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisableContextualDisplaying = !enabled
}

// fileWriter implements the RotatingWriter interface
type fileWriter struct {
	writer *bufio.Writer
	file   *os.File

	config Config
}

// Flush implements the RotatingWriter interface
func (fw *fileWriter) Flush() error {
	return fw.writer.Flush()
}

// Write implements the RotatingWriter interface
func (fw *fileWriter) Write(b []byte) (int, error) {
	return fw.writer.Write(b)
}

// WriteString implements the RotatingWriter interface
func (fw *fileWriter) WriteString(s string) (int, error) {
	return fw.writer.WriteString(s)
}

// Close implements the RotatingWriter interface
func (fw *fileWriter) Close() error {
	return fw.file.Close()
}

// Rotate implements the RotatingWriter interface
func (fw *fileWriter) Rotate() error {
	for i := fw.config.RotationSize - 1; i > 0; i-- {
		sourceFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log.%d", fw.config.LoggerName, i))
		destFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log.%d", fw.config.LoggerName, i+1))
		if _, err := os.Stat(sourceFilename); !os.IsNotExist(err) {
			if err := os.Rename(sourceFilename, destFilename); err != nil {
				return err
			}
		}
	}
	sourceFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log", fw.config.LoggerName))
	destFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log.1", fw.config.LoggerName))
	if err := os.Rename(sourceFilename, destFilename); err != nil {
		return err
	}
	writer, file, err := fw.create()
	if err != nil {
		return err
	}
	fw.file = file
	fw.writer = writer
	return nil
}

// Creates a file if it does not exist or opens it in append mode if it does
func (fw *fileWriter) create() (*bufio.Writer, *os.File, error) {
	filename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log", fw.config.LoggerName))
	file, err := os.OpenFile(filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return nil, nil, err
	}
	writer := bufio.NewWriter(file)
	return writer, file, nil
}

// Initialize implements the RotatingWriter interface
func (fw *fileWriter) Initialize(config Config) (int, error) {
	fw.config = config
	writer, file, err := fw.create()
	if err != nil {
		return 0, err
	}
	fw.writer = writer
	fw.file = file
	fileSize, err := file.Seek(0, io.SeekEnd)
	if err != nil {
		return 0, err
	}
	return int(fileSize), nil
}
<commit_msg>use file.Stat() instead of file.Seek() to get file size more safely<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package logging

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/ava-labs/avalanchego/utils/constants"
)

var (
	filePrefix = fmt.Sprintf("%s/", constants.AppName)
)

// Log implements the Logger interface
type Log struct {
	config Config

	messages []string
	size     int

	wg                               sync.WaitGroup
	flushLock, writeLock, configLock sync.Mutex
	needsFlush                       *sync.Cond

	closed bool

	writer RotatingWriter
}

// New returns a new logger set up according to [config]
func New(config Config) (*Log, error) {
	if err := os.MkdirAll(config.Directory, os.ModePerm); err != nil {
		return nil, err
	}
	l := &Log{
		config: config,
		writer: &fileWriter{},
	}
	l.needsFlush = sync.NewCond(&l.flushLock)

	l.wg.Add(1)

	go l.RecoverAndPanic(l.run)

	return l, nil
}

func (l *Log) run() {
	defer l.wg.Done()

	l.writeLock.Lock()
	defer l.writeLock.Unlock()

	currentSize, err := l.writer.Initialize(l.config)
	if err != nil {
		panic(err)
	}

	closed := false
	nextRotation := time.Now().Add(l.config.RotationInterval)

	for !closed {
		l.writeLock.Unlock()
		l.flushLock.Lock()
		for l.size < l.config.FlushSize && !l.closed {
			l.needsFlush.Wait()
		}
		closed = l.closed
		prevMessages := l.messages
		l.messages = nil
		l.size = 0
		l.flushLock.Unlock()
		l.writeLock.Lock()

		for _, msg := range prevMessages {
			n, _ := l.writer.WriteString(msg)
			currentSize += n
		}

		if !l.config.DisableFlushOnWrite {
			// attempt to flush after the write
			_ = l.writer.Flush()
		}

		if now := time.Now(); nextRotation.Before(now) || currentSize > l.config.FileSize {
			nextRotation = now.Add(l.config.RotationInterval)
			currentSize = 0
			// attempt to flush before closing
			_ = l.writer.Flush()
			// attempt to close the file
			_ = l.writer.Close()

			if err := l.writer.Rotate(); err != nil {
				panic(err)
			}
		}
	}
	// attempt to flush when exiting
	_ = l.writer.Flush()
	// attempt to close the file when exiting
	_ = l.writer.Close()
}

func (l *Log) Write(p []byte) (int, error) {
	if l == nil {
		return 0, nil
	}

	l.configLock.Lock()
	defer l.configLock.Unlock()

	if !l.config.DisableLogging {
		l.flushLock.Lock()
		l.messages = append(l.messages, string(p))
		l.size += len(p)
		l.needsFlush.Signal()
		l.flushLock.Unlock()
	}

	return len(p), nil
}

// Stop implements the Logger interface
func (l *Log) Stop() {
	l.flushLock.Lock()
	l.closed = true
	l.needsFlush.Signal()
	l.flushLock.Unlock()

	l.wg.Wait()
}

// Should only be called from [Level] functions.
func (l *Log) log(level Level, format string, args ...interface{}) {
	if l == nil {
		return
	}

	l.configLock.Lock()
	defer l.configLock.Unlock()

	shouldLog := !l.config.DisableLogging && level <= l.config.LogLevel
	shouldDisplay := (!l.config.DisableDisplaying && level <= l.config.DisplayLevel) || level == Fatal

	if !shouldLog && !shouldDisplay {
		return
	}

	output := l.format(level, format, args...)

	if shouldLog {
		l.flushLock.Lock()
		l.messages = append(l.messages, output)
		l.size += len(output)
		l.needsFlush.Signal()
		l.flushLock.Unlock()
	}

	if shouldDisplay {
		switch {
		case l.config.DisableContextualDisplaying:
			fmt.Println(fmt.Sprintf(format, args...))
		case l.config.DisplayHighlight == Plain:
			fmt.Print(output)
		default:
			fmt.Print(level.Color().Wrap(output))
		}
	}
}

func (l *Log) format(level Level, format string, args ...interface{}) string {
	loc := "?"
	if _, file, no, ok := runtime.Caller(3); ok {
		loc = fmt.Sprintf("%s#%d", file, no)
	}
	if i := strings.Index(loc, filePrefix); i != -1 {
		loc = loc[i+len(filePrefix):]
	}
	text := fmt.Sprintf("%s: %s", loc, fmt.Sprintf(format, args...))

	prefix := ""
	if l.config.MsgPrefix != "" {
		prefix = fmt.Sprintf(" <%s>", l.config.MsgPrefix)
	}

	return fmt.Sprintf("%s[%s]%s %s\n",
		level,
		time.Now().Format("01-02|15:04:05"),
		prefix,
		text)
}

// Fatal implements the Logger interface
func (l *Log) Fatal(format string, args ...interface{}) { l.log(Fatal, format, args...) }

// Error implements the Logger interface
func (l *Log) Error(format string, args ...interface{}) { l.log(Error, format, args...) }

// Warn implements the Logger interface
func (l *Log) Warn(format string, args ...interface{}) { l.log(Warn, format, args...) }

// Info implements the Logger interface
func (l *Log) Info(format string, args ...interface{}) { l.log(Info, format, args...) }

// Debug implements the Logger interface
func (l *Log) Debug(format string, args ...interface{}) { l.log(Debug, format, args...) }

// Verbo implements the Logger interface
func (l *Log) Verbo(format string, args ...interface{}) { l.log(Verbo, format, args...) }

// AssertNoError implements the Logger interface
func (l *Log) AssertNoError(err error) {
	if err != nil {
		l.log(Fatal, "%s", err)
	}
	if l.config.Assertions && err != nil {
		l.Stop()
		panic(err)
	}
}

// AssertTrue implements the Logger interface
func (l *Log) AssertTrue(b bool, format string, args ...interface{}) {
	if !b {
		l.log(Fatal, format, args...)
	}
	if l.config.Assertions && !b {
		l.Stop()
		panic(fmt.Sprintf(format, args...))
	}
}

// AssertDeferredTrue implements the Logger interface
func (l *Log) AssertDeferredTrue(f func() bool, format string, args ...interface{}) {
	// Note, the logger will only be notified here if assertions are enabled
	if l.config.Assertions && !f() {
		err := fmt.Sprintf(format, args...)
		l.log(Fatal, err)
		l.Stop()
		panic(err)
	}
}

// AssertDeferredNoError implements the Logger interface
func (l *Log) AssertDeferredNoError(f func() error) {
	if l.config.Assertions {
		err := f()
		if err != nil {
			l.log(Fatal, "%s", err)
		}
		if l.config.Assertions && err != nil {
			l.Stop()
			panic(err)
		}
	}
}

// StopOnPanic implements the Logger interface
func (l *Log) StopOnPanic() {
	if r := recover(); r != nil {
		l.Fatal("Panicking due to:\n%s\nFrom:\n%s", r, Stacktrace{})
		l.Stop()
		panic(r)
	}
}

// RecoverAndPanic implements the Logger interface
func (l *Log) RecoverAndPanic(f func()) { defer l.StopOnPanic(); f() }

func (l *Log) stopAndExit(exit func()) {
	if r := recover(); r != nil {
		l.Fatal("Panicking due to:\n%s\nFrom:\n%s", r, Stacktrace{})
		l.Stop()
		exit()
	}
}

// RecoverAndExit implements the Logger interface
func (l *Log) RecoverAndExit(f, exit func()) { defer l.stopAndExit(exit); f() }

// SetLogLevel ...
func (l *Log) SetLogLevel(lvl Level) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.LogLevel = lvl
}

// SetDisplayLevel implements the Logger interface
func (l *Log) SetDisplayLevel(lvl Level) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisplayLevel = lvl
}

// SetPrefix implements the Logger interface
func (l *Log) SetPrefix(prefix string) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.MsgPrefix = prefix
}

// SetLoggingEnabled implements the Logger interface
func (l *Log) SetLoggingEnabled(enabled bool) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisableLogging = !enabled
}

// SetDisplayingEnabled implements the Logger interface
func (l *Log) SetDisplayingEnabled(enabled bool) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisableDisplaying = !enabled
}

// SetContextualDisplayingEnabled implements the Logger interface
func (l *Log) SetContextualDisplayingEnabled(enabled bool) {
	l.configLock.Lock()
	defer l.configLock.Unlock()

	l.config.DisableContextualDisplaying = !enabled
}

// fileWriter implements the RotatingWriter interface
type fileWriter struct {
	writer *bufio.Writer
	file   *os.File

	config Config
}

// Flush implements the RotatingWriter interface
func (fw *fileWriter) Flush() error {
	return fw.writer.Flush()
}

// Write implements the RotatingWriter interface
func (fw *fileWriter) Write(b []byte) (int, error) {
	return fw.writer.Write(b)
}

// WriteString implements the RotatingWriter interface
func (fw *fileWriter) WriteString(s string) (int, error) {
	return fw.writer.WriteString(s)
}

// Close implements the RotatingWriter interface
func (fw *fileWriter) Close() error {
	return fw.file.Close()
}

// Rotate implements the RotatingWriter interface
func (fw *fileWriter) Rotate() error {
	for i := fw.config.RotationSize - 1; i > 0; i-- {
		sourceFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log.%d", fw.config.LoggerName, i))
		destFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log.%d", fw.config.LoggerName, i+1))
		if _, err := os.Stat(sourceFilename); !os.IsNotExist(err) {
			if err := os.Rename(sourceFilename, destFilename); err != nil {
				return err
			}
		}
	}
	sourceFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log", fw.config.LoggerName))
	destFilename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log.1", fw.config.LoggerName))
	if err := os.Rename(sourceFilename, destFilename); err != nil {
		return err
	}
	writer, file, err := fw.create()
	if err != nil {
		return err
	}
	fw.file = file
	fw.writer = writer
	return nil
}

// Creates a file if it does not exist or opens it in append mode if it does
func (fw *fileWriter) create() (*bufio.Writer, *os.File, error) {
	filename := filepath.Join(fw.config.Directory, fmt.Sprintf("%s.log", fw.config.LoggerName))
	file, err := os.OpenFile(filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return nil, nil, err
	}
	writer := bufio.NewWriter(file)
	return writer, file, nil
}

// Initialize implements the RotatingWriter interface
func (fw *fileWriter) Initialize(config Config) (int, error) {
	fw.config = config
	writer, file, err := fw.create()
	if err != nil {
		return 0, err
	}
	fw.writer = writer
	fw.file = file
	fileinfo, err := file.Stat()
	if err != nil {
		return 0, err
	}
	fileSize := fileinfo.Size()
	return int(fileSize), nil
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package v2

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"syscall"
	"testing"

	"golang.org/x/sys/unix"
)

const defaultCgroup2Path = "/sys/fs/cgroup"

type testCgroup struct {
	cgv2Root  string
	groupPath string
}

func NewCgroupDir(group string, t *testing.T) *testCgroup {
	var st syscall.Statfs_t
	if err := syscall.Statfs(defaultCgroup2Path, &st); err != nil {
		t.Fatal("cannot statfs cgroup root")
	}
	isUnified := st.Type == unix.CGROUP2_SUPER_MAGIC
	if !isUnified {
		t.Skip("System running in hybrid or cgroupv1 mode")
	}
	testCgroupPath := filepath.Join(defaultCgroup2Path, fmt.Sprintf("%s-%d", group, os.Getpid()))
	err := os.Mkdir(testCgroupPath, 0755)
	if err != nil {
		t.Fatal(err)
	}

	return &testCgroup{cgv2Root: defaultCgroup2Path, groupPath: testCgroupPath}
}

func (c *testCgroup) writeFileContents(fileContents map[string]string, t *testing.T) {
	for file, contents := range fileContents {
		err := writeFile(c.groupPath, file, contents)
		if err != nil {
			t.Fatal(err)
		}
	}
}

func writeFile(dir, file, data string) error {
	// Normally dir should not be empty, one case is that cgroup subsystem
	// is not mounted, we will get empty dir, and we want it fail here.
	if dir == "" {
		return fmt.Errorf("no such directory for %s", file)
	}
	if err := ioutil.WriteFile(filepath.Join(dir, file), []byte(data), 0700); err != nil {
		return fmt.Errorf("failed to write %v to %v: %v", data, file, err)
	}
	return nil
}
<commit_msg>GroupPath fix<commit_after>/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package v2

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"syscall"
	"testing"

	"golang.org/x/sys/unix"
)

const defaultCgroup2Path = "/sys/fs/cgroup"

type testCgroup struct {
	cgv2Root  string
	groupPath string
}

func NewCgroupDir(group string, t *testing.T) *testCgroup {
	var st syscall.Statfs_t
	if err := syscall.Statfs(defaultCgroup2Path, &st); err != nil {
		t.Fatal("cannot statfs cgroup root")
	}
	isUnified := st.Type == unix.CGROUP2_SUPER_MAGIC
	if !isUnified {
		t.Skip("System running in hybrid or cgroupv1 mode")
	}
	groupPath := fmt.Sprintf("%s-%d", group, os.Getpid())
	err := os.Mkdir(filepath.Join(defaultCgroup2Path, groupPath), 0755)
	if err != nil {
		t.Fatal(err)
	}

	return &testCgroup{cgv2Root: defaultCgroup2Path, groupPath: groupPath}
}

func (c *testCgroup) writeFileContents(fileContents map[string]string, t *testing.T) {
	for file, contents := range fileContents {
		err := writeFile(c.groupPath, file, contents)
		if err != nil {
			t.Fatal(err)
		}
	}
}

func writeFile(dir, file, data string) error {
	// Normally dir should not be empty, one case is that cgroup subsystem
	// is not mounted, we will get empty dir, and we want it fail here.
	if dir == "" {
		return fmt.Errorf("no such directory for %s", file)
	}
	if err := ioutil.WriteFile(filepath.Join(dir, file), []byte(data), 0700); err != nil {
		return fmt.Errorf("failed to write %v to %v: %v", data, file, err)
	}
	return nil
}
<|endoftext|>