<commit_before>package gremgo

import (
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

type dialer interface {
	connect() error
	write([]byte) error
	read() ([]byte, error)
}

/////
/*
WebSocket Connection
*/
/////

// Ws is the dialer for a WebSocket connection
type Ws struct {
	host string
	conn *websocket.Conn
}

func (ws *Ws) connect() (err error) {
	d := websocket.Dialer{}
	ws.conn, _, err = d.Dial(ws.host, http.Header{})
	return
}

func (ws *Ws) write(msg []byte) (err error) {
	err = ws.conn.WriteMessage(2, msg)
	return
}

func (ws *Ws) read() (msg []byte, err error) {
	_, msg, err = ws.conn.ReadMessage()
	return
}

/////

func (c *Client) writeWorker() { // writeWorker works on a loop and dispatches messages as soon as it recieves them
	for {
		select {
		case msg := <-c.requests: // Wait for message send request
			err := c.conn.write(msg) // Write message
			if err != nil {          // TODO: Fix error handling here
				log.Fatal(err)
			}
		default:
		}
	}
}

func (c *Client) readWorker() { // readWorker works on a loop and sorts messages as soon as it recieves them
	for {
		msg, err := c.conn.read()
		if err != nil {
			log.Fatal(err)
		}
		if msg != nil {
			// TODO: Make this multithreaded
			c.handleResponse(msg) // Send message for sorting and retrieval on a separate thread
		}
	}
}
<commit_msg>Fixed read and write buffers<commit_after>package gremgo

import (
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

type dialer interface {
	connect() error
	write([]byte) error
	read() ([]byte, error)
}

/////
/*
WebSocket Connection
*/
/////

// Ws is the dialer for a WebSocket connection
type Ws struct {
	host string
	conn *websocket.Conn
}

func (ws *Ws) connect() (err error) {
	d := websocket.Dialer{
		WriteBufferSize: 8192,
		ReadBufferSize:  8192,
	}
	ws.conn, _, err = d.Dial(ws.host, http.Header{})
	return
}

func (ws *Ws) write(msg []byte) (err error) {
	err = ws.conn.WriteMessage(2, msg)
	return
}

func (ws *Ws) read() (msg []byte, err error) {
	_, msg, err = ws.conn.ReadMessage()
	return
}

/////

func (c *Client) writeWorker() { // writeWorker works on a loop and dispatches messages as soon as it recieves them
	for {
		select {
		case msg := <-c.requests: // Wait for message send request
			err := c.conn.write(msg) // Write message
			if err != nil {          // TODO: Fix error handling here
				log.Fatal(err)
			}
		default:
		}
	}
}

func (c *Client) readWorker() { // readWorker works on a loop and sorts messages as soon as it recieves them
	for {
		msg, err := c.conn.read()
		if err != nil {
			log.Fatal(err)
		}
		if msg != nil {
			// TODO: Make this multithreaded
			c.handleResponse(msg) // Send message for sorting and retrieval on a separate thread
		}
	}
}
<|endoftext|>


<commit_before>package pop

import (
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/markbates/going/defaults"
	"github.com/markbates/going/randx"
)

// Connections contains all of the available connections
var Connections = map[string]*Connection{}

// Connection represents all of the necessary details for
// talking with a datastore
type Connection struct {
	ID      string
	Store   Store
	Dialect Dialect
	Timings []time.Duration
}

func (c *Connection) String() string {
	return c.Dialect.URL()
}

func (c *Connection) URL() string {
	return c.String()
}

// NewConnection creates a new connection, and sets it's `Dialect`
// appropriately based on the `ConnectionDetails` passed into it.
func NewConnection(deets *ConnectionDetails) *Connection {
	c := &Connection{
		ID:      randx.String(30),
		Timings: []time.Duration{},
	}
	switch deets.Dialect {
	case "postgres":
		c.Dialect = NewPostgreSQL(deets)
	case "mysql":
		c.Dialect = NewMySQL(deets)
		// case "sqlite3":
		// 	c.Dialect = NewSQLite(deets)
	}
	return c
}

// Connect takes the name of a connection, default is "development", and will
// return that connection from the available `Connections`. If a connection with
// that name can not be found an error will be returned. If a connection is
// found, and it has yet to open a connection with its underlying datastore,
// a connection to that store will be opened.
func Connect(e string) (*Connection, error) {
	e = defaults.String(e, "development")
	c := Connections[e]
	if c == nil {
		return c, fmt.Errorf("Could not find connection named %s!", e)
	}
	if c.Store != nil {
		return c, nil
	}
	db, err := sqlx.Open(c.Dialect.Details().Dialect, c.Dialect.URL())
	if err == nil {
		c.Store = &dB{db}
	}
	return c, nil
}

// Transaction will start a new transaction on the connection. If the inner function
// returns an error then the transaction will be rolled back, otherwise the transaction
// will automatically commit at the end.
func (c *Connection) Transaction(fn func(tx *Connection) error) error {
	tx, err := c.Store.Transaction()
	if err != nil {
		return err
	}
	cn := &Connection{
		ID:      randx.String(30),
		Store:   tx,
		Dialect: c.Dialect,
		Timings: []time.Duration{},
	}
	err = fn(cn)
	if err != nil {
		err = tx.Rollback()
	} else {
		err = tx.Commit()
	}
	return err
}

// Rollback will open a new transaction and automatically rollback that transaction
// when the inner function returns, regardless. This can be useful for tests, etc...
func (c *Connection) Rollback(fn func(tx *Connection)) error {
	tx, err := c.Store.Transaction()
	if err != nil {
		return err
	}
	cn := &Connection{
		ID:      randx.String(30),
		Store:   tx,
		Dialect: c.Dialect,
		Timings: []time.Duration{},
	}
	fn(cn)
	return tx.Rollback()
}

// Q creates a new "empty" query for the current connection.
func (c *Connection) Q() *Query {
	return Q(c)
}

func (c *Connection) timeFunc(name string, fn func() error) error {
	now := time.Now()
	err := fn()
	c.Timings = append(c.Timings, time.Now().Sub(now))
	return err
}
<commit_msg>if the connection is already a transaction, just use that<commit_after>package pop

import (
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/markbates/going/defaults"
	"github.com/markbates/going/randx"
)

// Connections contains all of the available connections
var Connections = map[string]*Connection{}

// Connection represents all of the necessary details for
// talking with a datastore
type Connection struct {
	ID      string
	Store   Store
	Dialect Dialect
	Timings []time.Duration
	tX      *tX
}

func (c *Connection) String() string {
	return c.Dialect.URL()
}

func (c *Connection) URL() string {
	return c.String()
}

// NewConnection creates a new connection, and sets it's `Dialect`
// appropriately based on the `ConnectionDetails` passed into it.
func NewConnection(deets *ConnectionDetails) *Connection {
	c := &Connection{
		ID:      randx.String(30),
		Timings: []time.Duration{},
	}
	switch deets.Dialect {
	case "postgres":
		c.Dialect = NewPostgreSQL(deets)
	case "mysql":
		c.Dialect = NewMySQL(deets)
		// case "sqlite3":
		// 	c.Dialect = NewSQLite(deets)
	}
	return c
}

// Connect takes the name of a connection, default is "development", and will
// return that connection from the available `Connections`. If a connection with
// that name can not be found an error will be returned. If a connection is
// found, and it has yet to open a connection with its underlying datastore,
// a connection to that store will be opened.
func Connect(e string) (*Connection, error) {
	e = defaults.String(e, "development")
	c := Connections[e]
	if c == nil {
		return c, fmt.Errorf("Could not find connection named %s!", e)
	}
	if c.Store != nil {
		return c, nil
	}
	db, err := sqlx.Open(c.Dialect.Details().Dialect, c.Dialect.URL())
	if err == nil {
		c.Store = &dB{db}
	}
	return c, nil
}

// Transaction will start a new transaction on the connection. If the inner function
// returns an error then the transaction will be rolled back, otherwise the transaction
// will automatically commit at the end.
func (c *Connection) Transaction(fn func(tx *Connection) error) error {
	var cn *Connection
	if c.tX == nil {
		tx, err := c.Store.Transaction()
		if err != nil {
			return err
		}
		cn = &Connection{
			ID:      randx.String(30),
			Store:   tx,
			Dialect: c.Dialect,
			Timings: []time.Duration{},
			tX:      tx,
		}
	} else {
		cn = c
	}
	err := fn(cn)
	if err != nil {
		err = cn.tX.Rollback()
	} else {
		err = cn.tX.Commit()
	}
	return err
}

// Rollback will open a new transaction and automatically rollback that transaction
// when the inner function returns, regardless. This can be useful for tests, etc...
func (c *Connection) Rollback(fn func(tx *Connection)) error {
	var cn *Connection
	if c.tX == nil {
		tx, err := c.Store.Transaction()
		if err != nil {
			return err
		}
		cn = &Connection{
			ID:      randx.String(30),
			Store:   tx,
			Dialect: c.Dialect,
			Timings: []time.Duration{},
			tX:      tx,
		}
	} else {
		cn = c
	}
	fn(cn)
	return cn.tX.Rollback()
}

// Q creates a new "empty" query for the current connection.
func (c *Connection) Q() *Query {
	return Q(c)
}

func (c *Connection) timeFunc(name string, fn func() error) error {
	now := time.Now()
	err := fn()
	c.Timings = append(c.Timings, time.Now().Sub(now))
	return err
}
<|endoftext|>


<commit_before>package tarantool

import (
	"bufio"
	"errors"
	"io"
	"net"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var (
	ErrEmptyDefaultSpace = errors.New("zero-length default space or unnecessary slash in dsn.path")
	ErrSyncFailed        = errors.New("SYNC failed")
)

type Options struct {
	ConnectTimeout time.Duration
	QueryTimeout   time.Duration
	DefaultSpace   string
	User           string
	Password       string
	UUID           string
	ReplicaSetUUID string
	Perf           PerfCount
}

type Greeting struct {
	Version []byte
	Auth    []byte
}

type Connection struct {
	requestID uint64
	requests  *requestMap
	writeChan chan *BinaryPacket // packed messages with header
	closeOnce sync.Once
	exit      chan bool
	closed    chan bool
	tcpConn   net.Conn

	ccr io.Reader
	ccw io.Writer

	// options
	queryTimeout   time.Duration
	greeting       *Greeting
	packData       *packData
	remoteAddr     string
	firstError     error
	firstErrorLock *sync.Mutex
	perf           PerfCount
}

// Connect to tarantool instance with options.
// Returned Connection could be used to execute queries.
func Connect(dsnString string, options *Options) (conn *Connection, err error) {
	dsn, opts, err := parseOptions(dsnString, options)
	if err != nil {
		return nil, err
	}

	conn, err = newConn(dsn.Scheme, dsn.Host, opts)
	if err != nil {
		return
	}

	// set schema pulling deadline
	deadline := time.Now().Add(opts.ConnectTimeout)
	conn.tcpConn.SetDeadline(deadline)

	err = conn.pullSchema()
	if err != nil {
		conn.tcpConn.Close()
		conn = nil
		return
	}

	// remove deadline
	conn.tcpConn.SetDeadline(time.Time{})

	go conn.worker()

	return
}

func newConn(scheme, addr string, opts Options) (conn *Connection, err error) {

	defer func() { // close opened connection if error
		if err != nil && conn != nil {
			if conn.tcpConn != nil {
				conn.tcpConn.Close()
			}
			conn = nil
		}
	}()

	conn = &Connection{
		remoteAddr:     addr,
		requests:       newRequestMap(),
		writeChan:      make(chan *BinaryPacket, 256),
		exit:           make(chan bool),
		closed:         make(chan bool),
		firstErrorLock: &sync.Mutex{},
		packData:       newPackData(opts.DefaultSpace),
		queryTimeout:   opts.QueryTimeout,
		perf:           opts.Perf,
	}

	conn.tcpConn, err = net.DialTimeout(scheme, conn.remoteAddr, opts.ConnectTimeout)
	if err != nil {
		return nil, err
	}

	if conn.perf.NetRead != nil {
		conn.ccr = NewCountedReader(conn.tcpConn, conn.perf.NetRead)
	} else {
		conn.ccr = conn.tcpConn
	}

	if conn.perf.NetWrite != nil {
		conn.ccw = NewCountedWriter(conn.tcpConn, conn.perf.NetWrite)
	} else {
		conn.ccw = conn.tcpConn
	}

	greeting := make([]byte, 128)

	connectDeadline := time.Now().Add(opts.ConnectTimeout)
	conn.tcpConn.SetDeadline(connectDeadline)
	// removing deadline deferred
	defer conn.tcpConn.SetDeadline(time.Time{})

	_, err = io.ReadFull(conn.ccr, greeting)
	if err != nil {
		return
	}

	conn.greeting = &Greeting{
		Version: greeting[:64],
		Auth:    greeting[64:108],
	}

	// try to authenticate if user have been provided
	if len(opts.User) > 0 {
		requestID := conn.nextID()

		pp := packetPool.GetWithID(requestID)

		err = pp.packMsg(&Auth{
			User:         opts.User,
			Password:     opts.Password,
			GreetingAuth: conn.greeting.Auth,
		}, conn.packData)
		if err != nil {
			pp.Release()
			return
		}

		_, err = pp.WriteTo(conn.ccw)
		pp.Release()
		if err != nil {
			return
		}

		pp = packetPool.Get()
		defer pp.Release()

		if err = pp.readPacket(conn.ccr); err != nil {
			return
		}

		authResponse := &pp.packet
		if authResponse.requestID != requestID {
			err = ErrSyncFailed
			return
		}

		if authResponse.Result != nil && authResponse.Result.Error != nil {
			err = authResponse.Result.Error
			return
		}
	}

	return
}

func parseOptions(dsnString string, options *Options) (*url.URL, Options, error) {
	if options == nil {
		options = &Options{}
	}
	opts := *options // copy to new object

	// remove schema, if present

	// === for backward compatibility (only tcp despite of user wishes :)
	dsnString = strings.TrimPrefix(dsnString, "unix:")
	// ===

	// tcp is the default scheme
	switch {
	case strings.HasPrefix(dsnString, "tcp://"):
	case strings.HasPrefix(dsnString, "//"):
		dsnString = "tcp:" + dsnString
	default:
		dsnString = "tcp://" + dsnString
	}
	dsn, err := url.Parse(dsnString)
	if err != nil {
		return dsn, opts, err
	}

	if opts.ConnectTimeout.Nanoseconds() == 0 {
		opts.ConnectTimeout = DefaultConnectTimeout
	}
	if opts.QueryTimeout.Nanoseconds() == 0 {
		opts.QueryTimeout = DefaultQueryTimeout
	}

	if len(opts.User) == 0 {
		if user := dsn.User; user != nil {
			opts.User = user.Username()
			opts.Password, _ = user.Password()
		}
	}

	if len(opts.DefaultSpace) == 0 && len(dsn.Path) > 0 {
		path := strings.TrimPrefix(dsn.Path, "/")
		// check it if it is necessary
		switch {
		case len(path) == 0:
			return nil, opts, ErrEmptyDefaultSpace
		//case strings.IndexAny(path, "/ ,") != -1:
		//	return nil, opts, ErrBadDSNPath
		default:
			opts.DefaultSpace = path
		}
	}

	return dsn, opts, nil
}

func (conn *Connection) pullSchema() (err error) {
	// select space and index schema
	request := func(q Query) (*Result, error) {
		var err error

		requestID := conn.nextID()

		pp := packetPool.GetWithID(requestID)
		if err = pp.packMsg(q, conn.packData); err != nil {
			pp.Release()
			return nil, err
		}

		_, err = pp.WriteTo(conn.ccw)
		pp.Release()
		if err != nil {
			return nil, err
		}

		pp = packetPool.Get()
		defer pp.Release()

		if err = pp.readPacket(conn.ccr); err != nil {
			return nil, err
		}

		response := &pp.packet
		if response.requestID != requestID {
			return nil, errors.New("Bad response requestID")
		}

		if response.Result == nil {
			return nil, errors.New("Nil response result")
		}

		if response.Result.Error != nil {
			return nil, response.Result.Error
		}

		return response.Result, nil
	}

	res, err := request(&Select{
		Space:    ViewSpace,
		Key:      0,
		Iterator: IterAll,
	})
	if err != nil {
		return
	}

	for _, space := range res.Data {
		spaceID, _ := conn.packData.spaceNo(space[0])
		conn.packData.spaceMap[space[2].(string)] = spaceID
	}

	res, err = request(&Select{
		Space:    ViewIndex,
		Key:      0,
		Iterator: IterAll,
	})
	if err != nil {
		return
	}

	for _, index := range res.Data {
		spaceID, _ := conn.packData.fieldNo(index[0])
		indexID, _ := conn.packData.fieldNo(index[1])
		indexName := index[2].(string)
		indexAttr := index[4].(map[string]interface{}) // e.g: {"unique": true}
		indexFields := index[5].([]interface{})        // e.g: [[0 num] [1 str]]

		indexSpaceMap, exists := conn.packData.indexMap[spaceID]
		if !exists {
			indexSpaceMap = make(map[string]uint64)
			conn.packData.indexMap[spaceID] = indexSpaceMap
		}
		indexSpaceMap[indexName] = indexID

		// build list of primary key field numbers for this space, if the PK is detected
		if indexAttr != nil && indexID == 0 {
			if unique, ok := indexAttr["unique"]; ok && unique.(bool) {
				pk := make([]int, len(indexFields))
				for i := range indexFields {
					descr := indexFields[i].([]interface{})
					f, _ := conn.packData.fieldNo(descr[0])
					pk[i] = int(f)
				}
				conn.packData.primaryKeyMap[spaceID] = pk
			}
		}
	}

	return
}

func (conn *Connection) nextID() uint64 {
	return atomic.AddUint64(&conn.requestID, 1)
}

func (conn *Connection) stop() {
	conn.closeOnce.Do(func() {
		// debug.PrintStack()
		close(conn.exit)
		conn.tcpConn.Close()
		runtime.GC()
	})
}

func (conn *Connection) GetPrimaryKeyFields(space interface{}) ([]int, bool) {
	var spaceID uint64
	var err error

	if conn.packData == nil {
		return nil, false
	}
	if spaceID, err = conn.packData.spaceNo(space); err != nil {
		return nil, false
	}

	f, ok := conn.packData.primaryKeyMap[spaceID]
	return f, ok
}

func (conn *Connection) Close() {
	conn.stop()
	<-conn.closed
}

func (conn *Connection) String() string {
	return conn.remoteAddr
}

func (conn *Connection) IsClosed() bool {
	select {
	case <-conn.exit:
		return true
	default:
		return false
	}
}

func (conn *Connection) getError() error {
	conn.firstErrorLock.Lock()
	defer conn.firstErrorLock.Unlock()
	return conn.firstError
}

func (conn *Connection) setError(err error) {
	if err != nil && err != io.EOF {
		conn.firstErrorLock.Lock()
		if conn.firstError == nil {
			conn.firstError = err
		}
		conn.firstErrorLock.Unlock()
	}
}

func (conn *Connection) worker() {
	var wg sync.WaitGroup

	wg.Add(2)

	go func() {
		err := conn.writer()
		conn.setError(err)
		conn.stop()
		wg.Done()
	}()

	go func() {
		err := conn.reader()
		conn.setError(err)
		conn.stop()
		wg.Done()
	}()

	wg.Wait()

	// release all pending packets
	writeChan := conn.writeChan
	conn.writeChan = nil

CLEANUP_LOOP:
	for {
		select {
		case pp := <-writeChan:
			pp.Release()
		default:
			break CLEANUP_LOOP
		}
	}

	// send error reply to all pending requests
	conn.requests.CleanUp(func(req *request) {
		select {
		case req.replyChan <- &AsyncResult{
			Error:     ConnectionClosedError(conn),
			ErrorCode: ErrNoConnection,
			Opaque:    req.opaque,
		}:
		default:
		}
	})

	close(conn.closed)
}

func (conn *Connection) writer() (err error) {
	writeChan := conn.writeChan
	stopChan := conn.exit
	w := bufio.NewWriterSize(conn.ccw, DefaultWriterBufSize)

	wp := func(w io.Writer, packet *BinaryPacket) error {
		if conn.perf.NetPacketsOut != nil {
			conn.perf.NetPacketsOut.Add(1)
		}
		_, err := packet.WriteTo(w)
		packet.Release()
		return err
	}

WRITER_LOOP:
	for {
		select {
		case packet, ok := <-writeChan:
			if !ok {
				break WRITER_LOOP
			}
			if err = wp(w, packet); err != nil {
				break WRITER_LOOP
			}
		case <-stopChan:
			break WRITER_LOOP
		default:
			if err = w.Flush(); err != nil {
				break WRITER_LOOP
			}

			// same without flush
			select {
			case packet, ok := <-writeChan:
				if !ok {
					break WRITER_LOOP
				}
				if err = wp(w, packet); err != nil {
					break WRITER_LOOP
				}
			case <-stopChan:
				break WRITER_LOOP
			}
		}
	}

	return
}

func (conn *Connection) reader() (err error) {
	var pp *BinaryPacket
	var requestID uint64

	r := bufio.NewReaderSize(conn.ccr, DefaultReaderBufSize)

READER_LOOP:
	for {
		pp := packetPool.Get()
		if requestID, err = pp.readRawPacket(r); err != nil {
			break READER_LOOP
		}

		if conn.perf.NetPacketsIn != nil {
			conn.perf.NetPacketsIn.Add(1)
		}

		req := conn.requests.Pop(requestID)
		if req == nil {
			pp.Release()
			pp = nil
			continue
		}

		select {
		case req.replyChan <- &AsyncResult{0, nil, pp, conn, req.opaque}:
			pp = nil
		default:
		}
	}

	if pp != nil {
		pp.Release()
	}
	return
}
<commit_msg>Add GetPerf method to Connection<commit_after>package tarantool

import (
	"bufio"
	"errors"
	"io"
	"net"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var (
	ErrEmptyDefaultSpace = errors.New("zero-length default space or unnecessary slash in dsn.path")
	ErrSyncFailed        = errors.New("SYNC failed")
)

type Options struct {
	ConnectTimeout time.Duration
	QueryTimeout   time.Duration
	DefaultSpace   string
	User           string
	Password       string
	UUID           string
	ReplicaSetUUID string
	Perf           PerfCount
}

type Greeting struct {
	Version []byte
	Auth    []byte
}

type Connection struct {
	requestID uint64
	requests  *requestMap
	writeChan chan *BinaryPacket // packed messages with header
	closeOnce sync.Once
	exit      chan bool
	closed    chan bool
	tcpConn   net.Conn

	ccr io.Reader
	ccw io.Writer

	// options
	queryTimeout   time.Duration
	greeting       *Greeting
	packData       *packData
	remoteAddr     string
	firstError     error
	firstErrorLock *sync.Mutex
	perf           PerfCount
}

// Connect to tarantool instance with options.
// Returned Connection could be used to execute queries.
func Connect(dsnString string, options *Options) (conn *Connection, err error) {
	dsn, opts, err := parseOptions(dsnString, options)
	if err != nil {
		return nil, err
	}

	conn, err = newConn(dsn.Scheme, dsn.Host, opts)
	if err != nil {
		return
	}

	// set schema pulling deadline
	deadline := time.Now().Add(opts.ConnectTimeout)
	conn.tcpConn.SetDeadline(deadline)

	err = conn.pullSchema()
	if err != nil {
		conn.tcpConn.Close()
		conn = nil
		return
	}

	// remove deadline
	conn.tcpConn.SetDeadline(time.Time{})

	go conn.worker()

	return
}

func newConn(scheme, addr string, opts Options) (conn *Connection, err error) {

	defer func() { // close opened connection if error
		if err != nil && conn != nil {
			if conn.tcpConn != nil {
				conn.tcpConn.Close()
			}
			conn = nil
		}
	}()

	conn = &Connection{
		remoteAddr:     addr,
		requests:       newRequestMap(),
		writeChan:      make(chan *BinaryPacket, 256),
		exit:           make(chan bool),
		closed:         make(chan bool),
		firstErrorLock: &sync.Mutex{},
		packData:       newPackData(opts.DefaultSpace),
		queryTimeout:   opts.QueryTimeout,
		perf:           opts.Perf,
	}

	conn.tcpConn, err = net.DialTimeout(scheme, conn.remoteAddr, opts.ConnectTimeout)
	if err != nil {
		return nil, err
	}

	if conn.perf.NetRead != nil {
		conn.ccr = NewCountedReader(conn.tcpConn, conn.perf.NetRead)
	} else {
		conn.ccr = conn.tcpConn
	}

	if conn.perf.NetWrite != nil {
		conn.ccw = NewCountedWriter(conn.tcpConn, conn.perf.NetWrite)
	} else {
		conn.ccw = conn.tcpConn
	}

	greeting := make([]byte, 128)

	connectDeadline := time.Now().Add(opts.ConnectTimeout)
	conn.tcpConn.SetDeadline(connectDeadline)
	// removing deadline deferred
	defer conn.tcpConn.SetDeadline(time.Time{})

	_, err = io.ReadFull(conn.ccr, greeting)
	if err != nil {
		return
	}

	conn.greeting = &Greeting{
		Version: greeting[:64],
		Auth:    greeting[64:108],
	}

	// try to authenticate if user have been provided
	if len(opts.User) > 0 {
		requestID := conn.nextID()

		pp := packetPool.GetWithID(requestID)

		err = pp.packMsg(&Auth{
			User:         opts.User,
			Password:     opts.Password,
			GreetingAuth: conn.greeting.Auth,
		}, conn.packData)
		if err != nil {
			pp.Release()
			return
		}

		_, err = pp.WriteTo(conn.ccw)
		pp.Release()
		if err != nil {
			return
		}

		pp = packetPool.Get()
		defer pp.Release()

		if err = pp.readPacket(conn.ccr); err != nil {
			return
		}

		authResponse := &pp.packet
		if authResponse.requestID != requestID {
			err = ErrSyncFailed
			return
		}

		if authResponse.Result != nil && authResponse.Result.Error != nil {
			err = authResponse.Result.Error
			return
		}
	}

	return
}

func parseOptions(dsnString string, options *Options) (*url.URL, Options, error) {
	if options == nil {
		options = &Options{}
	}
	opts := *options // copy to new object

	// remove schema, if present

	// === for backward compatibility (only tcp despite of user wishes :)
	dsnString = strings.TrimPrefix(dsnString, "unix:")
	// ===

	// tcp is the default scheme
	switch {
	case strings.HasPrefix(dsnString, "tcp://"):
	case strings.HasPrefix(dsnString, "//"):
		dsnString = "tcp:" + dsnString
	default:
		dsnString = "tcp://" + dsnString
	}
	dsn, err := url.Parse(dsnString)
	if err != nil {
		return dsn, opts, err
	}

	if opts.ConnectTimeout.Nanoseconds() == 0 {
		opts.ConnectTimeout = DefaultConnectTimeout
	}
	if opts.QueryTimeout.Nanoseconds() == 0 {
		opts.QueryTimeout = DefaultQueryTimeout
	}

	if len(opts.User) == 0 {
		if user := dsn.User; user != nil {
			opts.User = user.Username()
			opts.Password, _ = user.Password()
		}
	}

	if len(opts.DefaultSpace) == 0 && len(dsn.Path) > 0 {
		path := strings.TrimPrefix(dsn.Path, "/")
		// check it if it is necessary
		switch {
		case len(path) == 0:
			return nil, opts, ErrEmptyDefaultSpace
		//case strings.IndexAny(path, "/ ,") != -1:
		//	return nil, opts, ErrBadDSNPath
		default:
			opts.DefaultSpace = path
		}
	}

	return dsn, opts, nil
}

func (conn *Connection) pullSchema() (err error) {
	// select space and index schema
	request := func(q Query) (*Result, error) {
		var err error

		requestID := conn.nextID()

		pp := packetPool.GetWithID(requestID)
		if err = pp.packMsg(q, conn.packData); err != nil {
			pp.Release()
			return nil, err
		}

		_, err = pp.WriteTo(conn.ccw)
		pp.Release()
		if err != nil {
			return nil, err
		}

		pp = packetPool.Get()
		defer pp.Release()

		if err = pp.readPacket(conn.ccr); err != nil {
			return nil, err
		}

		response := &pp.packet
		if response.requestID != requestID {
			return nil, errors.New("Bad response requestID")
		}

		if response.Result == nil {
			return nil, errors.New("Nil response result")
		}

		if response.Result.Error != nil {
			return nil, response.Result.Error
		}

		return response.Result, nil
	}

	res, err := request(&Select{
		Space:    ViewSpace,
		Key:      0,
		Iterator: IterAll,
	})
	if err != nil {
		return
	}

	for _, space := range res.Data {
		spaceID, _ := conn.packData.spaceNo(space[0])
		conn.packData.spaceMap[space[2].(string)] = spaceID
	}

	res, err = request(&Select{
		Space:    ViewIndex,
		Key:      0,
		Iterator: IterAll,
	})
	if err != nil {
		return
	}

	for _, index := range res.Data {
		spaceID, _ := conn.packData.fieldNo(index[0])
		indexID, _ := conn.packData.fieldNo(index[1])
		indexName := index[2].(string)
		indexAttr := index[4].(map[string]interface{}) // e.g: {"unique": true}
		indexFields := index[5].([]interface{})        // e.g: [[0 num] [1 str]]

		indexSpaceMap, exists := conn.packData.indexMap[spaceID]
		if !exists {
			indexSpaceMap = make(map[string]uint64)
			conn.packData.indexMap[spaceID] = indexSpaceMap
		}
		indexSpaceMap[indexName] = indexID

		// build list of primary key field numbers for this space, if the PK is detected
		if indexAttr != nil && indexID == 0 {
			if unique, ok := indexAttr["unique"]; ok && unique.(bool) {
				pk := make([]int, len(indexFields))
				for i := range indexFields {
					descr := indexFields[i].([]interface{})
					f, _ := conn.packData.fieldNo(descr[0])
					pk[i] = int(f)
				}
				conn.packData.primaryKeyMap[spaceID] = pk
			}
		}
	}

	return
}

func (conn *Connection) nextID() uint64 {
	return atomic.AddUint64(&conn.requestID, 1)
}

func (conn *Connection) stop() {
	conn.closeOnce.Do(func() {
		// debug.PrintStack()
		close(conn.exit)
		conn.tcpConn.Close()
		runtime.GC()
	})
}

func (conn *Connection) GetPerf() PerfCount {
	return conn.perf
}

func (conn *Connection) GetPrimaryKeyFields(space interface{}) ([]int, bool) {
	var spaceID uint64
	var err error

	if conn.packData == nil {
		return nil, false
	}
	if spaceID, err = conn.packData.spaceNo(space); err != nil {
		return nil, false
	}

	f, ok := conn.packData.primaryKeyMap[spaceID]
	return f, ok
}

func (conn *Connection) Close() {
	conn.stop()
	<-conn.closed
}

func (conn *Connection) String() string {
	return conn.remoteAddr
}

func (conn *Connection) IsClosed() bool {
	select {
	case <-conn.exit:
		return true
	default:
		return false
	}
}

func (conn *Connection) getError() error {
	conn.firstErrorLock.Lock()
	defer conn.firstErrorLock.Unlock()
	return conn.firstError
}

func (conn *Connection) setError(err error) {
	if err != nil && err != io.EOF {
		conn.firstErrorLock.Lock()
		if conn.firstError == nil {
			conn.firstError = err
		}
		conn.firstErrorLock.Unlock()
	}
}

func (conn *Connection) worker() {
	var wg sync.WaitGroup

	wg.Add(2)

	go func() {
		err := conn.writer()
		conn.setError(err)
		conn.stop()
		wg.Done()
	}()

	go func() {
		err := conn.reader()
		conn.setError(err)
		conn.stop()
		wg.Done()
	}()

	wg.Wait()

	// release all pending packets
	writeChan := conn.writeChan
	conn.writeChan = nil

CLEANUP_LOOP:
	for {
		select {
		case pp := <-writeChan:
			pp.Release()
		default:
			break CLEANUP_LOOP
		}
	}

	// send error reply to all pending requests
	conn.requests.CleanUp(func(req *request) {
		select {
		case req.replyChan <- &AsyncResult{
			Error:     ConnectionClosedError(conn),
			ErrorCode: ErrNoConnection,
			Opaque:    req.opaque,
		}:
		default:
		}
	})

	close(conn.closed)
}

func (conn *Connection) writer() (err error) {
	writeChan := conn.writeChan
	stopChan := conn.exit
	w := bufio.NewWriterSize(conn.ccw, DefaultWriterBufSize)

	wp := func(w io.Writer, packet *BinaryPacket) error {
		if conn.perf.NetPacketsOut != nil {
			conn.perf.NetPacketsOut.Add(1)
		}
		_, err := packet.WriteTo(w)
		packet.Release()
		return err
	}

WRITER_LOOP:
	for {
		select {
		case packet, ok := <-writeChan:
			if !ok {
				break WRITER_LOOP
			}
			if err = wp(w, packet); err != nil {
				break WRITER_LOOP
			}
		case <-stopChan:
			break WRITER_LOOP
		default:
			if err = w.Flush(); err != nil {
				break WRITER_LOOP
			}

			// same without flush
			select {
			case packet, ok := <-writeChan:
				if !ok {
					break WRITER_LOOP
				}
				if err = wp(w, packet); err != nil {
					break WRITER_LOOP
				}
			case <-stopChan:
				break WRITER_LOOP
			}
		}
	}

	return
}

func (conn *Connection) reader() (err error) {
	var pp *BinaryPacket
	var requestID uint64

	r := bufio.NewReaderSize(conn.ccr, DefaultReaderBufSize)

READER_LOOP:
	for {
		pp := packetPool.Get()
		if requestID, err = pp.readRawPacket(r); err != nil {
			break READER_LOOP
		}

		if conn.perf.NetPacketsIn != nil {
			conn.perf.NetPacketsIn.Add(1)
		}

		req := conn.requests.Pop(requestID)
		if req == nil {
			pp.Release()
			pp = nil
			continue
		}

		select {
		case req.replyChan <- &AsyncResult{0, nil, pp, conn, req.opaque}:
			pp = nil
		default:
		}
	}

	if pp != nil {
		pp.Release()
	}
	return
}
<|endoftext|>


<commit_before>package reminderplugin

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/iopred/bruxism"
)

// A Reminder holds data about a specific reminder.
type Reminder struct {
	StartTime time.Time
	Time      time.Time
	Requester string
	Target    string
	Message   string
	IsPrivate bool
}

// ReminderPlugin is a plugin that reminds users.
type ReminderPlugin struct {
	sync.RWMutex
	bot            *bruxism.Bot
	Reminders      []*Reminder
	TotalReminders int
}

var randomTimes = []string{
	"1 minute",
	"10 minutes",
	"1 hour",
	"4 hours",
	"tomorrow",
	"next week",
}

var randomMessages = []string{
	"walk the dog",
	"take pizza out of the oven",
	"check my email",
	"feed the baby",
	"play some quake",
}

func (p *ReminderPlugin) random(list []string) string {
	return list[rand.Intn(len(list))]
}

func (p *ReminderPlugin) randomReminder(service bruxism.Service) string {
	ticks := ""
	if service.Name() == bruxism.DiscordServiceName {
		ticks = "`"
	}

	return fmt.Sprintf("%s%sreminder %s %s%s", ticks, service.CommandPrefix(), p.random(randomTimes), p.random(randomMessages), ticks)
}

// Help returns a list of help strings that are printed when the user requests them.
func (p *ReminderPlugin) Help(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, detailed bool) []string {
	help := []string{
		bruxism.CommandHelp(service, "reminder", "<time> <reminder>", "Sets a reminder that is sent after the provided time.")[0],
	}
	if detailed {
		help = append(help, []string{
			"Examples: ",
			p.randomReminder(service),
			p.randomReminder(service),
		}...)
	}
	return help
}

func (p *ReminderPlugin) parseReminder(parts []string) (time.Time, string, error) {
	if parts[0] == "tomorrow" {
		return time.Now().Add(1 * time.Hour * 24), strings.Join(parts[1:], " "), nil
	}

	if parts[0] == "next" {
		switch parts[1] {
		case "week":
			return time.Now().Add(1 * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
		case "month":
			return time.Now().Add(1 * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
		case "year":
			return time.Now().Add(1 * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
		default:
			return time.Time{}, "", errors.New("Invalid next.")
		}
	}

	i, err := strconv.Atoi(parts[0])
	if err != nil {
		return time.Time{}, "", err
	}

	switch {
	case strings.HasPrefix(parts[1], "sec"):
		return time.Now().Add(time.Duration(i) * time.Second), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "min"):
		return time.Now().Add(time.Duration(i) * time.Minute), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "hour"):
		return time.Now().Add(time.Duration(i) * time.Hour), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "day"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "week"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "month"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "year"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
	}

	return time.Time{}, "", errors.New("Invalid string.")
}

// AddReminder adds a reminder.
func (p *ReminderPlugin) AddReminder(reminder *Reminder) error {
	p.Lock()
	defer p.Unlock()

	i := 0
	for _, r := range p.Reminders {
		if r.Requester == reminder.Requester {
			i++
			if i > 5 {
				return errors.New("You have too many reminders already.")
			}
		}
	}

	i = 0
	for _, r := range p.Reminders {
		if r.Time.After(reminder.Time) {
			break
		}
		i++
	}

	p.Reminders = append(p.Reminders, reminder)
	copy(p.Reminders[i+1:], p.Reminders[i:])
	p.Reminders[i] = reminder
	p.TotalReminders++

	return nil
}

func (p *ReminderPlugin) Message(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) {
	defer bruxism.MessageRecover()

	if service.IsMe(message) {
		return
	}

	if !bruxism.MatchesCommand(service, "remind", message) && !bruxism.MatchesCommand(service, "reminder", message) {
		return
	}

	_, parts := bruxism.ParseCommand(service, message)

	if len(parts) < 2 {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no time or message. eg: %s", p.randomReminder(service)))
		return
	}

	t, r, err := p.parseReminder(parts)

	now := time.Now()

	if err != nil || t.Before(now) || t.After(now.Add(time.Hour*24*365+time.Hour)) {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid time. eg: %s", strings.Join(randomTimes, ", ")))
		return
	}

	if r == "" {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no message. eg: %s", p.randomReminder(service)))
		return
	}

	requester := message.UserName()
	if service.Name() == bruxism.DiscordServiceName {
		requester = fmt.Sprintf("<@%s>", message.UserID())
	}

	err = p.AddReminder(&Reminder{
		StartTime: now,
		Time:      t,
		Requester: requester,
		Target:    message.Channel(),
		Message:   r,
		IsPrivate: service.IsPrivate(message),
	})
	if err != nil {
		service.SendMessage(message.Channel(), err.Error())
		return
	}

	service.SendMessage(message.Channel(), fmt.Sprintf("Reminder set for %s.", humanize.Time(t)))
}

// SendReminder sends a reminder.
func (p *ReminderPlugin) SendReminder(service bruxism.Service, reminder *Reminder) {
	if reminder.IsPrivate {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s you set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Message))
	} else {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s %s set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Requester, reminder.Message))
	}
}

// Run will block until a reminder needs to be fired and then fire it.
func (p *ReminderPlugin) Run(bot *bruxism.Bot, service bruxism.Service) {
	for {
		p.RLock()

		if len(p.Reminders) > 0 {
			reminder := p.Reminders[0]
			if time.Now().After(reminder.Time) {
				p.RUnlock()

				p.SendReminder(service, reminder)

				p.Lock()
				p.Reminders = p.Reminders[1:]
				p.Unlock()

				continue
			}
		}

		p.RUnlock()
		time.Sleep(500 * time.Millisecond)
	}
}

// Load will load plugin state from a byte array.
func (p *ReminderPlugin) Load(bot *bruxism.Bot, service bruxism.Service, data []byte) error {
	if data != nil {
		if err := json.Unmarshal(data, p); err != nil {
			log.Println("Error loading data", err)
		}
	}
	if len(p.Reminders) > p.TotalReminders {
		p.TotalReminders = len(p.Reminders)
	}

	go p.Run(bot, service)
	return nil
}

// Save will save plugin state to a byte array.
func (p *ReminderPlugin) Save() ([]byte, error) {
	return json.Marshal(p)
}

// Stats will return the stats for a plugin.
func (p *ReminderPlugin) Stats(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) []string {
	return []string{fmt.Sprintf("Reminders: \t%d\n", p.TotalReminders)}
}

// Name returns the name of the plugin.
func (p *ReminderPlugin) Name() string {
	return "Reminder"
}

// New will create a new Reminder plugin.
func New() bruxism.Plugin {
	return &ReminderPlugin{
		Reminders: []*Reminder{},
	}
}
<commit_msg>Prevent Sep from @everyone.<commit_after>package reminderplugin

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/iopred/bruxism"
)

// A Reminder holds data about a specific reminder.
type Reminder struct {
	StartTime time.Time
	Time      time.Time
	Requester string
	Target    string
	Message   string
	IsPrivate bool
}

// ReminderPlugin is a plugin that reminds users.
type ReminderPlugin struct {
	sync.RWMutex
	bot            *bruxism.Bot
	Reminders      []*Reminder
	TotalReminders int
}

var randomTimes = []string{
	"1 minute",
	"10 minutes",
	"1 hour",
	"4 hours",
	"tomorrow",
	"next week",
}

var randomMessages = []string{
	"walk the dog",
	"take pizza out of the oven",
	"check my email",
	"feed the baby",
	"play some quake",
}

func (p *ReminderPlugin) random(list []string) string {
	return list[rand.Intn(len(list))]
}

func (p *ReminderPlugin) randomReminder(service bruxism.Service) string {
	ticks := ""
	if service.Name() == bruxism.DiscordServiceName {
		ticks = "`"
	}

	return fmt.Sprintf("%s%sreminder %s %s%s", ticks, service.CommandPrefix(), p.random(randomTimes), p.random(randomMessages), ticks)
}

// Help returns a list of help strings that are printed when the user requests them.
func (p *ReminderPlugin) Help(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, detailed bool) []string {
	help := []string{
		bruxism.CommandHelp(service, "reminder", "<time> <reminder>", "Sets a reminder that is sent after the provided time.")[0],
	}
	if detailed {
		help = append(help, []string{
			"Examples: ",
			p.randomReminder(service),
			p.randomReminder(service),
		}...)
	}
	return help
}

func (p *ReminderPlugin) parseReminder(parts []string) (time.Time, string, error) {
	if parts[0] == "tomorrow" {
		return time.Now().Add(1 * time.Hour * 24), strings.Join(parts[1:], " "), nil
	}

	if parts[0] == "next" {
		switch parts[1] {
		case "week":
			return time.Now().Add(1 * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
		case "month":
			return time.Now().Add(1 * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
		case "year":
			return time.Now().Add(1 * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
		default:
			return time.Time{}, "", errors.New("Invalid next.")
		}
	}

	i, err := strconv.Atoi(parts[0])
	if err != nil {
		return time.Time{}, "", err
	}

	switch {
	case strings.HasPrefix(parts[1], "sec"):
		return time.Now().Add(time.Duration(i) * time.Second), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "min"):
		return time.Now().Add(time.Duration(i) * time.Minute), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "hour"):
		return time.Now().Add(time.Duration(i) * time.Hour), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "day"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "week"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "month"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "year"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
	}

	return time.Time{}, "", errors.New("Invalid string.")
}

// AddReminder adds a reminder.
func (p *ReminderPlugin) AddReminder(reminder *Reminder) error {
	p.Lock()
	defer p.Unlock()

	i := 0
	for _, r := range p.Reminders {
		if r.Requester == reminder.Requester {
			i++
			if i > 5 {
				return errors.New("You have too many reminders already.")
			}
		}
	}

	i = 0
	for _, r := range p.Reminders {
		if r.Time.After(reminder.Time) {
			break
		}
		i++
	}

	p.Reminders = append(p.Reminders, reminder)
	copy(p.Reminders[i+1:], p.Reminders[i:])
	p.Reminders[i] = reminder
	p.TotalReminders++

	return nil
}

func (p *ReminderPlugin) Message(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) {
	defer bruxism.MessageRecover()

	if service.IsMe(message) {
		return
	}

	if !bruxism.MatchesCommand(service, "remind", message) && !bruxism.MatchesCommand(service, "reminder", message) {
		return
	}

	_, parts := bruxism.ParseCommand(service, message)

	if len(parts) < 2 {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no time or message. eg: %s", p.randomReminder(service)))
		return
	}

	t, r, err := p.parseReminder(parts)

	now := time.Now()

	if err != nil || t.Before(now) || t.After(now.Add(time.Hour*24*365+time.Hour)) {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid time. eg: %s", strings.Join(randomTimes, ", ")))
		return
	}

	if r == "" {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no message. eg: %s", p.randomReminder(service)))
		return
	}

	requester := message.UserName()
	if service.Name() == bruxism.DiscordServiceName {
		requester = fmt.Sprintf("<@%s>", message.UserID())

		if strings.Index(r, "<@") != -1 {
			service.SendMessage(message.Channel(), "Invalid reminder, no mentions, sorry.")
			return
		}
	}

	err = p.AddReminder(&Reminder{
		StartTime: now,
		Time:      t,
		Requester: requester,
		Target:    message.Channel(),
		Message:   r,
		IsPrivate: service.IsPrivate(message),
	})
	if err != nil {
		service.SendMessage(message.Channel(), err.Error())
		return
	}

	service.SendMessage(message.Channel(), fmt.Sprintf("Reminder set for %s.", humanize.Time(t)))
}

// SendReminder sends a reminder.
func (p *ReminderPlugin) SendReminder(service bruxism.Service, reminder *Reminder) {
	if strings.Index(reminder.Message, "<@") != -1 {
		return
	}

	if reminder.IsPrivate {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s you set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Message))
	} else {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s %s set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Requester, reminder.Message))
	}
}

// Run will block until a reminder needs to be fired and then fire it.
func (p *ReminderPlugin) Run(bot *bruxism.Bot, service bruxism.Service) {
	for {
		p.RLock()

		if len(p.Reminders) > 0 {
			reminder := p.Reminders[0]
			if time.Now().After(reminder.Time) {
				p.RUnlock()

				p.SendReminder(service, reminder)

				p.Lock()
				p.Reminders = p.Reminders[1:]
				p.Unlock()

				continue
			}
		}

		p.RUnlock()
		time.Sleep(500 * time.Millisecond)
	}
}

// Load will load plugin state from a byte array.
func (p *ReminderPlugin) Load(bot *bruxism.Bot, service bruxism.Service, data []byte) error {
	if data != nil {
		if err := json.Unmarshal(data, p); err != nil {
			log.Println("Error loading data", err)
		}
	}
	if len(p.Reminders) > p.TotalReminders {
		p.TotalReminders = len(p.Reminders)
	}

	go p.Run(bot, service)
	return nil
}

// Save will save plugin state to a byte array.
func (p *ReminderPlugin) Save() ([]byte, error) {
	return json.Marshal(p)
}

// Stats will return the stats for a plugin.
func (p *ReminderPlugin) Stats(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) []string {
	return []string{fmt.Sprintf("Reminders: \t%d\n", p.TotalReminders)}
}

// Name returns the name of the plugin.
func (p *ReminderPlugin) Name() string {
	return "Reminder"
}

// New will create a new Reminder plugin.
func New() bruxism.Plugin {
	return &ReminderPlugin{
		Reminders: []*Reminder{},
	}
}
<|endoftext|>


<commit_before>package reminderplugin

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/iopred/bruxism"
)

// A Reminder holds data about a specific reminder.
type Reminder struct {
	StartTime time.Time
	Time      time.Time
	Requester string
	Target    string
	Message   string
	IsPrivate bool
}

// ReminderPlugin is a plugin that reminds users.
type ReminderPlugin struct {
	sync.RWMutex
	bot            *bruxism.Bot
	Reminders      []*Reminder
	TotalReminders int
}

var randomTimes = []string{
	"1 minute",
	"10 minutes",
	"1 hour",
	"4 hours",
	"tomorrow",
	"next week",
}

var randomMessages = []string{
	"walk the dog",
	"take pizza out of the oven",
	"check my email",
	"feed the baby",
	"play some quake",
}

func (p *ReminderPlugin) random(list []string) string {
	return list[rand.Intn(len(list))]
}

func (p *ReminderPlugin) randomReminder(service bruxism.Service) string {
	ticks := ""
	if service.Name() == bruxism.DiscordServiceName {
		ticks = "`"
	}

	return fmt.Sprintf("%s%sreminder %s %s%s", ticks, service.CommandPrefix(), p.random(randomTimes), p.random(randomMessages), ticks)
}

// Help returns a list of help strings that are printed when the user requests them.
func (p *ReminderPlugin) Help(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, detailed bool) []string {
	help := []string{
		bruxism.CommandHelp(service, "reminder", "<time> <reminder>", "Sets a reminder that is sent after the provided time.")[0],
		bruxism.CommandHelp(service, "reminderlist", "", "List all active reminders.")[0],
		bruxism.CommandHelp(service, "reminderdelete", "<index>", "Deletes a reminder by index. eg: reminderdelete 0")[0],
	}
	if detailed {
		help = append(help, []string{
			"Examples: ",
			p.randomReminder(service),
			p.randomReminder(service),
		}...)
	}
	return help
}

func (p *ReminderPlugin) parseReminder(parts []string) (time.Time, string, error) {
	if parts[0] == "tomorrow" {
		return time.Now().Add(1 * time.Hour * 24), strings.Join(parts[1:], " "), nil
	}

	if parts[0] == "next" {
		switch parts[1] {
		case "week":
			return time.Now().Add(1 * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
		case "month":
			return time.Now().Add(1 * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
		case "year":
			return time.Now().Add(1 * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
		default:
			return time.Time{}, "", errors.New("Invalid next.")
		}
	}

	i, err := strconv.Atoi(parts[0])
	if err != nil {
		return time.Time{}, "", err
	}

	switch {
	case strings.HasPrefix(parts[1], "sec"):
		return time.Now().Add(time.Duration(i) * time.Second), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "min"):
		return time.Now().Add(time.Duration(i) * time.Minute), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "hour"):
		return time.Now().Add(time.Duration(i) * time.Hour), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "day"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "week"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "month"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "year"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
	}

	return time.Time{}, "", errors.New("Invalid string.")
}

// AddReminder adds a reminder.
func (p *ReminderPlugin) AddReminder(reminder *Reminder) error {
	p.Lock()
	defer p.Unlock()

	i := 0
	for _, r := range p.Reminders {
		if r.Requester == reminder.Requester {
			i++
			if i > 5 {
				return errors.New("You have too many reminders already.")
			}
		}
	}

	i = 0
	for _, r := range p.Reminders {
		if r.Time.After(reminder.Time) {
			break
		}
		i++
	}

	p.Reminders = append(p.Reminders, reminder)
	copy(p.Reminders[i+1:], p.Reminders[i:])
	p.Reminders[i] = reminder
	p.TotalReminders++

	return nil
}

func (p *ReminderPlugin) Message(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) {
	defer bruxism.MessageRecover()

	if service.IsMe(message) {
		return
	}

	requester := message.UserName()
	if service.Name() == bruxism.DiscordServiceName {
		requester = fmt.Sprintf("<@%s>", message.UserID())
	}

	if bruxism.MatchesCommand(service, "remindlist", message) || bruxism.MatchesCommand(service, "reminderlist", message) {
		reminders := []string{}
		i := 0
		for _, r := range p.Reminders {
			if r.Requester == requester {
				reminders = append(reminders, fmt.Sprintf("%d - %s: %s", i, humanize.Time(r.Time), r.Message))
				i++
			}
		}
		if len(reminders) > 0 {
			if service.SupportsMultiline() {
				service.SendMessage(message.Channel(), fmt.Sprintf("Your reminders:\n%s", strings.Join(reminders, "\n")))
			} else {
				service.SendMessage(message.Channel(), fmt.Sprintf("Your reminders: %s", strings.Join(reminders, ". ")))
			}
		} else {
			service.SendMessage(message.Channel(), "You have no reminders.")
		}
		return
	}

	if bruxism.MatchesCommand(service, "reminddelete", message) || bruxism.MatchesCommand(service, "reminderdelete", message) {
		index := 0
		indexString, parts := bruxism.ParseCommand(service, message)
		var err error
		if len(parts) > 0 {
			index, err = strconv.Atoi(indexString)
			if err != nil {
				service.SendMessage(message.Channel(), "Invalid reminder.")
				return
			}
		}

		j := 0
		for i, r := range p.Reminders {
			if r.Requester == requester {
				if j == index {
					p.Lock()
					p.Reminders = append(p.Reminders[:j], p.Reminders[i+j:]...)
					p.Unlock()
					service.SendMessage(message.Channel(), "Reminder deleted.")
					return
				}
				j++
			}
		}
		return
	}

	if !bruxism.MatchesCommand(service, "remind", message) && !bruxism.MatchesCommand(service, "reminder", message) {
		return
	}

	_, parts := bruxism.ParseCommand(service, message)

	if len(parts) < 2 {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no time or message. eg: %s", p.randomReminder(service)))
		return
	}

	t, r, err := p.parseReminder(parts)

	now := time.Now()

	if err != nil || t.Before(now) || t.After(now.Add(time.Hour*24*365+time.Hour)) {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid time. eg: %s", strings.Join(randomTimes, ", ")))
		return
	}

	if r == "" {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no message. eg: %s", p.randomReminder(service)))
		return
	}

	if service.Name() == bruxism.DiscordServiceName {
		if strings.Index(r, "<@") != -1 || strings.Index(strings.ToLower(r), "@everyone") != -1 || strings.Index(strings.ToLower(r), "@here") != -1 {
			service.SendMessage(message.Channel(), "Invalid reminder, no mentions, sorry.")
			return
		}
	}

	err = p.AddReminder(&Reminder{
		StartTime: now,
		Time:      t,
		Requester: requester,
		Target:    message.Channel(),
		Message:   r,
		IsPrivate: service.IsPrivate(message),
	})
	if err != nil {
		service.SendMessage(message.Channel(), err.Error())
		return
	}

	service.SendMessage(message.Channel(), fmt.Sprintf("Reminder set for %s.", humanize.Time(t)))
}

// SendReminder sends a reminder.
func (p *ReminderPlugin) SendReminder(service bruxism.Service, reminder *Reminder) {
	if strings.Index(reminder.Message, "<@") != -1 || strings.Index(strings.ToLower(reminder.Message), "@everyone") != -1 || strings.Index(strings.ToLower(reminder.Message), "@here") != -1 {
		return
	}

	if reminder.IsPrivate {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s you set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Message))
	} else {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s %s set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Requester, reminder.Message))
	}
}

// Run will block until a reminder needs to be fired and then fire it.
func (p *ReminderPlugin) Run(bot *bruxism.Bot, service bruxism.Service) {
	for {
		p.RLock()

		if len(p.Reminders) > 0 {
			reminder := p.Reminders[0]
			if time.Now().After(reminder.Time) {
				p.RUnlock()

				p.SendReminder(service, reminder)

				p.Lock()
				p.Reminders = p.Reminders[1:]
				p.Unlock()

				continue
			}
		}

		p.RUnlock()
		time.Sleep(500 * time.Millisecond)
	}
}

// Load will load plugin state from a byte array.
func (p *ReminderPlugin) Load(bot *bruxism.Bot, service bruxism.Service, data []byte) error {
	if data != nil {
		if err := json.Unmarshal(data, p); err != nil {
			log.Println("Error loading data", err)
		}
	}
	if len(p.Reminders) > p.TotalReminders {
		p.TotalReminders = len(p.Reminders)
	}

	go p.Run(bot, service)
	return nil
}

// Save will save plugin state to a byte array.
func (p *ReminderPlugin) Save() ([]byte, error) {
	return json.Marshal(p)
}

// Stats will return the stats for a plugin.
func (p *ReminderPlugin) Stats(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) []string {
	return []string{fmt.Sprintf("Reminders: \t%d\n", p.TotalReminders)}
}

// Name returns the name of the plugin.
func (p *ReminderPlugin) Name() string {
	return "Reminder"
}

// New will create a new Reminder plugin.
func New() bruxism.Plugin {
	return &ReminderPlugin{
		Reminders: []*Reminder{},
	}
}
<commit_msg>I am really bad at development.<commit_after>package reminderplugin

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/iopred/bruxism"
)

// A Reminder holds data about a specific reminder.
type Reminder struct {
	StartTime time.Time
	Time      time.Time
	Requester string
	Target    string
	Message   string
	IsPrivate bool
}

// ReminderPlugin is a plugin that reminds users.
type ReminderPlugin struct {
	sync.RWMutex
	bot            *bruxism.Bot
	Reminders      []*Reminder
	TotalReminders int
}

var randomTimes = []string{
	"1 minute",
	"10 minutes",
	"1 hour",
	"4 hours",
	"tomorrow",
	"next week",
}

var randomMessages = []string{
	"walk the dog",
	"take pizza out of the oven",
	"check my email",
	"feed the baby",
	"play some quake",
}

func (p *ReminderPlugin) random(list []string) string {
	return list[rand.Intn(len(list))]
}

func (p *ReminderPlugin) randomReminder(service bruxism.Service) string {
	ticks := ""
	if service.Name() == bruxism.DiscordServiceName {
		ticks = "`"
	}

	return fmt.Sprintf("%s%sreminder %s %s%s", ticks, service.CommandPrefix(), p.random(randomTimes), p.random(randomMessages), ticks)
}

// Help returns a list of help strings that are printed when the user requests them.
func (p *ReminderPlugin) Help(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, detailed bool) []string {
	help := []string{
		bruxism.CommandHelp(service, "reminder", "<time> <reminder>", "Sets a reminder that is sent after the provided time.")[0],
		bruxism.CommandHelp(service, "reminderlist", "", "List all active reminders.")[0],
		bruxism.CommandHelp(service, "reminderdelete", "<index>", "Deletes a reminder by index. eg: reminderdelete 0")[0],
	}
	if detailed {
		help = append(help, []string{
			"Examples: ",
			p.randomReminder(service),
			p.randomReminder(service),
		}...)
	}
	return help
}

func (p *ReminderPlugin) parseReminder(parts []string) (time.Time, string, error) {
	if parts[0] == "tomorrow" {
		return time.Now().Add(1 * time.Hour * 24), strings.Join(parts[1:], " "), nil
	}

	if parts[0] == "next" {
		switch parts[1] {
		case "week":
			return time.Now().Add(1 * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
		case "month":
			return time.Now().Add(1 * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
		case "year":
			return time.Now().Add(1 * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
		default:
			return time.Time{}, "", errors.New("Invalid next.")
		}
	}

	i, err := strconv.Atoi(parts[0])
	if err != nil {
		return time.Time{}, "", err
	}

	switch {
	case strings.HasPrefix(parts[1], "sec"):
		return time.Now().Add(time.Duration(i) * time.Second), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "min"):
		return time.Now().Add(time.Duration(i) * time.Minute), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "hour"):
		return time.Now().Add(time.Duration(i) * time.Hour), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "day"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "week"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "month"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 7 * 4), strings.Join(parts[2:], " "), nil
	case strings.HasPrefix(parts[1], "year"):
		return time.Now().Add(time.Duration(i) * time.Hour * 24 * 365), strings.Join(parts[2:], " "), nil
	}

	return time.Time{}, "", errors.New("Invalid string.")
}

// AddReminder adds a reminder.
func (p *ReminderPlugin) AddReminder(reminder *Reminder) error {
	p.Lock()
	defer p.Unlock()

	i := 0
	for _, r := range p.Reminders {
		if r.Requester == reminder.Requester {
			i++
			if i > 5 {
				return errors.New("You have too many reminders already.")
			}
		}
	}

	i = 0
	for _, r := range p.Reminders {
		if r.Time.After(reminder.Time) {
			break
		}
		i++
	}

	p.Reminders = append(p.Reminders, reminder)
	copy(p.Reminders[i+1:], p.Reminders[i:])
	p.Reminders[i] = reminder
	p.TotalReminders++

	return nil
}

func (p *ReminderPlugin) Message(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) {
	defer bruxism.MessageRecover()

	if service.IsMe(message) {
		return
	}

	requester := message.UserName()
	if service.Name() == bruxism.DiscordServiceName {
		requester = fmt.Sprintf("<@%s>", message.UserID())
	}

	if bruxism.MatchesCommand(service, "remindlist", message) || bruxism.MatchesCommand(service, "reminderlist", message) {
		reminders := []string{}
		i := 0
		for _, r := range p.Reminders {
			if r.Requester == requester {
				reminders = append(reminders, fmt.Sprintf("%d - %s: %s", i, humanize.Time(r.Time), r.Message))
				i++
			}
		}
		if len(reminders) > 0 {
			if service.SupportsMultiline() {
				service.SendMessage(message.Channel(), fmt.Sprintf("Your reminders:\n%s", strings.Join(reminders, "\n")))
			} else {
				service.SendMessage(message.Channel(), fmt.Sprintf("Your reminders: %s", strings.Join(reminders, ". ")))
			}
		} else {
			service.SendMessage(message.Channel(), "You have no reminders.")
		}
		return
	}

	if bruxism.MatchesCommand(service, "reminddelete", message) || bruxism.MatchesCommand(service, "reminderdelete", message) {
		index := 0
		indexString, parts := bruxism.ParseCommand(service, message)
		var err error
		if len(parts) > 0 {
			index, err = strconv.Atoi(indexString)
			if err != nil {
				service.SendMessage(message.Channel(), "Invalid reminder.")
				return
			}
		}

		j := 0
		for i, r := range p.Reminders {
			if r.Requester == requester {
				if j == index {
					p.Lock()
					p.Reminders = append(p.Reminders[:i], p.Reminders[i+1:]...)
					p.Unlock()
					service.SendMessage(message.Channel(), "Reminder deleted.")
					return
				}
				j++
			}
		}
		return
	}

	if !bruxism.MatchesCommand(service, "remind", message) && !bruxism.MatchesCommand(service, "reminder", message) {
		return
	}

	_, parts := bruxism.ParseCommand(service, message)

	if len(parts) < 2 {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no time or message. eg: %s", p.randomReminder(service)))
		return
	}

	t, r, err := p.parseReminder(parts)

	now := time.Now()

	if err != nil || t.Before(now) || t.After(now.Add(time.Hour*24*365+time.Hour)) {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid time. eg: %s", strings.Join(randomTimes, ", ")))
		return
	}

	if r == "" {
		service.SendMessage(message.Channel(), fmt.Sprintf("Invalid reminder, no message. eg: %s", p.randomReminder(service)))
		return
	}

	if service.Name() == bruxism.DiscordServiceName {
		if strings.Index(r, "<@") != -1 || strings.Index(strings.ToLower(r), "@everyone") != -1 || strings.Index(strings.ToLower(r), "@here") != -1 {
			service.SendMessage(message.Channel(), "Invalid reminder, no mentions, sorry.")
			return
		}
	}

	err = p.AddReminder(&Reminder{
		StartTime: now,
		Time:      t,
		Requester: requester,
		Target:    message.Channel(),
		Message:   r,
		IsPrivate: service.IsPrivate(message),
	})
	if err != nil {
		service.SendMessage(message.Channel(), err.Error())
		return
	}

	service.SendMessage(message.Channel(), fmt.Sprintf("Reminder set for %s.", humanize.Time(t)))
}

// SendReminder sends a reminder.
func (p *ReminderPlugin) SendReminder(service bruxism.Service, reminder *Reminder) {
	if strings.Index(reminder.Message, "<@") != -1 || strings.Index(strings.ToLower(reminder.Message), "@everyone") != -1 || strings.Index(strings.ToLower(reminder.Message), "@here") != -1 {
		return
	}

	if reminder.IsPrivate {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s you set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Message))
	} else {
		service.SendMessage(reminder.Target, fmt.Sprintf("%s %s set a reminder: %s", humanize.Time(reminder.StartTime), reminder.Requester, reminder.Message))
	}
}

// Run will block until a reminder needs to be fired and then fire it.
func (p *ReminderPlugin) Run(bot *bruxism.Bot, service bruxism.Service) {
	for {
		p.RLock()

		if len(p.Reminders) > 0 {
			reminder := p.Reminders[0]
			if time.Now().After(reminder.Time) {
				p.RUnlock()

				p.SendReminder(service, reminder)

				p.Lock()
				p.Reminders = p.Reminders[1:]
				p.Unlock()

				continue
			}
		}

		p.RUnlock()
		time.Sleep(500 * time.Millisecond)
	}
}

// Load will load plugin state from a byte array.
func (p *ReminderPlugin) Load(bot *bruxism.Bot, service bruxism.Service, data []byte) error {
	if data != nil {
		if err := json.Unmarshal(data, p); err != nil {
			log.Println("Error loading data", err)
		}
	}
	if len(p.Reminders) > p.TotalReminders {
		p.TotalReminders = len(p.Reminders)
	}

	go p.Run(bot, service)
	return nil
}

// Save will save plugin state to a byte array.
func (p *ReminderPlugin) Save() ([]byte, error) {
	return json.Marshal(p)
}

// Stats will return the stats for a plugin.
func (p *ReminderPlugin) Stats(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) []string {
	return []string{fmt.Sprintf("Reminders: \t%d\n", p.TotalReminders)}
}

// Name returns the name of the plugin.
func (p *ReminderPlugin) Name() string {
	return "Reminder"
}

// New will create a new Reminder plugin.
func New() bruxism.Plugin {
	return &ReminderPlugin{
		Reminders: []*Reminder{},
	}
}
<|endoftext|>


<commit_before>package samples_test

import (
	"fmt"
	"io/ioutil"
	"path"
	"sort"
	"testing"
	"time"

	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/tft"
	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/utils"
)

const (
	setupPath = "../setup"
	sampleDir = "../../"
)

// testGroups represents a collection of samples matched to a projectID.
type testGroups struct {
	projectID string
	group     int
	samples   []string
}

// Retry if these errors are encountered.
var retryErrors = map[string]string{
	// IAM for Eventarc service agent is eventually consistent
	".*Permission denied while using the Eventarc Service Agent.*": "Eventarc Service Agent IAM is eventually consistent",
}

func TestSamples(t *testing.T) {
	// common test env
	testEnv := map[string]string{
		"GOOGLE_REGION": "us-central1",
		"GOOGLE_ZONE":   "us-central1-a",
	}
	for k, v := range testEnv {
		utils.SetEnv(t, k, v)
	}

	// This initial blueprint test is to extract output info
	// so we only have to set the env vars once.
	setup := tft.NewTFBlueprintTest(t,
		tft.WithTFDir(sampleDir),
		tft.WithSetupPath(setupPath),
	)
	testProjectIDs := setup.GetTFSetupOutputListVal("project_ids")
	// number of groups is determined by length of testProjectIDs slice
	testGroups := discoverTestCaseGroups(t, testProjectIDs)

	for _, tg := range testGroups {
		for _, sample := range tg.samples {
			testName := fmt.Sprintf("%d/%s", tg.group, sample)
			t.Run(testName, func(t *testing.T) {
				utils.SetEnv(t, "GOOGLE_PROJECT", tg.projectID)
				sampleTest := tft.NewTFBlueprintTest(t,
					tft.WithTFDir(path.Join(sampleDir, sample)),
					tft.WithSetupPath(setupPath),
					tft.WithRetryableTerraformErrors(retryErrors, 10, time.Minute),
				)
				sampleTest.Test()
			})
		}
	}
}

// skipDiscoverDirs are directories that are skipped when discovering test cases.
var skipDiscoverDirs = map[string]bool{
	"test":  true,
	"build": true,
	".git":  true,
}

// discoverTestCaseGroups discovers individual sample directories in the parent directory
// and assigns them to a test group based on projects provided.
// It also skips known directories that are not samples.
func discoverTestCaseGroups(t *testing.T, projects []string) []*testGroups {
	t.Helper()
	samples := []string{}
	dirs, err := ioutil.ReadDir(sampleDir)
	if err != nil {
		t.Fatal(err)
	}
	for _, f := range dirs {
		if !f.IsDir() || skipDiscoverDirs[f.Name()] {
			continue
		}
		samples = append(samples, f.Name())
	}
	sort.Strings(samples)

	// One test group is associated to one project.
	groups := []*testGroups{}
	for i, project := range projects {
		groups = append(groups, &testGroups{projectID: project, samples: []string{}, group: i})
	}
	// Rather than chunking we assign them in a round robin fashion as some samples like sql takes more time.
	// Chunking would result in all sql* assigned to a single project.
	// We sort the sample slice beforehand so assignments should be stable for a given run.
	for i, sample := range samples {
		groupIndex := i % len(projects)
		groups[groupIndex].samples = append(groups[groupIndex].samples, sample)
	}
	return groups
}
<commit_msg>chore: retry service activation, IAM (#181)<commit_after>package samples_test

import (
	"fmt"
	"io/ioutil"
	"path"
	"sort"
	"testing"
	"time"

	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/tft"
	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/utils"
)

const (
	setupPath = "../setup"
	sampleDir = "../../"
)

// testGroups represents a collection of samples matched to a projectID.
type testGroups struct {
	projectID string
	group     int
	samples   []string
}

// Retry if these errors are encountered.
var retryErrors = map[string]string{
	// IAM for Eventarc service agent is eventually consistent
	".*Permission denied while using the Eventarc Service Agent.*": "Eventarc Service Agent IAM is eventually consistent",
	// API activation is eventually consistent.
	".*SERVICE_DISABLED.*": "Service enablement is eventually consistent",
	// Retry function GCS access
        ".*does not have storage.objects.get access.*": "GCS IAM is eventually consistent",
}

func TestSamples(t *testing.T) {
	// common test env
	testEnv := map[string]string{
		"GOOGLE_REGION": "us-central1",
		"GOOGLE_ZONE":   "us-central1-a",
	}
	for k, v := range testEnv {
		utils.SetEnv(t, k, v)
	}

	// This initial blueprint test is to extract output info
	// so we only have to set the env vars once.
	setup := tft.NewTFBlueprintTest(t,
		tft.WithTFDir(sampleDir),
		tft.WithSetupPath(setupPath),
	)
	testProjectIDs := setup.GetTFSetupOutputListVal("project_ids")
	// number of groups is determined by length of testProjectIDs slice
	testGroups := discoverTestCaseGroups(t, testProjectIDs)

	for _, tg := range testGroups {
		for _, sample := range tg.samples {
			testName := fmt.Sprintf("%d/%s", tg.group, sample)
			t.Run(testName, func(t *testing.T) {
				utils.SetEnv(t, "GOOGLE_PROJECT", tg.projectID)
				sampleTest := tft.NewTFBlueprintTest(t,
					tft.WithTFDir(path.Join(sampleDir, sample)),
					tft.WithSetupPath(setupPath),
					tft.WithRetryableTerraformErrors(retryErrors, 10, time.Minute),
				)
				sampleTest.Test()
			})
		}
	}
}

// skipDiscoverDirs are directories that are skipped when discovering test cases.
var skipDiscoverDirs = map[string]bool{
	"test":  true,
	"build": true,
	".git":  true,
}

// discoverTestCaseGroups discovers individual sample directories in the parent directory
// and assigns them to a test group based on projects provided.
// It also skips known directories that are not samples.
func discoverTestCaseGroups(t *testing.T, projects []string) []*testGroups {
	t.Helper()
	samples := []string{}
	dirs, err := ioutil.ReadDir(sampleDir)
	if err != nil {
		t.Fatal(err)
	}
	for _, f := range dirs {
		if !f.IsDir() || skipDiscoverDirs[f.Name()] {
			continue
		}
		samples = append(samples, f.Name())
	}
	sort.Strings(samples)

	// One test group is associated to one project.
	groups := []*testGroups{}
	for i, project := range projects {
		groups = append(groups, &testGroups{projectID: project, samples: []string{}, group: i})
	}
	// Rather than chunking we assign them in a round robin fashion as some samples like sql takes more time.
	// Chunking would result in all sql* assigned to a single project.
	// We sort the sample slice beforehand so assignments should be stable for a given run.
	for i, sample := range samples {
		groupIndex := i % len(projects)
		groups[groupIndex].samples = append(groups[groupIndex].samples, sample)
	}
	return groups
}
<|endoftext|>


<commit_before>package consul

import (
	"errors"
	"strings"
	"time"

	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/consul/structs"
)

const (
	// aclNotFound indicates there is no matching ACL
	aclNotFound = "ACL not found"

	// anonymousToken is the token ID we re-write to if there
	// is no token ID provided
	anonymousToken = "anonymous"
)

// aclCacheEntry is used to cache non-authoritative ACL's
// If non-authoritative, then we must respect a TTL
type aclCacheEntry struct {
	ACL     acl.ACL
	Expires time.Time
	ETag    string
}

// aclFault is used to fault in the rules for an ACL if we take a miss
func (s *Server) aclFault(id string) (string, error) {
	state := s.fsm.State()
	_, acl, err := state.ACLGet(id)
	if err != nil {
		return "", err
	}
	if acl == nil {
		return "", errors.New(aclNotFound)
	}
	return acl.Rules, nil
}

// resolveToken is used to resolve an ACL is any is appropriate
func (s *Server) resolveToken(id string) (acl.ACL, error) {
	// Check if there is no ACL datacenter (ACL's disabled)
	authDC := s.config.ACLDatacenter
	if len(authDC) == 0 {
		return nil, nil
	}

	// Handle the anonymous token
	if len(id) == 0 {
		id = anonymousToken
	}

	// Check if we are the ACL datacenter and the leader, use the
	// authoritative cache
	if s.config.Datacenter == authDC && s.IsLeader() {
		return s.aclAuthCache.GetACL(id)
	}

	// Use our non-authoritative cache
	return s.lookupACL(id, authDC)
}

// lookupACL is used when we are non-authoritative, and need
// to resolve an ACL
func (s *Server) lookupACL(id, authDC string) (acl.ACL, error) {
	// Check the cache for the ACL
	var cached *aclCacheEntry
	raw, ok := s.aclCache.Get(id)
	if ok {
		cached = raw.(*aclCacheEntry)
	}

	// Check for live cache
	if cached != nil && time.Now().Before(cached.Expires) {
		return cached.ACL, nil
	}

	// Attempt to refresh the policy
	args := structs.ACLPolicyRequest{
		Datacenter: authDC,
		ACL:        id,
	}
	var out structs.ACLPolicy
	err := s.RPC("ACL.GetPolicy", &args, &out)

	// Handle the happy path
	if err == nil {
		return s.useACLPolicy(id, cached, &out)
	}

	// Check for not-found
	if strings.Contains(err.Error(), aclNotFound) {
		return nil, errors.New(aclNotFound)
	} else {
		s.logger.Printf("[ERR] consul.acl: Failed to get policy for '%s': %v", id, err)
	}

	// Unable to refresh, apply the down policy
	switch s.config.ACLDownPolicy {
	case "allow":
		return acl.AllowAll(), nil
	case "extend-cache":
		if cached != nil {
			return cached.ACL, nil
		}
		fallthrough
	default:
		return acl.DenyAll(), nil
	}
}

// useACLPolicy handles an ACLPolicy response
func (s *Server) useACLPolicy(id string, cached *aclCacheEntry, p *structs.ACLPolicy) (acl.ACL, error) {
	// Check if we can used the cached policy
	if cached != nil && cached.ETag == p.ETag {
		if p.TTL > 0 {
			cached.Expires = time.Now().Add(p.TTL)
		}
		return cached.ACL, nil
	}

	// Check for a cached compiled policy
	var compiled acl.ACL
	raw, ok := s.aclPolicyCache.Get(p.ETag)
	if ok {
		compiled = raw.(acl.ACL)
	} else {
		// Determine the root policy
		var root acl.ACL
		switch p.Root {
		case "allow":
			root = acl.AllowAll()
		default:
			root = acl.DenyAll()
		}

		// Compile the ACL
		acl, err := acl.New(root, p.Policy)
		if err != nil {
			return nil, err
		}

		// Cache the policy
		s.aclPolicyCache.Add(p.ETag, acl)
		compiled = acl
	}

	// Cache the ACL
	cached = &aclCacheEntry{
		ACL:  compiled,
		ETag: p.ETag,
	}
	if p.TTL > 0 {
		cached.Expires = time.Now().Add(p.TTL)
	}
	s.aclCache.Add(id, cached)
	return compiled, nil
}
<commit_msg>consul: Adding some metrics for ACL usage<commit_after>package consul

import (
	"errors"
	"strings"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/consul/structs"
)

const (
	// aclNotFound indicates there is no matching ACL
	aclNotFound = "ACL not found"

	// anonymousToken is the token ID we re-write to if there
	// is no token ID provided
	anonymousToken = "anonymous"
)

// aclCacheEntry is used to cache non-authoritative ACL's
// If non-authoritative, then we must respect a TTL
type aclCacheEntry struct {
	ACL     acl.ACL
	Expires time.Time
	ETag    string
}

// aclFault is used to fault in the rules for an ACL if we take a miss
func (s *Server) aclFault(id string) (string, error) {
	defer metrics.MeasureSince([]string{"consul", "acl", "fault"}, time.Now())
	state := s.fsm.State()
	_, acl, err := state.ACLGet(id)
	if err != nil {
		return "", err
	}
	if acl == nil {
		return "", errors.New(aclNotFound)
	}
	return acl.Rules, nil
}

// resolveToken is used to resolve an ACL is any is appropriate
func (s *Server) resolveToken(id string) (acl.ACL, error) {
	// Check if there is no ACL datacenter (ACL's disabled)
	authDC := s.config.ACLDatacenter
	if len(authDC) == 0 {
		return nil, nil
	}
	defer metrics.MeasureSince([]string{"consul", "acl", "resolveToken"}, time.Now())

	// Handle the anonymous token
	if len(id) == 0 {
		id = anonymousToken
	}

	// Check if we are the ACL datacenter and the leader, use the
	// authoritative cache
	if s.config.Datacenter == authDC && s.IsLeader() {
		return s.aclAuthCache.GetACL(id)
	}

	// Use our non-authoritative cache
	return s.lookupACL(id, authDC)
}

// lookupACL is used when we are non-authoritative, and need
// to resolve an ACL
func (s *Server) lookupACL(id, authDC string) (acl.ACL, error) {
	// Check the cache for the ACL
	var cached *aclCacheEntry
	raw, ok := s.aclCache.Get(id)
	if ok {
		cached = raw.(*aclCacheEntry)
	}

	// Check for live cache
	if cached != nil && time.Now().Before(cached.Expires) {
		metrics.IncrCounter([]string{"consul", "acl", "cache_hit"}, 1)
		return cached.ACL, nil
	} else {
		metrics.IncrCounter([]string{"consul", "acl", "cache_miss"}, 1)
	}

	// Attempt to refresh the policy
	args := structs.ACLPolicyRequest{
		Datacenter: authDC,
		ACL:        id,
	}
	var out structs.ACLPolicy
	err := s.RPC("ACL.GetPolicy", &args, &out)

	// Handle the happy path
	if err == nil {
		return s.useACLPolicy(id, cached, &out)
	}

	// Check for not-found
	if strings.Contains(err.Error(), aclNotFound) {
		return nil, errors.New(aclNotFound)
	} else {
		s.logger.Printf("[ERR] consul.acl: Failed to get policy for '%s': %v", id, err)
	}

	// Unable to refresh, apply the down policy
	switch s.config.ACLDownPolicy {
	case "allow":
		return acl.AllowAll(), nil
	case "extend-cache":
		if cached != nil {
			return cached.ACL, nil
		}
		fallthrough
	default:
		return acl.DenyAll(), nil
	}
}

// useACLPolicy handles an ACLPolicy response
func (s *Server) useACLPolicy(id string, cached *aclCacheEntry, p *structs.ACLPolicy) (acl.ACL, error) {
	// Check if we can used the cached policy
	if cached != nil && cached.ETag == p.ETag {
		if p.TTL > 0 {
			cached.Expires = time.Now().Add(p.TTL)
		}
		return cached.ACL, nil
	}

	// Check for a cached compiled policy
	var compiled acl.ACL
	raw, ok := s.aclPolicyCache.Get(p.ETag)
	if ok {
		compiled = raw.(acl.ACL)
	} else {
		// Determine the root policy
		var root acl.ACL
		switch p.Root {
		case "allow":
			root = acl.AllowAll()
		default:
			root = acl.DenyAll()
		}

		// Compile the ACL
		acl, err := acl.New(root, p.Policy)
		if err != nil {
			return nil, err
		}

		// Cache the policy
		s.aclPolicyCache.Add(p.ETag, acl)
		compiled = acl
	}

	// Cache the ACL
	cached = &aclCacheEntry{
		ACL:  compiled,
		ETag: p.ETag,
	}
	if p.TTL > 0 {
		cached.Expires = time.Now().Add(p.TTL)
	}
	s.aclCache.Add(id, cached)
	return compiled, nil
}
<|endoftext|>


<commit_before>package xlog

import "io"

// globalInstance stores the global logger.
var globalInstance *Logger

// Close releases any resources held by the global logger. The logger should
// not be used again after calling this method without re-configuring it, as
// this method sets the global instance to nil.
func Close() {
	instance().Close()
	globalInstance = nil
}

// SetName sets the name of the global logger.
func SetName(name string) {
	instance().SetName(name)
}

// Append adds a file to the global logger.
func Append(file string, level Level) {
	instance().Append(file, level)
}

// MultiAppend adds one or more files to the global logger.
func MultiAppend(files []string, level Level) {
	instance().MultiAppend(files, level)
}

// AppendWriter adds a writer to the global logger.
func AppendWriter(writer io.Writer, level Level) {
	instance().AppendWriter(writer, level)
}

// MultiAppendWriters adds one or more io.Writer instances to the global logger.
func MultiAppendWriters(writers []io.Writer, level Level) {
	instance().MultiAppendWriters(writers, level)
}

// Writable returns true when global logging is enabled, and the global logger
// hasn't been closed.
func Writable() bool {
	return instance().Writable()
}

// Log writes the message to each logger appended to the global logger at given level
// or higher.
func Log(level Level, v ...interface{}) {
	instance().Log(level, v...)
}

// Logf writes the message to each logger appended to the global logger given
// level or higher.
func Logf(level Level, format string, v ...interface{}) {
	instance().Logf(level, format, v...)
}

// Debug writes to the global logger at DebugLevel.
// Arguments are handled in the manner of fmt.Print.
func Debug(v ...interface{}) {
	instance().Debug(v...)
}

// Debugf writes to the global logger at DebugLevel.
// Arguments are handled in the manner of fmt.Printf.
func Debugf(format string, v ...interface{}) {
	instance().Debugf(format, v...)
}

// Info writes to the global logger at InfoLevel.
// Arguments are handled in the manner of fmt.Print.
func Info(v ...interface{}) {
	instance().Info(v...)
}

// Infof writes to the global logger at InfoLevel.
// Arguments are handled in the manner of fmt.Printf.
func Infof(format string, v ...interface{}) {
	instance().Infof(format, v...)
}

// Notice writes to the global logger at NoticeLevel.
// Arguments are handled in the manner of fmt.Print.
func Notice(v ...interface{}) {
	instance().Notice(v...)
}

// Noticef writes to the global logger at NoticeLevel.
// Arguments are handled in the manner of fmt.Printf.
func Noticef(format string, v ...interface{}) {
	instance().Noticef(format, v...)
}

// Warning writes to the global logger at WarningLevel.
// Arguments are handled in the manner of fmt.Print.
func Warning(v ...interface{}) {
	instance().Warning(v...)
}

// Warningf writes to the global logger at WarningLevel.
// Arguments are handled in the manner of fmt.Printf.
func Warningf(format string, v ...interface{}) {
	instance().Warningf(format, v...)
}

// Error writes to the global logger at ErrorLevel.
// Arguments are handled in the manner of fmt.Print.
func Error(v ...interface{}) {
	instance().Error(v...)
}

// Errorf writes to the global logger at ErrorLevel.
// Arguments are handled in the manner of fmt.Printf.
func Errorf(format string, v ...interface{}) {
	instance().Errorf(format, v...)
}

// Critical writes to the global logger at CriticalLevel.
// Arguments are handled in the manner of fmt.Print.
func Critical(v ...interface{}) {
	instance().Critical(v...)
}

// Criticalf writes to the global logger at CriticalLevel.
// Arguments are handled in the manner of fmt.Printf.
func Criticalf(format string, v ...interface{}) {
	instance().Criticalf(format, v...)
}

// Alert writes to the global logger at AlertLevel.
// Arguments are handled in the manner of fmt.Print.
func Alert(v ...interface{}) {
	instance().Alert(v...)
}

// Alertf writes to the global logger at AlertLevel.
// Arguments are handled in the manner of fmt.Printf.
func Alertf(format string, v ...interface{}) {
	instance().Alertf(format, v...)
}

// Emergency writes to the global logger at EmergencyLevel.
// Arguments are handled in the manner of fmt.Print.
func Emergency(v ...interface{}) {
	instance().Emergency(v...)
}

// Emergencyf writes to the global logger at EmergencyLevel.
// Arguments are handled in the manner of fmt.Printf.
func Emergencyf(format string, v ...interface{}) {
	instance().Emergencyf(format, v...)
}

// instance calls panic() when the global logger has not been configured.
func instance() *Logger {
	if globalInstance == nil {
		globalInstance = NewLogger("xlog")
	}
	return globalInstance
}
<commit_msg>Added some global functions<commit_after>package xlog

import "io"

// globalInstance stores the global logger.
var globalInstance *Logger

// Close releases any resources held by the global logger. The logger should
// not be used again after calling this method without re-configuring it, as
// this method sets the global instance to nil.
func Close() {
	instance().Close()
	globalInstance = nil
}

// SetName sets the name of the global logger.
func SetName(name string) {
	instance().SetName(name)
}

// SetFormatter sets the formatter used by the global logger.
func SetFormatter(formatter Formatter) {
	instance().Formatter = formatter
}

// SetLoggerMap sets the logger map used by the global logger.
func SetLoggerMap(lm LoggerMap) {
	instance().Loggers = lm
}

// Enabled returns whether the global logger is enabled.
func Enabled() bool {
	return instance().Enabled
}

// SetEnabled sets whether the global logger is enabled.
func SetEnabled(enabled bool) {
	instance().Enabled = enabled
}

// Append adds a file to the global logger.
func Append(file string, level Level) {
	instance().Append(file, level)
}

// MultiAppend adds one or more files to the global logger.
func MultiAppend(files []string, level Level) {
	instance().MultiAppend(files, level)
}

// AppendWriter adds a writer to the global logger.
func AppendWriter(writer io.Writer, level Level) {
	instance().AppendWriter(writer, level)
}

// MultiAppendWriters adds one or more io.Writer instances to the global logger.
func MultiAppendWriters(writers []io.Writer, level Level) {
	instance().MultiAppendWriters(writers, level)
}

// Writable returns true when global logging is enabled, and the global logger
// hasn't been closed.
func Writable() bool {
	return instance().Writable()
}

// Log writes the message to each logger appended to the global logger at given level
// or higher.
func Log(level Level, v ...interface{}) {
	instance().Log(level, v...)
}

// Logf writes the message to each logger appended to the global logger given
// level or higher.
func Logf(level Level, format string, v ...interface{}) {
	instance().Logf(level, format, v...)
}

// Debug writes to the global logger at DebugLevel.
// Arguments are handled in the manner of fmt.Print.
func Debug(v ...interface{}) {
	instance().Debug(v...)
}

// Debugf writes to the global logger at DebugLevel.
// Arguments are handled in the manner of fmt.Printf.
func Debugf(format string, v ...interface{}) {
	instance().Debugf(format, v...)
}

// Info writes to the global logger at InfoLevel.
// Arguments are handled in the manner of fmt.Print.
func Info(v ...interface{}) {
	instance().Info(v...)
}

// Infof writes to the global logger at InfoLevel.
// Arguments are handled in the manner of fmt.Printf.
func Infof(format string, v ...interface{}) {
	instance().Infof(format, v...)
}

// Notice writes to the global logger at NoticeLevel.
// Arguments are handled in the manner of fmt.Print.
func Notice(v ...interface{}) {
	instance().Notice(v...)
}

// Noticef writes to the global logger at NoticeLevel.
// Arguments are handled in the manner of fmt.Printf.
func Noticef(format string, v ...interface{}) {
	instance().Noticef(format, v...)
}

// Warning writes to the global logger at WarningLevel.
// Arguments are handled in the manner of fmt.Print.
func Warning(v ...interface{}) {
	instance().Warning(v...)
}

// Warningf writes to the global logger at WarningLevel.
// Arguments are handled in the manner of fmt.Printf.
func Warningf(format string, v ...interface{}) {
	instance().Warningf(format, v...)
}

// Error writes to the global logger at ErrorLevel.
// Arguments are handled in the manner of fmt.Print.
func Error(v ...interface{}) {
	instance().Error(v...)
}

// Errorf writes to the global logger at ErrorLevel.
// Arguments are handled in the manner of fmt.Printf.
func Errorf(format string, v ...interface{}) {
	instance().Errorf(format, v...)
}

// Critical writes to the global logger at CriticalLevel.
// Arguments are handled in the manner of fmt.Print.
func Critical(v ...interface{}) {
	instance().Critical(v...)
}

// Criticalf writes to the global logger at CriticalLevel.
// Arguments are handled in the manner of fmt.Printf.
func Criticalf(format string, v ...interface{}) {
	instance().Criticalf(format, v...)
}

// Alert writes to the global logger at AlertLevel.
// Arguments are handled in the manner of fmt.Print.
func Alert(v ...interface{}) {
	instance().Alert(v...)
}

// Alertf writes to the global logger at AlertLevel.
// Arguments are handled in the manner of fmt.Printf.
func Alertf(format string, v ...interface{}) {
	instance().Alertf(format, v...)
}

// Emergency writes to the global logger at EmergencyLevel.
// Arguments are handled in the manner of fmt.Print.
func Emergency(v ...interface{}) {
	instance().Emergency(v...)
}

// Emergencyf writes to the global logger at EmergencyLevel.
// Arguments are handled in the manner of fmt.Printf.
func Emergencyf(format string, v ...interface{}) {
	instance().Emergencyf(format, v...)
}

// instance calls panic() when the global logger has not been configured.
func instance() *Logger {
	if globalInstance == nil {
		globalInstance = NewLogger("xlog")
	}
	return globalInstance
}
<|endoftext|>


<commit_before>// Copyright 2017 Tomas Machalek <tomas.machalek@gmail.com>
//
// 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

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/tomachalek/gloomy/index/builder"
	"github.com/tomachalek/gloomy/index/extras"
	"github.com/tomachalek/gloomy/index/gconf"
	"github.com/tomachalek/gloomy/service"
)

const (
	createIndexAction   = "create-index"
	extractNgramsAction = "extract-ngrams"
	searchServiceAction = "search-service"
	searchAction        = "search"
)

func help() {
	fmt.Println("HELP:")
}

func createIndex(conf *gconf.IndexBuilderConf, ngramSize int) {
	if conf.VerticalFilePath == "" {
		fmt.Println("Vertical file not specified")
		os.Exit(1)
	}
	log.Println("Importing vertical file ", conf.VerticalFilePath)
	if conf.OutDirectory == "" {
		conf.OutDirectory = filepath.Dir(conf.VerticalFilePath)
	}
	fmt.Println("Output directory: ", conf.OutDirectory)
	t0 := time.Now()
	builder.CreateGloomyIndex(conf, ngramSize)
	fmt.Printf("DONE in %s\n", time.Since(t0))
}

func extractNgrams(conf *gconf.IndexBuilderConf, ngramSize int) {
	if conf.VerticalFilePath == "" {
		fmt.Println("Vertical file not specified")
		os.Exit(1)
	}
	log.Println("Processing vertical file ", conf.VerticalFilePath)
	if conf.OutDirectory == "" {
		conf.OutDirectory = filepath.Dir(conf.VerticalFilePath)
	}
	fmt.Println("Output directory: ", conf.OutDirectory)
	t0 := time.Now()
	extras.ExtractUniqueNgrams(conf, ngramSize)
	fmt.Printf("DONE in %s\n", time.Since(t0))
}

func loadSearchConf(confBasePath string) *gconf.SearchConf {
	if confBasePath == "" {
		var err error
		confBasePath, err = os.Getwd()
		if err != nil {
			panic(err)
		}
		confBasePath = filepath.Join(confBasePath, "gloomy.json")
	}
	return gconf.LoadSearchConf(confBasePath)
}

func searchCLI(confBasePath string, query string) {
	conf := loadSearchConf(confBasePath)
	ans, err := service.Search(conf.DataPath, "susanne", query)
	if err != nil {
		log.Printf("Srch error: %s", err)
	}
	for i := 0; ans.HasNext(); i++ {
		v := ans.Next()
		log.Printf("res[%d]: %s", i, v)
	}
}

func startSearchService(confBasePath string) {
	conf := loadSearchConf(confBasePath)
	service.Serve(conf)
}

func main() {
	ngramSize := flag.Int("ngram-size", 2, "N-gram size, 2: bigram (default), ...")
	srchConfPath := flag.String("conf-path", "", "Path to the gloomy.conf (by default, working dir is used")
	flag.Parse()
	if len(flag.Args()) == 0 {
		fmt.Println("Missing action, try -h for help")
		os.Exit(1)

	} else {
		switch flag.Arg(0) {
		case "help":
			help()
		case createIndexAction:
			conf := gconf.LoadIndexBuilderConf(flag.Arg(1))
			createIndex(conf, *ngramSize)
		case extractNgramsAction:
			conf := gconf.LoadIndexBuilderConf(flag.Arg(1))
			extractNgrams(conf, *ngramSize)
		case searchServiceAction:
			startSearchService(*srchConfPath)
		case searchAction:
			searchCLI(*srchConfPath, flag.Arg(1))
		default:
			panic("Unknown action")
		}
	}
}
<commit_msg>Remove hardcoded 'susanne' corpus<commit_after>// Copyright 2017 Tomas Machalek <tomas.machalek@gmail.com>
//
// 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

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/tomachalek/gloomy/index/builder"
	"github.com/tomachalek/gloomy/index/extras"
	"github.com/tomachalek/gloomy/index/gconf"
	"github.com/tomachalek/gloomy/service"
)

const (
	createIndexAction   = "create-index"
	extractNgramsAction = "extract-ngrams"
	searchServiceAction = "search-service"
	searchAction        = "search"
)

func help() {
	fmt.Println("HELP:")
}

func createIndex(conf *gconf.IndexBuilderConf, ngramSize int) {
	if conf.VerticalFilePath == "" {
		fmt.Println("Vertical file not specified")
		os.Exit(1)
	}
	log.Println("Importing vertical file ", conf.VerticalFilePath)
	if conf.OutDirectory == "" {
		conf.OutDirectory = filepath.Dir(conf.VerticalFilePath)
	}
	fmt.Println("Output directory: ", conf.OutDirectory)
	t0 := time.Now()
	builder.CreateGloomyIndex(conf, ngramSize)
	fmt.Printf("DONE in %s\n", time.Since(t0))
}

func extractNgrams(conf *gconf.IndexBuilderConf, ngramSize int) {
	if conf.VerticalFilePath == "" {
		fmt.Println("Vertical file not specified")
		os.Exit(1)
	}
	log.Println("Processing vertical file ", conf.VerticalFilePath)
	if conf.OutDirectory == "" {
		conf.OutDirectory = filepath.Dir(conf.VerticalFilePath)
	}
	fmt.Println("Output directory: ", conf.OutDirectory)
	t0 := time.Now()
	extras.ExtractUniqueNgrams(conf, ngramSize)
	fmt.Printf("DONE in %s\n", time.Since(t0))
}

func loadSearchConf(confBasePath string) *gconf.SearchConf {
	if confBasePath == "" {
		var err error
		confBasePath, err = os.Getwd()
		if err != nil {
			panic(err)
		}
		confBasePath = filepath.Join(confBasePath, "gloomy.json")
	}
	return gconf.LoadSearchConf(confBasePath)
}

func searchCLI(confBasePath string, corpus string, query string) {
	conf := loadSearchConf(confBasePath)
	ans, err := service.Search(conf.DataPath, corpus, query)
	if err != nil {
		log.Printf("Srch error: %s", err)
	}
	for i := 0; ans.HasNext(); i++ {
		v := ans.Next()
		log.Printf("res[%d]: %s", i, v)
	}
}

func startSearchService(confBasePath string) {
	conf := loadSearchConf(confBasePath)
	service.Serve(conf)
}

func main() {
	ngramSize := flag.Int("ngram-size", 2, "N-gram size, 2: bigram (default), ...")
	srchConfPath := flag.String("conf-path", "", "Path to the gloomy.conf (by default, working dir is used")
	flag.Parse()
	if len(flag.Args()) == 0 {
		fmt.Println("Missing action, try -h for help")
		os.Exit(1)

	} else {
		switch flag.Arg(0) {
		case "help":
			help()
		case createIndexAction:
			conf := gconf.LoadIndexBuilderConf(flag.Arg(1))
			createIndex(conf, *ngramSize)
		case extractNgramsAction:
			conf := gconf.LoadIndexBuilderConf(flag.Arg(1))
			extractNgrams(conf, *ngramSize)
		case searchServiceAction:
			startSearchService(*srchConfPath)
		case searchAction:
			if flag.Arg(1) == "" || flag.Arg(2) == "" {
				log.Fatal("Missing argument (both corpus and query must be specified)")
			}
			searchCLI(*srchConfPath, flag.Arg(1), flag.Arg(2))
		default:
			panic("Unknown action")
		}
	}
}
<|endoftext|>