<commit_before>package git

import (
	"io"
	"os"
	"testing"

	"gopkg.in/src-d/go-git.v3/clients/common"
	"gopkg.in/src-d/go-git.v3/core"

	. "gopkg.in/check.v1"
)

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

type MockGitUploadPackService struct {
	Auth common.AuthMethod
	RC   io.ReadCloser
}

func (s *MockGitUploadPackService) Connect(url common.Endpoint) error {
	return nil
}

func (s *MockGitUploadPackService) ConnectWithAuth(url common.Endpoint, auth common.AuthMethod) error {
	s.Auth = auth
	return nil
}

func (s *MockGitUploadPackService) Info() (*common.GitUploadPackInfo, error) {
	hash := core.NewHash("6ecf0ef2c2dffb796033e5a02219af86ec6584e5")

	cap := common.NewCapabilities()
	cap.Decode("6ecf0ef2c2dffb796033e5a02219af86ec6584e5 HEADmulti_ack thin-pack side-band side-band-64k ofs-delta shallow no-progress include-tag multi_ack_detailed no-done symref=HEAD:refs/heads/master agent=git/2:2.4.8~dbussink-fix-enterprise-tokens-compilation-1167-gc7006cf")

	return &common.GitUploadPackInfo{
		Capabilities: cap,
		Head:         hash,
		Refs:         map[string]core.Hash{"refs/heads/master": hash},
	}, nil
}

func (s *MockGitUploadPackService) Fetch(*common.GitUploadPackRequest) (io.ReadCloser, error) {
	var err error
	s.RC, err = os.Open("formats/packfile/fixtures/git-fixture.ref-delta")
	return s.RC, err
}

var fixtureRepos = [...]struct {
	url      string
	packfile string
}{
	{"https://github.com/tyba/git-fixture.git", "formats/packfile/fixtures/git-fixture.ofs-delta"},
	{"https://github.com/jamesob/desk.git", "formats/packfile/fixtures/jamesob-desk.pack"},
	{"https://github.com/spinnaker/spinnaker.git", "formats/packfile/fixtures/spinnaker-spinnaker.pack"},
}

type SuiteCommon struct{}

var _ = Suite(&SuiteCommon{})

var countLinesTests = [...]struct {
	i string // the string we want to count lines from
	e int    // the expected number of lines in i
}{
	{"", 0},
	{"a", 1},
	{"a\n", 1},
	{"a\nb", 2},
	{"a\nb\n", 2},
	{"a\nb\nc", 3},
	{"a\nb\nc\n", 3},
	{"a\n\n\nb\n", 4},
	{"first line\n\tsecond line\nthird line\n", 3},
}

func (s *SuiteCommon) TestCountLines(c *C) {
	for i, t := range countLinesTests {
		o := countLines(t.i)
		c.Assert(o, Equals, t.e, Commentf("subtest %d, input=%q", i, t.i))
	}
}
<commit_msg>Added helper function for test fixture unpacking<commit_after>package git

import (
	"io"
	"os"
	"testing"

	"gopkg.in/src-d/go-git.v3/clients/common"
	"gopkg.in/src-d/go-git.v3/core"
	"gopkg.in/src-d/go-git.v3/formats/packfile"

	. "gopkg.in/check.v1"
)

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

type MockGitUploadPackService struct {
	Auth common.AuthMethod
	RC   io.ReadCloser
}

func (s *MockGitUploadPackService) Connect(url common.Endpoint) error {
	return nil
}

func (s *MockGitUploadPackService) ConnectWithAuth(url common.Endpoint, auth common.AuthMethod) error {
	s.Auth = auth
	return nil
}

func (s *MockGitUploadPackService) Info() (*common.GitUploadPackInfo, error) {
	hash := core.NewHash("6ecf0ef2c2dffb796033e5a02219af86ec6584e5")

	cap := common.NewCapabilities()
	cap.Decode("6ecf0ef2c2dffb796033e5a02219af86ec6584e5 HEADmulti_ack thin-pack side-band side-band-64k ofs-delta shallow no-progress include-tag multi_ack_detailed no-done symref=HEAD:refs/heads/master agent=git/2:2.4.8~dbussink-fix-enterprise-tokens-compilation-1167-gc7006cf")

	return &common.GitUploadPackInfo{
		Capabilities: cap,
		Head:         hash,
		Refs:         map[string]core.Hash{"refs/heads/master": hash},
	}, nil
}

func (s *MockGitUploadPackService) Fetch(*common.GitUploadPackRequest) (io.ReadCloser, error) {
	var err error
	s.RC, err = os.Open("formats/packfile/fixtures/git-fixture.ref-delta")
	return s.RC, err
}

type packedFixture struct {
	url      string
	packfile string
}

var fixtureRepos = []packedFixture{
	{"https://github.com/tyba/git-fixture.git", "formats/packfile/fixtures/git-fixture.ofs-delta"},
	{"https://github.com/jamesob/desk.git", "formats/packfile/fixtures/jamesob-desk.pack"},
	{"https://github.com/spinnaker/spinnaker.git", "formats/packfile/fixtures/spinnaker-spinnaker.pack"},
}

func unpackFixtures(c *C, fixtures ...[]packedFixture) map[string]*Repository {
	repos := make(map[string]*Repository, 0)
	for _, group := range fixtures {
		for _, fixture := range group {
			if _, existing := repos[fixture.url]; existing {
				continue
			}
			repos[fixture.url] = NewPlainRepository()

			d, err := os.Open(fixture.packfile)
			c.Assert(err, IsNil)

			r := packfile.NewReader(d)
			r.Format = packfile.OFSDeltaFormat // TODO: how to know the format of a pack file ahead of time?

			_, err = r.Read(repos[fixture.url].Storage)
			c.Assert(err, IsNil)

			c.Assert(d.Close(), IsNil)
		}
	}
	return repos
}

type SuiteCommon struct{}

var _ = Suite(&SuiteCommon{})

var countLinesTests = [...]struct {
	i string // the string we want to count lines from
	e int    // the expected number of lines in i
}{
	{"", 0},
	{"a", 1},
	{"a\n", 1},
	{"a\nb", 2},
	{"a\nb\n", 2},
	{"a\nb\nc", 3},
	{"a\nb\nc\n", 3},
	{"a\n\n\nb\n", 4},
	{"first line\n\tsecond line\nthird line\n", 3},
}

func (s *SuiteCommon) TestCountLines(c *C) {
	for i, t := range countLinesTests {
		o := countLines(t.i)
		c.Assert(o, Equals, t.e, Commentf("subtest %d, input=%q", i, t.i))
	}
}
<|endoftext|>


<commit_before>package capn_test

import (
	"bytes"

	capn "github.com/glycerine/go-capnproto"
)

func zdateFilledSegment(n int, packed bool) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	z := NewRootZ(seg)
	list := NewZdateList(seg, n)
	// hand added a Set() method to messages_test.go, so plist not needed
	// plist := capn.PointerList(list)

	for i := 0; i < n; i++ {
		d := NewZdate(seg)
		d.SetMonth(12)
		d.SetDay(7)
		d.SetYear(int16(2004 + i))
		//plist.Set(i, capn.Object(d))
		list.Set(i, d)
	}
	z.SetZdatevec(list)

	buf := bytes.Buffer{}
	if packed {
		seg.WriteToPacked(&buf)
	} else {
		seg.WriteTo(&buf)
	}
	return seg, buf.Bytes()
}

func zdateReader(n int, packed bool) *bytes.Reader {
	_, byteSlice := zdateFilledSegment(n, packed)
	return bytes.NewReader(byteSlice)
}

func zdataFilledSegment(n int) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	d := NewRootZdata(seg)
	b := make([]byte, n)
	for i := 0; i < len(b); i++ {
		b[i] = byte(i)
	}
	d.SetData(b)
	buf := bytes.Buffer{}
	seg.WriteTo(&buf)
	return seg, buf.Bytes()
}

func zdataReader(n int) *bytes.Reader {
	_, byteSlice := zdataFilledSegment(n)
	return bytes.NewReader(byteSlice)
}
<commit_msg>zdataFilledSegment returns a Z with a Zdata inside<commit_after>package capn_test

import (
	"bytes"

	capn "github.com/glycerine/go-capnproto"
)

func zdateFilledSegment(n int, packed bool) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	z := NewRootZ(seg)
	list := NewZdateList(seg, n)
	// hand added a Set() method to messages_test.go, so plist not needed
	// plist := capn.PointerList(list)

	for i := 0; i < n; i++ {
		d := NewZdate(seg)
		d.SetMonth(12)
		d.SetDay(7)
		d.SetYear(int16(2004 + i))
		//plist.Set(i, capn.Object(d))
		list.Set(i, d)
	}
	z.SetZdatevec(list)

	buf := bytes.Buffer{}
	if packed {
		seg.WriteToPacked(&buf)
	} else {
		seg.WriteTo(&buf)
	}
	return seg, buf.Bytes()
}

func zdateReader(n int, packed bool) *bytes.Reader {
	_, byteSlice := zdateFilledSegment(n, packed)
	return bytes.NewReader(byteSlice)
}

func zdataFilledSegment(n int) (*capn.Segment, []byte) {
	seg := capn.NewBuffer(nil)
	z := NewRootZ(seg)
	d := NewZdata(seg)

	b := make([]byte, n)
	for i := 0; i < len(b); i++ {
		b[i] = byte(i)
	}
	d.SetData(b)
	z.SetZdata(d)

	buf := bytes.Buffer{}
	seg.WriteTo(&buf)
	return seg, buf.Bytes()
}

func zdataReader(n int) *bytes.Reader {
	_, byteSlice := zdataFilledSegment(n)
	return bytes.NewReader(byteSlice)
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2019 Bitmark Inc.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package messagebus_test

import (
	"math/rand"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/bitmark-inc/bitmarkd/messagebus"
)

func setup(t *testing.T) {
	t.Logf("running %s\n", t.Name())
}

func teardown(t *testing.T) {
	messagebus.Bus.Announce.Release()
	messagebus.Bus.Blockstore.Release()
	messagebus.Bus.Connector.Release()
	messagebus.Bus.TestQueue.Release()
	messagebus.Bus.Broadcast.Release()

	messagebus.Bus.Broadcast.DropCache()

}

func TestQueue(t *testing.T) {

	setup(t)
	defer teardown(t)

	items := []messagebus.Message{
		{
			Command:    "c1",
			Parameters: nil,
		},
		{
			Command:    "c2",
			Parameters: nil,
		},
		{
			Command:    "c3",
			Parameters: nil,
		},
	}

	for _, item := range items {
		messagebus.Bus.TestQueue.Send(item.Command)
	}

	queue := messagebus.Bus.TestQueue.Chan()
	for _, item := range items {
		received := <-queue
		if received.Command != item.Command {
			t.Errorf("actual: %q  expected: %q", received.Command, item.Command)
		}
	}

}

func TestBroadcast(t *testing.T) {

	setup(t)
	defer teardown(t)

	items := []messagebus.Message{
		{
			Command:    "c1",
			Parameters: nil,
		},
		{
			Command:    "c2",
			Parameters: nil,
		},
		{
			Command:    "c3",
			Parameters: nil,
		},
	}

	// nothing listening so these messages should be dropped
	for _, item := range items {
		messagebus.Bus.Broadcast.Send("ignored:" + item.Command)
	}

	// allow background to run
	time.Sleep(20 * time.Millisecond)

	// create some listeners
	const listeners = 5

	var l [listeners]int
	var wg sync.WaitGroup

	for i := 0; i < listeners; i += 1 {
		wg.Add(1)
		go func(n int) {
			queue := messagebus.Bus.Broadcast.Chan(0)
			for _, item := range items {
				received := <-queue
				if received.Command != item.Command {
					t.Errorf("actual: %q  expected: %q", received.Command, item.Command)
				} else {
					l[n] += 1
				}
			}
			wg.Done()
		}(i)

	}

	// all listening so these messages should be received
	for _, item := range items {
		time.Sleep(20 * time.Millisecond)
		messagebus.Bus.Broadcast.Send(item.Command)
	}
	for _, item := range items {
		time.Sleep(20 * time.Millisecond)
		messagebus.Bus.Broadcast.Send(item.Command)
	}

	// wait for completion
	wg.Wait()
	for i, n := range l {
		if n != len(items) {
			t.Errorf("listener[%d] received: %d  expected: %d", i, n, len(items))
		}
	}
}

func TestCache(t *testing.T) {

	setup(t)
	defer teardown(t)

	cacheableCmd := []string{"assets", "issues", "transfer", "proof", "block"}
	uncacheableCmd := []string{"rpc", "peer"}
	c1 := cacheableCmd[rand.Intn(len(cacheableCmd))]
	c2 := uncacheableCmd[rand.Intn(len(uncacheableCmd))]
	c := []string{c1, c2}
	p := []byte{0x05, 0xFA, 0xFE}

	// declare listener
	queue := messagebus.Bus.Broadcast.Chan(50)

	// send a message has not been cached before
	for _, cmd := range c {
		messagebus.Bus.Broadcast.Send(cmd, p)
	}

	// wait for background
	time.Sleep(20 * time.Millisecond)

	s := make([]string, 0)
	for len(queue) > 0 {
		m := <-queue
		s = append(s, m.Command)
	}

	// verify received values from queue are the ones sent before
	if len(c) != len(s) {
		t.Errorf("actual count: %d, expected: %d", len(s), len(c))
	}

	for i := 0; i < len(c); i += 1 {
		if c[i] != s[i] {
			t.Errorf("actual command: %s, expected: %s", s[i], c[i])
		}
	}

	// func to check whether a string is contained in an array string
	f := func(a []string, i string) bool {
		for _, item := range a {
			if item == i {
				return true
			}
		}
		return false
	}

	// recreate s
	s = make([]string, 0)
	for _, cmd := range c {
		messagebus.Bus.Broadcast.Send(cmd, p)
	}
	// wait for background
	time.Sleep(20 * time.Millisecond)

	for len(queue) > 0 {
		m := <-queue
		s = append(s, m.Command)
	}

	// verify the queue did not contains cache value
	if len(c) == len(s) {
		t.Errorf("actual count: %d, expected: %d", len(s), len(c))
	}

	for i := 0; i < len(s); i += 1 {
		if !f(uncacheableCmd, s[i]) {
			t.Errorf("actual: %q, expected in %q", s[i], uncacheableCmd)
		}
	}

	// drop cache and resend it
	params := [][]byte{p}
	messagebus.Bus.Broadcast.DropCache(messagebus.Message{Command: c1, Parameters: params})
	messagebus.Bus.Broadcast.Send(c1, p)
	// wait for background
	time.Sleep(20 * time.Millisecond)

	select {
	case received := <-queue:
		if received.Command != c1 {
			t.Errorf("actual command: %q, expected: %q", received.Command, c1)
		}
	default:
		t.Errorf("actual nothing but expected is %q", c1)
	}

}

func TestQueueOverflow(t *testing.T) {

	setup(t)
	defer teardown(t)

	const queueSize = 10
	cmd := []string{"assets", "issues", "transfer", "proof", "rpc", "peer"}
	p1 := []byte{0x11, 0x99}
	p2 := []byte{0xAA, 0xFF}
	queue := messagebus.Bus.Broadcast.Chan(queueSize)

	// put at least 1 block message
	messagebus.Bus.Broadcast.Send("block", p1)
	for i := 0; i < queueSize-1; i += 1 {
		c := cmd[rand.Intn(len(cmd))]
		p := make([]byte, rand.Intn(1024))
		rand.Read(p)
		messagebus.Bus.Broadcast.Send(c, p)
	}

	time.Sleep(20 * time.Millisecond)

	// verify the queue is full
	if queueSize != len(queue) {
		t.Errorf("len(queue) actual: %d, expected: %d", len(queue), queueSize)
	}

	// put one more block message into queue
	messagebus.Bus.Broadcast.Send("block", p2)
	time.Sleep(20 * time.Millisecond)

	// verify there are 2 block messages in queue
	count := 0
	for len(queue) > 0 {
		m := <-queue
		if "block" == m.Command {
			count += 1
		}
	}
	if 2 != count {
		t.Errorf("block message count actual: %d, expected: %d", count, 2)
	}

	// make the queue is full
	for i := 0; i < queueSize; i += 1 {
		c := cmd[rand.Intn(len(cmd))]
		p := make([]byte, rand.Intn(1024))
		rand.Read(p)
		messagebus.Bus.Broadcast.Send(c, p)
	}

	time.Sleep(20 * time.Millisecond)

	// verify the queue is full
	if queueSize != len(queue) {
		t.Errorf("len(queue) actual: %d, expected: %d", len(queue), queueSize)
	}

	// put one non-block message to the queue
	c := cmd[rand.Intn(len(cmd))]
	messagebus.Bus.Broadcast.Send(c, p2)
	time.Sleep(20 * time.Millisecond)

	// verify the new message is just dropped
	for len(queue) > 0 {
		m := <-queue
		if c == m.Command && 1 == len(m.Parameters) && reflect.DeepEqual(p2, m.Parameters[0]) {
			t.Error("Expected new message won't be in queue but actually it's")
		}
	}

}

func TestCacheStateQueueOverflow(t *testing.T) {

	setup(t)
	defer teardown(t)

	const queueSize = 3
	cmd := []string{"assets", "issues", "transfer", "proof", "rpc", "peer"}

	// listener
	queue := messagebus.Bus.Broadcast.Chan(queueSize)

	// make the queue is overflow and continue send one item
	for i := 0; i < queueSize; i++ {
		c := cmd[rand.Intn(len(cmd))]
		p := make([]byte, rand.Intn(1024))
		rand.Read(p)
		messagebus.Bus.Broadcast.Send(c, p)
	}

	p := []byte{0x0A, 0x9F}
	c := cmd[0]

	// continue to send one more
	messagebus.Bus.Broadcast.Send(c, p)

	time.Sleep(20 * time.Millisecond)

	// drop one item
	<-queue

	// then push the last sent item
	messagebus.Bus.Broadcast.Send(c, p)

	time.Sleep(20 * time.Millisecond)

	// verify it already contains in queue
	found := false

searchLoop:
	for len(queue) > 0 {
		m := <-queue
		if c == m.Command && len(m.Parameters) == 1 && reflect.DeepEqual(p, m.Parameters[0]) {
			found = true
			break searchLoop
		}
	}

	if !found {
		t.Errorf("Expected the message command %q is contained in queue, but actually not", c)
	}

}
<commit_msg>[messagebus] update test for new queue dedup method<commit_after>// Copyright (c) 2014-2019 Bitmark Inc.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package messagebus_test

import (
	"math/rand"
	"sync"
	"testing"

	"github.com/bitmark-inc/bitmarkd/messagebus"
)

// command to exit goroutines
const DONE = "***DONE***"

func setup(t *testing.T) {
	t.Logf("running %s\n", t.Name())
}

func teardown(t *testing.T) {
	messagebus.Bus.Announce.Release()
	messagebus.Bus.Blockstore.Release()
	messagebus.Bus.Connector.Release()
	messagebus.Bus.TestQueue.Release()
	messagebus.Bus.Broadcast.Release()
}

func TestQueue(t *testing.T) {

	setup(t)
	defer teardown(t)

	items := []messagebus.Message{
		{
			Command:    "c1",
			Parameters: nil,
		},
		{
			Command:    "c2",
			Parameters: nil,
		},
		{
			Command:    "c3",
			Parameters: nil,
		},
	}

	for _, item := range items {
		messagebus.Bus.TestQueue.Send(item.Command)
	}

	queue := messagebus.Bus.TestQueue.Chan()
	for _, item := range items {
		received := <-queue
		if received.Command != item.Command {
			t.Errorf("actual: %q  expected: %q", received.Command, item.Command)
		}
	}

}

func TestBroadcast(t *testing.T) {

	setup(t)
	defer teardown(t)

	items := []messagebus.Message{
		{
			Command:    "c1",
			Parameters: nil,
		},
		{
			Command:    "c2",
			Parameters: nil,
		},
		{
			Command:    "c3",
			Parameters: nil,
		},
	}

	// nothing listening so these messages should be dropped
	for _, item := range items {
		messagebus.Bus.Broadcast.Send("ignored:" + item.Command)
	}

	// create some listeners
	const listeners = 5

	var l [listeners]int
	var wg sync.WaitGroup

	for i := 0; i < listeners; i += 1 {
		wg.Add(1)

		// queue created outside to avoid having spurious sleeps
		// to wait for goroutines to start
		// ensure queue is of sufficient size to prevent deadlock
		queue := messagebus.Bus.Broadcast.Chan(len(items) + 1)

		go func(n int, queue <-chan messagebus.Message) {
		loop:
			for {
				for _, item := range items {
					received := <-queue
					if DONE == received.Command {
						break loop
					}
					if received.Command != item.Command {
						t.Errorf("%d: actual: %q  expected: %q", n, received.Command, item.Command)
					} else {
						l[n] += 1
					}
				}
			}
			wg.Done()
		}(i, queue)

	}

	// all listening so one copy of each messages should be received
	for _, item := range items {
		for i := 0; i < 10; i += 1 {
			messagebus.Bus.Broadcast.Send(item.Command)
		}
	}

	for _, item := range items {
		messagebus.Bus.Broadcast.Send(item.Command)
	}

	messagebus.Bus.Broadcast.Send(DONE)

	// wait for completion
	wg.Wait()

	// check right number of items received
	for i, n := range l {
		if n != len(items) {
			t.Errorf("listener[%d] received: %d  expected: %d", i, n, len(items))
		}
	}
}

func TestQueueOverflow(t *testing.T) {

	setup(t)
	defer teardown(t)

	const queueSize = 15

	cmd := []string{"assets", "issues", "transfer", "proof", "rpc", "peer"}

	queue := messagebus.Bus.Broadcast.Chan(queueSize)

	// fill the queue
	for i := 0; i < 2*queueSize; i += 1 {
		c := cmd[rand.Intn(len(cmd))]
		p := make([]byte, rand.Intn(1024))
		rand.Read(p)
		messagebus.Bus.Broadcast.Send(c, p)
	}

	if len(queue) >= queueSize {
		t.Fatal("queue was filled by normal messages")
	}

	messagebus.Bus.Broadcast.Send("block", []byte{0x11, 0x99})

	if len(queue) != queueSize {
		t.Fatal("queue could not accept block")
	}

	// verify block message in queue
	count := 0

	for len(queue) > 0 {
		m := <-queue
		if "block" == m.Command {
			count += 1
		}
	}
	if 0 == count {
		t.Fatal("no block message in queue")
	}
}
<|endoftext|>


<commit_before>package main
import "fmt"

func main()  {
  fmt.Println("Hello World")
}
<commit_msg>gofmt file<commit_after>package main

import "fmt"

func main() {
	fmt.Println("Hello World")
}
<|endoftext|>


<commit_before>package noise

import (
	"encoding/binary"
	"fmt"
	"io"
)

// Each encrypted transport message must be <= 65,535 bytes, including 16
// bytes of authentication data. To write larger plaintexts, we split them
// into fragments of maxPlaintextLength before encrypting.
const maxPlaintextLength = 65519

// Read reads from the secure connection, filling `buf` with plaintext data.
// May read less than len(buf) if data is available immediately.
func (s *secureSession) Read(buf []byte) (int, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	l := len(buf)

	// if we have previously unread bytes, and they fit into the buf, copy them over and return
	if l <= len(s.msgBuffer) {
		copy(buf, s.msgBuffer)
		s.msgBuffer = s.msgBuffer[l:]
		return l, nil
	}

	readChunk := func(buf []byte) (int, error) {
		// read length of encrypted message
		ciphertext, err := s.readMsgInsecure()
		if err != nil {
			return 0, err
		}

		plaintext, err := s.Decrypt(ciphertext)
		if err != nil {
			return 0, err
		}

		// append plaintext to message buffer, copy over what can fit in the buf
		// then advance message buffer to remove what was copied
		s.msgBuffer = append(s.msgBuffer, plaintext...)
		c := copy(buf, s.msgBuffer)
		s.msgBuffer = s.msgBuffer[c:]
		return c, nil
	}

	total := 0
	for i := 0; i < len(buf); i += maxPlaintextLength {
		end := i + maxPlaintextLength
		if end > len(buf) {
			end = len(buf)
		}

		c, err := readChunk(buf[i:end])
		total += c
		if err != nil {
			return total, err
		}
	}

	return total, nil
}

// Write encrypts the plaintext `in` data and sends it on the
// secure connection.
func (s *secureSession) Write(in []byte) (int, error) {
	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	writeChunk := func(in []byte) (int, error) {
		ciphertext, err := s.Encrypt(in)
		if err != nil {
			return 0, err
		}

		err = s.writeMsgInsecure(ciphertext)
		if err != nil {
			return 0, err
		}
		return len(in), err
	}

	written := 0
	for i := 0; i < len(in); i += maxPlaintextLength {
		end := i + maxPlaintextLength
		if end > len(in) {
			end = len(in)
		}

		l, err := writeChunk(in[i:end])
		written += l
		if err != nil {
			return written, err
		}
	}
	return written, nil
}

// readMsgInsecure reads a message from the insecure channel.
// it first reads the message length, then consumes that many bytes
// from the insecure conn.
func (s *secureSession) readMsgInsecure() ([]byte, error) {
	buf := make([]byte, 2)
	_, err := io.ReadFull(s.insecure, buf)
	if err != nil {
		return nil, fmt.Errorf("error reading length prefix: %s", err)
	}
	size := int(binary.BigEndian.Uint16(buf))
	buf = make([]byte, size)
	_, err = io.ReadFull(s.insecure, buf)
	return buf, err
}

// writeMsgInsecure writes to the insecure conn.
// data will be prefixed with its length in bytes, written as a 16-bit uint in network order.
func (s *secureSession) writeMsgInsecure(data []byte) error {
	buf := make([]byte, 2)
	binary.BigEndian.PutUint16(buf, uint16(len(data)))
	_, err := s.insecure.Write(buf)
	if err != nil {
		return fmt.Errorf("error writing length prefix: %s", err)
	}
	_, err = s.insecure.Write(data)
	return err
}
<commit_msg>fix EOF handling in Read<commit_after>package noise

import (
	"encoding/binary"
	"fmt"
	"io"
)

// Each encrypted transport message must be <= 65,535 bytes, including 16
// bytes of authentication data. To write larger plaintexts, we split them
// into fragments of maxPlaintextLength before encrypting.
const maxPlaintextLength = 65519

// Read reads from the secure connection, filling `buf` with plaintext data.
// May read less than len(buf) if data is available immediately.
func (s *secureSession) Read(buf []byte) (int, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	l := len(buf)

	// if we have previously unread bytes, and they fit into the buf, copy them over and return
	if l <= len(s.msgBuffer) {
		copy(buf, s.msgBuffer)
		s.msgBuffer = s.msgBuffer[l:]
		return l, nil
	}

	readChunk := func(buf []byte) (int, error) {
		// read length of encrypted message
		ciphertext, err := s.readMsgInsecure()
		if err != nil {
			return 0, err
		}

		plaintext, err := s.Decrypt(ciphertext)
		if err != nil {
			return 0, err
		}

		// append plaintext to message buffer, copy over what can fit in the buf
		// then advance message buffer to remove what was copied
		s.msgBuffer = append(s.msgBuffer, plaintext...)
		c := copy(buf, s.msgBuffer)
		s.msgBuffer = s.msgBuffer[c:]
		return c, nil
	}

	total := 0
	for i := 0; i < len(buf); i += maxPlaintextLength {
		end := i + maxPlaintextLength
		if end > len(buf) {
			end = len(buf)
		}

		c, err := readChunk(buf[i:end])
		total += c
		if err != nil {
			return total, err
		}
	}

	return total, nil
}

// Write encrypts the plaintext `in` data and sends it on the
// secure connection.
func (s *secureSession) Write(in []byte) (int, error) {
	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	writeChunk := func(in []byte) (int, error) {
		ciphertext, err := s.Encrypt(in)
		if err != nil {
			return 0, err
		}

		err = s.writeMsgInsecure(ciphertext)
		if err != nil {
			return 0, err
		}
		return len(in), err
	}

	written := 0
	for i := 0; i < len(in); i += maxPlaintextLength {
		end := i + maxPlaintextLength
		if end > len(in) {
			end = len(in)
		}

		l, err := writeChunk(in[i:end])
		written += l
		if err != nil {
			return written, err
		}
	}
	return written, nil
}

// readMsgInsecure reads a message from the insecure channel.
// it first reads the message length, then consumes that many bytes
// from the insecure conn.
func (s *secureSession) readMsgInsecure() ([]byte, error) {
	buf := make([]byte, 2)
	_, err := io.ReadFull(s.insecure, buf)
	if err != nil {
		return nil, err
	}
	size := int(binary.BigEndian.Uint16(buf))
	buf = make([]byte, size)
	_, err = io.ReadFull(s.insecure, buf)
	return buf, err
}

// writeMsgInsecure writes to the insecure conn.
// data will be prefixed with its length in bytes, written as a 16-bit uint in network order.
func (s *secureSession) writeMsgInsecure(data []byte) error {
	buf := make([]byte, 2)
	binary.BigEndian.PutUint16(buf, uint16(len(data)))
	_, err := s.insecure.Write(buf)
	if err != nil {
		return fmt.Errorf("error writing length prefix: %s", err)
	}
	_, err = s.insecure.Write(data)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"sync"
	"time"

	"github.com/grafana/grafana/pkg/log"
	//"github.com/dgryski/go-tsz"
	"github.com/raintank/go-tsz"
)

var serverStart uint32
var statsPeriod time.Duration

func init() {
	serverStart = uint32(time.Now().Unix())
	statsPeriod = time.Duration(1) * time.Second
}

// AggMetric takes in new values, updates the in-memory data and streams the points to aggregators
// it uses a circular buffer of chunks
// each chunk starts at their respective t0
// a t0 is a timestamp divisible by chunkSpan without a remainder (e.g. 2 hour boundaries)
// firstT0's data is held at index 0, indexes go up and wrap around from numChunks-1 to 0
// in addition, keep in mind that the last chunk is always a work in progress and not useable for aggregation
// AggMetric is concurrency-safe
type AggMetric struct {
	sync.RWMutex
	Key             string
	CurrentChunkPos int    // element in []Chunks that is active. All others are either finished or nil.
	NumChunks       uint32 // size of the circular buffer
	ChunkSpan       uint32 // span of individual chunks in seconds
	Chunks          []*Chunk
	aggregators     []*Aggregator
}

type aggMetricOnDisk struct {
	Key             string
	CurrentChunkPos int
	NumChunks       uint32
	ChunkSpan       uint32
	Chunks          []*Chunk
}

// re-order the chunks with the oldest at start of the list and newest at the end.
// this is to support increaseing the chunkspan in future at startup.
func (a *aggMetricOnDisk) orderedChunks() []*Chunk {
	orderdChunks := make([]*Chunk, len(a.Chunks))
	if uint32(len(a.Chunks)) < a.NumChunks {
		// the circular buffer is not yet full.
		orderdChunks = a.Chunks
	} else {
		// start by writting the oldest chunk first, then each chunk in turn.
		pos := a.CurrentChunkPos - 1
		if pos < 0 {
			pos += len(a.Chunks)
		}
		for i := 0; i < len(a.Chunks); i++ {
			orderdChunks[i] = a.Chunks[pos]
			pos++
			if pos >= len(a.Chunks) {
				pos = 0
			}
		}
	}
	return orderdChunks
}

func (a AggMetric) GobEncode() ([]byte, error) {
	aOnDisk := aggMetricOnDisk{
		Key:             a.Key,
		CurrentChunkPos: a.CurrentChunkPos,
		NumChunks:       a.NumChunks,
		ChunkSpan:       a.ChunkSpan,
		Chunks:          a.Chunks,
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	err := enc.Encode(aOnDisk)

	return b.Bytes(), err
}

func (a *AggMetric) GobDecode(data []byte) error {
	r := bytes.NewReader(data)
	dec := gob.NewDecoder(r)
	aOnDisk := &aggMetricOnDisk{}
	err := dec.Decode(aOnDisk)
	if err != nil {
		return err
	}
	a.Key = aOnDisk.Key

	a.NumChunks = aOnDisk.NumChunks
	a.ChunkSpan = aOnDisk.ChunkSpan
	a.Chunks = aOnDisk.orderedChunks()
	//the current chunk is the last chunk.
	a.CurrentChunkPos = len(a.Chunks) - 1
	return nil
}

// NewAggMetric creates a metric with given key, it retains the given number of chunks each chunkSpan seconds long
// it optionally also creates aggregations with the given settings
func NewAggMetric(key string, chunkSpan, numChunks uint32, aggsetting ...aggSetting) *AggMetric {
	m := AggMetric{
		Key:       key,
		ChunkSpan: chunkSpan,
		NumChunks: numChunks,
		Chunks:    make([]*Chunk, 0),
	}
	for _, as := range aggsetting {
		m.aggregators = append(m.aggregators, NewAggregator(key, as.span, as.chunkSpan, as.numChunks))
	}
	go m.stats()
	return &m
}

func (a *AggMetric) stats() {
	for range time.Tick(statsPeriod) {
		sum := 0
		a.RLock()
		for _, chunk := range a.Chunks {
			if chunk != nil {
				sum += int(chunk.NumPoints)
			}
		}
		a.RUnlock()
		points.Update(int64(sum))
	}
}

func (a *AggMetric) getChunk(pos int) *Chunk {
	if pos >= len(a.Chunks) {
		return nil
	}
	return a.Chunks[pos]
}

// Get all data between the requested time rages. From is inclusive, to is exclusive. from <= x < to
func (a *AggMetric) Get(from, to uint32) (uint32, []*tsz.Iter) {
	log.Debug("GET: %s from: %d to:%d", a.Key, from, to)
	if from >= to {
		panic("invalid request. to must > from")
	}
	a.RLock()
	defer a.RUnlock()
	firstT0 := from - (from % a.ChunkSpan)
	lastT0 := (to - 1) - ((to - 1) % a.ChunkSpan)

	newestPos := a.CurrentChunkPos

	newestChunk := a.getChunk(a.CurrentChunkPos)

	if newestChunk == nil {
		// we dont have any data yet.
		log.Debug("no data for requested range.")
		return 0, make([]*tsz.Iter, 0)
	}
	if firstT0 > newestChunk.T0 {
		// we have no data in the requested range.
		log.Debug("no data for requested range.")
		return 0, make([]*tsz.Iter, 0)
	}

	// get the oldest chunk we have.
	// eg if we have 5 chunks, N is the current chunk and n-4 is the oldest chunk.
	// -----------------------------
	// | n | n-4 | n-3 | n-2 | n-1 |  CurrentChunkPos = 0
	// -----------------------------
	// -----------------------------
	// | n-2 | n-1 | n | n-4 | n-3 |  CurrentChunkPos = 3
	// -----------------------------
	oldestPos := a.CurrentChunkPos + 1
	if oldestPos >= len(a.Chunks) {
		oldestPos = 0
	}

	oldestChunk := a.getChunk(oldestPos)

	// Find the oldest Chunk that the "from" ts falls in.  If from extends before the oldest
	// chunk, then we just use the oldest chunk.
	for oldestChunk != nil && firstT0 > oldestChunk.T0 {
		oldestPos++
		if oldestPos >= len(a.Chunks) {
			oldestPos = 0
		}
		oldestChunk = a.getChunk(oldestPos)
	}

	firstT0 = oldestChunk.T0

	// find the newest Chunk that "to" falls in.  If "to" extends to after the newest data
	// then just return the newest chunk.
	for lastT0 < newestChunk.T0 {
		newestPos--
		if newestPos < 0 {
			newestPos += int(a.NumChunks)
		}
		newestChunk = a.getChunk(newestPos)
		if newestChunk == nil {
			// the requested time range ends before we collected data.
			return 0, make([]*tsz.Iter, 0)
		}
	}

	// now just start at oldestPos and move through the Chunks circular Buffer to newestPos
	iters := make([]*tsz.Iter, 0)
	for oldestPos != newestPos {
		iters = append(iters, a.getChunk(oldestPos).Iter())
		oldestPos++
		if oldestPos >= int(a.NumChunks) {
			oldestPos = 0
		}
	}
	// add the last chunk
	iters = append(iters, a.getChunk(oldestPos).Iter())

	return firstT0, iters
}

// this function must only be called while holding the lock
func (a *AggMetric) addAggregators(ts uint32, val float64) {
	for _, agg := range a.aggregators {
		agg.Add(ts, val)
	}
}

func (a *AggMetric) Persist(c *Chunk) {
	go func() {
		log.Debug("starting to save %v", c)
		a.RLock()
		data := c.Series.Bytes()
		a.RUnlock()
		chunkSizeAtSave.Value(int64(len(data)))
		err := InsertMetric(a.Key, c.T0, data, *metricTTL)
		if err == nil {
			a.Lock()
			c.Saved = true
			a.Unlock()
			log.Debug("save complete. %v", c)
			chunkSaveOk.Inc(1)
		} else {
			log.Error(1, "failed to save metric to cassandra. %v, %s", c, err)
			chunkSaveFail.Inc(1)
			// TODO
		}
	}()
}

// don't ever call with a ts of 0, cause we use 0 to mean not initialized!
func (a *AggMetric) Add(ts uint32, val float64) {
	a.Lock()
	defer a.Unlock()

	t0 := ts - (ts % a.ChunkSpan)

	currentChunk := a.getChunk(a.CurrentChunkPos)
	if currentChunk == nil {
		chunkCreate.Inc(1)
		if len(a.Chunks) < int(a.NumChunks) {
			log.Debug("adding new chunk to cirular Buffer. now %d chunks", a.CurrentChunkPos+1)
			a.Chunks = append(a.Chunks, NewChunk(t0))
		} else {
			a.Chunks[a.CurrentChunkPos] = NewChunk(t0)
		}

		if err := a.Chunks[a.CurrentChunkPos].Push(ts, val); err != nil {
			log.Error(1, "failed to add metric to chunk. %s", err)
			return
		}

		log.Debug("created new chunk. %v", a.Chunks[0])
		a.addAggregators(ts, val)
		return
	}

	if t0 == currentChunk.T0 {

		if currentChunk.Saved {
			//TODO(awoods): allow the chunk to be re-opened.
			log.Error(3, "cant write to chunk that has already been saved.", nil)
			return
		}
		// last prior data was in same chunk as new point
		if err := a.Chunks[a.CurrentChunkPos].Push(ts, val); err != nil {
			log.Error(3, "failed to add metric to chunk. %s", err)
			return
		}
	} else {
		currentChunk.Finish()
		a.Persist(currentChunk)

		a.CurrentChunkPos++
		if a.CurrentChunkPos >= int(a.NumChunks) {
			a.CurrentChunkPos = 0
		}

		chunkCreate.Inc(1)
		if len(a.Chunks) < int(a.NumChunks) {
			log.Debug("adding new chunk to cirular Buffer. now %d chunks", a.CurrentChunkPos+1)
			a.Chunks = append(a.Chunks, NewChunk(t0))
		} else {
			chunkClear.Inc(1)
			log.Debug("numChunks: %d  currentPos: %d", len(a.Chunks), a.CurrentChunkPos)
			log.Debug("clearing chunk from circular buffer. %v", a.Chunks[a.CurrentChunkPos])
			a.Chunks[a.CurrentChunkPos] = NewChunk(t0)
		}
		log.Debug("created new chunk. %v", a.Chunks[a.CurrentChunkPos])

		if err := a.Chunks[a.CurrentChunkPos].Push(ts, val); err != nil {
			log.Error(1, "failed to push metric to chunk. %s", err)
			return
		}
	}
	a.addAggregators(ts, val)
}

func (a *AggMetric) GC(minTs uint32) bool {
	a.Lock()
	defer a.Unlock()
	currentChunk := a.getChunk(a.CurrentChunkPos)
	if currentChunk == nil {
		return false
	}

	if currentChunk.T0 < minTs {
		if currentChunk.Saved {
			// already saved. lets check if we should just delete the metric from memory.
			if currentChunk.T0 < (minTs - (a.ChunkSpan * a.NumChunks)) {
				return true
			}
		}
		// chunk has not been written to in a while. Lets persist it.
		log.Info("Found stale Chunk, persisting it to Cassandra.")
		currentChunk.Finish()
		a.Persist(currentChunk)
	}
	return false
}
<commit_msg>log the metric key and T0 when GC persists chunk<commit_after>package main

import (
	"bytes"
	"encoding/gob"
	"sync"
	"time"

	"github.com/grafana/grafana/pkg/log"
	//"github.com/dgryski/go-tsz"
	"github.com/raintank/go-tsz"
)

var serverStart uint32
var statsPeriod time.Duration

func init() {
	serverStart = uint32(time.Now().Unix())
	statsPeriod = time.Duration(1) * time.Second
}

// AggMetric takes in new values, updates the in-memory data and streams the points to aggregators
// it uses a circular buffer of chunks
// each chunk starts at their respective t0
// a t0 is a timestamp divisible by chunkSpan without a remainder (e.g. 2 hour boundaries)
// firstT0's data is held at index 0, indexes go up and wrap around from numChunks-1 to 0
// in addition, keep in mind that the last chunk is always a work in progress and not useable for aggregation
// AggMetric is concurrency-safe
type AggMetric struct {
	sync.RWMutex
	Key             string
	CurrentChunkPos int    // element in []Chunks that is active. All others are either finished or nil.
	NumChunks       uint32 // size of the circular buffer
	ChunkSpan       uint32 // span of individual chunks in seconds
	Chunks          []*Chunk
	aggregators     []*Aggregator
}

type aggMetricOnDisk struct {
	Key             string
	CurrentChunkPos int
	NumChunks       uint32
	ChunkSpan       uint32
	Chunks          []*Chunk
}

// re-order the chunks with the oldest at start of the list and newest at the end.
// this is to support increaseing the chunkspan in future at startup.
func (a *aggMetricOnDisk) orderedChunks() []*Chunk {
	orderdChunks := make([]*Chunk, len(a.Chunks))
	if uint32(len(a.Chunks)) < a.NumChunks {
		// the circular buffer is not yet full.
		orderdChunks = a.Chunks
	} else {
		// start by writting the oldest chunk first, then each chunk in turn.
		pos := a.CurrentChunkPos - 1
		if pos < 0 {
			pos += len(a.Chunks)
		}
		for i := 0; i < len(a.Chunks); i++ {
			orderdChunks[i] = a.Chunks[pos]
			pos++
			if pos >= len(a.Chunks) {
				pos = 0
			}
		}
	}
	return orderdChunks
}

func (a AggMetric) GobEncode() ([]byte, error) {
	aOnDisk := aggMetricOnDisk{
		Key:             a.Key,
		CurrentChunkPos: a.CurrentChunkPos,
		NumChunks:       a.NumChunks,
		ChunkSpan:       a.ChunkSpan,
		Chunks:          a.Chunks,
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	err := enc.Encode(aOnDisk)

	return b.Bytes(), err
}

func (a *AggMetric) GobDecode(data []byte) error {
	r := bytes.NewReader(data)
	dec := gob.NewDecoder(r)
	aOnDisk := &aggMetricOnDisk{}
	err := dec.Decode(aOnDisk)
	if err != nil {
		return err
	}
	a.Key = aOnDisk.Key

	a.NumChunks = aOnDisk.NumChunks
	a.ChunkSpan = aOnDisk.ChunkSpan
	a.Chunks = aOnDisk.orderedChunks()
	//the current chunk is the last chunk.
	a.CurrentChunkPos = len(a.Chunks) - 1
	return nil
}

// NewAggMetric creates a metric with given key, it retains the given number of chunks each chunkSpan seconds long
// it optionally also creates aggregations with the given settings
func NewAggMetric(key string, chunkSpan, numChunks uint32, aggsetting ...aggSetting) *AggMetric {
	m := AggMetric{
		Key:       key,
		ChunkSpan: chunkSpan,
		NumChunks: numChunks,
		Chunks:    make([]*Chunk, 0),
	}
	for _, as := range aggsetting {
		m.aggregators = append(m.aggregators, NewAggregator(key, as.span, as.chunkSpan, as.numChunks))
	}
	go m.stats()
	return &m
}

func (a *AggMetric) stats() {
	for range time.Tick(statsPeriod) {
		sum := 0
		a.RLock()
		for _, chunk := range a.Chunks {
			if chunk != nil {
				sum += int(chunk.NumPoints)
			}
		}
		a.RUnlock()
		points.Update(int64(sum))
	}
}

func (a *AggMetric) getChunk(pos int) *Chunk {
	if pos >= len(a.Chunks) {
		return nil
	}
	return a.Chunks[pos]
}

// Get all data between the requested time rages. From is inclusive, to is exclusive. from <= x < to
func (a *AggMetric) Get(from, to uint32) (uint32, []*tsz.Iter) {
	log.Debug("GET: %s from: %d to:%d", a.Key, from, to)
	if from >= to {
		panic("invalid request. to must > from")
	}
	a.RLock()
	defer a.RUnlock()
	firstT0 := from - (from % a.ChunkSpan)
	lastT0 := (to - 1) - ((to - 1) % a.ChunkSpan)

	newestPos := a.CurrentChunkPos

	newestChunk := a.getChunk(a.CurrentChunkPos)

	if newestChunk == nil {
		// we dont have any data yet.
		log.Debug("no data for requested range.")
		return 0, make([]*tsz.Iter, 0)
	}
	if firstT0 > newestChunk.T0 {
		// we have no data in the requested range.
		log.Debug("no data for requested range.")
		return 0, make([]*tsz.Iter, 0)
	}

	// get the oldest chunk we have.
	// eg if we have 5 chunks, N is the current chunk and n-4 is the oldest chunk.
	// -----------------------------
	// | n | n-4 | n-3 | n-2 | n-1 |  CurrentChunkPos = 0
	// -----------------------------
	// -----------------------------
	// | n-2 | n-1 | n | n-4 | n-3 |  CurrentChunkPos = 3
	// -----------------------------
	oldestPos := a.CurrentChunkPos + 1
	if oldestPos >= len(a.Chunks) {
		oldestPos = 0
	}

	oldestChunk := a.getChunk(oldestPos)

	// Find the oldest Chunk that the "from" ts falls in.  If from extends before the oldest
	// chunk, then we just use the oldest chunk.
	for oldestChunk != nil && firstT0 > oldestChunk.T0 {
		oldestPos++
		if oldestPos >= len(a.Chunks) {
			oldestPos = 0
		}
		oldestChunk = a.getChunk(oldestPos)
	}

	firstT0 = oldestChunk.T0

	// find the newest Chunk that "to" falls in.  If "to" extends to after the newest data
	// then just return the newest chunk.
	for lastT0 < newestChunk.T0 {
		newestPos--
		if newestPos < 0 {
			newestPos += int(a.NumChunks)
		}
		newestChunk = a.getChunk(newestPos)
		if newestChunk == nil {
			// the requested time range ends before we collected data.
			return 0, make([]*tsz.Iter, 0)
		}
	}

	// now just start at oldestPos and move through the Chunks circular Buffer to newestPos
	iters := make([]*tsz.Iter, 0)
	for oldestPos != newestPos {
		iters = append(iters, a.getChunk(oldestPos).Iter())
		oldestPos++
		if oldestPos >= int(a.NumChunks) {
			oldestPos = 0
		}
	}
	// add the last chunk
	iters = append(iters, a.getChunk(oldestPos).Iter())

	return firstT0, iters
}

// this function must only be called while holding the lock
func (a *AggMetric) addAggregators(ts uint32, val float64) {
	for _, agg := range a.aggregators {
		agg.Add(ts, val)
	}
}

func (a *AggMetric) Persist(c *Chunk) {
	go func() {
		log.Debug("starting to save %v", c)
		a.RLock()
		data := c.Series.Bytes()
		a.RUnlock()
		chunkSizeAtSave.Value(int64(len(data)))
		err := InsertMetric(a.Key, c.T0, data, *metricTTL)
		if err == nil {
			a.Lock()
			c.Saved = true
			a.Unlock()
			log.Debug("save complete. %v", c)
			chunkSaveOk.Inc(1)
		} else {
			log.Error(1, "failed to save metric to cassandra. %v, %s", c, err)
			chunkSaveFail.Inc(1)
			// TODO
		}
	}()
}

// don't ever call with a ts of 0, cause we use 0 to mean not initialized!
func (a *AggMetric) Add(ts uint32, val float64) {
	a.Lock()
	defer a.Unlock()

	t0 := ts - (ts % a.ChunkSpan)

	currentChunk := a.getChunk(a.CurrentChunkPos)
	if currentChunk == nil {
		chunkCreate.Inc(1)
		if len(a.Chunks) < int(a.NumChunks) {
			log.Debug("adding new chunk to cirular Buffer. now %d chunks", a.CurrentChunkPos+1)
			a.Chunks = append(a.Chunks, NewChunk(t0))
		} else {
			a.Chunks[a.CurrentChunkPos] = NewChunk(t0)
		}

		if err := a.Chunks[a.CurrentChunkPos].Push(ts, val); err != nil {
			log.Error(1, "failed to add metric to chunk. %s", err)
			return
		}

		log.Debug("created new chunk. %v", a.Chunks[0])
		a.addAggregators(ts, val)
		return
	}

	if t0 == currentChunk.T0 {

		if currentChunk.Saved {
			//TODO(awoods): allow the chunk to be re-opened.
			log.Error(3, "cant write to chunk that has already been saved. %s T0:%d", a.Key, currentChunk.T0)
			return
		}
		// last prior data was in same chunk as new point
		if err := a.Chunks[a.CurrentChunkPos].Push(ts, val); err != nil {
			log.Error(3, "failed to add metric to chunk. %s", err)
			return
		}
	} else {
		currentChunk.Finish()
		a.Persist(currentChunk)

		a.CurrentChunkPos++
		if a.CurrentChunkPos >= int(a.NumChunks) {
			a.CurrentChunkPos = 0
		}

		chunkCreate.Inc(1)
		if len(a.Chunks) < int(a.NumChunks) {
			log.Debug("adding new chunk to cirular Buffer. now %d chunks", a.CurrentChunkPos+1)
			a.Chunks = append(a.Chunks, NewChunk(t0))
		} else {
			chunkClear.Inc(1)
			log.Debug("numChunks: %d  currentPos: %d", len(a.Chunks), a.CurrentChunkPos)
			log.Debug("clearing chunk from circular buffer. %v", a.Chunks[a.CurrentChunkPos])
			a.Chunks[a.CurrentChunkPos] = NewChunk(t0)
		}
		log.Debug("created new chunk. %v", a.Chunks[a.CurrentChunkPos])

		if err := a.Chunks[a.CurrentChunkPos].Push(ts, val); err != nil {
			log.Error(1, "failed to push metric to chunk. %s", err)
			return
		}
	}
	a.addAggregators(ts, val)
}

func (a *AggMetric) GC(minTs uint32) bool {
	a.Lock()
	defer a.Unlock()
	currentChunk := a.getChunk(a.CurrentChunkPos)
	if currentChunk == nil {
		return false
	}

	if currentChunk.T0 < minTs {
		if currentChunk.Saved {
			// already saved. lets check if we should just delete the metric from memory.
			if currentChunk.T0 < (minTs - (a.ChunkSpan * a.NumChunks)) {
				return true
			}
		}
		// chunk has not been written to in a while. Lets persist it.
		log.Info("Found stale Chunk, persisting it to Cassandra. key: %s T0: %d", a.Key, currentChunk.T0)
		currentChunk.Finish()
		a.Persist(currentChunk)
	}
	return false
}
<|endoftext|>


<commit_before>package middleware

import (
	"log"
	"net/http"
	"strings"

	"github.com/herald-it/goncord/models"

	. "github.com/julienschmidt/httprouter"
)

// CheckPermission middleware to check
// allowed hosts.
// The list of allowed hosts specify
// in the configuration file.
// A sample configuration file:
//	router:
//  	  register:
//          path: /reg
//	    allowedhost:
//	      - localhost:8000
func CheckPermission(next Handle) Handle {
	return func(w http.ResponseWriter, r *http.Request, p Params) {
		path := strings.Split(r.RequestURI, "?")[0]
		flagPermission := true

		remoteAddr := getIP(r.RemoteAddr)

		switch path {
		case models.Set.Router.Register.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Register.AllowedHost)
		case models.Set.Router.Login.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Login.AllowedHost)
		case models.Set.Router.Validate.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Validate.AllowedHost)
		case models.Set.Router.Logout.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Logout.AllowedHost)
		}

		if !flagPermission {
			log.Println("Remote address: ", remoteAddr, " request rejected.")
			log.Println("Insufficient permissions.")

			http.Error(w, "Insufficient rights.", 500)
			return
		}

		next(w, r, p)
	}
}

func getIP(addr string) string {
	tmp := strings.Split(addr, ":")

	if len(tmp) >= 1 {
		return tmp[0]
	}

	return "Invalid remote address."
}

func contains(s string, arr []string) bool {
	if len(arr) == 0 {
		return true
	}

	for _, v := range arr {
		if v == s {
			return true
		}
	}

	return false
}
<commit_msg>Fix comments<commit_after>package middleware

import (
	"log"
	"net/http"
	"strings"

	"github.com/herald-it/goncord/models"

	. "github.com/julienschmidt/httprouter"
)

// CheckPermission middleware to check
// allowed hosts.
// The list of allowed hosts specify
// in the configuration file.
// A sample configuration file:
//	router:
//  	  register:
//          path: /reg
//	    allowedhost:
//	      - localhost
func CheckPermission(next Handle) Handle {
	return func(w http.ResponseWriter, r *http.Request, p Params) {
		path := strings.Split(r.RequestURI, "?")[0]
		flagPermission := true

		remoteAddr := getIP(r.RemoteAddr)

		switch path {
		case models.Set.Router.Register.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Register.AllowedHost)
		case models.Set.Router.Login.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Login.AllowedHost)
		case models.Set.Router.Validate.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Validate.AllowedHost)
		case models.Set.Router.Logout.Path:
			flagPermission = flagPermission && contains(
				remoteAddr,
				models.Set.Router.Logout.AllowedHost)
		}

		if !flagPermission {
			log.Println("Remote address: ", remoteAddr, " request rejected.")
			log.Println("Insufficient permissions.")

			http.Error(w, "Insufficient rights.", 500)
			return
		}

		next(w, r, p)
	}
}

func getIP(addr string) string {
	tmp := strings.Split(addr, ":")

	if len(tmp) >= 1 {
		return tmp[0]
	}

	return "Invalid remote address."
}

func contains(s string, arr []string) bool {
	if len(arr) == 0 {
		return true
	}

	for _, v := range arr {
		if v == s {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>package manifest_test

import (
	. "cf/manifest"
	"fmt"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"path/filepath"
)

var _ = Describe("ManifestDiskRepository", func() {
	var repo ManifestRepository

	BeforeEach(func() {
		repo = NewManifestDiskRepository()
	})

	Describe("given a directory containing a file called 'manifest.yml'", func() {
		It("reads that file", func() {
			m, errs := repo.ReadManifest("../../fixtures/manifests")

			Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(filepath.Clean("../../fixtures/manifests/manifest.yml")))

			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-default-manifest"))
		})
	})

	Describe("given a directory that doesn't contain a file called 'manifest.y{a}ml'", func() {
		It("returns an error", func() {
			m, errs := repo.ReadManifest("../../fixtures")

			Expect(errs).NotTo(BeEmpty())
			Expect(m.Path).To(BeEmpty())
		})
	})

	Describe("given a directory that contains a file called 'manifest.yaml'", func() {
		It("reads that file", func() {
			m, errs := repo.ReadManifest("../../fixtures/manifests/only_yaml")

			fmt.Printf("\n errors: %v", errs)
			//Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(filepath.Clean("../../fixtures/manifests/only_yaml/manifest.yaml")))

			applications, errs := m.Applications()
			//Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-default-manifest"))
		})
	})

	Describe("given a directory contains files called 'manifest.yml' and 'manifest.yaml'", func() {
		It("reads the file named 'manifest.yml'", func() {
			m, errs := repo.ReadManifest("../../fixtures/manifests/both_yaml_yml")

			Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(filepath.Clean("../../fixtures/manifests/both_yaml_yml/manifest.yml")))

			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-default-manifest"))
		})
	})

	Describe("given a path to a file", func() {
		var (
			inputPath string
			m         *Manifest
			errs      ManifestErrors
		)

		BeforeEach(func() {
			inputPath = filepath.Clean("../../fixtures/manifests/different-manifest.yml")
			m, errs = repo.ReadManifest(inputPath)
		})

		It("reads the file at that path", func() {
			Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(inputPath))

			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-different-manifest"))
		})

		It("passes the base directory to the manifest file", func() {
			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(len(applications)).To(Equal(1))
			Expect(*applications[0].Name).To(Equal("from-different-manifest"))

			appPath := filepath.Clean("../../fixtures/manifests")
			Expect(*applications[0].Path).To(Equal(appPath))
		})
	})

	Describe("given a path to a file that doesn't exist", func() {
		It("returns an error", func() {
			_, errs := repo.ReadManifest("some/path/that/doesnt/exist/manifest.yml")
			Expect(errs).NotTo(BeEmpty())
		})

		It("returns empty string for the manifest path", func() {
			m, _ := repo.ReadManifest("some/path/that/doesnt/exist/manifest.yml")
			Expect(m.Path).To(Equal(""))
		})
	})

	Describe("when the manifest is not valid", func() {
		It("returns an error", func() {
			_, errs := repo.ReadManifest("../../fixtures/manifests/empty-manifest.yml")

			fmt.Printf("\n errors: %v", errs)

			Expect(errs).NotTo(BeEmpty())
		})

		It("returns the path to the manifest", func() {
			inputPath := filepath.Clean("../../fixtures/manifests/empty-manifest.yml")
			m, _ := repo.ReadManifest(inputPath)
			Expect(m.Path).To(Equal(inputPath))
		})
	})

	It("converts nested maps to generic maps", func() {
		m, errs := repo.ReadManifest("../../fixtures/manifests/different-manifest.yml")
		Expect(errs).To(BeEmpty())

		applications, errs := m.Applications()
		Expect(errs).To(BeEmpty())
		Expect(*applications[0].EnvironmentVars).To(Equal(map[string]string{
			"LD_LIBRARY_PATH": "/usr/lib/somewhere",
		}))
	})

	It("merges manifests with their 'inherited' manifests", func() {
		m, errs := repo.ReadManifest("../../fixtures/manifests/inherited-manifest.yml")
		Expect(errs).To(BeEmpty())

		applications, errs := m.Applications()
		Expect(errs).To(BeEmpty())
		Expect(*applications[0].Name).To(Equal("base-app"))
		Expect(*applications[0].ServicesToBind).To(Equal([]string{"base-service"}))
		Expect(*applications[0].EnvironmentVars).To(Equal(map[string]string{
			"foo":                "bar",
			"will-be-overridden": "my-value",
		}))

		Expect(*applications[1].Name).To(Equal("my-app"))

		env := *applications[1].EnvironmentVars
		Expect(env["will-be-overridden"]).To(Equal("my-value"))
		Expect(env["foo"]).To(Equal("bar"))

		services := *applications[1].ServicesToBind
		Expect(services).To(Equal([]string{"base-service", "foo-service"}))
	})
})
<commit_msg>Uncomment some assertions<commit_after>package manifest_test

import (
	. "cf/manifest"
	"fmt"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"path/filepath"
)

var _ = Describe("ManifestDiskRepository", func() {
	var repo ManifestRepository

	BeforeEach(func() {
		repo = NewManifestDiskRepository()
	})

	Describe("given a directory containing a file called 'manifest.yml'", func() {
		It("reads that file", func() {
			m, errs := repo.ReadManifest("../../fixtures/manifests")

			Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(filepath.Clean("../../fixtures/manifests/manifest.yml")))

			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-default-manifest"))
		})
	})

	Describe("given a directory that doesn't contain a file called 'manifest.y{a}ml'", func() {
		It("returns an error", func() {
			m, errs := repo.ReadManifest("../../fixtures")

			Expect(errs).NotTo(BeEmpty())
			Expect(m.Path).To(BeEmpty())
		})
	})

	Describe("given a directory that contains a file called 'manifest.yaml'", func() {
		It("reads that file", func() {
			m, errs := repo.ReadManifest("../../fixtures/manifests/only_yaml")

			Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(filepath.Clean("../../fixtures/manifests/only_yaml/manifest.yaml")))

			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-default-manifest"))
		})
	})

	Describe("given a directory contains files called 'manifest.yml' and 'manifest.yaml'", func() {
		It("reads the file named 'manifest.yml'", func() {
			m, errs := repo.ReadManifest("../../fixtures/manifests/both_yaml_yml")

			Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(filepath.Clean("../../fixtures/manifests/both_yaml_yml/manifest.yml")))

			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-default-manifest"))
		})
	})

	Describe("given a path to a file", func() {
		var (
			inputPath string
			m         *Manifest
			errs      ManifestErrors
		)

		BeforeEach(func() {
			inputPath = filepath.Clean("../../fixtures/manifests/different-manifest.yml")
			m, errs = repo.ReadManifest(inputPath)
		})

		It("reads the file at that path", func() {
			Expect(errs).To(BeEmpty())
			Expect(m.Path).To(Equal(inputPath))

			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(*applications[0].Name).To(Equal("from-different-manifest"))
		})

		It("passes the base directory to the manifest file", func() {
			applications, errs := m.Applications()
			Expect(errs).To(BeEmpty())
			Expect(len(applications)).To(Equal(1))
			Expect(*applications[0].Name).To(Equal("from-different-manifest"))

			appPath := filepath.Clean("../../fixtures/manifests")
			Expect(*applications[0].Path).To(Equal(appPath))
		})
	})

	Describe("given a path to a file that doesn't exist", func() {
		It("returns an error", func() {
			_, errs := repo.ReadManifest("some/path/that/doesnt/exist/manifest.yml")
			Expect(errs).NotTo(BeEmpty())
		})

		It("returns empty string for the manifest path", func() {
			m, _ := repo.ReadManifest("some/path/that/doesnt/exist/manifest.yml")
			Expect(m.Path).To(Equal(""))
		})
	})

	Describe("when the manifest is not valid", func() {
		It("returns an error", func() {
			_, errs := repo.ReadManifest("../../fixtures/manifests/empty-manifest.yml")

			fmt.Printf("\n errors: %v", errs)

			Expect(errs).NotTo(BeEmpty())
		})

		It("returns the path to the manifest", func() {
			inputPath := filepath.Clean("../../fixtures/manifests/empty-manifest.yml")
			m, _ := repo.ReadManifest(inputPath)
			Expect(m.Path).To(Equal(inputPath))
		})
	})

	It("converts nested maps to generic maps", func() {
		m, errs := repo.ReadManifest("../../fixtures/manifests/different-manifest.yml")
		Expect(errs).To(BeEmpty())

		applications, errs := m.Applications()
		Expect(errs).To(BeEmpty())
		Expect(*applications[0].EnvironmentVars).To(Equal(map[string]string{
			"LD_LIBRARY_PATH": "/usr/lib/somewhere",
		}))
	})

	It("merges manifests with their 'inherited' manifests", func() {
		m, errs := repo.ReadManifest("../../fixtures/manifests/inherited-manifest.yml")
		Expect(errs).To(BeEmpty())

		applications, errs := m.Applications()
		Expect(errs).To(BeEmpty())
		Expect(*applications[0].Name).To(Equal("base-app"))
		Expect(*applications[0].ServicesToBind).To(Equal([]string{"base-service"}))
		Expect(*applications[0].EnvironmentVars).To(Equal(map[string]string{
			"foo":                "bar",
			"will-be-overridden": "my-value",
		}))

		Expect(*applications[1].Name).To(Equal("my-app"))

		env := *applications[1].EnvironmentVars
		Expect(env["will-be-overridden"]).To(Equal("my-value"))
		Expect(env["foo"]).To(Equal("bar"))

		services := *applications[1].ServicesToBind
		Expect(services).To(Equal([]string{"base-service", "foo-service"}))
	})
})
<|endoftext|>


<commit_before>// Copyright (c) 2019 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 index

import (
	"errors"
	"sync"

	"github.com/m3db/m3/src/m3ninx/index"
	"github.com/m3db/m3/src/m3ninx/index/segment"
	"github.com/m3db/m3/src/m3ninx/postings"

	"github.com/pborman/uuid"
)

var (
	errCantGetReaderFromClosedSegment = errors.New("cant get reader from closed segment")
	errCantCloseClosedSegment         = errors.New("cant close closed segment")
)

// ReadThroughSegment wraps a segment with a postings list cache so that
// queries can be transparently cached in a read through manner. In addition,
// the postings lists returned by the segments may not be safe to use once the
// underlying segments are closed due to the postings lists pointing into the
// segments mmap'd region. As a result, the close method of the ReadThroughSegment
// will make sure that the cache is purged of all the segments postings lists before
// the segment itself is closed.
type ReadThroughSegment struct {
	segment.Segment
	sync.RWMutex

	uuid              uuid.UUID
	postingsListCache *PostingsListCache

	opts ReadThroughSegmentOptions

	closed bool
}

// ReadThroughSegmentOptions is the options struct for the
// ReadThroughSegment.
type ReadThroughSegmentOptions struct {
	// Whether the postings list for regexp queries should be cached.
	CacheRegexp bool
	// Whether the postings list for term queries should be cached.
	CacheTerms bool
}

// NewReadThroughSegment creates a new read through segment.
func NewReadThroughSegment(
	seg segment.Segment,
	cache *PostingsListCache,
	opts ReadThroughSegmentOptions,
) segment.Segment {
	return &ReadThroughSegment{
		Segment: seg,

		opts:              opts,
		uuid:              uuid.NewUUID(),
		postingsListCache: cache,
	}
}

// Reader returns a read through reader for the read through segment.
func (r *ReadThroughSegment) Reader() (index.Reader, error) {
	r.RLock()
	defer r.RUnlock()
	if r.closed {
		return nil, errCantGetReaderFromClosedSegment
	}

	reader, err := r.Segment.Reader()
	if err != nil {
		return nil, err
	}
	return newReadThroughSegmentReader(
		reader, r.uuid, r.postingsListCache, r.opts), nil
}

// Close purges all entries in the cache associated with this segment,
// and then closes the underlying segment.
func (r *ReadThroughSegment) Close() error {
	r.Lock()
	defer r.Unlock()
	if r.closed {
		return errCantCloseClosedSegment
	}

	r.closed = true

	if r.postingsListCache != nil {
		// Purge segments from the cache before closing the segment to avoid
		// temporarily having postings lists in the cache whose underlying
		// bytes are no longer mmap'd.
		r.postingsListCache.PurgeSegment(r.uuid)
	}
	return r.Segment.Close()
}

type readThroughSegmentReader struct {
	index.Reader

	opts              ReadThroughSegmentOptions
	uuid              uuid.UUID
	postingsListCache *PostingsListCache
}

func newReadThroughSegmentReader(
	reader index.Reader,
	uuid uuid.UUID,
	cache *PostingsListCache,
	opts ReadThroughSegmentOptions,
) index.Reader {
	return &readThroughSegmentReader{
		Reader: reader,

		opts:              opts,
		uuid:              uuid,
		postingsListCache: cache,
	}
}

// MatchRegexp returns a cached posting list or queries the underlying
// segment if their is a cache miss.
func (s *readThroughSegmentReader) MatchRegexp(
	field []byte,
	c index.CompiledRegex,
) (postings.List, error) {
	if s.postingsListCache == nil || !s.opts.CacheRegexp {
		return s.Reader.MatchRegexp(field, c)
	}

	// TODO(rartoul): Would be nice to not allocate strings here.
	fieldStr := string(field)
	patternStr := c.FSTSyntax.String()
	pl, ok := s.postingsListCache.GetRegexp(s.uuid, fieldStr, patternStr)
	if ok {
		return pl, nil
	}

	pl, err := s.Reader.MatchRegexp(field, c)
	if err == nil {
		s.postingsListCache.PutRegexp(s.uuid, fieldStr, patternStr, pl)
	}
	return pl, err
}

// MatchTerm returns a cached posting list or queries the underlying
// segment if their is a cache miss.
func (s *readThroughSegmentReader) MatchTerm(
	field []byte, term []byte,
) (postings.List, error) {
	if s.postingsListCache == nil || !s.opts.CacheTerms {
		return s.Reader.MatchTerm(field, term)
	}

	// TODO(rartoul): Would be nice to not allocate strings here.
	fieldStr := string(field)
	patternStr := string(term)
	pl, ok := s.postingsListCache.GetTerm(s.uuid, fieldStr, patternStr)
	if ok {
		return pl, nil
	}

	pl, err := s.Reader.MatchTerm(field, term)
	if err == nil {
		s.postingsListCache.PutTerm(s.uuid, fieldStr, patternStr, pl)
	}
	return pl, err
}
<commit_msg>[m3ninx] Explicitly implement index.Reader on read through segment reader (#1522)<commit_after>// Copyright (c) 2019 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 index

import (
	"errors"
	"sync"

	"github.com/m3db/m3/src/m3ninx/doc"
	"github.com/m3db/m3/src/m3ninx/index"
	"github.com/m3db/m3/src/m3ninx/index/segment"
	"github.com/m3db/m3/src/m3ninx/postings"

	"github.com/pborman/uuid"
)

var (
	errCantGetReaderFromClosedSegment = errors.New("cant get reader from closed segment")
	errCantCloseClosedSegment         = errors.New("cant close closed segment")
)

// ReadThroughSegment wraps a segment with a postings list cache so that
// queries can be transparently cached in a read through manner. In addition,
// the postings lists returned by the segments may not be safe to use once the
// underlying segments are closed due to the postings lists pointing into the
// segments mmap'd region. As a result, the close method of the ReadThroughSegment
// will make sure that the cache is purged of all the segments postings lists before
// the segment itself is closed.
type ReadThroughSegment struct {
	sync.RWMutex

	segment segment.Segment

	uuid              uuid.UUID
	postingsListCache *PostingsListCache

	opts ReadThroughSegmentOptions

	closed bool
}

// ReadThroughSegmentOptions is the options struct for the
// ReadThroughSegment.
type ReadThroughSegmentOptions struct {
	// Whether the postings list for regexp queries should be cached.
	CacheRegexp bool
	// Whether the postings list for term queries should be cached.
	CacheTerms bool
}

// NewReadThroughSegment creates a new read through segment.
func NewReadThroughSegment(
	seg segment.Segment,
	cache *PostingsListCache,
	opts ReadThroughSegmentOptions,
) segment.Segment {
	return &ReadThroughSegment{
		segment:           seg,
		opts:              opts,
		uuid:              uuid.NewUUID(),
		postingsListCache: cache,
	}
}

// Reader returns a read through reader for the read through segment.
func (r *ReadThroughSegment) Reader() (index.Reader, error) {
	r.RLock()
	defer r.RUnlock()
	if r.closed {
		return nil, errCantGetReaderFromClosedSegment
	}

	reader, err := r.segment.Reader()
	if err != nil {
		return nil, err
	}
	return newReadThroughSegmentReader(
		reader, r.uuid, r.postingsListCache, r.opts), nil
}

// Close purges all entries in the cache associated with this segment,
// and then closes the underlying segment.
func (r *ReadThroughSegment) Close() error {
	r.Lock()
	defer r.Unlock()
	if r.closed {
		return errCantCloseClosedSegment
	}

	r.closed = true

	if r.postingsListCache != nil {
		// Purge segments from the cache before closing the segment to avoid
		// temporarily having postings lists in the cache whose underlying
		// bytes are no longer mmap'd.
		r.postingsListCache.PurgeSegment(r.uuid)
	}
	return r.segment.Close()
}

// FieldsIterable is a pass through call to the segment, since there's no
// postings lists to cache for queries.
func (r *ReadThroughSegment) FieldsIterable() segment.FieldsIterable {
	return r.segment.FieldsIterable()
}

// TermsIterable is a pass through call to the segment, since there's no
// postings lists to cache for queries.
func (r *ReadThroughSegment) TermsIterable() segment.TermsIterable {
	return r.segment.TermsIterable()
}

// ContainsID is a pass through call to the segment, since there's no
// postings lists to cache for queries.
func (r *ReadThroughSegment) ContainsID(id []byte) (bool, error) {
	return r.segment.ContainsID(id)
}

// Size is a pass through call to the segment, since there's no
// postings lists to cache for queries.
func (r *ReadThroughSegment) Size() int64 {
	return r.segment.Size()
}

type readThroughSegmentReader struct {
	// reader is explicitly not embedded at the top level
	// of the struct to force new methods added to index.Reader
	// to be explicitly supported by the read through cache.
	reader            index.Reader
	opts              ReadThroughSegmentOptions
	uuid              uuid.UUID
	postingsListCache *PostingsListCache
}

func newReadThroughSegmentReader(
	reader index.Reader,
	uuid uuid.UUID,
	cache *PostingsListCache,
	opts ReadThroughSegmentOptions,
) index.Reader {
	return &readThroughSegmentReader{
		reader:            reader,
		opts:              opts,
		uuid:              uuid,
		postingsListCache: cache,
	}
}

// MatchRegexp returns a cached posting list or queries the underlying
// segment if their is a cache miss.
func (s *readThroughSegmentReader) MatchRegexp(
	field []byte,
	c index.CompiledRegex,
) (postings.List, error) {
	if s.postingsListCache == nil || !s.opts.CacheRegexp {
		return s.reader.MatchRegexp(field, c)
	}

	// TODO(rartoul): Would be nice to not allocate strings here.
	fieldStr := string(field)
	patternStr := c.FSTSyntax.String()
	pl, ok := s.postingsListCache.GetRegexp(s.uuid, fieldStr, patternStr)
	if ok {
		return pl, nil
	}

	pl, err := s.reader.MatchRegexp(field, c)
	if err == nil {
		s.postingsListCache.PutRegexp(s.uuid, fieldStr, patternStr, pl)
	}
	return pl, err
}

// MatchTerm returns a cached posting list or queries the underlying
// segment if their is a cache miss.
func (s *readThroughSegmentReader) MatchTerm(
	field []byte, term []byte,
) (postings.List, error) {
	if s.postingsListCache == nil || !s.opts.CacheTerms {
		return s.reader.MatchTerm(field, term)
	}

	// TODO(rartoul): Would be nice to not allocate strings here.
	fieldStr := string(field)
	patternStr := string(term)
	pl, ok := s.postingsListCache.GetTerm(s.uuid, fieldStr, patternStr)
	if ok {
		return pl, nil
	}

	pl, err := s.reader.MatchTerm(field, term)
	if err == nil {
		s.postingsListCache.PutTerm(s.uuid, fieldStr, patternStr, pl)
	}
	return pl, err
}

// MatchAll is a pass through call, since there's no postings list to cache.
// NB(r): The postings list returned by match all is just an iterator
// from zero to the maximum document number indexed by the segment and as such
// causes no allocations to compute and construct.
func (s *readThroughSegmentReader) MatchAll() (postings.MutableList, error) {
	return s.reader.MatchAll()
}

// AllDocs is a pass through call, since there's no postings list to cache.
func (s *readThroughSegmentReader) AllDocs() (index.IDDocIterator, error) {
	return s.reader.AllDocs()
}

// Doc is a pass through call, since there's no postings list to cache.
func (s *readThroughSegmentReader) Doc(id postings.ID) (doc.Document, error) {
	return s.reader.Doc(id)
}

// Docs is a pass through call, since there's no postings list to cache.
func (s *readThroughSegmentReader) Docs(pl postings.List) (doc.Iterator, error) {
	return s.reader.Docs(pl)
}

// Close is a pass through call.
func (s *readThroughSegmentReader) Close() error {
	return s.reader.Close()
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/transfer"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsTransferServer() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsTransferServerCreate,
		Read:   resourceAwsTransferServerRead,
		Update: resourceAwsTransferServerUpdate,
		Delete: resourceAwsTransferServerDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"endpoint": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"invocation_role": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateArn,
			},

			"url": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"identity_provider_type": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  transfer.IdentityProviderTypeServiceManaged,
				ValidateFunc: validation.StringInSlice([]string{
					transfer.IdentityProviderTypeServiceManaged,
					transfer.IdentityProviderTypeApiGateway,
				}, false),
			},

			"logging_role": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateArn,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsTransferServerCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	tags := tagsFromMapTransferServer(d.Get("tags").(map[string]interface{}))

	createOpts := &transfer.CreateServerInput{
		Tags: tags,
	}

	var identityProviderDetails *transfer.IdentityProviderDetails
	if attr, ok := d.GetOk("invocation_role"); ok {
		identityProviderDetails.SetInvocationRole(attr.(string))
	}

	if attr, ok := d.GetOk("url"); ok {
		identityProviderDetails.SetUrl(attr.(string))
	}
	if identityProviderDetails != nil {
		createOpts.SetIdentityProviderDetails(identityProviderDetails)
	}

	if attr, ok := d.GetOk("identity_provider_type"); ok {
		createOpts.SetIdentityProviderType(attr.(string))
	}

	if attr, ok := d.GetOk("logging_role"); ok {
		createOpts.SetIdentityProviderType(attr.(string))
	}

	log.Printf("[DEBUG] Create Transfer Server Option: %#v", createOpts)

	resp, err := conn.CreateServer(createOpts)
	if err != nil {
		return fmt.Errorf("Error creating Transfer Server: %s", err)
	}

	d.SetId(*resp.ServerId)

	return resourceAwsTransferServerRead(d, meta)
}

func resourceAwsTransferServerRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn

	descOpts := &transfer.DescribeServerInput{
		ServerId: aws.String(d.Id()),
	}

	log.Printf("[DEBUG] Describe Transfer Server Option: %#v", descOpts)

	resp, err := conn.DescribeServer(descOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			log.Printf("[WARN] Transfer Server (%s) not found, removing from state", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	endpoint := fmt.Sprintf("%s.server.transfer.%s.amazonaws.com", d.Id(), meta.(*AWSClient).region)
	log.Printf("[DEBUG] managed endpoint : %s", endpoint)

	d.Set("arn", resp.Server.Arn)
	d.Set("endpoint", endpoint)
	d.Set("invocation_role", "")
	d.Set("url", "")
	if resp.Server.IdentityProviderDetails != nil {
		d.Set("invocation_role", aws.StringValue(resp.Server.IdentityProviderDetails.InvocationRole))
		d.Set("url", aws.StringValue(resp.Server.IdentityProviderDetails.Url))
	}
	d.Set("identity_provider_type", resp.Server.IdentityProviderType)
	d.Set("logging_role", resp.Server.LoggingRole)

	if err := d.Set("tags", tagsToMapTransferServer(resp.Server.Tags)); err != nil {
		return fmt.Errorf("Error setting tags: %s", err)
	}
	return nil
}

func resourceAwsTransferServerUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	updateFlag := false
	updateOpts := &transfer.UpdateServerInput{
		ServerId: aws.String(d.Id()),
	}

	if d.HasChange("logging_role") {
		updateFlag = true
		updateOpts.SetLoggingRole(d.Get("logging_role").(string))
	}

	if d.HasChange("invocation_role") || d.HasChange("url") {
		var identityProviderDetails *transfer.IdentityProviderDetails
		updateFlag = true
		if attr, ok := d.GetOk("invocation_role"); ok {
			identityProviderDetails.SetInvocationRole(attr.(string))
		}

		if attr, ok := d.GetOk("url"); ok {
			identityProviderDetails.SetUrl(attr.(string))
		}
		updateOpts.SetIdentityProviderDetails(identityProviderDetails)
	}

	if updateFlag {
		_, err := conn.UpdateServer(updateOpts)
		if err != nil {
			if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
				log.Printf("[WARN] Transfer Server (%s) not found, removing from state", d.Id())
				d.SetId("")
				return nil
			}
			return fmt.Errorf("error updating Transfer Server (%s): %s", d.Id(), err)
		}
	}

	if err := setTagsTransferServer(conn, d); err != nil {
		return fmt.Errorf("Error update tags: %s", err)
	}

	return resourceAwsTransferServerRead(d, meta)
}

func resourceAwsTransferServerDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn

	delOpts := &transfer.DeleteServerInput{
		ServerId: aws.String(d.Id()),
	}

	log.Printf("[DEBUG] Delete Transfer Server Option: %#v", delOpts)

	_, err := conn.DeleteServer(delOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			log.Printf("[WARN] Transfer Server (%s) not found, removing from state", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	return nil
}
<commit_msg>fix to follow coding policy<commit_after>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/transfer"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsTransferServer() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsTransferServerCreate,
		Read:   resourceAwsTransferServerRead,
		Update: resourceAwsTransferServerUpdate,
		Delete: resourceAwsTransferServerDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"endpoint": {
				Type:     schema.TypeString,
				Computed: true,
			},

			"invocation_role": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateArn,
			},

			"url": {
				Type:     schema.TypeString,
				Optional: true,
			},

			"identity_provider_type": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  transfer.IdentityProviderTypeServiceManaged,
				ValidateFunc: validation.StringInSlice([]string{
					transfer.IdentityProviderTypeServiceManaged,
					transfer.IdentityProviderTypeApiGateway,
				}, false),
			},

			"logging_role": {
				Type:         schema.TypeString,
				Optional:     true,
				ValidateFunc: validateArn,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsTransferServerCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	tags := tagsFromMapTransferServer(d.Get("tags").(map[string]interface{}))

	createOpts := &transfer.CreateServerInput{
		Tags: tags,
	}

	var identityProviderDetails *transfer.IdentityProviderDetails
	if attr, ok := d.GetOk("invocation_role"); ok {
		identityProviderDetails.InvocationRole = aws.String(attr.(string))
	}

	if attr, ok := d.GetOk("url"); ok {
		identityProviderDetails.Url = aws.String(attr.(string))
	}
	if identityProviderDetails != nil {
		createOpts.IdentityProviderDetails = identityProviderDetails
	}

	if attr, ok := d.GetOk("identity_provider_type"); ok {
		createOpts.IdentityProviderType = aws.String(attr.(string))
	}

	if attr, ok := d.GetOk("logging_role"); ok {
		createOpts.LoggingRole = aws.String(attr.(string))
	}

	log.Printf("[DEBUG] Create Transfer Server Option: %#v", createOpts)

	resp, err := conn.CreateServer(createOpts)
	if err != nil {
		return fmt.Errorf("Error creating Transfer Server: %s", err)
	}

	d.SetId(*resp.ServerId)

	return resourceAwsTransferServerRead(d, meta)
}

func resourceAwsTransferServerRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn

	descOpts := &transfer.DescribeServerInput{
		ServerId: aws.String(d.Id()),
	}

	log.Printf("[DEBUG] Describe Transfer Server Option: %#v", descOpts)

	resp, err := conn.DescribeServer(descOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			log.Printf("[WARN] Transfer Server (%s) not found, removing from state", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	endpoint := fmt.Sprintf("%s.server.transfer.%s.amazonaws.com", d.Id(), meta.(*AWSClient).region)
	log.Printf("[DEBUG] managed endpoint : %s", endpoint)

	d.Set("arn", resp.Server.Arn)
	d.Set("endpoint", endpoint)
	d.Set("invocation_role", "")
	d.Set("url", "")
	if resp.Server.IdentityProviderDetails != nil {
		d.Set("invocation_role", aws.StringValue(resp.Server.IdentityProviderDetails.InvocationRole))
		d.Set("url", aws.StringValue(resp.Server.IdentityProviderDetails.Url))
	}
	d.Set("identity_provider_type", resp.Server.IdentityProviderType)
	d.Set("logging_role", resp.Server.LoggingRole)

	if err := d.Set("tags", tagsToMapTransferServer(resp.Server.Tags)); err != nil {
		return fmt.Errorf("Error setting tags: %s", err)
	}
	return nil
}

func resourceAwsTransferServerUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn
	updateFlag := false
	updateOpts := &transfer.UpdateServerInput{
		ServerId: aws.String(d.Id()),
	}

	if d.HasChange("logging_role") {
		updateFlag = true
		updateOpts.LoggingRole = aws.String(d.Get("logging_role").(string))
	}

	if d.HasChange("invocation_role") || d.HasChange("url") {
		var identityProviderDetails *transfer.IdentityProviderDetails
		updateFlag = true
		if attr, ok := d.GetOk("invocation_role"); ok {
			identityProviderDetails.InvocationRole = aws.String(attr.(string))
		}

		if attr, ok := d.GetOk("url"); ok {
			identityProviderDetails.Url = aws.String(attr.(string))
		}
		updateOpts.IdentityProviderDetails = identityProviderDetails
	}

	if updateFlag {
		_, err := conn.UpdateServer(updateOpts)
		if err != nil {
			if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
				log.Printf("[WARN] Transfer Server (%s) not found, removing from state", d.Id())
				d.SetId("")
				return nil
			}
			return fmt.Errorf("error updating Transfer Server (%s): %s", d.Id(), err)
		}
	}

	if err := setTagsTransferServer(conn, d); err != nil {
		return fmt.Errorf("Error update tags: %s", err)
	}

	return resourceAwsTransferServerRead(d, meta)
}

func resourceAwsTransferServerDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).transferconn

	delOpts := &transfer.DeleteServerInput{
		ServerId: aws.String(d.Id()),
	}

	log.Printf("[DEBUG] Delete Transfer Server Option: %#v", delOpts)

	_, err := conn.DeleteServer(delOpts)
	if err != nil {
		if isAWSErr(err, transfer.ErrCodeResourceNotFoundException, "") {
			log.Printf("[WARN] Transfer Server (%s) not found, removing from state", d.Id())
			d.SetId("")
			return nil
		}
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package adaptor

import (
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/compose/transporter/pkg/message"
	"github.com/compose/transporter/pkg/pipe"
	gorethink "gopkg.in/dancannon/gorethink.v0"
)

// Rethinkdb is an adaptor that writes metrics to rethinkdb (http://rethinkdb.com/)
// An open-source distributed database
type Rethinkdb struct {
	// pull these in from the config
	uri *url.URL

	// save time by setting these once
	database string
	table    string

	debug bool
	tail  bool

	//
	pipe *pipe.Pipe
	path string

	// rethinkdb connection and options
	client *gorethink.Session
}

// RethinkdbConfig provides custom configuration options for the RethinkDB adapter
type RethinkdbConfig struct {
	URI       string `json:"uri" doc:"the uri to connect to, in the form rethink://user:password@host.example:28015/database"`
	Namespace string `json:"namespace" doc:"rethink namespace to read/write, in the form database.table"`
	Debug     bool   `json:"debug" doc:"if true, verbose debugging information is displayed"`
	Tail      bool   `json:"tail" doc:"if true, the RethinkDB table will be monitored for changes after copying the namespace"`
}

// NewRethinkdb creates a new Rethinkdb database adaptor
func NewRethinkdb(p *pipe.Pipe, path string, extra Config) (StopStartListener, error) {
	var (
		conf RethinkdbConfig
		err  error
	)
	if err = extra.Construct(&conf); err != nil {
		return nil, err
	}

	u, err := url.Parse(conf.URI)
	if err != nil {
		return nil, err
	}

	r := &Rethinkdb{
		uri:  u,
		pipe: p,
		path: path,
		tail: conf.Tail,
	}

	r.database, r.table, err = extra.splitNamespace()
	if err != nil {
		return r, err
	}
	r.debug = conf.Debug

	if r.debug {
		fmt.Printf("Connecting to %s (database %s, table %s)\n", r.uri.Host, r.database, r.table)
	}
	r.client, err = gorethink.Connect(gorethink.ConnectOpts{
		Address: r.uri.Host,
		MaxIdle: 10,
		Timeout: time.Second * 10,
	})
	if err != nil {
		return nil, fmt.Errorf("unable to connect: %s", err)
	}
	r.client.Use(r.database)

	return r, nil
}

// Start the adaptor as a source (not implemented)
func (r *Rethinkdb) Start() error {
	if r.debug {
		fmt.Printf("getting a changes cursor\n")
	}

	var ccursor *gorethink.Cursor
	ccursor, err := gorethink.Table(r.table).Changes().Run(r.client)
	if err != nil {
		r.pipe.Err <- err
		return err
	}

	// Send all of the rows currently in the table
	if err := r.sendAllDocuments(); err != nil {
		r.pipe.Err <- err
		return err
	}

	// Monitor for changes
	if r.tail {
		if err := r.sendChanges(ccursor); err != nil {
			r.pipe.Err <- err
			return err
		}
	}

	return nil
}

func (r *Rethinkdb) sendAllDocuments() error {
	if r.debug {
		fmt.Printf("sending all documents\n")
	}

	cursor, err := gorethink.Table(r.table).Run(r.client)
	if err != nil {
		return err
	}

	var doc map[string]interface{}
	for cursor.Next(&doc) {
		if stop := r.pipe.Stopped; stop {
			return nil
		}

		msg := message.NewMsg(message.Insert, r.prepareDocument(doc))
		r.pipe.Send(msg)
	}

	if err := cursor.Err(); err != nil {
		return err
	}

	return nil
}

func (r *Rethinkdb) sendChanges(ccursor *gorethink.Cursor) error {
	if r.debug {
		fmt.Printf("sending changes\n")
	}

	// { "new_val": {...}, "old_val": {...} }
	var change map[string]map[string]interface{}
	for ccursor.Next(&change) {
		if stop := r.pipe.Stopped; stop {
			return nil
		}

		var msg *message.Msg
		if change["old_val"] != nil && change["new_val"] != nil {
			msg = message.NewMsg(message.Update, r.prepareDocument(change["new_val"]))
		} else if change["new_val"] != nil {
			msg = message.NewMsg(message.Insert, r.prepareDocument(change["new_val"]))
		} else if change["old_val"] != nil {
			msg = message.NewMsg(message.Delete, r.prepareDocument(change["old_val"]))
		}

		if msg != nil {
			r.pipe.Send(msg)
		}
	}

	if err := ccursor.Err(); err != nil {
		r.pipe.Err <- err
		return err
	}

	return nil
}

// prepareDocument munges the RethinkDB document slightly to play well with the
// rest of the system. For instance, it copies the "id" field to "_id" which is
// expected by some other databases.
func (r *Rethinkdb) prepareDocument(document map[string]interface{}) map[string]interface{} {
	document["_id"] = document["id"]
	return document
}

// Listen start's the adaptor's listener
func (r *Rethinkdb) Listen() (err error) {
	r.recreateTable()
	return r.pipe.Listen(r.applyOp)
}

// Stop the adaptor
func (r *Rethinkdb) Stop() error {
	r.pipe.Stop()
	return nil
}

// applyOp applies one operation to the database
func (r *Rethinkdb) applyOp(msg *message.Msg) (*message.Msg, error) {
	var (
		resp gorethink.WriteResponse
		err  error
	)

	if !msg.IsMap() {
		r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (document must be a json document)", msg.Data)
		return msg, nil
	}
	doc := msg.Map()

	switch msg.Op {
	case message.Delete:
		id, err := msg.IDString("id")
		if err != nil {
			r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (cannot delete an object with a nil id)", msg.Data)
			return msg, nil
		}
		resp, err = gorethink.Table(r.table).Get(id).Delete().RunWrite(r.client)
	case message.Insert:
		resp, err = gorethink.Table(r.table).Insert(doc).RunWrite(r.client)
	case message.Update:
		resp, err = gorethink.Table(r.table).Insert(doc, gorethink.InsertOpts{Conflict: "replace"}).RunWrite(r.client)
	}
	if err != nil {
		r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (%s)", err)
		return msg, nil
	}

	err = r.handleResponse(&resp)
	if err != nil {
		r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (%s)", err)
	}

	return msg, nil
}

func (r *Rethinkdb) recreateTable() {
	if r.debug {
		fmt.Printf("dropping and creating table '%s' on database '%s'\n", r.table, r.database)
	}
	gorethink.Db(r.database).TableDrop(r.table).RunWrite(r.client)
	gorethink.Db(r.database).TableCreate(r.table).RunWrite(r.client)
}

// handleresponse takes the rethink response and turn it into something we can consume elsewhere
func (r *Rethinkdb) handleResponse(resp *gorethink.WriteResponse) error {
	if resp.Errors != 0 {
		if !strings.Contains(resp.FirstError, "Duplicate primary key") { // we don't care about this error
			if r.debug {
				fmt.Printf("Reported %d errors\n", resp.Errors)
			}
			return fmt.Errorf("%s\n%s", "problem inserting docs", resp.FirstError)
		}
	}
	return nil
}
<commit_msg>This method is implemented now<commit_after>package adaptor

import (
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/compose/transporter/pkg/message"
	"github.com/compose/transporter/pkg/pipe"
	gorethink "gopkg.in/dancannon/gorethink.v0"
)

// Rethinkdb is an adaptor that writes metrics to rethinkdb (http://rethinkdb.com/)
// An open-source distributed database
type Rethinkdb struct {
	// pull these in from the config
	uri *url.URL

	// save time by setting these once
	database string
	table    string

	debug bool
	tail  bool

	//
	pipe *pipe.Pipe
	path string

	// rethinkdb connection and options
	client *gorethink.Session
}

// RethinkdbConfig provides custom configuration options for the RethinkDB adapter
type RethinkdbConfig struct {
	URI       string `json:"uri" doc:"the uri to connect to, in the form rethink://user:password@host.example:28015/database"`
	Namespace string `json:"namespace" doc:"rethink namespace to read/write, in the form database.table"`
	Debug     bool   `json:"debug" doc:"if true, verbose debugging information is displayed"`
	Tail      bool   `json:"tail" doc:"if true, the RethinkDB table will be monitored for changes after copying the namespace"`
}

// NewRethinkdb creates a new Rethinkdb database adaptor
func NewRethinkdb(p *pipe.Pipe, path string, extra Config) (StopStartListener, error) {
	var (
		conf RethinkdbConfig
		err  error
	)
	if err = extra.Construct(&conf); err != nil {
		return nil, err
	}

	u, err := url.Parse(conf.URI)
	if err != nil {
		return nil, err
	}

	r := &Rethinkdb{
		uri:  u,
		pipe: p,
		path: path,
		tail: conf.Tail,
	}

	r.database, r.table, err = extra.splitNamespace()
	if err != nil {
		return r, err
	}
	r.debug = conf.Debug

	if r.debug {
		fmt.Printf("Connecting to %s (database %s, table %s)\n", r.uri.Host, r.database, r.table)
	}
	r.client, err = gorethink.Connect(gorethink.ConnectOpts{
		Address: r.uri.Host,
		MaxIdle: 10,
		Timeout: time.Second * 10,
	})
	if err != nil {
		return nil, fmt.Errorf("unable to connect: %s", err)
	}
	r.client.Use(r.database)

	return r, nil
}

// Start the adaptor as a source
func (r *Rethinkdb) Start() error {
	if r.debug {
		fmt.Printf("getting a changes cursor\n")
	}

	var ccursor *gorethink.Cursor
	ccursor, err := gorethink.Table(r.table).Changes().Run(r.client)
	if err != nil {
		r.pipe.Err <- err
		return err
	}

	// Send all of the rows currently in the table
	if err := r.sendAllDocuments(); err != nil {
		r.pipe.Err <- err
		return err
	}

	// Monitor for changes
	if r.tail {
		if err := r.sendChanges(ccursor); err != nil {
			r.pipe.Err <- err
			return err
		}
	}

	return nil
}

func (r *Rethinkdb) sendAllDocuments() error {
	if r.debug {
		fmt.Printf("sending all documents\n")
	}

	cursor, err := gorethink.Table(r.table).Run(r.client)
	if err != nil {
		return err
	}

	var doc map[string]interface{}
	for cursor.Next(&doc) {
		if stop := r.pipe.Stopped; stop {
			return nil
		}

		msg := message.NewMsg(message.Insert, r.prepareDocument(doc))
		r.pipe.Send(msg)
	}

	if err := cursor.Err(); err != nil {
		return err
	}

	return nil
}

func (r *Rethinkdb) sendChanges(ccursor *gorethink.Cursor) error {
	if r.debug {
		fmt.Printf("sending changes\n")
	}

	// { "new_val": {...}, "old_val": {...} }
	var change map[string]map[string]interface{}
	for ccursor.Next(&change) {
		if stop := r.pipe.Stopped; stop {
			return nil
		}

		var msg *message.Msg
		if change["old_val"] != nil && change["new_val"] != nil {
			msg = message.NewMsg(message.Update, r.prepareDocument(change["new_val"]))
		} else if change["new_val"] != nil {
			msg = message.NewMsg(message.Insert, r.prepareDocument(change["new_val"]))
		} else if change["old_val"] != nil {
			msg = message.NewMsg(message.Delete, r.prepareDocument(change["old_val"]))
		}

		if msg != nil {
			r.pipe.Send(msg)
		}
	}

	if err := ccursor.Err(); err != nil {
		r.pipe.Err <- err
		return err
	}

	return nil
}

// prepareDocument munges the RethinkDB document slightly to play well with the
// rest of the system. For instance, it copies the "id" field to "_id" which is
// expected by some other databases.
func (r *Rethinkdb) prepareDocument(document map[string]interface{}) map[string]interface{} {
	document["_id"] = document["id"]
	return document
}

// Listen start's the adaptor's listener
func (r *Rethinkdb) Listen() (err error) {
	r.recreateTable()
	return r.pipe.Listen(r.applyOp)
}

// Stop the adaptor
func (r *Rethinkdb) Stop() error {
	r.pipe.Stop()
	return nil
}

// applyOp applies one operation to the database
func (r *Rethinkdb) applyOp(msg *message.Msg) (*message.Msg, error) {
	var (
		resp gorethink.WriteResponse
		err  error
	)

	if !msg.IsMap() {
		r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (document must be a json document)", msg.Data)
		return msg, nil
	}
	doc := msg.Map()

	switch msg.Op {
	case message.Delete:
		id, err := msg.IDString("id")
		if err != nil {
			r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (cannot delete an object with a nil id)", msg.Data)
			return msg, nil
		}
		resp, err = gorethink.Table(r.table).Get(id).Delete().RunWrite(r.client)
	case message.Insert:
		resp, err = gorethink.Table(r.table).Insert(doc).RunWrite(r.client)
	case message.Update:
		resp, err = gorethink.Table(r.table).Insert(doc, gorethink.InsertOpts{Conflict: "replace"}).RunWrite(r.client)
	}
	if err != nil {
		r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (%s)", err)
		return msg, nil
	}

	err = r.handleResponse(&resp)
	if err != nil {
		r.pipe.Err <- NewError(ERROR, r.path, "rethinkdb error (%s)", err)
	}

	return msg, nil
}

func (r *Rethinkdb) recreateTable() {
	if r.debug {
		fmt.Printf("dropping and creating table '%s' on database '%s'\n", r.table, r.database)
	}
	gorethink.Db(r.database).TableDrop(r.table).RunWrite(r.client)
	gorethink.Db(r.database).TableCreate(r.table).RunWrite(r.client)
}

// handleresponse takes the rethink response and turn it into something we can consume elsewhere
func (r *Rethinkdb) handleResponse(resp *gorethink.WriteResponse) error {
	if resp.Errors != 0 {
		if !strings.Contains(resp.FirstError, "Duplicate primary key") { // we don't care about this error
			if r.debug {
				fmt.Printf("Reported %d errors\n", resp.Errors)
			}
			return fmt.Errorf("%s\n%s", "problem inserting docs", resp.FirstError)
		}
	}
	return nil
}
<|endoftext|>