<commit_before>// Copyright (C) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 main

import (
	"context"
	"flag"
	"fmt"
	"image"
	"image/png"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/app/flags"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/client"
	"github.com/google/gapid/gapis/service"
	"github.com/google/gapid/gapis/service/path"

	img "github.com/google/gapid/core/image"
)

type screenshotVerb struct{ ScreenshotFlags }

func init() {
	verb := &screenshotVerb{
		ScreenshotFlags{
			At:    []flags.U64Slice{},
			Frame: []int{},
			Out:   "screenshot.png",
			NoOpt: false,
		},
	}

	app.AddVerb(&app.Verb{
		Name:      "screenshot",
		ShortHelp: "Produce a screenshot at a particular command from a .gfxtrace file",
		Action:    verb,
	})
}

func (verb *screenshotVerb) Run(ctx context.Context, flags flag.FlagSet) error {
	if flags.NArg() != 1 {
		app.Usage(ctx, "Exactly one gfx trace file expected, got %d", flags.NArg())
		return nil
	}

	client, capture, err := getGapisAndLoadCapture(ctx, verb.Gapis, verb.Gapir, flags.Arg(0), verb.CaptureFileFlags)
	if err != nil {
		return err
	}
	defer client.Close()

	device, err := getDevice(ctx, client, capture, verb.Gapir)
	if err != nil {
		return err
	}

	var commands []*path.Command
	if len(verb.At) > 0 {
		for _, at := range verb.At {
			commands = append(commands, capture.Command(at[0], at[1:]...))
		}
	} else if verb.ExecutedDraws > 0 {
		commands, err = verb.executedDrawCommands(ctx, capture, client, verb.ExecutedDraws)
		if err != nil {
			return err
		}
	} else {
		commands, err = verb.frameCommands(ctx, capture, client)
		if err != nil {
			return err
		}
	}

	// Submit requests in parallel, so that gapis will batch them.
	var wg sync.WaitGroup
	c := make(chan error)
	multi := len(commands) > 1
	for idx, command := range commands {
		wg.Add(1)
		go func(idx int, command *path.Command) {
			defer wg.Done()

			frame, err := verb.getSingleFrame(ctx, command, device, client)
			if err == nil {
				err = verb.writeSingleFrame(flipImg(frame), formatOut(verb.Out, idx, multi))
			}
			c <- err
		}(idx, command)
	}
	go func() {
		wg.Wait()
		close(c)
	}()
	for err := range c {
		if err != nil {
			return err
		}
	}
	return nil
}

func formatOut(out string, idx int, multi bool) string {
	if strings.Contains(out, "%d") {
		return fmt.Sprintf(out, idx)
	} else if !multi {
		return out
	} else if p := strings.LastIndex(out, "."); p != -1 {
		return fmt.Sprintf("%s_%d%s", out[:p], idx, out[p:])
	}
	return fmt.Sprintf("%s_%d", out, idx)

}

func (verb *screenshotVerb) writeSingleFrame(frame image.Image, fn string) error {
	out, err := os.Create(fn)
	if err != nil {
		return err
	}
	defer out.Close()
	return png.Encode(out, frame)
}

func (verb *screenshotVerb) getSingleFrame(ctx context.Context, cmd *path.Command, device *path.Device, client service.Service) (*image.NRGBA, error) {
	ctx = log.V{"cmd": cmd.Indices}.Bind(ctx)
	settings := &path.RenderSettings{
		MaxWidth:                  uint32(0xFFFFFFFF),
		MaxHeight:                 uint32(0xFFFFFFFF),
		DisableReplayOptimization: verb.NoOpt,
		DisplayToSurface:          verb.DisplayToSurface,
	}
	if verb.Overdraw {
		settings.DrawMode = path.DrawMode_OVERDRAW
	}

	attachment, err := verb.getAttachment(ctx)
	if err != nil {
		return nil, log.Errf(ctx, err, "Get color attachment failed")
	}
	fbPath := &path.FramebufferAttachment{
		After:          cmd,
		Index:          attachment,
		RenderSettings: settings,
		Hints:          nil,
	}
	iip, err := client.Get(ctx, fbPath.Path(), &path.ResolveConfig{ReplayDevice: device})
	if err != nil {
		return nil, log.Errf(ctx, err, "GetFramebufferAttachment failed")
	}
	iio, err := client.Get(ctx, iip.(*service.FramebufferAttachment).GetImageInfo().Path(), nil)
	if err != nil {
		return nil, log.Errf(ctx, err, "Get frame image.Info failed")
	}
	ii := iio.(*img.Info)
	dataO, err := client.Get(ctx, path.NewBlob(ii.Bytes.ID()).Path(), nil)
	if err != nil {
		return nil, log.Errf(ctx, err, "Get frame image data failed")
	}
	w, h, data := int(ii.Width), int(ii.Height), dataO.([]byte)

	ctx = log.V{
		"width":  w,
		"height": h,
		"format": ii.Format,
	}.Bind(ctx)
	if ii.Width == 0 || ii.Height == 0 {
		return nil, log.Err(ctx, nil, "Framebuffer has zero dimensions")
	}
	format := ii.Format
	if verb.Overdraw {
		format = img.Gray_U8_NORM
		rescaleBytes(ctx, data, verb.Max.Overdraw)
	}
	data, err = img.Convert(data, w, h, 1, format, img.RGBA_U8_NORM)
	if err != nil {
		return nil, log.Err(ctx, err, "Failed to convert frame to RGBA")
	}
	stride := w * 4
	return &image.NRGBA{
		Rect:   image.Rect(0, 0, w, h),
		Stride: stride,
		Pix:    data,
	}, nil
}

func (verb *screenshotVerb) frameCommands(ctx context.Context, capture *path.Capture, client service.Service) ([]*path.Command, error) {
	filter, err := verb.CommandFilterFlags.commandFilter(ctx, client, capture)
	if err != nil {
		return nil, log.Err(ctx, err, "Couldn't get filter")
	}

	requestEvents := path.Events{
		Capture:     capture,
		LastInFrame: true,
		DrawCalls:   verb.Draws,
		Filter:      filter,
	}

	// Get the end-of-frame and possibly draw call events.
	events, err := getEvents(ctx, client, &requestEvents)
	if err != nil {
		return nil, log.Err(ctx, err, "Couldn't get frame events")
	}

	// Compute an index of frame to event idx.
	frameIdx := map[int]int{}
	lastFrame := 0
	for i, e := range events {
		if e.Kind == service.EventKind_LastInFrame {
			lastFrame++
			frameIdx[lastFrame] = i
		}
	}

	if len(verb.Frame) == 0 {
		verb.Frame = []int{lastFrame}
	}

	var commands []*path.Command
	for _, frame := range verb.Frame {
		last, ok := frameIdx[frame]
		if !ok {
			return nil, fmt.Errorf("Invalid frame number %d (last frame is %d)", frame, lastFrame)
		}

		first := last
		if verb.Draws {
			if frame == 1 {
				first = 0
			} else {
				first = frameIdx[frame-1]
			}
		}
		for idx := first; idx <= last; idx++ {
			commands = append(commands, events[idx].Command)
		}
	}
	return commands, nil
}

func (verb *screenshotVerb) executedDrawCommands(ctx context.Context, capture *path.Capture, client client.Client, maxAmount int) ([]*path.Command, error) {
	filter, err := verb.CommandFilterFlags.commandFilter(ctx, client, capture)
	if err != nil {
		return nil, log.Err(ctx, err, "Couldn't get filter")
	}
	filter.OnlyExecutedDraws = true

	treePath := capture.CommandTree(filter)

	boxedTree, err := client.Get(ctx, treePath.Path(), nil)
	if err != nil {
		return nil, log.Err(ctx, err, "Failed to load the command tree")
	}

	tree := boxedTree.(*service.CommandTree)

	var allDrawCommands []*path.Command
	traverseCommandTree(ctx, client, tree.Root, func(n *service.CommandTreeNode, prefix string) error {
		if n.Group != "" || n.NumChildren > 0 {
			return nil
		}
		allDrawCommands = append(allDrawCommands, n.Commands.First())
		return nil
	}, "", true)

	if len(allDrawCommands) > maxAmount {
		commands := make([]*path.Command, maxAmount)
		// We use a float step so that we have a fairly even distribution when ratio < 2 (e.g. 5/3)
		step := float32(len(allDrawCommands)) / float32(maxAmount)

		for i := 0; i < len(commands); i++ {
			commands[i] = allDrawCommands[int32(float32(i)*step)]
		}

		return commands, nil
	} else {
		return allDrawCommands, nil
	}
}

// rescaleBytes scales the values in `data` from [0, `max`] to [0, 255].  If
// `max <= 0`, then the maximum value found in data is used as `max` instead.
func rescaleBytes(ctx context.Context, data []byte, max int) {
	if max <= 0 {
		for _, b := range data {
			if int(b) > max {
				max = int(b)
			}
		}
	}
	log.I(ctx, "Max overdraw: %v", max)
	if max < 1 {
		max = 1
	}

	for i, b := range data {
		if int(b) >= max {
			data[i] = 255
		} else {
			data[i] = byte(int(data[i]) * 255 / max)
		}
	}
}

func (verb *screenshotVerb) getAttachment(ctx context.Context) (uint32, error) {
	if verb.Attachment == "" {
		return 0, nil
	}

	// TODO: Add-back ability to type "depth" to get the depth attachment
	i, err := strconv.ParseUint(verb.Attachment, 10, 32)
	if err != nil {
		return 0, log.Errf(ctx, nil, "Invalid attachment %v", verb.Attachment)
	}
	return uint32(i), nil
}
<commit_msg>Use the first attachment for screenshot test instead of one with index = 0 (#515)<commit_after>// Copyright (C) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 main

import (
	"context"
	"flag"
	"fmt"
	"image"
	"image/png"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/app/flags"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/gapis/client"
	"github.com/google/gapid/gapis/service"
	"github.com/google/gapid/gapis/service/path"

	img "github.com/google/gapid/core/image"
)

type screenshotVerb struct{ ScreenshotFlags }

func init() {
	verb := &screenshotVerb{
		ScreenshotFlags{
			At:    []flags.U64Slice{},
			Frame: []int{},
			Out:   "screenshot.png",
			NoOpt: false,
		},
	}

	app.AddVerb(&app.Verb{
		Name:      "screenshot",
		ShortHelp: "Produce a screenshot at a particular command from a .gfxtrace file",
		Action:    verb,
	})
}

func (verb *screenshotVerb) Run(ctx context.Context, flags flag.FlagSet) error {
	if flags.NArg() != 1 {
		app.Usage(ctx, "Exactly one gfx trace file expected, got %d", flags.NArg())
		return nil
	}

	client, capture, err := getGapisAndLoadCapture(ctx, verb.Gapis, verb.Gapir, flags.Arg(0), verb.CaptureFileFlags)
	if err != nil {
		return err
	}
	defer client.Close()

	device, err := getDevice(ctx, client, capture, verb.Gapir)
	if err != nil {
		return err
	}

	var commands []*path.Command
	if len(verb.At) > 0 {
		for _, at := range verb.At {
			commands = append(commands, capture.Command(at[0], at[1:]...))
		}
	} else if verb.ExecutedDraws > 0 {
		commands, err = verb.executedDrawCommands(ctx, capture, client, verb.ExecutedDraws)
		if err != nil {
			return err
		}
	} else {
		commands, err = verb.frameCommands(ctx, capture, client)
		if err != nil {
			return err
		}
	}

	// Submit requests in parallel, so that gapis will batch them.
	var wg sync.WaitGroup
	c := make(chan error)
	multi := len(commands) > 1
	for idx, command := range commands {
		wg.Add(1)
		go func(idx int, command *path.Command) {
			defer wg.Done()

			frame, err := verb.getSingleFrame(ctx, command, device, client)
			if err == nil {
				err = verb.writeSingleFrame(flipImg(frame), formatOut(verb.Out, idx, multi))
			}
			c <- err
		}(idx, command)
	}
	go func() {
		wg.Wait()
		close(c)
	}()
	for err := range c {
		if err != nil {
			return err
		}
	}
	return nil
}

func formatOut(out string, idx int, multi bool) string {
	if strings.Contains(out, "%d") {
		return fmt.Sprintf(out, idx)
	} else if !multi {
		return out
	} else if p := strings.LastIndex(out, "."); p != -1 {
		return fmt.Sprintf("%s_%d%s", out[:p], idx, out[p:])
	}
	return fmt.Sprintf("%s_%d", out, idx)

}

func (verb *screenshotVerb) writeSingleFrame(frame image.Image, fn string) error {
	out, err := os.Create(fn)
	if err != nil {
		return err
	}
	defer out.Close()
	return png.Encode(out, frame)
}

func (verb *screenshotVerb) getSingleFrame(ctx context.Context, cmd *path.Command, device *path.Device, client service.Service) (*image.NRGBA, error) {
	ctx = log.V{"cmd": cmd.Indices}.Bind(ctx)
	settings := &path.RenderSettings{
		MaxWidth:                  uint32(0xFFFFFFFF),
		MaxHeight:                 uint32(0xFFFFFFFF),
		DisableReplayOptimization: verb.NoOpt,
		DisplayToSurface:          verb.DisplayToSurface,
	}
	if verb.Overdraw {
		settings.DrawMode = path.DrawMode_OVERDRAW
	}

	attachment, err := verb.getAttachment(ctx, cmd, device, client)
	if err != nil {
		return nil, log.Errf(ctx, err, "Get color attachment failed")
	}
	fbPath := &path.FramebufferAttachment{
		After:          cmd,
		Index:          attachment,
		RenderSettings: settings,
		Hints:          nil,
	}
	iip, err := client.Get(ctx, fbPath.Path(), &path.ResolveConfig{ReplayDevice: device})
	if err != nil {
		return nil, log.Errf(ctx, err, "GetFramebufferAttachment failed")
	}
	iio, err := client.Get(ctx, iip.(*service.FramebufferAttachment).GetImageInfo().Path(), nil)
	if err != nil {
		return nil, log.Errf(ctx, err, "Get frame image.Info failed")
	}
	ii := iio.(*img.Info)
	dataO, err := client.Get(ctx, path.NewBlob(ii.Bytes.ID()).Path(), nil)
	if err != nil {
		return nil, log.Errf(ctx, err, "Get frame image data failed")
	}
	w, h, data := int(ii.Width), int(ii.Height), dataO.([]byte)

	ctx = log.V{
		"width":  w,
		"height": h,
		"format": ii.Format,
	}.Bind(ctx)
	if ii.Width == 0 || ii.Height == 0 {
		return nil, log.Err(ctx, nil, "Framebuffer has zero dimensions")
	}
	format := ii.Format
	if verb.Overdraw {
		format = img.Gray_U8_NORM
		rescaleBytes(ctx, data, verb.Max.Overdraw)
	}
	data, err = img.Convert(data, w, h, 1, format, img.RGBA_U8_NORM)
	if err != nil {
		return nil, log.Err(ctx, err, "Failed to convert frame to RGBA")
	}
	stride := w * 4
	return &image.NRGBA{
		Rect:   image.Rect(0, 0, w, h),
		Stride: stride,
		Pix:    data,
	}, nil
}

func (verb *screenshotVerb) frameCommands(ctx context.Context, capture *path.Capture, client service.Service) ([]*path.Command, error) {
	filter, err := verb.CommandFilterFlags.commandFilter(ctx, client, capture)
	if err != nil {
		return nil, log.Err(ctx, err, "Couldn't get filter")
	}

	requestEvents := path.Events{
		Capture:     capture,
		LastInFrame: true,
		DrawCalls:   verb.Draws,
		Filter:      filter,
	}

	// Get the end-of-frame and possibly draw call events.
	events, err := getEvents(ctx, client, &requestEvents)
	if err != nil {
		return nil, log.Err(ctx, err, "Couldn't get frame events")
	}

	// Compute an index of frame to event idx.
	frameIdx := map[int]int{}
	lastFrame := 0
	for i, e := range events {
		if e.Kind == service.EventKind_LastInFrame {
			lastFrame++
			frameIdx[lastFrame] = i
		}
	}

	if len(verb.Frame) == 0 {
		verb.Frame = []int{lastFrame}
	}

	var commands []*path.Command
	for _, frame := range verb.Frame {
		last, ok := frameIdx[frame]
		if !ok {
			return nil, fmt.Errorf("Invalid frame number %d (last frame is %d)", frame, lastFrame)
		}

		first := last
		if verb.Draws {
			if frame == 1 {
				first = 0
			} else {
				first = frameIdx[frame-1]
			}
		}
		for idx := first; idx <= last; idx++ {
			commands = append(commands, events[idx].Command)
		}
	}
	return commands, nil
}

func (verb *screenshotVerb) executedDrawCommands(ctx context.Context, capture *path.Capture, client client.Client, maxAmount int) ([]*path.Command, error) {
	filter, err := verb.CommandFilterFlags.commandFilter(ctx, client, capture)
	if err != nil {
		return nil, log.Err(ctx, err, "Couldn't get filter")
	}
	filter.OnlyExecutedDraws = true

	treePath := capture.CommandTree(filter)

	boxedTree, err := client.Get(ctx, treePath.Path(), nil)
	if err != nil {
		return nil, log.Err(ctx, err, "Failed to load the command tree")
	}

	tree := boxedTree.(*service.CommandTree)

	var allDrawCommands []*path.Command
	traverseCommandTree(ctx, client, tree.Root, func(n *service.CommandTreeNode, prefix string) error {
		if n.Group != "" || n.NumChildren > 0 {
			return nil
		}
		allDrawCommands = append(allDrawCommands, n.Commands.First())
		return nil
	}, "", true)

	if len(allDrawCommands) > maxAmount {
		commands := make([]*path.Command, maxAmount)
		// We use a float step so that we have a fairly even distribution when ratio < 2 (e.g. 5/3)
		step := float32(len(allDrawCommands)) / float32(maxAmount)

		for i := 0; i < len(commands); i++ {
			commands[i] = allDrawCommands[int32(float32(i)*step)]
		}

		return commands, nil
	} else {
		return allDrawCommands, nil
	}
}

// rescaleBytes scales the values in `data` from [0, `max`] to [0, 255].  If
// `max <= 0`, then the maximum value found in data is used as `max` instead.
func rescaleBytes(ctx context.Context, data []byte, max int) {
	if max <= 0 {
		for _, b := range data {
			if int(b) > max {
				max = int(b)
			}
		}
	}
	log.I(ctx, "Max overdraw: %v", max)
	if max < 1 {
		max = 1
	}

	for i, b := range data {
		if int(b) >= max {
			data[i] = 255
		} else {
			data[i] = byte(int(data[i]) * 255 / max)
		}
	}
}

func (verb *screenshotVerb) getAttachment(ctx context.Context, cmd *path.Command, device *path.Device, client service.Service) (uint32, error) {
	if verb.Attachment == "" {
		fbsPath := &path.FramebufferAttachments{
			After: cmd,
		}
		fbs, err := client.Get(ctx, fbsPath.Path(), &path.ResolveConfig{ReplayDevice: device})
		if err != nil {
			return 0, log.Errf(ctx, err, "GetFramebufferAttachments failed")
		}
		attachments := fbs.(*service.FramebufferAttachments).GetAttachments()
		if len(attachments) == 0 {
			return 0, log.Errf(ctx, err, "No Framebuffer Attachments")
		}

		return attachments[0].GetIndex(), nil
	}

	// TODO: Add-back ability to type "depth" to get the depth attachment
	i, err := strconv.ParseUint(verb.Attachment, 10, 32)
	if err != nil {
		return 0, log.Errf(ctx, nil, "Invalid attachment %v", verb.Attachment)
	}
	return uint32(i), nil
}
<|endoftext|>


<commit_before>package promtail

import (
	"context"
	"os"
	"strings"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/discovery"
	sd_config "github.com/prometheus/prometheus/discovery/config"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	"github.com/prometheus/prometheus/relabel"
)

const (
	pathLabel = "__path__"
	hostLabel = "__host__"
)

// NewTargetFunc is the type of function that creates a new target.
type NewTargetFunc func(path string, labels model.LabelSet) (*Target, error)

// TargetManager manages a set of targets.
type TargetManager struct {
	syncers map[string]*syncer
	manager *discovery.Manager

	log  log.Logger
	quit context.CancelFunc
}

func hostname() (string, error) {
	hostname := os.Getenv("HOSTNAME")
	if hostname != "" {
		return hostname, nil
	}

	return os.Hostname()
}

// NewTargetManager creates a new TargetManager.
func NewTargetManager(
	logger log.Logger,
	scrapeConfig []ScrapeConfig,
	fn NewTargetFunc,
) (*TargetManager, error) {
	ctx, quit := context.WithCancel(context.Background())
	tm := &TargetManager{
		log:     logger,
		syncers: map[string]*syncer{},
		manager: discovery.NewManager(ctx, log.With(logger, "component", "discovery")),
		quit:    quit,
	}

	hostname, err := hostname()
	if err != nil {
		return nil, err
	}

	config := map[string]sd_config.ServiceDiscoveryConfig{}
	for _, cfg := range scrapeConfig {
		s := &syncer{
			log:           logger,
			newTarget:     fn,
			relabelConfig: cfg.RelabelConfigs,
			targets:       map[string]*Target{},
			hostname:      hostname,
		}
		tm.syncers[cfg.JobName] = s
		config[cfg.JobName] = cfg.ServiceDiscoveryConfig
	}

	go tm.run()
	go tm.manager.Run()

	return tm, tm.manager.ApplyConfig(config)
}

func (tm *TargetManager) run() {
	for targetGoups := range tm.manager.SyncCh() {
		for jobName, groups := range targetGoups {
			tm.syncers[jobName].Sync(groups)
		}
	}
}

// Stop the TargetManager.
func (tm *TargetManager) Stop() {
	tm.quit()

	for _, s := range tm.syncers {
		s.stop()
	}
}

type syncer struct {
	log           log.Logger
	newTarget     NewTargetFunc
	targets       map[string]*Target
	relabelConfig []*config.RelabelConfig
	hostname      string
}

func (s *syncer) Sync(groups []*targetgroup.Group) {
	targets := map[string]struct{}{}

	for _, group := range groups {
		for _, t := range group.Targets {
			level.Debug(s.log).Log("msg", "new target", "labels", t)

			labels := group.Labels.Merge(t)
			labels = relabel.Process(labels, s.relabelConfig...)
			// Drop empty targets (drop in relabeling).
			if labels == nil {
				level.Debug(s.log).Log("msg", "dropping target, no labels")
				continue
			}

			host, ok := labels[hostLabel]
			if ok && string(host) != s.hostname {
				level.Debug(s.log).Log("msg", "ignoring target, wrong host", "labels", labels.String(), "hostname", s.hostname)
				continue
			}

			path, ok := labels[pathLabel]
			if !ok {
				level.Info(s.log).Log("msg", "no path for target", "labels", labels.String())
				continue
			}

			for k := range labels {
				if strings.HasPrefix(string(k), "__") {
					delete(labels, k)
				}
			}

			key := labels.String()
			targets[key] = struct{}{}
			if _, ok := s.targets[key]; ok {
				level.Debug(s.log).Log("msg", "ignoring target, already exists", "labels", labels.String())
				continue
			}

			level.Info(s.log).Log("msg", "Adding target", "key", key)
			t, err := s.newTarget(string(path), labels)
			if err != nil {
				level.Error(s.log).Log("msg", "Failed to create target", "key", key, "error", err)
				continue
			}

			s.targets[key] = t
		}
	}

	for key, target := range s.targets {
		if _, ok := targets[key]; !ok {
			level.Info(s.log).Log("msg", "Removing target", "key", key)
			target.Stop()
			delete(s.targets, key)
		}
	}
}

func (s *syncer) stop() {
	for key, target := range s.targets {
		level.Info(s.log).Log("msg", "Removing target", "key", key)
		target.Stop()
		delete(s.targets, key)
	}
}
<commit_msg>Don't ignore the error<commit_after>package promtail

import (
	"context"
	"os"
	"strings"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/grafana/tempo/pkg/helpers"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/discovery"
	sd_config "github.com/prometheus/prometheus/discovery/config"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	"github.com/prometheus/prometheus/relabel"
)

const (
	pathLabel = "__path__"
	hostLabel = "__host__"
)

// NewTargetFunc is the type of function that creates a new target.
type NewTargetFunc func(path string, labels model.LabelSet) (*Target, error)

// TargetManager manages a set of targets.
type TargetManager struct {
	syncers map[string]*syncer
	manager *discovery.Manager

	log  log.Logger
	quit context.CancelFunc
}

func hostname() (string, error) {
	hostname := os.Getenv("HOSTNAME")
	if hostname != "" {
		return hostname, nil
	}

	return os.Hostname()
}

// NewTargetManager creates a new TargetManager.
func NewTargetManager(
	logger log.Logger,
	scrapeConfig []ScrapeConfig,
	fn NewTargetFunc,
) (*TargetManager, error) {
	ctx, quit := context.WithCancel(context.Background())
	tm := &TargetManager{
		log:     logger,
		syncers: map[string]*syncer{},
		manager: discovery.NewManager(ctx, log.With(logger, "component", "discovery")),
		quit:    quit,
	}

	hostname, err := hostname()
	if err != nil {
		return nil, err
	}

	config := map[string]sd_config.ServiceDiscoveryConfig{}
	for _, cfg := range scrapeConfig {
		s := &syncer{
			log:           logger,
			newTarget:     fn,
			relabelConfig: cfg.RelabelConfigs,
			targets:       map[string]*Target{},
			hostname:      hostname,
		}
		tm.syncers[cfg.JobName] = s
		config[cfg.JobName] = cfg.ServiceDiscoveryConfig
	}

	go tm.run()
	go helpers.LogError("running target manager", tm.manager.Run)

	return tm, tm.manager.ApplyConfig(config)
}

func (tm *TargetManager) run() {
	for targetGoups := range tm.manager.SyncCh() {
		for jobName, groups := range targetGoups {
			tm.syncers[jobName].Sync(groups)
		}
	}
}

// Stop the TargetManager.
func (tm *TargetManager) Stop() {
	tm.quit()

	for _, s := range tm.syncers {
		s.stop()
	}
}

type syncer struct {
	log           log.Logger
	newTarget     NewTargetFunc
	targets       map[string]*Target
	relabelConfig []*config.RelabelConfig
	hostname      string
}

func (s *syncer) Sync(groups []*targetgroup.Group) {
	targets := map[string]struct{}{}

	for _, group := range groups {
		for _, t := range group.Targets {
			level.Debug(s.log).Log("msg", "new target", "labels", t)

			labels := group.Labels.Merge(t)
			labels = relabel.Process(labels, s.relabelConfig...)
			// Drop empty targets (drop in relabeling).
			if labels == nil {
				level.Debug(s.log).Log("msg", "dropping target, no labels")
				continue
			}

			host, ok := labels[hostLabel]
			if ok && string(host) != s.hostname {
				level.Debug(s.log).Log("msg", "ignoring target, wrong host", "labels", labels.String(), "hostname", s.hostname)
				continue
			}

			path, ok := labels[pathLabel]
			if !ok {
				level.Info(s.log).Log("msg", "no path for target", "labels", labels.String())
				continue
			}

			for k := range labels {
				if strings.HasPrefix(string(k), "__") {
					delete(labels, k)
				}
			}

			key := labels.String()
			targets[key] = struct{}{}
			if _, ok := s.targets[key]; ok {
				level.Debug(s.log).Log("msg", "ignoring target, already exists", "labels", labels.String())
				continue
			}

			level.Info(s.log).Log("msg", "Adding target", "key", key)
			t, err := s.newTarget(string(path), labels)
			if err != nil {
				level.Error(s.log).Log("msg", "Failed to create target", "key", key, "error", err)
				continue
			}

			s.targets[key] = t
		}
	}

	for key, target := range s.targets {
		if _, ok := targets[key]; !ok {
			level.Info(s.log).Log("msg", "Removing target", "key", key)
			target.Stop()
			delete(s.targets, key)
		}
	}
}

func (s *syncer) stop() {
	for key, target := range s.targets {
		level.Info(s.log).Log("msg", "Removing target", "key", key)
		target.Stop()
		delete(s.targets, key)
	}
}
<|endoftext|>


<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 adapter

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	cloudevents "github.com/cloudevents/sdk-go/v2"
	"github.com/google/uuid"
	"go.uber.org/zap"
	"gopkg.in/go-playground/webhooks.v5/gitlab"
	"knative.dev/eventing/pkg/adapter/v2"
	"knative.dev/pkg/logging"
)

const (
	glHeaderEvent = "X-Gitlab-Event"
)

type envConfig struct {
	adapter.EnvConfig

	// Environment variable containing Gitlab secret token
	EnvSecret string `envconfig:"GITLAB_SECRET_TOKEN" required:"true"`
	// Port to listen incoming connections
	Port string `envconfig:"PORT" default:"8080"`
}

// gitLabReceiveAdapter converts incoming GitLab webhook events to
// CloudEvents and then sends them to the specified Sink
type gitLabReceiveAdapter struct {
	logger      *zap.SugaredLogger
	client      cloudevents.Client
	secretToken string
	port        string
}

// NewEnvConfig function reads env variables defined in envConfig structure and
// returns accessor interface
func NewEnvConfig() adapter.EnvConfigAccessor {
	return &envConfig{}
}

// NewAdapter returns the instance of gitLabReceiveAdapter that implements adapter.Adapter interface
func NewAdapter(ctx context.Context, processed adapter.EnvConfigAccessor, ceClient cloudevents.Client) adapter.Adapter {
	logger := logging.FromContext(ctx)
	env := processed.(*envConfig)

	return &gitLabReceiveAdapter{
		logger:      logger,
		client:      ceClient,
		port:        env.Port,
		secretToken: env.EnvSecret,
	}
}

// Start implements adapter.Adapter
func (ra *gitLabReceiveAdapter) Start(ctx context.Context) error {
	return ra.start(ctx.Done())
}

func (ra *gitLabReceiveAdapter) start(stopCh <-chan struct{}) error {
	hook, err := gitlab.New(gitlab.Options.Secret(ra.secretToken))
	if err != nil {
		return fmt.Errorf("cannot create gitlab hook: %v", err)
	}

	server := &http.Server{
		Addr:    ":" + ra.port,
		Handler: ra.newRouter(hook),
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		gracefulShutdown(server, ra.logger, stopCh)
	}()

	ra.logger.Infof("Server is ready to handle requests at %s", server.Addr)
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		return fmt.Errorf("could not listen on %s: %v", server.Addr, err)
	}

	wg.Wait()
	ra.logger.Infof("Server stopped")
	return nil
}

func gracefulShutdown(server *http.Server, logger *zap.SugaredLogger, stopCh <-chan struct{}) {
	<-stopCh
	logger.Info("Server is shutting down...")

	// Try to graceful shutdown
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	server.SetKeepAlivesEnabled(false)
	if err := server.Shutdown(ctx); err != nil {
		logger.Fatalf("Could not gracefully shutdown the server: %v", err)
	}
}

func (ra *gitLabReceiveAdapter) newRouter(hook *gitlab.Webhook) *http.ServeMux {
	router := http.NewServeMux()
	router.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		payload, err := hook.Parse(r,
			gitlab.PushEvents,
			gitlab.TagEvents,
			gitlab.IssuesEvents,
			gitlab.ConfidentialIssuesEvents,
			gitlab.CommentEvents,
			gitlab.MergeRequestEvents,
			gitlab.WikiPageEvents,
			gitlab.PipelineEvents,
			gitlab.BuildEvents,
		)
		if err != nil {
			if err == gitlab.ErrEventNotFound {
				w.Write([]byte("event not registered"))
				return
			}
			ra.logger.Errorf("hook parser error: %v", err)
			w.WriteHeader(400)
			w.Write([]byte(err.Error()))
			return
		}
		err = ra.handleEvent(payload, r.Header)
		if err != nil {
			ra.logger.Errorf("event handler error: %v", err)
			w.WriteHeader(400)
			w.Write([]byte(err.Error()))
			return
		}
		ra.logger.Infof("event processed")
		w.WriteHeader(202)
		w.Write([]byte("accepted"))
	})
	return router
}

func (ra *gitLabReceiveAdapter) handleEvent(payload interface{}, header http.Header) error {
	gitLabEventType := header.Get(glHeaderEvent)
	if gitLabEventType == "" {
		return fmt.Errorf("%q header is not set", glHeaderEvent)
	}
	extensions := map[string]interface{}{
		glHeaderEvent: gitLabEventType,
	}

	uuid, err := uuid.NewRandom()
	if err != nil {
		return fmt.Errorf("can't generate event ID: %v", err)
	}
	eventID := uuid.String()

	cloudEventType := fmt.Sprintf("%s.%s", "dev.knative.sources.gitlabsource", gitLabEventType)
	source := sourceFromGitLabEvent(gitlab.Event(gitLabEventType), payload)

	return ra.postMessage(payload, source, cloudEventType, eventID, extensions)
}

func (ra *gitLabReceiveAdapter) postMessage(payload interface{}, source, eventType, eventID string,
	extensions map[string]interface{}) error {
	event := cloudevents.NewEvent(cloudevents.VersionV03)
	event.SetID(eventID)
	event.SetType(eventType)
	event.SetSource(source)
	err := event.SetData(cloudevents.ApplicationJSON, payload)
	if err != nil {
		return fmt.Errorf("failed to marshal event data: %w", err)
	}

	result := ra.client.Send(context.Background(), event)
	if !cloudevents.IsACK(result) {
		return result
	}
	return nil
}

func sourceFromGitLabEvent(gitLabEvent gitlab.Event, payload interface{}) string {
	switch gitLabEvent {
	case gitlab.PushEvents:
		pe := payload.(gitlab.PushEventPayload)
		return pe.Project.HTTPURL
	case gitlab.TagEvents:
		te := payload.(gitlab.TagEventPayload)
		return te.Project.HTTPURL
	case gitlab.IssuesEvents:
		ie := payload.(gitlab.IssueEventPayload)
		return ie.ObjectAttributes.URL
	case gitlab.ConfidentialIssuesEvents:
		cie := payload.(gitlab.ConfidentialIssueEventPayload)
		return cie.ObjectAttributes.URL
	case gitlab.CommentEvents:
		ce := payload.(gitlab.CommentEventPayload)
		return ce.ObjectAttributes.URL
	case gitlab.MergeRequestEvents:
		mre := payload.(gitlab.MergeRequestEventPayload)
		return mre.ObjectAttributes.URL
	case gitlab.WikiPageEvents:
		wpe := payload.(gitlab.WikiPageEventPayload)
		return wpe.ObjectAttributes.URL
	case gitlab.PipelineEvents:
		pe := payload.(gitlab.PipelineEventPayload)
		return pe.ObjectAttributes.URL
	case gitlab.BuildEvents:
		be := payload.(gitlab.BuildEventPayload)
		return be.Repository.Homepage
	}
	return ""
}
<commit_msg>Use CloudEvents 1.0 in GitLabSource (#1335)<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 adapter

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	cloudevents "github.com/cloudevents/sdk-go/v2"
	"github.com/google/uuid"
	"go.uber.org/zap"
	"gopkg.in/go-playground/webhooks.v5/gitlab"
	"knative.dev/eventing/pkg/adapter/v2"
	"knative.dev/pkg/logging"
)

const (
	glHeaderEvent = "X-Gitlab-Event"
)

type envConfig struct {
	adapter.EnvConfig

	// Environment variable containing Gitlab secret token
	EnvSecret string `envconfig:"GITLAB_SECRET_TOKEN" required:"true"`
	// Port to listen incoming connections
	Port string `envconfig:"PORT" default:"8080"`
}

// gitLabReceiveAdapter converts incoming GitLab webhook events to
// CloudEvents and then sends them to the specified Sink
type gitLabReceiveAdapter struct {
	logger      *zap.SugaredLogger
	client      cloudevents.Client
	secretToken string
	port        string
}

// NewEnvConfig function reads env variables defined in envConfig structure and
// returns accessor interface
func NewEnvConfig() adapter.EnvConfigAccessor {
	return &envConfig{}
}

// NewAdapter returns the instance of gitLabReceiveAdapter that implements adapter.Adapter interface
func NewAdapter(ctx context.Context, processed adapter.EnvConfigAccessor, ceClient cloudevents.Client) adapter.Adapter {
	logger := logging.FromContext(ctx)
	env := processed.(*envConfig)

	return &gitLabReceiveAdapter{
		logger:      logger,
		client:      ceClient,
		port:        env.Port,
		secretToken: env.EnvSecret,
	}
}

// Start implements adapter.Adapter
func (ra *gitLabReceiveAdapter) Start(ctx context.Context) error {
	return ra.start(ctx.Done())
}

func (ra *gitLabReceiveAdapter) start(stopCh <-chan struct{}) error {
	hook, err := gitlab.New(gitlab.Options.Secret(ra.secretToken))
	if err != nil {
		return fmt.Errorf("cannot create gitlab hook: %v", err)
	}

	server := &http.Server{
		Addr:    ":" + ra.port,
		Handler: ra.newRouter(hook),
	}

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		gracefulShutdown(server, ra.logger, stopCh)
	}()

	ra.logger.Infof("Server is ready to handle requests at %s", server.Addr)
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		return fmt.Errorf("could not listen on %s: %v", server.Addr, err)
	}

	wg.Wait()
	ra.logger.Infof("Server stopped")
	return nil
}

func gracefulShutdown(server *http.Server, logger *zap.SugaredLogger, stopCh <-chan struct{}) {
	<-stopCh
	logger.Info("Server is shutting down...")

	// Try to graceful shutdown
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	server.SetKeepAlivesEnabled(false)
	if err := server.Shutdown(ctx); err != nil {
		logger.Fatalf("Could not gracefully shutdown the server: %v", err)
	}
}

func (ra *gitLabReceiveAdapter) newRouter(hook *gitlab.Webhook) *http.ServeMux {
	router := http.NewServeMux()
	router.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		payload, err := hook.Parse(r,
			gitlab.PushEvents,
			gitlab.TagEvents,
			gitlab.IssuesEvents,
			gitlab.ConfidentialIssuesEvents,
			gitlab.CommentEvents,
			gitlab.MergeRequestEvents,
			gitlab.WikiPageEvents,
			gitlab.PipelineEvents,
			gitlab.BuildEvents,
		)
		if err != nil {
			if err == gitlab.ErrEventNotFound {
				w.Write([]byte("event not registered"))
				return
			}
			ra.logger.Errorf("hook parser error: %v", err)
			w.WriteHeader(400)
			w.Write([]byte(err.Error()))
			return
		}
		err = ra.handleEvent(payload, r.Header)
		if err != nil {
			ra.logger.Errorf("event handler error: %v", err)
			w.WriteHeader(400)
			w.Write([]byte(err.Error()))
			return
		}
		ra.logger.Infof("event processed")
		w.WriteHeader(202)
		w.Write([]byte("accepted"))
	})
	return router
}

func (ra *gitLabReceiveAdapter) handleEvent(payload interface{}, header http.Header) error {
	gitLabEventType := header.Get(glHeaderEvent)
	if gitLabEventType == "" {
		return fmt.Errorf("%q header is not set", glHeaderEvent)
	}
	extensions := map[string]interface{}{
		glHeaderEvent: gitLabEventType,
	}

	uuid, err := uuid.NewRandom()
	if err != nil {
		return fmt.Errorf("can't generate event ID: %v", err)
	}
	eventID := uuid.String()

	cloudEventType := fmt.Sprintf("%s.%s", "dev.knative.sources.gitlabsource", gitLabEventType)
	source := sourceFromGitLabEvent(gitlab.Event(gitLabEventType), payload)

	return ra.postMessage(payload, source, cloudEventType, eventID, extensions)
}

func (ra *gitLabReceiveAdapter) postMessage(payload interface{}, source, eventType, eventID string,
	extensions map[string]interface{}) error {
	event := cloudevents.NewEvent(cloudevents.VersionV1)
	event.SetID(eventID)
	event.SetType(eventType)
	event.SetSource(source)
	err := event.SetData(cloudevents.ApplicationJSON, payload)
	if err != nil {
		return fmt.Errorf("failed to marshal event data: %w", err)
	}

	result := ra.client.Send(context.Background(), event)
	if !cloudevents.IsACK(result) {
		return result
	}
	return nil
}

func sourceFromGitLabEvent(gitLabEvent gitlab.Event, payload interface{}) string {
	switch gitLabEvent {
	case gitlab.PushEvents:
		pe := payload.(gitlab.PushEventPayload)
		return pe.Project.HTTPURL
	case gitlab.TagEvents:
		te := payload.(gitlab.TagEventPayload)
		return te.Project.HTTPURL
	case gitlab.IssuesEvents:
		ie := payload.(gitlab.IssueEventPayload)
		return ie.ObjectAttributes.URL
	case gitlab.ConfidentialIssuesEvents:
		cie := payload.(gitlab.ConfidentialIssueEventPayload)
		return cie.ObjectAttributes.URL
	case gitlab.CommentEvents:
		ce := payload.(gitlab.CommentEventPayload)
		return ce.ObjectAttributes.URL
	case gitlab.MergeRequestEvents:
		mre := payload.(gitlab.MergeRequestEventPayload)
		return mre.ObjectAttributes.URL
	case gitlab.WikiPageEvents:
		wpe := payload.(gitlab.WikiPageEventPayload)
		return wpe.ObjectAttributes.URL
	case gitlab.PipelineEvents:
		pe := payload.(gitlab.PipelineEventPayload)
		return pe.ObjectAttributes.URL
	case gitlab.BuildEvents:
		be := payload.(gitlab.BuildEventPayload)
		return be.Repository.Homepage
	}
	return ""
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"math"
	"strings"

	"code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"DROP TABLE IF EXISTS Inventories;",
		"DROP TABLE IF EXISTS TimeList;",
		"CREATE TABLE Inventories (SimID TEXT,ResID INTEGER,AgentID INTEGER,StartTime INTEGER,EndTime INTEGER,StateID INTEGER,Quantity REAL);",
		"CREATE TABLE TimeList AS SELECT DISTINCT Time FROM Transactions;",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimID", "ID", "StateID"),
		query.Index("Compositions", "SimID", "ID", "IsoID"),
		query.Index("Transactions", "SimID", "ResourceID"),
		query.Index("Transactions", "ID"),
		query.Index("ResCreators", "SimID", "ResID"),
		query.Index("Agents", "SimID", "Prototype"),
	}
	postExecStmts = []string{
		query.Index("Inventories", "SimID", "AgentID"),
		query.Index("Inventories", "SimID", "ResID", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ID,TimeCreated,StateID,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverID, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceID = ? AND tr.SimID = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ID,res.TimeCreated,rc.ModelID,res.StateID,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ID = rc.ResID
				  WHERE res.SimID = ? AND rc.SimID = ?;`
)

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating indexes and inventory table...")
	for _, sql := range preExecStmts {
		if err := conn.Exec(sql); err != nil {
			fmt.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating inventory indexes...")
	for _, sql := range postExecStmts {
		if err := conn.Exec(sql); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	StateId   int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sqlite3.Conn
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       string
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sqlite3.Stmt
	dumpStmt    *sqlite3.Stmt
	ownerStmt   *sqlite3.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(conn *sqlite3.Conn, simid string, history chan string) *Context {
	return &Context{
		Conn:  conn,
		Simid: simid,
	}
}

func (c *Context) init() {
	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// create temp res table without simid
	fmt.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + strings.Replace(c.Simid, "-", "_", -1)
	err := c.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql := "CREATE TABLE " + c.tmpResTbl + " AS SELECT ID,TimeCreated,Parent1,Parent2,StateID,Quantity FROM Resources WHERE SimID = ?;"
	err = c.Exec(sql, c.Simid)
	panicif(err)

	fmt.Println("Indexing temporary resource table...")
	err = c.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	err = c.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	fmt.Printf("--- Building inventories for simid %v ---\n", c.Simid)
	c.init()

	fmt.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	fmt.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		fmt.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	fmt.Println("Dropping temporary resource table...")
	err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimID = ?"
	stmt, err := c.Query(sql, c.Simid)
	panicif(err)

	n := 0
	err = stmt.Scan(&n)
	panicif(err)
	stmt.Reset()

	roots = make([]*Node, 0, n)
	for stmt, err = c.Query(rootsSql, c.Simid, c.Simid); err == nil; err = stmt.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := stmt.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.StateId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	if err != io.EOF {
		panic(err)
	}
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	err := c.tmpResStmt.Query(node.ResId, node.ResId)
	for ; err == nil; err = c.tmpResStmt.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := c.tmpResStmt.Scan(&child.ResId, &child.StartTime, &child.StateId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	if err != io.EOF {
		panic(err)
	}

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				StateId:   node.StateId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	err := c.ownerStmt.Query(id, c.Simid)
	for ; err == nil; err = c.ownerStmt.Next() {
		err := c.ownerStmt.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	if err != io.EOF {
		panic(err)
	}
	return owners, times
}

func (c *Context) dumpNodes() {
	fmt.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	err := c.Exec("BEGIN TRANSACTION;")
	panicif(err)

	for _, n := range c.nodes {
		err = c.dumpStmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.StateId, n.Quantity)
		panicif(err)
	}
	err = c.Exec("END TRANSACTION;")
	panicif(err)

	c.nodes = c.nodes[:0]
}
<commit_msg>updated inventory post-process to new db schema. And adding combining of agententry and agentexit tables into agents table. If post-processing has already been done for certain simids, they are skipped and not redone.<commit_after>package main

import (
	"fmt"
	"io"
	"math"
	"strings"

	"code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"DROP TABLE IF EXISTS TimeList;",
		"CREATE TABLE IF NOT EXISTS Agents (SimId TEXT,Kind TEXT,Implementation TEXT,Prototype TEXT,ParentId INTEGER,EnterTime INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Inventories (SimId TEXT,ResourceId INTEGER,AgentId INTEGER,StartTime INTEGER,EndTime INTEGER,StateId INTEGER,Quantity REAL);",
		"CREATE TABLE TimeList AS SELECT DISTINCT Time FROM Transactions;",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimId", "ResourceId", "StateId"),
		query.Index("Compositions", "SimId", "StateId", "NucId"),
		query.Index("Transactions", "SimId", "ResourceId"),
		query.Index("Transactions", "TransactionId"),
		query.Index("ResCreators", "SimId", "ResourceId"),
	}
	postExecStmts = []string{
		query.Index("Agents", "SimId", "Prototype"),
		query.Index("Inventories", "SimId", "AgentId"),
		query.Index("Inventories", "SimId", "ResourceId", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ResourceId,TimeCreated,StateId,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverId, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceId = ? AND tr.SimId = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ResourceId,res.TimeCreated,rc.AgentId,res.StateId,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ResourceId = rc.ResourceId
				  WHERE res.SimId = ? AND rc.SimId = ?;`
)

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating indexes and inventory table...")
	for _, sql := range preExecStmts {
		if err := conn.Exec(sql); err != nil {
			fmt.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating inventory indexes...")
	for _, sql := range postExecStmts {
		if err := conn.Exec(sql); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	StateId   int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sqlite3.Conn
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       string
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sqlite3.Stmt
	dumpStmt    *sqlite3.Stmt
	ownerStmt   *sqlite3.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(conn *sqlite3.Conn, simid string, history chan string) *Context {
	return &Context{
		Conn:  conn,
		Simid: simid,
	}
}

func (c *Context) init() {
	// skip if the post processing already exists for this simid in the db
	_, err := c.Query("SELECT * FROM Inventories WHERE SimId = ? LIMIT 1", c.Simid)
	if err == nil {
		panic(fmt.Errorf("SimId %v is already post-processed. Skipping.\n", c.Simid))
	} else if err != io.EOF {
		panicif(err)
	}

	// build Agents table
	sql := `INSERT INTO Agents
				SELECT n.SimId,n.Kind,n.Implementation,n.Prototype,n.ParentId,n.EnterTime,x.ExitTime
				FROM
					AgentEntry AS n
					INNER JOIN AgentExit AS x ON n.AgentId = x.AgentId
				WHERE
					n.SimId = x.SimId AND n.SimId = ?;`
	err = c.Exec(sql, c.Simid)
	panicif(err)

	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// create temp res table without simid
	fmt.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + strings.Replace(c.Simid, "-", "_", -1)
	err = c.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql = "CREATE TABLE " + c.tmpResTbl + " AS SELECT ResourceId,TimeCreated,Parent1,Parent2,StateId,Quantity FROM Resources WHERE SimId = ?;"
	err = c.Exec(sql, c.Simid)
	panicif(err)

	fmt.Println("Indexing temporary resource table...")
	err = c.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	err = c.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	fmt.Printf("--- Building inventories for simid %v ---\n", c.Simid)
	c.init()

	fmt.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	fmt.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		fmt.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	fmt.Println("Dropping temporary resource table...")
	err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimId = ?"
	stmt, err := c.Query(sql, c.Simid)
	panicif(err)

	n := 0
	err = stmt.Scan(&n)
	panicif(err)
	stmt.Reset()

	roots = make([]*Node, 0, n)
	for stmt, err = c.Query(rootsSql, c.Simid, c.Simid); err == nil; err = stmt.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := stmt.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.StateId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	if err != io.EOF {
		panic(err)
	}
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	err := c.tmpResStmt.Query(node.ResId, node.ResId)
	for ; err == nil; err = c.tmpResStmt.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := c.tmpResStmt.Scan(&child.ResId, &child.StartTime, &child.StateId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	if err != io.EOF {
		panic(err)
	}

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				StateId:   node.StateId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	err := c.ownerStmt.Query(id, c.Simid)
	for ; err == nil; err = c.ownerStmt.Next() {
		err := c.ownerStmt.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	if err != io.EOF {
		panic(err)
	}
	return owners, times
}

func (c *Context) dumpNodes() {
	fmt.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	err := c.Exec("BEGIN TRANSACTION;")
	panicif(err)

	for _, n := range c.nodes {
		err = c.dumpStmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.StateId, n.Quantity)
		panicif(err)
	}
	err = c.Exec("END TRANSACTION;")
	panicif(err)

	c.nodes = c.nodes[:0]
}
<|endoftext|>


<commit_before>package circular

import "fmt"

// Implement a circular buffer of bytes supporting both overflow-checked writes
// and unconditional, possibly overwriting, writes.
//
// We chose the provided API so that Buffer implements io.ByteReader
// and io.ByteWriter and can be used (size permitting) as a drop in
// replacement for anything using that interface.

// Define the Buffer type here.
type Buffer struct {
	capacity    int
	store       []byte
	length      int // tracks the number of units in the buffer
	insertIndex int
	readIndex   int
}

func NewBuffer(capacity int) *Buffer {
	store := make([]byte, capacity)
	return &Buffer{capacity: capacity, store: store}
}

func (b *Buffer) ReadByte() (byte, error) {
	fmt.Printf("Reading %v", b)
	if b.len() == 0 {
		return 0, fmt.Errorf("error reading from empty buffer")
	}
	result := b.store[b.readIndex]
	b.store[b.readIndex] = 0
	b.readIndex = nextIndex(b.readIndex, b.capacity)
	b.length--
	return result, nil
}

func (b *Buffer) WriteByte(c byte) error {
	if b.len() == b.capacity {
		return fmt.Errorf("error writing to full buffer")
	}
	b.store[b.insertIndex] = c
	b.insertIndex = nextIndex(b.insertIndex, b.capacity)
	b.length++
	return nil
}

func (b *Buffer) Overwrite(c byte) {
	if b.store[b.insertIndex] == 0 {
		b.length++
	}
	b.store[b.insertIndex] = c
	b.insertIndex = nextIndex(b.insertIndex, b.capacity)
	b.readIndex = nextIndex(b.readIndex, b.capacity)
}

func (b *Buffer) Reset() {
	b.store = make([]byte, b.capacity)
	b.length = 0
	b.insertIndex = 0
	b.readIndex = 0
}

func (b *Buffer) String() string {
	return fmt.Sprintf("store %v, capacity %v, length %v, readIndex %v, writeIndex %v\n", b.store, b.capacity, b.length, b.readIndex, b.insertIndex)
}

func (b *Buffer) len() int {
	return b.length
}

func nextIndex(current int, capacity int) int {
	// If we are at the last index in the buffer, loop around back to the 0
	// index
	if current == capacity-1 {
		return 0
	}
	return current + 1
}
<commit_msg>Solve circular buffer<commit_after>package circular

import "fmt"

// Implement a circular buffer of bytes supporting both overflow-checked writes
// and unconditional, possibly overwriting, writes.
//
// We chose the provided API so that Buffer implements io.ByteReader
// and io.ByteWriter and can be used (size permitting) as a drop in
// replacement for anything using that interface.

// Define the Buffer type here.
type Buffer struct {
	capacity    int
	store       []byte
	length      int // tracks the number of units in the buffer
	insertIndex int
	readIndex   int
}

func NewBuffer(capacity int) *Buffer {
	store := make([]byte, capacity)
	return &Buffer{capacity: capacity, store: store}
}

func (b *Buffer) ReadByte() (byte, error) {
	// fmt.Printf("Reading %v", b)
	if b.len() == 0 {
		return 0, fmt.Errorf("error reading from empty buffer")
	}
	result := b.store[b.readIndex]
	b.store[b.readIndex] = 0
	b.readIndex = nextIndex(b.readIndex, b.capacity)
	b.length--
	return result, nil
}

func (b *Buffer) WriteByte(c byte) error {
	if b.len() == b.capacity {
		return fmt.Errorf("error writing to full buffer")
	}
	b.store[b.insertIndex] = c
	b.insertIndex = nextIndex(b.insertIndex, b.capacity)
	b.length++
	return nil
}

func (b *Buffer) Overwrite(c byte) {
	if b.store[b.insertIndex] == 0 {
		b.length++
	} else {
		b.readIndex = nextIndex(b.readIndex, b.capacity)
	}
	b.store[b.insertIndex] = c
	b.insertIndex = nextIndex(b.insertIndex, b.capacity)
}

func (b *Buffer) Reset() {
	b.store = make([]byte, b.capacity)
	b.length = 0
	b.insertIndex = 0
	b.readIndex = 0
}

func (b *Buffer) String() string {
	return fmt.Sprintf("store %v, capacity %v, length %v, readIndex %v, writeIndex %v\n", b.store, b.capacity, b.length, b.readIndex, b.insertIndex)
}

func (b *Buffer) len() int {
	return b.length
}

func nextIndex(current int, capacity int) int {
	// If we are at the last index in the buffer, loop around back to the 0
	// index
	if current == capacity-1 {
		return 0
	}
	return current + 1
}
<|endoftext|>


<commit_before>// Copyright 2015 Reborndb Org. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package utils

import (
	"fmt"
	"os"
	"path"
	"testing"

	"github.com/reborndb/go/bytesize"
	"github.com/reborndb/qdb/pkg/engine/goleveldb"
	"github.com/reborndb/qdb/pkg/service"
	"github.com/reborndb/qdb/pkg/store"
	. "gopkg.in/check.v1"
)

func TestT(t *testing.T) {
	TestingT(t)
}

var _ = Suite(&testUtilsSuite{})

type testUtilsSuite struct {
	s *testServer

	auth string
}

type testServer struct {
	addr   string
	store  *store.Store
	server *service.Server
}

func (s *testServer) Close() {
	if s.server != nil {
		s.server.Close()
		s.server = nil
	}
}

func (s *testUtilsSuite) SetUpSuite(c *C) {
	s.auth = "abc"
	s.s = s.testCreateServer(c, 16380, s.auth)
}

func (s *testUtilsSuite) testCreateServer(c *C, port int, auth string) *testServer {
	base := fmt.Sprintf("/tmp/test_reborn/test_proxy_utils/%d", port)

	err := os.MkdirAll(base, 0700)
	c.Assert(err, IsNil)

	conf := goleveldb.NewDefaultConfig()
	testdb, err := goleveldb.Open(path.Join(base, "db"), conf, false)
	c.Assert(err, IsNil)

	cfg := service.NewDefaultConfig()
	cfg.Listen = fmt.Sprintf("127.0.0.1:%d", port)
	cfg.PidFile = fmt.Sprintf(base, "qdb.pid")
	cfg.DumpPath = path.Join(base, "rdb.dump")
	cfg.SyncFilePath = path.Join(base, "sync.pipe")
	cfg.Auth = auth
	cfg.ReplBacklogSize = bytesize.MB

	store := store.New(testdb)
	server, err := service.NewServer(cfg, store)
	c.Assert(err, IsNil)
	go server.Serve()

	ss := new(testServer)
	ss.addr = cfg.Listen
	ss.store = store
	ss.server = server

	c.Assert(ss, NotNil)

	return ss
}

func (s *testUtilsSuite) TearDownSuite(c *C) {
	if s.s != nil {
		s.s.Close()
	}
}

func (s *testUtilsSuite) TestPing(c *C) {
	err := Ping(s.s.addr, s.auth)
	c.Assert(err, IsNil)
}

func (s *testUtilsSuite) TestGetInfo(c *C) {
	_, err := GetRedisInfo(s.s.addr, "", s.auth)
	c.Assert(err, IsNil)
}
<commit_msg>Modify qdb server port for ci building<commit_after>// Copyright 2015 Reborndb Org. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package utils

import (
	"fmt"
	"os"
	"path"
	"testing"

	"github.com/reborndb/go/bytesize"
	"github.com/reborndb/qdb/pkg/engine/goleveldb"
	"github.com/reborndb/qdb/pkg/service"
	"github.com/reborndb/qdb/pkg/store"
	. "gopkg.in/check.v1"
)

func TestT(t *testing.T) {
	TestingT(t)
}

var _ = Suite(&testUtilsSuite{})

type testUtilsSuite struct {
	s *testServer

	auth string
}

type testServer struct {
	addr   string
	store  *store.Store
	server *service.Server
}

func (s *testServer) Close() {
	if s.server != nil {
		s.server.Close()
		s.server = nil
	}
}

func (s *testUtilsSuite) SetUpSuite(c *C) {
	s.auth = "abc"
	s.s = s.testCreateServer(c, 36380, s.auth)
}

func (s *testUtilsSuite) testCreateServer(c *C, port int, auth string) *testServer {
	base := fmt.Sprintf("/tmp/test_reborn/test_proxy_utils/%d", port)

	err := os.MkdirAll(base, 0700)
	c.Assert(err, IsNil)

	conf := goleveldb.NewDefaultConfig()
	testdb, err := goleveldb.Open(path.Join(base, "db"), conf, false)
	c.Assert(err, IsNil)

	cfg := service.NewDefaultConfig()
	cfg.Listen = fmt.Sprintf("127.0.0.1:%d", port)
	cfg.PidFile = fmt.Sprintf(base, "qdb.pid")
	cfg.DumpPath = path.Join(base, "rdb.dump")
	cfg.SyncFilePath = path.Join(base, "sync.pipe")
	cfg.Auth = auth
	cfg.ReplBacklogSize = bytesize.MB

	store := store.New(testdb)
	server, err := service.NewServer(cfg, store)
	c.Assert(err, IsNil)
	go server.Serve()

	ss := new(testServer)
	ss.addr = cfg.Listen
	ss.store = store
	ss.server = server

	c.Assert(ss, NotNil)

	return ss
}

func (s *testUtilsSuite) TearDownSuite(c *C) {
	if s.s != nil {
		s.s.Close()
	}
}

func (s *testUtilsSuite) TestPing(c *C) {
	err := Ping(s.s.addr, s.auth)
	c.Assert(err, IsNil)
}

func (s *testUtilsSuite) TestGetInfo(c *C) {
	_, err := GetRedisInfo(s.s.addr, "", s.auth)
	c.Assert(err, IsNil)
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package fc

import (
	"os"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/util/mount"
	"k8s.io/kubernetes/pkg/volume"
)

// Abstract interface to disk operations.
type diskManager interface {
	MakeGlobalPDName(disk fcDisk) string
	// Attaches the disk to the kubelet's host machine.
	AttachDisk(b fcDiskMounter) (string, error)
	// Detaches the disk from the kubelet's host machine.
	DetachDisk(disk fcDiskUnmounter, devName string) error
}

// utility to mount a disk based filesystem
func diskSetUp(manager diskManager, b fcDiskMounter, volPath string, mounter mount.Interface, fsGroup *int64) error {
	globalPDPath := manager.MakeGlobalPDName(*b.fcDisk)
	// TODO: handle failed mounts here.
	noMnt, err := mounter.IsLikelyNotMountPoint(volPath)

	if err != nil && !os.IsNotExist(err) {
		glog.Errorf("cannot validate mountpoint: %s", volPath)
		return err
	}
	if !noMnt {
		return nil
	}
	if err := os.MkdirAll(volPath, 0750); err != nil {
		glog.Errorf("failed to mkdir:%s", volPath)
		return err
	}
	// Perform a bind mount to the full path to allow duplicate mounts of the same disk.
	options := []string{"bind"}
	if b.readOnly {
		options = append(options, "ro")
	}
	err = mounter.Mount(globalPDPath, volPath, "", options)
	if err != nil {
		glog.Errorf("failed to bind mount:%s", globalPDPath)
		return err
	}

	if !b.readOnly {
		volume.SetVolumeOwnership(&b, fsGroup)
	}

	return nil
}

// utility to tear down a disk based filesystem
func diskTearDown(manager diskManager, c fcDiskUnmounter, volPath string, mounter mount.Interface) error {
	noMnt, err := mounter.IsLikelyNotMountPoint(volPath)
	if err != nil {
		glog.Errorf("cannot validate mountpoint %s", volPath)
		return err
	}
	if noMnt {
		return os.Remove(volPath)
	}

	if err := mounter.Unmount(volPath); err != nil {
		glog.Errorf("failed to unmount %s", volPath)
		return err
	}

	noMnt, mntErr := mounter.IsLikelyNotMountPoint(volPath)
	if mntErr != nil {
		glog.Errorf("isMountpoint check failed: %v", mntErr)
		return err
	}
	if noMnt {
		if err := os.Remove(volPath); err != nil {
			return err
		}
	}
	return nil

}
<commit_msg>handle failed mounts for fc volumes<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package fc

import (
	"os"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/util/mount"
	"k8s.io/kubernetes/pkg/volume"
)

// Abstract interface to disk operations.
type diskManager interface {
	MakeGlobalPDName(disk fcDisk) string
	// Attaches the disk to the kubelet's host machine.
	AttachDisk(b fcDiskMounter) (string, error)
	// Detaches the disk from the kubelet's host machine.
	DetachDisk(disk fcDiskUnmounter, devName string) error
}

// utility to mount a disk based filesystem
func diskSetUp(manager diskManager, b fcDiskMounter, volPath string, mounter mount.Interface, fsGroup *int64) error {
	globalPDPath := manager.MakeGlobalPDName(*b.fcDisk)
	noMnt, err := mounter.IsLikelyNotMountPoint(volPath)

	if err != nil && !os.IsNotExist(err) {
		glog.Errorf("cannot validate mountpoint: %s", volPath)
		return err
	}
	if !noMnt {
		return nil
	}
	if err := os.MkdirAll(volPath, 0750); err != nil {
		glog.Errorf("failed to mkdir:%s", volPath)
		return err
	}
	// Perform a bind mount to the full path to allow duplicate mounts of the same disk.
	options := []string{"bind"}
	if b.readOnly {
		options = append(options, "ro")
	}
	err = mounter.Mount(globalPDPath, volPath, "", options)
	if err != nil {
		glog.Errorf("Failed to bind mount: source:%s, target:%s, err:%v", globalPDPath, volPath, err)
		noMnt, mntErr := b.mounter.IsLikelyNotMountPoint(volPath)
		if mntErr != nil {
			glog.Errorf("IsLikelyNotMountPoint check failed: %v", mntErr)
			return err
		}
		if !noMnt {
			if mntErr = b.mounter.Unmount(volPath); mntErr != nil {
				glog.Errorf("Failed to unmount: %v", mntErr)
				return err
			}
			noMnt, mntErr = b.mounter.IsLikelyNotMountPoint(volPath)
			if mntErr != nil {
				glog.Errorf("IsLikelyNotMountPoint check failed: %v", mntErr)
				return err
			}
			if !noMnt {
				//  will most likely retry on next sync loop.
				glog.Errorf("%s is still mounted, despite call to unmount().  Will try again next sync loop.", volPath)
				return err
			}
		}
		os.Remove(volPath)

		return err
	}

	if !b.readOnly {
		volume.SetVolumeOwnership(&b, fsGroup)
	}

	return nil
}

// utility to tear down a disk based filesystem
func diskTearDown(manager diskManager, c fcDiskUnmounter, volPath string, mounter mount.Interface) error {
	noMnt, err := mounter.IsLikelyNotMountPoint(volPath)
	if err != nil {
		glog.Errorf("cannot validate mountpoint %s", volPath)
		return err
	}
	if noMnt {
		return os.Remove(volPath)
	}

	if err := mounter.Unmount(volPath); err != nil {
		glog.Errorf("failed to unmount %s", volPath)
		return err
	}

	noMnt, mntErr := mounter.IsLikelyNotMountPoint(volPath)
	if mntErr != nil {
		glog.Errorf("isMountpoint check failed: %v", mntErr)
		return err
	}
	if noMnt {
		if err := os.Remove(volPath); err != nil {
			return err
		}
	}
	return nil

}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"

	restserver "github.com/restic/rest-server"
	"github.com/spf13/cobra"
)

// cmdRoot is the base command when no other command has been specified.
var cmdRoot = &cobra.Command{
	Use:           "rest-server",
	Short:         "Run a REST server for use with restic",
	SilenceErrors: true,
	SilenceUsage:  true,
	RunE:          runRoot,
}

func init() {
	flags := cmdRoot.Flags()
	flags.StringVar(&restserver.Config.CPUProfile, "cpuprofile", restserver.Config.CPUProfile, "write CPU profile to file")
	flags.BoolVar(&restserver.Config.Debug, "debug", restserver.Config.Debug, "output debug messages")
	flags.StringVar(&restserver.Config.Listen, "listen", restserver.Config.Listen, "listen address")
	flags.StringVar(&restserver.Config.Log, "log", restserver.Config.Log, "log HTTP requests in the combined log format")
	flags.StringVar(&restserver.Config.Path, "path", restserver.Config.Path, "data directory")
	flags.BoolVar(&restserver.Config.TLS, "tls", restserver.Config.TLS, "turn on TLS support")
	flags.BoolVar(&restserver.Config.AppendOnly, "append-only", restserver.Config.AppendOnly, "enable append only mode")
	flags.BoolVar(&restserver.Config.Prometheus, "prometheus", restserver.Config.AppendOnly, "enable Prometheus metrics")
}

var version = "manually"

func runRoot(cmd *cobra.Command, args []string) error {
	log.SetFlags(0)

	log.Printf("rest-server %s compiled with %v on %v/%v\n", version, runtime.Version(), runtime.GOOS, runtime.GOARCH)
	log.Printf("Data directory: %s", restserver.Config.Path)

	if restserver.Config.CPUProfile != "" {
		f, err := os.Create(restserver.Config.CPUProfile)
		if err != nil {
			return err
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			return err
		}
		log.Println("CPU profiling enabled")
		defer pprof.StopCPUProfile()
	}

	mux := restserver.NewMux()

	var handler http.Handler
	htpasswdFile, err := restserver.NewHtpasswdFromFile(filepath.Join(restserver.Config.Path, ".htpasswd"))
	if err != nil {
		handler = mux
		log.Println("Authentication disabled")
	} else {
		handler = restserver.AuthHandler(htpasswdFile, mux)
		log.Println("Authentication enabled")
	}

	if !restserver.Config.TLS {
		log.Printf("Starting server on %s\n", restserver.Config.Listen)
		err = http.ListenAndServe(restserver.Config.Listen, handler)
	} else {
		privateKey := filepath.Join(restserver.Config.Path, "private_key")
		publicKey := filepath.Join(restserver.Config.Path, "public_key")
		log.Println("TLS enabled")
		log.Printf("Private key: %s", privateKey)
		log.Printf("Public key: %s", publicKey)
		log.Printf("Starting server on %s\n", restserver.Config.Listen)
		err = http.ListenAndServeTLS(restserver.Config.Listen, publicKey, privateKey, handler)
	}

	return err
}

func main() {
	if err := cmdRoot.Execute(); err != nil {
		log.Fatalf("error: %v", err)
	}
}
<commit_msg>fix prometheus config flag<commit_after>package main

import (
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"

	restserver "github.com/restic/rest-server"
	"github.com/spf13/cobra"
)

// cmdRoot is the base command when no other command has been specified.
var cmdRoot = &cobra.Command{
	Use:           "rest-server",
	Short:         "Run a REST server for use with restic",
	SilenceErrors: true,
	SilenceUsage:  true,
	RunE:          runRoot,
}

func init() {
	flags := cmdRoot.Flags()
	flags.StringVar(&restserver.Config.CPUProfile, "cpuprofile", restserver.Config.CPUProfile, "write CPU profile to file")
	flags.BoolVar(&restserver.Config.Debug, "debug", restserver.Config.Debug, "output debug messages")
	flags.StringVar(&restserver.Config.Listen, "listen", restserver.Config.Listen, "listen address")
	flags.StringVar(&restserver.Config.Log, "log", restserver.Config.Log, "log HTTP requests in the combined log format")
	flags.StringVar(&restserver.Config.Path, "path", restserver.Config.Path, "data directory")
	flags.BoolVar(&restserver.Config.TLS, "tls", restserver.Config.TLS, "turn on TLS support")
	flags.BoolVar(&restserver.Config.AppendOnly, "append-only", restserver.Config.AppendOnly, "enable append only mode")
	flags.BoolVar(&restserver.Config.Prometheus, "prometheus", restserver.Config.Prometheus, "enable Prometheus metrics")
}

var version = "manually"

func runRoot(cmd *cobra.Command, args []string) error {
	log.SetFlags(0)

	log.Printf("rest-server %s compiled with %v on %v/%v\n", version, runtime.Version(), runtime.GOOS, runtime.GOARCH)
	log.Printf("Data directory: %s", restserver.Config.Path)

	if restserver.Config.CPUProfile != "" {
		f, err := os.Create(restserver.Config.CPUProfile)
		if err != nil {
			return err
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			return err
		}
		log.Println("CPU profiling enabled")
		defer pprof.StopCPUProfile()
	}

	mux := restserver.NewMux()

	var handler http.Handler
	htpasswdFile, err := restserver.NewHtpasswdFromFile(filepath.Join(restserver.Config.Path, ".htpasswd"))
	if err != nil {
		handler = mux
		log.Println("Authentication disabled")
	} else {
		handler = restserver.AuthHandler(htpasswdFile, mux)
		log.Println("Authentication enabled")
	}

	if !restserver.Config.TLS {
		log.Printf("Starting server on %s\n", restserver.Config.Listen)
		err = http.ListenAndServe(restserver.Config.Listen, handler)
	} else {
		privateKey := filepath.Join(restserver.Config.Path, "private_key")
		publicKey := filepath.Join(restserver.Config.Path, "public_key")
		log.Println("TLS enabled")
		log.Printf("Private key: %s", privateKey)
		log.Printf("Public key: %s", publicKey)
		log.Printf("Starting server on %s\n", restserver.Config.Listen)
		err = http.ListenAndServeTLS(restserver.Config.Listen, publicKey, privateKey, handler)
	}

	return err
}

func main() {
	if err := cmdRoot.Execute(); err != nil {
		log.Fatalf("error: %v", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"time"

	"github.com/restic/restic/internal/cache"
	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/fs"
	"github.com/restic/restic/internal/ui/table"
	"github.com/spf13/cobra"
)

var cmdCache = &cobra.Command{
	Use:   "cache",
	Short: "Operate on local cache directories",
	Long: `
The "cache" command allows listing and cleaning local cache directories.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runCache(cacheOptions, globalOptions, args)
	},
}

// CacheOptions bundles all options for the snapshots command.
type CacheOptions struct {
	Cleanup bool
	MaxAge  uint
	NoSize  bool
}

var cacheOptions CacheOptions

func init() {
	cmdRoot.AddCommand(cmdCache)

	f := cmdCache.Flags()
	f.BoolVar(&cacheOptions.Cleanup, "cleanup", false, "remove old cache directories")
	f.UintVar(&cacheOptions.MaxAge, "max-age", 30, "max age in `days` for cache directories to be considered old")
	f.BoolVar(&cacheOptions.NoSize, "no-size", false, "do not output the size of the cache directories")
}

func runCache(opts CacheOptions, gopts GlobalOptions, args []string) error {
	if len(args) > 0 {
		return errors.Fatal("the cache command has no arguments")
	}

	if gopts.NoCache {
		return errors.Fatal("Refusing to do anything, the cache is disabled")
	}

	var (
		cachedir = gopts.CacheDir
		err      error
	)

	if cachedir == "" {
		cachedir, err = cache.DefaultDir()
		if err != nil {
			return err
		}
	}

	if opts.Cleanup || gopts.CleanupCache {
		oldDirs, err := cache.OlderThan(cachedir, time.Duration(opts.MaxAge)*24*time.Hour)
		if err != nil {
			return err
		}

		if len(oldDirs) == 0 {
			Verbosef("no old cache dirs found\n")
			return nil
		}

		Verbosef("remove %d old cache directories\n", len(oldDirs))

		for _, item := range oldDirs {
			dir := filepath.Join(cachedir, item.Name())
			err = fs.RemoveAll(dir)
			if err != nil {
				Warnf("unable to remove %v: %v\n", dir, err)
			}
		}

		return nil
	}

	tab := table.New()

	type data struct {
		ID   string
		Last string
		Old  string
		Size string
	}

	tab.AddColumn("Repo ID", "{{ .ID }}")
	tab.AddColumn("Last Used", "{{ .Last }}")
	tab.AddColumn("Old", "{{ .Old }}")

	if !opts.NoSize {
		tab.AddColumn("Size", "{{ .Size }}")
	}

	dirs, err := cache.All(cachedir)
	if err != nil {
		return err
	}

	if len(dirs) == 0 {
		Printf("no cache dirs found, basedir is %v\n", cachedir)
		return nil
	}

	sort.Slice(dirs, func(i, j int) bool {
		return dirs[i].ModTime().Before(dirs[j].ModTime())
	})

	for _, entry := range dirs {
		var old string
		if cache.IsOld(entry.ModTime(), time.Duration(opts.MaxAge)*24*time.Hour) {
			old = "yes"
		}

		var size string
		if !opts.NoSize {
			bytes, err := dirSize(filepath.Join(cachedir, entry.Name()))
			if err != nil {
				return err
			}
			size = fmt.Sprintf("%11s", formatBytes(uint64(bytes)))
		}

		tab.AddRow(data{
			entry.Name()[:10],
			fmt.Sprintf("%d days ago", uint(time.Since(entry.ModTime()).Hours()/24)),
			old,
			size,
		})
	}

	tab.Write(gopts.stdout)
	Printf("%d cache dirs in %s\n", len(dirs), cachedir)

	return nil
}

func dirSize(path string) (int64, error) {
	var size int64
	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			size += info.Size()
		}
		return err
	})
	return size, err
}
<commit_msg>Catch errors when walking cache directories<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"time"

	"github.com/restic/restic/internal/cache"
	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/fs"
	"github.com/restic/restic/internal/ui/table"
	"github.com/spf13/cobra"
)

var cmdCache = &cobra.Command{
	Use:   "cache",
	Short: "Operate on local cache directories",
	Long: `
The "cache" command allows listing and cleaning local cache directories.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runCache(cacheOptions, globalOptions, args)
	},
}

// CacheOptions bundles all options for the snapshots command.
type CacheOptions struct {
	Cleanup bool
	MaxAge  uint
	NoSize  bool
}

var cacheOptions CacheOptions

func init() {
	cmdRoot.AddCommand(cmdCache)

	f := cmdCache.Flags()
	f.BoolVar(&cacheOptions.Cleanup, "cleanup", false, "remove old cache directories")
	f.UintVar(&cacheOptions.MaxAge, "max-age", 30, "max age in `days` for cache directories to be considered old")
	f.BoolVar(&cacheOptions.NoSize, "no-size", false, "do not output the size of the cache directories")
}

func runCache(opts CacheOptions, gopts GlobalOptions, args []string) error {
	if len(args) > 0 {
		return errors.Fatal("the cache command has no arguments")
	}

	if gopts.NoCache {
		return errors.Fatal("Refusing to do anything, the cache is disabled")
	}

	var (
		cachedir = gopts.CacheDir
		err      error
	)

	if cachedir == "" {
		cachedir, err = cache.DefaultDir()
		if err != nil {
			return err
		}
	}

	if opts.Cleanup || gopts.CleanupCache {
		oldDirs, err := cache.OlderThan(cachedir, time.Duration(opts.MaxAge)*24*time.Hour)
		if err != nil {
			return err
		}

		if len(oldDirs) == 0 {
			Verbosef("no old cache dirs found\n")
			return nil
		}

		Verbosef("remove %d old cache directories\n", len(oldDirs))

		for _, item := range oldDirs {
			dir := filepath.Join(cachedir, item.Name())
			err = fs.RemoveAll(dir)
			if err != nil {
				Warnf("unable to remove %v: %v\n", dir, err)
			}
		}

		return nil
	}

	tab := table.New()

	type data struct {
		ID   string
		Last string
		Old  string
		Size string
	}

	tab.AddColumn("Repo ID", "{{ .ID }}")
	tab.AddColumn("Last Used", "{{ .Last }}")
	tab.AddColumn("Old", "{{ .Old }}")

	if !opts.NoSize {
		tab.AddColumn("Size", "{{ .Size }}")
	}

	dirs, err := cache.All(cachedir)
	if err != nil {
		return err
	}

	if len(dirs) == 0 {
		Printf("no cache dirs found, basedir is %v\n", cachedir)
		return nil
	}

	sort.Slice(dirs, func(i, j int) bool {
		return dirs[i].ModTime().Before(dirs[j].ModTime())
	})

	for _, entry := range dirs {
		var old string
		if cache.IsOld(entry.ModTime(), time.Duration(opts.MaxAge)*24*time.Hour) {
			old = "yes"
		}

		var size string
		if !opts.NoSize {
			bytes, err := dirSize(filepath.Join(cachedir, entry.Name()))
			if err != nil {
				return err
			}
			size = fmt.Sprintf("%11s", formatBytes(uint64(bytes)))
		}

		tab.AddRow(data{
			entry.Name()[:10],
			fmt.Sprintf("%d days ago", uint(time.Since(entry.ModTime()).Hours()/24)),
			old,
			size,
		})
	}

	tab.Write(gopts.stdout)
	Printf("%d cache dirs in %s\n", len(dirs), cachedir)

	return nil
}

func dirSize(path string) (int64, error) {
	var size int64
	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if err != nil || info == nil {
			return err
		}

		if !info.IsDir() {
			size += info.Size()
		}

		return nil
	})
	return size, err
}
<|endoftext|>


<commit_before>package datadog

import (
	"encoding/json"
	"github.com/DataDog/dd-trace-go/tracer"
	"github.com/flachnetz/dd-zipkin-proxy/proxy"
	"github.com/sirupsen/logrus"
	"os"
	"time"
)

var log = logrus.WithField("prefix", "datadog")
var logTraces = os.Getenv("DD_LOG_TRACES") == "true"

const flushInterval = 2 * time.Second
const flushSpanCount = 1000

// Create a new default transport.
func DefaultTransport(hostname, port string) tracer.Transport {
	return tracer.NewTransport(hostname, port)
}

func submitTraces(transport tracer.Transport, spansByTrace <-chan map[uint64][]*tracer.Span) {
	for buffer := range spansByTrace {
		count := 0

		// the transport expects a list of list, where each sub-list contains only
		// spans of the same trace.
		var traces [][]*tracer.Span
		for _, spans := range buffer {
			count += len(spans)
			traces = append(traces, spans)
		}

		// if we got traces, send them!
		if len(traces) > 0 {
			log.Infof("Sending %d spans in traces %d traces", count, len(traces))

			if logTraces {
				val, _ := json.MarshalIndent(traces, "", "  ")
				log.Info(string(val))
			} else {
				if _, err := transport.SendTraces(traces); err != nil {
					log.WithError(err).Warn("Error reporting spans to datadog")
				}
			}
		}
	}
}

func Sink(transport tracer.Transport, spans <-chan proxy.Span) {
	ticker := time.NewTicker(flushInterval)
	defer ticker.Stop()

	count := 0
	byTrace := make(map[uint64][]*tracer.Span)

	groupedSpans := make(chan map[uint64][]*tracer.Span, 8)
	defer close(groupedSpans)

	// send the spans in background
	go submitTraces(transport, groupedSpans)

	var ddSpans []tracer.Span

	var lastFlushTime time.Time

	for {
		var flush bool

		select {
		case span, ok := <-spans:
			if !ok {
				log.Info("Channel closed, stopping sender")
				return
			}

			// set lower bound on time
			duration := span.Duration
			if duration < 1*time.Microsecond {
				duration = 1 * time.Microsecond
			}

			// use fallback if name is empty
			resource := span.Name
			if resource == "" {
				resource = "(resource empty)"
			}

			// get a buffer of spans
			if len(spans) < 1 {
				ddSpans = make([]tracer.Span, 4*1024)
			}

			// get a pointer to a free span
			converted := &ddSpans[0]
			ddSpans = ddSpans[1:]

			*converted = tracer.Span{
				Resource: resource,

				// use span.Service as the datadog Service and Name
				Name:    span.Service,
				Service: span.Service,

				Start:    span.Timestamp.ToTime().UnixNano(),
				Duration: duration.Nanoseconds(),

				SpanID:   span.Id.Uint64(),
				TraceID:  span.Trace.Uint64(),
				ParentID: span.Parent.Uint64(),

				Meta:    span.Tags,
				Sampled: true,
			}

			count++
			byTrace[converted.TraceID] = append(byTrace[converted.TraceID], converted)
			flush = count >= flushSpanCount

		case <-ticker.C:
			// only flush if we didn't automatically flush shortly before
			flush = time.Since(lastFlushTime) >= 90*flushInterval/100
		}

		if flush && count > 0 {
			select {
			case groupedSpans <- byTrace:
			default:
				log.Warnf("Discarding %d traces, sending to datadog would block.", len(byTrace))
			}

			// reset collection
			count = 0
			byTrace = make(map[uint64][]*tracer.Span)
			lastFlushTime = time.Now()
		}
	}
}
<commit_msg>Fix check on ddSpans slice<commit_after>package datadog

import (
	"encoding/json"
	"github.com/DataDog/dd-trace-go/tracer"
	"github.com/flachnetz/dd-zipkin-proxy/proxy"
	"github.com/sirupsen/logrus"
	"os"
	"time"
)

var log = logrus.WithField("prefix", "datadog")
var logTraces = os.Getenv("DD_LOG_TRACES") == "true"

const flushInterval = 2 * time.Second
const flushSpanCount = 1000

// Create a new default transport.
func DefaultTransport(hostname, port string) tracer.Transport {
	return tracer.NewTransport(hostname, port)
}

func submitTraces(transport tracer.Transport, spansByTrace <-chan map[uint64][]*tracer.Span) {
	for buffer := range spansByTrace {
		count := 0

		// the transport expects a list of list, where each sub-list contains only
		// spans of the same trace.
		var traces [][]*tracer.Span
		for _, spans := range buffer {
			count += len(spans)
			traces = append(traces, spans)
		}

		// if we got traces, send them!
		if len(traces) > 0 {
			log.Infof("Sending %d spans in traces %d traces", count, len(traces))

			if logTraces {
				val, _ := json.MarshalIndent(traces, "", "  ")
				log.Info(string(val))
			} else {
				if _, err := transport.SendTraces(traces); err != nil {
					log.WithError(err).Warn("Error reporting spans to datadog")
				}
			}
		}
	}
}

func Sink(transport tracer.Transport, spans <-chan proxy.Span) {
	ticker := time.NewTicker(flushInterval)
	defer ticker.Stop()

	count := 0
	byTrace := make(map[uint64][]*tracer.Span)

	groupedSpans := make(chan map[uint64][]*tracer.Span, 8)
	defer close(groupedSpans)

	// send the spans in background
	go submitTraces(transport, groupedSpans)

	var ddSpans []tracer.Span

	var lastFlushTime time.Time

	for {
		var flush bool

		select {
		case span, ok := <-spans:
			if !ok {
				log.Info("Channel closed, stopping sender")
				return
			}

			// set lower bound on time
			duration := span.Duration
			if duration < 1*time.Microsecond {
				duration = 1 * time.Microsecond
			}

			// use fallback if name is empty
			resource := span.Name
			if resource == "" {
				resource = "(resource empty)"
			}

			// get a buffer of spans
			if len(ddSpans) < 1 {
				ddSpans = make([]tracer.Span, 4*1024)
			}

			// get a pointer to a free span
			converted := &ddSpans[0]
			ddSpans = ddSpans[1:]

			*converted = tracer.Span{
				Resource: resource,

				// use span.Service as the datadog Service and Name
				Name:    span.Service,
				Service: span.Service,

				Start:    span.Timestamp.ToTime().UnixNano(),
				Duration: duration.Nanoseconds(),

				SpanID:   span.Id.Uint64(),
				TraceID:  span.Trace.Uint64(),
				ParentID: span.Parent.Uint64(),

				Meta:    span.Tags,
				Sampled: true,
			}

			count++
			byTrace[converted.TraceID] = append(byTrace[converted.TraceID], converted)
			flush = count >= flushSpanCount

		case <-ticker.C:
			// only flush if we didn't automatically flush shortly before
			flush = time.Since(lastFlushTime) >= 90*flushInterval/100
		}

		if flush && count > 0 {
			select {
			case groupedSpans <- byTrace:
			default:
				log.Warnf("Discarding %d traces, sending to datadog would block.", len(byTrace))
			}

			// reset collection
			count = 0
			byTrace = make(map[uint64][]*tracer.Span)
			lastFlushTime = time.Now()
		}
	}
}
<|endoftext|>


<commit_before>// update-ocsp reads all-in-one bundle files (whose names are passed as
// command-line argument) and sends queries to the OCSP responders, storing the
// responses in *.ocsp files next to the input files.
// The argument can also identify a directory, in which case all files in the
// directory (with the exception of those ending in .ocsp, .issuer, or .sctl
// –for HAProxy compatibility–, or .key –for compatibility with almost anything
// else, storing private keys separately–) are treated as input files.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ocsp"

	"github.com/tbroyer/ocspd"
)

var interval = flag.Duration("interval", 24*time.Hour, "indicative interval between invocations of this tool")
var hookCmd = flag.String("hook", "", "optional program to run if all goes well")

var exitCode = 0

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Fprintf(os.Stderr, "Missing certificate filename(s)")
		flag.Usage()
		os.Exit(2)
	}
	var names = fileNames()
	for _, certBundleFileName := range names {
		cert, issuer, err := ocspd.ParsePEMCertificateBundle(certBundleFileName)
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		// check existing/cached OCSP response before querying the responder
		ocspFileName := certBundleFileName + ".ocsp"
		needsRefresh, err := ocspd.NeedsRefreshFile(ocspFileName, issuer, *interval)
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		if !needsRefresh {
			// cached response is "fresh" enough, don't refresh it
			continue
		}
		data, err := ocspd.Update(cert, issuer, "")
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		resp, err := ocsp.ParseResponse(data, issuer)
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		fmt.Printf("%v: %v\n", certBundleFileName, statusString(resp.Status))
		fmt.Printf("\tThis Update: %v\n", resp.ThisUpdate)
		if !resp.NextUpdate.IsZero() {
			fmt.Printf("\tNext Update: %v\n", resp.NextUpdate)
		}
		if resp.Status == ocsp.Revoked {
			fmt.Printf("\tReason: %v\n", revocationReasonString(resp.RevocationReason))
			fmt.Printf("\tRevocation Time: %v\n", resp.RevokedAt)
		}
		if err = ioutil.WriteFile(ocspFileName, data, 0644); err != nil {
			log.Print(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		if len(*hookCmd) > 0 {
			if err = ocspd.RunHookCmd(*hookCmd, data, os.Stdout, os.Stderr); err != nil {
				log.Println(certBundleFileName, ": ", err)
				exitCode = 1
				continue
			}
		}
	}
	os.Exit(exitCode)
}

func fileNames() (names []string) {
	for _, arg := range flag.Args() {
		stats, err := os.Stat(arg)
		if os.IsNotExist(err) {
			continue
		}
		if err != nil {
			log.Fatal(err)
		}
		if stats.IsDir() {
			f, err := os.Open(arg)
			if os.IsNotExist(err) {
				// dir has disappeared between Stat and Open,
				// let's do as if it never existed
				continue
			}
			if err != nil {
				log.Fatal(err)
			}
			ns, err := f.Readdirnames(-1)
			f.Close()
			if err != nil {
				log.Fatal(err)
			}
			for _, n := range ns {
				if strings.HasSuffix(n, ".issuer") || strings.HasSuffix(n, ".ocsp") || strings.HasSuffix(n, ".sctl") || strings.HasSuffix(n, ".key") {
					continue
				}
				n = filepath.Join(arg, n)
				stats, err := os.Stat(n)
				if os.IsNotExist(err) {
					// n has disappeared between Readdirnames and Open,
					// let's do as if it never existed
					continue
				}
				if err != nil {
					log.Fatal(err)
				}
				if stats.Mode().IsRegular() {
					names = append(names, n)
				}
			}
		} else if stats.Mode().IsRegular() {
			names = append(names, arg)
		}
	}
	return
}

func statusString(status int) string {
	switch status {
	case ocsp.Good:
		return "good"
	case ocsp.Unknown:
		return "unknown"
	case ocsp.Revoked:
		return "revoked"
	default:
		log.Panicf("Unknown status %v", status)
		return ""
	}
}

func revocationReasonString(revocationReason int) string {
	switch revocationReason {
	case ocsp.Unspecified:
		return "unspecified"
	case ocsp.KeyCompromise:
		return "keyCompromise"
	case ocsp.CACompromise:
		return "cACompromise"
	case ocsp.AffiliationChanged:
		return "affiliationChanged"
	case ocsp.Superseded:
		return "superseded"
	case ocsp.CessationOfOperation:
		return "cessationOfOperation"
	case ocsp.CertificateHold:
		return "certificateHold"
	case ocsp.RemoveFromCRL:
		return "removeFromCRL"
	case ocsp.PrivilegeWithdrawn:
		return "privilegeWithdrawn"
	case ocsp.AACompromise:
		return "aACompromise"
	default:
		log.Panicf("Unknown revocation reason %v", revocationReason)
		return ""
	}
}
<commit_msg>Add shorthand CLI arguments to update-ocsp tool<commit_after>// update-ocsp reads all-in-one bundle files (whose names are passed as
// command-line argument) and sends queries to the OCSP responders, storing the
// responses in *.ocsp files next to the input files.
// The argument can also identify a directory, in which case all files in the
// directory (with the exception of those ending in .ocsp, .issuer, or .sctl
// –for HAProxy compatibility–, or .key –for compatibility with almost anything
// else, storing private keys separately–) are treated as input files.
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ocsp"

	"github.com/tbroyer/ocspd"
)

var interval time.Duration
var hookCmd string

func init() {
	const (
		defaultInterval = 24 * time.Hour
		intervalUsage   = "indicative interval between invocations of this tool"
		hookUsage       = "optional program to run if all goes well"
	)
	flag.DurationVar(&interval, "interval", defaultInterval, intervalUsage)
	flag.DurationVar(&interval, "i", defaultInterval, intervalUsage+" (shorthand)")

	flag.StringVar(&hookCmd, "hook", "", hookUsage)
	flag.StringVar(&hookCmd, "h", "", hookUsage+" (shorthand)")
}

var exitCode = 0

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Fprintf(os.Stderr, "Missing certificate filename(s)")
		flag.Usage()
		os.Exit(2)
	}
	var names = fileNames()
	for _, certBundleFileName := range names {
		cert, issuer, err := ocspd.ParsePEMCertificateBundle(certBundleFileName)
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		// check existing/cached OCSP response before querying the responder
		ocspFileName := certBundleFileName + ".ocsp"
		needsRefresh, err := ocspd.NeedsRefreshFile(ocspFileName, issuer, interval)
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		if !needsRefresh {
			// cached response is "fresh" enough, don't refresh it
			continue
		}
		data, err := ocspd.Update(cert, issuer, "")
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		resp, err := ocsp.ParseResponse(data, issuer)
		if err != nil {
			log.Println(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		fmt.Printf("%v: %v\n", certBundleFileName, statusString(resp.Status))
		fmt.Printf("\tThis Update: %v\n", resp.ThisUpdate)
		if !resp.NextUpdate.IsZero() {
			fmt.Printf("\tNext Update: %v\n", resp.NextUpdate)
		}
		if resp.Status == ocsp.Revoked {
			fmt.Printf("\tReason: %v\n", revocationReasonString(resp.RevocationReason))
			fmt.Printf("\tRevocation Time: %v\n", resp.RevokedAt)
		}
		if err = ioutil.WriteFile(ocspFileName, data, 0644); err != nil {
			log.Print(certBundleFileName, ": ", err)
			exitCode = 1
			continue
		}
		if hookCmd != "" {
			if err = ocspd.RunHookCmd(hookCmd, data, os.Stdout, os.Stderr); err != nil {
				log.Println(certBundleFileName, ": ", err)
				exitCode = 1
				continue
			}
		}
	}
	os.Exit(exitCode)
}

func fileNames() (names []string) {
	for _, arg := range flag.Args() {
		stats, err := os.Stat(arg)
		if os.IsNotExist(err) {
			continue
		}
		if err != nil {
			log.Fatal(err)
		}
		if stats.IsDir() {
			f, err := os.Open(arg)
			if os.IsNotExist(err) {
				// dir has disappeared between Stat and Open,
				// let's do as if it never existed
				continue
			}
			if err != nil {
				log.Fatal(err)
			}
			ns, err := f.Readdirnames(-1)
			f.Close()
			if err != nil {
				log.Fatal(err)
			}
			for _, n := range ns {
				if strings.HasSuffix(n, ".issuer") || strings.HasSuffix(n, ".ocsp") || strings.HasSuffix(n, ".sctl") || strings.HasSuffix(n, ".key") {
					continue
				}
				n = filepath.Join(arg, n)
				stats, err := os.Stat(n)
				if os.IsNotExist(err) {
					// n has disappeared between Readdirnames and Open,
					// let's do as if it never existed
					continue
				}
				if err != nil {
					log.Fatal(err)
				}
				if stats.Mode().IsRegular() {
					names = append(names, n)
				}
			}
		} else if stats.Mode().IsRegular() {
			names = append(names, arg)
		}
	}
	return
}

func statusString(status int) string {
	switch status {
	case ocsp.Good:
		return "good"
	case ocsp.Unknown:
		return "unknown"
	case ocsp.Revoked:
		return "revoked"
	default:
		log.Panicf("Unknown status %v", status)
		return ""
	}
}

func revocationReasonString(revocationReason int) string {
	switch revocationReason {
	case ocsp.Unspecified:
		return "unspecified"
	case ocsp.KeyCompromise:
		return "keyCompromise"
	case ocsp.CACompromise:
		return "cACompromise"
	case ocsp.AffiliationChanged:
		return "affiliationChanged"
	case ocsp.Superseded:
		return "superseded"
	case ocsp.CessationOfOperation:
		return "cessationOfOperation"
	case ocsp.CertificateHold:
		return "certificateHold"
	case ocsp.RemoveFromCRL:
		return "removeFromCRL"
	case ocsp.PrivilegeWithdrawn:
		return "privilegeWithdrawn"
	case ocsp.AACompromise:
		return "aACompromise"
	default:
		log.Panicf("Unknown revocation reason %v", revocationReason)
		return ""
	}
}
<|endoftext|>