<commit_before>package ufop

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/qiniu/api/auth/digest"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
)

var (
	ufopPrefix  string
	jobHandlers map[string]UfopJobHandler
	unzipper    *UnZipper
)

type UfopJobHandler interface {
	Do(ufopReq UfopRequest) (interface{}, error)
}

type UfopServer struct {
	cfg *UfopConfig
}

func NewServer(cfg *UfopConfig) *UfopServer {
	serv := UfopServer{}
	serv.cfg = cfg
	return &serv
}

func (this *UfopServer) registerJobHandlers() {
	mac := digest.Mac{
		AccessKey: this.cfg.AccessKey,
		SecretKey: []byte(this.cfg.SecretKey),
	}
	jobHandlers = make(map[string]UfopJobHandler, 0)
	ufopPrefix = this.cfg.UfopPrefix
	//unzipper
	unzipper = &UnZipper{
		mac: &mac,
	}

	jobHandlers[ufopPrefix+"unzip"] = unzipper
}

func (this *UfopServer) Listen() {
	//register
	this.registerJobHandlers()

	//define handler
	http.HandleFunc("/uop", serveUfop)

	//bind and listen
	endPoint := fmt.Sprintf("%s:%d", this.cfg.ListenHost, this.cfg.ListenPort)
	ufopServer := &http.Server{
		Addr:           endPoint,
		ReadTimeout:    time.Duration(this.cfg.ReadTimeout) * time.Second,
		WriteTimeout:   time.Duration(this.cfg.WriteTimeout) * time.Second,
		MaxHeaderBytes: this.cfg.MaxHeaderBytes,
	}

	listenErr := ufopServer.ListenAndServe()
	if listenErr != nil {
		log.Println(listenErr)
	}
}

func serveUfop(w http.ResponseWriter, req *http.Request) {
	//check method
	if req.Method != "POST" {
		writeJsonError(w, 405, "method not allowed")
		return
	}

	defer req.Body.Close()
	var err error
	var ufopReq UfopRequest
	var ufopResult interface{}

	ufopReqData, err := ioutil.ReadAll(req.Body)
	if err != nil {
		writeJsonError(w, 500, "read ufop request body error")
		return
	}

	err = json.Unmarshal(ufopReqData, &ufopReq)
	if err != nil {
		writeJsonError(w, 500, "parse ufop request body error")
		return
	}

	ufopResult, err = handleJob(ufopReq)
	if err != nil {
		writeJsonError(w, 500, err.Error())
	} else {
		writeJsonResult(w, 200, ufopResult)
	}
}

func handleJob(ufopReq UfopRequest) (interface{}, error) {
	var ufopResult interface{}
	var err error
	cmd := ufopReq.Cmd

	items := strings.SplitN(cmd, "/", 2)
	fop := items[0]
	if jobHandler, ok := jobHandlers[fop]; ok {
		ufopReq.Cmd = strings.TrimPrefix(ufopReq.Cmd, ufopPrefix)
		ufopResult, err = jobHandler.Do(ufopReq)
	} else {
		err = errors.New("no fop available for the request")
	}
	return ufopResult, err
}

func writeJsonError(w http.ResponseWriter, statusCode int, message string) {
	log.Println(message)
	w.WriteHeader(statusCode)
	w.Header().Add("Content-Type", "application/json")
	io.WriteString(w, fmt.Sprintf(`{"error": "%s"}`, message))
}

func writeJsonResult(w http.ResponseWriter, statusCode int, result interface{}) {
	w.WriteHeader(statusCode)
	w.Header().Add("Content-Type", "application/json")
	data, err := json.Marshal(result)
	if err != nil {
		log.Println("Encode ufop result error,", err)
		writeJsonResult(w, 500, "Encode ufop result error")
	} else {
		io.WriteString(w, string(data))
	}
}
<commit_msg>Change processing error to 200.<commit_after>package ufop

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/qiniu/api/auth/digest"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"
)

var (
	ufopPrefix  string
	jobHandlers map[string]UfopJobHandler
	unzipper    *UnZipper
)

type UfopJobHandler interface {
	Do(ufopReq UfopRequest) (interface{}, error)
}

type UfopServer struct {
	cfg *UfopConfig
}

func NewServer(cfg *UfopConfig) *UfopServer {
	serv := UfopServer{}
	serv.cfg = cfg
	return &serv
}

func (this *UfopServer) registerJobHandlers() {
	mac := digest.Mac{
		AccessKey: this.cfg.AccessKey,
		SecretKey: []byte(this.cfg.SecretKey),
	}
	jobHandlers = make(map[string]UfopJobHandler, 0)
	ufopPrefix = this.cfg.UfopPrefix
	//unzipper
	unzipper = &UnZipper{
		mac: &mac,
	}

	jobHandlers[ufopPrefix+"unzip"] = unzipper
}

func (this *UfopServer) Listen() {
	//register
	this.registerJobHandlers()

	//define handler
	http.HandleFunc("/uop", serveUfop)

	//bind and listen
	endPoint := fmt.Sprintf("%s:%d", this.cfg.ListenHost, this.cfg.ListenPort)
	ufopServer := &http.Server{
		Addr:           endPoint,
		ReadTimeout:    time.Duration(this.cfg.ReadTimeout) * time.Second,
		WriteTimeout:   time.Duration(this.cfg.WriteTimeout) * time.Second,
		MaxHeaderBytes: this.cfg.MaxHeaderBytes,
	}

	listenErr := ufopServer.ListenAndServe()
	if listenErr != nil {
		log.Println(listenErr)
	}
}

func serveUfop(w http.ResponseWriter, req *http.Request) {
	//check method
	if req.Method != "POST" {
		writeJsonError(w, 405, "method not allowed")
		return
	}

	defer req.Body.Close()
	var err error
	var ufopReq UfopRequest
	var ufopResult interface{}

	ufopReqData, err := ioutil.ReadAll(req.Body)
	if err != nil {
		writeJsonError(w, 500, "read ufop request body error")
		return
	}

	err = json.Unmarshal(ufopReqData, &ufopReq)
	if err != nil {
		writeJsonError(w, 500, "parse ufop request body error")
		return
	}

	ufopResult, err = handleJob(ufopReq)
	if err != nil {
		writeJsonError(w, 200, err.Error())
	} else {
		writeJsonResult(w, 200, ufopResult)
	}
}

func handleJob(ufopReq UfopRequest) (interface{}, error) {
	var ufopResult interface{}
	var err error
	cmd := ufopReq.Cmd

	items := strings.SplitN(cmd, "/", 2)
	fop := items[0]
	if jobHandler, ok := jobHandlers[fop]; ok {
		ufopReq.Cmd = strings.TrimPrefix(ufopReq.Cmd, ufopPrefix)
		ufopResult, err = jobHandler.Do(ufopReq)
	} else {
		err = errors.New("no fop available for the request")
	}
	return ufopResult, err
}

func writeJsonError(w http.ResponseWriter, statusCode int, message string) {
	log.Println(message)
	w.WriteHeader(statusCode)
	w.Header().Add("Content-Type", "application/json")
	io.WriteString(w, fmt.Sprintf(`{"error": "%s"}`, message))
}

func writeJsonResult(w http.ResponseWriter, statusCode int, result interface{}) {
	w.WriteHeader(statusCode)
	w.Header().Add("Content-Type", "application/json")
	data, err := json.Marshal(result)
	if err != nil {
		log.Println("Encode ufop result error,", err)
		writeJsonResult(w, 500, "Encode ufop result error")
	} else {
		io.WriteString(w, string(data))
	}
}
<|endoftext|>


<commit_before>package message

import "github.com/tychoish/grip/level"

// Composer defines an interface with a "Resolve()" method that
// returns the message in string format. Objects that implement this
// interface, in combination to the Compose[*] operations, the
// Resolve() method is only caled if the priority of the method is
// greater than the threshold priority. This makes it possible to
// defer building log messages (that may be somewhat expensive to
// generate) until it's certain that we're going to be outputting the
// message.
type Composer interface {
	// Returns the content of the message as a string for use in
	// line-printing logging engines.
	Resolve() string

	// A "raw" format of the logging output for use by some Sender
	// implementations that write logged items to interfaces that
	// accept JSON or another structured format.
	Raw() interface{}

	// Returns "true" when the message has content and should be
	// logged, and false otherwise. When false, the sender can
	// (and should!) ignore messages even if they are otherwise
	// above the logging threshold.
	Loggable() bool

	// Priority returns the priority of the message.
	Priority() level.Priority
	SetPriority(level.Priority) error
}

// ConvertToComposer can coerce unknown objects into Composer
// instances, as possible.
func ConvertToComposer(p level.Priority, message interface{}) Composer {
	switch message := message.(type) {
	case Composer:
		_ = message.SetPriority(p)
		return message
	case string:
		return NewDefaultMessage(p, message)
	case []string, []interface{}:
		return NewLineMessage(p, message)
	case error, nil:
		return NewErrorMessage(p, message)
	default:
		return NewFormattedMessage(p, "%+v", message)
	}
}
<commit_msg>fix logging for nils<commit_after>package message

import "github.com/tychoish/grip/level"

// Composer defines an interface with a "Resolve()" method that
// returns the message in string format. Objects that implement this
// interface, in combination to the Compose[*] operations, the
// Resolve() method is only caled if the priority of the method is
// greater than the threshold priority. This makes it possible to
// defer building log messages (that may be somewhat expensive to
// generate) until it's certain that we're going to be outputting the
// message.
type Composer interface {
	// Returns the content of the message as a string for use in
	// line-printing logging engines.
	Resolve() string

	// A "raw" format of the logging output for use by some Sender
	// implementations that write logged items to interfaces that
	// accept JSON or another structured format.
	Raw() interface{}

	// Returns "true" when the message has content and should be
	// logged, and false otherwise. When false, the sender can
	// (and should!) ignore messages even if they are otherwise
	// above the logging threshold.
	Loggable() bool

	// Priority returns the priority of the message.
	Priority() level.Priority
	SetPriority(level.Priority) error
}

// ConvertToComposer can coerce unknown objects into Composer
// instances, as possible.
func ConvertToComposer(p level.Priority, message interface{}) Composer {
	switch message := message.(type) {
	case Composer:
		_ = message.SetPriority(p)
		return message
	case string:
		return NewDefaultMessage(p, message)
	case []string, []interface{}:
		return NewLineMessage(p, message)
	case error:
		return NewErrorMessage(p, message)
	case nil:
		return NewDefaultMessage(p, "")
	default:
		return NewFormattedMessage(p, "%+v", message)
	}
}
<|endoftext|>


<commit_before>package rest_test

import (
	"bufio"
	"io/ioutil"
	"net/url"
	"os"
	"os/exec"
	"restic"
	"strings"
	"testing"

	"restic/backend/rest"
	"restic/backend/test"
	. "restic/test"
)

func runRESTServer(t testing.TB, dir string) func() {
	srv, err := exec.LookPath("rest-server")
	if err != nil {
		t.Skip(err)
	}

	cmd := exec.Command(srv, "--path", dir)
	cmd.Stdout = os.Stdout
	stderr, err := cmd.StderrPipe()
	if err != nil {
		t.Fatal(err)
	}

	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}

	sc := bufio.NewScanner(stderr)
	for sc.Scan() {
		if strings.HasPrefix(sc.Text(), "Starting server") {
			break
		}
	}

	return func() {
		if err := cmd.Process.Kill(); err != nil {
			t.Fatal(err)
		}

		// ignore errors, we've killed the process
		_ = cmd.Wait()
	}
}

func TestBackend(t *testing.T) {
	dir, cleanup := TempDir(t)
	defer cleanup()

	cleanup = runRESTServer(t, dir)
	defer cleanup()

	suite := test.Suite{
		// NewConfig returns a config for a new temporary backend that will be used in tests.
		NewConfig: func() (interface{}, error) {
			dir, err := ioutil.TempDir(TestTempDir, "restic-test-rest-")
			if err != nil {
				t.Fatal(err)
			}

			t.Logf("create new backend at %v", dir)

			url, err := url.Parse("http://localhost:8000/restic-test")
			if err != nil {
				t.Fatal(err)
			}

			cfg := rest.Config{
				URL: url,
			}
			return cfg, nil
		},

		// CreateFn is a function that creates a temporary repository for the tests.
		Create: func(config interface{}) (restic.Backend, error) {
			cfg := config.(rest.Config)
			return rest.Create(cfg)
		},

		// OpenFn is a function that opens a previously created temporary repository.
		Open: func(config interface{}) (restic.Backend, error) {
			cfg := config.(rest.Config)
			return rest.Open(cfg)
		},

		// CleanupFn removes data created during the tests.
		Cleanup: func(config interface{}) error {
			return nil
		},
	}

	suite.RunTests(t)
}
<commit_msg>rest: Add context<commit_after>package rest_test

import (
	"bufio"
	"context"
	"io/ioutil"
	"net/url"
	"os"
	"os/exec"
	"restic"
	"strings"
	"testing"

	"restic/backend/rest"
	"restic/backend/test"
	. "restic/test"
)

func runRESTServer(ctx context.Context, t testing.TB, dir string) func() {
	srv, err := exec.LookPath("rest-server")
	if err != nil {
		t.Skip(err)
	}

	cmd := exec.CommandContext(ctx, srv, "--path", dir)
	cmd.Stdout = os.Stdout
	stderr, err := cmd.StderrPipe()
	if err != nil {
		t.Fatal(err)
	}

	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}

	sc := bufio.NewScanner(stderr)
	for sc.Scan() {
		if strings.HasPrefix(sc.Text(), "Starting server") {
			break
		}
	}

	return func() {
		if err := cmd.Process.Kill(); err != nil {
			t.Fatal(err)
		}

		// ignore errors, we've killed the process
		_ = cmd.Wait()
	}
}

func TestBackend(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	dir, cleanup := TempDir(t)
	defer cleanup()

	cleanup = runRESTServer(ctx, t, dir)
	defer cleanup()

	suite := test.Suite{
		// NewConfig returns a config for a new temporary backend that will be used in tests.
		NewConfig: func() (interface{}, error) {
			dir, err := ioutil.TempDir(TestTempDir, "restic-test-rest-")
			if err != nil {
				t.Fatal(err)
			}

			t.Logf("create new backend at %v", dir)

			url, err := url.Parse("http://localhost:8000/restic-test")
			if err != nil {
				t.Fatal(err)
			}

			cfg := rest.Config{
				URL: url,
			}
			return cfg, nil
		},

		// CreateFn is a function that creates a temporary repository for the tests.
		Create: func(config interface{}) (restic.Backend, error) {
			cfg := config.(rest.Config)
			return rest.Create(cfg)
		},

		// OpenFn is a function that opens a previously created temporary repository.
		Open: func(config interface{}) (restic.Backend, error) {
			cfg := config.(rest.Config)
			return rest.Open(cfg)
		},

		// CleanupFn removes data created during the tests.
		Cleanup: func(config interface{}) error {
			return nil
		},
	}

	suite.RunTests(t)
}
<|endoftext|>


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

package driver

import (
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/d4l3k/messagediff"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"

	"istio.io/proxy/test/envoye2e/env"
)

type Stats struct {
	AdminPort uint16
	Matchers  map[string]StatMatcher
}

type StatMatcher interface {
	Matches(*Params, *dto.MetricFamily) error
}

var _ Step = &Stats{}

func (s *Stats) Run(p *Params) error {
	for i := 0; i < 15; i++ {
		_, body, err := env.HTTPGet(fmt.Sprintf("http://127.0.0.1:%d/stats/prometheus", s.AdminPort))
		if err != nil {
			return err
		}
		reader := strings.NewReader(body)
		metrics, err := (&expfmt.TextParser{}).TextToMetricFamilies(reader)
		if err != nil {
			return err
		}
		count := 0
		for _, metric := range metrics {
			matcher, found := s.Matchers[metric.GetName()]
			if !found {
				continue
			}
			if err := matcher.Matches(p, metric); err == nil {
				log.Printf("matched metric %q", metric.GetName())
				count++
			} else {
				log.Printf("metric %q did not match: %v\n", metric.GetName(), err)
			}
		}
		if count == len(s.Matchers) {
			return nil
		}
		log.Printf("failed to match all metrics: want %#v", s.Matchers)
		time.Sleep(1 * time.Second)
	}
	return errors.New("failed to match all stats")
}

func (s *Stats) Cleanup() {}

type ExactStat struct {
	Metric string
}

func (me *ExactStat) Matches(params *Params, that *dto.MetricFamily) error {
	metric := &dto.MetricFamily{}
	params.LoadTestProto(me.Metric, metric)

	if s, same := messagediff.PrettyDiff(metric, that); !same {
		return fmt.Errorf("diff: %v, got: %v, want: %v", s, that, metric)
	}
	return nil
}

var _ StatMatcher = &ExactStat{}

type PartialStat struct {
	Metric string
}

func (me *PartialStat) Matches(params *Params, that *dto.MetricFamily) error {
	metric := &dto.MetricFamily{}
	params.LoadTestProto(me.Metric, metric)
	for _, wm := range metric.Metric {
		found := false
		for _, gm := range that.Metric {
			if _, same := messagediff.PrettyDiff(wm, gm); !same {
				continue
			}
			found = true
			break
		}
		if !found {
			return fmt.Errorf("cannot find metric, got: %v, want: %v", that.Metric, wm)
		}
	}
	return nil
}

var _ StatMatcher = &PartialStat{}
<commit_msg>test: enhance metrics check failure message. (#3494)<commit_after>// Copyright 2019 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package driver

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/d4l3k/messagediff"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/common/expfmt"

	"istio.io/proxy/test/envoye2e/env"
)

type Stats struct {
	AdminPort uint16
	Matchers  map[string]StatMatcher
}

type StatMatcher interface {
	Matches(*Params, *dto.MetricFamily) error
}

var _ Step = &Stats{}

func (s *Stats) Run(p *Params) error {
	var metrics map[string]*dto.MetricFamily
	for i := 0; i < 15; i++ {
		_, body, err := env.HTTPGet(fmt.Sprintf("http://127.0.0.1:%d/stats/prometheus", s.AdminPort))
		if err != nil {
			return err
		}
		reader := strings.NewReader(body)
		metrics, err = (&expfmt.TextParser{}).TextToMetricFamilies(reader)
		if err != nil {
			return err
		}
		count := 0
		for _, metric := range metrics {
			matcher, found := s.Matchers[metric.GetName()]
			if !found {
				continue
			}
			if err := matcher.Matches(p, metric); err == nil {
				log.Printf("matched metric %q", metric.GetName())
				count++
			} else {
				log.Printf("metric %q did not match: %v\n", metric.GetName(), err)
			}
		}
		if count == len(s.Matchers) {
			return nil
		}
		time.Sleep(1 * time.Second)
	}
	return fmt.Errorf("failed to match all metrics: want %v, but got %v", s.Matchers, metrics)
}

func (s *Stats) Cleanup() {}

type ExactStat struct {
	Metric string
}

func (me *ExactStat) Matches(params *Params, that *dto.MetricFamily) error {
	metric := &dto.MetricFamily{}
	params.LoadTestProto(me.Metric, metric)

	if s, same := messagediff.PrettyDiff(metric, that); !same {
		return fmt.Errorf("diff: %v, got: %v, want: %v", s, that, metric)
	}
	return nil
}

var _ StatMatcher = &ExactStat{}

type PartialStat struct {
	Metric string
}

func (me *PartialStat) Matches(params *Params, that *dto.MetricFamily) error {
	metric := &dto.MetricFamily{}
	params.LoadTestProto(me.Metric, metric)
	for _, wm := range metric.Metric {
		found := false
		for _, gm := range that.Metric {
			if _, same := messagediff.PrettyDiff(wm, gm); !same {
				continue
			}
			found = true
			break
		}
		if !found {
			return fmt.Errorf("cannot find metric, got: %v, want: %v", that.Metric, wm)
		}
	}
	return nil
}

var _ StatMatcher = &PartialStat{}
<|endoftext|>


<commit_before>package balancer

import (
    // "log"
    "github.com/trendrr/cheshire-golang/cheshire"
    shards "github.com/dustismo/cheshire-shards/shards"
    "fmt"
    // "io/ioutil"
    "time"
)

// All the operations necessary for rebalance and topology changes

// Moves data from one server to another
// Does the following:
// 0. Update router table
// 1. Lock all necessary clients
// 2. Move Data
// 3. Update router table on servers
// 4. Unlock
func MovePartition(services *Services, routerTable *shards.RouterTable, partition int, from, to *shards.RouterEntry) error {

    //update router table if needed
    routerTable, _ = RouterTableUpdate(services, routerTable, 3)

    entries, err := routerTable.PartitionEntries(partition)
    if err != nil {
        return err
    }

    // Lock partitions
    for _, e := range(entries) {
        response, err := cheshire.HttpApiCallSync(
            fmt.Sprintf("%s:%d", e.Address, e.HttpPort),
            cheshire.NewRequest(shards.PARTITION_LOCK, "GET"),
            5 * time.Second)
        if err != nil {
            services.Logger.Printf("ERROR While contacting %s -- %s", e.Address, err)
            continue
        }
        if response.StatusCode() != 200 {
            services.Logger.Printf("ERROR While locking partition %s -- %s", e.Address, response.StatusMessage())
        }
    }

    //Move the data!

    //create a new json connection
    fromClient := cheshire.NewJsonClient(from.Address, from.JsonPort)
    err = fromClient.Connect()
    if err != nil {
        return err
    }
    defer fromClient.Close()

    toClient := cheshire.NewJsonClient(to.Address, to.JsonPort)
    err = toClient.Connect()
    if err != nil {
        return err
    }
    defer toClient.Close()

    request := cheshire.NewRequest(shards.DATA_PULL, "GET")
    request.Params().Put("partition", partition)

    responseChan := make(chan *cheshire.Response, 10)
    errorChan := make(chan error)

    fromClient.ApiCall(
        request,
        responseChan,
        errorChan,
    )

    moved := 0

    toResponseChan := make(chan *cheshire.Response, 10)
    toErrorChan := make(chan error)

    for {
        select {
        case response := <- responseChan :
            //TODO: send the data to the toClient
            request := cheshire.NewRequest(shards.DATA_PUSH, "PUT")
            d, ok := response.GetDynMap("data")
            if !ok {
                services.Logger.Printf("ERROR: packet missing data :: %s ", request)
                continue
            }
            request.Params().Put("data", d)
            request.Params().Put("partition", partition)
            toClient.ApiCall(
                request,
                toResponseChan,
                toErrorChan,
            )

            // keep a log of items moved
            moved++
            if moved % 100 == 0 {
                services.Logger.Printf("Moving partition %d... Moved %d objects so far", partition, moved)
            }

            //check for completion
            if response.TxnStatus() == "complete" {
                // FINISHED!
                services.Logger.Printf("SUCCESSFULLY Moved partition %d. Moved %d objects!", partition, moved)
                break
            }

        case err := <- errorChan :
            services.Logger.Printf("ERROR While Moving data from %s -- %s", from.Address, err)
            return err

        case response := <- toResponseChan :
            if response.StatusCode() != 200 {
                services.Logger.Printf("ERROR While Moving data from %s -- %s.  \n Continuing...", from.Address, response.StatusMessage())
            }

            //do nothing,
        case err := <- toErrorChan :
            services.Logger.Printf("ERROR While Moving data to %s -- %s", to.Address, err)
            return err
        }
    }

    // Now make sure we got responses for all the PUT data ops
    count := 0
    for toClient.CurrentInFlight() > 1 {
        services.Logger.Printf("NOW Waiting for success messages to complete")
        select {
            case response := <- toResponseChan :
            //do nothing,
                if response.StatusCode() != 200 {
                    services.Logger.Printf("ERROR While Moving data to %s -- %s.  \n Continuing...", to.Address, response.StatusMessage())
                }

            case err := <- toErrorChan :
                services.Logger.Printf("ERROR While Moving data to %s -- %s", to.Address, err)
                return err
            default :
                if count > 30 {
                   services.Logger.Printf("GAH. Waited 30 seconds for completion.  there seems to be a problem,.")
                   return fmt.Errorf("GAH. Waited 30 seconds for completion.  there seems to be a problem.")
                }
                time.Sleep(1*time.Second)
        }
        count++
    }

    //Delete the data on the from server.
    request = cheshire.NewRequest(shards.PARTITION_DELETE, "DELETE")
    request.Params().Put("partition", partition)

    services.Logger.Printf("NOW Deleting the partition from the origin server %s", from.Address)
    response, err := fromClient.ApiCallSync(request, 300*time.Second)
    if response.StatusCode() != 200 {
        services.Logger.Printf("ERROR While deleting partition: %s", response.StatusMessage())
    }

    return nil
}

// Checks with entries to see if an updated router table is available.
func RouterTableUpdate(services *Services, routerTable *shards.RouterTable, maxChecks int) (*shards.RouterTable, bool) {
    checks := 0
    updated := false
    for _,e := range(routerTable.Entries) {
        if checks >= maxChecks {
            break
        }

        checks++
        // make sure our routertable is up to date.
        response, err := cheshire.HttpApiCallSync(
            fmt.Sprintf("%s:%d", e.Address, e.HttpPort),
            cheshire.NewRequest(shards.CHECKIN, "GET"),
            5 * time.Second)
        if err != nil {
            services.Logger.Printf("ERROR While contacting %s -- %s", e.Address, err)
            continue
        }

        rev := response.MustInt64("rt_revision", 0)
        if rev <= routerTable.Revision {
            continue
        }
        services.Logger.Printf("Found updated router table at: %s", e.Address)
        //get the new routertable.
        response, err = cheshire.HttpApiCallSync(
            fmt.Sprintf("%s:%d", e.Address, e.HttpPort),
            cheshire.NewRequest(shards.ROUTERTABLE_GET, "GET"),
            5 * time.Second)
        if err != nil {
            services.Logger.Printf("ERROR While contacting %s -- %s", e.Address, err)
            continue
        }
        mp, ok := response.GetDynMap("router_table")
        if !ok {
            services.Logger.Printf("ERROR from %s -- BAD ROUTER TABLE RESPONSE %s", e.Address, response)
            continue
        }

        rt, err := shards.ToRouterTable(mp)
        if err != nil {
            services.Logger.Printf("ERROR While parsing router table %s -- %s", e.Address, err)
            continue
        }

        services.Logger.Printf("SUCCESSFULLY update router table to revision %d", rt.Revision)
        updated = true
        routerTable = rt
    }
    return routerTable, updated
}

<commit_msg>add locking and unlocking functions<commit_after>package balancer

import (
    // "log"
    "github.com/trendrr/cheshire-golang/cheshire"
    shards "github.com/dustismo/cheshire-shards/shards"
    "fmt"
    // "io/ioutil"
    "time"
)

// All the operations necessary for rebalance and topology changes

// Locks the specified partition.
// Note, this does NOT update the routertable
// it is assumed the router table is up to date.
func LockPartition(services *Services, routerTable *shards.RouterTable, partition int) error {
    return locking(shards.PARTITION_LOCK, services, routerTable, partition)
}

// Unlocks the specified partition.
// Note, this does NOT update the routertable
// it is assumed the router table is up to date.
func UnlockPartition(services *Services, routerTable *shards.RouterTable, partition int) error {
    return locking(shards.PARTITION_UNLOCK, services, routerTable, partition)
}

// internal lock code shared by unlock and lock. (different endpoints)
func locking(endpoint string, services *Services, routerTable *shards.RouterTable, partition int) error {
    entries, err := routerTable.PartitionEntries(partition)
    if err != nil {
        return err
    }

    // Lock partitions
    for _, e := range(entries) {
        response, err := cheshire.HttpApiCallSync(
            fmt.Sprintf("%s:%d", e.Address, e.HttpPort),
            cheshire.NewRequest(endpoint, "GET"),
            5 * time.Second)
        if err != nil {
            services.Logger.Printf("ERROR While contacting %s -- %s", e.Address, err)
            //TODO: retry?
            continue
        }
        if response.StatusCode() != 200 {
            //TODO: retry?
            services.Logger.Printf("ERROR While locking partition %s -- %s", e.Address, response.StatusMessage())
        }
    }
    return nil
}

// Moves data from one server to another
// Does the following:
// 0. Update router table
// 1. Lock all necessary clients
// 2. Move Data
// 3. Update router table on servers
// 4. Unlock
func MovePartition(services *Services, routerTable *shards.RouterTable, partition int, from, to *shards.RouterEntry) error {

    //update router table if needed
    routerTable, _ = RouterTableUpdate(services, routerTable, 3)

    err := LockPartition(services, routerTable, partition)
    if err != nil {
        return err
    }

    //Unlock partition no matter what
    defer UnlockPartition(services, routerTable, partition)

    //Move the data!

    //create a new json connection
    fromClient := cheshire.NewJsonClient(from.Address, from.JsonPort)
    err = fromClient.Connect()
    if err != nil {
        return err
    }
    defer fromClient.Close()

    toClient := cheshire.NewJsonClient(to.Address, to.JsonPort)
    err = toClient.Connect()
    if err != nil {
        return err
    }
    defer toClient.Close()

    request := cheshire.NewRequest(shards.DATA_PULL, "GET")
    request.Params().Put("partition", partition)

    responseChan := make(chan *cheshire.Response, 10)
    errorChan := make(chan error)

    fromClient.ApiCall(
        request,
        responseChan,
        errorChan,
    )

    moved := 0

    toResponseChan := make(chan *cheshire.Response, 10)
    toErrorChan := make(chan error)

    for {
        select {
        case response := <- responseChan :
            //TODO: send the data to the toClient
            request := cheshire.NewRequest(shards.DATA_PUSH, "PUT")
            d, ok := response.GetDynMap("data")
            if !ok {
                services.Logger.Printf("ERROR: packet missing data :: %s ", request)
                continue
            }
            request.Params().Put("data", d)
            request.Params().Put("partition", partition)
            toClient.ApiCall(
                request,
                toResponseChan,
                toErrorChan,
            )

            // keep a log of items moved
            moved++
            if moved % 100 == 0 {
                services.Logger.Printf("Moving partition %d... Moved %d objects so far", partition, moved)
            }

            //check for completion
            if response.TxnStatus() == "complete" {
                // FINISHED!
                services.Logger.Printf("SUCCESSFULLY Moved partition %d. Moved %d objects!", partition, moved)
                break
            }

        case err := <- errorChan :
            services.Logger.Printf("ERROR While Moving data from %s -- %s", from.Address, err)
            return err

        case response := <- toResponseChan :
            if response.StatusCode() != 200 {
                services.Logger.Printf("ERROR While Moving data from %s -- %s.  \n Continuing...", from.Address, response.StatusMessage())
            }

            //do nothing,
        case err := <- toErrorChan :
            services.Logger.Printf("ERROR While Moving data to %s -- %s", to.Address, err)
            return err
        }
    }

    // Now make sure we got responses for all the PUT data ops
    count := 0
    for toClient.CurrentInFlight() > 1 {
        services.Logger.Printf("NOW Waiting for success messages to complete")
        select {
            case response := <- toResponseChan :
            //do nothing,
                if response.StatusCode() != 200 {
                    services.Logger.Printf("ERROR While Moving data to %s -- %s.  \n Continuing...", to.Address, response.StatusMessage())
                }

            case err := <- toErrorChan :
                services.Logger.Printf("ERROR While Moving data to %s -- %s", to.Address, err)
                return err
            default :
                if count > 30 {
                   services.Logger.Printf("GAH. Waited 30 seconds for completion.  there seems to be a problem,.")
                   return fmt.Errorf("GAH. Waited 30 seconds for completion.  there seems to be a problem.")
                }
                time.Sleep(1*time.Second)
        }
        count++
    }

    //update the router table.

    //Delete the data on the from server.
    request = cheshire.NewRequest(shards.PARTITION_DELETE, "DELETE")
    request.Params().Put("partition", partition)

    services.Logger.Printf("NOW Deleting the partition from the origin server %s", from.Address)
    response, err := fromClient.ApiCallSync(request, 300*time.Second)
    if response.StatusCode() != 200 {
        services.Logger.Printf("ERROR While deleting partition: %s", response.StatusMessage())
    }

    return nil
}

// Checks with entries to see if an updated router table is available.
func RouterTableUpdate(services *Services, routerTable *shards.RouterTable, maxChecks int) (*shards.RouterTable, bool) {
    checks := 0
    updated := false
    for _,e := range(routerTable.Entries) {
        if checks >= maxChecks {
            break
        }

        checks++
        // make sure our routertable is up to date.
        response, err := cheshire.HttpApiCallSync(
            fmt.Sprintf("%s:%d", e.Address, e.HttpPort),
            cheshire.NewRequest(shards.CHECKIN, "GET"),
            5 * time.Second)
        if err != nil {
            services.Logger.Printf("ERROR While contacting %s -- %s", e.Address, err)
            continue
        }

        rev := response.MustInt64("rt_revision", 0)
        if rev <= routerTable.Revision {
            continue
        }
        services.Logger.Printf("Found updated router table at: %s", e.Address)
        //get the new routertable.
        response, err = cheshire.HttpApiCallSync(
            fmt.Sprintf("%s:%d", e.Address, e.HttpPort),
            cheshire.NewRequest(shards.ROUTERTABLE_GET, "GET"),
            5 * time.Second)
        if err != nil {
            services.Logger.Printf("ERROR While contacting %s -- %s", e.Address, err)
            continue
        }
        mp, ok := response.GetDynMap("router_table")
        if !ok {
            services.Logger.Printf("ERROR from %s -- BAD ROUTER TABLE RESPONSE %s", e.Address, response)
            continue
        }

        rt, err := shards.ToRouterTable(mp)
        if err != nil {
            services.Logger.Printf("ERROR While parsing router table %s -- %s", e.Address, err)
            continue
        }

        services.Logger.Printf("SUCCESSFULLY update router table to revision %d", rt.Revision)
        updated = true
        routerTable = rt
    }
    return routerTable, updated
}

<|endoftext|>


<commit_before>package main

import (
	flags "github.com/jessevdk/go-flags"
	web "github.com/qmsk/go-web"
	server "github.com/qmsk/pngtile/go/server"

	"log"
)

var options struct {
	Web web.Options `group:"Web"`

	Path string `long:"pngtile-path"`
}

func main() {
	if args, err := flags.Parse(&options); err != nil {
		log.Fatalf("flags.Parse: %v", err)
	} else if len(args) > 0 {
		log.Fatalf("Usage: [options]")
	} else {

	}

	var config = server.Config{
		Path: options.Path,
	}

	if server, err := config.MakeServer(); err != nil {
		log.Fatalf("server:Config.MakeServer: %v", err)
	} else {
		options.Web.Server(
			options.Web.RouteStatic("/static/"),
			options.Web.Route("/", server),
		)
	}
}
<commit_msg>go pngtile-server: --pngtile-debug/quiet<commit_after>package main

import (
	flags "github.com/jessevdk/go-flags"
	web "github.com/qmsk/go-web"
	pngtile "github.com/qmsk/pngtile/go"
	server "github.com/qmsk/pngtile/go/server"

	"log"
)

var options struct {
	Web web.Options `group:"Web"`

	Path  string `long:"pngtile-path"`
	Debug bool   `long:"pngtile-debug"`
	Quiet bool   `long:"pngtile-quiet"`
}

func main() {
	if args, err := flags.Parse(&options); err != nil {
		log.Fatalf("flags.Parse: %v", err)
	} else if len(args) > 0 {
		log.Fatalf("Usage: [options]")
	} else {

	}

	pngtile.LogDebug(options.Debug)
	pngtile.LogWarn(!options.Quiet)

	var config = server.Config{
		Path: options.Path,
	}

	if server, err := config.MakeServer(); err != nil {
		log.Fatalf("server:Config.MakeServer: %v", err)
	} else {
		options.Web.Server(
			options.Web.RouteStatic("/static/"),
			options.Web.Route("/", server),
		)
	}
}
<|endoftext|>


<commit_before>package rest

import (
	"encoding/json"
	"fmt"
	"net/http"
	"socialapi/models"
)

func CreatePost(channelId, accountId int64) (*models.ChannelMessage, error) {
	return CreatePostWithBody(channelId, accountId, "create a message")
}

func CreatePostWithBody(channelId, accountId int64, body string) (*models.ChannelMessage, error) {
	cm := models.NewChannelMessage()
	cm.Body = body
	cm.AccountId = accountId

	return createPostRequest(channelId, cm, http.Header{})
}

func CreatePostWithHeader(channelId, accountId int64, header http.Header) (*models.ChannelMessage, error) {
	cm := models.NewChannelMessage()
	cm.Body = "Text1Text2"
	cm.AccountId = accountId

	return createPostRequest(channelId, cm, header)
}

func GetPost(id int64, accountId int64, groupName string) (*models.ChannelMessage, error) {
	url := fmt.Sprintf("/message/%d?accountId=%d&groupName=%s", id, accountId, groupName)
	cm := models.NewChannelMessage()
	cmI, err := sendModel("GET", url, cm)
	if err != nil {
		return nil, err
	}
	return cmI.(*models.ChannelMessage), nil
}

func GetPostBySlug(slug string, accountId int64) (*models.ChannelMessageContainer, error) {
	url := fmt.Sprintf("/message/slug/%s?accountId=%d", slug, accountId)
	cmc := models.NewChannelMessageContainer()
	cmI, err := sendModel("GET", url, cmc)
	if err != nil {
		return nil, err
	}
	return cmI.(*models.ChannelMessageContainer), nil
}

func DeletePost(id int64, accountId int64, groupName string) error {
	url := fmt.Sprintf("/message/%d?accountId=%d&groupName=%s", id, accountId, groupName)
	_, err := sendRequest("DELETE", url, nil)
	return err
}

func UpdatePost(cm *models.ChannelMessage) (*models.ChannelMessage, error) {
	cm.Body = "after update"

	url := fmt.Sprintf("/message/%d", cm.Id)
	cmI, err := sendModel("POST", url, cm)
	if err != nil {
		return nil, err
	}

	return cmI.(*models.ChannelMessage), nil
}

type PayloadRequest struct {
	Body      string                 `json:"body"`
	AccountId int64                  `json:"accountId,string"`
	Payload   map[string]interface{} `json:"payload"`
}

func CreatePostWithPayload(channelId, accountId int64, payload map[string]interface{}) (*models.ChannelMessage, error) {
	pr := PayloadRequest{}
	pr.Body = "message with payload"
	pr.AccountId = accountId
	pr.Payload = payload

	return createPostRequest(channelId, pr, http.Header{})
}

func createPostRequest(channelId int64, model interface{}, h http.Header) (*models.ChannelMessage, error) {
	url := fmt.Sprintf("/channel/%d/message", channelId)
	res, err := marshallAndSendRequestWithHeader("POST", url, model, h)
	if err != nil {
		return nil, err
	}

	container := models.NewChannelMessageContainer()
	err = json.Unmarshal(res, container)
	if err != nil {
		return nil, err
	}

	return container.Message, nil
}
<commit_msg>tests: add session data for requests related with messages<commit_after>package rest

import (
	"encoding/json"
	"fmt"
	"net/http"
	"socialapi/models"
)

func CreatePost(channelId, accountId int64, token string) (*models.ChannelMessage, error) {
	return CreatePostWithBodyAndAuth(channelId, accountId, "create a message", token)
}

func CreatePostWithBodyAndAuth(channelId, accountId int64, body, token string) (*models.ChannelMessage, error) {
	url := fmt.Sprintf("/channel/%d/message", channelId)
	cm := models.NewChannelMessage()
	cm.Body = body
	cm.AccountId = accountId
	res, err := marshallAndSendRequestWithAuth("POST", url, cm, token)
	if err != nil {
		return nil, err
	}

	container := models.NewChannelMessageContainer()
	err = json.Unmarshal(res, container)
	if err != nil {
		return nil, err
	}

	return container.Message, nil

	// return res.(*models.ChannelMessage), nil
}

func CreatePostWithBody(channelId, accountId int64, body string) (*models.ChannelMessage, error) {
	cm := models.NewChannelMessage()
	cm.Body = body
	cm.AccountId = accountId

	return createPostRequest(channelId, cm, http.Header{})
}

func CreatePostWithHeader(channelId, accountId int64, header http.Header) (*models.ChannelMessage, error) {
	cm := models.NewChannelMessage()
	cm.Body = "Text1Text2"
	cm.AccountId = accountId

	return createPostRequest(channelId, cm, header)
}

func GetPost(id int64, accountId int64, groupName string) (*models.ChannelMessage, error) {
	url := fmt.Sprintf("/message/%d?accountId=%d&groupName=%s", id, accountId, groupName)
	cm := models.NewChannelMessage()
	cmI, err := sendModel("GET", url, cm)
	if err != nil {
		return nil, err
	}
	return cmI.(*models.ChannelMessage), nil
}

func GetPostBySlug(slug string, accountId int64) (*models.ChannelMessageContainer, error) {
	url := fmt.Sprintf("/message/slug/%s?accountId=%d", slug, accountId)
	cmc := models.NewChannelMessageContainer()
	cmI, err := sendModel("GET", url, cmc)
	if err != nil {
		return nil, err
	}
	return cmI.(*models.ChannelMessageContainer), nil
}

func DeletePost(id int64, accountId int64, groupName, token string) error {
	url := fmt.Sprintf("/message/%d?accountId=%d&groupName=%s", id, accountId, groupName)
	_, err := sendRequestWithAuth("DELETE", url, nil, token)
	return err
}

func UpdatePost(cm *models.ChannelMessage, token string) (*models.ChannelMessage, error) {
	cm.Body = "after update"

	url := fmt.Sprintf("/message/%d", cm.Id)
	// cmI, err := sendModel("POST", url, cm)
	cmI, err := sendModelWithAuth("POST", url, cm, token)
	if err != nil {
		return nil, err
	}

	return cmI.(*models.ChannelMessage), nil
}

type PayloadRequest struct {
	Body      string                 `json:"body"`
	AccountId int64                  `json:"accountId,string"`
	Payload   map[string]interface{} `json:"payload"`
}

func CreatePostWithPayload(channelId, accountId int64, payload map[string]interface{}) (*models.ChannelMessage, error) {
	pr := PayloadRequest{}
	pr.Body = "message with payload"
	pr.AccountId = accountId
	pr.Payload = payload

	return createPostRequest(channelId, pr, http.Header{})
}

func createPostRequest(channelId int64, model interface{}, h http.Header) (*models.ChannelMessage, error) {
	url := fmt.Sprintf("/channel/%d/message", channelId)
	res, err := marshallAndSendRequestWithHeader("POST", url, model, h)
	if err != nil {
		return nil, err
	}

	container := models.NewChannelMessageContainer()
	err = json.Unmarshal(res, container)
	if err != nil {
		return nil, err
	}

	return container.Message, nil
}
<|endoftext|>


<commit_before>package algoliasearch

import "errors"

type IndexIterator struct {
	cursor string
	index  *Index
	page   BrowseRes
	params map[string]interface{}
	pos    int
}

func NewIndexIterator(index *Index, params map[string]interface{}) (it IndexIterator, err error) {
	it = IndexIterator{
		index:  index,
		params: duplicateMap(params),
		pos:    0,
	}
	err = it.loadNextPage()
	return
}

func (it *IndexIterator) Next() (res map[string]interface{}, err error) {
	// Abort if the user call `Next()` on a IndexIterator that has been
	// initialized without being able to load the first page.
	if len(it.page.Hits) == 0 {
		err = errors.New("No more hits")
		return
	}

	// If the last element of the page has been reached, the next one is loaded
	if it.pos == len(it.page.Hits) {
		if it.cursor == "" {
			err = errors.New("No more hits")

		} else {
			err = it.loadNextPage()
		}

		if err != nil {
			return
		}
	}

	res = it.page.Hits[it.pos]
	it.pos++

	return
}

func (it *IndexIterator) loadNextPage() (err error) {
	// Update the cursor for each new page except for the first one
	if it.cursor != "" {
		it.params["cursor"] = it.cursor
	}

	if it.page, err = it.index.Browse(it.params); err != nil {
		return
	}

	// Return an error if the newly loaded pages contains no results
	if len(it.page.Hits) == 0 {
		err = errors.New("No more hits")
		return
	}

	it.cursor = it.page.Cursor
	it.pos = 0
	return
}
<commit_msg>Remove extra line<commit_after>package algoliasearch

import "errors"

type IndexIterator struct {
	cursor string
	index  *Index
	page   BrowseRes
	params map[string]interface{}
	pos    int
}

func NewIndexIterator(index *Index, params map[string]interface{}) (it IndexIterator, err error) {
	it = IndexIterator{
		index:  index,
		params: duplicateMap(params),
		pos:    0,
	}
	err = it.loadNextPage()
	return
}

func (it *IndexIterator) Next() (res map[string]interface{}, err error) {
	// Abort if the user call `Next()` on a IndexIterator that has been
	// initialized without being able to load the first page.
	if len(it.page.Hits) == 0 {
		err = errors.New("No more hits")
		return
	}

	// If the last element of the page has been reached, the next one is loaded
	if it.pos == len(it.page.Hits) {
		if it.cursor == "" {
			err = errors.New("No more hits")
		} else {
			err = it.loadNextPage()
		}

		if err != nil {
			return
		}
	}

	res = it.page.Hits[it.pos]
	it.pos++

	return
}

func (it *IndexIterator) loadNextPage() (err error) {
	// Update the cursor for each new page except for the first one
	if it.cursor != "" {
		it.params["cursor"] = it.cursor
	}

	if it.page, err = it.index.Browse(it.params); err != nil {
		return
	}

	// Return an error if the newly loaded pages contains no results
	if len(it.page.Hits) == 0 {
		err = errors.New("No more hits")
		return
	}

	it.cursor = it.page.Cursor
	it.pos = 0
	return
}
<|endoftext|>


<commit_before>package system

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/qiniu/logkit/metric"
	. "github.com/qiniu/logkit/utils/models"
)

const (
	TypeMetricNet   = "net"
	MetricNetUsages = "网络设备状态(net)"

	// TypeMetricNet 信息中的字段
	KeyNetBytesSent       = "net_bytes_sent"
	KeyNetBytesSentPerSec = "net_bytes_sent_per_sec"
	KeyNetBytesRecv       = "net_bytes_recv"
	KeyNetBytesRecvPerSec = "net_bytes_recv_per_sec"
	KeyNetPacketsSent     = "net_packets_sent"
	KeyNetPacketsRecv     = "net_packets_recv"
	KeyNetErrIn           = "net_err_in"
	KeyNetErrOut          = "net_err_out"
	KeyNetDropIn          = "net_drop_in"
	KeyNetDropOut         = "net_drop_out"
	KeyNetInterface       = "net_interface"
)

// KeyNetUsages TypeMetricNet 中的字段名称
var KeyNetUsages = KeyValueSlice{
	{KeyNetBytesSent, "网卡发包总数(bytes)", ""},
	{KeyNetBytesSentPerSec, "网卡发包速率(bytes/s)", ""},
	{KeyNetBytesRecv, "网卡收包总数(bytes)", ""},
	{KeyNetBytesRecvPerSec, "网卡收包速率(bytes/s)", ""},
	{KeyNetPacketsSent, "网卡发包数量", ""},
	{KeyNetPacketsRecv, "网卡收包数量", ""},
	{KeyNetErrIn, "网卡收包错误数量", ""},
	{KeyNetErrOut, "网卡发包错误数量", ""},
	{KeyNetDropIn, "网卡收 丢包数量", ""},
	{KeyNetDropOut, "网卡发 丢包数量", ""},
	{KeyNetInterface, "网卡设备名称", ""},
}

type CollectInfo struct {
	timestamp time.Time
	BytesSent uint64
	BytesRecv uint64
}

type NetIOStats struct {
	ps          PS
	lastCollect map[string]CollectInfo

	skipChecks     bool
	skipProtoState bool     `json:"skip_protocols_state"`
	Interfaces     []string `json:"interfaces"`
}

func (_ *NetIOStats) Name() string {
	return TypeMetricNet
}

func (_ *NetIOStats) Usages() string {
	return MetricNetUsages
}

func (_ *NetIOStats) Tags() []string {
	return []string{KeyNetInterface}
}

func (_ *NetIOStats) Config() map[string]interface{} {
	configOption := []Option{
		{
			KeyName:      "interfaces",
			ChooseOnly:   false,
			Default:      "",
			DefaultNoUse: false,
			Description:  "收集特定网卡的信息,用','分隔(interfaces)",
			Type:         metric.ConfigTypeArray,
		},
		{
			KeyName:       "skip_protocols_state",
			Element:       Radio,
			ChooseOnly:    true,
			ChooseOptions: []interface{}{"true", "false"},
			Default:       "true",
			DefaultNoUse:  false,
			Description:   "是否忽略各个网络协议的状态信息",
			Type:          metric.ConfigTypeBool,
		},
	}
	config := map[string]interface{}{
		metric.OptionString:     configOption,
		metric.AttributesString: KeyNetUsages,
	}
	return config
}

func (s *NetIOStats) Collect() (datas []map[string]interface{}, err error) {
	netio, err := s.ps.NetIO()
	if err != nil {
		return nil, fmt.Errorf("error getting net io info: %s", err)
	}

	for _, io := range netio {
		if len(s.Interfaces) != 0 {
			var found bool

			for _, name := range s.Interfaces {
				if name == io.Name {
					found = true
					break
				}
			}

			if !found {
				continue
			}
		} else if !s.skipChecks {
			iface, err := net.InterfaceByName(io.Name)
			if err != nil {
				continue
			}

			if iface.Flags&net.FlagLoopback == net.FlagLoopback {
				continue
			}

			if iface.Flags&net.FlagUp == 0 {
				continue
			}
		}

		fields := map[string]interface{}{
			KeyNetBytesSent:   io.BytesSent,
			KeyNetBytesRecv:   io.BytesRecv,
			KeyNetPacketsSent: io.PacketsSent,
			KeyNetPacketsRecv: io.PacketsRecv,
			KeyNetErrIn:       io.Errin,
			KeyNetErrOut:      io.Errout,
			KeyNetDropIn:      io.Dropin,
			KeyNetDropOut:     io.Dropout,
			KeyNetInterface:   io.Name,
		}
		thisTime := time.Now()
		if info, ok := s.lastCollect[io.Name]; ok {
			dur := thisTime.Sub(info.timestamp)
			sentBytesDur := io.BytesSent - info.BytesSent
			recvBytesDur := io.BytesRecv - info.BytesRecv
			secs := float64(dur) / float64(time.Second)
			if secs > 0 {
				fields[KeyNetBytesSentPerSec] = uint64(float64(sentBytesDur) / secs)
				fields[KeyNetBytesRecvPerSec] = uint64(float64(recvBytesDur) / secs)
			}
		}
		s.lastCollect[io.Name] = CollectInfo{
			timestamp: thisTime,
			BytesRecv: io.BytesRecv,
			BytesSent: io.BytesSent,
		}
		datas = append(datas, fields)
	}

	if !s.skipProtoState {
		// Get system wide stats for different network protocols
		// (ignore these stats if the call fails)
		netprotos, _ := s.ps.NetProto()
		fields := make(map[string]interface{})
		for _, proto := range netprotos {
			for stat, value := range proto.Stats {
				name := TypeMetricNet + "_" + strings.ToLower(proto.Protocol) + "_" + strings.ToLower(stat)
				fields[name] = value
			}
		}
		fields[KeyNetInterface] = "all"
		datas = append(datas, fields)
	}
	return
}

func init() {
	metric.Add(TypeMetricNet, func() metric.Collector {
		return &NetIOStats{
			ps:          newSystemPS(),
			lastCollect: make(map[string]CollectInfo),
		}
	})
}
<commit_msg>logkit 增加网卡设备，用于服务端区分虚拟网卡<commit_after>package system

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/qiniu/logkit/metric"
	. "github.com/qiniu/logkit/utils/models"
)

const (
	TypeMetricNet   = "net"
	MetricNetUsages = "网络设备状态(net)"

	// TypeMetricNet 信息中的字段
	KeyNetBytesSent         = "net_bytes_sent"
	KeyNetBytesSentPerSec   = "net_bytes_sent_per_sec"
	KeyNetBytesRecv         = "net_bytes_recv"
	KeyNetBytesRecvPerSec   = "net_bytes_recv_per_sec"
	KeyNetPacketsSent       = "net_packets_sent"
	KeyNetPacketsSentPerSec = "net_packets_sent_per_sec"
	KeyNetPacketsRecv       = "net_packets_recv"
	KeyNetPacketsRecvPerSec = "net_packets_recv_per_sec"
	KeyNetErrIn             = "net_err_in"
	KeyNetErrOut            = "net_err_out"
	KeyNetDropIn            = "net_drop_in"
	KeyNetDropOut           = "net_drop_out"
	KeyNetInterface         = "net_interface"
	KeyNetHWAddr            = "net_hw_addr"
)

// KeyNetUsages TypeMetricNet 中的字段名称
var KeyNetUsages = KeyValueSlice{
	{KeyNetBytesSent, "网卡发包总数(bytes)", ""},
	{KeyNetBytesSentPerSec, "网卡发包速率(bytes/s)", ""},
	{KeyNetBytesRecv, "网卡收包总数(bytes)", ""},
	{KeyNetBytesRecvPerSec, "网卡收包速率(bytes/s)", ""},
	{KeyNetPacketsSent, "网卡发包数量", ""},
	{KeyNetPacketsSentPerSec, "每秒网卡发包数量", ""},
	{KeyNetPacketsRecv, "网卡收包数量", ""},
	{KeyNetPacketsRecvPerSec, "每秒网卡收包数量", ""},
	{KeyNetErrIn, "网卡收包错误数量", ""},
	{KeyNetErrOut, "网卡发包错误数量", ""},
	{KeyNetDropIn, "网卡收 丢包数量", ""},
	{KeyNetDropOut, "网卡发 丢包数量", ""},
	{KeyNetInterface, "网卡设备名称", ""},
	{KeyNetHWAddr, "网卡设备地址", ""},
}

type CollectInfo struct {
	timestamp   time.Time
	BytesSent   uint64
	BytesRecv   uint64
	PacketsSent uint64
	PacketsRecv uint64
}

type NetIOStats struct {
	ps          PS
	lastCollect map[string]CollectInfo

	skipChecks     bool
	skipProtoState bool     `json:"skip_protocols_state"`
	Interfaces     []string `json:"interfaces"`
	InterfacesMap  map[string]net.Interface
}

func (_ *NetIOStats) Name() string {
	return TypeMetricNet
}

func (_ *NetIOStats) Usages() string {
	return MetricNetUsages
}

func (_ *NetIOStats) Tags() []string {
	return []string{KeyNetInterface}
}

func (_ *NetIOStats) Config() map[string]interface{} {
	configOption := []Option{
		{
			KeyName:      "interfaces",
			ChooseOnly:   false,
			Default:      "",
			DefaultNoUse: false,
			Description:  "收集特定网卡的信息,用','分隔(interfaces)",
			Type:         metric.ConfigTypeArray,
		},
		{
			KeyName:       "skip_protocols_state",
			Element:       Radio,
			ChooseOnly:    true,
			ChooseOptions: []interface{}{"true", "false"},
			Default:       "true",
			DefaultNoUse:  false,
			Description:   "是否忽略各个网络协议的状态信息",
			Type:          metric.ConfigTypeBool,
		},
	}
	config := map[string]interface{}{
		metric.OptionString:     configOption,
		metric.AttributesString: KeyNetUsages,
	}
	return config
}

func (s *NetIOStats) Collect() (datas []map[string]interface{}, err error) {
	netio, err := s.ps.NetIO()
	if err != nil {
		return nil, fmt.Errorf("error getting net io info: %s", err)
	}

	for _, io := range netio {
		hwAddr := ""
		if _, ok := s.InterfacesMap[io.Name]; !ok {
			s.initInterfaces()
		}

		if iface, ok := s.InterfacesMap[io.Name]; ok {
			hwAddr = fmt.Sprintf("%v", iface.HardwareAddr)
		}

		if len(s.Interfaces) != 0 {
			var found bool

			for _, name := range s.Interfaces {
				if name == io.Name {
					found = true
					break
				}
			}

			if !found {
				continue
			}
		} else if !s.skipChecks {
			iface, err := net.InterfaceByName(io.Name)
			if err != nil {
				continue
			}

			if iface.Flags&net.FlagLoopback == net.FlagLoopback {
				continue
			}

			if iface.Flags&net.FlagUp == 0 {
				continue
			}
		}

		fields := map[string]interface{}{
			KeyNetBytesSent:   io.BytesSent,
			KeyNetBytesRecv:   io.BytesRecv,
			KeyNetPacketsSent: io.PacketsSent,
			KeyNetPacketsRecv: io.PacketsRecv,
			KeyNetErrIn:       io.Errin,
			KeyNetErrOut:      io.Errout,
			KeyNetDropIn:      io.Dropin,
			KeyNetDropOut:     io.Dropout,
			KeyNetInterface:   io.Name,
			KeyNetHWAddr:      hwAddr,
		}
		thisTime := time.Now()
		if info, ok := s.lastCollect[io.Name]; ok {
			dur := thisTime.Sub(info.timestamp)
			sentBytesDur := io.BytesSent - info.BytesSent
			recvBytesDur := io.BytesRecv - info.BytesRecv
			secs := float64(dur) / float64(time.Second)
			if secs > 0 {
				fields[KeyNetBytesSentPerSec] = uint64(float64(sentBytesDur) / secs)
				fields[KeyNetBytesRecvPerSec] = uint64(float64(recvBytesDur) / secs)
				fields[KeyNetPacketsRecvPerSec] = uint64(float64(io.PacketsRecv-info.PacketsRecv) / secs)
				fields[KeyNetPacketsSentPerSec] = uint64(float64(io.PacketsSent-info.PacketsSent) / secs)
			}
		}
		s.lastCollect[io.Name] = CollectInfo{
			timestamp: thisTime,
			BytesRecv: io.BytesRecv,
			BytesSent: io.BytesSent,
		}
		datas = append(datas, fields)
	}

	if !s.skipProtoState {
		// Get system wide stats for different network protocols
		// (ignore these stats if the call fails)
		netprotos, _ := s.ps.NetProto()
		fields := make(map[string]interface{})
		for _, proto := range netprotos {
			for stat, value := range proto.Stats {
				name := TypeMetricNet + "_" + strings.ToLower(proto.Protocol) + "_" + strings.ToLower(stat)
				fields[name] = value
			}
		}
		fields[KeyNetInterface] = "all"
		datas = append(datas, fields)
	}
	return
}

func (s *NetIOStats) initInterfaces() {
	if interfaces, err := net.Interfaces(); err == nil {
		for _, _interface := range interfaces {
			s.InterfacesMap[_interface.Name] = _interface
		}
	}
}

func init() {
	netIoStats := &NetIOStats{
		ps:            newSystemPS(),
		lastCollect:   make(map[string]CollectInfo),
		InterfacesMap: make(map[string]net.Interface, 0),
	}
	netIoStats.initInterfaces()
	metric.Add(TypeMetricNet, func() metric.Collector {
		return netIoStats
	})
}
<|endoftext|>


<commit_before>package testing

import (
	"bytes"
	"fmt"
	"testing"
	"time"

	"go.pedge.io/proto/stream"
	"go.pedge.io/protolog"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pkg/require"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
)

const (
	defaultJobTimeoutSec = 30
)

var (
	baseSimpleTest = &simpleTest{
		inputFilePathToContent: map[string][]byte{
			"1": []byte("test1"),
			"2": []byte("test2"),
			"3": []byte("test3"),
		},
		outputFilePathToContent: map[string][]byte{
			"1.output": []byte("test1"),
			"2.output": []byte("test2"),
			"3.output": []byte("test3"),
		},
		transform: &pps.Transform{
			Image: "ubuntu:14.04",
			Cmd: []string{
				fmt.Sprintf("for file in /pfs/*; do cp \"${file}\" /pfs/$(basename \"${file}\").output; done"),
			},
		},
	}
)

func TestSimpleJob(t *testing.T) {
	t.Skip()
	runSimpleJobTest(t, baseSimpleTest.copy())
}

func TestSimplePipeline(t *testing.T) {
	t.Skip()
	runSimplePipelineTest(t, baseSimpleTest.copy())
}

func runSimpleJobTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runJob)
}

func runSimplePipelineTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runPipeline)
}

type simpleTest struct {
	inputRepoName           string
	outputRepoName          string
	inputFilePathToContent  map[string][]byte
	outputFilePathToContent map[string][]byte
	transform               *pps.Transform
	pipelineName            string
	jobTimeoutSec           int
	expectError             bool
}

func (s *simpleTest) copy() *simpleTest {
	c := &simpleTest{
		inputRepoName:           s.inputRepoName,
		outputRepoName:          s.outputRepoName,
		inputFilePathToContent:  make(map[string][]byte),
		outputFilePathToContent: make(map[string][]byte),
		pipelineName:            s.pipelineName,
		jobTimeoutSec:           s.jobTimeoutSec,
		expectError:             s.expectError,
	}
	for key, value := range s.inputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.inputFilePathToContent[key] = buf
	}
	for key, value := range s.outputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.outputFilePathToContent[key] = buf
	}
	if s.transform != nil {
		c.transform = &pps.Transform{
			Image: s.transform.Image,
			Cmd:   make([]string, len(s.transform.Cmd)),
		}
		for i, value := range s.transform.Cmd {
			c.transform.Cmd[i] = value
		}
	}
	return c
}

func (s *simpleTest) setup() {
	if s.inputRepoName == "" {
		s.inputRepoName = uuid.NewWithoutDashes()
	}
	if s.outputRepoName == "" {
		s.outputRepoName = uuid.NewWithoutDashes()
	}
	if s.pipelineName == "" {
		s.pipelineName = uuid.NewWithoutDashes()
	}
	if s.jobTimeoutSec == 0 {
		s.jobTimeoutSec = defaultJobTimeoutSec
	}
}

func (s *simpleTest) runJob(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	inputCommit := setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	outputParentCommit := setupPFSOutputParentCommit(t, pfsAPIClient, outputRepo)
	job := createJob(t, jobAPIClient, s.transform, inputCommit, outputParentCommit)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.OutputCommit, s.outputFilePathToContent)
}

func (s *simpleTest) runPipeline(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	pipeline := createPipeline(t, pipelineAPIClient, s.pipelineName, s.transform, inputRepo, outputRepo)
	setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	// TODO: all the waits heh
	time.Sleep(5 * time.Second)
	job := getJobForPipeline(t, jobAPIClient, pipeline)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.OutputCommit, s.outputFilePathToContent)
}

func setupPFSRepo(t *testing.T, pfsAPIClient pfs.APIClient, repoName string) *pfs.Repo {
	repo := &pfs.Repo{
		Name: repoName,
	}
	_, err := pfsAPIClient.CreateRepo(
		context.Background(),
		&pfs.CreateRepoRequest{
			Repo: repo,
		},
	)
	require.NoError(t, err)
	return repo
}

// TODO: handle directories in filePathToContent
func setupPFSInputCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo, filePathToContent map[string][]byte) *pfs.Commit {
	commit, err := pfsAPIClient.StartCommit(
		context.Background(),
		&pfs.StartCommitRequest{
			Parent: &pfs.Commit{
				Repo: repo,
				Id:   pfs.InitialCommitID,
			},
		},
	)
	require.NoError(t, err)
	for filePath, content := range filePathToContent {
		apiPutFileClient, err := pfsAPIClient.PutFile(
			context.Background(),
		)
		require.NoError(t, err)
		err = apiPutFileClient.Send(
			&pfs.PutFileRequest{
				File: &pfs.File{
					Commit: commit,
					Path:   filePath,
				},
				FileType: pfs.FileType_FILE_TYPE_REGULAR,
				Value:    content,
			},
		)
		require.NoError(t, err)
		_, _ = apiPutFileClient.CloseAndRecv()
	}
	_, err = pfsAPIClient.FinishCommit(
		context.Background(),
		&pfs.FinishCommitRequest{
			Commit: commit,
		},
	)
	require.NoError(t, err)
	return commit
}

func setupPFSOutputParentCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo) *pfs.Commit {
	return &pfs.Commit{Repo: repo, Id: pfs.InitialCommitID}
}

func createJob(t *testing.T, jobAPIClient pps.JobAPIClient, transform *pps.Transform, inputCommit *pfs.Commit, outputParentCommit *pfs.Commit) *pps.Job {
	job, err := jobAPIClient.CreateJob(
		context.Background(),
		&pps.CreateJobRequest{
			Spec: &pps.CreateJobRequest_Transform{
				Transform: transform,
			},
			Input:        inputCommit,
			OutputParent: outputParentCommit,
		},
	)
	require.NoError(t, err)
	return job
}

func waitForJob(t *testing.T, jobAPIClient pps.JobAPIClient, job *pps.Job, timeoutSec int, expectError bool) *pps.JobInfo {
	for i := 0; i < timeoutSec; i++ {
		time.Sleep(1 * time.Second)
		jobInfo, err := jobAPIClient.InspectJob(
			context.Background(),
			&pps.InspectJobRequest{
				Job: job,
			},
		)
		require.NoError(t, err)
		if len(jobInfo.JobStatus) == 0 {
			continue
		}
		jobStatus := jobInfo.JobStatus[0]
		protolog.Infof("status of job %s at %d seconds: %v", job.Id, i+1, jobInfo.JobStatus)
		switch jobStatus.Type {
		case pps.JobStatusType_JOB_STATUS_TYPE_ERROR:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if !expectError {
				t.Fatalf("job %s had error", job.Id)
			}
			return jobInfo
		case pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if expectError {
				t.Fatalf("job %s did not have error", job.Id)
			}
			return jobInfo
		}
	}
	t.Fatalf("job %s did not finish in %d seconds", job.Id, timeoutSec)
	return nil
}

func handleEndJobStatus(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo) {
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDOUT)
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDERR)
}

func handleEndJobStatusOutputStream(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo, outputStream pps.OutputStream) {
	jobAPIGetJobLogsClient, err := jobAPIClient.GetJobLogs(
		context.Background(),
		&pps.GetJobLogsRequest{
			Job:          jobInfo.Job,
			OutputStream: outputStream,
		},
	)
	require.NoError(t, err)
	require.NoError(t, protostream.WriteFromStreamingBytesClient(jobAPIGetJobLogsClient, protolog.Writer()))
}

func createPipeline(t *testing.T, pipelineAPIClient pps.PipelineAPIClient, pipelineName string, transform *pps.Transform, inputRepo *pfs.Repo, outputRepo *pfs.Repo) *pps.Pipeline {
	pipeline := &pps.Pipeline{
		Name: pipelineName,
	}
	_, err := pipelineAPIClient.CreatePipeline(
		context.Background(),
		&pps.CreatePipelineRequest{
			Pipeline:  pipeline,
			Transform: transform,
			Input:     inputRepo,
			Output:    outputRepo,
		},
	)
	require.NoError(t, err)
	return pipeline
}

func getJobForPipeline(t *testing.T, jobAPIClient pps.JobAPIClient, pipeline *pps.Pipeline) *pps.Job {
	jobInfos, err := jobAPIClient.ListJob(
		context.Background(),
		&pps.ListJobRequest{
			Pipeline: pipeline,
		},
	)
	require.NoError(t, err)
	require.NotNil(t, jobInfos)
	require.Equal(t, 1, len(jobInfos.JobInfo))
	return jobInfos.JobInfo[0].Job
}

func checkPFSOutput(t *testing.T, pfsAPIClient pfs.APIClient, outputCommit *pfs.Commit, filePathToContent map[string][]byte) {
	for filePath, content := range filePathToContent {
		getContent, err := getPFSContent(pfsAPIClient, outputCommit, filePath)
		require.NoError(t, err)
		require.Equal(t, content, getContent)
	}
}

func getPFSContent(pfsAPIClient pfs.APIClient, commit *pfs.Commit, filePath string) ([]byte, error) {
	apiGetFileClient, err := pfsAPIClient.GetFile(
		context.Background(),
		&pfs.GetFileRequest{
			File: &pfs.File{
				Commit: commit,
				Path:   filePath,
			},
		},
	)
	if err != nil {
		return nil, err
	}
	buffer := bytes.NewBuffer(nil)
	if err := protostream.WriteFromStreamingBytesClient(apiGetFileClient, buffer); err != nil {
		return nil, err
	}
	return buffer.Bytes(), nil
}
<commit_msg>Make pps tests compile.<commit_after>package testing

import (
	"bytes"
	"fmt"
	"testing"
	"time"

	"go.pedge.io/proto/stream"
	"go.pedge.io/protolog"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pkg/require"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
)

const (
	defaultJobTimeoutSec = 30
)

var (
	baseSimpleTest = &simpleTest{
		inputFilePathToContent: map[string][]byte{
			"1": []byte("test1"),
			"2": []byte("test2"),
			"3": []byte("test3"),
		},
		outputFilePathToContent: map[string][]byte{
			"1.output": []byte("test1"),
			"2.output": []byte("test2"),
			"3.output": []byte("test3"),
		},
		transform: &pps.Transform{
			Image: "ubuntu:14.04",
			Cmd: []string{
				fmt.Sprintf("for file in /pfs/*; do cp \"${file}\" /pfs/$(basename \"${file}\").output; done"),
			},
		},
	}
)

func TestSimpleJob(t *testing.T) {
	t.Skip()
	runSimpleJobTest(t, baseSimpleTest.copy())
}

func TestSimplePipeline(t *testing.T) {
	t.Skip()
	runSimplePipelineTest(t, baseSimpleTest.copy())
}

func runSimpleJobTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runJob)
}

func runSimplePipelineTest(t *testing.T, simpleTest *simpleTest) {
	RunTest(t, simpleTest.runPipeline)
}

type simpleTest struct {
	inputRepoName           string
	outputRepoName          string
	inputFilePathToContent  map[string][]byte
	outputFilePathToContent map[string][]byte
	transform               *pps.Transform
	pipelineName            string
	jobTimeoutSec           int
	expectError             bool
}

func (s *simpleTest) copy() *simpleTest {
	c := &simpleTest{
		inputRepoName:           s.inputRepoName,
		outputRepoName:          s.outputRepoName,
		inputFilePathToContent:  make(map[string][]byte),
		outputFilePathToContent: make(map[string][]byte),
		pipelineName:            s.pipelineName,
		jobTimeoutSec:           s.jobTimeoutSec,
		expectError:             s.expectError,
	}
	for key, value := range s.inputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.inputFilePathToContent[key] = buf
	}
	for key, value := range s.outputFilePathToContent {
		buf := make([]byte, len(value))
		copy(buf, value)
		c.outputFilePathToContent[key] = buf
	}
	if s.transform != nil {
		c.transform = &pps.Transform{
			Image: s.transform.Image,
			Cmd:   make([]string, len(s.transform.Cmd)),
		}
		for i, value := range s.transform.Cmd {
			c.transform.Cmd[i] = value
		}
	}
	return c
}

func (s *simpleTest) setup() {
	if s.inputRepoName == "" {
		s.inputRepoName = uuid.NewWithoutDashes()
	}
	if s.outputRepoName == "" {
		s.outputRepoName = uuid.NewWithoutDashes()
	}
	if s.pipelineName == "" {
		s.pipelineName = uuid.NewWithoutDashes()
	}
	if s.jobTimeoutSec == 0 {
		s.jobTimeoutSec = defaultJobTimeoutSec
	}
}

func (s *simpleTest) runJob(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	inputCommit := setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	outputParentCommit := setupPFSOutputParentCommit(t, pfsAPIClient, outputRepo)
	job := createJob(t, jobAPIClient, s.transform, inputCommit, outputParentCommit)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.OutputCommit, s.outputFilePathToContent)
}

func (s *simpleTest) runPipeline(
	t *testing.T,
	pfsAPIClient pfs.APIClient,
	jobAPIClient pps.JobAPIClient,
	pipelineAPIClient pps.PipelineAPIClient,
) {
	s.setup()
	inputRepo := setupPFSRepo(t, pfsAPIClient, s.inputRepoName)
	outputRepo := setupPFSRepo(t, pfsAPIClient, s.outputRepoName)
	pipeline := createPipeline(t, pipelineAPIClient, s.pipelineName, s.transform, inputRepo, outputRepo)
	setupPFSInputCommit(t, pfsAPIClient, inputRepo, s.inputFilePathToContent)
	// TODO: all the waits heh
	time.Sleep(5 * time.Second)
	job := getJobForPipeline(t, jobAPIClient, pipeline)
	jobInfo := waitForJob(t, jobAPIClient, job, s.jobTimeoutSec, s.expectError)
	checkPFSOutput(t, pfsAPIClient, jobInfo.OutputCommit, s.outputFilePathToContent)
}

func setupPFSRepo(t *testing.T, pfsAPIClient pfs.APIClient, repoName string) *pfs.Repo {
	repo := &pfs.Repo{
		Name: repoName,
	}
	_, err := pfsAPIClient.CreateRepo(
		context.Background(),
		&pfs.CreateRepoRequest{
			Repo: repo,
		},
	)
	require.NoError(t, err)
	return repo
}

// TODO: handle directories in filePathToContent
func setupPFSInputCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo, filePathToContent map[string][]byte) *pfs.Commit {
	commit, err := pfsAPIClient.StartCommit(
		context.Background(),
		&pfs.StartCommitRequest{
			Parent: &pfs.Commit{
				Repo: repo,
				Id:   pfs.InitialCommitID,
			},
		},
	)
	require.NoError(t, err)
	for filePath, content := range filePathToContent {
		apiPutFileClient, err := pfsAPIClient.PutFile(
			context.Background(),
		)
		require.NoError(t, err)
		err = apiPutFileClient.Send(
			&pfs.PutFileRequest{
				File: &pfs.File{
					Commit: commit,
					Path:   filePath,
				},
				FileType: pfs.FileType_FILE_TYPE_REGULAR,
				Value:    content,
			},
		)
		require.NoError(t, err)
		_, _ = apiPutFileClient.CloseAndRecv()
	}
	_, err = pfsAPIClient.FinishCommit(
		context.Background(),
		&pfs.FinishCommitRequest{
			Commit: commit,
		},
	)
	require.NoError(t, err)
	return commit
}

func setupPFSOutputParentCommit(t *testing.T, pfsAPIClient pfs.APIClient, repo *pfs.Repo) *pfs.Commit {
	return &pfs.Commit{Repo: repo, Id: pfs.InitialCommitID}
}

func createJob(t *testing.T, jobAPIClient pps.JobAPIClient, transform *pps.Transform, inputCommit *pfs.Commit, outputParentCommit *pfs.Commit) *pps.Job {
	job, err := jobAPIClient.CreateJob(
		context.Background(),
		&pps.CreateJobRequest{
			Spec: &pps.CreateJobRequest_Transform{
				Transform: transform,
			},
			InputCommit:  []*pfs.Commit{inputCommit},
			OutputParent: outputParentCommit,
		},
	)
	require.NoError(t, err)
	return job
}

func waitForJob(t *testing.T, jobAPIClient pps.JobAPIClient, job *pps.Job, timeoutSec int, expectError bool) *pps.JobInfo {
	for i := 0; i < timeoutSec; i++ {
		time.Sleep(1 * time.Second)
		jobInfo, err := jobAPIClient.InspectJob(
			context.Background(),
			&pps.InspectJobRequest{
				Job: job,
			},
		)
		require.NoError(t, err)
		if len(jobInfo.JobStatus) == 0 {
			continue
		}
		jobStatus := jobInfo.JobStatus[0]
		protolog.Infof("status of job %s at %d seconds: %v", job.Id, i+1, jobInfo.JobStatus)
		switch jobStatus.Type {
		case pps.JobStatusType_JOB_STATUS_TYPE_ERROR:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if !expectError {
				t.Fatalf("job %s had error", job.Id)
			}
			return jobInfo
		case pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS:
			handleEndJobStatus(t, jobAPIClient, jobInfo)
			if expectError {
				t.Fatalf("job %s did not have error", job.Id)
			}
			return jobInfo
		}
	}
	t.Fatalf("job %s did not finish in %d seconds", job.Id, timeoutSec)
	return nil
}

func handleEndJobStatus(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo) {
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDOUT)
	handleEndJobStatusOutputStream(t, jobAPIClient, jobInfo, pps.OutputStream_OUTPUT_STREAM_STDERR)
}

func handleEndJobStatusOutputStream(t *testing.T, jobAPIClient pps.JobAPIClient, jobInfo *pps.JobInfo, outputStream pps.OutputStream) {
	jobAPIGetJobLogsClient, err := jobAPIClient.GetJobLogs(
		context.Background(),
		&pps.GetJobLogsRequest{
			Job:          jobInfo.Job,
			OutputStream: outputStream,
		},
	)
	require.NoError(t, err)
	require.NoError(t, protostream.WriteFromStreamingBytesClient(jobAPIGetJobLogsClient, protolog.Writer()))
}

func createPipeline(t *testing.T, pipelineAPIClient pps.PipelineAPIClient, pipelineName string, transform *pps.Transform, inputRepo *pfs.Repo, outputRepo *pfs.Repo) *pps.Pipeline {
	pipeline := &pps.Pipeline{
		Name: pipelineName,
	}
	_, err := pipelineAPIClient.CreatePipeline(
		context.Background(),
		&pps.CreatePipelineRequest{
			Pipeline:   pipeline,
			Transform:  transform,
			InputRepo:  []*pfs.Repo{inputRepo},
			OutputRepo: outputRepo,
		},
	)
	require.NoError(t, err)
	return pipeline
}

func getJobForPipeline(t *testing.T, jobAPIClient pps.JobAPIClient, pipeline *pps.Pipeline) *pps.Job {
	jobInfos, err := jobAPIClient.ListJob(
		context.Background(),
		&pps.ListJobRequest{
			Pipeline: pipeline,
		},
	)
	require.NoError(t, err)
	require.NotNil(t, jobInfos)
	require.Equal(t, 1, len(jobInfos.JobInfo))
	return jobInfos.JobInfo[0].Job
}

func checkPFSOutput(t *testing.T, pfsAPIClient pfs.APIClient, outputCommit *pfs.Commit, filePathToContent map[string][]byte) {
	for filePath, content := range filePathToContent {
		getContent, err := getPFSContent(pfsAPIClient, outputCommit, filePath)
		require.NoError(t, err)
		require.Equal(t, content, getContent)
	}
}

func getPFSContent(pfsAPIClient pfs.APIClient, commit *pfs.Commit, filePath string) ([]byte, error) {
	apiGetFileClient, err := pfsAPIClient.GetFile(
		context.Background(),
		&pfs.GetFileRequest{
			File: &pfs.File{
				Commit: commit,
				Path:   filePath,
			},
		},
	)
	if err != nil {
		return nil, err
	}
	buffer := bytes.NewBuffer(nil)
	if err := protostream.WriteFromStreamingBytesClient(apiGetFileClient, buffer); err != nil {
		return nil, err
	}
	return buffer.Bytes(), nil
}
<|endoftext|>


<commit_before>package middleware

import (
	"crypto/tls"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"strconv"

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

type logger struct {
	logger lager.Logger
}

func NewLogger(l lager.Logger) Middleware {
	return logger{
		logger: l,
	}
}

func (l logger) Wrap(next http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		loggingResponseWriter := responseWriter{
			rw,
			[]byte{},
			0,
			0,
		}
		next.ServeHTTP(&loggingResponseWriter, req)

		requestCopy := *req
		requestCopy.Header["Authorization"] = nil

		response := map[string]interface{}{
			"Header":     loggingResponseWriter.Header(),
			"StatusCode": loggingResponseWriter.statusCode,
			"Size":       loggingResponseWriter.size,
		}

		l.logger.Debug("", lager.Data{
			"request":  fromHTTPRequest(requestCopy),
			"response": response,
		})
	})
}

type responseWriter struct {
	http.ResponseWriter
	body       []byte
	statusCode int
	size       int
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if rw.statusCode == 0 {
		rw.WriteHeader(http.StatusOK)
	}

	size, err := rw.ResponseWriter.Write(b)
	rw.body = append(rw.body, b...)
	rw.size += size

	rw.Header().Set("Content-Length", strconv.Itoa(rw.size))

	return size, err
}

func (rw *responseWriter) WriteHeader(s int) {
	rw.statusCode = s
	rw.ResponseWriter.WriteHeader(s)
}

// Golang 1.5 introduces a http.Request.Cancel field,
// of type <-chan struct{} which the lager library fails to deal with.
// We introduced loggableHTTPRequest as a way of handling this
// until such time as lager can deal with it.
//
// Once lager can handle the request directly, remove this struct.
// #101259402
type LoggableHTTPRequest struct {
	Method           string
	URL              *url.URL
	Proto            string
	ProtoMajor       int
	ProtoMinor       int
	Header           http.Header
	Body             io.ReadCloser
	ContentLength    int64
	TransferEncoding []string
	Close            bool
	Host             string
	Form             url.Values
	PostForm         url.Values
	MultipartForm    *multipart.Form
	Trailer          http.Header
	RemoteAddr       string
	RequestURI       string
	TLS              *tls.ConnectionState
}

func fromHTTPRequest(req http.Request) LoggableHTTPRequest {
	var form, postForm url.Values
	if req.Form != nil {
		form = sanitizeCredentialsFromForm(req.Form)
	}

	if req.PostForm != nil {
		postForm = sanitizeCredentialsFromForm(req.PostForm)
	}

	return LoggableHTTPRequest{
		Method:           req.Method,
		URL:              req.URL,
		Proto:            req.Proto,
		ProtoMajor:       req.ProtoMajor,
		ProtoMinor:       req.ProtoMinor,
		Header:           req.Header,
		Body:             req.Body,
		ContentLength:    req.ContentLength,
		TransferEncoding: req.TransferEncoding,
		Close:            req.Close,
		Host:             req.Host,
		Form:             form,
		PostForm:         postForm,
		MultipartForm:    req.MultipartForm,
		Trailer:          req.Trailer,
		RemoteAddr:       req.RemoteAddr,
		RequestURI:       req.RequestURI,
		TLS:              req.TLS,
	}
}

func sanitizeCredentialsFromForm(form url.Values) url.Values {
	form.Set("password", "***")
	return form
}
<commit_msg>Remove comment about removing LoggableHTTPRequest<commit_after>package middleware

import (
	"crypto/tls"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"strconv"

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

type logger struct {
	logger lager.Logger
}

func NewLogger(l lager.Logger) Middleware {
	return logger{
		logger: l,
	}
}

func (l logger) Wrap(next http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		loggingResponseWriter := responseWriter{
			rw,
			[]byte{},
			0,
			0,
		}
		next.ServeHTTP(&loggingResponseWriter, req)

		requestCopy := *req
		requestCopy.Header["Authorization"] = nil

		response := map[string]interface{}{
			"Header":     loggingResponseWriter.Header(),
			"StatusCode": loggingResponseWriter.statusCode,
			"Size":       loggingResponseWriter.size,
		}

		l.logger.Debug("", lager.Data{
			"request":  fromHTTPRequest(requestCopy),
			"response": response,
		})
	})
}

type responseWriter struct {
	http.ResponseWriter
	body       []byte
	statusCode int
	size       int
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	if rw.statusCode == 0 {
		rw.WriteHeader(http.StatusOK)
	}

	size, err := rw.ResponseWriter.Write(b)
	rw.body = append(rw.body, b...)
	rw.size += size

	rw.Header().Set("Content-Length", strconv.Itoa(rw.size))

	return size, err
}

func (rw *responseWriter) WriteHeader(s int) {
	rw.statusCode = s
	rw.ResponseWriter.WriteHeader(s)
}

type LoggableHTTPRequest struct {
	Method           string
	URL              *url.URL
	Proto            string
	ProtoMajor       int
	ProtoMinor       int
	Header           http.Header
	Body             io.ReadCloser
	ContentLength    int64
	TransferEncoding []string
	Close            bool
	Host             string
	Form             url.Values
	PostForm         url.Values
	MultipartForm    *multipart.Form
	Trailer          http.Header
	RemoteAddr       string
	RequestURI       string
	TLS              *tls.ConnectionState
}

func fromHTTPRequest(req http.Request) LoggableHTTPRequest {
	var form, postForm url.Values
	if req.Form != nil {
		form = sanitizeCredentialsFromForm(req.Form)
	}

	if req.PostForm != nil {
		postForm = sanitizeCredentialsFromForm(req.PostForm)
	}

	return LoggableHTTPRequest{
		Method:           req.Method,
		URL:              req.URL,
		Proto:            req.Proto,
		ProtoMajor:       req.ProtoMajor,
		ProtoMinor:       req.ProtoMinor,
		Header:           req.Header,
		Body:             req.Body,
		ContentLength:    req.ContentLength,
		TransferEncoding: req.TransferEncoding,
		Close:            req.Close,
		Host:             req.Host,
		Form:             form,
		PostForm:         postForm,
		MultipartForm:    req.MultipartForm,
		Trailer:          req.Trailer,
		RemoteAddr:       req.RemoteAddr,
		RequestURI:       req.RequestURI,
		TLS:              req.TLS,
	}
}

func sanitizeCredentialsFromForm(form url.Values) url.Values {
	form.Set("password", "***")
	return form
}
<|endoftext|>


<commit_before>package middleware

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/Clever/kayvee-go.v6/logger"
)

var globalRollupRouter *RollupRouter

// RollupLogger will log info / error rollups depending on status code.
type RollupLogger interface {
	InfoD(title string, data map[string]interface{})
	ErrorD(title string, data map[string]interface{})
}

// EnableRollups turns on rollups for kv middleware logs.
func EnableRollups(ctx context.Context, logger RollupLogger, reportingInterval time.Duration) {
	globalRollupRouter = NewRollupRouter(ctx, logger, reportingInterval)
}

// RollupRouter rolls up log lines and periodically logs them as one log line.
type RollupRouter struct {
	logger         RollupLogger
	reportingDelay time.Duration
	ctx            context.Context
	ctxDone        bool

	// create a rollup object per unique (status-code, path) pair
	rollupsMu sync.Mutex
	rollups   map[string]*logRollup
}

// NewRollupRouter creates a new log rollup output.
// Rollups will stop when the context is canceled.
func NewRollupRouter(ctx context.Context, logger RollupLogger, reportingDelay time.Duration) *RollupRouter {
	l := &RollupRouter{
		logger:         logger,
		reportingDelay: reportingDelay,
		rollups:        map[string]*logRollup{},
		ctx:            ctx,
		ctxDone:        false,
	}
	go func() {
		select {
		case <-ctx.Done():
			l.rollupsMu.Lock()
			l.rollups = map[string]*logRollup{}
			l.ctxDone = true
			l.rollupsMu.Unlock()
		}
	}()
	return l
}

// ShouldRollup returns true when a log msg meets the criteria for rollup.
func (r *RollupRouter) ShouldRollup(logmsg map[string]interface{}) bool {
	statusCode, ok := logmsg["status-code"].(int)
	if !ok {
		return false
	}

	responseTime, ok := logmsg["response-time"].(time.Duration)
	if !ok {
		return false
	}

	// in the future allow more configurability, for now default to 200s and < 500ms
	if statusCode == 200 && responseTime < 500*time.Millisecond {
		return true
	}

	return false
}

// Process rolls up a log message.
func (r *RollupRouter) Process(logmsg map[string]interface{}) {
	if r.ctxDone {
		return
	}

	statusCode, ok := logmsg["status-code"].(int)
	if !ok {
		return
	}
	path, ok := logmsg["path"].(string)
	if !ok {
		return
	}
	canary, ok := logmsg["canary"].(bool)
	if !ok {
		return
	}
	r.findOrCreate(statusCode, path, canary).add(logmsg)
}

func (r *RollupRouter) findOrCreate(statusCode int, path string, canary bool) *logRollup {
	r.rollupsMu.Lock()
	defer r.rollupsMu.Unlock()
	rollupKey := fmt.Sprintf("%d-%s", statusCode, path)
	if canary {
		rollupKey += "-canary"
	}
	if rollup, ok := r.rollups[rollupKey]; ok {
		return rollup
	}
	rollup := &logRollup{
		Logger:           r.logger,
		ReportingDelayNs: (r.reportingDelay).Nanoseconds(),
		StatusCode:       statusCode,
		Path:             path,
		Canary:           canary,
	}
	r.rollups[rollupKey] = rollup
	go rollup.schedule(r.ctx)
	return rollup
}

// logRollup represents a single rollup.
type logRollup struct {
	Logger           RollupLogger
	ReportingDelayNs int64
	StatusCode       int
	Path             string
	Canary           bool

	rollupMu                sync.Mutex
	rollupMsg               map[string]interface{}
	rollupResponseTimeMsSum int64
}

func (r *logRollup) report() {
	r.rollupMu.Lock()
	defer r.rollupMu.Unlock()
	if r.rollupMsg != nil {
		r.rollupMsg["response-time-ms-sum"] = r.rollupResponseTimeMsSum
		r.rollupMsg["response-time-ms"] = r.rollupResponseTimeMsSum / r.rollupMsg["count"].(int64)
		switch logLevelFromStatus(r.StatusCode) {
		case logger.Error:
			r.Logger.ErrorD("request-finished-rollup", r.rollupMsg)
		default:
			r.Logger.InfoD("request-finished-rollup", r.rollupMsg)
		}
		r.rollupMsg = nil
		r.rollupResponseTimeMsSum = 0
	}
}

func (r *logRollup) schedule(ctx context.Context) {
	lastReport := time.Now()
	for {
		reportingDelay := time.Duration(atomic.LoadInt64(&r.ReportingDelayNs))
		wakeupTime := lastReport.Add(reportingDelay)
		now := time.Now()
		if now.After(wakeupTime) {
			wakeupTime = now.Add(reportingDelay)
		}
		sleepTime := wakeupTime.Sub(now)

		select {
		case <-ctx.Done():
			return
		case <-time.After(sleepTime):
			lastReport = time.Now()
			r.report()
		}
	}
}

func (r *logRollup) add(logmsg map[string]interface{}) {
	r.rollupMu.Lock()
	defer r.rollupMu.Unlock()

	if r.rollupMsg == nil {
		r.rollupMsg = map[string]interface{}{
			"status-code": r.StatusCode,
			"path":        r.Path,
			"count":       int64(0),
			"canary":      r.Canary,
			"via":         "kayvee-middleware",
		}
	}

	r.rollupMsg["count"] = r.rollupMsg["count"].(int64) + 1
	r.rollupResponseTimeMsSum += int64(logmsg["response-time"].(time.Duration) / time.Millisecond)
}
<commit_msg>hotfix response times<commit_after>package middleware

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/Clever/kayvee-go.v6/logger"
)

var globalRollupRouter *RollupRouter

// RollupLogger will log info / error rollups depending on status code.
type RollupLogger interface {
	InfoD(title string, data map[string]interface{})
	ErrorD(title string, data map[string]interface{})
}

// EnableRollups turns on rollups for kv middleware logs.
func EnableRollups(ctx context.Context, logger RollupLogger, reportingInterval time.Duration) {
	globalRollupRouter = NewRollupRouter(ctx, logger, reportingInterval)
}

// RollupRouter rolls up log lines and periodically logs them as one log line.
type RollupRouter struct {
	logger         RollupLogger
	reportingDelay time.Duration
	ctx            context.Context
	ctxDone        bool

	// create a rollup object per unique (status-code, path) pair
	rollupsMu sync.Mutex
	rollups   map[string]*logRollup
}

// NewRollupRouter creates a new log rollup output.
// Rollups will stop when the context is canceled.
func NewRollupRouter(ctx context.Context, logger RollupLogger, reportingDelay time.Duration) *RollupRouter {
	l := &RollupRouter{
		logger:         logger,
		reportingDelay: reportingDelay,
		rollups:        map[string]*logRollup{},
		ctx:            ctx,
		ctxDone:        false,
	}
	go func() {
		select {
		case <-ctx.Done():
			l.rollupsMu.Lock()
			l.rollups = map[string]*logRollup{}
			l.ctxDone = true
			l.rollupsMu.Unlock()
		}
	}()
	return l
}

// ShouldRollup returns true when a log msg meets the criteria for rollup.
func (r *RollupRouter) ShouldRollup(logmsg map[string]interface{}) bool {
	statusCode, ok := logmsg["status-code"].(int)
	if !ok {
		return false
	}

	responseTime, ok := logmsg["response-time"].(time.Duration)
	if !ok {
		return false
	}

	// in the future allow more configurability, for now default to 200s and < 500ms
	if statusCode == 200 && responseTime < 500*time.Millisecond {
		return true
	}

	return false
}

// Process rolls up a log message.
func (r *RollupRouter) Process(logmsg map[string]interface{}) {
	if r.ctxDone {
		return
	}

	statusCode, ok := logmsg["status-code"].(int)
	if !ok {
		return
	}
	path, ok := logmsg["path"].(string)
	if !ok {
		return
	}
	canary, ok := logmsg["canary"].(bool)
	if !ok {
		return
	}
	r.findOrCreate(statusCode, path, canary).add(logmsg)
}

func (r *RollupRouter) findOrCreate(statusCode int, path string, canary bool) *logRollup {
	r.rollupsMu.Lock()
	defer r.rollupsMu.Unlock()
	rollupKey := fmt.Sprintf("%d-%s", statusCode, path)
	if canary {
		rollupKey += "-canary"
	}
	if rollup, ok := r.rollups[rollupKey]; ok {
		return rollup
	}
	rollup := &logRollup{
		Logger:           r.logger,
		ReportingDelayNs: (r.reportingDelay).Nanoseconds(),
		StatusCode:       statusCode,
		Path:             path,
		Canary:           canary,
	}
	r.rollups[rollupKey] = rollup
	go rollup.schedule(r.ctx)
	return rollup
}

// logRollup represents a single rollup.
type logRollup struct {
	Logger           RollupLogger
	ReportingDelayNs int64
	StatusCode       int
	Path             string
	Canary           bool

	rollupMu                sync.Mutex
	rollupMsg               map[string]interface{}
	rollupResponseTimeNsSum int64
}

func (r *logRollup) report() {
	r.rollupMu.Lock()
	defer r.rollupMu.Unlock()
	if r.rollupMsg != nil {
		sum := r.rollupResponseTimeNsSum / int64(time.Millisecond)
		r.rollupMsg["response-time-ms-sum"] = sum
		r.rollupMsg["response-time-ms"] = sum / r.rollupMsg["count"].(int64)
		switch logLevelFromStatus(r.StatusCode) {
		case logger.Error:
			r.Logger.ErrorD("request-finished-rollup", r.rollupMsg)
		default:
			r.Logger.InfoD("request-finished-rollup", r.rollupMsg)
		}
		r.rollupMsg = nil
		r.rollupResponseTimeNsSum = 0
	}
}

func (r *logRollup) schedule(ctx context.Context) {
	lastReport := time.Now()
	for {
		reportingDelay := time.Duration(atomic.LoadInt64(&r.ReportingDelayNs))
		wakeupTime := lastReport.Add(reportingDelay)
		now := time.Now()
		if now.After(wakeupTime) {
			wakeupTime = now.Add(reportingDelay)
		}
		sleepTime := wakeupTime.Sub(now)

		select {
		case <-ctx.Done():
			return
		case <-time.After(sleepTime):
			lastReport = time.Now()
			r.report()
		}
	}
}

func (r *logRollup) add(logmsg map[string]interface{}) {
	r.rollupMu.Lock()
	defer r.rollupMu.Unlock()

	if r.rollupMsg == nil {
		r.rollupMsg = map[string]interface{}{
			"status-code": r.StatusCode,
			"path":        r.Path,
			"count":       int64(0),
			"canary":      r.Canary,
			"via":         "kayvee-middleware",
		}
	}

	r.rollupMsg["count"] = r.rollupMsg["count"].(int64) + 1
	r.rollupResponseTimeNsSum += logmsg["response-time"].(time.Duration).Nanoseconds()
}
<|endoftext|>


<commit_before>package instance

import (
	"database/sql"
	"fmt"

	_ "github.com/go-sql-driver/mysql"
	"github.com/scouter-argos/agent/monsql"
)

type Instance struct {
	Type      string `json:"db.type"`
	IP        string `json:"db.ip"`
	Port      string `json:"db.port"`
	User      string `json:"db.user"`
	Password  string `json:"db.password"`
	Slowquery string `json:"db.slowquery"`
}

func (inst *Instance) StartMonitor() {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s", inst.User, inst.Password, inst.IP, inst.Port, "mysql")

	db, err := sql.Open("mysql", dsn) //// Open doesn't open a connection. Validate DSN data
	if err != nil {
		panic(err.Error()) //todo : change error handler instead of panic.
	}
	defer db.Close()
	err = db.Ping()
	if err != nil {
		//Write log.
		return
	}

	go inst.run()
}

func (inst *Instance) run() {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s", inst.User, inst.Password, inst.IP, inst.Port, "mysql")
	db, err := sql.Open("mysql", dsn) //// Open doesn't open a connection. Validate DSN data
	if err != nil {
		return
	}
	for {
		resultRow, err := db.Query(monsql.GlobalStatusSQL)
		defer resultRow.Close()
		if err != nil {
			return
		}

		var statusName string
		var statusValue int64
		for resultRow.Next() {
			// get RawBytes from data
			err = resultRow.Scan(&statusName, &statusValue)
			if err != nil {
				panic(err.Error()) // proper error handling instead of panic in your app
			}

		}
	}
}
<commit_msg>add channel<commit_after>package instance

import (
	"database/sql"
	"fmt"

	_ "github.com/go-sql-driver/mysql"
	"github.com/scouter-argos/agent/monsql"
)

type Instance struct {
	Type      string `json:"db.type"`
	IP        string `json:"db.ip"`
	Port      string `json:"db.port"`
	User      string `json:"db.user"`
	Password  string `json:"db.password"`
	Slowquery string `json:"db.slowquery"`

}

func (inst *Instance) StartMonitor() {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s", inst.User, inst.Password, inst.IP, inst.Port, "mysql")

	db, err := sql.Open("mysql", dsn) //// Open doesn't open a connection. Validate DSN data
	if err != nil {
		panic(err.Error()) //todo : change error handler instead of panic.
	}
	defer db.Close()
	err = db.Ping()
	if err != nil {
		//Write log.
		return
	}

	go inst.run()
}

func (inst *Instance) run() {
	var wg sync.WaitGroup
	wg.Add(1)
	ch := make(chan int)
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s", inst.User, inst.Password, inst.IP, inst.Port, "mysql")
	db, err := sql.Open("mysql", dsn) //// Open doesn't open a connection. Validate DSN data
	if err != nil {
		return
	}
	for {
		resultRow, err := db.Query(monsql.GlobalStatusSQL)
		defer resultRow.Close()
		if err != nil {
			return
		}

		var statusName string
		var statusValue int64

		statusMap := make(map[string] int64)
		for resultRow.Next() {

			// get RawBytes from data
			err = resultRow.Scan(&statusName, &statusValue)
			if err != nil {
				panic(err.Error()) // proper error handling instead of panic in your app
			}
			map[statusName]= statusValue

		}
	}
}
<|endoftext|>