<commit_before>package iocap

import (
	"sync"
	"time"
)

// bucket is a simple "leaky bucket" abstraction to provide a way to
// limit the number of operations (in this case, byte reads/writes)
// allowed within a given interval.
type bucket struct {
	opts    RateOpts
	drained time.Time

	// Tokens is the number of tokens present in the bucket. A simple int is
	// used to allow for faster token acquisition, rather than a channel.
	// Arguably, due to the blocking nature of iocap, a channel may be
	// theoretically more appropriate for this use. The reality pitfall is
	// that billions of channel reads are far more expensive than taking a
	// lock and doing basic math.
	tokens int

	sync.RWMutex
}

// newBucket creates a new bucket to use for readers and writers.
func newBucket(opts RateOpts) *bucket {
	return &bucket{
		opts: opts,
	}
}

// insert performs a best-effort token insert of n tokens. v contains
// the number of tokens inserted, which will differ from n if the
// bucket overflows. insert will block until at least one token is
// successfully inserted.
func (b *bucket) insert(n int) (v int) {
	// Call a non-blocking drain up-front to make room for tokens.
	b.drain(false)

INSERT:
	var remain int

	b.RLock()
	tokens := b.tokens
	opts := b.opts
	b.RUnlock()

	switch {
	case opts == Unlimited:
		// No limit should be applied.
		return n

	case tokens == opts.Size:
		// Bucket is full. Call a blocking drain to wait for the next
		// drain interval (earliest we can insert more tokens).
		b.drain(true)
		goto INSERT

	case tokens+n > opts.Size:
		// Some tokens, but not all, were inserted. The bucket is now
		// full and subsequent inserts will overflow and block.
		v = opts.Size - tokens
		remain = opts.Size

	default:
		// All tokens inserted successfully.
		v = n
		remain = tokens + n
	}

	b.Lock()

	// Check if the token count was modified before the lock
	// was acquired.
	if b.tokens != tokens {
		b.Unlock()
		goto INSERT
	}

	b.tokens = remain
	b.Unlock()
	return
}

// drain is used to drain the bucket of tokens. If wait is true, drain
// will wait until the next drain cycle and then continue. Otherwise,
// drain only drains the bucket if it is due.
//
// This implementation is heavy-handed in that it brackets "leaking" tokens
// to the full duration of the configured interval. In other words, the
// bucket leaks not in single drops, but rather multiples, and only when the
// token drain window has elapsed. This side-steps near-hot-looping with
// dense token expiration (short interval + high size) and heavy lock
// contention. A possible enhancement would be to make this more granular.
func (b *bucket) drain(wait bool) {
	b.RLock()
	last := b.drained
	interval := b.opts.Interval
	b.RUnlock()

	switch {
	case time.Since(last) >= interval:
		b.Lock()
		defer b.Unlock()

		// Make sure the timestamp was not updated; prevents a time-of-
		// check vs. time-of-use error.
		if !b.drained.Equal(last) {
			return
		}

		// Drain the bucket.
		b.tokens = 0

		// Update the drain timestamp.
		b.drained = time.Now()

	case wait:
		delay := last.Add(interval).Sub(time.Now())
		time.Sleep(delay)
		b.drain(false)
	}
}

// setRate safely replaces the RateOpts on the bucket.
func (b *bucket) setRate(opts RateOpts) {
	b.Lock()
	b.opts = opts
	b.Unlock()
}
<commit_msg>Make mutex private on bucket<commit_after>package iocap

import (
	"sync"
	"time"
)

// bucket is a simple "leaky bucket" abstraction to provide a way to
// limit the number of operations (in this case, byte reads/writes)
// allowed within a given interval.
type bucket struct {
	opts    RateOpts
	drained time.Time

	// Tokens is the number of tokens present in the bucket. A simple int is
	// used to allow for faster token acquisition, rather than a channel.
	// Arguably, due to the blocking nature of iocap, a channel may be
	// theoretically more appropriate for this use. The reality pitfall is
	// that billions of channel reads are far more expensive than taking a
	// lock and doing basic math.
	tokens int

	l sync.RWMutex
}

// newBucket creates a new bucket to use for readers and writers.
func newBucket(opts RateOpts) *bucket {
	return &bucket{
		opts: opts,
	}
}

// insert performs a best-effort token insert of n tokens. v contains
// the number of tokens inserted, which will differ from n if the
// bucket overflows. insert will block until at least one token is
// successfully inserted.
func (b *bucket) insert(n int) (v int) {
	// Call a non-blocking drain up-front to make room for tokens.
	b.drain(false)

INSERT:
	var remain int

	b.l.RLock()
	tokens := b.tokens
	opts := b.opts
	b.l.RUnlock()

	switch {
	case opts == Unlimited:
		// No limit should be applied.
		return n

	case tokens == opts.Size:
		// Bucket is full. Call a blocking drain to wait for the next
		// drain interval (earliest we can insert more tokens).
		b.drain(true)
		goto INSERT

	case tokens+n > opts.Size:
		// Some tokens, but not all, were inserted. The bucket is now
		// full and subsequent inserts will overflow and block.
		v = opts.Size - tokens
		remain = opts.Size

	default:
		// All tokens inserted successfully.
		v = n
		remain = tokens + n
	}

	b.l.Lock()

	// Check if the token count was modified before the lock
	// was acquired.
	if b.tokens != tokens {
		b.l.Unlock()
		goto INSERT
	}

	b.tokens = remain
	b.l.Unlock()
	return
}

// drain is used to drain the bucket of tokens. If wait is true, drain
// will wait until the next drain cycle and then continue. Otherwise,
// drain only drains the bucket if it is due.
//
// This implementation is heavy-handed in that it brackets "leaking" tokens
// to the full duration of the configured interval. In other words, the
// bucket leaks not in single drops, but rather multiples, and only when the
// token drain window has elapsed. This side-steps near-hot-looping with
// dense token expiration (short interval + high size) and heavy lock
// contention. A possible enhancement would be to make this more granular.
func (b *bucket) drain(wait bool) {
	b.l.RLock()
	last := b.drained
	interval := b.opts.Interval
	b.l.RUnlock()

	switch {
	case time.Since(last) >= interval:
		b.l.Lock()
		defer b.l.Unlock()

		// Make sure the timestamp was not updated; prevents a time-of-
		// check vs. time-of-use error.
		if !b.drained.Equal(last) {
			return
		}

		// Drain the bucket.
		b.tokens = 0

		// Update the drain timestamp.
		b.drained = time.Now()

	case wait:
		delay := last.Add(interval).Sub(time.Now())
		time.Sleep(delay)
		b.drain(false)
	}
}

// setRate safely replaces the RateOpts on the bucket.
func (b *bucket) setRate(opts RateOpts) {
	b.l.Lock()
	b.opts = opts
	b.l.Unlock()
}
<|endoftext|>


<commit_before>package logged

import (
	"bytes"
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestInfo(t *testing.T) {
	var (
		buf bytes.Buffer
		l   = New(&Config{Writer: &buf})
		out entry

		msg  = "a test 123"
		data = Data{"test": "123", "test2": "1111"}
	)

	l.Info(msg, data)

	assert.NoError(t, json.NewDecoder(&buf).Decode(&out))
	assert.Equal(t, msg, out.Message)
	assert.Equal(t, Info, out.Level)
	assert.NotEmpty(t, out.Timestamp)
	assert.Equal(t, data, out.Data)
}

func TestDebug(t *testing.T) {
	var (
		buf bytes.Buffer
		l   = New(&Config{Writer: &buf})
		out entry

		msg  = "a test 123"
		data = Data{"test": "123", "test2": "1111"}
	)

	l.Debug(msg, data)

	assert.NoError(t, json.NewDecoder(&buf).Decode(&out))
	assert.Equal(t, msg, out.Message)
	assert.Equal(t, Debug, out.Level)
	assert.NotEmpty(t, out.Timestamp)
	assert.Equal(t, data, out.Data)
}

func TestMergedDataNoDefaults(t *testing.T) {
	var (
		data = Data{"one": "1"}
		l    = &log{}
	)

	assert.Nil(t, l.mergedData(nil))
	assert.Equal(t, data, l.mergedData(data))
}

func TestMergedDataNoData(t *testing.T) {
	var (
		defaults = Data{"one": "1"}
		l        = &log{defaults: defaults}
	)

	assert.Equal(t, defaults, l.mergedData(nil))
}

func TestMergedData(t *testing.T) {
	var (
		defaults = Data{"one": "1"}
		data     = Data{"two": "2"}
		out      = Data{"one": "1", "two": "2"}
		l        = &log{defaults: defaults}
	)

	assert.Equal(t, out, l.mergedData(data))
}

func TestMergedDataOverwriteDefaults(t *testing.T) {
	var (
		defaults = Data{"one": "1"}
		data     = Data{"one": "111"}
		l        = &log{defaults: defaults}
	)

	assert.Equal(t, data, l.mergedData(data))
}
<commit_msg>Fix tests<commit_after>package logged

import (
	"bytes"
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestInfo(t *testing.T) {
	var (
		buf bytes.Buffer
		l   = New(&Config{Writer: &buf})
		out entry

		msg  = "a test 123"
		data = Data{"test": "123", "test2": "1111"}
	)

	l.InfoEx(msg, data)

	assert.NoError(t, json.NewDecoder(&buf).Decode(&out))
	assert.Equal(t, msg, out.Message)
	assert.Equal(t, Info, out.Level)
	assert.NotEmpty(t, out.Timestamp)
	assert.Equal(t, data, out.Data)
}

func TestDebug(t *testing.T) {
	var (
		buf bytes.Buffer
		l   = New(&Config{Writer: &buf})
		out entry

		msg  = "a test 123"
		data = Data{"test": "123", "test2": "1111"}
	)

	l.DebugEx(msg, data)

	assert.NoError(t, json.NewDecoder(&buf).Decode(&out))
	assert.Equal(t, msg, out.Message)
	assert.Equal(t, Debug, out.Level)
	assert.NotEmpty(t, out.Timestamp)
	assert.Equal(t, data, out.Data)
}

func TestMergedDataNoDefaults(t *testing.T) {
	var (
		data = Data{"one": "1"}
		l    = &log{}
	)

	assert.Nil(t, l.mergedData(nil))
	assert.Equal(t, data, l.mergedData(data))
}

func TestMergedDataNoData(t *testing.T) {
	var (
		defaults = Data{"one": "1"}
		l        = &log{defaults: defaults}
	)

	assert.Equal(t, defaults, l.mergedData(nil))
}

func TestMergedData(t *testing.T) {
	var (
		defaults = Data{"one": "1"}
		data     = Data{"two": "2"}
		out      = Data{"one": "1", "two": "2"}
		l        = &log{defaults: defaults}
	)

	assert.Equal(t, out, l.mergedData(data))
}

func TestMergedDataOverwriteDefaults(t *testing.T) {
	var (
		defaults = Data{"one": "1"}
		data     = Data{"one": "111"}
		l        = &log{defaults: defaults}
	)

	assert.Equal(t, data, l.mergedData(data))
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"errors"
	"fmt"
	"reflect"
	"runtime"
	"testing"
	"time"
)

// Time returned in calling now(), setup and test in init.
var t1 = time.Date(2015, 9, 1, 14, 22, 36, 0, time.UTC)

func init() {
	margin := time.Millisecond
	t, t2 := now(), time.Now()
	if !t.Truncate(margin).Equal(t2.Truncate(margin)) {
		panic(fmt.Sprintf("now() doesn't return time.Now()! Expected %s, got %s",
			t2, t))
	}

	now = func() time.Time {
		return t1
	}
}

// A data type to be used in calling Log.
type user struct {
	ID   int
	Name string
}

// EventWriter that collects the events and errors.
type eventWriter struct {
	events []Event
	errors []error
	closed bool
}

func (ew *eventWriter) Write(event Event) error {
	ew.events = append(ew.events, event)
	return nil
}

func (ew *eventWriter) HandleError(err error) {
	ew.errors = append(ew.errors, err)
}

func (ew *eventWriter) Close() error {
	ew.closed = true
	return nil
}

func TestLog(t *testing.T) {
	defer reset()
	var ew eventWriter
	Start(&ew)

	tags := Tags{"my", "tags"}
	eventType := NewEventType("my-event-type")
	data := user{1, "Thomas"}
	event := Event{
		Type:    eventType,
		Tags:    tags,
		Message: "My event",
		Data:    data,
	}

	Debug(tags, "Debug message")
	Debugf(tags, "Debug %s message", "formatted")
	Info(tags, "Info message")
	Infof(tags, "Info %s message", "formatted")
	Warn(tags, "Warn message")
	Warnf(tags, "Warn %s message", "formatted")
	Error(tags, errors.New("Error message"))
	Errorf(tags, "Error %s message", "formatted")

	defer func() {
		recv := recover()
		Fatal(tags, recv)
		testThumstone(tags)
		Log(event)

		if err := Close(); err != nil {
			t.Fatal("Unexpected error closing: " + err.Error())
		}

		if len(ew.errors) != 0 {
			t.Fatalf("Unexpected error(s): %v", ew.errors)
		}

		pc, file, _, _ := runtime.Caller(0)
		fn := runtime.FuncForPC(pc)

		expected := []Event{
			Event{Type: DebugEvent, Tags: tags, Message: "Debug message"},
			Event{Type: DebugEvent, Tags: tags, Message: "Debug formatted message"},
			Event{Type: InfoEvent, Tags: tags, Message: "Info message"},
			Event{Type: InfoEvent, Tags: tags, Message: "Info formatted message"},
			Event{Type: WarnEvent, Tags: tags, Message: "Warn message"},
			Event{Type: WarnEvent, Tags: tags, Message: "Warn formatted message"},
			Event{Type: ErrorEvent, Tags: tags, Message: "Error message"},
			Event{Type: ErrorEvent, Tags: tags, Message: "Error formatted message"},
			Event{Type: FatalEvent, Tags: tags, Message: "Fatal message"},
			Event{Type: ThumbEvent, Tags: tags, Message: "Function testThumstone called by " + fn.Name() + ", from file " + file +
				" on line 87"},
			event,
		}

		if len(ew.events) != len(expected) {
			t.Fatalf("Expected to have %d events, but got %d",
				len(expected), len(ew.events))
		}

		for i, event := range ew.events {
			expectedEvent := expected[i]

			if expected, got := expectedEvent.Type, event.Type; expected != got {
				t.Errorf("Expected event #%d to have type %q, but got %q",
					i, expected, got)
			}

			if expected, got := t1, event.Timestamp; !expected.Equal(got) {
				t.Errorf("Expected event #%d to have timestamp %q, but got %q",
					i, expected, got)
			}

			if expected, got := expectedEvent.Tags, event.Tags; !reflect.DeepEqual(expected, got) {
				t.Errorf("Expected event #%d to have tags %q, but got %q",
					i, expected, got)
			}

			if expected, got := expectedEvent.Message, event.Message; expected != got {
				t.Errorf("Expected event #%d to have message %q, but got %q",
					i, expected, got)
			}

			// todo: test if we get a stacktrace with calling Fatal.
			if expectedEvent.Type == FatalEvent {
				event.Data = nil
			}

			if expected, got := expectedEvent.Data, event.Data; !reflect.DeepEqual(expected, got) {
				t.Errorf("Expected event #%d to have data %v, but got %v",
					i, expected, got)
			}
		}
	}()
	panic("Fatal message")
}

func TestStartTwice(t *testing.T) {
	defer reset()
	var ew eventWriter
	Start(&ew)
	if err := Close(); err != nil {
		t.Fatal("Unexpected error closing initial log: " + err.Error())
	}

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a second call to Start to panic, but it didn't")
		}

		const expected = "logger: can only Start once"
		got := interfaceToString(recv)
		if expected != got {
			t.Fatalf("Expected to panic with %q, but paniced with %q", expected, got)
		}
	}()
	Start(&ew)
}

func TestStartNoEventWriter(t *testing.T) {
	defer reset()
	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a call to Start without any EventWriters to panic, but it didn't")
		}

		const expected = "logger: need atleast a single EventWriter to write to"
		got := interfaceToString(recv)
		if expected != got {
			t.Fatalf("Expected to panic with %q, but paniced with %q", expected, got)
		}
	}()
	Start()
}

// EventWriter that always returns a write error with the event message in it.
type errorEventWriter struct {
	closeError error
	errors     []error
}

func (eew *errorEventWriter) Write(event Event) error {
	return errors.New("Write error: " + event.Message)
}

func (eew *errorEventWriter) HandleError(err error) {
	eew.errors = append(eew.errors, err)
}

func (eew *errorEventWriter) Close() error {
	return eew.closeError
}

func TestErrorEventWriter(t *testing.T) {
	closeError := errors.New("Close error")

	defer reset()
	eew := errorEventWriter{
		closeError: closeError,
	}
	Start(&eew)

	tags := Tags{"my", "tags"}
	Info(tags, "Info message1")
	Info(tags, "Info message2")
	Info(tags, "Info message3")
	Info(tags, "Info message4")
	Info(tags, "Info message5")
	Info(tags, "Won't be written to the writer")

	if err := Close(); err == nil {
		t.Fatal("Expected a closing error, but didn't get one")
	} else if err != closeError {
		t.Fatalf("Expceted the closing error to be %q, but got %q",
			closeError.Error(), err.Error())
	}

	// 10 = 5 bad write errors + 4 rewrite error of the first event + 1 bad
	// EventWriter error.
	if expected, got := 10, len(eew.errors); got != expected {
		t.Fatalf("Expected %d errors, but only got %d", expected, got)
	}

	// Expected errors:
	// 0 - write event 1
	// 1 - rewrite event 1
	// 2 - write event 2
	// 3 - rewrite event 1
	// 4 - write event 3
	// 5 - rewrite event 1
	// 6 - write event 4
	// 7 - rewrite event 1
	// 8 - write event 5
	// 9 - EventWriter is bad, more then 5 faulty writes, EventWriter will be dropped

	for i, got := range eew.errors {
		var expected error
		if i == 9 {
			expected = ErrBadEventWriter
		} else {
			d := 1
			switch i {
			case 2:
				d = 2
			case 4:
				d = 3
			case 6:
				d = 4
			case 8:
				d = 5
			}
			expected = fmt.Errorf("Write error: Info message%d", d)
		}

		if got.Error() != expected.Error() {
			t.Errorf("Expected error #%d to be %q, but got %q",
				i, expected.Error(), got.Error())
		}
	}
}

// todo: test with flaky writter, return nil, error, nil, error etc.
// todo: test with both good and bad Writer
// todo: test Thumbstone called from main, create a new sub proces and calling
// it and log to the stdout. Then parse the event and check it.

func testThumstone(tags Tags) {
	Thumbstone(tags, "testThumstone")
}

func reset() {
	eventChannel = make(chan Event, 1)
	eventChannelClosed = make(chan struct{}, 1)
	eventWriters = []EventWriter{}
	started = false
}
<commit_msg>Test stack trace of Fatal event<commit_after>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"bytes"
	"errors"
	"fmt"
	"reflect"
	"runtime"
	"testing"
	"time"
)

// Time returned in calling now(), setup and test in init.
var t1 = time.Date(2015, 9, 1, 14, 22, 36, 0, time.UTC)

func init() {
	margin := time.Millisecond
	t, t2 := now(), time.Now()
	if !t.Truncate(margin).Equal(t2.Truncate(margin)) {
		panic(fmt.Sprintf("now() doesn't return time.Now()! Expected %s, got %s",
			t2, t))
	}

	now = func() time.Time {
		return t1
	}
}

// A data type to be used in calling Log.
type user struct {
	ID   int
	Name string
}

// EventWriter that collects the events and errors.
type eventWriter struct {
	events []Event
	errors []error
	closed bool
}

func (ew *eventWriter) Write(event Event) error {
	ew.events = append(ew.events, event)
	return nil
}

func (ew *eventWriter) HandleError(err error) {
	ew.errors = append(ew.errors, err)
}

func (ew *eventWriter) Close() error {
	ew.closed = true
	return nil
}

func TestLog(t *testing.T) {
	defer reset()
	var ew eventWriter
	Start(&ew)

	tags := Tags{"my", "tags"}
	eventType := NewEventType("my-event-type")
	data := user{1, "Thomas"}
	event := Event{
		Type:    eventType,
		Tags:    tags,
		Message: "My event",
		Data:    data,
	}

	Debug(tags, "Debug message")
	Debugf(tags, "Debug %s message", "formatted")
	Info(tags, "Info message")
	Infof(tags, "Info %s message", "formatted")
	Warn(tags, "Warn message")
	Warnf(tags, "Warn %s message", "formatted")
	Error(tags, errors.New("Error message"))
	Errorf(tags, "Error %s message", "formatted")

	defer func() {
		recv := recover()
		Fatal(tags, recv)
		testThumstone(tags)
		Log(event)

		if err := Close(); err != nil {
			t.Fatal("Unexpected error closing: " + err.Error())
		}

		if len(ew.errors) != 0 {
			t.Fatalf("Unexpected error(s): %v", ew.errors)
		}

		pc, file, _, _ := runtime.Caller(0)
		fn := runtime.FuncForPC(pc)

		expected := []Event{
			Event{Type: DebugEvent, Tags: tags, Message: "Debug message"},
			Event{Type: DebugEvent, Tags: tags, Message: "Debug formatted message"},
			Event{Type: InfoEvent, Tags: tags, Message: "Info message"},
			Event{Type: InfoEvent, Tags: tags, Message: "Info formatted message"},
			Event{Type: WarnEvent, Tags: tags, Message: "Warn message"},
			Event{Type: WarnEvent, Tags: tags, Message: "Warn formatted message"},
			Event{Type: ErrorEvent, Tags: tags, Message: "Error message"},
			Event{Type: ErrorEvent, Tags: tags, Message: "Error formatted message"},
			Event{Type: FatalEvent, Tags: tags, Message: "Fatal message"},
			Event{Type: ThumbEvent, Tags: tags, Message: "Function testThumstone called by " + fn.Name() + ", from file " + file +
				" on line 87"},
			event,
		}

		if len(ew.events) != len(expected) {
			t.Fatalf("Expected to have %d events, but got %d",
				len(expected), len(ew.events))
		}

		for i, event := range ew.events {
			expectedEvent := expected[i]

			if expected, got := expectedEvent.Type, event.Type; expected != got {
				t.Errorf("Expected event #%d to have type %q, but got %q",
					i, expected, got)
			}

			if expected, got := t1, event.Timestamp; !expected.Equal(got) {
				t.Errorf("Expected event #%d to have timestamp %q, but got %q",
					i, expected, got)
			}

			if expected, got := expectedEvent.Tags, event.Tags; !reflect.DeepEqual(expected, got) {
				t.Errorf("Expected event #%d to have tags %q, but got %q",
					i, expected, got)
			}

			if expected, got := expectedEvent.Message, event.Message; expected != got {
				t.Errorf("Expected event #%d to have message %q, but got %q",
					i, expected, got)
			}

			if expectedEvent.Type == FatalEvent {
				// sortof test the stacktrace, best we can do.
				stacktrace := event.Data.([]byte)
				if !bytes.HasPrefix(stacktrace, []byte("goroutine")) {
					t.Errorf("Expected a stacktrace as data for a Fatal event, but got %s ",
						string(stacktrace))
				}
				event.Data = nil
			}

			if expected, got := expectedEvent.Data, event.Data; !reflect.DeepEqual(expected, got) {
				t.Errorf("Expected event #%d to have data %v, but got %v",
					i, expected, got)
			}
		}
	}()
	panic("Fatal message")
}

func TestStartTwice(t *testing.T) {
	defer reset()
	var ew eventWriter
	Start(&ew)
	if err := Close(); err != nil {
		t.Fatal("Unexpected error closing initial log: " + err.Error())
	}

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a second call to Start to panic, but it didn't")
		}

		const expected = "logger: can only Start once"
		got := interfaceToString(recv)
		if expected != got {
			t.Fatalf("Expected to panic with %q, but paniced with %q", expected, got)
		}
	}()
	Start(&ew)
}

func TestStartNoEventWriter(t *testing.T) {
	defer reset()
	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a call to Start without any EventWriters to panic, but it didn't")
		}

		const expected = "logger: need atleast a single EventWriter to write to"
		got := interfaceToString(recv)
		if expected != got {
			t.Fatalf("Expected to panic with %q, but paniced with %q", expected, got)
		}
	}()
	Start()
}

// EventWriter that always returns a write error with the event message in it.
type errorEventWriter struct {
	closeError error
	errors     []error
}

func (eew *errorEventWriter) Write(event Event) error {
	return errors.New("Write error: " + event.Message)
}

func (eew *errorEventWriter) HandleError(err error) {
	eew.errors = append(eew.errors, err)
}

func (eew *errorEventWriter) Close() error {
	return eew.closeError
}

func TestErrorEventWriter(t *testing.T) {
	closeError := errors.New("Close error")

	defer reset()
	eew := errorEventWriter{
		closeError: closeError,
	}
	Start(&eew)

	tags := Tags{"my", "tags"}
	Info(tags, "Info message1")
	Info(tags, "Info message2")
	Info(tags, "Info message3")
	Info(tags, "Info message4")
	Info(tags, "Info message5")
	Info(tags, "Won't be written to the writer")

	if err := Close(); err == nil {
		t.Fatal("Expected a closing error, but didn't get one")
	} else if err != closeError {
		t.Fatalf("Expceted the closing error to be %q, but got %q",
			closeError.Error(), err.Error())
	}

	// 10 = 5 bad write errors + 4 rewrite error of the first event + 1 bad
	// EventWriter error.
	if expected, got := 10, len(eew.errors); got != expected {
		t.Fatalf("Expected %d errors, but only got %d", expected, got)
	}

	// Expected errors:
	// 0 - write event 1
	// 1 - rewrite event 1
	// 2 - write event 2
	// 3 - rewrite event 1
	// 4 - write event 3
	// 5 - rewrite event 1
	// 6 - write event 4
	// 7 - rewrite event 1
	// 8 - write event 5
	// 9 - EventWriter is bad, more then 5 faulty writes, EventWriter will be dropped

	for i, got := range eew.errors {
		var expected error
		if i == 9 {
			expected = ErrBadEventWriter
		} else {
			d := 1
			switch i {
			case 2:
				d = 2
			case 4:
				d = 3
			case 6:
				d = 4
			case 8:
				d = 5
			}
			expected = fmt.Errorf("Write error: Info message%d", d)
		}

		if got.Error() != expected.Error() {
			t.Errorf("Expected error #%d to be %q, but got %q",
				i, expected.Error(), got.Error())
		}
	}
}

// todo: test with flaky writter, return nil, error, nil, error etc.
// todo: test with both good and bad Writer
// todo: test Thumbstone called from main, create a new sub proces and calling
// it and log to the stdout. Then parse the event and check it.

func testThumstone(tags Tags) {
	Thumbstone(tags, "testThumstone")
}

func reset() {
	eventChannel = make(chan Event, 1)
	eventChannelClosed = make(chan struct{}, 1)
	eventWriters = []EventWriter{}
	started = false
}
<|endoftext|>


<commit_before>package logstash

import (
	"encoding/json"
	"errors"
	"log"
	"net"

	"github.com/gliderlabs/logspout/router"
)

func init() {
	router.AdapterFactories.Register(NewLogstashAdapter, "logstash")
}

// LogstashAdapter is an adapter that streams UDP JSON to Logstash.
type LogstashAdapter struct {
	conn  net.Conn
	route *router.Route
}

// NewLogstashAdapter creates a LogstashAdapter with UDP as the default transport.
func NewLogstashAdapter(route *router.Route) (router.LogAdapter, error) {
	transport, found := router.AdapterTransports.Lookup(route.AdapterTransport("udp"))
	if !found {
		return nil, errors.New("unable to find adapter: " + route.Adapter)
	}

	conn, err := transport.Dial(route.Address, route.Options)
	if err != nil {
		return nil, err
	}

	return &LogstashAdapter{
		route: route,
		conn:  conn,
	}, nil
}

// Stream implements the router.LogAdapter interface.
func (a *LogstashAdapter) Stream(logstream chan *router.Message) {
	for m := range logstream {
		msg := LogstashMessage{
			Message:  m.Data,
			Name:     m.Container.Name,
			ID:       m.Container.ID,
			Image:    m.Container.Config.Image,
			Hostname: m.Container.Config.Hostname,
		}
		js, err := json.Marshal(msg)
		if err != nil {
			log.Println("logstash:", err)
			continue
		}
		_, err = a.conn.Write(js)
		if err != nil {
			log.Println("logstash:", err)
			continue
		}
	}
}

// LogstashMessage is a simple JSON input to Logstash.
type LogstashMessage struct {
	Message  string `json:"message"`
	Name     string `json:"container_name"`
	ID       string `json:"container_id"`
	Image    string `json:"image_name"`
	Hostname string `json:"host"`
}
<commit_msg>Use looplab/logspout-logstash#9 for streams.<commit_after>package logstash

import (
	"encoding/json"
	"errors"
	"log"
	"net"

	"github.com/gliderlabs/logspout/router"
)

func init() {
	router.AdapterFactories.Register(NewLogstashAdapter, "logstash")
}

// LogstashAdapter is an adapter that streams UDP JSON to Logstash.
type LogstashAdapter struct {
	conn  net.Conn
	route *router.Route
}

// NewLogstashAdapter creates a LogstashAdapter with UDP as the default transport.
func NewLogstashAdapter(route *router.Route) (router.LogAdapter, error) {
	transport, found := router.AdapterTransports.Lookup(route.AdapterTransport("udp"))
	if !found {
		return nil, errors.New("unable to find adapter: " + route.Adapter)
	}

	conn, err := transport.Dial(route.Address, route.Options)
	if err != nil {
		return nil, err
	}

	return &LogstashAdapter{
		route: route,
		conn:  conn,
	}, nil
}

// Stream implements the router.LogAdapter interface.
func (a *LogstashAdapter) Stream(logstream chan *router.Message) {
	for m := range logstream {
		msg := LogstashMessage{
			Message:  m.Data,
			Name:     m.Container.Name,
			ID:       m.Container.ID,
			Image:    m.Container.Config.Image,
			Hostname: m.Container.Config.Hostname,
			Stream:   m.Source,
		}
		js, err := json.Marshal(msg)
		if err != nil {
			log.Println("logstash:", err)
			continue
		}
		_, err = a.conn.Write(js)
		if err != nil {
			log.Println("logstash:", err)
			continue
		}
	}
}

// LogstashMessage is a simple JSON input to Logstash.
type LogstashMessage struct {
	Message  string `json:"message"`
	Name     string `json:"container_name"`
	ID       string `json:"container_id"`
	Image    string `json:"image_name"`
	Hostname string `json:"host"`
	Stream   string `json:"stream"`
}
<|endoftext|>


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

package poller

import (
	"context"
	"fmt"
	"strings"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/types/known/timestamppb"

	gerritutil "go.chromium.org/luci/common/api/gerrit"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gerritpb "go.chromium.org/luci/common/proto/gerrit"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/grpc/grpcutil"

	"go.chromium.org/luci/cv/internal/changelist"
	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/gerrit"
	"go.chromium.org/luci/cv/internal/gerrit/updater"
	"go.chromium.org/luci/cv/internal/prjmanager"
)

const (
	// fullPollInterval is between querying Gerrit for all changes relevant to CV as
	// if from scratch.
	fullPollInterval = time.Hour

	// incrementalPollOverlap is safety overlap of time range of Change.Updated
	// between two successive polls.
	//
	// While this doesn't guarantee that CV won't miss changes in-between
	// incremental polls, it should mitigate the most common reasons:
	//   * time skew between CV and Gerrit clocks,
	//   * hopping between potentially out of sync Gerrit mirrors.
	incrementalPollOverlap = time.Minute

	// changesPerPoll is how many changes CV will process per poll.
	//
	// A value that's too low here will first affect full polls, since they have
	// to (re-)process all interesting changes watched by a LUCI project.
	//
	// 10k is OK to fetch sequentially and keep in RAM without OOM-ing,
	// and currently enough for each of the LUCI projects.
	//
	// Higher values may need smarter full polling techniques.
	changesPerPoll = 10000

	// pageSize is how many changes to request in a single ListChangesRequest.
	pageSize = 1000

	// moreChangesTrustFactor controls when CV must not trust false value of
	// ListChangesResponse.MoreChanges.
	//
	// Value of 0.5 combined with pageSize of 1000 means that CV will trust
	// MoreChanges iff Gerrit returns <= 500 CLs.
	//
	// For more info, see corresponding field in
	// https://godoc.org/go.chromium.org/luci/common/api/gerrit#PagingListChangesOptions
	moreChangesTrustFactor = 0.5
)

// subpoll queries Gerrit and updates the state of individual SubPoller.
func subpoll(ctx context.Context, luciProject string, sp *SubPoller) error {
	q := singleQuery{
		luciProject: luciProject,
		sp:          sp,
	}
	var err error
	if q.client, err = gerrit.CurrentClient(ctx, sp.GetHost(), luciProject); err != nil {
		return err
	}
	if sp.GetLastFullTime() == nil {
		return q.full(ctx)
	}
	nextFullAt := sp.GetLastFullTime().AsTime().Add(fullPollInterval)
	if clock.Now(ctx).Before(nextFullAt) {
		return q.incremental(ctx)
	}
	return q.full(ctx)
}

type singleQuery struct {
	luciProject string
	sp          *SubPoller
	client      gerrit.QueryClient
}

func (q *singleQuery) full(ctx context.Context) error {
	ctx = logging.SetFields(ctx, logging.Fields{
		"luciProject": q.luciProject,
		"poll":        "full",
	})
	started := clock.Now(ctx)
	after := started.Add(-common.MaxTriggerAge)
	changes, err := q.fetch(ctx, after, buildQuery(q.sp, queryLimited))
	// There can be partial result even if err != nil.
	switch err2 := q.scheduleTasks(ctx, changes, true); {
	case err != nil:
		return err
	case err2 != nil:
		return err2
	}

	cur := uniqueSortedIDsOf(changes)
	if diff := common.DifferenceSorted(q.sp.Changes, cur); len(diff) != 0 {
		// `diff` changes are no longer matching the limited query,
		// so they probably updated since.
		if err := scheduleRefreshTasks(ctx, q.luciProject, q.sp.GetHost(), diff); err != nil {
			return err
		}
	}

	q.sp.Changes = cur
	q.sp.LastFullTime = timestamppb.New(started)
	q.sp.LastIncrTime = nil
	return nil
}

func (q *singleQuery) incremental(ctx context.Context) error {
	ctx = logging.SetFields(ctx, logging.Fields{
		"luciProject": q.luciProject,
		"poll":        "incremental",
	})
	started := clock.Now(ctx)

	lastInc := q.sp.GetLastIncrTime()
	if lastInc == nil {
		if lastInc = q.sp.GetLastFullTime(); lastInc == nil {
			panic("must have been a full poll")
		}
	}
	after := lastInc.AsTime().Add(-incrementalPollOverlap)
	// Unlike the full poll, query for all changes regardless of status or CQ
	// vote. This ensures that CV notices quickly when previously NEW & CQ-ed
	// change has either CQ vote removed OR status changed (e.g. submitted or
	// abandoned).
	changes, err := q.fetch(ctx, after, buildQuery(q.sp, queryAll))
	// There can be partial result even if err != nil.
	switch err2 := q.scheduleTasks(ctx, changes, false); {
	case err != nil:
		return err
	case err2 != nil:
		return err2
	}

	q.sp.Changes = common.UnionSorted(q.sp.Changes, uniqueSortedIDsOf(changes))
	q.sp.LastIncrTime = timestamppb.New(started)
	return nil
}

func (q *singleQuery) fetch(ctx context.Context, after time.Time, query string) ([]*gerritpb.ChangeInfo, error) {
	opts := gerritutil.PagingListChangesOptions{
		Limit:                  changesPerPoll,
		PageSize:               pageSize,
		MoreChangesTrustFactor: moreChangesTrustFactor,
		UpdatedAfter:           after,
	}
	req := gerritpb.ListChangesRequest{
		Options: []gerritpb.QueryOption{
			gerritpb.QueryOption_SKIP_MERGEABLE,
		},
		Query: query,
	}
	resp, err := gerritutil.PagingListChanges(ctx, q.client, &req, opts)
	switch grpcutil.Code(err) {
	case codes.OK:
		if resp.GetMoreChanges() {
			logging.Errorf(ctx, "Ignoring oldest changes because reached max (%d) allowed to process per poll", changesPerPoll)
		}
		return resp.GetChanges(), nil
	// TODO(tandrii): handle 403 and 404 if CV lacks access to entire host.
	default:
		// NOTE: resp may be set if there was partial success in fetching changes
		// followed by a typically transient error.
		return resp.GetChanges(), gerrit.UnhandledError(ctx, err, "PagingListChanges failed")
	}
}

// maxLoadCLBatchSize limits how many CL entities are loaded at once for
// notifying PM.
const maxLoadCLBatchSize = 100

func (q *singleQuery) scheduleTasks(ctx context.Context, changes []*gerritpb.ChangeInfo, forceNotifyPM bool) error {
	// TODO(tandrii): optimize by checking if CV is interested in the
	// (host,project,ref) these changes come from before triggering tasks.
	logging.Debugf(ctx, "scheduling %d CLUpdate tasks", len(changes))

	var clids []common.CLID
	if forceNotifyPM {
		// Objective: make PM aware of all CLs.
		// Optimization: avoid RefreshGerritCL with forceNotifyPM=true if possible,
		// as this removes ability to de-duplicate.
		var err error
		eids := make([]changelist.ExternalID, len(changes))
		for i, c := range changes {
			if eids[i], err = changelist.GobID(q.sp.GetHost(), c.GetNumber()); err != nil {
				return err
			}
		}
		clids, err = changelist.Lookup(ctx, eids)
		if err != nil {
			return err
		}
		if err := notifyPMifKnown(ctx, q.luciProject, clids, maxLoadCLBatchSize); err != nil {
			return err
		}
	}

	errs := parallel.WorkPool(10, func(work chan<- func() error) {
		for i, c := range changes {
			payload := &updater.RefreshGerritCL{
				LuciProject:   q.luciProject,
				Host:          q.sp.GetHost(),
				Change:        c.GetNumber(),
				UpdatedHint:   c.GetUpdated(),
				ForceNotifyPm: forceNotifyPM && (clids[i] == 0),
			}
			work <- func() error {
				return updater.Schedule(ctx, payload)
			}
		}
	})
	return common.MostSevereError(errs)
}

func scheduleRefreshTasks(ctx context.Context, luciProject, host string, changes []int64) error {
	logging.Debugf(ctx, "scheduling %d CLUpdate tasks for no longer matched CLs", len(changes))
	var err error
	eids := make([]changelist.ExternalID, len(changes))
	for i, c := range changes {
		if eids[i], err = changelist.GobID(host, c); err != nil {
			return err
		}
	}
	// Objective: make PM aware of all CLs.
	// Optimization: avoid RefreshGerritCL with forceNotifyPM=true if possible,
	// as this removes ability to de-duplicate.
	// Since all no longer matched CLs are typically already stored in the
	// Datastore, get their internal CL IDs and notify PM directly.
	clids, err := changelist.Lookup(ctx, eids)
	if err != nil {
		return err
	}

	if err := notifyPMifKnown(ctx, luciProject, clids, maxLoadCLBatchSize); err != nil {
		return err
	}

	errs := parallel.WorkPool(10, func(work chan<- func() error) {
		for i, c := range changes {
			payload := &updater.RefreshGerritCL{
				LuciProject:   luciProject,
				Host:          host,
				Change:        c,
				ForceNotifyPm: 0 == clids[i], // notify iff CL ID isn't yet known.
			}
			work <- func() error {
				return updater.Schedule(ctx, payload)
			}
		}
	})
	return common.MostSevereError(errs)
}

// notifyPMifKnown notifies PM to update its CLs for each non-0 CLID.
//
// Obtains EVersion of each CL before notify a PM. Unfortunately, this loads a
// lot of information we don't need, such as Snapshot. So, loads CLs in batches
// to avoid large memory footprint.
//
// In practice, most of these CLs would be already dscache-ed, so loading them
// is fast.
func notifyPMifKnown(ctx context.Context, luciProject string, clids []common.CLID, maxBatchSize int) error {
	cls := make([]*changelist.CL, 0, maxBatchSize)
	flush := func() error {
		if err := datastore.Get(ctx, cls); err != nil {
			return errors.Annotate(common.MostSevereError(err), "failed to load CLs").Tag(transient.Tag).Err()
		}
		return prjmanager.NotifyCLsUpdated(ctx, luciProject, cls)
	}

	for _, clid := range clids {
		switch l := len(cls); {
		case clid == 0:
		case l == maxBatchSize:
			if err := flush(); err != nil {
				return err
			}
			cls = cls[:0]
		default:
			cls = append(cls, &changelist.CL{ID: clid})
		}
	}
	if len(cls) > 0 {
		return flush()
	}
	return nil
}

type queryKind int

const (
	queryLimited queryKind = iota
	queryAll
)

// buildQuery returns query string.
//
// If queryLimited, unlike queryAll, searches for NEW CLs with CQ vote.
func buildQuery(sp *SubPoller, kind queryKind) string {
	buf := strings.Builder{}
	switch kind {
	case queryLimited:
		buf.WriteString("status:NEW ")
		// TODO(tandrii): make label optional to support Tricium use-case.
		buf.WriteString("label:Commit-Queue>0 ")
	case queryAll:
	default:
		panic(fmt.Errorf("unknown queryKind %d", kind))
	}
	// TODO(crbug/1163177): specify `branch:` search term to restrict search to
	// specific refs. This requires changing partitioning poller into subpollers,
	// but will provide more targeted queries, reducing load on CV & Gerrit.

	emitProjectValue := func(p string) {
		// Even though it appears to work without, Gerrit doc says project names
		// containing / must be surrounded by "" or {}:
		// https://gerrit-review.googlesource.com/Documentation/user-search.html#_argument_quoting
		buf.WriteRune('"')
		buf.WriteString(p)
		buf.WriteRune('"')
	}

	// One of .OrProjects or .CommonProjectPrefix must be set.
	switch prs := sp.GetOrProjects(); len(prs) {
	case 0:
		if sp.GetCommonProjectPrefix() == "" {
			panic("partitionConfig function should have ensured this")
		}
		// project*s* means find matching projects by prefix
		buf.WriteString("projects:")
		emitProjectValue(sp.GetCommonProjectPrefix())
	case 1:
		buf.WriteString("project:")
		emitProjectValue(prs[0])
	default:
		buf.WriteRune('(')
		for i, p := range prs {
			if i > 0 {
				buf.WriteString(" OR ")
			}
			buf.WriteString("project:")
			emitProjectValue(p)
		}
		buf.WriteRune(')')
	}
	return buf.String()
}

func uniqueSortedIDsOf(changes []*gerritpb.ChangeInfo) []int64 {
	if len(changes) == 0 {
		return nil
	}

	out := make([]int64, len(changes))
	for i, c := range changes {
		out[i] = c.GetNumber()
	}
	return common.UniqueSorted(out)
}
<commit_msg>cv: reduce full poll interval by half to 30 minutes.<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package poller

import (
	"context"
	"fmt"
	"strings"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/types/known/timestamppb"

	gerritutil "go.chromium.org/luci/common/api/gerrit"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	gerritpb "go.chromium.org/luci/common/proto/gerrit"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/grpc/grpcutil"

	"go.chromium.org/luci/cv/internal/changelist"
	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/gerrit"
	"go.chromium.org/luci/cv/internal/gerrit/updater"
	"go.chromium.org/luci/cv/internal/prjmanager"
)

const (
	// fullPollInterval is between querying Gerrit for all changes relevant to CV as
	// if from scratch.
	fullPollInterval = 30 * time.Minute

	// incrementalPollOverlap is safety overlap of time range of Change.Updated
	// between two successive polls.
	//
	// While this doesn't guarantee that CV won't miss changes in-between
	// incremental polls, it should mitigate the most common reasons:
	//   * time skew between CV and Gerrit clocks,
	//   * hopping between potentially out of sync Gerrit mirrors.
	incrementalPollOverlap = time.Minute

	// changesPerPoll is how many changes CV will process per poll.
	//
	// A value that's too low here will first affect full polls, since they have
	// to (re-)process all interesting changes watched by a LUCI project.
	//
	// 10k is OK to fetch sequentially and keep in RAM without OOM-ing,
	// and currently enough for each of the LUCI projects.
	//
	// Higher values may need smarter full polling techniques.
	changesPerPoll = 10000

	// pageSize is how many changes to request in a single ListChangesRequest.
	pageSize = 1000

	// moreChangesTrustFactor controls when CV must not trust false value of
	// ListChangesResponse.MoreChanges.
	//
	// Value of 0.5 combined with pageSize of 1000 means that CV will trust
	// MoreChanges iff Gerrit returns <= 500 CLs.
	//
	// For more info, see corresponding field in
	// https://godoc.org/go.chromium.org/luci/common/api/gerrit#PagingListChangesOptions
	moreChangesTrustFactor = 0.5
)

// subpoll queries Gerrit and updates the state of individual SubPoller.
func subpoll(ctx context.Context, luciProject string, sp *SubPoller) error {
	q := singleQuery{
		luciProject: luciProject,
		sp:          sp,
	}
	var err error
	if q.client, err = gerrit.CurrentClient(ctx, sp.GetHost(), luciProject); err != nil {
		return err
	}
	if sp.GetLastFullTime() == nil {
		return q.full(ctx)
	}
	nextFullAt := sp.GetLastFullTime().AsTime().Add(fullPollInterval)
	if clock.Now(ctx).Before(nextFullAt) {
		return q.incremental(ctx)
	}
	return q.full(ctx)
}

type singleQuery struct {
	luciProject string
	sp          *SubPoller
	client      gerrit.QueryClient
}

func (q *singleQuery) full(ctx context.Context) error {
	ctx = logging.SetFields(ctx, logging.Fields{
		"luciProject": q.luciProject,
		"poll":        "full",
	})
	started := clock.Now(ctx)
	after := started.Add(-common.MaxTriggerAge)
	changes, err := q.fetch(ctx, after, buildQuery(q.sp, queryLimited))
	// There can be partial result even if err != nil.
	switch err2 := q.scheduleTasks(ctx, changes, true); {
	case err != nil:
		return err
	case err2 != nil:
		return err2
	}

	cur := uniqueSortedIDsOf(changes)
	if diff := common.DifferenceSorted(q.sp.Changes, cur); len(diff) != 0 {
		// `diff` changes are no longer matching the limited query,
		// so they probably updated since.
		if err := scheduleRefreshTasks(ctx, q.luciProject, q.sp.GetHost(), diff); err != nil {
			return err
		}
	}

	q.sp.Changes = cur
	q.sp.LastFullTime = timestamppb.New(started)
	q.sp.LastIncrTime = nil
	return nil
}

func (q *singleQuery) incremental(ctx context.Context) error {
	ctx = logging.SetFields(ctx, logging.Fields{
		"luciProject": q.luciProject,
		"poll":        "incremental",
	})
	started := clock.Now(ctx)

	lastInc := q.sp.GetLastIncrTime()
	if lastInc == nil {
		if lastInc = q.sp.GetLastFullTime(); lastInc == nil {
			panic("must have been a full poll")
		}
	}
	after := lastInc.AsTime().Add(-incrementalPollOverlap)
	// Unlike the full poll, query for all changes regardless of status or CQ
	// vote. This ensures that CV notices quickly when previously NEW & CQ-ed
	// change has either CQ vote removed OR status changed (e.g. submitted or
	// abandoned).
	changes, err := q.fetch(ctx, after, buildQuery(q.sp, queryAll))
	// There can be partial result even if err != nil.
	switch err2 := q.scheduleTasks(ctx, changes, false); {
	case err != nil:
		return err
	case err2 != nil:
		return err2
	}

	q.sp.Changes = common.UnionSorted(q.sp.Changes, uniqueSortedIDsOf(changes))
	q.sp.LastIncrTime = timestamppb.New(started)
	return nil
}

func (q *singleQuery) fetch(ctx context.Context, after time.Time, query string) ([]*gerritpb.ChangeInfo, error) {
	opts := gerritutil.PagingListChangesOptions{
		Limit:                  changesPerPoll,
		PageSize:               pageSize,
		MoreChangesTrustFactor: moreChangesTrustFactor,
		UpdatedAfter:           after,
	}
	req := gerritpb.ListChangesRequest{
		Options: []gerritpb.QueryOption{
			gerritpb.QueryOption_SKIP_MERGEABLE,
		},
		Query: query,
	}
	resp, err := gerritutil.PagingListChanges(ctx, q.client, &req, opts)
	switch grpcutil.Code(err) {
	case codes.OK:
		if resp.GetMoreChanges() {
			logging.Errorf(ctx, "Ignoring oldest changes because reached max (%d) allowed to process per poll", changesPerPoll)
		}
		return resp.GetChanges(), nil
	// TODO(tandrii): handle 403 and 404 if CV lacks access to entire host.
	default:
		// NOTE: resp may be set if there was partial success in fetching changes
		// followed by a typically transient error.
		return resp.GetChanges(), gerrit.UnhandledError(ctx, err, "PagingListChanges failed")
	}
}

// maxLoadCLBatchSize limits how many CL entities are loaded at once for
// notifying PM.
const maxLoadCLBatchSize = 100

func (q *singleQuery) scheduleTasks(ctx context.Context, changes []*gerritpb.ChangeInfo, forceNotifyPM bool) error {
	// TODO(tandrii): optimize by checking if CV is interested in the
	// (host,project,ref) these changes come from before triggering tasks.
	logging.Debugf(ctx, "scheduling %d CLUpdate tasks", len(changes))

	var clids []common.CLID
	if forceNotifyPM {
		// Objective: make PM aware of all CLs.
		// Optimization: avoid RefreshGerritCL with forceNotifyPM=true if possible,
		// as this removes ability to de-duplicate.
		var err error
		eids := make([]changelist.ExternalID, len(changes))
		for i, c := range changes {
			if eids[i], err = changelist.GobID(q.sp.GetHost(), c.GetNumber()); err != nil {
				return err
			}
		}
		clids, err = changelist.Lookup(ctx, eids)
		if err != nil {
			return err
		}
		if err := notifyPMifKnown(ctx, q.luciProject, clids, maxLoadCLBatchSize); err != nil {
			return err
		}
	}

	errs := parallel.WorkPool(10, func(work chan<- func() error) {
		for i, c := range changes {
			payload := &updater.RefreshGerritCL{
				LuciProject:   q.luciProject,
				Host:          q.sp.GetHost(),
				Change:        c.GetNumber(),
				UpdatedHint:   c.GetUpdated(),
				ForceNotifyPm: forceNotifyPM && (clids[i] == 0),
			}
			work <- func() error {
				return updater.Schedule(ctx, payload)
			}
		}
	})
	return common.MostSevereError(errs)
}

func scheduleRefreshTasks(ctx context.Context, luciProject, host string, changes []int64) error {
	logging.Debugf(ctx, "scheduling %d CLUpdate tasks for no longer matched CLs", len(changes))
	var err error
	eids := make([]changelist.ExternalID, len(changes))
	for i, c := range changes {
		if eids[i], err = changelist.GobID(host, c); err != nil {
			return err
		}
	}
	// Objective: make PM aware of all CLs.
	// Optimization: avoid RefreshGerritCL with forceNotifyPM=true if possible,
	// as this removes ability to de-duplicate.
	// Since all no longer matched CLs are typically already stored in the
	// Datastore, get their internal CL IDs and notify PM directly.
	clids, err := changelist.Lookup(ctx, eids)
	if err != nil {
		return err
	}

	if err := notifyPMifKnown(ctx, luciProject, clids, maxLoadCLBatchSize); err != nil {
		return err
	}

	errs := parallel.WorkPool(10, func(work chan<- func() error) {
		for i, c := range changes {
			payload := &updater.RefreshGerritCL{
				LuciProject:   luciProject,
				Host:          host,
				Change:        c,
				ForceNotifyPm: 0 == clids[i], // notify iff CL ID isn't yet known.
			}
			work <- func() error {
				return updater.Schedule(ctx, payload)
			}
		}
	})
	return common.MostSevereError(errs)
}

// notifyPMifKnown notifies PM to update its CLs for each non-0 CLID.
//
// Obtains EVersion of each CL before notify a PM. Unfortunately, this loads a
// lot of information we don't need, such as Snapshot. So, loads CLs in batches
// to avoid large memory footprint.
//
// In practice, most of these CLs would be already dscache-ed, so loading them
// is fast.
func notifyPMifKnown(ctx context.Context, luciProject string, clids []common.CLID, maxBatchSize int) error {
	cls := make([]*changelist.CL, 0, maxBatchSize)
	flush := func() error {
		if err := datastore.Get(ctx, cls); err != nil {
			return errors.Annotate(common.MostSevereError(err), "failed to load CLs").Tag(transient.Tag).Err()
		}
		return prjmanager.NotifyCLsUpdated(ctx, luciProject, cls)
	}

	for _, clid := range clids {
		switch l := len(cls); {
		case clid == 0:
		case l == maxBatchSize:
			if err := flush(); err != nil {
				return err
			}
			cls = cls[:0]
		default:
			cls = append(cls, &changelist.CL{ID: clid})
		}
	}
	if len(cls) > 0 {
		return flush()
	}
	return nil
}

type queryKind int

const (
	queryLimited queryKind = iota
	queryAll
)

// buildQuery returns query string.
//
// If queryLimited, unlike queryAll, searches for NEW CLs with CQ vote.
func buildQuery(sp *SubPoller, kind queryKind) string {
	buf := strings.Builder{}
	switch kind {
	case queryLimited:
		buf.WriteString("status:NEW ")
		// TODO(tandrii): make label optional to support Tricium use-case.
		buf.WriteString("label:Commit-Queue>0 ")
	case queryAll:
	default:
		panic(fmt.Errorf("unknown queryKind %d", kind))
	}
	// TODO(crbug/1163177): specify `branch:` search term to restrict search to
	// specific refs. This requires changing partitioning poller into subpollers,
	// but will provide more targeted queries, reducing load on CV & Gerrit.

	emitProjectValue := func(p string) {
		// Even though it appears to work without, Gerrit doc says project names
		// containing / must be surrounded by "" or {}:
		// https://gerrit-review.googlesource.com/Documentation/user-search.html#_argument_quoting
		buf.WriteRune('"')
		buf.WriteString(p)
		buf.WriteRune('"')
	}

	// One of .OrProjects or .CommonProjectPrefix must be set.
	switch prs := sp.GetOrProjects(); len(prs) {
	case 0:
		if sp.GetCommonProjectPrefix() == "" {
			panic("partitionConfig function should have ensured this")
		}
		// project*s* means find matching projects by prefix
		buf.WriteString("projects:")
		emitProjectValue(sp.GetCommonProjectPrefix())
	case 1:
		buf.WriteString("project:")
		emitProjectValue(prs[0])
	default:
		buf.WriteRune('(')
		for i, p := range prs {
			if i > 0 {
				buf.WriteString(" OR ")
			}
			buf.WriteString("project:")
			emitProjectValue(p)
		}
		buf.WriteRune(')')
	}
	return buf.String()
}

func uniqueSortedIDsOf(changes []*gerritpb.ChangeInfo) []int64 {
	if len(changes) == 0 {
		return nil
	}

	out := make([]int64, len(changes))
	for i, c := range changes {
		out[i] = c.GetNumber()
	}
	return common.UniqueSorted(out)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/miquella/ask"

	"github.com/miquella/vaulted/lib"
)

type Spawn struct {
	SessionOptions

	Command       []string
	DisplayStatus bool
}

func (s *Spawn) Run(store vaulted.Store) error {
	session, err := GetSessionWithOptions(store, &s.SessionOptions)
	if err != nil {
		return err
	}

	timeRemaining := session.Expiration.Sub(time.Now())
	timeRemaining = time.Second * time.Duration(timeRemaining.Seconds())
	if s.DisplayStatus {
		ask.Print(fmt.Sprintf("%s — expires: %s (%s remaining)\n", session.Name, session.Expiration.Format("2 Jan 2006 15:04 MST"), timeRemaining))
	}

	code, err := session.Spawn(s.Command)
	if err != nil {
		return ErrorWithExitCode{err, 2}
	} else if *code != 0 {
		return ErrorWithExitCode{ErrNoError, *code}
	}

	return nil
}
<commit_msg>Replace non-breaking space with regular space<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/miquella/ask"

	"github.com/miquella/vaulted/lib"
)

type Spawn struct {
	SessionOptions

	Command       []string
	DisplayStatus bool
}

func (s *Spawn) Run(store vaulted.Store) error {
	session, err := GetSessionWithOptions(store, &s.SessionOptions)
	if err != nil {
		return err
	}

	timeRemaining := session.Expiration.Sub(time.Now())
	timeRemaining = time.Second * time.Duration(timeRemaining.Seconds())
	if s.DisplayStatus {
		ask.Print(fmt.Sprintf("%s — expires: %s (%s remaining)\n", session.Name, session.Expiration.Format("2 Jan 2006 15:04 MST"), timeRemaining))
	}

	code, err := session.Spawn(s.Command)
	if err != nil {
		return ErrorWithExitCode{err, 2}
	} else if *code != 0 {
		return ErrorWithExitCode{ErrNoError, *code}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	// std
	"encoding/json"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
	// base
	"github.com/Preetam/onecontactlink/internal-api/client"
	"github.com/Preetam/onecontactlink/schema"
	"github.com/Preetam/onecontactlink/web/linktoken"
	// vendor
	"github.com/VividCortex/siesta"
)

func serveGetRequest(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	linkStr := params.String("link", "", "link code")
	err := params.Parse(r.Form)

	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid OneContactLink",
		})
		return
	}

	if err != nil || !strings.Contains(*linkStr, "-") {
		invalidLink()
		return
	}
	parts := strings.Split(*linkStr, "-")
	if len(parts) != 2 {
		invalidLink()
		return
	}

	// get request link
	requestLink, err := internalAPIClient.GetRequestLinkByCode(parts[1])
	if err != nil {
		invalidLink()
		return
	}

	// get user
	user, err := internalAPIClient.GetUser(requestLink.User)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong",
		})
		return
	}

	if user.Code != parts[0] {
		invalidLink()
		return
	}

	templ.ExecuteTemplate(w, "request", map[string]string{
		"Name": user.Name,
	})
}

func servePostRequest(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	nameStr := params.String("name", "", "name")
	emailStr := params.String("email", "", "email")
	linkStr := params.String("link", "", "link code")
	err := params.Parse(r.Form)

	invalidLink := func() {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid OneContactLink",
		})
		return
	}

	if err != nil || !strings.Contains(*linkStr, "-") {
		invalidLink()
	}
	parts := strings.Split(*linkStr, "-")
	if len(parts) != 2 {
		invalidLink()
	}

	// get request link
	requestLink, err := internalAPIClient.GetRequestLinkByCode(parts[1])
	if err != nil {
		invalidLink()
		return
	}

	// get user
	user, err := internalAPIClient.GetUser(requestLink.User)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "Something went wrong",
		})
		return
	}

	if user.Code != parts[0] {
		invalidLink()
		return
	}

	if *nameStr == "" || *emailStr == "" {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "You must provide a name and email address",
		})
		return
	}

	toUser := user.ID
	fromUser := 0

	email, err := internalAPIClient.GetEmail(*emailStr)
	if err != nil {
		if err == client.ErrNotFound {
			// Email doesn't exist. Create a new user with the email.
			user, err := internalAPIClient.CreateUser(schema.NewUser(*nameStr, *emailStr))
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				templ.ExecuteTemplate(w, "request", map[string]string{
					"Name":    user.Name,
					"Warning": "Something went wrong. Please try again.",
				})
				return
			}
			fromUser = user.ID
		} else {
			w.WriteHeader(http.StatusInternalServerError)
			templ.ExecuteTemplate(w, "request", map[string]string{
				"Name":    user.Name,
				"Warning": "Something went wrong. Please try again.",
			})
			return
		}
	} else {
		fromUser = email.User
	}

	requestID, err := internalAPIClient.CreateRequest(fromUser, toUser)
	if err != nil {
		if err == client.ErrConflict {
			templ.ExecuteTemplate(w, "request", map[string]string{
				"Name": user.Name,
				"Info": "Looks like you already made this request. If " + user.Name +
					" has already approved your request, we'll send you their latest contact info.",
			})
			// Try to send another request email. This is idempotent.
			internalAPIClient.SendRequestEmail(requestID)
			internalAPIClient.SendContactInfoEmail(requestID)
			return
		}
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Name":    user.Name,
			"Warning": "Something went wrong. Please try again.",
		})
		return
	}
	err = internalAPIClient.SendRequestEmail(requestID)
	if err != nil {
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Name":    user.Name,
			"Warning": "Something went wrong. Please try again.",
		})
		return
	}

	templ.ExecuteTemplate(w, "success", map[string]string{
		"Success": "Request sent!",
	})
}

func serveManageRequest(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	linkStr := params.String("link", "", "link token")
	actionStr := params.String("action", "", "action")
	err := params.Parse(r.Form)

	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid link",
		})
		return
	}

	if err != nil {
		invalidLink()
		return
	}

	linkToken, err := tokenCodec.DecodeToken(*linkStr, new(linktoken.RequestTokenData))
	if err != nil {
		invalidLink()
		return
	}

	if *actionStr != "approve" && *actionStr != "reject" {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Invalid action.",
		})
		return
	}

	// Check if token expired
	if linkToken.Expires <= int(time.Now().Unix()) {
		// Token expired.
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "This link has expired.",
		})
		return
	}

	// extract request ID
	requestID := linkToken.Data.(*linktoken.RequestTokenData).Request
	err = internalAPIClient.ManageRequest(requestID, *actionStr)
	if err != nil {
		if serverErr, ok := err.(client.ServerError); ok {
			if int(serverErr) == http.StatusConflict {
				templ.ExecuteTemplate(w, "invalid", map[string]string{
					"Warning": "Oops, you already responded to this request. We can't change" +
						" anything right now.",
				})
				return
			}
		}

		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		log.Println(err)
		return

	}

	if *actionStr == "approve" {
		err = internalAPIClient.SendContactInfoEmail(requestID)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			templ.ExecuteTemplate(w, "invalid", map[string]string{
				"Error": "Something went wrong. Please try again.",
			})
			log.Println(err)
			return
		}

		templ.ExecuteTemplate(w, "success", map[string]string{
			"Success": "Approved! We'll send them an email with your contact information.",
		})
	} else {
		templ.ExecuteTemplate(w, "success", map[string]string{
			"Success": "Rejected. That email won't be able to send you any more requests.",
		})
	}
}

func serveAuth(c siesta.Context, w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	linkStr := params.String("link", "", "link token")
	err := params.Parse(r.Form)
	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid link",
		})
		return
	}
	if err != nil {
		invalidLink()
		return
	}

	linkToken, err := tokenCodec.DecodeToken(*linkStr, new(linktoken.UserTokenData))
	if err != nil {
		invalidLink()
		return
	}

	// Check if token expired
	if linkToken.Expires <= int(time.Now().Unix()) {
		// Token expired.
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "This link has expired.",
		})
		return
	}

	// extract user ID
	userID := linkToken.Data.(*linktoken.UserTokenData).User

	// get user information
	_, err = internalAPIClient.GetUser(userID)
	if err != nil {
		// Token expired.
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	// Update the expiration and set a cookie
	linkToken.Expires = int(time.Now().Unix() + 86400)

	token, err := tokenCodec.EncodeToken(linkToken)
	if err != nil {
		// Token expired.
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     "ocl",
		Value:    token,
		Domain:   CookieDomain,
		Path:     "/",
		HttpOnly: true,
		Secure:   !DevMode,
	})

	w.Header().Add("Refresh", "2; /app")

	templ.ExecuteTemplate(w, "success", map[string]string{
		"Success": "Logged in! Redirecting you to the app...",
	})
}

func servePostLogin(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	emailStr := params.String("email", "", "email")
	err := params.Parse(r.Form)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "Invalid parameters.",
		})
		return
	}

	internalAPIClient.SendAuth(*emailStr)
	templ.ExecuteTemplate(w, "success", map[string]string{
		"Info": "We've sent a login link to '" + *emailStr +
			"' if it's associated with a valid account.",
	})
}

func serveDevModeAuth(c siesta.Context, w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	userID := params.Int("user", 0, "user ID")
	err := params.Parse(r.Form)
	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid link",
		})
		return
	}
	if err != nil {
		invalidLink()
		return
	}

	linkToken := linktoken.NewLinkToken(&linktoken.UserTokenData{
		User: *userID,
	}, int(time.Now().Unix()+86400))

	// get user information
	_, err = internalAPIClient.GetUser(*userID)
	if err != nil {
		log.Println(err)
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	token, err := tokenCodec.EncodeToken(linkToken)
	if err != nil {
		log.Println(err)
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     "ocl",
		Value:    token,
		Domain:   CookieDomain,
		Path:     "/",
		HttpOnly: true,
		Secure:   !DevMode,
	})

	templ.ExecuteTemplate(w, "success", map[string]string{
		"Success": "Logged in!",
	})
}

func verifyCaptcha(c siesta.Context, r *http.Request, w http.ResponseWriter, q func()) {
	if DevMode {
		return
	}

	params := &siesta.Params{}
	recaptchaResponse := params.String("g-recaptcha-response", "", "reCAPTCHA response")
	err := params.Parse(r.Form)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "Missing CAPTCHA parameter.",
		})
		q()
		return
	}

	if *recaptchaResponse == "" {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Error": "Bad CAPTCHA",
		})
		q()
		return
	}

	// verify CAPTCHA
	resp, err := http.PostForm("https://www.google.com/recaptcha/api/siteverify", url.Values{
		"secret":   []string{RecaptchaSecret},
		"response": []string{*recaptchaResponse},
	})
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Warning": "Something went wrong. Please try again.",
		})
		q()
		return
	}
	recaptchaAPIResponse := struct {
		Success bool `json:"success"`
	}{}

	defer resp.Body.Close()
	err = json.NewDecoder(resp.Body).Decode(&recaptchaAPIResponse)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Warning": "Something went wrong. Please try again.",
		})
		q()
		return
	}
	if !recaptchaAPIResponse.Success {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Error": "Couldn't verify CAPTCHA. Please try again.",
		})
		q()
		return
	}
}
<commit_msg>fix verifyCaptcha handler<commit_after>package main

import (
	// std
	"encoding/json"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"
	// base
	"github.com/Preetam/onecontactlink/internal-api/client"
	"github.com/Preetam/onecontactlink/schema"
	"github.com/Preetam/onecontactlink/web/linktoken"
	// vendor
	"github.com/VividCortex/siesta"
)

func serveGetRequest(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	linkStr := params.String("link", "", "link code")
	err := params.Parse(r.Form)

	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid OneContactLink",
		})
		return
	}

	if err != nil || !strings.Contains(*linkStr, "-") {
		invalidLink()
		return
	}
	parts := strings.Split(*linkStr, "-")
	if len(parts) != 2 {
		invalidLink()
		return
	}

	// get request link
	requestLink, err := internalAPIClient.GetRequestLinkByCode(parts[1])
	if err != nil {
		invalidLink()
		return
	}

	// get user
	user, err := internalAPIClient.GetUser(requestLink.User)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong",
		})
		return
	}

	if user.Code != parts[0] {
		invalidLink()
		return
	}

	templ.ExecuteTemplate(w, "request", map[string]string{
		"Name": user.Name,
	})
}

func servePostRequest(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	nameStr := params.String("name", "", "name")
	emailStr := params.String("email", "", "email")
	linkStr := params.String("link", "", "link code")
	err := params.Parse(r.Form)

	invalidLink := func() {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid OneContactLink",
		})
		return
	}

	if err != nil || !strings.Contains(*linkStr, "-") {
		invalidLink()
	}
	parts := strings.Split(*linkStr, "-")
	if len(parts) != 2 {
		invalidLink()
	}

	// get request link
	requestLink, err := internalAPIClient.GetRequestLinkByCode(parts[1])
	if err != nil {
		invalidLink()
		return
	}

	// get user
	user, err := internalAPIClient.GetUser(requestLink.User)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "Something went wrong",
		})
		return
	}

	if user.Code != parts[0] {
		invalidLink()
		return
	}

	if *nameStr == "" || *emailStr == "" {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "You must provide a name and email address",
		})
		return
	}

	toUser := user.ID
	fromUser := 0

	email, err := internalAPIClient.GetEmail(*emailStr)
	if err != nil {
		if err == client.ErrNotFound {
			// Email doesn't exist. Create a new user with the email.
			user, err := internalAPIClient.CreateUser(schema.NewUser(*nameStr, *emailStr))
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				templ.ExecuteTemplate(w, "request", map[string]string{
					"Name":    user.Name,
					"Warning": "Something went wrong. Please try again.",
				})
				return
			}
			fromUser = user.ID
		} else {
			w.WriteHeader(http.StatusInternalServerError)
			templ.ExecuteTemplate(w, "request", map[string]string{
				"Name":    user.Name,
				"Warning": "Something went wrong. Please try again.",
			})
			return
		}
	} else {
		fromUser = email.User
	}

	requestID, err := internalAPIClient.CreateRequest(fromUser, toUser)
	if err != nil {
		if err == client.ErrConflict {
			templ.ExecuteTemplate(w, "request", map[string]string{
				"Name": user.Name,
				"Info": "Looks like you already made this request. If " + user.Name +
					" has already approved your request, we'll send you their latest contact info.",
			})
			// Try to send another request email. This is idempotent.
			internalAPIClient.SendRequestEmail(requestID)
			internalAPIClient.SendContactInfoEmail(requestID)
			return
		}
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Name":    user.Name,
			"Warning": "Something went wrong. Please try again.",
		})
		return
	}
	err = internalAPIClient.SendRequestEmail(requestID)
	if err != nil {
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Name":    user.Name,
			"Warning": "Something went wrong. Please try again.",
		})
		return
	}

	templ.ExecuteTemplate(w, "success", map[string]string{
		"Success": "Request sent!",
	})
}

func serveManageRequest(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	linkStr := params.String("link", "", "link token")
	actionStr := params.String("action", "", "action")
	err := params.Parse(r.Form)

	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid link",
		})
		return
	}

	if err != nil {
		invalidLink()
		return
	}

	linkToken, err := tokenCodec.DecodeToken(*linkStr, new(linktoken.RequestTokenData))
	if err != nil {
		invalidLink()
		return
	}

	if *actionStr != "approve" && *actionStr != "reject" {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Invalid action.",
		})
		return
	}

	// Check if token expired
	if linkToken.Expires <= int(time.Now().Unix()) {
		// Token expired.
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "This link has expired.",
		})
		return
	}

	// extract request ID
	requestID := linkToken.Data.(*linktoken.RequestTokenData).Request
	err = internalAPIClient.ManageRequest(requestID, *actionStr)
	if err != nil {
		if serverErr, ok := err.(client.ServerError); ok {
			if int(serverErr) == http.StatusConflict {
				templ.ExecuteTemplate(w, "invalid", map[string]string{
					"Warning": "Oops, you already responded to this request. We can't change" +
						" anything right now.",
				})
				return
			}
		}

		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		log.Println(err)
		return

	}

	if *actionStr == "approve" {
		err = internalAPIClient.SendContactInfoEmail(requestID)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			templ.ExecuteTemplate(w, "invalid", map[string]string{
				"Error": "Something went wrong. Please try again.",
			})
			log.Println(err)
			return
		}

		templ.ExecuteTemplate(w, "success", map[string]string{
			"Success": "Approved! We'll send them an email with your contact information.",
		})
	} else {
		templ.ExecuteTemplate(w, "success", map[string]string{
			"Success": "Rejected. That email won't be able to send you any more requests.",
		})
	}
}

func serveAuth(c siesta.Context, w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	linkStr := params.String("link", "", "link token")
	err := params.Parse(r.Form)
	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid link",
		})
		return
	}
	if err != nil {
		invalidLink()
		return
	}

	linkToken, err := tokenCodec.DecodeToken(*linkStr, new(linktoken.UserTokenData))
	if err != nil {
		invalidLink()
		return
	}

	// Check if token expired
	if linkToken.Expires <= int(time.Now().Unix()) {
		// Token expired.
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "This link has expired.",
		})
		return
	}

	// extract user ID
	userID := linkToken.Data.(*linktoken.UserTokenData).User

	// get user information
	_, err = internalAPIClient.GetUser(userID)
	if err != nil {
		// Token expired.
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	// Update the expiration and set a cookie
	linkToken.Expires = int(time.Now().Unix() + 86400)

	token, err := tokenCodec.EncodeToken(linkToken)
	if err != nil {
		// Token expired.
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     "ocl",
		Value:    token,
		Domain:   CookieDomain,
		Path:     "/",
		HttpOnly: true,
		Secure:   !DevMode,
	})

	w.Header().Add("Refresh", "2; /app")

	templ.ExecuteTemplate(w, "success", map[string]string{
		"Success": "Logged in! Redirecting you to the app...",
	})
}

func servePostLogin(w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	emailStr := params.String("email", "", "email")
	err := params.Parse(r.Form)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "Invalid parameters.",
		})
		return
	}

	internalAPIClient.SendAuth(*emailStr)
	templ.ExecuteTemplate(w, "success", map[string]string{
		"Info": "We've sent a login link to '" + *emailStr +
			"' if it's associated with a valid account.",
	})
}

func serveDevModeAuth(c siesta.Context, w http.ResponseWriter, r *http.Request) {
	params := &siesta.Params{}
	userID := params.Int("user", 0, "user ID")
	err := params.Parse(r.Form)
	invalidLink := func() {
		w.WriteHeader(http.StatusNotFound)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Not a valid link",
		})
		return
	}
	if err != nil {
		invalidLink()
		return
	}

	linkToken := linktoken.NewLinkToken(&linktoken.UserTokenData{
		User: *userID,
	}, int(time.Now().Unix()+86400))

	// get user information
	_, err = internalAPIClient.GetUser(*userID)
	if err != nil {
		log.Println(err)
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	token, err := tokenCodec.EncodeToken(linkToken)
	if err != nil {
		log.Println(err)
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "invalid", map[string]string{
			"Error": "Something went wrong. Please try again.",
		})
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     "ocl",
		Value:    token,
		Domain:   CookieDomain,
		Path:     "/",
		HttpOnly: true,
		Secure:   !DevMode,
	})

	templ.ExecuteTemplate(w, "success", map[string]string{
		"Success": "Logged in!",
	})
}

func verifyCaptcha(c siesta.Context, w http.ResponseWriter, r *http.Request, q func()) {
	if DevMode {
		return
	}

	params := &siesta.Params{}
	recaptchaResponse := params.String("g-recaptcha-response", "", "reCAPTCHA response")
	err := params.Parse(r.Form)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "request", map[string]string{
			"Error": "Missing CAPTCHA parameter.",
		})
		q()
		return
	}

	if *recaptchaResponse == "" {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Error": "Bad CAPTCHA",
		})
		q()
		return
	}

	// verify CAPTCHA
	resp, err := http.PostForm("https://www.google.com/recaptcha/api/siteverify", url.Values{
		"secret":   []string{RecaptchaSecret},
		"response": []string{*recaptchaResponse},
	})
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Warning": "Something went wrong. Please try again.",
		})
		q()
		return
	}
	recaptchaAPIResponse := struct {
		Success bool `json:"success"`
	}{}

	defer resp.Body.Close()
	err = json.NewDecoder(resp.Body).Decode(&recaptchaAPIResponse)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Warning": "Something went wrong. Please try again.",
		})
		q()
		return
	}
	if !recaptchaAPIResponse.Success {
		w.WriteHeader(http.StatusBadRequest)
		templ.ExecuteTemplate(w, "login", map[string]string{
			"Error": "Couldn't verify CAPTCHA. Please try again.",
		})
		q()
		return
	}
}
<|endoftext|>


<commit_before>// Package webprobe implements a HTTP probe.
package webprobe

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

	"github.com/hkjn/prober"
	"github.com/hkjn/probes"
)

const (
	MaxResponseBytes int64 = 10e5 // largest response size accepted
	defaultName            = "WebProber"
)

// WebProber probes a target's HTTP response.
type WebProber struct {
	Target         string // URL to probe
	Method         string // GET, POST, PUT, etc.
	Name           string // name of the prober
	Body           io.Reader
	wantCode       int
	wantInResponse string
	wantHeaders    map[string]string
}

// Name sets the name for the prober.
func Name(name string) func(*WebProber) {
	return func(p *WebProber) {
		p.Name = fmt.Sprintf("%s_%s", defaultName, name)
	}
}

// Body sets the HTTP request body for the prober.
func Body(body io.Reader) func(*WebProber) {
	return func(p *WebProber) {
		p.Body = body
	}
}

// InResponse applies the option that the prober wants given string in the HTTP response.
func InResponse(str string) func(*WebProber) {
	return func(p *WebProber) {
		p.wantInResponse = str
	}
}

// ResponseHeaders adds expected header names and values for the response.
func ResponseHeaders(headers map[string]string) func(*WebProber) {
	return func(p *WebProber) {
		p.wantHeaders = headers
	}
}

// New returns a new instance of the web probe with specified options.
func New(target, method string, code int, options ...func(*WebProber)) *prober.Probe {
	return NewWithGeneric(target, method, code, []prober.Option{}, options...)
}

// NewWithGeneric returns a new instance of the web probe with specified options.
//
// NewWithGeneric passes through specified prober.Options, after
// applying the webprobe-specific options.
func NewWithGeneric(target, method string, code int, genericOpts []prober.Option, options ...func(*WebProber)) *prober.Probe {
	name := defaultName
	p := &WebProber{
		Target:      target,
		Name:        name,
		Method:      method,
		wantCode:    code,
		wantHeaders: make(map[string]string),
	}
	for _, opt := range options {
		opt(p)
	}
	return prober.NewProbe(p, p.Name, fmt.Sprintf("Probes HTTP response of %s", target), genericOpts...)
}

// Probe verifies that the target's HTTP response is as expected.
func (p WebProber) Probe() prober.Result {
	req, err := http.NewRequest(p.Method, p.Target, p.Body)
	if err != nil {
		return prober.FailedWith(fmt.Errorf("failed to create HTTP request: %v", err))
	}
	// Inform the server that we'd like the connection to be closed once
	// we're done:
	// http://craigwickesser.com/2015/01/golang-http-to-many-open-files/
	req.Header.Set("Connection", "close")
	t := http.Transport{}
	resp, err := t.RoundTrip(req)
	if err != nil {
		return prober.FailedWith(fmt.Errorf("failed to send HTTP request: %v", err))
	}
	defer resp.Body.Close()

	for h, want := range p.wantHeaders {
		got := resp.Header.Get(h)
		if want != got {
			return prober.FailedWith(fmt.Errorf("want header %q=%q in response, got %q", h, want, got))
		}
	}

	if resp.StatusCode != p.wantCode {
		return prober.FailedWith(fmt.Errorf("bad HTTP response status; want %d, got %d", p.wantCode, resp.StatusCode))
	}
	body, err := ioutil.ReadAll(io.LimitReader(resp.Body, MaxResponseBytes))
	if err != nil {
		return prober.FailedWith(fmt.Errorf("failed to read HTTP response: %v", err))
	}
	sb := string(body)
	if !strings.Contains(sb, p.wantInResponse) {
		return prober.FailedWith(fmt.Errorf("response doesn't contain %q: \n%v\n", p.wantInResponse, sb))
	}
	return prober.Passed()
}

// Alert sends an alert notification via email.
func (p *WebProber) Alert(name, desc string, badness int, records prober.Records) error {
	return probes.SendAlertEmail(name, desc, badness, records)
}
<commit_msg>Allow for custom alert funcs to be set with Alert()<commit_after>// Package webprobe implements a HTTP probe.
package webprobe

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

	"github.com/hkjn/prober"
	"github.com/hkjn/probes"
)

const (
	MaxResponseBytes int64 = 10e5 // largest response size accepted
	defaultName            = "WebProber"
)

// WebProber probes a target's HTTP response.
type WebProber struct {
	Target         string // URL to probe
	Method         string // GET, POST, PUT, etc.
	Name           string // name of the prober
	Body           io.Reader
	wantCode       int
	wantInResponse string
	wantHeaders    map[string]string
	alertFn        prober.AlertFn
}

// Name sets the name for the prober.
func Name(name string) func(*WebProber) {
	return func(p *WebProber) {
		p.Name = fmt.Sprintf("%s_%s", defaultName, name)
	}
}

// Body sets the HTTP request body for the prober.
func Body(body io.Reader) func(*WebProber) {
	return func(p *WebProber) {
		p.Body = body
	}
}

// InResponse applies the option that the prober wants given string in the HTTP response.
func InResponse(str string) func(*WebProber) {
	return func(p *WebProber) {
		p.wantInResponse = str
	}
}

// ResponseHeaders adds expected header names and values for the response.
func ResponseHeaders(headers map[string]string) func(*WebProber) {
	return func(p *WebProber) {
		p.wantHeaders = headers
	}
}

// Alert sets a custom alerting function.
//
// If Alert is not called, the probes.SendAlertEmail function is called.
func Alert(fn prober.AlertFn) func(*WebProber) {
	return func(p *WebProber) {
		p.alertFn = fn
	}
}

// New returns a new instance of the web probe with specified options.
func New(target, method string, code int, options ...func(*WebProber)) *prober.Probe {
	return NewWithGeneric(target, method, code, []prober.Option{}, options...)
}

// NewWithGeneric returns a new instance of the web probe with specified options.
//
// NewWithGeneric passes through specified prober.Options, after
// applying the webprobe-specific options.
func NewWithGeneric(target, method string, code int, genericOpts []prober.Option, options ...func(*WebProber)) *prober.Probe {
	name := defaultName
	p := &WebProber{
		Target:      target,
		Name:        name,
		Method:      method,
		wantCode:    code,
		wantHeaders: make(map[string]string),
		alertFn:     probes.SendAlertEmail,
	}
	for _, opt := range options {
		opt(p)
	}
	return prober.NewProbe(p, p.Name, fmt.Sprintf("Probes HTTP response of %s", target), genericOpts...)
}

// Probe verifies that the target's HTTP response is as expected.
func (p WebProber) Probe() prober.Result {
	req, err := http.NewRequest(p.Method, p.Target, p.Body)
	if err != nil {
		return prober.FailedWith(fmt.Errorf("failed to create HTTP request: %v", err))
	}
	// Inform the server that we'd like the connection to be closed once
	// we're done:
	// http://craigwickesser.com/2015/01/golang-http-to-many-open-files/
	req.Header.Set("Connection", "close")
	t := http.Transport{}
	resp, err := t.RoundTrip(req)
	if err != nil {
		return prober.FailedWith(fmt.Errorf("failed to send HTTP request: %v", err))
	}
	defer resp.Body.Close()

	for h, want := range p.wantHeaders {
		got := resp.Header.Get(h)
		if want != got {
			return prober.FailedWith(fmt.Errorf("want header %q=%q in response, got %q", h, want, got))
		}
	}

	if resp.StatusCode != p.wantCode {
		return prober.FailedWith(fmt.Errorf("bad HTTP response status; want %d, got %d", p.wantCode, resp.StatusCode))
	}
	body, err := ioutil.ReadAll(io.LimitReader(resp.Body, MaxResponseBytes))
	if err != nil {
		return prober.FailedWith(fmt.Errorf("failed to read HTTP response: %v", err))
	}
	sb := string(body)
	if !strings.Contains(sb, p.wantInResponse) {
		return prober.FailedWith(fmt.Errorf("response doesn't contain %q: \n%v\n", p.wantInResponse, sb))
	}
	return prober.Passed()
}

// Alert calls the prober.AlertFn for the prober.
//
// If no prober.AlertFn was set with the Alert() option,
// probes.SendAlertEmail is used by default.
func (p *WebProber) Alert(name, desc string, badness int, records prober.Records) error {
	return p.alertFn(name, desc, badness, records)
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/m3db/m3/src/cmd/tools"
	"github.com/m3db/m3/src/dbnode/encoding"
	"github.com/m3db/m3/src/dbnode/encoding/m3tsz"
	"github.com/m3db/m3/src/dbnode/persist"
	"github.com/m3db/m3/src/dbnode/persist/fs"
	"github.com/m3db/m3/src/x/ident"

	"github.com/pborman/getopt"
	"go.uber.org/zap"
)

const snapshotType = "snapshot"
const flushType = "flush"

func main() {
	var (
		optPathPrefix  = getopt.StringLong("path-prefix", 'p', "", "Path prefix [e.g. /var/lib/m3db]")
		optNamespace   = getopt.StringLong("namespace", 'n', "", "Namespace [e.g. metrics]")
		optShard       = getopt.Uint32Long("shard", 's', 0, "Shard [expected format uint32]")
		optBlockstart  = getopt.Int64Long("block-start", 'b', 0, "Block Start Time [in nsec]")
		volume         = getopt.Int64Long("volume", 'v', 0, "Volume number")
		fileSetTypeArg = getopt.StringLong("fileset-type", 't', flushType, fmt.Sprintf("%s|%s", flushType, snapshotType))
		idFilter       = getopt.StringLong("id-filter", 'f', "", "ID Contains Filter (optional)")
	)
	getopt.Parse()

	rawLogger, err := zap.NewDevelopment()
	if err != nil {
		log.Fatalf("unable to create logger: %+v", err)
	}
	log := rawLogger.Sugar()

	if *optPathPrefix == "" ||
		*optNamespace == "" ||
		*optShard < 0 ||
		*optBlockstart <= 0 ||
		*volume < 0 ||
		(*fileSetTypeArg != snapshotType && *fileSetTypeArg != flushType) {
		getopt.Usage()
		os.Exit(1)
	}

	var fileSetType persist.FileSetType
	switch *fileSetTypeArg {
	case flushType:
		fileSetType = persist.FileSetFlushType
	case snapshotType:
		fileSetType = persist.FileSetSnapshotType
	default:
		log.Fatalf("unknown fileset type: %s", *fileSetTypeArg)
	}

	bytesPool := tools.NewCheckedBytesPool()
	bytesPool.Init()

	encodingOpts := encoding.NewOptions().SetBytesPool(bytesPool)

	fsOpts := fs.NewOptions().SetFilePathPrefix(*optPathPrefix)
	reader, err := fs.NewReader(bytesPool, fsOpts)
	if err != nil {
		log.Fatalf("could not create new reader: %v", err)
	}

	openOpts := fs.DataReaderOpenOptions{
		Identifier: fs.FileSetFileIdentifier{
			Namespace:   ident.StringID(*optNamespace),
			Shard:       *optShard,
			BlockStart:  time.Unix(0, *optBlockstart),
			VolumeIndex: int(*volume),
		},
		FileSetType: fileSetType,
	}

	err = reader.Open(openOpts)
	if err != nil {
		log.Fatalf("unable to open reader: %v", err)
	}

	for {
		id, _, data, _, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("err reading metadata: %v", err)
		}

		if *idFilter != "" && !strings.Contains(id.String(), *idFilter) {
			continue
		}

		data.IncRef()
		iter := m3tsz.NewReaderIterator(bytes.NewReader(data.Bytes()), true, encodingOpts)
		for iter.Next() {
			dp, _, annotation := iter.Current()
			// Use fmt package so it goes to stdout instead of stderr
			fmt.Printf("{id: %s, dp: %+v", id.String(), dp)
			if len(annotation) > 0 {
				fmt.Printf(", annotation: %s", hex.EncodeToString(annotation))
			}
			fmt.Println("}")
		}
		if err := iter.Err(); err != nil {
			log.Fatalf("unable to iterate original data: %v", err)
		}
		iter.Close()

		data.DecRef()
		data.Finalize()
	}
}
<commit_msg>[tools] Output annotations as base64 (#2743)<commit_after>// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/m3db/m3/src/cmd/tools"
	"github.com/m3db/m3/src/dbnode/encoding"
	"github.com/m3db/m3/src/dbnode/encoding/m3tsz"
	"github.com/m3db/m3/src/dbnode/persist"
	"github.com/m3db/m3/src/dbnode/persist/fs"
	"github.com/m3db/m3/src/x/ident"

	"github.com/pborman/getopt"
	"go.uber.org/zap"
)

const snapshotType = "snapshot"
const flushType = "flush"

func main() {
	var (
		optPathPrefix  = getopt.StringLong("path-prefix", 'p', "", "Path prefix [e.g. /var/lib/m3db]")
		optNamespace   = getopt.StringLong("namespace", 'n', "", "Namespace [e.g. metrics]")
		optShard       = getopt.Uint32Long("shard", 's', 0, "Shard [expected format uint32]")
		optBlockstart  = getopt.Int64Long("block-start", 'b', 0, "Block Start Time [in nsec]")
		volume         = getopt.Int64Long("volume", 'v', 0, "Volume number")
		fileSetTypeArg = getopt.StringLong("fileset-type", 't', flushType, fmt.Sprintf("%s|%s", flushType, snapshotType))
		idFilter       = getopt.StringLong("id-filter", 'f', "", "ID Contains Filter (optional)")
	)
	getopt.Parse()

	rawLogger, err := zap.NewDevelopment()
	if err != nil {
		log.Fatalf("unable to create logger: %+v", err)
	}
	log := rawLogger.Sugar()

	if *optPathPrefix == "" ||
		*optNamespace == "" ||
		*optShard < 0 ||
		*optBlockstart <= 0 ||
		*volume < 0 ||
		(*fileSetTypeArg != snapshotType && *fileSetTypeArg != flushType) {
		getopt.Usage()
		os.Exit(1)
	}

	var fileSetType persist.FileSetType
	switch *fileSetTypeArg {
	case flushType:
		fileSetType = persist.FileSetFlushType
	case snapshotType:
		fileSetType = persist.FileSetSnapshotType
	default:
		log.Fatalf("unknown fileset type: %s", *fileSetTypeArg)
	}

	bytesPool := tools.NewCheckedBytesPool()
	bytesPool.Init()

	encodingOpts := encoding.NewOptions().SetBytesPool(bytesPool)

	fsOpts := fs.NewOptions().SetFilePathPrefix(*optPathPrefix)
	reader, err := fs.NewReader(bytesPool, fsOpts)
	if err != nil {
		log.Fatalf("could not create new reader: %v", err)
	}

	openOpts := fs.DataReaderOpenOptions{
		Identifier: fs.FileSetFileIdentifier{
			Namespace:   ident.StringID(*optNamespace),
			Shard:       *optShard,
			BlockStart:  time.Unix(0, *optBlockstart),
			VolumeIndex: int(*volume),
		},
		FileSetType: fileSetType,
	}

	err = reader.Open(openOpts)
	if err != nil {
		log.Fatalf("unable to open reader: %v", err)
	}

	for {
		id, _, data, _, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("err reading metadata: %v", err)
		}

		if *idFilter != "" && !strings.Contains(id.String(), *idFilter) {
			continue
		}

		data.IncRef()
		iter := m3tsz.NewReaderIterator(bytes.NewReader(data.Bytes()), true, encodingOpts)
		for iter.Next() {
			dp, _, annotation := iter.Current()
			// Use fmt package so it goes to stdout instead of stderr
			fmt.Printf("{id: %s, dp: %+v", id.String(), dp)
			if len(annotation) > 0 {
				fmt.Printf(", annotation: %s", base64.StdEncoding.EncodeToString(annotation))
			}
			fmt.Println("}")
		}
		if err := iter.Err(); err != nil {
			log.Fatalf("unable to iterate original data: %v", err)
		}
		iter.Close()

		data.DecRef()
		data.Finalize()
	}
}
<|endoftext|>