<commit_before>// Copyright 2009 Michael Stephens.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mongo

import (
	"os";
	"io";
	"net";
	"fmt";
	"rand";
	"bytes";
	"encoding/binary";
	"container/vector";
)

var last_req int32

const (
	_OP_REPLY		= 1;
	_OP_MSG			= 1000;
	_OP_UPDATE		= 2001;
	_OP_INSERT		= 2002;
	_OP_GET_BY_OID		= 2003;
	_OP_QUERY		= 2004;
	_OP_GET_MORE		= 2005;
	_OP_DELETE		= 2006;
	_OP_KILL_CURSORS	= 2007;
)

type message interface {
	Bytes() []byte;
	RequestID() int32;
	OpCode() int32;
}

type Connection struct {
	host	string;
	port	int;
	conn	*net.TCPConn;
}

func Connect(host string, port int) (*Connection, os.Error) {
	laddr, _ := net.ResolveTCPAddr("localhost");
	addr, _ := net.ResolveTCPAddr(fmt.Sprintf("%s:%d", host, port));
	conn, err := net.DialTCP("tcp", laddr, addr);

	if err != nil {
		return nil, err
	}

	return &Connection{host, port, conn}, nil;
}

func header(length, reqID, respTo, opCode int32) []byte {
	b := make([]byte, 16);
	binary.LittleEndian.PutUint32(b[0:4], uint32(length));
	binary.LittleEndian.PutUint32(b[4:8], uint32(reqID));
	binary.LittleEndian.PutUint32(b[8:12], uint32(respTo));
	binary.LittleEndian.PutUint32(b[12:16], uint32(opCode));
	return b;
}

func (c *Connection) writeMessage(m message) os.Error {
	body := m.Bytes();
	hb := header(int32(len(body)+16), m.RequestID(), 0, m.OpCode());
	msg := bytes.Add(hb, body);

	_, err := c.conn.Write(msg);

	last_req = m.RequestID();
	return err;
}

func (c *Connection) readReply() (*replyMsg, os.Error) {
	size_bits, _ := io.ReadAll(io.LimitReader(c.conn, 4));
	size := binary.LittleEndian.Uint32(size_bits);
	rest, _ := io.ReadAll(io.LimitReader(c.conn, int64(size)-4));
	reply := parseReply(rest);
	return reply, nil;
}

type Database struct {
	conn	*Connection;
	name	string;
}

func (c *Connection) GetDB(name string) *Database {
	return &Database{c, name}
}

func (db *Database) Drop() os.Error {
	cmd, err := Marshal(map[string]int{"dropDatabase": 1});
	if err != nil {
		return err
	}

	_, err = db.Command(cmd);
	return err;
}

func (db *Database) Repair(preserveClonedFilesOnFailure, backupOriginalFiles bool) os.Error {
	cmd := &_Object{map[string]BSON{"repairDatabase": &_Number{1, _Null{}}, "preserveClonedFilesOnFailure": &_Boolean{preserveClonedFilesOnFailure, _Null{}}, "backupOriginalFiles": &_Boolean{backupOriginalFiles, _Null{}}}, _Null{}};
	_, err := db.Command(cmd);
	return err;
}

type Collection struct {
	db	*Database;
	name	string;
}

func (db *Database) GetCollection(name string) *Collection {
	return &Collection{db, name}
}

type Cursor struct {
	collection	*Collection;
	id		int64;
	pos		int;
	docs		*vector.Vector;
}

func (c *Cursor) HasMore() bool {
	if c.pos < c.docs.Len() {
		return true
	}

	err := c.GetMore();
	if err != nil {
		return false
	}

	return c.pos < c.docs.Len();
}

func (c *Cursor) GetNext() (BSON, os.Error) {
	if c.HasMore() {
		return c.docs.At(c.pos).(BSON), nil
	}
	return Null, os.NewError("cursor failure");
}

func (c *Cursor) GetMore() os.Error {
	if c.id == 0 {
		return os.NewError("no cursorID")
	}

	gm := &getMoreMsg{c.collection.fullName(), 0, c.id, rand.Int31()};
	conn := c.collection.db.conn;
	err := conn.writeMessage(gm);
	if err != nil {
		return err
	}

	reply, err := conn.readReply();
	if err != nil {
		return err
	}

	c.pos = 0;
	c.docs = reply.docs;

	return nil;
}

func (c *Collection) fullName() string	{ return c.db.name + "." + c.name }

type indexDesc struct {
	Name	string;
	Ns	string;
	Key	map[string]int;
}

func (c *Collection) EnsureIndex(name string, index map[string]int) os.Error {
	coll := c.db.GetCollection("system.indexes");
	id := &indexDesc{name, c.fullName(), index};
	desc, err := Marshal(id);
	if err != nil {
		return err
	}
	return coll.Insert(desc);
}

func (c *Collection) DropIndexes() os.Error	{ return c.DropIndex("*") }

func (c *Collection) DropIndex(name string) os.Error {
	cmdm := map[string]string{"deleteIndexes": c.fullName(), "index": name};
	cmd, err := Marshal(cmdm);
	if err != nil {
		return err
	}

	_, err = c.db.Command(cmd);
	return err;
}

func (c *Collection) Drop() os.Error {
	cmdm := map[string]string{"drop": c.fullName()};
	cmd, err := Marshal(cmdm);
	if err != nil {
		return err
	}

	_, err = c.db.Command(cmd);
	return err;
}

func (c *Collection) Insert(doc BSON) os.Error {
	im := &insertMsg{c.fullName(), doc, rand.Int31()};
	return c.db.conn.writeMessage(im);
}

func (c *Collection) Remove(selector BSON) os.Error {
	dm := &deleteMsg{c.fullName(), selector, rand.Int31()};
	return c.db.conn.writeMessage(dm);
}

func (coll *Collection) Query(query BSON, skip, limit int) (*Cursor, os.Error) {
	req_id := rand.Int31();
	conn := coll.db.conn;
	qm := &queryMsg{0, coll.fullName(), int32(skip), int32(limit), query, req_id};

	err := conn.writeMessage(qm);
	if err != nil {
		return nil, err
	}

	reply, err := conn.readReply();
	if err != nil {
		return nil, err
	}
	if reply.responseTo != req_id {
		return nil, os.NewError("wrong responseTo code")
	}

	return &Cursor{coll, reply.cursorID, 0, reply.docs}, nil;
}

func (coll *Collection) FindAll(query BSON) (*Cursor, os.Error) {
	return coll.Query(query, 0, 0)
}

func (coll *Collection) FindOne(query BSON) (BSON, os.Error) {
	cursor, err := coll.Query(query, 0, 1);
	if err != nil {
		return nil, err
	}
	return cursor.GetNext();
}

func (coll *Collection) update(um *updateMsg) os.Error {
	um.requestID = rand.Int31();
	conn := coll.db.conn;
	return conn.writeMessage(um);
}

func (coll *Collection) Update(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 0, selector, document, 0})
}

func (coll *Collection) Upsert(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 1, selector, document, 0})
}

func (coll *Collection) UpdateAll(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 2, selector, document, 0})
}

func (coll *Collection) UpsertAll(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 3, selector, document, 0})
}

func (db *Database) Command(cmd BSON) (BSON, os.Error) {
	coll := db.GetCollection("$cmd");
	return coll.FindOne(cmd);
}

type queryMsg struct {
	opts			int32;
	fullCollectionName	string;
	numberToSkip		int32;
	numberToReturn		int32;
	query			BSON;
	requestID		int32;
}

func (q *queryMsg) OpCode() int32	{ return _OP_QUERY }
func (q *queryMsg) RequestID() int32	{ return q.requestID }
func (q *queryMsg) Bytes() []byte {
	b := make([]byte, 4);
	binary.LittleEndian.PutUint32(b, uint32(q.opts));

	buf := bytes.NewBuffer(b);
	buf.WriteString(q.fullCollectionName);
	buf.WriteByte(0);

	binary.LittleEndian.PutUint32(b, uint32(q.numberToSkip));
	buf.Write(b);

	binary.LittleEndian.PutUint32(b, uint32(q.numberToReturn));
	buf.Write(b);

	buf.Write(q.query.Bytes());
	return buf.Bytes();
}

type insertMsg struct {
	fullCollectionName	string;
	doc			BSON;
	requestID		int32;
}

func (i *insertMsg) OpCode() int32	{ return _OP_INSERT }
func (i *insertMsg) RequestID() int32	{ return i.requestID }
func (i *insertMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4));
	buf.WriteString(i.fullCollectionName);
	buf.WriteByte(0);
	buf.Write(i.doc.Bytes());
	return buf.Bytes();
}

type deleteMsg struct {
	fullCollectionName	string;
	selector		BSON;
	requestID		int32;
}

func (d *deleteMsg) OpCode() int32	{ return _OP_DELETE }
func (d *deleteMsg) RequestID() int32	{ return d.requestID }
func (d *deleteMsg) Bytes() []byte {
	zero := make([]byte, 4);
	buf := bytes.NewBuffer(zero);
	buf.WriteString(d.fullCollectionName);
	buf.WriteByte(0);
	buf.Write(zero);
	buf.Write(d.selector.Bytes());
	return buf.Bytes();

}

type getMoreMsg struct {
	fullCollectionName	string;
	numberToReturn		int32;
	cursorID		int64;
	requestID		int32;
}

func (g *getMoreMsg) OpCode() int32	{ return _OP_GET_MORE }
func (g *getMoreMsg) RequestID() int32	{ return g.requestID }
func (g *getMoreMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4));
	buf.WriteString(g.fullCollectionName);
	buf.WriteByte(0);

	b := make([]byte, 4);
	binary.LittleEndian.PutUint32(b, uint32(g.numberToReturn));
	buf.Write(b);

	b = make([]byte, 8);
	binary.LittleEndian.PutUint64(b, uint64(g.cursorID));
	buf.Write(b);

	return buf.Bytes();
}

func (db *Database) GetCollectionNames() *vector.StringVector {
	return vector.NewStringVector(0)
}

type replyMsg struct {
	responseTo	int32;
	responseFlag	int32;
	cursorID	int64;
	startingFrom	int32;
	numberReturned	int32;
	docs		*vector.Vector;
}

func parseReply(b []byte) *replyMsg {
	r := new(replyMsg);
	r.responseTo = int32(binary.LittleEndian.Uint32(b[4:8]));
	r.responseFlag = int32(binary.LittleEndian.Uint32(b[12:16]));
	r.cursorID = int64(binary.LittleEndian.Uint64(b[16:24]));
	r.startingFrom = int32(binary.LittleEndian.Uint32(b[24:28]));
	r.numberReturned = int32(binary.LittleEndian.Uint32(b[28:32]));
	r.docs = vector.New(0);

	if r.numberReturned > 0 {
		buf := bytes.NewBuffer(b[36:len(b)]);
		for i := 0; int32(i) < r.numberReturned; i++ {
			var bson BSON;
			bb := new(_BSONBuilder);
			bb.ptr = &bson;
			bb.Object();
			Parse(buf, bb);
			r.docs.Push(bson);
			io.ReadAll(io.LimitReader(buf, 4));
		}
	}

	return r;
}

type updateMsg struct {
	fullCollectionName	string;
	flags			int32;
	selector, document	BSON;
	requestID		int32;
}

func (u *updateMsg) OpCode() int32	{ return _OP_UPDATE }
func (u *updateMsg) RequestID() int32	{ return u.requestID }
func (u *updateMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4));
	buf.WriteString(u.fullCollectionName);
	buf.WriteByte(0);

	b := make([]byte, 4);
	binary.LittleEndian.PutUint32(b, uint32(u.flags));
	buf.Write(b);

	buf.Write(u.selector.Bytes());
	buf.Write(u.document.Bytes());

	return buf.Bytes();
}
<commit_msg>Cursor.Close()<commit_after>// Copyright 2009 Michael Stephens.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mongo

import (
	"os";
	"io";
	"net";
	"fmt";
	"rand";
	"bytes";
	"encoding/binary";
	"container/vector";
)

var last_req int32

const (
	_OP_REPLY		= 1;
	_OP_MSG			= 1000;
	_OP_UPDATE		= 2001;
	_OP_INSERT		= 2002;
	_OP_GET_BY_OID		= 2003;
	_OP_QUERY		= 2004;
	_OP_GET_MORE		= 2005;
	_OP_DELETE		= 2006;
	_OP_KILL_CURSORS	= 2007;
)

type message interface {
	Bytes() []byte;
	RequestID() int32;
	OpCode() int32;
}

type Connection struct {
	host	string;
	port	int;
	conn	*net.TCPConn;
}

func Connect(host string, port int) (*Connection, os.Error) {
	laddr, _ := net.ResolveTCPAddr("localhost");
	addr, _ := net.ResolveTCPAddr(fmt.Sprintf("%s:%d", host, port));
	conn, err := net.DialTCP("tcp", laddr, addr);

	if err != nil {
		return nil, err
	}

	return &Connection{host, port, conn}, nil;
}

func header(length, reqID, respTo, opCode int32) []byte {
	b := make([]byte, 16);
	binary.LittleEndian.PutUint32(b[0:4], uint32(length));
	binary.LittleEndian.PutUint32(b[4:8], uint32(reqID));
	binary.LittleEndian.PutUint32(b[8:12], uint32(respTo));
	binary.LittleEndian.PutUint32(b[12:16], uint32(opCode));
	return b;
}

func (c *Connection) writeMessage(m message) os.Error {
	body := m.Bytes();
	hb := header(int32(len(body)+16), m.RequestID(), 0, m.OpCode());
	msg := bytes.Add(hb, body);

	_, err := c.conn.Write(msg);

	last_req = m.RequestID();
	return err;
}

func (c *Connection) readReply() (*replyMsg, os.Error) {
	size_bits, _ := io.ReadAll(io.LimitReader(c.conn, 4));
	size := binary.LittleEndian.Uint32(size_bits);
	rest, _ := io.ReadAll(io.LimitReader(c.conn, int64(size)-4));
	reply := parseReply(rest);
	return reply, nil;
}

type Database struct {
	conn	*Connection;
	name	string;
}

func (c *Connection) GetDB(name string) *Database {
	return &Database{c, name}
}

func (db *Database) Drop() os.Error {
	cmd, err := Marshal(map[string]int{"dropDatabase": 1});
	if err != nil {
		return err
	}

	_, err = db.Command(cmd);
	return err;
}

func (db *Database) Repair(preserveClonedFilesOnFailure, backupOriginalFiles bool) os.Error {
	cmd := &_Object{map[string]BSON{"repairDatabase": &_Number{1, _Null{}}, "preserveClonedFilesOnFailure": &_Boolean{preserveClonedFilesOnFailure, _Null{}}, "backupOriginalFiles": &_Boolean{backupOriginalFiles, _Null{}}}, _Null{}};
	_, err := db.Command(cmd);
	return err;
}

type Collection struct {
	db	*Database;
	name	string;
}

func (db *Database) GetCollection(name string) *Collection {
	return &Collection{db, name}
}

type Cursor struct {
	collection	*Collection;
	id		int64;
	pos		int;
	docs		*vector.Vector;
}

func (c *Cursor) HasMore() bool {
	if c.pos < c.docs.Len() {
		return true
	}

	err := c.GetMore();
	if err != nil {
		return false
	}

	return c.pos < c.docs.Len();
}

func (c *Cursor) GetNext() (BSON, os.Error) {
	if c.HasMore() {
		return c.docs.At(c.pos).(BSON), nil
	}
	return Null, os.NewError("cursor failure");
}

func (c *Cursor) GetMore() os.Error {
	if c.id == 0 {
		return os.NewError("no cursorID")
	}

	gm := &getMoreMsg{c.collection.fullName(), 0, c.id, rand.Int31()};
	conn := c.collection.db.conn;
	err := conn.writeMessage(gm);
	if err != nil {
		return err
	}

	reply, err := conn.readReply();
	if err != nil {
		return err
	}

	c.pos = 0;
	c.docs = reply.docs;

	return nil;
}

func (c *Cursor) Close() os.Error {
	if c.id == 0 {
		// not open on server
		return nil
	}

	req_id := rand.Int31();
	km := &killMsg{1, []int64{c.id}, req_id};
	conn := c.collection.db.conn;
	return conn.writeMessage(km);
}

func (c *Collection) fullName() string	{ return c.db.name + "." + c.name }

type indexDesc struct {
	Name	string;
	Ns	string;
	Key	map[string]int;
}

func (c *Collection) EnsureIndex(name string, index map[string]int) os.Error {
	coll := c.db.GetCollection("system.indexes");
	id := &indexDesc{name, c.fullName(), index};
	desc, err := Marshal(id);
	if err != nil {
		return err
	}
	return coll.Insert(desc);
}

func (c *Collection) DropIndexes() os.Error	{ return c.DropIndex("*") }

func (c *Collection) DropIndex(name string) os.Error {
	cmdm := map[string]string{"deleteIndexes": c.fullName(), "index": name};
	cmd, err := Marshal(cmdm);
	if err != nil {
		return err
	}

	_, err = c.db.Command(cmd);
	return err;
}

func (c *Collection) Drop() os.Error {
	cmdm := map[string]string{"drop": c.fullName()};
	cmd, err := Marshal(cmdm);
	if err != nil {
		return err
	}

	_, err = c.db.Command(cmd);
	return err;
}

func (c *Collection) Insert(doc BSON) os.Error {
	im := &insertMsg{c.fullName(), doc, rand.Int31()};
	return c.db.conn.writeMessage(im);
}

func (c *Collection) Remove(selector BSON) os.Error {
	dm := &deleteMsg{c.fullName(), selector, rand.Int31()};
	return c.db.conn.writeMessage(dm);
}

func (coll *Collection) Query(query BSON, skip, limit int) (*Cursor, os.Error) {
	req_id := rand.Int31();
	conn := coll.db.conn;
	qm := &queryMsg{0, coll.fullName(), int32(skip), int32(limit), query, req_id};

	err := conn.writeMessage(qm);
	if err != nil {
		return nil, err
	}

	reply, err := conn.readReply();
	if err != nil {
		return nil, err
	}
	if reply.responseTo != req_id {
		return nil, os.NewError("wrong responseTo code")
	}

	return &Cursor{coll, reply.cursorID, 0, reply.docs}, nil;
}

func (coll *Collection) FindAll(query BSON) (*Cursor, os.Error) {
	return coll.Query(query, 0, 0)
}

func (coll *Collection) FindOne(query BSON) (BSON, os.Error) {
	cursor, err := coll.Query(query, 0, 1);
	if err != nil {
		return nil, err
	}
	return cursor.GetNext();
}

func (coll *Collection) update(um *updateMsg) os.Error {
	um.requestID = rand.Int31();
	conn := coll.db.conn;
	return conn.writeMessage(um);
}

func (coll *Collection) Update(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 0, selector, document, 0})
}

func (coll *Collection) Upsert(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 1, selector, document, 0})
}

func (coll *Collection) UpdateAll(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 2, selector, document, 0})
}

func (coll *Collection) UpsertAll(selector, document BSON) os.Error {
	return coll.update(&updateMsg{coll.fullName(), 3, selector, document, 0})
}

func (db *Database) Command(cmd BSON) (BSON, os.Error) {
	coll := db.GetCollection("$cmd");
	return coll.FindOne(cmd);
}

type queryMsg struct {
	opts			int32;
	fullCollectionName	string;
	numberToSkip		int32;
	numberToReturn		int32;
	query			BSON;
	requestID		int32;
}

func (q *queryMsg) OpCode() int32	{ return _OP_QUERY }
func (q *queryMsg) RequestID() int32	{ return q.requestID }
func (q *queryMsg) Bytes() []byte {
	b := make([]byte, 4);
	binary.LittleEndian.PutUint32(b, uint32(q.opts));

	buf := bytes.NewBuffer(b);
	buf.WriteString(q.fullCollectionName);
	buf.WriteByte(0);

	binary.LittleEndian.PutUint32(b, uint32(q.numberToSkip));
	buf.Write(b);

	binary.LittleEndian.PutUint32(b, uint32(q.numberToReturn));
	buf.Write(b);

	buf.Write(q.query.Bytes());
	return buf.Bytes();
}

type insertMsg struct {
	fullCollectionName	string;
	doc			BSON;
	requestID		int32;
}

func (i *insertMsg) OpCode() int32	{ return _OP_INSERT }
func (i *insertMsg) RequestID() int32	{ return i.requestID }
func (i *insertMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4));
	buf.WriteString(i.fullCollectionName);
	buf.WriteByte(0);
	buf.Write(i.doc.Bytes());
	return buf.Bytes();
}

type deleteMsg struct {
	fullCollectionName	string;
	selector		BSON;
	requestID		int32;
}

func (d *deleteMsg) OpCode() int32	{ return _OP_DELETE }
func (d *deleteMsg) RequestID() int32	{ return d.requestID }
func (d *deleteMsg) Bytes() []byte {
	zero := make([]byte, 4);
	buf := bytes.NewBuffer(zero);
	buf.WriteString(d.fullCollectionName);
	buf.WriteByte(0);
	buf.Write(zero);
	buf.Write(d.selector.Bytes());
	return buf.Bytes();

}

type getMoreMsg struct {
	fullCollectionName	string;
	numberToReturn		int32;
	cursorID		int64;
	requestID		int32;
}

func (g *getMoreMsg) OpCode() int32	{ return _OP_GET_MORE }
func (g *getMoreMsg) RequestID() int32	{ return g.requestID }
func (g *getMoreMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4));
	buf.WriteString(g.fullCollectionName);
	buf.WriteByte(0);

	b := make([]byte, 4);
	binary.LittleEndian.PutUint32(b, uint32(g.numberToReturn));
	buf.Write(b);

	b = make([]byte, 8);
	binary.LittleEndian.PutUint64(b, uint64(g.cursorID));
	buf.Write(b);

	return buf.Bytes();
}

func (db *Database) GetCollectionNames() *vector.StringVector {
	return vector.NewStringVector(0)
}

type replyMsg struct {
	responseTo	int32;
	responseFlag	int32;
	cursorID	int64;
	startingFrom	int32;
	numberReturned	int32;
	docs		*vector.Vector;
}

func parseReply(b []byte) *replyMsg {
	r := new(replyMsg);
	r.responseTo = int32(binary.LittleEndian.Uint32(b[4:8]));
	r.responseFlag = int32(binary.LittleEndian.Uint32(b[12:16]));
	r.cursorID = int64(binary.LittleEndian.Uint64(b[16:24]));
	r.startingFrom = int32(binary.LittleEndian.Uint32(b[24:28]));
	r.numberReturned = int32(binary.LittleEndian.Uint32(b[28:32]));
	r.docs = vector.New(0);

	if r.numberReturned > 0 {
		buf := bytes.NewBuffer(b[36:len(b)]);
		for i := 0; int32(i) < r.numberReturned; i++ {
			var bson BSON;
			bb := new(_BSONBuilder);
			bb.ptr = &bson;
			bb.Object();
			Parse(buf, bb);
			r.docs.Push(bson);
			io.ReadAll(io.LimitReader(buf, 4));
		}
	}

	return r;
}

type updateMsg struct {
	fullCollectionName	string;
	flags			int32;
	selector, document	BSON;
	requestID		int32;
}

func (u *updateMsg) OpCode() int32	{ return _OP_UPDATE }
func (u *updateMsg) RequestID() int32	{ return u.requestID }
func (u *updateMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4));
	buf.WriteString(u.fullCollectionName);
	buf.WriteByte(0);

	b := make([]byte, 4);
	binary.LittleEndian.PutUint32(b, uint32(u.flags));
	buf.Write(b);

	buf.Write(u.selector.Bytes());
	buf.Write(u.document.Bytes());

	return buf.Bytes();
}

type killMsg struct {
	numberOfCursorIDs	int32;
	cursorIDs		[]int64;
	requestID		int32;
}

func (k *killMsg) OpCode() int32	{ return _OP_KILL_CURSORS }
func (k *killMsg) RequestID() int32	{ return k.requestID }
func (k *killMsg) Bytes() []byte {
	buf := bytes.NewBuffer(make([]byte, 4));

	b := make([]byte, 4);
	binary.LittleEndian.PutUint32(b, uint32(k.numberOfCursorIDs));
	buf.Write(b);

	b = make([]byte, 8);
	for _, id := range k.cursorIDs {
		binary.LittleEndian.PutUint64(b, uint64(id));
		buf.Write(b);
	}

	return buf.Bytes();
}
<|endoftext|>


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

// Package main is the main entry point for the app.
package main

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/proto/access"
	"go.chromium.org/luci/config/server/cfgmodule"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server"
	"go.chromium.org/luci/server/cron"
	"go.chromium.org/luci/server/gaeemulation"
	"go.chromium.org/luci/server/module"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/tq"

	// Enable datastore transactional tasks support.
	_ "go.chromium.org/luci/server/tq/txn/datastore"

	"go.chromium.org/luci/buildbucket/appengine/internal/config"
	"go.chromium.org/luci/buildbucket/appengine/rpc"
	pb "go.chromium.org/luci/buildbucket/proto"
)

// isBeefy returns whether the request was intended for the beefy service.
func isBeefy(req *http.Request) bool {
	return strings.Contains(req.Host, "beefy")
}

// isDev returns whether the request was intended for the dev instance.
func isDev(req *http.Request) bool {
	return strings.HasSuffix(req.Host, "-dev.appspot.com")
}

func main() {
	mods := []module.Module{
		cfgmodule.NewModuleFromFlags(),
		cron.NewModuleFromFlags(),
		gaeemulation.NewModuleFromFlags(),
		tq.NewModuleFromFlags(),
	}

	server.Main(nil, mods, func(srv *server.Server) error {
		// Proxy buildbucket.v2.Builds pRPC requests back to the Python
		// service in order to achieve a programmatic traffic split.
		// Because of the way dispatch routes work, requests are proxied
		// to a copy of the Python service hosted at a different path.
		// TODO(crbug/1042991): Remove the proxy once the go service handles all traffic.
		pythonURL, err := url.Parse(fmt.Sprintf("https://default-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		beefyURL, err := url.Parse(fmt.Sprintf("https://beefy-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		prx := httputil.NewSingleHostReverseProxy(pythonURL)
		prx.Director = func(req *http.Request) {
			target := pythonURL
			if isBeefy(req) {
				target = beefyURL
			}
			// According to net.Request documentation, setting Host is unnecessary
			// because URL.Host is supposed to be used for outbound requests.
			// However, on GAE, it seems that req.Host is incorrectly used.
			req.Host = target.Host
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			req.URL.Path = fmt.Sprintf("%s%s", target.Path, req.URL.Path)
		}
		// makeOverride returns a prpc.Override which allows the given percentage of requests
		// through to this service, proxying the remainder to Python.
		makeOverride := func(prodPct, devPct int) func(*router.Context) bool {
			return func(ctx *router.Context) bool {
				pct := prodPct
				if isDev(ctx.Request) {
					pct = devPct
				}
				switch val := ctx.Request.Header.Get("Should-Proxy"); val {
				case "true":
					pct = 0
					logging.Debugf(ctx.Context, "request demanded to be proxied")
				case "false":
					pct = 100
					logging.Debugf(ctx.Context, "request demanded not to be proxied")
				}
				if mathrand.Intn(ctx.Context, 100) < pct {
					return false
				}
				target := pythonURL
				if isBeefy(ctx.Request) {
					target = beefyURL
				}
				logging.Debugf(ctx.Context, "proxying request to %s", target)
				prx.ServeHTTP(ctx.Writer, ctx.Request)
				return true
			}
		}

		srv.PRPC.AccessControl = prpc.AllowOriginAll
		access.RegisterAccessServer(srv.PRPC, &access.UnimplementedAccessServer{})
		pb.RegisterBuildsServer(srv.PRPC, rpc.NewBuilds())
		pb.RegisterBuildersServer(srv.PRPC, rpc.NewBuilders())
		// TODO(crbug/1082369): Remove this workaround once field masks can be decoded.
		srv.PRPC.HackFixFieldMasksForJSON = true

		// Percentage of prod ScheduleBuild requests to serve from Go.
		pct := 25
		// Traffic split for unbatched ScheduleBuild requests. makeOverride(prod % -> Go, dev % -> Go).
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "ScheduleBuild", makeOverride(pct, 100))
		// Traffic split for batched ScheduleBuild requests in prod (see rpc/batch.go).
		srv.Context = rpc.WithTrafficSplit(srv.Context, pct)

		cron.RegisterHandler("update_config", config.UpdateSettingsCfg)
		return nil
	})
}
<commit_msg>[buildbucket] Serve 50% of batched/unbatched prod ScheduleBuild requests<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package main is the main entry point for the app.
package main

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/proto/access"
	"go.chromium.org/luci/config/server/cfgmodule"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server"
	"go.chromium.org/luci/server/cron"
	"go.chromium.org/luci/server/gaeemulation"
	"go.chromium.org/luci/server/module"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/tq"

	// Enable datastore transactional tasks support.
	_ "go.chromium.org/luci/server/tq/txn/datastore"

	"go.chromium.org/luci/buildbucket/appengine/internal/config"
	"go.chromium.org/luci/buildbucket/appengine/rpc"
	pb "go.chromium.org/luci/buildbucket/proto"
)

// isBeefy returns whether the request was intended for the beefy service.
func isBeefy(req *http.Request) bool {
	return strings.Contains(req.Host, "beefy")
}

// isDev returns whether the request was intended for the dev instance.
func isDev(req *http.Request) bool {
	return strings.HasSuffix(req.Host, "-dev.appspot.com")
}

func main() {
	mods := []module.Module{
		cfgmodule.NewModuleFromFlags(),
		cron.NewModuleFromFlags(),
		gaeemulation.NewModuleFromFlags(),
		tq.NewModuleFromFlags(),
	}

	server.Main(nil, mods, func(srv *server.Server) error {
		// Proxy buildbucket.v2.Builds pRPC requests back to the Python
		// service in order to achieve a programmatic traffic split.
		// Because of the way dispatch routes work, requests are proxied
		// to a copy of the Python service hosted at a different path.
		// TODO(crbug/1042991): Remove the proxy once the go service handles all traffic.
		pythonURL, err := url.Parse(fmt.Sprintf("https://default-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		beefyURL, err := url.Parse(fmt.Sprintf("https://beefy-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		prx := httputil.NewSingleHostReverseProxy(pythonURL)
		prx.Director = func(req *http.Request) {
			target := pythonURL
			if isBeefy(req) {
				target = beefyURL
			}
			// According to net.Request documentation, setting Host is unnecessary
			// because URL.Host is supposed to be used for outbound requests.
			// However, on GAE, it seems that req.Host is incorrectly used.
			req.Host = target.Host
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			req.URL.Path = fmt.Sprintf("%s%s", target.Path, req.URL.Path)
		}
		// makeOverride returns a prpc.Override which allows the given percentage of requests
		// through to this service, proxying the remainder to Python.
		makeOverride := func(prodPct, devPct int) func(*router.Context) bool {
			return func(ctx *router.Context) bool {
				pct := prodPct
				if isDev(ctx.Request) {
					pct = devPct
				}
				switch val := ctx.Request.Header.Get("Should-Proxy"); val {
				case "true":
					pct = 0
					logging.Debugf(ctx.Context, "request demanded to be proxied")
				case "false":
					pct = 100
					logging.Debugf(ctx.Context, "request demanded not to be proxied")
				}
				if mathrand.Intn(ctx.Context, 100) < pct {
					return false
				}
				target := pythonURL
				if isBeefy(ctx.Request) {
					target = beefyURL
				}
				logging.Debugf(ctx.Context, "proxying request to %s", target)
				prx.ServeHTTP(ctx.Writer, ctx.Request)
				return true
			}
		}

		srv.PRPC.AccessControl = prpc.AllowOriginAll
		access.RegisterAccessServer(srv.PRPC, &access.UnimplementedAccessServer{})
		pb.RegisterBuildsServer(srv.PRPC, rpc.NewBuilds())
		pb.RegisterBuildersServer(srv.PRPC, rpc.NewBuilders())
		// TODO(crbug/1082369): Remove this workaround once field masks can be decoded.
		srv.PRPC.HackFixFieldMasksForJSON = true

		// Percentage of prod ScheduleBuild requests to serve from Go.
		pct := 50
		// Traffic split for unbatched ScheduleBuild requests. makeOverride(prod % -> Go, dev % -> Go).
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "ScheduleBuild", makeOverride(pct, 100))
		// Traffic split for batched ScheduleBuild requests in prod (see rpc/batch.go).
		srv.Context = rpc.WithTrafficSplit(srv.Context, pct)

		cron.RegisterHandler("update_config", config.UpdateSettingsCfg)
		return nil
	})
}
<|endoftext|>


<commit_before>package party

import (
	"bytes"
	"html/template"
	"net/http"
	"path/filepath"
	"sync"
)

const (
	templateDir = "./template"
	partialDir  = templateDir + "/partial"
)

type templateHandler interface {
	String() (string, error)
	ServeHTTP(w http.ResponseWriter, r *http.Request)
	Data(key string) (interface{}, bool)
	SetData(key string, value interface{}) templateHandler
	AddDataMap(data map[string]interface{}) templateHandler
	ClearData() templateHandler
}

// New create a templateHandler. If development is true, returns a dynamic template
// which provide more development flexibility. Otherwise returns a static template
// that maximize the speed
func New(filename string, data map[string]interface{}, development bool) templateHandler {
	if data == nil {
		data = make(map[string]interface{})
	}
	data["IsDevelopment"] = devlopment
	if development {
		return newDynamicTemplateHandler(filename, data)
	} else {
		return newStaticTemplateHandler(filename, data)
	}
}

///////////////////////////////////////////////////////////////////////
// dynamicTemplateHandler - Flexible but slower, used in development.
///////////////////////////////////////////////////////////////////////

type dynamicTemplateHandler struct {
	once     sync.Once
	funcs    template.FuncMap
	filename string
	data     map[string]interface{}
}

func newDynamicTemplateHandler(filename string, data map[string]interface{}) *dynamicTemplateHandler {
	t := &dynamicTemplateHandler{
		filename: filename,
		data:     data,
	}
	t.once.Do(func() {
		t.funcs = template.FuncMap{
			"include": func(part string) (template.HTML, error) {
				buf := bytes.NewBuffer(nil)
				pt := template.Must(template.ParseFiles(filepath.Join(partialDir, part)))
				err := pt.Execute(buf, t.data)
				return template.HTML(buf.String()), err
			},
		}
	})
	return t
}

func (t *dynamicTemplateHandler) String() (string, error) {
	buf := bytes.NewBuffer(nil)
	tpl := template.Must(template.New(t.filename).Funcs(t.funcs).ParseFiles(filepath.Join(templateDir, t.filename)))
	if err := tpl.ExecuteTemplate(buf, t.filename, t.data); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func (t *dynamicTemplateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	tpl := template.Must(template.New(t.filename).
		Funcs(t.funcs).
		ParseFiles(filepath.Join(templateDir, t.filename)))
	if err := tpl.ExecuteTemplate(w, t.filename, t.data); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Data get template data
func (t *dynamicTemplateHandler) Data(key string) (interface{}, bool) {
	v, found := t.data[key]
	return v, found
}

// SetData set template data
func (t *dynamicTemplateHandler) SetData(key string, value interface{}) templateHandler {
	t.data[key] = value
	return t
}

// AddDataMap add template data map
func (t *dynamicTemplateHandler) AddDataMap(m map[string]interface{}) templateHandler {
	if len(t.data) == 0 {
		t.data = m
	} else {
		for k, v := range m {
			t.data[k] = v
		}
	}
	return t
}

// ClearData delete template data
func (t *dynamicTemplateHandler) ClearData() templateHandler {
	t.data = map[string]interface{}{}
	return t
}

//////////////////////////////////////////////////////////
// staticTemplateHandler - Faster, used in production.
//////////////////////////////////////////////////////////

// staticTemplateHandler handles HTTP request by rendering a template
type staticTemplateHandler struct {
	once     sync.Once
	filename string
	templ    *template.Template
	data     map[string]interface{}
}

// newStaticStaticTemplateHandler create a new Template Handler
func newStaticTemplateHandler(filename string, data map[string]interface{}) *staticTemplateHandler {
	if data == nil {
		data = map[string]interface{}{}
	}

	t := &staticTemplateHandler{
		filename: filename,
		data:     data,
	}

	t.once.Do(func() {
		funcs := template.FuncMap{
			"include": func(part string) (template.HTML, error) {
				buf := bytes.NewBuffer(nil)
				pt := template.Must(template.ParseFiles(filepath.Join(partialDir, part)))
				err := pt.Execute(buf, t.data)
				return template.HTML(buf.String()), err
			},
		}
		t.templ = template.Must(template.New(t.filename).Funcs(funcs).
			ParseFiles(filepath.Join(templateDir, t.filename)))
	})
	return t
}

func (t *staticTemplateHandler) String() (string, error) {
	buf := bytes.NewBuffer(nil)
	if err := t.templ.Execute(buf, t.data); err != nil {
		return "", err
	}
	return buf.String(), nil
}

// ServeHTTP provide HTTP service for staticTemplateHandler
func (t *staticTemplateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if err := t.templ.Execute(w, t.data); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Data get template data
func (t *staticTemplateHandler) Data(key string) (interface{}, bool) {
	v, found := t.data[key]
	return v, found
}

// SetData set template data
func (t *staticTemplateHandler) SetData(key string, value interface{}) templateHandler {
	t.data[key] = value
	return t
}

// AddDataMap add template data map
func (t *staticTemplateHandler) AddDataMap(m map[string]interface{}) templateHandler {
	if len(t.data) == 0 {
		t.data = m
	} else {
		for k, v := range m {
			t.data[k] = v
		}
	}
	return t
}

// ClearData delete template data
func (t *staticTemplateHandler) ClearData() templateHandler {
	t.data = map[string]interface{}{}
	return t
}
<commit_msg>Fixed a typo<commit_after>package party

import (
	"bytes"
	"html/template"
	"net/http"
	"path/filepath"
	"sync"
)

const (
	templateDir = "./template"
	partialDir  = templateDir + "/partial"
)

type templateHandler interface {
	String() (string, error)
	ServeHTTP(w http.ResponseWriter, r *http.Request)
	Data(key string) (interface{}, bool)
	SetData(key string, value interface{}) templateHandler
	AddDataMap(data map[string]interface{}) templateHandler
	ClearData() templateHandler
}

// New create a templateHandler. If development is true, returns a dynamic template
// which provide more development flexibility. Otherwise returns a static template
// that maximize the speed
func New(filename string, data map[string]interface{}, development bool) templateHandler {
	if data == nil {
		data = make(map[string]interface{})
	}
	data["IsDevelopment"] = development
	if development {
		return newDynamicTemplateHandler(filename, data)
	} else {
		return newStaticTemplateHandler(filename, data)
	}
}

///////////////////////////////////////////////////////////////////////
// dynamicTemplateHandler - Flexible but slower, used in development.
///////////////////////////////////////////////////////////////////////

type dynamicTemplateHandler struct {
	once     sync.Once
	funcs    template.FuncMap
	filename string
	data     map[string]interface{}
}

func newDynamicTemplateHandler(filename string, data map[string]interface{}) *dynamicTemplateHandler {
	t := &dynamicTemplateHandler{
		filename: filename,
		data:     data,
	}
	t.once.Do(func() {
		t.funcs = template.FuncMap{
			"include": func(part string) (template.HTML, error) {
				buf := bytes.NewBuffer(nil)
				pt := template.Must(template.ParseFiles(filepath.Join(partialDir, part)))
				err := pt.Execute(buf, t.data)
				return template.HTML(buf.String()), err
			},
		}
	})
	return t
}

func (t *dynamicTemplateHandler) String() (string, error) {
	buf := bytes.NewBuffer(nil)
	tpl := template.Must(template.New(t.filename).Funcs(t.funcs).ParseFiles(filepath.Join(templateDir, t.filename)))
	if err := tpl.ExecuteTemplate(buf, t.filename, t.data); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func (t *dynamicTemplateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	tpl := template.Must(template.New(t.filename).
		Funcs(t.funcs).
		ParseFiles(filepath.Join(templateDir, t.filename)))
	if err := tpl.ExecuteTemplate(w, t.filename, t.data); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Data get template data
func (t *dynamicTemplateHandler) Data(key string) (interface{}, bool) {
	v, found := t.data[key]
	return v, found
}

// SetData set template data
func (t *dynamicTemplateHandler) SetData(key string, value interface{}) templateHandler {
	t.data[key] = value
	return t
}

// AddDataMap add template data map
func (t *dynamicTemplateHandler) AddDataMap(m map[string]interface{}) templateHandler {
	if len(t.data) == 0 {
		t.data = m
	} else {
		for k, v := range m {
			t.data[k] = v
		}
	}
	return t
}

// ClearData delete template data
func (t *dynamicTemplateHandler) ClearData() templateHandler {
	t.data = map[string]interface{}{}
	return t
}

//////////////////////////////////////////////////////////
// staticTemplateHandler - Faster, used in production.
//////////////////////////////////////////////////////////

// staticTemplateHandler handles HTTP request by rendering a template
type staticTemplateHandler struct {
	once     sync.Once
	filename string
	templ    *template.Template
	data     map[string]interface{}
}

// newStaticStaticTemplateHandler create a new Template Handler
func newStaticTemplateHandler(filename string, data map[string]interface{}) *staticTemplateHandler {
	if data == nil {
		data = map[string]interface{}{}
	}

	t := &staticTemplateHandler{
		filename: filename,
		data:     data,
	}

	t.once.Do(func() {
		funcs := template.FuncMap{
			"include": func(part string) (template.HTML, error) {
				buf := bytes.NewBuffer(nil)
				pt := template.Must(template.ParseFiles(filepath.Join(partialDir, part)))
				err := pt.Execute(buf, t.data)
				return template.HTML(buf.String()), err
			},
		}
		t.templ = template.Must(template.New(t.filename).Funcs(funcs).
			ParseFiles(filepath.Join(templateDir, t.filename)))
	})
	return t
}

func (t *staticTemplateHandler) String() (string, error) {
	buf := bytes.NewBuffer(nil)
	if err := t.templ.Execute(buf, t.data); err != nil {
		return "", err
	}
	return buf.String(), nil
}

// ServeHTTP provide HTTP service for staticTemplateHandler
func (t *staticTemplateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if err := t.templ.Execute(w, t.data); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Data get template data
func (t *staticTemplateHandler) Data(key string) (interface{}, bool) {
	v, found := t.data[key]
	return v, found
}

// SetData set template data
func (t *staticTemplateHandler) SetData(key string, value interface{}) templateHandler {
	t.data[key] = value
	return t
}

// AddDataMap add template data map
func (t *staticTemplateHandler) AddDataMap(m map[string]interface{}) templateHandler {
	if len(t.data) == 0 {
		t.data = m
	} else {
		for k, v := range m {
			t.data[k] = v
		}
	}
	return t
}

// ClearData delete template data
func (t *staticTemplateHandler) ClearData() templateHandler {
	t.data = map[string]interface{}{}
	return t
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"fmt"

	"strings"

	"strconv"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	. "github.com/claudetech/loggo/default"
	"golang.org/x/net/context"
)

// Mount the fuse volume
func Mount(client *Drive, mountpoint string, mountOptions []string, uid, gid uint32, umask os.FileMode) error {
	Log.Infof("Mounting path %v", mountpoint)

	if _, err := os.Stat(mountpoint); os.IsNotExist(err) {
		Log.Debugf("Mountpoint doesn't exist, creating...")
		if err := os.MkdirAll(mountpoint, 0644); nil != err {
			Log.Debugf("%v", err)
			return fmt.Errorf("Could not create mount directory %v", mountpoint)
		}
	}

	fuse.Debug = func(msg interface{}) {
		Log.Tracef("FUSE %v", msg)
	}

	// Set mount options
	options := []fuse.MountOption{
		fuse.NoAppleDouble(),
		fuse.NoAppleXattr(),
	}
	for _, option := range mountOptions {
		if "allow_other" == option {
			options = append(options, fuse.AllowOther())
		} else if "allow_root" == option {
			options = append(options, fuse.AllowRoot())
		} else if "allow_dev" == option {
			options = append(options, fuse.AllowDev())
		} else if "allow_non_empty_mount" == option {
			options = append(options, fuse.AllowNonEmptyMount())
		} else if "allow_suid" == option {
			options = append(options, fuse.AllowSUID())
		} else if strings.Contains(option, "max_readahead=") {
			data := strings.Split(option, "=")
			value, err := strconv.ParseUint(data[1], 10, 32)
			if nil != err {
				Log.Debugf("%v", err)
				return fmt.Errorf("Could not parse max_readahead value")
			}
			options = append(options, fuse.MaxReadahead(uint32(value)))
		} else if "default_permissions" == option {
			options = append(options, fuse.DefaultPermissions())
		} else if "excl_create" == option {
			options = append(options, fuse.ExclCreate())
		} else if strings.Contains(option, "fs_name") {
			data := strings.Split(option, "=")
			options = append(options, fuse.FSName(data[1]))
		} else if "local_volume" == option {
			options = append(options, fuse.LocalVolume())
		} else if "writeback_cache" == option {
			options = append(options, fuse.WritebackCache())
		} else if strings.Contains(option, "volume_name") {
			data := strings.Split(option, "=")
			options = append(options, fuse.VolumeName(data[1]))
		} else if "read_only" == option {
			options = append(options, fuse.ReadOnly())
		} else {
			Log.Warningf("Fuse option %v is not supported, yet", option)
		}
	}

	c, err := fuse.Mount(mountpoint, options...)
	if err != nil {
		return err
	}
	defer c.Close()

	filesys := &FS{
		client: client,
		uid:    uid,
		gid:    gid,
		umask:  umask,
	}
	if err := fs.Serve(c, filesys); err != nil {
		return err
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Error mounting FUSE")
	}

	return Unmount(mountpoint, true)
}

// Unmount unmounts the mountpoint
func Unmount(mountpoint string, notify bool) error {
	if notify {
		Log.Infof("Unmounting path %v", mountpoint)
	}
	fuse.Unmount(mountpoint)
	return nil
}

// FS the fuse filesystem
type FS struct {
	client *Drive
	uid    uint32
	gid    uint32
	umask  os.FileMode
}

// Root returns the root path
func (f *FS) Root() (fs.Node, error) {
	object, err := f.client.GetRoot()
	if nil != err {
		Log.Warningf("%v", err)
		return nil, fmt.Errorf("Could not get root object")
	}
	return &Object{
		client: f.client,
		object: object,
		uid:    f.uid,
		gid:    f.gid,
		umask:  f.umask,
	}, nil
}

// Object represents one drive object
type Object struct {
	client *Drive
	object *APIObject
	buffer *Buffer
	uid    uint32
	gid    uint32
	umask  os.FileMode
}

// Attr returns the attributes for a directory
func (o *Object) Attr(ctx context.Context, attr *fuse.Attr) error {
	if o.object.IsDir {
		if o.umask > 0 {
			attr.Mode = os.ModeDir | o.umask
		} else {
			attr.Mode = os.ModeDir | 0755
		}
		attr.Size = 0
	} else {
		if o.umask > 0 {
			attr.Mode = o.umask
		} else {
			attr.Mode = 0644
		}
		attr.Size = o.object.Size
	}

	attr.Uid = uint32(o.uid)
	attr.Gid = uint32(o.gid)

	attr.Mtime = o.object.LastModified
	attr.Crtime = o.object.LastModified
	attr.Ctime = o.object.LastModified

	attr.Blocks = (attr.Size + 511) / 512

	return nil
}

// ReadDirAll shows all files in the current directory
func (o *Object) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	objects, err := o.client.GetObjectsByParent(o.object.ObjectID)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fuse.ENOENT
	}

	dirs := []fuse.Dirent{}
	for _, object := range objects {
		if object.IsDir {
			dirs = append(dirs, fuse.Dirent{
				Name: object.Name,
				Type: fuse.DT_Dir,
			})
		} else {
			dirs = append(dirs, fuse.Dirent{
				Name: object.Name,
				Type: fuse.DT_File,
			})
		}
	}
	return dirs, nil
}

// Lookup tests if a file is existent in the current directory
func (o *Object) Lookup(ctx context.Context, name string) (fs.Node, error) {
	object, err := o.client.GetObjectByParentAndName(o.object.ObjectID, name)
	if nil != err {
		Log.Tracef("%v", err)
		return nil, fuse.ENOENT
	}

	return &Object{
		client: o.client,
		object: object,
		uid:    o.uid,
		gid:    o.gid,
		umask:  o.umask,
	}, nil
}

// Open opens a file for reading
func (o *Object) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	if req.Dir {
		return o, nil
	}

	buffer, err := o.client.Open(o.object)
	if nil != err {
		Log.Warningf("%v", err)
		return o, fuse.ENOENT
	}
	o.buffer = buffer

	return o, nil
}

// Release a stream
func (o *Object) Release(ctx context.Context, req *fuse.ReleaseRequest) error {
	if nil != o.buffer {
		if err := o.buffer.Close(); nil != err {
			Log.Debugf("%v", err)
			Log.Warningf("Could not close buffer stream")
		}
	}
	return nil
}

// Read reads some bytes or the whole file
func (o *Object) Read(ctx context.Context, req *fuse.ReadRequest, resp *fuse.ReadResponse) error {
	buf, err := o.buffer.ReadBytes(req.Offset, int64(req.Size), false)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	resp.Data = buf[:]
	return nil
}

// Remove deletes an element
func (o *Object) Remove(ctx context.Context, req *fuse.RemoveRequest) error {
	obj, err := o.client.GetObjectByParentAndName(o.object.ObjectID, req.Name)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	err = o.client.Remove(obj, o.object.ObjectID)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	return nil
}
<commit_msg>Add Rename<commit_after>package main

import (
	"os"

	"fmt"

	"strings"

	"strconv"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	. "github.com/claudetech/loggo/default"
	"golang.org/x/net/context"
)

// Mount the fuse volume
func Mount(client *Drive, mountpoint string, mountOptions []string, uid, gid uint32, umask os.FileMode) error {
	Log.Infof("Mounting path %v", mountpoint)

	if _, err := os.Stat(mountpoint); os.IsNotExist(err) {
		Log.Debugf("Mountpoint doesn't exist, creating...")
		if err := os.MkdirAll(mountpoint, 0644); nil != err {
			Log.Debugf("%v", err)
			return fmt.Errorf("Could not create mount directory %v", mountpoint)
		}
	}

	fuse.Debug = func(msg interface{}) {
		Log.Tracef("FUSE %v", msg)
	}

	// Set mount options
	options := []fuse.MountOption{
		fuse.NoAppleDouble(),
		fuse.NoAppleXattr(),
	}
	for _, option := range mountOptions {
		if "allow_other" == option {
			options = append(options, fuse.AllowOther())
		} else if "allow_root" == option {
			options = append(options, fuse.AllowRoot())
		} else if "allow_dev" == option {
			options = append(options, fuse.AllowDev())
		} else if "allow_non_empty_mount" == option {
			options = append(options, fuse.AllowNonEmptyMount())
		} else if "allow_suid" == option {
			options = append(options, fuse.AllowSUID())
		} else if strings.Contains(option, "max_readahead=") {
			data := strings.Split(option, "=")
			value, err := strconv.ParseUint(data[1], 10, 32)
			if nil != err {
				Log.Debugf("%v", err)
				return fmt.Errorf("Could not parse max_readahead value")
			}
			options = append(options, fuse.MaxReadahead(uint32(value)))
		} else if "default_permissions" == option {
			options = append(options, fuse.DefaultPermissions())
		} else if "excl_create" == option {
			options = append(options, fuse.ExclCreate())
		} else if strings.Contains(option, "fs_name") {
			data := strings.Split(option, "=")
			options = append(options, fuse.FSName(data[1]))
		} else if "local_volume" == option {
			options = append(options, fuse.LocalVolume())
		} else if "writeback_cache" == option {
			options = append(options, fuse.WritebackCache())
		} else if strings.Contains(option, "volume_name") {
			data := strings.Split(option, "=")
			options = append(options, fuse.VolumeName(data[1]))
		} else if "read_only" == option {
			options = append(options, fuse.ReadOnly())
		} else {
			Log.Warningf("Fuse option %v is not supported, yet", option)
		}
	}

	c, err := fuse.Mount(mountpoint, options...)
	if err != nil {
		return err
	}
	defer c.Close()

	filesys := &FS{
		client: client,
		uid:    uid,
		gid:    gid,
		umask:  umask,
	}
	if err := fs.Serve(c, filesys); err != nil {
		return err
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Error mounting FUSE")
	}

	return Unmount(mountpoint, true)
}

// Unmount unmounts the mountpoint
func Unmount(mountpoint string, notify bool) error {
	if notify {
		Log.Infof("Unmounting path %v", mountpoint)
	}
	fuse.Unmount(mountpoint)
	return nil
}

// FS the fuse filesystem
type FS struct {
	client *Drive
	uid    uint32
	gid    uint32
	umask  os.FileMode
}

// Root returns the root path
func (f *FS) Root() (fs.Node, error) {
	object, err := f.client.GetRoot()
	if nil != err {
		Log.Warningf("%v", err)
		return nil, fmt.Errorf("Could not get root object")
	}
	return &Object{
		client: f.client,
		object: object,
		uid:    f.uid,
		gid:    f.gid,
		umask:  f.umask,
	}, nil
}

// Object represents one drive object
type Object struct {
	client *Drive
	object *APIObject
	buffer *Buffer
	uid    uint32
	gid    uint32
	umask  os.FileMode
}

// Attr returns the attributes for a directory
func (o *Object) Attr(ctx context.Context, attr *fuse.Attr) error {
	if o.object.IsDir {
		if o.umask > 0 {
			attr.Mode = os.ModeDir | o.umask
		} else {
			attr.Mode = os.ModeDir | 0755
		}
		attr.Size = 0
	} else {
		if o.umask > 0 {
			attr.Mode = o.umask
		} else {
			attr.Mode = 0644
		}
		attr.Size = o.object.Size
	}

	attr.Uid = uint32(o.uid)
	attr.Gid = uint32(o.gid)

	attr.Mtime = o.object.LastModified
	attr.Crtime = o.object.LastModified
	attr.Ctime = o.object.LastModified

	attr.Blocks = (attr.Size + 511) / 512

	return nil
}

// ReadDirAll shows all files in the current directory
func (o *Object) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	objects, err := o.client.GetObjectsByParent(o.object.ObjectID)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fuse.ENOENT
	}

	dirs := []fuse.Dirent{}
	for _, object := range objects {
		if object.IsDir {
			dirs = append(dirs, fuse.Dirent{
				Name: object.Name,
				Type: fuse.DT_Dir,
			})
		} else {
			dirs = append(dirs, fuse.Dirent{
				Name: object.Name,
				Type: fuse.DT_File,
			})
		}
	}
	return dirs, nil
}

// Lookup tests if a file is existent in the current directory
func (o *Object) Lookup(ctx context.Context, name string) (fs.Node, error) {
	object, err := o.client.GetObjectByParentAndName(o.object.ObjectID, name)
	if nil != err {
		Log.Tracef("%v", err)
		return nil, fuse.ENOENT
	}

	return &Object{
		client: o.client,
		object: object,
		uid:    o.uid,
		gid:    o.gid,
		umask:  o.umask,
	}, nil
}

// Open opens a file for reading
func (o *Object) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	if req.Dir {
		return o, nil
	}

	buffer, err := o.client.Open(o.object)
	if nil != err {
		Log.Warningf("%v", err)
		return o, fuse.ENOENT
	}
	o.buffer = buffer

	return o, nil
}

// Release a stream
func (o *Object) Release(ctx context.Context, req *fuse.ReleaseRequest) error {
	if nil != o.buffer {
		if err := o.buffer.Close(); nil != err {
			Log.Debugf("%v", err)
			Log.Warningf("Could not close buffer stream")
		}
	}
	return nil
}

// Read reads some bytes or the whole file
func (o *Object) Read(ctx context.Context, req *fuse.ReadRequest, resp *fuse.ReadResponse) error {
	buf, err := o.buffer.ReadBytes(req.Offset, int64(req.Size), false)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	resp.Data = buf[:]
	return nil
}

// Remove deletes an element
func (o *Object) Remove(ctx context.Context, req *fuse.RemoveRequest) error {
	obj, err := o.client.GetObjectByParentAndName(o.object.ObjectID, req.Name)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	err = o.client.Remove(obj, o.object.ObjectID)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	return nil
}

// Rename renames an element
func (o *Object) Rename(ctx context.Context, req *fuse.RenameRequest, newDir fs.Node) error {
	obj, err := o.client.GetObjectByParentAndName(o.object.ObjectID, req.OldName)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	destDir, ok := newDir.(*Object)
	if !ok {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	err = o.client.Rename(obj, destDir.object.ObjectID, req.NewName)
	if nil != err {
		Log.Warningf("%v", err)
		return fuse.EIO
	}

	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSIAMGroup_basic(t *testing.T) {
	var conf iam.GetGroupOutput

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSGroupDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSGroupConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupExists("aws_iam_group.group", &conf),
					testAccCheckAWSGroupAttributes(&conf),
				),
			},
			resource.TestStep{
				Config: testAccAWSGroupConfig2,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupExists("aws_iam_group.group", &conf),
					testAccCheckAWSGroupAttributes(&conf),
				),
			},
		},
	})
}

func testAccCheckAWSGroupDestroy(s *terraform.State) error {
	iamconn := testAccProvider.Meta().(*AWSClient).iamconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_iam_group" {
			continue
		}

		// Try to get group
		_, err := iamconn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(rs.Primary.ID),
		})
		if err == nil {
			return fmt.Errorf("still exist.")
		}

		// Verify the error is what we want
		ec2err, ok := err.(awserr.Error)
		if !ok {
			return err
		}
		if ec2err.Code() != "NoSuchEntity" {
			return err
		}
	}

	return nil
}

func testAccCheckAWSGroupExists(n string, res *iam.GetGroupOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No Group name is set")
		}

		iamconn := testAccProvider.Meta().(*AWSClient).iamconn

		resp, err := iamconn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(rs.Primary.ID),
		})
		if err != nil {
			return err
		}

		*res = *resp

		return nil
	}
}

func testAccCheckAWSGroupAttributes(group *iam.GetGroupOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *group.Group.GroupName != "test-group" {
			return fmt.Errorf("Bad name: %s", *group.Group.GroupName)
		}

		if *group.Group.Path != "/" {
			return fmt.Errorf("Bad path: %s", *group.Group.Path)
		}

		return nil
	}
}

const testAccAWSGroupConfig = `
resource "aws_iam_group" "group" {
	name = "test-group"
	path = "/"
}
`
const testAccAWSGroupConfig2 = `
resource "aws_iam_group" "group" {
	name = "test-group2"
	path = "/funnypath/"
}
`
<commit_msg>update aws_iam_group tests to check proper paths and name<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSIAMGroup_basic(t *testing.T) {
	var conf iam.GetGroupOutput

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSGroupDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSGroupConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupExists("aws_iam_group.group", &conf),
					testAccCheckAWSGroupAttributes(&conf, "test-group", "/"),
				),
			},
			resource.TestStep{
				Config: testAccAWSGroupConfig2,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupExists("aws_iam_group.group", &conf),
					testAccCheckAWSGroupAttributes(&conf, "test-group2", "/funnypath/"),
				),
			},
		},
	})
}

func testAccCheckAWSGroupDestroy(s *terraform.State) error {
	iamconn := testAccProvider.Meta().(*AWSClient).iamconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_iam_group" {
			continue
		}

		// Try to get group
		_, err := iamconn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(rs.Primary.ID),
		})
		if err == nil {
			return fmt.Errorf("still exist.")
		}

		// Verify the error is what we want
		ec2err, ok := err.(awserr.Error)
		if !ok {
			return err
		}
		if ec2err.Code() != "NoSuchEntity" {
			return err
		}
	}

	return nil
}

func testAccCheckAWSGroupExists(n string, res *iam.GetGroupOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No Group name is set")
		}

		iamconn := testAccProvider.Meta().(*AWSClient).iamconn

		resp, err := iamconn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(rs.Primary.ID),
		})
		if err != nil {
			return err
		}

		*res = *resp

		return nil
	}
}

func testAccCheckAWSGroupAttributes(group *iam.GetGroupOutput, name string, path string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if *group.Group.GroupName != name {
			return fmt.Errorf("Bad name: %s when %s was expected", *group.Group.GroupName, name)
		}

		if *group.Group.Path != path {
			return fmt.Errorf("Bad path: %s when %s was expected", *group.Group.Path, path)
		}

		return nil
	}
}

const testAccAWSGroupConfig = `
resource "aws_iam_group" "group" {
	name = "test-group"
	path = "/"
}
`
const testAccAWSGroupConfig2 = `
resource "aws_iam_group" "group" {
	name = "test-group2"
	path = "/funnypath/"
}
`
<|endoftext|>


<commit_before>package perfM

import (
	"log"
	"time"

	hist "github.com/VividCortex/gohistogram"
)

type Job interface {
	Done(PerfMonitor) error //count the cost about this job and add to the perfmonitor count channel
}

type job struct {
	start time.Time //set for every single request start time
}

func (j *job) Done(p *perfMonitor) error {
	cost := time.Since(j.start)
	p.GlobalChannel <- cost
	return nil
}

type PerfMonitor interface {
	Start() error    //start the perf monitor
	Stop() error     //stop the perf montior
	Do() (Job, erro) //set a timer to count the single request's cost
}

type perfMonitor struct {
	done           chan bool              //stor the perfM
	counter        int                    //count the sum of the request
	startTime      time.Time              //keep the start time
	timer          time.Timer             //the frequency sampling timer
	GlobalChannel  chan time.Duration     //get the request cost from every done()
	localCount     int                    //count for the number in the sampling times
	localTimeCount time.Duration          //count for the sampling time total costs
	histogram      *hist.NumericHistogram //used to print the histogram
}

func New(conf Config) {
	return &perfMonitor{
		counter:       0,
		startTime:     time.Time,
		timer:         time.NewTimer(time.Second * int64(conf.Frequency)),
		GlobalChannel: make(chan time.Duration, conf.BufferSize),
		histogram:     hist.NewHistogram(conf.BinsNumber),
	}
}

func (p *perfMonitor) Start() error {
	p.startTime = time.Now()
	for {
		select {
		case cost := <-p.GlobalChannel:
			p.counter++
			p.localCount++
			p.localTimeCount += cost
			p.histogram.Add(float64(cost))
		case <-p.timer:
			log.Println("Qps: ", p.localCount, "Avg Latency: ", p.localTimeCount/int64(p.localCount))
			p.localCount = 0
			p.localTimeCount = 0
		case <-p.done:
			return nil
		}
	}
	return nil
}

func (p *perfMonitor) Stop() error {
	//TODO:show the info of the performence test
	p.done <- true
	//show the summery

	// here show the histogram
	log.Println(p.histogram.String())
	return nil
}

func (p *perfMonitor) Do() (*Job, error) {
	presentJob := new(job)
	presentJob.start = time.Now()
	return presentJob, nil
}
<commit_msg>fast fix<commit_after>package perfM

import (
	"log"
	"time"

	hist "github.com/VividCortex/gohistogram"
)

type Job interface {
	Done(PerfMonitor) //count the cost about this job and add to the perfmonitor count channel
}

type job struct {
	start time.Time //set for every single request start time
}

func (j *job) Done(p PerfMonitor) {
	cost := time.Since(j.start)
	p.collect(cost)
}

type PerfMonitor interface {
	Start()                //start the perf monitor
	Stop()                 //stop the perf montior
	collect(time.Duration) //send the cost to the channel
	Do() Job               //set a timer to count the single request's cost
}

type perfMonitor struct {
	done           chan bool              //stor the perfM
	counter        int                    //count the sum of the request
	startTime      time.Time              //keep the start time
	timer          <-chan time.Time       //the frequency sampling timer
	Collector      chan time.Duration     //get the request cost from every done()
	localCount     int                    //count for the number in the sampling times
	localTimeCount time.Duration          //count for the sampling time total costs
	histogram      *hist.NumericHistogram //used to print the histogram
}

func New(conf Config) PerfMonitor {
	if conf.BinsNumber == 0 {
		conf.BinsNumber = 10
		conf.BufferSize = 1000000
		conf.Frequency = 1
	}
	return &perfMonitor{
		counter:   0,
		startTime: time.Now(),
		timer:     time.Tick(time.Duration(int64(1000000000 * conf.Frequency))),
		Collector: make(chan time.Duration, conf.BufferSize),
		histogram: hist.NewHistogram(conf.BinsNumber),
	}
}

func (p *perfMonitor) Start() {
	p.startTime = time.Now()
	for {
		select {
		case cost := <-p.Collector:
			p.counter++
			p.localCount++
			p.localTimeCount += cost
			p.histogram.Add(float64(cost))
		case <-p.timer:
			if p.localCount == 0 {
				continue
			}
			log.Println("Qps: ", p.localCount, "Avg Latency: ", p.localTimeCount.Nanoseconds()/int64(p.localCount)/1000000)
			p.localCount = 0
			p.localTimeCount = 0
		case <-p.done:
			return
		}
	}
}

func (p *perfMonitor) Stop() {
	//TODO:show the info of the performence test
	p.done <- true
	//show the summery
	log.Println("Total Recv: ", p.counter)

	// here show the histogram
	log.Println(p.histogram.String())
}

func (p *perfMonitor) Do() Job {
	presentJob := new(job)
	presentJob.start = time.Now()
	return presentJob
}

func (p *perfMonitor) collect(t time.Duration) {
	p.Collector <- t
}
<|endoftext|>


<commit_before>package streammux

import (
	"io"
	"net"
	"time"
)

// Stream is a bidirectional io pipe within a connection.
type Stream interface {
	io.Reader
	io.Writer

	// Close closes the stream for writing. Reading will still work (that
	// is, the remote side can still write).
	io.Closer

	// Reset closes both ends of the stream. Use this to tell the remote
	// side to hang up and go away.
	Reset() error

	SetDeadline(time.Time) error
	SetReadDeadline(time.Time) error
	SetWriteDeadline(time.Time) error
}

// NoOpHandler do nothing. close streams as soon as they are opened.
var NoOpHandler = func(s Stream) { s.Close() }

// Conn is a stream-multiplexing connection to a remote peer.
type Conn interface {
	io.Closer

	// IsClosed returns whether a connection is fully closed, so it can
	// be garbage collected.
	IsClosed() bool

	// OpenStream creates a new stream.
	OpenStream() (Stream, error)

	// AcceptStream accepts a stream opened by the other side.
	AcceptStream() (Stream, error)
}

// Transport constructs go-stream-muxer compatible connections.
type Transport interface {

	// NewConn constructs a new connection
	NewConn(c net.Conn, isServer bool) (Conn, error)
}
<commit_msg>make the NoOpHandler reset streams instead of closing them.<commit_after>package streammux

import (
	"io"
	"net"
	"time"
)

// Stream is a bidirectional io pipe within a connection.
type Stream interface {
	io.Reader
	io.Writer

	// Close closes the stream for writing. Reading will still work (that
	// is, the remote side can still write).
	io.Closer

	// Reset closes both ends of the stream. Use this to tell the remote
	// side to hang up and go away.
	Reset() error

	SetDeadline(time.Time) error
	SetReadDeadline(time.Time) error
	SetWriteDeadline(time.Time) error
}

// NoOpHandler do nothing. Resets streams as soon as they are opened.
var NoOpHandler = func(s Stream) { s.Reset() }

// Conn is a stream-multiplexing connection to a remote peer.
type Conn interface {
	io.Closer

	// IsClosed returns whether a connection is fully closed, so it can
	// be garbage collected.
	IsClosed() bool

	// OpenStream creates a new stream.
	OpenStream() (Stream, error)

	// AcceptStream accepts a stream opened by the other side.
	AcceptStream() (Stream, error)
}

// Transport constructs go-stream-muxer compatible connections.
type Transport interface {

	// NewConn constructs a new connection
	NewConn(c net.Conn, isServer bool) (Conn, error)
}
<|endoftext|>


<commit_before>// Package pgenv provides the ability to get default postgres
// credentials from the users env
package pgenv

import (
	"os"
	"strings"
)

// ConnStr string map to relate the key values pairs of a pg connection string
type ConnStr map[string]string

/*
Options are from: http://www.postgresql.org/docs/current/static/libpq-envars.html

Future plans are to support more options.
*/

func getVar(envVal string, def string) string {
	var ret string
	ret = os.Getenv(envVal)
	if ret == "" {
		ret = def
	}
	return ret
}

// SetDefaults sets a few common default variables based on os env
func (c ConnStr) SetDefaults() ConnStr {
	c["user"] = getVar("PGUSER", "postgres")
	c["dbname"] = getVar("PGDATABASE", "postgres")
	c["host"] = getVar("PGHOST", "localhost")
	c["port"] = getVar("PGPORT", "5432")
	c["password"] = getVar("PGPASSWD", "")
	c["sslmode"] = getVar("PGSSLMODE", "disable")

	return c
}

// ToString returns a string that can be dropped into sql.Open
func (c ConnStr) ToString() string {
	var s []string
	for key, val := range c {
		if val != "" {
			s = append(s, key+"="+val)
		}
	}

	return strings.Join(s, " ")
}
<commit_msg>documentation<commit_after>// Package pgenv provides a quick way to create db connection strings
// from the running users environment variables.
package pgenv

import (
	"os"
	"strings"
)

// ConnStr string map to relate the key values pairs of a pg connection string
type ConnStr map[string]string

/*
Options are from: http://www.postgresql.org/docs/current/static/libpq-envars.html

Future plans are to support more options.
*/

func getVar(envVal string, def string) string {
	var ret string
	ret = os.Getenv(envVal)
	if ret == "" {
		ret = def
	}
	return ret
}

// SetDefaults pulls common environment variables and sets them for later use.
func (c ConnStr) SetDefaults() ConnStr {
	c["user"] = getVar("PGUSER", "postgres")
	c["dbname"] = getVar("PGDATABASE", "postgres")
	c["host"] = getVar("PGHOST", "localhost")
	c["port"] = getVar("PGPORT", "5432")
	c["password"] = getVar("PGPASSWD", "")
	c["sslmode"] = getVar("PGSSLMODE", "disable")

	return c
}

// ToString converts the string map to a sql.Open() compatable connection string.
func (c ConnStr) ToString() string {
	var s []string
	for key, val := range c {
		if val != "" {
			s = append(s, key+"="+val)
		}
	}

	return strings.Join(s, " ")
}
<|endoftext|>


<commit_before>package apparmor

import (
	"fmt"
	"strings"
	"text/template"

	"github.com/lxc/lxd/lxd/sys"
	"github.com/lxc/lxd/shared"
)

var dnsmasqProfileTpl = template.Must(template.New("dnsmasqProfile").Parse(`#include <tunables/global>
profile "{{ .name }}" flags=(attach_disconnected,mediate_deleted) {
  #include <abstractions/base>
  #include <abstractions/dbus>
  #include <abstractions/nameservice>

  # Capabilities
  capability chown,
  capability net_bind_service,
  capability setgid,
  capability setuid,
  capability dac_override,
  capability dac_read_search,
  capability net_admin,         # for DHCP server
  capability net_raw,           # for DHCP server ping checks

  # Network access
  network inet raw,
  network inet6 raw,

  # Network-specific paths
  {{ .varPath }}/networks/{{ .networkName }}/dnsmasq.hosts/{,*} r,
  {{ .varPath }}/networks/{{ .networkName }}/dnsmasq.leases rw,
  {{ .varPath }}/networks/{{ .networkName }}/dnsmasq.raw r,

  # Additional system files
  @{PROC}/sys/net/ipv6/conf/*/mtu r,
  @{PROC}/@{pid}/fd/ r,
  {{ .rootPath }}/etc/localtime  r,
  {{ .rootPath }}/usr/share/zoneinfo/**  r,

  # System configuration access
  {{ .rootPath }}/etc/gai.conf           r,
  {{ .rootPath }}/etc/group              r,
  {{ .rootPath }}/etc/host.conf          r,
  {{ .rootPath }}/etc/hosts              r,
  {{ .rootPath }}/etc/nsswitch.conf      r,
  {{ .rootPath }}/etc/passwd             r,
  {{ .rootPath }}/etc/protocols          r,

  {{ .rootPath }}/etc/resolv.conf        r,
  {{ .rootPath }}/etc/resolvconf/run/resolv.conf r,

  {{ .rootPath }}/run/{resolvconf,NetworkManager,systemd/resolve,connman,netconfig}/resolv.conf r,
  {{ .rootPath }}/run/systemd/resolve/stub-resolv.conf r,

{{- if .snap }}

  # The binary itself (for nesting)
  /snap/lxd/*/bin/dnsmasq                 mr,

  # Snap-specific libraries
  /snap/lxd/*/lib/**.so*                  mr,
{{- end }}
}
`))

// dnsmasqProfile generates the AppArmor profile template from the given network.
func dnsmasqProfile(sysOS *sys.OS, n network) (string, error) {
	rootPath := ""
	if shared.InSnap() {
		rootPath = "/var/lib/snapd/hostfs"
	}

	// Render the profile.
	var sb *strings.Builder = &strings.Builder{}
	err := dnsmasqProfileTpl.Execute(sb, map[string]any{
		"name":        DnsmasqProfileName(n),
		"networkName": n.Name(),
		"varPath":     shared.VarPath(""),
		"rootPath":    rootPath,
		"snap":        shared.InSnap(),
	})
	if err != nil {
		return "", err
	}

	return sb.String(), nil
}

// DnsmasqProfileName returns the AppArmor profile name.
func DnsmasqProfileName(n network) string {
	path := shared.VarPath("")
	name := fmt.Sprintf("%s_<%s>", n.Name(), path)
	return profileName("dnsmasq", name)
}

// dnsmasqProfileFilename returns the name of the on-disk profile name.
func dnsmasqProfileFilename(n network) string {
	return profileName("dnsmasq", n.Name())
}
<commit_msg>lxd/apparmor/dnsmasq: Support non-snap nesting<commit_after>package apparmor

import (
	"fmt"
	"strings"
	"text/template"

	"github.com/lxc/lxd/lxd/sys"
	"github.com/lxc/lxd/shared"
)

var dnsmasqProfileTpl = template.Must(template.New("dnsmasqProfile").Parse(`#include <tunables/global>
profile "{{ .name }}" flags=(attach_disconnected,mediate_deleted) {
  #include <abstractions/base>
  #include <abstractions/dbus>
  #include <abstractions/nameservice>

  # Capabilities
  capability chown,
  capability net_bind_service,
  capability setgid,
  capability setuid,
  capability dac_override,
  capability dac_read_search,
  capability net_admin,         # for DHCP server
  capability net_raw,           # for DHCP server ping checks

  # Network access
  network inet raw,
  network inet6 raw,

  # Network-specific paths
  {{ .varPath }}/networks/{{ .networkName }}/dnsmasq.hosts/{,*} r,
  {{ .varPath }}/networks/{{ .networkName }}/dnsmasq.leases rw,
  {{ .varPath }}/networks/{{ .networkName }}/dnsmasq.raw r,

  # Additional system files
  @{PROC}/sys/net/ipv6/conf/*/mtu r,
  @{PROC}/@{pid}/fd/ r,
  {{ .rootPath }}/etc/localtime  r,
  {{ .rootPath }}/usr/share/zoneinfo/**  r,

  # System configuration access
  {{ .rootPath }}/etc/gai.conf           r,
  {{ .rootPath }}/etc/group              r,
  {{ .rootPath }}/etc/host.conf          r,
  {{ .rootPath }}/etc/hosts              r,
  {{ .rootPath }}/etc/nsswitch.conf      r,
  {{ .rootPath }}/etc/passwd             r,
  {{ .rootPath }}/etc/protocols          r,

  {{ .rootPath }}/etc/resolv.conf        r,
  {{ .rootPath }}/etc/resolvconf/run/resolv.conf r,

  {{ .rootPath }}/run/{resolvconf,NetworkManager,systemd/resolve,connman,netconfig}/resolv.conf r,
  {{ .rootPath }}/run/systemd/resolve/stub-resolv.conf r,

{{- if .snap }}

  # The binary itself (for nesting)
  /snap/lxd/*/bin/dnsmasq                 mr,

  # Snap-specific libraries
  /snap/lxd/*/lib/**.so*                  mr,
{{ else }}
  # The binary itself (for nesting)
  /{,usr/}sbin/dnsmasq                    mr,
{{- end }}
}
`))

// dnsmasqProfile generates the AppArmor profile template from the given network.
func dnsmasqProfile(sysOS *sys.OS, n network) (string, error) {
	rootPath := ""
	if shared.InSnap() {
		rootPath = "/var/lib/snapd/hostfs"
	}

	// Render the profile.
	var sb *strings.Builder = &strings.Builder{}
	err := dnsmasqProfileTpl.Execute(sb, map[string]any{
		"name":        DnsmasqProfileName(n),
		"networkName": n.Name(),
		"varPath":     shared.VarPath(""),
		"rootPath":    rootPath,
		"snap":        shared.InSnap(),
	})
	if err != nil {
		return "", err
	}

	return sb.String(), nil
}

// DnsmasqProfileName returns the AppArmor profile name.
func DnsmasqProfileName(n network) string {
	path := shared.VarPath("")
	name := fmt.Sprintf("%s_<%s>", n.Name(), path)
	return profileName("dnsmasq", name)
}

// dnsmasqProfileFilename returns the name of the on-disk profile name.
func dnsmasqProfileFilename(n network) string {
	return profileName("dnsmasq", n.Name())
}
<|endoftext|>


<commit_before>package copystructure

import (
	"errors"
	"reflect"
	"sync"

	"github.com/mitchellh/reflectwalk"
)

// Copy returns a deep copy of v.
func Copy(v interface{}) (interface{}, error) {
	return Config{}.Copy(v)
}

// CopierFunc is a function that knows how to deep copy a specific type.
// Register these globally with the Copiers variable.
type CopierFunc func(interface{}) (interface{}, error)

// Copiers is a map of types that behave specially when they are copied.
// If a type is found in this map while deep copying, this function
// will be called to copy it instead of attempting to copy all fields.
//
// The key should be the type, obtained using: reflect.TypeOf(value with type).
//
// It is unsafe to write to this map after Copies have started. If you
// are writing to this map while also copying, wrap all modifications to
// this map as well as to Copy in a mutex.
var Copiers map[reflect.Type]CopierFunc = make(map[reflect.Type]CopierFunc)

// Must is a helper that wraps a call to a function returning
// (interface{}, error) and panics if the error is non-nil. It is intended
// for use in variable initializations and should only be used when a copy
// error should be a crashing case.
func Must(v interface{}, err error) interface{} {
	if err != nil {
		panic("copy error: " + err.Error())
	}

	return v
}

var errPointerRequired = errors.New("Copy argument must be a pointer when Lock is true")

type Config struct {
	// Lock any types that are a sync.Locker and are not a mutex while copying.
	// If there is an RLocker method, use that to get the sync.Locker.
	Lock bool

	// Copiers is a map of types associated with a CopierFunc. Use the global
	// Copiers map if this is nil.
	Copiers map[reflect.Type]CopierFunc
}

func (c Config) Copy(v interface{}) (interface{}, error) {
	if c.Lock && reflect.ValueOf(v).Kind() != reflect.Ptr {
		return nil, errPointerRequired
	}

	w := new(walker)
	if c.Lock {
		w.useLocks = true
	}

	if c.Copiers == nil {
		c.Copiers = Copiers
	}

	err := reflectwalk.Walk(v, w)
	if err != nil {
		return nil, err
	}

	// Get the result. If the result is nil, then we want to turn it
	// into a typed nil if we can.
	result := w.Result
	if result == nil {
		val := reflect.ValueOf(v)
		result = reflect.Indirect(reflect.New(val.Type())).Interface()
	}

	return result, nil
}

// Return the key used to index interfaces types we've seen. Store the number
// of pointers in the upper 32bits, and the depth in the lower 32bits. This is
// easy to calculate, easy to match a key with our current depth, and we don't
// need to deal with initializing and cleaning up nested maps or slices.
func ifaceKey(pointers, depth int) uint64 {
	return uint64(pointers)<<32 | uint64(depth)
}

type walker struct {
	Result interface{}

	depth       int
	ignoreDepth int
	vals        []reflect.Value
	cs          []reflect.Value

	// This stores the number of pointers we've walked over, indexed by depth.
	ps []int

	// If an interface is indirected by a pointer, we need to know the type of
	// interface to create when creating the new value.  Store the interface
	// types here, indexed by both the walk depth and the number of pointers
	// already seen at that depth. Use ifaceKey to calculate the proper uint64
	// value.
	ifaceTypes map[uint64]reflect.Type

	// any locks we've taken, indexed by depth
	locks []sync.Locker
	// take locks while walking the structure
	useLocks bool
}

func (w *walker) Enter(l reflectwalk.Location) error {
	w.depth++

	// ensure we have enough elements to index via w.depth
	for w.depth >= len(w.locks) {
		w.locks = append(w.locks, nil)
	}

	for len(w.ps) < w.depth+1 {
		w.ps = append(w.ps, 0)
	}

	return nil
}

func (w *walker) Exit(l reflectwalk.Location) error {
	locker := w.locks[w.depth]
	w.locks[w.depth] = nil
	if locker != nil {
		defer locker.Unlock()
	}

	// clear out pointers and interfaces as we exit the stack
	w.ps[w.depth] = 0

	for k := range w.ifaceTypes {
		mask := uint64(^uint32(0))
		if k&mask == uint64(w.depth) {
			delete(w.ifaceTypes, k)
		}
	}

	w.depth--
	if w.ignoreDepth > w.depth {
		w.ignoreDepth = 0
	}

	if w.ignoring() {
		return nil
	}

	switch l {
	case reflectwalk.Map:
		fallthrough
	case reflectwalk.Slice:
		// Pop map off our container
		w.cs = w.cs[:len(w.cs)-1]
	case reflectwalk.MapValue:
		// Pop off the key and value
		mv := w.valPop()
		mk := w.valPop()
		m := w.cs[len(w.cs)-1]

		// If mv is the zero value, SetMapIndex deletes the key form the map,
		// or in this case never adds it. We need to create a properly typed
		// zero value so that this key can be set.
		if !mv.IsValid() {
			mv = reflect.Zero(m.Type().Elem())
		}
		m.SetMapIndex(mk, mv)
	case reflectwalk.SliceElem:
		// Pop off the value and the index and set it on the slice
		v := w.valPop()
		if v.IsValid() {
			i := w.valPop().Interface().(int)
			s := w.cs[len(w.cs)-1]
			se := s.Index(i)
			if se.CanSet() {
				se.Set(v)
			}
		}
	case reflectwalk.Struct:
		w.replacePointerMaybe()

		// Remove the struct from the container stack
		w.cs = w.cs[:len(w.cs)-1]
	case reflectwalk.StructField:
		// Pop off the value and the field
		v := w.valPop()
		f := w.valPop().Interface().(reflect.StructField)
		if v.IsValid() {
			s := w.cs[len(w.cs)-1]
			sf := reflect.Indirect(s).FieldByName(f.Name)

			if sf.CanSet() {
				sf.Set(v)
			}
		}
	case reflectwalk.WalkLoc:
		// Clear out the slices for GC
		w.cs = nil
		w.vals = nil
	}

	return nil
}

func (w *walker) Map(m reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(m)

	// Create the map. If the map itself is nil, then just make a nil map
	var newMap reflect.Value
	if m.IsNil() {
		newMap = reflect.Indirect(reflect.New(m.Type()))
	} else {
		newMap = reflect.MakeMap(m.Type())
	}

	w.cs = append(w.cs, newMap)
	w.valPush(newMap)
	return nil
}

func (w *walker) MapElem(m, k, v reflect.Value) error {
	return nil
}

func (w *walker) PointerEnter(v bool) error {
	if v {
		w.ps[w.depth]++
	}
	return nil
}

func (w *walker) PointerExit(v bool) error {
	if v {
		w.ps[w.depth]--
	}
	return nil
}

func (w *walker) Interface(v reflect.Value) error {
	if !v.IsValid() {
		return nil
	}
	if w.ifaceTypes == nil {
		w.ifaceTypes = make(map[uint64]reflect.Type)
	}

	w.ifaceTypes[ifaceKey(w.ps[w.depth], w.depth)] = v.Type()
	return nil
}

func (w *walker) Primitive(v reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(v)

	// IsValid verifies the v is non-zero and CanInterface verifies
	// that we're allowed to read this value (unexported fields).
	var newV reflect.Value
	if v.IsValid() && v.CanInterface() {
		newV = reflect.New(v.Type())
		newV.Elem().Set(v)
	}

	w.valPush(newV)
	w.replacePointerMaybe()
	return nil
}

func (w *walker) Slice(s reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(s)

	var newS reflect.Value
	if s.IsNil() {
		newS = reflect.Indirect(reflect.New(s.Type()))
	} else {
		newS = reflect.MakeSlice(s.Type(), s.Len(), s.Cap())
	}

	w.cs = append(w.cs, newS)
	w.valPush(newS)
	return nil
}

func (w *walker) SliceElem(i int, elem reflect.Value) error {
	if w.ignoring() {
		return nil
	}

	// We don't write the slice here because elem might still be
	// arbitrarily complex. Just record the index and continue on.
	w.valPush(reflect.ValueOf(i))

	return nil
}

func (w *walker) Struct(s reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(s)

	var v reflect.Value
	if c, ok := Copiers[s.Type()]; ok {
		// We have a Copier for this struct, so we use that copier to
		// get the copy, and we ignore anything deeper than this.
		w.ignoreDepth = w.depth

		dup, err := c(s.Interface())
		if err != nil {
			return err
		}

		v = reflect.ValueOf(dup)
	} else {
		// No copier, we copy ourselves and allow reflectwalk to guide
		// us deeper into the structure for copying.
		v = reflect.New(s.Type())
	}

	// Push the value onto the value stack for setting the struct field,
	// and add the struct itself to the containers stack in case we walk
	// deeper so that its own fields can be modified.
	w.valPush(v)
	w.cs = append(w.cs, v)

	return nil
}

func (w *walker) StructField(f reflect.StructField, v reflect.Value) error {
	if w.ignoring() {
		return nil
	}

	// If PkgPath is non-empty, this is a private (unexported) field.
	// We do not set this unexported since the Go runtime doesn't allow us.
	if f.PkgPath != "" {
		return reflectwalk.SkipEntry
	}

	// Push the field onto the stack, we'll handle it when we exit
	// the struct field in Exit...
	w.valPush(reflect.ValueOf(f))
	return nil
}

// ignore causes the walker to ignore any more values until we exit this on
func (w *walker) ignore() {
	w.ignoreDepth = w.depth
}

func (w *walker) ignoring() bool {
	return w.ignoreDepth > 0 && w.depth >= w.ignoreDepth
}

func (w *walker) pointerPeek() bool {
	return w.ps[w.depth] > 0
}

func (w *walker) valPop() reflect.Value {
	result := w.vals[len(w.vals)-1]
	w.vals = w.vals[:len(w.vals)-1]

	// If we're out of values, that means we popped everything off. In
	// this case, we reset the result so the next pushed value becomes
	// the result.
	if len(w.vals) == 0 {
		w.Result = nil
	}

	return result
}

func (w *walker) valPush(v reflect.Value) {
	w.vals = append(w.vals, v)

	// If we haven't set the result yet, then this is the result since
	// it is the first (outermost) value we're seeing.
	if w.Result == nil && v.IsValid() {
		w.Result = v.Interface()
	}
}

func (w *walker) replacePointerMaybe() {
	// Determine the last pointer value. If it is NOT a pointer, then
	// we need to push that onto the stack.
	if !w.pointerPeek() {
		w.valPush(reflect.Indirect(w.valPop()))
		return
	}

	v := w.valPop()
	for i := 1; i < w.ps[w.depth]; i++ {
		if iType, ok := w.ifaceTypes[ifaceKey(w.ps[w.depth]-i, w.depth)]; ok {
			iface := reflect.New(iType).Elem()
			iface.Set(v)
			v = iface
		}

		p := reflect.New(v.Type())
		p.Elem().Set(v)
		v = p
	}

	w.valPush(v)
}

// if this value is a Locker, lock it and add it to the locks slice
func (w *walker) lock(v reflect.Value) {
	if !w.useLocks {
		return
	}

	if !v.IsValid() || !v.CanInterface() {
		return
	}

	type rlocker interface {
		RLocker() sync.Locker
	}

	var locker sync.Locker

	// We can't call Interface() on a value directly, since that requires
	// a copy. This is OK, since the pointer to a value which is a sync.Locker
	// is also a sync.Locker.
	if v.Kind() == reflect.Ptr {
		switch l := v.Interface().(type) {
		case rlocker:
			// don't lock a mutex directly
			if _, ok := l.(*sync.RWMutex); !ok {
				locker = l.RLocker()
			}
		case sync.Locker:
			locker = l
		}
	} else if v.CanAddr() {
		switch l := v.Addr().Interface().(type) {
		case rlocker:
			// don't lock a mutex directly
			if _, ok := l.(*sync.RWMutex); !ok {
				locker = l.RLocker()
			}
		case sync.Locker:
			locker = l
		}
	}

	// still no callable locker
	if locker == nil {
		return
	}

	// don't lock a mutex directly
	switch locker.(type) {
	case *sync.Mutex, *sync.RWMutex:
		return
	}

	locker.Lock()
	w.locks[w.depth] = locker
}
<commit_msg>Always pop the slice index off the stack<commit_after>package copystructure

import (
	"errors"
	"reflect"
	"sync"

	"github.com/mitchellh/reflectwalk"
)

// Copy returns a deep copy of v.
func Copy(v interface{}) (interface{}, error) {
	return Config{}.Copy(v)
}

// CopierFunc is a function that knows how to deep copy a specific type.
// Register these globally with the Copiers variable.
type CopierFunc func(interface{}) (interface{}, error)

// Copiers is a map of types that behave specially when they are copied.
// If a type is found in this map while deep copying, this function
// will be called to copy it instead of attempting to copy all fields.
//
// The key should be the type, obtained using: reflect.TypeOf(value with type).
//
// It is unsafe to write to this map after Copies have started. If you
// are writing to this map while also copying, wrap all modifications to
// this map as well as to Copy in a mutex.
var Copiers map[reflect.Type]CopierFunc = make(map[reflect.Type]CopierFunc)

// Must is a helper that wraps a call to a function returning
// (interface{}, error) and panics if the error is non-nil. It is intended
// for use in variable initializations and should only be used when a copy
// error should be a crashing case.
func Must(v interface{}, err error) interface{} {
	if err != nil {
		panic("copy error: " + err.Error())
	}

	return v
}

var errPointerRequired = errors.New("Copy argument must be a pointer when Lock is true")

type Config struct {
	// Lock any types that are a sync.Locker and are not a mutex while copying.
	// If there is an RLocker method, use that to get the sync.Locker.
	Lock bool

	// Copiers is a map of types associated with a CopierFunc. Use the global
	// Copiers map if this is nil.
	Copiers map[reflect.Type]CopierFunc
}

func (c Config) Copy(v interface{}) (interface{}, error) {
	if c.Lock && reflect.ValueOf(v).Kind() != reflect.Ptr {
		return nil, errPointerRequired
	}

	w := new(walker)
	if c.Lock {
		w.useLocks = true
	}

	if c.Copiers == nil {
		c.Copiers = Copiers
	}

	err := reflectwalk.Walk(v, w)
	if err != nil {
		return nil, err
	}

	// Get the result. If the result is nil, then we want to turn it
	// into a typed nil if we can.
	result := w.Result
	if result == nil {
		val := reflect.ValueOf(v)
		result = reflect.Indirect(reflect.New(val.Type())).Interface()
	}

	return result, nil
}

// Return the key used to index interfaces types we've seen. Store the number
// of pointers in the upper 32bits, and the depth in the lower 32bits. This is
// easy to calculate, easy to match a key with our current depth, and we don't
// need to deal with initializing and cleaning up nested maps or slices.
func ifaceKey(pointers, depth int) uint64 {
	return uint64(pointers)<<32 | uint64(depth)
}

type walker struct {
	Result interface{}

	depth       int
	ignoreDepth int
	vals        []reflect.Value
	cs          []reflect.Value

	// This stores the number of pointers we've walked over, indexed by depth.
	ps []int

	// If an interface is indirected by a pointer, we need to know the type of
	// interface to create when creating the new value.  Store the interface
	// types here, indexed by both the walk depth and the number of pointers
	// already seen at that depth. Use ifaceKey to calculate the proper uint64
	// value.
	ifaceTypes map[uint64]reflect.Type

	// any locks we've taken, indexed by depth
	locks []sync.Locker
	// take locks while walking the structure
	useLocks bool
}

func (w *walker) Enter(l reflectwalk.Location) error {
	w.depth++

	// ensure we have enough elements to index via w.depth
	for w.depth >= len(w.locks) {
		w.locks = append(w.locks, nil)
	}

	for len(w.ps) < w.depth+1 {
		w.ps = append(w.ps, 0)
	}

	return nil
}

func (w *walker) Exit(l reflectwalk.Location) error {
	locker := w.locks[w.depth]
	w.locks[w.depth] = nil
	if locker != nil {
		defer locker.Unlock()
	}

	// clear out pointers and interfaces as we exit the stack
	w.ps[w.depth] = 0

	for k := range w.ifaceTypes {
		mask := uint64(^uint32(0))
		if k&mask == uint64(w.depth) {
			delete(w.ifaceTypes, k)
		}
	}

	w.depth--
	if w.ignoreDepth > w.depth {
		w.ignoreDepth = 0
	}

	if w.ignoring() {
		return nil
	}

	switch l {
	case reflectwalk.Map:
		fallthrough
	case reflectwalk.Slice:
		// Pop map off our container
		w.cs = w.cs[:len(w.cs)-1]
	case reflectwalk.MapValue:
		// Pop off the key and value
		mv := w.valPop()
		mk := w.valPop()
		m := w.cs[len(w.cs)-1]

		// If mv is the zero value, SetMapIndex deletes the key form the map,
		// or in this case never adds it. We need to create a properly typed
		// zero value so that this key can be set.
		if !mv.IsValid() {
			mv = reflect.Zero(m.Type().Elem())
		}
		m.SetMapIndex(mk, mv)
	case reflectwalk.SliceElem:
		// Pop off the value and the index and set it on the slice
		v := w.valPop()
		i := w.valPop().Interface().(int)
		if v.IsValid() {
			s := w.cs[len(w.cs)-1]
			se := s.Index(i)
			if se.CanSet() {
				se.Set(v)
			}
		}
	case reflectwalk.Struct:
		w.replacePointerMaybe()

		// Remove the struct from the container stack
		w.cs = w.cs[:len(w.cs)-1]
	case reflectwalk.StructField:
		// Pop off the value and the field
		v := w.valPop()
		f := w.valPop().Interface().(reflect.StructField)
		if v.IsValid() {
			s := w.cs[len(w.cs)-1]
			sf := reflect.Indirect(s).FieldByName(f.Name)

			if sf.CanSet() {
				sf.Set(v)
			}
		}
	case reflectwalk.WalkLoc:
		// Clear out the slices for GC
		w.cs = nil
		w.vals = nil
	}

	return nil
}

func (w *walker) Map(m reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(m)

	// Create the map. If the map itself is nil, then just make a nil map
	var newMap reflect.Value
	if m.IsNil() {
		newMap = reflect.Indirect(reflect.New(m.Type()))
	} else {
		newMap = reflect.MakeMap(m.Type())
	}

	w.cs = append(w.cs, newMap)
	w.valPush(newMap)
	return nil
}

func (w *walker) MapElem(m, k, v reflect.Value) error {
	return nil
}

func (w *walker) PointerEnter(v bool) error {
	if v {
		w.ps[w.depth]++
	}
	return nil
}

func (w *walker) PointerExit(v bool) error {
	if v {
		w.ps[w.depth]--
	}
	return nil
}

func (w *walker) Interface(v reflect.Value) error {
	if !v.IsValid() {
		return nil
	}
	if w.ifaceTypes == nil {
		w.ifaceTypes = make(map[uint64]reflect.Type)
	}

	w.ifaceTypes[ifaceKey(w.ps[w.depth], w.depth)] = v.Type()
	return nil
}

func (w *walker) Primitive(v reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(v)

	// IsValid verifies the v is non-zero and CanInterface verifies
	// that we're allowed to read this value (unexported fields).
	var newV reflect.Value
	if v.IsValid() && v.CanInterface() {
		newV = reflect.New(v.Type())
		newV.Elem().Set(v)
	}

	w.valPush(newV)
	w.replacePointerMaybe()
	return nil
}

func (w *walker) Slice(s reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(s)

	var newS reflect.Value
	if s.IsNil() {
		newS = reflect.Indirect(reflect.New(s.Type()))
	} else {
		newS = reflect.MakeSlice(s.Type(), s.Len(), s.Cap())
	}

	w.cs = append(w.cs, newS)
	w.valPush(newS)
	return nil
}

func (w *walker) SliceElem(i int, elem reflect.Value) error {
	if w.ignoring() {
		return nil
	}

	// We don't write the slice here because elem might still be
	// arbitrarily complex. Just record the index and continue on.
	w.valPush(reflect.ValueOf(i))

	return nil
}

func (w *walker) Struct(s reflect.Value) error {
	if w.ignoring() {
		return nil
	}
	w.lock(s)

	var v reflect.Value
	if c, ok := Copiers[s.Type()]; ok {
		// We have a Copier for this struct, so we use that copier to
		// get the copy, and we ignore anything deeper than this.
		w.ignoreDepth = w.depth

		dup, err := c(s.Interface())
		if err != nil {
			return err
		}

		v = reflect.ValueOf(dup)
	} else {
		// No copier, we copy ourselves and allow reflectwalk to guide
		// us deeper into the structure for copying.
		v = reflect.New(s.Type())
	}

	// Push the value onto the value stack for setting the struct field,
	// and add the struct itself to the containers stack in case we walk
	// deeper so that its own fields can be modified.
	w.valPush(v)
	w.cs = append(w.cs, v)

	return nil
}

func (w *walker) StructField(f reflect.StructField, v reflect.Value) error {
	if w.ignoring() {
		return nil
	}

	// If PkgPath is non-empty, this is a private (unexported) field.
	// We do not set this unexported since the Go runtime doesn't allow us.
	if f.PkgPath != "" {
		return reflectwalk.SkipEntry
	}

	// Push the field onto the stack, we'll handle it when we exit
	// the struct field in Exit...
	w.valPush(reflect.ValueOf(f))
	return nil
}

// ignore causes the walker to ignore any more values until we exit this on
func (w *walker) ignore() {
	w.ignoreDepth = w.depth
}

func (w *walker) ignoring() bool {
	return w.ignoreDepth > 0 && w.depth >= w.ignoreDepth
}

func (w *walker) pointerPeek() bool {
	return w.ps[w.depth] > 0
}

func (w *walker) valPop() reflect.Value {
	result := w.vals[len(w.vals)-1]
	w.vals = w.vals[:len(w.vals)-1]

	// If we're out of values, that means we popped everything off. In
	// this case, we reset the result so the next pushed value becomes
	// the result.
	if len(w.vals) == 0 {
		w.Result = nil
	}

	return result
}

func (w *walker) valPush(v reflect.Value) {
	w.vals = append(w.vals, v)

	// If we haven't set the result yet, then this is the result since
	// it is the first (outermost) value we're seeing.
	if w.Result == nil && v.IsValid() {
		w.Result = v.Interface()
	}
}

func (w *walker) replacePointerMaybe() {
	// Determine the last pointer value. If it is NOT a pointer, then
	// we need to push that onto the stack.
	if !w.pointerPeek() {
		w.valPush(reflect.Indirect(w.valPop()))
		return
	}

	v := w.valPop()
	for i := 1; i < w.ps[w.depth]; i++ {
		if iType, ok := w.ifaceTypes[ifaceKey(w.ps[w.depth]-i, w.depth)]; ok {
			iface := reflect.New(iType).Elem()
			iface.Set(v)
			v = iface
		}

		p := reflect.New(v.Type())
		p.Elem().Set(v)
		v = p
	}

	w.valPush(v)
}

// if this value is a Locker, lock it and add it to the locks slice
func (w *walker) lock(v reflect.Value) {
	if !w.useLocks {
		return
	}

	if !v.IsValid() || !v.CanInterface() {
		return
	}

	type rlocker interface {
		RLocker() sync.Locker
	}

	var locker sync.Locker

	// We can't call Interface() on a value directly, since that requires
	// a copy. This is OK, since the pointer to a value which is a sync.Locker
	// is also a sync.Locker.
	if v.Kind() == reflect.Ptr {
		switch l := v.Interface().(type) {
		case rlocker:
			// don't lock a mutex directly
			if _, ok := l.(*sync.RWMutex); !ok {
				locker = l.RLocker()
			}
		case sync.Locker:
			locker = l
		}
	} else if v.CanAddr() {
		switch l := v.Addr().Interface().(type) {
		case rlocker:
			// don't lock a mutex directly
			if _, ok := l.(*sync.RWMutex); !ok {
				locker = l.RLocker()
			}
		case sync.Locker:
			locker = l
		}
	}

	// still no callable locker
	if locker == nil {
		return
	}

	// don't lock a mutex directly
	switch locker.(type) {
	case *sync.Mutex, *sync.RWMutex:
		return
	}

	locker.Lock()
	w.locks[w.depth] = locker
}
<|endoftext|>