<commit_before>package knx

import (
	"errors"
	"time"
	"context"
)

//
type Client struct {
	sock    *Socket
	reaper  context.CancelFunc

	//
	Inbound <-chan []byte
}

//
var (
	ErrConnTimeout = errors.New("Gateway did not response to connection request")
	ErrConnRejected = errors.New("Gateway rejected connection request")
)

//
func NewClient(gatewayAddress string) (*Client, error) {
	sock, err := NewClientSocket(gatewayAddress)
	if err != nil {
		return nil, err
	}

	req := &ConnectionRequest{}

	err = sock.Send(req)
	if err != nil {
		return nil, err
	}

	resChan := awaitConnectionResponse(sock)

	// Connection cycle
	for i := 0; i < 5; i++ {
		select {
		case res := <-resChan:
			if res.Status == 0 {
				// Connection is established.

				Logger.Printf("Client[%v]: Connection has been established on channel %v",
				              sock.conn.RemoteAddr(), res.Channel)

				return makeClient(sock, res.Channel), nil
			} else {
				// Connection attempt was rejected.

				Logger.Printf("Client[%v]: Connection attempt was rejected",
				              sock.conn.RemoteAddr())

				sock.Close()
				return nil, ErrConnRejected
			}

		case <-time.After(time.Second):
			// Resend the connection request, if we haven't received a response from the gateway
			// after 1 second.
			err := sock.Send(req)
			if err != nil {
				return nil, err
			}
		}
	}

	Logger.Printf("Client[%v]: Connection attempts timed out", sock.conn.RemoteAddr())

	sock.Close()
	return nil, ErrConnTimeout
}

//
func (client *Client) Close() {
	client.reaper()
	client.sock.Close()
}

//
func awaitConnectionResponse(sock *Socket) <-chan *ConnectionResponse {
	resChan := make(chan *ConnectionResponse)

	go func() {
		for payload := range sock.Inbound {
			res, ok := payload.(*ConnectionResponse)
			if ok {
				resChan <- res
				return
			}
		}
	}()

	return resChan
}

//
func makeClient(sock *Socket, channel byte) *Client {
	ctx, reaper := context.WithCancel(context.Background())

	inbound := make(chan []byte)
	go clientInboundWorker(ctx, reaper, sock, channel, inbound)

	return &Client{sock, reaper, inbound}
}

//
func clientInboundWorker(
	ctx     context.Context,
	reaper  context.CancelFunc,
	sock    *Socket,
	channel byte,
	inbound chan<- []byte,
) {
	Logger.Printf("Client[%v]: Started inbound worker", sock.conn.RemoteAddr())
	defer Logger.Printf("Client[%v]: Stopped inbound worker", sock.conn.RemoteAddr())

	defer close(inbound)

	heartbeatTrigger := make(chan struct{})
	stateResChan := make(chan *ConnectionStateResponse)
	go clientHeartbeatWorker(ctx, reaper, sock, channel, heartbeatTrigger, stateResChan)

	for {
		select {
		// Goroutine exit has been requested
		case <-ctx.Done():
			return

		// 10 seconds without communication, time for a heartbeat
		case <-time.After(10 * time.Second):
			Logger.Printf("Client[%v]: Triggering heartbeat", sock.conn.RemoteAddr())

			select {
			case <-ctx.Done():
				return

			case heartbeatTrigger <- struct{}{}:
			}

		// Incoming packets
		case payload, open := <-sock.Inbound:
			// If the socket inbound channel is closed, this goroutine has no purpose.
			if !open {
				Logger.Printf("Client[%v]: Inbound channel has been closed", sock.conn.RemoteAddr())
				reaper()
				return
			}

			switch payload.(type) {
			case *ConnectionStateResponse:
				res := payload.(*ConnectionStateResponse)

				if res.Channel == channel {
					stateResChan <- res
				}
			}
		}
	}
}

//
func clientHeartbeatWorker(
	ctx     context.Context,
	reaper  context.CancelFunc,
	sock    *Socket,
	channel byte,
	trigger <-chan struct{},
	resChan <-chan *ConnectionStateResponse,
) {
	Logger.Printf("Client[%v]: Started heartbeat worker", sock.conn.RemoteAddr())
	defer Logger.Printf("Client[%v]: Stopped heartbeat worker", sock.conn.RemoteAddr())

	// Make sure we tell the others to exit
	defer reaper()

	outerLoop:
	for {
		select {
		// Gorouting has been asked to exit
		case <-ctx.Done():
			return

		// Inbound worker has triggered a heartbeat
		case <-trigger:
			req := &ConnectionStateRequest{channel, 0, HostInfo{}}

			err := sock.Send(req)
			if err != nil {
				Logger.Printf("Client[%v]: Error while sending heartbeat: %v",
				              sock.conn.RemoteAddr(), err)
				return
			}

			// Heartbeat cycle
			for i := 0; i < 5; i++ {
				select {
				case <-ctx.Done():
					return

				case res := <-resChan:
					if res.Status == 0 {
						Logger.Printf("Client[%v]: Heartbeat successful", sock.conn.RemoteAddr())
						continue outerLoop
					} else {
						Logger.Printf("Client[%v]: Gateway rejected heartbeat",
						              sock.conn.RemoteAddr())
						return
					}

				case <-time.After(time.Second):
					err := sock.Send(req)
					if err != nil {
						Logger.Printf("Client[%v]: Error while sending heartbeat: %v",
						              sock.conn.RemoteAddr(), err)
						return
					}
				}
			}

			// We get here, if the gateway did not respond

			Logger.Printf("Client[%v]: Gateway timed out during heartbeat", sock.conn.RemoteAddr())
			return

		case <-resChan:
			// Discard any connection state response that appears out-of-cycle
		}
	}
}
<commit_msg>Add send and receive functionality to Client<commit_after>package knx

import (
	"errors"
	"time"
	"context"
)

//
type Client struct {
	sock    *Socket
	channel byte
	reaper  context.CancelFunc
	seq     <-chan byte
	ack     <-chan byte

	//
	Inbound <-chan []byte
}

//
var (
	ErrConnTimeout = errors.New("Gateway did not response to connection request")
	ErrConnRejected = errors.New("Gateway rejected connection request")
)

//
func NewClient(gatewayAddress string) (*Client, error) {
	sock, err := NewClientSocket(gatewayAddress)
	if err != nil {
		return nil, err
	}

	req := &ConnectionRequest{}

	err = sock.Send(req)
	if err != nil {
		return nil, err
	}

	resChan := awaitConnectionResponse(sock)

	// Connection cycle
	for i := 0; i < 5; i++ {
		select {
		case res := <-resChan:
			if res.Status == 0 {
				// Connection is established.

				Logger.Printf("Client[%v]: Connection has been established on channel %v",
				              sock.conn.RemoteAddr(), res.Channel)

				return makeClient(sock, res.Channel), nil
			} else {
				// Connection attempt was rejected.

				Logger.Printf("Client[%v]: Connection attempt was rejected",
				              sock.conn.RemoteAddr())

				sock.Close()
				return nil, ErrConnRejected
			}

		case <-time.After(time.Second):
			// Resend the connection request, if we haven't received a response from the gateway
			// after 1 second.
			err := sock.Send(req)
			if err != nil {
				return nil, err
			}
		}
	}

	Logger.Printf("Client[%v]: Connection attempts timed out", sock.conn.RemoteAddr())

	sock.Close()
	return nil, ErrConnTimeout
}

//
func (client *Client) Close() {
	client.reaper()
	client.sock.Close()
}

var (
	ErrSendClosed   = errors.New("Outbound worker has terminated")
	ErrSendRejected = errors.New("Gateway rejected tunnel request")
	ErrSendTimeout  = errors.New("Gateway did not acknowledge tunnel request in time")
)

//
func (client *Client) Send(data []byte) error {
	seqNumber, open := <-client.seq
	if !open {
		return ErrSendClosed
	}

	req := &TunnelRequest{client.channel, seqNumber, data}
	err := client.sock.Send(req)
	if err != nil {
		return err
	}

	for i := 0; i < 5; i++ {
		select {
		case status, open := <-client.ack:
			if !open {
				return ErrSendClosed
			}

			if status == 0 {
				return nil
			} else {
				return ErrSendRejected
			}

		case <-time.After(time.Second):
			err := client.sock.Send(req)
			if err != nil {
				return err
			}
		}
	}

	return ErrSendTimeout
}

//
func awaitConnectionResponse(sock *Socket) <-chan *ConnectionResponse {
	resChan := make(chan *ConnectionResponse)

	go func() {
		for payload := range sock.Inbound {
			res, ok := payload.(*ConnectionResponse)
			if ok {
				resChan <- res
				return
			}
		}
	}()

	return resChan
}

//
func makeClient(sock *Socket, channel byte) *Client {
	ctx, reaper := context.WithCancel(context.Background())

	seq := make(chan byte)
	ack := make(chan byte)
	tunRes := make(chan *TunnelResponse)
	go clientOutboundWorker(ctx, reaper, tunRes, seq, ack)

	inbound := make(chan []byte)
	go clientInboundWorker(ctx, reaper, sock, channel, tunRes, inbound)

	return &Client{sock, channel, reaper, seq, ack, inbound}
}

//
func clientInboundWorker(
	ctx     context.Context,
	reaper  context.CancelFunc,
	sock    *Socket,
	channel byte,
	tunRes  chan<- *TunnelResponse,
	inbound chan<- []byte,
) {
	Logger.Printf("Client[%v]: Started inbound worker", sock.conn.RemoteAddr())
	defer Logger.Printf("Client[%v]: Stopped inbound worker", sock.conn.RemoteAddr())

	defer close(inbound)
	defer reaper()

	heartbeatTrigger := make(chan struct{})
	stateResChan := make(chan *ConnectionStateResponse)
	go clientHeartbeatWorker(ctx, reaper, sock, channel, heartbeatTrigger, stateResChan)

	var incomingSeqNumber byte = 0

	for {
		select {
		// Goroutine exit has been requested
		case <-ctx.Done():
			return

		// 10 seconds without communication, time for a heartbeat
		case <-time.After(10 * time.Second):
			Logger.Printf("Client[%v]: Triggering heartbeat", sock.conn.RemoteAddr())

			select {
			case <-ctx.Done():
				return

			case heartbeatTrigger <- struct{}{}:
			}

		// Incoming packets
		case payload, open := <-sock.Inbound:
			// If the socket inbound channel is closed, this goroutine has no purpose.
			if !open {
				Logger.Printf("Client[%v]: Inbound channel has been closed", sock.conn.RemoteAddr())
				return
			}

			switch payload.(type) {
			case *ConnectionStateResponse:
				res := payload.(*ConnectionStateResponse)

				if res.Channel == channel {
					select {
					case <-ctx.Done():
						return

					case stateResChan <- res:
					}
				}

			case *TunnelRequest:
				req := payload.(*TunnelRequest)

				if req.Channel == channel {
					// Acknowledge tunnel request
					err := sock.Send(&TunnelResponse{channel, req.SeqNumber, 0})
					if err != nil {
						Logger.Printf("Client[%v]: Error while sending tunnel response: %v",
						              sock.conn.RemoteAddr(), err)
						return
					}

					// Relay to user if it fits the sequence
					if req.SeqNumber == incomingSeqNumber {
						Logger.Printf("Client[%v]: Inbound tunnel request: %v",
						              sock.conn.RemoteAddr(), req.Payload)

						select {
						case <-ctx.Done():
							return

						case inbound <- req.Payload:
						}

						incomingSeqNumber++
					}
				}

			case *TunnelResponse:
				res := payload.(*TunnelResponse)

				if res.Channel == channel {
					select {
					case <-ctx.Done():
						return

					case tunRes <- res:
					}
				}
			}
		}
	}
}

//
func clientOutboundWorker(
	ctx     context.Context,
	reaper  context.CancelFunc,
	tunRes  <-chan *TunnelResponse,
	seq     chan<- byte,
	ack     chan<- byte,
) {
	defer reaper()
	defer close(seq)
	defer close(ack)

	var seqNumber byte = 0

	outerLoop:
	for {
		select {
		case <-ctx.Done():
			return

		// Client requests a sequence number because it wants to send something
		case seq <- seqNumber:

			for {
				select {
				case <-ctx.Done():
					return

				// Await tunnel response
				case res := <-tunRes:
					// We're only interested in the ones that match our sequence number
					if res.SeqNumber == seqNumber {
						select {
						case <-ctx.Done():
							return

						// Send result of the tunnel request to the sender
						case ack <- res.Status:
							seqNumber++
							continue outerLoop
						}
					}
				}
			}

		// Discard out-of-cycle tunnel responses
		case <-tunRes:
		}
	}
}

//
func clientHeartbeatWorker(
	ctx     context.Context,
	reaper  context.CancelFunc,
	sock    *Socket,
	channel byte,
	trigger <-chan struct{},
	resChan <-chan *ConnectionStateResponse, // Really? Why not just 'byte'?
) {
	Logger.Printf("Client[%v]: Started heartbeat worker", sock.conn.RemoteAddr())
	defer Logger.Printf("Client[%v]: Stopped heartbeat worker", sock.conn.RemoteAddr())

	// Make sure we tell the others to exit
	defer reaper()

	outerLoop:
	for {
		select {
		// Gorouting has been asked to exit
		case <-ctx.Done():
			return

		// Inbound worker has triggered a heartbeat
		case <-trigger:
			req := &ConnectionStateRequest{channel, 0, HostInfo{}}

			err := sock.Send(req)
			if err != nil {
				Logger.Printf("Client[%v]: Error while sending heartbeat: %v",
				              sock.conn.RemoteAddr(), err)
				return
			}

			// Heartbeat cycle
			for i := 0; i < 5; i++ {
				select {
				case <-ctx.Done():
					return

				case res := <-resChan:
					if res.Status == 0 {
						Logger.Printf("Client[%v]: Heartbeat successful", sock.conn.RemoteAddr())
						continue outerLoop
					} else {
						Logger.Printf("Client[%v]: Gateway rejected heartbeat",
						              sock.conn.RemoteAddr())
						return
					}

				case <-time.After(time.Second):
					err := sock.Send(req)
					if err != nil {
						Logger.Printf("Client[%v]: Error while sending heartbeat: %v",
						              sock.conn.RemoteAddr(), err)
						return
					}
				}
			}

			// We get here, if the gateway did not respond

			Logger.Printf("Client[%v]: Gateway timed out during heartbeat", sock.conn.RemoteAddr())
			return

		case <-resChan:
			// Discard any connection state response that appears out-of-cycle
		}
	}
}
<|endoftext|>


<commit_before>package kvite

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"
)

func withDB(t *testing.T, fn func(db *DB, t *testing.T)) {
	file := tempfile()
	db, err := Open(file, "testing")
	ok(t, err)
	defer os.Remove(file)
	defer db.Close()
	fn(db, t)
}

func TestOpen(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
	})
}

func TestBegin(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		if err != nil {
			t.Fatal(err)
		}
		defer tx.Rollback()
	})
}

func TestRollback(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		err = tx.Rollback()
		ok(t, err)
	})
}

func TestCommit(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestCreateBucket(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer tx.Rollback()
		_, err = tx.CreateBucket("test")
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestCreateBucketIfNotExists(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer tx.Rollback()
		_, err = tx.CreateBucketIfNotExists("test")
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestPut(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer tx.Rollback()
		b, err := tx.CreateBucket("test")
		ok(t, err)

		err = b.Put("foo", []byte("bar"))
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestGet(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer tx.Rollback()
		b, err := tx.CreateBucket("test")
		ok(t, err)

		err = b.Put("foo", []byte("bar"))
		ok(t, err)

		val, err := b.Get("foo")
		ok(t, err)

		equals(t, string(val), "bar")

		err = tx.Commit()
		ok(t, err)
	})
}

func TestDelete(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer tx.Rollback()
		b, err := tx.CreateBucket("test")
		ok(t, err)

		err = b.Put("foo", []byte("bar"))
		ok(t, err)

		val, err := b.Get("foo")
		ok(t, err)

		equals(t, string(val), "bar")

		err = b.Delete("foo")
		ok(t, err)

		val, err = b.Get("foo")
		equals(t, []byte(nil), val)
		ok(t, err)

		err = tx.Commit()
		ok(t, err)
	})
}

func TestTransaction(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			b, err := tx.CreateBucket("test")
			ok(t, err)

			err = b.Put("foo", []byte("bar"))
			ok(t, err)

			val, err := b.Get("foo")
			ok(t, err)

			equals(t, string(val), "bar")

			return nil
		})

		if err != nil {
			t.Fatal(err)
		}
	})
}

func TestForEach(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			b, err := tx.CreateBucket("test")
			ok(t, err)

			err = b.Put("foo", []byte("bar"))
			ok(t, err)

			err = b.Put("baz", []byte("stuff"))
			ok(t, err)

			var items []string
			err = b.ForEach(func(k string, v []byte) error {
				items = append(items, k)
				return nil
			})
			ok(t, err)

			if len(items) != 2 {
				return fmt.Errorf("length does not match")
			}

			return nil
		})

		ok(t, err)
	})
}

func TestBuckets(t *testing.T) {
	buckets := []string{"one", "two", "three"}
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			for _, name := range buckets {
				b, err := tx.CreateBucket(name)
				ok(t, err)
				err = b.Put("foo", []byte("bar"))
				ok(t, err)
			}
			return nil
		})

		names, err := db.Buckets()
		ok(t, err)
		equals(t, buckets, names)
	})
}

func TestUnique(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			b, err := tx.CreateBucket("test")
			ok(t, err)

			err = b.Put("foo", []byte("bar"))
			ok(t, err)
			err = b.Put("foo", []byte("baz"))

			err = b.ForEach(func(k string, v []byte) error {
				equals(t, "baz", string(v))
				return nil
			})
			ok(t, err)
			return nil
		})
		ok(t, err)
	})
}

func BenchmarkPutGet(bm *testing.B) {
	file := tempfile()
	db, err := Open(file, "testing")

	if err != nil {
		bm.Fatal(err)
	}
	defer os.Remove(file)
	defer db.Close()

	err = db.Transaction(func(tx *Tx) error {
		b, err := tx.CreateBucket("test")
		if err != nil {
			return err
		}

		for n := 0; n < bm.N; n++ {
			err = b.Put("foo", []byte("bar"))
			if err != nil {
				return err
			}
			_, err := b.Get("foo")
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		bm.Fatal(err)
	}
}

// tempfile returns a temporary file path.
func tempfile() string {
	f, _ := ioutil.TempFile("", "kvite-")
	f.Close()
	os.Remove(f.Name())
	return f.Name()
}

// Thanks to https://github.com/benbjohnson/testing

// assert fails the test if the condition is false.
func assert(tb testing.TB, condition bool, msg string, v ...interface{}) {
	if !condition {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: "+msg+"\033[39m\n\n", append([]interface{}{filepath.Base(file), line}, v...)...)
		tb.FailNow()
	}
}

// ok fails the test if an err is not nil.
func ok(tb testing.TB, err error) {
	if err != nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: unexpected error: %s\033[39m\n\n", filepath.Base(file), line, err.Error())
		tb.FailNow()
	}
}

// equals fails the test if exp is not equal to act.
func equals(tb testing.TB, exp, act interface{}) {
	if !reflect.DeepEqual(exp, act) {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n", filepath.Base(file), line, exp, act)
		tb.FailNow()
	}
}
<commit_msg>Log deferred errors instead of ignoring<commit_after>package kvite

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"
)

func withDB(t *testing.T, fn func(db *DB, t *testing.T)) {
	file := tempfile()
	db, err := Open(file, "testing")
	ok(t, err)
	defer removeFileAndLogError(file)
	defer logErr(db.Close, "database close")
	fn(db, t)
}

func TestOpen(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
	})
}

func TestBegin(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		if err != nil {
			t.Fatal(err)
		}
		defer logErr(tx.Rollback, "Transaction Rollback")
	})
}

func TestRollback(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		err = tx.Rollback()
		ok(t, err)
	})
}

func TestCommit(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestCreateBucket(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer logErr(tx.Rollback, "Transaction Rollback")
		_, err = tx.CreateBucket("test")
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestCreateBucketIfNotExists(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer logErr(tx.Rollback, "Transaction Rollback")
		_, err = tx.CreateBucketIfNotExists("test")
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestPut(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer logErr(tx.Rollback, "Transaction Rollback")
		b, err := tx.CreateBucket("test")
		ok(t, err)

		err = b.Put("foo", []byte("bar"))
		ok(t, err)
		err = tx.Commit()
		ok(t, err)
	})
}

func TestGet(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer logErr(tx.Rollback, "Transaction Rollback")
		b, err := tx.CreateBucket("test")
		ok(t, err)

		err = b.Put("foo", []byte("bar"))
		ok(t, err)

		val, err := b.Get("foo")
		ok(t, err)

		equals(t, string(val), "bar")

		err = tx.Commit()
		ok(t, err)
	})
}

func TestDelete(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		tx, err := db.Begin()
		ok(t, err)
		defer logErr(tx.Rollback, "Transaction Rollback")
		b, err := tx.CreateBucket("test")
		ok(t, err)

		err = b.Put("foo", []byte("bar"))
		ok(t, err)

		val, err := b.Get("foo")
		ok(t, err)

		equals(t, string(val), "bar")

		err = b.Delete("foo")
		ok(t, err)

		val, err = b.Get("foo")
		equals(t, []byte(nil), val)
		ok(t, err)

		err = tx.Commit()
		ok(t, err)
	})
}

func TestTransaction(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			b, err := tx.CreateBucket("test")
			ok(t, err)

			err = b.Put("foo", []byte("bar"))
			ok(t, err)

			val, err := b.Get("foo")
			ok(t, err)

			equals(t, string(val), "bar")

			return nil
		})

		if err != nil {
			t.Fatal(err)
		}
	})
}

func TestForEach(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			b, err := tx.CreateBucket("test")
			ok(t, err)

			err = b.Put("foo", []byte("bar"))
			ok(t, err)

			err = b.Put("baz", []byte("stuff"))
			ok(t, err)

			var items []string
			err = b.ForEach(func(k string, v []byte) error {
				items = append(items, k)
				return nil
			})
			ok(t, err)

			if len(items) != 2 {
				return fmt.Errorf("length does not match")
			}

			return nil
		})

		ok(t, err)
	})
}

func TestBuckets(t *testing.T) {
	buckets := []string{"one", "two", "three"}
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			for _, name := range buckets {
				b, err := tx.CreateBucket(name)
				ok(t, err)
				err = b.Put("foo", []byte("bar"))
				ok(t, err)
			}
			return nil
		})

		names, err := db.Buckets()
		ok(t, err)
		equals(t, buckets, names)
	})
}

func TestUnique(t *testing.T) {
	withDB(t, func(db *DB, t *testing.T) {
		err := db.Transaction(func(tx *Tx) error {
			b, err := tx.CreateBucket("test")
			ok(t, err)

			err = b.Put("foo", []byte("bar"))
			ok(t, err)
			err = b.Put("foo", []byte("baz"))

			err = b.ForEach(func(k string, v []byte) error {
				equals(t, "baz", string(v))
				return nil
			})
			ok(t, err)
			return nil
		})
		ok(t, err)
	})
}

func BenchmarkPutGet(bm *testing.B) {
	file := tempfile()
	db, err := Open(file, "testing")

	if err != nil {
		bm.Fatal(err)
	}
	defer removeFileAndLogError(file)
	defer logErr(db.Close, "database close")

	err = db.Transaction(func(tx *Tx) error {
		b, err := tx.CreateBucket("test")
		if err != nil {
			return err
		}

		for n := 0; n < bm.N; n++ {
			err = b.Put("foo", []byte("bar"))
			if err != nil {
				return err
			}
			_, err := b.Get("foo")
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		bm.Fatal(err)
	}
}

// tempfile returns a temporary file path.
func tempfile() string {
	f, _ := ioutil.TempFile("", "kvite-")
	logErr(f.Close, "temp file close")
	removeFileAndLogError(f.Name())
	return f.Name()
}

// Thanks to https://github.com/benbjohnson/testing

// assert fails the test if the condition is false.
func assert(tb testing.TB, condition bool, msg string, v ...interface{}) {
	if !condition {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: "+msg+"\033[39m\n\n", append([]interface{}{filepath.Base(file), line}, v...)...)
		tb.FailNow()
	}
}

// ok fails the test if an err is not nil.
func ok(tb testing.TB, err error) {
	if err != nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d: unexpected error: %s\033[39m\n\n", filepath.Base(file), line, err.Error())
		tb.FailNow()
	}
}

// equals fails the test if exp is not equal to act.
func equals(tb testing.TB, exp, act interface{}) {
	if !reflect.DeepEqual(exp, act) {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n", filepath.Base(file), line, exp, act)
		tb.FailNow()
	}
}

func logErr(fn func() error, message string) {
	if err := fn(); err != nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("%s:%d: Error: %s: %s", filepath.Base(file), line, message, err.Error())
	}
}

func removeFileAndLogError(file string) {
	if err := os.Remove(file); err != nil {
		_, file, line, _ := runtime.Caller(1)
		fmt.Printf("%s:%d: Error removing file '%s': %s", filepath.Base(file), line, file, err.Error())
	}
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"time"

	"github.com/hailocab/elastigo/api"
	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/output"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 5
	esMaxConns         = 20
	esRecvBuffer       = 100
	esSendBuffer       = 100
)

type Indexer struct {
	events      int
	buffer      *bytes.Buffer
	indexPrefix string
	indexType   string
}

type Config struct {
	Hosts       []string `json:"hosts"`
	IndexPrefix string   `json:"index"`
	IndexType   string   `json:"indexType"`
}

type ESServer struct {
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
}

func init() {
	output.Register("elasticsearch", &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	})
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func bulkSend(buf *bytes.Buffer) error {
	_, err := api.DoCommand("POST", "/_bulk", buf)
	if err != nil {
		return err
	}

	return nil
}

func indexDoc(ev *buffer.Event) *map[string]interface{} {
	return &*ev.Fields
}

func (i *Indexer) writeBulk(index string, _type string, data interface{}) error {
	w := `{"index":{"_index":"%s","_type":"%s"}}`

	i.buffer.WriteString(fmt.Sprintf(w, index, _type))
	i.buffer.WriteByte('\n')

	switch v := data.(type) {
	case *bytes.Buffer:
		io.Copy(i.buffer, v)
	case []byte:
		i.buffer.Write(v)
	case string:
		i.buffer.WriteString(v)
	default:
		body, err := json.Marshal(data)
		if err != nil {
			log.Printf("Error writing bulk data: %v", err)
			return err
		}
		i.buffer.Write(body)
	}
	i.buffer.WriteByte('\n')
	return nil
}

func (i *Indexer) flush() {
	if i.events == 0 {
		return
	}

	log.Printf("Flushing %d event(s) to elasticsearch", i.events)
	for j := 0; j < 3; j++ {
		if err := bulkSend(i.buffer); err != nil {
			log.Printf("Failed to index event (will retry): %v", err)
			time.Sleep(time.Duration(50) * time.Millisecond)
			continue
		}
		break
	}

	i.buffer.Reset()
	i.events = 0
}

func (i *Indexer) index(ev *buffer.Event) {
	doc := indexDoc(ev)
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	i.events++
	i.writeBulk(idx, typ, doc)

	if i.events < esSendBuffer {
		return
	}

	log.Printf("Flushing %d event(s) to elasticsearch", i.events)
	for j := 0; j < 3; j++ {
		if err := bulkSend(i.buffer); err != nil {
			log.Printf("Failed to index event (will retry): %v", err)
			time.Sleep(time.Duration(50) * time.Millisecond)
			continue
		}
		break
	}

	i.buffer.Reset()
	i.events = 0
}

func (e *ESServer) Init(config json.RawMessage, b buffer.Sender) error {
	var esConfig *Config
	if err := json.Unmarshal(config, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	return nil
}

func (es *ESServer) Start() error {
	api.SetHosts(es.hosts)

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	// Create indexer
	idx := &Indexer{0, bytes.NewBuffer(nil), es.config.IndexPrefix, es.config.IndexType}

	// Loop events and publish to elasticsearch
	tick := time.NewTicker(time.Duration(esFlushInterval) * time.Second)

	for {
		select {
		case ev := <-receiveChan:
			idx.index(ev)
		case <-tick.C:
			idx.flush()
		case <-es.term:
			tick.Stop()
			log.Println("Elasticsearch received term signal")
			break
		}
	}

	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<commit_msg>Switch to elastic for Elasticsearch client for better timeout control<commit_after>package elasticsearch

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/output"
	"gopkg.in/olivere/elastic.v2"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 5
	esMaxConns         = 20
	esRecvBuffer       = 100
	esSendBuffer       = 100
)

type Indexer struct {
	bulkService *elastic.BulkService
	indexPrefix string
	indexType   string
}

type Config struct {
	Hosts       []string `json:"hosts"`
	IndexPrefix string   `json:"index"`
	IndexType   string   `json:"indexType"`
}

type ESServer struct {
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
}

func init() {
	output.Register("elasticsearch", &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	})
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func indexDoc(ev *buffer.Event) *map[string]interface{} {
	return &*ev.Fields
}

func (i *Indexer) flush() {
	log.Printf("Flushing %d event(s) to elasticsearch", i.bulkService.NumberOfActions())
	i.bulkService.Do()
}

func (i *Indexer) index(ev *buffer.Event) {
	log.Printf("Received event here")
	doc := indexDoc(ev)
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	request := elastic.NewBulkIndexRequest().Index(idx).Type(typ).Doc(doc)
	i.bulkService.Add(request)
	numEvents := i.bulkService.NumberOfActions()

	if numEvents < esSendBuffer {
		return
	}

	i.flush()
}

func (e *ESServer) Init(config json.RawMessage, b buffer.Sender) error {
	var esConfig *Config
	if err := json.Unmarshal(config, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	return nil
}

func (es *ESServer) Start() error {
	client, err := elastic.NewClient(elastic.SetURL(es.hosts...))

	if err != nil {
		log.Printf("Error starting Elasticsearch: %s", err)
		return err
	}

	service := elastic.NewBulkService(client)

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	// Create indexer
	idx := &Indexer{service, es.config.IndexPrefix, es.config.IndexType}

	// Loop events and publish to elasticsearch
	tick := time.NewTicker(time.Duration(esFlushInterval) * time.Second)

	for {
		select {
		case ev := <-receiveChan:
			idx.index(ev)
		case <-tick.C:
			idx.flush()
		case <-es.term:
			tick.Stop()
			log.Println("Elasticsearch received term signal")
			break
		}
	}

	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<|endoftext|>


<commit_before>package peerstore

import (
	"context"
	"sync"
	"time"

	ds "github.com/ipfs/go-datastore"
	"github.com/ipfs/go-datastore/query"
	"github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
	mh "github.com/multiformats/go-multihash"
)

// Number of times to retry transactional writes
var dsWriteRetries = 5

// DatastoreAddrManager is an address manager backed by a Datastore with both an
// in-memory TTL manager and an in-memory address stream manager.
type DatastoreAddrManager struct {
	ds          ds.Batching
	ttlManager  *ttlmanager
	subsManager *AddrSubManager
}

// NewDatastoreAddrManager initializes a new DatastoreAddrManager given a
// Datastore instance, a context for managing the TTL manager, and the interval
// at which the TTL manager should sweep the Datastore.
func NewDatastoreAddrManager(ctx context.Context, ds ds.Batching, ttlInterval time.Duration) *DatastoreAddrManager {
	mgr := &DatastoreAddrManager{
		ds:          ds,
		ttlManager:  newTTLManager(ctx, ds, ttlInterval),
		subsManager: NewAddrSubManager(),
	}
	return mgr
}

// Stop will signal the TTL manager to stop and block until it returns.
func (mgr *DatastoreAddrManager) Stop() {
	mgr.ttlManager.cancel()
}

func peerAddressKey(p *peer.ID, addr *ma.Multiaddr) (ds.Key, error) {
	hash, err := mh.Sum((*addr).Bytes(), mh.MURMUR3, -1)
	if err != nil {
		return ds.Key{}, nil
	}
	return ds.NewKey(peer.IDB58Encode(*p)).ChildString(hash.B58String()), nil
}

func peerIDFromKey(key ds.Key) (peer.ID, error) {
	idstring := key.Parent().Name()
	return peer.IDB58Decode(idstring)
}

// AddAddr will add a new address if it's not already in the AddrBook.
func (mgr *DatastoreAddrManager) AddAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mgr.AddAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// AddAddrs will add many new addresses if they're not already in the AddrBook.
func (mgr *DatastoreAddrManager) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	if ttl <= 0 {
		return
	}

	mgr.setAddrs(p, addrs, ttl, true)
}

// SetAddr will add or update the TTL of an address in the AddrBook.
func (mgr *DatastoreAddrManager) SetAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mgr.SetAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// SetAddrs will add or update the TTLs of addresses in the AddrBook.
func (mgr *DatastoreAddrManager) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mgr.setAddrs(p, addrs, ttl, false)
}

func (mgr *DatastoreAddrManager) setAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration, add bool) {
	for i := 0; i < dsWriteRetries; i++ {
		// keys to add to the TTL manager
		var keys []ds.Key
		batch, err := mgr.ds.Batch()
		if err != nil {
			log.Error(err)
			return
		}

		for _, addr := range addrs {
			if addr == nil {
				continue
			}

			key, err := peerAddressKey(&p, &addr)
			if err != nil {
				log.Error(err)
				continue
			}
			keys = append(keys, key)

			if ttl <= 0 {
				batch.Delete(key)
				continue
			}
			has, err := mgr.ds.Has(key)
			if err != nil || !has {
				mgr.subsManager.BroadcastAddr(p, addr)
			}

			// Allows us to support AddAddr and SetAddr in one function
			if !has {
				if err := batch.Put(key, addr.Bytes()); err != nil {
					log.Error(err)
				}
			}
		}
		if err := batch.Commit(); err != nil {
			log.Errorf("failed to write addresses for peer %s: %s\n", p.Pretty(), err)
			continue
		}
		mgr.ttlManager.setTTLs(keys, ttl, add)
		return
	}
	log.Errorf("failed to avoid write conflict for peer %s after %d retries\n", p.Pretty(), dsWriteRetries)
}

// UpdateAddrs will update any addresses for a given peer and TTL combination to
// have a new TTL.
func (mgr *DatastoreAddrManager) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration) {
	prefix := ds.NewKey(p.Pretty())
	mgr.ttlManager.updateTTLs(prefix, oldTTL, newTTL)
}

// Addrs Returns all of the non-expired addresses for a given peer.
func (mgr *DatastoreAddrManager) Addrs(p peer.ID) []ma.Multiaddr {
	prefix := ds.NewKey(p.Pretty())
	q := query.Query{Prefix: prefix.String()}
	results, err := mgr.ds.Query(q)
	if err != nil {
		log.Error(err)
		return nil
	}

	var addrs []ma.Multiaddr
	for result := range results.Next() {
		addrbytes := result.Value.([]byte)
		addr, err := ma.NewMultiaddrBytes(addrbytes)
		if err != nil {
			log.Error(err)
			continue
		}
		addrs = append(addrs, addr)
	}

	return addrs
}

// Peers returns all of the peer IDs for which the AddrBook has addresses.
func (mgr *DatastoreAddrManager) Peers() []peer.ID {
	q := query.Query{}
	results, err := mgr.ds.Query(q)
	if err != nil {
		log.Error(err)
		return []peer.ID{}
	}

	idset := make(map[peer.ID]struct{})
	for result := range results.Next() {
		key := ds.RawKey(result.Key)
		id, err := peerIDFromKey(key)
		if err != nil {
			continue
		}
		idset[id] = struct{}{}
	}

	ids := make([]peer.ID, 0, len(idset))
	for id := range idset {
		ids = append(ids, id)
	}
	return ids
}

// AddrStream returns a channel on which all new addresses discovered for a
// given peer ID will be published.
func (mgr *DatastoreAddrManager) AddrStream(ctx context.Context, p peer.ID) <-chan ma.Multiaddr {
	initial := mgr.Addrs(p)
	return mgr.subsManager.AddrStream(ctx, p, initial)
}

// ClearAddrs will delete all known addresses for a peer ID.
func (mgr *DatastoreAddrManager) ClearAddrs(p peer.ID) {
	prefix := ds.NewKey(p.Pretty())
	for i := 0; i < dsWriteRetries; i++ {
		q := query.Query{Prefix: prefix.String()}
		results, err := mgr.ds.Query(q)
		if err != nil {
			log.Error(err)
			return
		}
		batch, err := mgr.ds.Batch()
		if err != nil {
			log.Error(err)
			return
		}

		for result := range results.Next() {
			err := batch.Delete(ds.NewKey(result.Key))
			if err != nil {
				// From inspectin badger, errors here signify a problem with
				// the transaction as a whole, so we can log and abort.
				log.Error(err)
				return
			}
		}
		if err = batch.Commit(); err != nil {
			log.Errorf("failed to clear addresses for peer %s: %s\n", p.Pretty(), err)
			continue
		}
		mgr.ttlManager.clear(ds.NewKey(p.Pretty()))
		return
	}
	log.Errorf("failed to clear addresses for peer %s after %d attempts\n", p.Pretty(), dsWriteRetries)
}

// ttlmanager

type ttlentry struct {
	TTL       time.Duration
	ExpiresAt time.Time
}

type ttlmanager struct {
	sync.RWMutex
	entries map[ds.Key]*ttlentry

	ctx    context.Context
	cancel context.CancelFunc
	ticker *time.Ticker
	ds     ds.Batching
}

func newTTLManager(parent context.Context, d ds.Datastore, tick time.Duration) *ttlmanager {
	ctx, cancel := context.WithCancel(parent)
	batching, ok := d.(ds.Batching)
	if !ok {
		panic("must construct ttlmanager with batching datastore")
	}
	mgr := &ttlmanager{
		entries: make(map[ds.Key]*ttlentry),
		ctx:     ctx,
		cancel:  cancel,
		ticker:  time.NewTicker(tick),
		ds:      batching,
	}

	go func() {
		for {
			select {
			case <-mgr.ctx.Done():
				mgr.ticker.Stop()
				return
			case <-mgr.ticker.C:
				mgr.tick()
			}
		}
	}()

	return mgr
}

// To be called by TTL manager's coroutine only.
func (mgr *ttlmanager) tick() {
	mgr.RLock()
	defer mgr.RUnlock()

	now := time.Now()
	batch, err := mgr.ds.Batch()
	if err != nil {
		log.Error(err)
		return
	}
	for key, entry := range mgr.entries {
		if entry.ExpiresAt.Before(now) {
			if err := batch.Delete(key); err != nil {
				log.Error(err)
			}
			delete(mgr.entries, key)
		}
	}
	err = batch.Commit()
	if err != nil {
		log.Error(err)
	}
}

func (mgr *ttlmanager) setTTLs(keys []ds.Key, ttl time.Duration, add bool) {
	mgr.Lock()
	defer mgr.Unlock()

	expiration := time.Now().Add(ttl)
	for _, key := range keys {
		update := true
		if add {
			if entry, ok := mgr.entries[key]; ok {
				if entry.ExpiresAt.After(expiration) {
					update = false
				}
			}
		}
		if update {
			if ttl <= 0 {
				delete(mgr.entries, key)
			} else {
				mgr.entries[key] = &ttlentry{TTL: ttl, ExpiresAt: expiration}
			}
		}
	}
}

func (mgr *ttlmanager) updateTTLs(prefix ds.Key, oldTTL, newTTL time.Duration) {
	mgr.Lock()
	defer mgr.Unlock()

	now := time.Now()
	var keys []ds.Key
	for key, entry := range mgr.entries {
		if key.IsDescendantOf(prefix) && entry.TTL == oldTTL {
			keys = append(keys, key)
			entry.TTL = newTTL
			entry.ExpiresAt = now.Add(newTTL)
		}
	}
}

func (mgr *ttlmanager) clear(prefix ds.Key) {
	mgr.Lock()
	defer mgr.Unlock()

	for key := range mgr.entries {
		if key.IsDescendantOf(prefix) {
			delete(mgr.entries, key)
		}
	}
}
<commit_msg>Acquire full lock when ticking ttlmanager<commit_after>package peerstore

import (
	"context"
	"sync"
	"time"

	ds "github.com/ipfs/go-datastore"
	"github.com/ipfs/go-datastore/query"
	"github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"
	mh "github.com/multiformats/go-multihash"
)

// Number of times to retry transactional writes
var dsWriteRetries = 5

// DatastoreAddrManager is an address manager backed by a Datastore with both an
// in-memory TTL manager and an in-memory address stream manager.
type DatastoreAddrManager struct {
	ds          ds.Batching
	ttlManager  *ttlmanager
	subsManager *AddrSubManager
}

// NewDatastoreAddrManager initializes a new DatastoreAddrManager given a
// Datastore instance, a context for managing the TTL manager, and the interval
// at which the TTL manager should sweep the Datastore.
func NewDatastoreAddrManager(ctx context.Context, ds ds.Batching, ttlInterval time.Duration) *DatastoreAddrManager {
	mgr := &DatastoreAddrManager{
		ds:          ds,
		ttlManager:  newTTLManager(ctx, ds, ttlInterval),
		subsManager: NewAddrSubManager(),
	}
	return mgr
}

// Stop will signal the TTL manager to stop and block until it returns.
func (mgr *DatastoreAddrManager) Stop() {
	mgr.ttlManager.cancel()
}

func peerAddressKey(p *peer.ID, addr *ma.Multiaddr) (ds.Key, error) {
	hash, err := mh.Sum((*addr).Bytes(), mh.MURMUR3, -1)
	if err != nil {
		return ds.Key{}, nil
	}
	return ds.NewKey(peer.IDB58Encode(*p)).ChildString(hash.B58String()), nil
}

func peerIDFromKey(key ds.Key) (peer.ID, error) {
	idstring := key.Parent().Name()
	return peer.IDB58Decode(idstring)
}

// AddAddr will add a new address if it's not already in the AddrBook.
func (mgr *DatastoreAddrManager) AddAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mgr.AddAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// AddAddrs will add many new addresses if they're not already in the AddrBook.
func (mgr *DatastoreAddrManager) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	if ttl <= 0 {
		return
	}

	mgr.setAddrs(p, addrs, ttl, true)
}

// SetAddr will add or update the TTL of an address in the AddrBook.
func (mgr *DatastoreAddrManager) SetAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mgr.SetAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// SetAddrs will add or update the TTLs of addresses in the AddrBook.
func (mgr *DatastoreAddrManager) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mgr.setAddrs(p, addrs, ttl, false)
}

func (mgr *DatastoreAddrManager) setAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration, add bool) {
	for i := 0; i < dsWriteRetries; i++ {
		// keys to add to the TTL manager
		var keys []ds.Key
		batch, err := mgr.ds.Batch()
		if err != nil {
			log.Error(err)
			return
		}

		for _, addr := range addrs {
			if addr == nil {
				continue
			}

			key, err := peerAddressKey(&p, &addr)
			if err != nil {
				log.Error(err)
				continue
			}
			keys = append(keys, key)

			if ttl <= 0 {
				batch.Delete(key)
				continue
			}
			has, err := mgr.ds.Has(key)
			if err != nil || !has {
				mgr.subsManager.BroadcastAddr(p, addr)
			}

			// Allows us to support AddAddr and SetAddr in one function
			if !has {
				if err := batch.Put(key, addr.Bytes()); err != nil {
					log.Error(err)
				}
			}
		}
		if err := batch.Commit(); err != nil {
			log.Errorf("failed to write addresses for peer %s: %s\n", p.Pretty(), err)
			continue
		}
		mgr.ttlManager.setTTLs(keys, ttl, add)
		return
	}
	log.Errorf("failed to avoid write conflict for peer %s after %d retries\n", p.Pretty(), dsWriteRetries)
}

// UpdateAddrs will update any addresses for a given peer and TTL combination to
// have a new TTL.
func (mgr *DatastoreAddrManager) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration) {
	prefix := ds.NewKey(p.Pretty())
	mgr.ttlManager.updateTTLs(prefix, oldTTL, newTTL)
}

// Addrs Returns all of the non-expired addresses for a given peer.
func (mgr *DatastoreAddrManager) Addrs(p peer.ID) []ma.Multiaddr {
	prefix := ds.NewKey(p.Pretty())
	q := query.Query{Prefix: prefix.String()}
	results, err := mgr.ds.Query(q)
	if err != nil {
		log.Error(err)
		return nil
	}

	var addrs []ma.Multiaddr
	for result := range results.Next() {
		addrbytes := result.Value.([]byte)
		addr, err := ma.NewMultiaddrBytes(addrbytes)
		if err != nil {
			log.Error(err)
			continue
		}
		addrs = append(addrs, addr)
	}

	return addrs
}

// Peers returns all of the peer IDs for which the AddrBook has addresses.
func (mgr *DatastoreAddrManager) Peers() []peer.ID {
	q := query.Query{}
	results, err := mgr.ds.Query(q)
	if err != nil {
		log.Error(err)
		return []peer.ID{}
	}

	idset := make(map[peer.ID]struct{})
	for result := range results.Next() {
		key := ds.RawKey(result.Key)
		id, err := peerIDFromKey(key)
		if err != nil {
			continue
		}
		idset[id] = struct{}{}
	}

	ids := make([]peer.ID, 0, len(idset))
	for id := range idset {
		ids = append(ids, id)
	}
	return ids
}

// AddrStream returns a channel on which all new addresses discovered for a
// given peer ID will be published.
func (mgr *DatastoreAddrManager) AddrStream(ctx context.Context, p peer.ID) <-chan ma.Multiaddr {
	initial := mgr.Addrs(p)
	return mgr.subsManager.AddrStream(ctx, p, initial)
}

// ClearAddrs will delete all known addresses for a peer ID.
func (mgr *DatastoreAddrManager) ClearAddrs(p peer.ID) {
	prefix := ds.NewKey(p.Pretty())
	for i := 0; i < dsWriteRetries; i++ {
		q := query.Query{Prefix: prefix.String()}
		results, err := mgr.ds.Query(q)
		if err != nil {
			log.Error(err)
			return
		}
		batch, err := mgr.ds.Batch()
		if err != nil {
			log.Error(err)
			return
		}

		for result := range results.Next() {
			err := batch.Delete(ds.NewKey(result.Key))
			if err != nil {
				// From inspectin badger, errors here signify a problem with
				// the transaction as a whole, so we can log and abort.
				log.Error(err)
				return
			}
		}
		if err = batch.Commit(); err != nil {
			log.Errorf("failed to clear addresses for peer %s: %s\n", p.Pretty(), err)
			continue
		}
		mgr.ttlManager.clear(ds.NewKey(p.Pretty()))
		return
	}
	log.Errorf("failed to clear addresses for peer %s after %d attempts\n", p.Pretty(), dsWriteRetries)
}

// ttlmanager

type ttlentry struct {
	TTL       time.Duration
	ExpiresAt time.Time
}

type ttlmanager struct {
	sync.RWMutex
	entries map[ds.Key]*ttlentry

	ctx    context.Context
	cancel context.CancelFunc
	ticker *time.Ticker
	ds     ds.Batching
}

func newTTLManager(parent context.Context, d ds.Datastore, tick time.Duration) *ttlmanager {
	ctx, cancel := context.WithCancel(parent)
	batching, ok := d.(ds.Batching)
	if !ok {
		panic("must construct ttlmanager with batching datastore")
	}
	mgr := &ttlmanager{
		entries: make(map[ds.Key]*ttlentry),
		ctx:     ctx,
		cancel:  cancel,
		ticker:  time.NewTicker(tick),
		ds:      batching,
	}

	go func() {
		for {
			select {
			case <-mgr.ctx.Done():
				mgr.ticker.Stop()
				return
			case <-mgr.ticker.C:
				mgr.tick()
			}
		}
	}()

	return mgr
}

// To be called by TTL manager's coroutine only.
func (mgr *ttlmanager) tick() {
	mgr.Lock()
	defer mgr.Unlock()

	now := time.Now()
	batch, err := mgr.ds.Batch()
	if err != nil {
		log.Error(err)
		return
	}
	for key, entry := range mgr.entries {
		if entry.ExpiresAt.Before(now) {
			if err := batch.Delete(key); err != nil {
				log.Error(err)
			}
			delete(mgr.entries, key)
		}
	}
	err = batch.Commit()
	if err != nil {
		log.Error(err)
	}
}

func (mgr *ttlmanager) setTTLs(keys []ds.Key, ttl time.Duration, add bool) {
	mgr.Lock()
	defer mgr.Unlock()

	expiration := time.Now().Add(ttl)
	for _, key := range keys {
		update := true
		if add {
			if entry, ok := mgr.entries[key]; ok {
				if entry.ExpiresAt.After(expiration) {
					update = false
				}
			}
		}
		if update {
			if ttl <= 0 {
				delete(mgr.entries, key)
			} else {
				mgr.entries[key] = &ttlentry{TTL: ttl, ExpiresAt: expiration}
			}
		}
	}
}

func (mgr *ttlmanager) updateTTLs(prefix ds.Key, oldTTL, newTTL time.Duration) {
	mgr.Lock()
	defer mgr.Unlock()

	now := time.Now()
	var keys []ds.Key
	for key, entry := range mgr.entries {
		if key.IsDescendantOf(prefix) && entry.TTL == oldTTL {
			keys = append(keys, key)
			entry.TTL = newTTL
			entry.ExpiresAt = now.Add(newTTL)
		}
	}
}

func (mgr *ttlmanager) clear(prefix ds.Key) {
	mgr.Lock()
	defer mgr.Unlock()

	for key := range mgr.entries {
		if key.IsDescendantOf(prefix) {
			delete(mgr.entries, key)
		}
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"bufio"
	"bytes"
	"os"
	"regexp"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"

	"gopkg.in/h2non/gock.v1"
)

var expected_vars = map[string]string{
	"NAV_TRUSTSTORE_KEYSTOREALIAS": "app-key",
	"APP_DB_URL":                   "jdbc:oracle:thin:@//testdatabase.local:1521/app_db",
	"APP_DB_USERNAME":              "generic_database_username",
	"SOME_API_REST_URL":            "https://external-application.nais.preprod.local/rest/v1/very_useful_call",
}

func runCommand(t *testing.T, args []string) ([]byte, int) {
	defer gock.Off()

	gock.New("https://fasit.local").
		Get("/api/v2/scopedresource").
		MatchParam("alias", "app_db").
		MatchParam("application", "test-application").
		Reply(200).File("../../api/testdata/fasit_response_datasource.json")

	gock.New("https://fasit.local").
		Get("/api/v2/scopedresource").
		MatchParam("alias", "nav_truststore").
		MatchParam("application", "test-application").
		Reply(200).File("../../api/testdata/fasitTruststoreResponse.json")

	gock.New("https://fasit.local").
		Get("/api/v2/scopedresource").
		MatchParam("alias", "some_api").
		MatchParam("application", "test-application").
		Reply(200).File("../../api/testdata/fasit_response_restservice.json")

	gock.New("https://fasit.local").
		Get("/api/v2/resources/3024713/file/keystore").
		Reply(200).Body(bytes.NewReader([]byte("very secure certificate file :)")))

	// Since we're testing the output to console we need to save the default os.Stdout
	stdout := os.Stdout

	// To read from os.Stdout we replace it with a pipe
	r, w, err := os.Pipe()
	assert.Nil(t, err)
	os.Stdout = w

	reader := bufio.NewReader(r)

	// Execute the command
	RootCmd.SetArgs(args)
	err = environmentCommand.Execute()

	assert.Nil(t, err)

	// Extract the console output from our pipe
	commandResult := make([]byte, 512)
	commandResultLength, err := reader.Read(commandResult)

	assert.Nil(t, err)

	// When the command is done executing we revert to the original os.Stdout
	os.Stdout = stdout

	return commandResult, commandResultLength
}

func TestExportFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	// Make sure the export format has a trailing newline
	assert.Equal(t, commandResult[commandResultLength-1], byte('\n'))

	// Remove the trailing newline and then split it by newlines
	resultStrings := strings.Split(string(commandResult[:commandResultLength-1]), "\n")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("export ([A-Z_]+)='(.+)'")
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}
func TestMultilineFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-o", "multiline", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	// Make sure the export format has a trailing newline
	assert.Equal(t, commandResult[commandResultLength-1], byte('\n'))

	// Remove the trailing newline and then split it by newlines
	resultStrings := strings.Split(string(commandResult[:commandResultLength-1]), "\n")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("([A-Z_]+)='(.+)'")
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}

func TestDockerFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-o", "docker", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	// Remove the first -e, so we can split by the environment flag
	resultStrings := strings.Split(string(commandResult[2:]), " -e")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("([A-Z_]+)='(.+)'") // The strings we got after splitting should have NAME='value' format
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}

func TestJavaFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-o", "java", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	/* For java we're using -DNAME=variable, and in the current test data we don't expect any
	variables to have any spaces, so we can just split on space. If we later want to test
	with newlines we need to redo how we test induvidual variables. */
	resultStrings := strings.Split(string(commandResult[:]), " ")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("-D([A-Z_]+)='(.+)'")
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}
<commit_msg>Add missing test for the inline format<commit_after>package cmd

import (
	"bufio"
	"bytes"
	"os"
	"regexp"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"

	"gopkg.in/h2non/gock.v1"
)

var expected_vars = map[string]string{
	"NAV_TRUSTSTORE_KEYSTOREALIAS": "app-key",
	"APP_DB_URL":                   "jdbc:oracle:thin:@//testdatabase.local:1521/app_db",
	"APP_DB_USERNAME":              "generic_database_username",
	"SOME_API_REST_URL":            "https://external-application.nais.preprod.local/rest/v1/very_useful_call",
}

func runCommand(t *testing.T, args []string) ([]byte, int) {
	defer gock.Off()

	gock.New("https://fasit.local").
		Get("/api/v2/scopedresource").
		MatchParam("alias", "app_db").
		MatchParam("application", "test-application").
		Reply(200).File("../../api/testdata/fasit_response_datasource.json")

	gock.New("https://fasit.local").
		Get("/api/v2/scopedresource").
		MatchParam("alias", "nav_truststore").
		MatchParam("application", "test-application").
		Reply(200).File("../../api/testdata/fasitTruststoreResponse.json")

	gock.New("https://fasit.local").
		Get("/api/v2/scopedresource").
		MatchParam("alias", "some_api").
		MatchParam("application", "test-application").
		Reply(200).File("../../api/testdata/fasit_response_restservice.json")

	gock.New("https://fasit.local").
		Get("/api/v2/resources/3024713/file/keystore").
		Reply(200).Body(bytes.NewReader([]byte("very secure certificate file :)")))

	// Since we're testing the output to console we need to save the default os.Stdout
	stdout := os.Stdout

	// To read from os.Stdout we replace it with a pipe
	r, w, err := os.Pipe()
	assert.Nil(t, err)
	os.Stdout = w

	reader := bufio.NewReader(r)

	// Execute the command
	RootCmd.SetArgs(args)
	err = environmentCommand.Execute()

	assert.Nil(t, err)

	// Extract the console output from our pipe
	commandResult := make([]byte, 512)
	commandResultLength, err := reader.Read(commandResult)

	assert.Nil(t, err)

	// When the command is done executing we revert to the original os.Stdout
	os.Stdout = stdout

	return commandResult, commandResultLength
}

func TestExportFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	// Make sure the export format has a trailing newline
	assert.Equal(t, commandResult[commandResultLength-1], byte('\n'))

	// Remove the trailing newline and then split it by newlines
	resultStrings := strings.Split(string(commandResult[:commandResultLength-1]), "\n")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("export ([A-Z_]+)='(.+)'")
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}
func TestMultilineFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-o", "multiline", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	// Make sure the export format has a trailing newline
	assert.Equal(t, commandResult[commandResultLength-1], byte('\n'))

	// Remove the trailing newline and then split it by newlines
	resultStrings := strings.Split(string(commandResult[:commandResultLength-1]), "\n")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("([A-Z_]+)='(.+)'")
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}

func TestDockerFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-o", "docker", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	// Remove the first -e, so we can split by the environment flag
	resultStrings := strings.Split(string(commandResult[2:]), " -e")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("([A-Z_]+)='(.+)'") // The strings we got after splitting should have NAME='value' format
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}

func TestJavaFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-o", "java", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	/* For java we're using -DNAME=variable, and in the current test data we don't expect any
	variables to have any spaces, so we can just split on space. If we later want to test
	with newlines we need to redo how we test induvidual variables. */
	resultStrings := strings.Split(string(commandResult[:]), " ")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("-D([A-Z_]+)='(.+)'")
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}

func TestInlineFormat(t *testing.T) {
	commandResult, commandResultLength := runCommand(t, []string{"env", "-u", "https://fasit.local", "-o", "inline", "-f", "../../api/testdata/nais_used_resources.yaml", "test-application"})

	assert.NotEqual(t, 0, commandResultLength)

	// This should be fine as long as we don't have spaces in the variables
	resultStrings := strings.Split(string(commandResult[:]), " ")

	assert.Equal(t, len(expected_vars), len(resultStrings))

	exportRegex, _ := regexp.Compile("([A-Z_]+)='(.+)'")
	for _, value := range resultStrings {
		matches := exportRegex.FindStringSubmatch(value)

		assert.NotNil(t, matches)

		assert.Contains(t, expected_vars, matches[1])
		assert.Equal(t, expected_vars[matches[1]], matches[2])
	}
}
<|endoftext|>


<commit_before>package shade

import (
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"time"
)

var (
	chunksize = flag.Int("chunksize", 16*1024*1024, "size of a chunk, in bytes")
)

// File represents the metadata of a file stored in Shade.  It is stored and
// retrieved by the drive.Client API, and boiled down
type File struct {
	// Filename is a fully qualified path, with no leading slash.
	Filename string
	Filesize int64 // Bytes

	// ModifiedTime represents the "commit" time of this File object.  A given
	// Filename is represented by the valid File with the latest ModifiedTime.
	ModifiedTime time.Time

	// Chunks represets an ordered list of the bytes in the file.
	Chunks []Chunk

	// Chunksize is the maximum size of each plaintext Chunk, in bytes.
	Chunksize int

	// LastChunksize is the size of the last chunk in the File.  Storing this
	// explicity avoids the need to fetch the last chunk to update the Filesize.
	LastChunksize int

	// Deleted indicates all previous versions of this file should be suppressed.
	Deleted bool

	// AesKey is a 256 bit key used to encrypt the Chunks with AES-GCM.  If no
	// key is provided, the blocks are not encrypted.  The GCM nonce is stored at
	// the front of the encrypted Chunk using gcm.Seal(); use gcm.Open() to
	// recover the Nonce when decrypting.  Nb: This increases the encrypted
	// Chunk's size by gcm.NonceSize(), currently 12 bytes.
	AesKey *[32]byte
}

func NewFile(filename string) *File {
	return &File{
		Filename:     filename,
		ModifiedTime: time.Now(),
		Chunksize:    *chunksize,
		AesKey:       NewSymmetricKey(),
	}
}

// Chunk represents a portion of the content of the File being stored.
type Chunk struct {
	Index  int
	Sha256 []byte
	Nonce  []byte // If encrypted, use this Nonce to store/retrieve the Sum.
}

func (f *File) String() string {
	out := fmt.Sprintf("{Filename: %q, Filesize: %d, Chunksize: %d, AesKey: %q, Chunks:", f.Filename, f.Filesize, f.Chunksize, f.AesKey)
	sep := ", "
	if len(f.Chunks) < 2 {
		out += " "
	} else {
		out += "\n"
		sep = ",\n"
	}
	for i, c := range f.Chunks {
		if i == len(f.Chunks) {
			out += c.String() + sep
		} else {
			out += c.String()
		}
	}
	return out
}

// ToJSON returns a JSON representation of the File struct.
func (f *File) ToJSON() ([]byte, error) {
	fj, err := json.Marshal(f)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal file %x: %s", f.Filename, err)
	}
	return fj, nil
}

// FromJSON populates the fields of this File struct from a JSON representation.
// It primarily provides a convenient error message if this fails.
func (f *File) FromJSON(fj []byte) error {
	if err := json.Unmarshal(fj, f); err != nil {
		return fmt.Errorf("failed to unmarshal sha256sum %x: %s", SumString(fj), err)
	}
	return nil
}

// UpdateFilesize calculates the size of the assocaited Chunks and sets the
// Filesize member of the struct.
func (f *File) UpdateFilesize() {
	f.Filesize = int64((len(f.Chunks) - 1) * f.Chunksize)
	f.Filesize += int64(f.LastChunksize)
}

func NewChunk() Chunk {
	return Chunk{Nonce: NewNonce()}
}

func (c *Chunk) String() string {
	return fmt.Sprintf("{Index: %d, Sha256: %x}", c.Index, c.Sha256)
}

// NewSymmetricKey generates a random 256-bit AES key for File{}s.
// It panics if the source of randomness fails.
func NewSymmetricKey() *[32]byte {
	key := [32]byte{}
	_, err := io.ReadFull(rand.Reader, key[:])
	if err != nil {
		panic(err)
	}
	return &key
}

// NewNonce generates a random Nonce for AES-GCM.
// It panics if the source of randomness fails.
func NewNonce() []byte {
	nonce := make([]byte, 12)
	_, err := io.ReadFull(rand.Reader, nonce)
	if err != nil {
		panic(err)
	}
	return nonce
}
<commit_msg>Resolve the pointer on AesKey and print as %v<commit_after>package shade

import (
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"time"
)

var (
	chunksize = flag.Int("chunksize", 16*1024*1024, "size of a chunk, in bytes")
)

// File represents the metadata of a file stored in Shade.  It is stored and
// retrieved by the drive.Client API, and boiled down
type File struct {
	// Filename is a fully qualified path, with no leading slash.
	Filename string
	Filesize int64 // Bytes

	// ModifiedTime represents the "commit" time of this File object.  A given
	// Filename is represented by the valid File with the latest ModifiedTime.
	ModifiedTime time.Time

	// Chunks represets an ordered list of the bytes in the file.
	Chunks []Chunk

	// Chunksize is the maximum size of each plaintext Chunk, in bytes.
	Chunksize int

	// LastChunksize is the size of the last chunk in the File.  Storing this
	// explicity avoids the need to fetch the last chunk to update the Filesize.
	LastChunksize int

	// Deleted indicates all previous versions of this file should be suppressed.
	Deleted bool

	// AesKey is a 256 bit key used to encrypt the Chunks with AES-GCM.  If no
	// key is provided, the blocks are not encrypted.  The GCM nonce is stored at
	// the front of the encrypted Chunk using gcm.Seal(); use gcm.Open() to
	// recover the Nonce when decrypting.  Nb: This increases the encrypted
	// Chunk's size by gcm.NonceSize(), currently 12 bytes.
	AesKey *[32]byte
}

func NewFile(filename string) *File {
	return &File{
		Filename:     filename,
		ModifiedTime: time.Now(),
		Chunksize:    *chunksize,
		AesKey:       NewSymmetricKey(),
	}
}

// Chunk represents a portion of the content of the File being stored.
type Chunk struct {
	Index  int
	Sha256 []byte
	Nonce  []byte // If encrypted, use this Nonce to store/retrieve the Sum.
}

func (f *File) String() string {
	out := fmt.Sprintf("{Filename: %q, Filesize: %d, Chunksize: %d, AesKey: %v, Chunks:", f.Filename, f.Filesize, f.Chunksize, *f.AesKey)
	sep := ", "
	if len(f.Chunks) < 2 {
		out += " "
	} else {
		out += "\n"
		sep = ",\n"
	}
	for i, c := range f.Chunks {
		if i == len(f.Chunks) {
			out += c.String() + sep
		} else {
			out += c.String()
		}
	}
	return out
}

// ToJSON returns a JSON representation of the File struct.
func (f *File) ToJSON() ([]byte, error) {
	fj, err := json.Marshal(f)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal file %x: %s", f.Filename, err)
	}
	return fj, nil
}

// FromJSON populates the fields of this File struct from a JSON representation.
// It primarily provides a convenient error message if this fails.
func (f *File) FromJSON(fj []byte) error {
	if err := json.Unmarshal(fj, f); err != nil {
		return fmt.Errorf("failed to unmarshal sha256sum %x: %s", SumString(fj), err)
	}
	return nil
}

// UpdateFilesize calculates the size of the assocaited Chunks and sets the
// Filesize member of the struct.
func (f *File) UpdateFilesize() {
	f.Filesize = int64((len(f.Chunks) - 1) * f.Chunksize)
	f.Filesize += int64(f.LastChunksize)
}

func NewChunk() Chunk {
	return Chunk{Nonce: NewNonce()}
}

func (c *Chunk) String() string {
	return fmt.Sprintf("{Index: %d, Sha256: %x}", c.Index, c.Sha256)
}

// NewSymmetricKey generates a random 256-bit AES key for File{}s.
// It panics if the source of randomness fails.
func NewSymmetricKey() *[32]byte {
	key := [32]byte{}
	_, err := io.ReadFull(rand.Reader, key[:])
	if err != nil {
		panic(err)
	}
	return &key
}

// NewNonce generates a random Nonce for AES-GCM.
// It panics if the source of randomness fails.
func NewNonce() []byte {
	nonce := make([]byte, 12)
	_, err := io.ReadFull(rand.Reader, nonce)
	if err != nil {
		panic(err)
	}
	return nonce
}
<|endoftext|>


<commit_before>package goimage

import (
	"io"
	"mime/multipart"
	"os"
	"path/filepath"

	"github.com/omar-h/goimage/utils"
)

// File holds a file.
type File struct {
	File      multipart.File
	Header    *multipart.FileHeader
	Basename  string // Without extension
	Fullname  string // With extension
	Extension string
	MIMEType  string
	Size      int
}

// NewFile will create a new file from a multipart.FileHeader.
func NewFile(file multipart.File, fileHeader *multipart.FileHeader) *File {
	return &File{
		File:      file,
		Header:    fileHeader,
		Basename:  utils.GetFileBasename(fileHeader.Filename),
		Fullname:  fileHeader.Filename,
		Extension: filepath.Ext(fileHeader.Filename),
		MIMEType:  fileHeader.Header["Content-Type"][0],
		Size:      int(fileHeader.Size),
	}
}

// Place will move the file onto a specific location.
// Returns os package errors. (os.ErrFileExist and os.ErrPermission)
func (f *File) Place(location string) error {
	_, err := os.Stat(location + f.Fullname)
	if !os.IsNotExist(err) {
		return os.ErrExist
	}

	file, err := os.OpenFile(location+f.Fullname, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	_, err = io.Copy(file, f.File)
	if err != nil {
		return err
	}

	return nil
}

// GenerateName will generate a new name with a given length.
func (f *File) GenerateName(len int) *File {
	return f.GiveName(utils.GenerateName(len))
}

// GiveName will give the File a new name, and update the basename and fullname.
func (f *File) GiveName(name string) *File {
	f.Basename = name
	f.Fullname = name + f.Extension
	return f
}

// Close will properly close the file.
func (f *File) Close() error {
	return f.File.Close()
}
<commit_msg>Make the file.Place method use filepath.Join to ensure path is correct.<commit_after>package goimage

import (
	"io"
	"mime/multipart"
	"os"
	"path/filepath"

	"github.com/omar-h/goimage/utils"
)

// File holds a file.
type File struct {
	File      multipart.File
	Header    *multipart.FileHeader
	Basename  string // Without extension
	Fullname  string // With extension
	Extension string
	MIMEType  string
	Size      int
}

// NewFile will create a new file from a multipart.FileHeader.
func NewFile(file multipart.File, fileHeader *multipart.FileHeader) *File {
	return &File{
		File:      file,
		Header:    fileHeader,
		Basename:  utils.GetFileBasename(fileHeader.Filename),
		Fullname:  fileHeader.Filename,
		Extension: filepath.Ext(fileHeader.Filename),
		MIMEType:  fileHeader.Header["Content-Type"][0],
		Size:      int(fileHeader.Size),
	}
}

// Place will move the file onto a specific location.
// Returns os package errors. (os.ErrFileExist and os.ErrPermission)
func (f *File) Place(location string) error {
	fullpath := filepath.Join(location, f.Fullname)
	_, err := os.Stat(fullpath)
	if !os.IsNotExist(err) {
		return os.ErrExist
	}

	file, err := os.OpenFile(fullpath, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}

	defer file.Close()

	_, err = io.Copy(file, f.File)
	if err != nil {
		return err
	}

	return nil
}

// GenerateName will generate a new name with a given length.
func (f *File) GenerateName(len int) *File {
	return f.GiveName(utils.GenerateName(len))
}

// GiveName will give the File a new name, and update the basename and fullname.
func (f *File) GiveName(name string) *File {
	f.Basename = name
	f.Fullname = name + f.Extension
	return f
}

// Close will properly close the file.
func (f *File) Close() error {
	return f.File.Close()
}
<|endoftext|>


<commit_before>package scipipe

import (
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"sync"
)

// ======= FileTarget ========

type FileTarget struct {
	path     string
	buffer   *bytes.Buffer
	doStream bool
	lock     *sync.Mutex
}

func NewFileTarget(path string) *FileTarget {
	ft := new(FileTarget)
	ft.path = path
	ft.lock = new(sync.Mutex)
	//Don't init buffer if not needed?
	//buf := make([]byte, 0, 128)
	//ft.buffer = bytes.NewBuffer(buf)
	return ft
}

func (ft *FileTarget) GetPath() string {
	return ft.path
}

func (ft *FileTarget) GetTempPath() string {
	return ft.path + ".tmp"
}

func (ft *FileTarget) GetFifoPath() string {
	return ft.path + ".fifo"
}

func (ft *FileTarget) Open() *os.File {
	f, err := os.Open(ft.GetPath())
	Check(err)
	return f
}

func (ft *FileTarget) Read() []byte {
	dat, err := ioutil.ReadFile(ft.GetPath())
	Check(err)
	return dat
}

func (ft *FileTarget) Write(dat []byte) {
	err := ioutil.WriteFile(ft.GetTempPath(), dat, 0644)
	ft.Atomize()
	Check(err)
}

func (ft *FileTarget) Atomize() {
	Debug.Println("FileTarget: Atomizing", ft.GetTempPath(), "->", ft.GetPath())
	ft.lock.Lock()
	err := os.Rename(ft.GetTempPath(), ft.path)
	Check(err)
	ft.lock.Unlock()
	Debug.Println("FileTarget: Done atomizing", ft.GetTempPath(), "->", ft.GetPath())
}

func (ft *FileTarget) CreateFifo() {
	ft.lock.Lock()
	cmd := "mkfifo " + ft.GetFifoPath()
	Debug.Println("Now creating FIFO with command:", cmd)
	_, err := exec.Command("bash", "-c", cmd).Output()
	Check(err)
	ft.lock.Unlock()
}

func (ft *FileTarget) RemoveFifo() {
	ft.lock.Lock()
	output, err := exec.Command("bash", "-c", "rm "+ft.GetFifoPath()).Output()
	Check(err)
	Debug.Println("Removed FIFO output: ", output)
	ft.lock.Unlock()
}

func (ft *FileTarget) Exists() bool {
	exists := false
	ft.lock.Lock()
	if _, err := os.Stat(ft.GetPath()); err == nil {
		exists = true
	}
	ft.lock.Unlock()
	return exists
}

// ======= FileQueue =======

type FileQueue struct {
	process
	Out       chan *FileTarget
	FilePaths []string
}

func FQ(fps ...string) (fq *FileQueue) {
	return NewFileQueue(fps...)
}

func NewFileQueue(fps ...string) (fq *FileQueue) {
	filePaths := []string{}
	for _, fp := range fps {
		filePaths = append(filePaths, fp)
	}
	fq = &FileQueue{
		Out:       make(chan *FileTarget, BUFSIZE),
		FilePaths: filePaths,
	}
	return
}

func (proc *FileQueue) Run() {
	defer close(proc.Out)
	for _, fp := range proc.FilePaths {
		proc.Out <- NewFileTarget(fp)
	}
}

// ======= Sink =======

type Sink struct {
	process
	In chan *FileTarget
}

func NewSink() (s *Sink) {
	return &Sink{}
}

func (proc *Sink) Run() {
	for ft := range proc.In {
		Debug.Println("Received file in sink: ", ft.GetPath())
	}
}
<commit_msg>Check if FIFO exists before creating<commit_after>package scipipe

import (
	"bytes"
	"io/ioutil"
	"os"
	"os/exec"
	"sync"
)

// ======= FileTarget ========

type FileTarget struct {
	path     string
	buffer   *bytes.Buffer
	doStream bool
	lock     *sync.Mutex
}

func NewFileTarget(path string) *FileTarget {
	ft := new(FileTarget)
	ft.path = path
	ft.lock = new(sync.Mutex)
	//Don't init buffer if not needed?
	//buf := make([]byte, 0, 128)
	//ft.buffer = bytes.NewBuffer(buf)
	return ft
}

func (ft *FileTarget) GetPath() string {
	return ft.path
}

func (ft *FileTarget) GetTempPath() string {
	return ft.path + ".tmp"
}

func (ft *FileTarget) GetFifoPath() string {
	return ft.path + ".fifo"
}

func (ft *FileTarget) Open() *os.File {
	f, err := os.Open(ft.GetPath())
	Check(err)
	return f
}

func (ft *FileTarget) Read() []byte {
	dat, err := ioutil.ReadFile(ft.GetPath())
	Check(err)
	return dat
}

func (ft *FileTarget) Write(dat []byte) {
	err := ioutil.WriteFile(ft.GetTempPath(), dat, 0644)
	ft.Atomize()
	Check(err)
}

func (ft *FileTarget) Atomize() {
	Debug.Println("FileTarget: Atomizing", ft.GetTempPath(), "->", ft.GetPath())
	ft.lock.Lock()
	err := os.Rename(ft.GetTempPath(), ft.path)
	Check(err)
	ft.lock.Unlock()
	Debug.Println("FileTarget: Done atomizing", ft.GetTempPath(), "->", ft.GetPath())
}

func (ft *FileTarget) CreateFifo() {
	ft.lock.Lock()
	cmd := "mkfifo " + ft.GetFifoPath()
	Debug.Println("Now creating FIFO with command:", cmd)

	if _, err := os.Stat(ft.GetFifoPath()); err == nil {
		Warn.Println("FIFO already exists, so cannot be created:", ft.GetFifoPath())
	} else {
		_, err := exec.Command("bash", "-c", cmd).Output()
		Check(err)
	}

	ft.lock.Unlock()
}

func (ft *FileTarget) RemoveFifo() {
	ft.lock.Lock()
	output, err := exec.Command("bash", "-c", "rm "+ft.GetFifoPath()).Output()
	Check(err)
	Debug.Println("Removed FIFO output: ", output)
	ft.lock.Unlock()
}

func (ft *FileTarget) Exists() bool {
	exists := false
	ft.lock.Lock()
	if _, err := os.Stat(ft.GetPath()); err == nil {
		exists = true
	}
	ft.lock.Unlock()
	return exists
}

// ======= FileQueue =======

type FileQueue struct {
	process
	Out       chan *FileTarget
	FilePaths []string
}

func FQ(fps ...string) (fq *FileQueue) {
	return NewFileQueue(fps...)
}

func NewFileQueue(fps ...string) (fq *FileQueue) {
	filePaths := []string{}
	for _, fp := range fps {
		filePaths = append(filePaths, fp)
	}
	fq = &FileQueue{
		Out:       make(chan *FileTarget, BUFSIZE),
		FilePaths: filePaths,
	}
	return
}

func (proc *FileQueue) Run() {
	defer close(proc.Out)
	for _, fp := range proc.FilePaths {
		proc.Out <- NewFileTarget(fp)
	}
}

// ======= Sink =======

type Sink struct {
	process
	In chan *FileTarget
}

func NewSink() (s *Sink) {
	return &Sink{}
}

func (proc *Sink) Run() {
	for ft := range proc.In {
		Debug.Println("Received file in sink: ", ft.GetPath())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/kisielk/whisper-go/whisper"
	"github.com/splatpm/subhuman"
	"github.com/urfave/cli"
)

var (
	DataChan  chan map[string]string
	FillerSig chan bool
	Count     int64
	Mutex     *sync.Mutex
)

type Dirstate struct {
	Location string
	Contents map[string][]string
}

type Overlap struct {
	Source      string
	Destination string
	Contents    map[string][]string
}

func Filler() {
	for {
		select {
		case msg := <-DataChan:
			Debug <- fmt.Sprintf("Backfill: %s -> %s", msg["Source"], msg["Destination"])

			// Open our filehandles
			sDb, sErr := whisper.Open(msg["Source"])
			dDb, dErr := whisper.Open(msg["Destination"])
			if !chkErr(sErr) {
				Error <- fmt.Sprintf("%s: %s", msg["Source"], sErr.Error())
			} else if !chkErr(dErr) {
				Error <- fmt.Sprintf("%s: %s", msg["Destination"], dErr.Error())
				os.Exit(1)
			}

			// Now for a series of checks, first to ensure that both
			// files have the same number of archives in them.
			if sDb.Header.Metadata.ArchiveCount != dDb.Header.Metadata.ArchiveCount {
				Error <- fmt.Sprintln("The files have a mismatched set of archives.")
			} else {
				// Now we'll start processing the archives, checking as we go to see if they
				// are matched. That way we at least fill in what we can, possibly....
				for i, a := range sDb.Header.Archives {
					// The offset
					if a.Offset == dDb.Header.Archives[i].Offset {
						// and the number of points
						if a.Points == dDb.Header.Archives[i].Points {
							// and finally the interval
							if a.SecondsPerPoint == dDb.Header.Archives[i].SecondsPerPoint {
								// ok, now let's get rolling through the archives
								sp, se := sDb.DumpArchive(i)
								if se != nil {
									Error <- fmt.Sprintf("%s: %s", msg["Source"], se.Error())
									os.Exit(1)
								}
								dp, de := dDb.DumpArchive(i)
								if de != nil {
									Error <- fmt.Sprintln(de.Error())
									os.Exit(1)
								}
								for idx := 0; idx < len(sp); idx++ {
									if sp[idx].Timestamp != 0 && sp[idx].Value != 0 {
										if dp[idx].Timestamp == 0 || dp[idx].Value == 0 {
											dp[idx].Timestamp = sp[idx].Timestamp
											dp[idx].Value = sp[idx].Value
											dDb.Update(dp[idx])
										}
									}
								}
							}
						}
					}
				}
			}

			// Defer their closings
			sDb.Close()
			dDb.Close()

			// and increment our counter
			Mutex.Lock()
			Count++
			Mutex.Unlock()
		case <-FillerSig:
			FillerSig <- true
			return
		}
	}
}

func BackFill(c *cli.Context) {
	if len(c.Args()) < 2 {
		Error <- "Invalid arguments. See 'flustro help fill' for more information."
		os.Exit(1)
	} else {
		// declare our variables
		var srcObj Dirstate
		var dstObj Dirstate
		srcDir := c.Args().Get(0)
		dstDir := c.Args().Get(1)

		// Now let's do the heavy lifting
		if isDir(srcDir) && isDir(dstDir) {
			// First let's get our dir contents
			srcObj = ListDir(srcDir)
			dstObj = ListDir(dstDir)
			// then spawn our worker pool, and get to processing
			for i := 0; i < c.Int("j"); i++ {
				go Filler()
			}
			// next we'll start processing through our srcObj and dstObj lists and
			// backfill everything that's present in both locations
			overlap := Overlap{
				Source:      srcObj.Location,
				Destination: dstObj.Location,
				Contents:    make(map[string][]string),
			}
			overlap_c := 0
			for k, _ := range srcObj.Contents {
				if _, ok := dstObj.Contents[k]; ok {
					for _, v := range srcObj.Contents[k] {
						for _, dv := range dstObj.Contents[k] {
							if v == dv {
								if _, ok := overlap.Contents[k]; ok {
									overlap.Contents[k] = append(overlap.Contents[k], v)
									overlap_c++
								} else {
									overlap.Contents[k] = []string{v}
									overlap_c++
								}
							}
						}
					}
				}
			}
			Info <- fmt.Sprintf("%d entries shared between %s and %s.",
				overlap_c,
				overlap.Source,
				overlap.Destination)
			time.Sleep(500 * time.Millisecond)
			start_t := int64(time.Now().Unix())
			for k, v := range overlap.Contents {
				for _, f := range v {
					DataChan <- map[string]string{
						"Source":      fmt.Sprintf("%s/%s/%s", overlap.Source, k, f),
						"Destination": fmt.Sprintf("%s/%s/%s", overlap.Destination, k, f),
					}
				}
			}
			for {
				runtime := (int64(time.Now().Unix()) - start_t)
				if len(DataChan) == 0 {
					time.Sleep(2 * time.Second)
					runrate := (float32(overlap_c) / float32(runtime))
					Info <- fmt.Sprintf("%d files processed in %d sec @ %.02f/sec.", overlap_c, runtime, runrate)
					time.Sleep(100 * time.Millisecond)
					return
				}
				runrate := (float32(Count) / float32(runtime))
				Progress <- fmt.Sprintf("%d files processed in %s @ %.02f/sec.", Count, subhuman.HumanTimeColon(runtime), runrate)
				time.Sleep(1 * time.Second)
			}
		} else if isFile(srcDir) && isFile(dstDir) {
			// we only need one worker for this job
			go Filler()
			DataChan <- map[string]string{
				"Source":      srcDir,
				"Destination": dstDir,
			}
		} else {
			Error <- fmt.Sprintf("SRC and DST must be either both files or both dirs.")
		}

		time.Sleep(500 * time.Millisecond)
	}
	return
}

func init() {
	DataChan = make(chan map[string]string, 8192)
	Mutex = &sync.Mutex{}
	Commands = append(Commands, cli.Command{
		Name:        "fill",
		Aliases:     []string{"f"},
		Usage:       "Backfill datapoints in the dst from the src",
		Description: "Backfill datapoints in the dst from the src",
		ArgsUsage:   "<src(File|Dir)> <dst(File|Dir)>",
		Flags: []cli.Flag{
			cli.IntFlag{
				Name:  "j",
				Usage: "Number of workers (for directory recursion)",
				Value: (runtime.GOMAXPROCS(0) * 2),
			},
		},
		SkipFlagParsing: false,
		HideHelp:        false,
		Hidden:          false,
		Action:          BackFill,
	})
}
<commit_msg>better rate limiting and progress output<commit_after>package main

import (
	"fmt"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/kisielk/whisper-go/whisper"
	"github.com/splatpm/subhuman"
	"github.com/urfave/cli"
)

var (
	DataChan  chan map[string]string
	FillerSig chan bool
	Count     int64
	Mutex     *sync.Mutex
)

type Dirstate struct {
	Location string
	Contents map[string][]string
}

type Overlap struct {
	Source      string
	Destination string
	Contents    map[string][]string
}

func Filler() {
	for {
		select {
		case msg := <-DataChan:
			Debug <- fmt.Sprintf("Backfill: %s -> %s", msg["Source"], msg["Destination"])

			// Open our filehandles
			sDb, sErr := whisper.Open(msg["Source"])
			dDb, dErr := whisper.Open(msg["Destination"])
			if !chkErr(sErr) {
				Error <- fmt.Sprintf("%s: %s", msg["Source"], sErr.Error())
			} else if !chkErr(dErr) {
				Error <- fmt.Sprintf("%s: %s", msg["Destination"], dErr.Error())
				os.Exit(1)
			}

			// Now for a series of checks, first to ensure that both
			// files have the same number of archives in them.
			if sDb.Header.Metadata.ArchiveCount != dDb.Header.Metadata.ArchiveCount {
				Error <- fmt.Sprintln("The files have a mismatched set of archives.")
			} else {
				// Now we'll start processing the archives, checking as we go to see if they
				// are matched. That way we at least fill in what we can, possibly....
				for i, a := range sDb.Header.Archives {
					// The offset
					if a.Offset == dDb.Header.Archives[i].Offset {
						// and the number of points
						if a.Points == dDb.Header.Archives[i].Points {
							// and finally the interval
							if a.SecondsPerPoint == dDb.Header.Archives[i].SecondsPerPoint {
								// ok, now let's get rolling through the archives
								sp, se := sDb.DumpArchive(i)
								if se != nil {
									Error <- fmt.Sprintf("%s: %s", msg["Source"], se.Error())
									os.Exit(1)
								}
								dp, de := dDb.DumpArchive(i)
								if de != nil {
									Error <- fmt.Sprintln(de.Error())
									os.Exit(1)
								}
								for idx := 0; idx < len(sp); idx++ {
									if sp[idx].Timestamp != 0 && sp[idx].Value != 0 {
										if dp[idx].Timestamp == 0 || dp[idx].Value == 0 {
											dp[idx].Timestamp = sp[idx].Timestamp
											dp[idx].Value = sp[idx].Value
											dDb.Update(dp[idx])
										}
									}
								}
							}
						}
					}
				}
			}

			// Defer their closings
			sDb.Close()
			dDb.Close()

			// and increment our counter
			Mutex.Lock()
			Count++
			Mutex.Unlock()
		case <-FillerSig:
			FillerSig <- true
			return
		}
	}
}

func BackFill(c *cli.Context) {
	if len(c.Args()) < 2 {
		Error <- "Invalid arguments. See 'flustro help fill' for more information."
		os.Exit(1)
	} else {
		// declare our variables
		var srcObj Dirstate
		var dstObj Dirstate
		srcDir := c.Args().Get(0)
		dstDir := c.Args().Get(1)

		// Now let's do the heavy lifting
		if isDir(srcDir) && isDir(dstDir) {
			// First let's get our dir contents
			srcObj = ListDir(srcDir)
			dstObj = ListDir(dstDir)
			// then spawn our worker pool, and get to processing
			for i := 0; i < c.Int("j"); i++ {
				go Filler()
			}
			// next we'll start processing through our srcObj and dstObj lists and
			// backfill everything that's present in both locations
			overlap := Overlap{
				Source:      srcObj.Location,
				Destination: dstObj.Location,
				Contents:    make(map[string][]string),
			}
			overlap_c := 0
			for k, _ := range srcObj.Contents {
				if _, ok := dstObj.Contents[k]; ok {
					for _, v := range srcObj.Contents[k] {
						for _, dv := range dstObj.Contents[k] {
							if v == dv {
								if _, ok := overlap.Contents[k]; ok {
									overlap.Contents[k] = append(overlap.Contents[k], v)
									overlap_c++
								} else {
									overlap.Contents[k] = []string{v}
									overlap_c++
								}
							}
						}
					}
				}
			}
			Info <- fmt.Sprintf("%d entries shared between %s and %s.",
				overlap_c,
				overlap.Source,
				overlap.Destination)
			time.Sleep(500 * time.Millisecond)
			start_t := int64(time.Now().Unix())
			for k, v := range overlap.Contents {
				for _, f := range v {
					DataChan <- map[string]string{
						"Source":      fmt.Sprintf("%s/%s/%s", overlap.Source, k, f),
						"Destination": fmt.Sprintf("%s/%s/%s", overlap.Destination, k, f),
					}
					runtime := (int64(time.Now().Unix()) - start_t)
					if len(DataChan) != 0 {
						runrate := (float32(Count) / float32(runtime))
						Progress <- fmt.Sprintf("%d files processed in %s @ %.02f/sec.",
							Count, subhuman.HumanTimeColon(runtime), runrate)
					}
				}
			}
			for {
				runtime := (int64(time.Now().Unix()) - start_t)
				if len(DataChan) == 0 {
					time.Sleep(2 * time.Second)
					runrate := (float32(overlap_c) / float32(runtime))
					Info <- fmt.Sprintf("%d files processed in %s sec @ %.02f/sec.",
						overlap_c, subhuman.HumanTimeColon(runtime), runrate)
					time.Sleep(100 * time.Millisecond)
					return
				}
				runrate := (float32(Count) / float32(runtime))
				Progress <- fmt.Sprintf("%d files processed in %s @ %.02f/sec.",
					Count, subhuman.HumanTimeColon(runtime), runrate)
				time.Sleep(1 * time.Second)
			}
		} else if isFile(srcDir) && isFile(dstDir) {
			// we only need one worker for this job
			go Filler()
			DataChan <- map[string]string{
				"Source":      srcDir,
				"Destination": dstDir,
			}
		} else {
			Error <- fmt.Sprintf("SRC and DST must be either both files or both dirs.")
		}

		time.Sleep(500 * time.Millisecond)
	}
	return
}

func init() {
	DataChan = make(chan map[string]string, 128)
	Mutex = &sync.Mutex{}
	Commands = append(Commands, cli.Command{
		Name:        "fill",
		Aliases:     []string{"f"},
		Usage:       "Backfill datapoints in the dst from the src",
		Description: "Backfill datapoints in the dst from the src",
		ArgsUsage:   "<src(File|Dir)> <dst(File|Dir)>",
		Flags: []cli.Flag{
			cli.IntFlag{
				Name:  "j",
				Usage: "Number of workers (for directory recursion)",
				Value: (runtime.GOMAXPROCS(0) * 2),
			},
		},
		SkipFlagParsing: false,
		HideHelp:        false,
		Hidden:          false,
		Action:          BackFill,
	})
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
	"path"
	"runtime"

	. "github.com/fuzzy/gocolor"
	"github.com/kisielk/whisper-go/whisper"
	"github.com/urfave/cli"
)

func isDir(p string) bool {
	f, e := os.Stat(p)
	if e != nil {
		return false
	}
	return f.IsDir()
}

func chkErr(e error) bool {
	if e != nil {
		fmt.Printf("%s: %s\n", String("ERROR").Red().Bold(), e)
		return false
	} else {
		return true
	}
}

func backfillFile(s string, d string) bool {
	// Open our filehandles
	srcDb, srcErr := whisper.Open(s)
	chkErr(srcErr)
	dstDb, dstErr := whisper.Open(d)
	chkErr(dstErr)
	// Defer their closings
	defer srcDb.Close()
	defer dstDb.Close()

	// Now for a series of checks, first to ensure that both
	// files have the same number of archives in them.
	if srcDb.Header.Metadata.ArchiveCount != dstDb.Header.Metadata.ArchiveCount {
		fmt.Printf("%s: The files have a mismatched set of archives.\n", String("ERROR").Red().Bold())
		return false
	}

	// Now we'll start processing the archives, checking as we go to see if they are matched.
	// that way we at least fill in what we can, possibly....
	for i, a := range srcDb.Header.Archives {
		// The offset
		if a.Offset == dstDb.Header.Archives[i].Offset {
			// and the number of points
			if a.Points == dstDb.Header.Archives[i].Points {
				// and finally the interval
				if a.SecondsPerPoint == dstDb.Header.Archives[i].SecondsPerPoint {
					// ok, now let's get rolling through the archives
					fmt.Println("WE ARE GO, I REPEAT, WE ARE FUCKING GO!")
				}
			}
		}
	}
	return true
}

func Filler(c *cli.Context) error {
	if len(c.Args()) == 2 {
		args := c.Args()
		if isDir(args[0]) && isDir(args[1]) {
			e := fmt.Sprintf("%s: Dir comparison not complete yet", String("ERROR").Red().Bold())
			fmt.Println(e)
			return errors.New(e)
		} else {
			if !backfillFile(args[0], args[1]) {
				e := fmt.Sprintf("%s: There has been an error.", String("ERROR").Red().Bold())
				fmt.Println(e)
				return errors.New(e)
			}
		}
	} else {
		var e string
		if !c.Bool("c") {
			e = fmt.Sprintf("%s: Wrong number of paramters given.\n%s: Try '%s help fill' for more information",
				String("ERROR").Red().Bold(),
				String("ERROR").Red().Bold(),
				path.Base(os.Args[0]))
		} else {
			e = fmt.Sprintf("ERROR: Wrong number of parameters given.\nERROR: Try '%s help fill' for more information.",
				path.Base(os.Args[0]))
		}
		fmt.Println(e)
		return errors.New(e)
	}
	return nil
}

func init() {
	Commands = append(Commands, cli.Command{
		Name:        "fill",
		Aliases:     []string{"f"},
		Usage:       "Backfill datapoints in the dst(file|dir) from the src(file|dir)",
		Description: "Backfill datapoints in the dst(file|dir) from the src(file|dir)",
		ArgsUsage:   "<src(File|Dir)> <dst(File|Dir)>",
		Flags: []cli.Flag{
			cli.IntFlag{
				Name:  "j",
				Usage: "Number of workers (for directory recursion)",
				Value: runtime.GOMAXPROCS(0),
			},
			cli.BoolFlag{Name: "c", Usage: "Prevent colors from being used"},
		},
		SkipFlagParsing: false,
		HideHelp:        false,
		Hidden:          false,
		Action:          Filler,
	})
}
<commit_msg>comparison of points complete, just need to update the target whisper file, and add the directory walking bits<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"path"
	"runtime"

	. "github.com/fuzzy/gocolor"
	"github.com/kisielk/whisper-go/whisper"
	"github.com/urfave/cli"
)

func isDir(p string) bool {
	f, e := os.Stat(p)
	if e != nil {
		return false
	}
	return f.IsDir()
}

func chkErr(e error) bool {
	if e != nil {
		fmt.Printf("%s: %s\n", String("ERROR").Red().Bold(), e)
		return false
	} else {
		return true
	}
}

func backfillFile(s string, d string) bool {
	// Open our filehandles
	srcDb, srcErr := whisper.Open(s)
	chkErr(srcErr)
	dstDb, dstErr := whisper.Open(d)
	chkErr(dstErr)
	// Defer their closings
	defer srcDb.Close()
	defer dstDb.Close()

	// Now for a series of checks, first to ensure that both
	// files have the same number of archives in them.
	if srcDb.Header.Metadata.ArchiveCount != dstDb.Header.Metadata.ArchiveCount {
		fmt.Printf("%s: The files have a mismatched set of archives.\n", String("ERROR").Red().Bold())
		return false
	}

	// Now we'll start processing the archives, checking as we go to see if they are matched.
	// that way we at least fill in what we can, possibly....
	for i, a := range srcDb.Header.Archives {
		// The offset
		if a.Offset == dstDb.Header.Archives[i].Offset {
			// and the number of points
			if a.Points == dstDb.Header.Archives[i].Points {
				// and finally the interval
				if a.SecondsPerPoint == dstDb.Header.Archives[i].SecondsPerPoint {
					// ok, now let's get rolling through the archives
					fmt.Println("WE ARE GO, I REPEAT, WE ARE FUCKING GO!")
					sp, se := srcDb.DumpArchive(i)
					if se != nil {
						fmt.Printf("%s: %s\n", String("ERROR").Red().Bold(), se)
						os.Exit(1)
					}
					dp, de := dstDb.DumpArchive(i)
					if de != nil {
						fmt.Printf("%s: %s\n", String("ERROR").Red().Bold(), de)
						os.Exit(1)
					}
					for idx := 0; idx < len(sp); idx++ {
						if sp[idx].Timestamp != 0 && sp[idx].Value != 0 {
							if dp[idx].Timestamp == 0 || dp[idx].Value == 0 {
								fmt.Printf("SRC: %s %d %f\n", sp[idx].Time(), sp[idx].Timestamp, sp[idx].Value)
								fmt.Printf("DST: %s %d %f\n", dp[idx].Time(), dp[idx].Timestamp, dp[idx].Value)
							}
						}
					}
				}
			}
		}
	}
	return true
}

func Filler(c *cli.Context) error {
	if len(c.Args()) == 2 {
		args := c.Args()
		if isDir(args[0]) && isDir(args[1]) {
			e := fmt.Sprintf("%s: Dir comparison not complete yet", String("ERROR").Red().Bold())
			fmt.Println(e)
			return errors.New(e)
		} else {
			if !backfillFile(args[0], args[1]) {
				e := fmt.Sprintf("%s: There has been an error.", String("ERROR").Red().Bold())
				fmt.Println(e)
				return errors.New(e)
			}
		}
	} else {
		var e string
		if !c.Bool("c") {
			e = fmt.Sprintf("%s: Wrong number of paramters given.\n%s: Try '%s help fill' for more information",
				String("ERROR").Red().Bold(),
				String("ERROR").Red().Bold(),
				path.Base(os.Args[0]))
		} else {
			e = fmt.Sprintf("ERROR: Wrong number of parameters given.\nERROR: Try '%s help fill' for more information.",
				path.Base(os.Args[0]))
		}
		fmt.Println(e)
		return errors.New(e)
	}
	return nil
}

func init() {
	Commands = append(Commands, cli.Command{
		Name:        "fill",
		Aliases:     []string{"f"},
		Usage:       "Backfill datapoints in the dst(file|dir) from the src(file|dir)",
		Description: "Backfill datapoints in the dst(file|dir) from the src(file|dir)",
		ArgsUsage:   "<src(File|Dir)> <dst(File|Dir)>",
		Flags: []cli.Flag{
			cli.IntFlag{
				Name:  "j",
				Usage: "Number of workers (for directory recursion)",
				Value: runtime.GOMAXPROCS(0),
			},
			cli.BoolFlag{Name: "c", Usage: "Prevent colors from being used"},
		},
		SkipFlagParsing: false,
		HideHelp:        false,
		Hidden:          false,
		Action:          Filler,
	})
}
<|endoftext|>


<commit_before><commit_msg>Add a function to predict labels for testing.<commit_after><|endoftext|>