<commit_before>/*
Copyright 2020 The Knative 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 dispatcher

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/cloudevents/sdk-go/v2/binding"
	"github.com/cloudevents/sdk-go/v2/binding/test"
	"github.com/cloudevents/sdk-go/v2/binding/transformer"
	protocolhttp "github.com/cloudevents/sdk-go/v2/protocol/http"
	"go.uber.org/zap"
	eventingduck "knative.dev/eventing/pkg/apis/duck/v1beta1"
	"knative.dev/eventing/pkg/channel"
	"knative.dev/eventing/pkg/channel/fanout"
	"knative.dev/eventing/pkg/channel/multichannelfanout"
	"knative.dev/eventing/pkg/kncloudevents"
	"knative.dev/eventing/pkg/tracing"
	"knative.dev/pkg/apis"
	duckv1 "knative.dev/pkg/apis/duck/v1"

	"knative.dev/eventing-contrib/kafka/channel/pkg/utils"
)

// This dispatcher tests the full integration of the dispatcher code with Kafka.
// This test doesn't run on the CI because unit tests script doesn't start a Kafka cluster.
// Use it in emergency situations when you can't reproduce the e2e test failures and the failure might be
// in the dispatcher code.
// Start a kafka cluster with docker: docker run --rm --net=host -e ADV_HOST=localhost -e SAMPLEDATA=0 lensesio/fast-data-dev
// Keep also the port 8080 free for the MessageReceiver
func TestDispatcher(t *testing.T) {
	if os.Getenv("CI") == "true" {
		t.Skipf("This test can't run in CI")
	}

	logger, err := zap.NewDevelopment(zap.AddStacktrace(zap.WarnLevel))
	if err != nil {
		t.Fatal(err)
	}

	tracing.SetupStaticPublishing(logger.Sugar(), "localhost", tracing.AlwaysSample)

	dispatcherArgs := KafkaDispatcherArgs{
		KnCEConnectionArgs: nil,
		ClientID:           "testing",
		Brokers:            []string{"localhost:9092"},
		TopicFunc:          utils.TopicName,
		Logger:             logger,
	}

	// Create the dispatcher. At this point, if Kafka is not up, this thing fails
	dispatcher, err := NewDispatcher(context.Background(), &dispatcherArgs)
	if err != nil {
		t.Skipf("no dispatcher: %v", err)
	}

	// Start the dispatcher
	go func() {
		if err := dispatcher.Start(context.Background()); err != nil {
			t.Fatal(err)
		}
	}()

	time.Sleep(1 * time.Second)

	// We need a channelaproxy and channelbproxy for handling correctly the Host header
	channelAProxy := httptest.NewServer(createReverseProxy(t, "channela.svc"))
	defer channelAProxy.Close()
	channelBProxy := httptest.NewServer(createReverseProxy(t, "channelb.svc"))
	defer channelBProxy.Close()

	// Start a bunch of test servers to simulate the various services
	transformationsWg := sync.WaitGroup{}
	transformationsWg.Add(1)
	transformationsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer transformationsWg.Done()
		message := protocolhttp.NewMessageFromHttpRequest(r)
		defer message.Finish(nil)

		err := protocolhttp.WriteResponseWriter(context.Background(), message, 200, w, transformer.AddExtension("transformed", "true"))
		if err != nil {
			w.WriteHeader(500)
			t.Fatal(err)
		}
	}))
	defer transformationsServer.Close()

	receiverWg := sync.WaitGroup{}
	receiverWg.Add(1)
	receiverServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer receiverWg.Done()
		transformed := r.Header.Get("ce-transformed")
		if transformed != "true" {
			w.WriteHeader(500)
			t.Fatalf("Expecting ce-transformed: true, found %s", transformed)
		}
	}))
	defer receiverServer.Close()

	transformationsFailureWg := sync.WaitGroup{}
	transformationsFailureWg.Add(1)
	transformationsFailureServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer transformationsFailureWg.Done()
		w.WriteHeader(500)
	}))
	defer transformationsFailureServer.Close()

	deadLetterWg := sync.WaitGroup{}
	deadLetterWg.Add(1)
	deadLetterServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer deadLetterWg.Done()
		transformed := r.Header.Get("ce-transformed")
		if transformed != "" {
			w.WriteHeader(500)
			t.Fatalf("Not expecting ce-transformed, found %s", transformed)
		}
	}))
	defer deadLetterServer.Close()

	logger.Debug("Test servers",
		zap.String("transformations server", transformationsServer.URL),
		zap.String("transformations failure server", transformationsFailureServer.URL),
		zap.String("receiver server", receiverServer.URL),
		zap.String("dead letter server", deadLetterServer.URL),
	)

	// send -> channela -> sub with transformationServer and reply to channelb -> channelb -> sub with receiver -> receiver
	config := multichannelfanout.Config{
		ChannelConfigs: []multichannelfanout.ChannelConfig{
			{
				Namespace: "default",
				Name:      "channela",
				HostName:  "channela.svc",
				FanoutConfig: fanout.Config{
					AsyncHandler:     false,
					DispatcherConfig: channel.EventDispatcherConfig{},
					Subscriptions: []eventingduck.SubscriberSpec{{
						UID:           "aaaa",
						Generation:    1,
						SubscriberURI: mustParseUrl(t, transformationsServer.URL),
						ReplyURI:      mustParseUrl(t, channelBProxy.URL),
					}, {
						UID:           "cccc",
						Generation:    1,
						SubscriberURI: mustParseUrl(t, transformationsFailureServer.URL),
						ReplyURI:      mustParseUrl(t, channelBProxy.URL),
						Delivery: &eventingduck.DeliverySpec{
							DeadLetterSink: &duckv1.Destination{URI: mustParseUrl(t, deadLetterServer.URL)},
						},
					}},
				},
			},
			{
				Namespace: "default",
				Name:      "channelb",
				HostName:  "channelb.svc",
				FanoutConfig: fanout.Config{
					AsyncHandler:     false,
					DispatcherConfig: channel.EventDispatcherConfig{},
					Subscriptions: []eventingduck.SubscriberSpec{{
						UID:           "bbbb",
						Generation:    1,
						SubscriberURI: mustParseUrl(t, receiverServer.URL),
					}},
				},
			},
		},
	}

	err = dispatcher.UpdateHostToChannelMap(&config)
	if err != nil {
		t.Fatal(err)
	}

	failed, err := dispatcher.UpdateKafkaConsumers(&config)
	if err != nil {
		t.Fatal(err)
	}
	if len(failed) != 0 {
		t.Fatal(err)
	}

	time.Sleep(5 * time.Second)

	// Ok now everything should be ready to send the event
	httpsender, err := kncloudevents.NewHttpMessageSender(nil, channelAProxy.URL)
	if err != nil {
		t.Fatal(err)
	}

	req, err := httpsender.NewCloudEventRequest(context.Background())
	if err != nil {
		t.Fatal(err)
	}

	event := test.FullEvent()
	_ = protocolhttp.WriteRequest(context.Background(), binding.ToMessage(&event), req)

	res, err := httpsender.Send(req)
	if err != nil {
		t.Fatal(err)
	}

	if res.StatusCode != 202 {
		t.Fatalf("Expected 202, Have %d", res.StatusCode)
	}

	transformationsFailureWg.Wait()
	deadLetterWg.Wait()
	transformationsWg.Wait()
	receiverWg.Wait()
}

func createReverseProxy(t *testing.T, host string) *httputil.ReverseProxy {
	director := func(req *http.Request) {
		target := mustParseUrl(t, "http://localhost:8080")
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = target.Path
		req.Host = host
	}
	return &httputil.ReverseProxy{Director: director}
}

func mustParseUrl(t *testing.T, str string) *apis.URL {
	url, err := apis.ParseURL(str)
	if err != nil {
		t.Fatal(err)
	}
	return url
}
<commit_msg>changes from eventing master (#1150)<commit_after>/*
Copyright 2020 The Knative 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 dispatcher

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/cloudevents/sdk-go/v2/binding"
	"github.com/cloudevents/sdk-go/v2/binding/test"
	"github.com/cloudevents/sdk-go/v2/binding/transformer"
	protocolhttp "github.com/cloudevents/sdk-go/v2/protocol/http"
	"go.uber.org/zap"
	eventingduck "knative.dev/eventing/pkg/apis/duck/v1beta1"
	"knative.dev/eventing/pkg/channel/fanout"
	"knative.dev/eventing/pkg/channel/multichannelfanout"
	"knative.dev/eventing/pkg/kncloudevents"
	"knative.dev/eventing/pkg/tracing"
	"knative.dev/pkg/apis"
	duckv1 "knative.dev/pkg/apis/duck/v1"

	"knative.dev/eventing-contrib/kafka/channel/pkg/utils"
)

// This dispatcher tests the full integration of the dispatcher code with Kafka.
// This test doesn't run on the CI because unit tests script doesn't start a Kafka cluster.
// Use it in emergency situations when you can't reproduce the e2e test failures and the failure might be
// in the dispatcher code.
// Start a kafka cluster with docker: docker run --rm --net=host -e ADV_HOST=localhost -e SAMPLEDATA=0 lensesio/fast-data-dev
// Keep also the port 8080 free for the MessageReceiver
func TestDispatcher(t *testing.T) {
	if os.Getenv("CI") == "true" {
		t.Skipf("This test can't run in CI")
	}

	logger, err := zap.NewDevelopment(zap.AddStacktrace(zap.WarnLevel))
	if err != nil {
		t.Fatal(err)
	}

	tracing.SetupStaticPublishing(logger.Sugar(), "localhost", tracing.AlwaysSample)

	dispatcherArgs := KafkaDispatcherArgs{
		KnCEConnectionArgs: nil,
		ClientID:           "testing",
		Brokers:            []string{"localhost:9092"},
		TopicFunc:          utils.TopicName,
		Logger:             logger,
	}

	// Create the dispatcher. At this point, if Kafka is not up, this thing fails
	dispatcher, err := NewDispatcher(context.Background(), &dispatcherArgs)
	if err != nil {
		t.Skipf("no dispatcher: %v", err)
	}

	// Start the dispatcher
	go func() {
		if err := dispatcher.Start(context.Background()); err != nil {
			t.Fatal(err)
		}
	}()

	time.Sleep(1 * time.Second)

	// We need a channelaproxy and channelbproxy for handling correctly the Host header
	channelAProxy := httptest.NewServer(createReverseProxy(t, "channela.svc"))
	defer channelAProxy.Close()
	channelBProxy := httptest.NewServer(createReverseProxy(t, "channelb.svc"))
	defer channelBProxy.Close()

	// Start a bunch of test servers to simulate the various services
	transformationsWg := sync.WaitGroup{}
	transformationsWg.Add(1)
	transformationsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer transformationsWg.Done()
		message := protocolhttp.NewMessageFromHttpRequest(r)
		defer message.Finish(nil)

		err := protocolhttp.WriteResponseWriter(context.Background(), message, 200, w, transformer.AddExtension("transformed", "true"))
		if err != nil {
			w.WriteHeader(500)
			t.Fatal(err)
		}
	}))
	defer transformationsServer.Close()

	receiverWg := sync.WaitGroup{}
	receiverWg.Add(1)
	receiverServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer receiverWg.Done()
		transformed := r.Header.Get("ce-transformed")
		if transformed != "true" {
			w.WriteHeader(500)
			t.Fatalf("Expecting ce-transformed: true, found %s", transformed)
		}
	}))
	defer receiverServer.Close()

	transformationsFailureWg := sync.WaitGroup{}
	transformationsFailureWg.Add(1)
	transformationsFailureServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer transformationsFailureWg.Done()
		w.WriteHeader(500)
	}))
	defer transformationsFailureServer.Close()

	deadLetterWg := sync.WaitGroup{}
	deadLetterWg.Add(1)
	deadLetterServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer deadLetterWg.Done()
		transformed := r.Header.Get("ce-transformed")
		if transformed != "" {
			w.WriteHeader(500)
			t.Fatalf("Not expecting ce-transformed, found %s", transformed)
		}
	}))
	defer deadLetterServer.Close()

	logger.Debug("Test servers",
		zap.String("transformations server", transformationsServer.URL),
		zap.String("transformations failure server", transformationsFailureServer.URL),
		zap.String("receiver server", receiverServer.URL),
		zap.String("dead letter server", deadLetterServer.URL),
	)

	// send -> channela -> sub with transformationServer and reply to channelb -> channelb -> sub with receiver -> receiver
	config := multichannelfanout.Config{
		ChannelConfigs: []multichannelfanout.ChannelConfig{
			{
				Namespace: "default",
				Name:      "channela",
				HostName:  "channela.svc",
				FanoutConfig: fanout.Config{
					AsyncHandler: false,
					Subscriptions: []eventingduck.SubscriberSpec{{
						UID:           "aaaa",
						Generation:    1,
						SubscriberURI: mustParseUrl(t, transformationsServer.URL),
						ReplyURI:      mustParseUrl(t, channelBProxy.URL),
					}, {
						UID:           "cccc",
						Generation:    1,
						SubscriberURI: mustParseUrl(t, transformationsFailureServer.URL),
						ReplyURI:      mustParseUrl(t, channelBProxy.URL),
						Delivery: &eventingduck.DeliverySpec{
							DeadLetterSink: &duckv1.Destination{URI: mustParseUrl(t, deadLetterServer.URL)},
						},
					}},
				},
			},
			{
				Namespace: "default",
				Name:      "channelb",
				HostName:  "channelb.svc",
				FanoutConfig: fanout.Config{
					AsyncHandler: false,
					Subscriptions: []eventingduck.SubscriberSpec{{
						UID:           "bbbb",
						Generation:    1,
						SubscriberURI: mustParseUrl(t, receiverServer.URL),
					}},
				},
			},
		},
	}

	err = dispatcher.UpdateHostToChannelMap(&config)
	if err != nil {
		t.Fatal(err)
	}

	failed, err := dispatcher.UpdateKafkaConsumers(&config)
	if err != nil {
		t.Fatal(err)
	}
	if len(failed) != 0 {
		t.Fatal(err)
	}

	time.Sleep(5 * time.Second)

	// Ok now everything should be ready to send the event
	httpsender, err := kncloudevents.NewHttpMessageSender(nil, channelAProxy.URL)
	if err != nil {
		t.Fatal(err)
	}

	req, err := httpsender.NewCloudEventRequest(context.Background())
	if err != nil {
		t.Fatal(err)
	}

	event := test.FullEvent()
	_ = protocolhttp.WriteRequest(context.Background(), binding.ToMessage(&event), req)

	res, err := httpsender.Send(req)
	if err != nil {
		t.Fatal(err)
	}

	if res.StatusCode != 202 {
		t.Fatalf("Expected 202, Have %d", res.StatusCode)
	}

	transformationsFailureWg.Wait()
	deadLetterWg.Wait()
	transformationsWg.Wait()
	receiverWg.Wait()
}

func createReverseProxy(t *testing.T, host string) *httputil.ReverseProxy {
	director := func(req *http.Request) {
		target := mustParseUrl(t, "http://localhost:8080")
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = target.Path
		req.Host = host
	}
	return &httputil.ReverseProxy{Director: director}
}

func mustParseUrl(t *testing.T, str string) *apis.URL {
	url, err := apis.ParseURL(str)
	if err != nil {
		t.Fatal(err)
	}
	return url
}
<|endoftext|>


<commit_before>package reply

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"

	v3pb "github.com/go-graphite/protocol/carbonapi_v3_pb"
	"github.com/lomik/graphite-clickhouse/helper/point"
	"github.com/lomik/graphite-clickhouse/pkg/scope"
	"github.com/lomik/graphite-clickhouse/render/data"
	"go.uber.org/zap"
)

const (
	Repeated = 2
	Float32  = 5
)

type V3pb struct{}

func (*V3pb) ParseRequest(r *http.Request) (data.MultiTarget, error) {
	logger := scope.Logger(r.Context()).Named("render")

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		logger.Error("failed to read request", zap.Error(err))
		return nil, fmt.Errorf("failed to read request body: %w", err)
	}

	var pv3Request v3pb.MultiFetchRequest
	if err := pv3Request.Unmarshal(body); err != nil {
		logger.Error("failed to unmarshal request", zap.Error(err))
		return nil, fmt.Errorf("failed to unmarshal request: %w", err)
	}

	multiTarget := data.MFRToMultiTarget(&pv3Request)
	if len(pv3Request.Metrics) > 0 {
		for _, m := range pv3Request.Metrics {
			logger.Info(
				"pb3_target",
				zap.Int64("from", m.StartTime),
				zap.Int64("until", m.StopTime),
				zap.Int64("maxDataPoints", m.MaxDataPoints),
				zap.String("target", m.PathExpression),
			)
		}
	}

	if scope.Debug(r.Context(), "Output") {
		request, err := json.Marshal(pv3Request)
		if err == nil {
			logger.Info("v3pb_request", zap.ByteString("json", request))
		}
	}

	return multiTarget, nil
}

func (*V3pb) Reply(w http.ResponseWriter, r *http.Request, multiData data.CHResponses) {
	replyProtobuf(w, r, multiData, true)
}

func writePB3(mb, mb2 *bytes.Buffer, writer *bufio.Writer, target, name, function string, from, until, step uint32, points []point.Point) {
	start, stop, count, getValue := point.FillNulls(points, from, until, step)

	mb.Reset()

	// First chunk
	// name
	VarintWrite(mb, (1<<3)+Repeated) // tag
	VarintWrite(mb, uint64(len(name)))
	mb.WriteString(name)

	// pathExpression
	VarintWrite(mb, (2<<3)+Repeated) // tag
	VarintWrite(mb, uint64(len(target)))
	mb.WriteString(target)

	consolidationFunc := function
	// consolidationFunc
	VarintWrite(mb, (3<<3)+Repeated) // tag
	VarintWrite(mb, uint64(len(consolidationFunc)))
	mb.WriteString(consolidationFunc)

	// start
	VarintWrite(mb, 4<<3) // tag
	VarintWrite(mb, uint64(start))

	// stop
	VarintWrite(mb, 5<<3) // tag
	VarintWrite(mb, uint64(stop))

	// step
	VarintWrite(mb, 6<<3) // tag
	VarintWrite(mb, uint64(step))

	// xFilesFactor
	VarintWrite(mb, (7<<3)+Float32) // tag
	ProtobufWriteSingle(mb, 0.0)

	// highPrecisionTimestamps
	VarintWrite(mb, 8<<3) // tag
	mb.WriteByte('\x00')  // False

	// Values header
	VarintWrite(mb, (9<<3)+Repeated) // tag
	VarintWrite(mb, uint64(8*count))
	for {
		value, err := getValue()
		if err != nil {
			if errors.Is(err, point.ErrTimeGreaterStop) {
				break
			}
			// if err is not point.ErrTimeGreaterStop, the points are corrupted
			return
		}
		ProtobufWriteDouble(mb, value)
	}

	// rest fields, that goes after values

	// Fields with default values are skipped, so this should be uncommented if support for appliedFunctions will be
	// implemented
	// appliedFunctions
	//VarintWrite(mb2, (10<<3)+Repeated)  // tag
	//VarintWrite(mb2, VarintLen(0)) // currently not supported

	// requestStartTime
	VarintWrite(mb, 11<<3)
	VarintWrite(mb, uint64(from))

	// requestStopTime
	VarintWrite(mb, 12<<3)
	VarintWrite(mb, uint64(until))

	// start write to output
	// repeated FetchResponse metrics = 1;
	// write tag and len
	VarintWrite(writer, (1<<3)+2)
	VarintWrite(writer, uint64(mb.Len()))

	writer.Write(mb.Bytes())
}
<commit_msg>Fix logger name<commit_after>package reply

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"

	v3pb "github.com/go-graphite/protocol/carbonapi_v3_pb"
	"github.com/lomik/graphite-clickhouse/helper/point"
	"github.com/lomik/graphite-clickhouse/pkg/scope"
	"github.com/lomik/graphite-clickhouse/render/data"
	"go.uber.org/zap"
)

const (
	Repeated = 2
	Float32  = 5
)

type V3pb struct{}

func (*V3pb) ParseRequest(r *http.Request) (data.MultiTarget, error) {
	logger := scope.Logger(r.Context()).Named("pb3parser")

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		logger.Error("failed to read request", zap.Error(err))
		return nil, fmt.Errorf("failed to read request body: %w", err)
	}

	var pv3Request v3pb.MultiFetchRequest
	if err := pv3Request.Unmarshal(body); err != nil {
		logger.Error("failed to unmarshal request", zap.Error(err))
		return nil, fmt.Errorf("failed to unmarshal request: %w", err)
	}

	multiTarget := data.MFRToMultiTarget(&pv3Request)
	if len(pv3Request.Metrics) > 0 {
		for _, m := range pv3Request.Metrics {
			logger.Info(
				"pb3_target",
				zap.Int64("from", m.StartTime),
				zap.Int64("until", m.StopTime),
				zap.Int64("maxDataPoints", m.MaxDataPoints),
				zap.String("target", m.PathExpression),
			)
		}
	}

	if scope.Debug(r.Context(), "Output") {
		request, err := json.Marshal(pv3Request)
		if err == nil {
			logger.Info("v3pb_request", zap.ByteString("json", request))
		}
	}

	return multiTarget, nil
}

func (*V3pb) Reply(w http.ResponseWriter, r *http.Request, multiData data.CHResponses) {
	replyProtobuf(w, r, multiData, true)
}

func writePB3(mb, mb2 *bytes.Buffer, writer *bufio.Writer, target, name, function string, from, until, step uint32, points []point.Point) {
	start, stop, count, getValue := point.FillNulls(points, from, until, step)

	mb.Reset()

	// First chunk
	// name
	VarintWrite(mb, (1<<3)+Repeated) // tag
	VarintWrite(mb, uint64(len(name)))
	mb.WriteString(name)

	// pathExpression
	VarintWrite(mb, (2<<3)+Repeated) // tag
	VarintWrite(mb, uint64(len(target)))
	mb.WriteString(target)

	consolidationFunc := function
	// consolidationFunc
	VarintWrite(mb, (3<<3)+Repeated) // tag
	VarintWrite(mb, uint64(len(consolidationFunc)))
	mb.WriteString(consolidationFunc)

	// start
	VarintWrite(mb, 4<<3) // tag
	VarintWrite(mb, uint64(start))

	// stop
	VarintWrite(mb, 5<<3) // tag
	VarintWrite(mb, uint64(stop))

	// step
	VarintWrite(mb, 6<<3) // tag
	VarintWrite(mb, uint64(step))

	// xFilesFactor
	VarintWrite(mb, (7<<3)+Float32) // tag
	ProtobufWriteSingle(mb, 0.0)

	// highPrecisionTimestamps
	VarintWrite(mb, 8<<3) // tag
	mb.WriteByte('\x00')  // False

	// Values header
	VarintWrite(mb, (9<<3)+Repeated) // tag
	VarintWrite(mb, uint64(8*count))
	for {
		value, err := getValue()
		if err != nil {
			if errors.Is(err, point.ErrTimeGreaterStop) {
				break
			}
			// if err is not point.ErrTimeGreaterStop, the points are corrupted
			return
		}
		ProtobufWriteDouble(mb, value)
	}

	// rest fields, that goes after values

	// Fields with default values are skipped, so this should be uncommented if support for appliedFunctions will be
	// implemented
	// appliedFunctions
	//VarintWrite(mb2, (10<<3)+Repeated)  // tag
	//VarintWrite(mb2, VarintLen(0)) // currently not supported

	// requestStartTime
	VarintWrite(mb, 11<<3)
	VarintWrite(mb, uint64(from))

	// requestStopTime
	VarintWrite(mb, 12<<3)
	VarintWrite(mb, uint64(until))

	// start write to output
	// repeated FetchResponse metrics = 1;
	// write tag and len
	VarintWrite(writer, (1<<3)+2)
	VarintWrite(writer, uint64(mb.Len()))

	writer.Write(mb.Bytes())
}
<|endoftext|>


<commit_before>// Package s3 implements Storage based on an S3 bucket.
package s3

import (
	"bytes"
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/minio/minio-go/v7"
	"github.com/minio/minio-go/v7/pkg/credentials"
	"github.com/pkg/errors"

	"github.com/kopia/kopia/internal/clock"
	"github.com/kopia/kopia/internal/gather"
	"github.com/kopia/kopia/internal/iocopy"
	"github.com/kopia/kopia/repo/blob"
	"github.com/kopia/kopia/repo/blob/retrying"
)

const (
	s3storageType   = "s3"
	latestVersionID = ""
)

type s3Storage struct {
	Options

	cli *minio.Client

	storageConfig *StorageConfig
}

func (s *s3Storage) GetCapacity(ctx context.Context) (blob.Capacity, error) {
	return blob.Capacity{}, blob.ErrNotAVolume
}

func (s *s3Storage) GetBlob(ctx context.Context, b blob.ID, offset, length int64, output blob.OutputBuffer) error {
	return s.getBlobWithVersion(ctx, b, latestVersionID, offset, length, output)
}

// getBlobWithVersion returns full or partial contents of a blob with given ID and version.
func (s *s3Storage) getBlobWithVersion(ctx context.Context, b blob.ID, version string, offset, length int64, output blob.OutputBuffer) error {
	output.Reset()

	attempt := func() error {
		opt := minio.GetObjectOptions{VersionID: version}

		if length > 0 {
			if err := opt.SetRange(offset, offset+length-1); err != nil {
				return errors.Wrap(blob.ErrInvalidRange, "unable to set range")
			}
		}

		if length == 0 {
			// zero-length ranges require special handling, set non-zero range and
			// we won't be trying to read the response anyway.
			if err := opt.SetRange(0, 1); err != nil {
				return errors.Wrap(blob.ErrInvalidRange, "unable to set range")
			}
		}

		o, err := s.cli.GetObject(ctx, s.BucketName, s.getObjectNameString(b), opt)
		if err != nil {
			return errors.Wrap(err, "GetObject")
		}

		defer o.Close() //nolint:errcheck

		if length == 0 {
			return nil
		}

		// nolint:wrapcheck
		return iocopy.JustCopy(output, o)
	}

	if err := attempt(); err != nil {
		return translateError(err)
	}

	// nolint:wrapcheck
	return blob.EnsureLengthExactly(output.Length(), length)
}

func isInvalidCredentials(err error) bool {
	return err != nil && strings.Contains(err.Error(), blob.InvalidCredentialsErrStr)
}

func translateError(err error) error {
	var me minio.ErrorResponse

	if isInvalidCredentials(err) {
		return blob.ErrInvalidCredentials
	}

	if errors.As(err, &me) {
		switch me.StatusCode {
		case http.StatusOK:
			return nil

		case http.StatusNotFound:
			return blob.ErrBlobNotFound

		case http.StatusRequestedRangeNotSatisfiable:
			return blob.ErrInvalidRange
		}
	}

	return err
}

func (s *s3Storage) GetMetadata(ctx context.Context, b blob.ID) (blob.Metadata, error) {
	vm, err := s.getVersionMetadata(ctx, b, "")

	return vm.Metadata, err
}

func (s *s3Storage) getVersionMetadata(ctx context.Context, b blob.ID, version string) (versionMetadata, error) {
	opts := minio.GetObjectOptions{
		VersionID: version,
	}

	oi, err := s.cli.StatObject(ctx, s.BucketName, s.getObjectNameString(b), opts)
	if err != nil {
		return versionMetadata{}, errors.Wrap(translateError(err), "StatObject")
	}

	return infoToVersionMetadata(s.Prefix, &oi), nil
}

func (s *s3Storage) PutBlob(ctx context.Context, b blob.ID, data blob.Bytes, opts blob.PutOptions) error {
	switch {
	case opts.DoNotRecreate:
		return errors.Wrap(blob.ErrUnsupportedPutBlobOption, "do-not-recreate")
	case !opts.SetModTime.IsZero():
		return blob.ErrSetTimeUnsupported
	}

	_, err := s.putBlob(ctx, b, data, opts)

	if opts.GetModTime != nil {
		bm, err2 := s.GetMetadata(ctx, b)
		if err2 != nil {
			return err2
		}

		*opts.GetModTime = bm.Timestamp
	}

	return err
}

func (s *s3Storage) putBlob(ctx context.Context, b blob.ID, data blob.Bytes, opts blob.PutOptions) (versionMetadata, error) {
	var (
		storageClass    = s.storageConfig.getStorageClassForBlobID(b)
		retentionMode   minio.RetentionMode
		retainUntilDate time.Time
	)

	if opts.RetentionPeriod != 0 {
		retentionMode = minio.RetentionMode(opts.RetentionMode)
		if !retentionMode.IsValid() {
			return versionMetadata{}, errors.Errorf("invalid retention mode: %q", opts.RetentionMode)
		}

		retainUntilDate = clock.Now().Add(opts.RetentionPeriod).UTC()
	}

	uploadInfo, err := s.cli.PutObject(ctx, s.BucketName, s.getObjectNameString(b), data.Reader(), int64(data.Length()), minio.PutObjectOptions{
		ContentType: "application/x-kopia",
		// Kopia already splits snapshot contents into small blobs to improve
		// upload throughput. There is no need for further splitting
		// through multipart uploads.
		DisableMultipart: true,
		// The Content-MD5 header is required for any request to upload an object
		// with a retention period configured using Amazon S3 Object Lock.
		// Unconditionally computing the content MD5, potentially incurring
		// a slightly higher CPU overhead.
		SendContentMd5:  true,
		StorageClass:    storageClass,
		RetainUntilDate: retainUntilDate,
		Mode:            retentionMode,
	})

	if isInvalidCredentials(err) {
		return versionMetadata{}, blob.ErrInvalidCredentials
	}

	var er minio.ErrorResponse

	if errors.As(err, &er) && er.Code == "InvalidRequest" && strings.Contains(strings.ToLower(er.Message), "content-md5") {
		return versionMetadata{}, err // nolint:wrapcheck
	}

	if errors.Is(err, io.EOF) && uploadInfo.Size == 0 {
		// special case empty stream
		_, err = s.cli.PutObject(ctx, s.BucketName, s.getObjectNameString(b), bytes.NewBuffer(nil), 0, minio.PutObjectOptions{
			ContentType:     "application/x-kopia",
			StorageClass:    storageClass,
			RetainUntilDate: retainUntilDate,
			Mode:            retentionMode,
		})
	}

	if err != nil {
		return versionMetadata{}, err // nolint:wrapcheck
	}

	return versionMetadata{
		Metadata: blob.Metadata{
			BlobID:    b,
			Length:    uploadInfo.Size,
			Timestamp: uploadInfo.LastModified,
		},
		Version: uploadInfo.VersionID,
	}, nil
}

func (s *s3Storage) DeleteBlob(ctx context.Context, b blob.ID) error {
	err := translateError(s.cli.RemoveObject(ctx, s.BucketName, s.getObjectNameString(b), minio.RemoveObjectOptions{}))
	if errors.Is(err, blob.ErrBlobNotFound) {
		return nil
	}

	return err
}

func (s *s3Storage) getObjectNameString(b blob.ID) string {
	return s.Prefix + string(b)
}

func (s *s3Storage) ListBlobs(ctx context.Context, prefix blob.ID, callback func(blob.Metadata) error) error {
	ctx, cancel := context.WithCancel(ctx)

	defer cancel()

	oi := s.cli.ListObjects(ctx, s.BucketName, minio.ListObjectsOptions{
		Prefix: s.getObjectNameString(prefix),
	})
	for o := range oi {
		if err := o.Err; err != nil {
			if isInvalidCredentials(err) {
				return blob.ErrInvalidCredentials
			}

			return err
		}

		bm := blob.Metadata{
			BlobID:    blob.ID(o.Key[len(s.Prefix):]),
			Length:    o.Size,
			Timestamp: o.LastModified,
		}

		if bm.BlobID == ConfigName {
			continue
		}

		if err := callback(bm); err != nil {
			return err
		}
	}

	return nil
}

func (s *s3Storage) ConnectionInfo() blob.ConnectionInfo {
	return blob.ConnectionInfo{
		Type:   s3storageType,
		Config: &s.Options,
	}
}

func (s *s3Storage) Close(ctx context.Context) error {
	return nil
}

func (s *s3Storage) String() string {
	return fmt.Sprintf("s3://%v/%v", s.BucketName, s.Prefix)
}

func (s *s3Storage) DisplayName() string {
	return fmt.Sprintf("S3: %v %v", s.Endpoint, s.BucketName)
}

func (s *s3Storage) FlushCaches(ctx context.Context) error {
	return nil
}

func getCustomTransport(insecureSkipVerify bool) (transport *http.Transport) {
	// nolint:gosec
	customTransport := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: insecureSkipVerify}}
	return customTransport
}

// New creates new S3-backed storage with specified options:
//
// - the 'BucketName' field is required and all other parameters are optional.
func New(ctx context.Context, opt *Options) (blob.Storage, error) {
	st, err := newStorage(ctx, opt)
	if err != nil {
		return nil, err
	}

	s, err := maybePointInTimeStore(ctx, st, opt.PointInTime)
	if err != nil {
		return nil, err
	}

	return retrying.NewWrapper(s), nil
}

func newStorage(ctx context.Context, opt *Options) (*s3Storage, error) {
	return newStorageWithCredentials(ctx, credentials.NewStaticV4(opt.AccessKeyID, opt.SecretAccessKey, opt.SessionToken), opt)
}

func newStorageWithCredentials(ctx context.Context, creds *credentials.Credentials, opt *Options) (*s3Storage, error) {
	if opt.BucketName == "" {
		return nil, errors.New("bucket name must be specified")
	}

	minioOpts := &minio.Options{
		Creds:  creds,
		Secure: !opt.DoNotUseTLS,
		Region: opt.Region,
	}

	if opt.DoNotVerifyTLS {
		minioOpts.Transport = getCustomTransport(true)
	}

	cli, err := minio.New(opt.Endpoint, minioOpts)
	if err != nil {
		return nil, errors.Wrap(err, "unable to create client")
	}

	ok, err := cli.BucketExists(ctx, opt.BucketName)
	if err != nil {
		return nil, errors.Wrapf(err, "unable to determine if bucket %q exists", opt.BucketName)
	}

	if !ok {
		return nil, errors.Errorf("bucket %q does not exist", opt.BucketName)
	}

	s := s3Storage{
		Options:       *opt,
		cli:           cli,
		storageConfig: &StorageConfig{},
	}

	var scOutput gather.WriteBuffer

	if getBlobErr := s.GetBlob(ctx, ConfigName, 0, -1, &scOutput); getBlobErr == nil {
		if scErr := s.storageConfig.Load(scOutput.Bytes().Reader()); scErr != nil {
			return nil, errors.Wrapf(scErr, "error parsing storage config for bucket %q", opt.BucketName)
		}
	} else if !errors.Is(getBlobErr, blob.ErrBlobNotFound) {
		return nil, errors.Wrapf(getBlobErr, "error retrieving storage config from bucket %q", opt.BucketName)
	}

	return &s, nil
}

func init() {
	blob.AddSupportedStorage(
		s3storageType,
		func() interface{} {
			return &Options{}
		},
		func(ctx context.Context, o interface{}, isCreate bool) (blob.Storage, error) {
			return New(ctx, o.(*Options)) // nolint:forcetypeassert
		})
}
<commit_msg>feat(providers): add AWS IAM as a credential provider for s3 storage (#2213) (#2222)<commit_after>// Package s3 implements Storage based on an S3 bucket.
package s3

import (
	"bytes"
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/minio/minio-go/v7"
	"github.com/minio/minio-go/v7/pkg/credentials"
	"github.com/pkg/errors"

	"github.com/kopia/kopia/internal/clock"
	"github.com/kopia/kopia/internal/gather"
	"github.com/kopia/kopia/internal/iocopy"
	"github.com/kopia/kopia/repo/blob"
	"github.com/kopia/kopia/repo/blob/retrying"
)

const (
	s3storageType   = "s3"
	latestVersionID = ""
)

type s3Storage struct {
	Options

	cli *minio.Client

	storageConfig *StorageConfig
}

func (s *s3Storage) GetCapacity(ctx context.Context) (blob.Capacity, error) {
	return blob.Capacity{}, blob.ErrNotAVolume
}

func (s *s3Storage) GetBlob(ctx context.Context, b blob.ID, offset, length int64, output blob.OutputBuffer) error {
	return s.getBlobWithVersion(ctx, b, latestVersionID, offset, length, output)
}

// getBlobWithVersion returns full or partial contents of a blob with given ID and version.
func (s *s3Storage) getBlobWithVersion(ctx context.Context, b blob.ID, version string, offset, length int64, output blob.OutputBuffer) error {
	output.Reset()

	attempt := func() error {
		opt := minio.GetObjectOptions{VersionID: version}

		if length > 0 {
			if err := opt.SetRange(offset, offset+length-1); err != nil {
				return errors.Wrap(blob.ErrInvalidRange, "unable to set range")
			}
		}

		if length == 0 {
			// zero-length ranges require special handling, set non-zero range and
			// we won't be trying to read the response anyway.
			if err := opt.SetRange(0, 1); err != nil {
				return errors.Wrap(blob.ErrInvalidRange, "unable to set range")
			}
		}

		o, err := s.cli.GetObject(ctx, s.BucketName, s.getObjectNameString(b), opt)
		if err != nil {
			return errors.Wrap(err, "GetObject")
		}

		defer o.Close() //nolint:errcheck

		if length == 0 {
			return nil
		}

		// nolint:wrapcheck
		return iocopy.JustCopy(output, o)
	}

	if err := attempt(); err != nil {
		return translateError(err)
	}

	// nolint:wrapcheck
	return blob.EnsureLengthExactly(output.Length(), length)
}

func isInvalidCredentials(err error) bool {
	return err != nil && strings.Contains(err.Error(), blob.InvalidCredentialsErrStr)
}

func translateError(err error) error {
	var me minio.ErrorResponse

	if isInvalidCredentials(err) {
		return blob.ErrInvalidCredentials
	}

	if errors.As(err, &me) {
		switch me.StatusCode {
		case http.StatusOK:
			return nil

		case http.StatusNotFound:
			return blob.ErrBlobNotFound

		case http.StatusRequestedRangeNotSatisfiable:
			return blob.ErrInvalidRange
		}
	}

	return err
}

func (s *s3Storage) GetMetadata(ctx context.Context, b blob.ID) (blob.Metadata, error) {
	vm, err := s.getVersionMetadata(ctx, b, "")

	return vm.Metadata, err
}

func (s *s3Storage) getVersionMetadata(ctx context.Context, b blob.ID, version string) (versionMetadata, error) {
	opts := minio.GetObjectOptions{
		VersionID: version,
	}

	oi, err := s.cli.StatObject(ctx, s.BucketName, s.getObjectNameString(b), opts)
	if err != nil {
		return versionMetadata{}, errors.Wrap(translateError(err), "StatObject")
	}

	return infoToVersionMetadata(s.Prefix, &oi), nil
}

func (s *s3Storage) PutBlob(ctx context.Context, b blob.ID, data blob.Bytes, opts blob.PutOptions) error {
	switch {
	case opts.DoNotRecreate:
		return errors.Wrap(blob.ErrUnsupportedPutBlobOption, "do-not-recreate")
	case !opts.SetModTime.IsZero():
		return blob.ErrSetTimeUnsupported
	}

	_, err := s.putBlob(ctx, b, data, opts)

	if opts.GetModTime != nil {
		bm, err2 := s.GetMetadata(ctx, b)
		if err2 != nil {
			return err2
		}

		*opts.GetModTime = bm.Timestamp
	}

	return err
}

func (s *s3Storage) putBlob(ctx context.Context, b blob.ID, data blob.Bytes, opts blob.PutOptions) (versionMetadata, error) {
	var (
		storageClass    = s.storageConfig.getStorageClassForBlobID(b)
		retentionMode   minio.RetentionMode
		retainUntilDate time.Time
	)

	if opts.RetentionPeriod != 0 {
		retentionMode = minio.RetentionMode(opts.RetentionMode)
		if !retentionMode.IsValid() {
			return versionMetadata{}, errors.Errorf("invalid retention mode: %q", opts.RetentionMode)
		}

		retainUntilDate = clock.Now().Add(opts.RetentionPeriod).UTC()
	}

	uploadInfo, err := s.cli.PutObject(ctx, s.BucketName, s.getObjectNameString(b), data.Reader(), int64(data.Length()), minio.PutObjectOptions{
		ContentType: "application/x-kopia",
		// Kopia already splits snapshot contents into small blobs to improve
		// upload throughput. There is no need for further splitting
		// through multipart uploads.
		DisableMultipart: true,
		// The Content-MD5 header is required for any request to upload an object
		// with a retention period configured using Amazon S3 Object Lock.
		// Unconditionally computing the content MD5, potentially incurring
		// a slightly higher CPU overhead.
		SendContentMd5:  true,
		StorageClass:    storageClass,
		RetainUntilDate: retainUntilDate,
		Mode:            retentionMode,
	})

	if isInvalidCredentials(err) {
		return versionMetadata{}, blob.ErrInvalidCredentials
	}

	var er minio.ErrorResponse

	if errors.As(err, &er) && er.Code == "InvalidRequest" && strings.Contains(strings.ToLower(er.Message), "content-md5") {
		return versionMetadata{}, err // nolint:wrapcheck
	}

	if errors.Is(err, io.EOF) && uploadInfo.Size == 0 {
		// special case empty stream
		_, err = s.cli.PutObject(ctx, s.BucketName, s.getObjectNameString(b), bytes.NewBuffer(nil), 0, minio.PutObjectOptions{
			ContentType:     "application/x-kopia",
			StorageClass:    storageClass,
			RetainUntilDate: retainUntilDate,
			Mode:            retentionMode,
		})
	}

	if err != nil {
		return versionMetadata{}, err // nolint:wrapcheck
	}

	return versionMetadata{
		Metadata: blob.Metadata{
			BlobID:    b,
			Length:    uploadInfo.Size,
			Timestamp: uploadInfo.LastModified,
		},
		Version: uploadInfo.VersionID,
	}, nil
}

func (s *s3Storage) DeleteBlob(ctx context.Context, b blob.ID) error {
	err := translateError(s.cli.RemoveObject(ctx, s.BucketName, s.getObjectNameString(b), minio.RemoveObjectOptions{}))
	if errors.Is(err, blob.ErrBlobNotFound) {
		return nil
	}

	return err
}

func (s *s3Storage) getObjectNameString(b blob.ID) string {
	return s.Prefix + string(b)
}

func (s *s3Storage) ListBlobs(ctx context.Context, prefix blob.ID, callback func(blob.Metadata) error) error {
	ctx, cancel := context.WithCancel(ctx)

	defer cancel()

	oi := s.cli.ListObjects(ctx, s.BucketName, minio.ListObjectsOptions{
		Prefix: s.getObjectNameString(prefix),
	})
	for o := range oi {
		if err := o.Err; err != nil {
			if isInvalidCredentials(err) {
				return blob.ErrInvalidCredentials
			}

			return err
		}

		bm := blob.Metadata{
			BlobID:    blob.ID(o.Key[len(s.Prefix):]),
			Length:    o.Size,
			Timestamp: o.LastModified,
		}

		if bm.BlobID == ConfigName {
			continue
		}

		if err := callback(bm); err != nil {
			return err
		}
	}

	return nil
}

func (s *s3Storage) ConnectionInfo() blob.ConnectionInfo {
	return blob.ConnectionInfo{
		Type:   s3storageType,
		Config: &s.Options,
	}
}

func (s *s3Storage) Close(ctx context.Context) error {
	return nil
}

func (s *s3Storage) String() string {
	return fmt.Sprintf("s3://%v/%v", s.BucketName, s.Prefix)
}

func (s *s3Storage) DisplayName() string {
	return fmt.Sprintf("S3: %v %v", s.Endpoint, s.BucketName)
}

func (s *s3Storage) FlushCaches(ctx context.Context) error {
	return nil
}

func getCustomTransport(insecureSkipVerify bool) (transport *http.Transport) {
	// nolint:gosec
	customTransport := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: insecureSkipVerify}}
	return customTransport
}

// New creates new S3-backed storage with specified options:
//
// - the 'BucketName' field is required and all other parameters are optional.
func New(ctx context.Context, opt *Options) (blob.Storage, error) {
	st, err := newStorage(ctx, opt)
	if err != nil {
		return nil, err
	}

	s, err := maybePointInTimeStore(ctx, st, opt.PointInTime)
	if err != nil {
		return nil, err
	}

	return retrying.NewWrapper(s), nil
}

func newStorage(ctx context.Context, opt *Options) (*s3Storage, error) {
	creds := credentials.NewChainCredentials(
		[]credentials.Provider{
			&credentials.Static{
				Value: credentials.Value{
					AccessKeyID:     opt.AccessKeyID,
					SecretAccessKey: opt.SecretAccessKey,
					SessionToken:    opt.SessionToken,
					SignerType:      credentials.SignatureV4,
				},
			},
			&credentials.EnvAWS{},
			&credentials.IAM{
				Client: &http.Client{
					Transport: http.DefaultTransport,
				},
			},
		},
	)

	return newStorageWithCredentials(ctx, creds, opt)
}

func newStorageWithCredentials(ctx context.Context, creds *credentials.Credentials, opt *Options) (*s3Storage, error) {
	if opt.BucketName == "" {
		return nil, errors.New("bucket name must be specified")
	}

	minioOpts := &minio.Options{
		Creds:  creds,
		Secure: !opt.DoNotUseTLS,
		Region: opt.Region,
	}

	if opt.DoNotVerifyTLS {
		minioOpts.Transport = getCustomTransport(true)
	}

	cli, err := minio.New(opt.Endpoint, minioOpts)
	if err != nil {
		return nil, errors.Wrap(err, "unable to create client")
	}

	ok, err := cli.BucketExists(ctx, opt.BucketName)
	if err != nil {
		return nil, errors.Wrapf(err, "unable to determine if bucket %q exists", opt.BucketName)
	}

	if !ok {
		return nil, errors.Errorf("bucket %q does not exist", opt.BucketName)
	}

	s := s3Storage{
		Options:       *opt,
		cli:           cli,
		storageConfig: &StorageConfig{},
	}

	var scOutput gather.WriteBuffer

	if getBlobErr := s.GetBlob(ctx, ConfigName, 0, -1, &scOutput); getBlobErr == nil {
		if scErr := s.storageConfig.Load(scOutput.Bytes().Reader()); scErr != nil {
			return nil, errors.Wrapf(scErr, "error parsing storage config for bucket %q", opt.BucketName)
		}
	} else if !errors.Is(getBlobErr, blob.ErrBlobNotFound) {
		return nil, errors.Wrapf(getBlobErr, "error retrieving storage config from bucket %q", opt.BucketName)
	}

	return &s, nil
}

func init() {
	blob.AddSupportedStorage(
		s3storageType,
		func() interface{} {
			return &Options{}
		},
		func(ctx context.Context, o interface{}, isCreate bool) (blob.Storage, error) {
			return New(ctx, o.(*Options)) // nolint:forcetypeassert
		})
}
<|endoftext|>


<commit_before>package jobserver

import (
	"fmt"
	"sync"
	"time"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/persist"
	"go.pedge.io/google-protobuf"
	"go.pedge.io/proto/rpclog"
	"go.pedge.io/proto/time"
	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/apis/extensions"
	kube "k8s.io/kubernetes/pkg/client/unversioned"
)

var (
	trueVal = true
	suite   = "pachyderm"
)

type apiServer struct {
	protorpclog.Logger
	pfsAPIClient     pfs.APIClient
	persistAPIClient persist.APIClient
	kubeClient       *kube.Client
	startJobCounter  map[string]uint64
	startJobLock     sync.Mutex
	finishJobCounter map[string]uint64
	finishJobLock    sync.Mutex
}

func newAPIServer(
	pfsAPIClient pfs.APIClient,
	persistAPIClient persist.APIClient,
	kubeClient *kube.Client,
) *apiServer {
	return &apiServer{
		protorpclog.NewLogger("pachyderm.pps.JobAPI"),
		pfsAPIClient,
		persistAPIClient,
		kubeClient,
		make(map[string]uint64),
		sync.Mutex{},
		make(map[string]uint64),
		sync.Mutex{},
	}
}

func (a *apiServer) CreateJob(ctx context.Context, request *pps.CreateJobRequest) (response *pps.Job, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	if request.OutputParent == nil {
		return nil, fmt.Errorf("pachyderm.pps.jobserver: request.OutputParent cannot be nil")
	}
	if request.Shards == 0 {
		return nil, fmt.Errorf("pachyderm.pps.jobserver: request.Shards cannot be 0")
	}
	// TODO validate job to make sure input commits and output repo exist
	persistJobInfo := &persist.JobInfo{
		Shards:       request.Shards,
		Transform:    request.Transform,
		InputCommit:  request.InputCommit,
		OutputParent: request.OutputParent,
	}
	if request.Pipeline != nil {
		persistJobInfo.PipelineName = request.Pipeline.Name
	}
	persistJobInfo.JobId = uuid.NewWithoutDashes()
	persistJobInfo.CreatedAt = prototime.TimeToTimestamp(time.Now())
	if a.kubeClient == nil {
		return nil, fmt.Errorf("pachyderm.pps.jobserver: no job backend")
	}
	_, err := a.persistAPIClient.CreateJobInfo(ctx, persistJobInfo)
	if err != nil {
		return nil, err
	}
	if _, err := a.kubeClient.Jobs(api.NamespaceDefault).Create(job(persistJobInfo)); err != nil {
		return nil, err
	}
	return &pps.Job{
		Id: persistJobInfo.JobId,
	}, nil
}

func (a *apiServer) InspectJob(ctx context.Context, request *pps.InspectJobRequest) (response *pps.JobInfo, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	persistJobInfo, err := a.persistAPIClient.GetJobInfo(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	return a.persistJobInfoToJobInfo(ctx, persistJobInfo)
}

func (a *apiServer) ListJob(ctx context.Context, request *pps.ListJobRequest) (response *pps.JobInfos, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	persistJobInfos, err := a.persistAPIClient.ListJobInfos(ctx, request)
	if err != nil {
		return nil, err
	}
	jobInfos := make([]*pps.JobInfo, len(persistJobInfos.JobInfo))
	for i, persistJobInfo := range persistJobInfos.JobInfo {
		jobInfo, err := a.persistJobInfoToJobInfo(ctx, persistJobInfo)
		if err != nil {
			return nil, err
		}
		jobInfos[i] = jobInfo
	}
	return &pps.JobInfos{
		JobInfo: jobInfos,
	}, nil
}

func (a *apiServer) StartJob(ctx context.Context, request *pps.StartJobRequest) (response *pps.StartJobResponse, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	jobInfo, err := a.persistAPIClient.GetJobInfo(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	var shard uint64
	if err := func() error {
		a.startJobLock.Lock()
		defer a.startJobLock.Unlock()
		shard = a.startJobCounter[request.Job.Id]
		a.startJobCounter[request.Job.Id] = shard + 1
		if shard == 0 {
			commit, err := a.pfsAPIClient.StartCommit(ctx, &pfs.StartCommitRequest{
				Parent: jobInfo.OutputParent,
			})
			if err != nil {
				return err
			}
			if _, err := a.persistAPIClient.CreateJobOutput(
				ctx,
				&persist.JobOutput{
					JobId:        request.Job.Id,
					OutputCommit: commit,
				}); err != nil {
				return err
			}
		}
		return nil
	}(); err != nil {
		return nil, err
	}
	if jobInfo.GetTransform() == nil {
		return nil, fmt.Errorf("jobInfo.GetTransform() should not be nil (this is likely a bug)")
	}
	jobOutput, err := a.persistAPIClient.GetJobOutput(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	if jobOutput.OutputCommit == nil {
		return nil, fmt.Errorf("jobOutput.OutputCommit should not be nil (this is likely a bug)")
	}
	return &pps.StartJobResponse{
		Transform:    jobInfo.GetTransform(),
		InputCommit:  jobInfo.InputCommit,
		OutputCommit: jobOutput.OutputCommit,
		Shard: &pfs.Shard{
			Number:  shard,
			Modulus: jobInfo.Shards,
		},
	}, nil
}

func (a *apiServer) FinishJob(ctx context.Context, request *pps.FinishJobRequest) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	jobInfo, err := a.persistAPIClient.GetJobInfo(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	var finished uint64
	func() {
		a.finishJobLock.Lock()
		defer a.finishJobLock.Unlock()
		a.finishJobCounter[request.Job.Id] = a.finishJobCounter[request.Job.Id] + 1
		finished = a.finishJobCounter[request.Job.Id]
	}()
	if finished == jobInfo.Shards {
		// all of the shards have finished so we finish the commit
		jobOutput, err := a.persistAPIClient.GetJobOutput(ctx, request.Job)
		if err != nil {
			return nil, err
		}
		if jobOutput.OutputCommit == nil {
			return nil, fmt.Errorf("jobOutput.OutputCommit should not be nil (this is likely a bug)")
		}
		if _, err := a.pfsAPIClient.FinishCommit(ctx, &pfs.FinishCommitRequest{
			Commit: jobOutput.OutputCommit,
		}); err != nil {
			return nil, err
		}
	}
	return google_protobuf.EmptyInstance, nil
}

func (a *apiServer) persistJobInfoToJobInfo(ctx context.Context, persistJobInfo *persist.JobInfo) (*pps.JobInfo, error) {
	job := &pps.Job{Id: persistJobInfo.JobId}
	jobInfo := &pps.JobInfo{
		Job:         job,
		Transform:   persistJobInfo.Transform,
		Pipeline:    &pps.Pipeline{Name: persistJobInfo.PipelineName},
		Shards:      persistJobInfo.Shards,
		InputCommit: persistJobInfo.InputCommit,
	}
	persistJobOutput, err := a.persistAPIClient.GetJobOutput(ctx, job)
	if err == nil && persistJobOutput != nil {
		jobInfo.OutputCommit = persistJobOutput.OutputCommit
	}
	return jobInfo, nil
}

func job(jobInfo *persist.JobInfo) *extensions.Job {
	app := jobInfo.JobId
	shards := int(jobInfo.Shards)
	image := "pachyderm/job-shim"
	if jobInfo.GetTransform().Image != "" {
		image = jobInfo.GetTransform().Image
	}
	return &extensions.Job{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Job",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   jobInfo.JobId,
			Labels: labels(app),
		},
		Spec: extensions.JobSpec{
			Selector: &extensions.PodSelector{
				MatchLabels: labels(app),
			},
			Parallelism: &shards,
			Completions: &shards,
			Template: api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   jobInfo.JobId,
					Labels: labels(app),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:    "user",
							Image:   image,
							Command: []string{"/job-shim", jobInfo.JobId},
							SecurityContext: &api.SecurityContext{
								Privileged: &trueVal, // god is this dumb
							},
						},
					},
					RestartPolicy: "OnFailure",
				},
			},
		},
	}
}

func labels(app string) map[string]string {
	return map[string]string{
		"app":   app,
		"suite": suite,
	}
}
<commit_msg>No more GetTransform().<commit_after>package jobserver

import (
	"fmt"
	"sync"
	"time"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/persist"
	"go.pedge.io/google-protobuf"
	"go.pedge.io/proto/rpclog"
	"go.pedge.io/proto/time"
	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/apis/extensions"
	kube "k8s.io/kubernetes/pkg/client/unversioned"
)

var (
	trueVal = true
	suite   = "pachyderm"
)

type apiServer struct {
	protorpclog.Logger
	pfsAPIClient     pfs.APIClient
	persistAPIClient persist.APIClient
	kubeClient       *kube.Client
	startJobCounter  map[string]uint64
	startJobLock     sync.Mutex
	finishJobCounter map[string]uint64
	finishJobLock    sync.Mutex
}

func newAPIServer(
	pfsAPIClient pfs.APIClient,
	persistAPIClient persist.APIClient,
	kubeClient *kube.Client,
) *apiServer {
	return &apiServer{
		protorpclog.NewLogger("pachyderm.pps.JobAPI"),
		pfsAPIClient,
		persistAPIClient,
		kubeClient,
		make(map[string]uint64),
		sync.Mutex{},
		make(map[string]uint64),
		sync.Mutex{},
	}
}

func (a *apiServer) CreateJob(ctx context.Context, request *pps.CreateJobRequest) (response *pps.Job, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	if request.OutputParent == nil {
		return nil, fmt.Errorf("pachyderm.pps.jobserver: request.OutputParent cannot be nil")
	}
	if request.Shards == 0 {
		return nil, fmt.Errorf("pachyderm.pps.jobserver: request.Shards cannot be 0")
	}
	// TODO validate job to make sure input commits and output repo exist
	persistJobInfo := &persist.JobInfo{
		Shards:       request.Shards,
		Transform:    request.Transform,
		InputCommit:  request.InputCommit,
		OutputParent: request.OutputParent,
	}
	if request.Pipeline != nil {
		persistJobInfo.PipelineName = request.Pipeline.Name
	}
	persistJobInfo.JobId = uuid.NewWithoutDashes()
	persistJobInfo.CreatedAt = prototime.TimeToTimestamp(time.Now())
	if a.kubeClient == nil {
		return nil, fmt.Errorf("pachyderm.pps.jobserver: no job backend")
	}
	_, err := a.persistAPIClient.CreateJobInfo(ctx, persistJobInfo)
	if err != nil {
		return nil, err
	}
	if _, err := a.kubeClient.Jobs(api.NamespaceDefault).Create(job(persistJobInfo)); err != nil {
		return nil, err
	}
	return &pps.Job{
		Id: persistJobInfo.JobId,
	}, nil
}

func (a *apiServer) InspectJob(ctx context.Context, request *pps.InspectJobRequest) (response *pps.JobInfo, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	persistJobInfo, err := a.persistAPIClient.GetJobInfo(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	return a.persistJobInfoToJobInfo(ctx, persistJobInfo)
}

func (a *apiServer) ListJob(ctx context.Context, request *pps.ListJobRequest) (response *pps.JobInfos, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	persistJobInfos, err := a.persistAPIClient.ListJobInfos(ctx, request)
	if err != nil {
		return nil, err
	}
	jobInfos := make([]*pps.JobInfo, len(persistJobInfos.JobInfo))
	for i, persistJobInfo := range persistJobInfos.JobInfo {
		jobInfo, err := a.persistJobInfoToJobInfo(ctx, persistJobInfo)
		if err != nil {
			return nil, err
		}
		jobInfos[i] = jobInfo
	}
	return &pps.JobInfos{
		JobInfo: jobInfos,
	}, nil
}

func (a *apiServer) StartJob(ctx context.Context, request *pps.StartJobRequest) (response *pps.StartJobResponse, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	jobInfo, err := a.persistAPIClient.GetJobInfo(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	var shard uint64
	if err := func() error {
		a.startJobLock.Lock()
		defer a.startJobLock.Unlock()
		shard = a.startJobCounter[request.Job.Id]
		a.startJobCounter[request.Job.Id] = shard + 1
		if shard == 0 {
			commit, err := a.pfsAPIClient.StartCommit(ctx, &pfs.StartCommitRequest{
				Parent: jobInfo.OutputParent,
			})
			if err != nil {
				return err
			}
			if _, err := a.persistAPIClient.CreateJobOutput(
				ctx,
				&persist.JobOutput{
					JobId:        request.Job.Id,
					OutputCommit: commit,
				}); err != nil {
				return err
			}
		}
		return nil
	}(); err != nil {
		return nil, err
	}
	if jobInfo.Transform == nil {
		return nil, fmt.Errorf("jobInfo.Transform should not be nil (this is likely a bug)")
	}
	jobOutput, err := a.persistAPIClient.GetJobOutput(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	if jobOutput.OutputCommit == nil {
		return nil, fmt.Errorf("jobOutput.OutputCommit should not be nil (this is likely a bug)")
	}
	return &pps.StartJobResponse{
		Transform:    jobInfo.Transform,
		InputCommit:  jobInfo.InputCommit,
		OutputCommit: jobOutput.OutputCommit,
		Shard: &pfs.Shard{
			Number:  shard,
			Modulus: jobInfo.Shards,
		},
	}, nil
}

func (a *apiServer) FinishJob(ctx context.Context, request *pps.FinishJobRequest) (response *google_protobuf.Empty, retErr error) {
	defer func(start time.Time) { a.Log(request, response, retErr, time.Since(start)) }(time.Now())
	jobInfo, err := a.persistAPIClient.GetJobInfo(ctx, request.Job)
	if err != nil {
		return nil, err
	}
	var finished uint64
	func() {
		a.finishJobLock.Lock()
		defer a.finishJobLock.Unlock()
		a.finishJobCounter[request.Job.Id] = a.finishJobCounter[request.Job.Id] + 1
		finished = a.finishJobCounter[request.Job.Id]
	}()
	if finished == jobInfo.Shards {
		// all of the shards have finished so we finish the commit
		jobOutput, err := a.persistAPIClient.GetJobOutput(ctx, request.Job)
		if err != nil {
			return nil, err
		}
		if jobOutput.OutputCommit == nil {
			return nil, fmt.Errorf("jobOutput.OutputCommit should not be nil (this is likely a bug)")
		}
		if _, err := a.pfsAPIClient.FinishCommit(ctx, &pfs.FinishCommitRequest{
			Commit: jobOutput.OutputCommit,
		}); err != nil {
			return nil, err
		}
	}
	return google_protobuf.EmptyInstance, nil
}

func (a *apiServer) persistJobInfoToJobInfo(ctx context.Context, persistJobInfo *persist.JobInfo) (*pps.JobInfo, error) {
	job := &pps.Job{Id: persistJobInfo.JobId}
	jobInfo := &pps.JobInfo{
		Job:         job,
		Transform:   persistJobInfo.Transform,
		Pipeline:    &pps.Pipeline{Name: persistJobInfo.PipelineName},
		Shards:      persistJobInfo.Shards,
		InputCommit: persistJobInfo.InputCommit,
	}
	persistJobOutput, err := a.persistAPIClient.GetJobOutput(ctx, job)
	if err == nil && persistJobOutput != nil {
		jobInfo.OutputCommit = persistJobOutput.OutputCommit
	}
	return jobInfo, nil
}

func job(jobInfo *persist.JobInfo) *extensions.Job {
	app := jobInfo.JobId
	shards := int(jobInfo.Shards)
	image := "pachyderm/job-shim"
	if jobInfo.Transform.Image != "" {
		image = jobInfo.Transform.Image
	}
	return &extensions.Job{
		TypeMeta: unversioned.TypeMeta{
			Kind:       "Job",
			APIVersion: "v1",
		},
		ObjectMeta: api.ObjectMeta{
			Name:   jobInfo.JobId,
			Labels: labels(app),
		},
		Spec: extensions.JobSpec{
			Selector: &extensions.PodSelector{
				MatchLabels: labels(app),
			},
			Parallelism: &shards,
			Completions: &shards,
			Template: api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Name:   jobInfo.JobId,
					Labels: labels(app),
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:    "user",
							Image:   image,
							Command: []string{"/job-shim", jobInfo.JobId},
							SecurityContext: &api.SecurityContext{
								Privileged: &trueVal, // god is this dumb
							},
						},
					},
					RestartPolicy: "OnFailure",
				},
			},
		},
	}
}

func labels(app string) map[string]string {
	return map[string]string{
		"app":   app,
		"suite": suite,
	}
}
<|endoftext|>


<commit_before>// Gone Time Tracker -or- Where has my time gone?
package main

import (
	"encoding/gob"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/mewkiz/pkg/goutil"
)

var (
	goneDir       string
	dumpFileName  string
	logFileName   string
	indexFileName string
	tracks        Tracks
	zzz           bool
	logger        *log.Logger
	current       Window
)

func init() {
	var err error
	goneDir, err = goutil.SrcDir("github.com/dim13/gone")
	if err != nil {
		log.Fatal("init: ", err)
	}
	dumpFileName = filepath.Join(goneDir, "gone.gob")
	logFileName = filepath.Join(goneDir, "gone.log")
	indexFileName = filepath.Join(goneDir, "index.html")
}

type Tracker interface {
	Update(Window)
	Snooze(time.Duration)
	Wakeup()
}

type Tracks map[Window]Track

type Track struct {
	Seen  time.Time
	Spent time.Duration
}

type Window struct {
	Class string
	Name  string
}

func (t Track) String() string {
	return fmt.Sprintf("%s %s",
		t.Seen.Format("2006/01/02 15:04:05"), t.Spent)
}

func (w Window) String() string {
	return fmt.Sprintf("%s %s", w.Class, w.Name)
}

func (t Tracks) Snooze(idle time.Duration) {
	if zzz == false {
		log.Println("away from keyboard, idle for", idle)
		if c, ok := t[current]; ok {
			if idle > c.Spent && c.Spent > 0 {
				c.Spent -= idle
				t[current] = c
			}
		}
		zzz = true
	}
}

func (t Tracks) Wakeup() {
	if zzz == true {
		log.Println("back to keyboard")
		zzz = false
	}
}

func (t Tracks) Update(w Window) {
	if zzz == false {
		if c, ok := t[current]; ok {
			c.Spent += time.Since(c.Seen)
			t[current] = c
		}
	}

	if _, ok := t[w]; !ok {
		t[w] = Track{}
	}

	s := t[w]
	s.Seen = time.Now()
	t[w] = s

	current = w
}

func (t Tracks) Remove(d time.Duration) {
	for k, v := range t {
		if time.Since(v.Seen) > d {
			logger.Println(v, k)
			delete(t, k)
		}
	}
}

func Load(fname string) Tracks {
	t := make(Tracks)
	dump, err := os.Open(fname)
	if err != nil {
		log.Println(err)
		return t
	}
	defer dump.Close()
	dec := gob.NewDecoder(dump)
	err = dec.Decode(&t)
	if err != nil {
		log.Println(err)
	}
	return t
}

func (t Tracks) Store(fname string) {
	tmp := fname + ".tmp"
	dump, err := os.Create(tmp)
	if err != nil {
		log.Println(err)
		return
	}
	defer dump.Close()
	enc := gob.NewEncoder(dump)
	err = enc.Encode(t)
	if err != nil {
		log.Println(err)
		os.Remove(tmp)
		return
	}
	os.Rename(tmp, fname)
}

func (t Tracks) Cleanup() {
	for {
		tracks.Remove(8 * time.Hour)
		tracks.Store(dumpFileName)
		time.Sleep(time.Minute)
	}
}

func main() {
	X := Connect()
	defer X.Close()

	logfile, err := os.OpenFile(logFileName,
		os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		log.Fatal(err)
	}
	defer logfile.Close()
	logger = log.New(logfile, "", log.LstdFlags)

	tracks = Load(dumpFileName)

	go X.Collect(tracks)
	go tracks.Cleanup()

	webReporter("127.0.0.1:8001")
}
<commit_msg>reset seen time for current task on wakeup<commit_after>// Gone Time Tracker -or- Where has my time gone?
package main

import (
	"encoding/gob"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/mewkiz/pkg/goutil"
)

var (
	goneDir       string
	dumpFileName  string
	logFileName   string
	indexFileName string
	tracks        Tracks
	zzz           bool
	logger        *log.Logger
	current       Window
)

func init() {
	var err error
	goneDir, err = goutil.SrcDir("github.com/dim13/gone")
	if err != nil {
		log.Fatal("init: ", err)
	}
	dumpFileName = filepath.Join(goneDir, "gone.gob")
	logFileName = filepath.Join(goneDir, "gone.log")
	indexFileName = filepath.Join(goneDir, "index.html")
}

type Tracker interface {
	Update(Window)
	Snooze(time.Duration)
	Wakeup()
}

type Tracks map[Window]Track

type Track struct {
	Seen  time.Time
	Spent time.Duration
}

type Window struct {
	Class string
	Name  string
}

func (t Track) String() string {
	return fmt.Sprintf("%s %s",
		t.Seen.Format("2006/01/02 15:04:05"), t.Spent)
}

func (w Window) String() string {
	return fmt.Sprintf("%s %s", w.Class, w.Name)
}

func (t Tracks) Snooze(idle time.Duration) {
	if zzz == false {
		log.Println("away from keyboard, idle for", idle)
		if c, ok := t[current]; ok {
			if idle > c.Spent && c.Spent > 0 {
				c.Spent -= idle
				t[current] = c
			}
		}
		zzz = true
	}
}

func (t Tracks) Wakeup() {
	if zzz == true {
		log.Println("back to keyboard")
		if c, ok := t[current]; ok {
			c.Seen = time.Now()
			t[current] = c
		}
		zzz = false
	}
}

func (t Tracks) Update(w Window) {
	if zzz == false {
		if c, ok := t[current]; ok {
			c.Spent += time.Since(c.Seen)
			t[current] = c
		}
	}

	if _, ok := t[w]; !ok {
		t[w] = Track{}
	}

	s := t[w]
	s.Seen = time.Now()
	t[w] = s

	current = w
}

func (t Tracks) Remove(d time.Duration) {
	for k, v := range t {
		if time.Since(v.Seen) > d {
			logger.Println(v, k)
			delete(t, k)
		}
	}
}

func Load(fname string) Tracks {
	t := make(Tracks)
	dump, err := os.Open(fname)
	if err != nil {
		log.Println(err)
		return t
	}
	defer dump.Close()
	dec := gob.NewDecoder(dump)
	err = dec.Decode(&t)
	if err != nil {
		log.Println(err)
	}
	return t
}

func (t Tracks) Store(fname string) {
	tmp := fname + ".tmp"
	dump, err := os.Create(tmp)
	if err != nil {
		log.Println(err)
		return
	}
	defer dump.Close()
	enc := gob.NewEncoder(dump)
	err = enc.Encode(t)
	if err != nil {
		log.Println(err)
		os.Remove(tmp)
		return
	}
	os.Rename(tmp, fname)
}

func (t Tracks) Cleanup() {
	for {
		tracks.Remove(8 * time.Hour)
		tracks.Store(dumpFileName)
		time.Sleep(time.Minute)
	}
}

func main() {
	X := Connect()
	defer X.Close()

	logfile, err := os.OpenFile(logFileName,
		os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		log.Fatal(err)
	}
	defer logfile.Close()
	logger = log.New(logfile, "", log.LstdFlags)

	tracks = Load(dumpFileName)

	go X.Collect(tracks)
	go tracks.Cleanup()

	webReporter("127.0.0.1:8001")
}
<|endoftext|>


<commit_before>// Copyright 2016 Red Hat, Inc. and/or its affiliates
// and other contributors.
//
// 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 backend
package random

import (
	"fmt"
	"math/rand"

	"github.com/yaacov/mohawk/backend"
)

type Backend struct {
	Items []backend.Item
}

func (r Backend) Name() string {
	return "Backend-Random"
}

func (r *Backend) Open() {
	r.Items = make([]backend.Item, 0)

	seeds := []map[string]string{
		map[string]string{"type": "node", "group_id": "cpu/usage_rate", "units": "cpu", "issue": "42"},
		map[string]string{"type": "node", "group_id": "memory/usage_rate", "units": "byte"},
		map[string]string{"type": "node", "group_id": "cpu/usage_rate", "units": "cpu", "issue": "42"},
		map[string]string{"type": "node", "group_id": "memory/usage_rate", "units": "byte"},
		map[string]string{"type": "node", "group_id": "cpu/limit", "units": "cpu", "issue": "442"},
		map[string]string{"type": "node", "group_id": "memory/limit", "units": "byte", "issue": "442"},
		map[string]string{"type": "node", "group_id": "filesystem/usage_rate", "units": "byte"},
	}

	for i := 0; i < 120; i++ {
		seed := seeds[rand.Intn(len(seeds))]
		tags := map[string]string{
			"type":     seed["type"],
			"group_id": seed["group_id"],
			"units":    seed["units"],
			"issue":    seed["issue"],
			"hostname": fmt.Sprintf("example.%03d.com", i/4),
		}
		id := fmt.Sprintf("machine/%s/%s", tags["hostname"], tags["group_id"])

		validID := true
		for _, o := range r.Items {
			if id == o.Id {
				validID = false
			}
		}

		if validID {
			r.Items = append(r.Items, backend.Item{
				Id:   id,
				Type: "gauge",
				Tags: tags,
			})
		}
	}
}

func (r Backend) GetItemList(tags map[string]string) []backend.Item {
	res := r.Items

	if len(tags) > 0 {
		for key, value := range tags {
			res = backend.FilterItems(res, func(i backend.Item) bool { return i.Tags[key] == value })
		}
	}

	return res
}

func (r Backend) GetRawData(id string, end int64, start int64, limit int64, order string) []backend.DataItem {
	var i int64
	var res []backend.DataItem

	delta := int64(5 * 60 * 1000)
	base := 10 + rand.Intn(250)
	variant := 20 + rand.Intn(50)

	for i = 0; i < limit && (end-i*delta) > start; i++ {
		res = append(res, backend.DataItem{
			Timestamp: end - i*delta,
			Value:     float64(base + rand.Intn(variant)),
		})
	}

	if order == "ASC" {
		for i, j := 0, len(res)-1; i < j; i, j = i+1, j-1 {
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r Backend) GetStatData(id string, end int64, start int64, limit int64, order string, bucketDuration int64) []backend.StatItem {
	var i int64
	var res []backend.StatItem

	delta := bucketDuration * 1000
	base := 10 + rand.Intn(250)
	variant := 20 + rand.Intn(50)

	for i = 0; i < limit && (end-(i-1)*delta) > start; i++ {
		value := float64(base + rand.Intn(variant))
		res = append(res, backend.StatItem{
			Start:   end - (i-1)*delta,
			End:     end - i*delta,
			Empty:   false,
			Samples: 1,
			Min:     value,
			Max:     value,
			Avg:     value,
			Median:  value,
			Sum:     value,
		})
	}

	if order == "ASC" {
		for i, j := 0, len(res)-1; i < j; i, j = i+1, j-1 {
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r Backend) PostRawData(id string, t int64, v float64) bool {
	return false
}

func (r Backend) PutTags(id string, tags map[string]string) bool {
	return false
}
<commit_msg>more random metrics<commit_after>// Copyright 2016 Red Hat, Inc. and/or its affiliates
// and other contributors.
//
// 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 backend
package random

import (
	"fmt"
	"math/rand"

	"github.com/yaacov/mohawk/backend"
)

type Backend struct {
	Items []backend.Item
}

func (r Backend) Name() string {
	return "Backend-Random"
}

func (r *Backend) Open() {
	r.Items = make([]backend.Item, 0)

	seeds := []map[string]string{
		map[string]string{"type": "node", "group_id": "cpu/usage_rate", "units": "cpu", "issue": "42"},
		map[string]string{"type": "node", "group_id": "memory/usage_rate", "units": "byte"},
		map[string]string{"type": "node", "group_id": "cpu/usage_rate", "units": "cpu", "issue": "42"},
		map[string]string{"type": "node", "group_id": "memory/usage_rate", "units": "byte"},
		map[string]string{"type": "node", "group_id": "cpu/limit", "units": "cpu", "issue": "442"},
		map[string]string{"type": "node", "group_id": "memory/limit", "units": "byte", "issue": "442"},
		map[string]string{"type": "node", "group_id": "filesystem/usage_rate", "units": "byte"},
	}

	for i := 0; i < 2020; i++ {
		seed := seeds[rand.Intn(len(seeds))]
		tags := map[string]string{
			"type":     seed["type"],
			"group_id": seed["group_id"],
			"units":    seed["units"],
			"issue":    seed["issue"],
			"hostname": fmt.Sprintf("example.%04d.com", i/4),
		}
		id := fmt.Sprintf("machine/%s/%s", tags["hostname"], tags["group_id"])

		validID := true
		for _, o := range r.Items {
			if id == o.Id {
				validID = false
			}
		}

		if validID {
			r.Items = append(r.Items, backend.Item{
				Id:   id,
				Type: "gauge",
				Tags: tags,
			})
		}
	}
}

func (r Backend) GetItemList(tags map[string]string) []backend.Item {
	res := r.Items

	if len(tags) > 0 {
		for key, value := range tags {
			res = backend.FilterItems(res, func(i backend.Item) bool { return i.Tags[key] == value })
		}
	}

	return res
}

func (r Backend) GetRawData(id string, end int64, start int64, limit int64, order string) []backend.DataItem {
	var i int64
	var res []backend.DataItem

	delta := int64(5 * 60 * 1000)
	base := 10 + rand.Intn(250)
	variant := 20 + rand.Intn(50)

	for i = 0; i < limit && (end-i*delta) > start; i++ {
		res = append(res, backend.DataItem{
			Timestamp: end - i*delta,
			Value:     float64(base + rand.Intn(variant)),
		})
	}

	if order == "ASC" {
		for i, j := 0, len(res)-1; i < j; i, j = i+1, j-1 {
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r Backend) GetStatData(id string, end int64, start int64, limit int64, order string, bucketDuration int64) []backend.StatItem {
	var i int64
	var res []backend.StatItem

	delta := bucketDuration * 1000
	base := 10 + rand.Intn(250)
	variant := 20 + rand.Intn(50)

	for i = 0; i < limit && (end-(i-1)*delta) > start; i++ {
		value := float64(base + rand.Intn(variant))
		res = append(res, backend.StatItem{
			Start:   end - (i-1)*delta,
			End:     end - i*delta,
			Empty:   false,
			Samples: 1,
			Min:     value,
			Max:     value,
			Avg:     value,
			Median:  value,
			Sum:     value,
		})
	}

	if order == "ASC" {
		for i, j := 0, len(res)-1; i < j; i, j = i+1, j-1 {
			res[i], res[j] = res[j], res[i]
		}
	}

	return res
}

func (r Backend) PostRawData(id string, t int64, v float64) bool {
	return false
}

func (r Backend) PutTags(id string, tags map[string]string) bool {
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"text/template"
)

var binaryPath = flag.String("bin", "", "binary to be included into the package")

type Build struct {
	appName    string
	version    string
	output     string
	binaryPath string
}

func main() {
	flag.Parse()

	if *binaryPath == "" {
		fmt.Println("please specify application binary with --bin flag")
		os.Exit(1)
	}

	if !fileExist(*binaryPath) {
		fmt.Printf("specified binary doesn't exist: %s\n", *binaryPath)
		os.Exit(1)
	}

	// use binary name as appName
	appName := filepath.Base(*binaryPath)

	build := &Build{
		appName:    appName,
		version:    "0.0.1",
		binaryPath: *binaryPath,
	}

	build.do()
}

func (b *Build) do() {
	switch runtime.GOOS {
	case "darwin":
		b.darwin()
	default:
		fmt.Printf("not supported os: %s.\n", runtime.GOOS)
	}
}

// darwin is building a new .pkg installer for darwin based OS'es.
func (b *Build) darwin() {
	version := b.version
	if b.output == "" {
		b.output = fmt.Sprintf("koding-%s", b.appName)
	}

	scriptDir := "./darwin/scripts"
	installRoot := "./root"
	installRootUsr := filepath.Join(installRoot, "/usr/local/bin")

	os.MkdirAll(installRootUsr, 0755)
	err := copyFile(b.binaryPath, installRootUsr+"/"+b.appName)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	tempDest, err := ioutil.TempDir("", "tempDest")
	if err != nil {
		return
	}
	defer os.RemoveAll(tempDest)

	b.createScripts(scriptDir)
	b.createLaunchAgent(installRoot)

	cmdPkg := exec.Command("pkgbuild",
		"--identifier", fmt.Sprintf("com.koding.kite.%s.pkg", b.appName),
		"--version", version,
		"--scripts", scriptDir,
		"--root", installRoot,
		"--install-location", "/",
		fmt.Sprintf("%s/com.koding.kite.%s.pkg", tempDest, b.appName),
		// used for next step, also set up for distribution.xml
	)

	res, err := cmdPkg.CombinedOutput()
	if err != nil {
		fmt.Println("res, err", string(res), err)
		return
	}

	distributionFile := "./darwin/Distribution.xml"
	resources := "./darwin/Resources"
	targetFile := b.output + ".pkg"

	b.createDistribution(distributionFile)

	cmdBuild := exec.Command("productbuild",
		"--distribution", distributionFile,
		"--resources", resources,
		"--package-path", tempDest,
		targetFile,
	)

	res, err = cmdBuild.CombinedOutput()
	if err != nil {
		fmt.Println("res, err", string(res), err)
		return
	}

	fmt.Println("everything is ok")

}

func (b *Build) createLaunchAgent(rootDir string) {
	launchDir := fmt.Sprintf("%s/Library/LaunchAgents/", rootDir)
	os.MkdirAll(launchDir, 0700)

	launchFile := fmt.Sprintf("%s/com.koding.kite.%s.plist", launchDir, b.appName)

	lFile, err := os.Create(launchFile)
	if err != nil {
		log.Fatalln(err)
	}

	t := template.Must(template.New("launchAgent").Parse(launchAgent))
	t.Execute(lFile, b.appName)

}

func (b *Build) createDistribution(file string) {
	distFile, err := os.Create(file)
	if err != nil {
		log.Fatalln(err)
	}

	t := template.Must(template.New("distribution").Parse(distribution))
	t.Execute(distFile, b.appName)

}

func (b *Build) createScripts(scriptDir string) {
	os.MkdirAll(scriptDir, 0700) // does return nil if exists

	postInstallFile, err := os.Create(scriptDir + "/postInstall")
	if err != nil {
		log.Fatalln(err)
	}
	postInstallFile.Chmod(0755)

	preInstallFile, err := os.Create(scriptDir + "/preInstall")
	if err != nil {
		log.Fatalln(err)
	}
	preInstallFile.Chmod(0755)

	t := template.Must(template.New("postInstall").Parse(postInstall))
	t.Execute(postInstallFile, b.appName)

	t = template.Must(template.New("preInstall").Parse(preInstall))
	t.Execute(preInstallFile, b.appName)
}

func fileExist(dir string) bool {
	var err error
	_, err = os.Stat(dir)
	if err == nil {
		return true // file exist
	}

	if os.IsNotExist(err) {
		return false // file does not exist
	}

	panic(err) // permission errors or something else bad
}

func copyFile(src, dst string) error {
	sf, err := os.Open(src)
	if err != nil {
		return err
	}
	defer sf.Close()

	fi, err := sf.Stat()
	if err != nil {
		return err
	}

	if fi.IsDir() {
		return errors.New("src is a directory, please provide a file")
	}

	df, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, fi.Mode())
	if err != nil {
		return err
	}
	defer df.Close()

	if _, err := io.Copy(df, sf); err != nil {
		return err
	}

	return nil
}
<commit_msg>kd/build: fix cleaning old builds<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"text/template"
)

var binaryPath = flag.String("bin", "", "binary to be included into the package")

type Build struct {
	appName    string
	version    string
	output     string
	binaryPath string
}

func main() {
	flag.Parse()

	if *binaryPath == "" {
		fmt.Println("please specify application binary with --bin flag")
		os.Exit(1)
	}

	if !fileExist(*binaryPath) {
		fmt.Printf("specified binary doesn't exist: %s\n", *binaryPath)
		os.Exit(1)
	}

	// use binary name as appName
	appName := filepath.Base(*binaryPath)

	build := &Build{
		appName:    appName,
		version:    "0.0.1",
		binaryPath: *binaryPath,
	}

	build.do()
}

func (b *Build) do() {
	switch runtime.GOOS {
	case "darwin":
		b.darwin()
	default:
		fmt.Printf("not supported os: %s.\n", runtime.GOOS)
	}
}

// darwin is building a new .pkg installer for darwin based OS'es.
func (b *Build) darwin() {
	version := b.version
	if b.output == "" {
		b.output = fmt.Sprintf("koding-%s", b.appName)
	}

	scriptDir := "./darwin/scripts"
	installRoot := "./root" // TODO REMOVE

	os.RemoveAll(installRoot) // clean up old build before we continue

	installRootUsr := filepath.Join(installRoot, "/usr/local/bin")

	os.MkdirAll(installRootUsr, 0755)
	err = copyFile(b.binaryPath, installRootUsr+"/"+b.appName)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	tempDest, err := ioutil.TempDir("", "tempDest")
	if err != nil {
		return
	}
	defer os.RemoveAll(tempDest)

	b.createScripts(scriptDir)
	b.createLaunchAgent(installRoot)

	cmdPkg := exec.Command("pkgbuild",
		"--identifier", fmt.Sprintf("com.koding.kite.%s.pkg", b.appName),
		"--version", version,
		"--scripts", scriptDir,
		"--root", installRoot,
		"--install-location", "/",
		fmt.Sprintf("%s/com.koding.kite.%s.pkg", tempDest, b.appName),
		// used for next step, also set up for distribution.xml
	)

	res, err := cmdPkg.CombinedOutput()
	if err != nil {
		fmt.Println("res, err", string(res), err)
		return
	}

	distributionFile := "./darwin/Distribution.xml"
	resources := "./darwin/Resources"
	targetFile := b.output + ".pkg"

	b.createDistribution(distributionFile)

	cmdBuild := exec.Command("productbuild",
		"--distribution", distributionFile,
		"--resources", resources,
		"--package-path", tempDest,
		targetFile,
	)

	res, err = cmdBuild.CombinedOutput()
	if err != nil {
		fmt.Println("res, err", string(res), err)
		return
	}

	fmt.Println("everything is ok")

}

func (b *Build) createLaunchAgent(rootDir string) {
	launchDir := fmt.Sprintf("%s/Library/LaunchAgents/", rootDir)
	os.MkdirAll(launchDir, 0700)

	launchFile := fmt.Sprintf("%s/com.koding.kite.%s.plist", launchDir, b.appName)

	lFile, err := os.Create(launchFile)
	if err != nil {
		log.Fatalln(err)
	}

	t := template.Must(template.New("launchAgent").Parse(launchAgent))
	t.Execute(lFile, b.appName)

}

func (b *Build) createDistribution(file string) {
	distFile, err := os.Create(file)
	if err != nil {
		log.Fatalln(err)
	}

	t := template.Must(template.New("distribution").Parse(distribution))
	t.Execute(distFile, b.appName)

}

func (b *Build) createScripts(scriptDir string) {
	os.MkdirAll(scriptDir, 0700) // does return nil if exists

	postInstallFile, err := os.Create(scriptDir + "/postInstall")
	if err != nil {
		log.Fatalln(err)
	}
	postInstallFile.Chmod(0755)

	preInstallFile, err := os.Create(scriptDir + "/preInstall")
	if err != nil {
		log.Fatalln(err)
	}
	preInstallFile.Chmod(0755)

	t := template.Must(template.New("postInstall").Parse(postInstall))
	t.Execute(postInstallFile, b.appName)

	t = template.Must(template.New("preInstall").Parse(preInstall))
	t.Execute(preInstallFile, b.appName)
}

func fileExist(dir string) bool {
	var err error
	_, err = os.Stat(dir)
	if err == nil {
		return true // file exist
	}

	if os.IsNotExist(err) {
		return false // file does not exist
	}

	panic(err) // permission errors or something else bad
}

func copyFile(src, dst string) error {
	sf, err := os.Open(src)
	if err != nil {
		return err
	}
	defer sf.Close()

	fi, err := sf.Stat()
	if err != nil {
		return err
	}

	if fi.IsDir() {
		return errors.New("src is a directory, please provide a file")
	}

	df, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, fi.Mode())
	if err != nil {
		return err
	}
	defer df.Close()

	if _, err := io.Copy(df, sf); err != nil {
		return err
	}

	return nil
}
<|endoftext|>