<commit_before>package kcp

import (
	"encoding/binary"
	"sync/atomic"

	"github.com/klauspost/reedsolomon"
)

const (
	fecHeaderSize      = 6
	fecHeaderSizePlus2 = fecHeaderSize + 2 // plus 2B data size
	typeData           = 0xf1
	typeFEC            = 0xf2
)

type (
	// fecPacket is a decoded FEC packet
	fecPacket struct {
		seqid uint32
		flag  uint16
		data  []byte
	}

	// FECDecoder for decoding incoming packets
	FECDecoder struct {
		rxlimit      int // queue size limit
		dataShards   int
		parityShards int
		shardSize    int
		rx           []fecPacket // ordered receive queue

		// caches
		decodeCache [][]byte
		shardsflag  []bool

		// RS decoder
		enc reedsolomon.Encoder
	}
)

func newFECDecoder(rxlimit, dataShards, parityShards int) *FECDecoder {
	if dataShards <= 0 || parityShards <= 0 {
		return nil
	}
	if rxlimit < dataShards+parityShards {
		return nil
	}

	fec := new(FECDecoder)
	fec.rxlimit = rxlimit
	fec.dataShards = dataShards
	fec.parityShards = parityShards
	fec.shardSize = dataShards + parityShards
	enc, err := reedsolomon.New(dataShards, parityShards, reedsolomon.WithMaxGoroutines(1))
	if err != nil {
		return nil
	}
	fec.enc = enc
	fec.decodeCache = make([][]byte, fec.shardSize)
	fec.shardsflag = make([]bool, fec.shardSize)
	return fec
}

// decodeBytes a fec packet
func (dec *FECDecoder) decodeBytes(data []byte) fecPacket {
	var pkt fecPacket
	pkt.seqid = binary.LittleEndian.Uint32(data)
	pkt.flag = binary.LittleEndian.Uint16(data[4:])
	// allocate memory & copy
	buf := xmitBuf.Get().([]byte)[:len(data)-6]
	copy(buf, data[6:])
	pkt.data = buf
	return pkt
}

// Decode a fec packet
func (dec *FECDecoder) Decode(pkt fecPacket) (recovered [][]byte) {
	// insertion
	n := len(dec.rx) - 1
	insertIdx := 0
	for i := n; i >= 0; i-- {
		if pkt.seqid == dec.rx[i].seqid { // de-duplicate
			xmitBuf.Put(pkt.data)
			return nil
		} else if _itimediff(pkt.seqid, dec.rx[i].seqid) > 0 { // insertion
			insertIdx = i + 1
			break
		}
	}

	// insert into ordered rx queue
	if insertIdx == n+1 {
		dec.rx = append(dec.rx, pkt)
	} else {
		dec.rx = append(dec.rx, fecPacket{})
		copy(dec.rx[insertIdx+1:], dec.rx[insertIdx:]) // shift right
		dec.rx[insertIdx] = pkt
	}

	// shard range for current packet
	shardBegin := pkt.seqid - pkt.seqid%uint32(dec.shardSize)
	shardEnd := shardBegin + uint32(dec.shardSize) - 1

	// max search range in ordered queue for current shard
	searchBegin := insertIdx - int(pkt.seqid%uint32(dec.shardSize))
	if searchBegin < 0 {
		searchBegin = 0
	}
	searchEnd := searchBegin + dec.shardSize - 1
	if searchEnd >= len(dec.rx) {
		searchEnd = len(dec.rx) - 1
	}

	// re-construct datashards
	if searchEnd > searchBegin && searchEnd-searchBegin+1 >= dec.dataShards {
		numshard := 0
		numDataShard := 0
		first := -1
		maxlen := 0
		shards := dec.decodeCache
		shardsflag := dec.shardsflag
		for k := range dec.decodeCache {
			shards[k] = nil
			shardsflag[k] = false
		}

		for i := searchBegin; i <= searchEnd; i++ {
			seqid := dec.rx[i].seqid
			if _itimediff(seqid, shardEnd) > 0 {
				break
			} else if _itimediff(seqid, shardBegin) >= 0 {
				shards[seqid%uint32(dec.shardSize)] = dec.rx[i].data
				shardsflag[seqid%uint32(dec.shardSize)] = true
				numshard++
				if dec.rx[i].flag == typeData {
					numDataShard++
				}
				if numshard == 1 {
					first = i
				}
				if len(dec.rx[i].data) > maxlen {
					maxlen = len(dec.rx[i].data)
				}
			}
		}

		if numDataShard == dec.dataShards { // no lost
			for i := first; i < first+numshard; i++ { // free
				xmitBuf.Put(dec.rx[i].data)
			}
			copy(dec.rx[first:], dec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				dec.rx[len(dec.rx)-1-i] = fecPacket{}
			}
			dec.rx = dec.rx[:len(dec.rx)-numshard]
		} else if numshard >= dec.dataShards { // recoverable
			for k := range shards {
				if shards[k] != nil {
					dlen := len(shards[k])
					shards[k] = shards[k][:maxlen]
					xorBytes(shards[k][dlen:], shards[k][dlen:], shards[k][dlen:])
				}
			}
			if err := dec.enc.Reconstruct(shards); err == nil {
				for k := range shards[:dec.dataShards] {
					if !shardsflag[k] {
						recovered = append(recovered, shards[k])
					}
				}
			}

			for i := first; i < first+numshard; i++ { // free
				xmitBuf.Put(dec.rx[i].data)
			}
			copy(dec.rx[first:], dec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				dec.rx[len(dec.rx)-1-i] = fecPacket{}
			}
			dec.rx = dec.rx[:len(dec.rx)-numshard]
		}
	}

	// keep rxlimit
	if len(dec.rx) > dec.rxlimit {
		if dec.rx[0].flag == typeData { // record unrecoverable data
			atomic.AddUint64(&DefaultSnmp.FECShortShards, 1)
		}
		xmitBuf.Put(dec.rx[0].data) // free
		copy(dec.rx, dec.rx[1:])    // shift left
		dec.rx[len(dec.rx)-1] = fecPacket{}
		dec.rx = dec.rx[:len(dec.rx)-1]
	}
	return
}

type (
	// FECEncoder for encoding outgoing packets
	FECEncoder struct {
		dataShards   int
		parityShards int
		shardSize    int
		paws         uint32 // Protect Against Wrapped Sequence numbers
		next         uint32 // next seqid

		shardCount int // count the number of datashards collected
		maxSize    int // record maximum data length in datashard

		headerOffset  int // FEC header offset
		payloadOffset int // FEC payload offset

		// caches
		shardCache  [][]byte
		encodeCache [][]byte

		// RS encoder
		enc reedsolomon.Encoder
	}
)

func newFECEncoder(dataShards, parityShards, offset int) *FECEncoder {
	if dataShards <= 0 || parityShards <= 0 {
		return nil
	}
	fec := new(FECEncoder)
	fec.dataShards = dataShards
	fec.parityShards = parityShards
	fec.shardSize = dataShards + parityShards
	fec.paws = (0xffffffff/uint32(fec.shardSize) - 1) * uint32(fec.shardSize)
	fec.headerOffset = offset
	fec.payloadOffset = fec.headerOffset + fecHeaderSize

	enc, err := reedsolomon.New(dataShards, parityShards, reedsolomon.WithMaxGoroutines(1))
	if err != nil {
		return nil
	}
	fec.enc = enc

	// caches
	fec.encodeCache = make([][]byte, fec.shardSize)
	fec.shardCache = make([][]byte, fec.shardSize)
	for k := range fec.shardCache {
		fec.shardCache[k] = make([]byte, mtuLimit)
	}
	return fec
}

// Encode the packet, output parity shards if we have enough datashards
// the content of returned parityshards will change in next Encode
func (enc *FECEncoder) Encode(b []byte) (ps [][]byte) {
	enc.markData(b[enc.headerOffset:])
	binary.LittleEndian.PutUint16(b[enc.payloadOffset:], uint16(len(b[enc.payloadOffset:])))

	// copy data to fec datashards
	sz := len(b)
	enc.shardCache[enc.shardCount] = enc.shardCache[enc.shardCount][:sz]
	copy(enc.shardCache[enc.shardCount], b)
	enc.shardCount++

	// record max datashard length
	if sz > enc.maxSize {
		enc.maxSize = sz
	}

	//  calculate Reed-Solomon Erasure Code
	if enc.shardCount == enc.dataShards {
		// bzero each datashard's tail
		for i := 0; i < enc.dataShards; i++ {
			shard := enc.shardCache[i]
			slen := len(shard)
			xorBytes(shard[slen:enc.maxSize], shard[slen:enc.maxSize], shard[slen:enc.maxSize])
		}

		// construct equal-sized slice with stripped header
		cache := enc.encodeCache
		for k := range cache {
			cache[k] = enc.shardCache[k][enc.payloadOffset:enc.maxSize]
		}

		// rs encode
		if err := enc.enc.Encode(cache); err == nil {
			ps = enc.shardCache[enc.dataShards:]
			for k := range ps {
				enc.markFEC(ps[k][enc.headerOffset:])
				ps[k] = ps[k][:enc.maxSize]
			}
		}

		// reset counters to zero
		enc.shardCount = 0
		enc.maxSize = 0
	}

	return
}

func (enc *FECEncoder) markData(data []byte) {
	binary.LittleEndian.PutUint32(data, enc.next)
	binary.LittleEndian.PutUint16(data[4:], typeData)
	enc.next++
}

func (enc *FECEncoder) markFEC(data []byte) {
	binary.LittleEndian.PutUint32(data, enc.next)
	binary.LittleEndian.PutUint16(data[4:], typeFEC)
	enc.next = (enc.next + 1) % enc.paws
}
<commit_msg>GoRename<commit_after>package kcp

import (
	"encoding/binary"
	"sync/atomic"

	"github.com/klauspost/reedsolomon"
)

const (
	fecHeaderSize      = 6
	fecHeaderSizePlus2 = fecHeaderSize + 2 // plus 2B data size
	typeData           = 0xf1
	typeFEC            = 0xf2
)

type (
	// fecPacket is a decoded FEC packet
	fecPacket struct {
		seqid uint32
		flag  uint16
		data  []byte
	}

	// FECDecoder for decoding incoming packets
	FECDecoder struct {
		rxlimit      int // queue size limit
		dataShards   int
		parityShards int
		shardSize    int
		rx           []fecPacket // ordered receive queue

		// caches
		decodeCache [][]byte
		flagCache   []bool

		// RS decoder
		enc reedsolomon.Encoder
	}
)

func newFECDecoder(rxlimit, dataShards, parityShards int) *FECDecoder {
	if dataShards <= 0 || parityShards <= 0 {
		return nil
	}
	if rxlimit < dataShards+parityShards {
		return nil
	}

	fec := new(FECDecoder)
	fec.rxlimit = rxlimit
	fec.dataShards = dataShards
	fec.parityShards = parityShards
	fec.shardSize = dataShards + parityShards
	enc, err := reedsolomon.New(dataShards, parityShards, reedsolomon.WithMaxGoroutines(1))
	if err != nil {
		return nil
	}
	fec.enc = enc
	fec.decodeCache = make([][]byte, fec.shardSize)
	fec.flagCache = make([]bool, fec.shardSize)
	return fec
}

// decodeBytes a fec packet
func (dec *FECDecoder) decodeBytes(data []byte) fecPacket {
	var pkt fecPacket
	pkt.seqid = binary.LittleEndian.Uint32(data)
	pkt.flag = binary.LittleEndian.Uint16(data[4:])
	// allocate memory & copy
	buf := xmitBuf.Get().([]byte)[:len(data)-6]
	copy(buf, data[6:])
	pkt.data = buf
	return pkt
}

// Decode a fec packet
func (dec *FECDecoder) Decode(pkt fecPacket) (recovered [][]byte) {
	// insertion
	n := len(dec.rx) - 1
	insertIdx := 0
	for i := n; i >= 0; i-- {
		if pkt.seqid == dec.rx[i].seqid { // de-duplicate
			xmitBuf.Put(pkt.data)
			return nil
		} else if _itimediff(pkt.seqid, dec.rx[i].seqid) > 0 { // insertion
			insertIdx = i + 1
			break
		}
	}

	// insert into ordered rx queue
	if insertIdx == n+1 {
		dec.rx = append(dec.rx, pkt)
	} else {
		dec.rx = append(dec.rx, fecPacket{})
		copy(dec.rx[insertIdx+1:], dec.rx[insertIdx:]) // shift right
		dec.rx[insertIdx] = pkt
	}

	// shard range for current packet
	shardBegin := pkt.seqid - pkt.seqid%uint32(dec.shardSize)
	shardEnd := shardBegin + uint32(dec.shardSize) - 1

	// max search range in ordered queue for current shard
	searchBegin := insertIdx - int(pkt.seqid%uint32(dec.shardSize))
	if searchBegin < 0 {
		searchBegin = 0
	}
	searchEnd := searchBegin + dec.shardSize - 1
	if searchEnd >= len(dec.rx) {
		searchEnd = len(dec.rx) - 1
	}

	// re-construct datashards
	if searchEnd > searchBegin && searchEnd-searchBegin+1 >= dec.dataShards {
		numshard := 0
		numDataShard := 0
		first := -1
		maxlen := 0
		shards := dec.decodeCache
		shardsflag := dec.flagCache
		for k := range dec.decodeCache {
			shards[k] = nil
			shardsflag[k] = false
		}

		for i := searchBegin; i <= searchEnd; i++ {
			seqid := dec.rx[i].seqid
			if _itimediff(seqid, shardEnd) > 0 {
				break
			} else if _itimediff(seqid, shardBegin) >= 0 {
				shards[seqid%uint32(dec.shardSize)] = dec.rx[i].data
				shardsflag[seqid%uint32(dec.shardSize)] = true
				numshard++
				if dec.rx[i].flag == typeData {
					numDataShard++
				}
				if numshard == 1 {
					first = i
				}
				if len(dec.rx[i].data) > maxlen {
					maxlen = len(dec.rx[i].data)
				}
			}
		}

		if numDataShard == dec.dataShards { // no lost
			for i := first; i < first+numshard; i++ { // free
				xmitBuf.Put(dec.rx[i].data)
			}
			copy(dec.rx[first:], dec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				dec.rx[len(dec.rx)-1-i] = fecPacket{}
			}
			dec.rx = dec.rx[:len(dec.rx)-numshard]
		} else if numshard >= dec.dataShards { // recoverable
			for k := range shards {
				if shards[k] != nil {
					dlen := len(shards[k])
					shards[k] = shards[k][:maxlen]
					xorBytes(shards[k][dlen:], shards[k][dlen:], shards[k][dlen:])
				}
			}
			if err := dec.enc.Reconstruct(shards); err == nil {
				for k := range shards[:dec.dataShards] {
					if !shardsflag[k] {
						recovered = append(recovered, shards[k])
					}
				}
			}

			for i := first; i < first+numshard; i++ { // free
				xmitBuf.Put(dec.rx[i].data)
			}
			copy(dec.rx[first:], dec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				dec.rx[len(dec.rx)-1-i] = fecPacket{}
			}
			dec.rx = dec.rx[:len(dec.rx)-numshard]
		}
	}

	// keep rxlimit
	if len(dec.rx) > dec.rxlimit {
		if dec.rx[0].flag == typeData { // record unrecoverable data
			atomic.AddUint64(&DefaultSnmp.FECShortShards, 1)
		}
		xmitBuf.Put(dec.rx[0].data) // free
		copy(dec.rx, dec.rx[1:])    // shift left
		dec.rx[len(dec.rx)-1] = fecPacket{}
		dec.rx = dec.rx[:len(dec.rx)-1]
	}
	return
}

type (
	// FECEncoder for encoding outgoing packets
	FECEncoder struct {
		dataShards   int
		parityShards int
		shardSize    int
		paws         uint32 // Protect Against Wrapped Sequence numbers
		next         uint32 // next seqid

		shardCount int // count the number of datashards collected
		maxSize    int // record maximum data length in datashard

		headerOffset  int // FEC header offset
		payloadOffset int // FEC payload offset

		// caches
		shardCache  [][]byte
		encodeCache [][]byte

		// RS encoder
		enc reedsolomon.Encoder
	}
)

func newFECEncoder(dataShards, parityShards, offset int) *FECEncoder {
	if dataShards <= 0 || parityShards <= 0 {
		return nil
	}
	fec := new(FECEncoder)
	fec.dataShards = dataShards
	fec.parityShards = parityShards
	fec.shardSize = dataShards + parityShards
	fec.paws = (0xffffffff/uint32(fec.shardSize) - 1) * uint32(fec.shardSize)
	fec.headerOffset = offset
	fec.payloadOffset = fec.headerOffset + fecHeaderSize

	enc, err := reedsolomon.New(dataShards, parityShards, reedsolomon.WithMaxGoroutines(1))
	if err != nil {
		return nil
	}
	fec.enc = enc

	// caches
	fec.encodeCache = make([][]byte, fec.shardSize)
	fec.shardCache = make([][]byte, fec.shardSize)
	for k := range fec.shardCache {
		fec.shardCache[k] = make([]byte, mtuLimit)
	}
	return fec
}

// Encode the packet, output parity shards if we have enough datashards
// the content of returned parityshards will change in next Encode
func (enc *FECEncoder) Encode(b []byte) (ps [][]byte) {
	enc.markData(b[enc.headerOffset:])
	binary.LittleEndian.PutUint16(b[enc.payloadOffset:], uint16(len(b[enc.payloadOffset:])))

	// copy data to fec datashards
	sz := len(b)
	enc.shardCache[enc.shardCount] = enc.shardCache[enc.shardCount][:sz]
	copy(enc.shardCache[enc.shardCount], b)
	enc.shardCount++

	// record max datashard length
	if sz > enc.maxSize {
		enc.maxSize = sz
	}

	//  calculate Reed-Solomon Erasure Code
	if enc.shardCount == enc.dataShards {
		// bzero each datashard's tail
		for i := 0; i < enc.dataShards; i++ {
			shard := enc.shardCache[i]
			slen := len(shard)
			xorBytes(shard[slen:enc.maxSize], shard[slen:enc.maxSize], shard[slen:enc.maxSize])
		}

		// construct equal-sized slice with stripped header
		cache := enc.encodeCache
		for k := range cache {
			cache[k] = enc.shardCache[k][enc.payloadOffset:enc.maxSize]
		}

		// rs encode
		if err := enc.enc.Encode(cache); err == nil {
			ps = enc.shardCache[enc.dataShards:]
			for k := range ps {
				enc.markFEC(ps[k][enc.headerOffset:])
				ps[k] = ps[k][:enc.maxSize]
			}
		}

		// reset counters to zero
		enc.shardCount = 0
		enc.maxSize = 0
	}

	return
}

func (enc *FECEncoder) markData(data []byte) {
	binary.LittleEndian.PutUint32(data, enc.next)
	binary.LittleEndian.PutUint16(data[4:], typeData)
	enc.next++
}

func (enc *FECEncoder) markFEC(data []byte) {
	binary.LittleEndian.PutUint32(data, enc.next)
	binary.LittleEndian.PutUint16(data[4:], typeFEC)
	enc.next = (enc.next + 1) % enc.paws
}
<|endoftext|>


<commit_before>package treewatcher

import (
	"github.com/howeyc/fsnotify"
	"io/ioutil"
	"os"
	"path"
	"testing"
	"time"
)

var (
	tempDir string
)

func expectNoEvent(t *testing.T, tw *TreeWatcher) {
	time.Sleep(50 * time.Millisecond)
	for {
		select {
		case event := <-tw.Event:
			// Only fail if the error is for a file that exists. In some odd cases involving
			// move directories, we can get double notifications under the old and new directory names.
			_, err := os.Stat(event.Name)
			if err == nil {
				t.Error("Got unexpected event", event)
			} else {
				t.Log("Ignoring", event)
			}

		case err := <-tw.Error:
			t.Error("fswatcher returned error", err)

		default:
			return
		}
	}
}

func expectEvent(t *testing.T, tw *TreeWatcher) *fsnotify.FileEvent {
	noError := false
	for !noError {
		select {
		case err := <-tw.Error:
			t.Error("fswatcher returned error", err)
		default:
			noError = true
		}
	}
	select {
	case event := <-tw.Event:
		t.Log(event)
		// This sleep takes care of race conditions in fsnotify so that the test can be more
		// strict about what it expects.
		time.Sleep(50 * time.Millisecond)
		return event
	// Wait at most 1 second for the event to show up.
	case <-time.After(1 * time.Second):
		t.Fatal("Expected an event but found none.")
	}

	return nil
}

func writeFile(t *testing.T, tw *TreeWatcher, filename string, data string) {
	fullname := path.Join(tempDir, filename)
	_, err := os.Stat(fullname)
	create := err != nil

	var file *os.File
	if create {
		t.Log("Create file", fullname)
		file, err = os.Create(fullname)
		event := expectEvent(t, tw)
		if event.Name != fullname || !event.IsCreate() {
			t.Errorf("Expected create event on \"%s\" but got %s", filename, event.String())
		}
	} else {
		t.Log("Modify file", filename)
		file, err = os.OpenFile(fullname, os.O_WRONLY, 0666)
	}

	if err != nil {
		t.Fatal("Failed to open file", filename)
	}
	_, err = file.WriteString(data)
	if err != nil {
		t.Fatalf("Write file %s failed with error %s", filename, err)
	}
	err = file.Sync()
	if err != nil {
		t.Fatalf("Sync file %s failed with error %s", filename, err)
	}
	file.Close()

	event := expectEvent(t, tw)
	if event.Name != fullname || !event.IsModify() {
		t.Errorf("Expected modify event on \"%s\" but got %s", filename, event.String())
	}

	expectNoEvent(t, tw)
}

func mkDir(t *testing.T, tw *TreeWatcher, name string) {
	t.Log("Create directory", name)
	name = path.Join(tempDir, name)
	err := os.Mkdir(name, 0755)
	if err != nil {
		t.Fatalf("Mkdir %s failed with error %s", name, err)
	}
	event := expectEvent(t, tw)
	if !event.IsCreate() || event.Name != name {
		t.Errorf("Expected create event on \"%s\" but got %s", name, event.String())
	}

	expectNoEvent(t, tw)
}

func deleteFile(t *testing.T, tw *TreeWatcher, filename string) {
	t.Log("Remove file", filename)
	filename = path.Join(tempDir, filename)
	err := os.Remove(filename)
	if err != nil {
		t.Fatal("Remove file %s failed with error %s", filename, err)
	}
	event := expectEvent(t, tw)
	if !event.IsDelete() || event.Name != filename {
		t.Errorf("Expected create event on \"%s\" but got %s", filename, event.String())
	}
	expectNoEvent(t, tw)
}

func renameFile(t *testing.T, tw *TreeWatcher, oldName, newName string) {
	t.Logf("Rename %s to %s", oldName, newName)
	newName = path.Join(tempDir, newName)
	oldName = path.Join(tempDir, oldName)
	err := os.Rename(oldName, newName)
	if err != nil {
		t.Fatalf("Rename %s to %s failed with %s", oldName, newName, err)
	}

	event := expectEvent(t, tw)
	if !event.IsCreate() || event.Name != newName {
		t.Errorf("Expected rename event on \"%s\" but got %s", newName, event.String())
	}

	event = expectEvent(t, tw)
	if !event.IsRename() || event.Name != oldName {
		t.Errorf("Expected rename event on \"%s\" but got %s", oldName, event.String())
	}

	expectNoEvent(t, tw)
}

func TestTreeWatcher(t *testing.T) {
	var err error
	tempDir, err = ioutil.TempDir("", "treewatcher_test")
	if err != nil {
		t.Fatal("Could not create temp directory")
	}
	t.Log("Using temporary directory", tempDir)
	// defer os.RemoveAll(tempDir)

	tw, err := New()
	if err != nil {
		t.Fatal("Could not create watcher:", err)
	}
	defer tw.Close()

	tw.WatchTree(tempDir)

	writeFile(t, tw, "abc.txt", "abc")
	writeFile(t, tw, "def.txt", "def")
	writeFile(t, tw, "abc.txt", "def")
	deleteFile(t, tw, "def.txt")
	renameFile(t, tw, "abc.txt", "def.txt")
	mkDir(t, tw, "dir")
	writeFile(t, tw, "dir/abc.txt", "abc")
	mkDir(t, tw, "dir2")
	renameFile(t, tw, "dir/abc.txt", "dir2/abc.txt")
	mkDir(t, tw, "dir/a")
	mkDir(t, tw, "dir/a/b")
	writeFile(t, tw, "dir/a/b/c.txt", "jklsdf")
	renameFile(t, tw, "dir/a", "dir2/a")
	writeFile(t, tw, "dir2/a/b/c.txt", "ggg")

	t.Log("Creating unwatched tree ppp/qqq/rrr/a.txt")
	unwatchedDir, _ := ioutil.TempDir("", "ppp")
	defer os.RemoveAll(unwatchedDir)

	nestedDir := path.Join(unwatchedDir, "ppp", "qqq", "rrr")
	os.MkdirAll(nestedDir, 0755)
	ioutil.WriteFile(path.Join(nestedDir, "a.txt"), []byte("data"), 0666)
	expectNoEvent(t, tw)

	t.Log("Moving unwatched tree under watched tree")
	movedDir := path.Join(tempDir, "ppp")
	os.Rename(path.Join(unwatchedDir, "ppp"), movedDir)
	event := expectEvent(t, tw)
	if event.Name != movedDir || !event.IsCreate() {
		t.Fatalf("Expected create of %s, got %s", movedDir, event)
	}
	writeFile(t, tw, "ppp/qqq/rrr/new.txt", "data")
	writeFile(t, tw, "ppp/qqq/rrr/a.txt", "data")
	deleteFile(t, tw, "ppp/qqq/rrr/new.txt")
	expectNoEvent(t, tw)
}
<commit_msg>Update test to be more tolerant of rename event ordering<commit_after>package treewatcher

import (
	"github.com/howeyc/fsnotify"
	"io/ioutil"
	"os"
	"path"
	"testing"
	"time"
)

var (
	tempDir string
)

func expectNoEvent(t *testing.T, tw *TreeWatcher) {
	time.Sleep(50 * time.Millisecond)
	for {
		select {
		case event := <-tw.Event:
			// Only fail if the error is for a file that exists. In some odd cases involving
			// move directories, we can get double notifications under the old and new directory names.
			_, err := os.Stat(event.Name)
			if err == nil {
				t.Error("Got unexpected event", event)
			} else {
				t.Log("Ignoring", event)
			}

		case err := <-tw.Error:
			t.Error("fswatcher returned error", err)

		default:
			return
		}
	}
}

func expectEvent(t *testing.T, tw *TreeWatcher) *fsnotify.FileEvent {
	noError := false
	for !noError {
		select {
		case err := <-tw.Error:
			t.Error("fswatcher returned error", err)
		default:
			noError = true
		}
	}
	select {
	case event := <-tw.Event:
		t.Log(event)
		// This sleep takes care of race conditions in fsnotify so that the test can be more
		// strict about what it expects.
		time.Sleep(50 * time.Millisecond)
		return event
	// Wait at most 1 second for the event to show up.
	case <-time.After(1 * time.Second):
		t.Fatal("Expected an event but found none.")
	}

	return nil
}

func writeFile(t *testing.T, tw *TreeWatcher, filename string, data string) {
	fullname := path.Join(tempDir, filename)
	_, err := os.Stat(fullname)
	create := err != nil

	var file *os.File
	if create {
		t.Log("Create file", fullname)
		file, err = os.Create(fullname)
		event := expectEvent(t, tw)
		if event.Name != fullname || !event.IsCreate() {
			t.Errorf("Expected create event on \"%s\" but got %s", filename, event.String())
		}
	} else {
		t.Log("Modify file", filename)
		file, err = os.OpenFile(fullname, os.O_WRONLY, 0666)
	}

	if err != nil {
		t.Fatal("Failed to open file", filename)
	}
	_, err = file.WriteString(data)
	if err != nil {
		t.Fatalf("Write file %s failed with error %s", filename, err)
	}
	err = file.Sync()
	if err != nil {
		t.Fatalf("Sync file %s failed with error %s", filename, err)
	}
	file.Close()

	event := expectEvent(t, tw)
	if event.Name != fullname || !event.IsModify() {
		t.Errorf("Expected modify event on \"%s\" but got %s", filename, event.String())
	}

	expectNoEvent(t, tw)
}

func mkDir(t *testing.T, tw *TreeWatcher, name string) {
	t.Log("Create directory", name)
	name = path.Join(tempDir, name)
	err := os.Mkdir(name, 0755)
	if err != nil {
		t.Fatalf("Mkdir %s failed with error %s", name, err)
	}
	event := expectEvent(t, tw)
	if !event.IsCreate() || event.Name != name {
		t.Errorf("Expected create event on \"%s\" but got %s", name, event.String())
	}

	expectNoEvent(t, tw)
}

func deleteFile(t *testing.T, tw *TreeWatcher, filename string) {
	t.Log("Remove file", filename)
	filename = path.Join(tempDir, filename)
	err := os.Remove(filename)
	if err != nil {
		t.Fatal("Remove file %s failed with error %s", filename, err)
	}
	event := expectEvent(t, tw)
	if !event.IsDelete() || event.Name != filename {
		t.Errorf("Expected create event on \"%s\" but got %s", filename, event.String())
	}
	expectNoEvent(t, tw)
}

func renameFile(t *testing.T, tw *TreeWatcher, oldName, newName string) {
	t.Logf("Rename %s to %s", oldName, newName)
	newName = path.Join(tempDir, newName)
	oldName = path.Join(tempDir, oldName)
	err := os.Rename(oldName, newName)
	if err != nil {
		t.Fatalf("Rename %s to %s failed with %s", oldName, newName, err)
	}

	events := make([]*fsnotify.FileEvent, 2)
	events[0] = expectEvent(t, tw)
	events[1] = expectEvent(t, tw)

	sawCreate := false
	sawRename := false

	for _, event := range events {
		if event.IsCreate() && event.Name == newName {
			sawCreate = true

		}

		if event.IsRename() && event.Name == oldName {
			sawRename = true
		}
	}

	if !sawCreate {
		t.Errorf("Expected create event on \"%s\" but got %v", newName, events)
	}

	if !sawRename {
		t.Errorf("Expected rename event on \"%s\" but got %v", oldName, events)
	}

	expectNoEvent(t, tw)
}

func TestTreeWatcher(t *testing.T) {
	var err error
	tempDir, err = ioutil.TempDir("", "treewatcher_test")
	if err != nil {
		t.Fatal("Could not create temp directory")
	}
	t.Log("Using temporary directory", tempDir)
	// defer os.RemoveAll(tempDir)

	tw, err := New()
	if err != nil {
		t.Fatal("Could not create watcher:", err)
	}
	defer tw.Close()

	tw.WatchTree(tempDir)

	writeFile(t, tw, "abc.txt", "abc")
	writeFile(t, tw, "def.txt", "def")
	writeFile(t, tw, "abc.txt", "def")
	deleteFile(t, tw, "def.txt")
	renameFile(t, tw, "abc.txt", "def.txt")
	mkDir(t, tw, "dir")
	writeFile(t, tw, "dir/abc.txt", "abc")
	mkDir(t, tw, "dir2")
	renameFile(t, tw, "dir/abc.txt", "dir2/abc.txt")
	mkDir(t, tw, "dir/a")
	mkDir(t, tw, "dir/a/b")
	writeFile(t, tw, "dir/a/b/c.txt", "jklsdf")
	renameFile(t, tw, "dir/a", "dir2/a")
	writeFile(t, tw, "dir2/a/b/c.txt", "ggg")

	t.Log("Creating unwatched tree ppp/qqq/rrr/a.txt")
	unwatchedDir, _ := ioutil.TempDir("", "ppp")
	defer os.RemoveAll(unwatchedDir)

	nestedDir := path.Join(unwatchedDir, "ppp", "qqq", "rrr")
	os.MkdirAll(nestedDir, 0755)
	ioutil.WriteFile(path.Join(nestedDir, "a.txt"), []byte("data"), 0666)
	expectNoEvent(t, tw)

	t.Log("Moving unwatched tree under watched tree")
	movedDir := path.Join(tempDir, "ppp")
	os.Rename(path.Join(unwatchedDir, "ppp"), movedDir)
	event := expectEvent(t, tw)
	if event.Name != movedDir || !event.IsCreate() {
		t.Fatalf("Expected create of %s, got %s", movedDir, event)
	}
	writeFile(t, tw, "ppp/qqq/rrr/new.txt", "data")
	writeFile(t, tw, "ppp/qqq/rrr/a.txt", "data")
	deleteFile(t, tw, "ppp/qqq/rrr/new.txt")
	expectNoEvent(t, tw)
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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.

// The ct_server binary runs the CT personality.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/etcd/clientv3"
	etcdnaming "github.com/coreos/etcd/clientv3/naming"
	"github.com/golang/glog"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/trillian/ctfe/configpb"
	"github.com/google/certificate-transparency-go/trillian/util"
	"github.com/google/trillian"
	"github.com/google/trillian/monitoring/opencensus"
	"github.com/google/trillian/monitoring/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/rs/cors"
	"github.com/tomasen/realip"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer/roundrobin"
	"google.golang.org/grpc/naming"

	// Register PEMKeyFile, PrivateKey and PKCS11Config ProtoHandlers
	_ "github.com/google/trillian/crypto/keys/der/proto"
	_ "github.com/google/trillian/crypto/keys/pem/proto"
	_ "github.com/google/trillian/crypto/keys/pkcs11/proto"
)

// Global flags that affect all log instances.
var (
	httpEndpoint       = flag.String("http_endpoint", "localhost:6962", "Endpoint for HTTP (host:port)")
	metricsEndpoint    = flag.String("metrics_endpoint", "localhost:6963", "Endpoint for serving metrics; if left empty, metrics will be visible on --http_endpoint")
	rpcBackend         = flag.String("log_rpc_server", "localhost:8090", "Backend specification; comma-separated list or etcd service name (if --etcd_servers specified). If unset backends are specified in config (as a LogMultiConfig proto)")
	rpcDeadline        = flag.Duration("rpc_deadline", time.Second*10, "Deadline for backend RPC requests")
	getSTHInterval     = flag.Duration("get_sth_interval", time.Second*180, "Interval between internal get-sth operations (0 to disable)")
	logConfig          = flag.String("log_config", "", "File holding log config in text proto format")
	maxGetEntries      = flag.Int64("max_get_entries", 0, "Max number of entries we allow in a get-entries request (0=>use default 1000)")
	etcdServers        = flag.String("etcd_servers", "", "A comma-separated list of etcd servers")
	etcdHTTPService    = flag.String("etcd_http_service", "trillian-ctfe-http", "Service name to announce our HTTP endpoint under")
	etcdMetricsService = flag.String("etcd_metrics_service", "trillian-ctfe-metrics-http", "Service name to announce our HTTP metrics endpoint under")
	tracing            = flag.Bool("tracing", false, "If true opencensus Stackdriver tracing will be enabled. See https://opencensus.io/.")
	tracingProjectID   = flag.String("tracing_project_id", "", "project ID to pass to stackdriver. Can be empty for GCP, consult docs for other platforms.")
	tracingPercent     = flag.Int("tracing_percent", 0, "Percent of requests to be traced. Zero is a special case to use the DefaultSampler")
	quotaRemote        = flag.Bool("quota_remote", true, "Enable requesting of quota for IP address sending incoming requests")
	quotaIntermediate  = flag.Bool("quota_intermediate", true, "Enable requesting of quota for intermediate certificates in sumbmitted chains")
	handlerPrefix      = flag.String("handler_prefix", "", "If set e.g. to '/logs' will prefix all handlers that don't define a custom prefix")
)

// nolint:staticcheck
func main() {
	flag.Parse()
	ctx := context.Background()

	if *maxGetEntries > 0 {
		ctfe.MaxGetEntriesAllowed = *maxGetEntries
	}

	var cfg *configpb.LogMultiConfig
	var err error
	// Get log config from file before we start. This is a different proto
	// type if we're using a multi backend configuration (no rpcBackend set
	// in flags). The single-backend config is converted to a multi config so
	// they can be treated the same.
	if len(*rpcBackend) > 0 {
		var cfgs []*configpb.LogConfig
		if cfgs, err = ctfe.LogConfigFromFile(*logConfig); err == nil {
			cfg = ctfe.ToMultiLogConfig(cfgs, *rpcBackend)
		}
	} else {
		cfg, err = ctfe.MultiLogConfigFromFile(*logConfig)
	}

	if err != nil {
		glog.Exitf("Failed to read config: %v", err)
	}

	beMap, err := ctfe.ValidateLogMultiConfig(cfg)
	if err != nil {
		glog.Exitf("Invalid config: %v", err)
	}

	glog.CopyStandardLogTo("WARNING")
	glog.Info("**** CT HTTP Server Starting ****")

	metricsAt := *metricsEndpoint
	if metricsAt == "" {
		metricsAt = *httpEndpoint
	}

	dialOpts := []grpc.DialOption{grpc.WithInsecure()}
	if len(*etcdServers) > 0 {
		// Use etcd to provide endpoint resolution.
		// TODO(daviddrysdale): re-enable dial timeout when upstream client code fixed
		// for https://github.com/grpc/grpc-go/pull/2733/files#r271705181
		cfg := clientv3.Config{
			Endpoints: strings.Split(*etcdServers, ","),
			// DialTimeout: 5 * time.Second,
		}
		client, err := clientv3.New(cfg)
		if err != nil {
			glog.Exitf("Failed to connect to etcd at %v: %v", *etcdServers, err)
		}
		etcdRes := &etcdnaming.GRPCResolver{Client: client}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(etcdRes)))

		// Also announce ourselves.
		updateHTTP := naming.Update{Op: naming.Add, Addr: *httpEndpoint}
		updateMetrics := naming.Update{Op: naming.Add, Addr: metricsAt}
		glog.Infof("Announcing our presence in %v with %+v", *etcdHTTPService, updateHTTP)
		etcdRes.Update(ctx, *etcdHTTPService, updateHTTP)
		glog.Infof("Announcing our presence in %v with %+v", *etcdMetricsService, updateMetrics)
		etcdRes.Update(ctx, *etcdMetricsService, updateMetrics)

		byeHTTP := naming.Update{Op: naming.Delete, Addr: *httpEndpoint}
		byeMetrics := naming.Update{Op: naming.Delete, Addr: metricsAt}
		defer func() {
			glog.Infof("Removing our presence in %v with %+v", *etcdHTTPService, byeHTTP)
			etcdRes.Update(ctx, *etcdHTTPService, byeHTTP)
			glog.Infof("Removing our presence in %v with %+v", *etcdMetricsService, byeMetrics)
			etcdRes.Update(ctx, *etcdMetricsService, byeMetrics)
		}()
	} else if strings.Contains(*rpcBackend, ",") {
		glog.Infof("Using FixedBackendResolver")
		// Use a fixed endpoint resolution that just returns the addresses configured on the command line.
		res := util.FixedBackendResolver{}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(res)))
	} else {
		glog.Infof("Using regular DNS resolver")
		dialOpts = append(dialOpts, grpc.WithBalancerName(roundrobin.Name))
	}

	// Dial all our log backends.
	clientMap := make(map[string]trillian.TrillianLogClient)
	for _, be := range beMap {
		glog.Infof("Dialling backend: %v", be)
		if len(beMap) == 1 {
			// If there's only one of them we use the blocking option as we can't
			// serve anything until connected.
			dialOpts = append(dialOpts, grpc.WithBlock())
		}
		conn, err := grpc.Dial(be.BackendSpec, dialOpts...)
		if err != nil {
			glog.Exitf("Could not dial RPC server: %v: %v", be, err)
		}
		defer conn.Close()
		clientMap[be.Name] = trillian.NewTrillianLogClient(conn)
	}

	// Allow cross-origin requests to all handlers registered on corsMux.
	// This is safe for CT log handlers because the log is public and
	// unauthenticated so cross-site scripting attacks are not a concern.
	corsMux := http.NewServeMux()
	corsHandler := cors.AllowAll().Handler(corsMux)
	http.Handle("/", corsHandler)

	// Register handlers for all the configured logs using the correct RPC
	// client.
	for _, c := range cfg.LogConfigs.Config {
		inst, err := setupAndRegister(ctx, clientMap[c.LogBackendName], *rpcDeadline, c, corsMux, *handlerPrefix)
		if err != nil {
			glog.Exitf("Failed to set up log instance for %+v: %v", cfg, err)
		}
		if *getSTHInterval > 0 {
			go inst.RunUpdateSTH(ctx, *getSTHInterval)
		}
	}

	// Return a 200 on the root, for GCE default health checking :/
	corsMux.HandleFunc("/", func(resp http.ResponseWriter, req *http.Request) {
		if req.URL.Path == "/" {
			resp.WriteHeader(http.StatusOK)
		} else {
			resp.WriteHeader(http.StatusNotFound)
		}
	})

	if metricsAt != *httpEndpoint {
		// Run a separate handler for metrics.
		go func() {
			mux := http.NewServeMux()
			mux.Handle("/metrics", promhttp.Handler())
			metricsServer := http.Server{Addr: metricsAt, Handler: mux}
			err := metricsServer.ListenAndServe()
			glog.Warningf("Metrics server exited: %v", err)
		}()
	} else {
		// Handle metrics on the DefaultServeMux.
		http.Handle("/metrics", promhttp.Handler())
	}

	// If we're enabling tracing we need to use an instrumented http.Handler.
	var handler http.Handler
	if *tracing {
		handler, err = opencensus.EnableHTTPServerTracing(*tracingProjectID, *tracingPercent)
		if err != nil {
			glog.Exitf("Failed to initialize stackdriver / opencensus tracing: %v", err)
		}
	}

	// Bring up the HTTP server and serve until we get a signal not to.
	srv := http.Server{Addr: *httpEndpoint, Handler: handler}
	shutdownWG := new(sync.WaitGroup)
	go awaitSignal(func() {
		shutdownWG.Add(1)
		defer shutdownWG.Done()
		// Allow 60s for any pending requests to finish then terminate any stragglers
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
		defer cancel()
		glog.Info("Shutting down HTTP server...")
		srv.Shutdown(ctx)
		glog.Info("HTTP server shutdown")
	})

	err = srv.ListenAndServe()
	if err != http.ErrServerClosed {
		glog.Warningf("Server exited: %v", err)
	}
	// Wait will only block if the function passed to awaitSignal was called,
	// in which case it'll block until the HTTP server has gracefully shutdown
	shutdownWG.Wait()
	glog.Flush()
}

// awaitSignal waits for standard termination signals, then runs the given
// function; it should be run as a separate goroutine.
func awaitSignal(doneFn func()) {
	// Arrange notification for the standard set of signals used to terminate a server
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	// Now block main and wait for a signal
	sig := <-sigs
	glog.Warningf("Signal received: %v", sig)
	glog.Flush()

	doneFn()
}

func setupAndRegister(ctx context.Context, client trillian.TrillianLogClient, deadline time.Duration, cfg *configpb.LogConfig, mux *http.ServeMux, globalHandlerPrefix string) (*ctfe.Instance, error) {
	vCfg, err := ctfe.ValidateLogConfig(cfg)
	if err != nil {
		return nil, err
	}

	opts := ctfe.InstanceOptions{
		Validated:     vCfg,
		Client:        client,
		Deadline:      deadline,
		MetricFactory: prometheus.MetricFactory{},
		RequestLog:    new(ctfe.DefaultRequestLog),
	}
	if *quotaRemote {
		glog.Info("Enabling quota for requesting IP")
		opts.RemoteQuotaUser = realip.FromRequest
	}
	if *quotaIntermediate {
		glog.Info("Enabling quota for intermediate certificates")
		opts.CertificateQuotaUser = ctfe.QuotaUserForCert
	}
	// Full handler pattern will be of the form "/logs/yyz/ct/v1/add-chain", where "/logs" is the
	// HandlerPrefix and "yyz" is the c.Prefix for this particular log. Use the default
	// HandlerPrefix unless the log config overrides it. The custom prefix in
	// the log configuration intended for use in migration scenarios where logs
	// have an existing URL path that differs from the global one. For example
	// if all new logs are served on "/logs/log/..." and a previously existing
	// log is at "/log/..." this is now supported.
	lhp := globalHandlerPrefix
	if ohPrefix := cfg.OverrideHandlerPrefix; len(ohPrefix) > 0 {
		glog.Infof("Log with prefix: %s is using a custom HandlerPrefix: %s", cfg.Prefix, ohPrefix)
		lhp = "/" + strings.Trim(ohPrefix, "/")
	}
	inst, err := ctfe.SetUpInstance(ctx, opts)
	if err != nil {
		return nil, err
	}
	for path, handler := range inst.Handlers {
		mux.Handle(lhp+path, handler)
	}
	return inst, nil
}
<commit_msg>Revert "Temporary workaround for upstream cancellation problem"<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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.

// The ct_server binary runs the CT personality.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/etcd/clientv3"
	etcdnaming "github.com/coreos/etcd/clientv3/naming"
	"github.com/golang/glog"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/trillian/ctfe/configpb"
	"github.com/google/certificate-transparency-go/trillian/util"
	"github.com/google/trillian"
	"github.com/google/trillian/monitoring/opencensus"
	"github.com/google/trillian/monitoring/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/rs/cors"
	"github.com/tomasen/realip"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer/roundrobin"
	"google.golang.org/grpc/naming"

	// Register PEMKeyFile, PrivateKey and PKCS11Config ProtoHandlers
	_ "github.com/google/trillian/crypto/keys/der/proto"
	_ "github.com/google/trillian/crypto/keys/pem/proto"
	_ "github.com/google/trillian/crypto/keys/pkcs11/proto"
)

// Global flags that affect all log instances.
var (
	httpEndpoint       = flag.String("http_endpoint", "localhost:6962", "Endpoint for HTTP (host:port)")
	metricsEndpoint    = flag.String("metrics_endpoint", "localhost:6963", "Endpoint for serving metrics; if left empty, metrics will be visible on --http_endpoint")
	rpcBackend         = flag.String("log_rpc_server", "localhost:8090", "Backend specification; comma-separated list or etcd service name (if --etcd_servers specified). If unset backends are specified in config (as a LogMultiConfig proto)")
	rpcDeadline        = flag.Duration("rpc_deadline", time.Second*10, "Deadline for backend RPC requests")
	getSTHInterval     = flag.Duration("get_sth_interval", time.Second*180, "Interval between internal get-sth operations (0 to disable)")
	logConfig          = flag.String("log_config", "", "File holding log config in text proto format")
	maxGetEntries      = flag.Int64("max_get_entries", 0, "Max number of entries we allow in a get-entries request (0=>use default 1000)")
	etcdServers        = flag.String("etcd_servers", "", "A comma-separated list of etcd servers")
	etcdHTTPService    = flag.String("etcd_http_service", "trillian-ctfe-http", "Service name to announce our HTTP endpoint under")
	etcdMetricsService = flag.String("etcd_metrics_service", "trillian-ctfe-metrics-http", "Service name to announce our HTTP metrics endpoint under")
	tracing            = flag.Bool("tracing", false, "If true opencensus Stackdriver tracing will be enabled. See https://opencensus.io/.")
	tracingProjectID   = flag.String("tracing_project_id", "", "project ID to pass to stackdriver. Can be empty for GCP, consult docs for other platforms.")
	tracingPercent     = flag.Int("tracing_percent", 0, "Percent of requests to be traced. Zero is a special case to use the DefaultSampler")
	quotaRemote        = flag.Bool("quota_remote", true, "Enable requesting of quota for IP address sending incoming requests")
	quotaIntermediate  = flag.Bool("quota_intermediate", true, "Enable requesting of quota for intermediate certificates in sumbmitted chains")
	handlerPrefix      = flag.String("handler_prefix", "", "If set e.g. to '/logs' will prefix all handlers that don't define a custom prefix")
)

// nolint:staticcheck
func main() {
	flag.Parse()
	ctx := context.Background()

	if *maxGetEntries > 0 {
		ctfe.MaxGetEntriesAllowed = *maxGetEntries
	}

	var cfg *configpb.LogMultiConfig
	var err error
	// Get log config from file before we start. This is a different proto
	// type if we're using a multi backend configuration (no rpcBackend set
	// in flags). The single-backend config is converted to a multi config so
	// they can be treated the same.
	if len(*rpcBackend) > 0 {
		var cfgs []*configpb.LogConfig
		if cfgs, err = ctfe.LogConfigFromFile(*logConfig); err == nil {
			cfg = ctfe.ToMultiLogConfig(cfgs, *rpcBackend)
		}
	} else {
		cfg, err = ctfe.MultiLogConfigFromFile(*logConfig)
	}

	if err != nil {
		glog.Exitf("Failed to read config: %v", err)
	}

	beMap, err := ctfe.ValidateLogMultiConfig(cfg)
	if err != nil {
		glog.Exitf("Invalid config: %v", err)
	}

	glog.CopyStandardLogTo("WARNING")
	glog.Info("**** CT HTTP Server Starting ****")

	metricsAt := *metricsEndpoint
	if metricsAt == "" {
		metricsAt = *httpEndpoint
	}

	dialOpts := []grpc.DialOption{grpc.WithInsecure()}
	if len(*etcdServers) > 0 {
		// Use etcd to provide endpoint resolution.
		cfg := clientv3.Config{Endpoints: strings.Split(*etcdServers, ","), DialTimeout: 5 * time.Second}
		client, err := clientv3.New(cfg)
		if err != nil {
			glog.Exitf("Failed to connect to etcd at %v: %v", *etcdServers, err)
		}
		etcdRes := &etcdnaming.GRPCResolver{Client: client}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(etcdRes)))

		// Also announce ourselves.
		updateHTTP := naming.Update{Op: naming.Add, Addr: *httpEndpoint}
		updateMetrics := naming.Update{Op: naming.Add, Addr: metricsAt}
		glog.Infof("Announcing our presence in %v with %+v", *etcdHTTPService, updateHTTP)
		etcdRes.Update(ctx, *etcdHTTPService, updateHTTP)
		glog.Infof("Announcing our presence in %v with %+v", *etcdMetricsService, updateMetrics)
		etcdRes.Update(ctx, *etcdMetricsService, updateMetrics)

		byeHTTP := naming.Update{Op: naming.Delete, Addr: *httpEndpoint}
		byeMetrics := naming.Update{Op: naming.Delete, Addr: metricsAt}
		defer func() {
			glog.Infof("Removing our presence in %v with %+v", *etcdHTTPService, byeHTTP)
			etcdRes.Update(ctx, *etcdHTTPService, byeHTTP)
			glog.Infof("Removing our presence in %v with %+v", *etcdMetricsService, byeMetrics)
			etcdRes.Update(ctx, *etcdMetricsService, byeMetrics)
		}()
	} else if strings.Contains(*rpcBackend, ",") {
		glog.Infof("Using FixedBackendResolver")
		// Use a fixed endpoint resolution that just returns the addresses configured on the command line.
		res := util.FixedBackendResolver{}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(res)))
	} else {
		glog.Infof("Using regular DNS resolver")
		dialOpts = append(dialOpts, grpc.WithBalancerName(roundrobin.Name))
	}

	// Dial all our log backends.
	clientMap := make(map[string]trillian.TrillianLogClient)
	for _, be := range beMap {
		glog.Infof("Dialling backend: %v", be)
		if len(beMap) == 1 {
			// If there's only one of them we use the blocking option as we can't
			// serve anything until connected.
			dialOpts = append(dialOpts, grpc.WithBlock())
		}
		conn, err := grpc.Dial(be.BackendSpec, dialOpts...)
		if err != nil {
			glog.Exitf("Could not dial RPC server: %v: %v", be, err)
		}
		defer conn.Close()
		clientMap[be.Name] = trillian.NewTrillianLogClient(conn)
	}

	// Allow cross-origin requests to all handlers registered on corsMux.
	// This is safe for CT log handlers because the log is public and
	// unauthenticated so cross-site scripting attacks are not a concern.
	corsMux := http.NewServeMux()
	corsHandler := cors.AllowAll().Handler(corsMux)
	http.Handle("/", corsHandler)

	// Register handlers for all the configured logs using the correct RPC
	// client.
	for _, c := range cfg.LogConfigs.Config {
		inst, err := setupAndRegister(ctx, clientMap[c.LogBackendName], *rpcDeadline, c, corsMux, *handlerPrefix)
		if err != nil {
			glog.Exitf("Failed to set up log instance for %+v: %v", cfg, err)
		}
		if *getSTHInterval > 0 {
			go inst.RunUpdateSTH(ctx, *getSTHInterval)
		}
	}

	// Return a 200 on the root, for GCE default health checking :/
	corsMux.HandleFunc("/", func(resp http.ResponseWriter, req *http.Request) {
		if req.URL.Path == "/" {
			resp.WriteHeader(http.StatusOK)
		} else {
			resp.WriteHeader(http.StatusNotFound)
		}
	})

	if metricsAt != *httpEndpoint {
		// Run a separate handler for metrics.
		go func() {
			mux := http.NewServeMux()
			mux.Handle("/metrics", promhttp.Handler())
			metricsServer := http.Server{Addr: metricsAt, Handler: mux}
			err := metricsServer.ListenAndServe()
			glog.Warningf("Metrics server exited: %v", err)
		}()
	} else {
		// Handle metrics on the DefaultServeMux.
		http.Handle("/metrics", promhttp.Handler())
	}

	// If we're enabling tracing we need to use an instrumented http.Handler.
	var handler http.Handler
	if *tracing {
		handler, err = opencensus.EnableHTTPServerTracing(*tracingProjectID, *tracingPercent)
		if err != nil {
			glog.Exitf("Failed to initialize stackdriver / opencensus tracing: %v", err)
		}
	}

	// Bring up the HTTP server and serve until we get a signal not to.
	srv := http.Server{Addr: *httpEndpoint, Handler: handler}
	shutdownWG := new(sync.WaitGroup)
	go awaitSignal(func() {
		shutdownWG.Add(1)
		defer shutdownWG.Done()
		// Allow 60s for any pending requests to finish then terminate any stragglers
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
		defer cancel()
		glog.Info("Shutting down HTTP server...")
		srv.Shutdown(ctx)
		glog.Info("HTTP server shutdown")
	})

	err = srv.ListenAndServe()
	if err != http.ErrServerClosed {
		glog.Warningf("Server exited: %v", err)
	}
	// Wait will only block if the function passed to awaitSignal was called,
	// in which case it'll block until the HTTP server has gracefully shutdown
	shutdownWG.Wait()
	glog.Flush()
}

// awaitSignal waits for standard termination signals, then runs the given
// function; it should be run as a separate goroutine.
func awaitSignal(doneFn func()) {
	// Arrange notification for the standard set of signals used to terminate a server
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	// Now block main and wait for a signal
	sig := <-sigs
	glog.Warningf("Signal received: %v", sig)
	glog.Flush()

	doneFn()
}

func setupAndRegister(ctx context.Context, client trillian.TrillianLogClient, deadline time.Duration, cfg *configpb.LogConfig, mux *http.ServeMux, globalHandlerPrefix string) (*ctfe.Instance, error) {
	vCfg, err := ctfe.ValidateLogConfig(cfg)
	if err != nil {
		return nil, err
	}

	opts := ctfe.InstanceOptions{
		Validated:     vCfg,
		Client:        client,
		Deadline:      deadline,
		MetricFactory: prometheus.MetricFactory{},
		RequestLog:    new(ctfe.DefaultRequestLog),
	}
	if *quotaRemote {
		glog.Info("Enabling quota for requesting IP")
		opts.RemoteQuotaUser = realip.FromRequest
	}
	if *quotaIntermediate {
		glog.Info("Enabling quota for intermediate certificates")
		opts.CertificateQuotaUser = ctfe.QuotaUserForCert
	}
	// Full handler pattern will be of the form "/logs/yyz/ct/v1/add-chain", where "/logs" is the
	// HandlerPrefix and "yyz" is the c.Prefix for this particular log. Use the default
	// HandlerPrefix unless the log config overrides it. The custom prefix in
	// the log configuration intended for use in migration scenarios where logs
	// have an existing URL path that differs from the global one. For example
	// if all new logs are served on "/logs/log/..." and a previously existing
	// log is at "/log/..." this is now supported.
	lhp := globalHandlerPrefix
	if ohPrefix := cfg.OverrideHandlerPrefix; len(ohPrefix) > 0 {
		glog.Infof("Log with prefix: %s is using a custom HandlerPrefix: %s", cfg.Prefix, ohPrefix)
		lhp = "/" + strings.Trim(ohPrefix, "/")
	}
	inst, err := ctfe.SetUpInstance(ctx, opts)
	if err != nil {
		return nil, err
	}
	for path, handler := range inst.Handlers {
		mux.Handle(lhp+path, handler)
	}
	return inst, nil
}
<|endoftext|>


<commit_before>package ui

import (
	"strings"

	"github.com/fatih/color"
)

func (ui *UI) DisplayInstancesTableForApp(table [][]string) {
	redColor := color.New(color.FgRed, color.Bold)
	trDown, trCrashed := ui.TranslateText("down"), ui.TranslateText("crashed")

	for i, row := range table {
		if row[1] == trDown || row[1] == trCrashed {
			table[i][1] = ui.modifyColor(row[1], redColor)
		}
	}
	ui.DisplayTableWithHeader("", table, 3)
}

func (ui *UI) DisplayKeyValueTableForApp(table [][]string) {
	runningInstances := strings.Split(table[2][1], "/")[0]
	state := table[1][1]

	if runningInstances == "0" && state != ui.TranslateText("stopped") {
		redColor := color.New(color.FgRed, color.Bold)
		table[1][1] = ui.modifyColor(table[1][1], redColor)
		table[2][1] = ui.modifyColor(table[2][1], redColor)
	}
	ui.DisplayKeyValueTable("", table, 3)
}

func (ui *UI) DisplayKeyValueTableForV3App(table [][]string, crashedProcesses []string) {
	if len(crashedProcesses) > 0 {
		redColor := color.New(color.FgRed, color.Bold)
		table[1][1] = ui.modifyColor(table[1][1], redColor)

		processes := strings.Split(table[2][1], ",")
		newProcesses := []string{}
		for _, process := range processes {
			parts := strings.Split(process, ":")
			isCrashedProcess := false
			for _, crashedProcess := range crashedProcesses {
				if parts[0] == crashedProcess {
					isCrashedProcess = true
					break
				}
			}

			if isCrashedProcess {
				println("modifying color")
				newProcesses = append(newProcesses, ui.modifyColor(process, redColor))
			} else {
				newProcesses = append(newProcesses, process)
			}
		}

		table[2][1] = strings.Join(newProcesses, ",")
	}

	ui.DisplayKeyValueTable("", table, 3)
}
<commit_msg>remove debugging statement<commit_after>package ui

import (
	"strings"

	"github.com/fatih/color"
)

func (ui *UI) DisplayInstancesTableForApp(table [][]string) {
	redColor := color.New(color.FgRed, color.Bold)
	trDown, trCrashed := ui.TranslateText("down"), ui.TranslateText("crashed")

	for i, row := range table {
		if row[1] == trDown || row[1] == trCrashed {
			table[i][1] = ui.modifyColor(row[1], redColor)
		}
	}
	ui.DisplayTableWithHeader("", table, 3)
}

func (ui *UI) DisplayKeyValueTableForApp(table [][]string) {
	runningInstances := strings.Split(table[2][1], "/")[0]
	state := table[1][1]

	if runningInstances == "0" && state != ui.TranslateText("stopped") {
		redColor := color.New(color.FgRed, color.Bold)
		table[1][1] = ui.modifyColor(table[1][1], redColor)
		table[2][1] = ui.modifyColor(table[2][1], redColor)
	}
	ui.DisplayKeyValueTable("", table, 3)
}

func (ui *UI) DisplayKeyValueTableForV3App(table [][]string, crashedProcesses []string) {
	if len(crashedProcesses) > 0 {
		redColor := color.New(color.FgRed, color.Bold)
		table[1][1] = ui.modifyColor(table[1][1], redColor)

		processes := strings.Split(table[2][1], ",")
		newProcesses := []string{}
		for _, process := range processes {
			parts := strings.Split(process, ":")
			isCrashedProcess := false
			for _, crashedProcess := range crashedProcesses {
				if parts[0] == crashedProcess {
					isCrashedProcess = true
					break
				}
			}

			if isCrashedProcess {
				newProcesses = append(newProcesses, ui.modifyColor(process, redColor))
			} else {
				newProcesses = append(newProcesses, process)
			}
		}

		table[2][1] = strings.Join(newProcesses, ",")
	}

	ui.DisplayKeyValueTable("", table, 3)
}
<|endoftext|>


<commit_before>package swf

import (
	"errors"
	"fmt"
	"log"
	"reflect"
)

// The marker name used when recording the current state and data of a workflow
const (
	STATE_MARKER        = "FSM.State"
	ERROR_SIGNAL        = "FSM.Error"
	SYSTEM_ERROR_SIGNAL = "FSM.SystemError"
)

// Decider decides an Outcome based on an event and the current data for an FSM
type Decider func(*FSM, HistoryEvent, interface{}) *Outcome

// EventDataType should return an empty struct of the correct type based on the event
// the FSM will unmarshal data from the event into this struct
type EventDataType func(HistoryEvent) interface{}

type ErrorHandler func(*FSM, *FSMState, HistoryEvent, []HistoryEvent, interface{}, error) *Outcome

// on error state marker FSM.Error, serialize ErrorState(HistoryEvent, []HistoryEvent, interface{[])
// signal self with signal = "ERROR"

// Outcome is created by Deciders
type Outcome struct {
	Data      interface{}
	NextState string
	Decisions []*Decision
}

// FSMState defines the behavior of one state of an FSM
type FSMState struct {
	Name    string
	Decider Decider
}

// FSM models the decision handling logic a workflow in SWF
type FSM struct {
	Name           string
	Domain         string
	TaskList       string
	Identity       string
	DecisionWorker *DecisionWorker
	Input          chan *PollForDecisionTaskResponse
	DataType       interface{}
	EventDataType  EventDataType
	ErrorHandler   ErrorHandler
	states         map[string]*FSMState
	initialState   *FSMState
	stop           chan bool
}

func (f *FSM) AddInitialState(state *FSMState) {
	f.AddState(state)
	f.initialState = state
}
func (f *FSM) AddState(state *FSMState) {
	if f.states == nil {
		f.states = make(map[string]*FSMState)
	}
	f.states[state.Name] = state
}

func (f *FSM) Start() {
	if f.initialState == nil {
		panic("No Initial State Defined For FSM")
	}
	if f.stop == nil {
		f.stop = make(chan bool)
	}

	go func() {
		poller := f.DecisionWorker.PollTaskList(f.Domain, f.Identity, f.TaskList, f.Input)
		for {
			select {
			case decisionTask, ok := <-f.Input:
				if ok {
					decisions := f.Tick(decisionTask)
					err := f.DecisionWorker.Decide(decisionTask.TaskToken, decisions)
					if err != nil {
						f.log("action=tick at=decide-request-failed error=%s", err.Error())
						//TODO Retry the Decide?
						poller.Stop()
						return
					}

				} else {
					poller.Stop()
					return
				}
			case <-f.stop:
				poller.Stop()
				return
			}
		}
	}()
}

func (f *FSM) Tick(decisionTask *PollForDecisionTaskResponse) []*Decision {
	execution := decisionTask.WorkflowExecution
	serializedState, err := f.findSerializedState(decisionTask.Events)

	if err != nil {
		return f.captureSystemError(execution, "FindSerializedStateError", decisionTask.Events, err)
	}

	f.log("action=tick at=find-current-state state=%s", serializedState.State)
	data := reflect.New(reflect.TypeOf(f.DataType)).Interface()
	err = f.Serializer().Deserialize(serializedState.Data, data)
	if err != nil {
		f.log("action=tick at=error=deserialize-state-failed")
		return f.captureSystemError(execution, "DeserializeStateError", decisionTask.Events, err)
	}

	f.log("action=tick at=find-current-data data=%v", data)
	lastEvents := f.findLastEvents(decisionTask.PreviousStartedEventId, decisionTask.Events)

	outcome := new(Outcome)
	outcome.Data = data
	outcome.NextState = serializedState.State

	//iterate through events oldest to newest, calling the decider for the current state.
	//if the outcome changes the state use the right FSMState
	for i := len(lastEvents) - 1; i >= 0; i-- {
		e := lastEvents[i]
		f.log("action=tick at=history id=%d type=%s", e.EventId, e.EventType)
		fsmState, ok := f.states[outcome.NextState]
		if ok {
			anOutcome, err := f.decide(fsmState, e, outcome.Data)
			if err != nil {
				f.log("at=error error=decision-execution-error state=%s next-state=%", fsmState.Name, outcome.NextState)
				return f.captureDecisionError(execution, i, lastEvents, outcome.NextState, outcome.Data, err)
			}

			f.log("action=tick at=decided-event state=%s next-state=%s decisions=%d", outcome.NextState, anOutcome.NextState, len(anOutcome.Decisions))
			outcome.Data = anOutcome.Data
			outcome.NextState = anOutcome.NextState
			outcome.Decisions = append(outcome.Decisions, anOutcome.Decisions...)
		} else {
			f.log("action=tick at=error error=marked-state-not-in-fsm state=%s", outcome.NextState)
			return f.captureSystemError(execution, "MissingFsmStateError", lastEvents[i:], errors.New(outcome.NextState))
		}
	}

	f.log("action=tick at=events-processed next-state=%s decisions=%d", outcome.NextState, len(outcome.Decisions))

	for _, d := range outcome.Decisions {
		f.log("action=tick at=decide next-state=%s decision=%s", outcome.NextState, d.DecisionType)
	}

	final, err := f.appendState(outcome)
	if err != nil {
		f.log("action=tick at=error error=state-serialization-error error-type=system")
		return append(outcome.Decisions, f.captureSystemError(execution, "StateSerializationError", []HistoryEvent{}, err)...)
	}
	return final
}

//if the outcome is good good if its an error, we capture the error state above

func (f *FSM) decide(state *FSMState, event HistoryEvent, data interface{}) (anOutcome *Outcome, anErr error) {
	defer func() {
		if r := recover(); r != nil {
			f.log("at=error error=decide-panic-recovery")
			anErr = errors.New("panic in decider, capture error state")
		}
	}()
	anOutcome = state.Decider(f, event, data)
	return
}

func (f *FSM) captureDecisionError(execution WorkflowExecution, event int, lastEvents []HistoryEvent, stateName string, stateData interface{}, err error) []*Decision {
	return f.captureError(ERROR_SIGNAL, execution, &SerializedDecisionError{
		ErrorEvent:        lastEvents[event],
		UnprocessedEvents: lastEvents[event+1:],
		StateName:         stateName,
		StateData:         stateData,
	})
}

func (f *FSM) captureSystemError(execution WorkflowExecution, errorType string, lastEvents []HistoryEvent, err error) []*Decision {
	return f.captureError(SYSTEM_ERROR_SIGNAL, execution, &SerializedSystemError{
		ErrorType:         errorType,
		UnprocessedEvents: lastEvents,
		Error:             err,
	})
}

func (f *FSM) captureError(signal string, execution WorkflowExecution, error interface{}) []*Decision {
	decisions := f.EmptyDecisions()
	r, err := f.DecisionWorker.RecordMarker(signal, error)
	if err != nil {
		//really bail
		panic("giving up, cant even create a RecordMarker decsion")
	}
	d := &Decision{
		DecisionType: DecisionTypeSignalExternalWorkflowExecution,
		SignalExternalWorkflowExecutionDecisionAttributes: &SignalExternalWorkflowExecutionDecisionAttributes{
			WorkflowId: execution.WorkflowId,
			RunId:      execution.RunId,
			SignalName: signal,
			Input:      r.RecordMarkerDecisionAttributes.Details,
		},
	}
	return append(decisions, d)
}

func (f *FSM) EventData(event HistoryEvent) interface{} {
	eventData := f.EventDataType(event)

	if eventData != nil {
		var serialized string
		switch event.EventType {
		case EventTypeActivityTaskCompleted:
			serialized = event.ActivityTaskCompletedEventAttributes.Result
		case EventTypeWorkflowExecutionCompleted:
			serialized = event.WorkflowExecutionCompletedEventAttributes.Result
		case EventTypeChildWorkflowExecutionCompleted:
			serialized = event.ChildWorkflowExecutionCompletedEventAttributes.Result
		case EventTypeWorkflowExecutionSignaled:
			serialized = event.WorkflowExecutionSignaledEventAttributes.Input
		case EventTypeWorkflowExecutionStarted:
			serialized = event.WorkflowExecutionStartedEventAttributes.Input
		case EventTypeWorkflowExecutionContinuedAsNew:
			serialized = event.WorkflowExecutionContinuedAsNewEventAttributes.Input
		}
		if serialized != "" {
			err := f.Serializer().Deserialize(serialized, eventData)
			if err != nil {
				f.log("action=EventData at=error error=unable-to-deserialize")
				panic("Unable to Deserialize Event Data")
			}
		}
	}

	return eventData

}

func (f *FSM) log(format string, data ...interface{}) {
	actualFormat := fmt.Sprintf("component=FSM name=%s %s", f.Name, format)
	log.Printf(actualFormat, data...)
}

func (f *FSM) findSerializedState(events []HistoryEvent) (*SerializedState, error) {
	for _, event := range events {
		if f.isStateMarker(event) {
			state := &SerializedState{}
			err := f.Serializer().Deserialize(event.MarkerRecordedEventAttributes.Details, state)
			return state, err
		} else if event.EventType == EventTypeWorkflowExecutionStarted {
			return &SerializedState{State: f.initialState.Name, Data: event.WorkflowExecutionStartedEventAttributes.Input}, nil
		}
	}
	return nil, errors.New("Cant Find Current Data")
}

func (f *FSM) findLastEvents(prevStarted int, events []HistoryEvent) []HistoryEvent {
	lastEvents := make([]HistoryEvent, 0)
	for _, event := range events {
		if event.EventId == prevStarted {
			return lastEvents
		} else {
			switch event.EventType {
			case EventTypeDecisionTaskCompleted, EventTypeDecisionTaskScheduled,
				EventTypeDecisionTaskStarted, EventTypeDecisionTaskTimedOut:
				//no-op
			case EventTypeMarkerRecorded:
				if !f.isStateMarker(event) {
					lastEvents = append(lastEvents, event)
				}
			default:
				lastEvents = append(lastEvents, event)
			}
		}
	}

	return lastEvents
}

func (f *FSM) appendState(outcome *Outcome) ([]*Decision, error) {

	serializedData, err := f.Serializer().Serialize(outcome.Data)

	state := SerializedState{
		State: outcome.NextState,
		Data:  serializedData,
	}

	d, err := f.DecisionWorker.RecordMarker(STATE_MARKER, state)
	if err != nil {
		return nil, err
	}
	decisions := f.EmptyDecisions()
	decisions = append(decisions, d)
	decisions = append(decisions, outcome.Decisions...)
	return decisions, nil
}

func (f *FSM) Stop() {
	f.stop <- true
}

func (f *FSM) isStateMarker(e HistoryEvent) bool {
	return e.EventType == EventTypeMarkerRecorded && e.MarkerRecordedEventAttributes.MarkerName == STATE_MARKER
}

func (f *FSM) Serializer() StateSerializer {
	return f.DecisionWorker.StateSerializer
}

func (f *FSM) EmptyDecisions() []*Decision {
	return make([]*Decision, 0)
}

type SerializedState struct {
	State string `json:"state"`
	Data  string `json:"data"`
}

type SerializedDecisionError struct {
	ErrorEvent        HistoryEvent   `json:"errorEvent"`
	UnprocessedEvents []HistoryEvent `json:"unprocessedEvents"`
	StateName         string         `json:"stateName"`
	StateData         interface{}    `json:"stateData`
}

type SerializedSystemError struct {
	ErrorType         string         `json:"errorType"`
	Error             interface{}    `json:"error"`
	UnprocessedEvents []HistoryEvent `json:"unprocessedEvents"`
}

type DecisionErrorPointer struct {
	Error error
}

/*tigertonic-like marhslling of data*/
type MarshalledDecider struct {
	v reflect.Value
}

func TypedDecider(decider interface{}) Decider {
	t := reflect.TypeOf(decider)
	if reflect.Func != t.Kind() {
		panic(fmt.Sprintf("kind was %v, not Func", t.Kind()))
	}
	if 3 != t.NumIn() {
		panic(fmt.Sprintf(
			"input arity was %v, not 3",
			t.NumIn(),
		))
	}
	if "*swf.FSM" != t.In(0).String() {
		panic(fmt.Sprintf(
			"type of first argument was %v, not *swf.FSM",
			t.In(0),
		))
	}
	if "swf.HistoryEvent" != t.In(1).String() {
		panic(fmt.Sprintf(
			"type of second argument was %v, not swf.HistoryEvent",
			t.In(1),
		))
	}

	if "*swf.Outcome" != t.Out(0).String() {
		panic(fmt.Sprintf(
			"type of return value was %v, not *swf.Outcome",
			t.Out(0),
		))
	}

	return MarshalledDecider{reflect.ValueOf(decider)}.Decide
}

func (m MarshalledDecider) Decide(f *FSM, h HistoryEvent, data interface{}) *Outcome {
	return m.v.Call([]reflect.Value{reflect.ValueOf(f), reflect.ValueOf(h), reflect.ValueOf(data)})[0].Interface().(*Outcome)
}
<commit_msg>serialize event ids, not the full event in Error States,  add the concept of an ErrorState and a default impl to FSM<commit_after>package swf

import (
	"errors"
	"fmt"
	"log"
	"reflect"
)

// The marker name used when recording the current state and data of a workflow
const (
	STATE_MARKER        = "FSM.State"
	ERROR_SIGNAL        = "FSM.Error"
	SYSTEM_ERROR_SIGNAL = "FSM.SystemError"
)

// Decider decides an Outcome based on an event and the current data for an FSM
// the interface{} parameter that is passed to the decider is safe to
// be asserted to be the type of the DataType field in the FSM
// Alternatively you can use the TypedDecider to avoid having to do the assertion.
type Decider func(*FSM, HistoryEvent, interface{}) *Outcome

// EventDataType should return an empty struct of the correct type based on the event
// the FSM will unmarshal data from the event into this struct
type EventDataType func(HistoryEvent) interface{}

type ErrorHandler func(*FSM, *FSMState, HistoryEvent, []HistoryEvent, interface{}, error) *Outcome

// on error state marker FSM.Error, serialize ErrorState(HistoryEvent, []HistoryEvent, interface{[])
// signal self with signal = "ERROR"

// Outcome is created by Deciders
type Outcome struct {
	Data      interface{}
	NextState string
	Decisions []*Decision
}

// FSMState defines the behavior of one state of an FSM
type FSMState struct {
	Name    string
	Decider Decider
}

// FSM models the decision handling logic a workflow in SWF
type FSM struct {
	Name           string
	Domain         string
	TaskList       string
	Identity       string
	DecisionWorker *DecisionWorker
	Input          chan *PollForDecisionTaskResponse
	DataType       interface{}
	EventDataType  EventDataType
	ErrorHandler   ErrorHandler
	states         map[string]*FSMState
	initialState   *FSMState
	errorState     *FSMState
	stop           chan bool
}

func (f *FSM) AddInitialState(state *FSMState) {
	f.AddState(state)
	f.initialState = state
}
func (f *FSM) AddState(state *FSMState) {
	if f.states == nil {
		f.states = make(map[string]*FSMState)
	}
	f.states[state.Name] = state
}

func (f *FSM) AddErrorState(state *FSMState) {
	f.AddState(state)
	f.errorState = state
}

func (f *FSM) DefaultErrorState() *FSMState {
	return &FSMState{
		Name: "error",
		Decider: func(f *FSM, h HistoryEvent, data interface{}) *Outcome {
			switch h.EventType {
			case EventTypeWorkflowExecutionSignaled:
				{
					switch h.WorkflowExecutionSignaledEventAttributes.SignalName {
					case ERROR_SIGNAL:
						err := &SerializedDecisionError{}
						f.Serializer().Deserialize(h.WorkflowExecutionSignaledEventAttributes.Input, err)
						f.log("action=default-handle-error at=handle-decision-error error=%+v", err)
						f.log("YOU SHOULD CREATE AN ERROR STATE FOR YOUR FSM, Workflow %s is Hung", h.WorkflowExecutionSignaledEventAttributes.ExternalWorkflowExecution.WorkflowId)
					case SYSTEM_ERROR_SIGNAL:
						err := &SerializedSystemError{}
						f.Serializer().Deserialize(h.WorkflowExecutionSignaledEventAttributes.Input, err)
						f.log("action=default-handle-error at=handle-system-error error=%+v", err)
						f.log("YOU SHOULD CREATE AN ERROR STATE FOR YOUR FSM, Workflow %s is Hung", h.WorkflowExecutionSignaledEventAttributes.ExternalWorkflowExecution.WorkflowId)
					}
				}
			}
			return &Outcome{NextState: "error", Data: data, Decisions: []*Decision{}}
		},
	}
}

func (f *FSM) Start() {
	if f.initialState == nil {
		panic("No Initial State Defined For FSM")
	}

	if f.errorState == nil {
		f.AddErrorState(f.DefaultErrorState())
	}

	if f.stop == nil {
		f.stop = make(chan bool)
	}

	go func() {
		poller := f.DecisionWorker.PollTaskList(f.Domain, f.Identity, f.TaskList, f.Input)
		for {
			select {
			case decisionTask, ok := <-f.Input:
				if ok {
					decisions := f.Tick(decisionTask)
					err := f.DecisionWorker.Decide(decisionTask.TaskToken, decisions)
					if err != nil {
						f.log("action=tick at=decide-request-failed error=%s", err.Error())
						//TODO Retry the Decide?
						poller.Stop()
						return
					}

				} else {
					poller.Stop()
					return
				}
			case <-f.stop:
				poller.Stop()
				return
			}
		}
	}()
}

func (f *FSM) Tick(decisionTask *PollForDecisionTaskResponse) []*Decision {
	execution := decisionTask.WorkflowExecution
	serializedState, err := f.findSerializedState(decisionTask.Events)

	if err != nil {
		return f.captureSystemError(execution, "FindSerializedStateError", decisionTask.Events, err)
	}

	f.log("action=tick at=find-current-state state=%s", serializedState.StateName)
	data := reflect.New(reflect.TypeOf(f.DataType)).Interface()
	err = f.Serializer().Deserialize(serializedState.StateData, data)
	if err != nil {
		f.log("action=tick at=error=deserialize-state-failed")
		return f.captureSystemError(execution, "DeserializeStateError", decisionTask.Events, err)
	}

	f.log("action=tick at=find-current-data data=%v", data)
	lastEvents := f.findLastEvents(decisionTask.PreviousStartedEventId, decisionTask.Events)

	outcome := new(Outcome)
	outcome.Data = data
	outcome.NextState = serializedState.StateName

	//iterate through events oldest to newest, calling the decider for the current state.
	//if the outcome changes the state use the right FSMState
	for i := len(lastEvents) - 1; i >= 0; i-- {
		e := lastEvents[i]
		f.log("action=tick at=history id=%d type=%s", e.EventId, e.EventType)
		fsmState, ok := f.states[outcome.NextState]
		if ok {
			anOutcome, err := f.panicSafeDecide(fsmState, e, outcome.Data)
			if err != nil {
				f.log("at=error error=decision-execution-error state=%s next-state=%", fsmState.Name, outcome.NextState)
				return f.captureDecisionError(execution, i, lastEvents, outcome.NextState, outcome.Data, err)
			}

			f.log("action=tick at=decided-event state=%s next-state=%s decisions=%d", outcome.NextState, anOutcome.NextState, len(anOutcome.Decisions))
			outcome.Data = anOutcome.Data
			outcome.NextState = anOutcome.NextState
			outcome.Decisions = append(outcome.Decisions, anOutcome.Decisions...)
		} else {
			f.log("action=tick at=error error=marked-state-not-in-fsm state=%s", outcome.NextState)
			return f.captureSystemError(execution, "MissingFsmStateError", lastEvents[i:], errors.New(outcome.NextState))
		}
	}

	f.log("action=tick at=events-processed next-state=%s decisions=%d", outcome.NextState, len(outcome.Decisions))

	for _, d := range outcome.Decisions {
		f.log("action=tick at=decide next-state=%s decision=%s", outcome.NextState, d.DecisionType)
	}

	final, err := f.appendState(outcome)
	if err != nil {
		f.log("action=tick at=error error=state-serialization-error error-type=system")
		return append(outcome.Decisions, f.captureSystemError(execution, "StateSerializationError", []HistoryEvent{}, err)...)
	}
	return final
}

//if the outcome is good good if its an error, we capture the error state above

func (f *FSM) panicSafeDecide(state *FSMState, event HistoryEvent, data interface{}) (anOutcome *Outcome, anErr error) {
	defer func() {
		if r := recover(); r != nil {
			f.log("at=error error=decide-panic-recovery")
			anErr = errors.New("panic in decider, capture error state")
		}
	}()
	anOutcome = state.Decider(f, event, data)
	return
}

func (f *FSM) captureDecisionError(execution WorkflowExecution, event int, lastEvents []HistoryEvent, stateName string, stateData interface{}, err error) []*Decision {
	return f.captureError(ERROR_SIGNAL, execution, &SerializedDecisionError{
		ErrorEventId:        lastEvents[event].EventId,
		UnprocessedEventIds: f.eventIds(lastEvents[event+1:]),
		StateName:           stateName,
		StateData:           stateData,
	})
}

func (f *FSM) captureSystemError(execution WorkflowExecution, errorType string, lastEvents []HistoryEvent, err error) []*Decision {
	return f.captureError(SYSTEM_ERROR_SIGNAL, execution, &SerializedSystemError{
		ErrorType:           errorType,
		UnprocessedEventIds: f.eventIds(lastEvents),
		Error:               err,
	})
}

func (f *FSM) eventIds(events []HistoryEvent) []int {
	ids := make([]int, len(events))
	for _, e := range events {
		ids = append(ids, e.EventId)
	}
	return ids
}

func (f *FSM) captureError(signal string, execution WorkflowExecution, error interface{}) []*Decision {
	decisions := f.EmptyDecisions()
	r, err := f.DecisionWorker.RecordMarker(signal, error)
	if err != nil {
		//really bail
		panic("giving up, cant even create a RecordMarker decsion")
	}
	d := &Decision{
		DecisionType: DecisionTypeSignalExternalWorkflowExecution,
		SignalExternalWorkflowExecutionDecisionAttributes: &SignalExternalWorkflowExecutionDecisionAttributes{
			WorkflowId: execution.WorkflowId,
			RunId:      execution.RunId,
			SignalName: signal,
			Input:      r.RecordMarkerDecisionAttributes.Details,
		},
	}
	return append(decisions, d)
}

func (f *FSM) EventData(event HistoryEvent) interface{} {
	eventData := f.EventDataType(event)

	if eventData != nil {
		var serialized string
		switch event.EventType {
		case EventTypeActivityTaskCompleted:
			serialized = event.ActivityTaskCompletedEventAttributes.Result
		case EventTypeWorkflowExecutionCompleted:
			serialized = event.WorkflowExecutionCompletedEventAttributes.Result
		case EventTypeChildWorkflowExecutionCompleted:
			serialized = event.ChildWorkflowExecutionCompletedEventAttributes.Result
		case EventTypeWorkflowExecutionSignaled:
			serialized = event.WorkflowExecutionSignaledEventAttributes.Input
		case EventTypeWorkflowExecutionStarted:
			serialized = event.WorkflowExecutionStartedEventAttributes.Input
		case EventTypeWorkflowExecutionContinuedAsNew:
			serialized = event.WorkflowExecutionContinuedAsNewEventAttributes.Input
		}
		if serialized != "" {
			err := f.Serializer().Deserialize(serialized, eventData)
			if err != nil {
				f.log("action=EventData at=error error=unable-to-deserialize")
				panic("Unable to Deserialize Event Data")
			}
		}
	}

	return eventData

}

func (f *FSM) log(format string, data ...interface{}) {
	actualFormat := fmt.Sprintf("component=FSM name=%s %s", f.Name, format)
	log.Printf(actualFormat, data...)
}

func (f *FSM) findSerializedState(events []HistoryEvent) (*SerializedState, error) {
	for _, event := range events {
		if f.isStateMarker(event) {
			state := &SerializedState{}
			err := f.Serializer().Deserialize(event.MarkerRecordedEventAttributes.Details, state)
			return state, err
		} else if event.EventType == EventTypeWorkflowExecutionStarted {
			return &SerializedState{StateName: f.initialState.Name, StateData: event.WorkflowExecutionStartedEventAttributes.Input}, nil
		}
	}
	return nil, errors.New("Cant Find Current Data")
}

func (f *FSM) findLastEvents(prevStarted int, events []HistoryEvent) []HistoryEvent {
	lastEvents := make([]HistoryEvent, 0)
	for _, event := range events {
		if event.EventId == prevStarted {
			return lastEvents
		} else {
			switch event.EventType {
			case EventTypeDecisionTaskCompleted, EventTypeDecisionTaskScheduled,
				EventTypeDecisionTaskStarted, EventTypeDecisionTaskTimedOut:
				//no-op
			case EventTypeMarkerRecorded:
				if !f.isStateMarker(event) {
					lastEvents = append(lastEvents, event)
				}
			default:
				lastEvents = append(lastEvents, event)
			}
		}
	}

	return lastEvents
}

func (f *FSM) appendState(outcome *Outcome) ([]*Decision, error) {

	serializedData, err := f.Serializer().Serialize(outcome.Data)

	state := SerializedState{
		StateName: outcome.NextState,
		StateData: serializedData,
	}

	d, err := f.DecisionWorker.RecordMarker(STATE_MARKER, state)
	if err != nil {
		return nil, err
	}
	decisions := f.EmptyDecisions()
	decisions = append(decisions, d)
	decisions = append(decisions, outcome.Decisions...)
	return decisions, nil
}

func (f *FSM) Stop() {
	f.stop <- true
}

func (f *FSM) isStateMarker(e HistoryEvent) bool {
	return e.EventType == EventTypeMarkerRecorded && e.MarkerRecordedEventAttributes.MarkerName == STATE_MARKER
}

func (f *FSM) Serializer() StateSerializer {
	return f.DecisionWorker.StateSerializer
}

func (f *FSM) EmptyDecisions() []*Decision {
	return make([]*Decision, 0)
}

type SerializedState struct {
	StateName string `json:"stateName"`
	StateData string `json:"stateData"`
}

type SerializedDecisionError struct {
	ErrorEventId        int         `json:"errorEventIds"`
	UnprocessedEventIds []int       `json:"unprocessedEventIds"`
	StateName           string      `json:"stateName"`
	StateData           interface{} `json:"stateData`
}

type SerializedSystemError struct {
	ErrorType           string      `json:"errorType"`
	Error               interface{} `json:"error"`
	UnprocessedEventIds []int       `json:"unprocessedEventIds"`
}

type DecisionErrorPointer struct {
	Error error
}

/*tigertonic-like marhslling of data*/
type MarshalledDecider struct {
	v reflect.Value
}

func TypedDecider(decider interface{}) Decider {
	t := reflect.TypeOf(decider)
	if reflect.Func != t.Kind() {
		panic(fmt.Sprintf("kind was %v, not Func", t.Kind()))
	}
	if 3 != t.NumIn() {
		panic(fmt.Sprintf(
			"input arity was %v, not 3",
			t.NumIn(),
		))
	}
	if "*swf.FSM" != t.In(0).String() {
		panic(fmt.Sprintf(
			"type of first argument was %v, not *swf.FSM",
			t.In(0),
		))
	}
	if "swf.HistoryEvent" != t.In(1).String() {
		panic(fmt.Sprintf(
			"type of second argument was %v, not swf.HistoryEvent",
			t.In(1),
		))
	}

	if "*swf.Outcome" != t.Out(0).String() {
		panic(fmt.Sprintf(
			"type of return value was %v, not *swf.Outcome",
			t.Out(0),
		))
	}

	return MarshalledDecider{reflect.ValueOf(decider)}.Decide
}

func (m MarshalledDecider) Decide(f *FSM, h HistoryEvent, data interface{}) *Outcome {
	return m.v.Call([]reflect.Value{reflect.ValueOf(f), reflect.ValueOf(h), reflect.ValueOf(data)})[0].Interface().(*Outcome)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"strings"
)

const (
	VERSION            = "0.0.1"
	defaultPlaceholder = "{{}}"
)

var originalSttyState bytes.Buffer
var placeholder string

var usage = `fzz allows you to run a command interactively.

Usage:

	fzz command

The command MUST include the placeholder '{{}}'.

Arguments:

	-v		Print version and exit
`

func printUsage() {
	fmt.Printf(usage)
}

func isPipe(f *os.File) bool {
	s, err := f.Stat()
	if err != nil {
		return false
	}

	return s.Mode()&os.ModeNamedPipe != 0
}

func containsPlaceholder(s []string, ph string) bool {
	for _, v := range s {
		if v == ph {
			return true
		}
	}
	return false
}

func stripTrailingNewline(b *bytes.Buffer) {
	s := b.Bytes()
	if s[len(s)-1] == '\n' {
		b.Truncate(b.Len()-1)
	}
}

func main() {
	flVersion := flag.Bool("v", false, "Print fzz version and quit")
	flag.Usage = printUsage
	flag.Parse()

	if *flVersion {
		fmt.Printf("fzz %s\n", VERSION)
		os.Exit(2)
	}

	if len(flag.Args()) < 2 {
		fmt.Fprintf(os.Stderr, usage)
		os.Exit(2)
	}

	if placeholder = os.Getenv("FZZ_PLACEHOLDER"); placeholder == "" {
		placeholder = defaultPlaceholder
	}

	if !containsPlaceholder(flag.Args(), placeholder) {
		fmt.Fprintf(os.Stderr, "No placeholder in arguments\n")
		os.Exit(2)
	}

	tty, err := NewTTY()
	if err != nil {
		log.Fatal(err)
	}

	err = tty.getSttyState(&originalSttyState)
	if err != nil {
		log.Fatal(err)
	}
	stripTrailingNewline(&originalSttyState)
	defer tty.setSttyState(&originalSttyState)

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		tty.setSttyState(&originalSttyState)
		os.Exit(1)
	}()

	tty.setSttyState(bytes.NewBufferString("cbreak"))
	tty.setSttyState(bytes.NewBufferString("-echo"))

	cmdTemplate := strings.Join(flag.Args(), " ")
	printer := NewPrinter(tty, tty.cols, tty.rows-3)
	runner := &Runner{
		printer:     printer,
		template:    cmdTemplate,
		placeholder: placeholder,
	}

	if isPipe(os.Stdin) {
		runner.stdinbuf = new(bytes.Buffer)
		io.Copy(runner.stdinbuf, os.Stdin)
	}

	input := make([]byte, 0)
	b := make([]byte, 1)

	for {
		tty.resetScreen()
		tty.printPrompt(input[:len(input)])

		if len(input) > 0 {
			runner.killCurrent()

			go func() {
				runner.runWithInput(input[:len(input)])
				tty.cursorAfterPrompt(len(input))
			}()
		}

		tty.Read(b)
		switch b[0] {
		case 8, 127:
			// Backspace, delete
			if len(input) > 1 {
				input = input[:len(input)-1]
			} else if len(input) == 1 {
				input = nil
			}
		case 4, 10, 13:
			// Ctrl-D, line feed, carriage return
			tty.resetScreen()
			runner.writeCmdStdout(os.Stdout)
			return
		default:
			// TODO: Default is wrong here. Only append printable characters to
			// input
			input = append(input, b...)
		}
	}
}
<commit_msg>Fix containsPlaceholder<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"strings"
)

const (
	VERSION            = "0.0.1"
	defaultPlaceholder = "{{}}"
)

var originalSttyState bytes.Buffer
var placeholder string

var usage = `fzz allows you to run a command interactively.

Usage:

	fzz command

The command MUST include the placeholder '{{}}'.

Arguments:

	-v		Print version and exit
`

func printUsage() {
	fmt.Printf(usage)
}

func isPipe(f *os.File) bool {
	s, err := f.Stat()
	if err != nil {
		return false
	}

	return s.Mode()&os.ModeNamedPipe != 0
}

func containsPlaceholder(s []string, ph string) bool {
	for _, v := range s {
		if strings.Contains(v, ph) {
			return true
		}
	}
	return false
}

func stripTrailingNewline(b *bytes.Buffer) {
	s := b.Bytes()
	if s[len(s)-1] == '\n' {
		b.Truncate(b.Len()-1)
	}
}

func main() {
	flVersion := flag.Bool("v", false, "Print fzz version and quit")
	flag.Usage = printUsage
	flag.Parse()

	if *flVersion {
		fmt.Printf("fzz %s\n", VERSION)
		os.Exit(2)
	}

	if len(flag.Args()) < 2 {
		fmt.Fprintf(os.Stderr, usage)
		os.Exit(2)
	}

	if placeholder = os.Getenv("FZZ_PLACEHOLDER"); placeholder == "" {
		placeholder = defaultPlaceholder
	}

	if !containsPlaceholder(flag.Args(), placeholder) {
		fmt.Fprintf(os.Stderr, "No placeholder in arguments\n")
		os.Exit(2)
	}

	tty, err := NewTTY()
	if err != nil {
		log.Fatal(err)
	}

	err = tty.getSttyState(&originalSttyState)
	if err != nil {
		log.Fatal(err)
	}
	stripTrailingNewline(&originalSttyState)
	defer tty.setSttyState(&originalSttyState)

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		tty.setSttyState(&originalSttyState)
		os.Exit(1)
	}()

	tty.setSttyState(bytes.NewBufferString("cbreak"))
	tty.setSttyState(bytes.NewBufferString("-echo"))

	cmdTemplate := strings.Join(flag.Args(), " ")
	printer := NewPrinter(tty, tty.cols, tty.rows-3)
	runner := &Runner{
		printer:     printer,
		template:    cmdTemplate,
		placeholder: placeholder,
	}

	if isPipe(os.Stdin) {
		runner.stdinbuf = new(bytes.Buffer)
		io.Copy(runner.stdinbuf, os.Stdin)
	}

	input := make([]byte, 0)
	b := make([]byte, 1)

	for {
		tty.resetScreen()
		tty.printPrompt(input[:len(input)])

		if len(input) > 0 {
			runner.killCurrent()

			go func() {
				runner.runWithInput(input[:len(input)])
				tty.cursorAfterPrompt(len(input))
			}()
		}

		tty.Read(b)
		switch b[0] {
		case 8, 127:
			// Backspace, delete
			if len(input) > 1 {
				input = input[:len(input)-1]
			} else if len(input) == 1 {
				input = nil
			}
		case 4, 10, 13:
			// Ctrl-D, line feed, carriage return
			tty.resetScreen()
			runner.writeCmdStdout(os.Stdout)
			return
		default:
			// TODO: Default is wrong here. Only append printable characters to
			// input
			input = append(input, b...)
		}
	}
}
<|endoftext|>


<commit_before>package ds

import (
	"context"
	"reflect"

	"cloud.google.com/go/datastore"
)

// GetByKey retrieves model from datastore by key
func (client *Client) GetByKey(ctx context.Context, key *datastore.Key, dst interface{}) error {
	if client.Cache != nil && client.Cache.Get(key, dst) == nil {
		return nil
	}
	err := client.Get(ctx, key, dst)
	SetKey(key, dst)
	if client.Cache != nil {
		client.Cache.Set(key, dst)
	}
	if err != nil {
		return err
	}
	return nil
}

// GetByKeys retrieves models from datastore by keys
func (client *Client) GetByKeys(ctx context.Context, keys []*datastore.Key, dst interface{}) error {
	// prepare slice if dst is pointer to 0 len slice
	if rf := reflect.ValueOf(dst); rf.Kind() == reflect.Ptr {
		rs := rf.Elem()
		if rs.Kind() == reflect.Slice && rs.Len() == 0 {
			l := len(keys)
			rs.Set(reflect.MakeSlice(rs.Type(), l, l))
		}
		dst = rs.Interface()
	}

	if client.Cache != nil {
		err := client.Cache.GetMulti(keys, dst)
		if err == nil {
			nfKeys := []*datastore.Key{}
			nfMap := []int{}
			rf := valueOf(dst)
			for i := 0; i < rf.Len(); i++ {
				if rf.Index(i).IsNil() {
					nfKeys = append(nfKeys, keys[i])
					nfMap = append(nfMap, i)
				}
			}
			l := len(nfKeys)
			nfDstRf := reflect.MakeSlice(rf.Type(), l, l)
			err := client.GetMulti(ctx, keys, nfDstRf.Interface())
			for i, k := range nfMap {
				rf.Index(k).Set(nfDstRf.Index(i))
			}
			SetKeys(keys, dst)
			if err != nil {
				return err
			}
			return nil
		}
	}

	err := client.GetMulti(ctx, keys, dst)
	SetKeys(keys, dst)
	if client.Cache != nil {
		client.Cache.SetMulti(keys, dst)
	}
	if err != nil {
		return err
	}
	return nil
}

// GetByModel retrieves model from datastore by key from model
func (client *Client) GetByModel(ctx context.Context, dst interface{}) error {
	key := ExtractKey(dst)
	return client.GetByKey(ctx, key, dst)
}

// GetByModels retrieves models from datastore by keys from models
func (client *Client) GetByModels(ctx context.Context, dst interface{}) error {
	keys := ExtractKeys(dst)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByID retrieves model from datastore by id
func (client *Client) GetByID(ctx context.Context, kind string, id int64, dst interface{}) error {
	return client.GetByKey(ctx, datastore.IDKey(kind, id, nil), dst)
}

// GetByIDs retrieves models from datastore by ids
func (client *Client) GetByIDs(ctx context.Context, kind string, ids []int64, dst interface{}) error {
	keys := BuildIDKeys(kind, ids)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByStringID retrieves model from datastore by string id
func (client *Client) GetByStringID(ctx context.Context, kind string, id string, dst interface{}) error {
	tid := parseID(id)
	if tid == 0 {
		return datastore.ErrInvalidKey
	}
	return client.GetByKey(ctx, datastore.IDKey(kind, tid, nil), dst)
}

// GetByStringIDs retrieves models from datastore by string ids
func (client *Client) GetByStringIDs(ctx context.Context, kind string, ids []string, dst interface{}) error {
	keys := BuildStringIDKeys(kind, ids)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByName retrieves model from datastore by name
func (client *Client) GetByName(ctx context.Context, kind string, name string, dst interface{}) error {
	return client.GetByKey(ctx, datastore.NameKey(kind, name, nil), dst)
}

// GetByNames retrieves models from datastore by names
func (client *Client) GetByNames(ctx context.Context, kind string, names []string, dst interface{}) error {
	keys := BuildNameKeys(kind, names)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByQuery retrieves model from datastore by datastore query
func (client *Client) GetByQuery(ctx context.Context, q *datastore.Query, dst interface{}) error {
	_, err := client.GetAll(ctx, q, dst)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>support get more than 1000 entity<commit_after>package ds

import (
	"context"
	"reflect"

	"cloud.google.com/go/datastore"
)

// GetByKey retrieves model from datastore by key
func (client *Client) GetByKey(ctx context.Context, key *datastore.Key, dst interface{}) error {
	if client.Cache != nil && client.Cache.Get(key, dst) == nil {
		return nil
	}
	err := client.Get(ctx, key, dst)
	SetKey(key, dst)
	if client.Cache != nil {
		client.Cache.Set(key, dst)
	}
	if err != nil {
		return err
	}
	return nil
}

// GetByKeys retrieves models from datastore by keys
func (client *Client) GetByKeys(ctx context.Context, keys []*datastore.Key, dst interface{}) error {
	// prepare slice if dst is pointer to 0 len slice
	if rf := reflect.ValueOf(dst); rf.Kind() == reflect.Ptr {
		rs := rf.Elem()
		if rs.Kind() == reflect.Slice && rs.Len() == 0 {
			l := len(keys)
			rs.Set(reflect.MakeSlice(rs.Type(), l, l))
		}
		dst = rs.Interface()
	}

	if client.Cache != nil {
		err := client.Cache.GetMulti(keys, dst)
		if err == nil {
			nfKeys := []*datastore.Key{}
			nfMap := []int{}
			rf := valueOf(dst)
			for i := 0; i < rf.Len(); i++ {
				if rf.Index(i).IsNil() {
					nfKeys = append(nfKeys, keys[i])
					nfMap = append(nfMap, i)
				}
			}
			l := len(nfKeys)
			nfDstRf := reflect.MakeSlice(rf.Type(), l, l)
			err := client.GetMulti(ctx, keys, nfDstRf.Interface())
			for i, k := range nfMap {
				rf.Index(k).Set(nfDstRf.Index(i))
			}
			SetKeys(keys, dst)
			if err != nil {
				return err
			}
			return nil
		}
	}

	var err error
	l := len(keys)
	p := 1000
	if l > p {
		rfDst := valueOf(dst)
		for i := 0; i < l/p+1; i++ {
			m := (i + 1) * p
			if l-m+1 < p {
				m = l
			}
			if i*p == m {
				break
			}
			e := client.GetMulti(ctx, keys[i*p:m], rfDst.Slice(i*p, m).Interface())
			if e != nil {
				if err == nil {
					err = e
				} else {
					if errs, ok := err.(datastore.MultiError); ok {
						err = append(errs, e)
					} else {
						err = datastore.MultiError{err, e}
					}
				}
			}
		}
	} else {
		err = client.GetMulti(ctx, keys, dst)
	}
	SetKeys(keys, dst)
	if client.Cache != nil {
		client.Cache.SetMulti(keys, dst)
	}
	if err != nil {
		return err
	}
	return nil
}

// GetByModel retrieves model from datastore by key from model
func (client *Client) GetByModel(ctx context.Context, dst interface{}) error {
	key := ExtractKey(dst)
	return client.GetByKey(ctx, key, dst)
}

// GetByModels retrieves models from datastore by keys from models
func (client *Client) GetByModels(ctx context.Context, dst interface{}) error {
	keys := ExtractKeys(dst)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByID retrieves model from datastore by id
func (client *Client) GetByID(ctx context.Context, kind string, id int64, dst interface{}) error {
	return client.GetByKey(ctx, datastore.IDKey(kind, id, nil), dst)
}

// GetByIDs retrieves models from datastore by ids
func (client *Client) GetByIDs(ctx context.Context, kind string, ids []int64, dst interface{}) error {
	keys := BuildIDKeys(kind, ids)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByStringID retrieves model from datastore by string id
func (client *Client) GetByStringID(ctx context.Context, kind string, id string, dst interface{}) error {
	tid := parseID(id)
	if tid == 0 {
		return datastore.ErrInvalidKey
	}
	return client.GetByKey(ctx, datastore.IDKey(kind, tid, nil), dst)
}

// GetByStringIDs retrieves models from datastore by string ids
func (client *Client) GetByStringIDs(ctx context.Context, kind string, ids []string, dst interface{}) error {
	keys := BuildStringIDKeys(kind, ids)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByName retrieves model from datastore by name
func (client *Client) GetByName(ctx context.Context, kind string, name string, dst interface{}) error {
	return client.GetByKey(ctx, datastore.NameKey(kind, name, nil), dst)
}

// GetByNames retrieves models from datastore by names
func (client *Client) GetByNames(ctx context.Context, kind string, names []string, dst interface{}) error {
	keys := BuildNameKeys(kind, names)
	return client.GetByKeys(ctx, keys, dst)
}

// GetByQuery retrieves model from datastore by datastore query
func (client *Client) GetByQuery(ctx context.Context, q *datastore.Query, dst interface{}) error {
	_, err := client.GetAll(ctx, q, dst)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>