<commit_before>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

func metricName(names ...string) string {
	return strings.Join(names, ".")
}

// FluentdPlugin mackerel plugin for Fluentd
type FluentdPlugin struct {
	Target          string
	Prefix          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string

	plugins []FluentdPluginMetrics
}

// FluentdMetrics is alias for backward compatibility.
type FluentdMetrics = FluentdPlugin

// MetricKeyPrefix interface for PluginWithPrefix
func (f FluentdPlugin) MetricKeyPrefix() string {
	if f.Prefix == "" {
		f.Prefix = "fluentd"
	}
	return f.Prefix
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (fpm FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(fpm.EmitRecords)
	case "emit_count":
		return float64(fpm.EmitCount)
	case "write_count":
		return float64(fpm.WriteCount)
	case "rollback_count":
		return float64(fpm.RollbackCount)
	case "slow_flush_count":
		return float64(fpm.SlowFlushCount)
	case "flush_time_count":
		return float64(fpm.FlushTimeCount)
	case "buffer_stage_length":
		return float64(fpm.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(fpm.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(fpm.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return fpm.BufferAvailableBufferSpaceRatios
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdPlugin) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics[metricName("retry_count", pid)] = float64(p.RetryCount)
		metrics[metricName("buffer_queue_length", pid)] = float64(p.BufferQueueLength)
		metrics[metricName("buffer_total_queued_size", pid)] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			metrics[metricName(name, pid)] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdPlugin) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

// FetchMetrics interface for mackerelplugin
func (f FluentdPlugin) FetchMetrics() (map[string]interface{}, error) {
	req, err := http.NewRequest(http.MethodGet, f.Target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

var defaultGraphs = map[string]mp.Graphs{
	"retry_count": {
		Label: "retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_length": {
		Label: "queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_total_queued_size": {
		Label: "buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"emit_records": {
		Label: "emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"emit_count": {
		Label: "emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"write_count": {
		Label: "write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"rollback_count": {
		Label: "rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"slow_flush_count": {
		Label: "slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"flush_time_count": {
		Label: "buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"buffer_stage_length": {
		Label: "length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_stage_byte_size": {
		Label: "bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_byte_size": {
		Label: "bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_available_buffer_space_ratios": {
		Label: "available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(f.Prefix)
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		graphs[key] = mp.Graphs{
			Label:   (labelPrefix + " " + g.Label),
			Unit:    g.Unit,
			Metrics: g.Metrics,
		}
	}
	for _, name := range f.extendedMetrics {
		fullName := metricName(name)
		if g, ok := extendedGraphs[fullName]; ok {
			graphs[fullName] = mp.Graphs{
				Label:   (labelPrefix + " " + g.Label),
				Unit:    g.Unit,
				Metrics: g.Metrics,
			}
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	prefix := flag.String("metric-key-prefix", "fluentd", "Metric key prefix")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, key)
		}
	case "":
	default:
		for _, name := range strings.Split(*extendedMetricNames, ",") {
			fullName := metricName(name)
			if _, exists := extendedGraphs[fullName]; !exists {
				fmt.Fprintf(os.Stderr, "extended_metrics %s is not supported. See also https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released\n", name)
				os.Exit(1)
			}
			extendedMetrics = append(extendedMetrics, name)
		}
	}
	f := FluentdPlugin{
		Target:          fmt.Sprintf("http://%s:%s/api/plugins.json", *host, *port),
		Prefix:          *prefix,
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<commit_msg>[fluentd] add workers option.<commit_after>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

func metricName(names ...string) string {
	return strings.Join(names, ".")
}

// FluentdPlugin mackerel plugin for Fluentd
type FluentdPlugin struct {
	Host            string
	Port            string
	Prefix          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string
	Workers         uint

	plugins []FluentdPluginMetrics
}

// FluentdMetrics is alias for backward compatibility.
type FluentdMetrics = FluentdPlugin

// MetricKeyPrefix interface for PluginWithPrefix
func (f FluentdPlugin) MetricKeyPrefix() string {
	if f.Prefix == "" {
		f.Prefix = "fluentd"
	}
	return f.Prefix
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (fpm FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(fpm.EmitRecords)
	case "emit_count":
		return float64(fpm.EmitCount)
	case "write_count":
		return float64(fpm.WriteCount)
	case "rollback_count":
		return float64(fpm.RollbackCount)
	case "slow_flush_count":
		return float64(fpm.SlowFlushCount)
	case "flush_time_count":
		return float64(fpm.FlushTimeCount)
	case "buffer_stage_length":
		return float64(fpm.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(fpm.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(fpm.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return fpm.BufferAvailableBufferSpaceRatios
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdPlugin) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics[metricName("retry_count", pid)] = float64(p.RetryCount)
		metrics[metricName("buffer_queue_length", pid)] = float64(p.BufferQueueLength)
		metrics[metricName("buffer_total_queued_size", pid)] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			metrics[metricName(name, pid)] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdPlugin) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

func (f *FluentdPlugin) fetchFluentdMetrics(host string, port int) (map[string]interface{}, error) {
	target := fmt.Sprintf("http://%s:%d/api/plugins.json", host, port)
	req, err := http.NewRequest(http.MethodGet, target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

// FetchMetrics interface for mackerelplugin
func (f FluentdPlugin) FetchMetrics() (map[string]interface{}, error) {
	port, _ := strconv.Atoi(f.Port)
	if f.Workers > 1 {
		metrics := make(map[string]interface{})
		for workerNumber := 0; workerNumber < int(f.Workers); workerNumber++ {
			m, e := f.fetchFluentdMetrics(f.Host, port+workerNumber)
			if e != nil {
				continue
			}

			workerName := fmt.Sprintf("worker%d", workerNumber)
			for k, v := range m {
				ks := strings.Split(k, ".")
				ks, last := ks[:len(ks)-1], ks[len(ks)-1]
				ks = append(ks, workerName)
				ks = append(ks, last)
				metrics[strings.Join(ks, ".")] = v
			}
		}
		if len(metrics) == 0 {
			err := fmt.Errorf("failed to connect to fluentd's monitor_agent")
			return metrics, err
		}
		return metrics, nil
	}
	return f.fetchFluentdMetrics(f.Host, port)
}

var defaultGraphs = map[string]mp.Graphs{
	"retry_count": {
		Label: "retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_length": {
		Label: "queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_total_queued_size": {
		Label: "buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"emit_records": {
		Label: "emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"emit_count": {
		Label: "emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"write_count": {
		Label: "write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"rollback_count": {
		Label: "rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"slow_flush_count": {
		Label: "slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"flush_time_count": {
		Label: "buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"buffer_stage_length": {
		Label: "length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_stage_byte_size": {
		Label: "bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_byte_size": {
		Label: "bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_available_buffer_space_ratios": {
		Label: "available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(f.Prefix)
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		graphs[key] = mp.Graphs{
			Label:   (labelPrefix + " " + g.Label),
			Unit:    g.Unit,
			Metrics: g.Metrics,
		}
	}
	for _, name := range f.extendedMetrics {
		fullName := metricName(name)
		if g, ok := extendedGraphs[fullName]; ok {
			graphs[fullName] = mp.Graphs{
				Label:   (labelPrefix + " " + g.Label),
				Unit:    g.Unit,
				Metrics: g.Metrics,
			}
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	prefix := flag.String("metric-key-prefix", "fluentd", "Metric key prefix")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	workers := flag.Uint("workers", 1, "specifying the number of Fluentd's multi-process workers")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, key)
		}
	case "":
	default:
		for _, name := range strings.Split(*extendedMetricNames, ",") {
			fullName := metricName(name)
			if _, exists := extendedGraphs[fullName]; !exists {
				fmt.Fprintf(os.Stderr, "extended_metrics %s is not supported. See also https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released\n", name)
				os.Exit(1)
			}
			extendedMetrics = append(extendedMetrics, name)
		}
	}
	f := FluentdPlugin{
		Host:            *host,
		Port:            *port,
		Prefix:          *prefix,
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
		Workers:         *workers,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<|endoftext|>


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

package common

import (
	"strconv"
	"strings"

	"istio.io/istio/pkg/config/protocol"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/components/echo/echoboot"
	"istio.io/istio/pkg/test/framework/components/istio"
	"istio.io/istio/pkg/test/framework/components/istio/ingress"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/resource"
	"istio.io/istio/pkg/test/util/tmpl"
)

type EchoDeployments struct {
	// Namespace echo apps will be deployed
	Namespace namespace.Instance
	// Namespace where external echo app will be deployed
	ExternalNamespace namespace.Instance

	// Ingressgateway instance
	Ingress ingress.Instance

	// Standard echo app to be used by tests
	PodA echo.Instances
	// Standard echo app to be used by tests
	PodB echo.Instances
	// Standard echo app to be used by tests
	PodC echo.Instances
	// Standard echo app with TPROXY interception mode to be used by tests
	PodTproxy echo.Instances
	// Headless echo app to be used by tests
	Headless echo.Instances
	// Echo app to be used by tests, with no sidecar injected
	Naked echo.Instances
	// A virtual machine echo app (only deployed to one cluster)
	VM echo.Instances

	// Echo app to be used by tests, with no sidecar injected
	External echo.Instances

	All echo.Instances
}

const (
	PodASvc      = "a"
	PodBSvc      = "b"
	PodCSvc      = "c"
	PodTproxySvc = "tproxy"
	VMSvc        = "vm"
	HeadlessSvc  = "headless"
	NakedSvc     = "naked"
	ExternalSvc  = "external"

	externalHostname = "fake.external.com"
)

var EchoPorts = []echo.Port{
	{Name: "http", Protocol: protocol.HTTP, ServicePort: 80, InstancePort: 18080},
	{Name: "grpc", Protocol: protocol.GRPC, ServicePort: 7070, InstancePort: 17070},
	{Name: "tcp", Protocol: protocol.TCP, ServicePort: 9090, InstancePort: 19090},
	{Name: "https", Protocol: protocol.HTTPS, ServicePort: 443, InstancePort: 18443, TLS: true},
	{Name: "tcp-server", Protocol: protocol.TCP, ServicePort: 9091, InstancePort: 16060, ServerFirst: true},
	{Name: "auto-tcp", Protocol: protocol.TCP, ServicePort: 9092, InstancePort: 19091},
	{Name: "auto-tcp-server", Protocol: protocol.TCP, ServicePort: 9093, InstancePort: 16061, ServerFirst: true},
	{Name: "auto-http", Protocol: protocol.HTTP, ServicePort: 81, InstancePort: 18081},
	{Name: "auto-grpc", Protocol: protocol.GRPC, ServicePort: 7071, InstancePort: 17071},
	{Name: "auto-https", Protocol: protocol.HTTPS, ServicePort: 9443, InstancePort: 19443},
	{Name: "http-instance", Protocol: protocol.HTTP, ServicePort: 82, InstancePort: 18082, InstanceIP: true},
	{Name: "http-localhost", Protocol: protocol.HTTP, ServicePort: 84, InstancePort: 18084, LocalhostIP: true},
}

var WorkloadPorts = []echo.WorkloadPort{
	{Protocol: protocol.TCP, Port: 19092},
	{Protocol: protocol.HTTP, Port: 18083},
}

func FindPortByName(name string) echo.Port {
	for _, p := range EchoPorts {
		if p.Name == name {
			return p
		}
	}
	return echo.Port{}
}

func serviceEntryPorts() []echo.Port {
	res := []echo.Port{}
	for _, p := range EchoPorts {
		if strings.HasPrefix(p.Name, "auto") {
			// The protocol needs to be set in EchoPorts to configure the echo deployment
			// But for service entry, we want to ensure we set it to "" which will use sniffing
			p.Protocol = ""
		}
		res = append(res, p)
	}
	return res
}

func SetupApps(t resource.Context, i istio.Instance, apps *EchoDeployments) error {
	var err error
	apps.Namespace, err = namespace.New(t, namespace.Config{
		Prefix: "echo",
		Inject: true,
	})
	if err != nil {
		return err
	}
	apps.ExternalNamespace, err = namespace.New(t, namespace.Config{
		Prefix: "external",
		Inject: false,
	})
	if err != nil {
		return err
	}

	apps.Ingress = i.IngressFor(t.Clusters().Default())

	// Headless services don't work with targetPort, set to same port
	headlessPorts := make([]echo.Port, len(EchoPorts))
	for i, p := range EchoPorts {
		p.ServicePort = p.InstancePort
		headlessPorts[i] = p
	}
	builder := echoboot.NewBuilder(t).
		WithClusters(t.Clusters()...).
		WithConfig(echo.Config{
			Service:           PodASvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			Subsets:           []echo.SubsetConfig{{}},
			Locality:          "region.zone.subzone",
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           PodBSvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           PodCSvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           HeadlessSvc,
			Headless:          true,
			Namespace:         apps.Namespace,
			Ports:             headlessPorts,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:   NakedSvc,
			Namespace: apps.Namespace,
			Ports:     EchoPorts,
			Subsets: []echo.SubsetConfig{
				{
					Annotations: map[echo.Annotation]*echo.AnnotationValue{
						echo.SidecarInject: {
							Value: strconv.FormatBool(false),
						},
					},
				},
			},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           ExternalSvc,
			Namespace:         apps.ExternalNamespace,
			DefaultHostHeader: externalHostname,
			Ports:             EchoPorts,
			Subsets: []echo.SubsetConfig{
				{
					Annotations: map[echo.Annotation]*echo.AnnotationValue{
						echo.SidecarInject: {
							Value: strconv.FormatBool(false),
						},
					},
				},
			},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:   PodTproxySvc,
			Namespace: apps.Namespace,
			Ports:     EchoPorts,
			Subsets: []echo.SubsetConfig{{
				Annotations: echo.NewAnnotations().Set(echo.SidecarInterceptionMode, "TPROXY"),
			}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           VMSvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			DeployAsVM:        true,
			AutoRegisterVM:    true,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		})

	echos, err := builder.Build()
	if err != nil {
		return err
	}
	apps.All = echos
	apps.PodA = echos.Match(echo.Service(PodASvc))
	apps.PodB = echos.Match(echo.Service(PodBSvc))
	apps.PodC = echos.Match(echo.Service(PodCSvc))
	apps.PodTproxy = echos.Match(echo.Service(PodTproxySvc))
	apps.Headless = echos.Match(echo.Service(HeadlessSvc))
	apps.Naked = echos.Match(echo.Service(NakedSvc))
	apps.External = echos.Match(echo.Service(ExternalSvc))
	if !t.Settings().SkipVM {
		apps.VM = echos.Match(echo.Service(VMSvc))
	}

	if err := t.Config().ApplyYAML(apps.Namespace.Name(), `
apiVersion: networking.istio.io/v1alpha3
kind: Sidecar
metadata:
  name: restrict-to-namespace
spec:
  egress:
  - hosts:
    - "./*"
    - "istio-system/*"
`); err != nil {
		return err
	}

	se, err := tmpl.Evaluate(`apiVersion: networking.istio.io/v1alpha3
kind: ServiceEntry
metadata:
  name: external-service
spec:
  hosts:
  - {{.Hostname}}
  location: MESH_EXTERNAL
  resolution: DNS
  endpoints:
  - address: external.{{.Namespace}}.svc.cluster.local
  ports:
  - name: http-tls-origination
    number: 8888
    protocol: http
    targetPort: 443
  - name: http2-tls-origination
    number: 8882
    protocol: http2
    targetPort: 443
{{- range $i, $p := .Ports }}
  - name: {{$p.Name}}
    number: {{$p.ServicePort}}
    protocol: "{{$p.Protocol}}"
{{- end }}
`, map[string]interface{}{"Namespace": apps.ExternalNamespace.Name(), "Hostname": externalHostname, "Ports": serviceEntryPorts()})
	if err != nil {
		return err
	}
	if err := t.Config().ApplyYAML(apps.Namespace.Name(), se); err != nil {
		return err
	}
	return nil
}

func (d EchoDeployments) IsMulticluster() bool {
	return d.All.Clusters().IsMulticluster()
}
<commit_msg>Fix "HTTPS" port that wasn't actually using TLS (#31792)<commit_after>// +build integ
// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package common

import (
	"strconv"
	"strings"

	"istio.io/istio/pkg/config/protocol"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/components/echo/echoboot"
	"istio.io/istio/pkg/test/framework/components/istio"
	"istio.io/istio/pkg/test/framework/components/istio/ingress"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/resource"
	"istio.io/istio/pkg/test/util/tmpl"
)

type EchoDeployments struct {
	// Namespace echo apps will be deployed
	Namespace namespace.Instance
	// Namespace where external echo app will be deployed
	ExternalNamespace namespace.Instance

	// Ingressgateway instance
	Ingress ingress.Instance

	// Standard echo app to be used by tests
	PodA echo.Instances
	// Standard echo app to be used by tests
	PodB echo.Instances
	// Standard echo app to be used by tests
	PodC echo.Instances
	// Standard echo app with TPROXY interception mode to be used by tests
	PodTproxy echo.Instances
	// Headless echo app to be used by tests
	Headless echo.Instances
	// Echo app to be used by tests, with no sidecar injected
	Naked echo.Instances
	// A virtual machine echo app (only deployed to one cluster)
	VM echo.Instances

	// Echo app to be used by tests, with no sidecar injected
	External echo.Instances

	All echo.Instances
}

const (
	PodASvc      = "a"
	PodBSvc      = "b"
	PodCSvc      = "c"
	PodTproxySvc = "tproxy"
	VMSvc        = "vm"
	HeadlessSvc  = "headless"
	NakedSvc     = "naked"
	ExternalSvc  = "external"

	externalHostname = "fake.external.com"
)

var EchoPorts = []echo.Port{
	{Name: "http", Protocol: protocol.HTTP, ServicePort: 80, InstancePort: 18080},
	{Name: "grpc", Protocol: protocol.GRPC, ServicePort: 7070, InstancePort: 17070},
	{Name: "tcp", Protocol: protocol.TCP, ServicePort: 9090, InstancePort: 19090},
	{Name: "https", Protocol: protocol.HTTPS, ServicePort: 443, InstancePort: 18443, TLS: true},
	{Name: "tcp-server", Protocol: protocol.TCP, ServicePort: 9091, InstancePort: 16060, ServerFirst: true},
	{Name: "auto-tcp", Protocol: protocol.TCP, ServicePort: 9092, InstancePort: 19091},
	{Name: "auto-tcp-server", Protocol: protocol.TCP, ServicePort: 9093, InstancePort: 16061, ServerFirst: true},
	{Name: "auto-http", Protocol: protocol.HTTP, ServicePort: 81, InstancePort: 18081},
	{Name: "auto-grpc", Protocol: protocol.GRPC, ServicePort: 7071, InstancePort: 17071},
	{Name: "auto-https", Protocol: protocol.HTTPS, ServicePort: 9443, InstancePort: 19443, TLS: true},
	{Name: "http-instance", Protocol: protocol.HTTP, ServicePort: 82, InstancePort: 18082, InstanceIP: true},
	{Name: "http-localhost", Protocol: protocol.HTTP, ServicePort: 84, InstancePort: 18084, LocalhostIP: true},
}

var WorkloadPorts = []echo.WorkloadPort{
	{Protocol: protocol.TCP, Port: 19092},
	{Protocol: protocol.HTTP, Port: 18083},
}

func FindPortByName(name string) echo.Port {
	for _, p := range EchoPorts {
		if p.Name == name {
			return p
		}
	}
	return echo.Port{}
}

func serviceEntryPorts() []echo.Port {
	res := []echo.Port{}
	for _, p := range EchoPorts {
		if strings.HasPrefix(p.Name, "auto") {
			// The protocol needs to be set in EchoPorts to configure the echo deployment
			// But for service entry, we want to ensure we set it to "" which will use sniffing
			p.Protocol = ""
		}
		res = append(res, p)
	}
	return res
}

func SetupApps(t resource.Context, i istio.Instance, apps *EchoDeployments) error {
	var err error
	apps.Namespace, err = namespace.New(t, namespace.Config{
		Prefix: "echo",
		Inject: true,
	})
	if err != nil {
		return err
	}
	apps.ExternalNamespace, err = namespace.New(t, namespace.Config{
		Prefix: "external",
		Inject: false,
	})
	if err != nil {
		return err
	}

	apps.Ingress = i.IngressFor(t.Clusters().Default())

	// Headless services don't work with targetPort, set to same port
	headlessPorts := make([]echo.Port, len(EchoPorts))
	for i, p := range EchoPorts {
		p.ServicePort = p.InstancePort
		headlessPorts[i] = p
	}
	builder := echoboot.NewBuilder(t).
		WithClusters(t.Clusters()...).
		WithConfig(echo.Config{
			Service:           PodASvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			Subsets:           []echo.SubsetConfig{{}},
			Locality:          "region.zone.subzone",
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           PodBSvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           PodCSvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           HeadlessSvc,
			Headless:          true,
			Namespace:         apps.Namespace,
			Ports:             headlessPorts,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:   NakedSvc,
			Namespace: apps.Namespace,
			Ports:     EchoPorts,
			Subsets: []echo.SubsetConfig{
				{
					Annotations: map[echo.Annotation]*echo.AnnotationValue{
						echo.SidecarInject: {
							Value: strconv.FormatBool(false),
						},
					},
				},
			},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           ExternalSvc,
			Namespace:         apps.ExternalNamespace,
			DefaultHostHeader: externalHostname,
			Ports:             EchoPorts,
			Subsets: []echo.SubsetConfig{
				{
					Annotations: map[echo.Annotation]*echo.AnnotationValue{
						echo.SidecarInject: {
							Value: strconv.FormatBool(false),
						},
					},
				},
			},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:   PodTproxySvc,
			Namespace: apps.Namespace,
			Ports:     EchoPorts,
			Subsets: []echo.SubsetConfig{{
				Annotations: echo.NewAnnotations().Set(echo.SidecarInterceptionMode, "TPROXY"),
			}},
			WorkloadOnlyPorts: WorkloadPorts,
		}).
		WithConfig(echo.Config{
			Service:           VMSvc,
			Namespace:         apps.Namespace,
			Ports:             EchoPorts,
			DeployAsVM:        true,
			AutoRegisterVM:    true,
			Subsets:           []echo.SubsetConfig{{}},
			WorkloadOnlyPorts: WorkloadPorts,
		})

	echos, err := builder.Build()
	if err != nil {
		return err
	}
	apps.All = echos
	apps.PodA = echos.Match(echo.Service(PodASvc))
	apps.PodB = echos.Match(echo.Service(PodBSvc))
	apps.PodC = echos.Match(echo.Service(PodCSvc))
	apps.PodTproxy = echos.Match(echo.Service(PodTproxySvc))
	apps.Headless = echos.Match(echo.Service(HeadlessSvc))
	apps.Naked = echos.Match(echo.Service(NakedSvc))
	apps.External = echos.Match(echo.Service(ExternalSvc))
	if !t.Settings().SkipVM {
		apps.VM = echos.Match(echo.Service(VMSvc))
	}

	if err := t.Config().ApplyYAML(apps.Namespace.Name(), `
apiVersion: networking.istio.io/v1alpha3
kind: Sidecar
metadata:
  name: restrict-to-namespace
spec:
  egress:
  - hosts:
    - "./*"
    - "istio-system/*"
`); err != nil {
		return err
	}

	se, err := tmpl.Evaluate(`apiVersion: networking.istio.io/v1alpha3
kind: ServiceEntry
metadata:
  name: external-service
spec:
  hosts:
  - {{.Hostname}}
  location: MESH_EXTERNAL
  resolution: DNS
  endpoints:
  - address: external.{{.Namespace}}.svc.cluster.local
  ports:
  - name: http-tls-origination
    number: 8888
    protocol: http
    targetPort: 443
  - name: http2-tls-origination
    number: 8882
    protocol: http2
    targetPort: 443
{{- range $i, $p := .Ports }}
  - name: {{$p.Name}}
    number: {{$p.ServicePort}}
    protocol: "{{$p.Protocol}}"
{{- end }}
`, map[string]interface{}{"Namespace": apps.ExternalNamespace.Name(), "Hostname": externalHostname, "Ports": serviceEntryPorts()})
	if err != nil {
		return err
	}
	if err := t.Config().ApplyYAML(apps.Namespace.Name(), se); err != nil {
		return err
	}
	return nil
}

func (d EchoDeployments) IsMulticluster() bool {
	return d.All.Clusters().IsMulticluster()
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package pd

import (
	"path"
	"strconv"
	"sync"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"golang.org/x/net/context"
)

const (
	pdRootPath        = "/pd"
	requestTimeout    = 3 * time.Second
	maxRetryGetLeader = 100
)

// Client is a PD (Placement Driver) client.
// It should not be used after calling Close().
type Client interface {
	// GetTS gets a timestamp from PD.
	GetTS() (int64, int64, error)
	// GetRegion gets a region and its leader Peer from PD by key.
	// The region may expire after split. Caller is responsible for caching and
	// taking care of region change.
	// Also it may return nil if PD finds no Region for the key temporarily,
	// client should retry later.
	GetRegion(key []byte) (*metapb.Region, *metapb.Peer, error)
	// GetStore gets a store from PD by store id.
	// The store may expire later. Caller is responsible for caching and taking care
	// of store change.
	GetStore(storeID uint64) (*metapb.Store, error)
	// Close closes the client.
	Close()
}

type client struct {
	clusterID   uint64
	etcdClient  *clientv3.Client
	workerMutex sync.RWMutex
	worker      *rpcWorker
	wg          sync.WaitGroup
	quit        chan struct{}
}

func getLeaderPath(clusterID uint64) string {
	return path.Join(pdRootPath, strconv.FormatUint(clusterID, 10), "leader")
}

// NewClient creates a PD client.
func NewClient(etcdAddrs []string, clusterID uint64) (Client, error) {
	log.Infof("[pd] create etcd client with endpoints %v", etcdAddrs)
	etcdClient, err := clientv3.New(clientv3.Config{
		Endpoints:   etcdAddrs,
		DialTimeout: requestTimeout,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}
	leaderPath := getLeaderPath(clusterID)

	var (
		leaderAddr string
		revision   int64
	)

	for i := 0; i < maxRetryGetLeader; i++ {
		leaderAddr, revision, err = getLeader(etcdClient, leaderPath)
		if err == nil {
			break
		}

		time.Sleep(50 * time.Millisecond)
	}

	if err != nil {
		return nil, errors.Trace(err)
	}

	client := &client{
		clusterID:  clusterID,
		etcdClient: etcdClient,
		worker:     newRPCWorker(leaderAddr, clusterID),
		quit:       make(chan struct{}),
	}

	client.wg.Add(1)
	go client.watchLeader(leaderPath, revision)

	return client, nil
}

func (c *client) Close() {
	c.etcdClient.Close()

	close(c.quit)
	// Must wait watchLeader done.
	c.wg.Wait()
	c.worker.stop(errors.New("[pd] pd-client closing"))
}

func (c *client) GetTS() (int64, int64, error) {
	req := &tsoRequest{
		done: make(chan error, 1),
	}
	c.workerMutex.RLock()
	c.worker.requests <- req
	c.workerMutex.RUnlock()
	err := <-req.done
	return req.physical, req.logical, err
}

func (c *client) GetRegion(key []byte) (*metapb.Region, *metapb.Peer, error) {
	req := &regionRequest{
		pbReq: &pdpb.GetRegionRequest{
			RegionKey: key,
		},
		done: make(chan error, 1),
	}
	c.workerMutex.RLock()
	c.worker.requests <- req
	c.workerMutex.RUnlock()
	err := <-req.done
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	return req.pbResp.GetRegion(), req.pbResp.GetLeader(), nil
}

func (c *client) GetStore(storeID uint64) (*metapb.Store, error) {
	req := &storeRequest{
		pbReq: &pdpb.GetStoreRequest{
			StoreId: proto.Uint64(storeID),
		},
		done: make(chan error, 1),
	}
	c.workerMutex.RLock()
	c.worker.requests <- req
	c.workerMutex.RUnlock()
	err := <-req.done
	if err != nil {
		return nil, errors.Trace(err)
	}
	store := req.pbResp.GetStore()
	if store == nil {
		return nil, errors.New("[pd] store field in rpc response not set")
	}
	return store, nil
}

// Use var here is for test changing.
// TODO: refactor this after etcd fixes https://github.com/coreos/etcd/issues/5985
var defaultWatchLeaderTimeout = 30 * time.Second

func (c *client) watchLeader(leaderPath string, revision int64) {
	defer c.wg.Done()

	for {
		log.Infof("[pd] start watch pd leader on path %v, revision %v", leaderPath, revision)
		ctx, cancel := context.WithTimeout(c.etcdClient.Ctx(), defaultWatchLeaderTimeout)
		rch := c.etcdClient.Watch(ctx, leaderPath, clientv3.WithRev(revision))

		for resp := range rch {
			if resp.Canceled {
				log.Warn("[pd] leader watcher canceled")
				break
			}

			// We don't watch any changed, no need to check leader again.
			if len(resp.Events) == 0 {
				break
			}

			leaderAddr, rev, err := getLeader(c.etcdClient, leaderPath)
			if err != nil {
				log.Warn(err)
				break
			}

			log.Infof("[pd] found new pd-server leader addr: %v", leaderAddr)
			c.workerMutex.Lock()
			c.worker.stop(errors.New("[pd] leader change"))
			c.worker = newRPCWorker(leaderAddr, c.clusterID)
			c.workerMutex.Unlock()
			revision = rev
		}

		cancel()

		select {
		case <-c.quit:
			return
		default:
		}
	}
}

func getLeader(etcdClient *clientv3.Client, path string) (string, int64, error) {
	kv := clientv3.NewKV(etcdClient)
	ctx, cancel := context.WithTimeout(etcdClient.Ctx(), requestTimeout)
	resp, err := kv.Get(ctx, path)
	cancel()
	if err != nil {
		return "", 0, errors.Trace(err)
	}
	if len(resp.Kvs) != 1 {
		return "", 0, errors.Errorf("invalid getLeader resp: %v", resp)
	}

	var leader pdpb.Leader
	if err = proto.Unmarshal(resp.Kvs[0].Value, &leader); err != nil {
		return "", 0, errors.Trace(err)
	}
	return leader.GetAddr(), resp.Header.Revision, nil
}
<commit_msg>pd-client: update pd client connect and request timeout. (#255)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package pd

import (
	"path"
	"strconv"
	"sync"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"golang.org/x/net/context"
)

const (
	pdRootPath        = "/pd"
	requestTimeout    = 10 * time.Second
	connectTimeout    = 30 * time.Second
	maxRetryGetLeader = 100
)

// Client is a PD (Placement Driver) client.
// It should not be used after calling Close().
type Client interface {
	// GetTS gets a timestamp from PD.
	GetTS() (int64, int64, error)
	// GetRegion gets a region and its leader Peer from PD by key.
	// The region may expire after split. Caller is responsible for caching and
	// taking care of region change.
	// Also it may return nil if PD finds no Region for the key temporarily,
	// client should retry later.
	GetRegion(key []byte) (*metapb.Region, *metapb.Peer, error)
	// GetStore gets a store from PD by store id.
	// The store may expire later. Caller is responsible for caching and taking care
	// of store change.
	GetStore(storeID uint64) (*metapb.Store, error)
	// Close closes the client.
	Close()
}

type client struct {
	clusterID   uint64
	etcdClient  *clientv3.Client
	workerMutex sync.RWMutex
	worker      *rpcWorker
	wg          sync.WaitGroup
	quit        chan struct{}
}

func getLeaderPath(clusterID uint64) string {
	return path.Join(pdRootPath, strconv.FormatUint(clusterID, 10), "leader")
}

// NewClient creates a PD client.
func NewClient(etcdAddrs []string, clusterID uint64) (Client, error) {
	log.Infof("[pd] create etcd client with endpoints %v", etcdAddrs)
	etcdClient, err := clientv3.New(clientv3.Config{
		Endpoints:   etcdAddrs,
		DialTimeout: connectTimeout,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}
	leaderPath := getLeaderPath(clusterID)

	var (
		leaderAddr string
		revision   int64
	)

	for i := 0; i < maxRetryGetLeader; i++ {
		leaderAddr, revision, err = getLeader(etcdClient, leaderPath)
		if err == nil {
			break
		}

		time.Sleep(50 * time.Millisecond)
	}

	if err != nil {
		return nil, errors.Trace(err)
	}

	client := &client{
		clusterID:  clusterID,
		etcdClient: etcdClient,
		worker:     newRPCWorker(leaderAddr, clusterID),
		quit:       make(chan struct{}),
	}

	client.wg.Add(1)
	go client.watchLeader(leaderPath, revision)

	return client, nil
}

func (c *client) Close() {
	c.etcdClient.Close()

	close(c.quit)
	// Must wait watchLeader done.
	c.wg.Wait()
	c.worker.stop(errors.New("[pd] pd-client closing"))
}

func (c *client) GetTS() (int64, int64, error) {
	req := &tsoRequest{
		done: make(chan error, 1),
	}
	c.workerMutex.RLock()
	c.worker.requests <- req
	c.workerMutex.RUnlock()
	err := <-req.done
	return req.physical, req.logical, err
}

func (c *client) GetRegion(key []byte) (*metapb.Region, *metapb.Peer, error) {
	req := &regionRequest{
		pbReq: &pdpb.GetRegionRequest{
			RegionKey: key,
		},
		done: make(chan error, 1),
	}
	c.workerMutex.RLock()
	c.worker.requests <- req
	c.workerMutex.RUnlock()
	err := <-req.done
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	return req.pbResp.GetRegion(), req.pbResp.GetLeader(), nil
}

func (c *client) GetStore(storeID uint64) (*metapb.Store, error) {
	req := &storeRequest{
		pbReq: &pdpb.GetStoreRequest{
			StoreId: proto.Uint64(storeID),
		},
		done: make(chan error, 1),
	}
	c.workerMutex.RLock()
	c.worker.requests <- req
	c.workerMutex.RUnlock()
	err := <-req.done
	if err != nil {
		return nil, errors.Trace(err)
	}
	store := req.pbResp.GetStore()
	if store == nil {
		return nil, errors.New("[pd] store field in rpc response not set")
	}
	return store, nil
}

// Use var here is for test changing.
// TODO: refactor this after etcd fixes https://github.com/coreos/etcd/issues/5985
var defaultWatchLeaderTimeout = 30 * time.Second

func (c *client) watchLeader(leaderPath string, revision int64) {
	defer c.wg.Done()

	for {
		log.Infof("[pd] start watch pd leader on path %v, revision %v", leaderPath, revision)
		ctx, cancel := context.WithTimeout(c.etcdClient.Ctx(), defaultWatchLeaderTimeout)
		rch := c.etcdClient.Watch(ctx, leaderPath, clientv3.WithRev(revision))

		for resp := range rch {
			if resp.Canceled {
				log.Warn("[pd] leader watcher canceled")
				break
			}

			// We don't watch any changed, no need to check leader again.
			if len(resp.Events) == 0 {
				break
			}

			leaderAddr, rev, err := getLeader(c.etcdClient, leaderPath)
			if err != nil {
				log.Warn(err)
				break
			}

			log.Infof("[pd] found new pd-server leader addr: %v", leaderAddr)
			c.workerMutex.Lock()
			c.worker.stop(errors.New("[pd] leader change"))
			c.worker = newRPCWorker(leaderAddr, c.clusterID)
			c.workerMutex.Unlock()
			revision = rev
		}

		cancel()

		select {
		case <-c.quit:
			return
		default:
		}
	}
}

func getLeader(etcdClient *clientv3.Client, path string) (string, int64, error) {
	kv := clientv3.NewKV(etcdClient)
	ctx, cancel := context.WithTimeout(etcdClient.Ctx(), requestTimeout)
	resp, err := kv.Get(ctx, path)
	cancel()
	if err != nil {
		return "", 0, errors.Trace(err)
	}
	if len(resp.Kvs) != 1 {
		return "", 0, errors.Errorf("invalid getLeader resp: %v", resp)
	}

	var leader pdpb.Leader
	if err = proto.Unmarshal(resp.Kvs[0].Value, &leader); err != nil {
		return "", 0, errors.Trace(err)
	}
	return leader.GetAddr(), resp.Header.Revision, nil
}
<|endoftext|>


<commit_before>package torrent

import (
	"encoding/base64"
	"errors"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/cenkalti/rain/internal/rpctypes"
)

var errTorrentNotFound = errors.New("torrent not found")

type rpcHandler struct {
	session *Session
}

func (h *rpcHandler) Version(args struct{}, reply *string) error {
	*reply = Version
	return nil
}

func (h *rpcHandler) ListTorrents(args *rpctypes.ListTorrentsRequest, reply *rpctypes.ListTorrentsResponse) error {
	torrents := h.session.ListTorrents()
	reply.Torrents = make([]rpctypes.Torrent, 0, len(torrents))
	for _, t := range torrents {
		reply.Torrents = append(reply.Torrents, newTorrent(t))
	}
	return nil
}

func (h *rpcHandler) AddTorrent(args *rpctypes.AddTorrentRequest, reply *rpctypes.AddTorrentResponse) error {
	r := base64.NewDecoder(base64.StdEncoding, strings.NewReader(args.Torrent))
	t, err := h.session.AddTorrent(r)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func (h *rpcHandler) AddURI(args *rpctypes.AddURIRequest, reply *rpctypes.AddURIResponse) error {
	t, err := h.session.AddURI(args.URI)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func newTorrent(t *Torrent) rpctypes.Torrent {
	return rpctypes.Torrent{
		ID:       t.ID(),
		Name:     t.Name(),
		InfoHash: t.InfoHash().String(),
		Port:     t.Port(),
		AddedAt:  rpctypes.Time{Time: t.AddedAt()},
	}
}

func (h *rpcHandler) RemoveTorrent(args *rpctypes.RemoveTorrentRequest, reply *rpctypes.RemoveTorrentResponse) error {
	h.session.RemoveTorrent(args.ID)
	return nil
}

func (h *rpcHandler) GetSessionStats(args *rpctypes.GetSessionStatsRequest, reply *rpctypes.GetSessionStatsResponse) error {
	s := h.session.Stats()
	var blocklistUpdatedAt *rpctypes.Time
	if !s.BlockListLastSuccessfulUpdate.IsZero() {
		blocklistUpdatedAt = &rpctypes.Time{Time: s.BlockListLastSuccessfulUpdate}
	}
	reply.Stats = rpctypes.SessionStats{
		Torrents:                      s.Torrents,
		AvailablePorts:                s.AvailablePorts,
		BlockListRules:                s.BlockListRules,
		BlockListLastSuccessfulUpdate: blocklistUpdatedAt,
		PieceCacheItems:               s.PieceCacheItems,
		PieceCacheSize:                s.PieceCacheSize,
		PieceCacheUtilization:         s.PieceCacheUtilization,
		ActivePieceBytes:              s.ActivePieceBytes,
		TorrentsPendingRAM:            s.TorrentsPendingRAM,
	}
	return nil
}

func (h *rpcHandler) GetTorrentStats(args *rpctypes.GetTorrentStatsRequest, reply *rpctypes.GetTorrentStatsResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	s := t.Stats()
	reply.Stats = rpctypes.Stats{
		Status: torrentStatusToString(s.Status),
		Pieces: struct {
			Checked   uint32
			Have      uint32
			Missing   uint32
			Available uint32
			Total     uint32
		}{
			Checked:   s.Pieces.Checked,
			Have:      s.Pieces.Have,
			Missing:   s.Pieces.Missing,
			Available: s.Pieces.Available,
			Total:     s.Pieces.Total,
		},
		Bytes: struct {
			Total      int64
			Allocated  int64
			Completed  int64
			Incomplete int64
			Downloaded int64
			Uploaded   int64
			Wasted     int64
		}{
			Total:      s.Bytes.Total,
			Allocated:  s.Bytes.Allocated,
			Completed:  s.Bytes.Completed,
			Incomplete: s.Bytes.Incomplete,
			Downloaded: s.Bytes.Downloaded,
			Uploaded:   s.Bytes.Uploaded,
			Wasted:     s.Bytes.Wasted,
		},
		Peers: struct {
			Total    int
			Incoming int
			Outgoing int
		}{
			Total:    s.Peers.Total,
			Incoming: s.Peers.Incoming,
			Outgoing: s.Peers.Outgoing,
		},
		Handshakes: struct {
			Total    int
			Incoming int
			Outgoing int
		}{
			Total:    s.Handshakes.Total,
			Incoming: s.Handshakes.Incoming,
			Outgoing: s.Handshakes.Outgoing,
		},
		Addresses: struct {
			Total   int
			Tracker int
			DHT     int
			PEX     int
		}{
			Total:   s.Addresses.Total,
			Tracker: s.Addresses.Tracker,
			DHT:     s.Addresses.DHT,
			PEX:     s.Addresses.PEX,
		},
		Downloads: struct {
			Total   int
			Running int
			Snubbed int
			Choked  int
		}{
			Total:   s.Downloads.Total,
			Running: s.Downloads.Running,
			Snubbed: s.Downloads.Snubbed,
			Choked:  s.Downloads.Choked,
		},
		MetadataDownloads: struct {
			Total   int
			Snubbed int
			Running int
		}{
			Total:   s.MetadataDownloads.Total,
			Snubbed: s.MetadataDownloads.Snubbed,
			Running: s.MetadataDownloads.Running,
		},
		Name:        s.Name,
		Private:     s.Private,
		PieceLength: s.PieceLength,
		SeededFor:   uint(s.SeededFor / time.Second),
		Speed: struct {
			Download uint
			Upload   uint
		}{
			Download: s.Speed.Download,
			Upload:   s.Speed.Upload,
		},
	}
	if s.Error != nil {
		errStr := s.Error.Error()
		reply.Stats.Error = &errStr
	}
	if s.ETA != nil {
		eta := uint(*s.ETA / time.Second)
		reply.Stats.ETA = &eta
	}
	return nil
}

func (h *rpcHandler) GetTorrentTrackers(args *rpctypes.GetTorrentTrackersRequest, reply *rpctypes.GetTorrentTrackersResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	trackers := t.Trackers()
	reply.Trackers = make([]rpctypes.Tracker, len(trackers))
	for i, t := range trackers {
		reply.Trackers[i] = rpctypes.Tracker{
			URL:      t.URL,
			Status:   trackerStatusToString(t.Status),
			Leechers: t.Leechers,
			Seeders:  t.Seeders,
		}
		if t.Error != nil {
			errStr := t.Error.Error()
			reply.Trackers[i].Error = &errStr
		}
	}
	return nil
}

func (h *rpcHandler) GetTorrentPeers(args *rpctypes.GetTorrentPeersRequest, reply *rpctypes.GetTorrentPeersResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	peers := t.Peers()
	reply.Peers = make([]rpctypes.Peer, len(peers))
	for i, p := range peers {
		var source string
		switch p.Source {
		case SourceTracker:
			source = "TRACKER"
		case SourceDHT:
			source = "DHT"
		case SourcePEX:
			source = "PEX"
		case SourceIncoming:
			source = "INCOMING"
		default:
			panic("unhandled peer source")
		}
		reply.Peers[i] = rpctypes.Peer{
			ID:                 string(p.ID[:]),
			Addr:               p.Addr.String(),
			Source:             source,
			ConnectedAt:        rpctypes.Time{Time: p.ConnectedAt.UTC()},
			Downloading:        p.Downloading,
			ClientInterested:   p.ClientInterested,
			ClientChoking:      p.ClientChoking,
			PeerInterested:     p.PeerInterested,
			PeerChoking:        p.PeerChoking,
			OptimisticUnchoked: p.OptimisticUnchoked,
			Snubbed:            p.Snubbed,
			EncryptedHandshake: p.EncryptedHandshake,
			EncryptedStream:    p.EncryptedStream,
			DownloadSpeed:      p.DownloadSpeed,
			UploadSpeed:        p.UploadSpeed,
		}
	}
	return nil
}

func (h *rpcHandler) StartTorrent(args *rpctypes.StartTorrentRequest, reply *rpctypes.StartTorrentResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	t.Start()
	return nil
}

func (h *rpcHandler) StopTorrent(args *rpctypes.StopTorrentRequest, reply *rpctypes.StopTorrentResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	t.Stop()
	return nil
}

func (h *rpcHandler) AddPeer(args *rpctypes.AddPeerRequest, reply *rpctypes.AddPeerResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	host, portString, err := net.SplitHostPort(args.Addr)
	ip := net.ParseIP(host).To4()
	if ip == nil {
		return errors.New("invalid v4 IP")
	}
	port, err := strconv.Atoi(portString)
	if err != nil {
		return err
	}
	if port == 0 {
		return errors.New("invalid port")
	}
	addr := &net.TCPAddr{
		IP:   ip,
		Port: port,
	}
	t.AddPeer(addr)
	return nil
}

func (h *rpcHandler) AddTracker(args *rpctypes.AddTrackerRequest, reply *rpctypes.AddTrackerResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	return t.AddTracker(args.URL)
}
<commit_msg>add missing methods to rpc handler<commit_after>package torrent

import (
	"encoding/base64"
	"errors"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/cenkalti/rain/internal/rpctypes"
)

var errTorrentNotFound = errors.New("torrent not found")

type rpcHandler struct {
	session *Session
}

func (h *rpcHandler) Version(args struct{}, reply *string) error {
	*reply = Version
	return nil
}

func (h *rpcHandler) ListTorrents(args *rpctypes.ListTorrentsRequest, reply *rpctypes.ListTorrentsResponse) error {
	torrents := h.session.ListTorrents()
	reply.Torrents = make([]rpctypes.Torrent, 0, len(torrents))
	for _, t := range torrents {
		reply.Torrents = append(reply.Torrents, newTorrent(t))
	}
	return nil
}

func (h *rpcHandler) AddTorrent(args *rpctypes.AddTorrentRequest, reply *rpctypes.AddTorrentResponse) error {
	r := base64.NewDecoder(base64.StdEncoding, strings.NewReader(args.Torrent))
	t, err := h.session.AddTorrent(r)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func (h *rpcHandler) AddURI(args *rpctypes.AddURIRequest, reply *rpctypes.AddURIResponse) error {
	t, err := h.session.AddURI(args.URI)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func newTorrent(t *Torrent) rpctypes.Torrent {
	return rpctypes.Torrent{
		ID:       t.ID(),
		Name:     t.Name(),
		InfoHash: t.InfoHash().String(),
		Port:     t.Port(),
		AddedAt:  rpctypes.Time{Time: t.AddedAt()},
	}
}

func (h *rpcHandler) RemoveTorrent(args *rpctypes.RemoveTorrentRequest, reply *rpctypes.RemoveTorrentResponse) error {
	h.session.RemoveTorrent(args.ID)
	return nil
}

func (h *rpcHandler) GetSessionStats(args *rpctypes.GetSessionStatsRequest, reply *rpctypes.GetSessionStatsResponse) error {
	s := h.session.Stats()
	var blocklistUpdatedAt *rpctypes.Time
	if !s.BlockListLastSuccessfulUpdate.IsZero() {
		blocklistUpdatedAt = &rpctypes.Time{Time: s.BlockListLastSuccessfulUpdate}
	}
	reply.Stats = rpctypes.SessionStats{
		Torrents:                      s.Torrents,
		AvailablePorts:                s.AvailablePorts,
		BlockListRules:                s.BlockListRules,
		BlockListLastSuccessfulUpdate: blocklistUpdatedAt,
		PieceCacheItems:               s.PieceCacheItems,
		PieceCacheSize:                s.PieceCacheSize,
		PieceCacheUtilization:         s.PieceCacheUtilization,
		ActivePieceBytes:              s.ActivePieceBytes,
		TorrentsPendingRAM:            s.TorrentsPendingRAM,
	}
	return nil
}

func (h *rpcHandler) GetTorrentStats(args *rpctypes.GetTorrentStatsRequest, reply *rpctypes.GetTorrentStatsResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	s := t.Stats()
	reply.Stats = rpctypes.Stats{
		Status: torrentStatusToString(s.Status),
		Pieces: struct {
			Checked   uint32
			Have      uint32
			Missing   uint32
			Available uint32
			Total     uint32
		}{
			Checked:   s.Pieces.Checked,
			Have:      s.Pieces.Have,
			Missing:   s.Pieces.Missing,
			Available: s.Pieces.Available,
			Total:     s.Pieces.Total,
		},
		Bytes: struct {
			Total      int64
			Allocated  int64
			Completed  int64
			Incomplete int64
			Downloaded int64
			Uploaded   int64
			Wasted     int64
		}{
			Total:      s.Bytes.Total,
			Allocated:  s.Bytes.Allocated,
			Completed:  s.Bytes.Completed,
			Incomplete: s.Bytes.Incomplete,
			Downloaded: s.Bytes.Downloaded,
			Uploaded:   s.Bytes.Uploaded,
			Wasted:     s.Bytes.Wasted,
		},
		Peers: struct {
			Total    int
			Incoming int
			Outgoing int
		}{
			Total:    s.Peers.Total,
			Incoming: s.Peers.Incoming,
			Outgoing: s.Peers.Outgoing,
		},
		Handshakes: struct {
			Total    int
			Incoming int
			Outgoing int
		}{
			Total:    s.Handshakes.Total,
			Incoming: s.Handshakes.Incoming,
			Outgoing: s.Handshakes.Outgoing,
		},
		Addresses: struct {
			Total   int
			Tracker int
			DHT     int
			PEX     int
		}{
			Total:   s.Addresses.Total,
			Tracker: s.Addresses.Tracker,
			DHT:     s.Addresses.DHT,
			PEX:     s.Addresses.PEX,
		},
		Downloads: struct {
			Total   int
			Running int
			Snubbed int
			Choked  int
		}{
			Total:   s.Downloads.Total,
			Running: s.Downloads.Running,
			Snubbed: s.Downloads.Snubbed,
			Choked:  s.Downloads.Choked,
		},
		MetadataDownloads: struct {
			Total   int
			Snubbed int
			Running int
		}{
			Total:   s.MetadataDownloads.Total,
			Snubbed: s.MetadataDownloads.Snubbed,
			Running: s.MetadataDownloads.Running,
		},
		Name:        s.Name,
		Private:     s.Private,
		PieceLength: s.PieceLength,
		SeededFor:   uint(s.SeededFor / time.Second),
		Speed: struct {
			Download uint
			Upload   uint
		}{
			Download: s.Speed.Download,
			Upload:   s.Speed.Upload,
		},
	}
	if s.Error != nil {
		errStr := s.Error.Error()
		reply.Stats.Error = &errStr
	}
	if s.ETA != nil {
		eta := uint(*s.ETA / time.Second)
		reply.Stats.ETA = &eta
	}
	return nil
}

func (h *rpcHandler) GetTorrentTrackers(args *rpctypes.GetTorrentTrackersRequest, reply *rpctypes.GetTorrentTrackersResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	trackers := t.Trackers()
	reply.Trackers = make([]rpctypes.Tracker, len(trackers))
	for i, t := range trackers {
		reply.Trackers[i] = rpctypes.Tracker{
			URL:      t.URL,
			Status:   trackerStatusToString(t.Status),
			Leechers: t.Leechers,
			Seeders:  t.Seeders,
		}
		if t.Error != nil {
			errStr := t.Error.Error()
			reply.Trackers[i].Error = &errStr
		}
	}
	return nil
}

func (h *rpcHandler) GetTorrentPeers(args *rpctypes.GetTorrentPeersRequest, reply *rpctypes.GetTorrentPeersResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	peers := t.Peers()
	reply.Peers = make([]rpctypes.Peer, len(peers))
	for i, p := range peers {
		var source string
		switch p.Source {
		case SourceTracker:
			source = "TRACKER"
		case SourceDHT:
			source = "DHT"
		case SourcePEX:
			source = "PEX"
		case SourceIncoming:
			source = "INCOMING"
		default:
			panic("unhandled peer source")
		}
		reply.Peers[i] = rpctypes.Peer{
			ID:                 string(p.ID[:]),
			Addr:               p.Addr.String(),
			Source:             source,
			ConnectedAt:        rpctypes.Time{Time: p.ConnectedAt.UTC()},
			Downloading:        p.Downloading,
			ClientInterested:   p.ClientInterested,
			ClientChoking:      p.ClientChoking,
			PeerInterested:     p.PeerInterested,
			PeerChoking:        p.PeerChoking,
			OptimisticUnchoked: p.OptimisticUnchoked,
			Snubbed:            p.Snubbed,
			EncryptedHandshake: p.EncryptedHandshake,
			EncryptedStream:    p.EncryptedStream,
			DownloadSpeed:      p.DownloadSpeed,
			UploadSpeed:        p.UploadSpeed,
		}
	}
	return nil
}

func (h *rpcHandler) StartTorrent(args *rpctypes.StartTorrentRequest, reply *rpctypes.StartTorrentResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	t.Start()
	return nil
}

func (h *rpcHandler) StopTorrent(args *rpctypes.StopTorrentRequest, reply *rpctypes.StopTorrentResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	t.Stop()
	return nil
}

func (h *rpcHandler) StartAllTorrents(args *rpctypes.StartAllTorrentsRequest, reply *rpctypes.StartAllTorrentsResponse) error {
	h.session.StartAll()
	return nil
}

func (h *rpcHandler) StopAllTorrents(args *rpctypes.StopAllTorrentsRequest, reply *rpctypes.StopAllTorrentsResponse) error {
	h.session.StopAll()
	return nil
}

func (h *rpcHandler) AddPeer(args *rpctypes.AddPeerRequest, reply *rpctypes.AddPeerResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	host, portString, err := net.SplitHostPort(args.Addr)
	ip := net.ParseIP(host).To4()
	if ip == nil {
		return errors.New("invalid v4 IP")
	}
	port, err := strconv.Atoi(portString)
	if err != nil {
		return err
	}
	if port == 0 {
		return errors.New("invalid port")
	}
	addr := &net.TCPAddr{
		IP:   ip,
		Port: port,
	}
	t.AddPeer(addr)
	return nil
}

func (h *rpcHandler) AddTracker(args *rpctypes.AddTrackerRequest, reply *rpctypes.AddTrackerResponse) error {
	t := h.session.GetTorrent(args.ID)
	if t == nil {
		return errTorrentNotFound
	}
	return t.AddTracker(args.URL)
}
<|endoftext|>


<commit_before>//go:build !js
// +build !js

package webrtc

import (
	"strings"
	"sync"

	"github.com/pion/rtp"
	"github.com/pion/webrtc/v3/internal/util"
	"github.com/pion/webrtc/v3/pkg/media"
)

// trackBinding is a single bind for a Track
// Bind can be called multiple times, this stores the
// result for a single bind call so that it can be used when writing
type trackBinding struct {
	id          string
	ssrc        SSRC
	payloadType PayloadType
	writeStream TrackLocalWriter
}

// TrackLocalStaticRTP  is a TrackLocal that has a pre-set codec and accepts RTP Packets.
// If you wish to send a media.Sample use TrackLocalStaticSample
type TrackLocalStaticRTP struct {
	mu                sync.RWMutex
	bindings          []trackBinding
	codec             RTPCodecCapability
	id, rid, streamID string
}

// NewTrackLocalStaticRTP returns a TrackLocalStaticRTP.
func NewTrackLocalStaticRTP(c RTPCodecCapability, id, streamID string, options ...func(*TrackLocalStaticRTP)) (*TrackLocalStaticRTP, error) {
	t := &TrackLocalStaticRTP{
		codec:    c,
		bindings: []trackBinding{},
		id:       id,
		streamID: streamID,
	}

	for _, option := range options {
		option(t)
	}

	return t, nil
}

// WithRTPStreamID sets the RTP stream ID for this TrackLocalStaticRTP.
func WithRTPStreamID(rid string) func(*TrackLocalStaticRTP) {
	return func(t *TrackLocalStaticRTP) {
		t.rid = rid
	}
}

// Bind is called by the PeerConnection after negotiation is complete
// This asserts that the code requested is supported by the remote peer.
// If so it setups all the state (SSRC and PayloadType) to have a call
func (s *TrackLocalStaticRTP) Bind(t TrackLocalContext) (RTPCodecParameters, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	parameters := RTPCodecParameters{RTPCodecCapability: s.codec}
	if codec, matchType := codecParametersFuzzySearch(parameters, t.CodecParameters()); matchType != codecMatchNone {
		s.bindings = append(s.bindings, trackBinding{
			ssrc:        t.SSRC(),
			payloadType: codec.PayloadType,
			writeStream: t.WriteStream(),
			id:          t.ID(),
		})
		return codec, nil
	}

	return RTPCodecParameters{}, ErrUnsupportedCodec
}

// Unbind implements the teardown logic when the track is no longer needed. This happens
// because a track has been stopped.
func (s *TrackLocalStaticRTP) Unbind(t TrackLocalContext) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	for i := range s.bindings {
		if s.bindings[i].id == t.ID() {
			s.bindings[i] = s.bindings[len(s.bindings)-1]
			s.bindings = s.bindings[:len(s.bindings)-1]
			return nil
		}
	}

	return ErrUnbindFailed
}

// ID is the unique identifier for this Track. This should be unique for the
// stream, but doesn't have to globally unique. A common example would be 'audio' or 'video'
// and StreamID would be 'desktop' or 'webcam'
func (s *TrackLocalStaticRTP) ID() string { return s.id }

// StreamID is the group this track belongs too. This must be unique
func (s *TrackLocalStaticRTP) StreamID() string { return s.streamID }

// RID is the RTP stream identifier.
func (s *TrackLocalStaticRTP) RID() string { return s.rid }

// Kind controls if this TrackLocal is audio or video
func (s *TrackLocalStaticRTP) Kind() RTPCodecType {
	switch {
	case strings.HasPrefix(s.codec.MimeType, "audio/"):
		return RTPCodecTypeAudio
	case strings.HasPrefix(s.codec.MimeType, "video/"):
		return RTPCodecTypeVideo
	default:
		return RTPCodecType(0)
	}
}

// Codec gets the Codec of the track
func (s *TrackLocalStaticRTP) Codec() RTPCodecCapability {
	return s.codec
}

// packetPool is a pool of packets used by WriteRTP and Write below
// nolint:gochecknoglobals
var rtpPacketPool = sync.Pool{
	New: func() interface{} {
		return &rtp.Packet{}
	},
}

// WriteRTP writes a RTP Packet to the TrackLocalStaticRTP
// If one PeerConnection fails the packets will still be sent to
// all PeerConnections. The error message will contain the ID of the failed
// PeerConnections so you can remove them
func (s *TrackLocalStaticRTP) WriteRTP(p *rtp.Packet) error {
	ipacket := rtpPacketPool.Get()
	packet := ipacket.(*rtp.Packet) //nolint:forcetypeassert
	defer func() {
		*packet = rtp.Packet{}
		rtpPacketPool.Put(ipacket)
	}()
	*packet = *p
	return s.writeRTP(packet)
}

// writeRTP is like WriteRTP, except that it may modify the packet p
func (s *TrackLocalStaticRTP) writeRTP(p *rtp.Packet) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	writeErrs := []error{}

	for _, b := range s.bindings {
		p.Header.SSRC = uint32(b.ssrc)
		p.Header.PayloadType = uint8(b.payloadType)
		if _, err := b.writeStream.WriteRTP(&p.Header, p.Payload); err != nil {
			writeErrs = append(writeErrs, err)
		}
	}

	return util.FlattenErrs(writeErrs)
}

// Write writes a RTP Packet as a buffer to the TrackLocalStaticRTP
// If one PeerConnection fails the packets will still be sent to
// all PeerConnections. The error message will contain the ID of the failed
// PeerConnections so you can remove them
func (s *TrackLocalStaticRTP) Write(b []byte) (n int, err error) {
	ipacket := rtpPacketPool.Get()
	packet := ipacket.(*rtp.Packet) //nolint:forcetypeassert
	defer func() {
		*packet = rtp.Packet{}
		rtpPacketPool.Put(ipacket)
	}()

	if err = packet.Unmarshal(b); err != nil {
		return 0, err
	}

	return len(b), s.writeRTP(packet)
}

// TrackLocalStaticSample is a TrackLocal that has a pre-set codec and accepts Samples.
// If you wish to send a RTP Packet use TrackLocalStaticRTP
type TrackLocalStaticSample struct {
	packetizer rtp.Packetizer
	sequencer  rtp.Sequencer
	rtpTrack   *TrackLocalStaticRTP
	clockRate  float64
}

// NewTrackLocalStaticSample returns a TrackLocalStaticSample
func NewTrackLocalStaticSample(c RTPCodecCapability, id, streamID string, options ...func(*TrackLocalStaticRTP)) (*TrackLocalStaticSample, error) {
	rtpTrack, err := NewTrackLocalStaticRTP(c, id, streamID, options...)
	if err != nil {
		return nil, err
	}

	return &TrackLocalStaticSample{
		rtpTrack: rtpTrack,
	}, nil
}

// ID is the unique identifier for this Track. This should be unique for the
// stream, but doesn't have to globally unique. A common example would be 'audio' or 'video'
// and StreamID would be 'desktop' or 'webcam'
func (s *TrackLocalStaticSample) ID() string { return s.rtpTrack.ID() }

// StreamID is the group this track belongs too. This must be unique
func (s *TrackLocalStaticSample) StreamID() string { return s.rtpTrack.StreamID() }

// RID is the RTP stream identifier.
func (s *TrackLocalStaticSample) RID() string { return s.rtpTrack.RID() }

// Kind controls if this TrackLocal is audio or video
func (s *TrackLocalStaticSample) Kind() RTPCodecType { return s.rtpTrack.Kind() }

// Codec gets the Codec of the track
func (s *TrackLocalStaticSample) Codec() RTPCodecCapability {
	return s.rtpTrack.Codec()
}

// Bind is called by the PeerConnection after negotiation is complete
// This asserts that the code requested is supported by the remote peer.
// If so it setups all the state (SSRC and PayloadType) to have a call
func (s *TrackLocalStaticSample) Bind(t TrackLocalContext) (RTPCodecParameters, error) {
	codec, err := s.rtpTrack.Bind(t)
	if err != nil {
		return codec, err
	}

	s.rtpTrack.mu.Lock()
	defer s.rtpTrack.mu.Unlock()

	// We only need one packetizer
	if s.packetizer != nil {
		return codec, nil
	}

	payloader, err := payloaderForCodec(codec.RTPCodecCapability)
	if err != nil {
		return codec, err
	}

	s.sequencer = rtp.NewRandomSequencer()
	s.packetizer = rtp.NewPacketizer(
		rtpOutboundMTU,
		0, // Value is handled when writing
		0, // Value is handled when writing
		payloader,
		s.sequencer,
		codec.ClockRate,
	)
	s.clockRate = float64(codec.RTPCodecCapability.ClockRate)
	return codec, nil
}

// Unbind implements the teardown logic when the track is no longer needed. This happens
// because a track has been stopped.
func (s *TrackLocalStaticSample) Unbind(t TrackLocalContext) error {
	return s.rtpTrack.Unbind(t)
}

// WriteSample writes a Sample to the TrackLocalStaticSample
// If one PeerConnection fails the packets will still be sent to
// all PeerConnections. The error message will contain the ID of the failed
// PeerConnections so you can remove them
func (s *TrackLocalStaticSample) WriteSample(sample media.Sample) error {
	s.rtpTrack.mu.RLock()
	p := s.packetizer
	clockRate := s.clockRate
	s.rtpTrack.mu.RUnlock()

	if p == nil {
		return nil
	}

	// skip packets by the number of previously dropped packets
	for i := uint16(0); i < sample.PrevDroppedPackets; i++ {
		s.sequencer.NextSequenceNumber()
	}

	samples := uint32(sample.Duration.Seconds() * clockRate)
	if sample.PrevDroppedPackets > 0 {
		p.SkipSamples(samples * uint32(sample.PrevDroppedPackets))
	}
	packets := p.Packetize(sample.Data, samples)

	writeErrs := []error{}
	for _, p := range packets {
		if err := s.rtpTrack.WriteRTP(p); err != nil {
			writeErrs = append(writeErrs, err)
		}
	}

	return util.FlattenErrs(writeErrs)
}
<commit_msg>Moved duplicate operation to function<commit_after>//go:build !js
// +build !js

package webrtc

import (
	"strings"
	"sync"

	"github.com/pion/rtp"
	"github.com/pion/webrtc/v3/internal/util"
	"github.com/pion/webrtc/v3/pkg/media"
)

// trackBinding is a single bind for a Track
// Bind can be called multiple times, this stores the
// result for a single bind call so that it can be used when writing
type trackBinding struct {
	id          string
	ssrc        SSRC
	payloadType PayloadType
	writeStream TrackLocalWriter
}

// TrackLocalStaticRTP  is a TrackLocal that has a pre-set codec and accepts RTP Packets.
// If you wish to send a media.Sample use TrackLocalStaticSample
type TrackLocalStaticRTP struct {
	mu                sync.RWMutex
	bindings          []trackBinding
	codec             RTPCodecCapability
	id, rid, streamID string
}

// NewTrackLocalStaticRTP returns a TrackLocalStaticRTP.
func NewTrackLocalStaticRTP(c RTPCodecCapability, id, streamID string, options ...func(*TrackLocalStaticRTP)) (*TrackLocalStaticRTP, error) {
	t := &TrackLocalStaticRTP{
		codec:    c,
		bindings: []trackBinding{},
		id:       id,
		streamID: streamID,
	}

	for _, option := range options {
		option(t)
	}

	return t, nil
}

// WithRTPStreamID sets the RTP stream ID for this TrackLocalStaticRTP.
func WithRTPStreamID(rid string) func(*TrackLocalStaticRTP) {
	return func(t *TrackLocalStaticRTP) {
		t.rid = rid
	}
}

// Bind is called by the PeerConnection after negotiation is complete
// This asserts that the code requested is supported by the remote peer.
// If so it setups all the state (SSRC and PayloadType) to have a call
func (s *TrackLocalStaticRTP) Bind(t TrackLocalContext) (RTPCodecParameters, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	parameters := RTPCodecParameters{RTPCodecCapability: s.codec}
	if codec, matchType := codecParametersFuzzySearch(parameters, t.CodecParameters()); matchType != codecMatchNone {
		s.bindings = append(s.bindings, trackBinding{
			ssrc:        t.SSRC(),
			payloadType: codec.PayloadType,
			writeStream: t.WriteStream(),
			id:          t.ID(),
		})
		return codec, nil
	}

	return RTPCodecParameters{}, ErrUnsupportedCodec
}

// Unbind implements the teardown logic when the track is no longer needed. This happens
// because a track has been stopped.
func (s *TrackLocalStaticRTP) Unbind(t TrackLocalContext) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	for i := range s.bindings {
		if s.bindings[i].id == t.ID() {
			s.bindings[i] = s.bindings[len(s.bindings)-1]
			s.bindings = s.bindings[:len(s.bindings)-1]
			return nil
		}
	}

	return ErrUnbindFailed
}

// ID is the unique identifier for this Track. This should be unique for the
// stream, but doesn't have to globally unique. A common example would be 'audio' or 'video'
// and StreamID would be 'desktop' or 'webcam'
func (s *TrackLocalStaticRTP) ID() string { return s.id }

// StreamID is the group this track belongs too. This must be unique
func (s *TrackLocalStaticRTP) StreamID() string { return s.streamID }

// RID is the RTP stream identifier.
func (s *TrackLocalStaticRTP) RID() string { return s.rid }

// Kind controls if this TrackLocal is audio or video
func (s *TrackLocalStaticRTP) Kind() RTPCodecType {
	switch {
	case strings.HasPrefix(s.codec.MimeType, "audio/"):
		return RTPCodecTypeAudio
	case strings.HasPrefix(s.codec.MimeType, "video/"):
		return RTPCodecTypeVideo
	default:
		return RTPCodecType(0)
	}
}

// Codec gets the Codec of the track
func (s *TrackLocalStaticRTP) Codec() RTPCodecCapability {
	return s.codec
}

// packetPool is a pool of packets used by WriteRTP and Write below
// nolint:gochecknoglobals
var rtpPacketPool = sync.Pool{
	New: func() interface{} {
		return &rtp.Packet{}
	},
}

func resetPacketPoolAllocation(localPacket *rtp.Packet) {
	*localPacket = rtp.Packet{}
	rtpPacketPool.Put(localPacket)
}

func getPacketAllocationFromPool() *rtp.Packet {
	ipacket := rtpPacketPool.Get()
	return ipacket.(*rtp.Packet) //nolint:forcetypeassert
}

// WriteRTP writes a RTP Packet to the TrackLocalStaticRTP
// If one PeerConnection fails the packets will still be sent to
// all PeerConnections. The error message will contain the ID of the failed
// PeerConnections so you can remove them
func (s *TrackLocalStaticRTP) WriteRTP(p *rtp.Packet) error {
	packet := getPacketAllocationFromPool()

	defer resetPacketPoolAllocation(packet)

	*packet = *p

	return s.writeRTP(packet)
}

// writeRTP is like WriteRTP, except that it may modify the packet p
func (s *TrackLocalStaticRTP) writeRTP(p *rtp.Packet) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	writeErrs := []error{}

	for _, b := range s.bindings {
		p.Header.SSRC = uint32(b.ssrc)
		p.Header.PayloadType = uint8(b.payloadType)
		if _, err := b.writeStream.WriteRTP(&p.Header, p.Payload); err != nil {
			writeErrs = append(writeErrs, err)
		}
	}

	return util.FlattenErrs(writeErrs)
}

// Write writes a RTP Packet as a buffer to the TrackLocalStaticRTP
// If one PeerConnection fails the packets will still be sent to
// all PeerConnections. The error message will contain the ID of the failed
// PeerConnections so you can remove them
func (s *TrackLocalStaticRTP) Write(b []byte) (n int, err error) {
	packet := getPacketAllocationFromPool()

	defer resetPacketPoolAllocation(packet)

	if err = packet.Unmarshal(b); err != nil {
		return 0, err
	}

	return len(b), s.writeRTP(packet)
}

// TrackLocalStaticSample is a TrackLocal that has a pre-set codec and accepts Samples.
// If you wish to send a RTP Packet use TrackLocalStaticRTP
type TrackLocalStaticSample struct {
	packetizer rtp.Packetizer
	sequencer  rtp.Sequencer
	rtpTrack   *TrackLocalStaticRTP
	clockRate  float64
}

// NewTrackLocalStaticSample returns a TrackLocalStaticSample
func NewTrackLocalStaticSample(c RTPCodecCapability, id, streamID string, options ...func(*TrackLocalStaticRTP)) (*TrackLocalStaticSample, error) {
	rtpTrack, err := NewTrackLocalStaticRTP(c, id, streamID, options...)
	if err != nil {
		return nil, err
	}

	return &TrackLocalStaticSample{
		rtpTrack: rtpTrack,
	}, nil
}

// ID is the unique identifier for this Track. This should be unique for the
// stream, but doesn't have to globally unique. A common example would be 'audio' or 'video'
// and StreamID would be 'desktop' or 'webcam'
func (s *TrackLocalStaticSample) ID() string { return s.rtpTrack.ID() }

// StreamID is the group this track belongs too. This must be unique
func (s *TrackLocalStaticSample) StreamID() string { return s.rtpTrack.StreamID() }

// RID is the RTP stream identifier.
func (s *TrackLocalStaticSample) RID() string { return s.rtpTrack.RID() }

// Kind controls if this TrackLocal is audio or video
func (s *TrackLocalStaticSample) Kind() RTPCodecType { return s.rtpTrack.Kind() }

// Codec gets the Codec of the track
func (s *TrackLocalStaticSample) Codec() RTPCodecCapability {
	return s.rtpTrack.Codec()
}

// Bind is called by the PeerConnection after negotiation is complete
// This asserts that the code requested is supported by the remote peer.
// If so it setups all the state (SSRC and PayloadType) to have a call
func (s *TrackLocalStaticSample) Bind(t TrackLocalContext) (RTPCodecParameters, error) {
	codec, err := s.rtpTrack.Bind(t)
	if err != nil {
		return codec, err
	}

	s.rtpTrack.mu.Lock()
	defer s.rtpTrack.mu.Unlock()

	// We only need one packetizer
	if s.packetizer != nil {
		return codec, nil
	}

	payloader, err := payloaderForCodec(codec.RTPCodecCapability)
	if err != nil {
		return codec, err
	}

	s.sequencer = rtp.NewRandomSequencer()
	s.packetizer = rtp.NewPacketizer(
		rtpOutboundMTU,
		0, // Value is handled when writing
		0, // Value is handled when writing
		payloader,
		s.sequencer,
		codec.ClockRate,
	)
	s.clockRate = float64(codec.RTPCodecCapability.ClockRate)
	return codec, nil
}

// Unbind implements the teardown logic when the track is no longer needed. This happens
// because a track has been stopped.
func (s *TrackLocalStaticSample) Unbind(t TrackLocalContext) error {
	return s.rtpTrack.Unbind(t)
}

// WriteSample writes a Sample to the TrackLocalStaticSample
// If one PeerConnection fails the packets will still be sent to
// all PeerConnections. The error message will contain the ID of the failed
// PeerConnections so you can remove them
func (s *TrackLocalStaticSample) WriteSample(sample media.Sample) error {
	s.rtpTrack.mu.RLock()
	p := s.packetizer
	clockRate := s.clockRate
	s.rtpTrack.mu.RUnlock()

	if p == nil {
		return nil
	}

	// skip packets by the number of previously dropped packets
	for i := uint16(0); i < sample.PrevDroppedPackets; i++ {
		s.sequencer.NextSequenceNumber()
	}

	samples := uint32(sample.Duration.Seconds() * clockRate)
	if sample.PrevDroppedPackets > 0 {
		p.SkipSamples(samples * uint32(sample.PrevDroppedPackets))
	}
	packets := p.Packetize(sample.Data, samples)

	writeErrs := []error{}
	for _, p := range packets {
		if err := s.rtpTrack.WriteRTP(p); err != nil {
			writeErrs = append(writeErrs, err)
		}
	}

	return util.FlattenErrs(writeErrs)
}
<|endoftext|>