<commit_before>package middlewares

import (
	"net/http"

	"github.com/solher/snakepit-seed/errs"

	"github.com/ansel1/merry"
	"github.com/pressly/chi"
	"github.com/solher/snakepit"
	"golang.org/x/net/context"
)

type Role string

const (
	Admin, Developer, User Role = "ADMIN", "DEVELOPER", "USER"
)

type Gate struct {
	json    *snakepit.JSON
	granter func(role Role) bool
}

func NewGate(j *snakepit.JSON, g func(role Role) bool) func(next chi.Handler) chi.Handler {
	gate := &Gate{json: j, granter: g}
	return gate.middleware
}

func (c *Gate) middleware(next chi.Handler) chi.Handler {
	return chi.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		session, err := GetCurrentSession(ctx)
		if err != nil {
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}
		if len(session.Role) == 0 {
			err := merry.New("empty role in context")
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}

		_, err = GetCurrentUser(ctx)
		if err != nil {
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}

		_, err = GetAccessToken(ctx)
		if err != nil {
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}

		if ok := c.granter(Role(session.Role)); !ok {
			err := merry.New("permission denied")
			c.json.RenderError(ctx, w, http.StatusForbidden, errs.APIForbidden, err)
			return
		}
		next.ServeHTTPC(ctx, w, r)
	})
}

func NewAdminOnly(j *snakepit.JSON) func(next chi.Handler) chi.Handler {
	gate := &Gate{
		json: j,
		granter: func(role Role) bool {
			return role == Admin
		},
	}
	return gate.middleware
}

func NewAuthenticatedOnly(j *snakepit.JSON) func(next chi.Handler) chi.Handler {
	gate := &Gate{
		json: j,
		granter: func(role Role) bool {
			return len(role) != 0 //&& role != Guest
		},
	}
	return gate.middleware
}

// func NewPublic(j *snakepit.JSON) func(next chi.Handler) chi.Handler {
// 	gate := &Gate{
// 		json: j,
// 		granter: func(role Role) bool {
// 			return len(role) != 0
// 		},
// 	}
// 	return gate.middleware
// }
<commit_msg>Old comments removed.<commit_after>package middlewares

import (
	"net/http"

	"github.com/solher/snakepit-seed/errs"

	"github.com/ansel1/merry"
	"github.com/pressly/chi"
	"github.com/solher/snakepit"
	"golang.org/x/net/context"
)

type Role string

const (
	Admin, Developer, User Role = "ADMIN", "DEVELOPER", "USER"
)

type Gate struct {
	json    *snakepit.JSON
	granter func(role Role) bool
}

func NewGate(j *snakepit.JSON, g func(role Role) bool) func(next chi.Handler) chi.Handler {
	gate := &Gate{json: j, granter: g}
	return gate.middleware
}

func (c *Gate) middleware(next chi.Handler) chi.Handler {
	return chi.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		session, err := GetCurrentSession(ctx)
		if err != nil {
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}
		if len(session.Role) == 0 {
			err := merry.New("empty role in context")
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}

		_, err = GetCurrentUser(ctx)
		if err != nil {
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}

		_, err = GetAccessToken(ctx)
		if err != nil {
			c.json.RenderError(ctx, w, http.StatusUnauthorized, errs.APIUnauthorized, err)
			return
		}

		if ok := c.granter(Role(session.Role)); !ok {
			err := merry.New("permission denied")
			c.json.RenderError(ctx, w, http.StatusForbidden, errs.APIForbidden, err)
			return
		}
		next.ServeHTTPC(ctx, w, r)
	})
}

func NewAdminOnly(j *snakepit.JSON) func(next chi.Handler) chi.Handler {
	gate := &Gate{
		json: j,
		granter: func(role Role) bool {
			return role == Admin
		},
	}
	return gate.middleware
}

func NewAuthenticatedOnly(j *snakepit.JSON) func(next chi.Handler) chi.Handler {
	gate := &Gate{
		json: j,
		granter: func(role Role) bool {
			return len(role) != 0
		},
	}
	return gate.middleware
}
<|endoftext|>


<commit_before>package builder

import (
	"log"

	"github.com/cloudfoundry-incubator/executor/log_streamer_factory"
	"github.com/cloudfoundry-incubator/garden/warden"

	"github.com/winston-ci/prole/api/builds"
)

type SourceFetcher interface {
	Fetch(source builds.BuildSource) (directory string, err error)
}

type ImageFetcher interface {
	Fetch(name string) (id string, err error)
}

type Builder struct {
	sourceFetcher      SourceFetcher
	wardenClient       warden.Client
	logStreamerFactory log_streamer_factory.LogStreamerFactory
}

func NewBuilder(
	sourceFetcher SourceFetcher,
	wardenClient warden.Client,
	logStreamerFactory log_streamer_factory.LogStreamerFactory,
) *Builder {
	return &Builder{
		sourceFetcher:      sourceFetcher,
		wardenClient:       wardenClient,
		logStreamerFactory: logStreamerFactory,
	}
}

func (builder *Builder) Build(build *builds.Build) (bool, error) {
	logStreamer := builder.logStreamerFactory(build.LogConfig)

	log.Println("fetching source")

	fetchedSource, err := builder.sourceFetcher.Fetch(build.Source)
	if err != nil {
		return false, err
	}

	log.Println("creating container")

	container, err := builder.wardenClient.Create(warden.ContainerSpec{
		RootFSPath: "image:" + build.Image,
	})
	if err != nil {
		return false, err
	}

	log.Println("copying in")

	err = container.CopyIn(fetchedSource+"/", build.Source.Path+"/")
	if err != nil {
		return false, err
	}

	log.Println("running")

	_, stream, err := container.Run(warden.ProcessSpec{Script: build.Script})
	if err != nil {
		return false, err
	}

	log.Println("streaming")

	succeeded := false

	for chunk := range stream {
		if chunk.ExitStatus != nil {
			succeeded = *chunk.ExitStatus == 0
		}

		switch chunk.Source {
		case warden.ProcessStreamSourceStdout:
			logStreamer.Stdout().Write(chunk.Data)
		case warden.ProcessStreamSourceStderr:
			logStreamer.Stderr().Write(chunk.Data)
		}

		logStreamer.Flush()
	}

	return succeeded, nil
}
<commit_msg>don't flush after every log message<commit_after>package builder

import (
	"log"

	"github.com/cloudfoundry-incubator/executor/log_streamer_factory"
	"github.com/cloudfoundry-incubator/garden/warden"

	"github.com/winston-ci/prole/api/builds"
)

type SourceFetcher interface {
	Fetch(source builds.BuildSource) (directory string, err error)
}

type ImageFetcher interface {
	Fetch(name string) (id string, err error)
}

type Builder struct {
	sourceFetcher      SourceFetcher
	wardenClient       warden.Client
	logStreamerFactory log_streamer_factory.LogStreamerFactory
}

func NewBuilder(
	sourceFetcher SourceFetcher,
	wardenClient warden.Client,
	logStreamerFactory log_streamer_factory.LogStreamerFactory,
) *Builder {
	return &Builder{
		sourceFetcher:      sourceFetcher,
		wardenClient:       wardenClient,
		logStreamerFactory: logStreamerFactory,
	}
}

func (builder *Builder) Build(build *builds.Build) (bool, error) {
	logStreamer := builder.logStreamerFactory(build.LogConfig)

	log.Println("fetching source")

	fetchedSource, err := builder.sourceFetcher.Fetch(build.Source)
	if err != nil {
		return false, err
	}

	log.Println("creating container")

	container, err := builder.wardenClient.Create(warden.ContainerSpec{
		RootFSPath: "image:" + build.Image,
	})
	if err != nil {
		return false, err
	}

	log.Println("copying in")

	err = container.CopyIn(fetchedSource+"/", build.Source.Path+"/")
	if err != nil {
		return false, err
	}

	log.Println("running")

	_, stream, err := container.Run(warden.ProcessSpec{Script: build.Script})
	if err != nil {
		return false, err
	}

	log.Println("streaming")

	succeeded := false

	for chunk := range stream {
		if chunk.ExitStatus != nil {
			logStreamer.Flush()
			succeeded = *chunk.ExitStatus == 0
			break
		}

		switch chunk.Source {
		case warden.ProcessStreamSourceStdout:
			logStreamer.Stdout().Write(chunk.Data)
		case warden.ProcessStreamSourceStderr:
			logStreamer.Stderr().Write(chunk.Data)
		}
	}

	return succeeded, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 The Prometheus 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 hetzner

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/pkg/errors"
	"github.com/prometheus/common/config"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/discovery/refresh"
	"github.com/prometheus/prometheus/discovery/targetgroup"
)

const (
	hetznerRobotLabelPrefix    = hetznerLabelPrefix + "robot_"
	hetznerLabelRobotProduct   = hetznerRobotLabelPrefix + "product"
	hetznerLabelRobotCancelled = hetznerRobotLabelPrefix + "cancelled"
)

// Discovery periodically performs Hetzner Robot requests. It implements
// the Discoverer interface.
type robotDiscovery struct {
	*refresh.Discovery
	client   *http.Client
	port     int
	endpoint string
}

// newRobotDiscovery returns a new robotDiscovery which periodically refreshes its targets.
func newRobotDiscovery(conf *SDConfig, logger log.Logger) (*robotDiscovery, error) {
	d := &robotDiscovery{
		port:     conf.Port,
		endpoint: conf.robotEndpoint,
	}

	rt, err := config.NewRoundTripperFromConfig(conf.HTTPClientConfig, "hetzner_sd", false, false)
	if err != nil {
		return nil, err
	}
	d.client = &http.Client{
		Transport: rt,
		Timeout:   time.Duration(conf.RefreshInterval),
	}

	return d, nil
}
func (d *robotDiscovery) refresh(ctx context.Context) ([]*targetgroup.Group, error) {
	resp, err := d.client.Get(d.endpoint + "/server")
	if err != nil {
		return nil, err
	}

	defer func() {
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}()

	if resp.StatusCode/100 != 2 {
		return nil, errors.Errorf("non 2xx status '%d' response during hetzner service discovery with role robot", resp.StatusCode)
	}

	var servers serversList
	err = json.NewDecoder(resp.Body).Decode(&servers)
	if err != nil {
		return nil, err
	}

	targets := make([]model.LabelSet, len(servers))
	for i, server := range servers {
		labels := model.LabelSet{
			hetznerLabelRole:           model.LabelValue(hetznerRoleRobot),
			hetznerLabelServerID:       model.LabelValue(strconv.Itoa(server.Server.ServerNumber)),
			hetznerLabelServerName:     model.LabelValue(server.Server.ServerName),
			hetznerLabelDatacenter:     model.LabelValue(strings.ToLower(server.Server.Dc)),
			hetznerLabelPublicIPv4:     model.LabelValue(server.Server.ServerIP),
			hetznerLabelServerStatus:   model.LabelValue(server.Server.Status),
			hetznerLabelRobotProduct:   model.LabelValue(server.Server.Product),
			hetznerLabelRobotCancelled: model.LabelValue(fmt.Sprintf("%t", server.Server.Canceled)),

			model.AddressLabel: model.LabelValue(net.JoinHostPort(server.Server.ServerIP, strconv.FormatUint(uint64(d.port), 10))),
		}
		for _, subnet := range server.Server.Subnet {
			ip := net.ParseIP(subnet.IP)
			if ip.To4() == nil {
				labels[hetznerLabelPublicIPv6Network] = model.LabelValue(fmt.Sprintf("%s/%s", subnet.IP, subnet.Mask))
				break
			}

		}
		targets[i] = labels
	}
	return []*targetgroup.Group{{Source: "hetzner", Targets: targets}}, nil
}

type serversList []struct {
	Server struct {
		ServerIP     string `json:"server_ip"`
		ServerNumber int    `json:"server_number"`
		ServerName   string `json:"server_name"`
		Dc           string `json:"dc"`
		Status       string `json:"status"`
		Product      string `json:"product"`
		Canceled     bool   `json:"cancelled"`
		Subnet       []struct {
			IP   string `json:"ip"`
			Mask string `json:"mask"`
		} `json:"subnet"`
	} `json:"server"`
}
<commit_msg>Add UserAgent to Hetzner Robot SD (#8008)<commit_after>// Copyright 2020 The Prometheus 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 hetzner

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/pkg/errors"
	"github.com/prometheus/common/config"
	"github.com/prometheus/common/model"
	"github.com/prometheus/common/version"
	"github.com/prometheus/prometheus/discovery/refresh"
	"github.com/prometheus/prometheus/discovery/targetgroup"
)

const (
	hetznerRobotLabelPrefix    = hetznerLabelPrefix + "robot_"
	hetznerLabelRobotProduct   = hetznerRobotLabelPrefix + "product"
	hetznerLabelRobotCancelled = hetznerRobotLabelPrefix + "cancelled"
)

var userAgent = fmt.Sprintf("Prometheus/%s", version.Version)

// Discovery periodically performs Hetzner Robot requests. It implements
// the Discoverer interface.
type robotDiscovery struct {
	*refresh.Discovery
	client   *http.Client
	port     int
	endpoint string
}

// newRobotDiscovery returns a new robotDiscovery which periodically refreshes its targets.
func newRobotDiscovery(conf *SDConfig, logger log.Logger) (*robotDiscovery, error) {
	d := &robotDiscovery{
		port:     conf.Port,
		endpoint: conf.robotEndpoint,
	}

	rt, err := config.NewRoundTripperFromConfig(conf.HTTPClientConfig, "hetzner_sd", false, false)
	if err != nil {
		return nil, err
	}
	d.client = &http.Client{
		Transport: rt,
		Timeout:   time.Duration(conf.RefreshInterval),
	}

	return d, nil
}
func (d *robotDiscovery) refresh(ctx context.Context) ([]*targetgroup.Group, error) {
	req, err := http.NewRequest("GET", d.endpoint+"/server", nil)
	if err != nil {
		return nil, err
	}

	req.Header.Add("User-Agent", userAgent)

	resp, err := d.client.Do(req)
	if err != nil {
		return nil, err
	}

	defer func() {
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}()

	if resp.StatusCode/100 != 2 {
		return nil, errors.Errorf("non 2xx status '%d' response during hetzner service discovery with role robot", resp.StatusCode)
	}

	var servers serversList
	err = json.NewDecoder(resp.Body).Decode(&servers)
	if err != nil {
		return nil, err
	}

	targets := make([]model.LabelSet, len(servers))
	for i, server := range servers {
		labels := model.LabelSet{
			hetznerLabelRole:           model.LabelValue(hetznerRoleRobot),
			hetznerLabelServerID:       model.LabelValue(strconv.Itoa(server.Server.ServerNumber)),
			hetznerLabelServerName:     model.LabelValue(server.Server.ServerName),
			hetznerLabelDatacenter:     model.LabelValue(strings.ToLower(server.Server.Dc)),
			hetznerLabelPublicIPv4:     model.LabelValue(server.Server.ServerIP),
			hetznerLabelServerStatus:   model.LabelValue(server.Server.Status),
			hetznerLabelRobotProduct:   model.LabelValue(server.Server.Product),
			hetznerLabelRobotCancelled: model.LabelValue(fmt.Sprintf("%t", server.Server.Canceled)),

			model.AddressLabel: model.LabelValue(net.JoinHostPort(server.Server.ServerIP, strconv.FormatUint(uint64(d.port), 10))),
		}
		for _, subnet := range server.Server.Subnet {
			ip := net.ParseIP(subnet.IP)
			if ip.To4() == nil {
				labels[hetznerLabelPublicIPv6Network] = model.LabelValue(fmt.Sprintf("%s/%s", subnet.IP, subnet.Mask))
				break
			}

		}
		targets[i] = labels
	}
	return []*targetgroup.Group{{Source: "hetzner", Targets: targets}}, nil
}

type serversList []struct {
	Server struct {
		ServerIP     string `json:"server_ip"`
		ServerNumber int    `json:"server_number"`
		ServerName   string `json:"server_name"`
		Dc           string `json:"dc"`
		Status       string `json:"status"`
		Product      string `json:"product"`
		Canceled     bool   `json:"cancelled"`
		Subnet       []struct {
			IP   string `json:"ip"`
			Mask string `json:"mask"`
		} `json:"subnet"`
	} `json:"server"`
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"io/ioutil"
	"log"
	"math"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// Version of tprotect, need to remember to update this
const version = "0.1"

type PidPageFault map[int]int

var pgfaults = make(PidPageFault)

type Config struct {
	SleepInterval            time.Duration  // In seconds
	FaultThreshold           int            // Number of faults per SleepInterval
	ProcessScanningThreshold int            // Number of pagefaults between each process scanning when the protectios doesn't kick in
	CmdWhitelist             map[string]int // Whitelisted processes
	UnfreezePopRatio         int            // Ratio of POP compared to GET
}

// Global config (TODO: Add Config file)
var cfg = new(Config)

func (c *Config) SetDefaults() {
	c.SleepInterval = 3
	c.FaultThreshold = 5
	c.ProcessScanningThreshold = cfg.FaultThreshold * 5
	c.CmdWhitelist = map[string]int{"init": 1, "sshd": 1, "bash": 1, "xinit": 1, "X": 1, "chromium-browser": 1}
	c.UnfreezePopRatio = 5
}

func ScanProcesses() (worstpid int) {
	stat_files, err := filepath.Glob("/proc/*/stat")
	if err != nil {
		log.Fatal("failed get files to stat")
	}

	for pfile := range stat_files {
		fs, err := ioutil.ReadFile(stat_files[pfile])
		if err != nil {
			//fmt.Print(stat_files[pfile])
			continue
		}
		buf0 := strings.Split(stat_files[pfile], "/")[2]
		pid, err := strconv.Atoi(buf0)
		if pid == 0 {
			continue
		}
		if err != nil {
			log.Println("PID ERR: ", pid, buf0)
			continue
			// self - hoppa över? dennna redan här?
		}
		stats := strings.Split(string(fs), " ")
		majflt, err := strconv.Atoi(stats[11])
		if err != nil {
			log.Printf("Could not convert %s to an int", stats[11])
			continue
		}
		buf1 := strings.Split(stats[1][1:], "/")[0]
		cmd := strings.Split(buf1, ")")[0]
		cmd = strings.TrimSpace(cmd) // just in case
		mypid := os.Getpid()
		max := 0

		if majflt > 0 {
			prev := pgfaults[pid]
			pgfaults[pid] = majflt
			if majflt-prev > max {
				if k := cfg.CmdWhitelist[cmd]; k == 1 {
					continue
				}
				if pid == mypid {
					continue
				}
				max = majflt - prev
				worstpid = pid
			}
		}
	}

	return worstpid

}

func GetPageFaults() (int, error) {
	file, err := os.Open("/proc/vmstat")
	if err != nil {
		log.Fatal("Could not open /proc/vmstat")
	}
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		buf := scanner.Text()
		if strings.HasPrefix(buf, "pgmajfault") {
			pf, err := strconv.Atoi(buf[12:])
			if err != nil {
				log.Fatal("Could not parse vmstat file")
			}
			return pf, nil
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatal(err)
	}

	return 0, errors.New("Unable to parse /proc/vmstat")
}

func freeze_something(frozen_pids []int, num_freezes int) ([]int, int) {
	pid_to_freeze := ScanProcesses()
	if pid_to_freeze == 0 {
		return frozen_pids, num_freezes
	}
	frozen_pids = append(frozen_pids, pid_to_freeze)
	num_freezes += 1

	ptf, err := os.FindProcess(pid_to_freeze)
	if err != nil {
		// Process wannished?
		return frozen_pids, num_freezes
	}
	log.Println("freezing pid: ", pid_to_freeze)
	err = ptf.Signal(syscall.SIGSTOP)
	if err != nil {
		log.Println("error sending SIGSTOP: ", err)
	}
	return frozen_pids, num_freezes
}

func unfreeze_something(frozen_pids []int, num_unfreezes int) ([]int, int) {

	var pid_to_unfreeze int

	if len(frozen_pids) > 0 {
		x := math.Remainder(float64(num_unfreezes), float64(cfg.UnfreezePopRatio))

		if x > 0 {
			// Trick to POP from an slice.. :-/
			pid_to_unfreeze, frozen_pids = frozen_pids[len(frozen_pids)-1], frozen_pids[:len(frozen_pids)-1]
		} else {
			pid_to_unfreeze = frozen_pids[0]
			frozen_pids = frozen_pids[1:]
		}

		ptuf, err := os.FindProcess(pid_to_unfreeze)
		if err != nil {
			// Process wannished?
			return frozen_pids, num_unfreezes
		}
		log.Println("Unfreezing pid: ", pid_to_unfreeze)
		err = ptuf.Signal(syscall.SIGCONT)
		if err != nil {
			log.Println("error sending SIGCONT: ", err)
		}

	}

	return frozen_pids, num_unfreezes
}

func MainLoop() {
	last_observed_pagefaults, _ := GetPageFaults()
	last_scan_pagefaults := 0
	var frozen_pids []int
	var num_freezes int = 0
	var num_unfreezes int = 0

	// Handle signals, what we are trying to do is when we quit
	// we unfreeze the freezed processes so we don't leave with
	// lots of freezed processes..
	// TODO: Need more testing.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	go func() {
		for sig := range c {
			log.Printf("Got %v, unfreezing freezed processes..", sig)
			for p := range frozen_pids {
				pu, err := os.FindProcess(frozen_pids[p])
				if err != nil {
					log.Println("Could not unfreeze pid: ", frozen_pids[p], err)
				}
				log.Printf("%d unfreezed", frozen_pids[p])
				pu.Signal(syscall.SIGCONT)
			}
			os.Exit(0)
		}
	}()

	for {
		current_pagefaults, err := GetPageFaults()
		if err != nil {
			log.Fatal(err)
		}
		switch {
		case current_pagefaults-last_observed_pagefaults > cfg.FaultThreshold:
			frozen_pids, num_freezes = freeze_something(frozen_pids, num_freezes)
		case current_pagefaults-last_observed_pagefaults == 0:
			frozen_pids, num_unfreezes = unfreeze_something(frozen_pids, num_unfreezes)
		}

		if current_pagefaults-last_scan_pagefaults > cfg.ProcessScanningThreshold {
			last_observed_pagefaults = current_pagefaults
		}

		time.Sleep(cfg.SleepInterval * time.Second)
	}
}

func main() {
	log.Printf("tprotect v%s start", version)
	cfg.SetDefaults()
	MainLoop()
}
<commit_msg>Cleaned up, fixed names etc to be more go-like<commit_after>package main

import (
	"bufio"
	"errors"
	"io/ioutil"
	"log"
	"math"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// Version of tprotect, need to remember to update this
const version = "0.1"

// Type Map that holds the Pid = Pagefault data
type pidPageFault map[int]int

var pgfaults = make(pidPageFault)

// Config Holds the configuration parameters
type Config struct {
	SleepInterval            time.Duration  // In seconds
	FaultThreshold           int            // Number of faults per SleepInterval
	ProcessScanningThreshold int            // Number of pagefaults between each process scanning when the protectios doesn't kick in
	CmdWhitelist             map[string]int // Whitelisted processes
	UnfreezePopRatio         int            // Ratio of POP compared to GET
}

// Global config (TODO: Add Config file)
var cfg = new(Config)

// SetDefaults sets the default Configuration
func (c *Config) SetDefaults() {
	c.SleepInterval = 3
	c.FaultThreshold = 5
	c.ProcessScanningThreshold = cfg.FaultThreshold * 5
	c.CmdWhitelist = map[string]int{"init": 1, "sshd": 1, "bash": 1, "xinit": 1, "X": 1, "chromium-browser": 1}
	c.UnfreezePopRatio = 5
}

func scanProcesses() (worstpid int) {
	statFiles, err := filepath.Glob("/proc/*/stat")
	if err != nil {
		log.Fatal("failed get files to stat")
	}

	for pfile := range statFiles {
		fs, err := ioutil.ReadFile(statFiles[pfile])
		if err != nil {
			//fmt.Print(statFiles[pfile])
			continue
		}
		buf0 := strings.Split(statFiles[pfile], "/")[2]
		pid, err := strconv.Atoi(buf0)
		if pid == 0 {
			continue
		}
		if err != nil {
			log.Println("PID ERR: ", pid, buf0)
			continue
			// self - hoppa över? dennna redan här?
		}
		stats := strings.Split(string(fs), " ")
		majflt, err := strconv.Atoi(stats[11])
		if err != nil {
			log.Printf("Could not convert %s to an int", stats[11])
			continue
		}
		buf1 := strings.Split(stats[1][1:], "/")[0]
		cmd := strings.Split(buf1, ")")[0]
		cmd = strings.TrimSpace(cmd) // just in case
		mypid := os.Getpid()
		max := 0

		if majflt > 0 {
			prev := pgfaults[pid]
			pgfaults[pid] = majflt
			if majflt-prev > max {
				if k := cfg.CmdWhitelist[cmd]; k == 1 {
					continue
				}
				if pid == mypid {
					continue
				}
				max = majflt - prev
				worstpid = pid
			}
		}
	}

	return worstpid

}

func getPageFaults() (int, error) {
	file, err := os.Open("/proc/vmstat")
	defer file.Close()
	if err != nil {
		log.Fatal("Could not open /proc/vmstat")
	}
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		buf := scanner.Text()
		if strings.HasPrefix(buf, "pgmajfault") {
			pf, err := strconv.Atoi(buf[12:])
			if err != nil {
				log.Fatal("Could not parse vmstat file")
			}
			return pf, nil
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatal(err)
	}

	return 0, errors.New("Unable to parse /proc/vmstat")
}

func freezeSomething(frozenPids []int, numFreezes int) ([]int, int) {
	pidToFreeze := scanProcesses()
	if pidToFreeze == 0 {
		return frozenPids, numFreezes
	}
	frozenPids = append(frozenPids, pidToFreeze)
	numFreezes += 1

	ptf, err := os.FindProcess(pidToFreeze)
	if err != nil {
		// Process wannished?
		return frozenPids, numFreezes
	}
	log.Println("freezing pid: ", pidToFreeze)
	err = ptf.Signal(syscall.SIGSTOP)
	if err != nil {
		log.Println("error sending SIGSTOP: ", err)
	}
	return frozenPids, numFreezes
}

func unfreezeSomething(frozenPids []int, numUnFreezes int) ([]int, int) {

	var pidToUnFreeze int

	if len(frozenPids) > 0 {
		x := math.Remainder(float64(numUnFreezes), float64(cfg.UnfreezePopRatio))

		if x > 0 {
			// Trick to POP from an slice.. :-/
			pidToUnFreeze, frozenPids = frozenPids[len(frozenPids)-1], frozenPids[:len(frozenPids)-1]
		} else {
			pidToUnFreeze = frozenPids[0]
			frozenPids = frozenPids[1:]
		}

		ptuf, err := os.FindProcess(pidToUnFreeze)
		if err != nil {
			// Process wannished?
			return frozenPids, numUnFreezes
		}
		log.Println("Unfreezing pid: ", pidToUnFreeze)
		err = ptuf.Signal(syscall.SIGCONT)
		if err != nil {
			log.Println("error sending SIGCONT: ", err)
		}

	}

	return frozenPids, numUnFreezes
}

func mainLoop() {
	lastObservedPagefaults, _ := getPageFaults()
	lastScanPagefaults := 0
	var frozenPids []int
	var numFreezes int
	var numUnFreezes int

	// Handle signals, what we are trying to do is when we quit
	// we unfreeze the freezed processes so we don't leave with
	// lots of freezed processes..
	// TODO: Need more testing.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	go func() {
		for sig := range c {
			log.Printf("Got %v, unfreezing freezed processes..", sig)
			for p := range frozenPids {
				pu, err := os.FindProcess(frozenPids[p])
				if err != nil {
					log.Println("Could not unfreeze pid: ", frozenPids[p], err)
				}
				log.Printf("%d unfreezed", frozenPids[p])
				pu.Signal(syscall.SIGCONT)
			}
			os.Exit(0)
		}
	}()

	for {
		currentPagefaults, err := getPageFaults()
		if err != nil {
			log.Fatal(err)
		}
		switch {
		case currentPagefaults-lastObservedPagefaults > cfg.FaultThreshold:
			frozenPids, numFreezes = freezeSomething(frozenPids, numFreezes)
		case currentPagefaults-lastObservedPagefaults == 0:
			frozenPids, numUnFreezes = unfreezeSomething(frozenPids, numUnFreezes)
		}

		if currentPagefaults-lastScanPagefaults > cfg.ProcessScanningThreshold {
			lastObservedPagefaults = currentPagefaults
		}

		time.Sleep(cfg.SleepInterval * time.Second)
	}
}

func main() {
	log.Printf("tprotect v%s start", version)
	cfg.SetDefaults()
	mainLoop()
}
<|endoftext|>


<commit_before>package main

import "github.com/jpillora/go-ogle-analytics"

func fireEvent( cid, ds, aid string,
	ai, av string,
	eType, eSubType, eId string,
	eState string,
	eLabel string, eVal int64 ) {

	client, err := ga.NewClient("UA-127388617-1")
	if err != nil {
		panic(err)
	}

	//client.ClientID("20eacc6c-e109-11e8-9f32-f2801f1b9fd1")
	client.ClientID(cid).
		DataSource(ds).
		ApplicationInstallerID( aid ).
		ApplicationID(ai).
		ApplicationVersion(av)

	client.ApplicationName(eSubType).
		DocumentTitle(eId)

	event := ga.NewEvent(eType, eState)
	event.
		Label(eLabel).
		Value(eVal)

	err = client.Send(event)
	//err = client.Send(ga.NewEvent("Foo", "Bar").Label("Bazz"))
	if err != nil {
		panic(err)
	}
	println("Event fired!")
}

func main() {

	/**
	//Install Events - Init or Running
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Install", "Linux-AMD64", "20eacc6c-0001-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	fireEvent( "20eacc6c-0002-0000-0000-f2801f1b9fd1", "AWS", "1.10",
			"OpenEBS", "0.8.0",
			"Install", "Linux-AMD64", "20eacc6c-0002-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	fireEvent( "20eacc6c-0003-0000-0000-f2801f1b9fd1", "GKE", "1.9.7",
			"OpenEBS", "0.7.1",
			"Install", "Linux-AMD64", "20eacc6c-0003-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	**/

	/**
	//Install Events - Running
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Install", "Linux-AMD64", "20eacc6c-0001-0000-0000-f2801f1b9fd1",
			"Running",
			"Nodes", 1 )
	**/

	/**
	//Install Events - Init (upgraded minikube openebs version)
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.3",
			"Install", "Linux-AMD64", "20eacc6c-0001-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	**/

	/**
	//Volume Provisioning Events
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Provisioned", "Jiva", "20eacc6c-0001-0001-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 10 )

	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Provisioned", "cStor", "20eacc6c-0001-0002-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 100 )

	//Volume Deprovisioning Events
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Deprovisioned", "Jiva", "20eacc6c-0001-0001-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 10 )

	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Deprovisioned", "cStor", "20eacc6c-0001-0002-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 100 )
	**/
}
<commit_msg>update openebs volume usage info<commit_after>package main

import "github.com/jpillora/go-ogle-analytics"

func fireVolumeEvent( clusterId, nodeType, k8sVersion string,
	volumeType, openebsVersion string,
	volumeEvent, pvName string,
	replicaCount string,
	capacityKey string, capacityVal int64 ) {

	client, err := ga.NewClient("UA-127388617-1")
	if err != nil {
		panic(err)
	}

	//client.ClientID("20eacc6c-e109-11e8-9f32-f2801f1b9fd1")
	client.ClientID(clusterId).
		CampaignSource("openebs-operator-kmova").
		CampaignContent(clusterId).
		CampaignID("campaign-id").
		CampaignName("campaign-name").
		CampaignKeyword("campaign-keyword").
		ApplicationID("OpenEBS").
		ApplicationVersion(openebsVersion).
		DataSource(nodeType).
		ApplicationName(volumeType).
		ApplicationInstallerID( k8sVersion ).
		DocumentTitle(pvName).
		DocumentHostName("hostname")

	event := ga.NewEvent(volumeEvent, replicaCount)
	event.Label(capacityKey)
	event.Value(capacityVal)

	err = client.Send(event)
	//err = client.Send(ga.NewEvent("Foo", "Bar").Label("Bazz"))
	if err != nil {
		panic(err)
	}
	println("Event fired!")
}

func main() {
	clusterId := "2ea928fa-b83c-463b-954c-d22b49c60da6"
	nodeType := "ubuntu 18.04.4 lts, 5.0.0-1032-gke"
	k8sVersion := "v1.15.11-gke.13"
	//k8sArch := ""
	openebsVersion := "1.11.0-f1152d3"

	volumeEvent := "volume-provision"
	volumeType := "jiva"
	pvName := "pvc-1827cb8d-8d45-4316-94ff-168ea312be61"
	replicaCount := "replica:3"

	fireVolumeEvent( clusterId, nodeType, k8sVersion,
			volumeType, openebsVersion,
			volumeEvent, pvName,
			replicaCount,
			"capacity", 2 )

	/**
	//Install Events - Init or Running
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Install", "Linux-AMD64", "20eacc6c-0001-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	fireEvent( "20eacc6c-0002-0000-0000-f2801f1b9fd1", "AWS", "1.10",
			"OpenEBS", "0.8.0",
			"Install", "Linux-AMD64", "20eacc6c-0002-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	fireEvent( "20eacc6c-0003-0000-0000-f2801f1b9fd1", "GKE", "1.9.7",
			"OpenEBS", "0.7.1",
			"Install", "Linux-AMD64", "20eacc6c-0003-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	**/

	/**
	//Install Events - Running
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Install", "Linux-AMD64", "20eacc6c-0001-0000-0000-f2801f1b9fd1",
			"Running",
			"Nodes", 1 )
	**/

	/**
	//Install Events - Init (upgraded minikube openebs version)
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.3",
			"Install", "Linux-AMD64", "20eacc6c-0001-0000-0000-f2801f1b9fd1",
			"Init",
			"Nodes", 1 )
	**/

	/**
	//Volume Provisioning Events
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Provisioned", "Jiva", "20eacc6c-0001-0001-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 10 )

	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Provisioned", "cStor", "20eacc6c-0001-0002-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 100 )

	//Volume Deprovisioning Events
	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Deprovisioned", "Jiva", "20eacc6c-0001-0001-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 10 )

	fireEvent( "20eacc6c-0001-0000-0000-f2801f1b9fd1", "Minikube", "1.12",
			"OpenEBS", "0.8.2",
			"Volume Deprovisioned", "cStor", "20eacc6c-0001-0002-0000-f2801f1b9fd1",
			"Success",
			"Capacity", 100 )
	**/
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/aarondl/pack"
	"log"
)

const (
	initialStackSize        = 20
	kidOffset               = 32
	stackIndexMask   uint64 = 0xFFFFFFFF
)

// bitFilter is a type used to filter versions based on their index in an array
// and an easy way to do cumulative versioning.
type bitFilter uint64

// Set sets a bit in the filter.
func (b bitFilter) Set(index uint) bitFilter {
	return bitFilter((1 << uint64(index)) | uint64(b))
}

// IsSet checks if a a bit in the filter is set.
func (b bitFilter) IsSet(index uint) bool {
	return 0 != ((1 << uint64(index)) & uint64(b))
}

// Clear turns off a bit in the filter.
func (b bitFilter) Clear(index uint) bitFilter {
	return bitFilter(^(1 << uint64(index)) & uint64(b))
}

// Add is the union of two bitFilters.
func (b bitFilter) Add(a bitFilter) bitFilter {
	return bitFilter(uint64(a) | uint64(b))
}

// versionProvider allows us to look up available versions for each package
// the array returned must be in sorted order for the best result from the
// solver as it assumes [0] is the latest version, and [1]... is less than that.
//
// GetVersions: Get the versions only from the dependency information.
// GetGraphs: Get a list of graphs showing dependency information for each
// version of the package.
type versionProvider interface {
	GetVersions(string) []*pack.Version
	GetGraph(string, *pack.Version) *depgraph
}

// stacknode helps to emulate recursion and perform safejumps.
type stacknode struct {
	kid     int
	version int
	current *depnode
	parent  *depnode
}

// savestate is the state of the algorithm at an activation point.
type savestate struct {
	kid     int
	version int
	index   int
	current *depnode
	stack   []stacknode
}

// activation is the details of a packages activation.
type activation struct {
	name    string
	version *pack.Version
	states  []*savestate
	filter  bitFilter
}

// String is used to debug activations.
func (a activation) String() string {
	var buf bytes.Buffer
	buf.WriteString(a.name)
	if a.version != nil {
		buf.WriteRune(space)
		buf.WriteString(a.version.String())
	}
	return buf.String()
}

/*
solve a dependency graph.

This algorithm is a depth first search with backjumping to resolve conflicts.

Possible optimization: don't attempt a new version of a package unless it's
dependencies have changed.
*/
func (g *depgraph) solve(vp versionProvider) error {
	if len(g.head.kids) == 0 {
		return nil
	}

	var current, parent *depnode = g.head, nil
	var stack = make([]stacknode, 0, initialStackSize) // Avoid allocations
	var si, kid = -1, 0
	var activations []*activation
	var versions = make(map[string][]*pack.Version)
	var active *activation
	var version *pack.Version
	var vs []*pack.Version
	var filter bitFilter
	var ok bool
	var verbose = true // Replace by flag.

	for i := 0; i < 10; i++ {
		name := current.d.Name
		if verbose {
			log.Println("Current:", current.d)
		}

		if current == g.head {
			if kid >= len(current.kids) {
				if verbose {
					log.Println("Success!")
				}
				return nil
			} else {
				goto NEXT
			}
		}

		// Skip Activation if we're on any child other than 0.
		if kid != 0 {
			goto NEXT
		}

		// Fetch Versions for current.
		vs = nil
		if vs, ok = versions[name]; !ok {
			vstmp := vp.GetVersions(current.d.Name)
			vs = make([]*pack.Version, len(vstmp))
			copy(vs, vstmp)
			versions[name] = vs
			if verbose {
				log.Printf("Fetched Versions")
			}
		}

		if verbose {
			log.Println("Versions:", vs)
		}

		// Weed out versions.
		filter = 0
		for j := 0; j < len(vs); j++ {
			for _, con := range current.d.Constraints {
				if !vs[j].Satisfies(con.Operator, con.Version) {
					if verbose {
						log.Println("Removing unacceptable version:", vs[j])
					}
					filter.Set(uint(j))
				}
			}
		}

		// Check for activeness.
		active = nil
		for j := 0; j < len(activations); j++ {
			if activations[j].name == name {
				active = activations[j]
				break
			}
		}

		version = nil
		if active != nil {
			if verbose {
				log.Println("Found activation:", active)
			}

			for _, con := range current.d.Constraints {
				if !active.version.Satisfies(con.Operator, con.Version) {
					// We've found a problem.
					log.Printf("Conflict (%v): %v fails new constraint: %v%v",
						name, active.version, con.Operator.String(),
						con.Version)
					return fmt.Errorf("Oh noes.")
				}
			}

			// Check that we comply with the currently active.
			active.states = append(active.states, &savestate{
				kid:     kid,
				version: 0, //Remove?
				index:   si,
				current: current,
				stack:   nil,
			})
			active.filter = active.filter.Add(filter)
		} else {
			if verbose {
				log.Println("Not activated:", name)
			}
			// Find a suitable version
			for j := 0; version == nil && j < len(vs); j++ {
				if active != nil && active.filter.IsSet(uint(j)) {
					continue
				}
				if len(current.d.Constraints) == 0 {
					version = vs[j]
					break
				}
				for _, con := range current.d.Constraints {
					if vs[j].Satisfies(con.Operator, con.Version) {
						version = vs[j]
						break
					}
				}
			}

			activations = append(activations, &activation{
				name:    name,
				version: version,
				filter:  filter,
				states: []*savestate{{
					kid:     kid,
					version: 0, //Remove?
					index:   si,
					current: current,
					stack:   nil,
				}},
			})
			current.v = version
		}

		if verbose {
			log.Printf("Added: %v %v to activations", name, version)
			log.Println("Activations:", activations)
		}

		if verbose {
			log.Println("Fetching Dependencies for:", name)
		}

		current.kids = vp.GetGraph(name, version).head.kids
		if verbose {
			var b bytes.Buffer
			for _, dep := range current.kids {
				b.WriteString(dep.d.String())
				b.WriteRune(space)
			}
			log.Println("Got dependencies:", b.String())
		}

		// Do the next node
	NEXT:
		// Pop off the stack back to parent.
		if kid < len(current.kids) {
			// Push us on to stack, go into child.
			if verbose {
				log.Println("Pushing:", name, kid)
			}
			stack = append(stack, stacknode{
				kid:     kid,
				version: 0, //remove?
				current: current,
				parent:  parent,
			})
			parent = current
			current = current.kids[kid]
			kid = 0
			si++
			continue
		}

		sn := stack[si]
		parent = sn.parent
		current = sn.current
		if verbose {
			log.Println("Popping:", current.d.Name)
		}
		kid = sn.kid
		kid++
		stack = stack[:len(stack)-1]
		si--

		// Try activating:
		// If activated:
		//  If conflict:
		//    backjump to first activated parent?
		// Else:
		//  Create activation with list of versions/dependencies.
		// Remove non-compatible versions with the current node.
		// Choose highest version as current activated.
	}

	return nil
}

/*
	var verbose = true // Move to flag

	var stack = make([]stacknode, 0, initialStackSize) // Avoid allocations
	var index = 0
	var kid, version int
	var backjump *savestate
	var current *depnode
	var vs []*pack.Version
	var active = make(map[string]*activation)

	current = g.head

	for i := 0; i < 40; i++ {
		if verbose {
			log.Printf("Eval: %s (%v, %v)\n", current.d.Name, kid, version)
		}

		// Have we run out of dependencies to resolve?
		if kid >= len(current.kids) {
			if verbose {
				log.Println("Ran out of children...")
			}
			// Jump up stack.
			if index > 0 {
				index--
				current = stack[index].depnode
				//version = stack[index].version
				stack[index].kid++
				kid = stack[index].kid
				if verbose {
					log.Printf("Pop: %s (%v, %v)\n", current.d.Name,
						kid, version)
				}
				continue
			}
			if verbose {
				log.Println("Nothing left to do, should be solved.")
			}
			// We did it!!!!
			return true
		}

		// Try to activate child
		curkid := current.kids[kid]
		name := curkid.d.Name
		log.Println("Attempting child activation:", name)

		// Check if already activated
		if act, ok := active[name]; ok && act.v != nil {
			var found bool
			for _, con := range curkid.d.Constraints {
				if act.v.Satisfies(con.Operator, con.Version) {
					found = true
					break
				}
			}

			if found {
				// Add ourselves to the activators list.
				save := &savestate{
					kid, version, index, current,
					make([]stacknode, index+1),
				}
				copy(save.stack, stack)
				act.states = append(act.states, save)

				if verbose {
					log.Println(name, "satisfied by previous activation:",
						act.v, act.states)
				}
			} else {
				// Backjump
				backjump = act.states[len(act.states)-1]
				if verbose {
					log.Printf("Found conflict: %s (%v)\n", name, act.v)
					log.Print("Previous states: ")
					for _, a := range act.states {
						log.Printf("%v ", a)
					}
				}
				act.states = act.states[:len(act.states)-1]
				act.v = nil

				current, kid, version, index = backjump.current,
					backjump.kid, backjump.version, backjump.index
				copy(stack, backjump.stack)
				if verbose {
					log.Println("Backjumping:", index, kid)
				}
				version++
				continue
			}
		}

		// Get versions
		vs = vp.GetVersions(name)
		if verbose {
			log.Printf("Versions: %s %v\n", name, vs)
			log.Println("Iterating from:", version)
		}
		// Each version
		for ; version < len(vs); version++ {
			// Each constraint
			ver := vs[version]
			for _, con := range curkid.d.Constraints {
				if ver.Satisfies(con.Operator, con.Version) {
					if verbose {
						log.Println("Satisfactory Version:", curkid.d, ver)
					}
					curkid.v = ver
				}
				if curkid.v != nil {
					log.Println("No need for more constraint checks... Breaking.")
					break
				}
			}
			if len(curkid.d.Constraints) == 0 {
				curkid.v = ver
			}
			if curkid.v != nil {
				log.Println("No need to check more versions... Breaking")
				break
			}
		}

		if curkid.v == nil {
			// No version found to satisfy.
			if verbose {
				log.Println("No versions available to satisfy:", curkid.d)
			}
			return false
		} else {
			// Activate
			if verbose {
				log.Printf("Activating: %s %v (%v, %v)\n",
					curkid.d.Name, curkid.v, version, index)
			}

			// Add save state info to activation
			save := &savestate{
				kid, version, index, current,
				make([]stacknode, index+1),
			}
			copy(save.stack, stack)

			if act, ok := active[name]; ok {
				act.states = []*savestate{save}
				act.v = curkid.v
			} else {
				active[name] = &activation{[]*savestate{save}, curkid.v}
			}

			// Pull in child dependencies on activation.
			graphs := vp.GetGraphs(name)
			found := false
			for i := 0; i < len(graphs); i++ {
				if graphs[i].head.v.Satisfies(pack.Equal, curkid.v) {
					found = true
					curkid.kids = graphs[i].head.kids
					if verbose {
						log.Println("Setting kids:", len(curkid.kids))
					}
					break
				}
			}

			if !found {
				if verbose {
					log.Println("Dependency graph missing for: %v %v",
						curkid.d.Name, curkid.v)
				}
				return false
			}
		}

		if len(curkid.kids) > 0 {
			// Push current on to stack, make the curkid the new current.
			if verbose {
				log.Println("Has kids:", len(curkid.kids))
				log.Printf("Push: %s (%v, %v)\n", current.d.Name, kid, version)
			}
			stack = append(stack, stacknode{kid, version, current})
			kid = 0
			version = 0
			current = curkid
			index++
		} else {
			// Continue through all kids.
			kid++
			if verbose {
				log.Println("Next kid:", kid)
			}
		}
	}

	return false
}*/
<commit_msg>Have conflict resolution flow in a good spot.<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/aarondl/pack"
	"log"
)

const (
	initialStackSize        = 20
	kidOffset               = 32
	stackIndexMask   uint64 = 0xFFFFFFFF
)

// bitFilter is a type used to filter versions based on their index in an array
// and an easy way to do cumulative versioning.
type bitFilter uint64

// Set sets a bit in the filter.
func (b bitFilter) Set(index uint) bitFilter {
	return bitFilter((1 << uint64(index)) | uint64(b))
}

// IsSet checks if a a bit in the filter is set.
func (b bitFilter) IsSet(index uint) bool {
	return 0 != ((1 << uint64(index)) & uint64(b))
}

// Clear turns off a bit in the filter.
func (b bitFilter) Clear(index uint) bitFilter {
	return bitFilter(^(1 << uint64(index)) & uint64(b))
}

// Add is the union of two bitFilters.
func (b bitFilter) Add(a bitFilter) bitFilter {
	return bitFilter(uint64(a) | uint64(b))
}

// versionProvider allows us to look up available versions for each package
// the array returned must be in sorted order for the best result from the
// solver as it assumes [0] is the latest version, and [1]... is less than that.
//
// GetVersions: Get the versions only from the dependency information.
// GetGraphs: Get a list of graphs showing dependency information for each
// version of the package.
type versionProvider interface {
	GetVersions(string) []*pack.Version
	GetGraph(string, *pack.Version) *depgraph
}

// stacknode helps to emulate recursion and perform safejumps.
type stacknode struct {
	kid     int
	version int
	ai      int
	current *depnode
	parent  *depnode
}

// savestate is the state of the algorithm at an activation point.
type savestate struct {
	kid     int
	version int
	si      int
	ai      int
	current *depnode
	stack   []stacknode
}

// activation is the details of a packages activation.
type activation struct {
	name    string
	version *pack.Version
	state   *savestate
	filter  bitFilter
}

// String is used to debug activations.
func (a activation) String() string {
	var buf bytes.Buffer
	buf.WriteString(a.name)
	if a.version != nil {
		buf.WriteRune(space)
		buf.WriteString(a.version.String())
	}
	return buf.String()
}

/*
solve a dependency graph. This algorithm is a depth first search with
backjumping to resolve conflicts.
*/
func (g *depgraph) solve(vp versionProvider) error {
	if len(g.head.kids) == 0 {
		return nil
	}

	var current, parent *depnode = g.head, nil
	var stack = make([]stacknode, 0, initialStackSize) // Avoid allocations
	var si, ai, kid = -1, -1, 0
	var activations []*activation
	var active *activation
	var versions = make(map[string][]*pack.Version)
	var version *pack.Version
	var vs []*pack.Version
	var vi int
	var noversions bool
	var filter bitFilter
	var ok bool
	var conflicts = make([]string, 0)

	var verbose = true // Replace by flag.

	for i := 0; i < 20; i++ {
		name := current.d.Name
		if verbose {
			log.Println("Current:", current.d)
		}

		if current == g.head {
			if kid >= len(current.kids) {
				if verbose {
					log.Println("Success!")
					log.Println(activations)
				}
				return nil
			} else {
				goto NEXT
			}
		}

		// Skip Activation if we're on any child other than 0.
		if kid != 0 {
			goto NEXT
		}

		// Fetch Versions for current.
		vs = nil
		if vs, ok = versions[name]; !ok {
			vstmp := vp.GetVersions(current.d.Name)
			vs = make([]*pack.Version, len(vstmp))
			copy(vs, vstmp)
			versions[name] = vs
			if verbose {
				log.Printf("Fetched Versions")
			}
		}

		if verbose {
			log.Println("Versions:", vs)
		}

		// Weed out versions.
		filter = 0
		noversions = true
		for j := 0; j < len(vs); j++ {
			for _, con := range current.d.Constraints {
				if !vs[j].Satisfies(con.Operator, con.Version) {
					if verbose {
						log.Println("Removing unacceptable version:", vs[j])
					}
					filter.Set(uint(j))
				} else {
					noversions = false
				}
			}
		}

		if noversions {
			if parent == g.head {
				return fmt.Errorf("No versions to satisfy root dependency: %v", current.d)
			} else {
				// Jump up the stack, conflict, something!
				return fmt.Errorf("No versions satisfy: %v", current.d)
			}
		}

		// Check for activeness. The first activation will always serve as the
		// main activation point, with the others simply being save points.
		active = nil
		for j := 0; j < len(activations); j++ {
			if activations[j].name == name {
				active = activations[j]
				break
			}
		}

		version = nil
		noversions = false
		if active != nil {
			if verbose {
				log.Println("Found activation:", active)
			}

			// Check that we comply with the currently active.
			for _, con := range current.d.Constraints {
				if !active.version.Satisfies(con.Operator, con.Version) {
					// We've found a problem.
					log.Printf("Conflict (%v): %v fails new constraint: %v%v",
						name, active.version, con.Operator.String(),
						con.Version)

					conflicts = append(conflicts, name)
				}
			}

			// If there's a conflict
			if len(conflicts) > 0 {
				if parent == g.head {
					return fmt.Errorf("Reached the top of the stack!")
				}

				// We can still climb the stack, try it.
				sn := stack[si]
				parent = sn.parent
				current = sn.current
				if verbose {
					log.Println("Popping:", current.d.Name)
				}
				ai = sn.ai
				vi = sn.version
				kid = 0
				stack = stack[:len(stack)-1]
				si--
				vi++
				if verbose {
					log.Println(activations)
					log.Println(activations[:ai])
					log.Println("Snipping activations back to:", ai) // TODO: REMOVE
				}
				activations = activations[:ai]
				continue
			}

			// Add the current filter to the primary activation.
			active.filter = active.filter.Add(filter)
		} else {
			if verbose {
				log.Println("Not activated:", name)
			}
			// Find a suitable version
			for ; version == nil && vi < len(vs); vi++ {
				if active != nil && active.filter.IsSet(uint(vi)) {
					continue
				}
				if len(current.d.Constraints) == 0 {
					version = vs[vi]
					break
				}
				for _, con := range current.d.Constraints {
					if vs[vi].Satisfies(con.Operator, con.Version) {
						version = vs[vi]
						break
					}
				}
			}

			if version == nil {
				// No version could be found, this is a conflict of sorts.
				return fmt.Errorf("No versions could be found for: %v", current.d)
			}
		}

		// Add ourselves to the list of activators.
		activations = append(activations, &activation{
			name:    name,
			version: version,
			filter:  filter,
			state: &savestate{
				kid:     kid,
				version: vi,
				si:      si,
				ai:      ai,
				current: current,
				stack:   nil,
			},
		})
		ai++
		current.v = version

		if verbose {
			log.Printf("Added: %v %v to activations", name, version)
			log.Println("Activations:", activations)
		}

		if verbose {
			log.Println("Fetching Dependencies for:", name)
		}

		current.kids = vp.GetGraph(name, version).head.kids
		if verbose {
			var b bytes.Buffer
			for _, dep := range current.kids {
				b.WriteString(dep.d.String())
				b.WriteRune(space)
			}
			log.Println("Got dependencies:", b.String())
		}

	NEXT:
		// Push current on to stack, go into child.
		if kid < len(current.kids) {
			if verbose {
				log.Println("Pushing:", name, kid)
			}
			stack = append(stack, stacknode{
				kid:     kid,
				version: vi,
				ai:      ai,
				current: current,
				parent:  parent,
			})
			parent = current
			current = current.kids[kid]
			ai = len(activations) - 1
			kid = 0
			vi = 0
			si++
			continue
		}

		// Pop off the stack back to parent.
		sn := stack[si]
		parent = sn.parent
		current = sn.current
		if verbose {
			log.Println("Popping:", current.d.Name)
		}
		kid = sn.kid
		vi = sn.version
		ai = sn.ai
		stack = stack[:len(stack)-1]
		kid++
		si--

		// Try activating:
		// If activated:
		//  If conflict:
		//    backjump to first activated parent?
		// Else:
		//  Create activation with list of versions/dependencies.
		// Remove non-compatible versions with the current node.
		// Choose highest version as current activated.
	}

	return nil
}

/*
	var verbose = true // Move to flag

	var stack = make([]stacknode, 0, initialStackSize) // Avoid allocations
	var index = 0
	var kid, version int
	var backjump *savestate
	var current *depnode
	var vs []*pack.Version
	var active = make(map[string]*activation)

	current = g.head

	for i := 0; i < 40; i++ {
		if verbose {
			log.Printf("Eval: %s (%v, %v)\n", current.d.Name, kid, version)
		}

		// Have we run out of dependencies to resolve?
		if kid >= len(current.kids) {
			if verbose {
				log.Println("Ran out of children...")
			}
			// Jump up stack.
			if index > 0 {
				index--
				current = stack[index].depnode
				//version = stack[index].version
				stack[index].kid++
				kid = stack[index].kid
				if verbose {
					log.Printf("Pop: %s (%v, %v)\n", current.d.Name,
						kid, version)
				}
				continue
			}
			if verbose {
				log.Println("Nothing left to do, should be solved.")
			}
			// We did it!!!!
			return true
		}

		// Try to activate child
		curkid := current.kids[kid]
		name := curkid.d.Name
		log.Println("Attempting child activation:", name)

		// Check if already activated
		if act, ok := active[name]; ok && act.v != nil {
			var found bool
			for _, con := range curkid.d.Constraints {
				if act.v.Satisfies(con.Operator, con.Version) {
					found = true
					break
				}
			}

			if found {
				// Add ourselves to the activators list.
				save := &savestate{
					kid, version, index, current,
					make([]stacknode, index+1),
				}
				copy(save.stack, stack)
				act.states = append(act.states, save)

				if verbose {
					log.Println(name, "satisfied by previous activation:",
						act.v, act.states)
				}
			} else {
				// Backjump
				backjump = act.states[len(act.states)-1]
				if verbose {
					log.Printf("Found conflict: %s (%v)\n", name, act.v)
					log.Print("Previous states: ")
					for _, a := range act.states {
						log.Printf("%v ", a)
					}
				}
				act.states = act.states[:len(act.states)-1]
				act.v = nil

				current, kid, version, index = backjump.current,
					backjump.kid, backjump.version, backjump.index
				copy(stack, backjump.stack)
				if verbose {
					log.Println("Backjumping:", index, kid)
				}
				version++
				continue
			}
		}

		// Get versions
		vs = vp.GetVersions(name)
		if verbose {
			log.Printf("Versions: %s %v\n", name, vs)
			log.Println("Iterating from:", version)
		}
		// Each version
		for ; version < len(vs); version++ {
			// Each constraint
			ver := vs[version]
			for _, con := range curkid.d.Constraints {
				if ver.Satisfies(con.Operator, con.Version) {
					if verbose {
						log.Println("Satisfactory Version:", curkid.d, ver)
					}
					curkid.v = ver
				}
				if curkid.v != nil {
					log.Println("No need for more constraint checks... Breaking.")
					break
				}
			}
			if len(curkid.d.Constraints) == 0 {
				curkid.v = ver
			}
			if curkid.v != nil {
				log.Println("No need to check more versions... Breaking")
				break
			}
		}

		if curkid.v == nil {
			// No version found to satisfy.
			if verbose {
				log.Println("No versions available to satisfy:", curkid.d)
			}
			return false
		} else {
			// Activate
			if verbose {
				log.Printf("Activating: %s %v (%v, %v)\n",
					curkid.d.Name, curkid.v, version, index)
			}

			// Add save state info to activation
			save := &savestate{
				kid, version, index, current,
				make([]stacknode, index+1),
			}
			copy(save.stack, stack)

			if act, ok := active[name]; ok {
				act.states = []*savestate{save}
				act.v = curkid.v
			} else {
				active[name] = &activation{[]*savestate{save}, curkid.v}
			}

			// Pull in child dependencies on activation.
			graphs := vp.GetGraphs(name)
			found := false
			for i := 0; i < len(graphs); i++ {
				if graphs[i].head.v.Satisfies(pack.Equal, curkid.v) {
					found = true
					curkid.kids = graphs[i].head.kids
					if verbose {
						log.Println("Setting kids:", len(curkid.kids))
					}
					break
				}
			}

			if !found {
				if verbose {
					log.Println("Dependency graph missing for: %v %v",
						curkid.d.Name, curkid.v)
				}
				return false
			}
		}

		if len(curkid.kids) > 0 {
			// Push current on to stack, make the curkid the new current.
			if verbose {
				log.Println("Has kids:", len(curkid.kids))
				log.Printf("Push: %s (%v, %v)\n", current.d.Name, kid, version)
			}
			stack = append(stack, stacknode{kid, version, current})
			kid = 0
			version = 0
			current = curkid
			index++
		} else {
			// Continue through all kids.
			kid++
			if verbose {
				log.Println("Next kid:", kid)
			}
		}
	}

	return false
}*/
<|endoftext|>


<commit_before>package virt

import (
	"errors"
	"labix.org/v2/mgo/bson"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

func (vm *VM) Start() error {
	if out, err := exec.Command("/usr/bin/lxc-start", "--name", vm.String(), "--daemon").CombinedOutput(); err != nil {
		return commandError("lxc-start failed.", err, out)
	}
	exec.Command("/sbin/ifconfig", vm.VEth(), "mtu", "1476").Start() // error ignored
	return vm.WaitForState("RUNNING", time.Second)
}

func (vm *VM) Stop() error {
	if out, err := exec.Command("/usr/bin/lxc-stop", "--name", vm.String()).CombinedOutput(); err != nil {
		return commandError("lxc-stop failed.", err, out)
	}
	return vm.WaitForState("STOPPED", time.Second)
}

func (vm *VM) Shutdown() error {
	if out, err := exec.Command("/usr/bin/lxc-shutdown", "--name", vm.String()).CombinedOutput(); err != nil {
		if vm.GetState() != "STOPPED" {
			return commandError("lxc-shutdown failed.", err, out)
		}
	}
	vm.WaitForState("STOPPED", 5*time.Second) // may time out, then vm is force stopped
	return vm.Stop()
}

func (vm *VM) AttachCommand(uid int, tty string, command ...string) *exec.Cmd {
	args := []string{"--name", vm.String()}
	if tty != "" {
		args = append(args, "--tty", tty)
	}
	args = append(args, "--", "/usr/bin/sudo", "-i", "-u", "#"+strconv.Itoa(uid), "--")
	args = append(args, command...)
	cmd := exec.Command("/usr/bin/lxc-attach", args...)
	cmd.Env = []string{"TERM=xterm-256color"}
	return cmd
}

func (vm *VM) GetState() string {
	out, err := exec.Command("/usr/bin/lxc-info", "--name", vm.String(), "--state").CombinedOutput()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(out)[6:])
}

func (vm *VM) WaitForState(state string, timeout time.Duration) error {
	tryUntil := time.Now().Add(timeout)
	for vm.GetState() != state {
		if time.Now().After(tryUntil) {
			return errors.New("Timeout while waiting for VM state.")
		}
		time.Sleep(time.Second / 10)
	}
	return nil
}

func SendMessageToVMUsers(vmId bson.ObjectId, message string) error {
	cmd := exec.Command("/usr/bin/lxc-attach", "--name", VMName(vmId), "--", "/usr/bin/wall", "--nobanner")
	cmd.Stdin = strings.NewReader(message)
	if out, err := cmd.CombinedOutput(); err != nil {
		return commandError("wall failed.", err, out)
	}
	return nil
}
<commit_msg>os kite: Fixed ifconfig zombies.<commit_after>package virt

import (
	"errors"
	"labix.org/v2/mgo/bson"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

func (vm *VM) Start() error {
	if out, err := exec.Command("/usr/bin/lxc-start", "--name", vm.String(), "--daemon").CombinedOutput(); err != nil {
		return commandError("lxc-start failed.", err, out)
	}
	exec.Command("/sbin/ifconfig", vm.VEth(), "mtu", "1476").Run() // error ignored
	return vm.WaitForState("RUNNING", time.Second)
}

func (vm *VM) Stop() error {
	if out, err := exec.Command("/usr/bin/lxc-stop", "--name", vm.String()).CombinedOutput(); err != nil {
		return commandError("lxc-stop failed.", err, out)
	}
	return vm.WaitForState("STOPPED", time.Second)
}

func (vm *VM) Shutdown() error {
	if out, err := exec.Command("/usr/bin/lxc-shutdown", "--name", vm.String()).CombinedOutput(); err != nil {
		if vm.GetState() != "STOPPED" {
			return commandError("lxc-shutdown failed.", err, out)
		}
	}
	vm.WaitForState("STOPPED", 5*time.Second) // may time out, then vm is force stopped
	return vm.Stop()
}

func (vm *VM) AttachCommand(uid int, tty string, command ...string) *exec.Cmd {
	args := []string{"--name", vm.String()}
	if tty != "" {
		args = append(args, "--tty", tty)
	}
	args = append(args, "--", "/usr/bin/sudo", "-i", "-u", "#"+strconv.Itoa(uid), "--")
	args = append(args, command...)
	cmd := exec.Command("/usr/bin/lxc-attach", args...)
	cmd.Env = []string{"TERM=xterm-256color"}
	return cmd
}

func (vm *VM) GetState() string {
	out, err := exec.Command("/usr/bin/lxc-info", "--name", vm.String(), "--state").CombinedOutput()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(out)[6:])
}

func (vm *VM) WaitForState(state string, timeout time.Duration) error {
	tryUntil := time.Now().Add(timeout)
	for vm.GetState() != state {
		if time.Now().After(tryUntil) {
			return errors.New("Timeout while waiting for VM state.")
		}
		time.Sleep(time.Second / 10)
	}
	return nil
}

func SendMessageToVMUsers(vmId bson.ObjectId, message string) error {
	cmd := exec.Command("/usr/bin/lxc-attach", "--name", VMName(vmId), "--", "/usr/bin/wall", "--nobanner")
	cmd.Stdin = strings.NewReader(message)
	if out, err := cmd.CombinedOutput(); err != nil {
		return commandError("wall failed.", err, out)
	}
	return nil
}
<|endoftext|>


<commit_before>package dnsr

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/miekg/dns"
)

// DNS Resolution configuration.
var (
	Timeout             = 2000 * time.Millisecond
	TypicalResponseTime = 100 * time.Millisecond
	MaxRecursion        = 10
	MaxNameservers      = 4
	MaxIPs              = 2
)

// Resolver errors.
var (
	NXDOMAIN = fmt.Errorf("NXDOMAIN")

	ErrMaxRecursion = fmt.Errorf("maximum recursion depth reached: %d", MaxRecursion)
	ErrMaxIPs       = fmt.Errorf("maximum name server IPs queried: %d", MaxIPs)
	ErrNoARecords   = fmt.Errorf("no A records found for name server")
	ErrNoResponse   = fmt.Errorf("no responses received")
	ErrTimeout      = fmt.Errorf("timeout expired") // TODO: Timeouter interface? e.g. func (e) Timeout() bool { return true }
)

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache   *cache
	expire  bool
	timeout time.Duration
}

// New initializes a Resolver with the specified cache size.
func New(capacity int) *Resolver {
	return NewWithTimeout(capacity, Timeout)
}

// NewWithTimeout initializes a Resolver with the specified cache size and resolution timeout.
func NewWithTimeout(capacity int, timeout time.Duration) *Resolver {
	r := &Resolver{
		cache:   newCache(capacity, false),
		expire:  false,
		timeout: timeout,
	}
	return r
}

// NewExpiring initializes an expiring Resolver with the specified cache size.
func NewExpiring(capacity int) *Resolver {
	return NewExpiringWithTimeout(capacity, Timeout)
}

// NewExpiringWithTimeout initializes an expiring Resolved with the specified cache size and resolution timeout.
func NewExpiringWithTimeout(capacity int, timeout time.Duration) *Resolver {
	r := &Resolver{
		cache:   newCache(capacity, true),
		expire:  true,
		timeout: timeout,
	}
	return r
}

// Resolve calls ResolveErr to find DNS records of type qtype for the domain qname.
// For nonexistent domains (NXDOMAIN), it will return an empty, non-nil slice.
func (r *Resolver) Resolve(qname, qtype string) RRs {
	rrs, err := r.ResolveErr(qname, qtype)
	if err == NXDOMAIN {
		return emptyRRs
	}
	if err != nil {
		return nil
	}
	return rrs
}

// ResolveErr finds DNS records of type qtype for the domain qname.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveErr(qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

// ResolveCtx finds DNS records of type qtype for the domain qname using
// the supplied context. Requests may time out earlier if timeout is
// shorter than a deadline set in ctx.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveCtx(ctx context.Context, qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(ctx, r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

func (r *Resolver) resolve(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		return nil, ErrMaxRecursion
	}
	rrs, err := r.cacheGet(ctx, qname, qtype)
	if err != nil {
		return nil, err
	}
	if len(rrs) > 0 {
		return rrs, nil
	}
	logResolveStart(qname, qtype, depth)
	start := time.Now()
	rrs, err = r.iterateParents(ctx, qname, qtype, depth)
	logResolveEnd(qname, qtype, rrs, depth, start, err)
	return rrs, err
}

func (r *Resolver) iterateParents(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	chanRRs := make(chan RRs, MaxNameservers)
	chanErrs := make(chan error, MaxNameservers)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		// If we’re looking for [foo.com,NS], then move on to the parent ([com,NS])
		if pname == qname && qtype == "NS" {
			continue
		}

		// Only query TLDs against the root nameservers
		if pname == "." && dns.CountLabel(qname) != 1 {
			// fmt.Fprintf(os.Stderr, "Warning: non-TLD query at root: dig +norecurse %s %s\n", qname, qtype)
			return nil, nil
		}

		// Get nameservers
		nrrs, err := r.resolve(ctx, pname, "NS", depth)
		if err == NXDOMAIN || err == ErrTimeout || err == context.DeadlineExceeded {
			return nil, err
		}
		if err != nil {
			continue
		}

		// Check cache for specific queries
		if len(nrrs) > 0 && qtype != "" {
			rrs, err := r.cacheGet(ctx, qname, qtype)
			if err != nil {
				return nil, err
			}
			if len(rrs) > 0 {
				return rrs, nil
			}
		}

		// Query all nameservers in parallel
		count := 0
		for i := 0; i < len(nrrs) && count < MaxNameservers; i++ {
			nrr := nrrs[i]
			if nrr.Type != "NS" {
				continue
			}

			go func(host string) {
				rrs, err := r.exchange(ctx, host, qname, qtype, depth)
				if err != nil {
					chanErrs <- err
				} else {
					chanRRs <- rrs
				}
			}(nrr.Value)

			count++
		}

		// Wait for answer, error, or cancellation
		for ; count > 0; count-- {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case rrs := <-chanRRs:
				for _, nrr := range nrrs {
					if nrr.Name == qname {
						rrs = append(rrs, nrr)
					}
				}
				cancel() // stop any other work here before recursing
				return r.resolveCNAMEs(ctx, qname, qtype, rrs, depth)
			case err = <-chanErrs:
				if err == NXDOMAIN {
					return nil, err
				}
			}

		}

		// NS queries naturally recurse, so stop further iteration
		if qtype == "NS" {
			return nil, err
		}
	}

	return nil, ErrNoResponse
}

func (r *Resolver) exchange(ctx context.Context, host, qname, qtype string, depth int) (RRs, error) {
	count := 0
	arrs, err := r.resolve(ctx, host, "A", depth)
	if err != nil {
		return nil, err
	}
	for _, arr := range arrs {
		// FIXME: support AAAA records?
		if arr.Type != "A" {
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return nil, ErrMaxIPs
		}

		rrs, err := r.exchangeIP(ctx, host, arr.Value, qname, qtype, depth)
		if err == nil || err == NXDOMAIN || err == ErrTimeout {
			return rrs, err
		}

		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
	}

	return nil, ErrNoARecords
}

func (r *Resolver) exchangeIP(ctx context.Context, host, ip, qname, qtype string, depth int) (RRs, error) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeANY
	}
	var qmsg dns.Msg
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Synchronously query this DNS server
	start := time.Now()
	timeout := r.timeout // belt and suspenders, since ctx has a deadline from ResolveErr
	if dl, ok := ctx.Deadline(); ok {
		if start.After(dl.Add(-TypicalResponseTime)) { // bail if we can't finish in time (start is too close to deadline)
			return nil, ErrTimeout
		}
		timeout = dl.Sub(start)
	}

	client := &dns.Client{Timeout: timeout} // client must finish within remaining timeout
	rmsg, dur, err := client.Exchange(&qmsg, ip+":53")
	select {
	case <-ctx.Done(): // Finished too late
		logCancellation(host, &qmsg, rmsg, depth, dur, timeout)
		return nil, ctx.Err()
	default:
		logExchange(host, &qmsg, rmsg, depth, dur, timeout, err) // Log hostname instead of IP
	}
	if err != nil {
		return nil, err
	}

	// FIXME: cache NXDOMAIN responses responsibly
	if rmsg.Rcode == dns.RcodeNameError {
		var hasSOA bool
		if qtype == "NS" {
			for _, drr := range rmsg.Ns {
				rr, ok := convertRR(drr, r.expire)
				if !ok {
					continue
				}
				if rr.Type == "SOA" {
					hasSOA = true
					break
				}
			}
		}
		if !hasSOA {
			r.cache.addNX(qname)
			return nil, NXDOMAIN
		}
	} else if rmsg.Rcode != dns.RcodeSuccess {
		return nil, errors.New(dns.RcodeToString[rmsg.Rcode]) // FIXME: should (*Resolver).exchange special-case this error?
	}

	// Cache records returned
	rrs := r.saveDNSRR(host, qname, append(append(rmsg.Answer, rmsg.Ns...), rmsg.Extra...))

	// Resolve IP addresses of TLD name servers if NS query doesn’t return additional section
	if qtype == "NS" {
		for _, rr := range rrs {
			if rr.Type != "NS" {
				continue
			}
			arrs, err := r.cacheGet(ctx, rr.Value, "A")
			if err == NXDOMAIN {
				continue
			}
			if err != nil {
				break
			}
			if len(arrs) == 0 {
				arrs, err = r.exchangeIP(ctx, host, ip, rr.Value, "A", depth+1)
				if err != nil {
					break
				}
			}
			rrs = append(rrs, arrs...)
		}
	}

	return rrs, nil
}

func (r *Resolver) resolveCNAMEs(ctx context.Context, qname, qtype string, crrs RRs, depth int) (RRs, error) {
	var rrs RRs
	for _, crr := range crrs {
		rrs = append(rrs, crr)
		if crr.Type != "CNAME" || crr.Name != qname {
			continue
		}
		logCNAME(crr.String(), depth)
		crrs, _ := r.resolve(ctx, crr.Value, qtype, depth)
		for _, rr := range crrs {
			r.cache.add(qname, rr)
			rrs = append(rrs, crr)
		}
	}
	return rrs, nil
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(host, qname string, drrs []dns.RR) RRs {
	var rrs RRs
	cl := dns.CountLabel(qname)
	for _, drr := range drrs {
		rr, ok := convertRR(drr, r.expire)
		if !ok {
			continue
		}
		if dns.CountLabel(rr.Name) < cl && dns.CompareDomainName(qname, rr.Name) < 2 {
			// fmt.Fprintf(os.Stderr, "Warning: potential poisoning from %s: %s -> %s\n", host, qname, drr.String())
			continue
		}
		r.cache.add(rr.Name, rr)
		if rr.Name != qname {
			continue
		}
		rrs = append(rrs, rr)
	}
	return rrs
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(ctx context.Context, qname, qtype string) (RRs, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}
	any := r.cache.get(qname)
	if any == nil {
		any = rootCache.get(qname)
	}
	if any == nil {
		return nil, nil
	}
	if len(any) == 0 {
		return nil, NXDOMAIN
	}
	rrs := make(RRs, 0, len(any))
	for _, rr := range any {
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, rr)
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil, nil
	}
	return rrs, nil
}
<commit_msg>Dial MaxNameservers down from 4 to 2<commit_after>package dnsr

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/miekg/dns"
)

// DNS Resolution configuration.
var (
	Timeout             = 2000 * time.Millisecond
	TypicalResponseTime = 100 * time.Millisecond
	MaxRecursion        = 10
	MaxNameservers      = 2
	MaxIPs              = 2
)

// Resolver errors.
var (
	NXDOMAIN = fmt.Errorf("NXDOMAIN")

	ErrMaxRecursion = fmt.Errorf("maximum recursion depth reached: %d", MaxRecursion)
	ErrMaxIPs       = fmt.Errorf("maximum name server IPs queried: %d", MaxIPs)
	ErrNoARecords   = fmt.Errorf("no A records found for name server")
	ErrNoResponse   = fmt.Errorf("no responses received")
	ErrTimeout      = fmt.Errorf("timeout expired") // TODO: Timeouter interface? e.g. func (e) Timeout() bool { return true }
)

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache   *cache
	expire  bool
	timeout time.Duration
}

// New initializes a Resolver with the specified cache size.
func New(capacity int) *Resolver {
	return NewWithTimeout(capacity, Timeout)
}

// NewWithTimeout initializes a Resolver with the specified cache size and resolution timeout.
func NewWithTimeout(capacity int, timeout time.Duration) *Resolver {
	r := &Resolver{
		cache:   newCache(capacity, false),
		expire:  false,
		timeout: timeout,
	}
	return r
}

// NewExpiring initializes an expiring Resolver with the specified cache size.
func NewExpiring(capacity int) *Resolver {
	return NewExpiringWithTimeout(capacity, Timeout)
}

// NewExpiringWithTimeout initializes an expiring Resolved with the specified cache size and resolution timeout.
func NewExpiringWithTimeout(capacity int, timeout time.Duration) *Resolver {
	r := &Resolver{
		cache:   newCache(capacity, true),
		expire:  true,
		timeout: timeout,
	}
	return r
}

// Resolve calls ResolveErr to find DNS records of type qtype for the domain qname.
// For nonexistent domains (NXDOMAIN), it will return an empty, non-nil slice.
func (r *Resolver) Resolve(qname, qtype string) RRs {
	rrs, err := r.ResolveErr(qname, qtype)
	if err == NXDOMAIN {
		return emptyRRs
	}
	if err != nil {
		return nil
	}
	return rrs
}

// ResolveErr finds DNS records of type qtype for the domain qname.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveErr(qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

// ResolveCtx finds DNS records of type qtype for the domain qname using
// the supplied context. Requests may time out earlier if timeout is
// shorter than a deadline set in ctx.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveCtx(ctx context.Context, qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(ctx, r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

func (r *Resolver) resolve(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		return nil, ErrMaxRecursion
	}
	rrs, err := r.cacheGet(ctx, qname, qtype)
	if err != nil {
		return nil, err
	}
	if len(rrs) > 0 {
		return rrs, nil
	}
	logResolveStart(qname, qtype, depth)
	start := time.Now()
	rrs, err = r.iterateParents(ctx, qname, qtype, depth)
	logResolveEnd(qname, qtype, rrs, depth, start, err)
	return rrs, err
}

func (r *Resolver) iterateParents(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	chanRRs := make(chan RRs, MaxNameservers)
	chanErrs := make(chan error, MaxNameservers)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		// If we’re looking for [foo.com,NS], then move on to the parent ([com,NS])
		if pname == qname && qtype == "NS" {
			continue
		}

		// Only query TLDs against the root nameservers
		if pname == "." && dns.CountLabel(qname) != 1 {
			// fmt.Fprintf(os.Stderr, "Warning: non-TLD query at root: dig +norecurse %s %s\n", qname, qtype)
			return nil, nil
		}

		// Get nameservers
		nrrs, err := r.resolve(ctx, pname, "NS", depth)
		if err == NXDOMAIN || err == ErrTimeout || err == context.DeadlineExceeded {
			return nil, err
		}
		if err != nil {
			continue
		}

		// Check cache for specific queries
		if len(nrrs) > 0 && qtype != "" {
			rrs, err := r.cacheGet(ctx, qname, qtype)
			if err != nil {
				return nil, err
			}
			if len(rrs) > 0 {
				return rrs, nil
			}
		}

		// Query all nameservers in parallel
		count := 0
		for i := 0; i < len(nrrs) && count < MaxNameservers; i++ {
			nrr := nrrs[i]
			if nrr.Type != "NS" {
				continue
			}

			go func(host string) {
				rrs, err := r.exchange(ctx, host, qname, qtype, depth)
				if err != nil {
					chanErrs <- err
				} else {
					chanRRs <- rrs
				}
			}(nrr.Value)

			count++
		}

		// Wait for answer, error, or cancellation
		for ; count > 0; count-- {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case rrs := <-chanRRs:
				for _, nrr := range nrrs {
					if nrr.Name == qname {
						rrs = append(rrs, nrr)
					}
				}
				cancel() // stop any other work here before recursing
				return r.resolveCNAMEs(ctx, qname, qtype, rrs, depth)
			case err = <-chanErrs:
				if err == NXDOMAIN {
					return nil, err
				}
			}

		}

		// NS queries naturally recurse, so stop further iteration
		if qtype == "NS" {
			return nil, err
		}
	}

	return nil, ErrNoResponse
}

func (r *Resolver) exchange(ctx context.Context, host, qname, qtype string, depth int) (RRs, error) {
	count := 0
	arrs, err := r.resolve(ctx, host, "A", depth)
	if err != nil {
		return nil, err
	}
	for _, arr := range arrs {
		// FIXME: support AAAA records?
		if arr.Type != "A" {
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return nil, ErrMaxIPs
		}

		rrs, err := r.exchangeIP(ctx, host, arr.Value, qname, qtype, depth)
		if err == nil || err == NXDOMAIN || err == ErrTimeout {
			return rrs, err
		}

		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
	}

	return nil, ErrNoARecords
}

func (r *Resolver) exchangeIP(ctx context.Context, host, ip, qname, qtype string, depth int) (RRs, error) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeANY
	}
	var qmsg dns.Msg
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Synchronously query this DNS server
	start := time.Now()
	timeout := r.timeout // belt and suspenders, since ctx has a deadline from ResolveErr
	if dl, ok := ctx.Deadline(); ok {
		if start.After(dl.Add(-TypicalResponseTime)) { // bail if we can't finish in time (start is too close to deadline)
			return nil, ErrTimeout
		}
		timeout = dl.Sub(start)
	}

	client := &dns.Client{Timeout: timeout} // client must finish within remaining timeout
	rmsg, dur, err := client.Exchange(&qmsg, ip+":53")
	select {
	case <-ctx.Done(): // Finished too late
		logCancellation(host, &qmsg, rmsg, depth, dur, timeout)
		return nil, ctx.Err()
	default:
		logExchange(host, &qmsg, rmsg, depth, dur, timeout, err) // Log hostname instead of IP
	}
	if err != nil {
		return nil, err
	}

	// FIXME: cache NXDOMAIN responses responsibly
	if rmsg.Rcode == dns.RcodeNameError {
		var hasSOA bool
		if qtype == "NS" {
			for _, drr := range rmsg.Ns {
				rr, ok := convertRR(drr, r.expire)
				if !ok {
					continue
				}
				if rr.Type == "SOA" {
					hasSOA = true
					break
				}
			}
		}
		if !hasSOA {
			r.cache.addNX(qname)
			return nil, NXDOMAIN
		}
	} else if rmsg.Rcode != dns.RcodeSuccess {
		return nil, errors.New(dns.RcodeToString[rmsg.Rcode]) // FIXME: should (*Resolver).exchange special-case this error?
	}

	// Cache records returned
	rrs := r.saveDNSRR(host, qname, append(append(rmsg.Answer, rmsg.Ns...), rmsg.Extra...))

	// Resolve IP addresses of TLD name servers if NS query doesn’t return additional section
	if qtype == "NS" {
		for _, rr := range rrs {
			if rr.Type != "NS" {
				continue
			}
			arrs, err := r.cacheGet(ctx, rr.Value, "A")
			if err == NXDOMAIN {
				continue
			}
			if err != nil {
				break
			}
			if len(arrs) == 0 {
				arrs, err = r.exchangeIP(ctx, host, ip, rr.Value, "A", depth+1)
				if err != nil {
					break
				}
			}
			rrs = append(rrs, arrs...)
		}
	}

	return rrs, nil
}

func (r *Resolver) resolveCNAMEs(ctx context.Context, qname, qtype string, crrs RRs, depth int) (RRs, error) {
	var rrs RRs
	for _, crr := range crrs {
		rrs = append(rrs, crr)
		if crr.Type != "CNAME" || crr.Name != qname {
			continue
		}
		logCNAME(crr.String(), depth)
		crrs, _ := r.resolve(ctx, crr.Value, qtype, depth)
		for _, rr := range crrs {
			r.cache.add(qname, rr)
			rrs = append(rrs, crr)
		}
	}
	return rrs, nil
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(host, qname string, drrs []dns.RR) RRs {
	var rrs RRs
	cl := dns.CountLabel(qname)
	for _, drr := range drrs {
		rr, ok := convertRR(drr, r.expire)
		if !ok {
			continue
		}
		if dns.CountLabel(rr.Name) < cl && dns.CompareDomainName(qname, rr.Name) < 2 {
			// fmt.Fprintf(os.Stderr, "Warning: potential poisoning from %s: %s -> %s\n", host, qname, drr.String())
			continue
		}
		r.cache.add(rr.Name, rr)
		if rr.Name != qname {
			continue
		}
		rrs = append(rrs, rr)
	}
	return rrs
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(ctx context.Context, qname, qtype string) (RRs, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}
	any := r.cache.get(qname)
	if any == nil {
		any = rootCache.get(qname)
	}
	if any == nil {
		return nil, nil
	}
	if len(any) == 0 {
		return nil, NXDOMAIN
	}
	rrs := make(RRs, 0, len(any))
	for _, rr := range any {
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, rr)
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil, nil
	}
	return rrs, nil
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package stats

import (
	"expvar"
	"reflect"
	"testing"
)

func TestCounters(t *testing.T) {
	clear()
	c := NewCounters("counter1")
	c.Add("c1", 1)
	c.Add("c2", 1)
	c.Add("c2", 1)
	want1 := `{"c1": 1, "c2": 2}`
	want2 := `{"c2": 2, "c1": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1"] != 1 {
		t.Errorf("want 1, got %d", counts["c1"])
	}
	if counts["c2"] != 2 {
		t.Errorf("want 2, got %d", counts["c2"])
	}
	f := CountersFunc(func() map[string]int64 {
		return map[string]int64{
			"c1": 1,
			"c2": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersTags(t *testing.T) {
	clear()
	c := NewCounters("counterTag1")
	want := map[string]int64{}
	got := c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}

	c = NewCounters("counterTag2", "tag1", "tag2")
	want = map[string]int64{"tag1": 0, "tag2": 0}
	got = c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}
}

func TestMultiCounters(t *testing.T) {
	clear()
	c := NewMultiCounters("mapCounter1", []string{"aaa", "bbb"})
	c.Add([]string{"c1a", "c1b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	want1 := `{"c1a.c1b": 1, "c2a.c2b": 2}`
	want2 := `{"c2a.c2b": 2, "c1a.c1b": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1a.c1b"] != 1 {
		t.Errorf("want 1, got %d", counts["c1a.c1b"])
	}
	if counts["c2a.c2b"] != 2 {
		t.Errorf("want 2, got %d", counts["c2a.c2b"])
	}
	f := NewMultiCountersFunc("", []string{"aaa", "bbb"}, func() map[string]int64 {
		return map[string]int64{
			"c1a.c1b": 1,
			"c2a.c2b": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersHook(t *testing.T) {
	var gotname string
	var gotv *Counters
	clear()
	Register(func(name string, v expvar.Var) {
		gotname = name
		gotv = v.(*Counters)
	})

	v := NewCounters("counter2")
	if gotname != "counter2" {
		t.Errorf("want counter2, got %s", gotname)
	}
	if gotv != v {
		t.Errorf("want %#v, got %#v", v, gotv)
	}
}
<commit_msg>Add benchmarks for stats.Counters.Add().<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package stats

import (
	"expvar"
	"reflect"
	"sort"
	"testing"
	"time"
)

func TestCounters(t *testing.T) {
	clear()
	c := NewCounters("counter1")
	c.Add("c1", 1)
	c.Add("c2", 1)
	c.Add("c2", 1)
	want1 := `{"c1": 1, "c2": 2}`
	want2 := `{"c2": 2, "c1": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1"] != 1 {
		t.Errorf("want 1, got %d", counts["c1"])
	}
	if counts["c2"] != 2 {
		t.Errorf("want 2, got %d", counts["c2"])
	}
	f := CountersFunc(func() map[string]int64 {
		return map[string]int64{
			"c1": 1,
			"c2": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersTags(t *testing.T) {
	clear()
	c := NewCounters("counterTag1")
	want := map[string]int64{}
	got := c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}

	c = NewCounters("counterTag2", "tag1", "tag2")
	want = map[string]int64{"tag1": 0, "tag2": 0}
	got = c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}
}

func TestMultiCounters(t *testing.T) {
	clear()
	c := NewMultiCounters("mapCounter1", []string{"aaa", "bbb"})
	c.Add([]string{"c1a", "c1b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	want1 := `{"c1a.c1b": 1, "c2a.c2b": 2}`
	want2 := `{"c2a.c2b": 2, "c1a.c1b": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1a.c1b"] != 1 {
		t.Errorf("want 1, got %d", counts["c1a.c1b"])
	}
	if counts["c2a.c2b"] != 2 {
		t.Errorf("want 2, got %d", counts["c2a.c2b"])
	}
	f := NewMultiCountersFunc("", []string{"aaa", "bbb"}, func() map[string]int64 {
		return map[string]int64{
			"c1a.c1b": 1,
			"c2a.c2b": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersHook(t *testing.T) {
	var gotname string
	var gotv *Counters
	clear()
	Register(func(name string, v expvar.Var) {
		gotname = name
		gotv = v.(*Counters)
	})

	v := NewCounters("counter2")
	if gotname != "counter2" {
		t.Errorf("want counter2, got %s", gotname)
	}
	if gotv != v {
		t.Errorf("want %#v, got %#v", v, gotv)
	}
}

var benchCounter = NewCounters("bench")

func BenchmarkCounters(b *testing.B) {
	clear()
	benchCounter.Add("c1", 1)
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			benchCounter.Add("c1", 1)
		}
	})
}

func BenchmarkCountersTailLatency(b *testing.B) {
	// For this one, ignore the time reported by 'go test'.
	// The 99th Percentile log line is all that matters.
	clear()
	benchCounter.Add("c1", 1)
	c := make(chan time.Duration, 100)
	done := make(chan struct{})
	go func() {
		all := make([]int, b.N)
		i := 0
		for dur := range c {
			all[i] = int(dur)
			i++
		}
		sort.Ints(all)
		p99 := time.Duration(all[b.N*99/100])
		b.Logf("99th Percentile (for N=%v): %v", b.N, p99)
		close(done)
	}()

	b.ResetTimer()
	b.SetParallelism(1000)
	b.RunParallel(func(pb *testing.PB) {
		var start time.Time

		for pb.Next() {
			start = time.Now()
			benchCounter.Add("c1", 1)
			c <- time.Since(start)
		}
	})
	b.StopTimer()

	close(c)
	<-done
}
<|endoftext|>


<commit_before>package restless

import (
    "encoding/json"
    "fmt"
    "github.com/gorilla/mux"
    "labix.org/v2/mgo"
    "labix.org/v2/mgo/bson"
    "net/http"
)

func GetAll(c *mgo.Collection, ip interface{}) {
    c.Find(nil).All(ip)
}

func Insert(c *mgo.Collection, i interface{}) (string, error) {
    info, err := c.Upsert(bson.M{"_id": nil}, i)
    id := info.UpsertedId.(bson.ObjectId)
    return id.Hex(), err
}

func RemoveId(c *mgo.Collection, id bson.ObjectId) error {
    return c.RemoveId(id)
}

func GetId(c *mgo.Collection, i interface{}, id bson.ObjectId) error {
    return c.FindId(id).One(i)
}

func UpdateId(c *mgo.Collection, i interface{}, id bson.ObjectId) error {
    return c.UpdateId(id, i)
}

func GetGenHandler(s *mgo.Session, dbName string, colName string, cns Constructor) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {

        var jdata []byte
        var err error

        ns := s.Clone()
        defer ns.Close()

        col := ns.DB(dbName).C(colName)

        switch r.Method {
        //TODO: Add ability to queary specifics
        case "GET":
            i := cns.Slice()
            GetAll(col, i)
            jdata, err = json.Marshal(i)

            w.Header().Add("Content-Type", "application/json")
            fmt.Fprintf(w, "%s", jdata)

        case "POST":
            var lastId string

            i := cns.Single()
            if err = r.ParseForm(); err != nil {
                http.Error(w, "Unable to parse form", http.StatusBadRequest)
                Log.Errorf("Parsing form : %s", err)
            }

            jString := []byte(r.PostForm.Get("json"))
            if err = json.Unmarshal(jString, i); err != nil {
                http.Error(w, "Unable to unmarshal data", http.StatusBadRequest)
                Log.Errorf("UnMarshal error : %s", err)
                return
            }

            if lastId, err = Insert(col, i); err != nil {
                http.Error(w, "Unable to unmarshal data", http.StatusInternalServerError)
                Log.Error("Insert Error : %#v", err)
                return
            }

            if jdata, err = json.Marshal(i); err != nil {
                http.Error(w, "Marshal error", http.StatusInternalServerError)
            }
            w.Header().Add("Location", fmt.Sprintf("%s/%s", r.URL, lastId))
            w.WriteHeader(http.StatusCreated)
        }
        return
    }
}

func GetIdHandler(s *mgo.Session, dbName string, colName string, cns Constructor) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {

        var jdata []byte
        var err error
        var ids string

        ns := s.Clone()
        defer ns.Close()

        vars := mux.Vars(r)
        ids = vars["id"]

        col := ns.DB(dbName).C(colName)

        if !bson.IsObjectIdHex(ids) {
            http.Error(w, "Provided ID is unknown", http.StatusNotFound)
            return
        }

        id := bson.ObjectIdHex(ids)
        i := cns.Single()

        if err = GetId(col, i, id); err != nil {
            http.Error(w, "Provided ID is unknown", http.StatusNotFound)
            return
        }

        if jdata, err = json.Marshal(i); err != nil {
            http.Error(w, "", http.StatusBadRequest)
        }

        switch r.Method {
        case "GET":
            w.Header().Add("Content-Type", "application/json")
            fmt.Fprintf(w, "%s", jdata)

        case "PUT":
            if r.ParseForm(); err != nil {
                http.Error(w, "", http.StatusBadRequest)
            }

            if err = json.Unmarshal([]byte(r.PostForm.Get("json")), i); err != nil {
                http.Error(w, "", http.StatusBadRequest)
                Log.Errorf("UnMarshal error : %s", err)
            }

            if err = UpdateId(col, i, id); err != nil {
                http.Error(w, "Failed to update provided ID", http.StatusInternalServerError)
                Log.Errorf("UnMarshal error : %s", err)
            }

        case "DELETE":
            if err = RemoveId(col, id); err != nil {
                http.Error(w, "Failed to remove provided ID", http.StatusInternalServerError)
                Log.Errorf("Failed to remove id %s; error : %s", id, err)
            }
        }
        return
    }
}
<commit_msg>Update<commit_after>package restless

import (
    "encoding/json"
    "fmt"
    "github.com/gorilla/mux"
    "labix.org/v2/mgo"
    "labix.org/v2/mgo/bson"
    "net/http"
)

func GetAll(c *mgo.Collection, ip interface{}) {
    c.Find(nil).All(ip)
}

func Insert(c *mgo.Collection, i interface{}) (string, error) {
    info, err := c.Upsert(bson.M{"_id": nil}, i)
    id := info.UpsertedId.(bson.ObjectId)
    return id.Hex(), err
}

func RemoveId(c *mgo.Collection, id bson.ObjectId) error {
    return c.RemoveId(id)
}

func GetId(c *mgo.Collection, i interface{}, id bson.ObjectId) error {
    return c.FindId(id).One(i)
}

func UpdateId(c *mgo.Collection, i interface{}, id bson.ObjectId) error {
    return c.UpdateId(id, i)
}

func GetGenHandler(s *mgo.Session, dbName string, colName string, cns Constructor) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {

        var jdata []byte
        var err error

        ns := s.Clone()
        defer ns.Close()

        col := ns.DB(dbName).C(colName)

        switch r.Method {
        //TODO: Add ability to queary specifics
        case "GET":
            i := cns.Slice()
            GetAll(col, i)
            jdata, err = json.Marshal(i)

            w.Header().Add("Content-Type", "application/json")
            fmt.Fprintf(w, "%s", jdata)

        case "POST":
            var lastId string

            i := cns.Single()
            if err = r.ParseForm(); err != nil {
                http.Error(w, "Unable to parse form", http.StatusBadRequest)
                Log.Errorf("Parsing form : %s", err)
            }

            jString := []byte(r.PostForm.Get("json"))
            if err = json.Unmarshal(jString, i); err != nil {
                http.Error(w, "Unable to unmarshal data", http.StatusBadRequest)
                Log.Errorf("UnMarshal error : %s", err)
                return
            }

            if lastId, err = Insert(col, i); err != nil {
                http.Error(w, "Unable to unmarshal data", http.StatusInternalServerError)
                Log.Errorf("Insert Error : %#v", err)
                return
            }

            if jdata, err = json.Marshal(i); err != nil {
                http.Error(w, "Marshal error", http.StatusInternalServerError)
            }
            w.Header().Add("Location", fmt.Sprintf("%s/%s", r.URL, lastId))
            w.WriteHeader(http.StatusCreated)
        }
        return
    }
}

func GetIdHandler(s *mgo.Session, dbName string, colName string, cns Constructor) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {

        var jdata []byte
        var err error
        var ids string

        ns := s.Clone()
        defer ns.Close()

        vars := mux.Vars(r)
        ids = vars["id"]

        col := ns.DB(dbName).C(colName)

        if !bson.IsObjectIdHex(ids) {
            http.Error(w, "Provided ID is unknown", http.StatusNotFound)
            return
        }

        id := bson.ObjectIdHex(ids)
        i := cns.Single()

        if err = GetId(col, i, id); err != nil {
            http.Error(w, "Provided ID is unknown", http.StatusNotFound)
            return
        }

        if jdata, err = json.Marshal(i); err != nil {
            http.Error(w, "", http.StatusBadRequest)
        }

        switch r.Method {
        case "GET":
            w.Header().Add("Content-Type", "application/json")
            fmt.Fprintf(w, "%s", jdata)

        case "PUT":
            if r.ParseForm(); err != nil {
                http.Error(w, "", http.StatusBadRequest)
            }

            if err = json.Unmarshal([]byte(r.PostForm.Get("json")), i); err != nil {
                http.Error(w, "", http.StatusBadRequest)
                Log.Errorf("UnMarshal error : %s", err)
            }

            if err = UpdateId(col, i, id); err != nil {
                http.Error(w, "Failed to update provided ID", http.StatusInternalServerError)
                Log.Errorf("UnMarshal error : %s", err)
            }

        case "DELETE":
            if err = RemoveId(col, id); err != nil {
                http.Error(w, "Failed to remove provided ID", http.StatusInternalServerError)
                Log.Errorf("Failed to remove id %s; error : %s", id, err)
            }
        }
        return
    }
}
<|endoftext|>