<commit_before>package bitread_test

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"github.com/markus-wa/gobitread"
	"math"
	"testing"
)

func TestReadBit(t *testing.T) {
	b := make([]byte, 0xff)
	for n := byte(0); n < byte(len(b)); n++ {
		b[n] = n
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	res := make([]byte, 8)
	var exp string
	for i := 0; i < len(b); i++ {
		for i := 0; i < 8; i++ {
			// Least significant bit first
			if br.ReadBit() {
				res[7-i] = "1"[0]
			} else {
				res[7-i] = "0"[0]
			}
		}

		exp = fmt.Sprintf("%b", b[i])
		// Pad cut off bits
		for len(exp) < 8 {
			exp = "0" + exp
		}
		if string(res) != exp {
			t.Fatalf("Expected %s got %s", exp, res)
		}
	}
}

func TestReadBytes(t *testing.T) {
	b := make([]byte, 1<<8)
	for n := 0; n < len(b); n++ {
		b[n] = byte(n)
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	r := br.ReadBytes(len(b))

	for i := 0; i < len(b); i++ {
		if b[i] != r[i] {
			t.Fatalf("Expected %b got %b", b[i], r[i])
		}
	}
}

func TestReadBytesBuffered(t *testing.T) {
	b := make([]byte, 1<<8)
	for n := 0; n < len(b); n++ {
		b[n] = byte(n)
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), len(b)<<1) // Make buffer large enough to hold all bytes

	r := br.ReadBytes(len(b))

	for i := 0; i < len(b); i++ {
		if b[i] != r[i] {
			t.Fatalf("Expected %b got %b", b[i], r[i])
		}
	}
}

func TestReadCString(t *testing.T) {
	s := "test"
	b := []byte(s)

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	r := br.ReadCString(len(s))
	if r != s {
		t.Fatalf("Expected %q got %q", s, r)
	}
}

func TestBitsToByte(t *testing.T) {
	b := make([]byte, 0x0f)
	for n := byte(0); n < byte(len(b)); n++ {
		b[n] = n
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	for i := byte(0); i < byte(len(b)); i++ {
		r := br.ReadBitsToByte(4)
		if r != b[i] {
			t.Fatalf("Expected %b got %b", b[i], r)
		}
		br.ReadBitsToByte(4)
	}
}

func TestReadBits(t *testing.T) {
	nums := []uint64{0xa3dc, 0x48c1}
	b := make([]byte, len(nums)<<3)
	for i := 0; i < len(nums); i++ {
		binary.LittleEndian.PutUint64(b[i<<3:], nums[i])
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	for i := 0; i < len(nums); i++ {
		r := binary.LittleEndian.Uint64(br.ReadBits(64))
		if r != nums[i] {
			t.Fatalf("Expected %d got %d", nums[0], r)
		}
	}
}

func TestReadInt(t *testing.T) {
	nums := []uint32{0, math.MaxUint32, 0x61cb83f0}
	b := make([]byte, len(nums)<<2)
	for i := 0; i < len(nums); i++ {
		binary.LittleEndian.PutUint32(b[i<<2:], nums[i])
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	for i := 0; i < len(nums); i++ {
		r := br.ReadInt(32)
		if r != uint(nums[i]) {
			t.Fatalf("Expected %q got %q", nums[i], r)
		}
	}
}

func TestReadSignedInt(t *testing.T) {
	nums := []int32{math.MaxInt32, math.MinInt32, 0, 0x4ac71bf}
	b := make([]byte, len(nums)<<2)
	for i := 0; i < len(nums); i++ {
		binary.LittleEndian.PutUint32(b[i<<2:], uint32(nums[i]))
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	for i := 0; i < len(nums); i++ {
		r := br.ReadSignedInt(32)
		if r != int(nums[i]) {
			t.Fatalf("Expected %q got %q", nums[i], r)
		}
	}
}

func TestPositions(t *testing.T) {
	b := []byte{0xac}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	// Positions
	if br.LazyPosition() != 0 {
		t.Fatalf("LazyPosition is %d expected 0", br.LazyPosition())
	}

	if br.ActualPosition() != 0 {
		t.Fatalf("ActualPostition is %d expected 0", br.ActualPosition())
	}

	br.ReadBit()

	if br.LazyPosition() != 0 {
		t.Fatalf("LazyPosition is %d expected 0", br.LazyPosition())
	}

	if br.ActualPosition() != 1 {
		t.Fatalf("ActualPostition is %d expected 1", br.ActualPosition())
	}
}

// Tests Open, Close & Chunks
func TestLifecycle(t *testing.T) {
	b := []byte("a")

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	br.BeginChunk(2)

	br.ReadBit()
	if br.ChunkFinished() {
		t.Error("ChunkFinished() should have returned false but returned true")
	}

	br.ReadBit()
	if !br.ChunkFinished() {
		t.Error("ChunkFinished() should have returned true but returned false")
	}

	br.EndChunk()
	br.Close()

	// Test reopen
	br.Open(bytes.NewReader(b), 32)
	r := br.ReadBytes(len(b))
	for i := 0; i < len(b); i++ {
		if r[i] != b[i] {
			t.Fatal("Expected", r[i], "got", b[i])
		}
	}
}

// Tests seeking at EndChunk()
func TestChunkSeek(t *testing.T) {
	b := make([]byte, 0xff)
	for n := byte(0); n < byte(len(b)); n++ {
		b[n] = n
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	br.BeginChunk(2)
	br.ReadBit()
	br.EndChunk()

	if br.ActualPosition() != 2 {
		t.Errorf("ActualPostition is %d expected 2", br.ActualPosition())
	}

	// Beyond buffer
	br.BeginChunk(16 << 3)
	br.EndChunk()

	// Check if lazyPos got updated when refilling the buffer
	if br.LazyPosition() != 16<<3 {
		t.Errorf("LazyPosition is %d expected %d", br.LazyPosition(), 16<<3)
	}

	if br.ActualPosition() != 2+(16<<3) {
		t.Errorf("ActualPostition is %d expected %d", br.ActualPosition(), 2+(16<<3))
	}

	// Get to an even position (in front of 18th byte)
	br.ReadBits(6)
	// Read 18th byte
	r := br.ReadSingleByte()
	if r != 17 {
		t.Fatalf("Expected 18th byte (17) got %d", r)
	}
}

func TestChunkExceeded(t *testing.T) {
	b := []byte("a")

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	br.BeginChunk(0)
	br.ReadBit()
	func() {
		defer func() {
			err := recover()
			if err == nil {
				t.Error("EndChunk() should have paniced after reading beyond chunk boundary")
			}
		}()
		br.EndChunk()
	}()
}

// TODO: The following type of tests
// Reading beyond a chunk boundary
// Not reading all the data in a chunk (seeking)
// With and without io.Seeker as underlying
// With data in sled
// Reopen after close

// TODO: Add Skip function (alias for private advance()) or make advance public? - BeginChunk followed by endchunk already does this tho - with more overhead
<commit_msg>Changed string to char in test<commit_after>package bitread_test

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"github.com/markus-wa/gobitread"
	"math"
	"testing"
)

func TestReadBit(t *testing.T) {
	b := make([]byte, 0xff)
	for n := byte(0); n < byte(len(b)); n++ {
		b[n] = n
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	res := make([]byte, 8)
	var exp string
	for i := 0; i < len(b); i++ {
		for i := 0; i < 8; i++ {
			// Least significant bit first
			if br.ReadBit() {
				res[7-i] = '1'
			} else {
				res[7-i] = '0'
			}
		}

		exp = fmt.Sprintf("%b", b[i])
		// Pad cut off bits
		for len(exp) < 8 {
			exp = "0" + exp
		}
		if string(res) != exp {
			t.Fatalf("Expected %s got %s", exp, res)
		}
	}
}

func TestReadBytes(t *testing.T) {
	b := make([]byte, 1<<8)
	for n := 0; n < len(b); n++ {
		b[n] = byte(n)
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	r := br.ReadBytes(len(b))

	for i := 0; i < len(b); i++ {
		if b[i] != r[i] {
			t.Fatalf("Expected %b got %b", b[i], r[i])
		}
	}
}

func TestReadBytesBuffered(t *testing.T) {
	b := make([]byte, 1<<8)
	for n := 0; n < len(b); n++ {
		b[n] = byte(n)
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), len(b)<<1) // Make buffer large enough to hold all bytes

	r := br.ReadBytes(len(b))

	for i := 0; i < len(b); i++ {
		if b[i] != r[i] {
			t.Fatalf("Expected %b got %b", b[i], r[i])
		}
	}
}

func TestReadCString(t *testing.T) {
	s := "test"
	b := []byte(s)

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	r := br.ReadCString(len(s))
	if r != s {
		t.Fatalf("Expected %q got %q", s, r)
	}
}

func TestBitsToByte(t *testing.T) {
	b := make([]byte, 0x0f)
	for n := byte(0); n < byte(len(b)); n++ {
		b[n] = n
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	for i := byte(0); i < byte(len(b)); i++ {
		r := br.ReadBitsToByte(4)
		if r != b[i] {
			t.Fatalf("Expected %b got %b", b[i], r)
		}
		br.ReadBitsToByte(4)
	}
}

func TestReadBits(t *testing.T) {
	nums := []uint64{0xa3dc, 0x48c1}
	b := make([]byte, len(nums)<<3)
	for i := 0; i < len(nums); i++ {
		binary.LittleEndian.PutUint64(b[i<<3:], nums[i])
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	for i := 0; i < len(nums); i++ {
		r := binary.LittleEndian.Uint64(br.ReadBits(64))
		if r != nums[i] {
			t.Fatalf("Expected %d got %d", nums[0], r)
		}
	}
}

func TestReadInt(t *testing.T) {
	nums := []uint32{0, math.MaxUint32, 0x61cb83f0}
	b := make([]byte, len(nums)<<2)
	for i := 0; i < len(nums); i++ {
		binary.LittleEndian.PutUint32(b[i<<2:], nums[i])
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	for i := 0; i < len(nums); i++ {
		r := br.ReadInt(32)
		if r != uint(nums[i]) {
			t.Fatalf("Expected %q got %q", nums[i], r)
		}
	}
}

func TestReadSignedInt(t *testing.T) {
	nums := []int32{math.MaxInt32, math.MinInt32, 0, 0x4ac71bf}
	b := make([]byte, len(nums)<<2)
	for i := 0; i < len(nums); i++ {
		binary.LittleEndian.PutUint32(b[i<<2:], uint32(nums[i]))
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	for i := 0; i < len(nums); i++ {
		r := br.ReadSignedInt(32)
		if r != int(nums[i]) {
			t.Fatalf("Expected %q got %q", nums[i], r)
		}
	}
}

func TestPositions(t *testing.T) {
	b := []byte{0xac}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	// Positions
	if br.LazyPosition() != 0 {
		t.Fatalf("LazyPosition is %d expected 0", br.LazyPosition())
	}

	if br.ActualPosition() != 0 {
		t.Fatalf("ActualPostition is %d expected 0", br.ActualPosition())
	}

	br.ReadBit()

	if br.LazyPosition() != 0 {
		t.Fatalf("LazyPosition is %d expected 0", br.LazyPosition())
	}

	if br.ActualPosition() != 1 {
		t.Fatalf("ActualPostition is %d expected 1", br.ActualPosition())
	}
}

// Tests Open, Close & Chunks
func TestLifecycle(t *testing.T) {
	b := []byte("a")

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	br.BeginChunk(2)

	br.ReadBit()
	if br.ChunkFinished() {
		t.Error("ChunkFinished() should have returned false but returned true")
	}

	br.ReadBit()
	if !br.ChunkFinished() {
		t.Error("ChunkFinished() should have returned true but returned false")
	}

	br.EndChunk()
	br.Close()

	// Test reopen
	br.Open(bytes.NewReader(b), 32)
	r := br.ReadBytes(len(b))
	for i := 0; i < len(b); i++ {
		if r[i] != b[i] {
			t.Fatal("Expected", r[i], "got", b[i])
		}
	}
}

// Tests seeking at EndChunk()
func TestChunkSeek(t *testing.T) {
	b := make([]byte, 0xff)
	for n := byte(0); n < byte(len(b)); n++ {
		b[n] = n
	}

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 16)

	br.BeginChunk(2)
	br.ReadBit()
	br.EndChunk()

	if br.ActualPosition() != 2 {
		t.Errorf("ActualPostition is %d expected 2", br.ActualPosition())
	}

	// Beyond buffer
	br.BeginChunk(16 << 3)
	br.EndChunk()

	// Check if lazyPos got updated when refilling the buffer
	if br.LazyPosition() != 16<<3 {
		t.Errorf("LazyPosition is %d expected %d", br.LazyPosition(), 16<<3)
	}

	if br.ActualPosition() != 2+(16<<3) {
		t.Errorf("ActualPostition is %d expected %d", br.ActualPosition(), 2+(16<<3))
	}

	// Get to an even position (in front of 18th byte)
	br.ReadBits(6)
	// Read 18th byte
	r := br.ReadSingleByte()
	if r != 17 {
		t.Fatalf("Expected 18th byte (17) got %d", r)
	}
}

func TestChunkExceeded(t *testing.T) {
	b := []byte("a")

	br := new(bitread.BitReader)
	br.Open(bytes.NewReader(b), 32)

	br.BeginChunk(0)
	br.ReadBit()
	func() {
		defer func() {
			err := recover()
			if err == nil {
				t.Error("EndChunk() should have paniced after reading beyond chunk boundary")
			}
		}()
		br.EndChunk()
	}()
}

// TODO: The following type of tests
// Reading beyond a chunk boundary
// Not reading all the data in a chunk (seeking)
// With and without io.Seeker as underlying
// With data in sled
// Reopen after close

// TODO: Add Skip function (alias for private advance()) or make advance public? - BeginChunk followed by endchunk already does this tho - with more overhead
<|endoftext|>


<commit_before>package renter

import (
	"path/filepath"
	"testing"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/modules/consensus"
	"github.com/NebulousLabs/Sia/modules/gateway"
	"github.com/NebulousLabs/Sia/modules/hostdb"
	"github.com/NebulousLabs/Sia/modules/tester"
	"github.com/NebulousLabs/Sia/modules/transactionpool"
	"github.com/NebulousLabs/Sia/modules/wallet"
)

// A RenterTester contains a consensus tester and a renter, and provides a set
// of helper functions for testing the renter without needing other modules
// such as the host.
type RenterTester struct {
	*consensus.ConsensusTester
	*Renter
}

// CreateHostTester initializes a HostTester.
func CreateRenterTester(name string, t *testing.T) (rt *RenterTester) {
	testdir := tester.TempDir("renter", name)
	cs, err := consensus.New(filepath.Join(testdir, modules.ConsensusDir))
	if err != nil {
		t.Fatal(err)
	}
	ct := consensus.NewConsensusTester(t, cs)
	g, err := gateway.New(":0", ct.State, filepath.Join(testdir, modules.GatewayDir))
	if err != nil {
		t.Fatal(err)
	}
	hdb, err := hostdb.New(ct.State, g)
	if err != nil {
		t.Fatal(err)
	}
	tp, err := transactionpool.New(ct.State, g)
	if err != nil {
		t.Fatal(err)
	}
	w, err := wallet.New(ct.State, tp, filepath.Join(testdir, modules.WalletDir))
	if err != nil {
		t.Fatal(err)
	}
	r, err := New(ct.State, g, hdb, w, filepath.Join(testdir, modules.RenterDir))
	if err != nil {
		t.Fatal(err)
	}

	rt = new(RenterTester)
	rt.ConsensusTester = ct
	rt.Renter = r
	return
}

// TestSaveLoad tests that saving and loading a Renter restores its data.
func TestSaveLoad(t *testing.T) {
	rt := CreateRenterTester("TestSaveLoad", t)
	err := rt.save()
	if err != nil {
		rt.Fatal(err)
	}
	err = rt.load()
	if err != nil {
		rt.Fatal(err)
	}
}
<commit_msg>remove all renter testing<commit_after><|endoftext|>


<commit_before>package firego

import (
	"bufio"
	"encoding/json"
	"log"
	"net/http"
	"strings"
)

const (
	// EventTypePut is the event type sent when new data is inserted to the
	// Firebase instance.
	EventTypePut = "put"
	// EventTypePatch is the event type sent when data at the Firebase instance is
	// updated.
	EventTypePatch = "patch"
	// EventTypeError is the event type sent when an unknown error is encountered.
	EventTypeError = "event_error"
	// EventTypeAuthRevoked is the event type sent when the supplied auth parameter
	// is no longer valid.
	EventTypeAuthRevoked = "auth_revoked"

	eventTypeKeepAlive  = "keep-alive"
	eventTypeCancel     = "cancel"
	eventTypeRulesDebug = "rules_debug"
)

// Event represents a notification received when watching a
// firebase reference.
type Event struct {
	// Type of event that was received
	Type string
	// Path to the data that changed
	Path string
	// Data that changed
	Data interface{}

	rawData string
}

// Value converts the raw payload of the event into the given interface.
func (e Event) Value(v interface{}) error {
	var tmp struct {
		Data interface{} `json:"data"`
	}
	tmp.Data = &v
	return json.Unmarshal([]byte(e.rawData), &tmp)
}

// StopWatching stops tears down all connections that are watching.
func (fb *Firebase) StopWatching() {
	if fb.isWatching() {
		// signal connection to terminal
		fb.stopWatching <- struct{}{}
		// flip the bit back to not watching
		fb.setWatching(false)
	}
}

func (fb *Firebase) isWatching() bool {
	fb.watchMtx.Lock()
	v := fb.watching
	fb.watchMtx.Unlock()
	return v
}

func (fb *Firebase) setWatching(v bool) {
	fb.watchMtx.Lock()
	fb.watching = v
	fb.watchMtx.Unlock()
}

// Watch listens for changes on a firebase instance and
// passes over to the given chan.
//
// Only one connection can be established at a time. The
// second call to this function without a call to fb.StopWatching
// will close the channel given and return nil immediately.
func (fb *Firebase) Watch(notifications chan Event) error {
	if fb.isWatching() {
		close(notifications)
		return nil
	}
	// set watching flag
	fb.setWatching(true)

	stop := make(chan struct{})
	events, err := fb.watch(stop)
	if err != nil {
		return err
	}

	var closedManually bool

	// monitor the stopWatching channel
	// if we're told to stop, close the response Body
	go func() {
		<-fb.stopWatching

		closedManually = true
		close(stop)
	}()

	go func() {
		for event := range events {
			if event.Type == EventTypeError && closedManually {
				break
			}

			notifications <- event
		}

		close(notifications)
	}()

	return nil
}

func (fb *Firebase) watch(stop chan struct{}) (chan Event, error) {
	// build SSE request
	req, err := http.NewRequest("GET", fb.String(), nil)
	if err != nil {
		fb.setWatching(false)
		return nil, err
	}
	req.Header.Add("Accept", "text/event-stream")

	// do request
	resp, err := fb.client.Do(req)
	if err != nil {
		fb.setWatching(false)
		return nil, err
	}

	notifications := make(chan Event)

	go func() {
		<-stop
		defer resp.Body.Close()
	}()

	// start parsing response body
	go func() {

		// build scanner for response body
		scanner := bufio.NewReader(resp.Body)
		var scanErr error

	scanning:
		for scanErr == nil {
			// split event string
			// 		event: put
			// 		data: {"path":"/","data":{"foo":"bar"}}

			var evt string
			var dat string

			// scan for 'event:'
			evt, scanErr = scanner.ReadString('\n')
			if scanErr != nil {
				break scanning
			}
			evt = strings.TrimSuffix(evt, "\n")

			// scan for 'data:'
			dat, scanErr = scanner.ReadString('\n')
			if scanErr != nil {
				break scanning
			}
			dat = strings.TrimSuffix(dat, "\n")

			// strip off last '\n'
			_, scanErr = scanner.ReadString('\n')
			if scanErr != nil {
				break scanning
			}

			// create a base event
			event := Event{
				Type:    strings.Replace(evt, "event: ", "", 1),
				rawData: strings.Replace(dat, "data: ", "", 1),
			}

			// should be reacting differently based off the type of event
			switch event.Type {
			case EventTypePut, EventTypePatch:
				// we've got extra data we've got to parse
				var data map[string]interface{}
				if err := json.Unmarshal([]byte(strings.Replace(dat, "data: ", "", 1)), &data); err != nil {
					scanErr = err
					break scanning
				}

				// set the extra fields
				event.Path = data["path"].(string)
				event.Data = data["data"]

				// ship it
				notifications <- event
			case eventTypeKeepAlive:
				// received ping - nothing to do here
			case eventTypeCancel:
				// The data for this event is null
				// This event will be sent if the Security and Firebase Rules
				// cause a read at the requested location to no longer be allowed

				// send the cancel event
				notifications <- event
				break scanning
			case EventTypeAuthRevoked:
				// The data for this event is a string indicating that a the credential has expired
				// This event will be sent when the supplied auth parameter is no longer valid
				event.Data = strings.Replace(dat, "data: ", "", 1)
				notifications <- event
				break scanning
			case eventTypeRulesDebug:
				log.Printf("Rules-Debug: %s\n%s\n", evt, dat)
			}
		}

		if scanErr != nil {
			notifications <- Event{
				Type: EventTypeError,
				Data: scanErr,
			}
		}

		// cleanup routines
		close(notifications)
	}()
	return notifications, nil
}
<commit_msg>cleaning up watch function a bit<commit_after>package firego

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"log"
	"net/http"
)

const (
	// EventTypePut is the event type sent when new data is inserted to the
	// Firebase instance.
	EventTypePut = "put"
	// EventTypePatch is the event type sent when data at the Firebase instance is
	// updated.
	EventTypePatch = "patch"
	// EventTypeError is the event type sent when an unknown error is encountered.
	EventTypeError = "event_error"
	// EventTypeAuthRevoked is the event type sent when the supplied auth parameter
	// is no longer valid.
	EventTypeAuthRevoked = "auth_revoked"

	eventTypeKeepAlive  = "keep-alive"
	eventTypeCancel     = "cancel"
	eventTypeRulesDebug = "rules_debug"
)

// Event represents a notification received when watching a
// firebase reference.
type Event struct {
	// Type of event that was received
	Type string
	// Path to the data that changed
	Path string
	// Data that changed
	Data interface{}

	rawData []byte
}

// Value converts the raw payload of the event into the given interface.
func (e Event) Value(v interface{}) error {
	var tmp struct {
		Data interface{} `json:"data"`
	}
	tmp.Data = &v
	return json.Unmarshal(e.rawData, &tmp)
}

// StopWatching stops tears down all connections that are watching.
func (fb *Firebase) StopWatching() {
	fb.watchMtx.Lock()
	defer fb.watchMtx.Unlock()

	if fb.watching {
		// flip the bit back to not watching
		fb.watching = false
		// signal connection to terminal
		fb.stopWatching <- struct{}{}
	}
}

func (fb *Firebase) setWatching(v bool) {
	fb.watchMtx.Lock()
	fb.watching = v
	fb.watchMtx.Unlock()
}

// Watch listens for changes on a firebase instance and
// passes over to the given chan.
//
// Only one connection can be established at a time. The
// second call to this function without a call to fb.StopWatching
// will close the channel given and return nil immediately.
func (fb *Firebase) Watch(notifications chan Event) error {
	fb.watchMtx.Lock()
	if fb.watching {
		fb.watchMtx.Unlock()
		close(notifications)
		return nil
	}
	fb.watching = true
	fb.watchMtx.Unlock()

	stop := make(chan struct{})
	events, err := fb.watch(stop)
	if err != nil {
		return err
	}

	var closedManually bool

	go func() {
		<-fb.stopWatching
		closedManually = true
		stop <- struct{}{}
	}()

	go func() {
		defer func() {
			close(notifications)
			close(stop)
		}()

		for event := range events {
			if closedManually {
				return
			}

			notifications <- event
		}
	}()

	return nil
}

func readLine(rdr *bufio.Reader, prefix string) ([]byte, error) {
	// read event: line
	line, err := rdr.ReadBytes('\n')
	if err != nil {
		return nil, err
	}

	// empty line check for empty prefix
	if len(prefix) == 0 {
		line = bytes.TrimSpace(line)
		if len(line) != 0 {
			return nil, errors.New("expected empty line")
		}
		return line, nil
	}

	// check line has event prefix
	if !bytes.HasPrefix(line, []byte(prefix)) {
		return nil, errors.New("missing prefix")
	}

	// trim space
	line = line[len(prefix):]
	return bytes.TrimSpace(line), nil
}

func (fb *Firebase) watch(stop chan struct{}) (chan Event, error) {
	// build SSE request
	req, err := http.NewRequest("GET", fb.String(), nil)
	if err != nil {
		fb.setWatching(false)
		return nil, err
	}
	req.Header.Add("Accept", "text/event-stream")

	// do request
	resp, err := fb.client.Do(req)
	if err != nil {
		fb.setWatching(false)
		return nil, err
	}

	notifications := make(chan Event)

	go func() {
		<-stop
		resp.Body.Close()
	}()

	// start parsing response body
	go func() {
		defer func() {
			resp.Body.Close()
			close(notifications)
		}()

		// build scanner for response body
		scanner := bufio.NewReader(resp.Body)
		sendError := func(err error) {
			notifications <- Event{
				Type: EventTypeError,
				Data: err,
			}
		}
		for {
			// scan for 'event:'
			evt, err := readLine(scanner, "event: ")
			if err != nil {
				sendError(err)
				return
			}

			// scan for 'data:'
			dat, err := readLine(scanner, "data: ")
			if err != nil {
				sendError(err)
				return
			}

			// read the empty line
			_, err = readLine(scanner, "")
			if err != nil {
				sendError(err)
				return
			}

			// create a base event
			event := Event{
				Type:    string(evt),
				Data:    string(dat),
				rawData: dat,
			}

			// should be reacting differently based off the type of event
			switch event.Type {
			case EventTypePut, EventTypePatch:
				// we've got extra data we've got to parse
				var data map[string]interface{}
				if err := json.Unmarshal(event.rawData, &data); err != nil {
					sendError(err)
					return
				}

				// set the extra fields
				event.Path = data["path"].(string)
				event.Data = data["data"]

				// ship it
				notifications <- event
			case eventTypeKeepAlive:
				// received ping - nothing to do here
			case eventTypeCancel:
				// The data for this event is null
				// This event will be sent if the Security and Firebase Rules
				// cause a read at the requested location to no longer be allowed

				// send the cancel event
				notifications <- event
				return
			case EventTypeAuthRevoked:
				// The data for this event is a string indicating that a the credential has expired
				// This event will be sent when the supplied auth parameter is no longer valid
				notifications <- event
				return
			case eventTypeRulesDebug:
				log.Printf("Rules-Debug: %s\n%s\n", evt, dat)
			}
		}
	}()
	return notifications, nil
}
<|endoftext|>


<commit_before>// TODO: avoid creating two instances of the fsnotify.Watcher struct
package tail

import (
	"github.com/howeyc/fsnotify"
	"os"
	"path/filepath"
	"time"
)

type FileWatcher interface {
	BlockUntilExists() error
	ChangeEvents() chan bool
}

// FileWatcher monitors file-level events
type InotifyFileWatcher struct {
	Filename string
}

func NewInotifyFileWatcher(filename string) *InotifyFileWatcher {
	fw := &InotifyFileWatcher{filename}
	return fw
}

// BlockUntilExists blocks until the file comes into existence. If the
// file already exists, then block until it is created again.
func (fw *InotifyFileWatcher) BlockUntilExists() error {
	w, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	defer w.Close()
	err = w.WatchFlags(filepath.Dir(fw.Filename), fsnotify.FSN_CREATE)
	if err != nil {
		return err
	}
	defer w.RemoveWatch(filepath.Dir(fw.Filename))
	<-w.Event
	return nil
}

// ChangeEvents returns a channel that gets updated when the file is ready to be read.
func (fw *InotifyFileWatcher) ChangeEvents() chan bool {
	w, err := fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
	err = w.Watch(fw.Filename)
	if err != nil {
		panic(err)
	}

	ch := make(chan bool)

	go func() {
		for {
			evt := <-w.Event
			switch {
			case evt.IsDelete():
				fallthrough

			case evt.IsRename():
				close(ch)
				w.RemoveWatch(fw.Filename)
				w.Close()
				return

			case evt.IsModify():
				// send only if channel is empty.
				select {
				case ch <- true:
				default:
				}
			}
		}
	}()

	return ch
}

// FileWatcher monitors file-level events
type PollingFileWatcher struct {
	Filename string
}

func NewPollingFileWatcher(filename string) *PollingFileWatcher {
	fw := &PollingFileWatcher{filename}
	return fw
}

// BlockUntilExists blocks until the file comes into existence. If the
// file already exists, then block until it is created again.
func (fw *PollingFileWatcher) BlockUntilExists() error {
	panic("not implemented")
	return nil
}

// ChangeEvents returns a channel that gets updated when the file is ready to be read.
func (fw *PollingFileWatcher) ChangeEvents() chan bool {
	ch := make(chan bool)
	stop := make(chan bool)
	every2Seconds := time.Tick(2 * time.Second)

	go func() {
		for {
			time.Sleep(250 * time.Millisecond)
			select {
			case ch <- true:
			case <-stop:
				return
			default:
			}
		}
	}()

	go func() {
		for {
			select {
			case <-every2Seconds:
				// NOTE: not using file descriptor.
				if _, err := os.Stat(fw.Filename); os.IsNotExist(err) {
					stop <- true
					close(ch)
					return
				}
			}
		}
	}()

	return ch
}
<commit_msg>fix: BlockUntilExists should only return if filename matches<commit_after>// TODO: avoid creating two instances of the fsnotify.Watcher struct
package tail

import (
	"github.com/howeyc/fsnotify"
	"os"
	"path/filepath"
	"time"
)

type FileWatcher interface {
	BlockUntilExists() error
	ChangeEvents() chan bool
}

// FileWatcher monitors file-level events
type InotifyFileWatcher struct {
	Filename string
}

func NewInotifyFileWatcher(filename string) *InotifyFileWatcher {
	fw := &InotifyFileWatcher{filename}
	return fw
}

// BlockUntilExists blocks until the file comes into existence. If the
// file already exists, then block until it is created again.
func (fw *InotifyFileWatcher) BlockUntilExists() error {
	w, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	defer w.Close()
	err = w.WatchFlags(filepath.Dir(fw.Filename), fsnotify.FSN_CREATE)
	if err != nil {
		return err
	}
	defer w.RemoveWatch(filepath.Dir(fw.Filename))
	for {
		evt := <-w.Event
		if evt.Name == fw.Filename {
			break
		}
	}
	return nil
}

// ChangeEvents returns a channel that gets updated when the file is ready to be read.
func (fw *InotifyFileWatcher) ChangeEvents() chan bool {
	w, err := fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
	err = w.Watch(fw.Filename)
	if err != nil {
		panic(err)
	}

	ch := make(chan bool)

	go func() {
		for {
			evt := <-w.Event
			switch {
			case evt.IsDelete():
				fallthrough

			case evt.IsRename():
				close(ch)
				w.RemoveWatch(fw.Filename)
				w.Close()
				return

			case evt.IsModify():
				// send only if channel is empty.
				select {
				case ch <- true:
				default:
				}
			}
		}
	}()

	return ch
}

// FileWatcher monitors file-level events
type PollingFileWatcher struct {
	Filename string
}

func NewPollingFileWatcher(filename string) *PollingFileWatcher {
	fw := &PollingFileWatcher{filename}
	return fw
}

// BlockUntilExists blocks until the file comes into existence. If the
// file already exists, then block until it is created again.
func (fw *PollingFileWatcher) BlockUntilExists() error {
	panic("not implemented")
	return nil
}

// ChangeEvents returns a channel that gets updated when the file is ready to be read.
func (fw *PollingFileWatcher) ChangeEvents() chan bool {
	ch := make(chan bool)
	stop := make(chan bool)
	every2Seconds := time.Tick(2 * time.Second)

	go func() {
		for {
			time.Sleep(250 * time.Millisecond)
			select {
			case ch <- true:
			case <-stop:
				return
			default:
			}
		}
	}()

	go func() {
		for {
			select {
			case <-every2Seconds:
				// NOTE: not using file descriptor.
				if _, err := os.Stat(fw.Filename); os.IsNotExist(err) {
					stop <- true
					close(ch)
					return
				}
			}
		}
	}()

	return ch
}
<|endoftext|>


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

package gcslock

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
)

type mutex struct {
	project string
	bucket  string
	object  string
	client  *http.Client
}

const (
	defaultStorageLockURL   = "https://www.googleapis.com/upload/storage/v1"
	defaultStorageUnlockURL = "https://www.googleapis.com/storage/v1"
)

var (
	// These vars are used in the requests below. Having separate default
	// values makes it easy to reset the standard config during testing.
	storageLockURL   = defaultStorageLockURL
	storageUnlockURL = defaultStorageUnlockURL
)

// Lock waits up to duration d for l.Lock() to succeed.
func Lock(l sync.Locker, d time.Duration) error {
	done := make(chan struct{}, 1)
	go func() {
		l.Lock()
		done <- struct{}{}
	}()
	select {
	case <-done:
		return nil
	case <-time.After(d):
		return errors.New("lock request timed out")
	}
}

// Unlock waits up to duration d for l.Unlock() to succeed.
func Unlock(l sync.Locker, d time.Duration) error {
	done := make(chan struct{}, 1)
	go func() {
		l.Unlock()
		done <- struct{}{}
	}()
	select {
	case <-done:
		return nil
	case <-time.After(d):
		return errors.New("unlock request timed out")
	}
}

// Lock waits indefinitely to acquire a global mutex lock.
func (m mutex) Lock() {
	q := url.Values{
		"name":              {m.object},
		"uploadType":        {"media"},
		"ifGenerationMatch": {"0"},
	}
	url := fmt.Sprintf("%s/b/%s/o?%s", storageLockURL, m.bucket, q.Encode())
	for i := 1; ; i *= 2 {
		res, err := m.client.Post(url, "plain/text", bytes.NewReader([]byte("1")))
		if err == nil {
			res.Body.Close()
			if res.StatusCode == 200 {
				return
			}
		}
		time.Sleep(time.Duration(i) * time.Millisecond)
	}
}

// Unlock waits indefinitely to relinquish a global mutex lock.
func (m mutex) Unlock() {
	url := fmt.Sprintf("%s/b/%s/o/%s?", storageUnlockURL, m.bucket, m.object)
	for i := 1; ; i *= 2 {
		req, err := http.NewRequest("DELETE", url, nil)
		if err == nil {
			res, err := m.client.Do(req)
			if err == nil {
				res.Body.Close()
				if res.StatusCode == 204 {
					return
				}
			}
		}
		time.Sleep(time.Duration(i) * time.Millisecond)
	}
}

// New creates a GCS-based sync.Locker.
// It uses Application Default Credentials to make authenticated requests
// to Google Cloud Storage. See the DefaultClient function of the
// golang.org/x/oauth2/google package for App Default Credentials details.
//
// If ctx argument is nil, context.Background is used.
//
func New(ctx context.Context, project, bucket, object string) (sync.Locker, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	scope := "https://www.googleapis.com/auth/devstorage.full_control"
	client, err := google.DefaultClient(ctx, scope)
	if err != nil {
		return nil, err
	}
	m := &mutex{
		project: project,
		bucket:  bucket,
		object:  object,
		client:  client,
	}
	return m, nil
}
<commit_msg>add missing pkgdoc<commit_after>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// gcslock is a scalable, distributed mutex that can be used to
// serialize computations anywhere on the global internet.
package gcslock

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
)

type mutex struct {
	project string
	bucket  string
	object  string
	client  *http.Client
}

const (
	defaultStorageLockURL   = "https://www.googleapis.com/upload/storage/v1"
	defaultStorageUnlockURL = "https://www.googleapis.com/storage/v1"
)

var (
	// These vars are used in the requests below. Having separate default
	// values makes it easy to reset the standard config during testing.
	storageLockURL   = defaultStorageLockURL
	storageUnlockURL = defaultStorageUnlockURL
)

// Lock waits up to duration d for l.Lock() to succeed.
func Lock(l sync.Locker, d time.Duration) error {
	done := make(chan struct{}, 1)
	go func() {
		l.Lock()
		done <- struct{}{}
	}()
	select {
	case <-done:
		return nil
	case <-time.After(d):
		return errors.New("lock request timed out")
	}
}

// Unlock waits up to duration d for l.Unlock() to succeed.
func Unlock(l sync.Locker, d time.Duration) error {
	done := make(chan struct{}, 1)
	go func() {
		l.Unlock()
		done <- struct{}{}
	}()
	select {
	case <-done:
		return nil
	case <-time.After(d):
		return errors.New("unlock request timed out")
	}
}

// Lock waits indefinitely to acquire a global mutex lock.
func (m mutex) Lock() {
	q := url.Values{
		"name":              {m.object},
		"uploadType":        {"media"},
		"ifGenerationMatch": {"0"},
	}
	url := fmt.Sprintf("%s/b/%s/o?%s", storageLockURL, m.bucket, q.Encode())
	for i := 1; ; i *= 2 {
		res, err := m.client.Post(url, "plain/text", bytes.NewReader([]byte("1")))
		if err == nil {
			res.Body.Close()
			if res.StatusCode == 200 {
				return
			}
		}
		time.Sleep(time.Duration(i) * time.Millisecond)
	}
}

// Unlock waits indefinitely to relinquish a global mutex lock.
func (m mutex) Unlock() {
	url := fmt.Sprintf("%s/b/%s/o/%s?", storageUnlockURL, m.bucket, m.object)
	for i := 1; ; i *= 2 {
		req, err := http.NewRequest("DELETE", url, nil)
		if err == nil {
			res, err := m.client.Do(req)
			if err == nil {
				res.Body.Close()
				if res.StatusCode == 204 {
					return
				}
			}
		}
		time.Sleep(time.Duration(i) * time.Millisecond)
	}
}

// New creates a GCS-based sync.Locker.
// It uses Application Default Credentials to make authenticated requests
// to Google Cloud Storage. See the DefaultClient function of the
// golang.org/x/oauth2/google package for App Default Credentials details.
//
// If ctx argument is nil, context.Background is used.
//
func New(ctx context.Context, project, bucket, object string) (sync.Locker, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	scope := "https://www.googleapis.com/auth/devstorage.full_control"
	client, err := google.DefaultClient(ctx, scope)
	if err != nil {
		return nil, err
	}
	m := &mutex{
		project: project,
		bucket:  bucket,
		object:  object,
		client:  client,
	}
	return m, nil
}
<|endoftext|>


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

package main

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
)

type LinuxStrategy struct {
}

// chown changes ownership of a path.
func (LinuxStrategy) chown(ctx context.Context, username, path string) error {
	user, err := user.Lookup(username)
	if err != nil {
		return err
	}
	uid, err := strconv.Atoi(user.Uid)
	if err != nil {
		return err
	}
	gid, err := strconv.Atoi(user.Gid)
	if err != nil {
		return err
	}
	return os.Chown(path, uid, gid)
}

// configureAutoMount mounts the specified disk and configures mount on startup.
//
// Assumes the disk is already formatted as ext4.
func (LinuxStrategy) configureAutoMount(ctx context.Context, disk string) error {
	// Configure auto-mount using fstab.
	f, err := os.OpenFile("/etc/fstab", os.O_RDWR, 0)
	if err != nil {
		return err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), disk) {
			logging.Infof(ctx, "Already mounted: %s.", disk)
			return nil
		}
	}
	logging.Infof(ctx, "Mounting: %s.", disk)
	// Ensure the cursor is at the end so the new line is appended.
	f.Seek(0, io.SeekEnd)
	line := []byte(fmt.Sprintf("%s /b ext4 defaults,nobootwait,nofail 0 2\n", disk))
	if _, err := f.Write(line); err != nil {
		return err
	}
	// Ensure the file is closed before mount reads from it.
	f.Close()
	return exec.Command("/bin/mount", "--all").Run()
}

// reboot reboots the machine.
func (LinuxStrategy) reboot(ctx context.Context) error {
	return exec.Command("/sbin/shutdown", "-r", "now").Run()
}

type SystemdStrategy struct {
	LinuxStrategy
}

// enableSwarming enables installed service.
func (SystemdStrategy) enableSwarming(ctx context.Context) error {
	if err := exec.Command("systemctl", "daemon-reload").Run(); err != nil {
		return err
	}
	return exec.Command("systemctl", "enable", "swarming-start-bot").Run()
}

// start starts the agent.
func (SystemdStrategy) start(ctx context.Context, _ string) error {
	if err := exec.Command("systemctl", "daemon-reload").Run(); err != nil {
		return err
	}
	if err := exec.Command("systemctl", "enable", "machine-provider-agent").Run(); err != nil {
		return err
	}
	return exec.Command("systemctl", "start", "machine-provider-agent").Run()
}

// stop stops all instances of the agent.
func (SystemdStrategy) stop(ctx context.Context) error {
	return exec.Command("systemctl", "stop", "machine-provider-agent").Run()
}

type UpstartStrategy struct {
	LinuxStrategy
}

// enableSwarming enables installed service.
func (UpstartStrategy) enableSwarming(ctx context.Context) error {
	return nil
}

// start starts the agent.
func (UpstartStrategy) start(ctx context.Context, _ string) error {
	if err := exec.Command("initctl", "reload-configuration").Run(); err != nil {
		return err
	}
	return exec.Command("start", "machine-provider-agent").Run()
}

// stop stops all instances of the agent.
func (UpstartStrategy) stop(ctx context.Context) error {
	return exec.Command("stop", "machine-provider-agent").Run()
}

// systemdFound returns true iff systemd is supported on this machine.
func systemdFound() bool {
	return exec.Command("which", "systemctl").Run() == nil
}

// upstartFound returns true iff upstart is supported on this machine.
func upstartFound() bool {
	return exec.Command("init", "--version").Run() == nil
}

// getAgent returns an agent which runs on Linux, depending on supported init systems.
func getAgent(ctx context.Context) (*Agent, error) {
	if systemdFound() {
		agent := Agent{
			agentAutoStartPath:        "/etc/systemd/system/machine-provider-agent.service",
			agentAutoStartTemplate:    "machine-provider-agent.service.tmpl",
			logsDir:                   "/var/log/machine-provider-agent",
			swarmingAutoStartPath:     "/etc/systemd/system/swarming-start-bot.service",
			swarmingAutoStartTemplate: "swarming-start-bot.service.tmpl",
			swarmingBotDir:            "/b/s",
			strategy:                  SystemdStrategy{},
		}
		logging.Infof(ctx, "Using systemd Linux agent.")
		return &agent, nil
	}

	if upstartFound() {
		agent := Agent{
			agentAutoStartPath:        "/etc/init/machine-provider-agent.conf",
			agentAutoStartTemplate:    "machine-provider-agent.conf.tmpl",
			logsDir:                   "/var/log/messages/machine-provider-agent",
			swarmingAutoStartPath:     "/etc/init/swarming-start-bot.conf",
			swarmingAutoStartTemplate: "swarming-start-bot.conf.tmpl",
			swarmingBotDir:            "/b/s",
			strategy:                  UpstartStrategy{},
		}
		logging.Infof(ctx, "Using Upstart Linux agent.")
		return &agent, nil
	}

	return nil, errors.New("unsupported init system, expected systemd or Upstart")
}
<commit_msg>[mpagent] Wait for disk to be attached before mounting<commit_after>// Copyright 2017 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"
	"syscall"
	"time"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
)

type LinuxStrategy struct {
}

// chown changes ownership of a path.
func (LinuxStrategy) chown(ctx context.Context, username, path string) error {
	user, err := user.Lookup(username)
	if err != nil {
		return err
	}
	uid, err := strconv.Atoi(user.Uid)
	if err != nil {
		return err
	}
	gid, err := strconv.Atoi(user.Gid)
	if err != nil {
		return err
	}
	return os.Chown(path, uid, gid)
}

// configureAutoMount mounts the specified disk and configures mount on startup.
//
// Assumes the disk is already formatted as ext4.
func (LinuxStrategy) configureAutoMount(ctx context.Context, disk string) error {
	// Wait for the disk to be attached.
	for {
		err := exec.Command("blkid", disk).Run()
		if err == nil {
			logging.Infof(ctx, "Attached: %s.", disk)
			break
		}
		s := err.(*exec.ExitError).Sys().(syscall.WaitStatus)
		if s.ExitStatus() != 2 {
			// 2 means the specified device wasn't found.
			// Keep waiting if 2, otherwise return an error.
			return err
		}
		time.Sleep(60 * time.Second)
		logging.Infof(ctx, "Waiting for disk to be attached...")
	}
	// Configure auto-mount using fstab.
	f, err := os.OpenFile("/etc/fstab", os.O_RDWR, 0)
	if err != nil {
		return err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), disk) {
			logging.Infof(ctx, "Already mounted: %s.", disk)
			return nil
		}
	}
	logging.Infof(ctx, "Mounting: %s.", disk)
	// Ensure the cursor is at the end so the new line is appended.
	f.Seek(0, io.SeekEnd)
	line := []byte(fmt.Sprintf("%s /b ext4 defaults,nobootwait,nofail 0 2\n", disk))
	if _, err := f.Write(line); err != nil {
		return err
	}
	// Ensure the file is closed before mount reads from it.
	f.Close()
	return exec.Command("/bin/mount", "--all").Run()
}

// reboot reboots the machine.
func (LinuxStrategy) reboot(ctx context.Context) error {
	return exec.Command("/sbin/shutdown", "-r", "now").Run()
}

type SystemdStrategy struct {
	LinuxStrategy
}

// enableSwarming enables installed service.
func (SystemdStrategy) enableSwarming(ctx context.Context) error {
	if err := exec.Command("systemctl", "daemon-reload").Run(); err != nil {
		return err
	}
	return exec.Command("systemctl", "enable", "swarming-start-bot").Run()
}

// start starts the agent.
func (SystemdStrategy) start(ctx context.Context, _ string) error {
	if err := exec.Command("systemctl", "daemon-reload").Run(); err != nil {
		return err
	}
	if err := exec.Command("systemctl", "enable", "machine-provider-agent").Run(); err != nil {
		return err
	}
	return exec.Command("systemctl", "start", "machine-provider-agent").Run()
}

// stop stops all instances of the agent.
func (SystemdStrategy) stop(ctx context.Context) error {
	return exec.Command("systemctl", "stop", "machine-provider-agent").Run()
}

type UpstartStrategy struct {
	LinuxStrategy
}

// enableSwarming enables installed service.
func (UpstartStrategy) enableSwarming(ctx context.Context) error {
	return nil
}

// start starts the agent.
func (UpstartStrategy) start(ctx context.Context, _ string) error {
	if err := exec.Command("initctl", "reload-configuration").Run(); err != nil {
		return err
	}
	return exec.Command("start", "machine-provider-agent").Run()
}

// stop stops all instances of the agent.
func (UpstartStrategy) stop(ctx context.Context) error {
	return exec.Command("stop", "machine-provider-agent").Run()
}

// systemdFound returns true iff systemd is supported on this machine.
func systemdFound() bool {
	return exec.Command("which", "systemctl").Run() == nil
}

// upstartFound returns true iff upstart is supported on this machine.
func upstartFound() bool {
	return exec.Command("init", "--version").Run() == nil
}

// getAgent returns an agent which runs on Linux, depending on supported init systems.
func getAgent(ctx context.Context) (*Agent, error) {
	if systemdFound() {
		agent := Agent{
			agentAutoStartPath:        "/etc/systemd/system/machine-provider-agent.service",
			agentAutoStartTemplate:    "machine-provider-agent.service.tmpl",
			logsDir:                   "/var/log/machine-provider-agent",
			swarmingAutoStartPath:     "/etc/systemd/system/swarming-start-bot.service",
			swarmingAutoStartTemplate: "swarming-start-bot.service.tmpl",
			swarmingBotDir:            "/b/s",
			strategy:                  SystemdStrategy{},
		}
		logging.Infof(ctx, "Using systemd Linux agent.")
		return &agent, nil
	}

	if upstartFound() {
		agent := Agent{
			agentAutoStartPath:        "/etc/init/machine-provider-agent.conf",
			agentAutoStartTemplate:    "machine-provider-agent.conf.tmpl",
			logsDir:                   "/var/log/messages/machine-provider-agent",
			swarmingAutoStartPath:     "/etc/init/swarming-start-bot.conf",
			swarmingAutoStartTemplate: "swarming-start-bot.conf.tmpl",
			swarmingBotDir:            "/b/s",
			strategy:                  UpstartStrategy{},
		}
		logging.Infof(ctx, "Using Upstart Linux agent.")
		return &agent, nil
	}

	return nil, errors.New("unsupported init system, expected systemd or Upstart")
}
<|endoftext|>


<commit_before>package gocbcore

import (
	"encoding/binary"
)

// Represents the state of a particular cluster snapshot.
type SnapshotState uint32

// Returns whether this snapshot is available in memory.
func (s SnapshotState) HasInMemory() bool {
	return uint32(s)&1 != 0
}

// Returns whether this snapshot is available on disk.
func (s SnapshotState) HasOnDisk() bool {
	return uint32(s)&2 != 0
}

// Represents a single entry in the server failover log.
type FailoverEntry struct {
	VbUuid VbUuid
	SeqNo  SeqNo
}

type StreamObserver interface {
	SnapshotMarker(startSeqNo, endSeqNo uint64, vbId uint16, snapshotType SnapshotState)
	Mutation(seqNo, revNo uint64, flags, expiry, lockTime uint32, cas uint64, datatype uint8, vbId uint16, key, value []byte)
	Deletion(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	Expiration(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	End(vbId uint16, err error)
}

type OpenStreamCallback func([]FailoverEntry, error)
type CloseStreamCallback func(error)
type GetFailoverLogCallback func([]FailoverEntry, error)
type GetVBucketSeqnosCallback func(uint16, SeqNo, error)

// **INTERNAL**
// Opens a DCP stream for a particular VBucket.
func (agent *Agent) OpenStream(vbId uint16, vbUuid VbUuid, startSeqNo, endSeqNo, snapStartSeqNo, snapEndSeqNo SeqNo, evtHandler StreamObserver, cb OpenStreamCallback) (PendingOp, error) {
	var req *memdQRequest
	handler := func(resp *memdQResponse, _ *memdQRequest, err error) {
		if resp.Magic == ResMagic {
			// This is the response to the open stream request.
			if err != nil {
				req.Cancel()

				// All client errors are handled by the StreamObserver
				cb(nil, err)
				return
			}

			numEntries := len(resp.Value) / 16
			entries := make([]FailoverEntry, numEntries)
			for i := 0; i < numEntries; i++ {
				entries[i] = FailoverEntry{
					VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
					SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
				}
			}

			cb(entries, nil)
			return
		}

		if err != nil {
			req.Cancel()
			evtHandler.End(vbId, err)
			return
		}

		// This is one of the stream events
		switch resp.Opcode {
		case CmdDcpSnapshotMarker:
			vbId := uint16(resp.Status)
			newStartSeqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			newEndSeqNo := binary.BigEndian.Uint64(resp.Extras[8:])
			snapshotType := binary.BigEndian.Uint32(resp.Extras[16:])
			evtHandler.SnapshotMarker(newStartSeqNo, newEndSeqNo, vbId, SnapshotState(snapshotType))
		case CmdDcpMutation:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			flags := binary.BigEndian.Uint32(resp.Extras[16:])
			expiry := binary.BigEndian.Uint32(resp.Extras[20:])
			lockTime := binary.BigEndian.Uint32(resp.Extras[24:])
			evtHandler.Mutation(seqNo, revNo, flags, expiry, lockTime, resp.Cas, resp.Datatype, vbId, resp.Key, resp.Value)
		case CmdDcpDeletion:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Deletion(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpExpiration:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Expiration(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpStreamEnd:
			vbId := uint16(resp.Status)
			code := StreamEndStatus(binary.BigEndian.Uint32(resp.Extras[0:]))
			evtHandler.End(vbId, getStreamEndError(code))
			req.Cancel()
		}
	}

	extraBuf := make([]byte, 48)
	binary.BigEndian.PutUint32(extraBuf[0:], 0)
	binary.BigEndian.PutUint32(extraBuf[4:], 0)
	binary.BigEndian.PutUint64(extraBuf[8:], uint64(startSeqNo))
	binary.BigEndian.PutUint64(extraBuf[16:], uint64(endSeqNo))
	binary.BigEndian.PutUint64(extraBuf[24:], uint64(vbUuid))
	binary.BigEndian.PutUint64(extraBuf[32:], uint64(snapStartSeqNo))
	binary.BigEndian.PutUint64(extraBuf[40:], uint64(snapEndSeqNo))

	req = &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdDcpStreamReq,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return agent.dispatchOp(req)
}

// **INTERNAL**
// Shuts down an open stream for the specified VBucket.
func (agent *Agent) CloseStream(vbId uint16, cb CloseStreamCallback) (PendingOp, error) {
	handler := func(_ *memdQResponse, _ *memdQRequest, err error) {
		cb(err)
	}

	req := &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdDcpCloseStream,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return agent.dispatchOp(req)
}

// **INTERNAL**
// Retrieves the failover log for a particular VBucket.  This is used
// to resume an interrupted stream after a node failover has occured.
func (agent *Agent) GetFailoverLog(vbId uint16, cb GetFailoverLogCallback) (PendingOp, error) {
	handler := func(resp *memdQResponse, _ *memdQRequest, err error) {
		if err != nil {
			cb(nil, err)
			return
		}

		numEntries := len(resp.Value) / 16
		entries := make([]FailoverEntry, numEntries)
		for i := 0; i < numEntries; i++ {
			entries[i] = FailoverEntry{
				VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
				SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
			}
		}
		cb(entries, nil)
	}

	req := &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdDcpGetFailoverLog,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return agent.dispatchOp(req)
}

// **INTERNAL**
// Returns the last checkpoint for a particular VBucket.  This is useful
// for starting a DCP stream from wherever the server currently is.
func (agent *Agent) GetVbucketSeqnos(serverIdx int, cb GetVBucketSeqnosCallback) (PendingOp, error) {
	handler := func(resp *memdQResponse, _ *memdQRequest, err error) {
		if err != nil {
			cb(0, 0, err)
			return
		}

		vbs := len(resp.Value) / 10
		for i := 0; i < vbs; i++ {
			vbid := binary.BigEndian.Uint16(resp.Value[i*10:])
			seqNo := SeqNo(binary.BigEndian.Uint64(resp.Value[i*10+2:]))
			cb(vbid, seqNo, nil)
		}
	}

	extraBuf := make([]byte, 4)
	binary.BigEndian.PutUint32(extraBuf[0:], uint32(VBucketStateActive))

	req := &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdGetAllVBSeqnos,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  0,
		},
		Callback:   handler,
		ReplicaIdx: -serverIdx,
		Persistent: false,
	}

	return agent.dispatchOp(req)
}
<commit_msg>GOCBC-137: Ensure we have response data before accessing its Magic.<commit_after>package gocbcore

import (
	"encoding/binary"
)

// Represents the state of a particular cluster snapshot.
type SnapshotState uint32

// Returns whether this snapshot is available in memory.
func (s SnapshotState) HasInMemory() bool {
	return uint32(s)&1 != 0
}

// Returns whether this snapshot is available on disk.
func (s SnapshotState) HasOnDisk() bool {
	return uint32(s)&2 != 0
}

// Represents a single entry in the server failover log.
type FailoverEntry struct {
	VbUuid VbUuid
	SeqNo  SeqNo
}

type StreamObserver interface {
	SnapshotMarker(startSeqNo, endSeqNo uint64, vbId uint16, snapshotType SnapshotState)
	Mutation(seqNo, revNo uint64, flags, expiry, lockTime uint32, cas uint64, datatype uint8, vbId uint16, key, value []byte)
	Deletion(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	Expiration(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	End(vbId uint16, err error)
}

type OpenStreamCallback func([]FailoverEntry, error)
type CloseStreamCallback func(error)
type GetFailoverLogCallback func([]FailoverEntry, error)
type GetVBucketSeqnosCallback func(uint16, SeqNo, error)

// **INTERNAL**
// Opens a DCP stream for a particular VBucket.
func (agent *Agent) OpenStream(vbId uint16, vbUuid VbUuid, startSeqNo, endSeqNo, snapStartSeqNo, snapEndSeqNo SeqNo, evtHandler StreamObserver, cb OpenStreamCallback) (PendingOp, error) {
	var req *memdQRequest
	handler := func(resp *memdQResponse, _ *memdQRequest, err error) {
		if resp != nil && resp.Magic == ResMagic {
			// This is the response to the open stream request.
			if err != nil {
				req.Cancel()

				// All client errors are handled by the StreamObserver
				cb(nil, err)
				return
			}

			numEntries := len(resp.Value) / 16
			entries := make([]FailoverEntry, numEntries)
			for i := 0; i < numEntries; i++ {
				entries[i] = FailoverEntry{
					VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
					SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
				}
			}

			cb(entries, nil)
			return
		}

		if err != nil {
			req.Cancel()
			evtHandler.End(vbId, err)
			return
		}

		// This is one of the stream events
		switch resp.Opcode {
		case CmdDcpSnapshotMarker:
			vbId := uint16(resp.Status)
			newStartSeqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			newEndSeqNo := binary.BigEndian.Uint64(resp.Extras[8:])
			snapshotType := binary.BigEndian.Uint32(resp.Extras[16:])
			evtHandler.SnapshotMarker(newStartSeqNo, newEndSeqNo, vbId, SnapshotState(snapshotType))
		case CmdDcpMutation:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			flags := binary.BigEndian.Uint32(resp.Extras[16:])
			expiry := binary.BigEndian.Uint32(resp.Extras[20:])
			lockTime := binary.BigEndian.Uint32(resp.Extras[24:])
			evtHandler.Mutation(seqNo, revNo, flags, expiry, lockTime, resp.Cas, resp.Datatype, vbId, resp.Key, resp.Value)
		case CmdDcpDeletion:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Deletion(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpExpiration:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Expiration(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpStreamEnd:
			vbId := uint16(resp.Status)
			code := StreamEndStatus(binary.BigEndian.Uint32(resp.Extras[0:]))
			evtHandler.End(vbId, getStreamEndError(code))
			req.Cancel()
		}
	}

	extraBuf := make([]byte, 48)
	binary.BigEndian.PutUint32(extraBuf[0:], 0)
	binary.BigEndian.PutUint32(extraBuf[4:], 0)
	binary.BigEndian.PutUint64(extraBuf[8:], uint64(startSeqNo))
	binary.BigEndian.PutUint64(extraBuf[16:], uint64(endSeqNo))
	binary.BigEndian.PutUint64(extraBuf[24:], uint64(vbUuid))
	binary.BigEndian.PutUint64(extraBuf[32:], uint64(snapStartSeqNo))
	binary.BigEndian.PutUint64(extraBuf[40:], uint64(snapEndSeqNo))

	req = &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdDcpStreamReq,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return agent.dispatchOp(req)
}

// **INTERNAL**
// Shuts down an open stream for the specified VBucket.
func (agent *Agent) CloseStream(vbId uint16, cb CloseStreamCallback) (PendingOp, error) {
	handler := func(_ *memdQResponse, _ *memdQRequest, err error) {
		cb(err)
	}

	req := &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdDcpCloseStream,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return agent.dispatchOp(req)
}

// **INTERNAL**
// Retrieves the failover log for a particular VBucket.  This is used
// to resume an interrupted stream after a node failover has occured.
func (agent *Agent) GetFailoverLog(vbId uint16, cb GetFailoverLogCallback) (PendingOp, error) {
	handler := func(resp *memdQResponse, _ *memdQRequest, err error) {
		if err != nil {
			cb(nil, err)
			return
		}

		numEntries := len(resp.Value) / 16
		entries := make([]FailoverEntry, numEntries)
		for i := 0; i < numEntries; i++ {
			entries[i] = FailoverEntry{
				VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
				SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
			}
		}
		cb(entries, nil)
	}

	req := &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdDcpGetFailoverLog,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return agent.dispatchOp(req)
}

// **INTERNAL**
// Returns the last checkpoint for a particular VBucket.  This is useful
// for starting a DCP stream from wherever the server currently is.
func (agent *Agent) GetVbucketSeqnos(serverIdx int, cb GetVBucketSeqnosCallback) (PendingOp, error) {
	handler := func(resp *memdQResponse, _ *memdQRequest, err error) {
		if err != nil {
			cb(0, 0, err)
			return
		}

		vbs := len(resp.Value) / 10
		for i := 0; i < vbs; i++ {
			vbid := binary.BigEndian.Uint16(resp.Value[i*10:])
			seqNo := SeqNo(binary.BigEndian.Uint64(resp.Value[i*10+2:]))
			cb(vbid, seqNo, nil)
		}
	}

	extraBuf := make([]byte, 4)
	binary.BigEndian.PutUint32(extraBuf[0:], uint32(VBucketStateActive))

	req := &memdQRequest{
		memdPacket: memdPacket{
			Magic:    ReqMagic,
			Opcode:   CmdGetAllVBSeqnos,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  0,
		},
		Callback:   handler,
		ReplicaIdx: -serverIdx,
		Persistent: false,
	}

	return agent.dispatchOp(req)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, The dcrdata developers
// See LICENSE for details.

package explorer

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"time"

	"golang.org/x/net/websocket"
)

// RootWebsocket is the websocket handler for all pages
func (exp *explorerUI) RootWebsocket(w http.ResponseWriter, r *http.Request) {
	wsHandler := websocket.Handler(func(ws *websocket.Conn) {
		// Create channel to signal updated data availability
		updateSig := make(hubSpoke, 3)
		// register websocket client with our signal channel
		exp.wsHub.RegisterClient(&updateSig)
		// unregister (and close signal channel) before return
		defer exp.wsHub.UnregisterClient(&updateSig)

		requestLimit := 1 << 20
		// set the max payload size to 1 MB
		ws.MaxPayloadBytes = requestLimit

		// Ticker for a regular ping
		ticker := time.NewTicker(pingInterval)
		defer ticker.Stop()

		// Periodically ping clients over websocket connection
		go func() {
			for range ticker.C {
				exp.wsHub.HubRelay <- sigPingAndUserCount
			}
		}()

		// Start listening for websocket messages from client with raw
		// transaction bytes (hex encoded) to decode or broadcast.
		go func() {
			for {
				// Wait to receive a message on the websocket
				msg := &WebSocketMessage{}
				ws.SetReadDeadline(time.Now().Add(wsReadTimeout))
				if err := websocket.JSON.Receive(ws, &msg); err != nil {
					log.Warnf("websocket client receive error: %v", err)
					exp.wsHub.UnregisterClient(&updateSig)
					return
				}

				// handle received message according to event ID
				var webData WebSocketMessage
				switch msg.EventId {
				case "decodetx":
					webData.EventId = msg.EventId + "Resp"
					if len(msg.Message) > requestLimit {
						log.Debug("Request size over limit")
						webData.Message = "Request too large"
						break
					}
					log.Debugf("Received decodetx signal for hex: %.40s...", msg.Message)
					tx, err := exp.blockData.DecodeRawTransaction(msg.Message)
					if err == nil {
						message, err := json.MarshalIndent(tx, "", "    ")
						if err != nil {
							log.Warn("Invalid JSON message: ", err)
							webData.Message = fmt.Sprintf("Error: Could not encode JSON message")
							break
						}
						webData.Message = string(message)
					} else {
						log.Debugf("Could not decode raw tx")
						webData.Message = fmt.Sprintf("Error: %v", err)
					}
				case "sendtx":
					webData.EventId = msg.EventId + "Resp"
					if len(msg.Message) > requestLimit {
						log.Debugf("Request size over limit")
						webData.Message = "Request too large"
						break
					}
					log.Debugf("Received sendtx signal for hex: %.40s...", msg.Message)
					txid, err := exp.blockData.SendRawTransaction(msg.Message)
					if err != nil {
						webData.Message = fmt.Sprintf("Error: %v", err)
					} else {
						webData.Message = fmt.Sprintf("Transaction sent: %s", txid)
					}
				case "getmempooltxs":
					webData.EventId = msg.EventId + "Resp"

					exp.MempoolData.Lock()
					msg, err := json.Marshal(exp.MempoolData)
					exp.MempoolData.Unlock()

					if err != nil {
						log.Warn("Invalid JSON message: ", err)
						webData.Message = fmt.Sprintf("Error: Could not encode JSON message")
						break
					}
					webData.Message = string(msg)
				case "ping":
					log.Tracef("We've been pinged: %.40s...", msg.Message)
					continue
				default:
					log.Warnf("Unrecognized event ID: %v", msg.EventId)
					continue
				}

				// send the response back on the websocket
				ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
				if err := websocket.JSON.Send(ws, webData); err != nil {
					log.Debugf("Failed to encode WebSocketMessage %s: %v",
						webData.EventId, err)
					// If the send failed, the client is probably gone, so close
					// the connection and quit.
					return
				}
			}
		}()

		// Ping and block update loop (send only)
	loop:
		for {
			// Wait for signal from the hub to update
			select {
			case sig, ok := <-updateSig:
				// Check if the update channel was closed. Either the websocket
				// hub will do it after unregistering the client, or forcibly in
				// response to (http.CloseNotifier).CloseNotify() and only then if
				// the hub has somehow lost track of the client.
				if !ok {
					//ws.WriteClose(1)
					exp.wsHub.UnregisterClient(&updateSig)
					break loop
				}

				if _, ok = eventIDs[sig]; !ok {
					break loop
				}

				log.Tracef("signaling client: %p", &updateSig)

				// Write block data to websocket client

				webData := WebSocketMessage{
					EventId: eventIDs[sig],
				}
				buff := new(bytes.Buffer)
				enc := json.NewEncoder(buff)
				switch sig {
				case sigNewBlock:
					exp.NewBlockDataMtx.Lock()
					enc.Encode(WebsocketBlock{
						Block: exp.NewBlockData,
						Extra: exp.ExtraInfo,
					})
					exp.NewBlockDataMtx.Unlock()
					webData.Message = buff.String()
				case sigMempoolUpdate:
					exp.MempoolData.Lock()
					enc.Encode(exp.MempoolData.MempoolShort)
					exp.MempoolData.Unlock()
					webData.Message = buff.String()
				case sigPingAndUserCount:
					// ping and send user count
					webData.Message = strconv.Itoa(exp.wsHub.NumClients())
				case sigNewTx:
					exp.wsHub.clients[&updateSig].Lock()
					enc.Encode(exp.wsHub.clients[&updateSig].newTxs)
					exp.wsHub.clients[&updateSig].Unlock()
					webData.Message = buff.String()
				}

				ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
				err := websocket.JSON.Send(ws, webData)

				if err != nil {
					log.Debugf("Failed to encode WebSocketMessage %v: %v", sig, err)
					// If the send failed, the client is probably gone, so close
					// the connection and quit.
					return
				}
			case <-exp.wsHub.quitWSHandler:
				break loop
			}
		}
	})

	wsHandler.ServeHTTP(w, r)
}
<commit_msg>explorer/websocket: avoid repeatedly unregistering a websocket client<commit_after>// Copyright (c) 2017, The dcrdata developers
// See LICENSE for details.

package explorer

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"time"

	"golang.org/x/net/websocket"
)

// RootWebsocket is the websocket handler for all pages
func (exp *explorerUI) RootWebsocket(w http.ResponseWriter, r *http.Request) {
	wsHandler := websocket.Handler(func(ws *websocket.Conn) {
		// Create channel to signal updated data availability
		updateSig := make(hubSpoke, 3)
		// register websocket client with our signal channel
		exp.wsHub.RegisterClient(&updateSig)
		// unregister (and close signal channel) before return
		defer exp.wsHub.UnregisterClient(&updateSig)

		requestLimit := 1 << 20
		// set the max payload size to 1 MB
		ws.MaxPayloadBytes = requestLimit

		// Ticker for a regular ping
		ticker := time.NewTicker(pingInterval)
		defer ticker.Stop()

		// Periodically ping clients over websocket connection
		go func() {
			for range ticker.C {
				exp.wsHub.HubRelay <- sigPingAndUserCount
			}
		}()

		// Start listening for websocket messages from client with raw
		// transaction bytes (hex encoded) to decode or broadcast.
		go func() {
			for {
				// Wait to receive a message on the websocket
				msg := &WebSocketMessage{}
				ws.SetReadDeadline(time.Now().Add(wsReadTimeout))
				if err := websocket.JSON.Receive(ws, &msg); err != nil {
					if err.Error() != "EOF" {
						log.Warnf("websocket client receive error: %v", err)
					}
					return
				}

				// handle received message according to event ID
				var webData WebSocketMessage
				switch msg.EventId {
				case "decodetx":
					webData.EventId = msg.EventId + "Resp"
					if len(msg.Message) > requestLimit {
						log.Debug("Request size over limit")
						webData.Message = "Request too large"
						break
					}
					log.Debugf("Received decodetx signal for hex: %.40s...", msg.Message)
					tx, err := exp.blockData.DecodeRawTransaction(msg.Message)
					if err == nil {
						message, err := json.MarshalIndent(tx, "", "    ")
						if err != nil {
							log.Warn("Invalid JSON message: ", err)
							webData.Message = fmt.Sprintf("Error: Could not encode JSON message")
							break
						}
						webData.Message = string(message)
					} else {
						log.Debugf("Could not decode raw tx")
						webData.Message = fmt.Sprintf("Error: %v", err)
					}
				case "sendtx":
					webData.EventId = msg.EventId + "Resp"
					if len(msg.Message) > requestLimit {
						log.Debugf("Request size over limit")
						webData.Message = "Request too large"
						break
					}
					log.Debugf("Received sendtx signal for hex: %.40s...", msg.Message)
					txid, err := exp.blockData.SendRawTransaction(msg.Message)
					if err != nil {
						webData.Message = fmt.Sprintf("Error: %v", err)
					} else {
						webData.Message = fmt.Sprintf("Transaction sent: %s", txid)
					}
				case "getmempooltxs":
					webData.EventId = msg.EventId + "Resp"

					exp.MempoolData.Lock()
					msg, err := json.Marshal(exp.MempoolData)
					exp.MempoolData.Unlock()

					if err != nil {
						log.Warn("Invalid JSON message: ", err)
						webData.Message = fmt.Sprintf("Error: Could not encode JSON message")
						break
					}
					webData.Message = string(msg)
				case "ping":
					log.Tracef("We've been pinged: %.40s...", msg.Message)
					continue
				default:
					log.Warnf("Unrecognized event ID: %v", msg.EventId)
					continue
				}

				// send the response back on the websocket
				ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
				if err := websocket.JSON.Send(ws, webData); err != nil {
					log.Debugf("Failed to encode WebSocketMessage %s: %v",
						webData.EventId, err)
					// If the send failed, the client is probably gone, so close
					// the connection and quit.
					return
				}
			}
		}()

		// Ping and block update loop (send only)
	loop:
		for {
			// Wait for signal from the hub to update
			select {
			case sig, ok := <-updateSig:
				// Check if the update channel was closed. Either the websocket
				// hub will do it after unregistering the client, or forcibly in
				// response to (http.CloseNotifier).CloseNotify() and only then if
				// the hub has somehow lost track of the client.
				if !ok {
					break loop
				}

				if _, ok = eventIDs[sig]; !ok {
					break loop
				}

				log.Tracef("signaling client: %p", &updateSig)

				// Write block data to websocket client

				webData := WebSocketMessage{
					EventId: eventIDs[sig],
				}
				buff := new(bytes.Buffer)
				enc := json.NewEncoder(buff)
				switch sig {
				case sigNewBlock:
					exp.NewBlockDataMtx.Lock()
					enc.Encode(WebsocketBlock{
						Block: exp.NewBlockData,
						Extra: exp.ExtraInfo,
					})
					exp.NewBlockDataMtx.Unlock()
					webData.Message = buff.String()
				case sigMempoolUpdate:
					exp.MempoolData.Lock()
					enc.Encode(exp.MempoolData.MempoolShort)
					exp.MempoolData.Unlock()
					webData.Message = buff.String()
				case sigPingAndUserCount:
					// ping and send user count
					webData.Message = strconv.Itoa(exp.wsHub.NumClients())
				case sigNewTx:
					exp.wsHub.clients[&updateSig].Lock()
					enc.Encode(exp.wsHub.clients[&updateSig].newTxs)
					exp.wsHub.clients[&updateSig].Unlock()
					webData.Message = buff.String()
				}

				ws.SetWriteDeadline(time.Now().Add(wsWriteTimeout))
				err := websocket.JSON.Send(ws, webData)

				if err != nil {
					log.Debugf("Failed to encode WebSocketMessage %v: %v", sig, err)
					// If the send failed, the client is probably gone, so close
					// the connection and quit.
					return
				}
			case <-exp.wsHub.quitWSHandler:
				break loop
			}
		}
	})

	wsHandler.ServeHTTP(w, r)
}
<|endoftext|>


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

package otlpexporter

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

	"google.golang.org/genproto/googleapis/rpc/errdetails"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"

	"go.opentelemetry.io/collector/config"
	"go.opentelemetry.io/collector/consumer/consumererror"
	"go.opentelemetry.io/collector/consumer/pdata"
	"go.opentelemetry.io/collector/exporter/exporterhelper"
	"go.opentelemetry.io/collector/internal"
	otlplogs "go.opentelemetry.io/collector/internal/data/protogen/collector/logs/v1"
	otlpmetrics "go.opentelemetry.io/collector/internal/data/protogen/collector/metrics/v1"
	otlptrace "go.opentelemetry.io/collector/internal/data/protogen/collector/trace/v1"
)

type exporterImp struct {
	// Input configuration.
	config *Config
	w      *grpcSender
}

// Crete new exporter and start it. The exporter will begin connecting but
// this function may return before the connection is established.
func newExporter(cfg config.Exporter) (*exporterImp, error) {
	oCfg := cfg.(*Config)

	if oCfg.Endpoint == "" {
		return nil, errors.New("OTLP exporter config requires an Endpoint")
	}

	e := &exporterImp{}
	e.config = oCfg
	w, err := newGrpcSender(oCfg)
	if err != nil {
		return nil, err
	}
	e.w = w
	return e, nil
}

func (e *exporterImp) shutdown(context.Context) error {
	return e.w.stop()
}

func (e *exporterImp) pushTraceData(ctx context.Context, td pdata.Traces) error {
	if err := e.w.exportTrace(ctx, internal.TracesToOtlp(td.InternalRep())); err != nil {
		return fmt.Errorf("failed to push trace data via OTLP exporter: %w", err)
	}
	return nil
}

func (e *exporterImp) pushMetricsData(ctx context.Context, md pdata.Metrics) error {
	req := internal.MetricsToOtlp(md.InternalRep())
	if err := e.w.exportMetrics(ctx, req); err != nil {
		return fmt.Errorf("failed to push metrics data via OTLP exporter: %w", err)
	}
	return nil
}

func (e *exporterImp) pushLogData(ctx context.Context, ld pdata.Logs) error {
	request := internal.LogsToOtlp(ld.InternalRep())
	if err := e.w.exportLogs(ctx, request); err != nil {
		return fmt.Errorf("failed to push log data via OTLP exporter: %w", err)
	}
	return nil
}

type grpcSender struct {
	// gRPC clients and connection.
	traceExporter  otlptrace.TraceServiceClient
	metricExporter otlpmetrics.MetricsServiceClient
	logExporter    otlplogs.LogsServiceClient
	grpcClientConn *grpc.ClientConn
	metadata       metadata.MD
	waitForReady   bool
}

func newGrpcSender(config *Config) (*grpcSender, error) {
	dialOpts, err := config.GRPCClientSettings.ToDialOptions()
	if err != nil {
		return nil, err
	}

	var clientConn *grpc.ClientConn
	if clientConn, err = grpc.Dial(config.GRPCClientSettings.Endpoint, dialOpts...); err != nil {
		return nil, err
	}

	gs := &grpcSender{
		traceExporter:  otlptrace.NewTraceServiceClient(clientConn),
		metricExporter: otlpmetrics.NewMetricsServiceClient(clientConn),
		logExporter:    otlplogs.NewLogsServiceClient(clientConn),
		grpcClientConn: clientConn,
		metadata:       metadata.New(config.GRPCClientSettings.Headers),
		waitForReady:   config.GRPCClientSettings.WaitForReady,
	}
	return gs, nil
}

func (gs *grpcSender) stop() error {
	return gs.grpcClientConn.Close()
}

func (gs *grpcSender) exportTrace(ctx context.Context, request *otlptrace.ExportTraceServiceRequest) error {
	_, err := gs.traceExporter.Export(gs.enhanceContext(ctx), request, grpc.WaitForReady(gs.waitForReady))
	return processError(err)
}

func (gs *grpcSender) exportMetrics(ctx context.Context, request *otlpmetrics.ExportMetricsServiceRequest) error {
	_, err := gs.metricExporter.Export(gs.enhanceContext(ctx), request, grpc.WaitForReady(gs.waitForReady))
	return processError(err)
}

func (gs *grpcSender) exportLogs(ctx context.Context, request *otlplogs.ExportLogsServiceRequest) error {
	_, err := gs.logExporter.Export(gs.enhanceContext(ctx), request, grpc.WaitForReady(gs.waitForReady))
	return processError(err)
}

func (gs *grpcSender) enhanceContext(ctx context.Context) context.Context {
	if gs.metadata.Len() > 0 {
		return metadata.NewOutgoingContext(ctx, gs.metadata)
	}
	return ctx
}

// Send a trace or metrics request to the server. "perform" function is expected to make
// the actual gRPC unary call that sends the request. This function implements the
// common OTLP logic around request handling such as retries and throttling.
func processError(err error) error {
	if err == nil {
		// Request is successful, we are done.
		return nil
	}

	// We have an error, check gRPC status code.

	st := status.Convert(err)
	if st.Code() == codes.OK {
		// Not really an error, still success.
		return nil
	}

	// Now, this is this a real error.

	if !shouldRetry(st.Code()) {
		// It is not a retryable error, we should not retry.
		return consumererror.Permanent(err)
	}

	// Need to retry.

	// Check if server returned throttling information.
	throttleDuration := getThrottleDuration(st)
	if throttleDuration != 0 {
		return exporterhelper.NewThrottleRetry(err, throttleDuration)
	}

	return err
}

func shouldRetry(code codes.Code) bool {
	switch code {
	case codes.OK:
		// Success. This function should not be called for this code, the best we
		// can do is tell the caller not to retry.
		return false

	case codes.Canceled,
		codes.DeadlineExceeded,
		codes.PermissionDenied,
		codes.Unauthenticated,
		codes.ResourceExhausted,
		codes.Aborted,
		codes.OutOfRange,
		codes.Unavailable,
		codes.DataLoss:
		// These are retryable errors.
		return true

	case codes.Unknown,
		codes.InvalidArgument,
		codes.NotFound,
		codes.AlreadyExists,
		codes.FailedPrecondition,
		codes.Unimplemented,
		codes.Internal:
		// These are fatal errors, don't retry.
		return false

	default:
		// Don't retry on unknown codes.
		return false
	}
}

func getThrottleDuration(status *status.Status) time.Duration {
	// See if throttling information is available.
	for _, detail := range status.Details() {
		if t, ok := detail.(*errdetails.RetryInfo); ok {
			if t.RetryDelay.Seconds > 0 || t.RetryDelay.Nanos > 0 {
				// We are throttled. Wait before retrying as requested by the server.
				return time.Duration(t.RetryDelay.Seconds)*time.Second + time.Duration(t.RetryDelay.Nanos)*time.Nanosecond
			}
			return 0
		}
	}
	return 0
}
<commit_msg>Avoid one unnecessary allocation in grpc OTLP exporter (#3122)<commit_after>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package otlpexporter

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

	"google.golang.org/genproto/googleapis/rpc/errdetails"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"

	"go.opentelemetry.io/collector/config"
	"go.opentelemetry.io/collector/consumer/consumererror"
	"go.opentelemetry.io/collector/consumer/pdata"
	"go.opentelemetry.io/collector/exporter/exporterhelper"
	"go.opentelemetry.io/collector/internal"
	otlplogs "go.opentelemetry.io/collector/internal/data/protogen/collector/logs/v1"
	otlpmetrics "go.opentelemetry.io/collector/internal/data/protogen/collector/metrics/v1"
	otlptrace "go.opentelemetry.io/collector/internal/data/protogen/collector/trace/v1"
)

type exporterImp struct {
	// Input configuration.
	config *Config
	w      *grpcSender
}

// Crete new exporter and start it. The exporter will begin connecting but
// this function may return before the connection is established.
func newExporter(cfg config.Exporter) (*exporterImp, error) {
	oCfg := cfg.(*Config)

	if oCfg.Endpoint == "" {
		return nil, errors.New("OTLP exporter config requires an Endpoint")
	}

	e := &exporterImp{}
	e.config = oCfg
	w, err := newGrpcSender(oCfg)
	if err != nil {
		return nil, err
	}
	e.w = w
	return e, nil
}

func (e *exporterImp) shutdown(context.Context) error {
	return e.w.stop()
}

func (e *exporterImp) pushTraceData(ctx context.Context, td pdata.Traces) error {
	if err := e.w.exportTrace(ctx, internal.TracesToOtlp(td.InternalRep())); err != nil {
		return fmt.Errorf("failed to push trace data via OTLP exporter: %w", err)
	}
	return nil
}

func (e *exporterImp) pushMetricsData(ctx context.Context, md pdata.Metrics) error {
	req := internal.MetricsToOtlp(md.InternalRep())
	if err := e.w.exportMetrics(ctx, req); err != nil {
		return fmt.Errorf("failed to push metrics data via OTLP exporter: %w", err)
	}
	return nil
}

func (e *exporterImp) pushLogData(ctx context.Context, ld pdata.Logs) error {
	request := internal.LogsToOtlp(ld.InternalRep())
	if err := e.w.exportLogs(ctx, request); err != nil {
		return fmt.Errorf("failed to push log data via OTLP exporter: %w", err)
	}
	return nil
}

type grpcSender struct {
	// gRPC clients and connection.
	traceExporter  otlptrace.TraceServiceClient
	metricExporter otlpmetrics.MetricsServiceClient
	logExporter    otlplogs.LogsServiceClient
	clientConn     *grpc.ClientConn
	metadata       metadata.MD
	callOptions    []grpc.CallOption
}

func newGrpcSender(config *Config) (*grpcSender, error) {
	dialOpts, err := config.GRPCClientSettings.ToDialOptions()
	if err != nil {
		return nil, err
	}

	var clientConn *grpc.ClientConn
	if clientConn, err = grpc.Dial(config.GRPCClientSettings.Endpoint, dialOpts...); err != nil {
		return nil, err
	}

	gs := &grpcSender{
		traceExporter:  otlptrace.NewTraceServiceClient(clientConn),
		metricExporter: otlpmetrics.NewMetricsServiceClient(clientConn),
		logExporter:    otlplogs.NewLogsServiceClient(clientConn),
		clientConn:     clientConn,
		metadata:       metadata.New(config.GRPCClientSettings.Headers),
		callOptions: []grpc.CallOption{
			grpc.WaitForReady(config.GRPCClientSettings.WaitForReady),
		},
	}
	return gs, nil
}

func (gs *grpcSender) stop() error {
	return gs.clientConn.Close()
}

func (gs *grpcSender) exportTrace(ctx context.Context, request *otlptrace.ExportTraceServiceRequest) error {
	_, err := gs.traceExporter.Export(gs.enhanceContext(ctx), request, gs.callOptions...)
	return processError(err)
}

func (gs *grpcSender) exportMetrics(ctx context.Context, request *otlpmetrics.ExportMetricsServiceRequest) error {
	_, err := gs.metricExporter.Export(gs.enhanceContext(ctx), request, gs.callOptions...)
	return processError(err)
}

func (gs *grpcSender) exportLogs(ctx context.Context, request *otlplogs.ExportLogsServiceRequest) error {
	_, err := gs.logExporter.Export(gs.enhanceContext(ctx), request, gs.callOptions...)
	return processError(err)
}

func (gs *grpcSender) enhanceContext(ctx context.Context) context.Context {
	if gs.metadata.Len() > 0 {
		return metadata.NewOutgoingContext(ctx, gs.metadata)
	}
	return ctx
}

// Send a trace or metrics request to the server. "perform" function is expected to make
// the actual gRPC unary call that sends the request. This function implements the
// common OTLP logic around request handling such as retries and throttling.
func processError(err error) error {
	if err == nil {
		// Request is successful, we are done.
		return nil
	}

	// We have an error, check gRPC status code.

	st := status.Convert(err)
	if st.Code() == codes.OK {
		// Not really an error, still success.
		return nil
	}

	// Now, this is this a real error.

	if !shouldRetry(st.Code()) {
		// It is not a retryable error, we should not retry.
		return consumererror.Permanent(err)
	}

	// Need to retry.

	// Check if server returned throttling information.
	throttleDuration := getThrottleDuration(st)
	if throttleDuration != 0 {
		return exporterhelper.NewThrottleRetry(err, throttleDuration)
	}

	return err
}

func shouldRetry(code codes.Code) bool {
	switch code {
	case codes.OK:
		// Success. This function should not be called for this code, the best we
		// can do is tell the caller not to retry.
		return false

	case codes.Canceled,
		codes.DeadlineExceeded,
		codes.PermissionDenied,
		codes.Unauthenticated,
		codes.ResourceExhausted,
		codes.Aborted,
		codes.OutOfRange,
		codes.Unavailable,
		codes.DataLoss:
		// These are retryable errors.
		return true

	case codes.Unknown,
		codes.InvalidArgument,
		codes.NotFound,
		codes.AlreadyExists,
		codes.FailedPrecondition,
		codes.Unimplemented,
		codes.Internal:
		// These are fatal errors, don't retry.
		return false

	default:
		// Don't retry on unknown codes.
		return false
	}
}

func getThrottleDuration(status *status.Status) time.Duration {
	// See if throttling information is available.
	for _, detail := range status.Details() {
		if t, ok := detail.(*errdetails.RetryInfo); ok {
			if t.RetryDelay.Seconds > 0 || t.RetryDelay.Nanos > 0 {
				// We are throttled. Wait before retrying as requested by the server.
				return time.Duration(t.RetryDelay.Seconds)*time.Second + time.Duration(t.RetryDelay.Nanos)*time.Nanosecond
			}
			return 0
		}
	}
	return 0
}
<|endoftext|>