<commit_before>package jpeg_test

import (
	"bufio"
	"fmt"
	"image"
	"image/color"
	nativeJPEG "image/jpeg"
	"os"
	"testing"

	"github.com/pixiv/go-libjpeg/jpeg"
	"github.com/pixiv/go-libjpeg/test/util"
)

var naturalImageFiles = []string{
	"cosmos.jpg",
	"kinkaku.jpg",
}

var subsampledImageFiles = []string{
	"checkerboard_444.jpg",
	"checkerboard_440.jpg",
	"checkerboard_422.jpg",
	"checkerboard_420.jpg",
}

func BenchmarkDecode(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, file := range naturalImageFiles {
			io := util.OpenFile(file)
			img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
			if img == nil {
				b.Error("Got nil")
			}
			if err != nil {
				b.Errorf("Got Error: %v", err)
			}
		}
	}
}

func BenchmarkDecodeIntoRGB(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, file := range naturalImageFiles {
			io := util.OpenFile(file)
			img, err := jpeg.DecodeIntoRGB(io, &jpeg.DecoderOptions{})
			if img == nil {
				b.Error("Got nil")
			}
			if err != nil {
				b.Errorf("Got Error: %v", err)
			}
		}
	}
}

func BenchmarkDecodeWithNativeJPEG(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, file := range naturalImageFiles {
			io := util.OpenFile(file)
			img, err := nativeJPEG.Decode(io)
			if img == nil {
				b.Error("Got nil")
			}
			if err != nil {
				b.Errorf("Got Error: %v", err)
			}
		}
	}
}

func TestDecode(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecode_%s.png", file))
	}
}

func TestDecodeScaled(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{ScaleTarget: image.Rect(0, 0, 100, 100)})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}
		if got := img.Bounds().Dx(); got != 256 {
			t.Errorf("Wrong scaled width: %v, expect: 128 (=1024/8)", got)
		}
		if got := img.Bounds().Dy(); got != 192 {
			t.Errorf("Wrong scaled height: %v, expect: 192 (=768/8)", got)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeScaled_%s.png", file))
	}
}

func TestDecodeIntoRGBA(t *testing.T) {
	if jpeg.SupportRGBA() != true {
		t.Skipf("This build is not support DecodeIntoRGBA.")
		return
	}
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.DecodeIntoRGBA(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Got Error: %v", err)
			continue
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeIntoRGBA_%s.png", file))
	}
}

func TestDecodeScaledIntoRGBA(t *testing.T) {
	if jpeg.SupportRGBA() != true {
		t.Skipf("This build is not support DecodeIntoRGBA.")
		return
	}
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.DecodeIntoRGBA(io, &jpeg.DecoderOptions{ScaleTarget: image.Rect(0, 0, 100, 100)})
		if err != nil {
			t.Errorf("Got Error: %v", err)
			continue
		}
		if got := img.Bounds().Dx(); got != 256 {
			t.Errorf("Wrong scaled width: %v, expect: 128 (=1024/8)", got)
		}
		if got := img.Bounds().Dy(); got != 192 {
			t.Errorf("Wrong scaled height: %v, expect: 192 (=768/8)", got)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeIntoRGBA_%s.png", file))
	}
}

func TestDecodeScaledIntoRGB(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.DecodeIntoRGB(io, &jpeg.DecoderOptions{ScaleTarget: image.Rect(0, 0, 100, 100)})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}
		if got := img.Bounds().Dx(); got != 256 {
			t.Errorf("Wrong scaled width: %v, expect: 128 (=1024/8)", got)
		}
		if got := img.Bounds().Dy(); got != 192 {
			t.Errorf("Wrong scaled height: %v, expect: 192 (=768/8)", got)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeIntoRGB_%s.png", file))
	}
}

func TestDecodeSubsampledImage(t *testing.T) {
	for _, file := range subsampledImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeSubsampledImage_%s.png", file))
	}
}

func TestDecodeAndEncode(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Decode returns error: %v", err)
		}

		// Create output file
		f, err := os.Create(util.GetOutFilePath(fmt.Sprintf("TestDecodeAndEncode_%s", file)))
		if err != nil {
			panic(err)
		}
		w := bufio.NewWriter(f)
		defer func() {
			w.Flush()
			f.Close()
		}()

		if err := jpeg.Encode(w, img, &jpeg.EncoderOptions{Quality: 90}); err != nil {
			t.Errorf("Encode returns error: %v", err)
		}
	}
}

func TestDecodeAndEncodeSubsampledImages(t *testing.T) {
	for _, file := range subsampledImageFiles {
		r := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(r, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Decode returns error: %v", err)
		}

		// Create output file
		f, err := os.Create(util.GetOutFilePath(fmt.Sprintf("TestDecodeAndEncodeSubsampledImages_%s", file)))
		if err != nil {
			panic(err)
		}
		w := bufio.NewWriter(f)
		defer func() {
			w.Flush()
			f.Close()
		}()

		if err := jpeg.Encode(w, img, &jpeg.EncoderOptions{Quality: 90}); err != nil {
			t.Errorf("Encode returns error: %v", err)
		}
	}
}

func TestDecodeConfig(t *testing.T) {
	for _, file := range naturalImageFiles {
		r := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		config, err := jpeg.DecodeConfig(r)
		if err != nil {
			t.Errorf("Got error: %v", err)
		}

		if got := config.ColorModel; got != color.YCbCrModel {
			t.Errorf("got wrong ColorModel: %v, expect: color.YCbCrModel", got)
		}
		if got := config.Width; got != 1024 {
			t.Errorf("got wrong width: %d, expect: 1024", got)
		}
		if got := config.Height; got != 768 {
			t.Errorf("got wrong height: %d, expect: 768", got)
		}
	}
}
<commit_msg>add test code for NewYCbcrAligned<commit_after>package jpeg_test

import (
	"bufio"
	"fmt"
	"image"
	"image/color"
	nativeJPEG "image/jpeg"
	"os"
	"testing"

	"github.com/pixiv/go-libjpeg/jpeg"
	"github.com/pixiv/go-libjpeg/test/util"
)

var naturalImageFiles = []string{
	"cosmos.jpg",
	"kinkaku.jpg",
}

var subsampledImageFiles = []string{
	"checkerboard_444.jpg",
	"checkerboard_440.jpg",
	"checkerboard_422.jpg",
	"checkerboard_420.jpg",
}

func BenchmarkDecode(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, file := range naturalImageFiles {
			io := util.OpenFile(file)
			img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
			if img == nil {
				b.Error("Got nil")
			}
			if err != nil {
				b.Errorf("Got Error: %v", err)
			}
		}
	}
}

func BenchmarkDecodeIntoRGB(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, file := range naturalImageFiles {
			io := util.OpenFile(file)
			img, err := jpeg.DecodeIntoRGB(io, &jpeg.DecoderOptions{})
			if img == nil {
				b.Error("Got nil")
			}
			if err != nil {
				b.Errorf("Got Error: %v", err)
			}
		}
	}
}

func BenchmarkDecodeWithNativeJPEG(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, file := range naturalImageFiles {
			io := util.OpenFile(file)
			img, err := nativeJPEG.Decode(io)
			if img == nil {
				b.Error("Got nil")
			}
			if err != nil {
				b.Errorf("Got Error: %v", err)
			}
		}
	}
}

func TestDecode(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecode_%s.png", file))
	}
}

func TestDecodeScaled(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{ScaleTarget: image.Rect(0, 0, 100, 100)})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}
		if got := img.Bounds().Dx(); got != 256 {
			t.Errorf("Wrong scaled width: %v, expect: 128 (=1024/8)", got)
		}
		if got := img.Bounds().Dy(); got != 192 {
			t.Errorf("Wrong scaled height: %v, expect: 192 (=768/8)", got)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeScaled_%s.png", file))
	}
}

func TestDecodeIntoRGBA(t *testing.T) {
	if jpeg.SupportRGBA() != true {
		t.Skipf("This build is not support DecodeIntoRGBA.")
		return
	}
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.DecodeIntoRGBA(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Got Error: %v", err)
			continue
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeIntoRGBA_%s.png", file))
	}
}

func TestDecodeScaledIntoRGBA(t *testing.T) {
	if jpeg.SupportRGBA() != true {
		t.Skipf("This build is not support DecodeIntoRGBA.")
		return
	}
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.DecodeIntoRGBA(io, &jpeg.DecoderOptions{ScaleTarget: image.Rect(0, 0, 100, 100)})
		if err != nil {
			t.Errorf("Got Error: %v", err)
			continue
		}
		if got := img.Bounds().Dx(); got != 256 {
			t.Errorf("Wrong scaled width: %v, expect: 128 (=1024/8)", got)
		}
		if got := img.Bounds().Dy(); got != 192 {
			t.Errorf("Wrong scaled height: %v, expect: 192 (=768/8)", got)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeIntoRGBA_%s.png", file))
	}
}

func TestDecodeScaledIntoRGB(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.DecodeIntoRGB(io, &jpeg.DecoderOptions{ScaleTarget: image.Rect(0, 0, 100, 100)})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}
		if got := img.Bounds().Dx(); got != 256 {
			t.Errorf("Wrong scaled width: %v, expect: 128 (=1024/8)", got)
		}
		if got := img.Bounds().Dy(); got != 192 {
			t.Errorf("Wrong scaled height: %v, expect: 192 (=768/8)", got)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeIntoRGB_%s.png", file))
	}
}

func TestDecodeSubsampledImage(t *testing.T) {
	for _, file := range subsampledImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Got Error: %v", err)
		}

		util.WritePNG(img, fmt.Sprintf("TestDecodeSubsampledImage_%s.png", file))
	}
}

func TestDecodeAndEncode(t *testing.T) {
	for _, file := range naturalImageFiles {
		io := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(io, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Decode returns error: %v", err)
		}

		// Create output file
		f, err := os.Create(util.GetOutFilePath(fmt.Sprintf("TestDecodeAndEncode_%s", file)))
		if err != nil {
			panic(err)
		}
		w := bufio.NewWriter(f)
		defer func() {
			w.Flush()
			f.Close()
		}()

		if err := jpeg.Encode(w, img, &jpeg.EncoderOptions{Quality: 90}); err != nil {
			t.Errorf("Encode returns error: %v", err)
		}
	}
}

func TestDecodeAndEncodeSubsampledImages(t *testing.T) {
	for _, file := range subsampledImageFiles {
		r := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		img, err := jpeg.Decode(r, &jpeg.DecoderOptions{})
		if err != nil {
			t.Errorf("Decode returns error: %v", err)
		}

		// Create output file
		f, err := os.Create(util.GetOutFilePath(fmt.Sprintf("TestDecodeAndEncodeSubsampledImages_%s", file)))
		if err != nil {
			panic(err)
		}
		w := bufio.NewWriter(f)
		defer func() {
			w.Flush()
			f.Close()
		}()

		if err := jpeg.Encode(w, img, &jpeg.EncoderOptions{Quality: 90}); err != nil {
			t.Errorf("Encode returns error: %v", err)
		}
	}
}

func TestDecodeConfig(t *testing.T) {
	for _, file := range naturalImageFiles {
		r := util.OpenFile(file)
		fmt.Printf(" - test: %s\n", file)

		config, err := jpeg.DecodeConfig(r)
		if err != nil {
			t.Errorf("Got error: %v", err)
		}

		if got := config.ColorModel; got != color.YCbCrModel {
			t.Errorf("got wrong ColorModel: %v, expect: color.YCbCrModel", got)
		}
		if got := config.Width; got != 1024 {
			t.Errorf("got wrong width: %d, expect: 1024", got)
		}
		if got := config.Height; got != 768 {
			t.Errorf("got wrong height: %d, expect: 768", got)
		}
	}
}

func TestNewYCbCrAlignedWithLandscape(t *testing.T) {
	got := jpeg.NewYCbCrAligned(image.Rect(0, 0, 125, 25), image.YCbCrSubsampleRatio444)

	if len(got.Y) != 6912 {
		t.Errorf("wrong array size Y: %d, expect: 6912", len(got.Y))
	}
	if len(got.Cb) != 6912 {
		t.Errorf("wrong array size Cb: %d, expect: 6912", len(got.Cb))
	}
	if len(got.Cr) != 6912 {
		t.Errorf("wrong array size Cr: %d, expect: 6912", len(got.Cr))
	}
	if got.YStride != 144 {
		t.Errorf("got wrong YStride: %d, expect: 128", got.YStride)
	}
	if got.CStride != 144 {
		t.Errorf("got wrong CStride: %d, expect: 128", got.CStride)
	}
}

func TestNewYCbCrAlignedWithPortrait(t *testing.T) {
	got := jpeg.NewYCbCrAligned(image.Rect(0, 0, 25, 125), image.YCbCrSubsampleRatio444)

	if len(got.Y) != 6912 {
		t.Errorf("wrong array size Y: %d, expect: 6912", len(got.Y))
	}
	if len(got.Cb) != 6912 {
		t.Errorf("wrong array size Cb: %d, expect: 6912", len(got.Cb))
	}
	if len(got.Cr) != 6912 {
		t.Errorf("wrong array size Cr: %d, expect: 6912", len(got.Cr))
	}
	if got.YStride != 48 {
		t.Errorf("got wrong YStride: %d, expect: 128", got.YStride)
	}
	if got.CStride != 48 {
		t.Errorf("got wrong CStride: %d, expect: 128", got.CStride)
	}
}
<|endoftext|>


<commit_before>/*
Wrapper around the HTTP trading API for type safety 'n' stuff.
*/
package tradeapi

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/Philipp15b/go-steam/community"
	"github.com/Philipp15b/go-steam/economy/inventory"
	"github.com/Philipp15b/go-steam/netutil"
	"github.com/Philipp15b/go-steam/steamid"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"time"
)

const tradeUrl = "https://steamcommunity.com/trade/%d/"

type Trade struct {
	client *http.Client
	other  steamid.SteamId

	LogPos  uint // not automatically updated
	Version uint // Incremented for each item change by Steam; not automatically updated.

	// the `sessionid` cookie is sent as a parameter/POST data for CSRF protection.
	sessionId string
	baseUrl   string
}

// Creates a new Trade based on the given cookies `sessionid`, `steamLogin`, `steamLoginSecure` and the trade partner's Steam ID.
func New(sessionId, steamLogin steamLoginSecure, string, other steamid.SteamId) *Trade {
	client := new(http.Client)
	client.Timeout = 10 * time.Second

	t := &Trade{
		client:    client,
		other:     other,
		sessionId: sessionId,
		baseUrl:   fmt.Sprintf(tradeUrl, other),
		Version:   1,
	}
	community.SetCookies(t.client, sessionId, steamLogin, steamLoginSecure)
	return t
}

type Main struct {
	PartnerOnProbation bool
}

var onProbationRegex = regexp.MustCompile(`var g_bTradePartnerProbation = (\w+);`)

// Fetches the main HTML page and parses it. Thread-safe.
func (t *Trade) GetMain() (*Main, error) {
	resp, err := t.client.Get(t.baseUrl)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	match := onProbationRegex.FindSubmatch(body)
	if len(match) == 0 {
		return nil, errors.New("tradeapi.GetMain: Could not find probation info")
	}

	return &Main{
		string(match[1]) == "true",
	}, nil
}

// Ajax POSTs to an API endpoint that should return a status
func (t *Trade) postWithStatus(url string, data map[string]string) (*Status, error) {
	status := new(Status)

	req := netutil.NewPostForm(url, netutil.ToUrlValues(data))
	// Tales of Madness and Pain, Episode 1: If you forget this, Steam will return an error
	// saying "missing required parameter", even though they are all there. IT WAS JUST THE HEADER, ARGH!
	req.Header.Add("Referer", t.baseUrl)

	resp, err := t.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(status)
	if err != nil {
		return nil, err
	}
	return status, nil
}

func (t *Trade) GetStatus() (*Status, error) {
	return t.postWithStatus(t.baseUrl+"tradestatus/", map[string]string{
		"sessionid": t.sessionId,
		"logpos":    strconv.FormatUint(uint64(t.LogPos), 10),
		"version":   strconv.FormatUint(uint64(t.Version), 10),
	})
}

// Thread-safe.
func (t *Trade) GetForeignInventory(contextId uint64, appId uint32, start *uint) (*inventory.PartialInventory, error) {
	data := map[string]string{
		"sessionid": t.sessionId,
		"steamid":   fmt.Sprintf("%d", t.other),
		"contextid": strconv.FormatUint(contextId, 10),
		"appid":     strconv.FormatUint(uint64(appId), 10),
	}
	if start != nil {
		data["start"] = strconv.FormatUint(uint64(*start), 10)
	}

	req, err := http.NewRequest("GET", t.baseUrl+"foreigninventory?"+netutil.ToUrlValues(data).Encode(), nil)
	if err != nil {
		panic(err)
	}
	req.Header.Add("Referer", t.baseUrl)

	return inventory.DoInventoryRequest(t.client, req)
}

// Thread-safe.
func (t *Trade) GetOwnInventory(contextId uint64, appId uint32) (*inventory.Inventory, error) {
	return inventory.GetOwnInventory(t.client, contextId, appId)
}

func (t *Trade) Chat(message string) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"chat", map[string]string{
		"sessionid": t.sessionId,
		"logpos":    strconv.FormatUint(uint64(t.LogPos), 10),
		"version":   strconv.FormatUint(uint64(t.Version), 10),
		"message":   message,
	})
}

func (t *Trade) AddItem(slot uint, itemId, contextId uint64, appId uint32) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"additem", map[string]string{
		"sessionid": t.sessionId,
		"slot":      strconv.FormatUint(uint64(slot), 10),
		"itemid":    strconv.FormatUint(itemId, 10),
		"contextid": strconv.FormatUint(contextId, 10),
		"appid":     strconv.FormatUint(uint64(appId), 10),
	})
}

func (t *Trade) RemoveItem(slot uint, itemId, contextId uint64, appId uint32) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"removeitem", map[string]string{
		"sessionid": t.sessionId,
		"slot":      strconv.FormatUint(uint64(slot), 10),
		"itemid":    strconv.FormatUint(itemId, 10),
		"contextid": strconv.FormatUint(contextId, 10),
		"appid":     strconv.FormatUint(uint64(appId), 10),
	})
}

func (t *Trade) SetCurrency(amount uint, currencyId, contextId uint64, appId uint32) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"setcurrency", map[string]string{
		"sessionid":  t.sessionId,
		"amount":     strconv.FormatUint(uint64(amount), 10),
		"currencyid": strconv.FormatUint(uint64(currencyId), 10),
		"contextid":  strconv.FormatUint(contextId, 10),
		"appid":      strconv.FormatUint(uint64(appId), 10),
	})
}

func (t *Trade) SetReady(ready bool) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"toggleready", map[string]string{
		"sessionid": t.sessionId,
		"version":   strconv.FormatUint(uint64(t.Version), 10),
		"ready":     fmt.Sprint(ready),
	})
}

func (t *Trade) Confirm() (*Status, error) {
	return t.postWithStatus(t.baseUrl+"confirm", map[string]string{
		"sessionid": t.sessionId,
		"version":   strconv.FormatUint(uint64(t.Version), 10),
	})
}

func (t *Trade) Cancel() (*Status, error) {
	return t.postWithStatus(t.baseUrl+"cancel", map[string]string{
		"sessionid": t.sessionId,
	})
}

func isSuccess(v interface{}) bool {
	if m, ok := v.(map[string]interface{}); ok {
		return m["success"] == true
	}
	return false
}
<commit_msg>Trading: Fix compilation<commit_after>/*
Wrapper around the HTTP trading API for type safety 'n' stuff.
*/
package tradeapi

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/Philipp15b/go-steam/community"
	"github.com/Philipp15b/go-steam/economy/inventory"
	"github.com/Philipp15b/go-steam/netutil"
	"github.com/Philipp15b/go-steam/steamid"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
	"time"
)

const tradeUrl = "https://steamcommunity.com/trade/%d/"

type Trade struct {
	client *http.Client
	other  steamid.SteamId

	LogPos  uint // not automatically updated
	Version uint // Incremented for each item change by Steam; not automatically updated.

	// the `sessionid` cookie is sent as a parameter/POST data for CSRF protection.
	sessionId string
	baseUrl   string
}

// Creates a new Trade based on the given cookies `sessionid`, `steamLogin`, `steamLoginSecure` and the trade partner's Steam ID.
func New(sessionId, steamLogin, steamLoginSecure string, other steamid.SteamId) *Trade {
	client := new(http.Client)
	client.Timeout = 10 * time.Second

	t := &Trade{
		client:    client,
		other:     other,
		sessionId: sessionId,
		baseUrl:   fmt.Sprintf(tradeUrl, other),
		Version:   1,
	}
	community.SetCookies(t.client, sessionId, steamLogin, steamLoginSecure)
	return t
}

type Main struct {
	PartnerOnProbation bool
}

var onProbationRegex = regexp.MustCompile(`var g_bTradePartnerProbation = (\w+);`)

// Fetches the main HTML page and parses it. Thread-safe.
func (t *Trade) GetMain() (*Main, error) {
	resp, err := t.client.Get(t.baseUrl)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	match := onProbationRegex.FindSubmatch(body)
	if len(match) == 0 {
		return nil, errors.New("tradeapi.GetMain: Could not find probation info")
	}

	return &Main{
		string(match[1]) == "true",
	}, nil
}

// Ajax POSTs to an API endpoint that should return a status
func (t *Trade) postWithStatus(url string, data map[string]string) (*Status, error) {
	status := new(Status)

	req := netutil.NewPostForm(url, netutil.ToUrlValues(data))
	// Tales of Madness and Pain, Episode 1: If you forget this, Steam will return an error
	// saying "missing required parameter", even though they are all there. IT WAS JUST THE HEADER, ARGH!
	req.Header.Add("Referer", t.baseUrl)

	resp, err := t.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(status)
	if err != nil {
		return nil, err
	}
	return status, nil
}

func (t *Trade) GetStatus() (*Status, error) {
	return t.postWithStatus(t.baseUrl+"tradestatus/", map[string]string{
		"sessionid": t.sessionId,
		"logpos":    strconv.FormatUint(uint64(t.LogPos), 10),
		"version":   strconv.FormatUint(uint64(t.Version), 10),
	})
}

// Thread-safe.
func (t *Trade) GetForeignInventory(contextId uint64, appId uint32, start *uint) (*inventory.PartialInventory, error) {
	data := map[string]string{
		"sessionid": t.sessionId,
		"steamid":   fmt.Sprintf("%d", t.other),
		"contextid": strconv.FormatUint(contextId, 10),
		"appid":     strconv.FormatUint(uint64(appId), 10),
	}
	if start != nil {
		data["start"] = strconv.FormatUint(uint64(*start), 10)
	}

	req, err := http.NewRequest("GET", t.baseUrl+"foreigninventory?"+netutil.ToUrlValues(data).Encode(), nil)
	if err != nil {
		panic(err)
	}
	req.Header.Add("Referer", t.baseUrl)

	return inventory.DoInventoryRequest(t.client, req)
}

// Thread-safe.
func (t *Trade) GetOwnInventory(contextId uint64, appId uint32) (*inventory.Inventory, error) {
	return inventory.GetOwnInventory(t.client, contextId, appId)
}

func (t *Trade) Chat(message string) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"chat", map[string]string{
		"sessionid": t.sessionId,
		"logpos":    strconv.FormatUint(uint64(t.LogPos), 10),
		"version":   strconv.FormatUint(uint64(t.Version), 10),
		"message":   message,
	})
}

func (t *Trade) AddItem(slot uint, itemId, contextId uint64, appId uint32) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"additem", map[string]string{
		"sessionid": t.sessionId,
		"slot":      strconv.FormatUint(uint64(slot), 10),
		"itemid":    strconv.FormatUint(itemId, 10),
		"contextid": strconv.FormatUint(contextId, 10),
		"appid":     strconv.FormatUint(uint64(appId), 10),
	})
}

func (t *Trade) RemoveItem(slot uint, itemId, contextId uint64, appId uint32) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"removeitem", map[string]string{
		"sessionid": t.sessionId,
		"slot":      strconv.FormatUint(uint64(slot), 10),
		"itemid":    strconv.FormatUint(itemId, 10),
		"contextid": strconv.FormatUint(contextId, 10),
		"appid":     strconv.FormatUint(uint64(appId), 10),
	})
}

func (t *Trade) SetCurrency(amount uint, currencyId, contextId uint64, appId uint32) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"setcurrency", map[string]string{
		"sessionid":  t.sessionId,
		"amount":     strconv.FormatUint(uint64(amount), 10),
		"currencyid": strconv.FormatUint(uint64(currencyId), 10),
		"contextid":  strconv.FormatUint(contextId, 10),
		"appid":      strconv.FormatUint(uint64(appId), 10),
	})
}

func (t *Trade) SetReady(ready bool) (*Status, error) {
	return t.postWithStatus(t.baseUrl+"toggleready", map[string]string{
		"sessionid": t.sessionId,
		"version":   strconv.FormatUint(uint64(t.Version), 10),
		"ready":     fmt.Sprint(ready),
	})
}

func (t *Trade) Confirm() (*Status, error) {
	return t.postWithStatus(t.baseUrl+"confirm", map[string]string{
		"sessionid": t.sessionId,
		"version":   strconv.FormatUint(uint64(t.Version), 10),
	})
}

func (t *Trade) Cancel() (*Status, error) {
	return t.postWithStatus(t.baseUrl+"cancel", map[string]string{
		"sessionid": t.sessionId,
	})
}

func isSuccess(v interface{}) bool {
	if m, ok := v.(map[string]interface{}); ok {
		return m["success"] == true
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/xeipuuv/gojsonreference"
)

type JsonSchemaBundler interface {
	Bundle()
	GetRootSchema() (JsonSchema)
	GetSchema(string) (JsonSchema)
}

type bundler struct {
	rootRef gojsonreference.JsonReference
	bundles map[string]bundle
}

type bundle struct {
	ref gojsonreference.JsonReference
	schema JsonSchema
}

func NewJsonSchemaBundler(rootReference gojsonreference.JsonReference) (JsonSchemaBundler) {
	b := bundler{ rootReference, make(map[string]bundle) }
	b.addBundle(b.rootRef.String(), b.rootRef)
	return &b
}

func (b *bundler) GetRootSchema() (JsonSchema) {
	return b.GetSchema(b.rootRef.String())
}

func (b *bundler) Bundle() {
	b.bundleRecursive(b.GetRootSchema())
}

func (b *bundler) bundleRecursive(schema JsonSchema) {
	if schema.Ref != "" {
		b.addBundle(schema.Ref, NewRelativeJsonReference(b.rootRef, NewJsonReference(schema.Ref)))
	}
	// object type
	for _, v := range schema.Properties {
		b.bundleRecursive(v)
	}
	for _, d := range schema.Dependencies {
		b.bundleRecursive(d)
	}
	// array type
	for _, i := range schema.GetItemList() {
		b.bundleRecursive(i)
	}
}

func (b *bundler) addBundle(path string, ref gojsonreference.JsonReference) {
	b.bundles[path] = bundle{
		ref,
		GetJsonSchemaLoader().Load(ref),
	}
}

func (b *bundler) GetSchema(path string) (JsonSchema) {
	return b.bundles[path].schema
}

<commit_msg>deleted bundle struct<commit_after>package main

import (
	"github.com/xeipuuv/gojsonreference"
)

type JsonSchemaBundler interface {
	Bundle()
	GetRootSchema() (JsonSchema)
	GetSchema(string) (JsonSchema)
}

type bundler struct {
	rootRef gojsonreference.JsonReference
	bundles map[string]JsonSchema
}

func NewJsonSchemaBundler(rootReference gojsonreference.JsonReference) (JsonSchemaBundler) {
	b := bundler{ rootReference, make(map[string]JsonSchema) }
	b.registerBundle(b.rootRef.String(), b.rootRef)
	return &b
}

func (b *bundler) GetRootSchema() (JsonSchema) {
	return b.GetSchema(b.rootRef.String())
}

func (b *bundler) Bundle() {
	b.bundleRecursive(b.GetRootSchema())
}

func (b *bundler) bundleRecursive(schema JsonSchema) {
	if schema.Ref != "" {
		b.registerBundle(schema.Ref, NewRelativeJsonReference(b.rootRef, NewJsonReference(schema.Ref)))
	}
	// object type
	for _, v := range schema.Properties {
		b.bundleRecursive(v)
	}
	for _, d := range schema.Dependencies {
		b.bundleRecursive(d)
	}
	// array type
	for _, i := range schema.GetItemList() {
		b.bundleRecursive(i)
	}
}

func (b *bundler) registerBundle(path string, ref gojsonreference.JsonReference) {
	b.bundles[path] = GetJsonSchemaLoader().Load(ref)
}

func (b *bundler) GetSchema(path string) (JsonSchema) {
	return b.bundles[path]
}

<|endoftext|>


<commit_before><commit_msg>Fix updateVersion commandline args in release tool (#572)<commit_after><|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"

	"github.com/confluentinc/confluent-kafka-go/kafka"
)

func consumer() {

	c, err := kafka.NewConsumer(&kafka.ConfigMap{
		"bootstrap.servers": "localhost",
		"group.id":          "myGroup",
		"auto.offset.reset": "earliest",
	})

	if err != nil {
		panic(err)
	}

	c.SubscribeTopics([]string{"myTopic", "^aRegex.*[Tt]opic"}, nil)

	for {
		msg, err := c.ReadMessage(-1)
		if err == nil {
			fmt.Printf("Message on %s: %s\n", msg.TopicPartition, string(msg.Value))
		} else {
			fmt.Printf("Consumer error: %v (%v)\n", err, msg)
			break
		}
	}

	c.Close()
}

func producer() {

	p, err := kafka.NewProducer(&kafka.ConfigMap{"bootstrap.servers": "localhost"})
	if err != nil {
		panic(err)
	}

	// Delivery report handler for produced messages
	go func() {
		for e := range p.Events() {
			switch ev := e.(type) {
			case *kafka.Message:
				if ev.TopicPartition.Error != nil {
					fmt.Printf("Delivery failed: %v\n", ev.TopicPartition)
				} else {
					fmt.Printf("Delivered message to %v\n", ev.TopicPartition)
				}
			}
		}
	}()

	// Produce messages to topic (asynchronously)
	topic := "myTopic"
	for _, word := range []string{"Welcome", "to", "the", "Confluent", "Kafka", "Golang", "client"} {
		p.Produce(&kafka.Message{
			TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny},
			Value:          []byte(word),
		}, nil)
	}

	// Wait for message deliveries
	p.Flush(15 * 1000)
}

func main() {
	c, err := kafka.NewConsumer(&kafka.ConfigMap{
		"bootstrap.servers": "192.168.240.86",
		"group.id":          "cclin",
		"auto.offset.reset": "earliest",
	})

	if err != nil {
		log.Fatal(err)
	}

	defer c.Close()

	err = c.Subscribe("foo", nil)

	if err != nil {
		log.Fatal(err)
	}

	//

	p, err := kafka.NewProducer(&kafka.ConfigMap{"bootstrap.servers": "192.168.240.86"})
	if err != nil {
		log.Fatal(err)
	}

	defer p.Close()

	//

	for {
		msg, err := c.ReadMessage(-1)
		if err == nil {
			log.Printf("Message on %s: %s\n", msg.TopicPartition, string(msg.Value))

			bar := "bar"
			err := p.Produce(&kafka.Message{
				TopicPartition: kafka.TopicPartition{Topic: &bar, Partition: kafka.PartitionAny},
				Value:          msg.Value,
			}, nil)

			if err != nil {
				log.Fatalf("Producer error: %v (%v)\n", err, msg)
			}

			break // just prosume one
		} else {
			log.Fatalf("Consumer error: %v (%v)\n", err, msg)
		}
	}

	// Wait for message deliveries
	p.Flush(15 * 1000)

	log.Println("End")
}
<commit_msg>mod prosumer/main.go<commit_after>package main

import (
	"log"

	"github.com/confluentinc/confluent-kafka-go/kafka"
)

func main() {
	c, err := kafka.NewConsumer(&kafka.ConfigMap{
		"bootstrap.servers": "192.168.240.86",
		"group.id":          "cclin",
		"auto.offset.reset": "earliest",
	})

	if err != nil {
		log.Fatal(err)
	}

	defer c.Close()

	err = c.Subscribe("foo", nil)

	if err != nil {
		log.Fatal(err)
	}

	//

	p, err := kafka.NewProducer(&kafka.ConfigMap{"bootstrap.servers": "192.168.240.86"})
	if err != nil {
		log.Fatal(err)
	}

	defer p.Close()

	// Wait for message deliveries
	defer p.Flush(15 * 1000)

	//

	for {
		msg, err := c.ReadMessage(-1)
		if err == nil {
			log.Printf("Message on %s: %s\n", msg.TopicPartition, string(msg.Value))

			bar := "bar"
			err := p.Produce(&kafka.Message{
				TopicPartition: kafka.TopicPartition{Topic: &bar, Partition: kafka.PartitionAny},
				Value:          msg.Value,
			}, nil)

			if err != nil {
				log.Fatalf("Producer error: %v (%v)\n", err, msg)
			}

			// break // just prosume one
		} else {
			log.Fatalf("Consumer error: %v (%v)\n", err, msg)
		}
	}
}
<|endoftext|>


<commit_before>package kafkahook

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/mailgun/holster/errors"
	"github.com/mailgun/holster/stack"
	"github.com/mailgun/logrus-hooks/common"
	"github.com/mailru/easyjson/jwriter"
	"github.com/sirupsen/logrus"
)

const bufferSize = 150

type KafkaHook struct {
	produce chan []byte
	conf    Config
	debug   bool

	// Process identity metadata
	hostName string
	appName  string
	cid      string
	pid      int

	// Sync stuff
	wg   sync.WaitGroup
	once sync.Once
}

type Config struct {
	Endpoints []string
	Topic     string
	Producer  sarama.AsyncProducer
}

func NewWithContext(ctx context.Context, conf Config) (hook *KafkaHook, err error) {
	done := make(chan struct{})
	go func() {
		hook, err = New(conf)
		close(done)
	}()

	select {
	case <-ctx.Done():
		return hook, fmt.Errorf("connect timeout while connecting to kafka peers %s",
			strings.Join(conf.Endpoints, ","))
	case <-done:
		return hook, err
	}
}

func New(conf Config) (*KafkaHook, error) {
	var err error

	kafkaConfig := sarama.NewConfig()
	kafkaConfig.Producer.RequiredAcks = sarama.WaitForAll
	kafkaConfig.Producer.Compression = sarama.CompressionSnappy
	kafkaConfig.Producer.Flush.Frequency = 200 * time.Millisecond
	kafkaConfig.Producer.Retry.Backoff = 10 * time.Second
	kafkaConfig.Producer.Retry.Max = 6
	kafkaConfig.Producer.Return.Errors = true

	// If the user failed to provide a producer create one
	if conf.Producer == nil {
		conf.Producer, err = sarama.NewAsyncProducer(conf.Endpoints, kafkaConfig)
		if err != nil {
			return nil, errors.Wrap(err, "kafka producer error")
		}
	}

	h := KafkaHook{
		produce: make(chan []byte, bufferSize),
		conf:    conf,
	}

	h.wg.Add(1)
	go func() {
		for {
			select {
			case prodErr := <-conf.Producer.Errors():
				fmt.Fprintf(os.Stderr, "[kafkahook] producer error '%s'; reconnecting producer...\n", prodErr.Err)

				// Close any previously open producers
				if conf.Producer != nil {
					// Since `Producer.Return.Errors` is true, `Producer.Close()` will send a shutdown error which
					// could cause `Close()` to block unless we read it. Also there could be other errors
					// that are waiting to be received and must be sent before shutdown is complete.
					go func() {
						for range conf.Producer.Errors() {}
					}()

					conf.Producer.Close()
					conf.Producer = nil
				}

				var backOff int
				for {
					// Attempt to reconnect to the brokers
					conf.Producer, err = sarama.NewAsyncProducer(conf.Endpoints, kafkaConfig)
					if err != nil {
						if backOff < 6 {
							backOff++
						}
						fmt.Fprintf(os.Stderr, "[kafkahook] reconnect error: %s; sleeping (%d)...\n", err, backOff)
						time.Sleep(time.Duration(backOff) * time.Second)
						continue
					}
					backOff = 0
					break
				}
			case buf, ok := <-h.produce:
				if !ok {
					h.wg.Done()
					// See comment above
					go func() {
						for range conf.Producer.Errors() {}
					}()
					if err := conf.Producer.Close(); err != nil {
						fmt.Fprintf(os.Stderr, "[kafkahook] producer close error: %s\n", err)
					}
					return
				}
				conf.Producer.Input() <- &sarama.ProducerMessage{
					Value: sarama.ByteEncoder(buf),
					Topic: conf.Topic,
					Key:   nil,
				}
			}
		}
	}()

	if h.hostName, err = os.Hostname(); err != nil {
		h.hostName = "unknown"
	}
	h.appName = filepath.Base(os.Args[0])
	if h.pid = os.Getpid(); h.pid == 1 {
		h.pid = 0
	}
	h.cid = getDockerCID()
	return &h, nil
}

func (h *KafkaHook) Fire(entry *logrus.Entry) error {
	var caller *stack.FrameInfo
	var err error

	caller = common.GetLogrusCaller()

	rec := &common.LogRecord{
		Category:  "logrus",
		AppName:   h.appName,
		HostName:  h.hostName,
		LogLevel:  strings.ToUpper(entry.Level.String()),
		FileName:  caller.File,
		FuncName:  caller.Func,
		LineNo:    caller.LineNo,
		Message:   entry.Message,
		Context:   nil,
		Timestamp: common.Number(float64(entry.Time.UnixNano()) / 1000000000),
		CID:       h.cid,
		PID:       h.pid,
	}
	rec.FromFields(entry.Data)

	var w jwriter.Writer
	rec.MarshalEasyJSON(&w)
	if w.Error != nil {
		return errors.Wrap(w.Error, "while marshalling json")
	}
	buf := w.Buffer.BuildBytes()

	if h.debug {
		fmt.Printf("%s\n", string(buf))
	}

	err = h.sendKafka(buf)
	if err != nil {
		return errors.Wrap(err, "while sending")
	}
	return nil
}

func (h *KafkaHook) sendKafka(buf []byte) error {
	select {
	case h.produce <- buf:
	default:
		// If the producer input channel buffer is full, then we better drop
		// a log record than block program execution.
		fmt.Fprintf(os.Stderr, "[kafkahook] buffer overflow: %s\n", string(buf))
	}
	return nil
}

// Given an io reader send the contents of the reader to udplog
func (h *KafkaHook) SendIO(input io.Reader) error {
	// Append our identifier
	buf := bytes.NewBuffer([]byte(""))
	_, err := buf.ReadFrom(input)
	if err != nil {
		return errors.Wrap(err, "while reading from IO")
	}

	if h.debug {
		fmt.Printf("%s\n", buf.String())
	}

	err = h.sendKafka(buf.Bytes())
	if err != nil {
		return errors.Wrap(err, "while sending")
	}
	return nil
}

// Levels returns the available logging levels.
func (h *KafkaHook) Levels() []logrus.Level {
	return logrus.AllLevels
}

func (h *KafkaHook) SetDebug(set bool) {
	h.debug = set
}

// Close the kakfa producer and flush any remaining logs
func (h *KafkaHook) Close() error {
	var err error
	h.once.Do(func() {
		close(h.produce)
		h.wg.Wait()
	})
	return err
}

func getDockerCID() string {
	f, err := os.Open("/proc/self/cgroup")
	if err != nil {
		return ""
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Split(line, "/docker/")
		if len(parts) != 2 {
			continue
		}

		fullDockerCID := parts[1]
		return fullDockerCID[:12]
	}
	return ""
}
<commit_msg>drop producer reconnect logic<commit_after>package kafkahook

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/mailgun/holster/errors"
	"github.com/mailgun/holster/stack"
	"github.com/mailgun/logrus-hooks/common"
	"github.com/mailru/easyjson/jwriter"
	"github.com/sirupsen/logrus"
)

const bufferSize = 150

type KafkaHook struct {
	produce chan []byte
	conf    Config
	debug   bool

	// Process identity metadata
	hostName string
	appName  string
	cid      string
	pid      int

	// Sync stuff
	wg   sync.WaitGroup
	once sync.Once
}

type Config struct {
	Endpoints []string
	Topic     string
	Producer  sarama.AsyncProducer
}

func NewWithContext(ctx context.Context, conf Config) (hook *KafkaHook, err error) {
	done := make(chan struct{})
	go func() {
		hook, err = New(conf)
		close(done)
	}()

	select {
	case <-ctx.Done():
		return hook, fmt.Errorf("connect timeout while connecting to kafka peers %s",
			strings.Join(conf.Endpoints, ","))
	case <-done:
		return hook, err
	}
}

func New(conf Config) (*KafkaHook, error) {
	var err error

	kafkaConfig := sarama.NewConfig()
	kafkaConfig.Producer.RequiredAcks = sarama.WaitForAll
	kafkaConfig.Producer.Compression = sarama.CompressionSnappy
	kafkaConfig.Producer.Flush.Frequency = 200 * time.Millisecond
	kafkaConfig.Producer.Retry.Backoff = 10 * time.Second
	kafkaConfig.Producer.Retry.Max = 6
	kafkaConfig.Producer.Return.Errors = true

	// If the user failed to provide a producer create one
	if conf.Producer == nil {
		conf.Producer, err = sarama.NewAsyncProducer(conf.Endpoints, kafkaConfig)
		if err != nil {
			return nil, errors.Wrap(err, "kafka producer error")
		}
	}

	h := KafkaHook{
		produce: make(chan []byte, bufferSize),
		conf:    conf,
	}

	h.wg.Add(1)
	go func() {
		for {
			select {
			case err := <-conf.Producer.Errors():
				msg, _ := err.Msg.Value.Encode()
				fmt.Fprintf(os.Stderr, "[kafkahook] produce error '%s' for: %s\n", err.Err, string(msg))

			case buf, ok := <-h.produce:
				if !ok {
					if err := conf.Producer.Close(); err != nil {
						fmt.Fprintf(os.Stderr, "[kafkahook] producer close error: %s\n", err)
					}
					h.wg.Done()
					return
				}
				conf.Producer.Input() <- &sarama.ProducerMessage{
					Value: sarama.ByteEncoder(buf),
					Topic: conf.Topic,
					Key:   nil,
				}
			}
		}
	}()

	if h.hostName, err = os.Hostname(); err != nil {
		h.hostName = "unknown"
	}
	h.appName = filepath.Base(os.Args[0])
	if h.pid = os.Getpid(); h.pid == 1 {
		h.pid = 0
	}
	h.cid = getDockerCID()
	return &h, nil
}

func (h *KafkaHook) Fire(entry *logrus.Entry) error {
	var caller *stack.FrameInfo
	var err error

	caller = common.GetLogrusCaller()

	rec := &common.LogRecord{
		Category:  "logrus",
		AppName:   h.appName,
		HostName:  h.hostName,
		LogLevel:  strings.ToUpper(entry.Level.String()),
		FileName:  caller.File,
		FuncName:  caller.Func,
		LineNo:    caller.LineNo,
		Message:   entry.Message,
		Context:   nil,
		Timestamp: common.Number(float64(entry.Time.UnixNano()) / 1000000000),
		CID:       h.cid,
		PID:       h.pid,
	}
	rec.FromFields(entry.Data)

	var w jwriter.Writer
	rec.MarshalEasyJSON(&w)
	if w.Error != nil {
		return errors.Wrap(w.Error, "while marshalling json")
	}
	buf := w.Buffer.BuildBytes()

	if h.debug {
		fmt.Printf("%s\n", string(buf))
	}

	err = h.sendKafka(buf)
	if err != nil {
		return errors.Wrap(err, "while sending")
	}
	return nil
}

func (h *KafkaHook) sendKafka(buf []byte) error {
	select {
	case h.produce <- buf:
	default:
		// If the producer input channel buffer is full, then we better drop
		// a log record than block program execution.
		fmt.Fprintf(os.Stderr, "[kafkahook] buffer overflow: %s\n", string(buf))
	}
	return nil
}

// Given an io reader send the contents of the reader to udplog
func (h *KafkaHook) SendIO(input io.Reader) error {
	// Append our identifier
	buf := bytes.NewBuffer([]byte(""))
	_, err := buf.ReadFrom(input)
	if err != nil {
		return errors.Wrap(err, "while reading from IO")
	}

	if h.debug {
		fmt.Printf("%s\n", buf.String())
	}

	err = h.sendKafka(buf.Bytes())
	if err != nil {
		return errors.Wrap(err, "while sending")
	}
	return nil
}

// Levels returns the available logging levels.
func (h *KafkaHook) Levels() []logrus.Level {
	return logrus.AllLevels
}

func (h *KafkaHook) SetDebug(set bool) {
	h.debug = set
}

// Close the kakfa producer and flush any remaining logs
func (h *KafkaHook) Close() error {
	var err error
	h.once.Do(func() {
		close(h.produce)
		h.wg.Wait()
	})
	return err
}

func getDockerCID() string {
	f, err := os.Open("/proc/self/cgroup")
	if err != nil {
		return ""
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Split(line, "/docker/")
		if len(parts) != 2 {
			continue
		}

		fullDockerCID := parts[1]
		return fullDockerCID[:12]
	}
	return ""
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/rusco/qunit"
	"github.com/soroushjp/humble/model"
	"reflect"
	"strconv"
)

type Todo struct {
	Id          int
	Title       string
	IsCompleted bool
}

func (t Todo) GetId() string {
	return strconv.Itoa(t.Id)
}

func (t Todo) RootURL() string {
	return "http://localhost:3000/todos"
}

func main() {
	qunit.AsyncTest("ReadAll", func() interface{} {
		qunit.Expect(2)
		expectedTodos := []*Todo{
			{
				Id:          0,
				Title:       "Write a frontend framework in Go",
				IsCompleted: false,
			},
			{
				Id:          1,
				Title:       "???",
				IsCompleted: false,
			},
			{
				Id:          2,
				Title:       "Profit!",
				IsCompleted: false,
			},
		}
		gotTodos := []*Todo{}
		err := model.ReadAll(&gotTodos)
		qunit.Ok(err == nil, fmt.Sprintf("model.ReadAll returned an error: %v", err))
		qunit.Ok(reflect.DeepEqual(gotTodos, expectedTodos), fmt.Sprintf("Expected: %v, Got: %v", expectedTodos, gotTodos))
		qunit.Start()
		return nil
	})
}
<commit_msg>Update karma/model_test.go to use the non-deprecated async function<commit_after>package main

import (
	"fmt"
	"github.com/rusco/qunit"
	"github.com/soroushjp/humble/model"
	"reflect"
	"strconv"
)

type Todo struct {
	Id          int
	Title       string
	IsCompleted bool
}

func (t Todo) GetId() string {
	return strconv.Itoa(t.Id)
}

func (t Todo) RootURL() string {
	return "http://localhost:3000/todos"
}

func main() {
	qunit.Test("ReadAll", func(assert qunit.QUnitAssert) {
		qunit.Expect(2)
		expectedTodos := []*Todo{
			{
				Id:          0,
				Title:       "Write a frontend framework in Go",
				IsCompleted: false,
			},
			{
				Id:          1,
				Title:       "???",
				IsCompleted: false,
			},
			{
				Id:          2,
				Title:       "Profit!",
				IsCompleted: false,
			},
		}
		done := assert.Call("async")
		go func() {
			gotTodos := []*Todo{}
			err := model.ReadAll(&gotTodos)
			assert.Ok(err == nil, fmt.Sprintf("model.ReadAll returned an error: %v", err))
			assert.Ok(reflect.DeepEqual(gotTodos, expectedTodos), fmt.Sprintf("Expected: %v, Got: %v", expectedTodos, gotTodos))
			done.Invoke()
		}()
	})
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/dustin/gojson"
	"github.com/dustin/seriesly/timelib"
	"github.com/mschoch/gouchstore"
)

func serverInfo(parts []string, w http.ResponseWriter, req *http.Request) {
	sinfo := map[string]string{
		"seriesly": "Why so series?", "version": "seriesly 0.0",
	}
	mustEncode(200, w, sinfo)
}

func listDatabases(parts []string, w http.ResponseWriter, req *http.Request) {
	mustEncode(200, w, dblist(*dbRoot))
}

func notImplemented(parts []string, w http.ResponseWriter, req *http.Request) {
	emitError(501, w, "not implemented", "TODO")
}

func createDB(parts []string, w http.ResponseWriter, req *http.Request) {
	path := dbPath(parts[0])
	err := dbcreate(path)
	if err == nil {
		w.WriteHeader(201)
	} else {
		emitError(500, w, "Server Error", err.Error())
	}
}

func checkDB(args []string, w http.ResponseWriter, req *http.Request) {
	dbname := args[0]
	if db, err := dbopen(dbname); err == nil {
		closeDBConn(db)
		w.WriteHeader(200)
	} else {
		w.WriteHeader(404)
	}
}

func newDocument(args []string, w http.ResponseWriter, req *http.Request) {
	var k, fk string
	form, err := url.ParseQuery(req.URL.RawQuery)
	if err == nil {
		fk = form.Get("ts")
	}

	if fk == "" {
		k = time.Now().UTC().Format(time.RFC3339Nano)
	} else {
		t, err := timelib.ParseTime(fk)
		if err != nil {
			emitError(400, w, "Bad time format", err.Error())
			return
		}
		k = t.UTC().Format(time.RFC3339Nano)
	}
	putDocument([]string{args[0], k}, w, req)
}

func putDocument(args []string, w http.ResponseWriter, req *http.Request) {
	dbname := args[0]
	k := args[1]
	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		emitError(400, w, "Bad Request",
			fmt.Sprintf("Error reading body: %v", err))
		return
	}

	err = json.Validate(body)
	if err != nil {
		emitError(400, w, "Error parsing JSON data", err.Error())
		return
	}

	err = dbstore(dbname, k, body)

	if err == nil {
		w.WriteHeader(201)
	} else {
		emitError(500, w, "Error storing data", err.Error())
	}
}

func cleanupRangeParam(in, def string) (string, error) {
	if in == "" {
		return def, nil
	}
	t, err := timelib.ParseTime(in)
	if err != nil {
		return in, err
	}
	return t.UTC().Format(time.RFC3339Nano), nil
}

func canGzip(req *http.Request) bool {
	acceptable := req.Header.Get("accept-encoding")
	return strings.Contains(acceptable, "gzip")
}

func query(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	group, err := strconv.Atoi(req.FormValue("group"))
	if err != nil {
		emitError(400, w, "Bad group value", err.Error())
		return
	}

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	ptrs := req.Form["ptr"]
	reds := make([]string, 0, len(ptrs))
	for _, r := range req.Form["reducer"] {
		_, ok := reducers[r]
		if !ok {
			emitError(400, w, "No such reducer", r)
			return
		}
		reds = append(reds, r)
	}

	if len(ptrs) < 1 {
		emitError(400, w, "Pointer required",
			"At least one ptr argument is required")
		return
	}

	if len(ptrs) != len(reds) {
		emitError(400, w, "Parameter mismatch",
			"Must supply the same number of pointers and reducers")
		return
	}

	filters := req.Form["f"]
	filtervals := req.Form["fv"]
	if len(filters) != len(filtervals) {
		emitError(400, w, "Parameter mismatch",
			"Must supply the same number of filters and filter values")
		return
	}

	q := executeQuery(args[0], from, to, group, ptrs, reds, filters, filtervals)
	defer close(q.out)
	defer close(q.cherr)

	output := io.Writer(w)

	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		output = gz
	} else {
		output = w
	}

	going := true
	finished := int32(0)
	started := false
	walkComplete := false
	for going {
		select {
		case po := <-q.out:
			if !started {
				started = true
				w.WriteHeader(200)
				output.Write([]byte{'{'})
			}
			if finished != 0 {
				output.Write([]byte{',', '\n'})
			}
			finished++

			_, err := fmt.Fprintf(output, `"%d": `, po.key/1e6)
			if err == nil {
				var d []byte
				d, err = json.Marshal(po.value)
				if err == nil {
					_, err = output.Write(d)
				}
			}
			if err != nil {
				log.Printf("Error sending item: %v", err)
				output = ioutil.Discard
				q.before = time.Time{}
			}
		case err = <-q.cherr:
			if err != nil {
				if !started {
					w.WriteHeader(500)
					w.Header().Set("Content-Type", "text/plain")
					fmt.Fprintf(output, "Error beginning traversal: %v", err)
				}
				log.Printf("Walk completed with err: %v", err)
				going = false
			}
			walkComplete = true
		}
		going = (q.started-finished > 0) || !walkComplete
	}

	if started {
		output.Write([]byte{'}'})
	}

	duration := time.Since(q.start)
	if duration > *minQueryLogDuration {
		log.Printf("Completed query processing in %v, %v keys, %v chunks",
			duration, humanize.Comma(int64(q.totalKeys)),
			humanize.Comma(int64(q.started)))
	}
}

func deleteBulk(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	compactAfter := strings.ToLower(req.FormValue("compact"))

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	db, err := dbopen(args[0])
	if err != nil {
		emitError(400, w, "Unable to open Database", err.Error())
		return
	}

	bulk := db.Bulk()
	deleteCount := 0
	commitThreshold := 10000

	err = dbwalkKeys(args[0], from, to, func(k string) error {

		bulk.Delete(gouchstore.NewDocumentInfo(k))
		deleteCount++
		if deleteCount >= commitThreshold {
			bulk.Commit()
			deleteCount = 0
		}
		return err
	})

	if deleteCount > 0 {
		bulk.Commit()
	}

	bulk.Close()

	if compactAfter == "true" {
		err = dbcompact(args[0])
	}

	w.WriteHeader(201)

}
func deleteDB(parts []string, w http.ResponseWriter, req *http.Request) {
	err := dbdelete(parts[0])
	if err == nil {
		mustEncode(200, w, map[string]interface{}{"ok": true})
	} else {
		emitError(500, w, "Error deleting DB", err.Error())
	}
}

func compact(parts []string, w http.ResponseWriter, req *http.Request) {
	err := dbcompact(parts[0])
	if err == nil {
		mustEncode(200, w, map[string]interface{}{"ok": true})
	} else {
		emitError(500, w, "Error compacting DB", err.Error())
	}
}

func allDocs(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	limit, err := strconv.Atoi(req.FormValue("limit"))
	if err != nil {
		limit = 2000000000
	}

	output := io.Writer(w)

	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		output = gz
	} else {
		output = w
	}
	w.WriteHeader(200)

	output.Write([]byte{'{'})
	defer output.Write([]byte{'}'})

	seenOne := false

	walked := 0
	err = dbwalk(args[0], from, to, func(k string, v []byte) error {
		if walked > limit {
			return io.EOF
		}
		walked++
		if seenOne {
			output.Write([]byte(",\n"))
		} else {
			seenOne = true
		}
		_, err := fmt.Fprintf(output, `"%s": `, k)
		if err != nil {
			return err
		}
		_, err = output.Write(v)
		return err
	})
}

func dumpDocs(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	limit, err := strconv.Atoi(req.FormValue("limit"))
	if err != nil {
		limit = 2000000000
	}

	output := io.Writer(w)

	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		output = gz
	} else {
		output = w
	}
	w.WriteHeader(200)

	walked := 0
	err = dbwalk(args[0], from, to, func(k string, v []byte) error {
		if walked > limit {
			return io.EOF
		}
		walked++
		_, err := fmt.Fprintf(output, `{"%s": `, k)
		if err != nil {
			return err
		}
		_, err = output.Write(v)
		output.Write([]byte{'}', '\n'})
		return err
	})
}

func getDocument(parts []string, w http.ResponseWriter, req *http.Request) {
	d, err := dbGetDoc(parts[0], parts[1])
	if err == nil {
		w.Write(d)
	} else {
		emitError(404, w, "Error retrieving value", err.Error())
	}
}

func dbInfo(args []string, w http.ResponseWriter, req *http.Request) {
	db, err := dbopen(args[0])
	if err != nil {
		emitError(500, w, "Error opening DB", err.Error())
		return
	}
	defer closeDBConn(db)

	inf, err := db.DatabaseInfo()
	if err == nil {
		mustEncode(200, w, map[string]interface{}{
			"last_seq":      inf.LastSeq,
			"doc_count":     inf.DocumentCount,
			"deleted_count": inf.DeletedCount,
			"space_used":    inf.SpaceUsed,
			"header_pos":    inf.HeaderPosition,
		})
	} else {
		emitError(500, w, "Error getting db info", err.Error())
	}
}

// TODO:

func dbChanges(parts []string, w http.ResponseWriter, req *http.Request) {
	notImplemented(parts, w, req)
}

func rmDocument(parts []string, w http.ResponseWriter, req *http.Request) {
	notImplemented(parts, w, req)
}
<commit_msg>common gzip stream stuff<commit_after>package main

import (
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/dustin/gojson"
	"github.com/dustin/seriesly/timelib"
	"github.com/mschoch/gouchstore"
)

func serverInfo(parts []string, w http.ResponseWriter, req *http.Request) {
	sinfo := map[string]string{
		"seriesly": "Why so series?", "version": "seriesly 0.0",
	}
	mustEncode(200, w, sinfo)
}

func listDatabases(parts []string, w http.ResponseWriter, req *http.Request) {
	mustEncode(200, w, dblist(*dbRoot))
}

func notImplemented(parts []string, w http.ResponseWriter, req *http.Request) {
	emitError(501, w, "not implemented", "TODO")
}

func createDB(parts []string, w http.ResponseWriter, req *http.Request) {
	path := dbPath(parts[0])
	err := dbcreate(path)
	if err == nil {
		w.WriteHeader(201)
	} else {
		emitError(500, w, "Server Error", err.Error())
	}
}

func checkDB(args []string, w http.ResponseWriter, req *http.Request) {
	dbname := args[0]
	if db, err := dbopen(dbname); err == nil {
		closeDBConn(db)
		w.WriteHeader(200)
	} else {
		w.WriteHeader(404)
	}
}

func newDocument(args []string, w http.ResponseWriter, req *http.Request) {
	var k, fk string
	form, err := url.ParseQuery(req.URL.RawQuery)
	if err == nil {
		fk = form.Get("ts")
	}

	if fk == "" {
		k = time.Now().UTC().Format(time.RFC3339Nano)
	} else {
		t, err := timelib.ParseTime(fk)
		if err != nil {
			emitError(400, w, "Bad time format", err.Error())
			return
		}
		k = t.UTC().Format(time.RFC3339Nano)
	}
	putDocument([]string{args[0], k}, w, req)
}

func putDocument(args []string, w http.ResponseWriter, req *http.Request) {
	dbname := args[0]
	k := args[1]
	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		emitError(400, w, "Bad Request",
			fmt.Sprintf("Error reading body: %v", err))
		return
	}

	err = json.Validate(body)
	if err != nil {
		emitError(400, w, "Error parsing JSON data", err.Error())
		return
	}

	err = dbstore(dbname, k, body)

	if err == nil {
		w.WriteHeader(201)
	} else {
		emitError(500, w, "Error storing data", err.Error())
	}
}

func cleanupRangeParam(in, def string) (string, error) {
	if in == "" {
		return def, nil
	}
	t, err := timelib.ParseTime(in)
	if err != nil {
		return in, err
	}
	return t.UTC().Format(time.RFC3339Nano), nil
}

func canGzip(req *http.Request) bool {
	acceptable := req.Header.Get("accept-encoding")
	return strings.Contains(acceptable, "gzip")
}

func query(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	group, err := strconv.Atoi(req.FormValue("group"))
	if err != nil {
		emitError(400, w, "Bad group value", err.Error())
		return
	}

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	ptrs := req.Form["ptr"]
	reds := make([]string, 0, len(ptrs))
	for _, r := range req.Form["reducer"] {
		_, ok := reducers[r]
		if !ok {
			emitError(400, w, "No such reducer", r)
			return
		}
		reds = append(reds, r)
	}

	if len(ptrs) < 1 {
		emitError(400, w, "Pointer required",
			"At least one ptr argument is required")
		return
	}

	if len(ptrs) != len(reds) {
		emitError(400, w, "Parameter mismatch",
			"Must supply the same number of pointers and reducers")
		return
	}

	filters := req.Form["f"]
	filtervals := req.Form["fv"]
	if len(filters) != len(filtervals) {
		emitError(400, w, "Parameter mismatch",
			"Must supply the same number of filters and filter values")
		return
	}

	q := executeQuery(args[0], from, to, group, ptrs, reds, filters, filtervals)
	defer close(q.out)
	defer close(q.cherr)

	output := newGzippingWriter(w, req)
	defer output.Close()

	going := true
	finished := int32(0)
	started := false
	walkComplete := false
	for going {
		select {
		case po := <-q.out:
			if !started {
				started = true
				w.WriteHeader(200)
				output.Write([]byte{'{'})
			}
			if finished != 0 {
				output.Write([]byte{',', '\n'})
			}
			finished++

			_, err := fmt.Fprintf(output, `"%d": `, po.key/1e6)
			if err == nil {
				var d []byte
				d, err = json.Marshal(po.value)
				if err == nil {
					_, err = output.Write(d)
				}
			}
			if err != nil {
				log.Printf("Error sending item: %v", err)
				output = ioutil.Discard
				q.before = time.Time{}
			}
		case err = <-q.cherr:
			if err != nil {
				if !started {
					w.WriteHeader(500)
					w.Header().Set("Content-Type", "text/plain")
					fmt.Fprintf(output, "Error beginning traversal: %v", err)
				}
				log.Printf("Walk completed with err: %v", err)
				going = false
			}
			walkComplete = true
		}
		going = (q.started-finished > 0) || !walkComplete
	}

	if started {
		output.Write([]byte{'}'})
	}

	duration := time.Since(q.start)
	if duration > *minQueryLogDuration {
		log.Printf("Completed query processing in %v, %v keys, %v chunks",
			duration, humanize.Comma(int64(q.totalKeys)),
			humanize.Comma(int64(q.started)))
	}
}

func deleteBulk(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	compactAfter := strings.ToLower(req.FormValue("compact"))

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	db, err := dbopen(args[0])
	if err != nil {
		emitError(400, w, "Unable to open Database", err.Error())
		return
	}

	bulk := db.Bulk()
	deleteCount := 0
	commitThreshold := 10000

	err = dbwalkKeys(args[0], from, to, func(k string) error {

		bulk.Delete(gouchstore.NewDocumentInfo(k))
		deleteCount++
		if deleteCount >= commitThreshold {
			bulk.Commit()
			deleteCount = 0
		}
		return err
	})

	if deleteCount > 0 {
		bulk.Commit()
	}

	bulk.Close()

	if compactAfter == "true" {
		err = dbcompact(args[0])
	}

	w.WriteHeader(201)

}
func deleteDB(parts []string, w http.ResponseWriter, req *http.Request) {
	err := dbdelete(parts[0])
	if err == nil {
		mustEncode(200, w, map[string]interface{}{"ok": true})
	} else {
		emitError(500, w, "Error deleting DB", err.Error())
	}
}

func compact(parts []string, w http.ResponseWriter, req *http.Request) {
	err := dbcompact(parts[0])
	if err == nil {
		mustEncode(200, w, map[string]interface{}{"ok": true})
	} else {
		emitError(500, w, "Error compacting DB", err.Error())
	}
}

type gzippingWriter struct {
	underlying http.ResponseWriter
	gz         *gzip.Writer
	output     io.Writer
}

func newGzippingWriter(w http.ResponseWriter, req *http.Request) *gzippingWriter {
	rv := &gzippingWriter{underlying: w, output: w}
	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		rv.gz = gzip.NewWriter(rv.underlying)
		rv.output = rv.gz
	}
	return rv
}

func (g *gzippingWriter) Write(b []byte) (int, error) {
	return g.output.Write(b)
}

func (g *gzippingWriter) Close() error {
	if g.gz != nil {
		return g.gz.Close()
	}
	return nil
}

func allDocs(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	limit, err := strconv.Atoi(req.FormValue("limit"))
	if err != nil {
		limit = 2000000000
	}

	output := newGzippingWriter(w, req)
	defer output.Close()
	w.WriteHeader(200)

	output.Write([]byte{'{'})
	defer output.Write([]byte{'}'})

	seenOne := false

	walked := 0
	err = dbwalk(args[0], from, to, func(k string, v []byte) error {
		if walked > limit {
			return io.EOF
		}
		walked++
		if seenOne {
			output.Write([]byte(",\n"))
		} else {
			seenOne = true
		}
		_, err := fmt.Fprintf(output, `"%s": `, k)
		if err != nil {
			return err
		}
		_, err = output.Write(v)
		return err
	})
}

func dumpDocs(args []string, w http.ResponseWriter, req *http.Request) {
	// Parse the params

	req.ParseForm()

	from, err := cleanupRangeParam(req.FormValue("from"), "")
	if err != nil {
		emitError(400, w, "Bad from value", err.Error())
		return
	}
	to, err := cleanupRangeParam(req.FormValue("to"), "")
	if err != nil {
		emitError(400, w, "Bad to value", err.Error())
		return
	}

	limit, err := strconv.Atoi(req.FormValue("limit"))
	if err != nil {
		limit = 2000000000
	}

	output := newGzippingWriter(w, req)
	defer output.Close()
	w.WriteHeader(200)

	walked := 0
	err = dbwalk(args[0], from, to, func(k string, v []byte) error {
		if walked > limit {
			return io.EOF
		}
		walked++
		_, err := fmt.Fprintf(output, `{"%s": `, k)
		if err != nil {
			return err
		}
		_, err = output.Write(v)
		output.Write([]byte{'}', '\n'})
		return err
	})
}

func getDocument(parts []string, w http.ResponseWriter, req *http.Request) {
	d, err := dbGetDoc(parts[0], parts[1])
	if err == nil {
		w.Write(d)
	} else {
		emitError(404, w, "Error retrieving value", err.Error())
	}
}

func dbInfo(args []string, w http.ResponseWriter, req *http.Request) {
	db, err := dbopen(args[0])
	if err != nil {
		emitError(500, w, "Error opening DB", err.Error())
		return
	}
	defer closeDBConn(db)

	inf, err := db.DatabaseInfo()
	if err == nil {
		mustEncode(200, w, map[string]interface{}{
			"last_seq":      inf.LastSeq,
			"doc_count":     inf.DocumentCount,
			"deleted_count": inf.DeletedCount,
			"space_used":    inf.SpaceUsed,
			"header_pos":    inf.HeaderPosition,
		})
	} else {
		emitError(500, w, "Error getting db info", err.Error())
	}
}

// TODO:

func dbChanges(parts []string, w http.ResponseWriter, req *http.Request) {
	notImplemented(parts, w, req)
}

func rmDocument(parts []string, w http.ResponseWriter, req *http.Request) {
	notImplemented(parts, w, req)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"path"
	"strings"

	"github.com/gorilla/mux"
)

const imageTypes = ".jpg .jpeg .png .gif"

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintln(w, "Hello world!")
}

func WikiHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	// Deny requests trying to traverse up the directory structure using
	// relative paths
	if strings.Contains(vars["filepath"], "..") {
		http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
		return
	}

	// Path to the file as it is on the the local file system
	realpath := fmt.Sprintf("%s/%s", options.Dir, vars["filepath"])

	// Serve (accepted) images
	for _, filext := range strings.Split(imageTypes, " ") {
		if path.Ext(r.URL.Path) == filext {
			http.ServeFile(w, r, realpath)
			return
		}
	}

	md, err := ioutil.ReadFile(realpath + ".md")
	if err != nil {
		http.NotFound(w, r)
		return
	}

	wiki := Wiki{
		Markdown: md,
		filepath: realpath,
		template: options.template,
	}

	wiki.Commits, err = Commits(vars["filepath"]+".md", 5)
	if err != nil {
		log.Println("ERROR", "Failed to get commits")
	}

	wiki.Write(w)
}
<commit_msg>Improve handler for landing page.<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"path"
	"strings"

	"github.com/gorilla/mux"
)

const imageTypes = ".jpg .jpeg .png .gif"

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	md, err := ioutil.ReadFile(options.Dir + "/home.md")
	if err != nil {
		log.Fatalln(err)
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}

	wiki := Wiki{Markdown: md, template: options.template}
	wiki.Write(w)
}

func WikiHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	// Deny requests trying to traverse up the directory structure using
	// relative paths
	if strings.Contains(vars["filepath"], "..") {
		http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
		return
	}

	// Path to the file as it is on the the local file system
	realpath := fmt.Sprintf("%s/%s", options.Dir, vars["filepath"])

	// Serve (accepted) images
	for _, filext := range strings.Split(imageTypes, " ") {
		if path.Ext(r.URL.Path) == filext {
			http.ServeFile(w, r, realpath)
			return
		}
	}

	md, err := ioutil.ReadFile(realpath + ".md")
	if err != nil {
		http.NotFound(w, r)
		return
	}

	wiki := Wiki{
		Markdown: md,
		filepath: realpath,
		template: options.template,
	}

	wiki.Commits, err = Commits(vars["filepath"]+".md", 5)
	if err != nil {
		log.Println("ERROR", "Failed to get commits")
	}

	wiki.Write(w)
}
<|endoftext|>


<commit_before>package rest

import (
	"github.com/zatiti/router"
	"net/http"
)

// write creates a http handler for creating or updating a document depending on the mode provided
func (s *Service) write(modelFactory ModelFactory, mode string) router.Handler {
	return func(w http.ResponseWriter, r *http.Request) {
		// model is request scoped
		model := modelFactory.New(r)
		// event is the name used to track metrics
		event := model.Name() + "_" + mode + "_one"
		// Track how long this function take to return
		stop := s.NewTimer(event)
		defer stop()
		// HTTP response status code
		var status int
		// HTTP response body
		var body []byte
		// Instanciate a value of the model being created from the request body
		err := model.Decode()
		if err == nil {
			// Validate what came through the wire
			verr := model.Validate()
			if verr == nil {
				var v interface{}
				if mode == "insert" {
					v, err = model.Create()
					status = http.StatusCreated
				}
				if mode == "update" {
					v, err = model.Update()
					status = http.StatusNoContent
				}
				if mode == "upsert" {
					v, err = model.Upsert()
					status = http.StatusAccepted
				}
				if err == nil {
					// The output from model.Create could be invalid
					body, _ = model.Encode(v)
					//if err == nil {
					// Allow event broker to be optional
					if s.Broker != nil {
						err = s.Broker.Publish(event, body)
					}
					if err == nil {
						// If a metrics client is defined use it
						if s.Metrics != nil {
							s.Metrics.Incr(event, 1)
						}
					} else {
						// Something wicked happened while publishing to the event stream
						s.Logger.Error(err)
						status, body = InternalServerErrorResponse()
					}
					//} else {
					//	s.Logger.Error(err)
					//	status, body = InternalServerErrorResponse()
					//}
				} else {
					s.Logger.Error(err)
					status, body = InternalServerErrorResponse()
				}
			} else {
				s.Logger.Error(verr.Message)
				status, body = verr.Code, []byte(verr.Message)
			}
		} else {
			s.Logger.Error(err)
			status, body = BadRequestResponse()
		}
		w.WriteHeader(status)
		w.Write(body)
	}
}

// Insert creates a http handler that will create a document in model's database.
func (s *Service) Insert(modelFactory ModelFactory) router.Handler {
	return s.write(modelFactory, "insert")
}

// Update creates a http handler that will updates a document by the model's update selector in model's database
func (s *Service) Update(modelFactory ModelFactory) router.Handler {
	return s.write(modelFactory, "update")
}

// Upsert creates a http handler that will upsert(create or update if it exists) a document selected by the model's upsert selector
func (s *Service) Upsert(modelFactory ModelFactory) router.Handler {
	return s.write(modelFactory, "upsert")
}

// find creates a http handler that will list documents or return one document from a model's database
func (s *Service) find(modelFactory ModelFactory, mode string) router.Handler {
	return func(w http.ResponseWriter, r *http.Request) {
		// model is request scoped
		model := modelFactory.New(r)
		// event is the name used to track metrics
		event := model.Name() + "_find_many"
		// HTTP response status code
		var status int
		// HTTP response body
		var body []byte
		// Track how long this function take to return
		stop := s.NewTimer(event)
		defer stop()
		// Validate
		verr := model.Validate()
		if verr == nil {
			var v interface{}
			var err error
			if mode == "one" {
				v, err = model.FindOne()
			}
			if mode == "many" {
				v, err = model.FindMany()
			}
			if err == nil {
				body, _ = model.Encode(v)
				// Notify other services, if an event broker exists
				if s.Broker != nil {
					err = s.Broker.Publish(event, body)
				}
				if err == nil {
					status = http.StatusOK
					if s.Metrics != nil {
						s.Metrics.Incr(event, 1)
					}
				} else {
					s.Logger.Error(err)
					status, body = InternalServerErrorResponse()
				}
			} else {
				// Something wicked happened while fetching document/s
				s.Logger.Error(err)
				status, body = InternalServerErrorResponse()
			}
		} else {
			s.Logger.Error(verr.Message)
			status, body = verr.Code, []byte(verr.Message)
		}
		w.WriteHeader(status)
		w.Write(body)
	}
}

// FindOne - creates a http handler that will return one document from a model's database if the id exists
func (s *Service) FindOne(modelFactory ModelFactory) router.Handler {
	return s.find(modelFactory, "one")
}

// FindMany - creates a http handler that will list documents from a model's database
func (s *Service) FindMany(modelFactory ModelFactory) router.Handler {
	return s.find(modelFactory, "many")
}

// Remove creates a http handler that will delete a document by remove selector specified in the model model
func (s *Service) Remove(modelFactory ModelFactory) router.Handler {
	return func(w http.ResponseWriter, r *http.Request) {
		model := modelFactory.New(r)
		// event is the name used to track metrics
		event := model.Name() + "_delete_one"
		// Track how long this function take to return
		stop := s.NewTimer(event)
		defer stop()
		// HTTP response status code
		var status int
		// HTTP response body
		var body []byte
		verr := model.Validate()
		if verr == nil {
			// Remove the item if it exists
			v, err := model.Remove()
			if err == nil {
				// Encode the output into []byte
				body, _ = model.Encode(v)
				//if err == nil {
				// Notify other services, if an event broker exists
				if s.Broker != nil {
					err = s.Broker.Publish(event, body)
				}
				if err == nil {
					status, body = NoContentResponse()
					// If a metrics client is defined use it
					if s.Metrics != nil {
						s.Metrics.Incr(event, 1)
					}
				} else {
					status, body = InternalServerErrorResponse()
					s.Logger.Error(err)
				}
				//} else {
				//	status, body = InternalServerErrorResponse()
				//	s.Logger.Error(err)
				//}
			} else {
				status, body = InternalServerErrorResponse()
				s.Logger.Error(err)
			}
		} else {
			status, body = verr.Code, []byte(verr.Message)
			s.Logger.Error(verr.Message)
		}
		w.WriteHeader(status)
		w.Write(body)
	}
}
<commit_msg>use versioned router lib<commit_after>package rest

import (
	"gopkg.in/zatiti/router.v1"
	"net/http"
)

// write creates a http handler for creating or updating a document depending on the mode provided
func (s *Service) write(modelFactory ModelFactory, mode string) router.Handler {
	return func(w http.ResponseWriter, r *http.Request) {
		// model is request scoped
		model := modelFactory.New(r)
		// event is the name used to track metrics
		event := model.Name() + "_" + mode + "_one"
		// Track how long this function take to return
		stop := s.NewTimer(event)
		defer stop()
		// HTTP response status code
		var status int
		// HTTP response body
		var body []byte
		// Instanciate a value of the model being created from the request body
		err := model.Decode()
		if err == nil {
			// Validate what came through the wire
			verr := model.Validate()
			if verr == nil {
				var v interface{}
				if mode == "insert" {
					v, err = model.Create()
					status = http.StatusCreated
				}
				if mode == "update" {
					v, err = model.Update()
					status = http.StatusNoContent
				}
				if mode == "upsert" {
					v, err = model.Upsert()
					status = http.StatusAccepted
				}
				if err == nil {
					// The output from model.Create could be invalid
					body, _ = model.Encode(v)
					//if err == nil {
					// Allow event broker to be optional
					if s.Broker != nil {
						err = s.Broker.Publish(event, body)
					}
					if err == nil {
						// If a metrics client is defined use it
						if s.Metrics != nil {
							s.Metrics.Incr(event, 1)
						}
					} else {
						// Something wicked happened while publishing to the event stream
						s.Logger.Error(err)
						status, body = InternalServerErrorResponse()
					}
					//} else {
					//	s.Logger.Error(err)
					//	status, body = InternalServerErrorResponse()
					//}
				} else {
					s.Logger.Error(err)
					status, body = InternalServerErrorResponse()
				}
			} else {
				s.Logger.Error(verr.Message)
				status, body = verr.Code, []byte(verr.Message)
			}
		} else {
			s.Logger.Error(err)
			status, body = BadRequestResponse()
		}
		w.WriteHeader(status)
		w.Write(body)
	}
}

// Insert creates a http handler that will create a document in model's database.
func (s *Service) Insert(modelFactory ModelFactory) router.Handler {
	return s.write(modelFactory, "insert")
}

// Update creates a http handler that will updates a document by the model's update selector in model's database
func (s *Service) Update(modelFactory ModelFactory) router.Handler {
	return s.write(modelFactory, "update")
}

// Upsert creates a http handler that will upsert(create or update if it exists) a document selected by the model's upsert selector
func (s *Service) Upsert(modelFactory ModelFactory) router.Handler {
	return s.write(modelFactory, "upsert")
}

// find creates a http handler that will list documents or return one document from a model's database
func (s *Service) find(modelFactory ModelFactory, mode string) router.Handler {
	return func(w http.ResponseWriter, r *http.Request) {
		// model is request scoped
		model := modelFactory.New(r)
		// event is the name used to track metrics
		event := model.Name() + "_find_many"
		// HTTP response status code
		var status int
		// HTTP response body
		var body []byte
		// Track how long this function take to return
		stop := s.NewTimer(event)
		defer stop()
		// Validate
		verr := model.Validate()
		if verr == nil {
			var v interface{}
			var err error
			if mode == "one" {
				v, err = model.FindOne()
			}
			if mode == "many" {
				v, err = model.FindMany()
			}
			if err == nil {
				body, _ = model.Encode(v)
				// Notify other services, if an event broker exists
				if s.Broker != nil {
					err = s.Broker.Publish(event, body)
				}
				if err == nil {
					status = http.StatusOK
					if s.Metrics != nil {
						s.Metrics.Incr(event, 1)
					}
				} else {
					s.Logger.Error(err)
					status, body = InternalServerErrorResponse()
				}
			} else {
				// Something wicked happened while fetching document/s
				s.Logger.Error(err)
				status, body = InternalServerErrorResponse()
			}
		} else {
			s.Logger.Error(verr.Message)
			status, body = verr.Code, []byte(verr.Message)
		}
		w.WriteHeader(status)
		w.Write(body)
	}
}

// FindOne - creates a http handler that will return one document from a model's database if the id exists
func (s *Service) FindOne(modelFactory ModelFactory) router.Handler {
	return s.find(modelFactory, "one")
}

// FindMany - creates a http handler that will list documents from a model's database
func (s *Service) FindMany(modelFactory ModelFactory) router.Handler {
	return s.find(modelFactory, "many")
}

// Remove creates a http handler that will delete a document by remove selector specified in the model model
func (s *Service) Remove(modelFactory ModelFactory) router.Handler {
	return func(w http.ResponseWriter, r *http.Request) {
		model := modelFactory.New(r)
		// event is the name used to track metrics
		event := model.Name() + "_delete_one"
		// Track how long this function take to return
		stop := s.NewTimer(event)
		defer stop()
		// HTTP response status code
		var status int
		// HTTP response body
		var body []byte
		verr := model.Validate()
		if verr == nil {
			// Remove the item if it exists
			v, err := model.Remove()
			if err == nil {
				// Encode the output into []byte
				body, _ = model.Encode(v)
				//if err == nil {
				// Notify other services, if an event broker exists
				if s.Broker != nil {
					err = s.Broker.Publish(event, body)
				}
				if err == nil {
					status, body = NoContentResponse()
					// If a metrics client is defined use it
					if s.Metrics != nil {
						s.Metrics.Incr(event, 1)
					}
				} else {
					status, body = InternalServerErrorResponse()
					s.Logger.Error(err)
				}
				//} else {
				//	status, body = InternalServerErrorResponse()
				//	s.Logger.Error(err)
				//}
			} else {
				status, body = InternalServerErrorResponse()
				s.Logger.Error(err)
			}
		} else {
			status, body = verr.Code, []byte(verr.Message)
			s.Logger.Error(verr.Message)
		}
		w.WriteHeader(status)
		w.Write(body)
	}
}
<|endoftext|>


<commit_before>
package main

import (
        "os"
        "os/signal"
        "syscall"
        "fmt"
        "log"
        "time"
        "flag"
        "runtime"
        "runtime/pprof"
        )

const VERSION = "0.1.5"

var (
    myId string
    cache *Cache
    runningCheckers = 0
    chanCheckers = make(map[string]chan int)
)

func addCheck(line string) {
    check, err := NewCheck(line)
    if err != nil {
        log.Printf("Warning: got invalid data on the \"dead\" " +
            "channel (\"%s\")\n", line)
        return
    }
    if check.BackendGroupLength <= 1 {
        // Add the check only if the frontend is scaled to several
        // backends (backend is part of a group)
        return
    }
    r, isMine := cache.LockBackend(check)
    if r == false {
        // Backend already locked
        if isMine == true {
            // If the backend is already monitored by the current process,
            // we advertise the goroutine (pingUrl) that we got a new dead
            // event (probably from a passive check inside Hipache)
            ch, exists := chanCheckers[check.BackendUrl]
            if exists {
                // Non-blocking send
                select {
                    case ch <- 1:
                    default:
                }
            }
        }
        return
    }
    // Create the routine's channel
    ch := make(chan int, 1)
    chanCheckers[check.BackendUrl] = ch
    // Set all the callbacks for the check. They will be called during
    // the PingUrl at different steps
    check.SetDeadCallback(func () {
        cache.MarkBackendDead(check)
    })
    check.SetAliveCallback(func () {
        cache.MarkBackendAlive(check)
    })
    check.SetCheckIfBreakCallback(func () bool {
        return cache.IsUnlockedBackend(check)
    })
    check.SetExitCallback(func () {
        runningCheckers -= 1
        cache.UnlockBackend(check)
        // clear the channel
        delete(chanCheckers, check.BackendUrl)
    })
    // Check the URL at a regular interval
    go check.PingUrl(ch)
    runningCheckers += 1
    log.Printf("Added check for: %s\n", check.BackendUrl)
}

/*
 * Prints some stats on runtime
 */
func printStats(cache *Cache) {
    go func () {
        const step = 10 // 10 seconds
        count := 0
        for {
            cache.PingAlive()
            time.Sleep(time.Duration(step) * time.Second)
            count += step
            if count >= 60 {
                // Every minute
                count = 0
                log.Printf("%d backend URLs are being tested\n", runningCheckers)
            }
        }
    }()
}

/*
 * Enables CPU profile
 */
func enableCPUProfile() {
    cwd, _ := os.Getwd()
    log.Printf("CPU profile will be written to \"%s/%s\"", cwd, "hchecker.prof")
    f, err := os.Create("hchecker.prof")
    if err != nil {
        log.Fatal("Cannot enable CPU profile:", err)
    }
    pprof.StartCPUProfile(f)
}

/*
 * Listens to signals
 */
func handleSignals() {
    c := make(chan os.Signal)
    signal.Notify(c, syscall.SIGINT)
    go func () {
        switch <-c {
            case syscall.SIGINT:
                pprof.StopCPUProfile()
                os.Exit(0)
        }
    }()
}

func parseFlags(cpuProfile *bool) {
    parseDuration := func (v *time.Duration, n string, def int, help string) {
        i := flag.Int(n, def, help)
        *v = time.Duration(*i) * time.Second
    }
    flag.StringVar(&httpMethod, "method", HTTP_METHOD,
        "HTTP method")
    parseDuration(&checkInterval, "interval", CHECK_INTERVAL,
        "Check interval (seconds)")
    parseDuration(&connectionTimeout, "connect", CONNECTION_TIMEOUT,
        "TCP connection timeout (seconds)")
    parseDuration(&ioTimeout, "io", IO_TIMEOUT,
        "Socket read/write timeout (seconds)")
    flag.StringVar(&userAgent, "agent", USER_AGENT,
        "HTTP User-Agent header")
    flag.StringVar(&redisAddress, "redis", REDIS_ADDRESS,
        "Network address of Redis")
    flag.BoolVar(cpuProfile, "cpuprofile", false,
        "Write CPU profile to \"hchecker.prof\" (current directory)")
    flag.Parse()
}

func main() {
    var (
        err error
        hostname string
        cpuProfile bool
        )
    for  _, arg := range os.Args {
        if !(arg == "-v" || arg == "--version" || arg == "-version") {
            continue
        }
        fmt.Println("hchecker version", VERSION)
        os.Exit(0)
    }
    parseFlags(&cpuProfile)
    // Force 1 CPU to reduce parallelism. If you want to use more CPUs, prefer
    // spawning several processes instead.
    runtime.GOMAXPROCS(1)
    hostname, _ = os.Hostname()
    myId = fmt.Sprintf("%s#%d", hostname, os.Getpid())
    // Prefix each line of log
    log.SetPrefix(myId + " ")
    if cpuProfile == true {
        enableCPUProfile()
    }
    handleSignals()
    cache, err = NewCache()
    if err != nil {
        log.Println(err.Error())
        os.Exit(1)
    }
    printStats(cache)
    err = cache.ListenToChannel("dead", addCheck)
    if err != nil {
        log.Println(err.Error())
    }
}
<commit_msg>Bumped version to 0.1.6<commit_after>
package main

import (
        "os"
        "os/signal"
        "syscall"
        "fmt"
        "log"
        "time"
        "flag"
        "runtime"
        "runtime/pprof"
        )

const VERSION = "0.1.6"

var (
    myId string
    cache *Cache
    runningCheckers = 0
    chanCheckers = make(map[string]chan int)
)

func addCheck(line string) {
    check, err := NewCheck(line)
    if err != nil {
        log.Printf("Warning: got invalid data on the \"dead\" " +
            "channel (\"%s\")\n", line)
        return
    }
    if check.BackendGroupLength <= 1 {
        // Add the check only if the frontend is scaled to several
        // backends (backend is part of a group)
        return
    }
    r, isMine := cache.LockBackend(check)
    if r == false {
        // Backend already locked
        if isMine == true {
            // If the backend is already monitored by the current process,
            // we advertise the goroutine (pingUrl) that we got a new dead
            // event (probably from a passive check inside Hipache)
            ch, exists := chanCheckers[check.BackendUrl]
            if exists {
                // Non-blocking send
                select {
                    case ch <- 1:
                    default:
                }
            }
        }
        return
    }
    // Create the routine's channel
    ch := make(chan int, 1)
    chanCheckers[check.BackendUrl] = ch
    // Set all the callbacks for the check. They will be called during
    // the PingUrl at different steps
    check.SetDeadCallback(func () {
        cache.MarkBackendDead(check)
    })
    check.SetAliveCallback(func () {
        cache.MarkBackendAlive(check)
    })
    check.SetCheckIfBreakCallback(func () bool {
        return cache.IsUnlockedBackend(check)
    })
    check.SetExitCallback(func () {
        runningCheckers -= 1
        cache.UnlockBackend(check)
        // clear the channel
        delete(chanCheckers, check.BackendUrl)
    })
    // Check the URL at a regular interval
    go check.PingUrl(ch)
    runningCheckers += 1
    log.Printf("Added check for: %s\n", check.BackendUrl)
}

/*
 * Prints some stats on runtime
 */
func printStats(cache *Cache) {
    go func () {
        const step = 10 // 10 seconds
        count := 0
        for {
            cache.PingAlive()
            time.Sleep(time.Duration(step) * time.Second)
            count += step
            if count >= 60 {
                // Every minute
                count = 0
                log.Printf("%d backend URLs are being tested\n", runningCheckers)
            }
        }
    }()
}

/*
 * Enables CPU profile
 */
func enableCPUProfile() {
    cwd, _ := os.Getwd()
    log.Printf("CPU profile will be written to \"%s/%s\"", cwd, "hchecker.prof")
    f, err := os.Create("hchecker.prof")
    if err != nil {
        log.Fatal("Cannot enable CPU profile:", err)
    }
    pprof.StartCPUProfile(f)
}

/*
 * Listens to signals
 */
func handleSignals() {
    c := make(chan os.Signal)
    signal.Notify(c, syscall.SIGINT)
    go func () {
        switch <-c {
            case syscall.SIGINT:
                pprof.StopCPUProfile()
                os.Exit(0)
        }
    }()
}

func parseFlags(cpuProfile *bool) {
    parseDuration := func (v *time.Duration, n string, def int, help string) {
        i := flag.Int(n, def, help)
        *v = time.Duration(*i) * time.Second
    }
    flag.StringVar(&httpMethod, "method", HTTP_METHOD,
        "HTTP method")
    parseDuration(&checkInterval, "interval", CHECK_INTERVAL,
        "Check interval (seconds)")
    parseDuration(&connectionTimeout, "connect", CONNECTION_TIMEOUT,
        "TCP connection timeout (seconds)")
    parseDuration(&ioTimeout, "io", IO_TIMEOUT,
        "Socket read/write timeout (seconds)")
    flag.StringVar(&userAgent, "agent", USER_AGENT,
        "HTTP User-Agent header")
    flag.StringVar(&redisAddress, "redis", REDIS_ADDRESS,
        "Network address of Redis")
    flag.BoolVar(cpuProfile, "cpuprofile", false,
        "Write CPU profile to \"hchecker.prof\" (current directory)")
    flag.Parse()
}

func main() {
    var (
        err error
        hostname string
        cpuProfile bool
        )
    for  _, arg := range os.Args {
        if !(arg == "-v" || arg == "--version" || arg == "-version") {
            continue
        }
        fmt.Println("hchecker version", VERSION)
        os.Exit(0)
    }
    parseFlags(&cpuProfile)
    // Force 1 CPU to reduce parallelism. If you want to use more CPUs, prefer
    // spawning several processes instead.
    runtime.GOMAXPROCS(1)
    hostname, _ = os.Hostname()
    myId = fmt.Sprintf("%s#%d", hostname, os.Getpid())
    // Prefix each line of log
    log.SetPrefix(myId + " ")
    if cpuProfile == true {
        enableCPUProfile()
    }
    handleSignals()
    cache, err = NewCache()
    if err != nil {
        log.Println(err.Error())
        os.Exit(1)
    }
    printStats(cache)
    err = cache.ListenToChannel("dead", addCheck)
    if err != nil {
        log.Println(err.Error())
    }
}
<|endoftext|>