<commit_before>/*
Copyright 2022 The Kubernetes 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.
*/

// aws-costexplorer-export
// exports data from AWS Cost Explorer and imports it into a GCS Bucket as JSON

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	awssdkconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/costexplorer"
	cetypes "github.com/aws/aws-sdk-go-v2/service/costexplorer/types"
	"gocloud.dev/blob"
	_ "gocloud.dev/blob/fileblob"
	_ "gocloud.dev/blob/gcsblob"
	_ "gocloud.dev/blob/memblob"
	_ "gocloud.dev/blob/s3blob"
)

// consts
const (
	// formats
	usageDateFormat  = "2006-01-02"
	resultDateFormat = "200601021504"

	// templates
	fileNameTemplate = "cncf-aws-infra-billing-and-usage-data-%v.json"
)
// default config for runtime
var (
	defaultConfig = AWSCostExplorerExportConfig{
		AWSRegion: "us-east-1",
		LocalOutputFile:"/tmp/local-cncf-aws-infra-billing-and-usage-data.json",
		LocalOutputFileEnable: false,
		BucketURI: "gs://cncf-aws-infra-cost-and-billing-data",
		AmountOfDaysToReportFrom: 365,
		PromoteToLatest: true,
	}
)

func marshalAsJSON(input interface{}) string {
	o, err := json.Marshal(input)
	if err != nil {
		log.Println(err)
		return ""
	}
	return string(o)
}

func writeFile(path string, contents string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(contents)
	if err != nil {
		return err
	}
	return nil
}

// AWSCostExplorerExportConfig stores configuration for the runtime
type AWSCostExplorerExportConfig struct {
	AWSRegion                string
	LocalOutputFile          string
	LocalOutputFileEnable    bool
	BucketURI                string
	AmountOfDaysToReportFrom int
	PromoteToLatest          bool
}

// usageClient stores the client for costexplorer
type usageClient struct {
	client *costexplorer.Client
	config AWSCostExplorerExportConfig
}

// GetInputForUsage returns an input for making the cost and usage data request
func (c usageClient) GetInputForUsage(nextPageToken *string) *costexplorer.GetCostAndUsageInput {
	start := time.Now().
		Add(-time.Duration(time.Hour * 24 * time.Duration(c.config.AmountOfDaysToReportFrom))).
		Format(usageDateFormat)
	end := time.Now().Format(usageDateFormat)
	input := &costexplorer.GetCostAndUsageInput{
		Filter: &cetypes.Expression{
			Not: &cetypes.Expression{
				Dimensions: &cetypes.DimensionValues{
					Key:          cetypes.DimensionPurchaseType,
					MatchOptions: []cetypes.MatchOption{cetypes.MatchOptionEquals},
					Values:       []string{"Refund", "Credit"},
				},
			},
		},
		Metrics:     []string{string(cetypes.MetricUnblendedCost)},
		Granularity: cetypes.GranularityDaily,
		GroupBy: []cetypes.GroupDefinition{{
			Type: cetypes.GroupDefinitionTypeDimension,
			Key:  aws.String(string(cetypes.DimensionLinkedAccount)),
		}},
		TimePeriod: &cetypes.DateInterval{
			Start: aws.String(start),
			End:   aws.String(end),
		},
		NextPageToken: nextPageToken,
	}
	return input
}

// GetUsage returns the cost and usage data collected without pages
func (c usageClient) GetUsage() (costAndUsageOutput *costexplorer.GetCostAndUsageOutput, err error) {
	costAndUsageOutput = &costexplorer.GetCostAndUsageOutput{}
	var nextPageToken *string
	for page := 0; true; page ++ {
		usage, err := c.client.GetCostAndUsage(context.TODO(), c.GetInputForUsage(nextPageToken))
		if err != nil {
			return &costexplorer.GetCostAndUsageOutput{}, fmt.Errorf("error with getting usage, %v", err)
		}
		if usage == nil {
			break
		}
		log.Printf("- page (%v): dimensions (%v); group definitions (%v); results by time (%v)\n", page, len(usage.DimensionValueAttributes), len(usage.GroupDefinitions), len(usage.ResultsByTime))
		costAndUsageOutput.DimensionValueAttributes = append(costAndUsageOutput.DimensionValueAttributes, usage.DimensionValueAttributes...)
		costAndUsageOutput.GroupDefinitions = append(costAndUsageOutput.GroupDefinitions, usage.GroupDefinitions...)
		costAndUsageOutput.ResultsByTime = append(costAndUsageOutput.ResultsByTime, usage.ResultsByTime...)

		if usage.NextPageToken == nil {
			break
		}
		nextPageToken = usage.NextPageToken
		costAndUsageOutput.ResultMetadata = usage.ResultMetadata
	}
	costAndUsageOutput.NextPageToken = nil
	if len(costAndUsageOutput.ResultsByTime) == 0 {
		return &costexplorer.GetCostAndUsageOutput{}, fmt.Errorf("error: empty dataset")
	}
	return costAndUsageOutput, nil
}

// BucketAccess stores config for accessing a bucket
type BucketAccess struct {
	URI    string
	Bucket *blob.Bucket
}

// Open returns access to a bucket
func (b BucketAccess) Open() (BucketAccess, error) {
	bucket, err := blob.OpenBucket(context.TODO(), b.URI)
	if err != nil {
		return BucketAccess{}, err
	}
	b.Bucket = bucket
	return b, nil
}

// WriteToFile writes a file in the bucket with access
func (b BucketAccess) WriteToFile(name string, data string) (err error) {
	w, err := b.Bucket.NewWriter(context.TODO(), name, nil)
	if err != nil {
		return err
	}
	_, writeErr := fmt.Fprintln(w, data)
	closeErr := w.Close()
	if writeErr != nil {
		return err
	}
	if closeErr != nil {
		return err
	}
	return nil
}

func main() {
	var config AWSCostExplorerExportConfig
	flag.StringVar(&config.AWSRegion, "aws-region", defaultConfig.AWSRegion, "specify an AWS region")
	flag.StringVar(&config.LocalOutputFile, "output-file", defaultConfig.LocalOutputFile, "specify a local file to write the usage JSON data to")
	flag.BoolVar(&config.LocalOutputFileEnable, "output-file-enable", defaultConfig.LocalOutputFileEnable, "specify whether the usage data is also written to disk locally")
	flag.StringVar(&config.BucketURI, "bucket-uri", defaultConfig.BucketURI, "specify a bucket to write to")
	flag.IntVar(&config.AmountOfDaysToReportFrom, "days-ago", defaultConfig.AmountOfDaysToReportFrom, "specify the amount of days back to report from today")
	flag.BoolVar(&config.PromoteToLatest, "promote-to-latest", defaultConfig.PromoteToLatest, "specifies whether to promote the cost and usage data to a latest JSON file")
	flag.Parse()

	log.Printf("Config: %#v\n", config)

	cfg, err := awssdkconfig.LoadDefaultConfig(context.TODO(),
		awssdkconfig.WithRegion(config.AWSRegion),
	)
	if err != nil {
		log.Printf("unable to load SDK config, %v", err)
		return
	}

	ceclient := costexplorer.NewFromConfig(cfg)

	log.Println("Fetching usage data")
	uc := usageClient{client: ceclient, config: config}
	costAndUsageOutput, err := uc.GetUsage()
	if err != nil {
		log.Printf("%v", err)
		return
	}
	log.Println("Fetched usage data")
	costAndUsageOutputJSON := marshalAsJSON(costAndUsageOutput)
	if config.LocalOutputFileEnable {
		log.Println("Writing usage data to file")
		err = writeFile(config.LocalOutputFile, costAndUsageOutputJSON)
		if err != nil {
			log.Printf("%v", err)
			return
		}
		log.Printf("Wrote usage data to file '%v'\n", config.LocalOutputFile)
	}
	log.Printf("Opening GCS bucket '%v'\n", config.BucketURI)
	ba := BucketAccess{URI: config.BucketURI}
	ba, err = ba.Open()
	if err != nil {
		log.Printf("%v", err)
		return
	}
	defer func() {
		log.Printf("Closing GCS bucket '%v'\n", config.BucketURI)
		ba.Bucket.Close()
	}()
	fileNames := []string{
		fmt.Sprintf(fileNameTemplate, time.Now().Format(resultDateFormat)),
	}
	if config.PromoteToLatest {
		fileNames = append(fileNames, fmt.Sprintf(fileNameTemplate, "latest"))
	}
	for _, fileName := range fileNames {
		log.Printf("Uploading '%v' to '%v/%v'", fileName, config.BucketURI, fileName)
		err = ba.WriteToFile(fileName, costAndUsageOutputJSON)
		if err != nil {
			log.Printf("%v", err)
			return
		}
		log.Printf("Uploaded '%v' to '%v/%v' successfully", fileName, config.BucketURI, fileName)
	}
}
<commit_msg>Update output into multiple files<commit_after>/*
Copyright 2022 The Kubernetes 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.
*/

// aws-costexplorer-export
// exports data from AWS Cost Explorer and imports it into a GCS Bucket as JSON

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"path"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	awssdkconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/costexplorer"
	cetypes "github.com/aws/aws-sdk-go-v2/service/costexplorer/types"
	"gocloud.dev/blob"
	_ "gocloud.dev/blob/fileblob"
	_ "gocloud.dev/blob/gcsblob"
	_ "gocloud.dev/blob/memblob"
	_ "gocloud.dev/blob/s3blob"
)

// consts
const (
	// formats
	usageDateFormat  = "2006-01-02"
	resultDateFormat = "200601021504"

	// templates
	fileNameTemplate = "cncf-aws-infra-billing-and-usage-data-%v-%v.json"
)

// default config for runtime
var (
	defaultConfig = AWSCostExplorerExportConfig{
		AWSRegion:                "us-east-1",
		LocalOutputFolder:        "/tmp/local-cncf-aws-infra-billing-and-usage-data",
		LocalOutputFolderEnable:  false,
		BucketURI:                "gs://cncf-aws-infra-cost-and-billing-data",
		AmountOfDaysToReportFrom: 365,
		PromoteToLatest:          true,
	}
)

func marshalAsJSON(input interface{}) string {
	o, err := json.Marshal(input)
	if err != nil {
		log.Println(err)
		return ""
	}
	return string(o)
}

func writeFile(path string, contents string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(contents)
	if err != nil {
		return err
	}
	return nil
}

type FileOutputs map[string]string

// AWSCostExplorerExportConfig stores configuration for the runtime
type AWSCostExplorerExportConfig struct {
	AWSRegion                string
	LocalOutputFolder        string
	LocalOutputFolderEnable  bool
	BucketURI                string
	AmountOfDaysToReportFrom int
	PromoteToLatest          bool
}

// usageClient stores the client for costexplorer
type usageClient struct {
	client *costexplorer.Client
	config AWSCostExplorerExportConfig
}

// GetInputForUsage returns an input for making the cost and usage data request
func (c usageClient) GetInputForUsage(nextPageToken *string) *costexplorer.GetCostAndUsageInput {
	start := time.Now().
		Add(-time.Duration(time.Hour * 24 * time.Duration(c.config.AmountOfDaysToReportFrom))).
		Format(usageDateFormat)
	end := time.Now().Format(usageDateFormat)
	input := &costexplorer.GetCostAndUsageInput{
		Filter: &cetypes.Expression{
			Not: &cetypes.Expression{
				Dimensions: &cetypes.DimensionValues{
					Key:          cetypes.DimensionPurchaseType,
					MatchOptions: []cetypes.MatchOption{cetypes.MatchOptionEquals},
					Values:       []string{"Refund", "Credit"},
				},
			},
		},
		Metrics:     []string{string(cetypes.MetricUnblendedCost)},
		Granularity: cetypes.GranularityDaily,
		GroupBy: []cetypes.GroupDefinition{{
			Type: cetypes.GroupDefinitionTypeDimension,
			Key:  aws.String(string(cetypes.DimensionLinkedAccount)),
		}},
		TimePeriod: &cetypes.DateInterval{
			Start: aws.String(start),
			End:   aws.String(end),
		},
		NextPageToken: nextPageToken,
	}
	return input
}

// GetUsage returns the cost and usage data collected without pages
func (c usageClient) GetUsage() (costAndUsageOutput *costexplorer.GetCostAndUsageOutput, err error) {
	costAndUsageOutput = &costexplorer.GetCostAndUsageOutput{}
	var nextPageToken *string
	for page := 0; true; page++ {
		input := c.GetInputForUsage(nextPageToken)
		usage, err := c.client.GetCostAndUsage(context.TODO(), input)
		if err != nil {
			return &costexplorer.GetCostAndUsageOutput{}, fmt.Errorf("error with getting usage, %v", err)
		}
		if usage == nil {
			break
		}
		log.Printf("- page (%v): dimensions (%v); group definitions (%v); results by time (%v)\n", page, len(usage.DimensionValueAttributes), len(usage.GroupDefinitions), len(usage.ResultsByTime))
		costAndUsageOutput.DimensionValueAttributes = append(costAndUsageOutput.DimensionValueAttributes, usage.DimensionValueAttributes...)
		costAndUsageOutput.GroupDefinitions = append(costAndUsageOutput.GroupDefinitions, usage.GroupDefinitions...)
		costAndUsageOutput.ResultsByTime = append(costAndUsageOutput.ResultsByTime, usage.ResultsByTime...)

		if usage.NextPageToken == nil {
			break
		}
		nextPageToken = usage.NextPageToken
		costAndUsageOutput.ResultMetadata = usage.ResultMetadata
	}
	costAndUsageOutput.NextPageToken = nil
	if len(costAndUsageOutput.ResultsByTime) == 0 {
		return &costexplorer.GetCostAndUsageOutput{}, fmt.Errorf("error: empty dataset")
	}
	return costAndUsageOutput, nil
}

type ResultByTime struct {
	cetypes.ResultByTime
	Total interface{} `json:"Total,omitempty"`
}

// BucketAccess stores config for accessing a bucket
type BucketAccess struct {
	URI    string
	Bucket *blob.Bucket
}

// Open returns access to a bucket
func (b BucketAccess) Open() (BucketAccess, error) {
	bucket, err := blob.OpenBucket(context.TODO(), b.URI)
	if err != nil {
		return BucketAccess{}, err
	}
	b.Bucket = bucket
	return b, nil
}

// WriteToFile writes a file in the bucket with access
func (b BucketAccess) WriteToFile(name string, data string) (err error) {
	w, err := b.Bucket.NewWriter(context.TODO(), name, nil)
	if err != nil {
		return err
	}
	_, writeErr := fmt.Fprintln(w, data)
	closeErr := w.Close()
	if writeErr != nil {
		return err
	}
	if closeErr != nil {
		return err
	}
	return nil
}

func main() {
	var config AWSCostExplorerExportConfig
	flag.StringVar(&config.AWSRegion, "aws-region", defaultConfig.AWSRegion, "specify an AWS region")
	flag.StringVar(&config.LocalOutputFolder, "output-file", defaultConfig.LocalOutputFolder, "specify a local file to write the usage JSON data to")
	flag.BoolVar(&config.LocalOutputFolderEnable, "output-file-enable", defaultConfig.LocalOutputFolderEnable, "specify whether the usage data is also written to disk locally")
	flag.StringVar(&config.BucketURI, "bucket-uri", defaultConfig.BucketURI, "specify a bucket to write to")
	flag.IntVar(&config.AmountOfDaysToReportFrom, "days-ago", defaultConfig.AmountOfDaysToReportFrom, "specify the amount of days back to report from today")
	flag.BoolVar(&config.PromoteToLatest, "promote-to-latest", defaultConfig.PromoteToLatest, "specifies whether to promote the cost and usage data to a latest JSON file")
	flag.Parse()

	log.Printf("Config: %#v\n", config)

	cfg, err := awssdkconfig.LoadDefaultConfig(context.TODO(),
		awssdkconfig.WithRegion(config.AWSRegion),
	)
	if err != nil {
		log.Printf("unable to load SDK config, %v", err)
		return
	}

	ceclient := costexplorer.NewFromConfig(cfg)

	log.Println("Fetching usage data")
	uc := usageClient{client: ceclient, config: config}
	costAndUsageOutput, err := uc.GetUsage()
	if err != nil {
		log.Printf("%v", err)
		return
	}
	log.Println("Fetched usage data")

	log.Println("Formatting data")
	fileOutputs := FileOutputs{}
	for _, value := range costAndUsageOutput.ResultsByTime {
		o := marshalAsJSON(ResultByTime{
			ResultByTime: value,
			Total: nil,
		})
		fileOutputs["ResultsByTime"] += o + "\n"
	}
	for _, value := range costAndUsageOutput.DimensionValueAttributes {
		o := marshalAsJSON(value)
		fileOutputs["DimensionValueAttributes"] += o + "\n"
	}
	for _, value := range costAndUsageOutput.GroupDefinitions {
		o := marshalAsJSON(value)
		fileOutputs["GroupDefinitions"] += o + "\n"
	}
	log.Println("Formatted data")

	if config.LocalOutputFolderEnable {
		_ = os.Mkdir(config.LocalOutputFolder, 0700)
		log.Printf("Writing usage data to files in '%v'\n", config.LocalOutputFolder)
		for name, content := range fileOutputs {
			log.Printf("- writing '%v'\n", name)
			err = writeFile(path.Join(config.LocalOutputFolder, name+".json"), content)
			if err != nil {
				log.Printf("%v", err)
				return
			}
		}
		log.Printf("Wrote usage data to file '%v'\n", config.LocalOutputFolder)
	}
	log.Printf("Opening GCS bucket '%v'\n", config.BucketURI)
	ba := BucketAccess{URI: config.BucketURI}
	ba, err = ba.Open()
	if err != nil {
		log.Printf("%v", err)
		return
	}
	defer func() {
		log.Printf("Closing GCS bucket '%v'\n", config.BucketURI)
		ba.Bucket.Close()
	}()
	for name, content := range fileOutputs {
		fileNames := []string{
			fmt.Sprintf(fileNameTemplate, name, time.Now().Format(resultDateFormat)),
		}
		if config.PromoteToLatest {
			fileNames = append(fileNames, fmt.Sprintf(fileNameTemplate, name, "latest"))
		}
		for _, fileName := range fileNames {
			log.Printf("Uploading '%v' to '%v/%v'", fileName, config.BucketURI, fileName)
			err = ba.WriteToFile(fileName, content)
			if err != nil {
				log.Printf("%v", err)
				return
			}
			log.Printf("Uploaded '%v' to '%v/%v' successfully", fileName, config.BucketURI, fileName)
		}
	}
	// bigquery stuff
	//   1. create dataset based on date
	//   2. load into dataset based on date
	//   3. if not promoting, exit
	//   4. delete latest dataset
	//   5. create latest dataset
	//   6. load into latest dataset
	// TODO declare schema to prevent errors
}
<|endoftext|>


<commit_before>package linkedql

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/cayleygraph/quad"
)

var (
	typeByName = make(map[string]reflect.Type)
	nameByType = make(map[reflect.Type]string)
)

// Register adds an Item type to the registry.
func Register(typ RegistryItem) {
	tp := reflect.TypeOf(typ)
	if tp.Kind() == reflect.Ptr {
		tp = tp.Elem()
	}
	if tp.Kind() != reflect.Struct {
		panic("only structs are allowed")
	}
	name := string(typ.Type())
	if _, ok := typeByName[name]; ok {
		panic("this name was already registered")
	}
	typeByName[name] = tp
	nameByType[tp] = name
}

var quadValue = reflect.TypeOf((*quad.Value)(nil)).Elem()
var quadSliceValue = reflect.TypeOf(([]quad.Value)(nil))

// Unmarshal attempts to unmarshal an Item or returns error.
func Unmarshal(data []byte) (RegistryItem, error) {
	var m map[string]json.RawMessage
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, err
	}
	var typ string
	if err := json.Unmarshal(m["@type"], &typ); err != nil {
		return nil, err
	}
	delete(m, "@type")
	tp, ok := typeByName[typ]
	if !ok {
		return nil, fmt.Errorf("unsupported item: %q", typ)
	}
	item := reflect.New(tp).Elem()
	for i := 0; i < tp.NumField(); i++ {
		f := tp.Field(i)
		name := f.Name
		tag := strings.SplitN(f.Tag.Get("json"), ",", 2)[0]
		if tag == "-" {
			continue
		} else if tag != "" {
			name = tag
		}
		v, ok := m[name]
		if !ok {
			continue
		}
		fv := item.Field(i)
		switch f.Type {
		case quadValue:
			var a interface{}
			err := json.Unmarshal(v, &a)
			if err != nil {
				return nil, err
			}
			value, err := parseValue(v)
			if err != nil {
				return nil, err
			}
			fv.Set(reflect.ValueOf(value))
			continue
		case quadSliceValue:
			var a []interface{}
			err := json.Unmarshal(v, &a)
			if err != nil {
				return nil, err
			}
			var values []quad.Value
			for _, item := range a {
				value, err := parseValue(item)
				if err != nil {
					return nil, err
				}
				values = append(values, value)
			}
			fv.Set(reflect.ValueOf(values))
			continue
		}
		switch f.Type.Kind() {
		case reflect.Interface:
			s, err := Unmarshal(v)
			if err != nil {
				return nil, err
			}
			fv.Set(reflect.ValueOf(s))
		case reflect.Slice:
			el := f.Type.Elem()
			if el.Kind() != reflect.Interface {
				err := json.Unmarshal(v, fv.Addr().Interface())
				if err != nil {
					return nil, err
				}
			} else {
				var arr []json.RawMessage
				if err := json.Unmarshal(v, &arr); err != nil {
					return nil, err
				}
				if arr != nil {
					va := reflect.MakeSlice(f.Type, len(arr), len(arr))
					for i, v := range arr {
						s, err := Unmarshal(v)
						if err != nil {
							return nil, err
						}
						va.Index(i).Set(reflect.ValueOf(s))
					}
					fv.Set(va)
				}
			}
		default:
			err := json.Unmarshal(v, fv.Addr().Interface())
			if err != nil {
				return nil, err
			}
		}
	}
	return item.Addr().Interface().(RegistryItem), nil
}

func parseValue(a interface{}) (quad.Value, error) {
	switch a := a.(type) {
	case string:
		return quad.String(a), nil
	case int64:
		return quad.Int(a), nil
	case float64:
		return quad.Float(a), nil
	case bool:
		return quad.Bool(a), nil
	case map[string]interface{}:
		id, ok := a["@id"].(string)
		if ok {
			if strings.HasPrefix(id, "_:") {
				return quad.BNode(id[2:]), nil
			}
			return quad.IRI(id), nil
		}
		value, ok := a["@value"].(string)
		if ok {
			if language, ok := a["@language"].(string); ok {
				return quad.LangString{Value: quad.String("value"), Lang: language}, nil
			}
			if _type, ok := a["@type"].(string); ok {
				return quad.TypedString{Value: quad.String(value), Type: quad.IRI(_type)}, nil
			}
		}
	}
	return nil, fmt.Errorf("cannot parse JSON-LD value: %#v", a)
}

// RegistryItem in the registry.
type RegistryItem interface {
	Type() quad.IRI
}
<commit_msg>linkedql: Support IRIs in registry<commit_after>package linkedql

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/cayleygraph/quad"
)

var (
	typeByName = make(map[string]reflect.Type)
	nameByType = make(map[reflect.Type]string)
)

// Register adds an Item type to the registry.
func Register(typ RegistryItem) {
	tp := reflect.TypeOf(typ)
	if tp.Kind() == reflect.Ptr {
		tp = tp.Elem()
	}
	if tp.Kind() != reflect.Struct {
		panic("only structs are allowed")
	}
	name := string(typ.Type())
	if _, ok := typeByName[name]; ok {
		panic("this name was already registered")
	}
	typeByName[name] = tp
	nameByType[tp] = name
}

var quadValue = reflect.TypeOf((*quad.Value)(nil)).Elem()
var quadSliceValue = reflect.TypeOf([]quad.Value{})
var quadIRI = reflect.TypeOf(quad.IRI(""))
var quadSliceIRI = reflect.TypeOf([]quad.IRI{})

// Unmarshal attempts to unmarshal an Item or returns error.
func Unmarshal(data []byte) (RegistryItem, error) {
	var m map[string]json.RawMessage
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, err
	}
	var typ string
	if err := json.Unmarshal(m["@type"], &typ); err != nil {
		return nil, err
	}
	delete(m, "@type")
	tp, ok := typeByName[typ]
	if !ok {
		return nil, fmt.Errorf("unsupported item: %q", typ)
	}
	item := reflect.New(tp).Elem()
	for i := 0; i < tp.NumField(); i++ {
		f := tp.Field(i)
		name := f.Name
		tag := strings.SplitN(f.Tag.Get("json"), ",", 2)[0]
		if tag == "-" {
			continue
		} else if tag != "" {
			name = tag
		}
		v, ok := m[name]
		if !ok {
			continue
		}
		fv := item.Field(i)
		switch f.Type {
		case quadValue:
			var a interface{}
			err := json.Unmarshal(v, &a)
			if err != nil {
				return nil, err
			}
			value, err := parseValue(v)
			if err != nil {
				return nil, err
			}
			fv.Set(reflect.ValueOf(value))
			continue
		case quadSliceValue:
			var a []interface{}
			err := json.Unmarshal(v, &a)
			if err != nil {
				return nil, err
			}
			var values []quad.Value
			for _, item := range a {
				value, err := parseValue(item)
				if err != nil {
					return nil, err
				}
				values = append(values, value)
			}
			fv.Set(reflect.ValueOf(values))
			continue
		case quadIRI:
			var a interface{}
			err := json.Unmarshal(v, &a)
			if err != nil {
				return nil, err
			}
			value, err := parseIRI(v)
			if err != nil {
				return nil, err
			}
			fv.Set(reflect.ValueOf(*value))
			continue
		case quadSliceIRI:
			var a []interface{}
			err := json.Unmarshal(v, &a)
			if err != nil {
				return nil, err
			}
			var values []quad.IRI
			for _, item := range a {
				value, err := parseIRI(item)
				if err != nil {
					return nil, err
				}
				values = append(values, *value)
			}
			fv.Set(reflect.ValueOf(values))
			continue
		}
		switch f.Type.Kind() {
		case reflect.Interface:
			s, err := Unmarshal(v)
			if err != nil {
				return nil, err
			}
			fv.Set(reflect.ValueOf(s))
		case reflect.Slice:
			el := f.Type.Elem()
			if el.Kind() != reflect.Interface {
				err := json.Unmarshal(v, fv.Addr().Interface())
				if err != nil {
					return nil, err
				}
			} else {
				var arr []json.RawMessage
				if err := json.Unmarshal(v, &arr); err != nil {
					return nil, err
				}
				if arr != nil {
					va := reflect.MakeSlice(f.Type, len(arr), len(arr))
					for i, v := range arr {
						s, err := Unmarshal(v)
						if err != nil {
							return nil, err
						}
						va.Index(i).Set(reflect.ValueOf(s))
					}
					fv.Set(va)
				}
			}
		default:
			err := json.Unmarshal(v, fv.Addr().Interface())
			if err != nil {
				return nil, err
			}
		}
	}
	return item.Addr().Interface().(RegistryItem), nil
}

func parseBNode(a interface{}) (*quad.BNode, error) {
	pointer, err := parseIdentifierString(a)
	id := *pointer
	if err != nil {
		return nil, err
	}
	if !strings.HasPrefix(id, "_:") {
		return nil, fmt.Errorf("blank node ID must start with \"_:\"")
	}
	bnode := quad.BNode(id[2:])
	return &bnode, nil
}

func parseIRI(a interface{}) (*quad.IRI, error) {
	pointer, err := parseIdentifierString(a)
	id := *pointer
	if err != nil {
		return nil, err
	}
	iri := quad.IRI(id)
	return &iri, nil
}

func parseIdentifierString(a interface{}) (*string, error) {
	m, ok := a.(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("Not a map")
	}
	id, ok := m["@id"].(string)
	if !ok {
		return nil, fmt.Errorf("No string @id key")
	}
	return &id, nil
}

func parseLiteral(a interface{}) (quad.Value, error) {
	switch a := a.(type) {
	case string:
		return quad.String(a), nil
	case int64:
		return quad.Int(a), nil
	case float64:
		return quad.Float(a), nil
	case bool:
		return quad.Bool(a), nil
	case map[string]interface{}:
		value, ok := a["@value"].(string)
		if ok {
			if language, ok := a["@language"].(string); ok {
				return quad.LangString{Value: quad.String("value"), Lang: language}, nil
			}
			if _type, ok := a["@type"].(string); ok {
				return quad.TypedString{Value: quad.String(value), Type: quad.IRI(_type)}, nil
			}
		}
	}
	return nil, fmt.Errorf("Can not parse %#v as literal", a)
}

func parseValue(a interface{}) (quad.Value, error) {
	bnode, err := parseBNode(a)
	if err == nil {
		return bnode, nil
	}
	iri, err := parseIRI(a)
	if err == nil {
		return iri, nil
	}
	literal, err := parseLiteral(a)
	if err == nil {
		return literal, nil
	}
	return nil, fmt.Errorf("Can not parse JSON-LD value: %#v", a)
}

// RegistryItem in the registry.
type RegistryItem interface {
	Type() quad.IRI
}
<|endoftext|>


<commit_before>package udptracker

// http://bittorrent.org/beps/bep_0015.html
// http://xbtt.sourceforge.net/udp_tracker_bt.html
// http://www.rasterbar.com/products/libtorrent/udp_tracker_bt.html

import (
	"bytes"
	"encoding/binary"
	"errors"
	"net"
	"net/url"
	"sync"
	"time"

	"github.com/cenkalti/backoff"

	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/tracker"
)

const connectionIDMagic = 0x41727101980
const connectionIDInterval = time.Minute

type UDPTracker struct {
	URL           *url.URL
	Log           logger.Logger
	conn          *net.UDPConn
	dialMutex     sync.Mutex
	connected     bool
	transactions  map[int32]*transaction
	transactionsM sync.Mutex
	writeC        chan *transaction
}

func New(u *url.URL, l logger.Logger) *UDPTracker {
	return &UDPTracker{
		URL:          u,
		Log:          l,
		transactions: make(map[int32]*transaction),
		writeC:       make(chan *transaction),
	}
}

func (t *UDPTracker) dial() error {
	serverAddr, err := net.ResolveUDPAddr("udp", t.URL.Host)
	if err != nil {
		return err
	}
	t.conn, err = net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		return err
	}
	go t.readLoop()
	go t.writeLoop()
	return nil
}

// Close the tracker connection.
func (t *UDPTracker) Close() error {
	close(t.writeC)
	if t.conn != nil {
		return t.conn.Close()
	}
	return nil
}

// readLoop reads datagrams from connection, finds the transaction and
// sends the bytes to the transaction's response channel.
func (t *UDPTracker) readLoop() {
	// Read buffer must be big enough to hold a UDP packet of maximum expected size.
	// Current value is: 320 = 20 + 50*6 (AnnounceResponse with 50 peers)
	buf := make([]byte, 320)
	for {
		n, err := t.conn.Read(buf)
		if err != nil {
			t.Log.Error(err)
			if nerr, ok := err.(net.Error); ok && !nerr.Temporary() {
				t.Log.Debug("End of tracker read loop")
				return
			}
			continue
		}
		t.Log.Debug("Read ", n, " bytes")

		var header udpMessageHeader
		if n < binary.Size(header) {
			t.Log.Error("response is too small")
			continue
		}

		err = binary.Read(bytes.NewReader(buf), binary.BigEndian, &header)
		if err != nil {
			t.Log.Error(err)
			continue
		}

		t.transactionsM.Lock()
		trx, ok := t.transactions[header.TransactionID]
		delete(t.transactions, header.TransactionID)
		t.transactionsM.Unlock()
		if !ok {
			t.Log.Errorln("unexpected transaction_id:", header.TransactionID)
			continue
		}

		// Tracker has sent and error.
		if header.Action == actionError {
			// The part after the header is the error message.
			trx.err = tracker.Error(buf[binary.Size(header):])
			trx.Done()
			continue
		}

		// Copy data into a new slice because buf will be overwritten at next read.
		trx.response = make([]byte, n)
		copy(trx.response, buf)
		trx.Done()
	}
}

// writeLoop receives a request from t.transactionC, sets a random TransactionID
// and sends it to the tracker.
func (t *UDPTracker) writeLoop() {
	var connectionID int64
	var connectionIDtime time.Time

	for trx := range t.writeC {
		if time.Since(connectionIDtime) > connectionIDInterval {
			connectionID = t.connect()
			connectionIDtime = time.Now()
		}
		trx.request.SetConnectionID(connectionID)

		t.writeTrx(trx)
	}
}

func (t *UDPTracker) writeTrx(trx *transaction) {
	t.Log.Debugln("Writing transaction. ID:", trx.ID())
	_, err := trx.request.WriteTo(t.conn)
	if err != nil {
		t.Log.Error(err)
	}
}

// connect sends a connectRequest and returns a ConnectionID given by the tracker.
// On error, it backs off with the algorithm described in BEP15 and retries.
// It does not return until tracker sends a ConnectionID.
func (t *UDPTracker) connect() int64 {
	req := new(connectRequest)
	req.SetAction(actionConnect)
	req.SetConnectionID(connectionIDMagic)

	trx := newTransaction(req)

	for {
		data, err := t.retryTransaction(t.writeTrx, trx, nil) // Does not return until transaction is completed.
		if err != nil {
			t.Log.Error(err)
			continue
		}

		var response connectResponse
		err = binary.Read(bytes.NewReader(data), binary.BigEndian, &response)
		if err != nil {
			t.Log.Error(err)
			continue
		}

		if response.Action != actionConnect {
			t.Log.Error("invalid action in connect response")
			continue
		}

		t.Log.Debugf("connect Response: %#v\n", response)
		return response.ConnectionID
	}
}

func (t *UDPTracker) retryTransaction(f func(*transaction), trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	t.transactionsM.Lock()
	t.transactions[trx.ID()] = trx
	t.transactionsM.Unlock()

	ticker := backoff.NewTicker(new(udpBackOff))
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			f(trx)
		case <-trx.done:
			// transaction is deleted in readLoop()
			return trx.response, trx.err
		case <-cancel:
			t.transactionsM.Lock()
			delete(t.transactions, trx.ID())
			t.transactionsM.Unlock()
			return nil, tracker.ErrRequestCancelled
		}
	}
}

func (t *UDPTracker) sendTransaction(trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	f := func(trx *transaction) { t.writeC <- trx }
	return t.retryTransaction(f, trx, cancel)
}

func (t *UDPTracker) Announce(transfer tracker.Transfer, e tracker.Event, cancel <-chan struct{}) (*tracker.AnnounceResponse, error) {
	t.dialMutex.Lock()
	if !t.connected {
		err := t.dial()
		if err != nil {
			t.dialMutex.Unlock()
			return nil, err
		}
		t.connected = true
	}
	t.dialMutex.Unlock()

	request := &announceRequest{
		InfoHash:   transfer.InfoHash(),
		PeerID:     transfer.PeerID(),
		Event:      e,
		IP:         0, // Tracker uses sender of this UDP packet.
		Key:        0, // TODO set it
		NumWant:    tracker.NumWant,
		Port:       uint16(transfer.Port()),
		Extensions: 0,
	}
	request.SetAction(actionAnnounce)

	request2 := &transferAnnounceRequest{
		transfer:        transfer,
		announceRequest: request,
		urlData:         t.URL.RequestURI(),
	}
	trx := newTransaction(request2)

	// t.request may block, that's why we pass cancel as argument.
	reply, err := t.sendTransaction(trx, cancel)
	if err != nil {
		if err, ok := err.(tracker.Error); ok {
			return &tracker.AnnounceResponse{Error: err}, nil
		}
		return nil, err
	}

	response, peers, err := t.parseAnnounceResponse(reply)
	if err != nil {
		return nil, err
	}
	t.Log.Debugf("Announce response: %#v", response)

	return &tracker.AnnounceResponse{
		Interval: time.Duration(response.Interval) * time.Second,
		Leechers: response.Leechers,
		Seeders:  response.Seeders,
		Peers:    peers,
	}, nil
}

func (t *UDPTracker) parseAnnounceResponse(data []byte) (*udpAnnounceResponse, []*net.TCPAddr, error) {
	response := new(udpAnnounceResponse)
	if len(data) < binary.Size(response) {
		return nil, nil, errors.New("response is too small")
	}

	reader := bytes.NewReader(data)

	err := binary.Read(reader, binary.BigEndian, response)
	if err != nil {
		return nil, nil, err
	}
	t.Log.Debugf("annouceResponse: %#v", response)

	if response.Action != actionAnnounce {
		return nil, nil, errors.New("invalid action")
	}

	peers, err := tracker.ParsePeersBinary(reader, t.Log)
	if err != nil {
		return nil, nil, err
	}

	return response, peers, nil
}
<commit_msg>unexport udp tracker fields<commit_after>package udptracker

// http://bittorrent.org/beps/bep_0015.html
// http://xbtt.sourceforge.net/udp_tracker_bt.html
// http://www.rasterbar.com/products/libtorrent/udp_tracker_bt.html

import (
	"bytes"
	"encoding/binary"
	"errors"
	"net"
	"net/url"
	"sync"
	"time"

	"github.com/cenkalti/backoff"

	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/tracker"
)

const connectionIDMagic = 0x41727101980
const connectionIDInterval = time.Minute

type UDPTracker struct {
	url           *url.URL
	log           logger.Logger
	conn          *net.UDPConn
	dialMutex     sync.Mutex
	connected     bool
	transactions  map[int32]*transaction
	transactionsM sync.Mutex
	writeC        chan *transaction
}

func New(u *url.URL, l logger.Logger) *UDPTracker {
	return &UDPTracker{
		url:          u,
		log:          l,
		transactions: make(map[int32]*transaction),
		writeC:       make(chan *transaction),
	}
}

func (t *UDPTracker) dial() error {
	serverAddr, err := net.ResolveUDPAddr("udp", t.url.Host)
	if err != nil {
		return err
	}
	t.conn, err = net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		return err
	}
	go t.readLoop()
	go t.writeLoop()
	return nil
}

// Close the tracker connection.
func (t *UDPTracker) Close() error {
	close(t.writeC)
	if t.conn != nil {
		return t.conn.Close()
	}
	return nil
}

// readLoop reads datagrams from connection, finds the transaction and
// sends the bytes to the transaction's response channel.
func (t *UDPTracker) readLoop() {
	// Read buffer must be big enough to hold a UDP packet of maximum expected size.
	// Current value is: 320 = 20 + 50*6 (AnnounceResponse with 50 peers)
	buf := make([]byte, 320)
	for {
		n, err := t.conn.Read(buf)
		if err != nil {
			t.log.Error(err)
			if nerr, ok := err.(net.Error); ok && !nerr.Temporary() {
				t.log.Debug("End of tracker read loop")
				return
			}
			continue
		}
		t.log.Debug("Read ", n, " bytes")

		var header udpMessageHeader
		if n < binary.Size(header) {
			t.log.Error("response is too small")
			continue
		}

		err = binary.Read(bytes.NewReader(buf), binary.BigEndian, &header)
		if err != nil {
			t.log.Error(err)
			continue
		}

		t.transactionsM.Lock()
		trx, ok := t.transactions[header.TransactionID]
		delete(t.transactions, header.TransactionID)
		t.transactionsM.Unlock()
		if !ok {
			t.log.Errorln("unexpected transaction_id:", header.TransactionID)
			continue
		}

		// Tracker has sent and error.
		if header.Action == actionError {
			// The part after the header is the error message.
			trx.err = tracker.Error(buf[binary.Size(header):])
			trx.Done()
			continue
		}

		// Copy data into a new slice because buf will be overwritten at next read.
		trx.response = make([]byte, n)
		copy(trx.response, buf)
		trx.Done()
	}
}

// writeLoop receives a request from t.transactionC, sets a random TransactionID
// and sends it to the tracker.
func (t *UDPTracker) writeLoop() {
	var connectionID int64
	var connectionIDtime time.Time

	for trx := range t.writeC {
		if time.Since(connectionIDtime) > connectionIDInterval {
			connectionID = t.connect()
			connectionIDtime = time.Now()
		}
		trx.request.SetConnectionID(connectionID)

		t.writeTrx(trx)
	}
}

func (t *UDPTracker) writeTrx(trx *transaction) {
	t.log.Debugln("Writing transaction. ID:", trx.ID())
	_, err := trx.request.WriteTo(t.conn)
	if err != nil {
		t.log.Error(err)
	}
}

// connect sends a connectRequest and returns a ConnectionID given by the tracker.
// On error, it backs off with the algorithm described in BEP15 and retries.
// It does not return until tracker sends a ConnectionID.
func (t *UDPTracker) connect() int64 {
	req := new(connectRequest)
	req.SetAction(actionConnect)
	req.SetConnectionID(connectionIDMagic)

	trx := newTransaction(req)

	for {
		data, err := t.retryTransaction(t.writeTrx, trx, nil) // Does not return until transaction is completed.
		if err != nil {
			t.log.Error(err)
			continue
		}

		var response connectResponse
		err = binary.Read(bytes.NewReader(data), binary.BigEndian, &response)
		if err != nil {
			t.log.Error(err)
			continue
		}

		if response.Action != actionConnect {
			t.log.Error("invalid action in connect response")
			continue
		}

		t.log.Debugf("connect Response: %#v\n", response)
		return response.ConnectionID
	}
}

func (t *UDPTracker) retryTransaction(f func(*transaction), trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	t.transactionsM.Lock()
	t.transactions[trx.ID()] = trx
	t.transactionsM.Unlock()

	ticker := backoff.NewTicker(new(udpBackOff))
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			f(trx)
		case <-trx.done:
			// transaction is deleted in readLoop()
			return trx.response, trx.err
		case <-cancel:
			t.transactionsM.Lock()
			delete(t.transactions, trx.ID())
			t.transactionsM.Unlock()
			return nil, tracker.ErrRequestCancelled
		}
	}
}

func (t *UDPTracker) sendTransaction(trx *transaction, cancel <-chan struct{}) ([]byte, error) {
	f := func(trx *transaction) { t.writeC <- trx }
	return t.retryTransaction(f, trx, cancel)
}

func (t *UDPTracker) Announce(transfer tracker.Transfer, e tracker.Event, cancel <-chan struct{}) (*tracker.AnnounceResponse, error) {
	t.dialMutex.Lock()
	if !t.connected {
		err := t.dial()
		if err != nil {
			t.dialMutex.Unlock()
			return nil, err
		}
		t.connected = true
	}
	t.dialMutex.Unlock()

	request := &announceRequest{
		InfoHash:   transfer.InfoHash(),
		PeerID:     transfer.PeerID(),
		Event:      e,
		IP:         0, // Tracker uses sender of this UDP packet.
		Key:        0, // TODO set it
		NumWant:    tracker.NumWant,
		Port:       uint16(transfer.Port()),
		Extensions: 0,
	}
	request.SetAction(actionAnnounce)

	request2 := &transferAnnounceRequest{
		transfer:        transfer,
		announceRequest: request,
		urlData:         t.url.RequestURI(),
	}
	trx := newTransaction(request2)

	// t.request may block, that's why we pass cancel as argument.
	reply, err := t.sendTransaction(trx, cancel)
	if err != nil {
		if err, ok := err.(tracker.Error); ok {
			return &tracker.AnnounceResponse{Error: err}, nil
		}
		return nil, err
	}

	response, peers, err := t.parseAnnounceResponse(reply)
	if err != nil {
		return nil, err
	}
	t.log.Debugf("Announce response: %#v", response)

	return &tracker.AnnounceResponse{
		Interval: time.Duration(response.Interval) * time.Second,
		Leechers: response.Leechers,
		Seeders:  response.Seeders,
		Peers:    peers,
	}, nil
}

func (t *UDPTracker) parseAnnounceResponse(data []byte) (*udpAnnounceResponse, []*net.TCPAddr, error) {
	response := new(udpAnnounceResponse)
	if len(data) < binary.Size(response) {
		return nil, nil, errors.New("response is too small")
	}

	reader := bytes.NewReader(data)

	err := binary.Read(reader, binary.BigEndian, response)
	if err != nil {
		return nil, nil, err
	}
	t.log.Debugf("annouceResponse: %#v", response)

	if response.Action != actionAnnounce {
		return nil, nil, errors.New("invalid action")
	}

	peers, err := tracker.ParsePeersBinary(reader, t.log)
	if err != nil {
		return nil, nil, err
	}

	return response, peers, nil
}
<|endoftext|>


<commit_before>package command

import (
	"context"
	"fmt"
	"github.com/golang/protobuf/jsonpb"
	jsoniter "github.com/json-iterator/go"
	"github.com/olivere/elastic/v7"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func init() {
	cmdFilerMetaTail.Run = runFilerMetaTail // break init cycle
}

var cmdFilerMetaTail = &Command{
	UsageLine: "filer.meta.tail [-filer=localhost:8888] [-target=/]",
	Short:     "see continuous changes on a filer",
	Long: `See continuous changes on a filer.

	weed filer.meta.tail -timeAgo=30h | grep truncate
	weed filer.meta.tail -timeAgo=30h | jq .
	weed filer.meta.tail -timeAgo=30h | jq .eventNotification.newEntry.name

  `,
}

var (
	tailFiler   = cmdFilerMetaTail.Flag.String("filer", "localhost:8888", "filer hostname:port")
	tailTarget  = cmdFilerMetaTail.Flag.String("pathPrefix", "/", "path to a folder or file, or common prefix for the folders or files on filer")
	tailStart   = cmdFilerMetaTail.Flag.Duration("timeAgo", 0, "start time before now. \"300ms\", \"1.5h\" or \"2h45m\". Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\"")
	tailPattern = cmdFilerMetaTail.Flag.String("pattern", "", "full path or just filename pattern, ex: \"/home/?opher\", \"*.pdf\", see https://golang.org/pkg/path/filepath/#Match ")
	esServers   = cmdFilerMetaTail.Flag.String("es", "", "comma-separated elastic servers http://<host:port>")
	esIndex     = cmdFilerMetaTail.Flag.String("es.index", "seaweedfs", "ES index name")
)

func runFilerMetaTail(cmd *Command, args []string) bool {

	grpcDialOption := security.LoadClientTLS(util.GetViper(), "grpc.client")

	var filterFunc func(dir, fname string) bool
	if *tailPattern != "" {
		if strings.Contains(*tailPattern, "/") {
			println("watch path pattern", *tailPattern)
			filterFunc = func(dir, fname string) bool {
				matched, err := filepath.Match(*tailPattern, dir+"/"+fname)
				if err != nil {
					fmt.Printf("error: %v", err)
				}
				return matched
			}
		} else {
			println("watch file pattern", *tailPattern)
			filterFunc = func(dir, fname string) bool {
				matched, err := filepath.Match(*tailPattern, fname)
				if err != nil {
					fmt.Printf("error: %v", err)
				}
				return matched
			}
		}
	}

	shouldPrint := func(resp *filer_pb.SubscribeMetadataResponse) bool {
		if filterFunc == nil {
			return true
		}
		if resp.EventNotification.OldEntry == nil && resp.EventNotification.NewEntry == nil {
			return false
		}
		if resp.EventNotification.OldEntry != nil && filterFunc(resp.Directory, resp.EventNotification.OldEntry.Name) {
			return true
		}
		if resp.EventNotification.NewEntry != nil && filterFunc(resp.EventNotification.NewParentPath, resp.EventNotification.NewEntry.Name) {
			return true
		}
		return false
	}

	jsonpbMarshaler := jsonpb.Marshaler{
		EmitDefaults: false,
	}
	eachEntryFunc := func(resp *filer_pb.SubscribeMetadataResponse) error {
		jsonpbMarshaler.Marshal(os.Stdout, resp)
		fmt.Fprintln(os.Stdout)
		return nil
	}
	if *esServers != "" {
		var err error
		eachEntryFunc, err = sendToElasticSearchFunc(*esServers, *esIndex)
		if err != nil {
			fmt.Printf("create elastic search client to %s: %+v\n", *esServers, err)
			return false
		}
	}

	tailErr := pb.WithFilerClient(*tailFiler, grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		stream, err := client.SubscribeMetadata(ctx, &filer_pb.SubscribeMetadataRequest{
			ClientName: "tail",
			PathPrefix: *tailTarget,
			SinceNs:    time.Now().Add(-*tailStart).UnixNano(),
		})
		if err != nil {
			return fmt.Errorf("listen: %v", err)
		}

		for {
			resp, listenErr := stream.Recv()
			if listenErr == io.EOF {
				return nil
			}
			if listenErr != nil {
				return listenErr
			}
			if !shouldPrint(resp) {
				continue
			}
			if err = eachEntryFunc(resp); err != nil {
				return err
			}
		}

	})
	if tailErr != nil {
		fmt.Printf("tail %s: %v\n", *tailFiler, tailErr)
	}

	return true
}

type EsDocument struct {
	Dir         string `json:"dir,omitempty"`
	Name        string `json:"name,omitempty"`
	IsDirectory bool   `json:"isDir,omitempty"`
	Size        uint64 `json:"size,omitempty"`
	Uid         uint32 `json:"uid,omitempty"`
	Gid         uint32 `json:"gid,omitempty"`
	UserName    string `json:"userName,omitempty"`
	Collection  string `json:"collection,omitempty"`
	Crtime      int64  `json:"crtime,omitempty"`
	Mtime       int64  `json:"mtime,omitempty"`
	Mime        string `json:"mime,omitempty"`
}

func toEsEntry(event *filer_pb.EventNotification) (*EsDocument, string) {
	entry := event.NewEntry
	dir, name := event.NewParentPath, entry.Name
	id := util.Md5String([]byte(util.NewFullPath(dir, name)))
	esEntry := &EsDocument{
		Dir:         dir,
		Name:        name,
		IsDirectory: entry.IsDirectory,
		Size:        entry.Attributes.FileSize,
		Uid:         entry.Attributes.Uid,
		Gid:         entry.Attributes.Gid,
		UserName:    entry.Attributes.UserName,
		Collection:  entry.Attributes.Collection,
		Crtime:      entry.Attributes.Crtime,
		Mtime:       entry.Attributes.Mtime,
		Mime:        entry.Attributes.Mime,
	}
	return esEntry, id
}

func sendToElasticSearchFunc(servers string, esIndex string) (func(resp *filer_pb.SubscribeMetadataResponse) error, error) {
	options := []elastic.ClientOptionFunc{}
	options = append(options, elastic.SetURL(strings.Split(servers, ",")...))
	options = append(options, elastic.SetSniff(false))
	options = append(options, elastic.SetHealthcheck(false))
	client, err := elastic.NewClient(options...)
	if err != nil {
		return nil, err
	}
	return func(resp *filer_pb.SubscribeMetadataResponse) error {
		event := resp.EventNotification
		if event.OldEntry != nil &&
			(event.NewEntry == nil || resp.Directory != event.NewParentPath || event.OldEntry.Name != event.NewEntry.Name) {
			// delete or not update the same file
			dir, name := resp.Directory, event.OldEntry.Name
			id := util.Md5String([]byte(util.NewFullPath(dir, name)))
			println("delete", id)
			_, err := client.Delete().Index(esIndex).Id(id).Do(context.Background())
			return err
		}
		if event.NewEntry != nil {
			// add a new file or update the same file
			esEntry, id := toEsEntry(event)
			value, err := jsoniter.Marshal(esEntry)
			if err != nil {
				return err
			}
			println(string(value))
			_, err = client.Index().Index(esIndex).Id(id).BodyJson(string(value)).Do(context.Background())
			return err
		}
		return nil
	}, nil
}
<commit_msg>adjust help message<commit_after>package command

import (
	"context"
	"fmt"
	"github.com/golang/protobuf/jsonpb"
	jsoniter "github.com/json-iterator/go"
	"github.com/olivere/elastic/v7"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func init() {
	cmdFilerMetaTail.Run = runFilerMetaTail // break init cycle
}

var cmdFilerMetaTail = &Command{
	UsageLine: "filer.meta.tail [-filer=localhost:8888] [-pathPrefix=/]",
	Short:     "see continuous changes on a filer",
	Long: `See continuous changes on a filer.

	weed filer.meta.tail -timeAgo=30h | grep truncate
	weed filer.meta.tail -timeAgo=30h | jq .
	weed filer.meta.tail -timeAgo=30h | jq .eventNotification.newEntry.name

  `,
}

var (
	tailFiler   = cmdFilerMetaTail.Flag.String("filer", "localhost:8888", "filer hostname:port")
	tailTarget  = cmdFilerMetaTail.Flag.String("pathPrefix", "/", "path to a folder or common prefix for the folders or files on filer")
	tailStart   = cmdFilerMetaTail.Flag.Duration("timeAgo", 0, "start time before now. \"300ms\", \"1.5h\" or \"2h45m\". Valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\"")
	tailPattern = cmdFilerMetaTail.Flag.String("pattern", "", "full path or just filename pattern, ex: \"/home/?opher\", \"*.pdf\", see https://golang.org/pkg/path/filepath/#Match ")
	esServers   = cmdFilerMetaTail.Flag.String("es", "", "comma-separated elastic servers http://<host:port>")
	esIndex     = cmdFilerMetaTail.Flag.String("es.index", "seaweedfs", "ES index name")
)

func runFilerMetaTail(cmd *Command, args []string) bool {

	grpcDialOption := security.LoadClientTLS(util.GetViper(), "grpc.client")

	var filterFunc func(dir, fname string) bool
	if *tailPattern != "" {
		if strings.Contains(*tailPattern, "/") {
			println("watch path pattern", *tailPattern)
			filterFunc = func(dir, fname string) bool {
				matched, err := filepath.Match(*tailPattern, dir+"/"+fname)
				if err != nil {
					fmt.Printf("error: %v", err)
				}
				return matched
			}
		} else {
			println("watch file pattern", *tailPattern)
			filterFunc = func(dir, fname string) bool {
				matched, err := filepath.Match(*tailPattern, fname)
				if err != nil {
					fmt.Printf("error: %v", err)
				}
				return matched
			}
		}
	}

	shouldPrint := func(resp *filer_pb.SubscribeMetadataResponse) bool {
		if filterFunc == nil {
			return true
		}
		if resp.EventNotification.OldEntry == nil && resp.EventNotification.NewEntry == nil {
			return false
		}
		if resp.EventNotification.OldEntry != nil && filterFunc(resp.Directory, resp.EventNotification.OldEntry.Name) {
			return true
		}
		if resp.EventNotification.NewEntry != nil && filterFunc(resp.EventNotification.NewParentPath, resp.EventNotification.NewEntry.Name) {
			return true
		}
		return false
	}

	jsonpbMarshaler := jsonpb.Marshaler{
		EmitDefaults: false,
	}
	eachEntryFunc := func(resp *filer_pb.SubscribeMetadataResponse) error {
		jsonpbMarshaler.Marshal(os.Stdout, resp)
		fmt.Fprintln(os.Stdout)
		return nil
	}
	if *esServers != "" {
		var err error
		eachEntryFunc, err = sendToElasticSearchFunc(*esServers, *esIndex)
		if err != nil {
			fmt.Printf("create elastic search client to %s: %+v\n", *esServers, err)
			return false
		}
	}

	tailErr := pb.WithFilerClient(*tailFiler, grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		stream, err := client.SubscribeMetadata(ctx, &filer_pb.SubscribeMetadataRequest{
			ClientName: "tail",
			PathPrefix: *tailTarget,
			SinceNs:    time.Now().Add(-*tailStart).UnixNano(),
		})
		if err != nil {
			return fmt.Errorf("listen: %v", err)
		}

		for {
			resp, listenErr := stream.Recv()
			if listenErr == io.EOF {
				return nil
			}
			if listenErr != nil {
				return listenErr
			}
			if !shouldPrint(resp) {
				continue
			}
			if err = eachEntryFunc(resp); err != nil {
				return err
			}
		}

	})
	if tailErr != nil {
		fmt.Printf("tail %s: %v\n", *tailFiler, tailErr)
	}

	return true
}

type EsDocument struct {
	Dir         string `json:"dir,omitempty"`
	Name        string `json:"name,omitempty"`
	IsDirectory bool   `json:"isDir,omitempty"`
	Size        uint64 `json:"size,omitempty"`
	Uid         uint32 `json:"uid,omitempty"`
	Gid         uint32 `json:"gid,omitempty"`
	UserName    string `json:"userName,omitempty"`
	Collection  string `json:"collection,omitempty"`
	Crtime      int64  `json:"crtime,omitempty"`
	Mtime       int64  `json:"mtime,omitempty"`
	Mime        string `json:"mime,omitempty"`
}

func toEsEntry(event *filer_pb.EventNotification) (*EsDocument, string) {
	entry := event.NewEntry
	dir, name := event.NewParentPath, entry.Name
	id := util.Md5String([]byte(util.NewFullPath(dir, name)))
	esEntry := &EsDocument{
		Dir:         dir,
		Name:        name,
		IsDirectory: entry.IsDirectory,
		Size:        entry.Attributes.FileSize,
		Uid:         entry.Attributes.Uid,
		Gid:         entry.Attributes.Gid,
		UserName:    entry.Attributes.UserName,
		Collection:  entry.Attributes.Collection,
		Crtime:      entry.Attributes.Crtime,
		Mtime:       entry.Attributes.Mtime,
		Mime:        entry.Attributes.Mime,
	}
	return esEntry, id
}

func sendToElasticSearchFunc(servers string, esIndex string) (func(resp *filer_pb.SubscribeMetadataResponse) error, error) {
	options := []elastic.ClientOptionFunc{}
	options = append(options, elastic.SetURL(strings.Split(servers, ",")...))
	options = append(options, elastic.SetSniff(false))
	options = append(options, elastic.SetHealthcheck(false))
	client, err := elastic.NewClient(options...)
	if err != nil {
		return nil, err
	}
	return func(resp *filer_pb.SubscribeMetadataResponse) error {
		event := resp.EventNotification
		if event.OldEntry != nil &&
			(event.NewEntry == nil || resp.Directory != event.NewParentPath || event.OldEntry.Name != event.NewEntry.Name) {
			// delete or not update the same file
			dir, name := resp.Directory, event.OldEntry.Name
			id := util.Md5String([]byte(util.NewFullPath(dir, name)))
			println("delete", id)
			_, err := client.Delete().Index(esIndex).Id(id).Do(context.Background())
			return err
		}
		if event.NewEntry != nil {
			// add a new file or update the same file
			esEntry, id := toEsEntry(event)
			value, err := jsoniter.Marshal(esEntry)
			if err != nil {
				return err
			}
			println(string(value))
			_, err = client.Index().Index(esIndex).Id(id).BodyJson(string(value)).Do(context.Background())
			return err
		}
		return nil
	}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"net"
	"path/filepath"
	"fmt"
	"strings"
);

func main() {
	logger := log.New(os.Stdout, "", log.LstdFlags);

	sigtoolPath := findSigtoolPath();
	logger.Printf("ClamAV executable sigtool found at path: %v", sigtoolPath);

	var mirrorDomain string = "current.cvd.clamav.net";
	var mirrorTxtRecord string = pullTxtRecord(mirrorDomain);
	logger.Printf("TXT record for [%v]: %v", mirrorDomain, mirrorTxtRecord);

	clamav, mainv, dailyv, x, y, z, safebrowsingv, bytecodev := parseTxtRecord(mirrorTxtRecord);
	logger.Printf("TXT record values parsed: " +
		"[clamav=%v,mainv=%v,dailyv=%v,x=%v,y=%v,z=%v,safebrowsingv=%v,bytecodev=%v",
		clamav, mainv, dailyv, x, y, z, safebrowsingv, bytecodev);
}

func pullTxtRecord(mirrorDomain string) (string) {
	mirrorTxtRecords, err := net.LookupTXT(mirrorDomain);

	if (err != nil) {
		msg := fmt.Sprintf("Unable to resolve TXT record for %v", mirrorDomain);
		log.Fatal(msg, err);
		os.Exit(1);
	}

	if (len(mirrorTxtRecords) < 1) {
		msg := fmt.Sprintf("No TXT records returned for %v", mirrorDomain);
		log.Fatal(msg);
		os.Exit(1);
	}

	return mirrorTxtRecords[0];
}

func parseTxtRecord(mirrorTxtRecord string) (string, string, string, string, string, string, string, string) {
	s := strings.SplitN(mirrorTxtRecord, ":", 8);
	return s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7];
}

func updateFile(file string) {

}

func findSigtoolPath() (path string) {
	var execName string = "sigtool";
	var separator string = string(os.PathSeparator);
	var envPathSeparator string = string(os.PathListSeparator);
	var envPath string = os.Getenv("PATH");

	var localPath string = "." + separator + execName;
	if (exists(localPath)) {
		execPath, err := filepath.Abs(localPath);

		if (err != nil) {
			msg := fmt.Sprintf("Error parsing absolute path for: %v", localPath);
			log.Fatal(msg, err);
			os.Exit(1);
		}

		return execPath;
	}

	for _, pathElement := range strings.Split(envPath, envPathSeparator) {
		var execPath string = pathElement + separator + execName;
		if (exists(execPath)) {
			return execPath;
		}
	}

	log.Fatal("The ClamAV executable sigtool was not found in the " +
		"current directory nor in the system path");
	os.Exit(1);

	return;
}

func exists(filePath string) (exists bool) {
	exists = true

	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		exists = false
	}

	return;
}<commit_msg>Added command line options support.<commit_after>package main

import (
	"log"
	"os"
	"net"
	"github.com/pborman/getopt"
	"path/filepath"
	"fmt"
	"strings"
);

func main() {
	logger := log.New(os.Stdout, "", log.LstdFlags);
	verboseMode := parseCliFlags();

	sigtoolPath := findSigtoolPath();
	if (verboseMode) {
		logger.Printf("ClamAV executable sigtool found at path: %v", sigtoolPath);
	}

	var mirrorDomain string = "current.cvd.clamav.net";
	var mirrorTxtRecord string = pullTxtRecord(mirrorDomain);

	if (verboseMode) {
		logger.Printf("TXT record for [%v]: %v", mirrorDomain, mirrorTxtRecord);
	}

	clamav, mainv, dailyv, x, y, z, safebrowsingv, bytecodev := parseTxtRecord(mirrorTxtRecord);

	if (verboseMode) {
		logger.Printf("TXT record values parsed: " +
			"[clamav=%v,mainv=%v,dailyv=%v,x=%v,y=%v,z=%v,safebrowsingv=%v,bytecodev=%v",
			clamav, mainv, dailyv, x, y, z, safebrowsingv, bytecodev);
	}
}

func parseCliFlags() (bool) {
	verbosePart := getopt.BoolLong("verbose", 'v', "Enable verbose mode with additional debugging information");
	getopt.Parse();

	var verboseMode bool = *verbosePart;

	return verboseMode;
}

func help() {
	getopt.Usage();
	os.Exit(0);
}

func pullTxtRecord(mirrorDomain string) (string) {
	mirrorTxtRecords, err := net.LookupTXT(mirrorDomain);

	if (err != nil) {
		msg := fmt.Sprintf("Unable to resolve TXT record for %v", mirrorDomain);
		log.Fatal(msg, err);
	}

	if (len(mirrorTxtRecords) < 1) {
		msg := fmt.Sprintf("No TXT records returned for %v", mirrorDomain);
		log.Fatal(msg);
	}

	return mirrorTxtRecords[0];
}

func parseTxtRecord(mirrorTxtRecord string) (string, string, string, string, string, string, string, string) {
	s := strings.SplitN(mirrorTxtRecord, ":", 8);
	return s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7];
}

func updateFile(file string) {

}

func findSigtoolPath() (path string) {
	var execName string = "sigtool";
	var separator string = string(os.PathSeparator);
	var envPathSeparator string = string(os.PathListSeparator);
	var envPath string = os.Getenv("PATH");

	var localPath string = "." + separator + execName;
	if (exists(localPath)) {
		execPath, err := filepath.Abs(localPath);

		if (err != nil) {
			msg := fmt.Sprintf("Error parsing absolute path for: %v", localPath);
			log.Fatal(msg, err);
		}

		return execPath;
	}

	for _, pathElement := range strings.Split(envPath, envPathSeparator) {
		var execPath string = pathElement + separator + execName;
		if (exists(execPath)) {
			return execPath;
		}
	}

	log.Fatal("The ClamAV executable sigtool was not found in the " +
		"current directory nor in the system path");

	return;
}

func exists(filePath string) (exists bool) {
	exists = true

	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		exists = false
	}

	return;
}<|endoftext|>


<commit_before>package libproxy

import (
	"bytes"
	"encoding/binary"
	"io"
	"log"
	"net"
	"sync"
)

// UDPListener defines a listener interface to read, write and close a UDP connection
type UDPListener interface {
	ReadFromUDP(b []byte) (int, *net.UDPAddr, error)
	WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
	Close() error
}

// udpEncapsulator encapsulates a UDP connection and listener
type udpEncapsulator struct {
	conn     *net.Conn
	listener net.Listener
	m        *sync.Mutex
	r        *sync.Mutex
	w        *sync.Mutex
}

func (u *udpEncapsulator) getConn() (net.Conn, error) {
	u.m.Lock()
	defer u.m.Unlock()
	if u.conn != nil {
		return *u.conn, nil
	}
	conn, err := u.listener.Accept()
	if err != nil {
		log.Printf("Failed to accept connection: %#v", err)
		return nil, err
	}
	u.conn = &conn
	return conn, nil
}

// ReadFromUDP reads the bytestream from a udpEncapsulator, returning the
// number of bytes read and the unpacked UDPAddr struct
func (u *udpEncapsulator) ReadFromUDP(b []byte) (int, *net.UDPAddr, error) {
	conn, err := u.getConn()
	if err != nil {
		return 0, nil, err
	}
	u.r.Lock()
	defer u.r.Unlock()
	datagram := &udpDatagram{payload: b}
	length, err := datagram.Unmarshal(conn)
	if err != nil {
		return 0, nil, err
	}
	udpAddr := net.UDPAddr{IP: *datagram.IP, Port: int(datagram.Port), Zone: datagram.Zone}
	return length, &udpAddr, nil
}

// WriteToUDP writes a bytestream to a specified UDPAddr, returning the number
// of bytes successfully written
func (u *udpEncapsulator) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
	conn, err := u.getConn()
	if err != nil {
		return 0, err
	}
	u.w.Lock()
	defer u.w.Unlock()
	datagram := &udpDatagram{payload: b, IP: &addr.IP, Port: uint16(addr.Port), Zone: addr.Zone}
	return len(b), datagram.Marshal(conn)
}

// Close closes the connection in the udpEncapsulator
func (u *udpEncapsulator) Close() error {
	if u.conn != nil {
		conn := *u.conn
		conn.Close()
	}
	u.listener.Close()
	return nil
}

// NewUDPConn initializes a new UDP connection
func NewUDPConn(conn net.Conn) UDPListener {
	var m sync.Mutex
	var r sync.Mutex
	var w sync.Mutex
	return &udpEncapsulator{
		conn:     &conn,
		listener: nil,
		m:        &m,
		r:        &r,
		w:        &w,
	}
}

// NewUDPListener initializes a new UDP listener
func NewUDPListener(listener net.Listener) UDPListener {
	var m sync.Mutex
	var r sync.Mutex
	var w sync.Mutex
	return &udpEncapsulator{
		conn:     nil,
		listener: listener,
		m:        &m,
		r:        &r,
		w:        &w,
	}
}

type udpDatagram struct {
	payload []byte
	IP      *net.IP
	Port    uint16
	Zone    string
}

// Marshal marshals data from the udpDatagram to the provided connection
func (u *udpDatagram) Marshal(conn net.Conn) error {
	// marshal the variable length header to a temporary buffer
	var header bytes.Buffer
	var length uint16
	length = uint16(len(*u.IP))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return err
	}

	if err := binary.Write(&header, binary.LittleEndian, u.IP); err != nil {
		return err
	}

	if err := binary.Write(&header, binary.LittleEndian, &u.Port); err != nil {
		return err
	}

	length = uint16(len(u.Zone))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return err
	}

	if err := binary.Write(&header, binary.LittleEndian, []byte(u.Zone)); err != nil {
		return nil
	}

	length = uint16(len(u.payload))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return nil
	}

	length = uint16(2 + header.Len() + len(u.payload))
	if err := binary.Write(conn, binary.LittleEndian, &length); err != nil {
		return nil
	}
	_, err := io.Copy(conn, &header)
	if err != nil {
		return err
	}
	payload := bytes.NewBuffer(u.payload)
	_, err = io.Copy(conn, payload)
	if err != nil {
		return err
	}
	return nil
}

// Unmarshal unmarshals data from the connection to the udpDatagram
func (u *udpDatagram) Unmarshal(conn net.Conn) (int, error) {
	var length uint16
	// frame length
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	var IP net.IP
	IP = make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &IP); err != nil {
		return 0, err
	}
	u.IP = &IP
	if err := binary.Read(conn, binary.LittleEndian, &u.Port); err != nil {
		return 0, err
	}
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	Zone := make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &Zone); err != nil {
		return 0, err
	}
	u.Zone = string(Zone)
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	_, err := io.ReadFull(conn, u.payload[0:length])
	if err != nil {
		return 0, err
	}
	return int(length), nil
}
<commit_msg>vpnkit-forwarder: remove dead code<commit_after>package libproxy

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"sync"
)

// UDPListener defines a listener interface to read, write and close a UDP connection
type UDPListener interface {
	ReadFromUDP(b []byte) (int, *net.UDPAddr, error)
	WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
	Close() error
}

// udpEncapsulator encapsulates a UDP connection and listener
type udpEncapsulator struct {
	conn net.Conn
	m    *sync.Mutex
	r    *sync.Mutex
	w    *sync.Mutex
}

// ReadFromUDP reads the bytestream from a udpEncapsulator, returning the
// number of bytes read and the unpacked UDPAddr struct
func (u *udpEncapsulator) ReadFromUDP(b []byte) (int, *net.UDPAddr, error) {
	u.r.Lock()
	defer u.r.Unlock()
	datagram := &udpDatagram{payload: b}
	length, err := datagram.Unmarshal(u.conn)
	if err != nil {
		return 0, nil, err
	}
	udpAddr := net.UDPAddr{IP: *datagram.IP, Port: int(datagram.Port), Zone: datagram.Zone}
	return length, &udpAddr, nil
}

// WriteToUDP writes a bytestream to a specified UDPAddr, returning the number
// of bytes successfully written
func (u *udpEncapsulator) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
	u.w.Lock()
	defer u.w.Unlock()
	datagram := &udpDatagram{payload: b, IP: &addr.IP, Port: uint16(addr.Port), Zone: addr.Zone}
	return len(b), datagram.Marshal(u.conn)
}

// Close closes the connection in the udpEncapsulator
func (u *udpEncapsulator) Close() error {
	return u.conn.Close()
}

// NewUDPConn initializes a new UDP connection
func NewUDPConn(conn net.Conn) UDPListener {
	var m sync.Mutex
	var r sync.Mutex
	var w sync.Mutex
	return &udpEncapsulator{
		conn: conn,
		m:    &m,
		r:    &r,
		w:    &w,
	}
}

type udpDatagram struct {
	payload []byte
	IP      *net.IP
	Port    uint16
	Zone    string
}

// Marshal marshals data from the udpDatagram to the provided connection
func (u *udpDatagram) Marshal(conn net.Conn) error {
	// marshal the variable length header to a temporary buffer
	var header bytes.Buffer
	var length uint16
	length = uint16(len(*u.IP))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return err
	}

	if err := binary.Write(&header, binary.LittleEndian, u.IP); err != nil {
		return err
	}

	if err := binary.Write(&header, binary.LittleEndian, &u.Port); err != nil {
		return err
	}

	length = uint16(len(u.Zone))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return err
	}

	if err := binary.Write(&header, binary.LittleEndian, []byte(u.Zone)); err != nil {
		return nil
	}

	length = uint16(len(u.payload))
	if err := binary.Write(&header, binary.LittleEndian, &length); err != nil {
		return nil
	}

	length = uint16(2 + header.Len() + len(u.payload))
	if err := binary.Write(conn, binary.LittleEndian, &length); err != nil {
		return nil
	}
	_, err := io.Copy(conn, &header)
	if err != nil {
		return err
	}
	payload := bytes.NewBuffer(u.payload)
	_, err = io.Copy(conn, payload)
	if err != nil {
		return err
	}
	return nil
}

// Unmarshal unmarshals data from the connection to the udpDatagram
func (u *udpDatagram) Unmarshal(conn net.Conn) (int, error) {
	var length uint16
	// frame length
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	var IP net.IP
	IP = make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &IP); err != nil {
		return 0, err
	}
	u.IP = &IP
	if err := binary.Read(conn, binary.LittleEndian, &u.Port); err != nil {
		return 0, err
	}
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	Zone := make([]byte, length)
	if err := binary.Read(conn, binary.LittleEndian, &Zone); err != nil {
		return 0, err
	}
	u.Zone = string(Zone)
	if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
		return 0, err
	}
	_, err := io.ReadFull(conn, u.payload[0:length])
	if err != nil {
		return 0, err
	}
	return int(length), nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"socialapi/config"
	"socialapi/db"
	"socialapi/models"
	"socialapi/workers/helper"
)

var (
	flagProfile = flag.String("c", "", "Configuration profile from file")
	flagDebug   = flag.Bool("d", false, "Debug mode")
)

func main() {
	flag.Parse()
	if *flagProfile == "" {
		fmt.Println("Please define config file with -c", "Exiting...")
		return
	}

	conf := config.Read(*flagProfile)
	// create logger for our package
	log := helper.CreateLogger("DB Creation Script", *flagDebug)

	helper.MustInitBongo(conf, log)
	db.DB.LogMode(true)
	db.DB.Exec("drop table channel_message_list;")
	db.DB.Exec("drop table channel_message;")
	db.DB.Exec("drop table message_reply;")
	db.DB.Exec("drop table channel_participant;")
	db.DB.Exec("drop table channel;")
	db.DB.Exec("drop table interaction;")
	db.DB.Exec("drop table account;")

	if err := db.DB.CreateTable(&models.ChannelMessage{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.MessageReply{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.Channel{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.ChannelMessageList{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.ChannelParticipant{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.Interaction{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.Account{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
}
<commit_msg>Social: add schema to the tables<commit_after>package main

import (
	"flag"
	"fmt"
	"socialapi/config"
	"socialapi/db"
	"socialapi/models"
	"socialapi/workers/helper"
)

var (
	flagProfile = flag.String("c", "", "Configuration profile from file")
	flagDebug   = flag.Bool("d", false, "Debug mode")
)

func main() {
	flag.Parse()
	if *flagProfile == "" {
		fmt.Println("Please define config file with -c", "Exiting...")
		return
	}

	conf := config.Read(*flagProfile)
	// create logger for our package
	log := helper.CreateLogger("DB Creation Script", *flagDebug)

	helper.MustInitBongo(conf, log)
	db.DB.LogMode(true)
	db.DB.Exec("drop table api.channel_message_list;")
	db.DB.Exec("drop table api.channel_message;")
	db.DB.Exec("drop table api.message_reply;")
	db.DB.Exec("drop table api.channel_participant;")
	db.DB.Exec("drop table api.channel;")
	db.DB.Exec("drop table api.interaction;")
	db.DB.Exec("drop table api.account;")

	if err := db.DB.CreateTable(&models.ChannelMessage{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.MessageReply{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.Channel{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.ChannelMessageList{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.ChannelParticipant{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.Interaction{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
	if err := db.DB.CreateTable(&models.Account{}).Error; err != nil {
		panic(fmt.Sprintf("No error should happen when create table, but got %+v", err))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"strings"
)

func main() {
	//go BlobServerStart()
	//hashfindwalk()
	in := bufio.NewReader(os.Stdin)
	input, err := in.ReadString('\n')
	if err != nil {
		panic(err)
	}
	fmt.Println(input)
	return
}

type hkid []byte

func get(objecthash hkid, path string) (b blob, err error) {
	typeString := "commit"
	//objecthash := hkid
	err = nil
	nameSegments := strings.SplitN(path, "/", 2)
	for {
		switch typeString {
		case "blob":
			if len(nameSegments) < 2 {
				b, err = GetBlob(objecthash)
				return
			}
		case "list":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			l, _ := GetList(objecthash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
		case "tag":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			t, err := GetTag(objecthash, nameSegments[0])
			if !t.Verifiy() {
				b = nil
				err = errors.New("Tag Verifiy Failed")
				return b, err
			}
			typeString = t.TypeString
			objecthash = t.HashBytes
		case "commit":
			c, err := GetCommit(objecthash)
			if !c.Verifiy() {
				b = nil
				err = errors.New("Commit Verifiy Failed")
				return b, err
			}
			l, err := GetList(c.listHash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
		}
		if err != nil{
			return nil, err
		}
	}
	b = nil
	return
}

func post(objecthash hkid, path string, b blob) (err error) {
	typeString := "commit"
	//objecthash := hkid
	err = nil
	nameSegments := strings.SplitN(path, "/", 2)
	regenlist := []blob{}
	regenpath := []string{}
	for {
		switch typeString {
		case "blob":
			if len(nameSegments) < 2 {
				b, err = GetBlob(objecthash)
				return
			}
		case "list":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			l, _ := GetList(objecthash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
			if err == nil {
				regenlist = append(regenlist, l.Bytes())
				regenpath = append(regenpath, nameSegments[0])
			} else {
				regen(regenlist, objecthash, b)
			}
		case "tag":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			t, _ := GetTag(objecthash, nameSegments[0])
			if !t.Verifiy() {
				return errors.New("Tag Verifiy Failed")
			}
			typeString = t.TypeString
			objecthash = t.HashBytes
		case "commit":
			c, _ := GetCommit(objecthash)
			if !c.Verifiy() {
				return errors.New("Commit Verifiy Failed")
			}
			l, _ := GetList(c.listHash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
		}
		if err != nil{
			return err
		}
		if typeString != "list"{
			regenlist = nil
			regenpath = []string{}
		}
	}
	b = nil
	return
}

func regen(objectsToRegen [][]byte, objecthash hkid, b blob) error {
	return nil
}

func initRepo(objecthash hkid, path string) (repoHkid []byte) {
	return nil
}

func initDomain(objecthash hkid, path string) (domainHkid []byte) {
	return nil
}
<commit_msg>testing golangs list<commit_after>package main

import (
	"bufio"
	golist "container/list"
	"errors"
	"fmt"
	"os"
	"strings"
)

func main() {
	//go BlobServerStart()
	//hashfindwalk()
	in := bufio.NewReader(os.Stdin)
	input, err := in.ReadString('\n')
	if err != nil {
		panic(err)
	}
	fmt.Println(input)
	return
}

type hkid []byte

func get(objecthash hkid, path string) (b blob, err error) {
	typeString := "commit"
	//objecthash := hkid
	err = nil
	nameSegments := strings.SplitN(path, "/", 2)
	for {
		switch typeString {
		case "blob":
			if len(nameSegments) < 2 {
				b, err = GetBlob(objecthash)
				return
			}
		case "list":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			l, _ := GetList(objecthash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
		case "tag":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			t, err := GetTag(objecthash, nameSegments[0])
			if !t.Verifiy() {
				b = nil
				err = errors.New("Tag Verifiy Failed")
				return b, err
			}
			typeString = t.TypeString
			objecthash = t.HashBytes
		case "commit":
			c, err := GetCommit(objecthash)
			if !c.Verifiy() {
				b = nil
				err = errors.New("Commit Verifiy Failed")
				return b, err
			}
			l, err := GetList(c.listHash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
		}
		if err != nil {
			return nil, err
		}
	}
	b = nil
	return
}

func post(objecthash hkid, path string, b blob) (err error) {
	typeString := "commit"
	//objecthash := hkid
	err = nil
	nameSegments := strings.SplitN(path, "/", 2)
	regenlist := golist.New()
	regenpath := []string{}
	for {
		switch typeString {
		case "blob":
			if len(nameSegments) < 2 {
				b, err = GetBlob(objecthash)
				return
			}
		case "list":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			l, _ := GetList(objecthash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
			if err == nil {
				regenlist.PushBack(l)
				regenpath = append(regenpath, nameSegments[0])
			} else {
				regen(regenlist, objecthash, b)
			}
		case "tag":
			nameSegments = strings.SplitN(nameSegments[1], "/", 2)
			t, _ := GetTag(objecthash, nameSegments[0])
			if !t.Verifiy() {
				return errors.New("Tag Verifiy Failed")
			}
			typeString = t.TypeString
			objecthash = t.HashBytes
			regenlist.Init()
			regenlist.PushBack(t)
		case "commit":
			c, _ := GetCommit(objecthash)
			if !c.Verifiy() {
				return errors.New("Commit Verifiy Failed")
			}
			regenlist.Init()
			regenlist.PushBack(c)
			l, _ := GetList(c.listHash)
			typeString, objecthash = l.hash_for_namesegment(nameSegments[0])
		}
		if err != nil {
			return err
		}
		if typeString != "list" {

			regenpath = []string{}
		}
	}
	b = nil
	return
}

func regen(objectsToRegen *golist.List, objecthash hkid, b blob) error {
	return nil
}

func initRepo(objecthash hkid, path string) (repoHkid []byte) {
	return nil
}

func initDomain(objecthash hkid, path string) (domainHkid []byte) {
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes 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 metrics

import (
	"context"
	"fmt"
	"sync"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/master/ports"
	e2epod "k8s.io/kubernetes/test/e2e/framework/pod"
	"k8s.io/kubernetes/test/e2e/system"

	"k8s.io/klog"
)

// Collection is metrics collection of components
type Collection struct {
	APIServerMetrics         APIServerMetrics
	ControllerManagerMetrics ControllerManagerMetrics
	KubeletMetrics           map[string]KubeletMetrics
	SchedulerMetrics         SchedulerMetrics
	ClusterAutoscalerMetrics ClusterAutoscalerMetrics
}

// Grabber provides functions which grab metrics from components
type Grabber struct {
	client                            clientset.Interface
	externalClient                    clientset.Interface
	grabFromAPIServer                 bool
	grabFromControllerManager         bool
	grabFromKubelets                  bool
	grabFromScheduler                 bool
	grabFromClusterAutoscaler         bool
	masterName                        string
	registeredMaster                  bool
	waitForControllerManagerReadyOnce sync.Once
}

// NewMetricsGrabber returns new metrics which are initialized.
func NewMetricsGrabber(c clientset.Interface, ec clientset.Interface, kubelets bool, scheduler bool, controllers bool, apiServer bool, clusterAutoscaler bool) (*Grabber, error) {
	registeredMaster := false
	masterName := ""
	nodeList, err := c.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, err
	}
	if len(nodeList.Items) < 1 {
		klog.Warning("Can't find any Nodes in the API server to grab metrics from")
	}
	for _, node := range nodeList.Items {
		if system.DeprecatedMightBeMasterNode(node.Name) {
			registeredMaster = true
			masterName = node.Name
			break
		}
	}
	if !registeredMaster {
		scheduler = false
		controllers = false
		clusterAutoscaler = ec != nil
		if clusterAutoscaler {
			klog.Warningf("Master node is not registered. Grabbing metrics from Scheduler, ControllerManager is disabled.")
		} else {
			klog.Warningf("Master node is not registered. Grabbing metrics from Scheduler, ControllerManager and ClusterAutoscaler is disabled.")
		}
	}

	return &Grabber{
		client:                    c,
		externalClient:            ec,
		grabFromAPIServer:         apiServer,
		grabFromControllerManager: controllers,
		grabFromKubelets:          kubelets,
		grabFromScheduler:         scheduler,
		grabFromClusterAutoscaler: clusterAutoscaler,
		masterName:                masterName,
		registeredMaster:          registeredMaster,
	}, nil
}

// HasRegisteredMaster returns if metrics grabber was able to find a master node
func (g *Grabber) HasRegisteredMaster() bool {
	return g.registeredMaster
}

// GrabFromKubelet returns metrics from kubelet
func (g *Grabber) GrabFromKubelet(nodeName string) (KubeletMetrics, error) {
	nodes, err := g.client.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{FieldSelector: fields.Set{"metadata.name": nodeName}.AsSelector().String()})
	if err != nil {
		return KubeletMetrics{}, err
	}
	if len(nodes.Items) != 1 {
		return KubeletMetrics{}, fmt.Errorf("Error listing nodes with name %v, got %v", nodeName, nodes.Items)
	}
	kubeletPort := nodes.Items[0].Status.DaemonEndpoints.KubeletEndpoint.Port
	return g.grabFromKubeletInternal(nodeName, int(kubeletPort))
}

func (g *Grabber) grabFromKubeletInternal(nodeName string, kubeletPort int) (KubeletMetrics, error) {
	if kubeletPort <= 0 || kubeletPort > 65535 {
		return KubeletMetrics{}, fmt.Errorf("Invalid Kubelet port %v. Skipping Kubelet's metrics gathering", kubeletPort)
	}
	output, err := g.getMetricsFromNode(nodeName, int(kubeletPort))
	if err != nil {
		return KubeletMetrics{}, err
	}
	return parseKubeletMetrics(output)
}

// GrabFromScheduler returns metrics from scheduler
func (g *Grabber) GrabFromScheduler() (SchedulerMetrics, error) {
	if !g.registeredMaster {
		return SchedulerMetrics{}, fmt.Errorf("Master's Kubelet is not registered. Skipping Scheduler's metrics gathering")
	}
	output, err := g.getMetricsFromPod(g.client, fmt.Sprintf("%v-%v", "kube-scheduler", g.masterName), metav1.NamespaceSystem, ports.InsecureSchedulerPort)
	if err != nil {
		return SchedulerMetrics{}, err
	}
	return parseSchedulerMetrics(output)
}

// GrabFromClusterAutoscaler returns metrics from cluster autoscaler
func (g *Grabber) GrabFromClusterAutoscaler() (ClusterAutoscalerMetrics, error) {
	if !g.registeredMaster && g.externalClient == nil {
		return ClusterAutoscalerMetrics{}, fmt.Errorf("Master's Kubelet is not registered. Skipping ClusterAutoscaler's metrics gathering")
	}
	var client clientset.Interface
	var namespace string
	if g.externalClient != nil {
		client = g.externalClient
		namespace = "kubemark"
	} else {
		client = g.client
		namespace = metav1.NamespaceSystem
	}
	output, err := g.getMetricsFromPod(client, "cluster-autoscaler", namespace, 8085)
	if err != nil {
		return ClusterAutoscalerMetrics{}, err
	}
	return parseClusterAutoscalerMetrics(output)
}

// GrabFromControllerManager returns metrics from controller manager
func (g *Grabber) GrabFromControllerManager() (ControllerManagerMetrics, error) {
	if !g.registeredMaster {
		return ControllerManagerMetrics{}, fmt.Errorf("Master's Kubelet is not registered. Skipping ControllerManager's metrics gathering")
	}

	var err error
	podName := fmt.Sprintf("%v-%v", "kube-controller-manager", g.masterName)
	g.waitForControllerManagerReadyOnce.Do(func() {
		if runningErr := e2epod.WaitForPodNameRunningInNamespace(g.client, podName, metav1.NamespaceSystem); runningErr != nil {
			err = fmt.Errorf("error waiting for controller manager pod to be running: %w", err)
			return
		}

		var lastMetricsFetchErr error
		if metricsWaitErr := wait.PollImmediate(time.Second, time.Minute, func() (bool, error) {
			_, lastMetricsFetchErr = g.getMetricsFromPod(g.client, podName, metav1.NamespaceSystem, ports.InsecureKubeControllerManagerPort)
			return lastMetricsFetchErr == nil, nil
		}); metricsWaitErr != nil {
			err = fmt.Errorf("error waiting for controller manager pod to expose metrics: %v; %v", err, lastMetricsFetchErr)
			return
		}
	})
	if err != nil {
		return ControllerManagerMetrics{}, err
	}

	output, err := g.getMetricsFromPod(g.client, podName, metav1.NamespaceSystem, ports.InsecureKubeControllerManagerPort)
	if err != nil {
		return ControllerManagerMetrics{}, err
	}
	return parseControllerManagerMetrics(output)
}

// GrabFromAPIServer returns metrics from API server
func (g *Grabber) GrabFromAPIServer() (APIServerMetrics, error) {
	output, err := g.getMetricsFromAPIServer()
	if err != nil {
		return APIServerMetrics{}, nil
	}
	return parseAPIServerMetrics(output)
}

// Grab returns metrics from corresponding component
func (g *Grabber) Grab() (Collection, error) {
	result := Collection{}
	var errs []error
	if g.grabFromAPIServer {
		metrics, err := g.GrabFromAPIServer()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.APIServerMetrics = metrics
		}
	}
	if g.grabFromScheduler {
		metrics, err := g.GrabFromScheduler()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.SchedulerMetrics = metrics
		}
	}
	if g.grabFromControllerManager {
		metrics, err := g.GrabFromControllerManager()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.ControllerManagerMetrics = metrics
		}
	}
	if g.grabFromClusterAutoscaler {
		metrics, err := g.GrabFromClusterAutoscaler()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.ClusterAutoscalerMetrics = metrics
		}
	}
	if g.grabFromKubelets {
		result.KubeletMetrics = make(map[string]KubeletMetrics)
		nodes, err := g.client.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
		if err != nil {
			errs = append(errs, err)
		} else {
			for _, node := range nodes.Items {
				kubeletPort := node.Status.DaemonEndpoints.KubeletEndpoint.Port
				metrics, err := g.grabFromKubeletInternal(node.Name, int(kubeletPort))
				if err != nil {
					errs = append(errs, err)
				}
				result.KubeletMetrics[node.Name] = metrics
			}
		}
	}
	if len(errs) > 0 {
		return result, fmt.Errorf("Errors while grabbing metrics: %v", errs)
	}
	return result, nil
}

func (g *Grabber) getMetricsFromPod(client clientset.Interface, podName string, namespace string, port int) (string, error) {
	rawOutput, err := client.CoreV1().RESTClient().Get().
		Namespace(namespace).
		Resource("pods").
		SubResource("proxy").
		Name(fmt.Sprintf("%v:%v", podName, port)).
		Suffix("metrics").
		Do(context.TODO()).Raw()
	if err != nil {
		return "", err
	}
	return string(rawOutput), nil
}
<commit_msg>fix logging on e2e metrics grabber test<commit_after>/*
Copyright 2015 The Kubernetes 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 metrics

import (
	"context"
	"fmt"
	"sync"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/pkg/master/ports"
	e2epod "k8s.io/kubernetes/test/e2e/framework/pod"
	"k8s.io/kubernetes/test/e2e/system"

	"k8s.io/klog"
)

// Collection is metrics collection of components
type Collection struct {
	APIServerMetrics         APIServerMetrics
	ControllerManagerMetrics ControllerManagerMetrics
	KubeletMetrics           map[string]KubeletMetrics
	SchedulerMetrics         SchedulerMetrics
	ClusterAutoscalerMetrics ClusterAutoscalerMetrics
}

// Grabber provides functions which grab metrics from components
type Grabber struct {
	client                            clientset.Interface
	externalClient                    clientset.Interface
	grabFromAPIServer                 bool
	grabFromControllerManager         bool
	grabFromKubelets                  bool
	grabFromScheduler                 bool
	grabFromClusterAutoscaler         bool
	masterName                        string
	registeredMaster                  bool
	waitForControllerManagerReadyOnce sync.Once
}

// NewMetricsGrabber returns new metrics which are initialized.
func NewMetricsGrabber(c clientset.Interface, ec clientset.Interface, kubelets bool, scheduler bool, controllers bool, apiServer bool, clusterAutoscaler bool) (*Grabber, error) {
	registeredMaster := false
	masterName := ""
	nodeList, err := c.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, err
	}
	if len(nodeList.Items) < 1 {
		klog.Warning("Can't find any Nodes in the API server to grab metrics from")
	}
	for _, node := range nodeList.Items {
		if system.DeprecatedMightBeMasterNode(node.Name) {
			registeredMaster = true
			masterName = node.Name
			break
		}
	}
	if !registeredMaster {
		scheduler = false
		controllers = false
		clusterAutoscaler = ec != nil
		if clusterAutoscaler {
			klog.Warningf("Master node is not registered. Grabbing metrics from Scheduler, ControllerManager is disabled.")
		} else {
			klog.Warningf("Master node is not registered. Grabbing metrics from Scheduler, ControllerManager and ClusterAutoscaler is disabled.")
		}
	}

	return &Grabber{
		client:                    c,
		externalClient:            ec,
		grabFromAPIServer:         apiServer,
		grabFromControllerManager: controllers,
		grabFromKubelets:          kubelets,
		grabFromScheduler:         scheduler,
		grabFromClusterAutoscaler: clusterAutoscaler,
		masterName:                masterName,
		registeredMaster:          registeredMaster,
	}, nil
}

// HasRegisteredMaster returns if metrics grabber was able to find a master node
func (g *Grabber) HasRegisteredMaster() bool {
	return g.registeredMaster
}

// GrabFromKubelet returns metrics from kubelet
func (g *Grabber) GrabFromKubelet(nodeName string) (KubeletMetrics, error) {
	nodes, err := g.client.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{FieldSelector: fields.Set{"metadata.name": nodeName}.AsSelector().String()})
	if err != nil {
		return KubeletMetrics{}, err
	}
	if len(nodes.Items) != 1 {
		return KubeletMetrics{}, fmt.Errorf("Error listing nodes with name %v, got %v", nodeName, nodes.Items)
	}
	kubeletPort := nodes.Items[0].Status.DaemonEndpoints.KubeletEndpoint.Port
	return g.grabFromKubeletInternal(nodeName, int(kubeletPort))
}

func (g *Grabber) grabFromKubeletInternal(nodeName string, kubeletPort int) (KubeletMetrics, error) {
	if kubeletPort <= 0 || kubeletPort > 65535 {
		return KubeletMetrics{}, fmt.Errorf("Invalid Kubelet port %v. Skipping Kubelet's metrics gathering", kubeletPort)
	}
	output, err := g.getMetricsFromNode(nodeName, int(kubeletPort))
	if err != nil {
		return KubeletMetrics{}, err
	}
	return parseKubeletMetrics(output)
}

// GrabFromScheduler returns metrics from scheduler
func (g *Grabber) GrabFromScheduler() (SchedulerMetrics, error) {
	if !g.registeredMaster {
		return SchedulerMetrics{}, fmt.Errorf("Master's Kubelet is not registered. Skipping Scheduler's metrics gathering")
	}
	output, err := g.getMetricsFromPod(g.client, fmt.Sprintf("%v-%v", "kube-scheduler", g.masterName), metav1.NamespaceSystem, ports.InsecureSchedulerPort)
	if err != nil {
		return SchedulerMetrics{}, err
	}
	return parseSchedulerMetrics(output)
}

// GrabFromClusterAutoscaler returns metrics from cluster autoscaler
func (g *Grabber) GrabFromClusterAutoscaler() (ClusterAutoscalerMetrics, error) {
	if !g.registeredMaster && g.externalClient == nil {
		return ClusterAutoscalerMetrics{}, fmt.Errorf("Master's Kubelet is not registered. Skipping ClusterAutoscaler's metrics gathering")
	}
	var client clientset.Interface
	var namespace string
	if g.externalClient != nil {
		client = g.externalClient
		namespace = "kubemark"
	} else {
		client = g.client
		namespace = metav1.NamespaceSystem
	}
	output, err := g.getMetricsFromPod(client, "cluster-autoscaler", namespace, 8085)
	if err != nil {
		return ClusterAutoscalerMetrics{}, err
	}
	return parseClusterAutoscalerMetrics(output)
}

// GrabFromControllerManager returns metrics from controller manager
func (g *Grabber) GrabFromControllerManager() (ControllerManagerMetrics, error) {
	if !g.registeredMaster {
		return ControllerManagerMetrics{}, fmt.Errorf("Master's Kubelet is not registered. Skipping ControllerManager's metrics gathering")
	}

	var err error
	podName := fmt.Sprintf("%v-%v", "kube-controller-manager", g.masterName)
	g.waitForControllerManagerReadyOnce.Do(func() {
		if runningErr := e2epod.WaitForPodNameRunningInNamespace(g.client, podName, metav1.NamespaceSystem); runningErr != nil {
			err = fmt.Errorf("error waiting for controller manager pod to be running: %w", runningErr)
			return
		}

		var lastMetricsFetchErr error
		if metricsWaitErr := wait.PollImmediate(time.Second, time.Minute, func() (bool, error) {
			_, lastMetricsFetchErr = g.getMetricsFromPod(g.client, podName, metav1.NamespaceSystem, ports.InsecureKubeControllerManagerPort)
			return lastMetricsFetchErr == nil, nil
		}); metricsWaitErr != nil {
			err = fmt.Errorf("error waiting for controller manager pod to expose metrics: %v; %v", metricsWaitErr, lastMetricsFetchErr)
			return
		}
	})
	if err != nil {
		return ControllerManagerMetrics{}, err
	}

	output, err := g.getMetricsFromPod(g.client, podName, metav1.NamespaceSystem, ports.InsecureKubeControllerManagerPort)
	if err != nil {
		return ControllerManagerMetrics{}, err
	}
	return parseControllerManagerMetrics(output)
}

// GrabFromAPIServer returns metrics from API server
func (g *Grabber) GrabFromAPIServer() (APIServerMetrics, error) {
	output, err := g.getMetricsFromAPIServer()
	if err != nil {
		return APIServerMetrics{}, nil
	}
	return parseAPIServerMetrics(output)
}

// Grab returns metrics from corresponding component
func (g *Grabber) Grab() (Collection, error) {
	result := Collection{}
	var errs []error
	if g.grabFromAPIServer {
		metrics, err := g.GrabFromAPIServer()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.APIServerMetrics = metrics
		}
	}
	if g.grabFromScheduler {
		metrics, err := g.GrabFromScheduler()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.SchedulerMetrics = metrics
		}
	}
	if g.grabFromControllerManager {
		metrics, err := g.GrabFromControllerManager()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.ControllerManagerMetrics = metrics
		}
	}
	if g.grabFromClusterAutoscaler {
		metrics, err := g.GrabFromClusterAutoscaler()
		if err != nil {
			errs = append(errs, err)
		} else {
			result.ClusterAutoscalerMetrics = metrics
		}
	}
	if g.grabFromKubelets {
		result.KubeletMetrics = make(map[string]KubeletMetrics)
		nodes, err := g.client.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
		if err != nil {
			errs = append(errs, err)
		} else {
			for _, node := range nodes.Items {
				kubeletPort := node.Status.DaemonEndpoints.KubeletEndpoint.Port
				metrics, err := g.grabFromKubeletInternal(node.Name, int(kubeletPort))
				if err != nil {
					errs = append(errs, err)
				}
				result.KubeletMetrics[node.Name] = metrics
			}
		}
	}
	if len(errs) > 0 {
		return result, fmt.Errorf("Errors while grabbing metrics: %v", errs)
	}
	return result, nil
}

func (g *Grabber) getMetricsFromPod(client clientset.Interface, podName string, namespace string, port int) (string, error) {
	rawOutput, err := client.CoreV1().RESTClient().Get().
		Namespace(namespace).
		Resource("pods").
		SubResource("proxy").
		Name(fmt.Sprintf("%v:%v", podName, port)).
		Suffix("metrics").
		Do(context.TODO()).Raw()
	if err != nil {
		return "", err
	}
	return string(rawOutput), nil
}
<|endoftext|>