<commit_before>package peerstream_transport_test

import (
	"bytes"
	crand "crypto/rand"
	"fmt"
	"io"
	mrand "math/rand"
	"net"
	"os"
	"reflect"
	"runtime"
	"sync"
	"testing"

	ps "github.com/jbenet/go-peerstream"
	pst "github.com/jbenet/go-peerstream/transport"
)

var randomness []byte
var nextPort = 20000

func init() {
	// read 1MB of randomness
	randomness = make([]byte, 1<<20)
	if _, err := crand.Read(randomness); err != nil {
		panic(err)
	}
}

func randBuf(size int) []byte {
	n := len(randomness) - size
	if size < 1 {
		panic(fmt.Errorf("requested too large buffer (%d). max is %d", size, len(randomness)))
	}

	start := mrand.Intn(n)
	return randomness[start : start+size]
}

func checkErr(t *testing.T, err error) {
	if err != nil {
		t.Fatal(err)
	}
}

func log(s string, v ...interface{}) {
	if testing.Verbose() {
		fmt.Fprintf(os.Stderr, "> "+s+"\n", v...)
	}
}

type echoSetup struct {
	swarm *ps.Swarm
	conns []*ps.Conn
}

func singleConn(t *testing.T, tr pst.Transport) echoSetup {
	swarm := ps.NewSwarm(tr)
	swarm.SetStreamHandler(func(s *ps.Stream) {
		defer s.Close()
		log("accepted stream")
		io.Copy(s, s) // echo everything
		log("closing stream")
	})

	log("listening at %s", "localhost:0")
	l, err := net.Listen("tcp", "localhost:0")
	checkErr(t, err)

	_, err = swarm.AddListener(l)
	checkErr(t, err)

	log("dialing to %s", l.Addr())
	nc1, err := net.Dial("tcp", l.Addr().String())
	checkErr(t, err)

	c1, err := swarm.AddConn(nc1)
	checkErr(t, err)

	return echoSetup{
		swarm: swarm,
		conns: []*ps.Conn{c1},
	}
}

func makeSwarm(t *testing.T, tr pst.Transport, nListeners int) *ps.Swarm {
	swarm := ps.NewSwarm(tr)
	swarm.SetStreamHandler(func(s *ps.Stream) {
		defer s.Close()
		log("accepted stream")
		io.Copy(s, s) // echo everything
		log("closing stream")
	})

	for i := 0; i < nListeners; i++ {
		log("%p listening at %s", swarm, "localhost:0")
		l, err := net.Listen("tcp", "localhost:0")
		checkErr(t, err)
		_, err = swarm.AddListener(l)
		checkErr(t, err)
	}

	return swarm
}

func makeSwarms(t *testing.T, tr pst.Transport, nSwarms, nListeners int) []*ps.Swarm {
	swarms := make([]*ps.Swarm, nSwarms)
	for i := 0; i < nSwarms; i++ {
		swarms[i] = makeSwarm(t, tr, nListeners)
	}
	return swarms
}

func SubtestConstructSwarm(t *testing.T, tr pst.Transport) {
	ps.NewSwarm(tr)
}

func SubtestSimpleWrite(t *testing.T, tr pst.Transport) {
	swarm := ps.NewSwarm(tr)
	defer swarm.Close()

	piper, pipew := io.Pipe()
	swarm.SetStreamHandler(func(s *ps.Stream) {
		defer s.Close()
		log("accepted stream")
		w := io.MultiWriter(s, pipew)
		io.Copy(w, s) // echo everything and write it to pipew
		log("closing stream")
	})

	log("listening at %s", "localhost:0")
	l, err := net.Listen("tcp", "localhost:0")
	checkErr(t, err)

	_, err = swarm.AddListener(l)
	checkErr(t, err)

	log("dialing to %s", l.Addr().String())
	nc1, err := net.Dial("tcp", l.Addr().String())
	checkErr(t, err)

	c1, err := swarm.AddConn(nc1)
	checkErr(t, err)
	defer c1.Close()

	log("creating stream")
	s1, err := c1.NewStream()
	checkErr(t, err)
	defer s1.Close()

	buf1 := randBuf(4096)
	log("writing %d bytes to stream", len(buf1))
	_, err = s1.Write(buf1)
	checkErr(t, err)

	buf2 := make([]byte, len(buf1))
	log("reading %d bytes from stream (echoed)", len(buf2))
	_, err = s1.Read(buf2)
	checkErr(t, err)
	if string(buf2) != string(buf1) {
		t.Error("buf1 and buf2 not equal: %s != %s", string(buf1), string(buf2))
	}

	buf3 := make([]byte, len(buf1))
	log("reading %d bytes from pipe (tee)", len(buf3))
	_, err = piper.Read(buf3)
	checkErr(t, err)
	if string(buf3) != string(buf1) {
		t.Error("buf1 and buf3 not equal: %s != %s", string(buf1), string(buf3))
	}
}

func SubtestSimpleWrite100msgs(t *testing.T, tr pst.Transport) {

	msgs := 100
	msgsize := 1 << 19
	es := singleConn(t, tr)

	log("creating stream")
	stream, err := es.conns[0].NewStream()
	checkErr(t, err)

	bufs := make(chan []byte, msgs)
	errs := make(chan error, msgs*100)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()

		for i := 0; i < msgs; i++ {
			buf := randBuf(msgsize)
			bufs <- buf
			log("writing %d bytes (message %d/%d #%x)", len(buf), i, msgs, buf[:3])
			if _, err := stream.Write(buf); err != nil {
				errs <- err
				continue
			}
		}
		close(bufs)
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		buf2 := make([]byte, msgsize)
		i := 0
		for buf1 := range bufs {
			log("reading %d bytes (message %d/%d #%x)", len(buf1), i, msgs, buf1[:3])
			i++

			if _, err := io.ReadFull(stream, buf2); err != nil {
				errs <- err
				continue
			}
			if !bytes.Equal(buf1, buf2) {
				errs <- fmt.Errorf("buffers not equal (%x != %x)", buf1[:3], buf2[:3])
			}
		}
	}()

	wg.Wait()
	close(errs)
	for err := range errs {
		t.Error(err)
	}
}

func SubtestStressNSwarmNConnNStreamNMsg(t *testing.T, tr pst.Transport, nSwarm, nConn, nStream, nMsg int) {

	msgsize := 1 << 11
	errs := make(chan error, nSwarm*nConn*nStream*nMsg*100) // dont block anything.

	rateLimitN := 5000
	rateLimitChan := make(chan struct{}, rateLimitN) // max of 5k funcs.
	for i := 0; i < rateLimitN; i++ {
		rateLimitChan <- struct{}{}
	}

	rateLimit := func(f func()) {
		<-rateLimitChan
		f()
		rateLimitChan <- struct{}{}
	}

	writeStream := func(s *ps.Stream, bufs chan<- []byte) {
		log("writeStream %p, %d nMsg", s, nMsg)

		for i := 0; i < nMsg; i++ {
			buf := randBuf(msgsize)
			bufs <- buf
			log("%p writing %d bytes (message %d/%d #%x)", s, len(buf), i, nMsg, buf[:3])
			if _, err := s.Write(buf); err != nil {
				errs <- err
				continue
			}
		}
	}

	readStream := func(s *ps.Stream, bufs <-chan []byte) {
		log("readStream %p, %d nMsg", s, nMsg)

		buf2 := make([]byte, msgsize)
		i := 0
		for buf1 := range bufs {
			log("%p reading %d bytes (message %d/%d #%x)", s, len(buf1), i, nMsg, buf1[:3])
			i++

			if _, err := io.ReadFull(s, buf2); err != nil {
				errs <- err
				continue
			}
			if !bytes.Equal(buf1, buf2) {
				errs <- fmt.Errorf("buffers not equal (%x != %x)", buf1[:3], buf2[:3])
			}
		}
	}

	openStreamAndRW := func(c *ps.Conn) {
		log("openStreamAndRW %p, %d nMsg", c, nMsg)

		s, err := c.NewStream()
		if err != nil {
			errs <- fmt.Errorf("Failed to create NewStream: %s", err)
			return
		}

		bufs := make(chan []byte, nMsg)
		go func() {
			writeStream(s, bufs)
			close(bufs)
		}()

		readStream(s, bufs)
		s.Close()
	}

	openConnAndRW := func(a, b *ps.Swarm) {
		log("openConnAndRW %p -> %p, %d nStream", a, b, nConn)

		ls := b.Listeners()
		l := ls[mrand.Intn(len(ls))]
		nl := l.NetListener()
		nla := nl.Addr()

		nc, err := net.Dial(nla.Network(), nla.String())
		if err != nil {
			errs <- err
			return
		}

		c, err := a.AddConn(nc)
		if err != nil {
			errs <- err
			return
		}

		var wg sync.WaitGroup
		for i := 0; i < nStream; i++ {
			wg.Add(1)
			go rateLimit(func() {
				defer wg.Done()
				openStreamAndRW(c)
			})
		}
		wg.Wait()
		c.Close()
	}

	openConnsAndRW := func(a, b *ps.Swarm) {
		log("openConnsAndRW %p -> %p, %d conns", a, b, nConn)

		var wg sync.WaitGroup
		for i := 0; i < nConn; i++ {
			wg.Add(1)
			go rateLimit(func() {
				defer wg.Done()
				openConnAndRW(a, b)
			})
		}
		wg.Wait()
	}

	connectSwarmsAndRW := func(swarms []*ps.Swarm) {
		log("connectSwarmsAndRW %d swarms", len(swarms))

		var wg sync.WaitGroup
		for _, a := range swarms {
			for _, b := range swarms {
				wg.Add(1)
				a := a // race
				b := b // race
				go rateLimit(func() {
					defer wg.Done()
					openConnsAndRW(a, b)
				})
			}
		}
		wg.Wait()
	}

	swarms := makeSwarms(t, tr, nSwarm, 3) // 3 listeners per swarm.

	go func() {
		connectSwarmsAndRW(swarms)
		close(errs) // done
	}()

	for err := range errs {
		t.Error(err)
	}

}

func SubtestStress1Swarm1Conn1Stream1Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 1, 1, 1)
}

func SubtestStress1Swarm1Conn1Stream100Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 1, 1, 100)
}

func SubtestStress1Swarm1Conn100Stream100Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 1, 100, 100)
}

func SubtestStress1Swarm10Conn50Stream50Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 10, 50, 50)
}

func SubtestStress5Swarm2Conn20Stream20Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 5, 2, 20, 20)
}

func SubtestStress10Swarm2Conn100Stream100Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 10, 2, 100, 100)
}

func SubtestAll(t *testing.T, tr pst.Transport) {

	tests := []TransportTest{
		SubtestConstructSwarm,
		SubtestSimpleWrite,
		SubtestSimpleWrite100msgs,
		SubtestStress1Swarm1Conn1Stream1Msg,
		SubtestStress1Swarm1Conn1Stream100Msg,
		SubtestStress1Swarm1Conn100Stream100Msg,
		SubtestStress1Swarm10Conn50Stream50Msg,
		SubtestStress5Swarm2Conn20Stream20Msg,
		// SubtestStress10Swarm2Conn100Stream100Msg, <-- this hoses the osx network stack...
	}

	for _, f := range tests {
		if testing.Verbose() {
			fmt.Fprintf(os.Stderr, "==== RUN %s\n", GetFunctionName(f))
		}
		f(t, tr)
	}
}

type TransportTest func(t *testing.T, tr pst.Transport)

func TestNoOp(t *testing.T) {}

func GetFunctionName(i interface{}) string {
	return runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
}
<commit_msg>test: close swarms<commit_after>package peerstream_transport_test

import (
	"bytes"
	crand "crypto/rand"
	"fmt"
	"io"
	mrand "math/rand"
	"net"
	"os"
	"reflect"
	"runtime"
	"sync"
	"testing"

	ps "github.com/jbenet/go-peerstream"
	pst "github.com/jbenet/go-peerstream/transport"
)

var randomness []byte
var nextPort = 20000

func init() {
	// read 1MB of randomness
	randomness = make([]byte, 1<<20)
	if _, err := crand.Read(randomness); err != nil {
		panic(err)
	}
}

func randBuf(size int) []byte {
	n := len(randomness) - size
	if size < 1 {
		panic(fmt.Errorf("requested too large buffer (%d). max is %d", size, len(randomness)))
	}

	start := mrand.Intn(n)
	return randomness[start : start+size]
}

func checkErr(t *testing.T, err error) {
	if err != nil {
		t.Fatal(err)
	}
}

func log(s string, v ...interface{}) {
	if testing.Verbose() {
		fmt.Fprintf(os.Stderr, "> "+s+"\n", v...)
	}
}

type echoSetup struct {
	swarm *ps.Swarm
	conns []*ps.Conn
}

func singleConn(t *testing.T, tr pst.Transport) echoSetup {
	swarm := ps.NewSwarm(tr)
	swarm.SetStreamHandler(func(s *ps.Stream) {
		defer s.Close()
		log("accepted stream")
		io.Copy(s, s) // echo everything
		log("closing stream")
	})

	log("listening at %s", "localhost:0")
	l, err := net.Listen("tcp", "localhost:0")
	checkErr(t, err)

	_, err = swarm.AddListener(l)
	checkErr(t, err)

	log("dialing to %s", l.Addr())
	nc1, err := net.Dial("tcp", l.Addr().String())
	checkErr(t, err)

	c1, err := swarm.AddConn(nc1)
	checkErr(t, err)

	return echoSetup{
		swarm: swarm,
		conns: []*ps.Conn{c1},
	}
}

func makeSwarm(t *testing.T, tr pst.Transport, nListeners int) *ps.Swarm {
	swarm := ps.NewSwarm(tr)
	swarm.SetStreamHandler(func(s *ps.Stream) {
		defer s.Close()
		log("accepted stream")
		io.Copy(s, s) // echo everything
		log("closing stream")
	})

	for i := 0; i < nListeners; i++ {
		log("%p listening at %s", swarm, "localhost:0")
		l, err := net.Listen("tcp", "localhost:0")
		checkErr(t, err)
		_, err = swarm.AddListener(l)
		checkErr(t, err)
	}

	return swarm
}

func makeSwarms(t *testing.T, tr pst.Transport, nSwarms, nListeners int) []*ps.Swarm {
	swarms := make([]*ps.Swarm, nSwarms)
	for i := 0; i < nSwarms; i++ {
		swarms[i] = makeSwarm(t, tr, nListeners)
	}
	return swarms
}

func SubtestConstructSwarm(t *testing.T, tr pst.Transport) {
	ps.NewSwarm(tr)
}

func SubtestSimpleWrite(t *testing.T, tr pst.Transport) {
	swarm := ps.NewSwarm(tr)
	defer swarm.Close()

	piper, pipew := io.Pipe()
	swarm.SetStreamHandler(func(s *ps.Stream) {
		defer s.Close()
		log("accepted stream")
		w := io.MultiWriter(s, pipew)
		io.Copy(w, s) // echo everything and write it to pipew
		log("closing stream")
	})

	log("listening at %s", "localhost:0")
	l, err := net.Listen("tcp", "localhost:0")
	checkErr(t, err)

	_, err = swarm.AddListener(l)
	checkErr(t, err)

	log("dialing to %s", l.Addr().String())
	nc1, err := net.Dial("tcp", l.Addr().String())
	checkErr(t, err)

	c1, err := swarm.AddConn(nc1)
	checkErr(t, err)
	defer c1.Close()

	log("creating stream")
	s1, err := c1.NewStream()
	checkErr(t, err)
	defer s1.Close()

	buf1 := randBuf(4096)
	log("writing %d bytes to stream", len(buf1))
	_, err = s1.Write(buf1)
	checkErr(t, err)

	buf2 := make([]byte, len(buf1))
	log("reading %d bytes from stream (echoed)", len(buf2))
	_, err = s1.Read(buf2)
	checkErr(t, err)
	if string(buf2) != string(buf1) {
		t.Error("buf1 and buf2 not equal: %s != %s", string(buf1), string(buf2))
	}

	buf3 := make([]byte, len(buf1))
	log("reading %d bytes from pipe (tee)", len(buf3))
	_, err = piper.Read(buf3)
	checkErr(t, err)
	if string(buf3) != string(buf1) {
		t.Error("buf1 and buf3 not equal: %s != %s", string(buf1), string(buf3))
	}
}

func SubtestSimpleWrite100msgs(t *testing.T, tr pst.Transport) {

	msgs := 100
	msgsize := 1 << 19
	es := singleConn(t, tr)

	log("creating stream")
	stream, err := es.conns[0].NewStream()
	checkErr(t, err)

	bufs := make(chan []byte, msgs)
	errs := make(chan error, msgs*100)
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		defer wg.Done()

		for i := 0; i < msgs; i++ {
			buf := randBuf(msgsize)
			bufs <- buf
			log("writing %d bytes (message %d/%d #%x)", len(buf), i, msgs, buf[:3])
			if _, err := stream.Write(buf); err != nil {
				errs <- err
				continue
			}
		}
		close(bufs)
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		buf2 := make([]byte, msgsize)
		i := 0
		for buf1 := range bufs {
			log("reading %d bytes (message %d/%d #%x)", len(buf1), i, msgs, buf1[:3])
			i++

			if _, err := io.ReadFull(stream, buf2); err != nil {
				errs <- err
				continue
			}
			if !bytes.Equal(buf1, buf2) {
				errs <- fmt.Errorf("buffers not equal (%x != %x)", buf1[:3], buf2[:3])
			}
		}
	}()

	wg.Wait()
	close(errs)
	for err := range errs {
		t.Error(err)
	}
}

func SubtestStressNSwarmNConnNStreamNMsg(t *testing.T, tr pst.Transport, nSwarm, nConn, nStream, nMsg int) {

	msgsize := 1 << 11
	errs := make(chan error, nSwarm*nConn*nStream*nMsg*100) // dont block anything.

	rateLimitN := 5000
	rateLimitChan := make(chan struct{}, rateLimitN) // max of 5k funcs.
	for i := 0; i < rateLimitN; i++ {
		rateLimitChan <- struct{}{}
	}

	rateLimit := func(f func()) {
		<-rateLimitChan
		f()
		rateLimitChan <- struct{}{}
	}

	writeStream := func(s *ps.Stream, bufs chan<- []byte) {
		log("writeStream %p, %d nMsg", s, nMsg)

		for i := 0; i < nMsg; i++ {
			buf := randBuf(msgsize)
			bufs <- buf
			log("%p writing %d bytes (message %d/%d #%x)", s, len(buf), i, nMsg, buf[:3])
			if _, err := s.Write(buf); err != nil {
				errs <- err
				continue
			}
		}
	}

	readStream := func(s *ps.Stream, bufs <-chan []byte) {
		log("readStream %p, %d nMsg", s, nMsg)

		buf2 := make([]byte, msgsize)
		i := 0
		for buf1 := range bufs {
			log("%p reading %d bytes (message %d/%d #%x)", s, len(buf1), i, nMsg, buf1[:3])
			i++

			if _, err := io.ReadFull(s, buf2); err != nil {
				errs <- err
				continue
			}
			if !bytes.Equal(buf1, buf2) {
				errs <- fmt.Errorf("buffers not equal (%x != %x)", buf1[:3], buf2[:3])
			}
		}
	}

	openStreamAndRW := func(c *ps.Conn) {
		log("openStreamAndRW %p, %d nMsg", c, nMsg)

		s, err := c.NewStream()
		if err != nil {
			errs <- fmt.Errorf("Failed to create NewStream: %s", err)
			return
		}

		bufs := make(chan []byte, nMsg)
		go func() {
			writeStream(s, bufs)
			close(bufs)
		}()

		readStream(s, bufs)
		s.Close()
	}

	openConnAndRW := func(a, b *ps.Swarm) {
		log("openConnAndRW %p -> %p, %d nStream", a, b, nConn)

		ls := b.Listeners()
		l := ls[mrand.Intn(len(ls))]
		nl := l.NetListener()
		nla := nl.Addr()

		nc, err := net.Dial(nla.Network(), nla.String())
		if err != nil {
			errs <- err
			return
		}

		c, err := a.AddConn(nc)
		if err != nil {
			errs <- err
			return
		}

		var wg sync.WaitGroup
		for i := 0; i < nStream; i++ {
			wg.Add(1)
			go rateLimit(func() {
				defer wg.Done()
				openStreamAndRW(c)
			})
		}
		wg.Wait()
		c.Close()
	}

	openConnsAndRW := func(a, b *ps.Swarm) {
		log("openConnsAndRW %p -> %p, %d conns", a, b, nConn)

		var wg sync.WaitGroup
		for i := 0; i < nConn; i++ {
			wg.Add(1)
			go rateLimit(func() {
				defer wg.Done()
				openConnAndRW(a, b)
			})
		}
		wg.Wait()
	}

	connectSwarmsAndRW := func(swarms []*ps.Swarm) {
		log("connectSwarmsAndRW %d swarms", len(swarms))

		var wg sync.WaitGroup
		for _, a := range swarms {
			for _, b := range swarms {
				wg.Add(1)
				a := a // race
				b := b // race
				go rateLimit(func() {
					defer wg.Done()
					openConnsAndRW(a, b)
				})
			}
		}
		wg.Wait()
	}

	swarms := makeSwarms(t, tr, nSwarm, 3) // 3 listeners per swarm.

	go func() {
		connectSwarmsAndRW(swarms)
		close(errs) // done
	}()

	for err := range errs {
		t.Error(err)
	}

	for _, s := range swarms {
		s.Close()
	}

}

func SubtestStress1Swarm1Conn1Stream1Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 1, 1, 1)
}

func SubtestStress1Swarm1Conn1Stream100Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 1, 1, 100)
}

func SubtestStress1Swarm1Conn100Stream100Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 1, 100, 100)
}

func SubtestStress1Swarm10Conn50Stream50Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 1, 10, 50, 50)
}

func SubtestStress5Swarm2Conn20Stream20Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 5, 2, 20, 20)
}

func SubtestStress10Swarm2Conn100Stream100Msg(t *testing.T, tr pst.Transport) {
	SubtestStressNSwarmNConnNStreamNMsg(t, tr, 10, 2, 100, 100)
}

func SubtestAll(t *testing.T, tr pst.Transport) {

	tests := []TransportTest{
		SubtestConstructSwarm,
		SubtestSimpleWrite,
		SubtestSimpleWrite100msgs,
		SubtestStress1Swarm1Conn1Stream1Msg,
		SubtestStress1Swarm1Conn1Stream100Msg,
		SubtestStress1Swarm1Conn100Stream100Msg,
		SubtestStress1Swarm10Conn50Stream50Msg,
		SubtestStress5Swarm2Conn20Stream20Msg,
		// SubtestStress10Swarm2Conn100Stream100Msg, <-- this hoses the osx network stack...
	}

	for _, f := range tests {
		if testing.Verbose() {
			fmt.Fprintf(os.Stderr, "==== RUN %s\n", GetFunctionName(f))
		}
		f(t, tr)
	}
}

type TransportTest func(t *testing.T, tr pst.Transport)

func TestNoOp(t *testing.T) {}

func GetFunctionName(i interface{}) string {
	return runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
}
<|endoftext|>


<commit_before>// Google Storage utility that contains methods for both CT master and worker
// scripts.
package util

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/skia-dev/glog"

	storage "code.google.com/p/google-api-go-client/storage/v1"
	"skia.googlesource.com/buildbot.git/go/auth"
	"skia.googlesource.com/buildbot.git/go/gs"
)

const (
	GOROUTINE_POOL_SIZE = 50
	MAX_CHANNEL_SIZE    = 100000
)

type GsUtil struct {
	// The client used to connect to Google Storage.
	client  *http.Client
	service *storage.Service
}

// NewGsUtil initializes and returns a utility for CT interations with Google
// Storage. If client is nil then auth.RunFlow is invoked. if client is nil then
// the client from GetOAuthClient is used.
func NewGsUtil(client *http.Client) (*GsUtil, error) {
	if client == nil {
		oauthClient, err := GetOAuthClient()
		if err != nil {
			return nil, err
		}
		client = oauthClient
	}
	service, err := storage.New(client)
	if err != nil {
		return nil, fmt.Errorf("Failed to create interface to Google Storage: %s", err)
	}
	return &GsUtil{client: client, service: service}, nil
}

func GetOAuthClient() (*http.Client, error) {
	config := auth.OAuthConfig(GSTokenPath, auth.SCOPE_READ_WRITE)
	return auth.RunFlow(config)
}

// Returns the response body of the specified GS object. Tries MAX_URI_GET_TRIES
// times if download is unsuccessful. Client must close the response body when
// finished with it.
func getRespBody(res *storage.Object, client *http.Client) (io.ReadCloser, error) {
	for i := 0; i < MAX_URI_GET_TRIES; i++ {
		glog.Infof("Fetching: %s", res.Name)
		request, err := gs.RequestForStorageURL(res.MediaLink)
		if err != nil {
			glog.Warningf("Unable to create Storage MediaURI request: %s\n", err)
			continue
		}

		resp, err := client.Do(request)
		if err != nil {
			glog.Warningf("Unable to retrieve Storage MediaURI: %s", err)
			continue
		}
		if resp.StatusCode != 200 {
			glog.Warningf("Failed to retrieve: %d  %s", resp.StatusCode, resp.Status)
			resp.Body.Close()
			continue
		}
		return resp.Body, nil
	}
	return nil, fmt.Errorf("Failed fetching file after %d attempts", MAX_URI_GET_TRIES)
}

// Returns the response body of the specified GS file. Client must close the
// response body when finished with it.
func (gs *GsUtil) GetRemoteFileContents(filePath string) (io.ReadCloser, error) {
	res, err := gs.service.Objects.Get(GS_BUCKET_NAME, filePath).Do()
	if err != nil {
		return nil, fmt.Errorf("Could not get %s from GS: %s", filePath, err)
	}
	return getRespBody(res, gs.client)
}

// AreTimeStampsEqual checks whether the TIMESTAMP in the local dir matches the
// TIMESTAMP in the remote Google Storage dir.
func (gs *GsUtil) AreTimeStampsEqual(localDir, gsDir string) (bool, error) {
	// Get timestamp from the local directory.
	localTimestampPath := filepath.Join(localDir, TIMESTAMP_FILE_NAME)
	fileContent, err := ioutil.ReadFile(localTimestampPath)
	if err != nil {
		return false, fmt.Errorf("Could not read %s: %s", localTimestampPath, err)
	}
	localTimestamp := strings.Trim(string(fileContent), "\n")

	// Get timestamp from the Google Storage directory.
	gsTimestampPath := filepath.Join(gsDir, TIMESTAMP_FILE_NAME)
	respBody, err := gs.GetRemoteFileContents(gsTimestampPath)
	if err != nil {
		return false, err
	}
	defer respBody.Close()
	resp, err := ioutil.ReadAll(respBody)
	if err != nil {
		return false, err
	}
	gsTimestamp := strings.Trim(string(resp), "\n")

	// Return the comparison of the two timestamps.
	return localTimestamp == gsTimestamp, nil
}

type filePathToStorageObject struct {
	storageObject *storage.Object
	filePath      string
}

// downloadRemoteDir downloads the specified Google Storage dir to the specified
// local dir. The local dir will be emptied and recreated. Handles multiple levels
// of directories.
func (gs *GsUtil) downloadRemoteDir(localDir, gsDir string) error {
	// Empty the local dir.
	os.RemoveAll(localDir)
	// Create the local dir.
	os.MkdirAll(localDir, 0700)
	// The channel where the storage objects to be deleted will be sent to.
	chStorageObjects := make(chan filePathToStorageObject, MAX_CHANNEL_SIZE)
	req := gs.service.Objects.List(GS_BUCKET_NAME).Prefix(gsDir + "/")
	for req != nil {
		resp, err := req.Do()
		if err != nil {
			return fmt.Errorf("Error occured while listing %s: %s", gsDir, err)
		}
		for _, result := range resp.Items {
			fileName := filepath.Base(result.Name)
			// If downloading from subdir then add it to the fileName.
			fileGsDir := filepath.Dir(result.Name)
			subDirs := strings.TrimPrefix(fileGsDir, gsDir)
			if subDirs != "" {
				dirTokens := strings.Split(subDirs, "/")
				for i := range dirTokens {
					fileName = filepath.Join(dirTokens[len(dirTokens)-i-1], fileName)
				}
				// Create the local directory.
				os.MkdirAll(filepath.Join(localDir, filepath.Dir(fileName)), 0700)
			}
			chStorageObjects <- filePathToStorageObject{storageObject: result, filePath: fileName}
		}
		if len(resp.NextPageToken) > 0 {
			req.PageToken(resp.NextPageToken)
		} else {
			req = nil
		}
	}
	close(chStorageObjects)

	// Kick off goroutines to download the storage objects.
	var wg sync.WaitGroup
	for i := 0; i < GOROUTINE_POOL_SIZE; i++ {
		wg.Add(1)
		go func(goroutineNum int) {
			defer wg.Done()
			for obj := range chStorageObjects {
				result := obj.storageObject
				filePath := obj.filePath
				respBody, err := getRespBody(result, gs.client)
				if err != nil {
					glog.Errorf("Could not fetch %s: %s", result.MediaLink, err)
					return
				}
				defer respBody.Close()
				outputFile := filepath.Join(localDir, filePath)
				out, err := os.Create(outputFile)
				if err != nil {
					glog.Errorf("Unable to create file %s: %s", outputFile, err)
					return
				}
				defer out.Close()
				if _, err = io.Copy(out, respBody); err != nil {
					glog.Error(err)
					return
				}
				glog.Infof("Downloaded gs://%s/%s to %s with goroutine#%d", GS_BUCKET_NAME, result.Name, outputFile, goroutineNum)
			}
		}(i + 1)
	}
	wg.Wait()
	return nil
}

// DownloadChromiumBuild downloads the specified Chromium build from Google
// Storage to a local dir.
func (gs *GsUtil) DownloadChromiumBuild(chromiumBuild string) error {
	localDir := filepath.Join(ChromiumBuildsDir, chromiumBuild)
	gsDir := filepath.Join(CHROMIUM_BUILDS_DIR_NAME, chromiumBuild)
	if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
		glog.Infof("Not downloading %s because TIMESTAMPS match", gsDir)
		return nil
	}
	glog.Infof("Timestamps between %s and %s are different. Downloading from Google Storage", localDir, gsDir)
	if err := gs.downloadRemoteDir(localDir, gsDir); err != nil {
		return fmt.Errorf("Error downloading %s into %s: %s", gsDir, localDir, err)
	}
	// Downloaded chrome binary needs to be set as an executable.
	os.Chmod(filepath.Join(localDir, "chrome"), 0777)

	return nil
}

// DownloadWorkerArtifacts downloads artifacts from Google Storage to a local dir.
func (gs *GsUtil) DownloadWorkerArtifacts(dirName, pagesetType string, workerNum int) error {
	localDir := filepath.Join(StorageDir, dirName, pagesetType)
	gsDir := filepath.Join(dirName, pagesetType, fmt.Sprintf("slave%d", workerNum))

	if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
		// No need to download artifacts they already exist locally.
		glog.Infof("Not downloading %s because TIMESTAMPS match", gsDir)
		return nil
	}
	glog.Infof("Timestamps between %s and %s are different. Downloading from Google Storage", localDir, gsDir)
	return gs.downloadRemoteDir(localDir, gsDir)
}

func (gs *GsUtil) deleteRemoteDir(gsDir string) error {
	// The channel where the GS filepaths to be deleted will be sent to.
	chFilePaths := make(chan string, MAX_CHANNEL_SIZE)
	req := gs.service.Objects.List(GS_BUCKET_NAME).Prefix(gsDir + "/")
	for req != nil {
		resp, err := req.Do()
		if err != nil {
			return fmt.Errorf("Error occured while listing %s: %s", gsDir, err)
		}
		for _, result := range resp.Items {
			chFilePaths <- result.Name
		}
		if len(resp.NextPageToken) > 0 {
			req.PageToken(resp.NextPageToken)
		} else {
			req = nil
		}
	}
	close(chFilePaths)

	// Kick off goroutines to delete the file paths.
	var wg sync.WaitGroup
	for i := 0; i < GOROUTINE_POOL_SIZE; i++ {
		wg.Add(1)
		go func(goroutineNum int) {
			defer wg.Done()
			for filePath := range chFilePaths {
				if err := gs.service.Objects.Delete(GS_BUCKET_NAME, filePath).Do(); err != nil {
					glog.Errorf("Goroutine#%d could not delete %s: %s", goroutineNum, filePath, err)
					return
				}
				glog.Infof("Deleted gs://%s/%s with goroutine#%d", GS_BUCKET_NAME, filePath, goroutineNum)
			}
		}(i + 1)
	}
	wg.Wait()
	return nil
}

// UploadFile uploads the specified file to the remote dir in Google Storage.
func (gs *GsUtil) UploadFile(fileName, localDir, gsDir string) error {
	localFile := filepath.Join(localDir, fileName)
	gsFile := filepath.Join(gsDir, fileName)
	object := &storage.Object{Name: gsFile}
	f, err := os.Open(localFile)
	if err != nil {
		return fmt.Errorf("Error opening %s: %s", localFile, err)
	}
	defer f.Close()
	if _, err := gs.service.Objects.Insert(GS_BUCKET_NAME, object).Media(f).Do(); err != nil {
		return fmt.Errorf("Objects.Insert failed: %s", err)
	}
	glog.Infof("Copied %s to %s", localFile, fmt.Sprintf("gs://%s/%s", GS_BUCKET_NAME, gsFile))
	return nil
}

// UploadWorkerArtifacts uploads artifacts from a local dir to Google Storage.
func (gs *GsUtil) UploadWorkerArtifacts(dirName, pagesetType string, workerNum int) error {
	localDir := filepath.Join(StorageDir, dirName, pagesetType)
	gsDir := filepath.Join(dirName, pagesetType, fmt.Sprintf("slave%d", workerNum))

	if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
		glog.Infof("Not uploading %s because TIMESTAMPS match", localDir)
		return nil
	}
	glog.Infof("Timestamps between %s and %s are different. Uploading to Google Storage", localDir, gsDir)
	return gs.UploadDir(localDir, gsDir)
}

// UploadDir uploads the specified local dir into the specified Google Storage dir.
func (gs *GsUtil) UploadDir(localDir, gsDir string) error {
	// Empty the remote dir.
	gs.deleteRemoteDir(gsDir)

	// Construct a dictionary of file paths to their file infos.
	pathsToFileInfos := map[string]os.FileInfo{}
	visit := func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			return nil
		}
		pathsToFileInfos[path] = f
		return nil
	}
	if err := filepath.Walk(localDir, visit); err != nil {
		return fmt.Errorf("Unable to read the local dir %s: %s", localDir, err)
	}

	// The channel where the filepaths to be uploaded will be sent to.
	chFilePaths := make(chan string, MAX_CHANNEL_SIZE)
	// File filepaths and send it to the above channel.
	for path, fileInfo := range pathsToFileInfos {
		fileName := fileInfo.Name()
		containingDir := strings.TrimSuffix(path, fileName)
		subDirs := strings.TrimPrefix(containingDir, localDir)
		if subDirs != "" {
			dirTokens := strings.Split(subDirs, "/")
			for i := range dirTokens {
				fileName = filepath.Join(dirTokens[len(dirTokens)-i-1], fileName)
			}
		}
		chFilePaths <- fileName
	}
	close(chFilePaths)

	// Kick off goroutines to upload the file paths.
	var wg sync.WaitGroup
	for i := 0; i < GOROUTINE_POOL_SIZE; i++ {
		wg.Add(1)
		go func(goroutineNum int) {
			defer wg.Done()
			for filePath := range chFilePaths {
				glog.Infof("Uploading %s to %s with goroutine#%d", filePath, gsDir, goroutineNum)
				if err := gs.UploadFile(filePath, localDir, gsDir); err != nil {
					glog.Errorf("Goroutine#%d could not upload %s to %s: %s", goroutineNum, filePath, localDir, err)
				}
			}
		}(i + 1)
	}
	wg.Wait()
	return nil
}
<commit_msg>All uploaded files must be readable by google.com accounts.<commit_after>// Google Storage utility that contains methods for both CT master and worker
// scripts.
package util

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/skia-dev/glog"

	storage "code.google.com/p/google-api-go-client/storage/v1"
	"skia.googlesource.com/buildbot.git/go/auth"
	"skia.googlesource.com/buildbot.git/go/gs"
)

const (
	GOROUTINE_POOL_SIZE = 50
	MAX_CHANNEL_SIZE    = 100000
)

type GsUtil struct {
	// The client used to connect to Google Storage.
	client  *http.Client
	service *storage.Service
}

// NewGsUtil initializes and returns a utility for CT interations with Google
// Storage. If client is nil then auth.RunFlow is invoked. if client is nil then
// the client from GetOAuthClient is used.
func NewGsUtil(client *http.Client) (*GsUtil, error) {
	if client == nil {
		oauthClient, err := GetOAuthClient()
		if err != nil {
			return nil, err
		}
		client = oauthClient
	}
	service, err := storage.New(client)
	if err != nil {
		return nil, fmt.Errorf("Failed to create interface to Google Storage: %s", err)
	}
	return &GsUtil{client: client, service: service}, nil
}

func GetOAuthClient() (*http.Client, error) {
	config := auth.OAuthConfig(GSTokenPath, auth.SCOPE_FULL_CONTROL)
	return auth.RunFlow(config)
}

// Returns the response body of the specified GS object. Tries MAX_URI_GET_TRIES
// times if download is unsuccessful. Client must close the response body when
// finished with it.
func getRespBody(res *storage.Object, client *http.Client) (io.ReadCloser, error) {
	for i := 0; i < MAX_URI_GET_TRIES; i++ {
		glog.Infof("Fetching: %s", res.Name)
		request, err := gs.RequestForStorageURL(res.MediaLink)
		if err != nil {
			glog.Warningf("Unable to create Storage MediaURI request: %s\n", err)
			continue
		}

		resp, err := client.Do(request)
		if err != nil {
			glog.Warningf("Unable to retrieve Storage MediaURI: %s", err)
			continue
		}
		if resp.StatusCode != 200 {
			glog.Warningf("Failed to retrieve: %d  %s", resp.StatusCode, resp.Status)
			resp.Body.Close()
			continue
		}
		return resp.Body, nil
	}
	return nil, fmt.Errorf("Failed fetching file after %d attempts", MAX_URI_GET_TRIES)
}

// Returns the response body of the specified GS file. Client must close the
// response body when finished with it.
func (gs *GsUtil) GetRemoteFileContents(filePath string) (io.ReadCloser, error) {
	res, err := gs.service.Objects.Get(GS_BUCKET_NAME, filePath).Do()
	if err != nil {
		return nil, fmt.Errorf("Could not get %s from GS: %s", filePath, err)
	}
	return getRespBody(res, gs.client)
}

// AreTimeStampsEqual checks whether the TIMESTAMP in the local dir matches the
// TIMESTAMP in the remote Google Storage dir.
func (gs *GsUtil) AreTimeStampsEqual(localDir, gsDir string) (bool, error) {
	// Get timestamp from the local directory.
	localTimestampPath := filepath.Join(localDir, TIMESTAMP_FILE_NAME)
	fileContent, err := ioutil.ReadFile(localTimestampPath)
	if err != nil {
		return false, fmt.Errorf("Could not read %s: %s", localTimestampPath, err)
	}
	localTimestamp := strings.Trim(string(fileContent), "\n")

	// Get timestamp from the Google Storage directory.
	gsTimestampPath := filepath.Join(gsDir, TIMESTAMP_FILE_NAME)
	respBody, err := gs.GetRemoteFileContents(gsTimestampPath)
	if err != nil {
		return false, err
	}
	defer respBody.Close()
	resp, err := ioutil.ReadAll(respBody)
	if err != nil {
		return false, err
	}
	gsTimestamp := strings.Trim(string(resp), "\n")

	// Return the comparison of the two timestamps.
	return localTimestamp == gsTimestamp, nil
}

type filePathToStorageObject struct {
	storageObject *storage.Object
	filePath      string
}

// downloadRemoteDir downloads the specified Google Storage dir to the specified
// local dir. The local dir will be emptied and recreated. Handles multiple levels
// of directories.
func (gs *GsUtil) downloadRemoteDir(localDir, gsDir string) error {
	// Empty the local dir.
	os.RemoveAll(localDir)
	// Create the local dir.
	os.MkdirAll(localDir, 0700)
	// The channel where the storage objects to be deleted will be sent to.
	chStorageObjects := make(chan filePathToStorageObject, MAX_CHANNEL_SIZE)
	req := gs.service.Objects.List(GS_BUCKET_NAME).Prefix(gsDir + "/")
	for req != nil {
		resp, err := req.Do()
		if err != nil {
			return fmt.Errorf("Error occured while listing %s: %s", gsDir, err)
		}
		for _, result := range resp.Items {
			fileName := filepath.Base(result.Name)
			// If downloading from subdir then add it to the fileName.
			fileGsDir := filepath.Dir(result.Name)
			subDirs := strings.TrimPrefix(fileGsDir, gsDir)
			if subDirs != "" {
				dirTokens := strings.Split(subDirs, "/")
				for i := range dirTokens {
					fileName = filepath.Join(dirTokens[len(dirTokens)-i-1], fileName)
				}
				// Create the local directory.
				os.MkdirAll(filepath.Join(localDir, filepath.Dir(fileName)), 0700)
			}
			chStorageObjects <- filePathToStorageObject{storageObject: result, filePath: fileName}
		}
		if len(resp.NextPageToken) > 0 {
			req.PageToken(resp.NextPageToken)
		} else {
			req = nil
		}
	}
	close(chStorageObjects)

	// Kick off goroutines to download the storage objects.
	var wg sync.WaitGroup
	for i := 0; i < GOROUTINE_POOL_SIZE; i++ {
		wg.Add(1)
		go func(goroutineNum int) {
			defer wg.Done()
			for obj := range chStorageObjects {
				result := obj.storageObject
				filePath := obj.filePath
				respBody, err := getRespBody(result, gs.client)
				if err != nil {
					glog.Errorf("Could not fetch %s: %s", result.MediaLink, err)
					return
				}
				defer respBody.Close()
				outputFile := filepath.Join(localDir, filePath)
				out, err := os.Create(outputFile)
				if err != nil {
					glog.Errorf("Unable to create file %s: %s", outputFile, err)
					return
				}
				defer out.Close()
				if _, err = io.Copy(out, respBody); err != nil {
					glog.Error(err)
					return
				}
				glog.Infof("Downloaded gs://%s/%s to %s with goroutine#%d", GS_BUCKET_NAME, result.Name, outputFile, goroutineNum)
			}
		}(i + 1)
	}
	wg.Wait()
	return nil
}

// DownloadChromiumBuild downloads the specified Chromium build from Google
// Storage to a local dir.
func (gs *GsUtil) DownloadChromiumBuild(chromiumBuild string) error {
	localDir := filepath.Join(ChromiumBuildsDir, chromiumBuild)
	gsDir := filepath.Join(CHROMIUM_BUILDS_DIR_NAME, chromiumBuild)
	if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
		glog.Infof("Not downloading %s because TIMESTAMPS match", gsDir)
		return nil
	}
	glog.Infof("Timestamps between %s and %s are different. Downloading from Google Storage", localDir, gsDir)
	if err := gs.downloadRemoteDir(localDir, gsDir); err != nil {
		return fmt.Errorf("Error downloading %s into %s: %s", gsDir, localDir, err)
	}
	// Downloaded chrome binary needs to be set as an executable.
	os.Chmod(filepath.Join(localDir, "chrome"), 0777)

	return nil
}

// DownloadWorkerArtifacts downloads artifacts from Google Storage to a local dir.
func (gs *GsUtil) DownloadWorkerArtifacts(dirName, pagesetType string, workerNum int) error {
	localDir := filepath.Join(StorageDir, dirName, pagesetType)
	gsDir := filepath.Join(dirName, pagesetType, fmt.Sprintf("slave%d", workerNum))

	if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
		// No need to download artifacts they already exist locally.
		glog.Infof("Not downloading %s because TIMESTAMPS match", gsDir)
		return nil
	}
	glog.Infof("Timestamps between %s and %s are different. Downloading from Google Storage", localDir, gsDir)
	return gs.downloadRemoteDir(localDir, gsDir)
}

func (gs *GsUtil) deleteRemoteDir(gsDir string) error {
	// The channel where the GS filepaths to be deleted will be sent to.
	chFilePaths := make(chan string, MAX_CHANNEL_SIZE)
	req := gs.service.Objects.List(GS_BUCKET_NAME).Prefix(gsDir + "/")
	for req != nil {
		resp, err := req.Do()
		if err != nil {
			return fmt.Errorf("Error occured while listing %s: %s", gsDir, err)
		}
		for _, result := range resp.Items {
			chFilePaths <- result.Name
		}
		if len(resp.NextPageToken) > 0 {
			req.PageToken(resp.NextPageToken)
		} else {
			req = nil
		}
	}
	close(chFilePaths)

	// Kick off goroutines to delete the file paths.
	var wg sync.WaitGroup
	for i := 0; i < GOROUTINE_POOL_SIZE; i++ {
		wg.Add(1)
		go func(goroutineNum int) {
			defer wg.Done()
			for filePath := range chFilePaths {
				if err := gs.service.Objects.Delete(GS_BUCKET_NAME, filePath).Do(); err != nil {
					glog.Errorf("Goroutine#%d could not delete %s: %s", goroutineNum, filePath, err)
					return
				}
				glog.Infof("Deleted gs://%s/%s with goroutine#%d", GS_BUCKET_NAME, filePath, goroutineNum)
			}
		}(i + 1)
	}
	wg.Wait()
	return nil
}

// UploadFile uploads the specified file to the remote dir in Google Storage. It
// also sets the appropriate ACLs on the uploaded file.
func (gs *GsUtil) UploadFile(fileName, localDir, gsDir string) error {
	localFile := filepath.Join(localDir, fileName)
	gsFile := filepath.Join(gsDir, fileName)
	object := &storage.Object{Name: gsFile}
	f, err := os.Open(localFile)
	if err != nil {
		return fmt.Errorf("Error opening %s: %s", localFile, err)
	}
	defer f.Close()
	if _, err := gs.service.Objects.Insert(GS_BUCKET_NAME, object).Media(f).Do(); err != nil {
		return fmt.Errorf("Objects.Insert failed: %s", err)
	}
	glog.Infof("Copied %s to %s", localFile, fmt.Sprintf("gs://%s/%s", GS_BUCKET_NAME, gsFile))

	// All objects uploaded to CT's bucket via this util must be readable by
	// the google.com domain. This will be fine tuned later if required.
	objectAcl := &storage.ObjectAccessControl{
		Bucket: GS_BUCKET_NAME, Entity: "domain-google.com", Object: gsFile, Role: "READER",
	}
	if _, err := gs.service.ObjectAccessControls.Insert(GS_BUCKET_NAME, gsFile, objectAcl).Do(); err != nil {
		return fmt.Errorf("Could not update ACL of %s: %s", object.Name, err)
	}
	glog.Infof("Updated ACL of %s", fmt.Sprintf("gs://%s/%s", GS_BUCKET_NAME, gsFile))

	return nil
}

// UploadWorkerArtifacts uploads artifacts from a local dir to Google Storage.
func (gs *GsUtil) UploadWorkerArtifacts(dirName, pagesetType string, workerNum int) error {
	localDir := filepath.Join(StorageDir, dirName, pagesetType)
	gsDir := filepath.Join(dirName, pagesetType, fmt.Sprintf("slave%d", workerNum))

	if equal, _ := gs.AreTimeStampsEqual(localDir, gsDir); equal {
		glog.Infof("Not uploading %s because TIMESTAMPS match", localDir)
		return nil
	}
	glog.Infof("Timestamps between %s and %s are different. Uploading to Google Storage", localDir, gsDir)
	return gs.UploadDir(localDir, gsDir)
}

// UploadDir uploads the specified local dir into the specified Google Storage dir.
func (gs *GsUtil) UploadDir(localDir, gsDir string) error {
	// Empty the remote dir.
	gs.deleteRemoteDir(gsDir)

	// Construct a dictionary of file paths to their file infos.
	pathsToFileInfos := map[string]os.FileInfo{}
	visit := func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			return nil
		}
		pathsToFileInfos[path] = f
		return nil
	}
	if err := filepath.Walk(localDir, visit); err != nil {
		return fmt.Errorf("Unable to read the local dir %s: %s", localDir, err)
	}

	// The channel where the filepaths to be uploaded will be sent to.
	chFilePaths := make(chan string, MAX_CHANNEL_SIZE)
	// File filepaths and send it to the above channel.
	for path, fileInfo := range pathsToFileInfos {
		fileName := fileInfo.Name()
		containingDir := strings.TrimSuffix(path, fileName)
		subDirs := strings.TrimPrefix(containingDir, localDir)
		if subDirs != "" {
			dirTokens := strings.Split(subDirs, "/")
			for i := range dirTokens {
				fileName = filepath.Join(dirTokens[len(dirTokens)-i-1], fileName)
			}
		}
		chFilePaths <- fileName
	}
	close(chFilePaths)

	// Kick off goroutines to upload the file paths.
	var wg sync.WaitGroup
	for i := 0; i < GOROUTINE_POOL_SIZE; i++ {
		wg.Add(1)
		go func(goroutineNum int) {
			defer wg.Done()
			for filePath := range chFilePaths {
				glog.Infof("Uploading %s to %s with goroutine#%d", filePath, gsDir, goroutineNum)
				if err := gs.UploadFile(filePath, localDir, gsDir); err != nil {
					glog.Errorf("Goroutine#%d could not upload %s to %s: %s", goroutineNum, filePath, localDir, err)
				}
			}
		}(i + 1)
	}
	wg.Wait()
	return nil
}
<|endoftext|>


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

// Package env makes it possible to track use of environment variables within a procress
// in order to generate documentation for these uses.
package env

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"strconv"
	"sync"
	"time"

	"istio.io/pkg/log"
)

// The type of a variable's value
type VarType byte

const (
	// Variable holds a free-form string.
	STRING VarType = iota
	// Variable holds a boolean value.
	BOOL
	// Variable holds a signed integer.
	INT
	// Variables holds a floating point value.
	FLOAT
	// Variable holds a time duration.
	DURATION
	// Variable holds a dynamic unknown type.
	OTHER
)

// Var describes a single environment variable
type Var struct {
	// The name of the environment variable.
	Name string

	// The optional default value of the environment variable.
	DefaultValue string

	// Description of the environment variable's purpose.
	Description string

	// Hide the existence of this variable when outputting usage information.
	Hidden bool

	// Mark this variable as deprecated when generating usage information.
	Deprecated bool

	// The type of the variable's value
	Type VarType

	// The underlying Go type of the variable
	GoType string
}

// StringVar represents a single string environment variable.
type StringVar struct {
	Var
}

// BoolVar represents a single boolean environment variable.
type BoolVar struct {
	Var
}

// IntVar represents a single integer environment variable.
type IntVar struct {
	Var
}

// FloatVar represents a single floating-point environment variable.
type FloatVar struct {
	Var
}

// DurationVar represents a single duration environment variable.
type DurationVar struct {
	Var
}

var (
	allVars = make(map[string]Var)
	mutex   sync.Mutex
)

// VarDescriptions returns a description of this process' environment variables, sorted by name.
func VarDescriptions() []Var {
	mutex.Lock()
	sorted := make([]Var, 0, len(allVars))
	for _, v := range allVars {
		sorted = append(sorted, v)
	}
	mutex.Unlock()

	sort.Slice(sorted, func(i, j int) bool {
		return sorted[i].Name < sorted[j].Name
	})

	return sorted
}

type Parseable interface {
	comparable
}

type GenericVar[T Parseable] struct {
	Var
	delegate specializedVar[T]
}

func Register[T Parseable](name string, defaultValue T, description string) GenericVar[T] {
	// Specialized cases
	// In the future, once only Register() remains, we can likely drop most of these.
	// however, time.Duration is needed still as it doesn't implement json
	switch d := any(defaultValue).(type) {
	case time.Duration:
		v := RegisterDurationVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case string:
		v := RegisterStringVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case float64:
		v := RegisterFloatVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case int:
		v := RegisterIntVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case bool:
		v := RegisterBoolVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	}
	b, _ := json.Marshal(defaultValue)
	v := Var{Name: name, DefaultValue: string(b), Description: description, Type: STRING, GoType: fmt.Sprintf("%T", defaultValue)}
	RegisterVar(v)
	return GenericVar[T]{getVar(name), nil}
}

// RegisterStringVar registers a new string environment variable.
func RegisterStringVar(name string, defaultValue string, description string) StringVar {
	v := Var{Name: name, DefaultValue: defaultValue, Description: description, Type: STRING}
	RegisterVar(v)
	return StringVar{getVar(name)}
}

// RegisterBoolVar registers a new boolean environment variable.
func RegisterBoolVar(name string, defaultValue bool, description string) BoolVar {
	v := Var{Name: name, DefaultValue: strconv.FormatBool(defaultValue), Description: description, Type: BOOL}
	RegisterVar(v)
	return BoolVar{getVar(name)}
}

// RegisterIntVar registers a new integer environment variable.
func RegisterIntVar(name string, defaultValue int, description string) IntVar {
	v := Var{Name: name, DefaultValue: strconv.FormatInt(int64(defaultValue), 10), Description: description, Type: INT}
	RegisterVar(v)
	return IntVar{getVar(name)}
}

// RegisterFloatVar registers a new floating-point environment variable.
func RegisterFloatVar(name string, defaultValue float64, description string) FloatVar {
	v := Var{Name: name, DefaultValue: strconv.FormatFloat(defaultValue, 'G', -1, 64), Description: description, Type: FLOAT}
	RegisterVar(v)
	return FloatVar{v}
}

// RegisterDurationVar registers a new duration environment variable.
func RegisterDurationVar(name string, defaultValue time.Duration, description string) DurationVar {
	v := Var{Name: name, DefaultValue: defaultValue.String(), Description: description, Type: DURATION}
	RegisterVar(v)
	return DurationVar{getVar(name)}
}

// RegisterVar registers a generic environment variable.
func RegisterVar(v Var) {
	mutex.Lock()

	if old, ok := allVars[v.Name]; ok {
		if v.Description != "" {
			allVars[v.Name] = v // last one with a description wins if the same variable name is registered multiple times
		}

		if old.Description != v.Description || old.DefaultValue != v.DefaultValue || old.Type != v.Type || old.Deprecated != v.Deprecated || old.Hidden != v.Hidden {
			log.Warnf("The environment variable %s was registered multiple times using different metadata: %v, %v", v.Name, old, v)
		}
	} else {
		allVars[v.Name] = v
	}

	mutex.Unlock()
}

func getVar(name string) Var {
	mutex.Lock()
	result := allVars[name]
	mutex.Unlock()

	return result
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v StringVar) Get() string {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v StringVar) Lookup() (string, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	return result, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v BoolVar) Get() bool {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v BoolVar) Lookup() (bool, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	b, err := strconv.ParseBool(result)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting true/false, defaulting to %v", result, v.DefaultValue)
		b, _ = strconv.ParseBool(v.DefaultValue)
	}

	return b, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v IntVar) Get() int {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v IntVar) Lookup() (int, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	i, err := strconv.Atoi(result)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting an integer, defaulting to %v", result, v.DefaultValue)
		i, _ = strconv.Atoi(v.DefaultValue)
	}

	return i, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v FloatVar) Get() float64 {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v FloatVar) Lookup() (float64, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	f, err := strconv.ParseFloat(result, 64)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting a floating-point value, defaulting to %v", result, v.DefaultValue)
		f, _ = strconv.ParseFloat(v.DefaultValue, 64)
	}

	return f, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v DurationVar) Get() time.Duration {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v DurationVar) Lookup() (time.Duration, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	d, err := time.ParseDuration(result)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting a duration, defaulting to %v", result, v.DefaultValue)
		d, _ = time.ParseDuration(v.DefaultValue)
	}

	return d, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v GenericVar[T]) Get() T {
	if v.delegate != nil {
		return v.delegate.Get()
	}
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v GenericVar[T]) Lookup() (T, bool) {
	if v.delegate != nil {
		return v.delegate.Lookup()
	}
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	res := new(T)

	if err := json.Unmarshal([]byte(result), res); err != nil {
		log.Warnf("Invalid environment variable value `%s` defaulting to %v: %v", result, v.DefaultValue, err)
		_ = json.Unmarshal([]byte(v.DefaultValue), res)
	}

	return *res, ok
}

// specializedVar represents a var that can Get/Lookup
type specializedVar[T any] interface {
	Lookup() (T, bool)
	Get() T
}
<commit_msg>Add interface to store collections of env vars (#619)<commit_after>// Copyright 2019 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package env makes it possible to track use of environment variables within a procress
// in order to generate documentation for these uses.
package env

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"strconv"
	"sync"
	"time"

	"istio.io/pkg/log"
)

// The type of a variable's value
type VarType byte

const (
	// Variable holds a free-form string.
	STRING VarType = iota
	// Variable holds a boolean value.
	BOOL
	// Variable holds a signed integer.
	INT
	// Variables holds a floating point value.
	FLOAT
	// Variable holds a time duration.
	DURATION
	// Variable holds a dynamic unknown type.
	OTHER
)

// Var describes a single environment variable
type Var struct {
	// The name of the environment variable.
	Name string

	// The optional default value of the environment variable.
	DefaultValue string

	// Description of the environment variable's purpose.
	Description string

	// Hide the existence of this variable when outputting usage information.
	Hidden bool

	// Mark this variable as deprecated when generating usage information.
	Deprecated bool

	// The type of the variable's value
	Type VarType

	// The underlying Go type of the variable
	GoType string
}

// StringVar represents a single string environment variable.
type StringVar struct {
	Var
}

// BoolVar represents a single boolean environment variable.
type BoolVar struct {
	Var
}

// IntVar represents a single integer environment variable.
type IntVar struct {
	Var
}

// FloatVar represents a single floating-point environment variable.
type FloatVar struct {
	Var
}

// DurationVar represents a single duration environment variable.
type DurationVar struct {
	Var
}

var (
	allVars = make(map[string]Var)
	mutex   sync.Mutex
)

// VarDescriptions returns a description of this process' environment variables, sorted by name.
func VarDescriptions() []Var {
	mutex.Lock()
	sorted := make([]Var, 0, len(allVars))
	for _, v := range allVars {
		sorted = append(sorted, v)
	}
	mutex.Unlock()

	sort.Slice(sorted, func(i, j int) bool {
		return sorted[i].Name < sorted[j].Name
	})

	return sorted
}

type Parseable interface {
	comparable
}

type GenericVar[T Parseable] struct {
	Var
	delegate specializedVar[T]
}

func Register[T Parseable](name string, defaultValue T, description string) GenericVar[T] {
	// Specialized cases
	// In the future, once only Register() remains, we can likely drop most of these.
	// however, time.Duration is needed still as it doesn't implement json
	switch d := any(defaultValue).(type) {
	case time.Duration:
		v := RegisterDurationVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case string:
		v := RegisterStringVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case float64:
		v := RegisterFloatVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case int:
		v := RegisterIntVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	case bool:
		v := RegisterBoolVar(name, d, description)
		return GenericVar[T]{v.Var, any(v).(specializedVar[T])}
	}
	b, _ := json.Marshal(defaultValue)
	v := Var{Name: name, DefaultValue: string(b), Description: description, Type: STRING, GoType: fmt.Sprintf("%T", defaultValue)}
	RegisterVar(v)
	return GenericVar[T]{getVar(name), nil}
}

// RegisterStringVar registers a new string environment variable.
func RegisterStringVar(name string, defaultValue string, description string) StringVar {
	v := Var{Name: name, DefaultValue: defaultValue, Description: description, Type: STRING}
	RegisterVar(v)
	return StringVar{getVar(name)}
}

// RegisterBoolVar registers a new boolean environment variable.
func RegisterBoolVar(name string, defaultValue bool, description string) BoolVar {
	v := Var{Name: name, DefaultValue: strconv.FormatBool(defaultValue), Description: description, Type: BOOL}
	RegisterVar(v)
	return BoolVar{getVar(name)}
}

// RegisterIntVar registers a new integer environment variable.
func RegisterIntVar(name string, defaultValue int, description string) IntVar {
	v := Var{Name: name, DefaultValue: strconv.FormatInt(int64(defaultValue), 10), Description: description, Type: INT}
	RegisterVar(v)
	return IntVar{getVar(name)}
}

// RegisterFloatVar registers a new floating-point environment variable.
func RegisterFloatVar(name string, defaultValue float64, description string) FloatVar {
	v := Var{Name: name, DefaultValue: strconv.FormatFloat(defaultValue, 'G', -1, 64), Description: description, Type: FLOAT}
	RegisterVar(v)
	return FloatVar{v}
}

// RegisterDurationVar registers a new duration environment variable.
func RegisterDurationVar(name string, defaultValue time.Duration, description string) DurationVar {
	v := Var{Name: name, DefaultValue: defaultValue.String(), Description: description, Type: DURATION}
	RegisterVar(v)
	return DurationVar{getVar(name)}
}

// RegisterVar registers a generic environment variable.
func RegisterVar(v Var) {
	mutex.Lock()

	if old, ok := allVars[v.Name]; ok {
		if v.Description != "" {
			allVars[v.Name] = v // last one with a description wins if the same variable name is registered multiple times
		}

		if old.Description != v.Description || old.DefaultValue != v.DefaultValue || old.Type != v.Type || old.Deprecated != v.Deprecated || old.Hidden != v.Hidden {
			log.Warnf("The environment variable %s was registered multiple times using different metadata: %v, %v", v.Name, old, v)
		}
	} else {
		allVars[v.Name] = v
	}

	mutex.Unlock()
}

func getVar(name string) Var {
	mutex.Lock()
	result := allVars[name]
	mutex.Unlock()

	return result
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v StringVar) Get() string {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v StringVar) Lookup() (string, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	return result, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v BoolVar) Get() bool {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v BoolVar) Lookup() (bool, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	b, err := strconv.ParseBool(result)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting true/false, defaulting to %v", result, v.DefaultValue)
		b, _ = strconv.ParseBool(v.DefaultValue)
	}

	return b, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v IntVar) Get() int {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v IntVar) Lookup() (int, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	i, err := strconv.Atoi(result)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting an integer, defaulting to %v", result, v.DefaultValue)
		i, _ = strconv.Atoi(v.DefaultValue)
	}

	return i, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v FloatVar) Get() float64 {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v FloatVar) Lookup() (float64, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	f, err := strconv.ParseFloat(result, 64)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting a floating-point value, defaulting to %v", result, v.DefaultValue)
		f, _ = strconv.ParseFloat(v.DefaultValue, 64)
	}

	return f, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v DurationVar) Get() time.Duration {
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v DurationVar) Lookup() (time.Duration, bool) {
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	d, err := time.ParseDuration(result)
	if err != nil {
		log.Warnf("Invalid environment variable value `%s`, expecting a duration, defaulting to %v", result, v.DefaultValue)
		d, _ = time.ParseDuration(v.DefaultValue)
	}

	return d, ok
}

// Get retrieves the value of the environment variable.
// It returns the value, which will be the default if the variable is not present.
// To distinguish between an empty value and an unset value, use Lookup.
func (v GenericVar[T]) Get() T {
	if v.delegate != nil {
		return v.delegate.Get()
	}
	result, _ := v.Lookup()
	return result
}

// Lookup retrieves the value of the environment variable. If the
// variable is present in the environment the
// value (which may be empty) is returned and the boolean is true.
// Otherwise the returned value will be the default and the boolean will
// be false.
func (v GenericVar[T]) Lookup() (T, bool) {
	if v.delegate != nil {
		return v.delegate.Lookup()
	}
	result, ok := os.LookupEnv(v.Name)
	if !ok {
		result = v.DefaultValue
	}

	res := new(T)

	if err := json.Unmarshal([]byte(result), res); err != nil {
		log.Warnf("Invalid environment variable value `%s` defaulting to %v: %v", result, v.DefaultValue, err)
		_ = json.Unmarshal([]byte(v.DefaultValue), res)
	}

	return *res, ok
}

func (v GenericVar[T]) IsSet() bool {
	_, ok := v.Lookup()
	return ok
}

func (v GenericVar[T]) GetName() string {
	return v.Var.Name
}

// specializedVar represents a var that can Get/Lookup
type specializedVar[T any] interface {
	Lookup() (T, bool)
	Get() T
}

// VariableInfo provides generic information about a variable. All Variables implement this interface.
// This is largely to workaround lack of covariance in Go.
type VariableInfo interface {
	GetName() string
	IsSet() bool
}
<|endoftext|>


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

package forgetfs

import (
	"fmt"
	"os"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
)

// Create a file system whose sole contents are a file named "foo" and a
// directory named "bar".
//
// The file "foo" may be opened for reading and/or writing, but reads and
// writes aren't supported. Additionally, any non-existent file or directory
// name may be created within any directory, but the resulting inode will
// appear to have been unlinked immediately.
//
// The file system maintains reference counts for the inodes involved. It will
// panic if a reference count becomes negative or if an inode ID is re-used
// after we expect it to be dead. Its Check method may be used to check that
// there are no inodes with non-zero reference counts remaining, after
// unmounting.
func NewFileSystem() (fs *ForgetFS) {
	// Set up the actual file system.
	impl := &fsImpl{
		inodes: map[fuseops.InodeID]*inode{
			cannedID_Root: &inode{
				lookupCount: 1,
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
			cannedID_Foo: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777,
				},
			},
			cannedID_Bar: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
		},
		nextInodeID: cannedID_Next,
	}

	impl.mu = syncutil.NewInvariantMutex(impl.checkInvariants)

	// Set up a wrapper that exposes only certain methods.
	fs = &ForgetFS{
		impl:   impl,
		server: fuseutil.NewFileSystemServer(impl),
	}

	return
}

////////////////////////////////////////////////////////////////////////
// ForgetFS
////////////////////////////////////////////////////////////////////////

type ForgetFS struct {
	impl   *fsImpl
	server fuse.Server
}

func (fs *ForgetFS) ServeOps(c *fuse.Connection) {
	fs.server.ServeOps(c)
}

// Panic if there are any inodes that have a non-zero reference count. For use
// after unmounting.
func (fs *ForgetFS) Check() {
	fs.impl.Check()
}

////////////////////////////////////////////////////////////////////////
// Actual implementation
////////////////////////////////////////////////////////////////////////

const (
	cannedID_Root = fuseops.RootInodeID + iota
	cannedID_Foo
	cannedID_Bar
	cannedID_Next
)

type fsImpl struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Mutable state
	/////////////////////////

	mu syncutil.InvariantMutex

	// An index of inode by ID, for all IDs we have issued.
	//
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	//
	// GUARDED_BY(mu)
	inodes map[fuseops.InodeID]*inode

	// The next ID to issue.
	//
	// INVARIANT: For each k in inodes, k < nextInodeID
	//
	// GUARDED_BY(mu)
	nextInodeID fuseops.InodeID
}

////////////////////////////////////////////////////////////////////////
// inode
////////////////////////////////////////////////////////////////////////

type inode struct {
	attributes fuseops.InodeAttributes

	// The current lookup count.
	lookupCount int

	// true if lookupCount has ever been positive.
	lookedUp bool
}

func (in *inode) Forgotten() bool {
	return in.lookedUp && in.lookupCount <= 0
}

func (in *inode) IncrementLookupCount() {
	in.lookupCount++
	in.lookedUp = true
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) checkInvariants() {
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	for _, v := range fs.inodes {
		if !(v.lookupCount >= 0) {
			panic("Negative lookup count")
		}
	}

	// INVARIANT: For each k in inodes, k < nextInodeID
	for k, _ := range fs.inodes {
		if !(k < fs.nextInodeID) {
			panic("Unexpectedly large inode ID")
		}
	}
}

// LOCKS_EXCLUDED(fs.mu)
func (fs *fsImpl) Check() {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	for k, v := range fs.inodes {
		if v.lookupCount != 0 {
			panic(fmt.Sprintf("Inode %v has lookup count %v", k, v.lookupCount))
		}
	}
}

// Look up the inode and verify it hasn't been forgotten.
//
// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) findInodeByID(id fuseops.InodeID) (in *inode) {
	in = fs.inodes[id]
	if in == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	if in.Forgotten() {
		panic(fmt.Sprintf("Forgotten inode: %v", id))
	}

	return
}

////////////////////////////////////////////////////////////////////////
// FileSystem methods
////////////////////////////////////////////////////////////////////////

func (fs *fsImpl) Init(
	op *fuseops.InitOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	return
}

func (fs *fsImpl) LookUpInode(
	op *fuseops.LookUpInodeOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Make sure the parent exists and has not been forgotten.
	_ = fs.findInodeByID(op.Parent)

	// Handle the names we support.
	var childID fuseops.InodeID
	switch {
	case op.Parent == cannedID_Root && op.Name == "foo":
		childID = cannedID_Foo

	case op.Parent == cannedID_Root && op.Name == "bar":
		childID = cannedID_Bar

	default:
		err = fuse.ENOENT
		return
	}

	// Look up the child.
	child := fs.findInodeByID(childID)
	child.IncrementLookupCount()

	// Return an appropriate entry.
	op.Entry = fuseops.ChildInodeEntry{
		Child:      childID,
		Attributes: child.attributes,
	}

	return
}

func (fs *fsImpl) GetInodeAttributes(
	op *fuseops.GetInodeAttributesOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Find the inode, verifying that it has not been forgotten.
	in := fs.findInodeByID(op.Inode)

	// Return appropriate attributes.
	op.Attributes = in.attributes

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

package forgetfs

import (
	"fmt"
	"os"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseops"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
)

// Create a file system whose sole contents are a file named "foo" and a
// directory named "bar".
//
// The file "foo" may be opened for reading and/or writing, but reads and
// writes aren't supported. Additionally, any non-existent file or directory
// name may be created within any directory, but the resulting inode will
// appear to have been unlinked immediately.
//
// The file system maintains reference counts for the inodes involved. It will
// panic if a reference count becomes negative or if an inode ID is re-used
// after we expect it to be dead. Its Check method may be used to check that
// there are no inodes with non-zero reference counts remaining, after
// unmounting.
func NewFileSystem() (fs *ForgetFS) {
	// Set up the actual file system.
	impl := &fsImpl{
		inodes: map[fuseops.InodeID]*inode{
			cannedID_Root: &inode{
				lookupCount: 1,
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
			cannedID_Foo: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777,
				},
			},
			cannedID_Bar: &inode{
				attributes: fuseops.InodeAttributes{
					Nlink: 1,
					Mode:  0777 | os.ModeDir,
				},
			},
		},
		nextInodeID: cannedID_Next,
	}

	impl.mu = syncutil.NewInvariantMutex(impl.checkInvariants)

	// Set up a wrapper that exposes only certain methods.
	fs = &ForgetFS{
		impl:   impl,
		server: fuseutil.NewFileSystemServer(impl),
	}

	return
}

////////////////////////////////////////////////////////////////////////
// ForgetFS
////////////////////////////////////////////////////////////////////////

type ForgetFS struct {
	impl   *fsImpl
	server fuse.Server
}

func (fs *ForgetFS) ServeOps(c *fuse.Connection) {
	fs.server.ServeOps(c)
}

// Panic if there are any inodes that have a non-zero reference count. For use
// after unmounting.
func (fs *ForgetFS) Check() {
	fs.impl.Check()
}

////////////////////////////////////////////////////////////////////////
// Actual implementation
////////////////////////////////////////////////////////////////////////

const (
	cannedID_Root = fuseops.RootInodeID + iota
	cannedID_Foo
	cannedID_Bar
	cannedID_Next
)

type fsImpl struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Mutable state
	/////////////////////////

	mu syncutil.InvariantMutex

	// An index of inode by ID, for all IDs we have issued.
	//
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	//
	// GUARDED_BY(mu)
	inodes map[fuseops.InodeID]*inode

	// The next ID to issue.
	//
	// INVARIANT: For each k in inodes, k < nextInodeID
	//
	// GUARDED_BY(mu)
	nextInodeID fuseops.InodeID
}

////////////////////////////////////////////////////////////////////////
// inode
////////////////////////////////////////////////////////////////////////

type inode struct {
	attributes fuseops.InodeAttributes

	// The current lookup count.
	lookupCount int

	// true if lookupCount has ever been positive.
	lookedUp bool
}

func (in *inode) Forgotten() bool {
	return in.lookedUp && in.lookupCount <= 0
}

func (in *inode) IncrementLookupCount() {
	in.lookupCount++
	in.lookedUp = true
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) checkInvariants() {
	// INVARIANT: For each v in inodes, v.lookupCount >= 0
	for _, v := range fs.inodes {
		if !(v.lookupCount >= 0) {
			panic("Negative lookup count")
		}
	}

	// INVARIANT: For each k in inodes, k < nextInodeID
	for k, _ := range fs.inodes {
		if !(k < fs.nextInodeID) {
			panic("Unexpectedly large inode ID")
		}
	}
}

// LOCKS_EXCLUDED(fs.mu)
func (fs *fsImpl) Check() {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	for k, v := range fs.inodes {
		if v.lookupCount != 0 {
			panic(fmt.Sprintf("Inode %v has lookup count %v", k, v.lookupCount))
		}
	}
}

// Look up the inode and verify it hasn't been forgotten.
//
// LOCKS_REQUIRED(fs.mu)
func (fs *fsImpl) findInodeByID(id fuseops.InodeID) (in *inode) {
	in = fs.inodes[id]
	if in == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	if in.Forgotten() {
		panic(fmt.Sprintf("Forgotten inode: %v", id))
	}

	return
}

////////////////////////////////////////////////////////////////////////
// FileSystem methods
////////////////////////////////////////////////////////////////////////

func (fs *fsImpl) Init(
	op *fuseops.InitOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	return
}

func (fs *fsImpl) LookUpInode(
	op *fuseops.LookUpInodeOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Make sure the parent exists and has not been forgotten.
	_ = fs.findInodeByID(op.Parent)

	// Handle the names we support.
	var childID fuseops.InodeID
	switch {
	case op.Parent == cannedID_Root && op.Name == "foo":
		childID = cannedID_Foo

	case op.Parent == cannedID_Root && op.Name == "bar":
		childID = cannedID_Bar

	default:
		err = fuse.ENOENT
		return
	}

	// Look up the child.
	child := fs.findInodeByID(childID)
	child.IncrementLookupCount()

	// Return an appropriate entry.
	op.Entry = fuseops.ChildInodeEntry{
		Child:      childID,
		Attributes: child.attributes,
	}

	return
}

func (fs *fsImpl) GetInodeAttributes(
	op *fuseops.GetInodeAttributesOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Find the inode, verifying that it has not been forgotten.
	in := fs.findInodeByID(op.Inode)

	// Return appropriate attributes.
	op.Attributes = in.attributes

	return
}

func (fs *fsImpl) OpenFile(
	op *fuseops.OpenFileOp) {
	var err error
	defer fuseutil.RespondToOp(op, &err)

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Verify that the inode has not been forgotten.
	_ = fs.findInodeByID(op.Inode)

	return
}
<|endoftext|>


<commit_before>package whois

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

type WhoisRecord struct {
	Domain string

	// times
	LastUpdatedAt time.Time
	CreatedAt time.Time
	ExpiresAt time.Time

	// contact points
	Registrar string
	ContactEmails []string
	ContactPhoneNumbers []string

	// tech details
	NameServers []url.URL
	DNSSECEnabled bool
}

func ParseWhoisResponse(resp string) (WhoisRecord, error) {
	record := WhoisRecord{}
	rp := &record

	last_updated_at, err := find_last_updated_at(resp)
	if err == nil {
		// From standard: 'Mon Jan 2 15:04:05 MST 2006'
		// incoming: Fri Nov 20 07:42:26 GMT 2015
		t, err := time.Parse("Mon Jan 2 15:04:05 MST 2006", last_updated_at)
		if err == nil {
			rp.LastUpdatedAt = t
		}
	}

	created_at, err := find_created_at(resp)

	if err == nil {
		t, err := time.Parse("Mon Jan 2 15:04:05 MST 2006", created_at)
		if err == nil {
			rp.CreatedAt = t
		}
	}

	expires_at, err := find_expires_at(resp)

	if err == nil {
		t, err := time.Parse("Mon Jan 2 15:04:05 MST 2006", expires_at)
		if err == nil {
			rp.ExpiresAt = t
		}
	}

	return record, nil
}

func find_last_updated_at(resp string) (string, error) {
	r, err := regexp.Compile(`(?im)Updated Date: \s+(.+)$`)

	if err != nil {
		return "", err
	}

	return find_date_time(resp, r, "Updated Date")
}

func find_created_at(resp string) (string, error) {
	r, err := regexp.Compile(`(?im)Registration Date: \s+(.+)$`)

	if err != nil {
		return "", err
	}

	return find_date_time(resp, r, "Registration Date")
}

func find_expires_at(resp string) (string, error) {
	r, err := regexp.Compile(`(?im)Expiration Date: \s+(.+)$`)

	if err != nil {
		return "", err
	}

	return find_date_time(resp, r, "Expiration Date")
}

func find_date_time(resp string, r *regexp.Regexp, key string) (string, error) {
	loc := r.FindStringIndex(resp)

	if len(loc) != 2 {
		return "", fmt.Errorf("error getting location of match")
	}

	old := resp[loc[0]:loc[1]]
	// fmt.Printf("old = '%s'\n", old)

	match := strings.TrimSpace(strings.Replace(old, key + ":", "", -1))

	// fmt.Printf("match = '%s'\n", match)

	if match == "" {
		return "", fmt.Errorf("no matches found for '%s' record", key)
	}

	return match, nil
}
<commit_msg>Simplify parsing code to handle more WHOIS formats<commit_after>package whois

import (
	"fmt"
	"time"
	"net/url"
	"regexp"
)

type WhoisRecord struct {
	Domain string

	// times
	LastUpdatedAt time.Time
	CreatedAt time.Time
	ExpiresAt time.Time

	// contact points
	Registrar string
	ContactEmails []string
	ContactPhoneNumbers []string

	// tech details
	NameServers []url.URL
	DNSSECEnabled bool
}

func ParseWhoisResponse(resp string) (WhoisRecord, error) {
	record := WhoisRecord{}
	rp := &record

	last_updated_at, err := find_last_updated_at(resp)
	if err == nil {
		rp.LastUpdatedAt = last_updated_at
	}

	created_at, err := find_created_at(resp)
	if err == nil {
		rp.CreatedAt = created_at
	}

	expires_at, err := find_expires_at(resp)
	if err == nil {
		rp.ExpiresAt = expires_at
	}

	return record, nil
}

func find_last_updated_at(resp string) (time.Time, error) {
	patterns_and_formats := map[*regexp.Regexp]string{
		regexp.MustCompile(`(?im)Updated Date: \s+(.+)$`): "Mon Jan 2 15:04:05 MST 2006",
	}
	return find_date_time(resp, patterns_and_formats, "LastUpdatedAt")
}

func find_created_at(resp string) (time.Time, error) {
	patterns_and_formats := map[*regexp.Regexp]string{
		regexp.MustCompile(`(?im)Registration Date: \s+(.+)$`): "Mon Jan 2 15:04:05 MST 2006",
	}
	return find_date_time(resp, patterns_and_formats, "CreatedAt")
}

func find_expires_at(resp string) (time.Time, error) {
	patterns_and_formats := map[*regexp.Regexp]string{
		regexp.MustCompile(`(?im)Expiration Date: \s+(.+)$`): "Mon Jan 2 15:04:05 MST 2006",
	}
	return find_date_time(resp, patterns_and_formats, "ExpiresAt")
}

func find_date_time(resp string, patterns_and_formats map[*regexp.Regexp]string, key string) (time.Time, error) {
	for r, format := range patterns_and_formats {
		res := r.FindStringSubmatch(resp)

		// Grab the first match
		if len(res) > 1 {
			t, err := time.Parse(format, res[1])
			if err == nil {
				return t, nil
			}
		}
	}
	return time.Now(), fmt.Errorf("unable to find patern for %s", key)
}
<|endoftext|>


<commit_before>package tsm1

// Int64 encoding uses two different strategies depending on the range of values in
// the uncompressed data.  Encoded values are first encoding used zig zag encoding.
// This interleaves positive and negative integers across a range of positive integers.
//
// For example, [-2,-1,0,1] becomes [3,1,0,2]. See
// https://developers.google.com/protocol-buffers/docs/encoding?hl=en#signed-integers
// for more information.
//
// If all the zig zag encoded values less than 1 << 60 - 1, they are compressed using
// simple8b encoding.  If any values is larger than 1 << 60 - 1, the values are stored uncompressed.
//
// Each encoded byte slice, contains a 1 byte header followed by multiple 8 byte packed integers
// or 8 byte uncompressed integers.  The 4 high bits of the first byte indicate the encoding type
// for the remaining bytes.
//
// There are currently two encoding types that can be used with room for 16 total.  These additional
// encoding slots are reserved for future use.  One improvement to be made is to use a patched
// encoding such as PFOR if only a small number of values exceed the max compressed value range.  This
// should improve compression ratios with very large integers near the ends of the int64 range.

import (
	"encoding/binary"
	"fmt"

	"github.com/jwilder/encoding/simple8b"
)

const (
	// intUncompressed is an uncompressed format using 8 bytes per point
	intUncompressed = 0
	// intCompressedSimple is a bit-packed format using simple8b encoding
	intCompressedSimple = 1
	// intCompressedRLE is a run-length encoding format
	intCompressedRLE = 2
)

// Int64Encoder encoders int64 into byte slices
type Int64Encoder interface {
	Write(v int64)
	Bytes() ([]byte, error)
}

// Int64Decoder decodes a byte slice into int64s
type Int64Decoder interface {
	Next() bool
	Read() int64
	Error() error
}

type int64Encoder struct {
	prev   int64
	rle    bool
	values []uint64
}

func NewInt64Encoder() Int64Encoder {
	return &int64Encoder{rle: true}
}

func (e *int64Encoder) Write(v int64) {
	// Delta-encode each value as it's written.  This happens before
	// ZigZagEncoding because the deltas could be negative.
	delta := v - e.prev
	e.prev = v
	enc := ZigZagEncode(delta)
	if len(e.values) > 1 {
		e.rle = e.rle && e.values[len(e.values)-1] == enc
	}

	e.values = append(e.values, enc)
}

func (e *int64Encoder) Bytes() ([]byte, error) {
	// Only run-length encode if it could be reduce storage size
	if e.rle && len(e.values) > 2 {
		return e.encodeRLE()
	}

	for _, v := range e.values {
		// Value is too large to encode using packed format
		if v > simple8b.MaxValue {
			return e.encodeUncompressed()
		}
	}

	return e.encodePacked()
}

func (e *int64Encoder) encodeRLE() ([]byte, error) {
	// Large varints can take up to 10 bytes
	b := make([]byte, 1+10*3)

	// 4 high bits used for the encoding type
	b[0] = byte(intCompressedRLE) << 4

	i := 1
	// The first value
	binary.BigEndian.PutUint64(b[i:], e.values[0])
	i += 8
	// The first delta
	i += binary.PutUvarint(b[i:], e.values[1])
	// The number of times the delta is repeated
	i += binary.PutUvarint(b[i:], uint64(len(e.values)-1))

	return b[:i], nil
}

func (e *int64Encoder) encodePacked() ([]byte, error) {
	if len(e.values) == 0 {
		return nil, nil
	}

	// Encode all but the first value.  Fist value is written unecoded
	// using 8 bytes.
	encoded, err := simple8b.EncodeAll(e.values[1:])
	if err != nil {
		return nil, err
	}

	b := make([]byte, 1+(len(encoded)+1)*8)
	// 4 high bits of first byte store the encoding type for the block
	b[0] = byte(intCompressedSimple) << 4

	// Write the first value since it's not part of the encoded values
	binary.BigEndian.PutUint64(b[1:9], e.values[0])

	// Write the encoded values
	for i, v := range encoded {
		binary.BigEndian.PutUint64(b[9+i*8:9+i*8+8], v)
	}
	return b, nil
}

func (e *int64Encoder) encodeUncompressed() ([]byte, error) {
	if len(e.values) == 0 {
		return nil, nil
	}

	b := make([]byte, 1+len(e.values)*8)
	// 4 high bits of first byte store the encoding type for the block
	b[0] = byte(intUncompressed) << 4

	for i, v := range e.values {
		binary.BigEndian.PutUint64(b[1+i*8:1+i*8+8], v)
	}
	return b, nil
}

type int64Decoder struct {
	values []uint64
	bytes  []byte
	i      int
	n      int
	prev   int64
	first  bool

	// The first value for a run-length encoded byte slice
	rleFirst uint64

	// The delta value for a run-length encoded byte slice
	rleDelta uint64
	encoding byte
	err      error
}

func NewInt64Decoder(b []byte) Int64Decoder {
	d := &int64Decoder{
		// 240 is the maximum number of values that can be encoded into a single uint64 using simple8b
		values: make([]uint64, 240),
	}

	d.SetBytes(b)
	return d
}

func (d *int64Decoder) SetBytes(b []byte) {
	if len(b) > 0 {
		d.encoding = b[0] >> 4
		d.bytes = b[1:]
	}
	d.first = true
	d.i = 0
	d.n = 0
}

func (d *int64Decoder) Next() bool {
	if d.i >= d.n && len(d.bytes) == 0 {
		return false
	}

	d.i += 1

	if d.i >= d.n {
		switch d.encoding {
		case intUncompressed:
			d.decodeUncompressed()
		case intCompressedSimple:
			d.decodePacked()
		case intCompressedRLE:
			d.decodeRLE()
		default:
			d.err = fmt.Errorf("unknown encoding %v", d.encoding)
		}
	}
	return d.i < d.n
}

func (d *int64Decoder) Error() error {
	return d.err
}

func (d *int64Decoder) Read() int64 {
	switch d.encoding {
	case intCompressedRLE:
		return ZigZagDecode(d.rleFirst + uint64(d.i)*d.rleDelta)
	default:
		v := ZigZagDecode(d.values[d.i])
		// v is the delta encoded value, we need to add the prior value to get the original
		v = v + d.prev
		d.prev = v
		return v

	}
}

func (d *int64Decoder) decodeRLE() {
	if len(d.bytes) == 0 {
		return
	}

	var i, n int

	// Next 8 bytes is the starting value
	first := binary.BigEndian.Uint64(d.bytes[i : i+8])
	i += 8

	// Next 1-10 bytes is the delta value
	value, n := binary.Uvarint(d.bytes[i:])

	i += n

	// Last 1-10 bytes is how many times the value repeats
	count, n := binary.Uvarint(d.bytes[i:])

	// Store the first value and delta value so we do not need to allocate
	// a large values slice.  We can comput the value at position d.i on
	// demand.
	d.rleFirst = first
	d.rleDelta = value
	d.n = int(count) + 1
	d.i = 0

	// We've process all the bytes
	d.bytes = nil
}

func (d *int64Decoder) decodePacked() {
	if len(d.bytes) == 0 {
		return
	}

	v := binary.BigEndian.Uint64(d.bytes[0:8])
	// The first value is always unencoded
	if d.first {
		d.first = false
		d.n = 1
		d.values[0] = v
	} else {
		n, err := simple8b.Decode(d.values, v)
		if err != nil {
			// Should never happen, only error that could be returned is if the the value to be decoded was not
			// actually encoded by simple8b encoder.
			d.err = fmt.Errorf("failed to decode value %v: %v", v, err)
		}

		d.n = n
	}
	d.i = 0
	d.bytes = d.bytes[8:]
}

func (d *int64Decoder) decodeUncompressed() {
	if len(d.bytes) == 0 {
		return
	}

	d.values[0] = binary.BigEndian.Uint64(d.bytes[0:8])
	d.i = 0
	d.n = 1
	d.bytes = d.bytes[8:]
}
<commit_msg>Fix typos<commit_after>package tsm1

// Int64 encoding uses two different strategies depending on the range of values in
// the uncompressed data.  Encoded values are first encoding used zig zag encoding.
// This interleaves positive and negative integers across a range of positive integers.
//
// For example, [-2,-1,0,1] becomes [3,1,0,2]. See
// https://developers.google.com/protocol-buffers/docs/encoding?hl=en#signed-integers
// for more information.
//
// If all the zig zag encoded values less than 1 << 60 - 1, they are compressed using
// simple8b encoding.  If any values is larger than 1 << 60 - 1, the values are stored uncompressed.
//
// Each encoded byte slice, contains a 1 byte header followed by multiple 8 byte packed integers
// or 8 byte uncompressed integers.  The 4 high bits of the first byte indicate the encoding type
// for the remaining bytes.
//
// There are currently two encoding types that can be used with room for 16 total.  These additional
// encoding slots are reserved for future use.  One improvement to be made is to use a patched
// encoding such as PFOR if only a small number of values exceed the max compressed value range.  This
// should improve compression ratios with very large integers near the ends of the int64 range.

import (
	"encoding/binary"
	"fmt"

	"github.com/jwilder/encoding/simple8b"
)

const (
	// intUncompressed is an uncompressed format using 8 bytes per point
	intUncompressed = 0
	// intCompressedSimple is a bit-packed format using simple8b encoding
	intCompressedSimple = 1
	// intCompressedRLE is a run-length encoding format
	intCompressedRLE = 2
)

// Int64Encoder encoders int64 into byte slices
type Int64Encoder interface {
	Write(v int64)
	Bytes() ([]byte, error)
}

// Int64Decoder decodes a byte slice into int64s
type Int64Decoder interface {
	Next() bool
	Read() int64
	Error() error
}

type int64Encoder struct {
	prev   int64
	rle    bool
	values []uint64
}

func NewInt64Encoder() Int64Encoder {
	return &int64Encoder{rle: true}
}

func (e *int64Encoder) Write(v int64) {
	// Delta-encode each value as it's written.  This happens before
	// ZigZagEncoding because the deltas could be negative.
	delta := v - e.prev
	e.prev = v
	enc := ZigZagEncode(delta)
	if len(e.values) > 1 {
		e.rle = e.rle && e.values[len(e.values)-1] == enc
	}

	e.values = append(e.values, enc)
}

func (e *int64Encoder) Bytes() ([]byte, error) {
	// Only run-length encode if it could be reduce storage size
	if e.rle && len(e.values) > 2 {
		return e.encodeRLE()
	}

	for _, v := range e.values {
		// Value is too large to encode using packed format
		if v > simple8b.MaxValue {
			return e.encodeUncompressed()
		}
	}

	return e.encodePacked()
}

func (e *int64Encoder) encodeRLE() ([]byte, error) {
	// Large varints can take up to 10 bytes
	b := make([]byte, 1+10*3)

	// 4 high bits used for the encoding type
	b[0] = byte(intCompressedRLE) << 4

	i := 1
	// The first value
	binary.BigEndian.PutUint64(b[i:], e.values[0])
	i += 8
	// The first delta
	i += binary.PutUvarint(b[i:], e.values[1])
	// The number of times the delta is repeated
	i += binary.PutUvarint(b[i:], uint64(len(e.values)-1))

	return b[:i], nil
}

func (e *int64Encoder) encodePacked() ([]byte, error) {
	if len(e.values) == 0 {
		return nil, nil
	}

	// Encode all but the first value.  Fist value is written unencoded
	// using 8 bytes.
	encoded, err := simple8b.EncodeAll(e.values[1:])
	if err != nil {
		return nil, err
	}

	b := make([]byte, 1+(len(encoded)+1)*8)
	// 4 high bits of first byte store the encoding type for the block
	b[0] = byte(intCompressedSimple) << 4

	// Write the first value since it's not part of the encoded values
	binary.BigEndian.PutUint64(b[1:9], e.values[0])

	// Write the encoded values
	for i, v := range encoded {
		binary.BigEndian.PutUint64(b[9+i*8:9+i*8+8], v)
	}
	return b, nil
}

func (e *int64Encoder) encodeUncompressed() ([]byte, error) {
	if len(e.values) == 0 {
		return nil, nil
	}

	b := make([]byte, 1+len(e.values)*8)
	// 4 high bits of first byte store the encoding type for the block
	b[0] = byte(intUncompressed) << 4

	for i, v := range e.values {
		binary.BigEndian.PutUint64(b[1+i*8:1+i*8+8], v)
	}
	return b, nil
}

type int64Decoder struct {
	values []uint64
	bytes  []byte
	i      int
	n      int
	prev   int64
	first  bool

	// The first value for a run-length encoded byte slice
	rleFirst uint64

	// The delta value for a run-length encoded byte slice
	rleDelta uint64
	encoding byte
	err      error
}

func NewInt64Decoder(b []byte) Int64Decoder {
	d := &int64Decoder{
		// 240 is the maximum number of values that can be encoded into a single uint64 using simple8b
		values: make([]uint64, 240),
	}

	d.SetBytes(b)
	return d
}

func (d *int64Decoder) SetBytes(b []byte) {
	if len(b) > 0 {
		d.encoding = b[0] >> 4
		d.bytes = b[1:]
	}
	d.first = true
	d.i = 0
	d.n = 0
}

func (d *int64Decoder) Next() bool {
	if d.i >= d.n && len(d.bytes) == 0 {
		return false
	}

	d.i += 1

	if d.i >= d.n {
		switch d.encoding {
		case intUncompressed:
			d.decodeUncompressed()
		case intCompressedSimple:
			d.decodePacked()
		case intCompressedRLE:
			d.decodeRLE()
		default:
			d.err = fmt.Errorf("unknown encoding %v", d.encoding)
		}
	}
	return d.i < d.n
}

func (d *int64Decoder) Error() error {
	return d.err
}

func (d *int64Decoder) Read() int64 {
	switch d.encoding {
	case intCompressedRLE:
		return ZigZagDecode(d.rleFirst + uint64(d.i)*d.rleDelta)
	default:
		v := ZigZagDecode(d.values[d.i])
		// v is the delta encoded value, we need to add the prior value to get the original
		v = v + d.prev
		d.prev = v
		return v

	}
}

func (d *int64Decoder) decodeRLE() {
	if len(d.bytes) == 0 {
		return
	}

	var i, n int

	// Next 8 bytes is the starting value
	first := binary.BigEndian.Uint64(d.bytes[i : i+8])
	i += 8

	// Next 1-10 bytes is the delta value
	value, n := binary.Uvarint(d.bytes[i:])

	i += n

	// Last 1-10 bytes is how many times the value repeats
	count, n := binary.Uvarint(d.bytes[i:])

	// Store the first value and delta value so we do not need to allocate
	// a large values slice.  We can compute the value at position d.i on
	// demand.
	d.rleFirst = first
	d.rleDelta = value
	d.n = int(count) + 1
	d.i = 0

	// We've process all the bytes
	d.bytes = nil
}

func (d *int64Decoder) decodePacked() {
	if len(d.bytes) == 0 {
		return
	}

	v := binary.BigEndian.Uint64(d.bytes[0:8])
	// The first value is always unencoded
	if d.first {
		d.first = false
		d.n = 1
		d.values[0] = v
	} else {
		n, err := simple8b.Decode(d.values, v)
		if err != nil {
			// Should never happen, only error that could be returned is if the the value to be decoded was not
			// actually encoded by simple8b encoder.
			d.err = fmt.Errorf("failed to decode value %v: %v", v, err)
		}

		d.n = n
	}
	d.i = 0
	d.bytes = d.bytes[8:]
}

func (d *int64Decoder) decodeUncompressed() {
	if len(d.bytes) == 0 {
		return
	}

	d.values[0] = binary.BigEndian.Uint64(d.bytes[0:8])
	d.i = 0
	d.n = 1
	d.bytes = d.bytes[8:]
}
<|endoftext|>