<commit_before>// returns data points from Win32_PerfRawData_PerfProc_Process
// https://msdn.microsoft.com/en-us/library/aa394323(v=vs.85).aspx - Win32_PerfRawData_PerfProc_Process class

// +build windows

package collector

import (
	"strconv"
	"strings"

	"github.com/StackExchange/wmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"gopkg.in/alecthomas/kingpin.v2"
)

func init() {
	Factories["process"] = NewProcessCollector
}

var (
	processWhereClause = kingpin.Flag(
		"collector.process.processes-where",
		"WQL 'where' clause to use in WMI metrics query. Limits the response to the processes you specify and reduces the size of the response.",
	).Default("").String()
)

// A ProcessCollector is a Prometheus collector for WMI Win32_PerfRawData_PerfProc_Process metrics
type ProcessCollector struct {
	StartTime         *prometheus.Desc
	CPUTimeTotal      *prometheus.Desc
	HandleCount       *prometheus.Desc
	IOBytesTotal      *prometheus.Desc
	IOOperationsTotal *prometheus.Desc
	PageFaultsTotal   *prometheus.Desc
	PageFileBytes     *prometheus.Desc
	PoolBytes         *prometheus.Desc
	PriorityBase      *prometheus.Desc
	PrivateBytes      *prometheus.Desc
	ThreadCount       *prometheus.Desc
	VirtualBytes      *prometheus.Desc
	WorkingSet        *prometheus.Desc

	queryWhereClause string
}

// NewProcessCollector ...
func NewProcessCollector() (Collector, error) {
	const subsystem = "process"

	if *processWhereClause == "" {
		log.Warn("No where-clause specified for process collector. This will generate a very large number of metrics!")
	}

	return &ProcessCollector{
		StartTime: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "start_time"),
			"Time of process start.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		CPUTimeTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "cpu_time_total"),
			"Returns elapsed time that all of the threads of this process used the processor to execute instructions by mode (privileged, user). An instruction is the basic unit of execution in a computer, a thread is the object that executes instructions, and a process is the object created when a program is run. Code executed to handle some hardware interrupts and trap conditions is included in this count.",
			[]string{"process", "process_id", "creating_process_id", "mode"},
			nil,
		),
		HandleCount: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "handle_count"),
			"Total number of handles the process has open. This number is the sum of the handles currently open by each thread in the process.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		IOBytesTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "io_bytes_total"),
			"Bytes issued to I/O operations in different modes (read, write, other). This property counts all I/O activity generated by the process to include file, network, and device I/Os. Read and write mode includes data operations; other mode includes those that do not involve data, such as control operations. ",
			[]string{"process", "process_id", "creating_process_id", "mode"},
			nil,
		),
		IOOperationsTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "io_operations_total"),
			"I/O operations issued in different modes (read, write, other). This property counts all I/O activity generated by the process to include file, network, and device I/Os. Read and write mode includes data operations; other mode includes those that do not involve data, such as control operations. ",
			[]string{"process", "process_id", "creating_process_id", "mode"},
			nil,
		),
		PageFaultsTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "page_faults_total"),
			"Page faults by the threads executing in this process. A page fault occurs when a thread refers to a virtual memory page that is not in its working set in main memory. This can cause the page not to be fetched from disk if it is on the standby list and hence already in main memory, or if it is in use by another process with which the page is shared.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		PageFileBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "page_file_bytes"),
			"Current number of bytes this process has used in the paging file(s). Paging files are used to store pages of memory used by the process that are not contained in other files. Paging files are shared by all processes, and lack of space in paging files can prevent other processes from allocating memory.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		PoolBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "pool_bytes"),
			"Pool Bytes is the last observed number of bytes in the paged or nonpaged pool. The nonpaged pool is an area of system memory (physical memory used by the operating system) for objects that cannot be written to disk, but must remain in physical memory as long as they are allocated. The paged pool is an area of system memory (physical memory used by the operating system) for objects that can be written to disk when they are not being used. Nonpaged pool bytes is calculated differently than paged pool bytes, so it might not equal the total of paged pool bytes.",
			[]string{"process", "process_id", "creating_process_id", "pool"},
			nil,
		),
		PriorityBase: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "priority_base"),
			"Current base priority of this process. Threads within a process can raise and lower their own base priority relative to the process base priority of the process.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		PrivateBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "private_bytes"),
			"Current number of bytes this process has allocated that cannot be shared with other processes.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		ThreadCount: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "thread_count"),
			"Number of threads currently active in this process. An instruction is the basic unit of execution in a processor, and a thread is the object that executes instructions. Every running process has at least one thread.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		VirtualBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "virtual_bytes"),
			"Current size, in bytes, of the virtual address space that the process is using. Use of virtual address space does not necessarily imply corresponding use of either disk or main memory pages. Virtual space is finite and, by using too much, the process can limit its ability to load libraries.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		WorkingSet: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "working_set"),
			"Maximum number of bytes in the working set of this process at any point in time. The working set is the set of memory pages touched recently by the threads in the process. If free memory in the computer is above a threshold, pages are left in the working set of a process even if they are not in use. When free memory falls below a threshold, pages are trimmed from working sets. If they are needed, they are then soft-faulted back into the working set before they leave main memory.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		queryWhereClause: *processWhereClause,
	}, nil
}

// Collect sends the metric values for each metric
// to the provided prometheus Metric channel.
func (c *ProcessCollector) Collect(ch chan<- prometheus.Metric) error {
	if desc, err := c.collect(ch); err != nil {
		log.Error("failed collecting process metrics:", desc, err)
		return err
	}
	return nil
}

type Win32_PerfRawData_PerfProc_Process struct {
	Name                    string
	CreatingProcessID       uint32
	ElapsedTime             uint64
	Frequency_Object        uint64
	HandleCount             uint32
	IDProcess               uint32
	IODataBytesPersec       uint64
	IODataOperationsPersec  uint64
	IOOtherBytesPersec      uint64
	IOOtherOperationsPersec uint64
	IOReadBytesPersec       uint64
	IOReadOperationsPersec  uint64
	IOWriteBytesPersec      uint64
	IOWriteOperationsPersec uint64
	PageFaultsPersec        uint32
	PageFileBytes           uint64
	PageFileBytesPeak       uint64
	PercentPrivilegedTime   uint64
	PercentProcessorTime    uint64
	PercentUserTime         uint64
	PoolNonpagedBytes       uint32
	PoolPagedBytes          uint32
	PriorityBase            uint32
	PrivateBytes            uint64
	ThreadCount             uint32
	Timestamp_Object        uint64
	VirtualBytes            uint64
	VirtualBytesPeak        uint64
	WorkingSet              uint64
	WorkingSetPeak          uint64
	WorkingSetPrivate       uint64
}

type WorkerProcess struct {
	AppPoolName string
	ProcessId   uint32
}

func (c *ProcessCollector) collect(ch chan<- prometheus.Metric) (*prometheus.Desc, error) {
	var dst []Win32_PerfRawData_PerfProc_Process
	q := queryAllWhere(&dst, c.queryWhereClause)
	if err := wmi.Query(q, &dst); err != nil {
		return nil, err
	}

	var dst_wp []WorkerProcess
	q_wp := queryAll(&dst_wp)
	wmi.QueryNamespace(q_wp, &dst_wp, "root\\WebAdministration")

	for _, process := range dst {

		if process.Name == "_Total" {
			continue
		}
		// Duplicate processes are suffixed # and an index number. Remove those.
		processName := strings.Split(process.Name, "#")[0]
		pid := strconv.FormatUint(uint64(process.IDProcess), 10)
		cpid := strconv.FormatUint(uint64(process.CreatingProcessID), 10)

		for _, wp := range dst_wp {
			if wp.ProcessId == process.IDProcess {
				processName = strings.Join([]string{processName, wp.AppPoolName}, "_")
				break
			}
		}

		ch <- prometheus.MustNewConstMetric(
			c.StartTime,
			prometheus.GaugeValue,
			// convert from Windows timestamp (1 jan 1601) to unix timestamp (1 jan 1970)
			float64(process.ElapsedTime-116444736000000000)/float64(process.Frequency_Object),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.HandleCount,
			prometheus.GaugeValue,
			float64(process.HandleCount),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.CPUTimeTotal,
			prometheus.CounterValue,
			float64(process.PercentPrivilegedTime)*ticksToSecondsScaleFactor,
			processName,
			pid,
			cpid,
			"privileged",
		)

		ch <- prometheus.MustNewConstMetric(
			c.CPUTimeTotal,
			prometheus.CounterValue,
			float64(process.PercentUserTime)*ticksToSecondsScaleFactor,
			processName,
			pid,
			cpid,
			"user",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOBytesTotal,
			prometheus.CounterValue,
			float64(process.IOOtherBytesPersec),
			processName,
			pid,
			cpid,
			"other",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOOperationsTotal,
			prometheus.CounterValue,
			float64(process.IOOtherOperationsPersec),
			processName,
			pid,
			cpid,
			"other",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOBytesTotal,
			prometheus.CounterValue,
			float64(process.IOReadBytesPersec),
			processName,
			pid,
			cpid,
			"read",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOOperationsTotal,
			prometheus.CounterValue,
			float64(process.IOReadOperationsPersec),
			processName,
			pid,
			cpid,
			"read",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOBytesTotal,
			prometheus.CounterValue,
			float64(process.IOWriteBytesPersec),
			processName,
			pid,
			cpid,
			"write",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOOperationsTotal,
			prometheus.CounterValue,
			float64(process.IOWriteOperationsPersec),
			processName,
			pid,
			cpid,
			"write",
		)

		ch <- prometheus.MustNewConstMetric(
			c.PageFaultsTotal,
			prometheus.CounterValue,
			float64(process.PageFaultsPersec),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PageFileBytes,
			prometheus.GaugeValue,
			float64(process.PageFileBytes),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PoolBytes,
			prometheus.GaugeValue,
			float64(process.PoolNonpagedBytes),
			processName,
			pid,
			cpid,
			"nonpaged",
		)

		ch <- prometheus.MustNewConstMetric(
			c.PoolBytes,
			prometheus.GaugeValue,
			float64(process.PoolPagedBytes),
			processName,
			pid,
			cpid,
			"paged",
		)

		ch <- prometheus.MustNewConstMetric(
			c.PriorityBase,
			prometheus.GaugeValue,
			float64(process.PriorityBase),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PrivateBytes,
			prometheus.GaugeValue,
			float64(process.PrivateBytes),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.ThreadCount,
			prometheus.GaugeValue,
			float64(process.ThreadCount),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.VirtualBytes,
			prometheus.GaugeValue,
			float64(process.VirtualBytes),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.WorkingSet,
			prometheus.GaugeValue,
			float64(process.WorkingSet),
			processName,
			pid,
			cpid,
		)
	}

	return nil, nil
}
<commit_msg>There was added the checking and logging of query worker processes<commit_after>// returns data points from Win32_PerfRawData_PerfProc_Process
// https://msdn.microsoft.com/en-us/library/aa394323(v=vs.85).aspx - Win32_PerfRawData_PerfProc_Process class

// +build windows

package collector

import (
	"strconv"
	"strings"

	"github.com/StackExchange/wmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"gopkg.in/alecthomas/kingpin.v2"
)

func init() {
	Factories["process"] = NewProcessCollector
}

var (
	processWhereClause = kingpin.Flag(
		"collector.process.processes-where",
		"WQL 'where' clause to use in WMI metrics query. Limits the response to the processes you specify and reduces the size of the response.",
	).Default("").String()
)

// A ProcessCollector is a Prometheus collector for WMI Win32_PerfRawData_PerfProc_Process metrics
type ProcessCollector struct {
	StartTime         *prometheus.Desc
	CPUTimeTotal      *prometheus.Desc
	HandleCount       *prometheus.Desc
	IOBytesTotal      *prometheus.Desc
	IOOperationsTotal *prometheus.Desc
	PageFaultsTotal   *prometheus.Desc
	PageFileBytes     *prometheus.Desc
	PoolBytes         *prometheus.Desc
	PriorityBase      *prometheus.Desc
	PrivateBytes      *prometheus.Desc
	ThreadCount       *prometheus.Desc
	VirtualBytes      *prometheus.Desc
	WorkingSet        *prometheus.Desc

	queryWhereClause string
}

// NewProcessCollector ...
func NewProcessCollector() (Collector, error) {
	const subsystem = "process"

	if *processWhereClause == "" {
		log.Warn("No where-clause specified for process collector. This will generate a very large number of metrics!")
	}

	return &ProcessCollector{
		StartTime: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "start_time"),
			"Time of process start.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		CPUTimeTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "cpu_time_total"),
			"Returns elapsed time that all of the threads of this process used the processor to execute instructions by mode (privileged, user). An instruction is the basic unit of execution in a computer, a thread is the object that executes instructions, and a process is the object created when a program is run. Code executed to handle some hardware interrupts and trap conditions is included in this count.",
			[]string{"process", "process_id", "creating_process_id", "mode"},
			nil,
		),
		HandleCount: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "handle_count"),
			"Total number of handles the process has open. This number is the sum of the handles currently open by each thread in the process.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		IOBytesTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "io_bytes_total"),
			"Bytes issued to I/O operations in different modes (read, write, other). This property counts all I/O activity generated by the process to include file, network, and device I/Os. Read and write mode includes data operations; other mode includes those that do not involve data, such as control operations. ",
			[]string{"process", "process_id", "creating_process_id", "mode"},
			nil,
		),
		IOOperationsTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "io_operations_total"),
			"I/O operations issued in different modes (read, write, other). This property counts all I/O activity generated by the process to include file, network, and device I/Os. Read and write mode includes data operations; other mode includes those that do not involve data, such as control operations. ",
			[]string{"process", "process_id", "creating_process_id", "mode"},
			nil,
		),
		PageFaultsTotal: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "page_faults_total"),
			"Page faults by the threads executing in this process. A page fault occurs when a thread refers to a virtual memory page that is not in its working set in main memory. This can cause the page not to be fetched from disk if it is on the standby list and hence already in main memory, or if it is in use by another process with which the page is shared.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		PageFileBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "page_file_bytes"),
			"Current number of bytes this process has used in the paging file(s). Paging files are used to store pages of memory used by the process that are not contained in other files. Paging files are shared by all processes, and lack of space in paging files can prevent other processes from allocating memory.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		PoolBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "pool_bytes"),
			"Pool Bytes is the last observed number of bytes in the paged or nonpaged pool. The nonpaged pool is an area of system memory (physical memory used by the operating system) for objects that cannot be written to disk, but must remain in physical memory as long as they are allocated. The paged pool is an area of system memory (physical memory used by the operating system) for objects that can be written to disk when they are not being used. Nonpaged pool bytes is calculated differently than paged pool bytes, so it might not equal the total of paged pool bytes.",
			[]string{"process", "process_id", "creating_process_id", "pool"},
			nil,
		),
		PriorityBase: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "priority_base"),
			"Current base priority of this process. Threads within a process can raise and lower their own base priority relative to the process base priority of the process.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		PrivateBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "private_bytes"),
			"Current number of bytes this process has allocated that cannot be shared with other processes.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		ThreadCount: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "thread_count"),
			"Number of threads currently active in this process. An instruction is the basic unit of execution in a processor, and a thread is the object that executes instructions. Every running process has at least one thread.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		VirtualBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "virtual_bytes"),
			"Current size, in bytes, of the virtual address space that the process is using. Use of virtual address space does not necessarily imply corresponding use of either disk or main memory pages. Virtual space is finite and, by using too much, the process can limit its ability to load libraries.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		WorkingSet: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "working_set"),
			"Maximum number of bytes in the working set of this process at any point in time. The working set is the set of memory pages touched recently by the threads in the process. If free memory in the computer is above a threshold, pages are left in the working set of a process even if they are not in use. When free memory falls below a threshold, pages are trimmed from working sets. If they are needed, they are then soft-faulted back into the working set before they leave main memory.",
			[]string{"process", "process_id", "creating_process_id"},
			nil,
		),
		queryWhereClause: *processWhereClause,
	}, nil
}

// Collect sends the metric values for each metric
// to the provided prometheus Metric channel.
func (c *ProcessCollector) Collect(ch chan<- prometheus.Metric) error {
	if desc, err := c.collect(ch); err != nil {
		log.Error("failed collecting process metrics:", desc, err)
		return err
	}
	return nil
}

type Win32_PerfRawData_PerfProc_Process struct {
	Name                    string
	CreatingProcessID       uint32
	ElapsedTime             uint64
	Frequency_Object        uint64
	HandleCount             uint32
	IDProcess               uint32
	IODataBytesPersec       uint64
	IODataOperationsPersec  uint64
	IOOtherBytesPersec      uint64
	IOOtherOperationsPersec uint64
	IOReadBytesPersec       uint64
	IOReadOperationsPersec  uint64
	IOWriteBytesPersec      uint64
	IOWriteOperationsPersec uint64
	PageFaultsPersec        uint32
	PageFileBytes           uint64
	PageFileBytesPeak       uint64
	PercentPrivilegedTime   uint64
	PercentProcessorTime    uint64
	PercentUserTime         uint64
	PoolNonpagedBytes       uint32
	PoolPagedBytes          uint32
	PriorityBase            uint32
	PrivateBytes            uint64
	ThreadCount             uint32
	Timestamp_Object        uint64
	VirtualBytes            uint64
	VirtualBytesPeak        uint64
	WorkingSet              uint64
	WorkingSetPeak          uint64
	WorkingSetPrivate       uint64
}

type WorkerProcess struct {
	AppPoolName string
	ProcessId   uint32
}

func (c *ProcessCollector) collect(ch chan<- prometheus.Metric) (*prometheus.Desc, error) {
	var dst []Win32_PerfRawData_PerfProc_Process
	q := queryAllWhere(&dst, c.queryWhereClause)
	if err := wmi.Query(q, &dst); err != nil {
		return nil, err
	}

	var dst_wp []WorkerProcess
	q_wp := queryAll(&dst_wp)
	if err := wmi.QueryNamespace(q_wp, &dst_wp, "root\\WebAdministration"); err != nil {
		log.Warnf("failed querying worker processes: %s", err)
	}

	for _, process := range dst {

		if process.Name == "_Total" {
			continue
		}
		// Duplicate processes are suffixed # and an index number. Remove those.
		processName := strings.Split(process.Name, "#")[0]
		pid := strconv.FormatUint(uint64(process.IDProcess), 10)
		cpid := strconv.FormatUint(uint64(process.CreatingProcessID), 10)

		for _, wp := range dst_wp {
			if wp.ProcessId == process.IDProcess {
				processName = strings.Join([]string{processName, wp.AppPoolName}, "_")
				break
			}
		}

		ch <- prometheus.MustNewConstMetric(
			c.StartTime,
			prometheus.GaugeValue,
			// convert from Windows timestamp (1 jan 1601) to unix timestamp (1 jan 1970)
			float64(process.ElapsedTime-116444736000000000)/float64(process.Frequency_Object),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.HandleCount,
			prometheus.GaugeValue,
			float64(process.HandleCount),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.CPUTimeTotal,
			prometheus.CounterValue,
			float64(process.PercentPrivilegedTime)*ticksToSecondsScaleFactor,
			processName,
			pid,
			cpid,
			"privileged",
		)

		ch <- prometheus.MustNewConstMetric(
			c.CPUTimeTotal,
			prometheus.CounterValue,
			float64(process.PercentUserTime)*ticksToSecondsScaleFactor,
			processName,
			pid,
			cpid,
			"user",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOBytesTotal,
			prometheus.CounterValue,
			float64(process.IOOtherBytesPersec),
			processName,
			pid,
			cpid,
			"other",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOOperationsTotal,
			prometheus.CounterValue,
			float64(process.IOOtherOperationsPersec),
			processName,
			pid,
			cpid,
			"other",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOBytesTotal,
			prometheus.CounterValue,
			float64(process.IOReadBytesPersec),
			processName,
			pid,
			cpid,
			"read",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOOperationsTotal,
			prometheus.CounterValue,
			float64(process.IOReadOperationsPersec),
			processName,
			pid,
			cpid,
			"read",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOBytesTotal,
			prometheus.CounterValue,
			float64(process.IOWriteBytesPersec),
			processName,
			pid,
			cpid,
			"write",
		)

		ch <- prometheus.MustNewConstMetric(
			c.IOOperationsTotal,
			prometheus.CounterValue,
			float64(process.IOWriteOperationsPersec),
			processName,
			pid,
			cpid,
			"write",
		)

		ch <- prometheus.MustNewConstMetric(
			c.PageFaultsTotal,
			prometheus.CounterValue,
			float64(process.PageFaultsPersec),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PageFileBytes,
			prometheus.GaugeValue,
			float64(process.PageFileBytes),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PoolBytes,
			prometheus.GaugeValue,
			float64(process.PoolNonpagedBytes),
			processName,
			pid,
			cpid,
			"nonpaged",
		)

		ch <- prometheus.MustNewConstMetric(
			c.PoolBytes,
			prometheus.GaugeValue,
			float64(process.PoolPagedBytes),
			processName,
			pid,
			cpid,
			"paged",
		)

		ch <- prometheus.MustNewConstMetric(
			c.PriorityBase,
			prometheus.GaugeValue,
			float64(process.PriorityBase),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.PrivateBytes,
			prometheus.GaugeValue,
			float64(process.PrivateBytes),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.ThreadCount,
			prometheus.GaugeValue,
			float64(process.ThreadCount),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.VirtualBytes,
			prometheus.GaugeValue,
			float64(process.VirtualBytes),
			processName,
			pid,
			cpid,
		)

		ch <- prometheus.MustNewConstMetric(
			c.WorkingSet,
			prometheus.GaugeValue,
			float64(process.WorkingSet),
			processName,
			pid,
			cpid,
		)
	}

	return nil, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2015 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2015 Andrii Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy

import (
    "math/rand"
    "time"
    "sync"

    "sippy/log"
    "sippy/utils"
)

type Timeout struct {
    callback        func()
    timeout         time.Duration
    logger          sippy_log.ErrorLogger
    shutdown_chan   chan struct{}
    shutdown        bool
    spread          time.Duration
    nticks          int
    lock            sync.Mutex
    cb_lock         sync.Locker
    started         bool
}

func StartTimeout(callback func(), cb_lock sync.Locker, _timeout time.Duration, nticks int, logger sippy_log.ErrorLogger) *Timeout {
    self := NewInactiveTimeout(callback, cb_lock, _timeout, nticks, logger)
    self.Start()
    return self
}

func NewInactiveTimeout(callback func(), cb_lock sync.Locker, _timeout time.Duration, nticks int, logger sippy_log.ErrorLogger) *Timeout {
    self := &Timeout{
        callback        : callback,
        timeout         : _timeout,
        nticks          : nticks,
        logger          : logger,
        shutdown_chan   : make(chan struct{}),
        shutdown        : false,
        spread          : 0,
        started         : false,
        cb_lock         : cb_lock,
    }
    return self
}

func (self *Timeout) Start() {
    self.lock.Lock()
    if ! self.started && self.callback != nil {
        self.started = true
        go self.run()
    }
    self.lock.Unlock()
}

func (self *Timeout) SpreadRuns(spread time.Duration) {
    self.spread = spread
}

func (self *Timeout) Cancel() {
    self.shutdown = true
    close(self.shutdown_chan)
}

func (self *Timeout) run() {
    for !self.shutdown {
        self._run()
    }
    self.callback = nil
    self.cb_lock = nil
}

func (self *Timeout) _run() {
    for !self.shutdown {
        if self.nticks == 0 {
            self.shutdown = true
            break
        }
        if self.nticks > 0 {
            self.nticks--
        }
        t := self.timeout
        if self.spread > 0 {
            t += time.Duration(float64(self.spread) * (rand.Float64() - 0.5))
        }
        select {
        case <-self.shutdown_chan:
            self.shutdown = true
        case <-time.After(t):
            sippy_utils.SafeCall(self.callback, self.cb_lock, self.logger)
        }
    }
}
<commit_msg>Properly shutdown the timer.<commit_after>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2015 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2015 Andrii Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy

import (
    "math/rand"
    "time"
    "sync"

    "sippy/log"
    "sippy/utils"
)

type Timeout struct {
    callback        func()
    timeout         time.Duration
    logger          sippy_log.ErrorLogger
    shutdown_chan   chan struct{}
    shutdown        bool
    spread          time.Duration
    nticks          int
    lock            sync.Mutex
    cb_lock         sync.Locker
    started         bool
}

func StartTimeout(callback func(), cb_lock sync.Locker, _timeout time.Duration, nticks int, logger sippy_log.ErrorLogger) *Timeout {
    self := NewInactiveTimeout(callback, cb_lock, _timeout, nticks, logger)
    self.Start()
    return self
}

func NewInactiveTimeout(callback func(), cb_lock sync.Locker, _timeout time.Duration, nticks int, logger sippy_log.ErrorLogger) *Timeout {
    self := &Timeout{
        callback        : callback,
        timeout         : _timeout,
        nticks          : nticks,
        logger          : logger,
        shutdown_chan   : make(chan struct{}),
        shutdown        : false,
        spread          : 0,
        started         : false,
        cb_lock         : cb_lock,
    }
    return self
}

func (self *Timeout) Start() {
    self.lock.Lock()
    if ! self.started && self.callback != nil {
        self.started = true
        go self.run()
    }
    self.lock.Unlock()
}

func (self *Timeout) SpreadRuns(spread time.Duration) {
    self.spread = spread
}

func (self *Timeout) Cancel() {
    self.shutdown = true
    close(self.shutdown_chan)
}

func (self *Timeout) run() {
    for !self.shutdown {
        self._run()
    }
    self.callback = nil
    self.cb_lock = nil
}

func (self *Timeout) _run() {
    for !self.shutdown {
        if self.nticks == 0 {
            self.shutdown = true
            break
        }
        if self.nticks > 0 {
            self.nticks--
        }
        t := self.timeout
        if self.spread > 0 {
            t += time.Duration(float64(self.spread) * (rand.Float64() - 0.5))
        }
        select {
        case <-self.shutdown_chan:
            self.shutdown = true
        case <-time.After(t):
            if ! self.shutdown {
                sippy_utils.SafeCall(self.callback, self.cb_lock, self.logger)
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package time

import (
	"sync"
)

// A Ticker holds a synchronous channel that delivers `ticks' of a clock
// at intervals.
type Ticker struct {
	C        <-chan int64 // The channel on which the ticks are delivered.
	c        chan<- int64 // The same channel, but the end we use.
	ns       int64
	shutdown chan bool // Buffered channel used to signal shutdown.
	nextTick int64
	next     *Ticker
}

// Stop turns off a ticker.  After Stop, no more ticks will be sent.
func (t *Ticker) Stop() {
	// Make it non-blocking so multiple Stops don't block.
	_ = t.shutdown <- true
}

// Tick is a convenience wrapper for NewTicker providing access to the ticking
// channel only.  Useful for clients that have no need to shut down the ticker.
func Tick(ns int64) <-chan int64 {
	if ns <= 0 {
		return nil
	}
	return NewTicker(ns).C
}

type alarmer struct {
	wakeUp   chan bool // wakeup signals sent/received here
	wakeMeAt chan int64
	wakeTime int64
}

// Set alarm to go off at time ns, if not already set earlier.
func (a *alarmer) set(ns int64) {
	switch {
	case a.wakeTime > ns:
		// Next tick we expect is too late; shut down the late runner
		// and (after fallthrough) start a new wakeLoop.
		a.wakeMeAt <- -1
		fallthrough
	case a.wakeMeAt == nil:
		// There's no wakeLoop, start one.
		a.wakeMeAt = make(chan int64, 10)
		go wakeLoop(a.wakeMeAt, a.wakeUp)
		fallthrough
	case a.wakeTime == 0:
		// Nobody else is waiting; it's just us.
		a.wakeTime = ns
		a.wakeMeAt <- ns
	default:
		// There's already someone scheduled.
	}
}

// Channel to notify tickerLoop of new Tickers being created.
var newTicker chan *Ticker

func startTickerLoop() {
	newTicker = make(chan *Ticker)
	go tickerLoop()
}

// wakeLoop delivers ticks at scheduled times, sleeping until the right moment.
// If another, earlier Ticker is created while it sleeps, tickerLoop() will start a new
// wakeLoop but they will share the wakeUp channel and signal that this one
// is done by giving it a negative time request.
func wakeLoop(wakeMeAt chan int64, wakeUp chan bool) {
	for {
		wakeAt := <-wakeMeAt
		if wakeAt < 0 { // tickerLoop has started another wakeLoop
			return
		}
		now := Nanoseconds()
		if wakeAt > now {
			Sleep(wakeAt - now)
			now = Nanoseconds()
		}
		wakeUp <- true
	}
}

// A single tickerLoop serves all ticks to Tickers.  It waits for two events:
// either the creation of a new Ticker or a tick from the alarm,
// signalling a time to wake up one or more Tickers.
func tickerLoop() {
	// Represents the next alarm to be delivered.
	var alarm alarmer
	// All wakeLoops deliver wakeups to this channel.
	alarm.wakeUp = make(chan bool, 10)
	var now, prevTime, wakeTime int64
	var tickers *Ticker
	for {
		select {
		case t := <-newTicker:
			// Add Ticker to list
			t.next = tickers
			tickers = t
			// Arrange for a new alarm if this one precedes the existing one.
			alarm.set(t.nextTick)
		case <-alarm.wakeUp:
			now = Nanoseconds()
			// Ignore an old time due to a dying wakeLoop
			if now < prevTime {
				continue
			}
			wakeTime = now + 1e15 // very long in the future
			var prev *Ticker = nil
			// Scan list of tickers, delivering updates to those
			// that need it and determining the next wake time.
			// TODO(r): list should be sorted in time order.
			for t := tickers; t != nil; t = t.next {
				if _, ok := <-t.shutdown; ok {
					// Ticker is done; remove it from list.
					if prev == nil {
						tickers = t.next
					} else {
						prev.next = t.next
					}
					continue
				}
				if t.nextTick <= now {
					if len(t.c) == 0 {
						// Only send if there's room.  We must not block.
						// The channel is allocated with a one-element
						// buffer, which is sufficient: if he hasn't picked
						// up the last tick, no point in sending more.
						t.c <- now
					}
					t.nextTick += t.ns
					if t.nextTick <= now {
						// Still behind; advance in one big step.
						t.nextTick += (now - t.nextTick + t.ns) / t.ns * t.ns
					}
				}
				if t.nextTick < wakeTime {
					wakeTime = t.nextTick
				}
				prev = t
			}
			if tickers != nil {
				// Please send wakeup at earliest required time.
				// If there are no tickers, don't bother.
				alarm.wakeMeAt <- wakeTime
			} else {
				alarm.wakeTime = 0
			}
		}
		prevTime = now
	}
}

var onceStartTickerLoop sync.Once

// NewTicker returns a new Ticker containing a channel that will
// send the time, in nanoseconds, every ns nanoseconds.  It adjusts the
// intervals to make up for pauses in delivery of the ticks.
func NewTicker(ns int64) *Ticker {
	if ns <= 0 {
		return nil
	}
	c := make(chan int64, 1) //  See comment on send in tickerLoop
	t := &Ticker{
		C:        c,
		c:        c,
		ns:       ns,
		shutdown: make(chan bool, 1),
		nextTick: Nanoseconds() + ns,
	}
	onceStartTickerLoop.Do(startTickerLoop)
	// must be run in background so global Tickers can be created
	go func() { newTicker <- t }()
	return t
}
<commit_msg>time.NewTicker: panic for intervals <= 0. Not absolutely certain it's right to do this, but since there's no error value coming back, it seems reasonable.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package time

import (
	"os"
	"sync"
)

// A Ticker holds a synchronous channel that delivers `ticks' of a clock
// at intervals.
type Ticker struct {
	C        <-chan int64 // The channel on which the ticks are delivered.
	c        chan<- int64 // The same channel, but the end we use.
	ns       int64
	shutdown chan bool // Buffered channel used to signal shutdown.
	nextTick int64
	next     *Ticker
}

// Stop turns off a ticker.  After Stop, no more ticks will be sent.
func (t *Ticker) Stop() {
	// Make it non-blocking so multiple Stops don't block.
	_ = t.shutdown <- true
}

// Tick is a convenience wrapper for NewTicker providing access to the ticking
// channel only.  Useful for clients that have no need to shut down the ticker.
func Tick(ns int64) <-chan int64 {
	if ns <= 0 {
		return nil
	}
	return NewTicker(ns).C
}

type alarmer struct {
	wakeUp   chan bool // wakeup signals sent/received here
	wakeMeAt chan int64
	wakeTime int64
}

// Set alarm to go off at time ns, if not already set earlier.
func (a *alarmer) set(ns int64) {
	switch {
	case a.wakeTime > ns:
		// Next tick we expect is too late; shut down the late runner
		// and (after fallthrough) start a new wakeLoop.
		a.wakeMeAt <- -1
		fallthrough
	case a.wakeMeAt == nil:
		// There's no wakeLoop, start one.
		a.wakeMeAt = make(chan int64, 10)
		go wakeLoop(a.wakeMeAt, a.wakeUp)
		fallthrough
	case a.wakeTime == 0:
		// Nobody else is waiting; it's just us.
		a.wakeTime = ns
		a.wakeMeAt <- ns
	default:
		// There's already someone scheduled.
	}
}

// Channel to notify tickerLoop of new Tickers being created.
var newTicker chan *Ticker

func startTickerLoop() {
	newTicker = make(chan *Ticker)
	go tickerLoop()
}

// wakeLoop delivers ticks at scheduled times, sleeping until the right moment.
// If another, earlier Ticker is created while it sleeps, tickerLoop() will start a new
// wakeLoop but they will share the wakeUp channel and signal that this one
// is done by giving it a negative time request.
func wakeLoop(wakeMeAt chan int64, wakeUp chan bool) {
	for {
		wakeAt := <-wakeMeAt
		if wakeAt < 0 { // tickerLoop has started another wakeLoop
			return
		}
		now := Nanoseconds()
		if wakeAt > now {
			Sleep(wakeAt - now)
			now = Nanoseconds()
		}
		wakeUp <- true
	}
}

// A single tickerLoop serves all ticks to Tickers.  It waits for two events:
// either the creation of a new Ticker or a tick from the alarm,
// signalling a time to wake up one or more Tickers.
func tickerLoop() {
	// Represents the next alarm to be delivered.
	var alarm alarmer
	// All wakeLoops deliver wakeups to this channel.
	alarm.wakeUp = make(chan bool, 10)
	var now, prevTime, wakeTime int64
	var tickers *Ticker
	for {
		select {
		case t := <-newTicker:
			// Add Ticker to list
			t.next = tickers
			tickers = t
			// Arrange for a new alarm if this one precedes the existing one.
			alarm.set(t.nextTick)
		case <-alarm.wakeUp:
			now = Nanoseconds()
			// Ignore an old time due to a dying wakeLoop
			if now < prevTime {
				continue
			}
			wakeTime = now + 1e15 // very long in the future
			var prev *Ticker = nil
			// Scan list of tickers, delivering updates to those
			// that need it and determining the next wake time.
			// TODO(r): list should be sorted in time order.
			for t := tickers; t != nil; t = t.next {
				if _, ok := <-t.shutdown; ok {
					// Ticker is done; remove it from list.
					if prev == nil {
						tickers = t.next
					} else {
						prev.next = t.next
					}
					continue
				}
				if t.nextTick <= now {
					if len(t.c) == 0 {
						// Only send if there's room.  We must not block.
						// The channel is allocated with a one-element
						// buffer, which is sufficient: if he hasn't picked
						// up the last tick, no point in sending more.
						t.c <- now
					}
					t.nextTick += t.ns
					if t.nextTick <= now {
						// Still behind; advance in one big step.
						t.nextTick += (now - t.nextTick + t.ns) / t.ns * t.ns
					}
				}
				if t.nextTick < wakeTime {
					wakeTime = t.nextTick
				}
				prev = t
			}
			if tickers != nil {
				// Please send wakeup at earliest required time.
				// If there are no tickers, don't bother.
				alarm.wakeMeAt <- wakeTime
			} else {
				alarm.wakeTime = 0
			}
		}
		prevTime = now
	}
}

var onceStartTickerLoop sync.Once

// NewTicker returns a new Ticker containing a channel that will
// send the time, in nanoseconds, every ns nanoseconds.  It adjusts the
// intervals to make up for pauses in delivery of the ticks. The value of
// ns must be greater than zero; if not, NewTicker will panic.
func NewTicker(ns int64) *Ticker {
	if ns <= 0 {
		panic(os.ErrorString("non-positive interval for NewTicker"))
	}
	c := make(chan int64, 1) //  See comment on send in tickerLoop
	t := &Ticker{
		C:        c,
		c:        c,
		ns:       ns,
		shutdown: make(chan bool, 1),
		nextTick: Nanoseconds() + ns,
	}
	onceStartTickerLoop.Do(startTickerLoop)
	// must be run in background so global Tickers can be created
	go func() { newTicker <- t }()
	return t
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"time"

	"github.com/bwmarrin/discordgo"
)

// Random - Picks a random integer.
// min: Minumum amount in the integer
// max: Maximum amount in the integer
func Random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}

// GetPageContents - Get page content based on URL.
// url: Valid url of image.
func GetPageContents(url string) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	return body, err
}

// IsManager Check to see if a user has ManageServer Permissions.
// s: The Current Session between the bot and discord
// m: The Message Object sent back from Discord.
func IsManager(s *discordgo.Session, GuildID string, AuthorID string) bool {
	// Check the user permissions of the guild.
	perms, err := s.State.UserChannelPermissions(AuthorID, GuildID)
	if err == nil {
		if (perms & discordgo.PermissionManageServer) > 0 {
			return true
		}
	} else {
		return false
	}
	return false
}

// Save - Saves Database to config.json
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// m: The Message Object sent back from Discord.
func (bot *Object) Save() {
	for {
		<-time.After(5 * time.Minute)
		js, err := json.MarshalIndent(bot, "", "  ")
		if err == nil {
			ioutil.WriteFile("config.json", js, 0777)
		}
	}
}

/* Not functional yet
func (bot *Object) PruneMessages() {
	for {
		<-time.After(1 * time.Hour)
		for _, m := range bot.System.Messages {
			if m.Timestamp < time.Now()-(3600*24*7) {

			}
		}
	}
}
*/

// GetRoleID - Grabs the Discord Role ID
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// role: The Discord role
func (bot *Object) GetRoleID(s *discordgo.Session, role string) string {
	var id string
	r, err := s.State.Guild(bot.Guild)
	if err == nil {
		for _, v := range r.Roles {
			if v.Name == role {
				id = v.ID
			}
		}
	}
	return id
}

// MemberHasRole - Checks to see if the user has a role.
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// role: The Discord role
func (bot *Object) MemberHasRole(s *discordgo.Session, AuthorID string, role string) bool {
	therole := bot.GetRoleID(s, role)
	z, err := s.State.Member(bot.Guild, AuthorID)
	if err != nil {
		z, err = s.GuildMember(bot.Guild, AuthorID)
		if err != nil {
			fmt.Println("Error ->", err)
			return false
		}
	}
	for r := range z.Roles {
		if therole == z.Roles[r] {
			return true
		}
	}
	return false
}

// Register - Register new object.
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// m: Message Object sent back from Discord.
func (bot *Object) Register(s *discordgo.Session, m *discordgo.MessageCreate) {
	// check and make sure the server already exists in my collection.
	if bot.System != nil {
		return
	}
	c, err := s.State.Channel(m.ChannelID)
	if err != nil {
		fmt.Println(err)
		return
	}

	bot.Guild = c.GuildID
	chn := &Channels{
		Autorole: "",
		Greeting: "",
		ByeMsg:   "",
	}

	// Create a new Info pointer.
	info := &System{
		Prefix:   ".",
		Autorole: "",
		Greeting: "",
		ByeMsg:   "",
		Channels: chn,
	}
	// Add our Info object to the bot map.
	bot.System = info
}

// Task - Store new messages to object.
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// role: The Discord role
func (bot *Object) Task(s *discordgo.Session, m *discordgo.MessageCreate) {
	// Don't track the bots messages.
	if m.Author.ID == s.State.User.ID {
		return
	}

	if bot.System == nil {
		return
	}
	// Create a new pointer.
	msg := &Messages{
		ID:        m.ID,
		Author:    m.Author.ID,
		Channel:   m.ChannelID,
		Timestamp: time.Now().Unix(),
	}
	// Add this Message to our Info object.
	bot.System.Messages = append(bot.System.Messages, msg)
}

// AddStatus - Adds a status string to the main Object.
// message: the status message(s)
func (bot *Object) AddStatus(message string) error {
	if bot.System == nil {
		return nil
	}
	for _, s := range bot.System.Status {
		if s == message {
			return errors.New("status exists already")
		}
	}
	bot.System.Status = append(bot.System.Status, message)
	return nil
}

// RemoveStatus - Removes a status string from the main Object.
// message: the status message
func (bot *Object) RemoveStatus(message string) error {
	if bot.System == nil {
		return errors.New("object doesn't exist")
	}
	var ti int
	for i, k := range bot.System.Status {
		if k == message {
			ti = i
		}
	}
	if ti == 0 {
		return errors.New("status doesn't exist in my collection")
	}
	bot.System.Status = append(bot.System.Status[:ti], bot.System.Status[ti+1:]...)
	return nil
}

// Intro - Displays introduction and information on startup.
func (bot *Object) Intro(s *discordgo.Session) {
	var ars map[string]string
	fmt.Println(`
  ______             _   _           _
 |  ____|           | | | |         | |
 | |__ _ __ ___  ___| |_| |__  _   _| |_ ___
 |  __| '__/ _ \/ __| __| '_ \| | | | __/ _ \
 | |  | | | (_) \__ \ |_| |_) | |_| | ||  __/
 |_|  |_|  \___/|___/\__|_.__/ \__, |\__\___|
                                __/ |
                               |___/         `)
	<-time.After(2 * time.Second)
	// Collect some information and display it!
	guild, err := s.State.Guild(bot.Guild)
	if err != nil {
		fmt.Println(err)
	} else {
		// Channel count for server
		fmt.Print(len(guild.Channels), " Channel(s), ")
		// Member count for server
		fmt.Print(len(guild.Members), " Member(s), ")
		// Role count for server
		fmt.Print(len(guild.Roles), " Role(s), ")
	}

	// Collect the A.R.S Count.
	io, err := ioutil.ReadFile("autoresponse.json")
	if err != nil {
		fmt.Println(err)
	} else {
		json.Unmarshal(io, &ars)
		// ARS Count for the bot
		fmt.Println(len(ars), " A.R.S Rule(s).")
	}
	fmt.Println(`---------------------------------------------`)
}
<commit_msg>fixing misspelling of Minimum<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"time"

	"github.com/bwmarrin/discordgo"
)

// Random - Picks a random integer.
// min: Minimum amount in the integer
// max: Maximum amount in the integer
func Random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}

// GetPageContents - Get page content based on URL.
// url: Valid url of image.
func GetPageContents(url string) ([]byte, error) {
	req, err := http.NewRequest("GET", url, nil)
	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	return body, err
}

// IsManager Check to see if a user has ManageServer Permissions.
// s: The Current Session between the bot and discord
// m: The Message Object sent back from Discord.
func IsManager(s *discordgo.Session, GuildID string, AuthorID string) bool {
	// Check the user permissions of the guild.
	perms, err := s.State.UserChannelPermissions(AuthorID, GuildID)
	if err == nil {
		if (perms & discordgo.PermissionManageServer) > 0 {
			return true
		}
	} else {
		return false
	}
	return false
}

// Save - Saves Database to config.json
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// m: The Message Object sent back from Discord.
func (bot *Object) Save() {
	for {
		<-time.After(5 * time.Minute)
		js, err := json.MarshalIndent(bot, "", "  ")
		if err == nil {
			ioutil.WriteFile("config.json", js, 0777)
		}
	}
}

/* Not functional yet
func (bot *Object) PruneMessages() {
	for {
		<-time.After(1 * time.Hour)
		for _, m := range bot.System.Messages {
			if m.Timestamp < time.Now()-(3600*24*7) {

			}
		}
	}
}
*/

// GetRoleID - Grabs the Discord Role ID
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// role: The Discord role
func (bot *Object) GetRoleID(s *discordgo.Session, role string) string {
	var id string
	r, err := s.State.Guild(bot.Guild)
	if err == nil {
		for _, v := range r.Roles {
			if v.Name == role {
				id = v.ID
			}
		}
	}
	return id
}

// MemberHasRole - Checks to see if the user has a role.
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// role: The Discord role
func (bot *Object) MemberHasRole(s *discordgo.Session, AuthorID string, role string) bool {
	therole := bot.GetRoleID(s, role)
	z, err := s.State.Member(bot.Guild, AuthorID)
	if err != nil {
		z, err = s.GuildMember(bot.Guild, AuthorID)
		if err != nil {
			fmt.Println("Error ->", err)
			return false
		}
	}
	for r := range z.Roles {
		if therole == z.Roles[r] {
			return true
		}
	}
	return false
}

// Register - Register new object.
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// m: Message Object sent back from Discord.
func (bot *Object) Register(s *discordgo.Session, m *discordgo.MessageCreate) {
	// check and make sure the server already exists in my collection.
	if bot.System != nil {
		return
	}
	c, err := s.State.Channel(m.ChannelID)
	if err != nil {
		fmt.Println(err)
		return
	}

	bot.Guild = c.GuildID
	chn := &Channels{
		Autorole: "",
		Greeting: "",
		ByeMsg:   "",
	}

	// Create a new Info pointer.
	info := &System{
		Prefix:   ".",
		Autorole: "",
		Greeting: "",
		ByeMsg:   "",
		Channels: chn,
	}
	// Add our Info object to the bot map.
	bot.System = info
}

// Task - Store new messages to object.
// bot: Main Object with all your settings.
// s: The Current Session between the bot and discord
// role: The Discord role
func (bot *Object) Task(s *discordgo.Session, m *discordgo.MessageCreate) {
	// Don't track the bots messages.
	if m.Author.ID == s.State.User.ID {
		return
	}

	if bot.System == nil {
		return
	}
	// Create a new pointer.
	msg := &Messages{
		ID:        m.ID,
		Author:    m.Author.ID,
		Channel:   m.ChannelID,
		Timestamp: time.Now().Unix(),
	}
	// Add this Message to our Info object.
	bot.System.Messages = append(bot.System.Messages, msg)
}

// AddStatus - Adds a status string to the main Object.
// message: the status message(s)
func (bot *Object) AddStatus(message string) error {
	if bot.System == nil {
		return nil
	}
	for _, s := range bot.System.Status {
		if s == message {
			return errors.New("status exists already")
		}
	}
	bot.System.Status = append(bot.System.Status, message)
	return nil
}

// RemoveStatus - Removes a status string from the main Object.
// message: the status message
func (bot *Object) RemoveStatus(message string) error {
	if bot.System == nil {
		return errors.New("object doesn't exist")
	}
	var ti int
	for i, k := range bot.System.Status {
		if k == message {
			ti = i
		}
	}
	if ti == 0 {
		return errors.New("status doesn't exist in my collection")
	}
	bot.System.Status = append(bot.System.Status[:ti], bot.System.Status[ti+1:]...)
	return nil
}

// Intro - Displays introduction and information on startup.
func (bot *Object) Intro(s *discordgo.Session) {
	var ars map[string]string
	fmt.Println(`
  ______             _   _           _
 |  ____|           | | | |         | |
 | |__ _ __ ___  ___| |_| |__  _   _| |_ ___
 |  __| '__/ _ \/ __| __| '_ \| | | | __/ _ \
 | |  | | | (_) \__ \ |_| |_) | |_| | ||  __/
 |_|  |_|  \___/|___/\__|_.__/ \__, |\__\___|
                                __/ |
                               |___/         `)
	<-time.After(2 * time.Second)
	// Collect some information and display it!
	guild, err := s.State.Guild(bot.Guild)
	if err != nil {
		fmt.Println(err)
	} else {
		// Channel count for server
		fmt.Print(len(guild.Channels), " Channel(s), ")
		// Member count for server
		fmt.Print(len(guild.Members), " Member(s), ")
		// Role count for server
		fmt.Print(len(guild.Roles), " Role(s), ")
	}

	// Collect the A.R.S Count.
	io, err := ioutil.ReadFile("autoresponse.json")
	if err != nil {
		fmt.Println(err)
	} else {
		json.Unmarshal(io, &ars)
		// ARS Count for the bot
		fmt.Println(len(ars), " A.R.S Rule(s).")
	}
	fmt.Println(`---------------------------------------------`)
}
<|endoftext|>


<commit_before>package v7_test

import (
	"code.cloudfoundry.org/cli/command/commandfakes"
	"code.cloudfoundry.org/cli/command/flag"
	"code.cloudfoundry.org/cli/command/translatableerror"
	. "code.cloudfoundry.org/cli/command/v7"
	v7 "code.cloudfoundry.org/cli/command/v7"
	"code.cloudfoundry.org/cli/command/v7/v7fakes"
	"code.cloudfoundry.org/cli/types"
	"code.cloudfoundry.org/cli/util/ui"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
)

var _ = Describe("ConfigCommand", func() {
	var (
		cmd             v7.ConfigCommand
		testUI          *ui.UI
		fakeConfig      *commandfakes.FakeConfig
		fakeSharedActor *commandfakes.FakeSharedActor
		fakeActor       *v7fakes.FakeActor
		executeErr      error
	)

	BeforeEach(func() {
		testUI = ui.NewTestUI(nil, NewBuffer(), NewBuffer())
		fakeConfig = new(commandfakes.FakeConfig)
		fakeSharedActor = new(commandfakes.FakeSharedActor)
		fakeActor = new(v7fakes.FakeActor)

		cmd = ConfigCommand{
			BaseCommand: BaseCommand{
				UI:          testUI,
				Config:      fakeConfig,
				SharedActor: fakeSharedActor,
				Actor:       fakeActor,
			},
		}
	})

	JustBeforeEach(func() {
		executeErr = cmd.Execute(nil)
	})

	When("there are no flags given", func() {
		It("returns an error", func() {
			Expect(executeErr).To(MatchError(
				translatableerror.IncorrectUsageError{Message: "at least one flag must be provided"},
			))
		})
	})

	When("using the async timeout flag", func() {
		BeforeEach(func() {
			cmd.AsyncTimeout = flag.Timeout{NullInt: types.NullInt{IsSet: true, Value: 2}}
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetAsyncTimeoutCallCount()).To(Equal(1))
			value := fakeConfig.SetAsyncTimeoutArgsForCall(0)
			Expect(value).To(Equal(2))
		})
	})

	When("using the color flag", func() {
		BeforeEach(func() {
			cmd.Color = flag.Color{IsSet: true, Value: "true"}
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetColorEnabledCallCount()).To(Equal(1))
			value := fakeConfig.SetColorEnabledArgsForCall(0)
			Expect(value).To(Equal("true"))
		})
	})

	When("using the locale flag", func() {
		BeforeEach(func() {
			cmd.Locale = flag.Locale{Locale: "en-US"}
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetLocaleCallCount()).To(Equal(1))
			value := fakeConfig.SetLocaleArgsForCall(0)
			Expect(value).To(Equal("en-US"))
		})
	})

	When("using the trace flag", func() {
		BeforeEach(func() {
			cmd.Trace = "my-trace-file"
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetTraceCallCount()).To(Equal(1))
			value := fakeConfig.SetTraceArgsForCall(0)
			Expect(value).To(Equal("my-trace-file"))
		})
	})
})
<commit_msg>V7: remove base command from config test [#172213932]<commit_after>package v7_test

import (
	"code.cloudfoundry.org/cli/command/commandfakes"
	"code.cloudfoundry.org/cli/command/flag"
	"code.cloudfoundry.org/cli/command/translatableerror"
	. "code.cloudfoundry.org/cli/command/v7"
	v7 "code.cloudfoundry.org/cli/command/v7"
	"code.cloudfoundry.org/cli/types"
	"code.cloudfoundry.org/cli/util/ui"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
)

var _ = Describe("ConfigCommand", func() {
	var (
		cmd        v7.ConfigCommand
		testUI     *ui.UI
		fakeConfig *commandfakes.FakeConfig
		executeErr error
	)

	BeforeEach(func() {
		testUI = ui.NewTestUI(nil, NewBuffer(), NewBuffer())
		fakeConfig = new(commandfakes.FakeConfig)

		cmd = ConfigCommand{
			UI:     testUI,
			Config: fakeConfig,
		}
	})

	JustBeforeEach(func() {
		executeErr = cmd.Execute(nil)
	})

	When("there are no flags given", func() {
		It("returns an error", func() {
			Expect(executeErr).To(MatchError(
				translatableerror.IncorrectUsageError{Message: "at least one flag must be provided"},
			))
		})
	})

	When("using the async timeout flag", func() {
		BeforeEach(func() {
			cmd.AsyncTimeout = flag.Timeout{NullInt: types.NullInt{IsSet: true, Value: 2}}
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetAsyncTimeoutCallCount()).To(Equal(1))
			value := fakeConfig.SetAsyncTimeoutArgsForCall(0)
			Expect(value).To(Equal(2))
		})
	})

	When("using the color flag", func() {
		BeforeEach(func() {
			cmd.Color = flag.Color{IsSet: true, Value: "true"}
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetColorEnabledCallCount()).To(Equal(1))
			value := fakeConfig.SetColorEnabledArgsForCall(0)
			Expect(value).To(Equal("true"))
		})
	})

	When("using the locale flag", func() {
		BeforeEach(func() {
			cmd.Locale = flag.Locale{Locale: "en-US"}
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetLocaleCallCount()).To(Equal(1))
			value := fakeConfig.SetLocaleArgsForCall(0)
			Expect(value).To(Equal("en-US"))
		})
	})

	When("using the trace flag", func() {
		BeforeEach(func() {
			cmd.Trace = "my-trace-file"
		})

		It("successfully updates the config", func() {
			Expect(executeErr).To(Not(HaveOccurred()))
			Expect(fakeConfig.SetTraceCallCount()).To(Equal(1))
			value := fakeConfig.SetTraceArgsForCall(0)
			Expect(value).To(Equal("my-trace-file"))
		})
	})
})
<|endoftext|>


<commit_before>package clusterstatus

import (
	"strings"

	"fmt"

	"github.com/coldbrewcloud/coldbrew-cli/aws"
	"github.com/coldbrewcloud/coldbrew-cli/console"
	"github.com/coldbrewcloud/coldbrew-cli/core"
	"github.com/coldbrewcloud/coldbrew-cli/flags"
	"github.com/coldbrewcloud/coldbrew-cli/utils"
	"github.com/coldbrewcloud/coldbrew-cli/utils/conv"
	"gopkg.in/alecthomas/kingpin.v2"
)

type Command struct {
	globalFlags    *flags.GlobalFlags
	commandFlags   *Flags
	awsClient      *aws.Client
	clusterNameArg *string
}

func (c *Command) Init(ka *kingpin.Application, globalFlags *flags.GlobalFlags) *kingpin.CmdClause {
	c.globalFlags = globalFlags

	cmd := ka.Command("cluster-status",
		"See: "+console.ColorFnHelpLink("https://github.com/coldbrewcloud/coldbrew-cli/wiki/CLI-Command:-cluster-status"))
	c.commandFlags = NewFlags(cmd)

	c.clusterNameArg = cmd.Arg("cluster-name", "Cluster name").Required().String()

	return cmd
}

func (c *Command) Run() error {
	c.awsClient = c.globalFlags.GetAWSClient()

	// AWS networking
	regionName, vpcID, subnetIDs, err := c.getAWSInfo()
	if err != nil {
		return console.ExitWithError(err)
	}

	// cluster name
	clusterName := strings.TrimSpace(conv.S(c.clusterNameArg))
	if !core.ClusterNameRE.MatchString(clusterName) {
		return console.ExitWithError(core.NewErrorExtraInfo(
			fmt.Errorf("Invalid cluster name [%s]", clusterName), "https://github.com/coldbrewcloud/coldbrew-cli/wiki/Configuration-File#cluster"))
	}

	console.Info("Cluster")
	console.DetailWithResource("Name", clusterName)

	// AWS env
	console.Info("AWS")
	console.DetailWithResource("Region", regionName)
	console.DetailWithResource("VPC", vpcID)
	console.DetailWithResource("Subnets", strings.Join(subnetIDs, " "))

	// ECS
	console.Info("ECS")
	showECSClusterDetails := false

	// ecs cluster
	ecsClusterName := core.DefaultECSClusterName(clusterName)
	ecsCluster, err := c.awsClient.ECS().RetrieveCluster(ecsClusterName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve ECS Cluster [%s]: %s", ecsClusterName, err.Error())
	}
	if ecsCluster == nil || conv.S(ecsCluster.Status) == "INACTIVE" {
		console.DetailWithResourceNote("ECS Cluster", ecsClusterName, "(not found)", true)
	} else {
		console.DetailWithResource("ECS Cluster", ecsClusterName)
		showECSClusterDetails = true
	}

	// ecs service role
	ecsServiceRoleName := core.DefaultECSServiceRoleName(clusterName)
	ecsServiceRole, err := c.awsClient.IAM().RetrieveRole(ecsServiceRoleName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve IAM Role [%s]: %s", ecsServiceRoleName, err.Error())
	}
	if ecsServiceRole == nil {
		console.DetailWithResourceNote("IAM Role for ECS Services", ecsServiceRoleName, "(not found)", true)
	} else {
		console.DetailWithResource("IAM Role for ECS Services", ecsServiceRoleName)
	}

	// ecs cluster details
	if showECSClusterDetails {
		console.DetailWithResource("ECS Services", conv.I64S(conv.I64(ecsCluster.ActiveServicesCount)))
		console.DetailWithResource("ECS Tasks (running/pending)",
			fmt.Sprintf("%d/%d",
				conv.I64(ecsCluster.RunningTasksCount),
				conv.I64(ecsCluster.PendingTasksCount)))
		console.DetailWithResource("ECS Container Instances", conv.I64S(conv.I64(ecsCluster.RegisteredContainerInstancesCount)))

	}

	// launch config and auto scaling group
	console.Info("Auto Scaling")

	// launch configuration
	launchConfigName := core.DefaultLaunchConfigurationName(clusterName)
	launchConfig, err := c.awsClient.AutoScaling().RetrieveLaunchConfiguration(launchConfigName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve Launch Configuration [%s]: %s", launchConfigName, err.Error())
	}
	if launchConfig == nil {
		console.DetailWithResourceNote("EC2 Launch Configuration", launchConfigName, "(not found)", true)
	} else {
		console.DetailWithResource("EC2 Launch Configuration", launchConfigName)

		instanceProfileARN := conv.S(launchConfig.IamInstanceProfile)
		console.DetailWithResource("  IAM Instance Profile", aws.GetIAMInstanceProfileNameFromARN(instanceProfileARN))

		console.DetailWithResource("  Instance Type", conv.S(launchConfig.InstanceType))
		console.DetailWithResource("  Image ID", conv.S(launchConfig.ImageId))
		console.DetailWithResource("  Key Pair", conv.S(launchConfig.KeyName))

		securityGroupIDs := []string{}
		for _, sg := range launchConfig.SecurityGroups {
			securityGroupIDs = append(securityGroupIDs, conv.S(sg))
		}
		securityGroups, err := c.awsClient.EC2().RetrieveSecurityGroups(securityGroupIDs)
		if err != nil {
			return console.ExitWithErrorString("Failed to retrieve Security Groups [%s]: %s", strings.Join(securityGroupIDs, ","), err.Error())
		}
		securityGroupNames := []string{}
		for _, sg := range securityGroups {
			securityGroupNames = append(securityGroupNames, conv.S(sg.GroupName))
		}
		console.DetailWithResource("  Security Groups", strings.Join(securityGroupNames, " "))
	}

	// auto scaling group
	autoScalingGroupName := core.DefaultAutoScalingGroupName(clusterName)
	autoScalingGroup, err := c.awsClient.AutoScaling().RetrieveAutoScalingGroup(autoScalingGroupName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve Auto Scaling Group [%s]: %s", autoScalingGroupName, err.Error())
	}
	if autoScalingGroup == nil {
		console.DetailWithResourceNote("EC2 Auto Scaling Group", autoScalingGroupName, "(not found)", true)
	} else if utils.IsBlank(conv.S(autoScalingGroup.Status)) {
		console.DetailWithResource("EC2 Auto Scaling Group", autoScalingGroupName)
		console.DetailWithResource("  Instances (current/desired/min/max)",
			fmt.Sprintf("%d/%d/%d/%d",
				len(autoScalingGroup.Instances),
				conv.I64(autoScalingGroup.DesiredCapacity),
				conv.I64(autoScalingGroup.MinSize),
				conv.I64(autoScalingGroup.MaxSize)))
	} else {
		console.DetailWithResourceNote("EC2 Auto Scaling Group", autoScalingGroupName, "(deleting)", true)
	}

	// ECS Container Instances
	if ecsCluster != nil && !conv.B(c.commandFlags.ExcludeContainerInstanceInfos) {
		containerInstanceARNs, err := c.awsClient.ECS().ListContainerInstanceARNs(ecsClusterName)
		if err != nil {
			return console.ExitWithErrorString("Failed to list ECS Container Instances: %s", err.Error())
		}
		containerInstances, err := c.awsClient.ECS().RetrieveContainerInstances(ecsClusterName, containerInstanceARNs)
		if err != nil {
			return console.ExitWithErrorString("Failed to retrieve ECS Container Instances: %s", err.Error())
		}

		// retrieve EC2 Instance info
		ec2InstanceIDs := []string{}
		for _, ci := range containerInstances {
			ec2InstanceIDs = append(ec2InstanceIDs, conv.S(ci.Ec2InstanceId))
		}
		ec2Instances, err := c.awsClient.EC2().RetrieveInstances(ec2InstanceIDs)
		if err != nil {
			return console.ExitWithErrorString("Failed to retrieve EC2 Instances: %s", err.Error())
		}

		for _, ci := range containerInstances {
			console.Info("ECS Container Instance")

			console.DetailWithResource("ID", aws.GetECSContainerInstanceIDFromARN(conv.S(ci.ContainerInstanceArn)))

			if conv.B(ci.AgentConnected) {
				console.DetailWithResource("Status", conv.S(ci.Status))
			} else {
				console.DetailWithResourceNote("Status", conv.S(ci.Status), "(agent not connected)", true)
			}

			console.DetailWithResource("Tasks (running/pending)", fmt.Sprintf("%d/%d",
				conv.I64(ci.RunningTasksCount),
				conv.I64(ci.PendingTasksCount)))

			var registeredCPU, registeredMemory, remainingCPU, remainingMemory int64
			for _, rr := range ci.RegisteredResources {
				switch strings.ToLower(conv.S(rr.Name)) {
				case "cpu":
					registeredCPU = conv.I64(rr.IntegerValue)
				case "memory":
					registeredMemory = conv.I64(rr.IntegerValue)
				}
			}
			for _, rr := range ci.RemainingResources {
				switch strings.ToLower(conv.S(rr.Name)) {
				case "cpu":
					remainingCPU = conv.I64(rr.IntegerValue)
				case "memory":
					remainingMemory = conv.I64(rr.IntegerValue)
				}
			}

			console.DetailWithResource("CPU (remaining/registered)", fmt.Sprintf("%.2f/%.2f",
				float64(remainingCPU)/1024.0, float64(registeredCPU)/1024.0))
			console.DetailWithResource("Memory (remaining/registered)", fmt.Sprintf("%dM/%dM,",
				remainingMemory, registeredMemory))

			console.DetailWithResource("EC2 Instance ID", conv.S(ci.Ec2InstanceId))
			for _, ei := range ec2Instances {
				if conv.S(ei.InstanceId) == conv.S(ci.Ec2InstanceId) {
					if !utils.IsBlank(conv.S(ei.PrivateIpAddress)) {
						console.DetailWithResource("  Private IP", conv.S(ei.PrivateIpAddress))
					}
					if !utils.IsBlank(conv.S(ei.PublicIpAddress)) {
						console.DetailWithResource("  Public IP", conv.S(ei.PublicIpAddress))
					}
					break
				}
			}
		}
	}

	return nil
}
<commit_msg>fix "cluster-status" minor bug<commit_after>package clusterstatus

import (
	"strings"

	"fmt"

	"github.com/coldbrewcloud/coldbrew-cli/aws"
	"github.com/coldbrewcloud/coldbrew-cli/console"
	"github.com/coldbrewcloud/coldbrew-cli/core"
	"github.com/coldbrewcloud/coldbrew-cli/flags"
	"github.com/coldbrewcloud/coldbrew-cli/utils"
	"github.com/coldbrewcloud/coldbrew-cli/utils/conv"
	"gopkg.in/alecthomas/kingpin.v2"
)

type Command struct {
	globalFlags    *flags.GlobalFlags
	commandFlags   *Flags
	awsClient      *aws.Client
	clusterNameArg *string
}

func (c *Command) Init(ka *kingpin.Application, globalFlags *flags.GlobalFlags) *kingpin.CmdClause {
	c.globalFlags = globalFlags

	cmd := ka.Command("cluster-status",
		"See: "+console.ColorFnHelpLink("https://github.com/coldbrewcloud/coldbrew-cli/wiki/CLI-Command:-cluster-status"))
	c.commandFlags = NewFlags(cmd)

	c.clusterNameArg = cmd.Arg("cluster-name", "Cluster name").Required().String()

	return cmd
}

func (c *Command) Run() error {
	c.awsClient = c.globalFlags.GetAWSClient()

	// AWS networking
	regionName, vpcID, subnetIDs, err := c.getAWSInfo()
	if err != nil {
		return console.ExitWithError(err)
	}

	// cluster name
	clusterName := strings.TrimSpace(conv.S(c.clusterNameArg))
	if !core.ClusterNameRE.MatchString(clusterName) {
		return console.ExitWithError(core.NewErrorExtraInfo(
			fmt.Errorf("Invalid cluster name [%s]", clusterName), "https://github.com/coldbrewcloud/coldbrew-cli/wiki/Configuration-File#cluster"))
	}

	console.Info("Cluster")
	console.DetailWithResource("Name", clusterName)

	// AWS env
	console.Info("AWS")
	console.DetailWithResource("Region", regionName)
	console.DetailWithResource("VPC", vpcID)
	console.DetailWithResource("Subnets", strings.Join(subnetIDs, " "))

	// ECS
	console.Info("ECS")
	showECSClusterDetails := false

	// ecs cluster
	ecsClusterName := core.DefaultECSClusterName(clusterName)
	ecsCluster, err := c.awsClient.ECS().RetrieveCluster(ecsClusterName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve ECS Cluster [%s]: %s", ecsClusterName, err.Error())
	}
	if ecsCluster == nil || conv.S(ecsCluster.Status) == "INACTIVE" {
		console.DetailWithResourceNote("ECS Cluster", ecsClusterName, "(not found)", true)
		ecsCluster = nil
	} else {
		console.DetailWithResource("ECS Cluster", ecsClusterName)
		showECSClusterDetails = true
	}

	// ecs service role
	ecsServiceRoleName := core.DefaultECSServiceRoleName(clusterName)
	ecsServiceRole, err := c.awsClient.IAM().RetrieveRole(ecsServiceRoleName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve IAM Role [%s]: %s", ecsServiceRoleName, err.Error())
	}
	if ecsServiceRole == nil {
		console.DetailWithResourceNote("IAM Role for ECS Services", ecsServiceRoleName, "(not found)", true)
	} else {
		console.DetailWithResource("IAM Role for ECS Services", ecsServiceRoleName)
	}

	// ecs cluster details
	if showECSClusterDetails {
		console.DetailWithResource("ECS Services", conv.I64S(conv.I64(ecsCluster.ActiveServicesCount)))
		console.DetailWithResource("ECS Tasks (running/pending)",
			fmt.Sprintf("%d/%d",
				conv.I64(ecsCluster.RunningTasksCount),
				conv.I64(ecsCluster.PendingTasksCount)))
		console.DetailWithResource("ECS Container Instances", conv.I64S(conv.I64(ecsCluster.RegisteredContainerInstancesCount)))

	}

	// launch config and auto scaling group
	console.Info("Auto Scaling")

	// launch configuration
	launchConfigName := core.DefaultLaunchConfigurationName(clusterName)
	launchConfig, err := c.awsClient.AutoScaling().RetrieveLaunchConfiguration(launchConfigName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve Launch Configuration [%s]: %s", launchConfigName, err.Error())
	}
	if launchConfig == nil {
		console.DetailWithResourceNote("EC2 Launch Configuration", launchConfigName, "(not found)", true)
	} else {
		console.DetailWithResource("EC2 Launch Configuration", launchConfigName)

		instanceProfileARN := conv.S(launchConfig.IamInstanceProfile)
		console.DetailWithResource("  IAM Instance Profile", aws.GetIAMInstanceProfileNameFromARN(instanceProfileARN))

		console.DetailWithResource("  Instance Type", conv.S(launchConfig.InstanceType))
		console.DetailWithResource("  Image ID", conv.S(launchConfig.ImageId))
		console.DetailWithResource("  Key Pair", conv.S(launchConfig.KeyName))

		securityGroupIDs := []string{}
		for _, sg := range launchConfig.SecurityGroups {
			securityGroupIDs = append(securityGroupIDs, conv.S(sg))
		}
		securityGroups, err := c.awsClient.EC2().RetrieveSecurityGroups(securityGroupIDs)
		if err != nil {
			return console.ExitWithErrorString("Failed to retrieve Security Groups [%s]: %s", strings.Join(securityGroupIDs, ","), err.Error())
		}
		securityGroupNames := []string{}
		for _, sg := range securityGroups {
			securityGroupNames = append(securityGroupNames, conv.S(sg.GroupName))
		}
		console.DetailWithResource("  Security Groups", strings.Join(securityGroupNames, " "))
	}

	// auto scaling group
	autoScalingGroupName := core.DefaultAutoScalingGroupName(clusterName)
	autoScalingGroup, err := c.awsClient.AutoScaling().RetrieveAutoScalingGroup(autoScalingGroupName)
	if err != nil {
		return console.ExitWithErrorString("Failed to retrieve Auto Scaling Group [%s]: %s", autoScalingGroupName, err.Error())
	}
	if autoScalingGroup == nil {
		console.DetailWithResourceNote("EC2 Auto Scaling Group", autoScalingGroupName, "(not found)", true)
	} else if utils.IsBlank(conv.S(autoScalingGroup.Status)) {
		console.DetailWithResource("EC2 Auto Scaling Group", autoScalingGroupName)
		console.DetailWithResource("  Instances (current/desired/min/max)",
			fmt.Sprintf("%d/%d/%d/%d",
				len(autoScalingGroup.Instances),
				conv.I64(autoScalingGroup.DesiredCapacity),
				conv.I64(autoScalingGroup.MinSize),
				conv.I64(autoScalingGroup.MaxSize)))
	} else {
		console.DetailWithResourceNote("EC2 Auto Scaling Group", autoScalingGroupName, "(deleting)", true)
	}

	// ECS Container Instances
	if ecsCluster != nil && !conv.B(c.commandFlags.ExcludeContainerInstanceInfos) {
		containerInstanceARNs, err := c.awsClient.ECS().ListContainerInstanceARNs(ecsClusterName)
		if err != nil {
			return console.ExitWithErrorString("Failed to list ECS Container Instances: %s", err.Error())
		}
		containerInstances, err := c.awsClient.ECS().RetrieveContainerInstances(ecsClusterName, containerInstanceARNs)
		if err != nil {
			return console.ExitWithErrorString("Failed to retrieve ECS Container Instances: %s", err.Error())
		}

		// retrieve EC2 Instance info
		ec2InstanceIDs := []string{}
		for _, ci := range containerInstances {
			ec2InstanceIDs = append(ec2InstanceIDs, conv.S(ci.Ec2InstanceId))
		}
		ec2Instances, err := c.awsClient.EC2().RetrieveInstances(ec2InstanceIDs)
		if err != nil {
			return console.ExitWithErrorString("Failed to retrieve EC2 Instances: %s", err.Error())
		}

		for _, ci := range containerInstances {
			console.Info("ECS Container Instance")

			console.DetailWithResource("ID", aws.GetECSContainerInstanceIDFromARN(conv.S(ci.ContainerInstanceArn)))

			if conv.B(ci.AgentConnected) {
				console.DetailWithResource("Status", conv.S(ci.Status))
			} else {
				console.DetailWithResourceNote("Status", conv.S(ci.Status), "(agent not connected)", true)
			}

			console.DetailWithResource("Tasks (running/pending)", fmt.Sprintf("%d/%d",
				conv.I64(ci.RunningTasksCount),
				conv.I64(ci.PendingTasksCount)))

			var registeredCPU, registeredMemory, remainingCPU, remainingMemory int64
			for _, rr := range ci.RegisteredResources {
				switch strings.ToLower(conv.S(rr.Name)) {
				case "cpu":
					registeredCPU = conv.I64(rr.IntegerValue)
				case "memory":
					registeredMemory = conv.I64(rr.IntegerValue)
				}
			}
			for _, rr := range ci.RemainingResources {
				switch strings.ToLower(conv.S(rr.Name)) {
				case "cpu":
					remainingCPU = conv.I64(rr.IntegerValue)
				case "memory":
					remainingMemory = conv.I64(rr.IntegerValue)
				}
			}

			console.DetailWithResource("CPU (remaining/registered)", fmt.Sprintf("%.2f/%.2f",
				float64(remainingCPU)/1024.0, float64(registeredCPU)/1024.0))
			console.DetailWithResource("Memory (remaining/registered)", fmt.Sprintf("%dM/%dM,",
				remainingMemory, registeredMemory))

			console.DetailWithResource("EC2 Instance ID", conv.S(ci.Ec2InstanceId))
			for _, ei := range ec2Instances {
				if conv.S(ei.InstanceId) == conv.S(ci.Ec2InstanceId) {
					if !utils.IsBlank(conv.S(ei.PrivateIpAddress)) {
						console.DetailWithResource("  Private IP", conv.S(ei.PrivateIpAddress))
					}
					if !utils.IsBlank(conv.S(ei.PublicIpAddress)) {
						console.DetailWithResource("  Public IP", conv.S(ei.PublicIpAddress))
					}
					break
				}
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// +build gofuzz

package fuzz

import (
	"context"
	"time"
)

var (
	ctx = context.Background()
	rdb *Client
)

func init() {
	rdb = NewClient(&Options{
		Addr:         ":6379",
		DialTimeout:  10 * time.Second,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 10 * time.Second,
		PoolSize:     10,
		PoolTimeout:  10 * time.Second,
	})
}

func Fuzz(data []byte) int {
	arrayLen := len(data)
	if arrayLen < 4 {
		return -1
	}
	maxIter := int(uint(data[0]))
	for i := 0; i < maxIter && i < arrayLen; i++ {
		n := i % arrayLen
		if n == 0 {
			_ = rdb.Set(ctx, string(data[i:]), string(data[i:]), 0).Err()
		} else if n == 1 {
			_, _ = rdb.Get(ctx, string(data[i:])).Result()
		} else if n == 2 {
			_, _ = rdb.Incr(ctx, string(data[i:])).Result()
		} else if n == 3 {
			var cursor uint64
			_, _, _ = rdb.Scan(ctx, cursor, string(data[i:]), 10).Result()
		}
	}
	return 1
}
<commit_msg>Fixed broken fuzzer<commit_after>// +build gofuzz

package fuzz

import (
	"context"
	"github.com/go-redis/redis"
	"time"
)

var (
	ctx = context.Background()
	rdb *redis.Client
)

func init() {
	rdb = redis.NewClient(&redis.Options{
		Addr:         ":6379",
		DialTimeout:  10 * time.Second,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 10 * time.Second,
		PoolSize:     10,
		PoolTimeout:  10 * time.Second,
	})
}

func Fuzz(data []byte) int {
	arrayLen := len(data)
	if arrayLen < 4 {
		return -1
	}
	maxIter := int(uint(data[0]))
	for i := 0; i < maxIter && i < arrayLen; i++ {
		n := i % arrayLen
		if n == 0 {
			_ = rdb.Set(ctx, string(data[i:]), string(data[i:]), 0).Err()
		} else if n == 1 {
			_, _ = rdb.Get(ctx, string(data[i:])).Result()
		} else if n == 2 {
			_, _ = rdb.Incr(ctx, string(data[i:])).Result()
		} else if n == 3 {
			var cursor uint64
			_, _, _ = rdb.Scan(ctx, cursor, string(data[i:]), 10).Result()
		}
	}
	return 1
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"log"
	"regexp"

	"github.com/github/hub/github"
	"github.com/github/hub/utils"
)

var cmdCheckout = &Command{
	Run:          checkout,
	GitExtension: true,
	Usage:        "checkout PULLREQ-URL [BRANCH]",
	Short:        "Switch the active branch to another branch",
	Long: `Checks out the head of the pull request as a local branch, to allow for
reviewing, rebasing and otherwise cleaning up the commits in the pull
request before merging. The name of the local branch can explicitly be
set with BRANCH.
`,
}

func init() {
	CmdRunner.Use(cmdCheckout)
}

/**
  $ gh checkout https://github.com/jingweno/gh/pull/73
  > git remote add -f -t feature git://github:com/foo/gh.git
  > git checkout --track -B foo-feature foo/feature

  $ gh checkout https://github.com/jingweno/gh/pull/73 custom-branch-name
**/
func checkout(command *Command, args *Args) {
	if !args.IsParamsEmpty() {
		err := transformCheckoutArgs(args)
		utils.Check(err)
	}
}

func transformCheckoutArgs(args *Args) error {
	words := args.Words()

	if len(words) == 0 {
		return nil
	}

	checkoutURL := words[0]
	var newBranchName string
	if len(words) > 1 {
		newBranchName = words[1]
	}

	url, err := github.ParseURL(checkoutURL)
	if err != nil {
		// not a valid GitHub URL
		return nil
	}

	pullURLRegex := regexp.MustCompile("^pull/(\\d+)")
	projectPath := url.ProjectPath()
	if !pullURLRegex.MatchString(projectPath) {
		// not a valid PR URL
		return nil
	}

	id := pullURLRegex.FindStringSubmatch(projectPath)[1]
	gh := github.NewClient(url.Project.Host)
	pullRequest, err := gh.PullRequest(url.Project, id)
	if err != nil {
		return err
	}

	if idx := args.IndexOfParam(newBranchName); idx >= 0 {
		args.RemoveParam(idx)
	}

	user, branch := parseUserBranchFromPR(pullRequest)
	if pullRequest.Head.Repo.ID == 0 {
		return fmt.Errorf("Error: %s's fork is not available anymore", user)
	}

	if newBranchName == "" {
		newBranchName = fmt.Sprintf("%s-%s", user, branch)
	}

	repo := github.LocalRepo()
	_, err = repo.RemoteByName(user)
	if err == nil {
		args.Before("git", "remote", "set-branches", "--add", user, branch)
		remoteURL := fmt.Sprintf("+refs/heads/%s:refs/remotes/%s/%s", branch, user, branch)
		args.Before("git", "fetch", user, remoteURL)
	} else {
		u := url.Project.GitURL("", user, pullRequest.Head.Repo.Private)
		args.Before("git", "remote", "add", "-f", "-t", branch, user, u)
	}

	idx := args.IndexOfParam(checkoutURL)
	args.RemoveParam(idx)
	args.InsertParam(idx, "--track", "-B", newBranchName, fmt.Sprintf("%s/%s", user, branch))

	return nil
}
<commit_msg>Don't import log<commit_after>package commands

import (
	"fmt"
	"regexp"

	"github.com/github/hub/github"
	"github.com/github/hub/utils"
)

var cmdCheckout = &Command{
	Run:          checkout,
	GitExtension: true,
	Usage:        "checkout PULLREQ-URL [BRANCH]",
	Short:        "Switch the active branch to another branch",
	Long: `Checks out the head of the pull request as a local branch, to allow for
reviewing, rebasing and otherwise cleaning up the commits in the pull
request before merging. The name of the local branch can explicitly be
set with BRANCH.
`,
}

func init() {
	CmdRunner.Use(cmdCheckout)
}

/**
  $ gh checkout https://github.com/jingweno/gh/pull/73
  > git remote add -f -t feature git://github:com/foo/gh.git
  > git checkout --track -B foo-feature foo/feature

  $ gh checkout https://github.com/jingweno/gh/pull/73 custom-branch-name
**/
func checkout(command *Command, args *Args) {
	if !args.IsParamsEmpty() {
		err := transformCheckoutArgs(args)
		utils.Check(err)
	}
}

func transformCheckoutArgs(args *Args) error {
	words := args.Words()

	if len(words) == 0 {
		return nil
	}

	checkoutURL := words[0]
	var newBranchName string
	if len(words) > 1 {
		newBranchName = words[1]
	}

	url, err := github.ParseURL(checkoutURL)
	if err != nil {
		// not a valid GitHub URL
		return nil
	}

	pullURLRegex := regexp.MustCompile("^pull/(\\d+)")
	projectPath := url.ProjectPath()
	if !pullURLRegex.MatchString(projectPath) {
		// not a valid PR URL
		return nil
	}

	id := pullURLRegex.FindStringSubmatch(projectPath)[1]
	gh := github.NewClient(url.Project.Host)
	pullRequest, err := gh.PullRequest(url.Project, id)
	if err != nil {
		return err
	}

	if idx := args.IndexOfParam(newBranchName); idx >= 0 {
		args.RemoveParam(idx)
	}

	user, branch := parseUserBranchFromPR(pullRequest)
	if pullRequest.Head.Repo.ID == 0 {
		return fmt.Errorf("Error: %s's fork is not available anymore", user)
	}

	if newBranchName == "" {
		newBranchName = fmt.Sprintf("%s-%s", user, branch)
	}

	repo := github.LocalRepo()
	_, err = repo.RemoteByName(user)
	if err == nil {
		args.Before("git", "remote", "set-branches", "--add", user, branch)
		remoteURL := fmt.Sprintf("+refs/heads/%s:refs/remotes/%s/%s", branch, user, branch)
		args.Before("git", "fetch", user, remoteURL)
	} else {
		u := url.Project.GitURL("", user, pullRequest.Head.Repo.Private)
		args.Before("git", "remote", "add", "-f", "-t", branch, user, u)
	}

	idx := args.IndexOfParam(checkoutURL)
	args.RemoveParam(idx)
	args.InsertParam(idx, "--track", "-B", newBranchName, fmt.Sprintf("%s/%s", user, branch))

	return nil
}
<|endoftext|>


<commit_before>package buf

import (
	"io"

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

// BytesToBufferReader is a Reader that adjusts its reading speed automatically.
type BytesToBufferReader struct {
	io.Reader
	buffer []byte
}

// NewBytesToBufferReader returns a new BytesToBufferReader.
func NewBytesToBufferReader(reader io.Reader) Reader {
	return &BytesToBufferReader{
		Reader: reader,
	}
}

func (r *BytesToBufferReader) readSmall() (MultiBuffer, error) {
	b := New()
	for i := 0; i < 64; i++ {
		err := b.Reset(ReadFrom(r.Reader))
		if b.IsFull() {
			r.buffer = newBytes(Size + 1)
		}
		if !b.IsEmpty() {
			return NewMultiBufferValue(b), nil
		}
		if err != nil {
			b.Release()
			return nil, err
		}
	}

	return nil, newError("Reader returns too many empty payloads.")
}

func (r *BytesToBufferReader) freeBuffer() {
	freeBytes(r.buffer)
	r.buffer = nil
}

// ReadMultiBuffer implements Reader.
func (r *BytesToBufferReader) ReadMultiBuffer() (MultiBuffer, error) {
	if r.buffer == nil {
		return r.readSmall()
	}

	nBytes, err := r.Reader.Read(r.buffer)
	if nBytes > 0 {
		mb := NewMultiBufferCap(int32(nBytes/Size) + 1)
		mb.Write(r.buffer[:nBytes])
		if nBytes == len(r.buffer) && nBytes < int(largeSize) {
			freeBytes(r.buffer)
			r.buffer = newBytes(int32(nBytes) + 1)
		} else if nBytes < Size {
			r.freeBuffer()
		}
		return mb, nil
	}

	r.freeBuffer()

	if err != nil {
		return nil, err
	}

	// Read() returns empty payload and nil err. We don't expect this to happen, but just in case.
	return r.readSmall()
}

// BufferedReader is a Reader that keeps its internal buffer.
type BufferedReader struct {
	stream   Reader
	leftOver MultiBuffer
	buffered bool
}

// NewBufferedReader returns a new BufferedReader.
func NewBufferedReader(reader Reader) *BufferedReader {
	return &BufferedReader{
		stream:   reader,
		buffered: true,
	}
}

// SetBuffered sets whether to keep the interal buffer.
func (r *BufferedReader) SetBuffered(f bool) {
	r.buffered = f
}

// IsBuffered returns true if internal buffer is used.
func (r *BufferedReader) IsBuffered() bool {
	return r.buffered
}

// BufferedBytes returns the number of bytes that is cached in this reader.
func (r *BufferedReader) BufferedBytes() int32 {
	return r.leftOver.Len()
}

// ReadByte implements io.ByteReader.
func (r *BufferedReader) ReadByte() (byte, error) {
	var b [1]byte
	_, err := r.Read(b[:])
	return b[0], err
}

// Read implements io.Reader. It reads from internal buffer first (if available) and then reads from the underlying reader.
func (r *BufferedReader) Read(b []byte) (int, error) {
	if r.leftOver != nil {
		nBytes, _ := r.leftOver.Read(b)
		if r.leftOver.IsEmpty() {
			r.leftOver.Release()
			r.leftOver = nil
		}
		return nBytes, nil
	}

	if !r.buffered {
		if reader, ok := r.stream.(io.Reader); ok {
			return reader.Read(b)
		}
	}

	mb, err := r.stream.ReadMultiBuffer()
	if mb != nil {
		nBytes, _ := mb.Read(b)
		if !mb.IsEmpty() {
			r.leftOver = mb
		}
		return nBytes, err
	}
	return 0, err
}

// ReadMultiBuffer implements Reader.
func (r *BufferedReader) ReadMultiBuffer() (MultiBuffer, error) {
	if r.leftOver != nil {
		mb := r.leftOver
		r.leftOver = nil
		return mb, nil
	}

	return r.stream.ReadMultiBuffer()
}

// ReadAtMost returns a MultiBuffer with at most size.
func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
	if r.leftOver == nil {
		mb, err := r.stream.ReadMultiBuffer()
		if mb.IsEmpty() && err != nil {
			return nil, err
		}
		r.leftOver = mb
	}

	mb := r.leftOver.SliceBySize(size)
	if r.leftOver.IsEmpty() {
		r.leftOver = nil
	}
	return mb, nil
}

func (r *BufferedReader) writeToInternal(writer io.Writer) (int64, error) {
	mbWriter := NewWriter(writer)
	totalBytes := int64(0)
	if r.leftOver != nil {
		totalBytes += int64(r.leftOver.Len())
		if err := mbWriter.WriteMultiBuffer(r.leftOver); err != nil {
			return 0, err
		}
		r.leftOver = nil
	}

	for {
		mb, err := r.stream.ReadMultiBuffer()
		if mb != nil {
			totalBytes += int64(mb.Len())
			if werr := mbWriter.WriteMultiBuffer(mb); werr != nil {
				return totalBytes, err
			}
		}
		if err != nil {
			return totalBytes, err
		}
	}
}

// WriteTo implements io.WriterTo.
func (r *BufferedReader) WriteTo(writer io.Writer) (int64, error) {
	nBytes, err := r.writeToInternal(writer)
	if errors.Cause(err) == io.EOF {
		return nBytes, nil
	}
	return nBytes, err
}
<commit_msg>add check for largeSize in case there is just one pool<commit_after>package buf

import (
	"io"

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

// BytesToBufferReader is a Reader that adjusts its reading speed automatically.
type BytesToBufferReader struct {
	io.Reader
	buffer []byte
}

// NewBytesToBufferReader returns a new BytesToBufferReader.
func NewBytesToBufferReader(reader io.Reader) Reader {
	return &BytesToBufferReader{
		Reader: reader,
	}
}

func (r *BytesToBufferReader) readSmall() (MultiBuffer, error) {
	b := New()
	for i := 0; i < 64; i++ {
		err := b.Reset(ReadFrom(r.Reader))
		if b.IsFull() && largeSize > Size {
			r.buffer = newBytes(Size + 1)
		}
		if !b.IsEmpty() {
			return NewMultiBufferValue(b), nil
		}
		if err != nil {
			b.Release()
			return nil, err
		}
	}

	return nil, newError("Reader returns too many empty payloads.")
}

func (r *BytesToBufferReader) freeBuffer() {
	freeBytes(r.buffer)
	r.buffer = nil
}

// ReadMultiBuffer implements Reader.
func (r *BytesToBufferReader) ReadMultiBuffer() (MultiBuffer, error) {
	if r.buffer == nil || largeSize == Size {
		return r.readSmall()
	}

	nBytes, err := r.Reader.Read(r.buffer)
	if nBytes > 0 {
		mb := NewMultiBufferCap(int32(nBytes/Size) + 1)
		mb.Write(r.buffer[:nBytes])
		if nBytes == len(r.buffer) && nBytes < int(largeSize) {
			freeBytes(r.buffer)
			r.buffer = newBytes(int32(nBytes) + 1)
		} else if nBytes < Size {
			r.freeBuffer()
		}
		return mb, nil
	}

	r.freeBuffer()

	if err != nil {
		return nil, err
	}

	// Read() returns empty payload and nil err. We don't expect this to happen, but just in case.
	return r.readSmall()
}

// BufferedReader is a Reader that keeps its internal buffer.
type BufferedReader struct {
	stream   Reader
	leftOver MultiBuffer
	buffered bool
}

// NewBufferedReader returns a new BufferedReader.
func NewBufferedReader(reader Reader) *BufferedReader {
	return &BufferedReader{
		stream:   reader,
		buffered: true,
	}
}

// SetBuffered sets whether to keep the interal buffer.
func (r *BufferedReader) SetBuffered(f bool) {
	r.buffered = f
}

// IsBuffered returns true if internal buffer is used.
func (r *BufferedReader) IsBuffered() bool {
	return r.buffered
}

// BufferedBytes returns the number of bytes that is cached in this reader.
func (r *BufferedReader) BufferedBytes() int32 {
	return r.leftOver.Len()
}

// ReadByte implements io.ByteReader.
func (r *BufferedReader) ReadByte() (byte, error) {
	var b [1]byte
	_, err := r.Read(b[:])
	return b[0], err
}

// Read implements io.Reader. It reads from internal buffer first (if available) and then reads from the underlying reader.
func (r *BufferedReader) Read(b []byte) (int, error) {
	if r.leftOver != nil {
		nBytes, _ := r.leftOver.Read(b)
		if r.leftOver.IsEmpty() {
			r.leftOver.Release()
			r.leftOver = nil
		}
		return nBytes, nil
	}

	if !r.buffered {
		if reader, ok := r.stream.(io.Reader); ok {
			return reader.Read(b)
		}
	}

	mb, err := r.stream.ReadMultiBuffer()
	if mb != nil {
		nBytes, _ := mb.Read(b)
		if !mb.IsEmpty() {
			r.leftOver = mb
		}
		return nBytes, err
	}
	return 0, err
}

// ReadMultiBuffer implements Reader.
func (r *BufferedReader) ReadMultiBuffer() (MultiBuffer, error) {
	if r.leftOver != nil {
		mb := r.leftOver
		r.leftOver = nil
		return mb, nil
	}

	return r.stream.ReadMultiBuffer()
}

// ReadAtMost returns a MultiBuffer with at most size.
func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
	if r.leftOver == nil {
		mb, err := r.stream.ReadMultiBuffer()
		if mb.IsEmpty() && err != nil {
			return nil, err
		}
		r.leftOver = mb
	}

	mb := r.leftOver.SliceBySize(size)
	if r.leftOver.IsEmpty() {
		r.leftOver = nil
	}
	return mb, nil
}

func (r *BufferedReader) writeToInternal(writer io.Writer) (int64, error) {
	mbWriter := NewWriter(writer)
	totalBytes := int64(0)
	if r.leftOver != nil {
		totalBytes += int64(r.leftOver.Len())
		if err := mbWriter.WriteMultiBuffer(r.leftOver); err != nil {
			return 0, err
		}
		r.leftOver = nil
	}

	for {
		mb, err := r.stream.ReadMultiBuffer()
		if mb != nil {
			totalBytes += int64(mb.Len())
			if werr := mbWriter.WriteMultiBuffer(mb); werr != nil {
				return totalBytes, err
			}
		}
		if err != nil {
			return totalBytes, err
		}
	}
}

// WriteTo implements io.WriterTo.
func (r *BufferedReader) WriteTo(writer io.Writer) (int64, error) {
	nBytes, err := r.writeToInternal(writer)
	if errors.Cause(err) == io.EOF {
		return nBytes, nil
	}
	return nBytes, err
}
<|endoftext|>


<commit_before>package compressor

import (
	"archive/zip"
	"io"
	"io/ioutil"
)

type ZipCompressor struct{}

func NewZipCompressor() *ZipCompressor {
	return &ZipCompressor{}
}

func (c *ZipCompressor) Compress(compressedFile io.Writer, targetDir string, files []string) error {
	w := zip.NewWriter(compressedFile)

	for _, file := range files {
		f, err := w.Create(file)
		if err != nil {
			return err
		}

		contents, _ := ioutil.ReadFile(targetDir + "/" + file)
		_, err = f.Write(contents)
		if err != nil {
			return err
		}
	}

	if err := w.Close(); err != nil {
		return err
	}

	return nil
}
<commit_msg>ZIPロジック変更<commit_after>package compressor

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

type ZipCompressor struct{}

func NewZipCompressor() *ZipCompressor {
	return &ZipCompressor{}
}

func (c *ZipCompressor) Compress(compressedFile io.Writer, targetDir string, files []string) error {
	w := zip.NewWriter(compressedFile)

	for _, filename := range files {
		filepath := fmt.Sprintf("%s/%s", targetDir, filename)
		info, err := os.Stat(filepath)
		if err != nil {
			return err
		}

		if info.IsDir() {
			continue
		}

		file, err := os.Open(filepath)
		if err != nil {
			return err
		}
		defer file.Close()

		hdr, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}

		hdr.Name = filename

		f, err := w.CreateHeader(hdr)
		if err != nil {
			return err
		}

		contents, _ := ioutil.ReadFile(filepath)
		_, err = f.Write(contents)
		if err != nil {
			return err
		}
	}

	if err := w.Close(); err != nil {
		return err
	}

	return nil
}
<|endoftext|>