<commit_before>/*
Package workerpool queues work to a limited number of goroutines.

The purpose of the worker pool is to limit the concurrency of the task
performed by the workers.  This is useful when performing a task requires
sufficient resources (CPU, memory, etc.), that running too many tasks at the
same time would exhaust resources.

Non-blocking task submission

A task is a function submitted to the worker pool for execution.  Submitting
tasks to this worker pool will not block, regardless of the number of tasks.
Tasks read from the input task queue are immediately dispatched to an available
worker.  If no worker is immediately available, then the task is passed to a go
routine which is started to wait for an available worker.  This clears the task
from the task queue immediately, whether or not a worker is currently
available, and will not block the submission of tasks.

The intent of the worker pool is to limit the concurrency of task execution,
not limit the number of tasks queued to be executed. Therefore, this unbounded
input of tasks is acceptable as the tasks cannot be discarded.  If the number
of inbound tasks if too many to even queue for pending processing, then this
solution is outside the scope of the worker pool, and should be solved by
distributing load over multiple systems, storing input that requires processing
in some intermediate storage (e.g. a database or file system).

Dispatcher

This worker pool uses a single dispatcher goroutine to read tasks from the
input task queue and dispatch them to a worker goroutine.  This allows for a
small input channel, and lets the dispatcher queue as many tasks as are
submitted when there are no available workers (using goroutines).
Additionally, the dispatcher can adjust the number of workers as appropriate
for the work load, without having to utilize locked counters and checks
incurred on task submission.

Usage note

It is advisable to use different worker pools for tasks that are bound by
different resources, or that have different resource use patterns.  For
example, tasks that use X Mb of memory may need different concurrency limits
than tasks that use Y Mb of memory.

Credits

This implementation builds on ideas from the following:

http://marcio.io/2015/07/handling-1-million-requests-per-minute-with-golang
http://nesv.github.io/golang/2014/02/25/worker-queues-in-go.html

*/
package workerpool

import "time"

const (
	// Size of queue to which tasks are submitted.  This can be small, no
	// matter how many tasks are submitted, because the dispatcher removes
	// tasks from this queue, scheduling each immediately to a ready worker, or
	// to a goroutine that will give it to the next ready worker.
	taskQueueSize = 16

	// If worker pool receives no new work for this period of time, then stop
	// a worker goroutine.
	idleTimeoutSec = 5

	// The minimum number of workers is always zero, because the time to start
	// new workers is insignificant.
	minWorkers = 0
)

type WorkerPool interface {
	// Submit enqueues a function for a worker to execute.
	//
	// Any external values needed by the task function must be captured in a
	// closure.  Any return values should be returned over a channel that is
	// captured in the task function closure.
	//
	// Submit will not block regardless of the number of tasks submitted.  Each
	// task is immediately given to an available worker or passed to a
	// goroutine to be given to the next available worker.  If there are no
	// available workers, the dispatcher adds a worker, until the maximum
	// number of workers is running.
	Submit(task func())

	// Stop stops the worker pool and waits for workers to complete.
	//
	// Since creating the worker pool starts at least one goroutine, for the
	// dispatcher, this function should be called when the worker pool is no
	// longer needed.
	Stop()

	// Stopped returns true if this worker pool has been stopped.
	Stopped() bool

	// WorkerCount returns the current number of worker goroutines.
	//
	// Do not rely on this value to be perfectly accurate as it may change by
	// the time the caller gets it.
	WorkerCount() int
}

// New creates and starts a pool of worker goroutines.
//
// The maxWorkers parameter specifies the maximum number of workers that will
// execute tasks concurrently.  After each timeout period, a worker goroutine
// is stopped until there are no remaining workers.
func New(maxWorkers int) (WorkerPool, error) {
	// There must be at least one worker.
	if maxWorkers < 1 {
		maxWorkers = 1
	}

	pool := &workerPool{
		taskQueue:    make(chan func(), taskQueueSize),
		maxWorkers:   maxWorkers,
		readyWorkers: make(chan chan func(), maxWorkers),
		timeout:      time.Second * idleTimeoutSec,
		stoppedChan:  make(chan struct{}),
	}

	// Start the task dispatcher.
	go pool.dispatch()

	return pool, nil
}

type workerPool struct {
	maxWorkers   int
	workerCount  int
	timeout      time.Duration
	taskQueue    chan func()
	readyWorkers chan chan func()
	stoppedChan  chan struct{}
}

// Stop stops the worker pool and waits for workers to complete.
func (p *workerPool) Stop() {
	if p.Stopped() {
		return
	}
	close(p.taskQueue)
	<-p.stoppedChan
}

// Stopped returns true if this worker pool has been stopped.
func (p *workerPool) Stopped() bool {
	select {
	case <-p.stoppedChan:
		return true
	default:
	}
	return false
}

// Submit enqueues a function for a worker to execute.
func (p *workerPool) Submit(task func()) {
	if task != nil {
		p.taskQueue <- task
	}
}

// WorkerCount returns the current number of worker goroutines.
func (p *workerPool) WorkerCount() int { return p.workerCount }

// dispatch sends the next queued task to an available worker.
func (p *workerPool) dispatch() {
	defer close(p.stoppedChan)
	timeout := time.NewTimer(p.timeout)
	var task func()
	var ok bool
	var workerTaskChan chan func()
shutdown:
	for {
		timeout.Reset(p.timeout)
		select {
		case task, ok = <-p.taskQueue:
			if !ok {
				break shutdown
			}
			// Got a task to do.
			select {
			case workerTaskChan = <-p.readyWorkers:
				// A worker is ready, so give task to worker.
				workerTaskChan <- task
			default:
				// No workers ready.
				// Create a new worker, if not at max.
				if p.workerCount < p.maxWorkers {
					p.workerCount++
					startWorker(p.readyWorkers)
				}
				// Start a goroutine to submit the task when a worker is ready.
				go func(t func()) {
					taskChan := <-p.readyWorkers
					taskChan <- t
				}(task)
			}
		case <-timeout.C:
			// Timed out waiting for work to arrive.  Kill a ready worker.
			if p.workerCount > minWorkers {
				select {
				case workerTaskChan = <-p.readyWorkers:
					// A worker is ready, so kill.
					close(workerTaskChan)
					p.workerCount--
				default:
					// No work, but no ready workers.  All workers are busy.
				}
			}
		}
	}

	// Stop all remaining workers as they become ready.
	for p.workerCount > 0 {
		workerTaskChan = <-p.readyWorkers
		close(workerTaskChan)
		p.workerCount--
	}
}

// startWorker starts a goroutine that executes tasks given by the dispatcher.
//
// A worker registers that is it available to do work by putting its task
// channel on the readyWorkers channel.  The dispatcher reads a worker's task
// channel from the readyWorkers channel, and writes a task to the worker over
// the worker's task channel.  To stop a worker, the dispatcher closes a
// worker's task channel, instead of writing a task to it.
func startWorker(readyWorkers chan chan func()) {
	go func() {
		taskChan := make(chan func())
		var task func()
		var ok bool
		for {
			// Register availability on readyWorkers channel.
			readyWorkers <- taskChan

			// Read task from dispatcher.
			task, ok = <-taskChan
			if !ok {
				// Dispatcher has told worker to stop.
				break
			}

			// Execute the task.
			task()
		}
	}()
}
<commit_msg>Remove countWorkers(). Remove unnecessary erro return.<commit_after>/*
Package workerpool queues work to a limited number of goroutines.

The purpose of the worker pool is to limit the concurrency of the task
performed by the workers.  This is useful when performing a task requires
sufficient resources (CPU, memory, etc.), that running too many tasks at the
same time would exhaust resources.

Non-blocking task submission

A task is a function submitted to the worker pool for execution.  Submitting
tasks to this worker pool will not block, regardless of the number of tasks.
Tasks read from the input task queue are immediately dispatched to an available
worker.  If no worker is immediately available, then the task is passed to a go
routine which is started to wait for an available worker.  This clears the task
from the task queue immediately, whether or not a worker is currently
available, and will not block the submission of tasks.

The intent of the worker pool is to limit the concurrency of task execution,
not limit the number of tasks queued to be executed. Therefore, this unbounded
input of tasks is acceptable as the tasks cannot be discarded.  If the number
of inbound tasks if too many to even queue for pending processing, then this
solution is outside the scope of the worker pool, and should be solved by
distributing load over multiple systems, storing input that requires processing
in some intermediate storage (e.g. a database or file system).

Dispatcher

This worker pool uses a single dispatcher goroutine to read tasks from the
input task queue and dispatch them to a worker goroutine.  This allows for a
small input channel, and lets the dispatcher queue as many tasks as are
submitted when there are no available workers (using goroutines).
Additionally, the dispatcher can adjust the number of workers as appropriate
for the work load, without having to utilize locked counters and checks
incurred on task submission.

Usage note

It is advisable to use different worker pools for tasks that are bound by
different resources, or that have different resource use patterns.  For
example, tasks that use X Mb of memory may need different concurrency limits
than tasks that use Y Mb of memory.

Credits

This implementation builds on ideas from the following:

http://marcio.io/2015/07/handling-1-million-requests-per-minute-with-golang
http://nesv.github.io/golang/2014/02/25/worker-queues-in-go.html

*/
package workerpool

import "time"

const (
	// Size of queue to which tasks are submitted.  This can be small, no
	// matter how many tasks are submitted, because the dispatcher removes
	// tasks from this queue, scheduling each immediately to a ready worker, or
	// to a goroutine that will give it to the next ready worker.
	taskQueueSize = 16

	// If worker pool receives no new work for this period of time, then stop
	// a worker goroutine.
	idleTimeoutSec = 5

	// The minimum number of workers is always zero, because the time to start
	// new workers is insignificant.
	minWorkers = 0
)

type WorkerPool interface {
	// Submit enqueues a function for a worker to execute.
	//
	// Any external values needed by the task function must be captured in a
	// closure.  Any return values should be returned over a channel that is
	// captured in the task function closure.
	//
	// Submit will not block regardless of the number of tasks submitted.  Each
	// task is immediately given to an available worker or passed to a
	// goroutine to be given to the next available worker.  If there are no
	// available workers, the dispatcher adds a worker, until the maximum
	// number of workers is running.
	Submit(task func())

	// Stop stops the worker pool and waits for workers to complete.
	//
	// Since creating the worker pool starts at least one goroutine, for the
	// dispatcher, this function should be called when the worker pool is no
	// longer needed.
	Stop()

	// Stopped returns true if this worker pool has been stopped.
	Stopped() bool
}

// New creates and starts a pool of worker goroutines.
//
// The maxWorkers parameter specifies the maximum number of workers that will
// execute tasks concurrently.  After each timeout period, a worker goroutine
// is stopped until there are no remaining workers.
func New(maxWorkers int) WorkerPool {
	// There must be at least one worker.
	if maxWorkers < 1 {
		maxWorkers = 1
	}

	pool := &workerPool{
		taskQueue:    make(chan func(), taskQueueSize),
		maxWorkers:   maxWorkers,
		readyWorkers: make(chan chan func(), maxWorkers),
		timeout:      time.Second * idleTimeoutSec,
		stoppedChan:  make(chan struct{}),
	}

	// Start the task dispatcher.
	go pool.dispatch()

	return pool
}

type workerPool struct {
	maxWorkers   int
	timeout      time.Duration
	taskQueue    chan func()
	readyWorkers chan chan func()
	stoppedChan  chan struct{}
}

// Stop stops the worker pool and waits for workers to complete.
func (p *workerPool) Stop() {
	if p.Stopped() {
		return
	}
	close(p.taskQueue)
	<-p.stoppedChan
}

// Stopped returns true if this worker pool has been stopped.
func (p *workerPool) Stopped() bool {
	select {
	case <-p.stoppedChan:
		return true
	default:
	}
	return false
}

// Submit enqueues a function for a worker to execute.
func (p *workerPool) Submit(task func()) {
	if task != nil {
		p.taskQueue <- task
	}
}

// dispatch sends the next queued task to an available worker.
func (p *workerPool) dispatch() {
	defer close(p.stoppedChan)
	timeout := time.NewTimer(p.timeout)
	var workerCount int
	var task func()
	var ok bool
	var workerTaskChan chan func()
shutdown:
	for {
		timeout.Reset(p.timeout)
		select {
		case task, ok = <-p.taskQueue:
			if !ok {
				break shutdown
			}
			// Got a task to do.
			select {
			case workerTaskChan = <-p.readyWorkers:
				// A worker is ready, so give task to worker.
				workerTaskChan <- task
			default:
				// No workers ready.
				// Create a new worker, if not at max.
				if workerCount < p.maxWorkers {
					workerCount++
					startWorker(p.readyWorkers)
				}
				// Start a goroutine to submit the task when a worker is ready.
				go func(t func()) {
					taskChan := <-p.readyWorkers
					taskChan <- t
				}(task)
			}
		case <-timeout.C:
			// Timed out waiting for work to arrive.  Kill a ready worker.
			if workerCount > minWorkers {
				select {
				case workerTaskChan = <-p.readyWorkers:
					// A worker is ready, so kill.
					close(workerTaskChan)
					workerCount--
				default:
					// No work, but no ready workers.  All workers are busy.
				}
			}
		}
	}

	// Stop all remaining workers as they become ready.
	for workerCount > 0 {
		workerTaskChan = <-p.readyWorkers
		close(workerTaskChan)
		workerCount--
	}
}

// startWorker starts a goroutine that executes tasks given by the dispatcher.
//
// A worker registers that is it available to do work by putting its task
// channel on the readyWorkers channel.  The dispatcher reads a worker's task
// channel from the readyWorkers channel, and writes a task to the worker over
// the worker's task channel.  To stop a worker, the dispatcher closes a
// worker's task channel, instead of writing a task to it.
func startWorker(readyWorkers chan chan func()) {
	go func() {
		taskChan := make(chan func())
		var task func()
		var ok bool
		for {
			// Register availability on readyWorkers channel.
			readyWorkers <- taskChan

			// Read task from dispatcher.
			task, ok = <-taskChan
			if !ok {
				// Dispatcher has told worker to stop.
				break
			}

			// Execute the task.
			task()
		}
	}()
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Google Inc. All rights reserved.
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 main

import (
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/skelterjohn/wgo/workspaces"
)

type workspace struct {
	workspaces.Workspace
}

func getCurrentWorkspace() (*workspace, error) {
	w, err := workspaces.GetCurrentWorkspace()
	return &workspace{*w}, err
}

func getWorkspace(start string) (*workspace, error) {
	w, err := workspaces.GetWorkspace(start)
	return &workspace{*w}, err
}

func guessGoCommand(args []string) string {
	if len(args) < 1 {
		return ""
	}
	return args[1]
}

func (w *workspace) shellOutToGo(args []string) {
	// we want to fetch new code directly into the workspace, for convenience
	gopath := w.Gopath(guessGoCommand(args) != "get")
	os.Setenv("GOPATH", gopath)
	log.Printf("using GOPATH=%s", gopath)
	shellOutToGo(args)
}

func (w *workspace) vendorRootSrc() string {
	firstGopath := "."
	if len(w.Gopaths) != 0 {
		firstGopath = w.Gopaths[0]
	}
	return filepath.Join(firstGopath, "src")
}

func shellOutToGo(args []string) {
	log.Printf("forking to go: %q", args[1:])
	cmd := exec.Command("go", args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if _, ok := err.(*exec.ExitError); ok {
		os.Exit(1)
	}
	os.Exit(0)
}
<commit_msg>Fixed bug with nil workspace.<commit_after>/*
Copyright 2015 Google Inc. All rights reserved.
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 main

import (
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/skelterjohn/wgo/workspaces"
)

type workspace struct {
	workspaces.Workspace
}

func getCurrentWorkspace() (*workspace, error) {
	if w, err := workspaces.GetCurrentWorkspace(); err != nil {
		return nil, err
	} else {
		return &workspace{*w}, nil
	}
}

func getWorkspace(start string) (*workspace, error) {
	if w, err := workspaces.GetWorkspace(start); err != nil {
		return nil, err
	} else {
		return &workspace{*w}, nil
	}
}

func guessGoCommand(args []string) string {
	if len(args) < 1 {
		return ""
	}
	return args[1]
}

func (w *workspace) shellOutToGo(args []string) {
	// we want to fetch new code directly into the workspace, for convenience
	gopath := w.Gopath(guessGoCommand(args) != "get")
	os.Setenv("GOPATH", gopath)
	log.Printf("using GOPATH=%s", gopath)
	shellOutToGo(args)
}

func (w *workspace) vendorRootSrc() string {
	firstGopath := "."
	if len(w.Gopaths) != 0 {
		firstGopath = w.Gopaths[0]
	}
	return filepath.Join(firstGopath, "src")
}

func shellOutToGo(args []string) {
	log.Printf("forking to go: %q", args[1:])
	cmd := exec.Command("go", args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if _, ok := err.(*exec.ExitError); ok {
		os.Exit(1)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>package beater

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/appcelerator/amp/api/rpc/logs"
	"github.com/appcelerator/amp/api/rpc/stats"
	"github.com/appcelerator/amp/cmd/ampbeat/config"
	"github.com/appcelerator/amp/pkg/nats-streaming"
	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/golang/protobuf/proto"
	"github.com/nats-io/go-nats-streaming"
)

// Ampbeat amp libbeat type
type Ampbeat struct {
	done          chan struct{}
	config        config.Config
	client        publisher.Client
	natsStreaming *ns.NatsStreaming
}

var bt = &Ampbeat{
	done: make(chan struct{}),
}

// New Creates beater
func New(b *beat.Beat, cfg *common.Config) (beat.Beater, error) {
	bt.config = config.DefaultConfig
	if err := cfg.Unpack(&bt.config); err != nil {
		return nil, fmt.Errorf("Error reading config file: %v", err)
	}

	// NATS Connect
	hostname, err := os.Hostname()
	if err != nil {
		return nil, fmt.Errorf("Unable to get hostname: %v", err)
	}
	bt.natsStreaming = ns.NewClient(ns.DefaultURL, ns.ClusterID, b.Name+"-"+hostname, time.Minute)
	if err = bt.natsStreaming.Connect(); err != nil {
		return nil, fmt.Errorf("Unable to connect to NATS: %v", err)
	}

	return bt, nil
}

//Run run the beat
func (bt *Ampbeat) Run(b *beat.Beat) error {
	logp.Info("ampbeat is running! Hit CTRL-C to stop it.")
	bt.client = b.Publisher.Connect()

	// logs subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.LogsSubject, logMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to logs subject")

	// metrics subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.MetricsSubject, metricsMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to metrics subject")

	select {
	case <-bt.done:
		return nil
	}
}

//Stop stop the beat
func (bt *Ampbeat) Stop() {
	bt.natsStreaming.Close()
	bt.client.Close()
	close(bt.done)
}

func logMessageHandler(msg *stan.Msg) {
	e := logs.LogEntry{}
	err := proto.Unmarshal(msg.Data, &e)
	if err != nil {
		logp.Err("Error unmarshalling log entry: %s", err.Error())
		return
	}
	timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
	if err != nil {
		timestamp = time.Now()
	}
	event := common.MapStr{
		"@timestamp":           common.Time(timestamp),
		"type":                 "logs",
		"container_id":         e.ContainerId,
		"container_name":       e.ContainerName,
		"container_short_name": e.ContainerShortName,
		"container_state":      e.ContainerState,
		"service_name":         e.ServiceName,
		"service_id":           e.ServiceId,
		"task_id":              e.TaskId,
		"stack_name":           e.StackName,
		"node_id":              e.NodeId,
		"msg":                  e.Msg,
	}
	if len(e.Labels) > 0 {
		event["labels"] = getLabels(e.Labels)
	}
	bt.client.PublishEvent(event)
}

func metricsMessageHandler(msg *stan.Msg) {
	e := stats.MetricsEntry{}
	err := proto.Unmarshal(msg.Data, &e)
	if err != nil {
		logp.Err("Error unmarshalling metrics entry: %s", err.Error())
		return
	}
	timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
	if err != nil {
		timestamp = time.Now()
	}
	event := common.MapStr{
		"@timestamp":           common.Time(timestamp),
		"type":                 "metrics",
		"container_id":         e.ContainerId,
		"container_name":       e.ContainerName,
		"container_short_name": e.ContainerShortName,
		"container_state":      e.ContainerState,
		"service_name":         e.ServiceName,
		"service_id":           e.ServiceId,
		"task_id":              e.TaskId,
		"stack_name":           e.StackName,
		"node_id":              e.NodeId,
	}
	if len(e.Labels) > 0 {
		event["labels"] = getLabels(e.Labels)
	}
	if e.Cpu != nil {
		event["cpu"] = common.MapStr{
			"total_usage":          e.Cpu.TotalUsage,
			"usage_in_kernel_mode": e.Cpu.UsageInKernelMode,
			"usage_in_user_mode":   e.Cpu.UsageInUserMode,
		}
	}
	if e.Io != nil {
		event["io"] = common.MapStr{
			"read":  e.Io.Read,
			"write": e.Io.Write,
			"total": e.Io.Total,
		}
	}
	if e.Mem != nil {
		event["mem"] = common.MapStr{
			"fail_count": e.Mem.Failcnt,
			"limit":      e.Mem.Limit,
			"max_usage":  e.Mem.Maxusage,
			"usage":      e.Mem.Usage,
			"usage_pct":  e.Mem.UsageP,
		}
	}
	if e.Net != nil {
		event["net"] = common.MapStr{
			"total_bytes": e.Net.TotalBytes,
			"rx_bytes":    e.Net.RxBytes,
			"rx_dropped":  e.Net.RxDropped,
			"rx_errors":   e.Net.RxErrors,
			"rx_packets":  e.Net.RxPackets,
			"tx_bytes":    e.Net.TxBytes,
			"tx_dropped":  e.Net.TxDropped,
			"tx_errors":   e.Net.TxErrors,
			"tx_packets":  e.Net.TxPackets,
		}
	}
	bt.client.PublishEvent(event)
}

func getLabels(labels map[string]string) common.MapStr {
	labmap := common.MapStr{}
	for name, value := range labels {
		if !strings.HasPrefix(name, "com.docker.") {
			name = strings.Replace(name, ".", "-", -1)
			labmap[name] = value
		}
	}
	return labmap
}
<commit_msg>comment update<commit_after>package beater

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/appcelerator/amp/api/rpc/logs"
	"github.com/appcelerator/amp/api/rpc/stats"
	"github.com/appcelerator/amp/cmd/ampbeat/config"
	"github.com/appcelerator/amp/pkg/nats-streaming"
	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/golang/protobuf/proto"
	"github.com/nats-io/go-nats-streaming"
)

// Ampbeat amp libbeat type
type Ampbeat struct {
	done          chan struct{}
	config        config.Config
	client        publisher.Client
	natsStreaming *ns.NatsStreaming
}

var bt = &Ampbeat{
	done: make(chan struct{}),
}

// New Creates beater
func New(b *beat.Beat, cfg *common.Config) (beat.Beater, error) {
	bt.config = config.DefaultConfig
	if err := cfg.Unpack(&bt.config); err != nil {
		return nil, fmt.Errorf("Error reading config file: %v", err)
	}

	// NATS Connect
	hostname, err := os.Hostname()
	if err != nil {
		return nil, fmt.Errorf("Unable to get hostname: %v", err)
	}
	bt.natsStreaming = ns.NewClient(ns.DefaultURL, ns.ClusterID, b.Name+"-"+hostname, time.Minute)
	if err = bt.natsStreaming.Connect(); err != nil {
		return nil, fmt.Errorf("Unable to connect to NATS: %v", err)
	}

	return bt, nil
}

//Run run the beat
func (bt *Ampbeat) Run(b *beat.Beat) error {
	logp.Info("ampbeat is running! Hit CTRL-C to stop it.")
	bt.client = b.Publisher.Connect()

	// logs subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.LogsSubject, logMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to logs subject")

	// metrics subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.MetricsSubject, metricsMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to metrics subject")

	select {
	case <-bt.done:
		return nil
	}
}

//Stop stop the beat
func (bt *Ampbeat) Stop() {
	bt.natsStreaming.Close()
	bt.client.Close()
	close(bt.done)
}

func logMessageHandler(msg *stan.Msg) {
	e := logs.LogEntry{}
	err := proto.Unmarshal(msg.Data, &e)
	if err != nil {
		logp.Err("Error unmarshalling log entry: %s", err.Error())
		return
	}
	timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
	if err != nil {
		timestamp = time.Now()
	}
	event := common.MapStr{
		"@timestamp":           common.Time(timestamp),
		"type":                 "logs",
		"container_id":         e.ContainerId,
		"container_name":       e.ContainerName,
		"container_short_name": e.ContainerShortName,
		"container_state":      e.ContainerState,
		"service_name":         e.ServiceName,
		"service_id":           e.ServiceId,
		"task_id":              e.TaskId,
		"stack_name":           e.StackName,
		"node_id":              e.NodeId,
		"msg":                  e.Msg,
	}
	if len(e.Labels) > 0 {
		event["labels"] = getLabels(e.Labels)
	}
	bt.client.PublishEvent(event)
}

func metricsMessageHandler(msg *stan.Msg) {
	e := stats.MetricsEntry{}
	err := proto.Unmarshal(msg.Data, &e)
	if err != nil {
		logp.Err("Error unmarshalling metrics entry: %s", err.Error())
		return
	}
	timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
	if err != nil {
		timestamp = time.Now()
	}
	event := common.MapStr{
		"@timestamp":           common.Time(timestamp),
		"type":                 "metrics",
		"container_id":         e.ContainerId,
		"container_name":       e.ContainerName,
		"container_short_name": e.ContainerShortName,
		"container_state":      e.ContainerState,
		"service_name":         e.ServiceName,
		"service_id":           e.ServiceId,
		"task_id":              e.TaskId,
		"stack_name":           e.StackName,
		"node_id":              e.NodeId,
	}
	if len(e.Labels) > 0 {
		event["labels"] = getLabels(e.Labels)
	}
	if e.Cpu != nil {
		event["cpu"] = common.MapStr{
			"total_usage":          e.Cpu.TotalUsage,
			"usage_in_kernel_mode": e.Cpu.UsageInKernelMode,
			"usage_in_user_mode":   e.Cpu.UsageInUserMode,
		}
	}
	if e.Io != nil {
		event["io"] = common.MapStr{
			"read":  e.Io.Read,
			"write": e.Io.Write,
			"total": e.Io.Total,
		}
	}
	if e.Mem != nil {
		event["mem"] = common.MapStr{
			"fail_count": e.Mem.Failcnt,
			"limit":      e.Mem.Limit,
			"max_usage":  e.Mem.Maxusage,
			"usage":      e.Mem.Usage,
			"usage_pct":  e.Mem.UsageP,
		}
	}
	if e.Net != nil {
		event["net"] = common.MapStr{
			"total_bytes": e.Net.TotalBytes,
			"rx_bytes":    e.Net.RxBytes,
			"rx_dropped":  e.Net.RxDropped,
			"rx_errors":   e.Net.RxErrors,
			"rx_packets":  e.Net.RxPackets,
			"tx_bytes":    e.Net.TxBytes,
			"tx_dropped":  e.Net.TxDropped,
			"tx_errors":   e.Net.TxErrors,
			"tx_packets":  e.Net.TxPackets,
		}
	}
	bt.client.PublishEvent(event)
}

// Need to replace '.'' by '-'', becuse '.' is used by ES for object operator it can't be part of a field name.
func getLabels(labels map[string]string) common.MapStr {
	labmap := common.MapStr{}
	for name, value := range labels {
		if !strings.HasPrefix(name, "com.docker.") {
			name = strings.Replace(name, ".", "-", -1)
			labmap[name] = value
		}
	}
	return labmap
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"time"

	"google.golang.org/api/googleapi"
	runapi "google.golang.org/api/run/v1"
)

// parseEnv parses K=V pairs into a map.
func parseEnv(envs []string) map[string]string {
	out := make(map[string]string)
	for _, v := range envs {
		p := strings.SplitN(v, "=", 2)
		out[p[0]] = p[1]
	}
	return out
}

// deploy reimplements the "gcloud run deploy" command, including setting IAM policy and
// waiting for Service to be Ready.
func deploy(project, name, image, region string, envs []string, options options) (string, error) {
	envVars := parseEnv(envs)

	client, err := runClient(region)
	if err != nil {
		return "", fmt.Errorf("failed to initialize Run API client: %w", err)
	}

	svc, err := getService(project, name, region)
	if err == nil {
		// existing service
		svc = patchService(svc, envVars, image)
		_, err = client.Namespaces.Services.ReplaceService("namespaces/"+project+"/services/"+name, svc).Do()
		if err != nil {
			if e, ok := err.(*googleapi.Error); ok {
				return "", fmt.Errorf("failed to deploy existing Service: code=%d message=%s -- %s", e.Code, e.Message, e.Body)
			}
			return "", fmt.Errorf("failed to deploy to existing Service: %w", err)
		}
	} else {
		// new service
		svc := newService(name, project, image, envVars)
		_, err = client.Namespaces.Services.Create("namespaces/"+project, svc).Do()
		if err != nil {
			if e, ok := err.(*googleapi.Error); ok {
				return "", fmt.Errorf("failed to deploy a new Service: code=%d message=%s -- %s", e.Code, e.Message, e.Body)
			}
			return "", fmt.Errorf("failed to deploy a new Service: %w", err)
		}
	}

	if options.AllowUnauthenticated == nil || *options.AllowUnauthenticated {
		if err := allowUnauthenticated(project, name, region); err != nil {
			return "", fmt.Errorf("failed to allow unauthenticated requests on the service: %w", err)
		}
	}

	if err := waitReady(project, name, region); err != nil {
		return "", err
	}

	out, err := getService(project, name, region)
	if err != nil {
		return "", fmt.Errorf("failed to get service after deploying: %w", err)
	}
	return out.Status.Url, nil
}

// newService initializes a new Knative Service object with given properties.
func newService(name, project, image string, envs map[string]string) *runapi.Service {
	var envVars []*runapi.EnvVar
	for k, v := range envs {
		envVars = append(envVars, &runapi.EnvVar{Name: k, Value: v})
	}

	svc := &runapi.Service{
		ApiVersion: "serving.knative.dev/v1",
		Kind:       "Service",
		Metadata: &runapi.ObjectMeta{
			Annotations: make(map[string]string),
			Name:        name,
			Namespace:   project,
		},
		Spec: &runapi.ServiceSpec{
			Template: &runapi.RevisionTemplate{
				Metadata: &runapi.ObjectMeta{
					Name:        generateRevisionName(name, 0),
					Annotations: make(map[string]string),
				},
				Spec: &runapi.RevisionSpec{
					Containers: []*runapi.Container{
						{
							Image: image,
							Env:   envVars,
						},
					},
				},
				ForceSendFields: nil,
				NullFields:      nil,
			},
		},
	}
	applyMeta(svc.Metadata.Annotations, image)
	applyMeta(svc.Spec.Template.Metadata.Annotations, image)

	return svc
}

// applyMeta applies optional annotations to the specified Metadata.Annotation field.
func applyMeta(meta map[string]string, userImage string) {
	meta["client.knative.dev/user-image"] = userImage
	meta["run.googleapis.com/client-name"] = "cloud-run-button"
}

// generateRevisionName attempts to generate a random revision name that is alphabetically increasing but also has
// a random suffix. objectGeneration is the current object generation.
func generateRevisionName(name string, objectGeneration int64) string {
	num := fmt.Sprintf("%05d", objectGeneration+1)
	out := name + "-" + num + "-"
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	for i := 0; i < 3; i++ {
		out += string(int('a') + r.Intn(26))
	}
	return out
}

// patchService modifies an existing Service with requested changes.
func patchService(svc *runapi.Service, envs map[string]string, image string) *runapi.Service {
	// merge env vars
	svc.Spec.Template.Spec.Containers[0].Env = mergeEnvs(svc.Spec.Template.Spec.Containers[0].Env, envs)

	// update container image
	svc.Spec.Template.Spec.Containers[0].Image = image

	// apply metadata annotations
	applyMeta(svc.Metadata.Annotations, image)
	applyMeta(svc.Spec.Template.Metadata.Annotations, image)

	// update revision name
	svc.Spec.Template.Metadata.Name = generateRevisionName(svc.Metadata.Name, svc.Metadata.Generation)

	return svc
}

// mergeEnvs updates variables in existing, and adds missing ones.
func mergeEnvs(existing []*runapi.EnvVar, env map[string]string) []*runapi.EnvVar {
	for i, ee := range existing {
		if v, ok := env[ee.Name]; ok {
			existing[i].Value = v
			delete(env, ee.Name)
		}
	}
	// add missing ones
	for k, v := range env {
		existing = append(existing, &runapi.EnvVar{Name: k, Value: v})
	}
	return existing
}

// waitReady waits until the specified service reaches Ready status
func waitReady(project, name, region string) error {
	wait := time.Second * 30
	deadline := time.Now().Add(wait)
	for time.Now().Before(deadline) {
		svc, err := getService(project, name, region)
		if err != nil {
			return fmt.Errorf("failed to query Service for readiness: %w", err)
		}

		for _, cond := range svc.Status.Conditions {
			if cond.Type == "Ready" {
				if cond.Status == "True" {
					return nil
				} else if cond.Status == "False" {
					return fmt.Errorf("reason=%s message=%s", cond.Reason, cond.Message)
				}
			}
		}
	}
	return fmt.Errorf("the service did not become ready in %s, check Cloud Console for logs to see why it failed", wait)
}

// allowUnauthenticated sets IAM policy on the specified Cloud Run service to give allUsers subject
// roles/run.invoker role.
func allowUnauthenticated(project, name, region string) error {
	client, err := runapi.NewService(context.TODO())
	if err != nil {
		return fmt.Errorf("failed to initialize Run API client: %w", err)
	}

	res := fmt.Sprintf("projects/%s/locations/%s/services/%s", project, region, name)
	policy, err := client.Projects.Locations.Services.GetIamPolicy(res).Do()
	if err != nil {
		return fmt.Errorf("failed to get IAM policy for Cloud Run Service: %w", err)
	}

	policy.Bindings = append(policy.Bindings, &runapi.Binding{
		Members: []string{"allUsers"},
		Role:    "roles/run.invoker",
	})

	_, err = client.Projects.Locations.Services.SetIamPolicy(res, &runapi.SetIamPolicyRequest{Policy: policy}).Do()
	if err != nil {
		var extra string
		e, ok := err.(*googleapi.Error)
		if ok {
			extra = fmt.Sprintf("code=%d, message=%s -- %s", e.Code, e.Message, e.Body)
		}
		return fmt.Errorf("failed to set IAM policy for Cloud Run Service: %w %s", err, extra)
	}
	return nil
}
<commit_msg>Initialize metadata.annotations map if empty (#165)<commit_after>package main

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"time"

	"google.golang.org/api/googleapi"
	runapi "google.golang.org/api/run/v1"
)

// parseEnv parses K=V pairs into a map.
func parseEnv(envs []string) map[string]string {
	out := make(map[string]string)
	for _, v := range envs {
		p := strings.SplitN(v, "=", 2)
		out[p[0]] = p[1]
	}
	return out
}

// deploy reimplements the "gcloud run deploy" command, including setting IAM policy and
// waiting for Service to be Ready.
func deploy(project, name, image, region string, envs []string, options options) (string, error) {
	envVars := parseEnv(envs)

	client, err := runClient(region)
	if err != nil {
		return "", fmt.Errorf("failed to initialize Run API client: %w", err)
	}

	svc, err := getService(project, name, region)
	if err == nil {
		// existing service
		svc = patchService(svc, envVars, image)
		_, err = client.Namespaces.Services.ReplaceService("namespaces/"+project+"/services/"+name, svc).Do()
		if err != nil {
			if e, ok := err.(*googleapi.Error); ok {
				return "", fmt.Errorf("failed to deploy existing Service: code=%d message=%s -- %s", e.Code, e.Message, e.Body)
			}
			return "", fmt.Errorf("failed to deploy to existing Service: %w", err)
		}
	} else {
		// new service
		svc := newService(name, project, image, envVars)
		_, err = client.Namespaces.Services.Create("namespaces/"+project, svc).Do()
		if err != nil {
			if e, ok := err.(*googleapi.Error); ok {
				return "", fmt.Errorf("failed to deploy a new Service: code=%d message=%s -- %s", e.Code, e.Message, e.Body)
			}
			return "", fmt.Errorf("failed to deploy a new Service: %w", err)
		}
	}

	if options.AllowUnauthenticated == nil || *options.AllowUnauthenticated {
		if err := allowUnauthenticated(project, name, region); err != nil {
			return "", fmt.Errorf("failed to allow unauthenticated requests on the service: %w", err)
		}
	}

	if err := waitReady(project, name, region); err != nil {
		return "", err
	}

	out, err := getService(project, name, region)
	if err != nil {
		return "", fmt.Errorf("failed to get service after deploying: %w", err)
	}
	return out.Status.Url, nil
}

// newService initializes a new Knative Service object with given properties.
func newService(name, project, image string, envs map[string]string) *runapi.Service {
	var envVars []*runapi.EnvVar
	for k, v := range envs {
		envVars = append(envVars, &runapi.EnvVar{Name: k, Value: v})
	}

	svc := &runapi.Service{
		ApiVersion: "serving.knative.dev/v1",
		Kind:       "Service",
		Metadata: &runapi.ObjectMeta{
			Annotations: make(map[string]string),
			Name:        name,
			Namespace:   project,
		},
		Spec: &runapi.ServiceSpec{
			Template: &runapi.RevisionTemplate{
				Metadata: &runapi.ObjectMeta{
					Name:        generateRevisionName(name, 0),
					Annotations: make(map[string]string),
				},
				Spec: &runapi.RevisionSpec{
					Containers: []*runapi.Container{
						{
							Image: image,
							Env:   envVars,
						},
					},
				},
				ForceSendFields: nil,
				NullFields:      nil,
			},
		},
	}
	applyMeta(svc.Metadata, image)
	applyMeta(svc.Spec.Template.Metadata, image)

	return svc
}

// applyMeta applies optional annotations to the specified Metadata.Annotation field.
func applyMeta(meta *runapi.ObjectMeta, userImage string) {
	if meta.Annotations == nil {
		meta.Annotations = make(map[string]string)
	}
	meta.Annotations["client.knative.dev/user-image"] = userImage
	meta.Annotations["run.googleapis.com/client-name"] = "cloud-run-button"
}

// generateRevisionName attempts to generate a random revision name that is alphabetically increasing but also has
// a random suffix. objectGeneration is the current object generation.
func generateRevisionName(name string, objectGeneration int64) string {
	num := fmt.Sprintf("%05d", objectGeneration+1)
	out := name + "-" + num + "-"
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	for i := 0; i < 3; i++ {
		out += string(int('a') + r.Intn(26))
	}
	return out
}

// patchService modifies an existing Service with requested changes.
func patchService(svc *runapi.Service, envs map[string]string, image string) *runapi.Service {
	// merge env vars
	svc.Spec.Template.Spec.Containers[0].Env = mergeEnvs(svc.Spec.Template.Spec.Containers[0].Env, envs)

	// update container image
	svc.Spec.Template.Spec.Containers[0].Image = image

	// apply metadata annotations
	applyMeta(svc.Metadata, image)
	applyMeta(svc.Spec.Template.Metadata, image)

	// update revision name
	svc.Spec.Template.Metadata.Name = generateRevisionName(svc.Metadata.Name, svc.Metadata.Generation)

	return svc
}

// mergeEnvs updates variables in existing, and adds missing ones.
func mergeEnvs(existing []*runapi.EnvVar, env map[string]string) []*runapi.EnvVar {
	for i, ee := range existing {
		if v, ok := env[ee.Name]; ok {
			existing[i].Value = v
			delete(env, ee.Name)
		}
	}
	// add missing ones
	for k, v := range env {
		existing = append(existing, &runapi.EnvVar{Name: k, Value: v})
	}
	return existing
}

// waitReady waits until the specified service reaches Ready status
func waitReady(project, name, region string) error {
	wait := time.Second * 30
	deadline := time.Now().Add(wait)
	for time.Now().Before(deadline) {
		svc, err := getService(project, name, region)
		if err != nil {
			return fmt.Errorf("failed to query Service for readiness: %w", err)
		}

		for _, cond := range svc.Status.Conditions {
			if cond.Type == "Ready" {
				if cond.Status == "True" {
					return nil
				} else if cond.Status == "False" {
					return fmt.Errorf("reason=%s message=%s", cond.Reason, cond.Message)
				}
			}
		}
	}
	return fmt.Errorf("the service did not become ready in %s, check Cloud Console for logs to see why it failed", wait)
}

// allowUnauthenticated sets IAM policy on the specified Cloud Run service to give allUsers subject
// roles/run.invoker role.
func allowUnauthenticated(project, name, region string) error {
	client, err := runapi.NewService(context.TODO())
	if err != nil {
		return fmt.Errorf("failed to initialize Run API client: %w", err)
	}

	res := fmt.Sprintf("projects/%s/locations/%s/services/%s", project, region, name)
	policy, err := client.Projects.Locations.Services.GetIamPolicy(res).Do()
	if err != nil {
		return fmt.Errorf("failed to get IAM policy for Cloud Run Service: %w", err)
	}

	policy.Bindings = append(policy.Bindings, &runapi.Binding{
		Members: []string{"allUsers"},
		Role:    "roles/run.invoker",
	})

	_, err = client.Projects.Locations.Services.SetIamPolicy(res, &runapi.SetIamPolicyRequest{Policy: policy}).Do()
	if err != nil {
		var extra string
		e, ok := err.(*googleapi.Error)
		if ok {
			extra = fmt.Sprintf("code=%d, message=%s -- %s", e.Code, e.Message, e.Body)
		}
		return fmt.Errorf("failed to set IAM policy for Cloud Run Service: %w %s", err, extra)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/cloudfoundry-incubator/guardian/guardiancmd"
	"github.com/concourse/atc"
	"github.com/concourse/baggageclaim/baggageclaimcmd"
	"github.com/concourse/baggageclaim/fs"
	"github.com/concourse/bin/bindata"
	"github.com/jessevdk/go-flags"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
)

const btrfsFSType = 0x9123683e

type GardenBackend guardiancmd.GuardianCommand

func (cmd WorkerCommand) lessenRequirements(command *flags.Command) {
	command.FindOptionByLongName("garden-depot").Required = false
	command.FindOptionByLongName("garden-graph").Required = false
	command.FindOptionByLongName("garden-runc-bin").Required = false
	command.FindOptionByLongName("garden-dadoo-bin").Required = false
	command.FindOptionByLongName("garden-init-bin").Required = false
	command.FindOptionByLongName("garden-iodaemon-bin").Required = false
	command.FindOptionByLongName("garden-kawasaki-bin").Required = false
	command.FindOptionByLongName("garden-nstar-bin").Required = false
	command.FindOptionByLongName("garden-tar-bin").Required = false
}

func (cmd *WorkerCommand) gardenRunner(logger lager.Logger, args []string) (atc.Worker, ifrit.Runner, error) {
	err := cmd.checkRoot()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	assetsDir, err := cmd.restoreVersionedAssets(logger.Session("unpack-assets"))
	if err != nil {
		return atc.Worker{}, nil, err
	}

	btrfsToolsDir := filepath.Join(assetsDir, "btrfs")
	iptablesDir := filepath.Join(assetsDir, "iptables")

	err = os.Setenv(
		"PATH",
		strings.Join(
			[]string{
				btrfsToolsDir,
				filepath.Join(iptablesDir, "sbin"),
				filepath.Join(iptablesDir, "bin"),
				os.Getenv("PATH"),
			},
			string(os.PathListSeparator),
		),
	)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	depotDir := filepath.Join(cmd.WorkDir, "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err = os.MkdirAll(depotDir, 0755)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	cmd.Garden.Server.BindIP = guardiancmd.IPFlag(cmd.BindIP)

	cmd.Garden.Containers.Dir = guardiancmd.DirFlag(depotDir)

	cmd.Garden.Bin.Runc = filepath.Join(assetsDir, "bin", "runc")
	cmd.Garden.Bin.Dadoo = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "dadoo"))
	cmd.Garden.Bin.Init = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "init"))
	cmd.Garden.Bin.IODaemon = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "iodaemon"))
	cmd.Garden.Bin.Kawasaki = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "kawasaki"))
	cmd.Garden.Bin.NSTar = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "nstar"))
	cmd.Garden.Bin.Tar = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "tar"))

	cmd.Garden.Network.AllowHostAccess = true

	worker := atc.Worker{
		Platform: "linux",
		Tags:     cmd.Tags,

		HTTPProxyURL:  cmd.HTTPProxy.String(),
		HTTPSProxyURL: cmd.HTTPSProxy.String(),
		NoProxy:       strings.Join(cmd.NoProxy, ","),
		StartTime:     time.Now().Unix(),
	}

	worker.ResourceTypes, err = cmd.extractResources(
		logger.Session("extract-resources"),
		assetsDir,
	)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	worker.Name, err = cmd.workerName()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	runner := guardiancmd.GuardianCommand(cmd.Garden)
	return worker, &runner, nil
}

func (cmd *WorkerCommand) restoreVersionedAssets(logger lager.Logger) (string, error) {
	assetsDir := filepath.Join(cmd.WorkDir, Version)

	restoredDir := filepath.Join(assetsDir, "linux")

	okMarker := filepath.Join(assetsDir, "ok")

	_, err := os.Stat(okMarker)
	if err == nil {
		logger.Info("already-done")
		return restoredDir, nil
	}

	logger.Info("unpacking")

	err = bindata.RestoreAssets(assetsDir, "linux")
	if err != nil {
		logger.Error("failed-to-unpack", err)
		return "", err
	}

	ok, err := os.Create(okMarker)
	if err != nil {
		logger.Error("failed-to-create-ok-marker", err)
		return "", err
	}

	err = ok.Close()
	if err != nil {
		logger.Error("failed-to-close-ok-marker", err)
		return "", err
	}

	logger.Info("done")

	return restoredDir, nil
}

func (cmd *WorkerCommand) baggageclaimRunner(logger lager.Logger) (ifrit.Runner, error) {
	volumesImage := filepath.Join(cmd.WorkDir, "volumes.img")
	volumesDir := filepath.Join(cmd.WorkDir, "volumes")

	err := os.MkdirAll(volumesDir, 0755)
	if err != nil {
		return nil, err
	}

	var fsStat syscall.Statfs_t
	err = syscall.Statfs(volumesDir, &fsStat)
	if err != nil {
		return nil, fmt.Errorf("failed to stat volumes filesystem: %s", err)
	}

	if fsStat.Type != btrfsFSType {
		filesystem := fs.New(logger.Session("fs"), volumesImage, volumesDir)

		err = filesystem.Create(fsStat.Blocks * uint64(fsStat.Bsize))
		if err != nil {
			logger.Error("falling-back-on-naive-driver", err)
			return cmd.naiveBaggageclaimRunner(logger)
		}
	}

	bc := &baggageclaimcmd.BaggageclaimCommand{
		BindIP:   baggageclaimcmd.IPFlag(cmd.Baggageclaim.BindIP.IP().String()),
		BindPort: cmd.Baggageclaim.BindPort,

		VolumesDir: baggageclaimcmd.DirFlag(volumesDir),

		Driver: "btrfs",

		ReapInterval: cmd.Baggageclaim.ReapInterval,

		Metrics: cmd.Metrics,
	}

	return bc.Runner(nil)
}

func (cmd *WorkerCommand) extractResources(logger lager.Logger, assetsDir string) ([]atc.WorkerResourceType, error) {
	var resourceTypes []atc.WorkerResourceType

	resourcesDir := filepath.Join(assetsDir, "resources")

	infos, err := ioutil.ReadDir(resourcesDir)
	if err != nil {
		logger.Error("failed-to-list-resource-assets", err)
		return nil, err
	}

	for _, info := range infos {
		resourceType := info.Name()

		workerResourceType, err := cmd.extractResource(
			logger.Session("extract", lager.Data{"resource-type": resourceType}),
			assetsDir,
			resourcesDir,
			resourceType,
		)
		if err != nil {
			logger.Error("failed-to-extract-resource", err)
			return nil, err
		}

		resourceTypes = append(resourceTypes, workerResourceType)
	}

	return resourceTypes, nil
}

func (cmd *WorkerCommand) extractResource(
	logger lager.Logger,
	assetsDir string,
	resourcesDir string,
	resourceType string,
) (atc.WorkerResourceType, error) {
	resourceImagesDir := filepath.Join(assetsDir, "resource-images")
	tarBin := filepath.Join(assetsDir, "bin", "tar")

	archive := filepath.Join(resourcesDir, resourceType, "rootfs.tar.gz")

	extractedDir := filepath.Join(resourceImagesDir, resourceType)

	rootfsDir := filepath.Join(extractedDir, "rootfs")
	okMarker := filepath.Join(extractedDir, "ok")

	var version string
	versionFile, err := os.Open(filepath.Join(resourcesDir, resourceType, "version"))
	if err != nil {
		logger.Error("failed-to-read-version", err)
		return atc.WorkerResourceType{}, err
	}

	_, err = fmt.Fscanf(versionFile, "%s", &version)
	if err != nil {
		logger.Error("failed-to-parse-version", err)
		return atc.WorkerResourceType{}, err
	}

	defer versionFile.Close()

	_, err = os.Stat(okMarker)
	if err == nil {
		logger.Info("already-extracted")
	} else {
		logger.Info("extracting")

		err := os.RemoveAll(rootfsDir)
		if err != nil {
			logger.Error("failed-to-clear-out-existing-rootfs", err)
			return atc.WorkerResourceType{}, err
		}

		err = os.MkdirAll(rootfsDir, 0755)
		if err != nil {
			logger.Error("failed-to-create-rootfs-dir", err)
			return atc.WorkerResourceType{}, err
		}

		tar := exec.Command(tarBin, "-zxf", archive, "-C", rootfsDir)

		output, err := tar.CombinedOutput()
		if err != nil {
			logger.Error("failed-to-extract-resource", err, lager.Data{
				"output": string(output),
			})
			return atc.WorkerResourceType{}, err
		}

		ok, err := os.Create(okMarker)
		if err != nil {
			logger.Error("failed-to-create-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}

		err = ok.Close()
		if err != nil {
			logger.Error("failed-to-close-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}
	}

	return atc.WorkerResourceType{
		Type:    resourceType,
		Image:   rootfsDir,
		Version: version,
	}, nil
}
<commit_msg>use garden flag to set iptables<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/cloudfoundry-incubator/guardian/guardiancmd"
	"github.com/concourse/atc"
	"github.com/concourse/baggageclaim/baggageclaimcmd"
	"github.com/concourse/baggageclaim/fs"
	"github.com/concourse/bin/bindata"
	"github.com/jessevdk/go-flags"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
)

const btrfsFSType = 0x9123683e

type GardenBackend guardiancmd.GuardianCommand

func (cmd WorkerCommand) lessenRequirements(command *flags.Command) {
	command.FindOptionByLongName("garden-depot").Required = false
	command.FindOptionByLongName("garden-graph").Required = false
	command.FindOptionByLongName("garden-runc-bin").Required = false
	command.FindOptionByLongName("garden-dadoo-bin").Required = false
	command.FindOptionByLongName("garden-init-bin").Required = false
	command.FindOptionByLongName("garden-iodaemon-bin").Required = false
	command.FindOptionByLongName("garden-kawasaki-bin").Required = false
	command.FindOptionByLongName("garden-nstar-bin").Required = false
	command.FindOptionByLongName("garden-tar-bin").Required = false
}

func (cmd *WorkerCommand) gardenRunner(logger lager.Logger, args []string) (atc.Worker, ifrit.Runner, error) {
	err := cmd.checkRoot()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	assetsDir, err := cmd.restoreVersionedAssets(logger.Session("unpack-assets"))
	if err != nil {
		return atc.Worker{}, nil, err
	}

	btrfsToolsDir := filepath.Join(assetsDir, "btrfs")

	err = os.Setenv(
		"PATH",
		strings.Join(
			[]string{
				btrfsToolsDir,
				os.Getenv("PATH"),
			},
			string(os.PathListSeparator),
		),
	)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	depotDir := filepath.Join(cmd.WorkDir, "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err = os.MkdirAll(depotDir, 0755)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	cmd.Garden.Server.BindIP = guardiancmd.IPFlag(cmd.BindIP)

	cmd.Garden.Containers.Dir = guardiancmd.DirFlag(depotDir)

	cmd.Garden.Bin.Runc = filepath.Join(assetsDir, "bin", "runc")
	cmd.Garden.Bin.Dadoo = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "dadoo"))
	cmd.Garden.Bin.Init = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "init"))
	cmd.Garden.Bin.IODaemon = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "iodaemon"))
	cmd.Garden.Bin.Kawasaki = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "kawasaki"))
	cmd.Garden.Bin.NSTar = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "nstar"))
	cmd.Garden.Bin.Tar = guardiancmd.FileFlag(filepath.Join(assetsDir, "bin", "tar"))

	iptablesDir := filepath.Join(assetsDir, "iptables")
	cmd.Garden.Bin.IPTables = guardiancmd.FileFlag(filepath.Join(iptablesDir, "sbin", "iptables"))

	cmd.Garden.Network.AllowHostAccess = true

	worker := atc.Worker{
		Platform: "linux",
		Tags:     cmd.Tags,

		HTTPProxyURL:  cmd.HTTPProxy.String(),
		HTTPSProxyURL: cmd.HTTPSProxy.String(),
		NoProxy:       strings.Join(cmd.NoProxy, ","),
		StartTime:     time.Now().Unix(),
	}

	worker.ResourceTypes, err = cmd.extractResources(
		logger.Session("extract-resources"),
		assetsDir,
	)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	worker.Name, err = cmd.workerName()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	runner := guardiancmd.GuardianCommand(cmd.Garden)
	return worker, &runner, nil
}

func (cmd *WorkerCommand) restoreVersionedAssets(logger lager.Logger) (string, error) {
	assetsDir := filepath.Join(cmd.WorkDir, Version)

	restoredDir := filepath.Join(assetsDir, "linux")

	okMarker := filepath.Join(assetsDir, "ok")

	_, err := os.Stat(okMarker)
	if err == nil {
		logger.Info("already-done")
		return restoredDir, nil
	}

	logger.Info("unpacking")

	err = bindata.RestoreAssets(assetsDir, "linux")
	if err != nil {
		logger.Error("failed-to-unpack", err)
		return "", err
	}

	ok, err := os.Create(okMarker)
	if err != nil {
		logger.Error("failed-to-create-ok-marker", err)
		return "", err
	}

	err = ok.Close()
	if err != nil {
		logger.Error("failed-to-close-ok-marker", err)
		return "", err
	}

	logger.Info("done")

	return restoredDir, nil
}

func (cmd *WorkerCommand) baggageclaimRunner(logger lager.Logger) (ifrit.Runner, error) {
	volumesImage := filepath.Join(cmd.WorkDir, "volumes.img")
	volumesDir := filepath.Join(cmd.WorkDir, "volumes")

	err := os.MkdirAll(volumesDir, 0755)
	if err != nil {
		return nil, err
	}

	var fsStat syscall.Statfs_t
	err = syscall.Statfs(volumesDir, &fsStat)
	if err != nil {
		return nil, fmt.Errorf("failed to stat volumes filesystem: %s", err)
	}

	if fsStat.Type != btrfsFSType {
		filesystem := fs.New(logger.Session("fs"), volumesImage, volumesDir)

		err = filesystem.Create(fsStat.Blocks * uint64(fsStat.Bsize))
		if err != nil {
			logger.Error("falling-back-on-naive-driver", err)
			return cmd.naiveBaggageclaimRunner(logger)
		}
	}

	bc := &baggageclaimcmd.BaggageclaimCommand{
		BindIP:   baggageclaimcmd.IPFlag(cmd.Baggageclaim.BindIP.IP().String()),
		BindPort: cmd.Baggageclaim.BindPort,

		VolumesDir: baggageclaimcmd.DirFlag(volumesDir),

		Driver: "btrfs",

		ReapInterval: cmd.Baggageclaim.ReapInterval,

		Metrics: cmd.Metrics,
	}

	return bc.Runner(nil)
}

func (cmd *WorkerCommand) extractResources(logger lager.Logger, assetsDir string) ([]atc.WorkerResourceType, error) {
	var resourceTypes []atc.WorkerResourceType

	resourcesDir := filepath.Join(assetsDir, "resources")

	infos, err := ioutil.ReadDir(resourcesDir)
	if err != nil {
		logger.Error("failed-to-list-resource-assets", err)
		return nil, err
	}

	for _, info := range infos {
		resourceType := info.Name()

		workerResourceType, err := cmd.extractResource(
			logger.Session("extract", lager.Data{"resource-type": resourceType}),
			assetsDir,
			resourcesDir,
			resourceType,
		)
		if err != nil {
			logger.Error("failed-to-extract-resource", err)
			return nil, err
		}

		resourceTypes = append(resourceTypes, workerResourceType)
	}

	return resourceTypes, nil
}

func (cmd *WorkerCommand) extractResource(
	logger lager.Logger,
	assetsDir string,
	resourcesDir string,
	resourceType string,
) (atc.WorkerResourceType, error) {
	resourceImagesDir := filepath.Join(assetsDir, "resource-images")
	tarBin := filepath.Join(assetsDir, "bin", "tar")

	archive := filepath.Join(resourcesDir, resourceType, "rootfs.tar.gz")

	extractedDir := filepath.Join(resourceImagesDir, resourceType)

	rootfsDir := filepath.Join(extractedDir, "rootfs")
	okMarker := filepath.Join(extractedDir, "ok")

	var version string
	versionFile, err := os.Open(filepath.Join(resourcesDir, resourceType, "version"))
	if err != nil {
		logger.Error("failed-to-read-version", err)
		return atc.WorkerResourceType{}, err
	}

	_, err = fmt.Fscanf(versionFile, "%s", &version)
	if err != nil {
		logger.Error("failed-to-parse-version", err)
		return atc.WorkerResourceType{}, err
	}

	defer versionFile.Close()

	_, err = os.Stat(okMarker)
	if err == nil {
		logger.Info("already-extracted")
	} else {
		logger.Info("extracting")

		err := os.RemoveAll(rootfsDir)
		if err != nil {
			logger.Error("failed-to-clear-out-existing-rootfs", err)
			return atc.WorkerResourceType{}, err
		}

		err = os.MkdirAll(rootfsDir, 0755)
		if err != nil {
			logger.Error("failed-to-create-rootfs-dir", err)
			return atc.WorkerResourceType{}, err
		}

		tar := exec.Command(tarBin, "-zxf", archive, "-C", rootfsDir)

		output, err := tar.CombinedOutput()
		if err != nil {
			logger.Error("failed-to-extract-resource", err, lager.Data{
				"output": string(output),
			})
			return atc.WorkerResourceType{}, err
		}

		ok, err := os.Create(okMarker)
		if err != nil {
			logger.Error("failed-to-create-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}

		err = ok.Close()
		if err != nil {
			logger.Error("failed-to-close-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}
	}

	return atc.WorkerResourceType{
		Type:    resourceType,
		Image:   rootfsDir,
		Version: version,
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus 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 main

import (
	"bufio"
	"bytes"
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/promlog/flag"
	"github.com/prometheus/common/version"
	"github.com/prometheus/exporter-toolkit/web"
	"github.com/prometheus/graphite_exporter/collector"
	"github.com/prometheus/statsd_exporter/pkg/mapper"
	"github.com/prometheus/statsd_exporter/pkg/mappercache/lru"
	"github.com/prometheus/statsd_exporter/pkg/mappercache/randomreplacement"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	listenAddress   = kingpin.Flag("web.listen-address", "Address on which to expose metrics.").Default(":9108").String()
	metricsPath     = kingpin.Flag("web.telemetry-path", "Path under which to expose Prometheus metrics.").Default("/metrics").String()
	graphiteAddress = kingpin.Flag("graphite.listen-address", "TCP and UDP address on which to accept samples.").Default(":9109").String()
	mappingConfig   = kingpin.Flag("graphite.mapping-config", "Metric mapping configuration file name.").Default("").String()
	sampleExpiry    = kingpin.Flag("graphite.sample-expiry", "How long a sample is valid for.").Default("5m").Duration()
	strictMatch     = kingpin.Flag("graphite.mapping-strict-match", "Only store metrics that match the mapping configuration.").Bool()
	cacheSize       = kingpin.Flag("graphite.cache-size", "Maximum size of your metric mapping cache. Relies on least recently used replacement policy if max size is reached.").Default("1000").Int()
	cacheType       = kingpin.Flag("graphite.cache-type", "Metric mapping cache type. Valid options are \"lru\" and \"random\"").Default("lru").Enum("lru", "random")
	dumpFSMPath     = kingpin.Flag("debug.dump-fsm", "The path to dump internal FSM generated for glob matching as Dot file.").Default("").String()
	checkConfig     = kingpin.Flag("check-config", "Check configuration and exit.").Default("false").Bool()
	configFile      = kingpin.Flag("web.config", "[EXPERIMENTAL] Path to config yaml file that can enable TLS or authentication.").Default("").String()
)

func init() {
	prometheus.MustRegister(version.NewCollector("graphite_exporter"))
}

func dumpFSM(mapper *mapper.MetricMapper, dumpFilename string, logger log.Logger) error {
	if mapper.FSM == nil {
		return fmt.Errorf("no FSM available to be dumped, possibly because the mapping contains regex patterns")
	}
	f, err := os.Create(dumpFilename)
	if err != nil {
		return err
	}
	level.Info(logger).Log("msg", "Start dumping FSM", "to", dumpFilename)
	w := bufio.NewWriter(f)
	mapper.FSM.DumpFSM(w)
	w.Flush()
	f.Close()
	level.Info(logger).Log("msg", "Finish dumping FSM")
	return nil
}

func main() {
	promlogConfig := &promlog.Config{}
	flag.AddFlags(kingpin.CommandLine, promlogConfig)
	kingpin.Version(version.Print("graphite_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()
	logger := promlog.New(promlogConfig)

	level.Info(logger).Log("msg", "Starting graphite_exporter", "version_info", version.Info())
	level.Info(logger).Log("build_context", version.BuildContext())

	http.Handle(*metricsPath, promhttp.Handler())
	c := collector.NewGraphiteCollector(logger, *strictMatch, *sampleExpiry)
	prometheus.MustRegister(c)

	metricMapper := &mapper.MetricMapper{}
	if *mappingConfig != "" {
		err := metricMapper.InitFromFile(*mappingConfig)
		if err != nil {
			level.Error(logger).Log("msg", "Error loading metric mapping config", "err", err)
			os.Exit(1)
		}
	}

	cache, err := getCache(*cacheSize, *cacheType, prometheus.DefaultRegisterer)
	if err != nil {
		level.Error(logger).Log("msg", "error initializing mapper cache", "err", err)
		os.Exit(1)
	}
	metricMapper.UseCache(cache)

	if *checkConfig {
		level.Info(logger).Log("msg", "Configuration check successful, exiting")
		return
	}

	if *dumpFSMPath != "" {
		err := dumpFSM(metricMapper, *dumpFSMPath, logger)
		if err != nil {
			level.Error(logger).Log("msg", "Error dumping FSM", "err", err)
			os.Exit(1)
		}
	}

	c.SetMapper(metricMapper)

	tcpSock, err := net.Listen("tcp", *graphiteAddress)
	if err != nil {
		level.Error(logger).Log("msg", "Error binding to TCP socket", "err", err)
		os.Exit(1)
	}
	go func() {
		for {
			conn, err := tcpSock.Accept()
			if err != nil {
				level.Error(logger).Log("msg", "Error accepting TCP connection", "err", err)
				continue
			}
			go func() {
				defer conn.Close()
				c.ProcessReader(conn)
			}()
		}
	}()

	udpAddress, err := net.ResolveUDPAddr("udp", *graphiteAddress)
	if err != nil {
		level.Error(logger).Log("msg", "Error resolving UDP address", "err", err)
		os.Exit(1)
	}
	udpSock, err := net.ListenUDP("udp", udpAddress)
	if err != nil {
		level.Error(logger).Log("msg", "Error listening to UDP address", "err", err)
		os.Exit(1)
	}
	go func() {
		defer udpSock.Close()
		for {
			buf := make([]byte, 65536)
			chars, srcAddress, err := udpSock.ReadFromUDP(buf)
			if err != nil {
				level.Error(logger).Log("msg", "Error reading UDP packet", "from", srcAddress, "err", err)
				continue
			}
			go c.ProcessReader(bytes.NewReader(buf[0:chars]))
		}
	}()

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}
		w.Write([]byte(`
      <title>Graphite Exporter</title>

Graphite Exporter

Accepting plaintext Graphite samples over TCP and UDP on ` + *graphiteAddress + `

<a href="` + *metricsPath + `">Metrics

      `))
	})

	level.Info(logger).Log("msg", "Listening on "+*listenAddress)
	server := &http.Server{Addr: *listenAddress}
	if err := web.ListenAndServe(server, *configFile, logger); err != nil {
		level.Error(logger).Log("err", err)
		os.Exit(1)
	}
}

// TODO(mr): this is copied verbatim from statsd_exporter/main.go. It should be a
// convenience function in mappercache, but that caused an import cycle.
func getCache(cacheSize int, cacheType string, registerer prometheus.Registerer) (mapper.MetricMapperCache, error) {
	var cache mapper.MetricMapperCache
	var err error
	if cacheSize == 0 {
		return nil, nil
	} else {
		switch cacheType {
		case "lru":
			cache, err = lru.NewMetricMapperLRUCache(registerer, cacheSize)
		case "random":
			cache, err = randomreplacement.NewMetricMapperRRCache(registerer, cacheSize)
		default:
			err = fmt.Errorf("unsupported cache type %q", cacheType)
		}

		if err != nil {
			return nil, err
		}
	}

	return cache, nil
}
<commit_msg>Pass a logger to the MetricMapper<commit_after>// Copyright 2015 The Prometheus 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 main

import (
	"bufio"
	"bytes"
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/promlog/flag"
	"github.com/prometheus/common/version"
	"github.com/prometheus/exporter-toolkit/web"
	"github.com/prometheus/graphite_exporter/collector"
	"github.com/prometheus/statsd_exporter/pkg/mapper"
	"github.com/prometheus/statsd_exporter/pkg/mappercache/lru"
	"github.com/prometheus/statsd_exporter/pkg/mappercache/randomreplacement"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	listenAddress   = kingpin.Flag("web.listen-address", "Address on which to expose metrics.").Default(":9108").String()
	metricsPath     = kingpin.Flag("web.telemetry-path", "Path under which to expose Prometheus metrics.").Default("/metrics").String()
	graphiteAddress = kingpin.Flag("graphite.listen-address", "TCP and UDP address on which to accept samples.").Default(":9109").String()
	mappingConfig   = kingpin.Flag("graphite.mapping-config", "Metric mapping configuration file name.").Default("").String()
	sampleExpiry    = kingpin.Flag("graphite.sample-expiry", "How long a sample is valid for.").Default("5m").Duration()
	strictMatch     = kingpin.Flag("graphite.mapping-strict-match", "Only store metrics that match the mapping configuration.").Bool()
	cacheSize       = kingpin.Flag("graphite.cache-size", "Maximum size of your metric mapping cache. Relies on least recently used replacement policy if max size is reached.").Default("1000").Int()
	cacheType       = kingpin.Flag("graphite.cache-type", "Metric mapping cache type. Valid options are \"lru\" and \"random\"").Default("lru").Enum("lru", "random")
	dumpFSMPath     = kingpin.Flag("debug.dump-fsm", "The path to dump internal FSM generated for glob matching as Dot file.").Default("").String()
	checkConfig     = kingpin.Flag("check-config", "Check configuration and exit.").Default("false").Bool()
	configFile      = kingpin.Flag("web.config", "[EXPERIMENTAL] Path to config yaml file that can enable TLS or authentication.").Default("").String()
)

func init() {
	prometheus.MustRegister(version.NewCollector("graphite_exporter"))
}

func dumpFSM(mapper *mapper.MetricMapper, dumpFilename string, logger log.Logger) error {
	if mapper.FSM == nil {
		return fmt.Errorf("no FSM available to be dumped, possibly because the mapping contains regex patterns")
	}
	f, err := os.Create(dumpFilename)
	if err != nil {
		return err
	}
	level.Info(logger).Log("msg", "Start dumping FSM", "to", dumpFilename)
	w := bufio.NewWriter(f)
	mapper.FSM.DumpFSM(w)
	w.Flush()
	f.Close()
	level.Info(logger).Log("msg", "Finish dumping FSM")
	return nil
}

func main() {
	promlogConfig := &promlog.Config{}
	flag.AddFlags(kingpin.CommandLine, promlogConfig)
	kingpin.Version(version.Print("graphite_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()
	logger := promlog.New(promlogConfig)

	level.Info(logger).Log("msg", "Starting graphite_exporter", "version_info", version.Info())
	level.Info(logger).Log("build_context", version.BuildContext())

	http.Handle(*metricsPath, promhttp.Handler())
	c := collector.NewGraphiteCollector(logger, *strictMatch, *sampleExpiry)
	prometheus.MustRegister(c)

	metricMapper := &mapper.MetricMapper{Logger: logger}
	if *mappingConfig != "" {
		err := metricMapper.InitFromFile(*mappingConfig)
		if err != nil {
			level.Error(logger).Log("msg", "Error loading metric mapping config", "err", err)
			os.Exit(1)
		}
	}

	cache, err := getCache(*cacheSize, *cacheType, prometheus.DefaultRegisterer)
	if err != nil {
		level.Error(logger).Log("msg", "error initializing mapper cache", "err", err)
		os.Exit(1)
	}
	metricMapper.UseCache(cache)

	if *checkConfig {
		level.Info(logger).Log("msg", "Configuration check successful, exiting")
		return
	}

	if *dumpFSMPath != "" {
		err := dumpFSM(metricMapper, *dumpFSMPath, logger)
		if err != nil {
			level.Error(logger).Log("msg", "Error dumping FSM", "err", err)
			os.Exit(1)
		}
	}

	c.SetMapper(metricMapper)

	tcpSock, err := net.Listen("tcp", *graphiteAddress)
	if err != nil {
		level.Error(logger).Log("msg", "Error binding to TCP socket", "err", err)
		os.Exit(1)
	}
	go func() {
		for {
			conn, err := tcpSock.Accept()
			if err != nil {
				level.Error(logger).Log("msg", "Error accepting TCP connection", "err", err)
				continue
			}
			go func() {
				defer conn.Close()
				c.ProcessReader(conn)
			}()
		}
	}()

	udpAddress, err := net.ResolveUDPAddr("udp", *graphiteAddress)
	if err != nil {
		level.Error(logger).Log("msg", "Error resolving UDP address", "err", err)
		os.Exit(1)
	}
	udpSock, err := net.ListenUDP("udp", udpAddress)
	if err != nil {
		level.Error(logger).Log("msg", "Error listening to UDP address", "err", err)
		os.Exit(1)
	}
	go func() {
		defer udpSock.Close()
		for {
			buf := make([]byte, 65536)
			chars, srcAddress, err := udpSock.ReadFromUDP(buf)
			if err != nil {
				level.Error(logger).Log("msg", "Error reading UDP packet", "from", srcAddress, "err", err)
				continue
			}
			go c.ProcessReader(bytes.NewReader(buf[0:chars]))
		}
	}()

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}
		w.Write([]byte(`
      <title>Graphite Exporter</title>

Graphite Exporter

Accepting plaintext Graphite samples over TCP and UDP on ` + *graphiteAddress + `

<a href="` + *metricsPath + `">Metrics

      `))
	})

	level.Info(logger).Log("msg", "Listening on "+*listenAddress)
	server := &http.Server{Addr: *listenAddress}
	if err := web.ListenAndServe(server, *configFile, logger); err != nil {
		level.Error(logger).Log("err", err)
		os.Exit(1)
	}
}

// TODO(mr): this is copied verbatim from statsd_exporter/main.go. It should be a
// convenience function in mappercache, but that caused an import cycle.
func getCache(cacheSize int, cacheType string, registerer prometheus.Registerer) (mapper.MetricMapperCache, error) {
	var cache mapper.MetricMapperCache
	var err error
	if cacheSize == 0 {
		return nil, nil
	} else {
		switch cacheType {
		case "lru":
			cache, err = lru.NewMetricMapperLRUCache(registerer, cacheSize)
		case "random":
			cache, err = randomreplacement.NewMetricMapperRRCache(registerer, cacheSize)
		default:
			err = fmt.Errorf("unsupported cache type %q", cacheType)
		}

		if err != nil {
			return nil, err
		}
	}

	return cache, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 cmd

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"os/exec"
	"os/user"
	"regexp"
	"strconv"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/addons"
	"k8s.io/minikube/pkg/minikube/assets"
	"k8s.io/minikube/pkg/minikube/style"

	"k8s.io/minikube/pkg/minikube/browser"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/mustload"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/proxy"
	"k8s.io/minikube/pkg/minikube/reason"
	"k8s.io/minikube/pkg/minikube/service"
	"k8s.io/minikube/pkg/util/retry"
)

var (
	dashboardURLMode     bool
	dashboardExposedPort int
	// Matches: 127.0.0.1:8001
	// TODO(tstromberg): Get kubectl to implement a stable supported output format.
	hostPortRe = regexp.MustCompile(`127.0.0.1:\d{4,}`)
)

// dashboardCmd represents the dashboard command
var dashboardCmd = &cobra.Command{
	Use:   "dashboard",
	Short: "Access the Kubernetes dashboard running within the minikube cluster",
	Long:  `Access the Kubernetes dashboard running within the minikube cluster`,
	Run: func(cmd *cobra.Command, args []string) {
		cname := ClusterFlagValue()
		co := mustload.Healthy(cname)

		for _, n := range co.Config.Nodes {
			if err := proxy.ExcludeIP(n.IP); err != nil {
				klog.Errorf("Error excluding IP from proxy: %s", err)
			}
		}

		if dashboardExposedPort < 0 || dashboardExposedPort > 65535 {
			exit.Message(reason.HostKubectlProxy, "Invalid port")
		}

		kubectlVersion := co.Config.KubernetesConfig.KubernetesVersion
		var err error

		// Check dashboard status before enabling it
		addon := assets.Addons["dashboard"]
		enabled := addon.IsEnabled(co.Config)

		if !enabled {
			// Send status messages to stderr for folks re-using this output.
			out.ErrT(style.Enabling, "Enabling dashboard ...")
			// Enable the dashboard add-on
			err = addons.SetAndSave(cname, "dashboard", "true")
			if err != nil {
				exit.Error(reason.InternalAddonEnable, "Unable to enable dashboard", err)
			}
		}

		ns := "kubernetes-dashboard"
		svc := "kubernetes-dashboard"
		out.ErrT(style.Verifying, "Verifying dashboard health ...")
		checkSVC := func() error { return service.CheckService(cname, ns, svc) }
		// for slow machines or parallels in CI to avoid #7503
		if err = retry.Expo(checkSVC, 100*time.Microsecond, time.Minute*10); err != nil {
			exit.Message(reason.SvcCheckTimeout, "dashboard service is not running: {{.error}}", out.V{"error": err})
		}

		out.ErrT(style.Launch, "Launching proxy ...")
		p, hostPort, err := kubectlProxy(kubectlVersion, co.Config.BinaryMirror, cname, dashboardExposedPort)
		if err != nil {
			exit.Error(reason.HostKubectlProxy, "kubectl proxy", err)
		}
		url := dashboardURL(hostPort, ns, svc)

		out.ErrT(style.Verifying, "Verifying proxy health ...")
		chkURL := func() error { return checkURL(url) }
		if err = retry.Expo(chkURL, 100*time.Microsecond, 10*time.Minute); err != nil {
			exit.Message(reason.SvcURLTimeout, "{{.url}} is not accessible: {{.error}}", out.V{"url": url, "error": err})
		}

		// check if current user is root
		user, err := user.Current()
		if err != nil {
			exit.Error(reason.HostCurrentUser, "Unable to get current user", err)
		}
		if dashboardURLMode || user.Uid == "0" {
			out.Ln(url)
		} else {
			out.Styled(style.Celebrate, "Opening {{.url}} in your default browser...", out.V{"url": url})
			if err = browser.OpenURL(url); err != nil {
				exit.Message(reason.HostBrowser, "failed to open browser: {{.error}}", out.V{"error": err})
			}
		}

		klog.Infof("Success! I will now quietly sit around until kubectl proxy exits!")
		if err = p.Wait(); err != nil {
			klog.Errorf("Wait: %v", err)
		}
	},
}

// kubectlProxy runs "kubectl proxy", returning host:port
func kubectlProxy(kubectlVersion string, binaryURL string, contextName string, port int) (*exec.Cmd, string, error) {
	// port=0 picks a random system port

	kubectlArgs := []string{"--context", contextName, "proxy", "--port", strconv.Itoa(port)}

	var cmd *exec.Cmd
	if kubectl, err := exec.LookPath("kubectl"); err == nil {
		cmd = exec.Command(kubectl, kubectlArgs...)
	} else if cmd, err = KubectlCommand(kubectlVersion, binaryURL, kubectlArgs...); err != nil {
		return nil, "", err
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return nil, "", errors.Wrap(err, "cmd stdout")
	}

	klog.Infof("Executing: %s %s", cmd.Path, cmd.Args)
	if err := cmd.Start(); err != nil {
		return nil, "", errors.Wrap(err, "proxy start")
	}

	klog.Infof("Waiting for kubectl to output host:port ...")
	reader := bufio.NewReader(stdoutPipe)

	var out []byte
	for {
		r, timedOut, err := readByteWithTimeout(reader, 5*time.Second)
		if err != nil {
			return cmd, "", fmt.Errorf("readByteWithTimeout: %v", err)
		}
		if r == byte('\n') {
			break
		}
		if timedOut {
			klog.Infof("timed out waiting for input: possibly due to an old kubectl version.")
			break
		}
		out = append(out, r)
	}
	klog.Infof("proxy stdout: %s", string(out))
	return cmd, hostPortRe.FindString(string(out)), nil
}

// readByteWithTimeout returns a byte from a reader or an indicator that a timeout has occurred.
func readByteWithTimeout(r io.ByteReader, timeout time.Duration) (byte, bool, error) {
	bc := make(chan byte, 1)
	ec := make(chan error, 1)
	go func() {
		b, err := r.ReadByte()
		if err != nil {
			ec <- err
		} else {
			bc <- b
		}
		close(bc)
		close(ec)
	}()
	select {
	case b := <-bc:
		return b, false, nil
	case err := <-ec:
		return byte(' '), false, err
	case <-time.After(timeout):
		return byte(' '), true, nil
	}
}

// dashboardURL generates a URL for accessing the dashboard service
func dashboardURL(proxy string, ns string, svc string) string {
	// Reference: https://github.com/kubernetes/dashboard/wiki/Accessing-Dashboard---1.7.X-and-above
	return fmt.Sprintf("http://%s/api/v1/namespaces/%s/services/http:%s:/proxy/", proxy, ns, svc)
}

// checkURL checks if a URL returns 200 HTTP OK
func checkURL(url string) error {
	resp, err := http.Get(url)
	klog.Infof("%s response: %v %+v", url, err, resp)
	if err != nil {
		return errors.Wrap(err, "checkURL")
	}
	if resp.StatusCode != http.StatusOK {
		return &retry.RetriableError{
			Err: fmt.Errorf("unexpected response code: %d", resp.StatusCode),
		}
	}
	return nil
}

func init() {
	dashboardCmd.Flags().BoolVar(&dashboardURLMode, "url", false, "Display dashboard URL instead of opening a browser")
	dashboardCmd.Flags().IntVar(&dashboardExposedPort, "port", 0, "Exposed port of the proxyfied dashboard. Set to 0 to pick a random port.")
}
<commit_msg>Accept space before port in DashboardProxy Regex<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 cmd

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"os/exec"
	"os/user"
	"regexp"
	"strconv"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/addons"
	"k8s.io/minikube/pkg/minikube/assets"
	"k8s.io/minikube/pkg/minikube/style"

	"k8s.io/minikube/pkg/minikube/browser"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/mustload"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/proxy"
	"k8s.io/minikube/pkg/minikube/reason"
	"k8s.io/minikube/pkg/minikube/service"
	"k8s.io/minikube/pkg/util/retry"
)

var (
	dashboardURLMode     bool
	dashboardExposedPort int
	// Matches: "127.0.0.1:8001" or "127.0.0.1 40012" etc.
	// TODO(tstromberg): Get kubectl to implement a stable supported output format.
	hostPortRe = regexp.MustCompile(`127.0.0.1(:| )\d{4,}`)
)

// dashboardCmd represents the dashboard command
var dashboardCmd = &cobra.Command{
	Use:   "dashboard",
	Short: "Access the Kubernetes dashboard running within the minikube cluster",
	Long:  `Access the Kubernetes dashboard running within the minikube cluster`,
	Run: func(cmd *cobra.Command, args []string) {
		cname := ClusterFlagValue()
		co := mustload.Healthy(cname)

		for _, n := range co.Config.Nodes {
			if err := proxy.ExcludeIP(n.IP); err != nil {
				klog.Errorf("Error excluding IP from proxy: %s", err)
			}
		}

		if dashboardExposedPort < 0 || dashboardExposedPort > 65535 {
			exit.Message(reason.HostKubectlProxy, "Invalid port")
		}

		kubectlVersion := co.Config.KubernetesConfig.KubernetesVersion
		var err error

		// Check dashboard status before enabling it
		addon := assets.Addons["dashboard"]
		enabled := addon.IsEnabled(co.Config)

		if !enabled {
			// Send status messages to stderr for folks re-using this output.
			out.ErrT(style.Enabling, "Enabling dashboard ...")
			// Enable the dashboard add-on
			err = addons.SetAndSave(cname, "dashboard", "true")
			if err != nil {
				exit.Error(reason.InternalAddonEnable, "Unable to enable dashboard", err)
			}
		}

		ns := "kubernetes-dashboard"
		svc := "kubernetes-dashboard"
		out.ErrT(style.Verifying, "Verifying dashboard health ...")
		checkSVC := func() error { return service.CheckService(cname, ns, svc) }
		// for slow machines or parallels in CI to avoid #7503
		if err = retry.Expo(checkSVC, 100*time.Microsecond, time.Minute*10); err != nil {
			exit.Message(reason.SvcCheckTimeout, "dashboard service is not running: {{.error}}", out.V{"error": err})
		}

		out.ErrT(style.Launch, "Launching proxy ...")
		p, hostPort, err := kubectlProxy(kubectlVersion, co.Config.BinaryMirror, cname, dashboardExposedPort)
		if err != nil {
			exit.Error(reason.HostKubectlProxy, "kubectl proxy", err)
		}
		url := dashboardURL(hostPort, ns, svc)

		out.ErrT(style.Verifying, "Verifying proxy health ...")
		chkURL := func() error { return checkURL(url) }
		if err = retry.Expo(chkURL, 100*time.Microsecond, 10*time.Minute); err != nil {
			exit.Message(reason.SvcURLTimeout, "{{.url}} is not accessible: {{.error}}", out.V{"url": url, "error": err})
		}

		// check if current user is root
		user, err := user.Current()
		if err != nil {
			exit.Error(reason.HostCurrentUser, "Unable to get current user", err)
		}
		if dashboardURLMode || user.Uid == "0" {
			out.Ln(url)
		} else {
			out.Styled(style.Celebrate, "Opening {{.url}} in your default browser...", out.V{"url": url})
			if err = browser.OpenURL(url); err != nil {
				exit.Message(reason.HostBrowser, "failed to open browser: {{.error}}", out.V{"error": err})
			}
		}

		klog.Infof("Success! I will now quietly sit around until kubectl proxy exits!")
		if err = p.Wait(); err != nil {
			klog.Errorf("Wait: %v", err)
		}
	},
}

// kubectlProxy runs "kubectl proxy", returning host:port
func kubectlProxy(kubectlVersion string, binaryURL string, contextName string, port int) (*exec.Cmd, string, error) {
	// port=0 picks a random system port

	kubectlArgs := []string{"--context", contextName, "proxy", "--port", strconv.Itoa(port)}

	var cmd *exec.Cmd
	if kubectl, err := exec.LookPath("kubectl"); err == nil {
		cmd = exec.Command(kubectl, kubectlArgs...)
	} else if cmd, err = KubectlCommand(kubectlVersion, binaryURL, kubectlArgs...); err != nil {
		return nil, "", err
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return nil, "", errors.Wrap(err, "cmd stdout")
	}

	klog.Infof("Executing: %s %s", cmd.Path, cmd.Args)
	if err := cmd.Start(); err != nil {
		return nil, "", errors.Wrap(err, "proxy start")
	}

	klog.Infof("Waiting for kubectl to output host:port ...")
	reader := bufio.NewReader(stdoutPipe)

	var out []byte
	for {
		r, timedOut, err := readByteWithTimeout(reader, 5*time.Second)
		if err != nil {
			return cmd, "", fmt.Errorf("readByteWithTimeout: %v", err)
		}
		if r == byte('\n') {
			break
		}
		if timedOut {
			klog.Infof("timed out waiting for input: possibly due to an old kubectl version.")
			break
		}
		out = append(out, r)
	}
	klog.Infof("proxy stdout: %s", string(out))
	return cmd, hostPortRe.FindString(string(out)), nil
}

// readByteWithTimeout returns a byte from a reader or an indicator that a timeout has occurred.
func readByteWithTimeout(r io.ByteReader, timeout time.Duration) (byte, bool, error) {
	bc := make(chan byte, 1)
	ec := make(chan error, 1)
	go func() {
		b, err := r.ReadByte()
		if err != nil {
			ec <- err
		} else {
			bc <- b
		}
		close(bc)
		close(ec)
	}()
	select {
	case b := <-bc:
		return b, false, nil
	case err := <-ec:
		return byte(' '), false, err
	case <-time.After(timeout):
		return byte(' '), true, nil
	}
}

// dashboardURL generates a URL for accessing the dashboard service
func dashboardURL(proxy string, ns string, svc string) string {
	// Reference: https://github.com/kubernetes/dashboard/wiki/Accessing-Dashboard---1.7.X-and-above
	return fmt.Sprintf("http://%s/api/v1/namespaces/%s/services/http:%s:/proxy/", proxy, ns, svc)
}

// checkURL checks if a URL returns 200 HTTP OK
func checkURL(url string) error {
	resp, err := http.Get(url)
	klog.Infof("%s response: %v %+v", url, err, resp)
	if err != nil {
		return errors.Wrap(err, "checkURL")
	}
	if resp.StatusCode != http.StatusOK {
		return &retry.RetriableError{
			Err: fmt.Errorf("unexpected response code: %d", resp.StatusCode),
		}
	}
	return nil
}

func init() {
	dashboardCmd.Flags().BoolVar(&dashboardURLMode, "url", false, "Display dashboard URL instead of opening a browser")
	dashboardCmd.Flags().IntVar(&dashboardExposedPort, "port", 0, "Exposed port of the proxyfied dashboard. Set to 0 to pick a random port.")
}
<|endoftext|>


<commit_before>package commands

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"path"
	"strconv"

	"github.com/go-openapi/loads"
	"github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/spec"
	"github.com/go-openapi/swag"
	"github.com/gorilla/handlers"
	"github.com/toqueteos/webbrowser"
)

// ServeCmd to serve a swagger spec with docs ui
type ServeCmd struct {
	BasePath string `long:"base-path" description:"the base path to serve the spec and UI at"`
	Flavor   string `short:"F" long:"flavor" description:"the flavor of docs, can be swagger or redoc" default:"redoc" choice:"redoc" choice:"swagger"`
	DocURL   string `long:"doc-url" description:"override the url which takes a url query param to render the doc ui"`
	NoOpen   bool   `long:"no-open" description:"when present won't open the the browser to show the url"`
	NoUI     bool   `long:"no-ui" description:"when present, only the swagger spec will be served"`
	Flatten  bool   `long:"flatten" description:"when present, flatten the swagger spec before serving it"`
	Port     int    `long:"port" short:"p" description:"the port to serve this site" env:"PORT"`
	Host     string `long:"host" description:"the interface to serve this site, defaults to 0.0.0.0" default:"0.0.0.0" env:"HOST"`
}

// Execute the serve command
func (s *ServeCmd) Execute(args []string) error {
	if len(args) == 0 {
		return errors.New("specify the spec to serve as argument to the serve command")
	}

	specDoc, err := loads.Spec(args[0])
	if err != nil {
		return err
	}

	if s.Flatten {
		specDoc, err = specDoc.Expanded(&spec.ExpandOptions{
			SkipSchemas:         false,
			ContinueOnError:     true,
			AbsoluteCircularRef: true,
		})

		if err != nil {
			return err
		}
	}

	b, err := json.MarshalIndent(specDoc.Spec(), "", "  ")
	if err != nil {
		return err
	}

	basePath := s.BasePath
	if basePath == "" {
		basePath = "/"
	}

	listener, err := net.Listen("tcp4", net.JoinHostPort(s.Host, strconv.Itoa(s.Port)))
	if err != nil {
		return err
	}
	sh, sp, err := swag.SplitHostPort(listener.Addr().String())
	if err != nil {
		return err
	}
	if sh == "0.0.0.0" {
		sh = "localhost"
	}

	visit := s.DocURL
	handler := http.NotFoundHandler()
	if !s.NoUI {
		if s.Flavor == "redoc" {
			handler = middleware.Redoc(middleware.RedocOpts{
				BasePath: basePath,
				SpecURL:  path.Join(basePath, "swagger.json"),
				Path:     "docs",
			}, handler)
			visit = fmt.Sprintf("http://%s:%d%s", sh, sp, path.Join(basePath, "docs"))
		} else if visit != "" || s.Flavor == "swagger" {
			handler = middleware.SwaggerUI(middleware.SwaggerUIOpts{
				BasePath: basePath,
				SpecURL:  path.Join(basePath, "swagger.json"),
				Path:     "docs",
			}, handler)
			visit = fmt.Sprintf("http://%s:%d%s", sh, sp, path.Join(basePath, "docs"))
		}
	}

	handler = handlers.CORS()(middleware.Spec(basePath, b, handler))
	errFuture := make(chan error)
	go func() {
		docServer := new(http.Server)
		docServer.SetKeepAlivesEnabled(true)
		docServer.Handler = handler

		errFuture <- docServer.Serve(listener)
	}()

	if !s.NoOpen && !s.NoUI {
		err := webbrowser.Open(visit)
		if err != nil {
			return err
		}
	}
	log.Println("serving docs at", visit)
	return <-errFuture
}
<commit_msg>add option to add your own path<commit_after>package commands

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"path"
	"strconv"

	"github.com/go-openapi/loads"
	"github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/spec"
	"github.com/go-openapi/swag"
	"github.com/gorilla/handlers"
	"github.com/toqueteos/webbrowser"
)

// ServeCmd to serve a swagger spec with docs ui
type ServeCmd struct {
	BasePath string `long:"base-path" description:"the base path to serve the spec and UI at"`
	Flavor   string `short:"F" long:"flavor" description:"the flavor of docs, can be swagger or redoc" default:"redoc" choice:"redoc" choice:"swagger"`
	DocURL   string `long:"doc-url" description:"override the url which takes a url query param to render the doc ui"`
	NoOpen   bool   `long:"no-open" description:"when present won't open the the browser to show the url"`
	NoUI     bool   `long:"no-ui" description:"when present, only the swagger spec will be served"`
	Flatten  bool   `long:"flatten" description:"when present, flatten the swagger spec before serving it"`
	Port     int    `long:"port" short:"p" description:"the port to serve this site" env:"PORT"`
	Host     string `long:"host" description:"the interface to serve this site, defaults to 0.0.0.0" default:"0.0.0.0" env:"HOST"`
	Path     string `long:"path" description:"the uri path at which the docs will be served" default:"docs"`
}

// Execute the serve command
func (s *ServeCmd) Execute(args []string) error {
	if len(args) == 0 {
		return errors.New("specify the spec to serve as argument to the serve command")
	}

	specDoc, err := loads.Spec(args[0])
	if err != nil {
		return err
	}

	if s.Flatten {
		specDoc, err = specDoc.Expanded(&spec.ExpandOptions{
			SkipSchemas:         false,
			ContinueOnError:     true,
			AbsoluteCircularRef: true,
		})

		if err != nil {
			return err
		}
	}

	b, err := json.MarshalIndent(specDoc.Spec(), "", "  ")
	if err != nil {
		return err
	}

	basePath := s.BasePath
	if basePath == "" {
		basePath = "/"
	}

	listener, err := net.Listen("tcp4", net.JoinHostPort(s.Host, strconv.Itoa(s.Port)))
	if err != nil {
		return err
	}
	sh, sp, err := swag.SplitHostPort(listener.Addr().String())
	if err != nil {
		return err
	}
	if sh == "0.0.0.0" {
		sh = "localhost"
	}

	visit := s.DocURL
	handler := http.NotFoundHandler()
	if !s.NoUI {
		if s.Flavor == "redoc" {
			handler = middleware.Redoc(middleware.RedocOpts{
				BasePath: basePath,
				SpecURL:  path.Join(basePath, "swagger.json"),
				Path:     s.Path,
			}, handler)
			visit = fmt.Sprintf("http://%s:%d%s", sh, sp, path.Join(basePath, "docs"))
		} else if visit != "" || s.Flavor == "swagger" {
			handler = middleware.SwaggerUI(middleware.SwaggerUIOpts{
				BasePath: basePath,
				SpecURL:  path.Join(basePath, "swagger.json"),
				Path:     s.Path,
			}, handler)
			visit = fmt.Sprintf("http://%s:%d%s", sh, sp, path.Join(basePath, s.Path))
		}
	}

	handler = handlers.CORS()(middleware.Spec(basePath, b, handler))
	errFuture := make(chan error)
	go func() {
		docServer := new(http.Server)
		docServer.SetKeepAlivesEnabled(true)
		docServer.Handler = handler

		errFuture <- docServer.Serve(listener)
	}()

	if !s.NoOpen && !s.NoUI {
		err := webbrowser.Open(visit)
		if err != nil {
			return err
		}
	}
	log.Println("serving docs at", visit)
	return <-errFuture
}
<|endoftext|>


<commit_before>package buf

import (
	"io"

	"v2ray.com/core/common/errors"
)

// BufferedWriter is an io.Writer with internal buffer. It writes to underlying writer when buffer is full or on demand.
// This type is not thread safe.
type BufferedWriter struct {
	writer   io.Writer
	buffer   *Buffer
	buffered bool
}

// NewBufferedWriter creates a new BufferedWriter.
func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
	return &BufferedWriter{
		writer:   rawWriter,
		buffer:   NewLocal(1024),
		buffered: true,
	}
}

// ReadFrom implements io.ReaderFrom.ReadFrom().
func (v *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
	totalBytes := int64(0)
	for {
		oriSize := v.buffer.Len()
		err := v.buffer.AppendSupplier(ReadFrom(reader))
		totalBytes += int64(v.buffer.Len() - oriSize)
		if err != nil {
			if errors.Cause(err) == io.EOF {
				return totalBytes, nil
			}
			return totalBytes, err
		}
		if err := v.Flush(); err != nil {
			return totalBytes, err
		}
	}
}

func (v *BufferedWriter) Write(b []byte) (int, error) {
	if !v.buffered || v.buffer == nil {
		return v.writer.Write(b)
	}
	nBytes, err := v.buffer.Write(b)
	if err != nil {
		return 0, err
	}
	if v.buffer.IsFull() {
		err := v.Flush()
		if err != nil {
			return 0, err
		}
		if nBytes < len(b) {
			if _, err := v.writer.Write(b[nBytes:]); err != nil {
				return nBytes, err
			}
		}
	}
	return len(b), nil
}

// Flush writes all buffered content into underlying writer, if any.
func (v *BufferedWriter) Flush() error {
	defer v.buffer.Clear()
	for !v.buffer.IsEmpty() {
		nBytes, err := v.writer.Write(v.buffer.Bytes())
		if err != nil {
			return err
		}
		v.buffer.SliceFrom(nBytes)
	}
	return nil
}

// IsBuffered returns true if this BufferedWriter holds a buffer.
func (v *BufferedWriter) IsBuffered() bool {
	return v.buffered
}

// SetBuffered controls whether the BufferedWriter holds a buffer for writing. If not buffered, any write() calls into underlying writer directly.
func (v *BufferedWriter) SetBuffered(cached bool) error {
	v.buffered = cached
	if !cached && !v.buffer.IsEmpty() {
		return v.Flush()
	}
	return nil
}
<commit_msg>remove unused code<commit_after>package buf

import "io"

// BufferedWriter is an io.Writer with internal buffer. It writes to underlying writer when buffer is full or on demand.
// This type is not thread safe.
type BufferedWriter struct {
	writer   io.Writer
	buffer   *Buffer
	buffered bool
}

// NewBufferedWriter creates a new BufferedWriter.
func NewBufferedWriter(rawWriter io.Writer) *BufferedWriter {
	return &BufferedWriter{
		writer:   rawWriter,
		buffer:   NewLocal(1024),
		buffered: true,
	}
}

func (v *BufferedWriter) Write(b []byte) (int, error) {
	if !v.buffered || v.buffer == nil {
		return v.writer.Write(b)
	}
	nBytes, err := v.buffer.Write(b)
	if err != nil {
		return 0, err
	}
	if v.buffer.IsFull() {
		err := v.Flush()
		if err != nil {
			return 0, err
		}
		if nBytes < len(b) {
			if _, err := v.writer.Write(b[nBytes:]); err != nil {
				return nBytes, err
			}
		}
	}
	return len(b), nil
}

// Flush writes all buffered content into underlying writer, if any.
func (v *BufferedWriter) Flush() error {
	defer v.buffer.Clear()
	for !v.buffer.IsEmpty() {
		nBytes, err := v.writer.Write(v.buffer.Bytes())
		if err != nil {
			return err
		}
		v.buffer.SliceFrom(nBytes)
	}
	return nil
}

// IsBuffered returns true if this BufferedWriter holds a buffer.
func (v *BufferedWriter) IsBuffered() bool {
	return v.buffered
}

// SetBuffered controls whether the BufferedWriter holds a buffer for writing. If not buffered, any write() calls into underlying writer directly.
func (v *BufferedWriter) SetBuffered(cached bool) error {
	v.buffered = cached
	if !cached && !v.buffer.IsEmpty() {
		return v.Flush()
	}
	return nil
}
<|endoftext|>