<commit_before>package opml

import (
	"bytes"
	"encoding/xml"
	log "github.com/golang/glog"
	"github.com/jrupac/goliath/models"
	"golang.org/x/net/html/charset"
	"io/ioutil"
	"os"
	"syscall"
	"time"
)

type header struct {
	Title       string `xml:"title"`
	DateCreated string `xml:"dateCreated"`
}

type outline struct {
	Text        string    `xml:"text,attr,omitempty"`
	Title       string    `xml:"title,attr,omitempty"`
	Description string    `xml:"description,attr,omitempty"`
	URL         string    `xml:"xmlUrl,attr,omitempty"`
	Folders     []outline `xml:"outline,omitempty"`
}

type internalOpmlType struct {
	XMLName xml.Name  `xml:"opml"`
	Header  header    `xml:"head"`
	Body    []outline `xml:"body>outline"`
}

// Opml is a nested tree of folders.
type Opml struct {
	Header  header
	Folders models.Folder
}

func createOpmlObject(oi *internalOpmlType) *Opml {
	return &Opml{
		Header:  oi.Header,
		Folders: parseOutline(oi.Body),
	}
}

// parseOutline converts "outline" objects into Folder and Feed objects.
func parseOutline(children []outline) models.Folder {
	folder := models.Folder{
		Name: models.RootFolder,
	}

	for _, o := range children {
		if o.URL != "" {
			// This entity has a URL so we presume it is a Feed and has no children.
			var name string
			switch {
			case o.Title != "":
				name = o.Title
			case o.Text != "":
				name = o.Text
			default:
				name = "<untitled>"
			}

			feed := models.Feed{
				Title:       name,
				Description: o.Description,
				URL:         o.URL,
			}
			folder.Feed = append(folder.Feed, feed)
		} else {
			// This entity has no URL so we presume it is a Folder.
			var name string
			switch {
			case o.Title != "":
				name = o.Title
			case o.Text != "":
				name = o.Text
			default:
				name = "<untitled>"
			}

			child := parseOutline(o.Folders)
			child.Name = name
			folder.Folders = append(folder.Folders, child)
		}
	}

	return folder
}

// parseFolders converts a rooted folder into a list of outline objects.
func parseFolders(folder models.Folder) []outline {
	var outlines []outline

	for _, feed := range folder.Feed {
		outlines = append(outlines, outline{Text: feed.Title, URL: feed.URL})
	}

	for _, child := range folder.Folders {
		o := outline{Text: child.Name}
		o.Folders = parseFolders(child)
		outlines = append(outlines, o)
	}

	return outlines
}

// ParseOpml open a file and returns a parsed OPML object.
func ParseOpml(filename string) (*Opml, error) {
	log.Infof("Loading OPML file from %s", filename)

	f, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	r := bytes.NewReader(f)
	d := xml.NewDecoder(r)
	d.CharsetReader = charset.NewReaderLabel

	oi := new(internalOpmlType)
	err = d.Decode(&oi)
	if err != nil {
		return nil, err
	}

	return createOpmlObject(oi), nil
}

// ExportOpml exports the given folder (with associated feeds and child folders)
// to a file of the given filename in OPML format.
// The file is created with 0777 mode if it does not exist.
func ExportOpml(tree *models.Folder, filename string) error {
	// OPML specifies that time fields conform to RFC822.
	exportTime := time.Now().Format(time.RFC822)
	export := &internalOpmlType{}
	export.Header = header{Title: "Goliath Feed Export", DateCreated: exportTime}
	export.Body = parseFolders(*tree)

	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, syscall.S_IRWXU|syscall.S_IRWXG|syscall.S_IRWXO)
	if err != nil {
		return err
	}
	defer f.Close()

	f.WriteString(xml.Header)
	e := xml.NewEncoder(f)
	e.Indent("", "  ")
	return e.Encode(export)
}
<commit_msg>core/opml: Make OPML export compliant to OPML 2.0 standard.<commit_after>package opml

import (
	"bytes"
	"encoding/xml"
	log "github.com/golang/glog"
	"github.com/jrupac/goliath/models"
	"golang.org/x/net/html/charset"
	"io/ioutil"
	"os"
	"syscall"
	"time"
)

type header struct {
	Title       string `xml:"title"`
	DateCreated string `xml:"dateCreated"`
}

type outline struct {
	Type        string    `xml:"type,attr,omitempty"`
	Text        string    `xml:"text,attr,omitempty"`
	Title       string    `xml:"title,attr,omitempty"`
	Description string    `xml:"description,attr,omitempty"`
	URL         string    `xml:"xmlUrl,attr,omitempty"`
	Folders     []outline `xml:"outline,omitempty"`
}

type internalOpmlType struct {
	XMLName xml.Name  `xml:"opml"`
	Version string    `xml:"version,attr"`
	Header  header    `xml:"head"`
	Body    []outline `xml:"body>outline"`
}

// Opml is a nested tree of folders.
type Opml struct {
	Header  header
	Folders models.Folder
}

func createOpmlObject(oi *internalOpmlType) *Opml {
	return &Opml{
		Header:  oi.Header,
		Folders: parseOutline(oi.Body),
	}
}

// parseOutline converts "outline" objects into Folder and Feed objects.
func parseOutline(children []outline) models.Folder {
	folder := models.Folder{
		Name: models.RootFolder,
	}

	for _, o := range children {
		if o.URL != "" {
			// This entity has a URL so we presume it is a Feed and has no children.
			var name string
			switch {
			case o.Title != "":
				name = o.Title
			case o.Text != "":
				name = o.Text
			default:
				name = "<untitled>"
			}

			feed := models.Feed{
				Title:       name,
				Description: o.Description,
				URL:         o.URL,
			}
			folder.Feed = append(folder.Feed, feed)
		} else {
			// This entity has no URL so we presume it is a Folder.
			var name string
			switch {
			case o.Title != "":
				name = o.Title
			case o.Text != "":
				name = o.Text
			default:
				name = "<untitled>"
			}

			child := parseOutline(o.Folders)
			child.Name = name
			folder.Folders = append(folder.Folders, child)
		}
	}

	return folder
}

// parseFolders converts a rooted folder into a list of outline objects.
func parseFolders(folder models.Folder) []outline {
	var outlines []outline

	for _, feed := range folder.Feed {
		outlines = append(outlines, outline{Type: "rss", Text: feed.Title, URL: feed.URL})
	}

	for _, child := range folder.Folders {
		o := outline{Text: child.Name}
		o.Folders = parseFolders(child)
		outlines = append(outlines, o)
	}

	return outlines
}

// ParseOpml open a file and returns a parsed OPML object.
func ParseOpml(filename string) (*Opml, error) {
	log.Infof("Loading OPML file from %s", filename)

	f, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	r := bytes.NewReader(f)
	d := xml.NewDecoder(r)
	d.CharsetReader = charset.NewReaderLabel

	oi := new(internalOpmlType)
	err = d.Decode(&oi)
	if err != nil {
		return nil, err
	}

	return createOpmlObject(oi), nil
}

// ExportOpml exports the given folder (with associated feeds and child folders)
// to a file of the given filename in OPML format.
// The file is created with 0777 mode if it does not exist.
func ExportOpml(tree *models.Folder, filename string) error {
	// OPML specifies that time fields conform to RFC822.
	exportTime := time.Now().Format(time.RFC822)
	export := &internalOpmlType{Version: "2.0"}
	export.Header = header{Title: "Goliath Feed Export", DateCreated: exportTime}
	export.Body = parseFolders(*tree)

	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, syscall.S_IRWXU|syscall.S_IRWXG|syscall.S_IRWXO)
	if err != nil {
		return err
	}
	defer f.Close()

	f.WriteString(xml.Header)
	e := xml.NewEncoder(f)
	e.Indent("", "  ")
	return e.Encode(export)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxReaders(math.MaxInt64)
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-badger", func(b *testing.B) {
		totalFound, totalNotFound, totalErr = 0, 0, 0
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64
			var item badger.KVItem
			for pb.Next() {
				key := newKey()
				err := bdb.Get(key, &item)
				if err != nil {
					errored++
				} else if item.Value() != nil {
					found++
				} else {
					notFound++
				}
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("badger %d keys had valid values.", totalFound)
	b.Logf("badger %d keys had no values", totalNotFound)
	b.Logf("badger %d keys had errors", totalErr)
	b.Logf("badger %d total keys looked at", totalFound+totalNotFound+totalErr)
	b.Logf("badger hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var totalCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if _, err := rdb.Get(key); err == nil {
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
	})
	b.Logf("rocks %d keys had valid values.", totalCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		totalFound, totalNotFound, totalErr = 0, 0, 0
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64

			for pb.Next() {
				key := newKey()
				err = lmdbEnv.View(func(txn *lmdb.Txn) error {
					_, err := txn.Get(lmdbDBI, key)
					if lmdb.IsNotFound(err) {
						notFound++
						return nil

					} else if err != nil {
						return err
					}
					found++
					return nil
				})
				if err != nil {
					errored++
				}
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("lmdb %d keys had valid values.", totalFound)
	b.Logf("lmdb %d keys had no values", totalNotFound)
	b.Logf("lmdb %d keys had errors", totalErr)
	b.Logf("lmdb %d total keys looked at", totalFound+totalNotFound+totalErr)
	b.Logf("lmdb hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<commit_msg>Set txn.RawRead to true in iterations<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxReaders(math.MaxInt64)
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-badger", func(b *testing.B) {
		totalFound, totalNotFound, totalErr = 0, 0, 0
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64
			var item badger.KVItem
			for pb.Next() {
				key := newKey()
				err := bdb.Get(key, &item)
				if err != nil {
					errored++
				} else if item.Value() != nil {
					found++
				} else {
					notFound++
				}
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("badger %d keys had valid values.", totalFound)
	b.Logf("badger %d keys had no values", totalNotFound)
	b.Logf("badger %d keys had errors", totalErr)
	b.Logf("badger %d total keys looked at", totalFound+totalNotFound+totalErr)
	b.Logf("badger hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var totalCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if _, err := rdb.Get(key); err == nil {
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
	})
	b.Logf("rocks %d keys had valid values.", totalCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		totalFound, totalNotFound, totalErr = 0, 0, 0
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64

			for pb.Next() {
				key := newKey()
				err = lmdbEnv.View(func(txn *lmdb.Txn) error {
					_, err := txn.Get(lmdbDBI, key)
					if lmdb.IsNotFound(err) {
						notFound++
						return nil

					} else if err != nil {
						return err
					}
					found++
					return nil
				})
				if err != nil {
					errored++
				}
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("lmdb %d keys had valid values.", totalFound)
	b.Logf("lmdb %d keys had no values", totalNotFound)
	b.Logf("lmdb %d keys had errors", totalErr)
	b.Logf("lmdb %d total keys looked at", totalFound+totalNotFound+totalErr)
	b.Logf("lmdb hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				txn.RawRead = true
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package nats

import (
	"testing"
	"time"
)

func BenchmarkPublishSpeed(b *testing.B) {
	b.StopTimer()
	server := startServer(b, DefaultPort, "")
	defer server.stopServer()
	nc, err := Connect(DefaultURL)
	if err != nil {
		b.Fatalf("Could not connect: %v\n", err)
	}
	defer nc.Close()
	b.StartTimer()
	b.SetBytes(1)

	msg := []byte("Hello World")

	for i := 0; i < b.N; i++ {
		err = nc.Publish("foo", msg)
		if err != nil {
			b.Fatalf("Error in benchmark during Publish: %v\n", err)
		}
	}
	// Make sure they are all processed.
	nc.Flush()
	b.StopTimer()
}

func BenchmarkPubSubSpeed(b *testing.B) {
	b.StopTimer()
	server := startServer(b, DefaultPort, "")
	defer server.stopServer()
	nc, err := Connect(DefaultURL)
	if err != nil {
		b.Fatalf("Could not connect: %v\n", err)
	}
	defer nc.Close()

	ch := make(chan bool)
	b.StartTimer()
	b.SetBytes(2)

	nc.Opts.AsyncErrorCB = func(nc *Conn, s *Subscription, err error) {
		b.Fatalf("Error : %v\n", err)
	}

	received := 0

	nc.Subscribe("foo", func(m *Msg) {
		received += 1
		if received >= b.N {
			ch <- true
		}
	})

	msg := []byte("Hello World")

	for i := 0; i < b.N; i++ {
		err = nc.Publish("foo", msg)
		if err != nil {
			b.Fatalf("Error in benchmark during Publish: %v\n", err)
		}
		// Don't overrun ourselves and be a slow consumer
		if i % 1000 == 0 {
			time.Sleep(10)
		}
	}

	// Make sure they are all processed.
	err = waitTime(ch, 10*time.Second)
	if err != nil {
		b.Fatal("Timed out waiting for messages")
	} else if received != b.N {
		b.Fatal(nc.LastError())
	}
	b.StopTimer()
}
<commit_msg>Removed helper setBytes<commit_after>package nats

import (
	"testing"
	"time"
)

func BenchmarkPublishSpeed(b *testing.B) {
	b.StopTimer()
	server := startServer(b, DefaultPort, "")
	defer server.stopServer()
	nc, err := Connect(DefaultURL)
	if err != nil {
		b.Fatalf("Could not connect: %v\n", err)
	}
	defer nc.Close()
	b.StartTimer()

	msg := []byte("Hello World")

	for i := 0; i < b.N; i++ {
		err = nc.Publish("foo", msg)
		if err != nil {
			b.Fatalf("Error in benchmark during Publish: %v\n", err)
		}
	}
	// Make sure they are all processed.
	nc.Flush()
	b.StopTimer()
}

func BenchmarkPubSubSpeed(b *testing.B) {
	b.StopTimer()
	server := startServer(b, DefaultPort, "")
	defer server.stopServer()
	nc, err := Connect(DefaultURL)
	if err != nil {
		b.Fatalf("Could not connect: %v\n", err)
	}
	defer nc.Close()

	ch := make(chan bool)
	b.StartTimer()

	nc.Opts.AsyncErrorCB = func(nc *Conn, s *Subscription, err error) {
		b.Fatalf("Error : %v\n", err)
	}

	received := 0

	nc.Subscribe("foo", func(m *Msg) {
		received += 1
		if received >= b.N {
			ch <- true
		}
	})

	msg := []byte("Hello World")

	for i := 0; i < b.N; i++ {
		err = nc.Publish("foo", msg)
		if err != nil {
			b.Fatalf("Error in benchmark during Publish: %v\n", err)
		}
		// Don't overrun ourselves and be a slow consumer
		if i % 1000 == 0 {
			time.Sleep(10)
		}
	}

	// Make sure they are all processed.
	err = waitTime(ch, 10*time.Second)
	if err != nil {
		b.Fatal("Timed out waiting for messages")
	} else if received != b.N {
		b.Fatal(nc.LastError())
	}
	b.StopTimer()
}
<|endoftext|>


<commit_before>package box

// TODO(ttacon):some of these fields pop up everywhere, make
// common struct and anonymously extend the others with it?
type Collaboration struct {
	Type           string  `json:"type"`
	ID             string  `json:"id"`
	CreatedBy      *Item   `json:"created_by"`  // TODO(ttacon): this should be user
	CreatedAt      string  `json:"created_at"`  // TODO(ttacon): transition this to time.Time
	ModifiedAt     string  `json:"modified_at"` // TODO(ttacon): transition to time.Time
	ExpiresAt      *string `json:"expires_at"`  // TODO(ttacon): *time.Time
	Status         string  `json:"status"`
	AccessibleBy   *Item   `json:"accessible_by"`   // TODO(ttacon): turn into user
	Role           string  `json:"role"`            // TODO(ttacon): enum
	AcknowledgedAt string  `json:"acknowledged_at"` // TODO(ttacon): time.Time
	Item           *Item   `json:"item"`            // TODO(ttacon): mini-folder struct
}

type Collaborations struct {
	TotalCount int `json:"total_count"`
	Entries    []*Collaboration
}
<commit_msg>Specify more<commit_after>package box

// TODO(ttacon):some of these fields pop up everywhere, make
// common struct and anonymously extend the others with it?
type Collaboration struct {
	Type           string  `json:"type"`
	ID             string  `json:"id"`
	CreatedBy      *Item   `json:"created_by"`  // TODO(ttacon): this should be user
	CreatedAt      string  `json:"created_at"`  // TODO(ttacon): transition this to time.Time
	ModifiedAt     string  `json:"modified_at"` // TODO(ttacon): transition to time.Time
	ExpiresAt      *string `json:"expires_at"`  // TODO(ttacon): *time.Time
	Status         string  `json:"status"`
	AccessibleBy   *Item   `json:"accessible_by"`   // TODO(ttacon): turn into user
	Role           string  `json:"role"`            // TODO(ttacon): enum (own file?)
	AcknowledgedAt string  `json:"acknowledged_at"` // TODO(ttacon): time.Time
	Item           *Item   `json:"item"`            // TODO(ttacon): mini-folder struct
}

type Collaborations struct {
	TotalCount int `json:"total_count"`
	Entries    []*Collaboration
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"log"
)

func main() {

    var (
        classID int
        ipSlice []string
    )

    db, err := sql.Open("sqlite3", "./watcher.db")
    if err != nil {
        log.Fatal(err)
    }
    defer db.Close()

    sqlClass := "SELECT class.id FROM class WHERE class.name = 'DNCC'"

    err = db.QueryRow(sqlClass).Scan(&classID)
    if err != nil {
        log.Fatal(err)
    }

    sqlStmt := "select hosts.name, hosts.ip from hosts where hosts.class = ?  "

    rows, err := db.Query(sqlStmt, classID)
    if err != nil {
        log.Fatal(err)
    }
    defer rows.Close()

    for rows.Next() {
        var name, ip string
        rows.Scan(&name, &ip)
        ipSlice = append(ipSlice, ip)
        printSlice("ipSlice", ipSlice)
        fmt.Println(name, ip)
    }

}

func printSlice(s string, x[] string) {
    fmt.Printf("%s len=%d cap=%d %v\n", s, len(x), cap(x), x)
}
<commit_msg>Most recent attempt to pull SNMP data from DNCCs<commit_after>package main

import (
	"database/sql"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
    "github.com/davecgh/go-spew/spew"
	"log"
  g "github.com/soniah/gosnmp"
    "runtime"
    "time"
    "strings"
)

type host struct {
    Name string
    Addr string
    SnmpVersion int
    Site int
}

var verbose bool = false

func main() {

    var (
        classID int
        ipSlice []string
        oidSlice []string
        oidMap map[string]string
        hostSlice []host
    )

    oidMap = make(map[string]string)
    db, err := sql.Open("sqlite3", "./watcher.db")
    if err != nil {
        log.Fatal(err)
    }
    defer db.Close()

    //
    // Extract the ClassID we want from the "class" table
    //
    sqlClass := "SELECT class.classID FROM class WHERE class.name = 'DNCC'"

    err = db.QueryRow(sqlClass).Scan(&classID)
    if err != nil {
        log.Fatal(err)
    }

    if verbose { fmt.Println("The class ID we want is ", classID) }

    //
    // Extract  the IP addresses belonging to the class in question
    //
    sqlStmt := "select host.name, host.ip, host.snmpVersion, host.site  from host where host.class = ?  "

    rows, err := db.Query(sqlStmt, classID)
    if err != nil {
        log.Fatal(err)
    }
    defer rows.Close()

    for rows.Next() {
        var name, ip string
        var version, site int
        rows.Scan(&name, &ip, &version, &site)
        h := host{ Name: name, Addr: ip, SnmpVersion: version, Site: site}
        hostSlice = append(hostSlice, h)
        if verbose {fmt.Printf("We are adding %s (%s)\n", name, ip) }
        ipSlice = append(ipSlice, ip)
        if verbose {
            printSlice("ipSlice", ipSlice)
            fmt.Println(name, ip)
        }
    }

    sqlOIDs := "select oid, name from test_master where class = ? "

    rows, err = db.Query(sqlOIDs, classID)
    if err != nil {
        log.Fatal(err)
    }
    defer rows.Close()

    if verbose { fmt.Println("We have read test_master. The error return value was ", err) }

    for rows.Next() {
        var oid, name string
        rows.Scan(&oid, &name)
        if verbose { fmt.Println("We got '", oid, "' (", name, ")") }
        oidSlice = append(oidSlice, oid)
        oidMap[oid] = name
    }

    if verbose {
	    fmt.Printf("After the database access, we have addresses of %v\n", ipSlice)
	    fmt.Printf("Each device will read OIDs: \n")
	    printSlice("oidSlice", oidSlice)
	    fmt.Println("The oidMap contains:\n")
	    spew.Dump(oidMap)
    }

    procs := 0
    for _, h := range hostSlice {
        procs++
        c := make(chan bool, 5)
        getData(h, oidSlice, c)
        d := <-c
        if d == true {
            procs--
        }
    }

    //time.Sleep(10 * time.Second)
}

func printSlice(s string, x[] string) {
    fmt.Printf("%s len=%d cap=%d %v\n", s, len(x), cap(x), x)
}

func getData(h host, oids []string, c chan bool) {
    fmt.Printf("This is getData, polling device %s. Oids are:\n", h.Addr)
    //spew.Dump(oids)
    var result []g.SnmpPDU
    result = make([]g.SnmpPDU, 1)

    g.Default.Target = h.Addr
    g.Default.Version = 0x1
    g.Default.Timeout = time.Duration(2) * time.Second
    fmt.Printf("DEBUG: community is %s\n", g.Default.Community)
    fmt.Printf("DEBUG: version is %v\n", g.Default.Version)
    err := g.Default.Connect()
    if err != nil {
        log.Fatalf("gosnmp.Connect err: %v", err)
        fmt.Printf("gosnmp.Connect err: %v", err)
    }
    runtime.Gosched()
    defer g.Default.Conn.Close()
    for _, o := range oids {
        fmt.Printf("About to BulkWalkAll %s on host %s\n", o, h.addr)
        result, err = g.Default.BulkWalkAll(o)
        if err != nil {
            log.Fatalf("gosnmp.BulkWalkAll failed: %v", err)
            fmt.Printf("gosnmp.BulkWalkAll failed: %v", err)
        } else {
            //fmt.Printf("Results for oid ", o)
            //spew.Dump(result)
            for _, r := range result {
                v := g.ToBigInt(r.Value)
                fmt.Println("oid ", r.Name, ", value ", v)
                ndxs := strings.Split(r.Name[len(o)+1:], ",")
                ndxCount := 0
                tags := map[string]string {
                    "host": h.Name,
                    "site": h.Site,
                    "name": oSlice[o].Name,
                }
                for i := range ndxs {
                    tagName = fmt.Sprintf("ndx%d", ndxCount)
                    ndxCount++
                    tags[tagname] = i
                }
                fields := map[string]interface{} {
                    "value": v,
                }
            }
        }
        //runtime.Gosched()
        //c <- true
    }
}
<|endoftext|>


<commit_before>package windows_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"github.com/cloudfoundry/bosh-agent/agent/action"
	"github.com/cloudfoundry/bosh-agent/integration/windows/utils"
	boshfileutil "github.com/cloudfoundry/bosh-utils/fileutil"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	agentGUID       = "123-456-789"
	agentID         = "agent." + agentGUID
	senderID        = "director.987-654-321"
	DefaultTimeout  = time.Second * 30
	DefaultInterval = time.Second
)

func natsIP() string {
	if ip := os.Getenv("NATS_PRIVATE_IP"); ip != "" {
		return ip
	}
	return ""
}

func natsURI() string {
	if VagrantProvider == "aws" {
		return fmt.Sprintf("nats://%s:4222", os.Getenv("NATS_ELASTIC_IP"))
	}
	return fmt.Sprintf("nats://%s:4222", natsIP())

}

func blobstoreURI() string {
	if VagrantProvider == "aws" {
		return fmt.Sprintf("http://%s:25250", os.Getenv("NATS_ELASTIC_IP"))
	}
	return fmt.Sprintf("http://%s:25250", natsIP())
}

func getNetworkProperty(key string, natsClient *NatsClient) string {
	message := fmt.Sprintf(`{"method":"get_state","arguments":["full"],"reply_to":"%s"}`, senderID)
	rawResponse, err := natsClient.SendRawMessage(message)
	Expect(err).NotTo(HaveOccurred())

	response := map[string]action.GetStateV1ApplySpec{}
	err = json.Unmarshal(rawResponse, &response)
	Expect(err).NotTo(HaveOccurred())

	for _, spec := range response["value"].NetworkSpecs {
		field, ok := spec.Fields[key]
		if !ok {
			return ""
		}
		if val, ok := field.(string); ok {
			return val
		}
	}
	return ""
}

var _ = Describe("An Agent running on Windows", func() {
	var (
		fs              boshsys.FileSystem
		natsClient      *NatsClient
		blobstoreClient utils.BlobClient
	)

	BeforeEach(func() {
		message := fmt.Sprintf(`{"method":"ping","arguments":[],"reply_to":"%s"}`, senderID)

		blobstoreClient = utils.NewBlobstore(blobstoreURI())

		logger := boshlog.NewLogger(boshlog.LevelNone)
		cmdRunner := boshsys.NewExecCmdRunner(logger)
		fs = boshsys.NewOsFileSystem(logger)
		compressor := boshfileutil.NewTarballCompressor(cmdRunner, fs)

		natsClient = NewNatsClient(compressor, blobstoreClient)
		err := natsClient.Setup()
		Expect(err).NotTo(HaveOccurred())

		testPing := func() (string, error) {
			response, err := natsClient.SendRawMessage(message)
			return string(response), err
		}

		Eventually(testPing, DefaultTimeout, DefaultInterval).Should(Equal(`{"value":"pong"}`))
	})

	AfterEach(func() {
		natsClient.Cleanup()
	})

	It("responds to 'get_state' message over NATS", func() {
		getStateSpecAgentID := func() string {
			message := fmt.Sprintf(`{"method":"get_state","arguments":[],"reply_to":"%s"}`, senderID)
			rawResponse, err := natsClient.SendRawMessage(message)

			response := map[string]action.GetStateV1ApplySpec{}
			err = json.Unmarshal(rawResponse, &response)
			Expect(err).NotTo(HaveOccurred())

			return response["value"].AgentID
		}

		Eventually(getStateSpecAgentID, DefaultTimeout, DefaultInterval).Should(Equal(agentGUID))
	})

	It("includes memory vitals in the 'get_state' response", func() {
		message := fmt.Sprintf(`{"method":"get_state","arguments":["full"],"reply_to":"%s"}`, senderID)

		checkVitals := func() error {
			b, err := natsClient.SendRawMessage(message)
			if err != nil {
				return err
			}

			var res map[string]action.GetStateV1ApplySpec
			if err := json.Unmarshal(b, &res); err != nil {
				return err
			}

			v := res["value"].Vitals
			if v == nil {
				return fmt.Errorf("nil Vitals")
			}
			if v.Mem.Kb == "" || v.Mem.Percent == "" {
				return fmt.Errorf("Empty Memory Vitals: %+v", v.Mem)
			}
			return nil
		}

		Eventually(checkVitals, DefaultTimeout, DefaultInterval).Should(Succeed())
	})

	It("can run a run_errand action", func() {
		natsClient.PrepareJob("say-hello")

		runErrandResponse, err := natsClient.RunErrand()
		Expect(err).NotTo(HaveOccurred())

		runErrandCheck := natsClient.CheckErrandResultStatus(runErrandResponse["value"]["agent_task_id"])
		Eventually(runErrandCheck, DefaultTimeout, DefaultInterval).Should(Equal(action.ErrandResult{
			Stdout:     "hello world\r\n",
			ExitStatus: 0,
		}))
	})

	It("can start a job", func() {
		natsClient.PrepareJob("say-hello")

		runStartResponse, err := natsClient.RunStart()
		Expect(err).NotTo(HaveOccurred())
		Expect(runStartResponse["value"]).To(Equal("started"))

		agentState := natsClient.GetState()
		Expect(agentState.JobState).To(Equal("running"))
	})

	It("can run a drain script", func() {
		natsClient.PrepareJob("say-hello")

		err := natsClient.RunDrain()
		Expect(err).NotTo(HaveOccurred())

		logsDir, err := fs.TempDir("windows-agent-drain-test")
		Expect(err).NotTo(HaveOccurred())
		defer fs.RemoveAll(logsDir)

		natsClient.FetchLogs(logsDir)

		drainLogContents, err := fs.ReadFileString(filepath.Join(logsDir, "say-hello", "drain.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(drainLogContents).To(ContainSubstring("Hello from drain"))
	})

	It("can unmonitor the job during drain script", func() {
		natsClient.PrepareJob("unmonitor-hello")

		runStartResponse, err := natsClient.RunStart()
		Expect(err).NotTo(HaveOccurred())
		Expect(runStartResponse["value"]).To(Equal("started"))

		agentState := natsClient.GetState()
		Expect(agentState.JobState).To(Equal("running"))

		err = natsClient.RunDrain()
		Expect(err).NotTo(HaveOccurred())

		logsDir, err := fs.TempDir("windows-agent-drain-test")
		Expect(err).NotTo(HaveOccurred())
		defer fs.RemoveAll(logsDir)

		natsClient.FetchLogs(logsDir)

		drainLogContents, err := fs.ReadFileString(filepath.Join(logsDir, "unmonitor-hello", "drain.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(drainLogContents).To(ContainSubstring("success"))
	})

	It("stops alerting failing jobs when job is stopped", func() {
		natsClient.PrepareJob("crashes-on-start")
		runStartResponse, err := natsClient.RunStart()
		Expect(err).NotTo(HaveOccurred())
		Expect(runStartResponse["value"]).To(Equal("started"))

		Eventually(func() string {
			return natsClient.GetState().JobState
		}, DefaultTimeout, DefaultInterval).Should(Equal("failing"))

		Eventually(func() (string, error) {
			alert, err := natsClient.GetNextAlert(10 * time.Second)
			if err != nil {
				return "", err
			}
			return alert.Title, nil
		}).Should(MatchRegexp(`crash-service \(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\) - pid failed - Start`))

		err = natsClient.RunStop()
		Expect(err).NotTo(HaveOccurred())

		_, err = natsClient.GetNextAlert(10 * time.Second)
		Expect(err).To(HaveOccurred())
		Expect(err).To(MatchError("nats: timeout"))
	})

	It("can run arbitrary user scripts", func() {
		natsClient.PrepareJob("say-hello")

		err := natsClient.RunScript("pre-start")
		Expect(err).NotTo(HaveOccurred())

		logsDir, err := fs.TempDir("windows-agent-prestart-test")
		Expect(err).NotTo(HaveOccurred())
		defer fs.RemoveAll(logsDir)

		natsClient.FetchLogs(logsDir)

		prestartStdoutContents, err := fs.ReadFileString(filepath.Join(logsDir, "say-hello", "pre-start.stdout.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(prestartStdoutContents).To(ContainSubstring("Hello from stdout"))

		prestartStderrContents, err := fs.ReadFileString(filepath.Join(logsDir, "say-hello", "pre-start.stderr.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(prestartStderrContents).To(ContainSubstring("Hello from stderr"))
	})

	It("can compile packages", func() {
		const (
			blobName     = "blob.tar"
			fileName     = "output.txt"
			fileContents = "i'm a compiled package!"
		)
		result, err := natsClient.CompilePackage("simple-package")
		Expect(err).NotTo(HaveOccurred())

		tempDir, err := fs.TempDir("windows-agent-compile-test")
		Expect(err).NotTo(HaveOccurred())

		path := filepath.Join(tempDir, blobName)
		Expect(blobstoreClient.Get(result.BlobstoreID, path)).To(Succeed())

		tarPath, err := ioutil.TempDir("", "")
		Expect(err).NotTo(HaveOccurred())
		defer os.Remove(tarPath)

		err = exec.Command("tar", "xf", path, "-C", tarPath).Run()
		Expect(err).NotTo(HaveOccurred())

		out, err := ioutil.ReadFile(filepath.Join(tarPath, fileName))
		Expect(err).NotTo(HaveOccurred())
		Expect(string(out)).To(ContainSubstring(fileContents))
	})

	It("Includes the default IP in the 'get_state' response", func() {
		getNetwork := func(key string) string {
			return getNetworkProperty(key, natsClient)
		}
		Eventually(getNetwork("ip"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
		Eventually(getNetwork("gateway"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
		Eventually(getNetwork("netmask"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
	})

	It("can compile longpath complex pakcage", func() {
		const (
			blobName     = "blob.tar"
			fileName     = "output.txt"
			fileContents = "i'm a compiled package!"
		)
		_, err := natsClient.CompilePackage("longpath-package")
		Expect(err).NotTo(HaveOccurred())
	})

	Context("when there are multiple networks", func() {
		BeforeEach(func() {
			if OsVersion == "2012R2" {
				Skip("Cannot create virtual interfaces in 2012R2")
			}
			natsClient.PrepareJob("add-network-interface")

			err := natsClient.RunScript("run")
			Expect(err).NotTo(HaveOccurred())
		})

		AfterEach(func() {
			natsClient.PrepareJob("remove-network-interface")

			err := natsClient.RunScript("run")
			Expect(err).NotTo(HaveOccurred())
		})

		It("chooses correct IP for default network", func() {
			getNetwork := func(key string) string {
				return getNetworkProperty(key, natsClient)
			}
			Eventually(getNetwork("ip"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
			Expect(getNetwork("ip")).ToNot(HavePrefix("172."))
		})
	})
})
<commit_msg>Increase the default timeout in windows tests to 60 seconds<commit_after>package windows_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"github.com/cloudfoundry/bosh-agent/agent/action"
	"github.com/cloudfoundry/bosh-agent/integration/windows/utils"
	boshfileutil "github.com/cloudfoundry/bosh-utils/fileutil"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const (
	agentGUID       = "123-456-789"
	agentID         = "agent." + agentGUID
	senderID        = "director.987-654-321"
	DefaultTimeout  = time.Second * 60
	DefaultInterval = time.Second
)

func natsIP() string {
	if ip := os.Getenv("NATS_PRIVATE_IP"); ip != "" {
		return ip
	}
	return ""
}

func natsURI() string {
	if VagrantProvider == "aws" {
		return fmt.Sprintf("nats://%s:4222", os.Getenv("NATS_ELASTIC_IP"))
	}
	return fmt.Sprintf("nats://%s:4222", natsIP())

}

func blobstoreURI() string {
	if VagrantProvider == "aws" {
		return fmt.Sprintf("http://%s:25250", os.Getenv("NATS_ELASTIC_IP"))
	}
	return fmt.Sprintf("http://%s:25250", natsIP())
}

func getNetworkProperty(key string, natsClient *NatsClient) string {
	message := fmt.Sprintf(`{"method":"get_state","arguments":["full"],"reply_to":"%s"}`, senderID)
	rawResponse, err := natsClient.SendRawMessage(message)
	Expect(err).NotTo(HaveOccurred())

	response := map[string]action.GetStateV1ApplySpec{}
	err = json.Unmarshal(rawResponse, &response)
	Expect(err).NotTo(HaveOccurred())

	for _, spec := range response["value"].NetworkSpecs {
		field, ok := spec.Fields[key]
		if !ok {
			return ""
		}
		if val, ok := field.(string); ok {
			return val
		}
	}
	return ""
}

var _ = Describe("An Agent running on Windows", func() {
	var (
		fs              boshsys.FileSystem
		natsClient      *NatsClient
		blobstoreClient utils.BlobClient
	)

	BeforeEach(func() {
		message := fmt.Sprintf(`{"method":"ping","arguments":[],"reply_to":"%s"}`, senderID)

		blobstoreClient = utils.NewBlobstore(blobstoreURI())

		logger := boshlog.NewLogger(boshlog.LevelNone)
		cmdRunner := boshsys.NewExecCmdRunner(logger)
		fs = boshsys.NewOsFileSystem(logger)
		compressor := boshfileutil.NewTarballCompressor(cmdRunner, fs)

		natsClient = NewNatsClient(compressor, blobstoreClient)
		err := natsClient.Setup()
		Expect(err).NotTo(HaveOccurred())

		testPing := func() (string, error) {
			response, err := natsClient.SendRawMessage(message)
			return string(response), err
		}

		Eventually(testPing, DefaultTimeout, DefaultInterval).Should(Equal(`{"value":"pong"}`))
	})

	AfterEach(func() {
		natsClient.Cleanup()
	})

	It("responds to 'get_state' message over NATS", func() {
		getStateSpecAgentID := func() string {
			message := fmt.Sprintf(`{"method":"get_state","arguments":[],"reply_to":"%s"}`, senderID)
			rawResponse, err := natsClient.SendRawMessage(message)

			response := map[string]action.GetStateV1ApplySpec{}
			err = json.Unmarshal(rawResponse, &response)
			Expect(err).NotTo(HaveOccurred())

			return response["value"].AgentID
		}

		Eventually(getStateSpecAgentID, DefaultTimeout, DefaultInterval).Should(Equal(agentGUID))
	})

	It("includes memory vitals in the 'get_state' response", func() {
		message := fmt.Sprintf(`{"method":"get_state","arguments":["full"],"reply_to":"%s"}`, senderID)

		checkVitals := func() error {
			b, err := natsClient.SendRawMessage(message)
			if err != nil {
				return err
			}

			var res map[string]action.GetStateV1ApplySpec
			if err := json.Unmarshal(b, &res); err != nil {
				return err
			}

			v := res["value"].Vitals
			if v == nil {
				return fmt.Errorf("nil Vitals")
			}
			if v.Mem.Kb == "" || v.Mem.Percent == "" {
				return fmt.Errorf("Empty Memory Vitals: %+v", v.Mem)
			}
			return nil
		}

		Eventually(checkVitals, DefaultTimeout, DefaultInterval).Should(Succeed())
	})

	It("can run a run_errand action", func() {
		natsClient.PrepareJob("say-hello")

		runErrandResponse, err := natsClient.RunErrand()
		Expect(err).NotTo(HaveOccurred())

		runErrandCheck := natsClient.CheckErrandResultStatus(runErrandResponse["value"]["agent_task_id"])
		Eventually(runErrandCheck, DefaultTimeout, DefaultInterval).Should(Equal(action.ErrandResult{
			Stdout:     "hello world\r\n",
			ExitStatus: 0,
		}))
	})

	It("can start a job", func() {
		natsClient.PrepareJob("say-hello")

		runStartResponse, err := natsClient.RunStart()
		Expect(err).NotTo(HaveOccurred())
		Expect(runStartResponse["value"]).To(Equal("started"))

		agentState := natsClient.GetState()
		Expect(agentState.JobState).To(Equal("running"))
	})

	It("can run a drain script", func() {
		natsClient.PrepareJob("say-hello")

		err := natsClient.RunDrain()
		Expect(err).NotTo(HaveOccurred())

		logsDir, err := fs.TempDir("windows-agent-drain-test")
		Expect(err).NotTo(HaveOccurred())
		defer fs.RemoveAll(logsDir)

		natsClient.FetchLogs(logsDir)

		drainLogContents, err := fs.ReadFileString(filepath.Join(logsDir, "say-hello", "drain.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(drainLogContents).To(ContainSubstring("Hello from drain"))
	})

	It("can unmonitor the job during drain script", func() {
		natsClient.PrepareJob("unmonitor-hello")

		runStartResponse, err := natsClient.RunStart()
		Expect(err).NotTo(HaveOccurred())
		Expect(runStartResponse["value"]).To(Equal("started"))

		agentState := natsClient.GetState()
		Expect(agentState.JobState).To(Equal("running"))

		err = natsClient.RunDrain()
		Expect(err).NotTo(HaveOccurred())

		logsDir, err := fs.TempDir("windows-agent-drain-test")
		Expect(err).NotTo(HaveOccurred())
		defer fs.RemoveAll(logsDir)

		natsClient.FetchLogs(logsDir)

		drainLogContents, err := fs.ReadFileString(filepath.Join(logsDir, "unmonitor-hello", "drain.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(drainLogContents).To(ContainSubstring("success"))
	})

	It("stops alerting failing jobs when job is stopped", func() {
		natsClient.PrepareJob("crashes-on-start")
		runStartResponse, err := natsClient.RunStart()
		Expect(err).NotTo(HaveOccurred())
		Expect(runStartResponse["value"]).To(Equal("started"))

		Eventually(func() string {
			return natsClient.GetState().JobState
		}, DefaultTimeout, DefaultInterval).Should(Equal("failing"))

		Eventually(func() (string, error) {
			alert, err := natsClient.GetNextAlert(10 * time.Second)
			if err != nil {
				return "", err
			}
			return alert.Title, nil
		}).Should(MatchRegexp(`crash-service \(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\) - pid failed - Start`))

		err = natsClient.RunStop()
		Expect(err).NotTo(HaveOccurred())

		_, err = natsClient.GetNextAlert(10 * time.Second)
		Expect(err).To(HaveOccurred())
		Expect(err).To(MatchError("nats: timeout"))
	})

	It("can run arbitrary user scripts", func() {
		natsClient.PrepareJob("say-hello")

		err := natsClient.RunScript("pre-start")
		Expect(err).NotTo(HaveOccurred())

		logsDir, err := fs.TempDir("windows-agent-prestart-test")
		Expect(err).NotTo(HaveOccurred())
		defer fs.RemoveAll(logsDir)

		natsClient.FetchLogs(logsDir)

		prestartStdoutContents, err := fs.ReadFileString(filepath.Join(logsDir, "say-hello", "pre-start.stdout.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(prestartStdoutContents).To(ContainSubstring("Hello from stdout"))

		prestartStderrContents, err := fs.ReadFileString(filepath.Join(logsDir, "say-hello", "pre-start.stderr.log"))
		Expect(err).NotTo(HaveOccurred())

		Expect(prestartStderrContents).To(ContainSubstring("Hello from stderr"))
	})

	It("can compile packages", func() {
		const (
			blobName     = "blob.tar"
			fileName     = "output.txt"
			fileContents = "i'm a compiled package!"
		)
		result, err := natsClient.CompilePackage("simple-package")
		Expect(err).NotTo(HaveOccurred())

		tempDir, err := fs.TempDir("windows-agent-compile-test")
		Expect(err).NotTo(HaveOccurred())

		path := filepath.Join(tempDir, blobName)
		Expect(blobstoreClient.Get(result.BlobstoreID, path)).To(Succeed())

		tarPath, err := ioutil.TempDir("", "")
		Expect(err).NotTo(HaveOccurred())
		defer os.Remove(tarPath)

		err = exec.Command("tar", "xf", path, "-C", tarPath).Run()
		Expect(err).NotTo(HaveOccurred())

		out, err := ioutil.ReadFile(filepath.Join(tarPath, fileName))
		Expect(err).NotTo(HaveOccurred())
		Expect(string(out)).To(ContainSubstring(fileContents))
	})

	It("Includes the default IP in the 'get_state' response", func() {
		getNetwork := func(key string) string {
			return getNetworkProperty(key, natsClient)
		}
		Eventually(getNetwork("ip"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
		Eventually(getNetwork("gateway"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
		Eventually(getNetwork("netmask"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
	})

	It("can compile longpath complex pakcage", func() {
		const (
			blobName     = "blob.tar"
			fileName     = "output.txt"
			fileContents = "i'm a compiled package!"
		)
		_, err := natsClient.CompilePackage("longpath-package")
		Expect(err).NotTo(HaveOccurred())
	})

	Context("when there are multiple networks", func() {
		BeforeEach(func() {
			if OsVersion == "2012R2" {
				Skip("Cannot create virtual interfaces in 2012R2")
			}
			natsClient.PrepareJob("add-network-interface")

			err := natsClient.RunScript("run")
			Expect(err).NotTo(HaveOccurred())
		})

		AfterEach(func() {
			natsClient.PrepareJob("remove-network-interface")

			err := natsClient.RunScript("run")
			Expect(err).NotTo(HaveOccurred())
		})

		It("chooses correct IP for default network", func() {
			getNetwork := func(key string) string {
				return getNetworkProperty(key, natsClient)
			}
			Eventually(getNetwork("ip"), DefaultTimeout, DefaultInterval).ShouldNot(BeEmpty())
			Expect(getNetwork("ip")).ToNot(HavePrefix("172."))
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/google/go-github/github"
)

type Contribution struct {
	Date time.Time
	Num  int
}

const dateFormat string = "2006-01-02"

func (c *Contribution) UnmarshalJSON(data []byte) error {
	var i []interface{}
	err := json.Unmarshal(data, &i)
	if err != nil {
		return err
	}

	if v, ok := i[0].(string); ok {
		c.Date, err = time.Parse(dateFormat, v)
	}
	if err != nil {
		return err
	}

	if v, ok := i[1].(float64); ok {
		c.Num = int(v)
	}

	return nil
}

func getContributions(user github.User) []Contribution {
	login := github.Stringify(user.Login)
	url := "https://github.com/users/" + login[1:len(login)-1] + "/contributions"
	var contribData []Contribution

	resp, err := http.Get(url)
	defer resp.Body.Close()
	if err != nil {
		log.Fatal(err)
	}

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	if data == nil {
		log.Fatalln("No data returned.")
	}

	err = json.Unmarshal(data, &contribData)
	if err != nil {
		log.Fatal(err)
	}

	return contribData
}

func getOrgMembers(orgName string) []github.User {
	client := github.NewClient(nil)

	users, _, err := client.Organizations.ListMembers(orgName, nil)
	if err != nil {
		log.Fatal(err)
	}

	return users
}

func main() {
	flag.Parse()
	var i interface{} = flag.Arg(0)
	var users []github.User

	var v, ok = i.(string)
	if ok == false || v == "" {
		log.Fatalln("Usage: orgstreak <orgName>")
	} else {
		users = getOrgMembers(v)
	}

	orgContribs := make(map[time.Time]int)
	userContributions := make(chan []Contribution, len(users))

	for _, user := range users {
		go func(u github.User) {
			userContributions <- getContributions(u)
		}(user)
	}

	for y := 0; y < len(users); y++ {
		select {
		case userContribution := <-userContributions:
			for _, contrib := range userContribution {
				orgContribs[contrib.Date] += contrib.Num
			}
		}
	}

	contribData := []interface{}{}
	for k, v := range orgContribs {
		contribData = append(contribData, []interface{}{k.Format(dateFormat), v})
	}
	j, err := json.Marshal(contribData)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s", j)
}
<commit_msg>fix data uri<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/google/go-github/github"
)

type Contribution struct {
	Date time.Time
	Num  int
}

const dateFormat string = "2006-01-02"

func (c *Contribution) UnmarshalJSON(data []byte) error {
	var i []interface{}
	err := json.Unmarshal(data, &i)
	if err != nil {
		return err
	}

	if v, ok := i[0].(string); ok {
		c.Date, err = time.Parse(dateFormat, v)
	}
	if err != nil {
		return err
	}

	if v, ok := i[1].(float64); ok {
		c.Num = int(v)
	}

	return nil
}

func getContributions(user github.User) []Contribution {
	login := github.Stringify(user.Login)
	url := "https://github.com/users/" + login[1:len(login)-1] + "/contributions.json"
	var contribData []Contribution

	resp, err := http.Get(url)
	defer resp.Body.Close()
	if err != nil {
		log.Fatal(err)
	}

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}
	if data == nil {
		log.Fatalln("No data returned.")
	}

	err = json.Unmarshal(data, &contribData)
	if err != nil {
		log.Fatal(err)
	}

	return contribData
}

func getOrgMembers(orgName string) []github.User {
	client := github.NewClient(nil)

	users, _, err := client.Organizations.ListMembers(orgName, nil)
	if err != nil {
		log.Fatal(err)
	}

	return users
}

func main() {
	flag.Parse()
	var i interface{} = flag.Arg(0)
	var users []github.User

	var v, ok = i.(string)
	if ok == false || v == "" {
		log.Fatalln("Usage: orgstreak <orgName>")
	} else {
		users = getOrgMembers(v)
	}

	orgContribs := make(map[time.Time]int)
	userContributions := make(chan []Contribution, len(users))

	for _, user := range users {
		go func(u github.User) {
			userContributions <- getContributions(u)
		}(user)
	}

	for y := 0; y < len(users); y++ {
		select {
		case userContribution := <-userContributions:
			for _, contrib := range userContribution {
				orgContribs[contrib.Date] += contrib.Num
			}
		}
	}

	contribData := []interface{}{}
	for k, v := range orgContribs {
		contribData = append(contribData, []interface{}{k.Format(dateFormat), v})
	}
	j, err := json.Marshal(contribData)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s", j)
}
<|endoftext|>


<commit_before>// Copyright 2014 Docker authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the DOCKER-LICENSE file.

package jsonmessage

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/Nvveen/Gotty"
	"github.com/docker/go-units"
	"github.com/fsouza/go-dockerclient/internal/term"
)

// RFC3339NanoFixed is time.RFC3339Nano with nanoseconds padded using zeros to
// ensure the formatted time isalways the same number of characters.
const RFC3339NanoFixed = "2006-01-02T15:04:05.000000000Z07:00"

// JSONError wraps a concrete Code and Message, `Code` is
// is an integer error code, `Message` is the error message.
type JSONError struct {
	Code    int    `json:"code,omitempty"`
	Message string `json:"message,omitempty"`
}

func (e *JSONError) Error() string {
	return e.Message
}

// JSONProgress describes a Progress. terminalFd is the fd of the current terminal,
// Start is the initial value for the operation. Current is the current status and
// value of the progress made towards Total. Total is the end value describing when
// we made 100% progress for an operation.
type JSONProgress struct {
	terminalFd uintptr
	Current    int64 `json:"current,omitempty"`
	Total      int64 `json:"total,omitempty"`
	Start      int64 `json:"start,omitempty"`
	// If true, don't show xB/yB
	HideCounts bool   `json:"hidecounts,omitempty"`
	Units      string `json:"units,omitempty"`
	nowFunc    func() time.Time
	winSize    int
}

func (p *JSONProgress) String() string {
	var (
		width       = p.width()
		pbBox       string
		numbersBox  string
		timeLeftBox string
	)
	if p.Current <= 0 && p.Total <= 0 {
		return ""
	}
	if p.Total <= 0 {
		switch p.Units {
		case "":
			current := units.HumanSize(float64(p.Current))
			return fmt.Sprintf("%8v", current)
		default:
			return fmt.Sprintf("%d %s", p.Current, p.Units)
		}
	}

	percentage := int(float64(p.Current)/float64(p.Total)*100) / 2
	if percentage > 50 {
		percentage = 50
	}
	if width > 110 {
		// this number can't be negative gh#7136
		numSpaces := 0
		if 50-percentage > 0 {
			numSpaces = 50 - percentage
		}
		pbBox = fmt.Sprintf("[%s>%s] ", strings.Repeat("=", percentage), strings.Repeat(" ", numSpaces))
	}

	switch {
	case p.HideCounts:
	case p.Units == "": // no units, use bytes
		current := units.HumanSize(float64(p.Current))
		total := units.HumanSize(float64(p.Total))

		numbersBox = fmt.Sprintf("%8v/%v", current, total)

		if p.Current > p.Total {
			// remove total display if the reported current is wonky.
			numbersBox = fmt.Sprintf("%8v", current)
		}
	default:
		numbersBox = fmt.Sprintf("%d/%d %s", p.Current, p.Total, p.Units)

		if p.Current > p.Total {
			// remove total display if the reported current is wonky.
			numbersBox = fmt.Sprintf("%d %s", p.Current, p.Units)
		}
	}

	if p.Current > 0 && p.Start > 0 && percentage < 50 {
		fromStart := p.now().Sub(time.Unix(p.Start, 0))
		perEntry := fromStart / time.Duration(p.Current)
		left := time.Duration(p.Total-p.Current) * perEntry
		left = (left / time.Second) * time.Second

		if width > 50 {
			timeLeftBox = " " + left.String()
		}
	}
	return pbBox + numbersBox + timeLeftBox
}

// shim for testing
func (p *JSONProgress) now() time.Time {
	if p.nowFunc == nil {
		p.nowFunc = func() time.Time {
			return time.Now().UTC()
		}
	}
	return p.nowFunc()
}

// shim for testing
func (p *JSONProgress) width() int {
	if p.winSize != 0 {
		return p.winSize
	}
	ws, err := term.GetWinsize(p.terminalFd)
	if err == nil {
		return int(ws.Width)
	}
	return 200
}

// JSONMessage defines a message struct. It describes
// the created time, where it from, status, ID of the
// message. It's used for docker events.
type JSONMessage struct {
	Stream          string        `json:"stream,omitempty"`
	Status          string        `json:"status,omitempty"`
	Progress        *JSONProgress `json:"progressDetail,omitempty"`
	ProgressMessage string        `json:"progress,omitempty"` //deprecated
	ID              string        `json:"id,omitempty"`
	From            string        `json:"from,omitempty"`
	Time            int64         `json:"time,omitempty"`
	TimeNano        int64         `json:"timeNano,omitempty"`
	Error           *JSONError    `json:"errorDetail,omitempty"`
	ErrorMessage    string        `json:"error,omitempty"` //deprecated
	// Aux contains out-of-band data, such as digests for push signing and image id after building.
	Aux *json.RawMessage `json:"aux,omitempty"`
}

/* Satisfied by gotty.TermInfo as well as noTermInfo from below */
type termInfo interface {
	Parse(attr string, params ...interface{}) (string, error)
}

type noTermInfo struct{} // canary used when no terminfo.

func (ti *noTermInfo) Parse(attr string, params ...interface{}) (string, error) {
	return "", fmt.Errorf("noTermInfo")
}

func clearLine(out io.Writer, ti termInfo) {
	// el2 (clear whole line) is not exposed by terminfo.

	// First clear line from beginning to cursor
	if attr, err := ti.Parse("el1"); err == nil {
		fmt.Fprintf(out, "%s", attr)
	} else {
		fmt.Fprintf(out, "\x1b[1K")
	}
	// Then clear line from cursor to end
	if attr, err := ti.Parse("el"); err == nil {
		fmt.Fprintf(out, "%s", attr)
	} else {
		fmt.Fprintf(out, "\x1b[K")
	}
}

func cursorUp(out io.Writer, ti termInfo, l int) {
	if l == 0 { // Should never be the case, but be tolerant
		return
	}
	if attr, err := ti.Parse("cuu", l); err == nil {
		fmt.Fprintf(out, "%s", attr)
	} else {
		fmt.Fprintf(out, "\x1b[%dA", l)
	}
}

func cursorDown(out io.Writer, ti termInfo, l int) {
	if l == 0 { // Should never be the case, but be tolerant
		return
	}
	if attr, err := ti.Parse("cud", l); err == nil {
		fmt.Fprintf(out, "%s", attr)
	} else {
		fmt.Fprintf(out, "\x1b[%dB", l)
	}
}

// Display displays the JSONMessage to `out`. `termInfo` is non-nil if `out`
// is a terminal. If this is the case, it will erase the entire current line
// when displaying the progressbar.
func (jm *JSONMessage) Display(out io.Writer, termInfo termInfo) error {
	if jm.Error != nil {
		if jm.Error.Code == 401 {
			return fmt.Errorf("authentication is required")
		}
		return jm.Error
	}
	var endl string
	if termInfo != nil && jm.Stream == "" && jm.Progress != nil {
		clearLine(out, termInfo)
		endl = "\r"
		fmt.Fprintf(out, endl)
	} else if jm.Progress != nil && jm.Progress.String() != "" { //disable progressbar in non-terminal
		return nil
	}
	if jm.TimeNano != 0 {
		fmt.Fprintf(out, "%s ", time.Unix(0, jm.TimeNano).Format(RFC3339NanoFixed))
	} else if jm.Time != 0 {
		fmt.Fprintf(out, "%s ", time.Unix(jm.Time, 0).Format(RFC3339NanoFixed))
	}
	if jm.ID != "" {
		fmt.Fprintf(out, "%s: ", jm.ID)
	}
	if jm.From != "" {
		fmt.Fprintf(out, "(from %s) ", jm.From)
	}
	if jm.Progress != nil && termInfo != nil {
		fmt.Fprintf(out, "%s %s%s", jm.Status, jm.Progress.String(), endl)
	} else if jm.ProgressMessage != "" { //deprecated
		fmt.Fprintf(out, "%s %s%s", jm.Status, jm.ProgressMessage, endl)
	} else if jm.Stream != "" {
		fmt.Fprintf(out, "%s%s", jm.Stream, endl)
	} else {
		fmt.Fprintf(out, "%s%s\n", jm.Status, endl)
	}
	return nil
}

// DisplayJSONMessagesStream displays a json message stream from `in` to `out`, `isTerminal`
// describes if `out` is a terminal. If this is the case, it will print `\n` at the end of
// each line and move the cursor while displaying.
func DisplayJSONMessagesStream(in io.Reader, out io.Writer, terminalFd uintptr, isTerminal bool, auxCallback func(JSONMessage)) error {
	var (
		dec = json.NewDecoder(in)
		ids = make(map[string]int)
	)

	var termInfo termInfo

	if isTerminal {
		term := os.Getenv("TERM")
		if term == "" {
			term = "vt102"
		}

		var err error
		if termInfo, err = gotty.OpenTermInfo(term); err != nil {
			termInfo = &noTermInfo{}
		}
	}

	for {
		diff := 0
		var jm JSONMessage
		if err := dec.Decode(&jm); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}

		if jm.Aux != nil {
			if auxCallback != nil {
				auxCallback(jm)
			}
			continue
		}

		if jm.Progress != nil {
			jm.Progress.terminalFd = terminalFd
		}
		if jm.ID != "" && (jm.Progress != nil || jm.ProgressMessage != "") {
			line, ok := ids[jm.ID]
			if !ok {
				// NOTE: This approach of using len(id) to
				// figure out the number of lines of history
				// only works as long as we clear the history
				// when we output something that's not
				// accounted for in the map, such as a line
				// with no ID.
				line = len(ids)
				ids[jm.ID] = line
				if termInfo != nil {
					fmt.Fprintf(out, "\n")
				}
			}
			diff = len(ids) - line
			if termInfo != nil {
				cursorUp(out, termInfo, diff)
			}
		} else {
			// When outputting something that isn't progress
			// output, clear the history of previous lines. We
			// don't want progress entries from some previous
			// operation to be updated (for example, pull -a
			// with multiple tags).
			ids = make(map[string]int)
		}
		err := jm.Display(out, termInfo)
		if jm.ID != "" && termInfo != nil {
			cursorDown(out, termInfo, diff)
		}
		if err != nil {
			return err
		}
	}
	return nil
}

type stream interface {
	io.Writer
	FD() uintptr
	IsTerminal() bool
}

// DisplayJSONMessagesToStream prints json messages to the output stream
func DisplayJSONMessagesToStream(in io.Reader, stream stream, auxCallback func(JSONMessage)) error {
	return DisplayJSONMessagesStream(in, stream, stream.FD(), stream.IsTerminal(), auxCallback)
}
<commit_msg>Surface errors that occur when performing writes on the io.Writer.<commit_after>// Copyright 2014 Docker authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the DOCKER-LICENSE file.

package jsonmessage

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/Nvveen/Gotty"
	"github.com/docker/go-units"
	"github.com/fsouza/go-dockerclient/internal/term"
)

// RFC3339NanoFixed is time.RFC3339Nano with nanoseconds padded using zeros to
// ensure the formatted time isalways the same number of characters.
const RFC3339NanoFixed = "2006-01-02T15:04:05.000000000Z07:00"

// JSONError wraps a concrete Code and Message, `Code` is
// is an integer error code, `Message` is the error message.
type JSONError struct {
	Code    int    `json:"code,omitempty"`
	Message string `json:"message,omitempty"`
}

func (e *JSONError) Error() string {
	return e.Message
}

// JSONProgress describes a Progress. terminalFd is the fd of the current terminal,
// Start is the initial value for the operation. Current is the current status and
// value of the progress made towards Total. Total is the end value describing when
// we made 100% progress for an operation.
type JSONProgress struct {
	terminalFd uintptr
	Current    int64 `json:"current,omitempty"`
	Total      int64 `json:"total,omitempty"`
	Start      int64 `json:"start,omitempty"`
	// If true, don't show xB/yB
	HideCounts bool   `json:"hidecounts,omitempty"`
	Units      string `json:"units,omitempty"`
	nowFunc    func() time.Time
	winSize    int
}

func (p *JSONProgress) String() string {
	var (
		width       = p.width()
		pbBox       string
		numbersBox  string
		timeLeftBox string
	)
	if p.Current <= 0 && p.Total <= 0 {
		return ""
	}
	if p.Total <= 0 {
		switch p.Units {
		case "":
			current := units.HumanSize(float64(p.Current))
			return fmt.Sprintf("%8v", current)
		default:
			return fmt.Sprintf("%d %s", p.Current, p.Units)
		}
	}

	percentage := int(float64(p.Current)/float64(p.Total)*100) / 2
	if percentage > 50 {
		percentage = 50
	}
	if width > 110 {
		// this number can't be negative gh#7136
		numSpaces := 0
		if 50-percentage > 0 {
			numSpaces = 50 - percentage
		}
		pbBox = fmt.Sprintf("[%s>%s] ", strings.Repeat("=", percentage), strings.Repeat(" ", numSpaces))
	}

	switch {
	case p.HideCounts:
	case p.Units == "": // no units, use bytes
		current := units.HumanSize(float64(p.Current))
		total := units.HumanSize(float64(p.Total))

		numbersBox = fmt.Sprintf("%8v/%v", current, total)

		if p.Current > p.Total {
			// remove total display if the reported current is wonky.
			numbersBox = fmt.Sprintf("%8v", current)
		}
	default:
		numbersBox = fmt.Sprintf("%d/%d %s", p.Current, p.Total, p.Units)

		if p.Current > p.Total {
			// remove total display if the reported current is wonky.
			numbersBox = fmt.Sprintf("%d %s", p.Current, p.Units)
		}
	}

	if p.Current > 0 && p.Start > 0 && percentage < 50 {
		fromStart := p.now().Sub(time.Unix(p.Start, 0))
		perEntry := fromStart / time.Duration(p.Current)
		left := time.Duration(p.Total-p.Current) * perEntry
		left = (left / time.Second) * time.Second

		if width > 50 {
			timeLeftBox = " " + left.String()
		}
	}
	return pbBox + numbersBox + timeLeftBox
}

// shim for testing
func (p *JSONProgress) now() time.Time {
	if p.nowFunc == nil {
		p.nowFunc = func() time.Time {
			return time.Now().UTC()
		}
	}
	return p.nowFunc()
}

// shim for testing
func (p *JSONProgress) width() int {
	if p.winSize != 0 {
		return p.winSize
	}
	ws, err := term.GetWinsize(p.terminalFd)
	if err == nil {
		return int(ws.Width)
	}
	return 200
}

// JSONMessage defines a message struct. It describes
// the created time, where it from, status, ID of the
// message. It's used for docker events.
type JSONMessage struct {
	Stream          string        `json:"stream,omitempty"`
	Status          string        `json:"status,omitempty"`
	Progress        *JSONProgress `json:"progressDetail,omitempty"`
	ProgressMessage string        `json:"progress,omitempty"` //deprecated
	ID              string        `json:"id,omitempty"`
	From            string        `json:"from,omitempty"`
	Time            int64         `json:"time,omitempty"`
	TimeNano        int64         `json:"timeNano,omitempty"`
	Error           *JSONError    `json:"errorDetail,omitempty"`
	ErrorMessage    string        `json:"error,omitempty"` //deprecated
	// Aux contains out-of-band data, such as digests for push signing and image id after building.
	Aux *json.RawMessage `json:"aux,omitempty"`
}

/* Satisfied by gotty.TermInfo as well as noTermInfo from below */
type termInfo interface {
	Parse(attr string, params ...interface{}) (string, error)
}

type noTermInfo struct{} // canary used when no terminfo.

func (ti *noTermInfo) Parse(attr string, params ...interface{}) (string, error) {
	return "", fmt.Errorf("noTermInfo")
}

func clearLine(out io.Writer, ti termInfo) error {
	// el2 (clear whole line) is not exposed by terminfo.

	// First clear line from beginning to cursor
	if attr, err := ti.Parse("el1"); err == nil {
		_, err = fmt.Fprintf(out, "%s", attr)
		if err != nil {
			return err
		}
	} else {
		_, err := fmt.Fprintf(out, "\x1b[1K")
		if err != nil {
			return err
		}
	}
	// Then clear line from cursor to end
	if attr, err := ti.Parse("el"); err == nil {
		_, err = fmt.Fprintf(out, "%s", attr)
		if err != nil {
			return err
		}
	} else {
		_, err := fmt.Fprintf(out, "\x1b[K")
		if err != nil {
			return err
		}
	}

	return nil
}

func cursorUp(out io.Writer, ti termInfo, l int) error {
	if l == 0 { // Should never be the case, but be tolerant
		return nil
	}
	if attr, err := ti.Parse("cuu", l); err == nil {
		_, err = fmt.Fprintf(out, "%s", attr)
		if err != nil {
			return err
		}
	} else {
		_, err := fmt.Fprintf(out, "\x1b[%dA", l)
		if err != nil {
			return err
		}
	}
	return nil
}

func cursorDown(out io.Writer, ti termInfo, l int) error {
	if l == 0 { // Should never be the case, but be tolerant
		return nil
	}
	if attr, err := ti.Parse("cud", l); err == nil {
		_, err = fmt.Fprintf(out, "%s", attr)
		if err != nil {
			return err
		}
	} else {
		_, err := fmt.Fprintf(out, "\x1b[%dB", l)
		if err != nil {
			return err
		}
	}

	return nil
}

// Display displays the JSONMessage to `out`. `termInfo` is non-nil if `out`
// is a terminal. If this is the case, it will erase the entire current line
// when displaying the progressbar.
func (jm *JSONMessage) Display(out io.Writer, termInfo termInfo) error {
	if jm.Error != nil {
		if jm.Error.Code == 401 {
			return fmt.Errorf("authentication is required")
		}
		return jm.Error
	}
	var endl string
	if termInfo != nil && jm.Stream == "" && jm.Progress != nil {
		clearLine(out, termInfo)
		endl = "\r"
		_, err := fmt.Fprintf(out, endl)
		if err != nil {
			return err
		}
	} else if jm.Progress != nil && jm.Progress.String() != "" { //disable progressbar in non-terminal
		return nil
	}
	if jm.TimeNano != 0 {
		_, err := fmt.Fprintf(out, "%s ", time.Unix(0, jm.TimeNano).Format(RFC3339NanoFixed))
		if err != nil {
			return err
		}
	} else if jm.Time != 0 {
		_, err := fmt.Fprintf(out, "%s ", time.Unix(jm.Time, 0).Format(RFC3339NanoFixed))
		if err != nil {
			return err
		}
	}
	if jm.ID != "" {
		_, err := fmt.Fprintf(out, "%s: ", jm.ID)
		if err != nil {
			return err
		}
	}
	if jm.From != "" {
		_, err := fmt.Fprintf(out, "(from %s) ", jm.From)
		if err != nil {
			return err
		}
	}
	if jm.Progress != nil && termInfo != nil {
		_, err := fmt.Fprintf(out, "%s %s%s", jm.Status, jm.Progress.String(), endl)
		if err != nil {
			return err
		}
	} else if jm.ProgressMessage != "" { //deprecated
		_, err := fmt.Fprintf(out, "%s %s%s", jm.Status, jm.ProgressMessage, endl)
		if err != nil {
			return err
		}
	} else if jm.Stream != "" {
		_, err := fmt.Fprintf(out, "%s%s", jm.Stream, endl)
		if err != nil {
			return err
		}
	} else {
		_, err := fmt.Fprintf(out, "%s%s\n", jm.Status, endl)
		if err != nil {
			return err
		}
	}
	return nil
}

// DisplayJSONMessagesStream displays a json message stream from `in` to `out`, `isTerminal`
// describes if `out` is a terminal. If this is the case, it will print `\n` at the end of
// each line and move the cursor while displaying.
func DisplayJSONMessagesStream(in io.Reader, out io.Writer, terminalFd uintptr, isTerminal bool, auxCallback func(JSONMessage)) error {
	var (
		dec = json.NewDecoder(in)
		ids = make(map[string]int)
	)

	var termInfo termInfo

	if isTerminal {
		term := os.Getenv("TERM")
		if term == "" {
			term = "vt102"
		}

		var err error
		if termInfo, err = gotty.OpenTermInfo(term); err != nil {
			termInfo = &noTermInfo{}
		}
	}

	for {
		diff := 0
		var jm JSONMessage
		if err := dec.Decode(&jm); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}

		if jm.Aux != nil {
			if auxCallback != nil {
				auxCallback(jm)
			}
			continue
		}

		if jm.Progress != nil {
			jm.Progress.terminalFd = terminalFd
		}
		if jm.ID != "" && (jm.Progress != nil || jm.ProgressMessage != "") {
			line, ok := ids[jm.ID]
			if !ok {
				// NOTE: This approach of using len(id) to
				// figure out the number of lines of history
				// only works as long as we clear the history
				// when we output something that's not
				// accounted for in the map, such as a line
				// with no ID.
				line = len(ids)
				ids[jm.ID] = line
				if termInfo != nil {
					_, err := fmt.Fprintf(out, "\n")
					if err != nil {
						return err
					}
				}
			}
			diff = len(ids) - line
			if termInfo != nil {
				if err := cursorUp(out, termInfo, diff); err != nil {
					return err
				}
			}
		} else {
			// When outputting something that isn't progress
			// output, clear the history of previous lines. We
			// don't want progress entries from some previous
			// operation to be updated (for example, pull -a
			// with multiple tags).
			ids = make(map[string]int)
		}
		err := jm.Display(out, termInfo)
		if jm.ID != "" && termInfo != nil {
			if err := cursorDown(out, termInfo, diff); err != nil {
				return err
			}
		}
		if err != nil {
			return err
		}
	}
	return nil
}

type stream interface {
	io.Writer
	FD() uintptr
	IsTerminal() bool
}

// DisplayJSONMessagesToStream prints json messages to the output stream
func DisplayJSONMessagesToStream(in io.Reader, stream stream, auxCallback func(JSONMessage)) error {
	return DisplayJSONMessagesStream(in, stream, stream.FD(), stream.IsTerminal(), auxCallback)
}
<|endoftext|>


<commit_before>package p

/*
#include "Python.h"

PyObject* idOrNone(PyObject* o)
{
  return o ? o : Py_BuildValue("");
}

void fetchPythonError(PyObject* excInfo)
{
  PyObject *type, *value, *traceback;

  PyErr_Fetch(&type, &value, &traceback);
  PyTuple_SetItem(excInfo, 0, idOrNone(type));
  PyTuple_SetItem(excInfo, 1, idOrNone(value));
  PyTuple_SetItem(excInfo, 2, idOrNone(traceback));
}
*/
import "C"
import (
	"errors"
	"fmt"
	"runtime"
	"strings"
	"unicode"
	"unsafe"
)

func init() {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	releaseGIL, err := initAndLockPython()
	if err != nil {
		panic(err)
	}
	defer releaseGIL()

	traceback, err := loadModule("traceback")
	if err != nil {
		panic(err)
	}
	defer traceback.decRef()
	if formatException, err := getPyFunc(traceback.p, "format_exception"); err != nil {
		panic(err)
	} else {
		tracebackFormatExceptionFunc = formatException
	}

	exceptions, err := loadModule("exceptions")
	if err != nil {
		panic(err)
	}
	defer exceptions.decRef()
	syntaxErrorCString := C.CString("SyntaxError")
	defer C.free(unsafe.Pointer(syntaxErrorCString))
	if syntaxError := C.PyObject_GetAttrString(exceptions.p, syntaxErrorCString); syntaxError == nil {
		panic("cannot load exceptions.SyntaxError")
	} else {
		syntaxErrorType.p = syntaxError
	}
}

func loadModule(name string) (mod ObjectModule, err error) {
	cName := C.CString(name)
	defer C.free(unsafe.Pointer(cName))
	if m := C.PyImport_ImportModule(cName); m == nil {
		err = fmt.Errorf("failed to import '%v'", name)
	} else {
		mod.p = m
	}
	return
}

var tracebackFormatExceptionFunc ObjectFunc
var syntaxErrorType Object

// tracebackFormatException calls traceback.format_exception().
func tracebackFormatException(excInfo Object) (Object, error) {
	formatted := C.PyObject_CallObject(tracebackFormatExceptionFunc.p, excInfo.p)
	if formatted == nil {
		return Object{p: nil}, errors.New("failed to extract error info from Python")
	}
	return Object{p: formatted}, nil
}

func isSyntaxError(p *C.PyObject) bool {
	return C.PyType_IsSubtype(pyObjectToPyTypeObject(p), pyObjectToPyTypeObject(syntaxErrorType.p)) != 0
}

func pyObjectToPyTypeObject(p *C.PyObject) *C.PyTypeObject {
	return (*C.PyTypeObject)(unsafe.Pointer(p))
}

func getPyErr() error {
	// TODO: consider to reserve excInfo
	excInfo := Object{p: C.PyTuple_New(3)}
	if excInfo.p == nil {
		return errors.New("cannot allocate python's tuple")
	}
	defer excInfo.decRef()
	C.fetchPythonError(excInfo.p)

	formatted, err := tracebackFormatException(excInfo)
	if err != nil {
		return err
	}
	defer formatted.decRef()

	ln := C.PyList_Size(formatted.p)
	mainMsg := strings.TrimSpace(extractLineFromFormattedErrorMessage(formatted, ln-1))
	syntaxErr := ""
	nTracebackLines := ln - 1
	if isSyntaxError(C.PyTuple_GetItem(excInfo.p, 0)) {
		firstLine := extractLineFromFormattedErrorMessage(formatted, ln-3)
		secondLine := extractLineFromFormattedErrorMessage(formatted, ln-2)
		syntaxErr = firstLine + secondLine
		nTracebackLines = ln - 3
	}
	stackTrace := ""
	for i := C.Py_ssize_t(0); i < nTracebackLines; i++ {
		stackTrace += extractLineFromFormattedErrorMessage(formatted, i)
	}
	stackTrace = strings.TrimRightFunc(stackTrace, unicode.IsSpace)
	return &pyErr{
		mainMsg:      mainMsg,
		syntaxErrMsg: syntaxErr,
		stackTrace:   stackTrace,
	}
}

func extractLineFromFormattedErrorMessage(formatted Object, n C.Py_ssize_t) string {
	line := C.PyList_GetItem(formatted.p, n)
	return C.GoString(C.PyString_AsString(line))
}

type pyErr struct {
	mainMsg      string
	syntaxErrMsg string
	stackTrace   string
}

// Error returns an error message string for pyErr.
// This string contains multiple lines for stacktrace.
func (e *pyErr) Error() string {
	return e.mainMsg + "\n" + e.syntaxErrMsg + e.stackTrace
}
<commit_msg>Fix getPyErr() to handle memory allocation error.<commit_after>package p

/*
#include "Python.h"

PyObject* idOrNone(PyObject* o)
{
  return o ? o : Py_BuildValue("");
}

void fetchPythonError(PyObject* excInfo)
{
  PyObject *type, *value, *traceback;

  PyErr_Fetch(&type, &value, &traceback);
  PyTuple_SetItem(excInfo, 0, idOrNone(type));
  PyTuple_SetItem(excInfo, 1, idOrNone(value));
  PyTuple_SetItem(excInfo, 2, idOrNone(traceback));
}
*/
import "C"
import (
	"errors"
	"fmt"
	"runtime"
	"strings"
	"unicode"
	"unsafe"
)

func init() {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	releaseGIL, err := initAndLockPython()
	if err != nil {
		panic(err)
	}
	defer releaseGIL()

	traceback, err := loadModule("traceback")
	if err != nil {
		panic(err)
	}
	defer traceback.decRef()
	if formatException, err := getPyFunc(traceback.p, "format_exception"); err != nil {
		panic(err)
	} else {
		tracebackFormatExceptionFunc = formatException
	}

	exceptions, err := loadModule("exceptions")
	if err != nil {
		panic(err)
	}
	defer exceptions.decRef()
	syntaxErrorCString := C.CString("SyntaxError")
	defer C.free(unsafe.Pointer(syntaxErrorCString))
	if syntaxError := C.PyObject_GetAttrString(exceptions.p, syntaxErrorCString); syntaxError == nil {
		panic("cannot load exceptions.SyntaxError")
	} else {
		syntaxErrorType.p = syntaxError
	}
}

func loadModule(name string) (mod ObjectModule, err error) {
	cName := C.CString(name)
	defer C.free(unsafe.Pointer(cName))
	if m := C.PyImport_ImportModule(cName); m == nil {
		err = fmt.Errorf("failed to import '%v'", name)
	} else {
		mod.p = m
	}
	return
}

var tracebackFormatExceptionFunc ObjectFunc
var syntaxErrorType Object

// tracebackFormatException calls traceback.format_exception().
func tracebackFormatException(excInfo Object) (Object, error) {
	formatted := C.PyObject_CallObject(tracebackFormatExceptionFunc.p, excInfo.p)
	if formatted == nil {
		return Object{p: nil}, errors.New("failed to extract error info from Python")
	}
	return Object{p: formatted}, nil
}

func isSyntaxError(p *C.PyObject) bool {
	return C.PyType_IsSubtype(pyObjectToPyTypeObject(p), pyObjectToPyTypeObject(syntaxErrorType.p)) != 0
}

func pyObjectToPyTypeObject(p *C.PyObject) *C.PyTypeObject {
	return (*C.PyTypeObject)(unsafe.Pointer(p))
}

func getPyErr() error {
	if isPyNoMemoryError() {
		// Fetching stacktrace requires some memory,
		// so just return an error without stacktrace.
		return pyNoMemoryError
	}

	// TODO: consider to reserve excInfo
	excInfo := Object{p: C.PyTuple_New(3)}
	if excInfo.p == nil {
		if isPyNoMemoryError() {
			return pyNoMemoryError
		}
		return getPyErr()
	}
	defer excInfo.decRef()
	C.fetchPythonError(excInfo.p)

	formatted, err := tracebackFormatException(excInfo)
	if err != nil {
		return err
	}
	defer formatted.decRef()

	ln := C.PyList_Size(formatted.p)
	mainMsg := strings.TrimSpace(extractLineFromFormattedErrorMessage(formatted, ln-1))
	syntaxErr := ""
	nTracebackLines := ln - 1
	if isSyntaxError(C.PyTuple_GetItem(excInfo.p, 0)) {
		firstLine := extractLineFromFormattedErrorMessage(formatted, ln-3)
		secondLine := extractLineFromFormattedErrorMessage(formatted, ln-2)
		syntaxErr = firstLine + secondLine
		nTracebackLines = ln - 3
	}
	stackTrace := ""
	for i := C.Py_ssize_t(0); i < nTracebackLines; i++ {
		stackTrace += extractLineFromFormattedErrorMessage(formatted, i)
	}
	stackTrace = strings.TrimRightFunc(stackTrace, unicode.IsSpace)
	return &pyErr{
		mainMsg:      mainMsg,
		syntaxErrMsg: syntaxErr,
		stackTrace:   stackTrace,
	}
}

func extractLineFromFormattedErrorMessage(formatted Object, n C.Py_ssize_t) string {
	line := C.PyList_GetItem(formatted.p, n)
	return C.GoString(C.PyString_AsString(line))
}

type pyErr struct {
	mainMsg      string
	syntaxErrMsg string
	stackTrace   string
}

// Error returns an error message string for pyErr.
// This string contains multiple lines for stacktrace.
func (e *pyErr) Error() string {
	return e.mainMsg + "\n" + e.syntaxErrMsg + e.stackTrace
}

func isPyNoMemoryError() bool {
	return C.PyErr_ExceptionMatches(C.PyExc_MemoryError) != 0
}

// pyNoMemoryError is an error value representing an allocation error on Python.
// This is not a pyErr because it is difficult to extract stacktrace from Python when the heap is exhaused.
var pyNoMemoryError = errors.New("python interpreter failed to allocate memory")
<|endoftext|>


<commit_before>package aws

import (
	"database/sql"
	"encoding/json"
	"errors"
	"math/rand"
	"net/http"

	"github.com/trackit/jsonlog"

	"github.com/trackit/trackit2/db"
	"github.com/trackit/trackit2/routes"
	"github.com/trackit/trackit2/users"
)

func init() {
	routes.Register(
		"/aws",
		awsAccount,
		routes.RequireMethod{"POST", "GET"},
		routes.RequireContentType{"application/json"},
		db.WithTransaction{db.Db},
		users.WithAuthenticatedUser{},
	)
}

func awsAccount(r *http.Request, a routes.Arguments) (int, interface{}) {
	switch r.Method {
	case "POST":
		return postAwsAccount(r, a)
	case "GET":
		return getAwsAccount(r, a)
	default:
		logger := jsonlog.LoggerFromContextOrDefault(r.Context())
		logger.Error("Bad method. Did 'RequireMethod' do its job?", r.Method)
		return 500, nil
	}
}

type postAwsAccountRequestBody struct {
	RoleArn  string `json:"roleArn"`
	External string `json:"external"`
}

func postAwsAccount(r *http.Request, a routes.Arguments) (int, interface{}) {
	var body postAwsAccountRequestBody
	err := decodeRequestBody(r, &body)
	if err == nil && isPostAwsAccountRequestBodyValid(body) {
		tx := a[db.Transaction].(*sql.Tx)
		u := a[users.AuthenticatedUser].(users.User)
		return postAwsAccountWithValidBody(r, tx, u, body)
	} else {
		return 400, errors.New("Body is invalid.")
	}
}

func postAwsAccountWithValidBody(r *http.Request, tx *sql.Tx, user users.User, body postAwsAccountRequestBody) (int, interface{}) {
	logger := jsonlog.LoggerFromContextOrDefault(r.Context())
	account := AwsAccount{
		RoleArn:  body.RoleArn,
		External: body.External,
		UserId:   user.Id,
	}
	_, err := GetTemporaryCredentials(account, "validityTest")
	if err == nil {
		err = account.CreateAwsAccount(r.Context(), tx)
		if err == nil {
			return 200, account
		} else {
			logger.Error("Failed to insert AWS account.", map[string]interface{}{
				"error":   err.Error(),
				"account": account,
				"user":    user,
			})
			return 500, errors.New("Failed to create AWS account.")
		}
	} else {
		return 400, errors.New("Could not validate role and external ID.")
	}
}

func isPostAwsAccountRequestBodyValid(body postAwsAccountRequestBody) bool {
	return body.RoleArn != ""
}

// decodeRequestBody decodes a JSON request body and returns nil in case it
// could do so.
func decodeRequestBody(request *http.Request, structuredBody interface{}) error {
	return json.NewDecoder(request.Body).Decode(structuredBody)
}

func getAwsAccount(r *http.Request, a routes.Arguments) (int, interface{}) {
	u := a[users.AuthenticatedUser].(users.User)
	tx := a[db.Transaction].(*sql.Tx)
	l := jsonlog.LoggerFromContextOrDefault(r.Context())
	awsAccounts, err := GetAwsAccountsFromUser(u, tx)
	if err == nil {
		return 200, awsAccounts
	} else {
		l.Error("Failed to get user's AWS accounts.", err.Error())
		return 500, errors.New("Failed to retrieve AWS accounts.")
	}
}

const (
	externalChars       = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_+=,.@-"
	externalCharsCount  = len(externalChars)
	externalBitsPerChar = 6 // 59 characters, so 6 bits are required to address them
	externalBitsMask    = ^(-1 << externalBitsPerChar)
	externalLength      = 40
)

// generateExternal generates an External Id for IAM roles. It is not supposed
// to be a secret thus we won't use a cryptographically secure random
// generator.
func generateExternal() string {
	var b [externalLength]byte
	var remainingBitsCount uint
	var remainingBits uint64
	for i := range b {
	top:
		if remainingBitsCount < externalBitsPerChar {
			remainingBits = rand.Uint64()
			remainingBitsCount = 64
		}
		bits := remainingBits & externalBitsMask
		remainingBits >>= externalBitsPerChar
		remainingBitsCount -= externalBitsPerChar
		if bits >= uint64(externalCharsCount) {
			goto top
		}
		b[i] = externalChars[bits]
	}
	return string(b[:])
}

func init() {
	routes.Register(
		"/aws/next",
		nextExternal,
		routes.RequireMethod{"GET"},
		db.WithTransaction{db.Db},
		users.WithAuthenticatedUser{},
	)
}

type nextExternalResponseBody struct {
	External  string `json:"external"`
	AccountId string `json:"accountId"`
}

func nextExternal(r *http.Request, a routes.Arguments) (int, interface{}) {
	user := a[users.AuthenticatedUser].(users.User)
	tx := a[db.Transaction].(*sql.Tx)
	ctx := r.Context()
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	if user.NextExternal == "" {
		user.NextExternal = generateExternal()
		err := user.UpdateNextExternal(ctx, tx)
		if err != nil {
			logger.Error("Failed to update external ID.", err.Error())
			return 500, "Failed get external ID."
		}
	}
	return 200, nextExternalResponseBody{
		External:  user.NextExternal,
		AccountId: AccountId(),
	}
}
<commit_msg>Check the external ID is the expected next ID.<commit_after>package aws

import (
	"database/sql"
	"encoding/json"
	"errors"
	"math/rand"
	"net/http"

	"github.com/trackit/jsonlog"

	"github.com/trackit/trackit2/db"
	"github.com/trackit/trackit2/routes"
	"github.com/trackit/trackit2/users"
)

func init() {
	routes.Register(
		"/aws",
		awsAccount,
		routes.RequireMethod{"POST", "GET"},
		routes.RequireContentType{"application/json"},
		db.WithTransaction{db.Db},
		users.WithAuthenticatedUser{},
	)
}

func awsAccount(r *http.Request, a routes.Arguments) (int, interface{}) {
	switch r.Method {
	case "POST":
		return postAwsAccount(r, a)
	case "GET":
		return getAwsAccount(r, a)
	default:
		logger := jsonlog.LoggerFromContextOrDefault(r.Context())
		logger.Error("Bad method. Did 'RequireMethod' do its job?", r.Method)
		return 500, nil
	}
}

type postAwsAccountRequestBody struct {
	RoleArn  string `json:"roleArn"`
	External string `json:"external"`
}

func postAwsAccount(r *http.Request, a routes.Arguments) (int, interface{}) {
	var body postAwsAccountRequestBody
	err := decodeRequestBody(r, &body)
	if err == nil && isPostAwsAccountRequestBodyValid(body) {
		tx := a[db.Transaction].(*sql.Tx)
		u := a[users.AuthenticatedUser].(users.User)
		return postAwsAccountWithValidBody(r, tx, u, body)
	} else {
		return 400, errors.New("Body is invalid.")
	}
}

func postAwsAccountWithValidBody(r *http.Request, tx *sql.Tx, user users.User, body postAwsAccountRequestBody) (int, interface{}) {
	ctx := r.Context()
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	account := AwsAccount{
		RoleArn:  body.RoleArn,
		External: body.External,
		UserId:   user.Id,
	}
	if account.External != user.NextExternal {
		logger.Warning("Tried to add AWS account with bad external.", account)
		return 400, errors.New("Incorrect external. Use /aws/next to get expected external.")
	}
	_, err := GetTemporaryCredentials(account, "validityTest")
	if err == nil {
		err = account.CreateAwsAccount(r.Context(), tx)
		if err == nil {
			user.NextExternal = ""
			err := user.UpdateNextExternal(ctx, tx)
			if err == nil {
				return 200, account
			} else {
				logger.Error("Failed to update external.", err.Error())
				return 500, errors.New("Failed to update external.")
			}
		} else {
			logger.Error("Failed to insert AWS account.", map[string]interface{}{
				"error":   err.Error(),
				"account": account,
				"user":    user,
			})
			return 500, errors.New("Failed to create AWS account.")
		}
	} else {
		return 400, errors.New("Could not validate role and external ID.")
	}
}

func isPostAwsAccountRequestBodyValid(body postAwsAccountRequestBody) bool {
	return body.RoleArn != ""
}

// decodeRequestBody decodes a JSON request body and returns nil in case it
// could do so.
func decodeRequestBody(request *http.Request, structuredBody interface{}) error {
	return json.NewDecoder(request.Body).Decode(structuredBody)
}

func getAwsAccount(r *http.Request, a routes.Arguments) (int, interface{}) {
	u := a[users.AuthenticatedUser].(users.User)
	tx := a[db.Transaction].(*sql.Tx)
	l := jsonlog.LoggerFromContextOrDefault(r.Context())
	awsAccounts, err := GetAwsAccountsFromUser(u, tx)
	if err == nil {
		return 200, awsAccounts
	} else {
		l.Error("Failed to get user's AWS accounts.", err.Error())
		return 500, errors.New("Failed to retrieve AWS accounts.")
	}
}

const (
	externalChars       = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_+=,.@-"
	externalCharsCount  = len(externalChars)
	externalBitsPerChar = 6 // 59 characters, so 6 bits are required to address them
	externalBitsMask    = ^(-1 << externalBitsPerChar)
	externalLength      = 40
)

// generateExternal generates an External Id for IAM roles. It is not supposed
// to be a secret thus we won't use a cryptographically secure random
// generator.
func generateExternal() string {
	var b [externalLength]byte
	var remainingBitsCount uint
	var remainingBits uint64
	for i := range b {
	top:
		if remainingBitsCount < externalBitsPerChar {
			remainingBits = rand.Uint64()
			remainingBitsCount = 64
		}
		bits := remainingBits & externalBitsMask
		remainingBits >>= externalBitsPerChar
		remainingBitsCount -= externalBitsPerChar
		if bits >= uint64(externalCharsCount) {
			goto top
		}
		b[i] = externalChars[bits]
	}
	return string(b[:])
}

func init() {
	routes.Register(
		"/aws/next",
		nextExternal,
		routes.RequireMethod{"GET"},
		db.WithTransaction{db.Db},
		users.WithAuthenticatedUser{},
	)
}

type nextExternalResponseBody struct {
	External  string `json:"external"`
	AccountId string `json:"accountId"`
}

func nextExternal(r *http.Request, a routes.Arguments) (int, interface{}) {
	user := a[users.AuthenticatedUser].(users.User)
	tx := a[db.Transaction].(*sql.Tx)
	ctx := r.Context()
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	if user.NextExternal == "" {
		user.NextExternal = generateExternal()
		err := user.UpdateNextExternal(ctx, tx)
		if err != nil {
			logger.Error("Failed to update external ID.", err.Error())
			return 500, "Failed get external ID."
		}
	}
	return 200, nextExternalResponseBody{
		External:  user.NextExternal,
		AccountId: AccountId(),
	}
}
<|endoftext|>


<commit_before>package babel

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"testing"

	"github.com/stvp/assert"
)

var opts = map[string]interface{}{
	"plugins": []string{
		"transform-react-jsx",
		"transform-block-scoping",
	},
}

var input = `let foo = 1;

		Hello JSX!
		The value of foo is {foo}.

`

func _TransformWithPool(t *testing.T, n int, p int) {
	expectedOutput := strings.Join([]string{
		"var foo = 1;",
		"",
		"/*#__PURE__*/",
		`React.createElement("div", null, "Hello JSX! The value of foo is ", foo, ".");`,
	}, "\n")
	if p > 0 {
		Init(p)
	}
	type result struct {
		Index  int
		Error  error
		Output string
	}
	outputCh := make(chan result)
	for i := 0; i < n; i++ { // make sure pool works by calling Transform multiple times
		go func(i int) {
			output, err := Transform(strings.NewReader(input), opts)
			if err != nil {
				outputCh <- result{
					Index: i,
					Error: err,
				}
				return
			}
			var outputBuf bytes.Buffer
			io.Copy(&outputBuf, output)
			outputCh <- result{
				Index:  i,
				Output: outputBuf.String(),
			}
		}(i)
	}

	for i := 0; i < n; i++ { // check outputs
		output := <-outputCh
		assert.Nil(t, output.Error, fmt.Sprintf("Transform(%d) failed", output.Index))
		assert.Equal(t, expectedOutput, output.Output, fmt.Sprintf("Transform(%d) failed", output.Index))
	}
}

func TestTransform(t *testing.T) {
	_TransformWithPool(t, 2, 0)
}

func TestTransformWithPool(t *testing.T) {
	_TransformWithPool(t, 4, 2)
}

func BenchmarkTransformString(b *testing.B) {
	for i := 0; i < b.N; i++ {
		TransformString(input, opts)
	}
}

func BenchmarkTransformStringWithSingletonPool(b *testing.B) {
	Init(1)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		TransformString(input, opts)
	}
}

func BenchmarkTransformStringWithLargePool(b *testing.B) {
	Init(4)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		TransformString(input, opts)
	}
}
<commit_msg>babel JSX output changed: adjust test to match<commit_after>package babel

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"testing"

	"github.com/stvp/assert"
)

var opts = map[string]interface{}{
	"plugins": []string{
		"transform-react-jsx",
		"transform-block-scoping",
	},
}

var input = `let foo = 1;

		Hello JSX!
		The value of foo is {foo}.

`

func _TransformWithPool(t *testing.T, n int, p int) {
	expectedOutput := strings.Join([]string{
		"var foo = 1;",
		`/*#__PURE__*/React.createElement("div", null, "Hello JSX! The value of foo is ", foo, ".");`,
	}, "\n")
	if p > 0 {
		Init(p)
	}
	type result struct {
		Index  int
		Error  error
		Output string
	}
	outputCh := make(chan result)
	for i := 0; i < n; i++ { // make sure pool works by calling Transform multiple times
		go func(i int) {
			output, err := Transform(strings.NewReader(input), opts)
			if err != nil {
				outputCh <- result{
					Index: i,
					Error: err,
				}
				return
			}
			var outputBuf bytes.Buffer
			io.Copy(&outputBuf, output)
			outputCh <- result{
				Index:  i,
				Output: outputBuf.String(),
			}
		}(i)
	}

	for i := 0; i < n; i++ { // check outputs
		output := <-outputCh
		assert.Nil(t, output.Error, fmt.Sprintf("Transform(%d) failed", output.Index))
		assert.Equal(t, expectedOutput, output.Output, fmt.Sprintf("Transform(%d) failed", output.Index))
	}
}

func TestTransform(t *testing.T) {
	_TransformWithPool(t, 2, 0)
}

func TestTransformWithPool(t *testing.T) {
	_TransformWithPool(t, 4, 2)
}

func BenchmarkTransformString(b *testing.B) {
	for i := 0; i < b.N; i++ {
		TransformString(input, opts)
	}
}

func BenchmarkTransformStringWithSingletonPool(b *testing.B) {
	Init(1)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		TransformString(input, opts)
	}
}

func BenchmarkTransformStringWithLargePool(b *testing.B) {
	Init(4)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		TransformString(input, opts)
	}
}
<|endoftext|>