<commit_before>package ydls

// TODO: test close reader prematurely

import (
	"context"
	"encoding/xml"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/wader/ydls/internal/ffmpeg"
	"github.com/wader/ydls/internal/leaktest"
	"github.com/wader/ydls/internal/rss"
	"github.com/wader/ydls/internal/stringprioset"
	"github.com/wader/ydls/internal/timerange"
	"github.com/wader/ydls/internal/youtubedl"
)

var youtubeTestVideoURL = "https://www.youtube.com/watch?v=C0DPdy98e4c"
var youtubeLongTestVideoURL = "https://www.youtube.com/watch?v=z7VYVjR_nwE"
var soundcloudTestAudioURL = "https://soundcloud.com/timsweeney/thedrifter"
var soundcloudTestPlaylistURL = "https://soundcloud.com/mattheis/sets/kindred-phenomena"

var testNetwork = os.Getenv("TEST_NETWORK") != ""
var testYoutubeldl = os.Getenv("TEST_YOUTUBEDL") != ""
var testFfmpeg = os.Getenv("TEST_FFMPEG") != ""

func stringsContains(strings []string, s string) bool {
	for _, ss := range strings {
		if ss == s {
			return true
		}
	}

	return false
}

func ydlsFromEnv(t *testing.T) YDLS {
	ydls, err := NewFromFile(os.Getenv("CONFIG"))
	if err != nil {
		t.Fatalf("failed to read config: %s", err)
	}

	return ydls
}

func TestSafeFilename(t *testing.T) {
	for _, c := range []struct {
		s      string
		expect string
	}{
		{"aba", "aba"},
		{"a/a", "a_a"},
		{"a\\a", "a_a"},
	} {
		actual := safeFilename(c.s)
		if actual != c.expect {
			t.Errorf("%s, got %v expected %v", c.s, actual, c.expect)
		}
	}
}

func TestForceCodec(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)
	forceCodecs := []string{"opus", "vp9"}

	// make sure codecs are not the perfered ones
	for _, s := range mkvFormat.Streams {
		for _, forceCodec := range forceCodecs {
			if c, ok := s.CodecNames.First(); ok && c == forceCodec {
				t.Errorf("test sanity check failed: codec already the preferred one")
				return
			}
		}
	}

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeTestVideoURL,
			Format:      &mkvFormat,
			Codecs:      forceCodecs,
		},
		nil)
	if err != nil {
		cancelFn()
		t.Errorf("%s: download failed: %s", youtubeTestVideoURL, err)
		return
	}

	pi, err := ffmpeg.Probe(ctx, ffmpeg.Reader{Reader: io.LimitReader(dr.Media, 10*1024*1024)}, nil, nil)
	dr.Media.Close()
	dr.Wait()
	cancelFn()
	if err != nil {
		t.Errorf("%s: probe failed: %s", youtubeTestVideoURL, err)
		return
	}

	if pi.FormatName() != "matroska" {
		t.Errorf("%s: force codec failed: found %s", youtubeTestVideoURL, pi)
		return
	}

	for i := 0; i < len(forceCodecs); i++ {
		if pi.Streams[i].CodecName != forceCodecs[i] {
			t.Errorf("%s: force codec failed: %s != %s", youtubeTestVideoURL, pi.Streams[i].CodecName, forceCodecs[i])
			return
		}
	}

}

func TestTimeRangeOption(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	timeRange, timeRangeErr := timerange.NewTimeRangeFromString("10s-15s")
	if timeRangeErr != nil {
		t.Fatalf("failed to parse time range")
	}

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeTestVideoURL,
			Format:      &mkvFormat,
			TimeRange:   timeRange,
		},
		nil)
	if err != nil {
		cancelFn()
		t.Fatalf("%s: download failed: %s", youtubeTestVideoURL, err)
	}

	pi, err := ffmpeg.Probe(ctx, ffmpeg.Reader{Reader: io.LimitReader(dr.Media, 10*1024*1024)}, nil, nil)
	dr.Media.Close()
	dr.Wait()
	cancelFn()
	if err != nil {
		t.Errorf("%s: probe failed: %s", youtubeTestVideoURL, err)
		return
	}

	if pi.Duration() != timeRange.Duration() {
		t.Errorf("%s: probed duration not %v, got %v", youtubeTestVideoURL, timeRange.Duration(), pi.Duration())
		return
	}
}

func TestMissingMediaStream(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	_, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: soundcloudTestAudioURL,
			Format:      &mkvFormat,
		},
		nil)
	cancelFn()
	if err == nil {
		t.Fatal("expected download to fail")
	}
}

func TestFindYDLFormat(t *testing.T) {
	ydlFormats := []youtubedl.Format{
		{FormatID: "1", Protocol: "http", NormACodec: "mp3", NormVCodec: "h264", NormBR: 1},
		{FormatID: "2", Protocol: "http", NormACodec: "", NormVCodec: "h264", NormBR: 2},
		{FormatID: "3", Protocol: "http", NormACodec: "aac", NormVCodec: "", NormBR: 3},
		{FormatID: "4", Protocol: "http", NormACodec: "vorbis", NormVCodec: "vp8", NormBR: 4},
		{FormatID: "5", Protocol: "http", NormACodec: "opus", NormVCodec: "vp9", NormBR: 5},
	}

	for i, c := range []struct {
		ydlFormats       []youtubedl.Format
		mediaType        mediaType
		codecs           stringprioset.Set
		expectedFormatID string
	}{
		{ydlFormats, MediaAudio, stringprioset.New([]string{"mp3"}), "1"},
		{ydlFormats, MediaAudio, stringprioset.New([]string{"aac"}), "3"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"h264"}), "2"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"h264"}), "2"},
		{ydlFormats, MediaAudio, stringprioset.New([]string{"vorbis"}), "4"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"vp8"}), "4"},
		{ydlFormats, MediaAudio, stringprioset.New([]string{"opus"}), "5"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"vp9"}), "5"},
	} {
		actualFormat, actaulFormatFound := findYDLFormat(c.ydlFormats, c.mediaType, c.codecs)
		if actaulFormatFound && actualFormat.FormatID != c.expectedFormatID {
			t.Errorf("%d: expected format %s, got %s", i, c.expectedFormatID, actualFormat)
		}
	}
}

func TestContextCloseProbe(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		time.Sleep(2 * time.Second)
		cancelFn()
		wg.Done()
	}()
	_, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeLongTestVideoURL,
			Format:      &mkvFormat,
		},
		nil)
	if err == nil {
		t.Error("expected error while probe")
	}
	cancelFn()
	wg.Wait()
}

func TestContextCloseDownload(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeLongTestVideoURL,
			Format:      &mkvFormat,
		},
		nil)
	if err != nil {
		t.Error("expected no error while probe")
	}
	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		time.Sleep(1 * time.Second)
		cancelFn()
		wg.Done()
	}()
	io.Copy(ioutil.Discard, dr.Media)
	cancelFn()
	wg.Wait()
}

func TestRSS(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "rss"
	rssFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: soundcloudTestPlaylistURL,
			Format:      &rssFormat,
			BaseURL:     &url.URL{Scheme: "http", Host: "dummy"},
		},
		nil)
	if err != nil {
		cancelFn()
		t.Fatalf("%s: download failed: %s", youtubeTestVideoURL, err)
	}
	defer cancelFn()

	if dr.Filename != "" {
		t.Errorf("expected no filename, got %s", dr.Filename)
	}
	expectedMIMEType := "text/xml"
	if dr.MIMEType != expectedMIMEType {
		t.Errorf("expected mimetype %s, got %s", expectedMIMEType, dr.MIMEType)
	}

	rssRoot := rss.RSS{}
	decoder := xml.NewDecoder(dr.Media)
	decodeErr := decoder.Decode(&rssRoot)
	if decodeErr != nil {
		t.Errorf("failed to parse rss: %s", decodeErr)
		return
	}
	dr.Media.Close()
	dr.Wait()

	expectedTitle := "Kindred Phenomena"
	if rssRoot.Channel.Title != expectedTitle {
		t.Errorf("expected title \"%s\" got \"%s\"", expectedTitle, rssRoot.Channel.Title)
	}

	// TODO: description, not returned by youtube-dl?

	expectedEntries := 8
	if len(rssRoot.Channel.Items) != expectedEntries {
		t.Errorf("expected %d entries got %d", expectedEntries, len(rssRoot.Channel.Items))
	}

	itemOne := rssRoot.Channel.Items[0]

	expectedItemTitle := "A1 Mattheis - Herds"
	if rssRoot.Channel.Items[0].Title != expectedItemTitle {
		t.Errorf("expected title \"%s\" got \"%s\"", expectedItemTitle, itemOne.Title)
	}

	expectedItemDescriptionPrefix := "Releasing my debut"
	if !strings.HasPrefix(rssRoot.Channel.Items[0].Description, expectedItemDescriptionPrefix) {
		t.Errorf("expected description prefix \"%s\" got \"%s\"", expectedItemDescriptionPrefix, itemOne.Description)
	}

	expectedItemGUID := "http://dummy/mp3/https://soundcloud.com/mattheis/sets/kindred-phenomena#293285002"
	if rssRoot.Channel.Items[0].GUID != expectedItemGUID {
		t.Errorf("expected guid \"%s\" got \"%s\"", expectedItemGUID, itemOne.GUID)
	}

	expectedItemURL := "http://dummy/media.mp3?format=mp3&url=https%3A%2F%2Fsoundcloud.com%2Fmattheis%2Fa1-mattheis-herds"
	if itemOne.Enclosure.URL != expectedItemURL {
		t.Errorf("expected enclousure url \"%s\" got \"%s\"", expectedItemURL, itemOne.Enclosure.URL)
	}

	expectedItemType := "audio/mpeg"
	if itemOne.Enclosure.Type != expectedItemType {
		t.Errorf("expected enclousure type \"%s\" got \"%s\"", expectedItemType, itemOne.Enclosure.Type)
	}
}
<commit_msg>Fix some grammar<commit_after>package ydls

// TODO: test close reader prematurely

import (
	"context"
	"encoding/xml"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/wader/ydls/internal/ffmpeg"
	"github.com/wader/ydls/internal/leaktest"
	"github.com/wader/ydls/internal/rss"
	"github.com/wader/ydls/internal/stringprioset"
	"github.com/wader/ydls/internal/timerange"
	"github.com/wader/ydls/internal/youtubedl"
)

var youtubeTestVideoURL = "https://www.youtube.com/watch?v=C0DPdy98e4c"
var youtubeLongTestVideoURL = "https://www.youtube.com/watch?v=z7VYVjR_nwE"
var soundcloudTestAudioURL = "https://soundcloud.com/timsweeney/thedrifter"
var soundcloudTestPlaylistURL = "https://soundcloud.com/mattheis/sets/kindred-phenomena"

var testNetwork = os.Getenv("TEST_NETWORK") != ""
var testYoutubeldl = os.Getenv("TEST_YOUTUBEDL") != ""
var testFfmpeg = os.Getenv("TEST_FFMPEG") != ""

func stringsContains(strings []string, s string) bool {
	for _, ss := range strings {
		if ss == s {
			return true
		}
	}

	return false
}

func ydlsFromEnv(t *testing.T) YDLS {
	ydls, err := NewFromFile(os.Getenv("CONFIG"))
	if err != nil {
		t.Fatalf("failed to read config: %s", err)
	}

	return ydls
}

func TestSafeFilename(t *testing.T) {
	for _, c := range []struct {
		s      string
		expect string
	}{
		{"aba", "aba"},
		{"a/a", "a_a"},
		{"a\\a", "a_a"},
	} {
		actual := safeFilename(c.s)
		if actual != c.expect {
			t.Errorf("%s, got %v expected %v", c.s, actual, c.expect)
		}
	}
}

func TestForceCodec(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)
	forceCodecs := []string{"opus", "vp9"}

	// make sure codecs are not the perferred ones
	for _, s := range mkvFormat.Streams {
		for _, forceCodec := range forceCodecs {
			if c, ok := s.CodecNames.First(); ok && c == forceCodec {
				t.Errorf("test sanity check failed: codec already the preferred one")
				return
			}
		}
	}

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeTestVideoURL,
			Format:      &mkvFormat,
			Codecs:      forceCodecs,
		},
		nil)
	if err != nil {
		cancelFn()
		t.Errorf("%s: download failed: %s", youtubeTestVideoURL, err)
		return
	}

	pi, err := ffmpeg.Probe(ctx, ffmpeg.Reader{Reader: io.LimitReader(dr.Media, 10*1024*1024)}, nil, nil)
	dr.Media.Close()
	dr.Wait()
	cancelFn()
	if err != nil {
		t.Errorf("%s: probe failed: %s", youtubeTestVideoURL, err)
		return
	}

	if pi.FormatName() != "matroska" {
		t.Errorf("%s: force codec failed: found %s", youtubeTestVideoURL, pi)
		return
	}

	for i := 0; i < len(forceCodecs); i++ {
		if pi.Streams[i].CodecName != forceCodecs[i] {
			t.Errorf("%s: force codec failed: %s != %s", youtubeTestVideoURL, pi.Streams[i].CodecName, forceCodecs[i])
			return
		}
	}

}

func TestTimeRangeOption(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	timeRange, timeRangeErr := timerange.NewTimeRangeFromString("10s-15s")
	if timeRangeErr != nil {
		t.Fatalf("failed to parse time range")
	}

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeTestVideoURL,
			Format:      &mkvFormat,
			TimeRange:   timeRange,
		},
		nil)
	if err != nil {
		cancelFn()
		t.Fatalf("%s: download failed: %s", youtubeTestVideoURL, err)
	}

	pi, err := ffmpeg.Probe(ctx, ffmpeg.Reader{Reader: io.LimitReader(dr.Media, 10*1024*1024)}, nil, nil)
	dr.Media.Close()
	dr.Wait()
	cancelFn()
	if err != nil {
		t.Errorf("%s: probe failed: %s", youtubeTestVideoURL, err)
		return
	}

	if pi.Duration() != timeRange.Duration() {
		t.Errorf("%s: probed duration not %v, got %v", youtubeTestVideoURL, timeRange.Duration(), pi.Duration())
		return
	}
}

func TestMissingMediaStream(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	_, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: soundcloudTestAudioURL,
			Format:      &mkvFormat,
		},
		nil)
	cancelFn()
	if err == nil {
		t.Fatal("expected download to fail")
	}
}

func TestFindYDLFormat(t *testing.T) {
	ydlFormats := []youtubedl.Format{
		{FormatID: "1", Protocol: "http", NormACodec: "mp3", NormVCodec: "h264", NormBR: 1},
		{FormatID: "2", Protocol: "http", NormACodec: "", NormVCodec: "h264", NormBR: 2},
		{FormatID: "3", Protocol: "http", NormACodec: "aac", NormVCodec: "", NormBR: 3},
		{FormatID: "4", Protocol: "http", NormACodec: "vorbis", NormVCodec: "vp8", NormBR: 4},
		{FormatID: "5", Protocol: "http", NormACodec: "opus", NormVCodec: "vp9", NormBR: 5},
	}

	for i, c := range []struct {
		ydlFormats       []youtubedl.Format
		mediaType        mediaType
		codecs           stringprioset.Set
		expectedFormatID string
	}{
		{ydlFormats, MediaAudio, stringprioset.New([]string{"mp3"}), "1"},
		{ydlFormats, MediaAudio, stringprioset.New([]string{"aac"}), "3"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"h264"}), "2"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"h264"}), "2"},
		{ydlFormats, MediaAudio, stringprioset.New([]string{"vorbis"}), "4"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"vp8"}), "4"},
		{ydlFormats, MediaAudio, stringprioset.New([]string{"opus"}), "5"},
		{ydlFormats, MediaVideo, stringprioset.New([]string{"vp9"}), "5"},
	} {
		actualFormat, actualFormatFound := findYDLFormat(c.ydlFormats, c.mediaType, c.codecs)
		if actualFormatFound && actualFormat.FormatID != c.expectedFormatID {
			t.Errorf("%d: expected format %s, got %s", i, c.expectedFormatID, actualFormat)
		}
	}
}

func TestContextCloseProbe(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		time.Sleep(2 * time.Second)
		cancelFn()
		wg.Done()
	}()
	_, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeLongTestVideoURL,
			Format:      &mkvFormat,
		},
		nil)
	if err == nil {
		t.Error("expected error while probe")
	}
	cancelFn()
	wg.Wait()
}

func TestContextCloseDownload(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "mkv"
	mkvFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: youtubeLongTestVideoURL,
			Format:      &mkvFormat,
		},
		nil)
	if err != nil {
		t.Error("expected no error while probe")
	}
	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		time.Sleep(1 * time.Second)
		cancelFn()
		wg.Done()
	}()
	io.Copy(ioutil.Discard, dr.Media)
	cancelFn()
	wg.Wait()
}

func TestRSS(t *testing.T) {
	if !testNetwork || !testFfmpeg || !testYoutubeldl {
		t.Skip("TEST_NETWORK, TEST_FFMPEG, TEST_YOUTUBEDL env not set")
	}

	defer leaktest.Check(t)()

	ydls := ydlsFromEnv(t)
	const formatName = "rss"
	rssFormat, _ := ydls.Config.Formats.FindByName(formatName)

	ctx, cancelFn := context.WithCancel(context.Background())

	dr, err := ydls.Download(ctx,
		DownloadOptions{
			MediaRawURL: soundcloudTestPlaylistURL,
			Format:      &rssFormat,
			BaseURL:     &url.URL{Scheme: "http", Host: "dummy"},
		},
		nil)
	if err != nil {
		cancelFn()
		t.Fatalf("%s: download failed: %s", youtubeTestVideoURL, err)
	}
	defer cancelFn()

	if dr.Filename != "" {
		t.Errorf("expected no filename, got %s", dr.Filename)
	}
	expectedMIMEType := "text/xml"
	if dr.MIMEType != expectedMIMEType {
		t.Errorf("expected mimetype %s, got %s", expectedMIMEType, dr.MIMEType)
	}

	rssRoot := rss.RSS{}
	decoder := xml.NewDecoder(dr.Media)
	decodeErr := decoder.Decode(&rssRoot)
	if decodeErr != nil {
		t.Errorf("failed to parse rss: %s", decodeErr)
		return
	}
	dr.Media.Close()
	dr.Wait()

	expectedTitle := "Kindred Phenomena"
	if rssRoot.Channel.Title != expectedTitle {
		t.Errorf("expected title \"%s\" got \"%s\"", expectedTitle, rssRoot.Channel.Title)
	}

	// TODO: description, not returned by youtube-dl?

	expectedEntries := 8
	if len(rssRoot.Channel.Items) != expectedEntries {
		t.Errorf("expected %d entries got %d", expectedEntries, len(rssRoot.Channel.Items))
	}

	itemOne := rssRoot.Channel.Items[0]

	expectedItemTitle := "A1 Mattheis - Herds"
	if rssRoot.Channel.Items[0].Title != expectedItemTitle {
		t.Errorf("expected title \"%s\" got \"%s\"", expectedItemTitle, itemOne.Title)
	}

	expectedItemDescriptionPrefix := "Releasing my debut"
	if !strings.HasPrefix(rssRoot.Channel.Items[0].Description, expectedItemDescriptionPrefix) {
		t.Errorf("expected description prefix \"%s\" got \"%s\"", expectedItemDescriptionPrefix, itemOne.Description)
	}

	expectedItemGUID := "http://dummy/mp3/https://soundcloud.com/mattheis/sets/kindred-phenomena#293285002"
	if rssRoot.Channel.Items[0].GUID != expectedItemGUID {
		t.Errorf("expected guid \"%s\" got \"%s\"", expectedItemGUID, itemOne.GUID)
	}

	expectedItemURL := "http://dummy/media.mp3?format=mp3&url=https%3A%2F%2Fsoundcloud.com%2Fmattheis%2Fa1-mattheis-herds"
	if itemOne.Enclosure.URL != expectedItemURL {
		t.Errorf("expected enclousure url \"%s\" got \"%s\"", expectedItemURL, itemOne.Enclosure.URL)
	}

	expectedItemType := "audio/mpeg"
	if itemOne.Enclosure.Type != expectedItemType {
		t.Errorf("expected enclousure type \"%s\" got \"%s\"", expectedItemType, itemOne.Enclosure.Type)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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.

//go:generate protoc -I ./model/cni --go_out=plugins=grpc:./model/cni ./model/cni/cni.proto

package contiv

import (
	"context"
	"git.fd.io/govpp.git/api"
	"github.com/contiv/vpp/plugins/contiv/containeridx"
	"github.com/contiv/vpp/plugins/contiv/model/cni"
	"github.com/contiv/vpp/plugins/kvdbproxy"
	"github.com/ligato/cn-infra/datasync/resync"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/rpc/grpc"
	"github.com/ligato/cn-infra/utils/safeclose"
	"github.com/ligato/vpp-agent/clientv1/linux"
	"github.com/ligato/vpp-agent/clientv1/linux/localclient"
	"github.com/ligato/vpp-agent/plugins/defaultplugins"
	"github.com/ligato/vpp-agent/plugins/govppmux"
)

// Plugin transforms GRPC requests into configuration for the VPP in order
// to connect a container into the network.
type Plugin struct {
	Deps
	govppCh *api.Channel

	configuredContainers *containeridx.ConfigIndex
	cniServer            *remoteCNIserver

	ctx           context.Context
	ctxCancelFunc context.CancelFunc
}

// Deps groups the dependencies of the Plugin.
type Deps struct {
	local.PluginInfraDeps
	GRPC   grpc.Server
	Proxy  *kvdbproxy.Plugin
	VPP    *defaultplugins.Plugin
	GoVPP  govppmux.API
	Resync resync.Subscriber
}

// Init initializes the grpc server handling the request from the CNI.
func (plugin *Plugin) Init() error {
	plugin.configuredContainers = containeridx.NewConfigIndex(plugin.Log, plugin.PluginName, "containers")

	plugin.ctx, plugin.ctxCancelFunc = context.WithCancel(context.Background())

	var err error
	plugin.govppCh, err = plugin.GoVPP.NewAPIChannel()
	if err != nil {
		return err
	}

	plugin.cniServer = newRemoteCNIServer(plugin.Log,
		func() linux.DataChangeDSL { return localclient.DataChangeRequest(plugin.PluginName) },
		plugin.Proxy,
		plugin.configuredContainers,
		plugin.govppCh,
		plugin.VPP.GetSwIfIndexes(),
		plugin.ServiceLabel.GetAgentLabel())
	cni.RegisterRemoteCNIServer(plugin.GRPC.Server(), plugin.cniServer)
	return nil
}

func (plugin *Plugin) AfterInit() error {
	if plugin.Resync != nil {
		reg := plugin.Resync.Register(string(plugin.PluginName))
		go plugin.handleResync(reg.StatusChan())
	}
	return nil
}

// Close cleans up the resources allocated by the plugin
func (plugin *Plugin) Close() error {
	plugin.ctxCancelFunc()
	plugin.cniServer.close()
	return safeclose.Close(plugin.govppCh)
}

// GetIfName looks up logical interface name that corresponds to the interface associated with the given pod.
func (plugin *Plugin) GetIfName(podNamespace string, podName string) (name string, exists bool) {
	podNamesMatch := plugin.configuredContainers.LookupPodName(podName)
	podNamespacesMatch := plugin.configuredContainers.LookupPodNamespace(podNamespace)

	for _, pod1 := range podNamespacesMatch {
		for _, pod2 := range podNamesMatch {
			if pod1 == pod2 {
				found, data := plugin.configuredContainers.LookupContainer(pod1)
				if found && data != nil && data.Afpacket != nil {
					return data.Afpacket.Name, true
				}
			}
		}
	}

	plugin.Log.WithFields(logging.Fields{"podNamespace": podNamespace, "podName": podName}).Warn("No matching result found")
	return "", false
}

func (plugin *Plugin) handleResync(resyncChan chan resync.StatusEvent) {
	for {
		select {
		case ev := <-resyncChan:
			status := ev.ResyncStatus()
			if status == resync.Started {
				err := plugin.cniServer.resync()
				if err != nil {
					plugin.Log.Error(err)
				}
			}
			ev.Ack()
		case <-plugin.ctx.Done():
			return
		}
	}
}
<commit_msg>Address go lint<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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.

//go:generate protoc -I ./model/cni --go_out=plugins=grpc:./model/cni ./model/cni/cni.proto

package contiv

import (
	"context"
	"git.fd.io/govpp.git/api"
	"github.com/contiv/vpp/plugins/contiv/containeridx"
	"github.com/contiv/vpp/plugins/contiv/model/cni"
	"github.com/contiv/vpp/plugins/kvdbproxy"
	"github.com/ligato/cn-infra/datasync/resync"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/rpc/grpc"
	"github.com/ligato/cn-infra/utils/safeclose"
	"github.com/ligato/vpp-agent/clientv1/linux"
	"github.com/ligato/vpp-agent/clientv1/linux/localclient"
	"github.com/ligato/vpp-agent/plugins/defaultplugins"
	"github.com/ligato/vpp-agent/plugins/govppmux"
)

// Plugin transforms GRPC requests into configuration for the VPP in order
// to connect a container into the network.
type Plugin struct {
	Deps
	govppCh *api.Channel

	configuredContainers *containeridx.ConfigIndex
	cniServer            *remoteCNIserver

	ctx           context.Context
	ctxCancelFunc context.CancelFunc
}

// Deps groups the dependencies of the Plugin.
type Deps struct {
	local.PluginInfraDeps
	GRPC   grpc.Server
	Proxy  *kvdbproxy.Plugin
	VPP    *defaultplugins.Plugin
	GoVPP  govppmux.API
	Resync resync.Subscriber
}

// Init initializes the grpc server handling the request from the CNI.
func (plugin *Plugin) Init() error {
	plugin.configuredContainers = containeridx.NewConfigIndex(plugin.Log, plugin.PluginName, "containers")

	plugin.ctx, plugin.ctxCancelFunc = context.WithCancel(context.Background())

	var err error
	plugin.govppCh, err = plugin.GoVPP.NewAPIChannel()
	if err != nil {
		return err
	}

	plugin.cniServer = newRemoteCNIServer(plugin.Log,
		func() linux.DataChangeDSL { return localclient.DataChangeRequest(plugin.PluginName) },
		plugin.Proxy,
		plugin.configuredContainers,
		plugin.govppCh,
		plugin.VPP.GetSwIfIndexes(),
		plugin.ServiceLabel.GetAgentLabel())
	cni.RegisterRemoteCNIServer(plugin.GRPC.Server(), plugin.cniServer)
	return nil
}

// AfterInit registers to the ResyncOrchestrator. The registration is done in this phase
// in order to trigger the resync for this plugin once the resync of defaultVPP plugins is finished.
func (plugin *Plugin) AfterInit() error {
	if plugin.Resync != nil {
		reg := plugin.Resync.Register(string(plugin.PluginName))
		go plugin.handleResync(reg.StatusChan())
	}
	return nil
}

// Close cleans up the resources allocated by the plugin
func (plugin *Plugin) Close() error {
	plugin.ctxCancelFunc()
	plugin.cniServer.close()
	return safeclose.Close(plugin.govppCh)
}

// GetIfName looks up logical interface name that corresponds to the interface associated with the given pod.
func (plugin *Plugin) GetIfName(podNamespace string, podName string) (name string, exists bool) {
	podNamesMatch := plugin.configuredContainers.LookupPodName(podName)
	podNamespacesMatch := plugin.configuredContainers.LookupPodNamespace(podNamespace)

	for _, pod1 := range podNamespacesMatch {
		for _, pod2 := range podNamesMatch {
			if pod1 == pod2 {
				found, data := plugin.configuredContainers.LookupContainer(pod1)
				if found && data != nil && data.Afpacket != nil {
					return data.Afpacket.Name, true
				}
			}
		}
	}

	plugin.Log.WithFields(logging.Fields{"podNamespace": podNamespace, "podName": podName}).Warn("No matching result found")
	return "", false
}

func (plugin *Plugin) handleResync(resyncChan chan resync.StatusEvent) {
	for {
		select {
		case ev := <-resyncChan:
			status := ev.ResyncStatus()
			if status == resync.Started {
				err := plugin.cniServer.resync()
				if err != nil {
					plugin.Log.Error(err)
				}
			}
			ev.Ack()
		case <-plugin.ctx.Done():
			return
		}
	}
}
<|endoftext|>


<commit_before>package worker

import (
	"encoding/json"
	"os"
	"time"

	"github.com/Shopify/sarama"
	"github.com/Sirupsen/logrus"
	"github.com/gocql/gocql"
	geoip2 "github.com/oschwald/geoip2-golang"
	"github.com/piotrkowalczuk/gonalytics-backend/lib"
	"github.com/piotrkowalczuk/gonalytics-backend/lib/models"
)

const (
	clientName        = "actions-worker-1"
	consumerGroupName = "actions-workers"
	topicName         = "actions"
)

// ActionsWorker ...
type ActionsWorker struct {
	Cassandra         *gocql.Session
	Config            *lib.ActionsWorkerConfig
	KafkaClient       *sarama.Client
	GeoIP             *geoip2.Reader
	Logger            *logrus.Logger
	RepositoryManager lib.RepositoryManager

	BrokerHost string

	startingOffset    int64
	broker            *sarama.Broker
	consumer          *sarama.Consumer
	partitionConsumer *sarama.PartitionConsumer
}

// ConsumedFunc ...
type ConsumedFunc func(trackRequest *models.TrackRequest) error

func (aw *ActionsWorker) panicIf(err error, message string) {
	if err != nil && err != sarama.ErrNoError {
		aw.Logger.Error(message)
		panic(err)
	}
}

func (aw *ActionsWorker) breakIf(err error) {
	if err != nil && err != sarama.ErrNoError {
		aw.Logger.Error(err.Error())
		os.Exit(1)
	}
}

// Start ...
func (aw *ActionsWorker) Start() {
	aw.initBroker()
	defer aw.broker.Close()

	aw.logNumberOfTopics()
	aw.checkIfAwailableOffsets()
	aw.fetchAndSetStartingOffset()

	aw.initConsumer()
	defer aw.partitionConsumer.Close()

	aw.consume(
		aw.saveToCassandra,
	)
}

func (aw *ActionsWorker) initBroker() {
	aw.broker = sarama.NewBroker(aw.BrokerHost)

	err := aw.broker.Open(nil)
	aw.panicIf(err, "Kafka broker initialization failed.")
}

func (aw *ActionsWorker) initConsumer() {
	consumerConfig := &sarama.ConsumerConfig{
		MinFetchSize: 1,
		MaxWaitTime:  250 * time.Millisecond,
	}

	partitionConsumerConfig := &sarama.PartitionConsumerConfig{
		DefaultFetchSize: 32768,
		OffsetMethod:     sarama.OffsetMethodManual,
		OffsetValue:      aw.startingOffset + 1,
	}

	consumer, err := sarama.NewConsumer(
		aw.KafkaClient,
		consumerConfig,
	)

	aw.panicIf(err, "Kafka consumer initialization failed.")

	partitionConsumer, err := consumer.ConsumePartition(
		topicName,
		0,
		partitionConsumerConfig,
	)

	aw.panicIf(err, "Kafka partition consumer initialization failed.")

	aw.Logger.WithFields(logrus.Fields{
		"offset": aw.startingOffset,
	}).Info("Kafka consumer initialized successfully.")

	aw.consumer = consumer
	aw.partitionConsumer = partitionConsumer
}

func (aw *ActionsWorker) logNumberOfTopics() {
	request := sarama.MetadataRequest{Topics: []string{}}
	response, err := aw.broker.GetMetadata(clientName, &request)
	aw.breakIf(err)

	aw.Logger.WithFields(logrus.Fields{
		"nbOfTopics": len(response.Topics),
	}).Info("Kafka metadata fetched successfully.")
}

func (aw *ActionsWorker) checkIfAwailableOffsets() {
	offsetsRequest := sarama.OffsetRequest{}
	offsetsRequest.AddBlock(topicName, 0, sarama.EarliestOffset, 255)
	offsetsResponse, err := aw.broker.GetAvailableOffsets(clientName, &offsetsRequest)

	aw.panicIf(err, "Kafka broker available offsets request failed.")
	aw.breakIf(offsetsResponse.Blocks[topicName][0].Err)

	aw.Logger.WithFields(logrus.Fields{
		"offsets": offsetsResponse.Blocks[topicName][0].Offsets,
	}).Info("Kafka available offsets fetched successfully.")
}

func (aw *ActionsWorker) fetchAndSetStartingOffset() {
	offsetFetchRequest := sarama.OffsetFetchRequest{ConsumerGroup: consumerGroupName}
	offsetFetchRequest.AddPartition(topicName, 0)
	offsetFetchResponse, err := aw.broker.FetchOffset(clientName, &offsetFetchRequest)

	aw.panicIf(err, "Kafka broker fetching current offset request failed.")

	err = offsetFetchResponse.Blocks[topicName][0].Err
	if err != sarama.ErrNoError {
		if err == sarama.ErrUnknownTopicOrPartition {
			offsetCommitRequest := &sarama.OffsetCommitRequest{ConsumerGroup: consumerGroupName}
			offsetCommitRequest.AddBlock(topicName, 0, 0, sarama.ReceiveTime, "")

			response, err := aw.broker.CommitOffset(clientName, offsetCommitRequest)

			if err != nil {
				aw.Logger.Error(err)
				return
			} else if response == nil {
				aw.Logger.Error("Brocker returns no response")
				return
			} else if response.Errors[topicName][0] != sarama.ErrNoError {
				aw.Logger.Error(response.Errors[topicName][0])
				return
			}
		} else {
			aw.breakIf(offsetFetchResponse.Blocks[topicName][0].Err)
		}
	}

	aw.startingOffset = offsetFetchResponse.Blocks[topicName][0].Offset

	aw.Logger.WithFields(logrus.Fields{
		"startingOffset": aw.startingOffset,
	}).Info("Kafka current offsets fetched successfully.")
}

func (aw *ActionsWorker) consume(callbacks ...ConsumedFunc) {
	for {
		select {
		case err := <-aw.partitionConsumer.Errors():
			if err != nil {
				panic(err)
			}
		case event := <-aw.partitionConsumer.Messages():
			trackRequest := &models.TrackRequest{}
			err := json.Unmarshal(event.Value, trackRequest)

			if err != nil {
				aw.Logger.Error(err)
				return
			}

			for _, fn := range callbacks {
				err := fn(trackRequest)
				if err != nil {
					aw.Logger.Error(err)
					return
				}
			}

			aw.Logger.WithFields(logrus.Fields{
				"url":     trackRequest.PageURL,
				"visitId": trackRequest.VisitID,
				"offset":  event.Offset,
			}).Debug("Action has been properly processed.")

			offsetCommitRequest := &sarama.OffsetCommitRequest{ConsumerGroup: consumerGroupName}
			offsetCommitRequest.AddBlock(topicName, 0, event.Offset, sarama.ReceiveTime, "")

			response, err := aw.broker.CommitOffset(clientName, offsetCommitRequest)
			if err != nil {
				aw.Logger.Error(err)
				return
			} else if response == nil {
				aw.Logger.Error("Brocker returns no response")
				return
			} else if response.Errors[topicName][0] != sarama.ErrNoError {
				aw.Logger.Error(response.Errors[topicName][0])
				return
			}

			aw.Logger.WithFields(logrus.Fields{
				"client":        clientName,
				"consumerGroup": consumerGroupName,
				"topic":         topicName,
				"offset":        event.Offset,
			}).Debug("New offset has been successfully submited to brocker.")
		}
	}
}

func (aw *ActionsWorker) saveToCassandra(trackRequest *models.TrackRequest) error {
	actionCreator := lib.NewActionCreator(aw.GeoIP)

	action, err := actionCreator.Create(trackRequest)
	if err != nil {
		return err
	}

	return aw.RepositoryManager.VisitActions.Insert(action)
}
<commit_msg>Sarama integration fix#2<commit_after>package worker

import (
	"encoding/json"
	"os"
	"time"

	"github.com/Shopify/sarama"
	"github.com/Sirupsen/logrus"
	"github.com/gocql/gocql"
	geoip2 "github.com/oschwald/geoip2-golang"
	"github.com/piotrkowalczuk/gonalytics-backend/lib"
	"github.com/piotrkowalczuk/gonalytics-backend/lib/models"
)

const (
	clientName        = "actions-worker-1"
	consumerGroupName = "actions-workers"
	topicName         = "actions"
)

// ActionsWorker ...
type ActionsWorker struct {
	Cassandra         *gocql.Session
	Config            *lib.ActionsWorkerConfig
	KafkaClient       *sarama.Client
	GeoIP             *geoip2.Reader
	Logger            *logrus.Logger
	RepositoryManager lib.RepositoryManager

	BrokerHost string

	startingOffset    int64
	broker            *sarama.Broker
	consumer          *sarama.Consumer
	partitionConsumer *sarama.PartitionConsumer
}

// ConsumedFunc ...
type ConsumedFunc func(trackRequest *models.TrackRequest) error

func (aw *ActionsWorker) panicIf(err error, message string) {
	if err != nil && err != sarama.ErrNoError {
		aw.Logger.Error(message)
		panic(err)
	}
}

func (aw *ActionsWorker) breakIf(err error) {
	if err != nil && err != sarama.ErrNoError {
		aw.Logger.Error(err.Error())
		os.Exit(1)
	}
}

// Start ...
func (aw *ActionsWorker) Start() {
	aw.initBroker()
	defer aw.broker.Close()

	aw.logNumberOfTopics()
	aw.checkIfAwailableOffsets()
	aw.fetchAndSetStartingOffset()

	aw.initConsumer()
	defer aw.partitionConsumer.Close()

	aw.consume(
		aw.saveToCassandra,
	)
}

func (aw *ActionsWorker) initBroker() {
	aw.broker = sarama.NewBroker(aw.BrokerHost)

	err := aw.broker.Open(nil)
	aw.panicIf(err, "Kafka broker initialization failed.")
}

func (aw *ActionsWorker) initConsumer() {
	consumerConfig := &sarama.ConsumerConfig{
		MinFetchSize: 1,
		MaxWaitTime:  250 * time.Millisecond,
	}

	partitionConsumerConfig := &sarama.PartitionConsumerConfig{
		DefaultFetchSize: 32768,
		OffsetMethod:     sarama.OffsetMethodManual,
		OffsetValue:      aw.startingOffset + 1,
	}

	consumer, err := sarama.NewConsumer(
		aw.KafkaClient,
		consumerConfig,
	)

	aw.panicIf(err, "Kafka consumer initialization failed.")

	partitionConsumer, err := consumer.ConsumePartition(
		topicName,
		0,
		partitionConsumerConfig,
	)

	aw.panicIf(err, "Kafka partition consumer initialization failed.")

	aw.Logger.WithFields(logrus.Fields{
		"offset": aw.startingOffset,
	}).Info("Kafka consumer initialized successfully.")

	aw.consumer = consumer
	aw.partitionConsumer = partitionConsumer
}

func (aw *ActionsWorker) logNumberOfTopics() {
	request := sarama.MetadataRequest{Topics: []string{}}
	response, err := aw.broker.GetMetadata(clientName, &request)
	aw.breakIf(err)

	aw.Logger.WithFields(logrus.Fields{
		"nbOfTopics": len(response.Topics),
	}).Info("Kafka metadata fetched successfully.")
}

func (aw *ActionsWorker) checkIfAwailableOffsets() {
	offsetsRequest := sarama.OffsetRequest{}
	offsetsRequest.AddBlock(topicName, 0, sarama.EarliestOffset, 255)
	offsetsResponse, err := aw.broker.GetAvailableOffsets(clientName, &offsetsRequest)

	aw.panicIf(err, "Kafka broker available offsets request failed.")
	aw.breakIf(offsetsResponse.Blocks[topicName][0].Err)

	aw.Logger.WithFields(logrus.Fields{
		"offsets": offsetsResponse.Blocks[topicName][0].Offsets,
	}).Info("Kafka available offsets fetched successfully.")
}

func (aw *ActionsWorker) fetchAndSetStartingOffset() {
	offsetFetchRequest := sarama.OffsetFetchRequest{ConsumerGroup: consumerGroupName}
	offsetFetchRequest.AddPartition(topicName, 0)
	offsetFetchResponse, err := aw.broker.FetchOffset(clientName, &offsetFetchRequest)

	aw.panicIf(err, "Kafka broker fetching current offset request failed.")

	err = offsetFetchResponse.Blocks[topicName][0].Err
	if err != sarama.ErrNoError {
		if err == sarama.ErrUnknownTopicOrPartition {
			offsetCommitRequest := &sarama.OffsetCommitRequest{ConsumerGroup: consumerGroupName}
			offsetCommitRequest.AddBlock(topicName, 0, 0, sarama.ReceiveTime, "")

			response, err := aw.broker.CommitOffset(clientName, offsetCommitRequest)

			if err != nil {
				aw.Logger.Error(err)
				return
			} else if response == nil {
				aw.Logger.Error("Brocker returns no response")
				return
			} else if response.Errors[topicName][0] != sarama.ErrNoError {
				aw.Logger.Error(response.Errors[topicName][0])
				return
			}

			aw.startingOffset = 0
		} else {
			aw.breakIf(offsetFetchResponse.Blocks[topicName][0].Err)
		}
	} else {
		aw.startingOffset = offsetFetchResponse.Blocks[topicName][0].Offset
	}

	aw.Logger.WithFields(logrus.Fields{
		"startingOffset": aw.startingOffset,
	}).Info("Kafka current offsets fetched successfully.")
}

func (aw *ActionsWorker) consume(callbacks ...ConsumedFunc) {
	for {
		select {
		case err := <-aw.partitionConsumer.Errors():
			if err != nil {
				panic(err)
			}
		case event := <-aw.partitionConsumer.Messages():
			trackRequest := &models.TrackRequest{}
			err := json.Unmarshal(event.Value, trackRequest)

			if err != nil {
				aw.Logger.Error(err)
				return
			}

			for _, fn := range callbacks {
				err := fn(trackRequest)
				if err != nil {
					aw.Logger.Error(err)
					return
				}
			}

			aw.Logger.WithFields(logrus.Fields{
				"url":     trackRequest.PageURL,
				"visitId": trackRequest.VisitID,
				"offset":  event.Offset,
			}).Debug("Action has been properly processed.")

			offsetCommitRequest := &sarama.OffsetCommitRequest{ConsumerGroup: consumerGroupName}
			offsetCommitRequest.AddBlock(topicName, 0, event.Offset, sarama.ReceiveTime, "")

			response, err := aw.broker.CommitOffset(clientName, offsetCommitRequest)
			if err != nil {
				aw.Logger.Error(err)
				return
			} else if response == nil {
				aw.Logger.Error("Brocker returns no response")
				return
			} else if response.Errors[topicName][0] != sarama.ErrNoError {
				aw.Logger.Error(response.Errors[topicName][0])
				return
			}

			aw.Logger.WithFields(logrus.Fields{
				"client":        clientName,
				"consumerGroup": consumerGroupName,
				"topic":         topicName,
				"offset":        event.Offset,
			}).Debug("New offset has been successfully submited to brocker.")
		}
	}
}

func (aw *ActionsWorker) saveToCassandra(trackRequest *models.TrackRequest) error {
	actionCreator := lib.NewActionCreator(aw.GeoIP)

	action, err := actionCreator.Create(trackRequest)
	if err != nil {
		return err
	}

	return aw.RepositoryManager.VisitActions.Insert(action)
}
<|endoftext|>


<commit_before>package main

import (
	"sync"
	"sync/atomic"
)

var wg sync.WaitGroup

type CallFrame struct {
	parent    *CallFrame
	var_table map[string]Object
	stack     []Object
	me        Object
}

type Transaction struct {
	startInst        Instruction
	instList         []Instruction
	objectSet        map[*RObject]*RObject
	transactionStack []*CallFrame
	env              *ThreadEnv
	inevitable       bool
	rev              int64
}

type ThreadEnv struct {
	instList      []Instruction
	threadStack   []*CallFrame
	transactionPC *Transaction
	id            int
}

type GobiesVM struct {
	instList       []Instruction
	callFrameStack []*CallFrame
	consts         map[string]*RObject
	symbols        map[string]int
	rev            int64
	globalLock     sync.RWMutex
	has_inevitable int32
}

func initVM() *GobiesVM {
	VM := &GobiesVM{}
	top := initCallFrame()
	VM.callFrameStack = append(VM.callFrameStack, top)
	VM.consts = make(map[string]*RObject)
	VM.initConsts()
	VM.symbols = make(map[string]int)
	top.me = initRKernel()
	return VM
}

func initCallFrame() *CallFrame {
	frame := &CallFrame{}
	frame.var_table = make(map[string]Object)
	return frame
}

func (VM *GobiesVM) initConsts() {
	VM.consts["RString"] = initRString()
	VM.consts["RFixnum"] = initRFixnum()
	VM.consts["RArray"] = initRArray()
	VM.consts["RHash"] = initRHash()
	VM.consts["IO"] = initRIO()
	VM.consts["Thread"] = initRThread()
}

func (obj *RObject) methodLookup(method_name string) *RMethod {
	if val, ok := obj.methods[method_name]; ok {
		return val
	}
	if obj.class != nil {
		return obj.class.methodLookup(method_name)
	}
	return nil
}

func (currentCallFrame *CallFrame) variableLookup(var_name string) Object {
	if obj, ok := currentCallFrame.var_table[var_name]; ok {
		return obj
	}
	if currentCallFrame.parent != nil {
		return currentCallFrame.parent.variableLookup(var_name)
	}
	return nil
}

func (t *Transaction) initTransaction(env *ThreadEnv, instList []Instruction) *Transaction {
	t.instList = []Instruction{}
	for _, inst := range instList {
		t.instList = append(t.instList, inst)
	}
	t.transactionStack = copyFrames(env.threadStack)
	t.objectSet = make(map[*RObject]*RObject)
	return t
}

func copyFrames(src []*CallFrame) []*CallFrame {
	newStack := []*CallFrame{}
	for _, frame := range src {
		newFrame := initCallFrame()
		newFrame.me = frame.me
		newFrame.parent = frame.parent

		// Copy every object "pointer" in old frame stack
		for _, obj := range frame.stack {
			newFrame.stack = append(newFrame.stack, obj)
		}

		// Copy every object "pointer" in instance variable table
		for key, obj := range frame.var_table {
			newFrame.var_table[key] = obj
		}

		newStack = append(newStack, newFrame)
	}
	return newStack
}

func (obj *RObject) copyObject(src *RObject) {
	// Copy instance variables
	ivars := src.ivars
	for k, v := range ivars {
		obj.ivars[k] = v
	}
}

func addRObjectToSet(obj *RObject, env *ThreadEnv) *RObject {
	if env == nil { // created during complication
		return obj
	}
	if _, ok := env.transactionPC.objectSet[obj]; !ok {
		env.transactionPC.objectSet[obj] = obj
	}
	return obj
}

func (VM *GobiesVM) execute() {
	// 	root := initTransaction(VM.instList)
	// 	root.inevitable = true

	VM.rev = 0

	// Execute root transaction which is inevitable
	VM.executeThread(VM.instList, nil)
	wg.Wait()
}

func (VM *GobiesVM) executeThread(instList []Instruction, parentScope *ThreadEnv) {
	// Create clean call frame without pushing back to VM stack
	currentCallFrame := initCallFrame()

	env := &ThreadEnv{instList: instList}
	if parentScope == nil {
		currentCallFrame.parent = VM.callFrameStack[len(VM.callFrameStack)-1]
	} else {
		env.threadStack = copyFrames(parentScope.threadStack)
		currentCallFrame.parent = env.threadStack[len(env.threadStack)-1]
		// currentCallFrame.parent = parentScope.threadStack[len(parentScope.threadStack)-1]
	}
	currentCallFrame.me = currentCallFrame.parent.me
	env.threadStack = append(env.threadStack, currentCallFrame)

	// t.inevitable = true
	VM.executeBytecodes(nil, env)
	if parentScope != nil {
		wg.Done()
	}
}

func (VM *GobiesVM) transactionBegin(env *ThreadEnv, inst []Instruction) *Transaction {
	t := &Transaction{}
	t.initTransaction(env, inst)
	t.rev = atomic.LoadInt64(&VM.rev)
	env.transactionPC = t
	t.env = env

	// Initialize environment
	if t.inevitable {
		VM.globalLock.Lock()
		atomic.StoreInt32(&VM.has_inevitable, 1)
	} else {
		VM.globalLock.RLock()
	}

	return t
}

func (VM *GobiesVM) transactionEnd(env *ThreadEnv) bool {
	t := env.transactionPC

	// Lock the write-set
	locked := []*RObject{}
	for orig_obj, new_obj := range t.objectSet {
		// fmt.Println(orig_obj, new_obj)
		if orig_obj != new_obj {
			if orig_obj.writeLock.TryLock() {
				locked = append(locked, orig_obj)
			} else { // Attempt to acquire lock failed
				// Release all locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY

			}
		}
	}

	// Increment global revision
	for !atomic.CompareAndSwapInt64(&VM.rev, VM.rev, VM.rev+1) {
	}

	// Validate the read-set
	for orig_obj, new_obj := range t.objectSet {
		if orig_obj == new_obj {
			if orig_obj.rev > t.rev {
				// Release write locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY
			}
		}
	}

	// Commit then release the locks
	for _, locked_obj := range locked {
		src := t.objectSet[locked_obj]
		locked_obj.copyObject(src)
	}
	for _, locked_obj := range locked {
		locked_obj.writeLock.Unlock()
	}

	env.threadStack = copyFrames(t.transactionStack)

	if t.inevitable {
		VM.globalLock.Unlock()
		atomic.StoreInt32(&VM.has_inevitable, 0)
	} else {
		VM.globalLock.RUnlock()
	}

	env.transactionPC = nil
	return true

TRANSACTION_RETRY:
	t.initTransaction(env, t.instList)
	t.rev = atomic.LoadInt64(&VM.rev)
	return false
}

func (VM *GobiesVM) executeBytecodes(instList []Instruction, env *ThreadEnv) {
	t := env.transactionPC

	if instList == nil {
		instList = env.instList
	}

	// SPECULATIVE_EXEC:
	// Speculative execution
	for i, v := range instList {
		t = env.transactionPC
		if t == nil {
			t = VM.transactionBegin(env, instList[i:])
		}
		currentCallFrame := t.transactionStack[len(t.transactionStack)-1]

		switch v.inst_type {
		case BC_PUTSELF:
			currentCallFrame.stack = append(currentCallFrame.stack, currentCallFrame.me)
		case BC_PUTNIL:
			currentCallFrame.stack = append(currentCallFrame.stack, nil)
		case BC_PUTOBJ:
			currentCallFrame.stack = append(currentCallFrame.stack, v.obj)
		case BC_PUTFIXNUM:
			currentCallFrame.stack = append(currentCallFrame.stack, RFixnum_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTSTRING:
			currentCallFrame.stack = append(currentCallFrame.stack, RString_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTTRUE:
		case BC_PUTFALSE:
		case BC_SETLOCAL:
			top := currentCallFrame.stack[len(currentCallFrame.stack)-1]
			currentCallFrame.var_table[v.obj.(string)] = top
			currentCallFrame.stack = currentCallFrame.stack[0 : len(currentCallFrame.stack)-1] // Pop object from stack
		case BC_GETLOCAL:
			obj := currentCallFrame.variableLookup(v.obj.(string)).(*RObject)
			if _, ok := t.objectSet[obj]; ok {
				obj = t.objectSet[obj]
			}
			currentCallFrame.stack = append(currentCallFrame.stack, obj)
			if VM.transactionEnd(env) == false {

			}
		case BC_SETGLOBAL:
		case BC_GETGLOBAL:
		case BC_SETSYMBOL:
		case BC_GETSYMBOL:
		case BC_SETCONST:
		case BC_GETCONST:
			currentCallFrame.stack = append(currentCallFrame.stack, VM.consts[v.obj.(string)])
		case BC_SETIVAR:
		case BC_GETIVAR:
		case BC_SETCVAR:
		case BC_GETCVAR:
		case BC_SEND:
			argLists := currentCallFrame.stack[len(currentCallFrame.stack)-(v.argc+1):] // argc + 1 ensures inclusion of receiver
			currentCallFrame.stack = currentCallFrame.stack[:len(currentCallFrame.stack)-(v.argc+1)]
			recv := argLists[0].(*RObject)
			argLists = argLists[1:]
			return_val := recv.methodLookup(v.obj.(string)).gofunc(VM, env, recv, argLists)
			// fmt.Println(env.transactionPC)
			// Update address since some functions might init new transaction
			currentCallFrame = env.transactionPC.transactionStack[len(env.transactionPC.transactionStack)-1]
			if return_val != nil {
				currentCallFrame.stack = append(currentCallFrame.stack, return_val)
			}
			if VM.transactionEnd(env) == false {

			}
		case BC_JUMP:

		}
	}
	// End transaction if any
	if env.transactionPC != nil {
		VM.transactionEnd(env)
	}
}

func (VM *GobiesVM) executeBlock(env *ThreadEnv, block *RObject, args []*RObject) {
	if env.transactionPC != nil { // Before
		VM.transactionEnd(env)
	}
	t := VM.transactionBegin(env, block.methods["def"].def)

	// Create clean call frame
	currentCallFrame := initCallFrame()
	currentCallFrame.parent = t.transactionStack[len(t.transactionStack)-1]
	currentCallFrame.me = currentCallFrame.parent.me
	t.transactionStack = append(t.transactionStack, currentCallFrame)

	// Fill in arguments to current call frame
	if block.ivars["params"] != nil {
		params := block.ivars["params"].(*RObject).ivars["array"].([]*RObject)
		for i, v := range params {
			var_name := v.val.str
			currentCallFrame.var_table[var_name] = args[i]
		}
	}

	// Execute block definition
	VM.executeBytecodes(block.methods["def"].def, env)

	VM.transactionBegin(env, []Instruction{})
}
<commit_msg>Define `Join` constant to simulate Thread.join<commit_after>package main

import (
	"sync"
	"sync/atomic"
)

var wg sync.WaitGroup

type CallFrame struct {
	parent    *CallFrame
	var_table map[string]Object
	stack     []Object
	me        Object
}

type Transaction struct {
	startInst        Instruction
	instList         []Instruction
	objectSet        map[*RObject]*RObject
	transactionStack []*CallFrame
	env              *ThreadEnv
	inevitable       bool
	rev              int64
}

type ThreadEnv struct {
	instList      []Instruction
	threadStack   []*CallFrame
	transactionPC *Transaction
	id            int
}

type GobiesVM struct {
	instList       []Instruction
	callFrameStack []*CallFrame
	consts         map[string]*RObject
	symbols        map[string]int
	rev            int64
	globalLock     sync.RWMutex
	has_inevitable int32
}

func initVM() *GobiesVM {
	VM := &GobiesVM{}
	top := initCallFrame()
	VM.callFrameStack = append(VM.callFrameStack, top)
	VM.consts = make(map[string]*RObject)
	VM.initConsts()
	VM.symbols = make(map[string]int)
	top.me = initRKernel()
	return VM
}

func initCallFrame() *CallFrame {
	frame := &CallFrame{}
	frame.var_table = make(map[string]Object)
	return frame
}

func (VM *GobiesVM) initConsts() {
	VM.consts["RString"] = initRString()
	VM.consts["RFixnum"] = initRFixnum()
	VM.consts["RArray"] = initRArray()
	VM.consts["RHash"] = initRHash()
	VM.consts["IO"] = initRIO()
	VM.consts["Thread"] = initRThread()
}

func (obj *RObject) methodLookup(method_name string) *RMethod {
	if val, ok := obj.methods[method_name]; ok {
		return val
	}
	if obj.class != nil {
		return obj.class.methodLookup(method_name)
	}
	return nil
}

func (currentCallFrame *CallFrame) variableLookup(var_name string) Object {
	if obj, ok := currentCallFrame.var_table[var_name]; ok {
		return obj
	}
	if currentCallFrame.parent != nil {
		return currentCallFrame.parent.variableLookup(var_name)
	}
	return nil
}

func (t *Transaction) initTransaction(env *ThreadEnv, instList []Instruction) *Transaction {
	t.instList = []Instruction{}
	for _, inst := range instList {
		t.instList = append(t.instList, inst)
	}
	t.transactionStack = copyFrames(env.threadStack)
	t.objectSet = make(map[*RObject]*RObject)
	return t
}

func copyFrames(src []*CallFrame) []*CallFrame {
	newStack := []*CallFrame{}
	for _, frame := range src {
		newFrame := initCallFrame()
		newFrame.me = frame.me
		newFrame.parent = frame.parent

		// Copy every object "pointer" in old frame stack
		for _, obj := range frame.stack {
			newFrame.stack = append(newFrame.stack, obj)
		}

		// Copy every object "pointer" in instance variable table
		for key, obj := range frame.var_table {
			newFrame.var_table[key] = obj
		}

		newStack = append(newStack, newFrame)
	}
	return newStack
}

func (obj *RObject) copyObject(src *RObject) {
	// Copy instance variables
	ivars := src.ivars
	for k, v := range ivars {
		obj.ivars[k] = v
	}
}

func addRObjectToSet(obj *RObject, env *ThreadEnv) *RObject {
	if env == nil { // created during complication
		return obj
	}
	if _, ok := env.transactionPC.objectSet[obj]; !ok {
		env.transactionPC.objectSet[obj] = obj
	}
	return obj
}

func (VM *GobiesVM) execute() {
	// 	root := initTransaction(VM.instList)
	// 	root.inevitable = true

	VM.rev = 0

	// Execute root transaction which is inevitable
	VM.executeThread(VM.instList, nil)
	wg.Wait()
}

func (VM *GobiesVM) executeThread(instList []Instruction, parentScope *ThreadEnv) {
	// Create clean call frame without pushing back to VM stack
	currentCallFrame := initCallFrame()

	env := &ThreadEnv{instList: instList}
	if parentScope == nil {
		currentCallFrame.parent = VM.callFrameStack[len(VM.callFrameStack)-1]
	} else {
		env.threadStack = copyFrames(parentScope.threadStack)
		currentCallFrame.parent = env.threadStack[len(env.threadStack)-1]
		// currentCallFrame.parent = parentScope.threadStack[len(parentScope.threadStack)-1]
	}
	currentCallFrame.me = currentCallFrame.parent.me
	env.threadStack = append(env.threadStack, currentCallFrame)

	// t.inevitable = true
	VM.executeBytecodes(nil, env)
	if parentScope != nil {
		wg.Done()
	}
}

func (VM *GobiesVM) transactionBegin(env *ThreadEnv, inst []Instruction) *Transaction {
	t := &Transaction{}
	t.initTransaction(env, inst)
	t.rev = atomic.LoadInt64(&VM.rev)
	env.transactionPC = t
	t.env = env

	// Initialize environment
	if t.inevitable {
		VM.globalLock.Lock()
		atomic.StoreInt32(&VM.has_inevitable, 1)
	} else {
		VM.globalLock.RLock()
	}

	return t
}

func (VM *GobiesVM) transactionEnd(env *ThreadEnv) bool {
	t := env.transactionPC

	// Lock the write-set
	locked := []*RObject{}
	for orig_obj, new_obj := range t.objectSet {
		// fmt.Println(orig_obj, new_obj)
		if orig_obj != new_obj {
			if orig_obj.writeLock.TryLock() {
				locked = append(locked, orig_obj)
			} else { // Attempt to acquire lock failed
				// Release all locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY

			}
		}
	}

	// Increment global revision
	for !atomic.CompareAndSwapInt64(&VM.rev, VM.rev, VM.rev+1) {
	}

	// Validate the read-set
	for orig_obj, new_obj := range t.objectSet {
		if orig_obj == new_obj {
			if orig_obj.rev > t.rev {
				// Release write locks
				for _, locked_obj := range locked {
					locked_obj.writeLock.Unlock()
				}
				// Retry
				goto TRANSACTION_RETRY
			}
		}
	}

	// Commit then release the locks
	for _, locked_obj := range locked {
		src := t.objectSet[locked_obj]
		locked_obj.copyObject(src)
	}
	for _, locked_obj := range locked {
		locked_obj.writeLock.Unlock()
	}

	env.threadStack = copyFrames(t.transactionStack)

	if t.inevitable {
		VM.globalLock.Unlock()
		atomic.StoreInt32(&VM.has_inevitable, 0)
	} else {
		VM.globalLock.RUnlock()
	}

	env.transactionPC = nil
	return true

TRANSACTION_RETRY:
	t.initTransaction(env, t.instList)
	t.rev = atomic.LoadInt64(&VM.rev)
	return false
}

func (VM *GobiesVM) executeBytecodes(instList []Instruction, env *ThreadEnv) {
	t := env.transactionPC

	if instList == nil {
		instList = env.instList
	}

	// SPECULATIVE_EXEC:
	// Speculative execution
	for i, v := range instList {
		t = env.transactionPC
		if t == nil {
			t = VM.transactionBegin(env, instList[i:])
		}
		currentCallFrame := t.transactionStack[len(t.transactionStack)-1]

		switch v.inst_type {
		case BC_PUTSELF:
			currentCallFrame.stack = append(currentCallFrame.stack, currentCallFrame.me)
		case BC_PUTNIL:
			currentCallFrame.stack = append(currentCallFrame.stack, nil)
		case BC_PUTOBJ:
			currentCallFrame.stack = append(currentCallFrame.stack, v.obj)
		case BC_PUTFIXNUM:
			currentCallFrame.stack = append(currentCallFrame.stack, RFixnum_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTSTRING:
			currentCallFrame.stack = append(currentCallFrame.stack, RString_new(VM, nil, nil, v.obj.([]Object)))
		case BC_PUTTRUE:
		case BC_PUTFALSE:
		case BC_SETLOCAL:
			top := currentCallFrame.stack[len(currentCallFrame.stack)-1]
			currentCallFrame.var_table[v.obj.(string)] = top
			currentCallFrame.stack = currentCallFrame.stack[0 : len(currentCallFrame.stack)-1] // Pop object from stack
		case BC_GETLOCAL:
			obj := currentCallFrame.variableLookup(v.obj.(string)).(*RObject)
			if _, ok := t.objectSet[obj]; ok {
				obj = t.objectSet[obj]
			}
			currentCallFrame.stack = append(currentCallFrame.stack, obj)
			if VM.transactionEnd(env) == false {

			}
		case BC_SETGLOBAL:
		case BC_GETGLOBAL:
		case BC_SETSYMBOL:
		case BC_GETSYMBOL:
		case BC_SETCONST:
		case BC_GETCONST:
			if v.obj.(string) == "Join" {
				wg.Wait()
			} else {
				currentCallFrame.stack = append(currentCallFrame.stack, VM.consts[v.obj.(string)])
			}
		case BC_SETIVAR:
		case BC_GETIVAR:
		case BC_SETCVAR:
		case BC_GETCVAR:
		case BC_SEND:
			argLists := currentCallFrame.stack[len(currentCallFrame.stack)-(v.argc+1):] // argc + 1 ensures inclusion of receiver
			currentCallFrame.stack = currentCallFrame.stack[:len(currentCallFrame.stack)-(v.argc+1)]
			recv := argLists[0].(*RObject)
			argLists = argLists[1:]
			return_val := recv.methodLookup(v.obj.(string)).gofunc(VM, env, recv, argLists)
			// fmt.Println(env.transactionPC)
			// Update address since some functions might init new transaction
			currentCallFrame = env.transactionPC.transactionStack[len(env.transactionPC.transactionStack)-1]
			if return_val != nil {
				currentCallFrame.stack = append(currentCallFrame.stack, return_val)
			}
			if VM.transactionEnd(env) == false {

			}
		case BC_JUMP:

		}
	}
	// End transaction if any
	if env.transactionPC != nil {
		VM.transactionEnd(env)
	}
}

func (VM *GobiesVM) executeBlock(env *ThreadEnv, block *RObject, args []*RObject) {
	if env.transactionPC != nil { // Before
		VM.transactionEnd(env)
	}
	t := VM.transactionBegin(env, block.methods["def"].def)

	// Create clean call frame
	currentCallFrame := initCallFrame()
	currentCallFrame.parent = t.transactionStack[len(t.transactionStack)-1]
	currentCallFrame.me = currentCallFrame.parent.me
	t.transactionStack = append(t.transactionStack, currentCallFrame)

	// Fill in arguments to current call frame
	if block.ivars["params"] != nil {
		params := block.ivars["params"].(*RObject).ivars["array"].([]*RObject)
		for i, v := range params {
			var_name := v.val.str
			currentCallFrame.var_table[var_name] = args[i]
		}
	}

	// Execute block definition
	VM.executeBytecodes(block.methods["def"].def, env)

	VM.transactionBegin(env, []Instruction{})
}
<|endoftext|>


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

package rpmutils

const (
	RPM_NULL_TYPE         = 0
	RPM_CHAR_TYPE         = 1
	RPM_INT8_TYPE         = 2
	RPM_INT16_TYPE        = 3
	RPM_INT32_TYPE        = 4
	RPM_INT64_TYPE        = 5
	RPM_STRING_TYPE       = 6
	RPM_BIN_TYPE          = 7
	RPM_STRING_ARRAY_TYPE = 8
	RPM_I18NSTRING_TYPE   = 9

	_GENERAL_TAG_BASE = 1000
	NAME              = 1000
	VERSION           = 1001
	RELEASE           = 1002
	EPOCH             = 1003
	SUMMARY           = 1004
	DESCRIPTION       = 1005
	BUILDTIME         = 1006
	BUILDHOST         = 1007
	SIZE              = 1009
	DISTRIBUTION      = 1010
	VENDOR            = 1011
	GIF               = 1012
	XPM               = 1013
	LICENSE           = 1014
	PACKAGER          = 1015
	GROUP             = 1016
	CHANGELOG         = 1017
	SOURCE            = 1018
	PATCH             = 1019
	URL               = 1020
	OS                = 1021
	ARCH              = 1022
	PREIN             = 1023
	POSTIN            = 1024
	PREUN             = 1025
	POSTUN            = 1026
	OLDFILENAMES      = 1027
	FILESIZES         = 1028
	FILEMODES         = 1030
	FILERDEVS         = 1033
	FILEMTIMES        = 1034
	FILEDIGESTS       = 1035 // AKA FILEMD5S
	FILELINKTOS       = 1036
	FILEFLAGS         = 1037 // bitmask: RPMFILE_* are bitmasks to interpret
	FILEUSERNAME      = 1039
	FILEGROUPNAME     = 1040
	ICON              = 1043
	SOURCERPM         = 1044
	FILEVERIFYFLAGS   = 1045 // bitmask: RPMVERIFY_* are bitmasks to interpret
	ARCHIVESIZE       = 1046
	PROVIDENAME       = 1047
	REQUIREFLAGS      = 1048
	REQUIRENAME       = 1049
	REQUIREVERSION    = 1050
	RPMVERSION        = 1064
	TRIGGERSCRIPTS    = 1065
	TRIGGERNAME       = 1066
	TRIGGERVERSION    = 1067
	TRIGGERFLAGS      = 1068 // bitmask: RPMSENSE_* are bitmasks to interpret
	TRIGGERINDEX      = 1069
	VERIFYSCRIPT      = 1079
	PREINPROG         = 1085
	POSTINPROG        = 1086
	PREUNPROG         = 1087
	POSTUNPROG        = 1088
	OBSOLETENAME      = 1090
	FILEDEVICES       = 1095
	FILEINODES        = 1096
	PROVIDEFLAGS      = 1112
	PROVIDEVERSION    = 1113
	OBSOLETEFLAGS     = 1114
	OBSOLETEVERSION   = 1115
	VERIFYSCRIPTPROG  = 1091
	TRIGGERSCRIPTPROG = 1092
	DIRINDEXES        = 1116
	BASENAMES         = 1117
	DIRNAMES          = 1118
	PAYLOADFORMAT     = 1124
	PAYLOADCOMPRESSOR = 1125
	FILECOLORS        = 1140
	// BLINK*, FLINK*, and TRIGGERPREIN included from SUSE fork of RPM
	BLINKPKGID   = 1164
	BLINKHDRID   = 1165
	BLINKNEVRA   = 1166
	FLINKPKGID   = 1167
	FLINKHDRID   = 1168
	FLINKNEVRA   = 1169
	TRIGGERPREIN = 1170

	LONGFILESIZES  = 5008
	LONGSIZE       = 5009
	FILECAPS       = 5010
	FILEDIGESTALGO = 5011
	BUGURL         = 5012
	VCS            = 5034
	ENCODING       = 5062

	SIG_BASE            = 256
	SIG_DSA             = SIG_BASE + 11 // DSA signature over header only
	SIG_RSA             = SIG_BASE + 12 // RSA signature over header only
	SIG_SHA1            = SIG_BASE + 13 // SHA1 over header only (hex)
	SIG_LONGSIZE        = SIG_BASE + 15 // header + compressed payload (uint64)
	SIG_LONGARCHIVESIZE = SIG_BASE + 15 // uncompressed payload bytes (uint64)

	// Given that there is overlap between signature tag headers and general tag
	// headers, we offset the signature ones by some amount
	_SIGHEADER_TAG_BASE = 16384
	SIG_SIZE            = _SIGHEADER_TAG_BASE + 1000 // Header + Payload size
	SIG_PGP             = _SIGHEADER_TAG_BASE + 1002 // Signature over header + payload
	SIG_MD5             = _SIGHEADER_TAG_BASE + 1004 // MD5SUM of header + payload
	SIG_GPG             = _SIGHEADER_TAG_BASE + 1005 // (same as SIG_PGP)
	SIG_PAYLOADSIZE     = _SIGHEADER_TAG_BASE + 1007 // uncompressed payload bytes (uint32)
	SIG_RESERVEDSPACE   = _SIGHEADER_TAG_BASE + 1008 // blank space that can be replaced by a signature

	// FILEFLAGS bitmask elements:
	RPMFILE_NONE      = 0
	RPMFILE_CONFIG    = 1 << 0
	RPMFILE_DOC       = 1 << 1
	RPMFILE_ICON      = 1 << 2
	RPMFILE_MISSINGOK = 1 << 3
	RPMFILE_NOREPLACE = 1 << 4
	RPMFILE_SPECFILE  = 1 << 5
	RPMFILE_GHOST     = 1 << 6
	RPMFILE_LICENSE   = 1 << 7
	RPMFILE_README    = 1 << 8
	RPMFILE_EXCLUDE   = 1 << 9
	RPMFILE_UNPATCHED = 1 << 10
	RPMFILE_PUBKEY    = 1 << 11
	RPMFILE_POLICY    = 1 << 12

	// FILEVERIFYFLAGS bitmask elements:
	RPMVERIFY_NONE       = 0
	RPMVERIFY_MD5        = 1 << 0
	RPMVERIFY_FILEDIGEST = 1 << 0
	RPMVERIFY_FILESIZE   = 1 << 1
	RPMVERIFY_LINKTO     = 1 << 2
	RPMVERIFY_USER       = 1 << 3
	RPMVERIFY_GROUP      = 1 << 4
	RPMVERIFY_MTIME      = 1 << 5
	RPMVERIFY_MODE       = 1 << 6
	RPMVERIFY_RDEV       = 1 << 7
	RPMVERIFY_CAPS       = 1 << 8
	RPMVERIFY_CONTEXTS   = 1 << 15

	// TRIGGERFLAGS bitmask elements -- not all rpmsenseFlags make sense
	// in TRIGGERFLAGS
	RPMSENSE_ANY           = 0
	RPMSENSE_LESS          = 1 << 1
	RPMSENSE_GREATER       = 1 << 2
	RPMSENSE_EQUAL         = 1 << 3
	RPMSENSE_TRIGGERIN     = 1 << 16
	RPMSENSE_TRIGGERUN     = 1 << 17
	RPMSENSE_TRIGGERPOSTUN = 1 << 18
	RPMSENSE_TRIGGERPREIN  = 1 << 25

	// Header region tags
	RPMTAG_HEADERSIGNATURES = 62
	RPMTAG_HEADERIMMUTABLE  = 63
	RPMTAG_HEADERREGIONS    = 64
)
<commit_msg>Add changelog-related tag constants.<commit_after>/*
 * Copyright (c) SAS Institute, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package rpmutils

const (
	RPM_NULL_TYPE         = 0
	RPM_CHAR_TYPE         = 1
	RPM_INT8_TYPE         = 2
	RPM_INT16_TYPE        = 3
	RPM_INT32_TYPE        = 4
	RPM_INT64_TYPE        = 5
	RPM_STRING_TYPE       = 6
	RPM_BIN_TYPE          = 7
	RPM_STRING_ARRAY_TYPE = 8
	RPM_I18NSTRING_TYPE   = 9

	_GENERAL_TAG_BASE = 1000
	NAME              = 1000
	VERSION           = 1001
	RELEASE           = 1002
	EPOCH             = 1003
	SUMMARY           = 1004
	DESCRIPTION       = 1005
	BUILDTIME         = 1006
	BUILDHOST         = 1007
	SIZE              = 1009
	DISTRIBUTION      = 1010
	VENDOR            = 1011
	GIF               = 1012
	XPM               = 1013
	LICENSE           = 1014
	PACKAGER          = 1015
	GROUP             = 1016
	CHANGELOG         = 1017
	SOURCE            = 1018
	PATCH             = 1019
	URL               = 1020
	OS                = 1021
	ARCH              = 1022
	PREIN             = 1023
	POSTIN            = 1024
	PREUN             = 1025
	POSTUN            = 1026
	OLDFILENAMES      = 1027
	FILESIZES         = 1028
	FILEMODES         = 1030
	FILERDEVS         = 1033
	FILEMTIMES        = 1034
	FILEDIGESTS       = 1035 // AKA FILEMD5S
	FILELINKTOS       = 1036
	FILEFLAGS         = 1037 // bitmask: RPMFILE_* are bitmasks to interpret
	FILEUSERNAME      = 1039
	FILEGROUPNAME     = 1040
	ICON              = 1043
	SOURCERPM         = 1044
	FILEVERIFYFLAGS   = 1045 // bitmask: RPMVERIFY_* are bitmasks to interpret
	ARCHIVESIZE       = 1046
	PROVIDENAME       = 1047
	REQUIREFLAGS      = 1048
	REQUIRENAME       = 1049
	REQUIREVERSION    = 1050
	RPMVERSION        = 1064
	TRIGGERSCRIPTS    = 1065
	TRIGGERNAME       = 1066
	TRIGGERVERSION    = 1067
	TRIGGERFLAGS      = 1068 // bitmask: RPMSENSE_* are bitmasks to interpret
	TRIGGERINDEX      = 1069
	VERIFYSCRIPT      = 1079
	CHANGELOGTIME     = 1080
	CHANGELOGNAME     = 1081
	CHANGELOGTEXT     = 1082
	PREINPROG         = 1085
	POSTINPROG        = 1086
	PREUNPROG         = 1087
	POSTUNPROG        = 1088
	OBSOLETENAME      = 1090
	FILEDEVICES       = 1095
	FILEINODES        = 1096
	PROVIDEFLAGS      = 1112
	PROVIDEVERSION    = 1113
	OBSOLETEFLAGS     = 1114
	OBSOLETEVERSION   = 1115
	VERIFYSCRIPTPROG  = 1091
	TRIGGERSCRIPTPROG = 1092
	DIRINDEXES        = 1116
	BASENAMES         = 1117
	DIRNAMES          = 1118
	PAYLOADFORMAT     = 1124
	PAYLOADCOMPRESSOR = 1125
	FILECOLORS        = 1140
	// BLINK*, FLINK*, and TRIGGERPREIN included from SUSE fork of RPM
	BLINKPKGID   = 1164
	BLINKHDRID   = 1165
	BLINKNEVRA   = 1166
	FLINKPKGID   = 1167
	FLINKHDRID   = 1168
	FLINKNEVRA   = 1169
	TRIGGERPREIN = 1170

	LONGFILESIZES  = 5008
	LONGSIZE       = 5009
	FILECAPS       = 5010
	FILEDIGESTALGO = 5011
	BUGURL         = 5012
	VCS            = 5034
	ENCODING       = 5062

	SIG_BASE            = 256
	SIG_DSA             = SIG_BASE + 11 // DSA signature over header only
	SIG_RSA             = SIG_BASE + 12 // RSA signature over header only
	SIG_SHA1            = SIG_BASE + 13 // SHA1 over header only (hex)
	SIG_LONGSIZE        = SIG_BASE + 15 // header + compressed payload (uint64)
	SIG_LONGARCHIVESIZE = SIG_BASE + 15 // uncompressed payload bytes (uint64)

	// Given that there is overlap between signature tag headers and general tag
	// headers, we offset the signature ones by some amount
	_SIGHEADER_TAG_BASE = 16384
	SIG_SIZE            = _SIGHEADER_TAG_BASE + 1000 // Header + Payload size
	SIG_PGP             = _SIGHEADER_TAG_BASE + 1002 // Signature over header + payload
	SIG_MD5             = _SIGHEADER_TAG_BASE + 1004 // MD5SUM of header + payload
	SIG_GPG             = _SIGHEADER_TAG_BASE + 1005 // (same as SIG_PGP)
	SIG_PAYLOADSIZE     = _SIGHEADER_TAG_BASE + 1007 // uncompressed payload bytes (uint32)
	SIG_RESERVEDSPACE   = _SIGHEADER_TAG_BASE + 1008 // blank space that can be replaced by a signature

	// FILEFLAGS bitmask elements:
	RPMFILE_NONE      = 0
	RPMFILE_CONFIG    = 1 << 0
	RPMFILE_DOC       = 1 << 1
	RPMFILE_ICON      = 1 << 2
	RPMFILE_MISSINGOK = 1 << 3
	RPMFILE_NOREPLACE = 1 << 4
	RPMFILE_SPECFILE  = 1 << 5
	RPMFILE_GHOST     = 1 << 6
	RPMFILE_LICENSE   = 1 << 7
	RPMFILE_README    = 1 << 8
	RPMFILE_EXCLUDE   = 1 << 9
	RPMFILE_UNPATCHED = 1 << 10
	RPMFILE_PUBKEY    = 1 << 11
	RPMFILE_POLICY    = 1 << 12

	// FILEVERIFYFLAGS bitmask elements:
	RPMVERIFY_NONE       = 0
	RPMVERIFY_MD5        = 1 << 0
	RPMVERIFY_FILEDIGEST = 1 << 0
	RPMVERIFY_FILESIZE   = 1 << 1
	RPMVERIFY_LINKTO     = 1 << 2
	RPMVERIFY_USER       = 1 << 3
	RPMVERIFY_GROUP      = 1 << 4
	RPMVERIFY_MTIME      = 1 << 5
	RPMVERIFY_MODE       = 1 << 6
	RPMVERIFY_RDEV       = 1 << 7
	RPMVERIFY_CAPS       = 1 << 8
	RPMVERIFY_CONTEXTS   = 1 << 15

	// TRIGGERFLAGS bitmask elements -- not all rpmsenseFlags make sense
	// in TRIGGERFLAGS
	RPMSENSE_ANY           = 0
	RPMSENSE_LESS          = 1 << 1
	RPMSENSE_GREATER       = 1 << 2
	RPMSENSE_EQUAL         = 1 << 3
	RPMSENSE_TRIGGERIN     = 1 << 16
	RPMSENSE_TRIGGERUN     = 1 << 17
	RPMSENSE_TRIGGERPOSTUN = 1 << 18
	RPMSENSE_TRIGGERPREIN  = 1 << 25

	// Header region tags
	RPMTAG_HEADERSIGNATURES = 62
	RPMTAG_HEADERIMMUTABLE  = 63
	RPMTAG_HEADERREGIONS    = 64
)
<|endoftext|>


<commit_before>package tasker

import (
	"context"
	"errors"
	"time"

	"github.com/NoneBorder/dora"
	"github.com/astaxie/beego/orm"
)

const (
	TaskStatPending = "pending"
	TaskStatRunning = "running"
	TaskStatRetry   = "retry"
	TaskStatFailed  = "failed"
	TaskStatSuccess = "success"
)

const (
	TaskCtxKeyLog       = "stdout"          // normal log record
	TaskCtxKeyExecAfter = "task-exec-after" // plan task exec time
)

// Task is the core object for tasker package.
type Task struct {
	Id         uint64
	Topic      string
	Status     string
	Timeout    int // 超时ms
	ExecAfter  time.Time `orm:"type(datetime);default(1971-01-01 00:00:00)"`
	Retry      int
	Input      string    `orm:"type(text)"`
	WorkerId   uint64
	WorkerFQDN string    `orm:"column(worker_fqdn);type(text)"`
	Created    time.Time `orm:"auto_now_add"`
	Updated    time.Time `orm:"auto_now"`
	Log        string    `orm:"type(text)"`
	Errlog     string    `orm:"type(text)"`
}

// ConsumeFn represents consume func definition.
type ConsumeFn func(ctx *context.Context, Input string, WorkerID uint64) (err error)

func (self *Task) TableIndex() [][]string {
	return [][]string{
		[]string{"Topic", "Status"},
	}
}

func (self *Task) TableEngine() string {
	return "INNODB"
}

// NewSimpleTask is new task with default settings.
func NewSimpleTask(topic, input string) *Task {
	return &Task{
		Topic:     topic,
		Status:    TaskStatPending,
		Timeout:   5000,
		ExecAfter: zeroTime,
		Retry:     3,
		Input:     input,
		WorkerId:  0,
	}
}

// Publish Task to Msg Queue.
func (self *Task) Publish() (err error) {
	o := orm.NewOrm()
	if _, err = o.Insert(self); err == nil {
		Stat(self.Topic, TaskStatPending, time.Duration(0))
	}
	return
}

func (self *Task) exec(ctx *context.Context, fn ConsumeFn) (err error) {
	errC := make(chan error)
	select {
	case errC <- fn(ctx, self.Input, self.WorkerId):
		err = <-errC

	case <-(*ctx).Done():
		err = (*ctx).Err()
	}

	return
}

func (self *Task) consume(fn ConsumeFn) bool {
	startT := time.Now().Local()
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(self.Timeout+1)*time.Microsecond)
	defer cancel()

	err := self.exec(&ctx, fn)
	success := err == nil

	if !success {
		// failed
		self.Errlog += err.Error() + "\n"
		self.Status = TaskStatFailed
		self.Retry -= 1
		if self.Retry > 0 {
			self.Status = TaskStatRetry

			// when the job will retry, you can update the next exec time
			if execAfter, ok := ctx.Value(TaskCtxKeyExecAfter).(time.Time); ok {
				self.ExecAfter = execAfter
			}
		}
	} else {
		self.Status = TaskStatSuccess
	}

	// stat task exec
	Stat(self.Topic, self.Status, time.Since(startT))

	// workerid set to 0
	self.WorkerId = 0

	// set stdout log from ctx
	if log, ok := ctx.Value(TaskCtxKeyLog).(string); ok && log != "" {
		self.Log = log
	}

	// update db record
	o := orm.NewOrm()
	for i := 0; i < 3; i++ {
		if _, err := o.Update(self); err != nil {
			dora.Error().Msgf("set task state failed: err=%s, task=%s", err.Error(), self)
		} else {
			break
		}
	}

	return success
}

// Consume a topic task.
func Consume(topic string, fn ConsumeFn) (int, error) {
	o := orm.NewOrm()
	if IsMaster {
		// 仅在 master 上对已经僵死的 task 进行重置，减轻 DB 压力
		if _, err := o.Raw(`UPDATE task SET status=?, retry=retry-1, worker_id=0 WHERE topic=? AND status=?
		AND TIMESTAMPDIFF(SECOND, updated, now())*1000-5000>timeout`,
			TaskStatRetry, topic, TaskStatRunning,
		).Exec(); err != nil {
			dora.Error().Msgf("update dead running task to waiting failed: %s", err.Error())
		}
	}

	maxConsumeTask := cap(RunningTaskChannel[topic]) - len(RunningTaskChannel[topic])
	if maxConsumeTask == 0 {
		return 0, errors.New("reached max running task limit")
	}

	workerID, err := UniqID.NextID()
	if err != nil {
		return 0, err
	}

	res, err := o.Raw(`UPDATE task SET status=?, worker_id=?, worker_fqdn=concat(worker_fqdn, ?) WHERE
		topic=? AND worker_id=0 AND status IN (?,?) AND exec_after > NOW() LIMIT ?`,
		TaskStatRunning, workerID, FQDN()+"\n", topic, TaskStatPending, TaskStatRetry, maxConsumeTask,
	).Exec()
	if err != nil {
		// not update db
		return 0, err
	}

	num, _ := res.RowsAffected()
	if num == 0 {
		// no task for consume
		return 0, nil
	}

	tasks := []*Task{}
	_, err = o.Raw("SELECT * FROM task WHERE `topic`=? AND `status`=? AND `worker_id`=?",
		topic, TaskStatRunning, workerID,
	).QueryRows(&tasks)
	if len(tasks) == 0 || err != nil {
		return 0, err
	}

	// TODO: 如何保证程序退出时正在运行的 task 可用？当前超时后 5s 僵尸 task 会被重新激活，但不能保证原子性
	for _, t := range tasks {
		RunningTaskChannel[topic] <- true
		go func(t *Task) {
			t.consume(fn)
			<-RunningTaskChannel[topic]
		}(t)
	}

	return len(tasks), nil
}

// ConsumeCtxStdout write stdout log to context for record
func ConsumeCtxStdout(ctx *context.Context, log string) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyLog, log)
}

// ConsumeCtxExecAfter update run time for next retry
func ConsumeCtxExecAfter(ctx *context.Context, t time.Time) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyExecAfter, t)
}
<commit_msg>fix sql bug<commit_after>package tasker

import (
	"context"
	"errors"
	"time"

	"github.com/NoneBorder/dora"
	"github.com/astaxie/beego/orm"
)

const (
	TaskStatPending = "pending"
	TaskStatRunning = "running"
	TaskStatRetry   = "retry"
	TaskStatFailed  = "failed"
	TaskStatSuccess = "success"
)

const (
	TaskCtxKeyLog       = "stdout"          // normal log record
	TaskCtxKeyExecAfter = "task-exec-after" // plan task exec time
)

// Task is the core object for tasker package.
type Task struct {
	Id         uint64
	Topic      string
	Status     string
	Timeout    int // 超时ms
	ExecAfter  time.Time `orm:"type(datetime);default(1971-01-01 00:00:00)"`
	Retry      int
	Input      string    `orm:"type(text)"`
	WorkerId   uint64
	WorkerFQDN string    `orm:"column(worker_fqdn);type(text)"`
	Created    time.Time `orm:"auto_now_add"`
	Updated    time.Time `orm:"auto_now"`
	Log        string    `orm:"type(text)"`
	Errlog     string    `orm:"type(text)"`
}

// ConsumeFn represents consume func definition.
type ConsumeFn func(ctx *context.Context, Input string, WorkerID uint64) (err error)

func (self *Task) TableIndex() [][]string {
	return [][]string{
		[]string{"Topic", "Status"},
	}
}

func (self *Task) TableEngine() string {
	return "INNODB"
}

// NewSimpleTask is new task with default settings.
func NewSimpleTask(topic, input string) *Task {
	return &Task{
		Topic:     topic,
		Status:    TaskStatPending,
		Timeout:   5000,
		ExecAfter: zeroTime,
		Retry:     3,
		Input:     input,
		WorkerId:  0,
	}
}

// Publish Task to Msg Queue.
func (self *Task) Publish() (err error) {
	o := orm.NewOrm()
	if _, err = o.Insert(self); err == nil {
		Stat(self.Topic, TaskStatPending, time.Duration(0))
	}
	return
}

func (self *Task) exec(ctx *context.Context, fn ConsumeFn) (err error) {
	errC := make(chan error)
	select {
	case errC <- fn(ctx, self.Input, self.WorkerId):
		err = <-errC

	case <-(*ctx).Done():
		err = (*ctx).Err()
	}

	return
}

func (self *Task) consume(fn ConsumeFn) bool {
	startT := time.Now().Local()
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(self.Timeout+1)*time.Microsecond)
	defer cancel()

	err := self.exec(&ctx, fn)
	success := err == nil

	if !success {
		// failed
		self.Errlog += err.Error() + "\n"
		self.Status = TaskStatFailed
		self.Retry -= 1
		if self.Retry > 0 {
			self.Status = TaskStatRetry

			// when the job will retry, you can update the next exec time
			if execAfter, ok := ctx.Value(TaskCtxKeyExecAfter).(time.Time); ok {
				self.ExecAfter = execAfter
			}
		}
	} else {
		self.Status = TaskStatSuccess
	}

	// stat task exec
	Stat(self.Topic, self.Status, time.Since(startT))

	// workerid set to 0
	self.WorkerId = 0

	// set stdout log from ctx
	if log, ok := ctx.Value(TaskCtxKeyLog).(string); ok && log != "" {
		self.Log = log
	}

	// update db record
	o := orm.NewOrm()
	for i := 0; i < 3; i++ {
		if _, err := o.Update(self); err != nil {
			dora.Error().Msgf("set task state failed: err=%s, task=%s", err.Error(), self)
		} else {
			break
		}
	}

	return success
}

// Consume a topic task.
func Consume(topic string, fn ConsumeFn) (int, error) {
	o := orm.NewOrm()
	if IsMaster {
		// 仅在 master 上对已经僵死的 task 进行重置，减轻 DB 压力
		if _, err := o.Raw(`UPDATE task SET status=?, retry=retry-1, worker_id=0 WHERE topic=? AND status=?
		AND TIMESTAMPDIFF(SECOND, updated, now())*1000-5000>timeout`,
			TaskStatRetry, topic, TaskStatRunning,
		).Exec(); err != nil {
			dora.Error().Msgf("update dead running task to waiting failed: %s", err.Error())
		}
	}

	maxConsumeTask := cap(RunningTaskChannel[topic]) - len(RunningTaskChannel[topic])
	if maxConsumeTask == 0 {
		return 0, errors.New("reached max running task limit")
	}

	workerID, err := UniqID.NextID()
	if err != nil {
		return 0, err
	}

	res, err := o.Raw(`UPDATE task SET status=?, worker_id=?, worker_fqdn=concat(worker_fqdn, ?) WHERE
		topic=? AND worker_id=0 AND status IN (?,?) AND exec_after <= NOW() LIMIT ?`,
		TaskStatRunning, workerID, FQDN()+"\n", topic, TaskStatPending, TaskStatRetry, maxConsumeTask,
	).Exec()
	if err != nil {
		// not update db
		return 0, err
	}

	num, _ := res.RowsAffected()
	if num == 0 {
		// no task for consume
		return 0, nil
	}

	tasks := []*Task{}
	_, err = o.Raw("SELECT * FROM task WHERE `topic`=? AND `status`=? AND `worker_id`=?",
		topic, TaskStatRunning, workerID,
	).QueryRows(&tasks)
	if len(tasks) == 0 || err != nil {
		return 0, err
	}

	// TODO: 如何保证程序退出时正在运行的 task 可用？当前超时后 5s 僵尸 task 会被重新激活，但不能保证原子性
	for _, t := range tasks {
		RunningTaskChannel[topic] <- true
		go func(t *Task) {
			t.consume(fn)
			<-RunningTaskChannel[topic]
		}(t)
	}

	return len(tasks), nil
}

// ConsumeCtxStdout write stdout log to context for record
func ConsumeCtxStdout(ctx *context.Context, log string) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyLog, log)
}

// ConsumeCtxExecAfter update run time for next retry
func ConsumeCtxExecAfter(ctx *context.Context, t time.Time) {
	*ctx = context.WithValue(*ctx, TaskCtxKeyExecAfter, t)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
)

func main() {
	http.Handle("/control/", handler)

	err := http.ListenAndServe(":8000", nil)
	if err != nil {
		panic(err.Error())
	}
}

func control(w http.ResponseWriter, req *http.Request) {
	fmt.Fprint(w, "Oh hi there!")
}
<commit_msg>stubbed out REST endpoints for move_x and move_y<commit_after>package main

import (
	"fmt"
	"github.com/gorilla/mux"
	"net/http"
	"strconv"
)

type movement struct {
	X_dif int
	Y_dif int
}

func main() {

	movements := make(chan movement, 500) //makes a channel for movements, with a depth of 500

	http.HandleFunc("/control/", control)
	http.HandleFunc("/move_x/{toMove}", func(w http.ResponseWriter, req *http.Request) {
		vars := mux.Vars(req)
		toMove, err := strconv.Atoi(vars["toMove"])
		if err != nil {
			w.WriteHeader(500)
			fmt.Fprintf(w, "Error: %s is probably not an int", vars["toMove"])
		} else {
			move_x(movements, toMove)
			fmt.Fprintf(w, "Moved %d on x", toMove)
		}
	})

	http.HandleFunc("/move_y/{toMove}", func(w http.ResponseWriter, req *http.Request) {
		vars := mux.Vars(req)
		toMove, err := strconv.Atoi(vars["toMove"])
		if err != nil {
			w.WriteHeader(500)
			fmt.Fprintf(w, "Error: %s is probably not an int", vars["toMove"])
		} else {
			move_y(movements, toMove)
			fmt.Fprintf(w, "Moved %d on y", toMove)
		}
	})

	err := http.ListenAndServe(":8000", nil)
	if err != nil {
		panic(err.Error())
	}
}

func control(w http.ResponseWriter, req *http.Request) {
	fmt.Fprint(w, "Oh hi there!")
}

func move_x(movements chan movement, toMove int) {

}

func move_y(movements chan movement, toMove int) {

}
<|endoftext|>


<commit_before>package radius

import (
	"errors"
	"net"
	"sync"
)

// Handler is a value that can handle a server's RADIUS packet event.
type Handler interface {
	ServeRadius(w ResponseWriter, p *Packet)
}

// HandlerFunc is a wrapper that allows ordinary functions to be used as a
// handler.
type HandlerFunc func(w ResponseWriter, p *Packet)

// ServeRadius calls h(w, p).
func (h HandlerFunc) ServeRadius(w ResponseWriter, p *Packet) {
	h(w, p)
}

// ResponseWriter is used by Handler when replying to a RADIUS packet.
type ResponseWriter interface {
	// LocalAddr returns the address of the local server that accepted the
	// packet.
	LocalAddr() net.Addr
	// RemoteAddr returns the address of the remote client that sent to packet.
	RemoteAddr() net.Addr

	// Write sends a packet to the sender.
	Write(packet *Packet) error

	// AccessAccept sends an Access-Accept packet to the sender that includes
	// the given attributes.
	AccessAccept(attributes ...*Attribute) error
	// AccessAccept sends an Access-Reject packet to the sender that includes
	// the given attributes.
	AccessReject(attributes ...*Attribute) error
	// AccessAccept sends an Access-Challenge packet to the sender that includes
	// the given attributes.
	AccessChallenge(attributes ...*Attribute) error
}

type responseWriter struct {
	// listener that received the packet
	conn *net.UDPConn
	// where the packet came from
	addr *net.UDPAddr
	// original packet
	packet *Packet
}

func (r *responseWriter) LocalAddr() net.Addr {
	return r.conn.LocalAddr()
}

func (r *responseWriter) RemoteAddr() net.Addr {
	return r.addr
}

func (r *responseWriter) accessRespond(code Code, attributes ...*Attribute) error {
	packet := Packet{
		Code:          code,
		Identifier:    r.packet.Identifier,
		Authenticator: r.packet.Authenticator,

		Secret: r.packet.Secret,

		Dictionary: r.packet.Dictionary,

		Attributes: attributes,
	}
	return r.Write(&packet)
}

func (r *responseWriter) AccessAccept(attributes ...*Attribute) error {
	// TOOD: do not send if packet was not Access-Request
	return r.accessRespond(CodeAccessAccept, attributes...)
}

func (r *responseWriter) AccessReject(attributes ...*Attribute) error {
	// TOOD: do not send if packet was not Access-Request
	return r.accessRespond(CodeAccessReject, attributes...)
}

func (r *responseWriter) AccessChallenge(attributes ...*Attribute) error {
	// TOOD: do not send if packet was not Access-Request
	return r.accessRespond(CodeAccessChallenge, attributes...)
}

func (r *responseWriter) Write(packet *Packet) error {
	raw, err := packet.Encode()
	if err != nil {
		return err
	}
	if _, err := r.conn.WriteToUDP(raw, r.addr); err != nil {
		return err
	}
	return nil
}

// Server is a server that listens for and handles RADIUS packets.
type Server struct {
	// Address to bind the server on. If empty, the address defaults to ":1812".
	Addr string
	// Network of the server. Valid values are "udp", "udp4", "udp6". If empty,
	// the network defaults to "udp".
	Network string
	// The shared secret between the client and server.
	Secret []byte

	// TODO: allow a secret function to be defined, which returned the secret
	// that should be used for the given client.

	// Dictionary used when decoding incoming packets.
	Dictionary *Dictionary

	// The packet handler that handles incoming, valid packets.
	Handler Handler

	listener *net.UDPConn
}

// ListenAndServe starts a RADIUS server on the address given in s.
func (s *Server) ListenAndServe() error {
	if s.listener != nil {
		return errors.New("radius: server already started")
	}

	if s.Handler == nil {
		return errors.New("radius: nil Handler")
	}

	addrStr := ":1812"
	if s.Addr != "" {
		addrStr = s.Addr
	}

	network := "udp"
	if s.Network != "" {
		network = s.Network
	}

	addr, err := net.ResolveUDPAddr(network, addrStr)
	if err != nil {
		return err
	}
	s.listener, err = net.ListenUDP(network, addr)
	if err != nil {
		return err
	}

	type activeKey struct {
		IP         string
		Identifier byte
	}

	var (
		activeLock sync.Mutex
		active     = map[activeKey]bool{}
	)

	for {
		buff := make([]byte, 4096)
		n, remoteAddr, err := s.listener.ReadFromUDP(buff)
		if err != nil && !err.(*net.OpError).Temporary() {
			break
		}
		if n == 0 {
			continue
		}
		buff = buff[:n]

		go func(conn *net.UDPConn, buff []byte, remoteAddr *net.UDPAddr) {
			packet, err := Parse(buff, s.Secret, s.Dictionary)
			if err != nil {
				return
			}
			key := activeKey{
				IP:         remoteAddr.String(),
				Identifier: packet.Identifier,
			}
			activeLock.Lock()
			if _, ok := active[key]; ok {
				activeLock.Unlock()
				return
			}
			active[key] = true
			activeLock.Unlock()

			response := responseWriter{
				conn:   conn,
				addr:   remoteAddr,
				packet: packet,
			}

			s.Handler.ServeRadius(&response, packet)

			activeLock.Lock()
			delete(active, key)
			activeLock.Unlock()
		}(s.listener, buff, remoteAddr)
	}
	// TODO: only return nil if s.Close was called
	s.listener = nil
	return nil
}

// Close stops listening for packets. Any packet that is currently being
// handled will not be able to respond to the sender.
func (s *Server) Close() error {
	if s.listener == nil {
		return nil
	}
	return s.listener.Close()
}
<commit_msg>rename ServeRadius to ServeRADIUS<commit_after>package radius

import (
	"errors"
	"net"
	"sync"
)

// Handler is a value that can handle a server's RADIUS packet event.
type Handler interface {
	ServeRADIUS(w ResponseWriter, p *Packet)
}

// HandlerFunc is a wrapper that allows ordinary functions to be used as a
// handler.
type HandlerFunc func(w ResponseWriter, p *Packet)

// ServeRADIUS calls h(w, p).
func (h HandlerFunc) ServeRADIUS(w ResponseWriter, p *Packet) {
	h(w, p)
}

var _ Handler = HandlerFunc(nil)

// ResponseWriter is used by Handler when replying to a RADIUS packet.
type ResponseWriter interface {
	// LocalAddr returns the address of the local server that accepted the
	// packet.
	LocalAddr() net.Addr
	// RemoteAddr returns the address of the remote client that sent to packet.
	RemoteAddr() net.Addr

	// Write sends a packet to the sender.
	Write(packet *Packet) error

	// AccessAccept sends an Access-Accept packet to the sender that includes
	// the given attributes.
	AccessAccept(attributes ...*Attribute) error
	// AccessAccept sends an Access-Reject packet to the sender that includes
	// the given attributes.
	AccessReject(attributes ...*Attribute) error
	// AccessAccept sends an Access-Challenge packet to the sender that includes
	// the given attributes.
	AccessChallenge(attributes ...*Attribute) error
}

type responseWriter struct {
	// listener that received the packet
	conn *net.UDPConn
	// where the packet came from
	addr *net.UDPAddr
	// original packet
	packet *Packet
}

func (r *responseWriter) LocalAddr() net.Addr {
	return r.conn.LocalAddr()
}

func (r *responseWriter) RemoteAddr() net.Addr {
	return r.addr
}

func (r *responseWriter) accessRespond(code Code, attributes ...*Attribute) error {
	packet := Packet{
		Code:          code,
		Identifier:    r.packet.Identifier,
		Authenticator: r.packet.Authenticator,

		Secret: r.packet.Secret,

		Dictionary: r.packet.Dictionary,

		Attributes: attributes,
	}
	return r.Write(&packet)
}

func (r *responseWriter) AccessAccept(attributes ...*Attribute) error {
	// TOOD: do not send if packet was not Access-Request
	return r.accessRespond(CodeAccessAccept, attributes...)
}

func (r *responseWriter) AccessReject(attributes ...*Attribute) error {
	// TOOD: do not send if packet was not Access-Request
	return r.accessRespond(CodeAccessReject, attributes...)
}

func (r *responseWriter) AccessChallenge(attributes ...*Attribute) error {
	// TOOD: do not send if packet was not Access-Request
	return r.accessRespond(CodeAccessChallenge, attributes...)
}

func (r *responseWriter) Write(packet *Packet) error {
	raw, err := packet.Encode()
	if err != nil {
		return err
	}
	if _, err := r.conn.WriteToUDP(raw, r.addr); err != nil {
		return err
	}
	return nil
}

// Server is a server that listens for and handles RADIUS packets.
type Server struct {
	// Address to bind the server on. If empty, the address defaults to ":1812".
	Addr string
	// Network of the server. Valid values are "udp", "udp4", "udp6". If empty,
	// the network defaults to "udp".
	Network string
	// The shared secret between the client and server.
	Secret []byte

	// TODO: allow a secret function to be defined, which returned the secret
	// that should be used for the given client.

	// Dictionary used when decoding incoming packets.
	Dictionary *Dictionary

	// The packet handler that handles incoming, valid packets.
	Handler Handler

	listener *net.UDPConn
}

// ListenAndServe starts a RADIUS server on the address given in s.
func (s *Server) ListenAndServe() error {
	if s.listener != nil {
		return errors.New("radius: server already started")
	}

	if s.Handler == nil {
		return errors.New("radius: nil Handler")
	}

	addrStr := ":1812"
	if s.Addr != "" {
		addrStr = s.Addr
	}

	network := "udp"
	if s.Network != "" {
		network = s.Network
	}

	addr, err := net.ResolveUDPAddr(network, addrStr)
	if err != nil {
		return err
	}
	s.listener, err = net.ListenUDP(network, addr)
	if err != nil {
		return err
	}

	type activeKey struct {
		IP         string
		Identifier byte
	}

	var (
		activeLock sync.Mutex
		active     = map[activeKey]bool{}
	)

	for {
		buff := make([]byte, 4096)
		n, remoteAddr, err := s.listener.ReadFromUDP(buff)
		if err != nil && !err.(*net.OpError).Temporary() {
			break
		}
		if n == 0 {
			continue
		}
		buff = buff[:n]

		go func(conn *net.UDPConn, buff []byte, remoteAddr *net.UDPAddr) {
			packet, err := Parse(buff, s.Secret, s.Dictionary)
			if err != nil {
				return
			}
			key := activeKey{
				IP:         remoteAddr.String(),
				Identifier: packet.Identifier,
			}
			activeLock.Lock()
			if _, ok := active[key]; ok {
				activeLock.Unlock()
				return
			}
			active[key] = true
			activeLock.Unlock()

			response := responseWriter{
				conn:   conn,
				addr:   remoteAddr,
				packet: packet,
			}

			s.Handler.ServeRADIUS(&response, packet)

			activeLock.Lock()
			delete(active, key)
			activeLock.Unlock()
		}(s.listener, buff, remoteAddr)
	}
	// TODO: only return nil if s.Close was called
	s.listener = nil
	return nil
}

// Close stops listening for packets. Any packet that is currently being
// handled will not be able to respond to the sender.
func (s *Server) Close() error {
	if s.listener == nil {
		return nil
	}
	return s.listener.Close()
}
<|endoftext|>