<commit_before>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package rec

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"

	"github.com/ikravets/errs"

	"my/ev/packet"
	"my/ev/packet/bats"
	"my/ev/packet/miax"
	"my/ev/packet/nasdaq"
	"my/ev/sim"
)

type SimLoggerConfig struct {
	EfhLoggerConfig
	SupernodeLevels int
}

type SimLogger struct {
	w               io.Writer
	tobOld, tobNew  []sim.PriceLevel
	efhLogger       EfhLogger
	supernodeLevels int
}

const SimLoggerDefaultSupernodeLevels = 256
const SimLoggerUpperSupernodeLevels = 8

func NewSimLogger(c SimLoggerConfig) *SimLogger {
	s := &SimLogger{
		w:               c.Writer,
		supernodeLevels: c.SupernodeLevels,
	}
	if s.supernodeLevels == 0 {
		s.supernodeLevels = SimLoggerDefaultSupernodeLevels
	}
	elc := c.EfhLoggerConfig
	elc.Printer = s
	elc.Writer = nil
	s.efhLogger = *NewEfhLogger(elc)
	return s
}
func (s *SimLogger) SetSupernodeLevels(levels int) {
	errs.Check(levels > 0)
	s.supernodeLevels = levels
}

func (s *SimLogger) printf(format string, vs ...interface{}) {
	_, err := fmt.Fprintf(s.w, format, vs...)
	errs.CheckE(err)
}
func (s *SimLogger) printfln(format string, vs ...interface{}) {
	f := format + "\n"
	s.printf(f, vs...)
}
func (s *SimLogger) MessageArrived(idm *sim.SimMessage) {
	outItto := func(name string, typ nasdaq.IttoMessageType, f string, vs ...interface{}) {
		s.printf("NORM %s %c ", name, typ)
		s.printfln(f, vs...)
	}
	outBats := func(f string, vs ...interface{}) {
		s.printf("NORM ORDER %02x ", idm.Pam.Layer().(bats.PitchMessage).Base().Type.ToInt())
		s.printfln(f, vs...)
	}
	outMiax := func(f string, vs ...interface{}) {
		s.printf("NORM TOM %02x ", idm.Pam.Layer().(miax.TomMessage).Base().Type.ToInt())
		s.printfln(f, vs...)
	}
	sideChar := func(s packet.MarketSide) byte {
		if s == packet.MarketSideAsk {
			return 'S'
		}
		return byte(s)
	}
	switch im := idm.Pam.Layer().(type) {
	case *nasdaq.IttoMessageAddOrder:
		outItto("ORDER", im.Type, "%c %012x %016x %08x %08x", sideChar(im.Side), im.OId.ToUint64(), im.RefNumD.ToUint32(), im.Size, im.Price)
	case *nasdaq.IttoMessageAddQuote:
		outItto("QBID", im.Type, "%012x %016x %08x %08x", im.OId.ToUint64(), im.Bid.RefNumD.ToUint32(), im.Bid.Size, im.Bid.Price)
		outItto("QASK", im.Type, "%012x %016x %08x %08x", im.OId.ToUint64(), im.Ask.RefNumD.ToUint32(), im.Ask.Size, im.Ask.Price)
	case *nasdaq.IttoMessageSingleSideExecuted:
		outItto("ORDER", im.Type, "%016x %08x", im.OrigRefNumD.ToUint32(), im.Size)
	case *nasdaq.IttoMessageSingleSideExecutedWithPrice:
		outItto("ORDER", im.Type, "%016x %08x", im.OrigRefNumD.ToUint32(), im.Size)
	case *nasdaq.IttoMessageOrderCancel:
		outItto("ORDER", im.Type, "%016x %08x", im.OrigRefNumD.ToUint32(), im.Size)
	case *nasdaq.IttoMessageSingleSideReplace:
		outItto("ORDER", im.Type, "%016x %016x %08x %08x", im.RefNumD.ToUint32(), im.OrigRefNumD.ToUint32(), im.Size, im.Price)
	case *nasdaq.IttoMessageSingleSideDelete:
		outItto("ORDER", im.Type, "%016x", im.OrigRefNumD.ToUint32())
	case *nasdaq.IttoMessageSingleSideUpdate:
		outItto("ORDER", im.Type, "%016x %08x %08x", im.RefNumD.ToUint32(), im.Size, im.Price)
	case *nasdaq.IttoMessageQuoteReplace:
		outItto("QBID", im.Type, "%016x %016x %08x %08x", im.Bid.RefNumD.ToUint32(), im.Bid.OrigRefNumD.ToUint32(), im.Bid.Size, im.Bid.Price)
		outItto("QASK", im.Type, "%016x %016x %08x %08x", im.Ask.RefNumD.ToUint32(), im.Ask.OrigRefNumD.ToUint32(), im.Ask.Size, im.Ask.Price)
	case *nasdaq.IttoMessageQuoteDelete:
		outItto("QBID", im.Type, "%016x", im.BidOrigRefNumD.ToUint32())
		outItto("QASK", im.Type, "%016x", im.AskOrigRefNumD.ToUint32())
	case *nasdaq.IttoMessageBlockSingleSideDelete:
		for _, r := range im.RefNumDs {
			outItto("ORDER", im.Type, "%016x", r.ToUint32())
		}

	case *bats.PitchMessageAddOrder:
		outBats("%c %012x %016x %08x %08x", sideChar(im.Side), im.Symbol.ToUint64(), im.OrderId.ToUint64(), im.Size, packet.PriceTo4Dec(im.Price))
	case *bats.PitchMessageDeleteOrder:
		outBats("%016x", im.OrderId.ToUint64())
	case *bats.PitchMessageOrderExecuted:
		outBats("%016x %08x", im.OrderId.ToUint64(), im.Size)
	case *bats.PitchMessageOrderExecutedAtPriceSize:
		outBats("%016x %08x", im.OrderId.ToUint64(), im.Size)
	case *bats.PitchMessageReduceSize:
		outBats("%016x %08x", im.OrderId.ToUint64(), im.Size)
	case *bats.PitchMessageModifyOrder:
		outBats("%016x %08x %08x", im.OrderId.ToUint64(), im.Size, packet.PriceTo4Dec(im.Price))
	case *miax.TomMessageTom:
		outMiax("%c %08x %08x %08x %08x", sideChar(im.Side), im.ProductId.ToUint32(), packet.PriceTo4Dec(im.Price), im.Size, im.PriorityCustomerSize)
	}
	s.efhLogger.MessageArrived(idm)
}
func (s *SimLogger) OperationAppliedToOrders(operation sim.SimOperation) {
	type ordrespLogInfo struct {
		notFound, addOp   int
		orderId           packet.OrderId
		optionId          packet.OptionId
		side, price, size int
		ordlSuffix        string
	}
	type orduLogInfo struct {
		orderId           packet.OrderId
		optionId          packet.OptionId
		side, price, size int
	}

	var or ordrespLogInfo
	var ou orduLogInfo
	if _, ok := operation.(*sim.OperationTop); ok {
		return
	} else if op, ok := operation.(*sim.OperationAdd); ok {
		var oid packet.OptionId
		if op.Independent() {
			oid = op.GetOptionId()
		}
		or = ordrespLogInfo{
			addOp:      1,
			orderId:    op.OrderId,
			optionId:   oid,
			ordlSuffix: fmt.Sprintf(" %012x", oid.ToUint64()),
		}
		ou = orduLogInfo{
			orderId:  or.orderId,
			optionId: op.GetOptionId(),
			price:    op.GetPrice(),
			size:     op.GetNewSize(sim.SizeKindDefault),
		}
		if op.GetSide() == packet.MarketSideAsk {
			ou.side = 1
		}
	} else {
		if operation.GetOptionId().Invalid() {
			or = ordrespLogInfo{notFound: 1}
		} else {
			newSize := operation.GetNewSize(sim.SizeKindDefault)
			or = ordrespLogInfo{
				optionId: operation.GetOptionId(),
				price:    operation.GetPrice(),
				size:     newSize - operation.GetDefaultSizeDelta(),
			}
			if operation.GetSide() == packet.MarketSideAsk {
				or.side = 1
			}
			if newSize != 0 {
				ou = orduLogInfo{
					optionId: or.optionId,
					side:     or.side,
					price:    or.price,
					size:     newSize,
				}
			}
		}
		or.orderId = operation.GetOrigOrderId()
		ou.orderId = or.orderId
	}
	s.printfln("ORDL %d %016x%s", or.addOp, or.orderId.ToUint64(), or.ordlSuffix)
	s.printfln("ORDRESP %d %d %d %08x %08x %012x %016x", or.notFound, or.addOp, or.side, or.size, or.price, or.optionId.ToUint64(), or.orderId.ToUint64())
	if operation.GetOptionId().Valid() {
		s.printfln("ORDU %016x %012x %d %08x %08x", ou.orderId.ToUint64(), ou.optionId.ToUint64(), ou.side, ou.price, ou.size)
	}
}
func (s *SimLogger) BeforeBookUpdate(book sim.Book, operation sim.SimOperation) {
	tobOld := book.GetTop(operation.GetOptionId(), operation.GetSide(), s.supernodeLevels)
	s.tobOld = make([]sim.PriceLevel, len(tobOld))
	for i, pl := range tobOld {
		s.tobOld[i] = pl.Clone()
	}
	s.efhLogger.BeforeBookUpdate(book, operation)
}
func (s *SimLogger) AfterBookUpdate(book sim.Book, operation sim.SimOperation) {
	if operation.GetOptionId().Valid() && s.supernodeLevels > 1 {
		var emptyPrice uint32
		if operation.GetSide() == packet.MarketSideAsk {
			emptyPrice -= 1
		}
		printablePriceLevel := func(pls []sim.PriceLevel, pos int) (price uint32, size int) {
			if pos < len(pls) {
				price = uint32(pls[pos].Price())
				size = pls[pos].Size(sim.SizeKindDefault)
			} else if operation.GetSide() == packet.MarketSideAsk {
				price = emptyPrice
			}
			return
		}
		s.tobNew = book.GetTop(operation.GetOptionId(), operation.GetSide(), s.supernodeLevels)
		for i := 0; i < s.accessedLevels(operation); i++ {
			priceOld, sizeOld := printablePriceLevel(s.tobOld, i)
			priceNew, sizeNew := printablePriceLevel(s.tobNew, i)
			s.printfln("SN_OLD_NEW %02d %08x %08x  %08x %08x", i, sizeOld, priceOld, sizeNew, priceNew)
		}
	}
	s.efhLogger.AfterBookUpdate(book, operation)
}
func (s *SimLogger) accessedLevels(operation sim.SimOperation) (levels int) {
	levels = SimLoggerUpperSupernodeLevels
	if s.supernodeLevels <= levels {
		return s.supernodeLevels
	}
	if operation.GetPrice() == 0 {
		// TODO hw can skip SN access at all
		return
	}
	lenOld, lenNew := len(s.tobOld), len(s.tobNew)
	if lenOld < levels {
		return
	}
	if lenOld == lenNew {
		for i := 0; i < levels; i++ {
			if !s.tobOld[i].Equals(s.tobNew[i]) {
				return
			}
		}
	}
	return s.supernodeLevels
}

func (s *SimLogger) PrintMessage(m efhMessage) (err error) {
	defer errs.PassE(&err)
	var bb bytes.Buffer
	errs.CheckE(binary.Write(&bb, binary.LittleEndian, m))
	if r := bb.Len() % 8; r > 0 {
		// pad to  multiple of 8 bytes
		z := make([]byte, 8)
		_, err = bb.Write(z[0 : 8-r])
		errs.CheckE(err)
	}

	for {
		var qw uint64
		if err := binary.Read(&bb, binary.LittleEndian, &qw); err != nil {
			if err == io.EOF {
				break
			}
			errs.CheckE(err)
		} else {
			s.printfln("DMATOHOST_DATA %016x", qw)
		}
	}
	s.printfln("DMATOHOST_TRAILER 00656e696c616b45")
	return
}
<commit_msg>rec:SimLogger: add TomMessageQuote<commit_after>// Copyright (c) Ilia Kravets, 2015. All rights reserved. PROVIDED "AS IS"
// WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED. See LICENSE file for details.

package rec

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"

	"github.com/ikravets/errs"

	"my/ev/packet"
	"my/ev/packet/bats"
	"my/ev/packet/miax"
	"my/ev/packet/nasdaq"
	"my/ev/sim"
)

type SimLoggerConfig struct {
	EfhLoggerConfig
	SupernodeLevels int
}

type SimLogger struct {
	w               io.Writer
	tobOld, tobNew  []sim.PriceLevel
	efhLogger       EfhLogger
	supernodeLevels int
}

const SimLoggerDefaultSupernodeLevels = 256
const SimLoggerUpperSupernodeLevels = 8

func NewSimLogger(c SimLoggerConfig) *SimLogger {
	s := &SimLogger{
		w:               c.Writer,
		supernodeLevels: c.SupernodeLevels,
	}
	if s.supernodeLevels == 0 {
		s.supernodeLevels = SimLoggerDefaultSupernodeLevels
	}
	elc := c.EfhLoggerConfig
	elc.Printer = s
	elc.Writer = nil
	s.efhLogger = *NewEfhLogger(elc)
	return s
}
func (s *SimLogger) SetSupernodeLevels(levels int) {
	errs.Check(levels > 0)
	s.supernodeLevels = levels
}

func (s *SimLogger) printf(format string, vs ...interface{}) {
	_, err := fmt.Fprintf(s.w, format, vs...)
	errs.CheckE(err)
}
func (s *SimLogger) printfln(format string, vs ...interface{}) {
	f := format + "\n"
	s.printf(f, vs...)
}
func (s *SimLogger) MessageArrived(idm *sim.SimMessage) {
	outItto := func(name string, typ nasdaq.IttoMessageType, f string, vs ...interface{}) {
		s.printf("NORM %s %c ", name, typ)
		s.printfln(f, vs...)
	}
	outBats := func(f string, vs ...interface{}) {
		s.printf("NORM ORDER %02x ", idm.Pam.Layer().(bats.PitchMessage).Base().Type.ToInt())
		s.printfln(f, vs...)
	}
	outMiax := func(name string, f string, vs ...interface{}) {
		s.printf("NORM %s %02x ", name, idm.Pam.Layer().(miax.TomMessage).Base().Type.ToInt())
		s.printfln(f, vs...)
	}
	sideChar := func(s packet.MarketSide) byte {
		if s == packet.MarketSideAsk {
			return 'S'
		}
		return byte(s)
	}
	switch im := idm.Pam.Layer().(type) {
	case *nasdaq.IttoMessageAddOrder:
		outItto("ORDER", im.Type, "%c %012x %016x %08x %08x", sideChar(im.Side), im.OId.ToUint64(), im.RefNumD.ToUint32(), im.Size, im.Price)
	case *nasdaq.IttoMessageAddQuote:
		outItto("QBID", im.Type, "%012x %016x %08x %08x", im.OId.ToUint64(), im.Bid.RefNumD.ToUint32(), im.Bid.Size, im.Bid.Price)
		outItto("QASK", im.Type, "%012x %016x %08x %08x", im.OId.ToUint64(), im.Ask.RefNumD.ToUint32(), im.Ask.Size, im.Ask.Price)
	case *nasdaq.IttoMessageSingleSideExecuted:
		outItto("ORDER", im.Type, "%016x %08x", im.OrigRefNumD.ToUint32(), im.Size)
	case *nasdaq.IttoMessageSingleSideExecutedWithPrice:
		outItto("ORDER", im.Type, "%016x %08x", im.OrigRefNumD.ToUint32(), im.Size)
	case *nasdaq.IttoMessageOrderCancel:
		outItto("ORDER", im.Type, "%016x %08x", im.OrigRefNumD.ToUint32(), im.Size)
	case *nasdaq.IttoMessageSingleSideReplace:
		outItto("ORDER", im.Type, "%016x %016x %08x %08x", im.RefNumD.ToUint32(), im.OrigRefNumD.ToUint32(), im.Size, im.Price)
	case *nasdaq.IttoMessageSingleSideDelete:
		outItto("ORDER", im.Type, "%016x", im.OrigRefNumD.ToUint32())
	case *nasdaq.IttoMessageSingleSideUpdate:
		outItto("ORDER", im.Type, "%016x %08x %08x", im.RefNumD.ToUint32(), im.Size, im.Price)
	case *nasdaq.IttoMessageQuoteReplace:
		outItto("QBID", im.Type, "%016x %016x %08x %08x", im.Bid.RefNumD.ToUint32(), im.Bid.OrigRefNumD.ToUint32(), im.Bid.Size, im.Bid.Price)
		outItto("QASK", im.Type, "%016x %016x %08x %08x", im.Ask.RefNumD.ToUint32(), im.Ask.OrigRefNumD.ToUint32(), im.Ask.Size, im.Ask.Price)
	case *nasdaq.IttoMessageQuoteDelete:
		outItto("QBID", im.Type, "%016x", im.BidOrigRefNumD.ToUint32())
		outItto("QASK", im.Type, "%016x", im.AskOrigRefNumD.ToUint32())
	case *nasdaq.IttoMessageBlockSingleSideDelete:
		for _, r := range im.RefNumDs {
			outItto("ORDER", im.Type, "%016x", r.ToUint32())
		}

	case *bats.PitchMessageAddOrder:
		outBats("%c %012x %016x %08x %08x", sideChar(im.Side), im.Symbol.ToUint64(), im.OrderId.ToUint64(), im.Size, packet.PriceTo4Dec(im.Price))
	case *bats.PitchMessageDeleteOrder:
		outBats("%016x", im.OrderId.ToUint64())
	case *bats.PitchMessageOrderExecuted:
		outBats("%016x %08x", im.OrderId.ToUint64(), im.Size)
	case *bats.PitchMessageOrderExecutedAtPriceSize:
		outBats("%016x %08x", im.OrderId.ToUint64(), im.Size)
	case *bats.PitchMessageReduceSize:
		outBats("%016x %08x", im.OrderId.ToUint64(), im.Size)
	case *bats.PitchMessageModifyOrder:
		outBats("%016x %08x %08x", im.OrderId.ToUint64(), im.Size, packet.PriceTo4Dec(im.Price))
	case *miax.TomMessageTom:
		outMiax("TOM", "%c %08x %08x %08x %08x", sideChar(im.Side), im.ProductId.ToUint32(), packet.PriceTo4Dec(im.Price), im.Size, im.PriorityCustomerSize)
	case *miax.TomMessageQuote:
		outMiax("QBID", "%08x %08x %08x %08x", im.ProductId.ToUint32(), packet.PriceTo4Dec(im.BidPrice), im.BidSize, im.BidPriorityCustomerSize)
		outMiax("QASK", "%08x %08x %08x %08x", im.ProductId.ToUint32(), packet.PriceTo4Dec(im.OfferPrice), im.OfferSize, im.OfferPriorityCustomerSize)
	}
	s.efhLogger.MessageArrived(idm)
}
func (s *SimLogger) OperationAppliedToOrders(operation sim.SimOperation) {
	type ordrespLogInfo struct {
		notFound, addOp   int
		orderId           packet.OrderId
		optionId          packet.OptionId
		side, price, size int
		ordlSuffix        string
	}
	type orduLogInfo struct {
		orderId           packet.OrderId
		optionId          packet.OptionId
		side, price, size int
	}

	var or ordrespLogInfo
	var ou orduLogInfo
	if _, ok := operation.(*sim.OperationTop); ok {
		return
	} else if op, ok := operation.(*sim.OperationAdd); ok {
		var oid packet.OptionId
		if op.Independent() {
			oid = op.GetOptionId()
		}
		or = ordrespLogInfo{
			addOp:      1,
			orderId:    op.OrderId,
			optionId:   oid,
			ordlSuffix: fmt.Sprintf(" %012x", oid.ToUint64()),
		}
		ou = orduLogInfo{
			orderId:  or.orderId,
			optionId: op.GetOptionId(),
			price:    op.GetPrice(),
			size:     op.GetNewSize(sim.SizeKindDefault),
		}
		if op.GetSide() == packet.MarketSideAsk {
			ou.side = 1
		}
	} else {
		if operation.GetOptionId().Invalid() {
			or = ordrespLogInfo{notFound: 1}
		} else {
			newSize := operation.GetNewSize(sim.SizeKindDefault)
			or = ordrespLogInfo{
				optionId: operation.GetOptionId(),
				price:    operation.GetPrice(),
				size:     newSize - operation.GetDefaultSizeDelta(),
			}
			if operation.GetSide() == packet.MarketSideAsk {
				or.side = 1
			}
			if newSize != 0 {
				ou = orduLogInfo{
					optionId: or.optionId,
					side:     or.side,
					price:    or.price,
					size:     newSize,
				}
			}
		}
		or.orderId = operation.GetOrigOrderId()
		ou.orderId = or.orderId
	}
	s.printfln("ORDL %d %016x%s", or.addOp, or.orderId.ToUint64(), or.ordlSuffix)
	s.printfln("ORDRESP %d %d %d %08x %08x %012x %016x", or.notFound, or.addOp, or.side, or.size, or.price, or.optionId.ToUint64(), or.orderId.ToUint64())
	if operation.GetOptionId().Valid() {
		s.printfln("ORDU %016x %012x %d %08x %08x", ou.orderId.ToUint64(), ou.optionId.ToUint64(), ou.side, ou.price, ou.size)
	}
}
func (s *SimLogger) BeforeBookUpdate(book sim.Book, operation sim.SimOperation) {
	tobOld := book.GetTop(operation.GetOptionId(), operation.GetSide(), s.supernodeLevels)
	s.tobOld = make([]sim.PriceLevel, len(tobOld))
	for i, pl := range tobOld {
		s.tobOld[i] = pl.Clone()
	}
	s.efhLogger.BeforeBookUpdate(book, operation)
}
func (s *SimLogger) AfterBookUpdate(book sim.Book, operation sim.SimOperation) {
	if operation.GetOptionId().Valid() && s.supernodeLevels > 1 {
		var emptyPrice uint32
		if operation.GetSide() == packet.MarketSideAsk {
			emptyPrice -= 1
		}
		printablePriceLevel := func(pls []sim.PriceLevel, pos int) (price uint32, size int) {
			if pos < len(pls) {
				price = uint32(pls[pos].Price())
				size = pls[pos].Size(sim.SizeKindDefault)
			} else if operation.GetSide() == packet.MarketSideAsk {
				price = emptyPrice
			}
			return
		}
		s.tobNew = book.GetTop(operation.GetOptionId(), operation.GetSide(), s.supernodeLevels)
		for i := 0; i < s.accessedLevels(operation); i++ {
			priceOld, sizeOld := printablePriceLevel(s.tobOld, i)
			priceNew, sizeNew := printablePriceLevel(s.tobNew, i)
			s.printfln("SN_OLD_NEW %02d %08x %08x  %08x %08x", i, sizeOld, priceOld, sizeNew, priceNew)
		}
	}
	s.efhLogger.AfterBookUpdate(book, operation)
}
func (s *SimLogger) accessedLevels(operation sim.SimOperation) (levels int) {
	levels = SimLoggerUpperSupernodeLevels
	if s.supernodeLevels <= levels {
		return s.supernodeLevels
	}
	if operation.GetPrice() == 0 {
		// TODO hw can skip SN access at all
		return
	}
	lenOld, lenNew := len(s.tobOld), len(s.tobNew)
	if lenOld < levels {
		return
	}
	if lenOld == lenNew {
		for i := 0; i < levels; i++ {
			if !s.tobOld[i].Equals(s.tobNew[i]) {
				return
			}
		}
	}
	return s.supernodeLevels
}

func (s *SimLogger) PrintMessage(m efhMessage) (err error) {
	defer errs.PassE(&err)
	var bb bytes.Buffer
	errs.CheckE(binary.Write(&bb, binary.LittleEndian, m))
	if r := bb.Len() % 8; r > 0 {
		// pad to  multiple of 8 bytes
		z := make([]byte, 8)
		_, err = bb.Write(z[0 : 8-r])
		errs.CheckE(err)
	}

	for {
		var qw uint64
		if err := binary.Read(&bb, binary.LittleEndian, &qw); err != nil {
			if err == io.EOF {
				break
			}
			errs.CheckE(err)
		} else {
			s.printfln("DMATOHOST_DATA %016x", qw)
		}
	}
	s.printfln("DMATOHOST_TRAILER 00656e696c616b45")
	return
}
<|endoftext|>


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

import (
	"bytes"
	"github.com/snowplow/sql-runner/playbook"
	"io/ioutil"
	"text/template"
	"time"
)

var (
	templFuncs = template.FuncMap{
		"nowWithFormat": func(format string) string {
			return time.Now().Format(format)
		},
		"awsChainCredentials":   awsChainCredentials,
		"awsEC2RoleCredentials": awsEC2RoleCredentials,
		"awsEnvCredentials":     awsEnvCredentials,
		"awsProfileCredentials": awsProfileCredentials,
	}
)

// Generalized interface to a database client
type Db interface {
	RunQuery(playbook.Query, string, map[string]interface{}) QueryStatus
	GetTarget() playbook.Target
}

// Reads the script and fills in the template
func prepareQuery(queryPath string, template bool, variables map[string]interface{}) (string, error) {

	script, err := readScript(queryPath)
	if err != nil {
		return "", err
	}

	if template {
		script, err = fillTemplate(script, variables) // Yech, mutate
		if err != nil {
			return "", err
		}
	}
	return script, nil
}

// Reads the file ready for executing
func readScript(file string) (string, error) {

	scriptBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return "", err
	}
	return string(scriptBytes), nil
}

// Fills in a script which is a template
func fillTemplate(script string, variables map[string]interface{}) (string, error) {

	t, err := template.New("playbook").Funcs(templFuncs).Parse(script)
	if err != nil {
		return "", err
	}

	var filled bytes.Buffer
	if err := t.Execute(&filled, variables); err != nil {
		return "", err
	}
	return filled.String(), nil
}
<commit_msg>Add an environment variable function (closes #30)<commit_after>//
// Copyright (c) 2015 Snowplow Analytics Ltd. All rights reserved.
//
// This program is licensed to you under the Apache License Version 2.0,
// and you may not use this file except in compliance with the Apache License Version 2.0.
// You may obtain a copy of the Apache License Version 2.0 at http://www.apache.org/licenses/LICENSE-2.0.
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the Apache License Version 2.0 is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the Apache License Version 2.0 for the specific language governing permissions and limitations there under.
//
package run

import (
	"bytes"
	"github.com/snowplow/sql-runner/playbook"
	"io/ioutil"
	"text/template"
	"time"
	"os"
)

var (
	templFuncs = template.FuncMap{
		"nowWithFormat": func(format string) string {
			return time.Now().Format(format)
		},
		"systemEnv": func(env string) string {
			return os.Getenv(env)
		},
		"awsChainCredentials":   awsChainCredentials,
		"awsEC2RoleCredentials": awsEC2RoleCredentials,
		"awsEnvCredentials":     awsEnvCredentials,
		"awsProfileCredentials": awsProfileCredentials,
	}
)

// Generalized interface to a database client
type Db interface {
	RunQuery(playbook.Query, string, map[string]interface{}) QueryStatus
	GetTarget() playbook.Target
}

// Reads the script and fills in the template
func prepareQuery(queryPath string, template bool, variables map[string]interface{}) (string, error) {

	script, err := readScript(queryPath)
	if err != nil {
		return "", err
	}

	if template {
		script, err = fillTemplate(script, variables) // Yech, mutate
		if err != nil {
			return "", err
		}
	}
	return script, nil
}

// Reads the file ready for executing
func readScript(file string) (string, error) {

	scriptBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return "", err
	}
	return string(scriptBytes), nil
}

// Fills in a script which is a template
func fillTemplate(script string, variables map[string]interface{}) (string, error) {

	t, err := template.New("playbook").Funcs(templFuncs).Parse(script)
	if err != nil {
		return "", err
	}

	var filled bytes.Buffer
	if err := t.Execute(&filled, variables); err != nil {
		return "", err
	}
	return filled.String(), nil
}
<|endoftext|>


<commit_before>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin dragonfly freebsd linux,!android netbsd openbsd
// +build cgo

// Note that this test does not work on Solaris: issue #22849.
// Don't run the test on Android because at least some versions of the
// C library do not define the posix_openpt function.

package signal_test

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal/internal/pty"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"
)

func TestTerminalSignal(t *testing.T) {
	const enteringRead = "test program entering read"
	if os.Getenv("GO_TEST_TERMINAL_SIGNALS") != "" {
		var b [1]byte
		fmt.Println(enteringRead)
		n, err := os.Stdin.Read(b[:])
		if n == 1 {
			if b[0] == '\n' {
				// This is what we expect
				fmt.Println("read newline")
			} else {
				fmt.Printf("read 1 byte: %q\n", b)
			}
		} else {
			fmt.Printf("read %d bytes\n", n)
		}
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	t.Parallel()

	// The test requires a shell that uses job control.
	bash, err := exec.LookPath("bash")
	if err != nil {
		t.Skipf("could not find bash: %v", err)
	}

	scale := 1
	if s := os.Getenv("GO_TEST_TIMEOUT_SCALE"); s != "" {
		if sc, err := strconv.Atoi(s); err == nil {
			scale = sc
		}
	}
	pause := time.Duration(scale) * 10 * time.Millisecond
	wait := time.Duration(scale) * 5 * time.Second

	// The test only fails when using a "slow device," in this
	// case a pseudo-terminal.

	master, sname, err := pty.Open()
	if err != nil {
		ptyErr := err.(*pty.PtyError)
		if ptyErr.FuncName == "posix_openpt" && ptyErr.Errno == syscall.EACCES {
			t.Skip("posix_openpt failed with EACCES, assuming chroot and skipping")
		}
		t.Fatal(err)
	}
	defer master.Close()
	slave, err := os.OpenFile(sname, os.O_RDWR, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer slave.Close()

	// Start an interactive shell.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, bash, "--norc", "--noprofile", "-i")
	cmd.Stdin = slave
	cmd.Stdout = slave
	cmd.Stderr = slave
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid:  true,
		Setctty: true,
		Ctty:    int(slave.Fd()),
	}

	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}

	if err := slave.Close(); err != nil {
		t.Errorf("closing slave: %v", err)
	}

	progReady := make(chan bool)
	sawPrompt := make(chan bool, 10)
	const prompt = "prompt> "

	// Read data from master in the background.
	go func() {
		input := bufio.NewReader(master)
		var line, handled []byte
		for {
			b, err := input.ReadByte()
			if err != nil {
				if len(line) > 0 || len(handled) > 0 {
					t.Logf("%q", append(handled, line...))
				}
				if perr, ok := err.(*os.PathError); ok {
					err = perr.Err
				}
				// EOF means master is closed.
				// EIO means child process is done.
				// "file already closed" means deferred close of master has happened.
				if err != io.EOF && err != syscall.EIO && !strings.Contains(err.Error(), "file already closed") {
					t.Logf("error reading from master: %v", err)
				}
				return
			}

			line = append(line, b)

			if b == '\n' {
				t.Logf("%q", append(handled, line...))
				line = nil
				handled = nil
				continue
			}

			if bytes.Contains(line, []byte(enteringRead)) {
				close(progReady)
				handled = append(handled, line...)
				line = nil
			} else if bytes.Contains(line, []byte(prompt)) && !bytes.Contains(line, []byte("PS1=")) {
				sawPrompt <- true
				handled = append(handled, line...)
				line = nil
			}
		}
	}()

	// Set the bash prompt so that we can see it.
	if _, err := master.Write([]byte("PS1='" + prompt + "'\n")); err != nil {
		t.Fatalf("setting prompt: %v", err)
	}
	select {
	case <-sawPrompt:
	case <-time.After(wait):
		t.Fatal("timed out waiting for shell prompt")
	}

	// Start a small program that reads from stdin
	// (namely the code at the top of this function).
	if _, err := master.Write([]byte("GO_TEST_TERMINAL_SIGNALS=1 " + os.Args[0] + " -test.run=TestTerminalSignal\n")); err != nil {
		t.Fatal(err)
	}

	// Wait for the program to print that it is starting.
	select {
	case <-progReady:
	case <-time.After(wait):
		t.Fatal("timed out waiting for program to start")
	}

	// Give the program time to enter the read call.
	// It doesn't matter much if we occasionally don't wait long enough;
	// we won't be testing what we want to test, but the overall test
	// will pass.
	time.Sleep(pause)

	// Send a ^Z to stop the program.
	if _, err := master.Write([]byte{26}); err != nil {
		t.Fatalf("writing ^Z to pty: %v", err)
	}

	// Wait for the program to stop and return to the shell.
	select {
	case <-sawPrompt:
	case <-time.After(wait):
		t.Fatal("timed out waiting for shell prompt")
	}

	// Restart the stopped program.
	if _, err := master.Write([]byte("fg\n")); err != nil {
		t.Fatalf("writing %q to pty: %v", "fg", err)
	}

	// Give the process time to restart.
	// This is potentially racy: if the process does not restart
	// quickly enough then the byte we send will go to bash rather
	// than the program. Unfortunately there isn't anything we can
	// look for to know that the program is running again.
	// bash will print the program name, but that happens before it
	// restarts the program.
	time.Sleep(10 * pause)

	// Write some data for the program to read,
	// which should cause it to exit.
	if _, err := master.Write([]byte{'\n'}); err != nil {
		t.Fatalf("writing %q to pty: %v", "\n", err)
	}

	// Wait for the program to exit.
	select {
	case <-sawPrompt:
	case <-time.After(wait):
		t.Fatal("timed out waiting for shell prompt")
	}

	// Exit the shell with the program's exit status.
	if _, err := master.Write([]byte("exit $?\n")); err != nil {
		t.Fatalf("writing %q to pty: %v", "exit", err)
	}

	if err = cmd.Wait(); err != nil {
		t.Errorf("subprogram failed: %v", err)
	}
}
<commit_msg>os/signal: disable loading of history during test<commit_after>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin dragonfly freebsd linux,!android netbsd openbsd
// +build cgo

// Note that this test does not work on Solaris: issue #22849.
// Don't run the test on Android because at least some versions of the
// C library do not define the posix_openpt function.

package signal_test

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal/internal/pty"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"
)

func TestTerminalSignal(t *testing.T) {
	const enteringRead = "test program entering read"
	if os.Getenv("GO_TEST_TERMINAL_SIGNALS") != "" {
		var b [1]byte
		fmt.Println(enteringRead)
		n, err := os.Stdin.Read(b[:])
		if n == 1 {
			if b[0] == '\n' {
				// This is what we expect
				fmt.Println("read newline")
			} else {
				fmt.Printf("read 1 byte: %q\n", b)
			}
		} else {
			fmt.Printf("read %d bytes\n", n)
		}
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	t.Parallel()

	// The test requires a shell that uses job control.
	bash, err := exec.LookPath("bash")
	if err != nil {
		t.Skipf("could not find bash: %v", err)
	}

	scale := 1
	if s := os.Getenv("GO_TEST_TIMEOUT_SCALE"); s != "" {
		if sc, err := strconv.Atoi(s); err == nil {
			scale = sc
		}
	}
	pause := time.Duration(scale) * 10 * time.Millisecond
	wait := time.Duration(scale) * 5 * time.Second

	// The test only fails when using a "slow device," in this
	// case a pseudo-terminal.

	master, sname, err := pty.Open()
	if err != nil {
		ptyErr := err.(*pty.PtyError)
		if ptyErr.FuncName == "posix_openpt" && ptyErr.Errno == syscall.EACCES {
			t.Skip("posix_openpt failed with EACCES, assuming chroot and skipping")
		}
		t.Fatal(err)
	}
	defer master.Close()
	slave, err := os.OpenFile(sname, os.O_RDWR, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer slave.Close()

	// Start an interactive shell.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, bash, "--norc", "--noprofile", "-i")
	// Clear HISTFILE so that we don't read or clobber the user's bash history.
	cmd.Env = append(os.Environ(), "HISTFILE=")
	cmd.Stdin = slave
	cmd.Stdout = slave
	cmd.Stderr = slave
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid:  true,
		Setctty: true,
		Ctty:    int(slave.Fd()),
	}

	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}

	if err := slave.Close(); err != nil {
		t.Errorf("closing slave: %v", err)
	}

	progReady := make(chan bool)
	sawPrompt := make(chan bool, 10)
	const prompt = "prompt> "

	// Read data from master in the background.
	go func() {
		input := bufio.NewReader(master)
		var line, handled []byte
		for {
			b, err := input.ReadByte()
			if err != nil {
				if len(line) > 0 || len(handled) > 0 {
					t.Logf("%q", append(handled, line...))
				}
				if perr, ok := err.(*os.PathError); ok {
					err = perr.Err
				}
				// EOF means master is closed.
				// EIO means child process is done.
				// "file already closed" means deferred close of master has happened.
				if err != io.EOF && err != syscall.EIO && !strings.Contains(err.Error(), "file already closed") {
					t.Logf("error reading from master: %v", err)
				}
				return
			}

			line = append(line, b)

			if b == '\n' {
				t.Logf("%q", append(handled, line...))
				line = nil
				handled = nil
				continue
			}

			if bytes.Contains(line, []byte(enteringRead)) {
				close(progReady)
				handled = append(handled, line...)
				line = nil
			} else if bytes.Contains(line, []byte(prompt)) && !bytes.Contains(line, []byte("PS1=")) {
				sawPrompt <- true
				handled = append(handled, line...)
				line = nil
			}
		}
	}()

	// Set the bash prompt so that we can see it.
	if _, err := master.Write([]byte("PS1='" + prompt + "'\n")); err != nil {
		t.Fatalf("setting prompt: %v", err)
	}
	select {
	case <-sawPrompt:
	case <-time.After(wait):
		t.Fatal("timed out waiting for shell prompt")
	}

	// Start a small program that reads from stdin
	// (namely the code at the top of this function).
	if _, err := master.Write([]byte("GO_TEST_TERMINAL_SIGNALS=1 " + os.Args[0] + " -test.run=TestTerminalSignal\n")); err != nil {
		t.Fatal(err)
	}

	// Wait for the program to print that it is starting.
	select {
	case <-progReady:
	case <-time.After(wait):
		t.Fatal("timed out waiting for program to start")
	}

	// Give the program time to enter the read call.
	// It doesn't matter much if we occasionally don't wait long enough;
	// we won't be testing what we want to test, but the overall test
	// will pass.
	time.Sleep(pause)

	// Send a ^Z to stop the program.
	if _, err := master.Write([]byte{26}); err != nil {
		t.Fatalf("writing ^Z to pty: %v", err)
	}

	// Wait for the program to stop and return to the shell.
	select {
	case <-sawPrompt:
	case <-time.After(wait):
		t.Fatal("timed out waiting for shell prompt")
	}

	// Restart the stopped program.
	if _, err := master.Write([]byte("fg\n")); err != nil {
		t.Fatalf("writing %q to pty: %v", "fg", err)
	}

	// Give the process time to restart.
	// This is potentially racy: if the process does not restart
	// quickly enough then the byte we send will go to bash rather
	// than the program. Unfortunately there isn't anything we can
	// look for to know that the program is running again.
	// bash will print the program name, but that happens before it
	// restarts the program.
	time.Sleep(10 * pause)

	// Write some data for the program to read,
	// which should cause it to exit.
	if _, err := master.Write([]byte{'\n'}); err != nil {
		t.Fatalf("writing %q to pty: %v", "\n", err)
	}

	// Wait for the program to exit.
	select {
	case <-sawPrompt:
	case <-time.After(wait):
		t.Fatal("timed out waiting for shell prompt")
	}

	// Exit the shell with the program's exit status.
	if _, err := master.Write([]byte("exit $?\n")); err != nil {
		t.Fatalf("writing %q to pty: %v", "exit", err)
	}

	if err = cmd.Wait(); err != nil {
		t.Errorf("subprogram failed: %v", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.
// Author: Robert Tingirica

package symlink

import (
	"errors"
	"os"
	"strings"
	"syscall"
	"unicode/utf16"
)

const (
	SYMBOLIC_LINK_FLAG_DIRECTORY = 1
)

//sys createSymbolicLink(symlinkname *uint16, targetname *uint16, flags uint32) (err error) = CreateSymbolicLinkW
//sys getFinalPathNameByHandle(handle Handle, buf *uint16, buflen uint32, flags uint32) (n uint32, err error) = GetFinalPathNameByHandleW

// New creates newname as a symbolic link to oldname.
// If there is an error, it will be of type *LinkError.
func New(oldname, newname string) error {
	fi, err := os.Stat(oldname)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}
	var flag uint32
	if fi.IsDir() {
		flag = SYMBOLIC_LINK_FLAG_DIRECTORY
	}

	targetp, err := getLongPath(oldname)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}

	linkp, err := syscall.UTF16PtrFromString(newname)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}

	err = createSymbolicLink(linkp, &targetp[0], flag)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}
	return nil
}

// Read returns the destination of the named symbolic link.
// If there is an error, it will be of type *PathError.
func Read(link string) (string, error) {
	linkp, err := getLongPath(link)
	if err != nil {
		return "", err
	}
	h, err := syscall.CreateFile(
		&linkp[0],
		syscall.GENERIC_READ,
		syscall.FILE_SHARE_READ,
		nil,
		syscall.OPEN_EXISTING,
		33554432, // for some reason, syscall.GENERIC_EXECUTE results in "Access Denied" error
		0)
	if err != nil {
		return "", &os.PathError{"readlink", link, err}
	}
	defer syscall.CloseHandle(h)

	pathw := make([]uint16, syscall.MAX_PATH)
	n, err := getFinalPathNameByHandle(h, &pathw[0], uint32(len(pathw)), 0)
	if err != nil {
		return "", &os.PathError{"readlink", link, err}
	}
	if n > uint32(len(pathw)) {
		pathw = make([]uint16, n)
		n, err = getFinalPathNameByHandle(h, &pathw[0], uint32(len(pathw)), 0)
		if err != nil {
			return "", &os.PathError{"readlink", link, err}
		}
		if n > uint32(len(pathw)) {
			return "", &os.PathError{"readlink", link, errors.New("link length too long")}
		}
	}
	ret := string(utf16.Decode(pathw[0:n]))

	if strings.HasPrefix(ret, `\\?\`) {
		return ret[4:], nil
	}

	retp, err := getLongPath(ret)
	if err != nil {
		return "", &os.PathError{"readlink", link, err}
	}
	return syscall.UTF16ToString(retp), nil
}

// getLongPath converts windows 8.1 short style paths (c:\Progra~1\foo) to full
// long paths.
func getLongPath(path string) ([]uint16, error) {
	pathp, err := syscall.UTF16FromString(path)
	if err != nil {
		return nil, err
	}

	longp := pathp
	n, err := syscall.GetLongPathName(&pathp[0], &longp[0], uint32(len(longp)))
	if err != nil {
		return nil, err
	}
	if n > uint32(len(longp)) {
		longp = make([]uint16, n)
		n, err = syscall.GetLongPathName(&pathp[0], &longp[0], uint32(len(longp)))
		if err != nil {
			return nil, err
		}
	}
	longp = longp[:n]

	return longp, nil
}
<commit_msg>Added GENERIC_EXECUTION constant<commit_after>// Copyright 2014 Cloudbase Solutions SRL
// Licensed under the AGPLv3, see LICENCE file for details.
// Author: Robert Tingirica

package symlink

import (
	"errors"
	"os"
	"strings"
	"syscall"
	"unicode/utf16"
)

const (
	SYMBOLIC_LINK_FLAG_DIRECTORY = 1
    GENERIC_EXECUTION = 33554432
)

//sys createSymbolicLink(symlinkname *uint16, targetname *uint16, flags uint32) (err error) = CreateSymbolicLinkW
//sys getFinalPathNameByHandle(handle Handle, buf *uint16, buflen uint32, flags uint32) (n uint32, err error) = GetFinalPathNameByHandleW

// New creates newname as a symbolic link to oldname.
// If there is an error, it will be of type *LinkError.
func New(oldname, newname string) error {
	fi, err := os.Stat(oldname)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}
	var flag uint32
	if fi.IsDir() {
		flag = SYMBOLIC_LINK_FLAG_DIRECTORY
	}

	targetp, err := getLongPath(oldname)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}

	linkp, err := syscall.UTF16PtrFromString(newname)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}

	err = createSymbolicLink(linkp, &targetp[0], flag)
	if err != nil {
		return &os.LinkError{"symlink", oldname, newname, err}
	}
	return nil
}

// Read returns the destination of the named symbolic link.
// If there is an error, it will be of type *PathError.
func Read(link string) (string, error) {
	linkp, err := getLongPath(link)
	if err != nil {
		return "", err
	}
	h, err := syscall.CreateFile(
		&linkp[0],
		syscall.GENERIC_READ,
		syscall.FILE_SHARE_READ,
		nil,
		syscall.OPEN_EXISTING,
		GENERIC_EXECUTION, // for some reason, syscall.GENERIC_EXECUTE results in "Access Denied" error
		0)
	if err != nil {
		return "", &os.PathError{"readlink", link, err}
	}
	defer syscall.CloseHandle(h)

	pathw := make([]uint16, syscall.MAX_PATH)
	n, err := getFinalPathNameByHandle(h, &pathw[0], uint32(len(pathw)), 0)
	if err != nil {
		return "", &os.PathError{"readlink", link, err}
	}
	if n > uint32(len(pathw)) {
		pathw = make([]uint16, n)
		n, err = getFinalPathNameByHandle(h, &pathw[0], uint32(len(pathw)), 0)
		if err != nil {
			return "", &os.PathError{"readlink", link, err}
		}
		if n > uint32(len(pathw)) {
			return "", &os.PathError{"readlink", link, errors.New("link length too long")}
		}
	}
	ret := string(utf16.Decode(pathw[0:n]))

	if strings.HasPrefix(ret, `\\?\`) {
		return ret[4:], nil
	}

	retp, err := getLongPath(ret)
	if err != nil {
		return "", &os.PathError{"readlink", link, err}
	}
	return syscall.UTF16ToString(retp), nil
}

// getLongPath converts windows 8.1 short style paths (c:\Progra~1\foo) to full
// long paths.
func getLongPath(path string) ([]uint16, error) {
	pathp, err := syscall.UTF16FromString(path)
	if err != nil {
		return nil, err
	}

	longp := pathp
	n, err := syscall.GetLongPathName(&pathp[0], &longp[0], uint32(len(longp)))
	if err != nil {
		return nil, err
	}
	if n > uint32(len(longp)) {
		longp = make([]uint16, n)
		n, err = syscall.GetLongPathName(&pathp[0], &longp[0], uint32(len(longp)))
		if err != nil {
			return nil, err
		}
	}
	longp = longp[:n]

	return longp, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/newrelic/go-agent"
	rk "github.com/svdberg/syncmysport-runkeeper/runkeeper"
	stv "github.com/svdberg/syncmysport-runkeeper/strava"
	sync "github.com/svdberg/syncmysport-runkeeper/sync"
	shared "github.com/svdberg/syncmysport-runkeeper/syncmysport-shared"

	log "github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/bgentry/que-go"
	"github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/jackc/pgx"
)

var (
	qc      *que.Client
	pgxpool *pgx.ConnPool
	app     newrelic.Application
)

const tsDelta = -45 //minutes

func init() {
	key := os.Getenv("NEW_RELIC_KEY")
	if key != "" {
		config := newrelic.NewConfig("SyncMySport-worker", key)
		var err error
		app, err = newrelic.NewApplication(config)
		if err != nil {
			log.Info("Error creating New Relic app instance, continuing...")
		}
	}
}

// syncTaskJob would do whatever syncing is necessary in the background
func syncTaskJob(j *que.Job) error {
	dbConnectionString := os.Getenv("CLEARDB_DATABASE_URL")
	log.Infof("syntask connection string: %s", dbConnectionString)
	repo := sync.CreateSyncDbRepo(dbConnectionString)

	defer j.Delete()
	defer j.Done()

	var txn newrelic.Transaction
	if app != nil {
		txn = app.StartTransaction("sync-task", nil, nil)
		txn.SetName("sync job")
		defer txn.End()
	}

	var synctask sync.SyncTask
	err := json.Unmarshal(j.Args, &synctask)
	if err != nil {
		log.WithField("args", string(j.Args)).Error("Unable to unmarshal job arguments into SyncTask")
		return err
	}

	log.WithField("SyncTask", synctask).Info("Processing Synctask!")

	stvClientImpl := stv.CreateStravaClient(synctask.StravaToken)
	rkClientImpl := rk.CreateRKClient(synctask.RunkeeperToken)
	itemsCreatedRk, totalItems, err := synctask.Sync(stvClientImpl, rkClientImpl, txn)
	if err != nil {
		if app != nil {
			app.RecordCustomEvent("sync_error_event", map[string]interface{}{"error": err})
		}
		if txn != nil {
			txn.NoticeError(err)
		}
		log.WithField("args", string(j.Args)).WithField("QueId", j.ID).Error("Error while syncing synctask.")
		return err //don't update Timestamp, so we will retry
	}
	//update last executed timestamp
	log.Print("Updating last seen timestamp")
	//subtract 45 minutes to prevent activites being missed
	synctask.LastSeenTimestamp = int(time.Now().Add(time.Duration(tsDelta) * time.Minute).Unix())
	rowsUpdated, err := repo.UpdateSyncTask(synctask)
	if err != nil || rowsUpdated != 1 {
		log.Errorf("Error updating the SyncTask record with a new timestamp: %e", err)
		return err
	}

	if app != nil {
		app.RecordCustomEvent("sync_items_created", map[string]interface{}{
			"rk-items-created": itemsCreatedRk,
			"total-items":      totalItems})
	}

	log.Infof("Executed synctask %d, now deleting it.", synctask.Uid)
	return nil
}

func main() {
	log.Print("Starting worker")

	var err error
	dbURL := os.Getenv("DATABASE_URL")
	pgxpool, qc, err = shared.Setup(dbURL)
	if err != nil {
		log.WithField("DATABASE_URL", dbURL).Fatal("Errors setting up the queue / database: ", err)
	}
	defer pgxpool.Close()

	wm := que.WorkMap{
		shared.SyncTaskJob: syncTaskJob,
	}

	// 2 worker go routines
	// LIMITED to 1 per worker.. otherwise we overload Runkeeper
	// and I guess create activties twice.
	workers := que.NewWorkerPool(qc, wm, 1)

	// Catch signal so we can shutdown gracefully
	sigCh := make(chan os.Signal)
	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)

	go workers.Start()

	// Wait for a signal
	sig := <-sigCh
	log.WithField("signal", sig).Info("Signal received. Shutting down.")

	workers.Shutdown()
}
<commit_msg>Added error reporting on task<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/newrelic/go-agent"
	rk "github.com/svdberg/syncmysport-runkeeper/runkeeper"
	stv "github.com/svdberg/syncmysport-runkeeper/strava"
	sync "github.com/svdberg/syncmysport-runkeeper/sync"
	shared "github.com/svdberg/syncmysport-runkeeper/syncmysport-shared"

	log "github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/bgentry/que-go"
	"github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/jackc/pgx"
)

var (
	qc      *que.Client
	pgxpool *pgx.ConnPool
	app     newrelic.Application
)

const tsDelta = -45 //minutes

func init() {
	key := os.Getenv("NEW_RELIC_KEY")
	if key != "" {
		config := newrelic.NewConfig("SyncMySport-worker", key)
		var err error
		app, err = newrelic.NewApplication(config)
		if err != nil {
			log.Info("Error creating New Relic app instance, continuing...")
		}
	}
}

// syncTaskJob would do whatever syncing is necessary in the background
func syncTaskJob(j *que.Job) error {
	dbConnectionString := os.Getenv("CLEARDB_DATABASE_URL")
	log.Infof("syntask connection string: %s", dbConnectionString)
	repo := sync.CreateSyncDbRepo(dbConnectionString)

	defer j.Done()

	var txn newrelic.Transaction
	if app != nil {
		txn = app.StartTransaction("sync-task", nil, nil)
		txn.SetName("sync job")
		defer txn.End()
	}

	var synctask sync.SyncTask
	err := json.Unmarshal(j.Args, &synctask)
	if err != nil {
		log.WithField("args", string(j.Args)).Error("Unable to unmarshal job arguments into SyncTask")
		return err
	}

	log.WithField("SyncTask", synctask).Info("Processing Synctask!")

	stvClientImpl := stv.CreateStravaClient(synctask.StravaToken)
	rkClientImpl := rk.CreateRKClient(synctask.RunkeeperToken)
	itemsCreatedRk, totalItems, err := synctask.Sync(stvClientImpl, rkClientImpl, txn)
	if err != nil {
		if app != nil {
			app.RecordCustomEvent("sync_error_event", map[string]interface{}{"error": err})
		}
		if txn != nil {
			txn.NoticeError(err)
		}
		log.WithField("args", string(j.Args)).WithField("QueId", j.ID).Error("Error while syncing synctask.")
		j.Error(fmt.Sprintf("Error in synctask: %e", err))
		return err //don't update Timestamp, so we will retry
	}
	//update last executed timestamp
	log.Print("Updating last seen timestamp")
	//subtract 45 minutes to prevent activites being missed
	synctask.LastSeenTimestamp = int(time.Now().Add(time.Duration(tsDelta) * time.Minute).Unix())
	rowsUpdated, err := repo.UpdateSyncTask(synctask)
	if err != nil || rowsUpdated != 1 {
		log.Errorf("Error updating the SyncTask record with a new timestamp: %e", err)
		return err
	}

	if app != nil {
		app.RecordCustomEvent("sync_items_created", map[string]interface{}{
			"rk-items-created": itemsCreatedRk,
			"total-items":      totalItems})
	}

	log.Infof("Executed synctask %d, now deleting it.", synctask.Uid)
	j.Delete()
	return nil
}

func main() {
	log.Print("Starting worker")

	var err error
	dbURL := os.Getenv("DATABASE_URL")
	pgxpool, qc, err = shared.Setup(dbURL)
	if err != nil {
		log.WithField("DATABASE_URL", dbURL).Fatal("Errors setting up the queue / database: ", err)
	}
	defer pgxpool.Close()

	wm := que.WorkMap{
		shared.SyncTaskJob: syncTaskJob,
	}

	// 2 worker go routines
	// LIMITED to 1 per worker.. otherwise we overload Runkeeper
	// and I guess create activties twice.
	workers := que.NewWorkerPool(qc, wm, 1)

	// Catch signal so we can shutdown gracefully
	sigCh := make(chan os.Signal)
	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)

	go workers.Start()

	// Wait for a signal
	sig := <-sigCh
	log.WithField("signal", sig).Info("Signal received. Shutting down.")

	workers.Shutdown()
}
<|endoftext|>


<commit_before>package generate

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/gobuffalo/makr"
	"github.com/pkg/errors"

	"github.com/markbates/going/defaults"
	"github.com/markbates/inflect"
	"github.com/markbates/pop"
	"github.com/spf13/cobra"
)

var skipMigration bool
var structTag string

func init() {
	ModelCmd.Flags().StringVarP(&structTag, "struct-tag", "", "json", "sets the struct tags for model (xml or json)")
	ModelCmd.Flags().BoolVarP(&skipMigration, "skip-migration", "s", false, "Skip creating a new fizz migration for this model.")
}

var nrx = regexp.MustCompile(`^nulls\.(.+)`)

type names struct {
	Original string
	Table    string
	Proper   string
	File     string
	Plural   string
	Char     string
}

func newName(name string) names {
	return names{
		Original: name,
		File:     inflect.Underscore(inflect.Singularize(name)),
		Table:    inflect.Tableize(inflect.Pluralize(name)),
		Proper:   inflect.ForeignKeyToAttribute(name),
		Plural:   inflect.Pluralize(inflect.Camelize(name)),
		Char:     strings.ToLower(string([]byte(name)[0])),
	}
}

type attribute struct {
	Names        names
	OriginalType string
	GoType       string
	Nullable     bool
}

func (a attribute) String() string {
	return fmt.Sprintf("\t%s %s `%s:\"%s\" db:\"%s\"`", a.Names.Proper, a.GoType, structTag, a.Names.Original, a.Names.Original)
}

type model struct {
	Package               string
	Imports               []string
	Names                 names
	Attributes            []attribute
	ValidatableAttributes []attribute
}

func (m model) Generate() error {
	g := makr.New()
	defer g.Fmt(".")
	ctx := makr.Data{}
	ctx["model"] = m
	ctx["plural_model_name"] = m.Names.Plural
	ctx["model_name"] = m.Names.Proper
	ctx["package_name"] = m.Package
	ctx["char"] = m.Names.Char
	ctx["encoding_type"] = structTag
	ctx["encoding_type_char"] = strings.ToLower(string([]byte(structTag)[0]))

	fname := filepath.Join(m.Package, m.Names.File+".go")
	g.Add(makr.NewFile(fname, modelTemplate))
	tfname := filepath.Join(m.Package, m.Names.File+"_test.go")
	g.Add(makr.NewFile(tfname, modelTestTemplate))
	return g.Run(".", ctx)
}

func (m model) Fizz() string {
	s := []string{fmt.Sprintf("create_table(\"%s\", func(t) {", m.Names.Table)}
	for _, a := range m.Attributes {
		switch a.Names.Original {
		case "created_at", "updated_at":
		case "id":
			s = append(s, fmt.Sprintf("\tt.Column(\"id\", \"%s\", {\"primary\": true})", fizzColType(a.OriginalType)))
		default:
			x := fmt.Sprintf("\tt.Column(\"%s\", \"%s\", {})", a.Names.Original, fizzColType(a.OriginalType))
			if a.Nullable {
				x = strings.Replace(x, "{}", `{"null": true}`, -1)
			}
			s = append(s, x)
		}
	}
	s = append(s, "})")
	return strings.Join(s, "\n")
}

func newModel(name string) model {
	m := model{
		Package: "models",
		Imports: []string{"time", "github.com/markbates/pop", "github.com/markbates/validate"},
		Names:   newName(name),
		Attributes: []attribute{
			{Names: newName("created_at"), OriginalType: "time.Time", GoType: "time.Time"},
			{Names: newName("updated_at"), OriginalType: "time.Time", GoType: "time.Time"},
		},
		ValidatableAttributes: []attribute{},
	}
	m.Names.Proper = inflect.Singularize(m.Names.Proper)
	return m
}

var ModelCmd = &cobra.Command{
	Use:     "model [name]",
	Aliases: []string{"m"},
	Short:   "Generates a model for your database",
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return errors.New("You must supply a name for your model!")
		}

		if structTag != "json" && structTag != "xml" {
			return errors.New("Invalid struct tags (use xml or json)!")
		}

		model := newModel(args[0])

		if structTag == "json" {
			model.Imports = append(model.Imports, "encoding/json")
		} else if structTag == "xml" {
			model.Imports = append(model.Imports, "encoding/xml")
		}

		hasId := false
		hasUuid := false
		hasNulls := false
		for _, def := range args[1:] {
			col := strings.Split(def, ":")
			if len(col) == 1 {
				col = append(col, "string")
			}

			nullable := nrx.MatchString(col[1])
			if !hasNulls && nullable {
				hasNulls = true
				model.Imports = append(model.Imports, "github.com/markbates/pop/nulls")
			}
			if strings.HasPrefix(col[1], "slices.") {
				model.Imports = append(model.Imports, "github.com/markbates/pop/slices")
			} else if !hasUuid && col[1] == "uuid" {
				hasUuid = true
				model.Imports = append(model.Imports, "github.com/satori/go.uuid")
			}

			a := attribute{
				Names:        newName(col[0]),
				OriginalType: col[1],
				GoType:       colType(col[1]),
				Nullable:     nullable,
			}
			if col[0] == "id" {
				// No need to create a default ID
				hasId = true
				// Ensure ID is the first attribute
				model.Attributes = append([]attribute{a}, model.Attributes...)
			} else {
				model.Attributes = append(model.Attributes, a)
			}
			if !a.Nullable {
				if a.GoType == "string" || a.GoType == "time.Time" || a.GoType == "int" {
					if a.GoType == "time.Time" {
						a.GoType = "Time"
					}
					model.ValidatableAttributes = append(model.ValidatableAttributes, a)
				}
			}
		}

		// Add a default UUID, if no custom ID is provided
		if !hasId {
			if !hasUuid {
				model.Imports = append(model.Imports, "github.com/satori/go.uuid")
			}
			id := newName("id")
			id.Proper = "ID"
			a := attribute{Names: id, OriginalType: "uuid.UUID", GoType: "uuid.UUID"}
			// Ensure ID is the first attribute
			model.Attributes = append([]attribute{a}, model.Attributes...)
		}

		err := os.MkdirAll(model.Package, 0766)
		if err != nil {
			return errors.Wrapf(err, "couldn't create folder %s", model.Package)
		}

		err = model.Generate()
		if err != nil {
			return err
		}

		if !skipMigration {
			cflag := cmd.Flag("path")
			migrationPath := defaults.String(cflag.Value.String(), "./migrations")
			err = pop.MigrationCreate(migrationPath, fmt.Sprintf("create_%s", model.Names.Table), "fizz", []byte(model.Fizz()), []byte(fmt.Sprintf("drop_table(\"%s\")", model.Names.Table)))
			if err != nil {
				return err
			}
		}

		return nil
	},
}

func colType(s string) string {
	switch strings.ToLower(s) {
	case "text":
		return "string"
	case "time", "timestamp", "datetime":
		return "time.Time"
	case "nulls.text":
		return "nulls.String"
	case "uuid":
		return "uuid.UUID"
	case "json", "jsonb":
		return "slices.Map"
	case "[]string":
		return "slices.String"
	case "[]int":
		return "slices.Int"
	case "slices.float", "[]float", "[]float32", "[]float64":
		return "slices.Float"
	case "decimal", "float":
		return "float64"
	default:
		return s
	}
}

func fizzColType(s string) string {
	switch strings.ToLower(s) {
	case "int":
		return "integer"
	case "time", "datetime":
		return "timestamp"
	case "uuid.uuid", "uuid":
		return "uuid"
	case "nulls.float32", "nulls.float64":
		return "float"
	case "slices.string", "slices.uuid", "[]string":
		return "varchar[]"
	case "slices.float", "[]float", "[]float32", "[]float64":
		return "numeric[]"
	case "slices.int":
		return "int[]"
	case "slices.map":
		return "jsonb"
	case "float32", "float64":
		return "decimal"
	default:
		if nrx.MatchString(s) {
			return fizzColType(strings.Replace(s, "nulls.", "", -1))
		}
		return strings.ToLower(s)
	}
}
<commit_msg>adding back float for decimal<commit_after>package generate

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/gobuffalo/makr"
	"github.com/pkg/errors"

	"github.com/markbates/going/defaults"
	"github.com/markbates/inflect"
	"github.com/markbates/pop"
	"github.com/spf13/cobra"
)

var skipMigration bool
var structTag string

func init() {
	ModelCmd.Flags().StringVarP(&structTag, "struct-tag", "", "json", "sets the struct tags for model (xml or json)")
	ModelCmd.Flags().BoolVarP(&skipMigration, "skip-migration", "s", false, "Skip creating a new fizz migration for this model.")
}

var nrx = regexp.MustCompile(`^nulls\.(.+)`)

type names struct {
	Original string
	Table    string
	Proper   string
	File     string
	Plural   string
	Char     string
}

func newName(name string) names {
	return names{
		Original: name,
		File:     inflect.Underscore(inflect.Singularize(name)),
		Table:    inflect.Tableize(inflect.Pluralize(name)),
		Proper:   inflect.ForeignKeyToAttribute(name),
		Plural:   inflect.Pluralize(inflect.Camelize(name)),
		Char:     strings.ToLower(string([]byte(name)[0])),
	}
}

type attribute struct {
	Names        names
	OriginalType string
	GoType       string
	Nullable     bool
}

func (a attribute) String() string {
	return fmt.Sprintf("\t%s %s `%s:\"%s\" db:\"%s\"`", a.Names.Proper, a.GoType, structTag, a.Names.Original, a.Names.Original)
}

type model struct {
	Package               string
	Imports               []string
	Names                 names
	Attributes            []attribute
	ValidatableAttributes []attribute
}

func (m model) Generate() error {
	g := makr.New()
	defer g.Fmt(".")
	ctx := makr.Data{}
	ctx["model"] = m
	ctx["plural_model_name"] = m.Names.Plural
	ctx["model_name"] = m.Names.Proper
	ctx["package_name"] = m.Package
	ctx["char"] = m.Names.Char
	ctx["encoding_type"] = structTag
	ctx["encoding_type_char"] = strings.ToLower(string([]byte(structTag)[0]))

	fname := filepath.Join(m.Package, m.Names.File+".go")
	g.Add(makr.NewFile(fname, modelTemplate))
	tfname := filepath.Join(m.Package, m.Names.File+"_test.go")
	g.Add(makr.NewFile(tfname, modelTestTemplate))
	return g.Run(".", ctx)
}

func (m model) Fizz() string {
	s := []string{fmt.Sprintf("create_table(\"%s\", func(t) {", m.Names.Table)}
	for _, a := range m.Attributes {
		switch a.Names.Original {
		case "created_at", "updated_at":
		case "id":
			s = append(s, fmt.Sprintf("\tt.Column(\"id\", \"%s\", {\"primary\": true})", fizzColType(a.OriginalType)))
		default:
			x := fmt.Sprintf("\tt.Column(\"%s\", \"%s\", {})", a.Names.Original, fizzColType(a.OriginalType))
			if a.Nullable {
				x = strings.Replace(x, "{}", `{"null": true}`, -1)
			}
			s = append(s, x)
		}
	}
	s = append(s, "})")
	return strings.Join(s, "\n")
}

func newModel(name string) model {
	m := model{
		Package: "models",
		Imports: []string{"time", "github.com/markbates/pop", "github.com/markbates/validate"},
		Names:   newName(name),
		Attributes: []attribute{
			{Names: newName("created_at"), OriginalType: "time.Time", GoType: "time.Time"},
			{Names: newName("updated_at"), OriginalType: "time.Time", GoType: "time.Time"},
		},
		ValidatableAttributes: []attribute{},
	}
	m.Names.Proper = inflect.Singularize(m.Names.Proper)
	return m
}

var ModelCmd = &cobra.Command{
	Use:     "model [name]",
	Aliases: []string{"m"},
	Short:   "Generates a model for your database",
	RunE: func(cmd *cobra.Command, args []string) error {
		if len(args) == 0 {
			return errors.New("You must supply a name for your model!")
		}

		if structTag != "json" && structTag != "xml" {
			return errors.New("Invalid struct tags (use xml or json)!")
		}

		model := newModel(args[0])

		if structTag == "json" {
			model.Imports = append(model.Imports, "encoding/json")
		} else if structTag == "xml" {
			model.Imports = append(model.Imports, "encoding/xml")
		}

		hasId := false
		hasUuid := false
		hasNulls := false
		for _, def := range args[1:] {
			col := strings.Split(def, ":")
			if len(col) == 1 {
				col = append(col, "string")
			}

			nullable := nrx.MatchString(col[1])
			if !hasNulls && nullable {
				hasNulls = true
				model.Imports = append(model.Imports, "github.com/markbates/pop/nulls")
			}
			if strings.HasPrefix(col[1], "slices.") {
				model.Imports = append(model.Imports, "github.com/markbates/pop/slices")
			} else if !hasUuid && col[1] == "uuid" {
				hasUuid = true
				model.Imports = append(model.Imports, "github.com/satori/go.uuid")
			}

			a := attribute{
				Names:        newName(col[0]),
				OriginalType: col[1],
				GoType:       colType(col[1]),
				Nullable:     nullable,
			}
			if col[0] == "id" {
				// No need to create a default ID
				hasId = true
				// Ensure ID is the first attribute
				model.Attributes = append([]attribute{a}, model.Attributes...)
			} else {
				model.Attributes = append(model.Attributes, a)
			}
			if !a.Nullable {
				if a.GoType == "string" || a.GoType == "time.Time" || a.GoType == "int" {
					if a.GoType == "time.Time" {
						a.GoType = "Time"
					}
					model.ValidatableAttributes = append(model.ValidatableAttributes, a)
				}
			}
		}

		// Add a default UUID, if no custom ID is provided
		if !hasId {
			if !hasUuid {
				model.Imports = append(model.Imports, "github.com/satori/go.uuid")
			}
			id := newName("id")
			id.Proper = "ID"
			a := attribute{Names: id, OriginalType: "uuid.UUID", GoType: "uuid.UUID"}
			// Ensure ID is the first attribute
			model.Attributes = append([]attribute{a}, model.Attributes...)
		}

		err := os.MkdirAll(model.Package, 0766)
		if err != nil {
			return errors.Wrapf(err, "couldn't create folder %s", model.Package)
		}

		err = model.Generate()
		if err != nil {
			return err
		}

		if !skipMigration {
			cflag := cmd.Flag("path")
			migrationPath := defaults.String(cflag.Value.String(), "./migrations")
			err = pop.MigrationCreate(migrationPath, fmt.Sprintf("create_%s", model.Names.Table), "fizz", []byte(model.Fizz()), []byte(fmt.Sprintf("drop_table(\"%s\")", model.Names.Table)))
			if err != nil {
				return err
			}
		}

		return nil
	},
}

func colType(s string) string {
	switch strings.ToLower(s) {
	case "text":
		return "string"
	case "time", "timestamp", "datetime":
		return "time.Time"
	case "nulls.text":
		return "nulls.String"
	case "uuid":
		return "uuid.UUID"
	case "json", "jsonb":
		return "slices.Map"
	case "[]string":
		return "slices.String"
	case "[]int":
		return "slices.Int"
	case "slices.float", "[]float", "[]float32", "[]float64":
		return "slices.Float"
	case "decimal", "float":
		return "float64"
	default:
		return s
	}
}

func fizzColType(s string) string {
	switch strings.ToLower(s) {
	case "int":
		return "integer"
	case "time", "datetime":
		return "timestamp"
	case "uuid.uuid", "uuid":
		return "uuid"
	case "nulls.float32", "nulls.float64":
		return "float"
	case "slices.string", "slices.uuid", "[]string":
		return "varchar[]"
	case "slices.float", "[]float", "[]float32", "[]float64":
		return "numeric[]"
	case "slices.int":
		return "int[]"
	case "slices.map":
		return "jsonb"
	case "float32", "float64", "float":
		return "decimal"
	default:
		if nrx.MatchString(s) {
			return fizzColType(strings.Replace(s, "nulls.", "", -1))
		}
		return strings.ToLower(s)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/glendc/cgreader"
	"math"
)

func main() {
	cgreader.SetTimeout(60); cgreader.RunAndValidateManualProgram(
		"../../input/horse_dual_3.txt",
		"../../output/horse_dual_3.txt",
		true,
		func(ch <-chan string) string {
			var n int
			fmt.Sscanf(<-ch, "%d", &n)

			horses := make([]int, n)

			D := math.MaxInt32
			for i := range horses {
				fmt.Sscanf(<-ch, "%d", &horses[i])
				for u := 0 ; u < i ; u++ {
					x := horses[u] - horses[i]

					if x < 0 {
						x *= -1
					}

					if x < D {
						D = x
					}
				}
			}

			return fmt.Sprintf("%d\n", D)
		})
}
<commit_msg>horse dual challenge is now correct<commit_after>package main

import (
	"fmt"
	"github.com/glendc/cgreader"
	"math"
	"sort"
)

type intArray []int

func (s intArray) Len() int {
	return len(s)
}

func (s intArray) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s intArray) Less(i, j int) bool {
	return s[i] < s[j]
}

func main() {
	cgreader.SetTimeout(2.0)
	cgreader.RunAndValidateManualProgram(
		"../../input/horse_dual_3.txt",
		"../../output/horse_dual_3.txt",
		true,
		func(ch <-chan string) string {
			var n int
			fmt.Sscanf(<-ch, "%d", &n)

			horses := make([]int, n)
			for i := range horses {
				fmt.Sscanf(<-ch, "%d", &horses[i])
			}

			sort.Sort(intArray(horses))

			D := math.MaxInt32
			for i := 1; i < n; i++ {
				x := horses[i-1] - horses[i]

				if x < 0 {
					x *= -1
				}

				if x < D {
					D = x
				}
			}

			return fmt.Sprintf("%d\n", D)
		})
}
<|endoftext|>


<commit_before>package asl

import (
	"parser"
	"tokenizer"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"os"
	"strings"
)

const (
    version = "1.0.0"
    extension = ".asl"
    sqfextension = ".sqf"
    PathSeparator = string(os.PathSeparator)
)

type ASLFile struct {
    in string
    out string
    newname string
}

var (
    recursive bool = false
    pretty bool = false
    exit bool = false
    aslFiles []ASLFile
    inDir string
)

func usage() {
	fmt.Println("Usage: asl [-v|-r|-pretty|--help] <input directory> <output directory>\n")
	fmt.Println("-v (optional) shows asl version")
	fmt.Println("-r (optional) recursivly compile all asl files in folder")
	fmt.Println("-pretty (optional) activates pretty printing\n")
	fmt.Println("--help (optional) shows usage\n")
	fmt.Println("<input directory> directory to compile")
	fmt.Println("<output directory> output directory, directory structure will be created corresponding to input directory")
}

func flags(flag string) bool {
    flag = strings.ToLower(flag)

    if flag[0] == '-' {
        if flag == "-v" {
            fmt.Println("asl version "+version)
            exit = true
        } else if flag == "-r" {
            recursive = true
        } else if flag == "-pretty" {
            pretty = true
        } else if flag == "--help" {
            usage()
            exit = true
        }

        return true
    }

    return false
}

func readAslFiles(path string) {
    dir, err := ioutil.ReadDir(path)

    if err != nil {
        fmt.Println("Error reading in directory!")
        return
    }

    for i := 0; i < len(dir); i++ {
        name := dir[i].Name()

        if dir[i].IsDir() && recursive {
            readAslFiles(filepath.FromSlash(path+PathSeparator+name))
            continue
        }

        if !dir[i].IsDir() && strings.ToLower(filepath.Ext(name)) == extension {
            in := filepath.FromSlash(path+PathSeparator+dir[i].Name())
            out := filepath.FromSlash("./"+path[len(inDir):len(path)])
            newname := name[:len(name)-len(filepath.Ext(name))]

            file := ASLFile{in, out, newname}
            aslFiles = append(aslFiles, file)
        }
    }
}

func compile(path string) {
    for i := 0; i < len(aslFiles); i++ {
        out := filepath.FromSlash(path+PathSeparator+aslFiles[i].out+PathSeparator+aslFiles[i].newname+sqfextension)
        fmt.Println(aslFiles[i].in+" -> "+out)
        code, err := ioutil.ReadFile(aslFiles[i].in)

        if err != nil {
            fmt.Println("Error reading file: "+aslFiles[i].in)
            continue
        }

    	token := tokenizer.Tokenize(code)
    	compiler := parser.Compiler{}
    	sqf := compiler.Parse(token, pretty)

    	os.MkdirAll(filepath.FromSlash(path+PathSeparator+aslFiles[i].out), 0777)
    	err = ioutil.WriteFile(out, []byte(sqf), 0666)

    	if err != nil {
    	    fmt.Println("Error writing file: "+aslFiles[i].out)
    	    fmt.Println(err)
    	}
    }
}

func main() {
	args := os.Args

	// flags
	if len(args) < 2 {
	    usage()
	    return
	}

	var i int
	for i = 1; i < len(args) && flags(args[i]); i++ {}

	if exit {
	    return
	}

	// in/out parameter
	out := ""

	if i < len(args) {
	    inDir = args[i]
	    i++
	} else {
	    return
	}

	if i < len(args) {
	    out = args[i]
	}

	readAslFiles(inDir)
	compile(out)
}
<commit_msg>Version number.<commit_after>package asl

import (
	"parser"
	"tokenizer"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"os"
	"strings"
)

const (
    version = "1.1.0"
    extension = ".asl"
    sqfextension = ".sqf"
    PathSeparator = string(os.PathSeparator)
)

type ASLFile struct {
    in string
    out string
    newname string
}

var (
    recursive bool = false
    pretty bool = false
    exit bool = false
    aslFiles []ASLFile
    inDir string
)

func usage() {
	fmt.Println("Usage: asl [-v|-r|-pretty|--help] <input directory> <output directory>\n")
	fmt.Println("-v (optional) shows asl version")
	fmt.Println("-r (optional) recursivly compile all asl files in folder")
	fmt.Println("-pretty (optional) activates pretty printing\n")
	fmt.Println("--help (optional) shows usage\n")
	fmt.Println("<input directory> directory to compile")
	fmt.Println("<output directory> output directory, directory structure will be created corresponding to input directory")
}

func flags(flag string) bool {
    flag = strings.ToLower(flag)

    if flag[0] == '-' {
        if flag == "-v" {
            fmt.Println("asl version "+version)
            exit = true
        } else if flag == "-r" {
            recursive = true
        } else if flag == "-pretty" {
            pretty = true
        } else if flag == "--help" {
            usage()
            exit = true
        }

        return true
    }

    return false
}

func readAslFiles(path string) {
    dir, err := ioutil.ReadDir(path)

    if err != nil {
        fmt.Println("Error reading in directory!")
        return
    }

    for i := 0; i < len(dir); i++ {
        name := dir[i].Name()

        if dir[i].IsDir() && recursive {
            readAslFiles(filepath.FromSlash(path+PathSeparator+name))
            continue
        }

        if !dir[i].IsDir() && strings.ToLower(filepath.Ext(name)) == extension {
            in := filepath.FromSlash(path+PathSeparator+dir[i].Name())
            out := filepath.FromSlash("./"+path[len(inDir):len(path)])
            newname := name[:len(name)-len(filepath.Ext(name))]

            file := ASLFile{in, out, newname}
            aslFiles = append(aslFiles, file)
        }
    }
}

func compile(path string) {
    for i := 0; i < len(aslFiles); i++ {
        out := filepath.FromSlash(path+PathSeparator+aslFiles[i].out+PathSeparator+aslFiles[i].newname+sqfextension)
        fmt.Println(aslFiles[i].in+" -> "+out)
        code, err := ioutil.ReadFile(aslFiles[i].in)

        if err != nil {
            fmt.Println("Error reading file: "+aslFiles[i].in)
            continue
        }

    	token := tokenizer.Tokenize(code)
    	compiler := parser.Compiler{}
    	sqf := compiler.Parse(token, pretty)

    	os.MkdirAll(filepath.FromSlash(path+PathSeparator+aslFiles[i].out), 0777)
    	err = ioutil.WriteFile(out, []byte(sqf), 0666)

    	if err != nil {
    	    fmt.Println("Error writing file: "+aslFiles[i].out)
    	    fmt.Println(err)
    	}
    }
}

func main() {
	args := os.Args

	// flags
	if len(args) < 2 {
	    usage()
	    return
	}

	var i int
	for i = 1; i < len(args) && flags(args[i]); i++ {}

	if exit {
	    return
	}

	// in/out parameter
	out := ""

	if i < len(args) {
	    inDir = args[i]
	    i++
	} else {
	    return
	}

	if i < len(args) {
	    out = args[i]
	}

	readAslFiles(inDir)
	compile(out)
}
<|endoftext|>


<commit_before>/*
 * util.go - Various helpers used throughout fscrypt
 *
 * Copyright 2017 Google Inc.
 * Author: Joe Richey (joerichey@google.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 util contains useful components for simplifying Go code.
//
// The package contains common error types (errors.go) and functions for
// converting arrays to pointers.
package util

import (
	"bufio"
	"math"
	"os"
	"unsafe"
)

// Ptr converts an Go byte array to a pointer to the start of the array.
func Ptr(slice []byte) unsafe.Pointer {
	if len(slice) == 0 {
		return nil
	}
	return unsafe.Pointer(&slice[0])
}

// ByteSlice takes a pointer to some data and views it as a slice of bytes.
// Note, indexing into this slice is unsafe.
func ByteSlice(ptr unsafe.Pointer) []byte {
	// Silce must fix in 32-bit address space to build on 32-bit platforms.
	return (*[math.MaxInt32]byte)(ptr)[:]
}

// PointerSlice takes a pointer to an array of pointers and views it as a slice
// of pointers. Note, indexing into this slice is unsafe.
func PointerSlice(ptr unsafe.Pointer) []unsafe.Pointer {
	// Silce must fix in 32-bit address space to build on 32-bit platforms.
	return (*[math.MaxInt32 / 4]unsafe.Pointer)(ptr)[:]
}

// Index returns the first index i such that inVal == inArray[i].
// ok is true if we find a match, false otherwise.
func Index(inVal int64, inArray []int64) (index int, ok bool) {
	for index, val := range inArray {
		if val == inVal {
			return index, true
		}
	}
	return 0, false
}

// Lookup finds inVal in inArray and returns the corresponding element in
// outArray. Specifically, if inVal == inArray[i], outVal == outArray[i].
// ok is true if we find a match, false otherwise.
func Lookup(inVal int64, inArray, outArray []int64) (outVal int64, ok bool) {
	index, ok := Index(inVal, inArray)
	if !ok {
		return 0, false
	}
	return outArray[index], true
}

// MinInt returns the lesser of a and b.
func MinInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// MaxInt returns the greater of a and b.
func MaxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// MinInt64 returns the lesser of a and b.
func MinInt64(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}

// ReadLine returns a line of input from standard input. An empty string is
// returned if the user didn't insert anything or on error.
func ReadLine() (string, error) {
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	return scanner.Text(), scanner.Err()
}
<commit_msg>util: Added parsing and effective user functions<commit_after>/*
 * util.go - Various helpers used throughout fscrypt
 *
 * Copyright 2017 Google Inc.
 * Author: Joe Richey (joerichey@google.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 util contains useful components for simplifying Go code.
//
// The package contains common error types (errors.go) and functions for
// converting arrays to pointers.
package util

import (
	"bufio"
	"math"
	"os"
	"os/user"
	"strconv"
	"unsafe"
)

// Ptr converts an Go byte array to a pointer to the start of the array.
func Ptr(slice []byte) unsafe.Pointer {
	if len(slice) == 0 {
		return nil
	}
	return unsafe.Pointer(&slice[0])
}

// ByteSlice takes a pointer to some data and views it as a slice of bytes.
// Note, indexing into this slice is unsafe.
func ByteSlice(ptr unsafe.Pointer) []byte {
	// Silce must fix in 32-bit address space to build on 32-bit platforms.
	return (*[math.MaxInt32]byte)(ptr)[:]
}

// PointerSlice takes a pointer to an array of pointers and views it as a slice
// of pointers. Note, indexing into this slice is unsafe.
func PointerSlice(ptr unsafe.Pointer) []unsafe.Pointer {
	// Silce must fix in 32-bit address space to build on 32-bit platforms.
	return (*[math.MaxInt32 / 4]unsafe.Pointer)(ptr)[:]
}

// Index returns the first index i such that inVal == inArray[i].
// ok is true if we find a match, false otherwise.
func Index(inVal int64, inArray []int64) (index int, ok bool) {
	for index, val := range inArray {
		if val == inVal {
			return index, true
		}
	}
	return 0, false
}

// Lookup finds inVal in inArray and returns the corresponding element in
// outArray. Specifically, if inVal == inArray[i], outVal == outArray[i].
// ok is true if we find a match, false otherwise.
func Lookup(inVal int64, inArray, outArray []int64) (outVal int64, ok bool) {
	index, ok := Index(inVal, inArray)
	if !ok {
		return 0, false
	}
	return outArray[index], true
}

// MinInt returns the lesser of a and b.
func MinInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// MaxInt returns the greater of a and b.
func MaxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// MinInt64 returns the lesser of a and b.
func MinInt64(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}

// ReadLine returns a line of input from standard input. An empty string is
// returned if the user didn't insert anything or on error.
func ReadLine() (string, error) {
	scanner := bufio.NewScanner(os.Stdin)
	scanner.Scan()
	return scanner.Text(), scanner.Err()
}

// AtoiOrPanic converts a string to an int or it panics. Should only be used in
// situations where the input MUST be a decimal number.
func AtoiOrPanic(input string) int {
	i, err := strconv.Atoi(input)
	if err != nil {
		panic(err)
	}
	return i
}

// EffectiveUser returns the user entry corresponding to the effective user.
func EffectiveUser() (*user.User, error) {
	return user.LookupId(strconv.Itoa(os.Geteuid()))
}

// IsUserRoot checks if the effective user is root.
func IsUserRoot() bool {
	return os.Geteuid() == 0
}
<|endoftext|>


<commit_before>// +build go1.7,!go1.8

package echo

import (
	"net/url"
)

// PathUnescape is wraps `url.QueryUnescape`
func PathUnescape(s string) (string, error) {
	return url.QueryUnescape(s)
}
<commit_msg>fix build constraint go1.8 (#993)<commit_after>// +build go1.7, !go1.8

package echo

import (
	"net/url"
)

// PathUnescape is wraps `url.QueryUnescape`
func PathUnescape(s string) (string, error) {
	return url.QueryUnescape(s)
}
<|endoftext|>


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

package integration

import (
	"math/rand"
	"sync"
	"testing"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/clientv3/concurrency"
	"github.com/coreos/etcd/contrib/recipes"
	"golang.org/x/net/context"
)

func TestMutexSingleNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	var clients []*clientv3.Client
	testMutex(t, 5, makeSingleNodeClients(t, clus.cluster, &clients))
	closeClients(t, clients)
}

func TestMutexMultiNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	var clients []*clientv3.Client
	testMutex(t, 5, makeMultiNodeClients(t, clus.cluster, &clients))
	closeClients(t, clients)
}

func testMutex(t *testing.T, waiters int, chooseClient func() *clientv3.Client) {
	// stream lock acquisitions
	lockedC := make(chan *concurrency.Mutex)
	for i := 0; i < waiters; i++ {
		go func() {
			session, err := concurrency.NewSession(chooseClient())
			if err != nil {
				t.Error(err)
			}
			m := concurrency.NewMutex(session, "test-mutex")
			if err := m.Lock(context.TODO()); err != nil {
				t.Fatalf("could not wait on lock (%v)", err)
			}
			lockedC <- m
		}()
	}
	// unlock locked mutexes
	timerC := time.After(time.Duration(waiters) * time.Second)
	for i := 0; i < waiters; i++ {
		select {
		case <-timerC:
			t.Fatalf("timed out waiting for lock %d", i)
		case m := <-lockedC:
			// lock acquired with m
			select {
			case <-lockedC:
				t.Fatalf("lock %d followers did not wait", i)
			default:
			}
			if err := m.Unlock(context.TODO()); err != nil {
				t.Fatalf("could not release lock (%v)", err)
			}
		}
	}
}

// TestMutexSessionRelock ensures that acquiring the same lock with the same
// session will not result in deadlock.
func TestMutexSessionRelock(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)
	session, err := concurrency.NewSession(clus.RandClient())
	if err != nil {
		t.Error(err)
	}

	m := concurrency.NewMutex(session, "test-mutex")
	if err := m.Lock(context.TODO()); err != nil {
		t.Fatal(err)
	}

	m2 := concurrency.NewMutex(session, "test-mutex")
	if err := m2.Lock(context.TODO()); err != nil {
		t.Fatal(err)
	}
}

func BenchmarkMutex4Waiters(b *testing.B) {
	// XXX switch tests to use TB interface
	clus := NewClusterV3(nil, &ClusterConfig{Size: 3})
	defer clus.Terminate(nil)
	for i := 0; i < b.N; i++ {
		testMutex(nil, 4, func() *clientv3.Client { return clus.RandClient() })
	}
}

func TestRWMutexSingleNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)
	testRWMutex(t, 5, func() *clientv3.Client { return clus.clients[0] })
}

func TestRWMutexMultiNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)
	testRWMutex(t, 5, func() *clientv3.Client { return clus.RandClient() })
}

func testRWMutex(t *testing.T, waiters int, chooseClient func() *clientv3.Client) {
	// stream rwlock acquistions
	rlockedC := make(chan *recipe.RWMutex, 1)
	wlockedC := make(chan *recipe.RWMutex, 1)
	for i := 0; i < waiters; i++ {
		go func() {
			session, err := concurrency.NewSession(chooseClient())
			if err != nil {
				t.Error(err)
			}
			rwm := recipe.NewRWMutex(session, "test-rwmutex")
			if rand.Intn(2) == 0 {
				if err := rwm.RLock(); err != nil {
					t.Fatalf("could not rlock (%v)", err)
				}
				rlockedC <- rwm
			} else {
				if err := rwm.Lock(); err != nil {
					t.Fatalf("could not lock (%v)", err)
				}
				wlockedC <- rwm
			}
		}()
	}
	// unlock locked rwmutexes
	timerC := time.After(time.Duration(waiters) * time.Second)
	for i := 0; i < waiters; i++ {
		select {
		case <-timerC:
			t.Fatalf("timed out waiting for lock %d", i)
		case wl := <-wlockedC:
			select {
			case <-rlockedC:
				t.Fatalf("rlock %d readers did not wait", i)
			default:
			}
			if err := wl.Unlock(); err != nil {
				t.Fatalf("could not release lock (%v)", err)
			}
		case rl := <-rlockedC:
			select {
			case <-wlockedC:
				t.Fatalf("rlock %d writers did not wait", i)
			default:
			}
			if err := rl.RUnlock(); err != nil {
				t.Fatalf("could not release rlock (%v)", err)
			}
		}
	}
}

func makeClients(t *testing.T, clients *[]*clientv3.Client, choose func() *member) func() *clientv3.Client {
	var mu sync.Mutex
	*clients = nil
	return func() *clientv3.Client {
		cli, err := NewClientV3(choose())
		if err != nil {
			t.Fatalf("cannot create client: %v", err)
		}
		mu.Lock()
		*clients = append(*clients, cli)
		mu.Unlock()
		return cli
	}
}

func makeSingleNodeClients(t *testing.T, clus *cluster, clients *[]*clientv3.Client) func() *clientv3.Client {
	return makeClients(t, clients, func() *member {
		return clus.Members[0]
	})
}

func makeMultiNodeClients(t *testing.T, clus *cluster, clients *[]*clientv3.Client) func() *clientv3.Client {
	return makeClients(t, clients, func() *member {
		return clus.Members[rand.Intn(len(clus.Members))]
	})
}

func closeClients(t *testing.T, clients []*clientv3.Client) {
	for _, cli := range clients {
		if err := cli.Close(); err != nil {
			t.Fatal(err)
		}
	}
}
<commit_msg>integration: add TestMutexWaitsOnCurrentHolder test<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package integration

import (
	"math/rand"
	"sync"
	"testing"
	"time"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/clientv3/concurrency"
	"github.com/coreos/etcd/contrib/recipes"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"github.com/coreos/etcd/pkg/testutil"
	"golang.org/x/net/context"
)

func TestMutexSingleNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	var clients []*clientv3.Client
	testMutex(t, 5, makeSingleNodeClients(t, clus.cluster, &clients))
	closeClients(t, clients)
}

func TestMutexMultiNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)

	var clients []*clientv3.Client
	testMutex(t, 5, makeMultiNodeClients(t, clus.cluster, &clients))
	closeClients(t, clients)
}

func testMutex(t *testing.T, waiters int, chooseClient func() *clientv3.Client) {
	// stream lock acquisitions
	lockedC := make(chan *concurrency.Mutex)
	for i := 0; i < waiters; i++ {
		go func() {
			session, err := concurrency.NewSession(chooseClient())
			if err != nil {
				t.Error(err)
			}
			m := concurrency.NewMutex(session, "test-mutex")
			if err := m.Lock(context.TODO()); err != nil {
				t.Fatalf("could not wait on lock (%v)", err)
			}
			lockedC <- m
		}()
	}
	// unlock locked mutexes
	timerC := time.After(time.Duration(waiters) * time.Second)
	for i := 0; i < waiters; i++ {
		select {
		case <-timerC:
			t.Fatalf("timed out waiting for lock %d", i)
		case m := <-lockedC:
			// lock acquired with m
			select {
			case <-lockedC:
				t.Fatalf("lock %d followers did not wait", i)
			default:
			}
			if err := m.Unlock(context.TODO()); err != nil {
				t.Fatalf("could not release lock (%v)", err)
			}
		}
	}
}

// TestMutexSessionRelock ensures that acquiring the same lock with the same
// session will not result in deadlock.
func TestMutexSessionRelock(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)
	session, err := concurrency.NewSession(clus.RandClient())
	if err != nil {
		t.Error(err)
	}

	m := concurrency.NewMutex(session, "test-mutex")
	if err := m.Lock(context.TODO()); err != nil {
		t.Fatal(err)
	}

	m2 := concurrency.NewMutex(session, "test-mutex")
	if err := m2.Lock(context.TODO()); err != nil {
		t.Fatal(err)
	}
}

// TestMutexWaitsOnCurrentHolder ensures a mutex is only acquired once all
// waiters older than the new owner are gone by testing the case where
// the waiter prior to the acquirer expires before the current holder.
func TestMutexWaitsOnCurrentHolder(t *testing.T) {
	defer testutil.AfterTest(t)

	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	cctx := context.Background()

	cli := clus.Client(0)

	firstOwnerSession, err := concurrency.NewSession(cli)
	if err != nil {
		t.Error(err)
	}
	defer firstOwnerSession.Close()
	firstOwnerMutex := concurrency.NewMutex(firstOwnerSession, "test-mutex")
	if err = firstOwnerMutex.Lock(cctx); err != nil {
		t.Fatal(err)
	}

	victimSession, err := concurrency.NewSession(cli)
	if err != nil {
		t.Error(err)
	}
	defer victimSession.Close()
	victimDonec := make(chan struct{})
	go func() {
		defer close(victimDonec)
		concurrency.NewMutex(victimSession, "test-mutex").Lock(cctx)
	}()

	// ensure mutexes associated with firstOwnerSession and victimSession waits before new owner
	wch := cli.Watch(cctx, "test-mutex", clientv3.WithPrefix(), clientv3.WithRev(1))
	putCounts := 0
	for putCounts < 2 {
		select {
		case wrp := <-wch:
			putCounts += len(wrp.Events)
		case <-time.After(time.Second):
			t.Fatal("failed to receive watch response")
		}
	}
	if putCounts != 2 {
		t.Fatalf("expect 2 put events, but got %v", putCounts)
	}

	newOwnerSession, err := concurrency.NewSession(cli)
	if err != nil {
		t.Error(err)
	}
	defer newOwnerSession.Close()
	newOwnerDonec := make(chan struct{})
	go func() {
		defer close(newOwnerDonec)
		concurrency.NewMutex(newOwnerSession, "test-mutex").Lock(cctx)
	}()

	select {
	case wrp := <-wch:
		if len(wrp.Events) != 1 {
			t.Fatalf("expect a event, but got %v events", len(wrp.Events))
		}
		if e := wrp.Events[0]; e.Type != mvccpb.PUT {
			t.Fatalf("expect a put event on prefix test-mutex, but got event type %v", e.Type)
		}
	case <-time.After(time.Second):
		t.Fatalf("failed to receive a watch response")
	}

	// simulate losing the client that's next in line to acquire the lock
	victimSession.Close()

	// ensures the deletion of victim waiter from server side.
	select {
	case wrp := <-wch:
		if len(wrp.Events) != 1 {
			t.Fatalf("expect a event, but got %v events", len(wrp.Events))
		}
		if e := wrp.Events[0]; e.Type != mvccpb.DELETE {
			t.Fatalf("expect a delete event on prefix test-mutex, but got event type %v", e.Type)
		}
	case <-time.After(time.Second):
		t.Fatal("failed to receive a watch response")
	}

	select {
	case <-newOwnerDonec:
		t.Fatal("new owner obtained lock before first owner unlocked")
	default:
	}

	if err := firstOwnerMutex.Unlock(cctx); err != nil {
		t.Fatal(err)
	}

	select {
	case <-newOwnerDonec:
	case <-time.After(time.Second):
		t.Fatal("new owner failed to obtain lock")
	}

	select {
	case <-victimDonec:
	case <-time.After(time.Second):
		t.Fatal("victim mutex failed to exit after first owner releases lock")
	}
}

func BenchmarkMutex4Waiters(b *testing.B) {
	// XXX switch tests to use TB interface
	clus := NewClusterV3(nil, &ClusterConfig{Size: 3})
	defer clus.Terminate(nil)
	for i := 0; i < b.N; i++ {
		testMutex(nil, 4, func() *clientv3.Client { return clus.RandClient() })
	}
}

func TestRWMutexSingleNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)
	testRWMutex(t, 5, func() *clientv3.Client { return clus.clients[0] })
}

func TestRWMutexMultiNode(t *testing.T) {
	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
	defer clus.Terminate(t)
	testRWMutex(t, 5, func() *clientv3.Client { return clus.RandClient() })
}

func testRWMutex(t *testing.T, waiters int, chooseClient func() *clientv3.Client) {
	// stream rwlock acquistions
	rlockedC := make(chan *recipe.RWMutex, 1)
	wlockedC := make(chan *recipe.RWMutex, 1)
	for i := 0; i < waiters; i++ {
		go func() {
			session, err := concurrency.NewSession(chooseClient())
			if err != nil {
				t.Error(err)
			}
			rwm := recipe.NewRWMutex(session, "test-rwmutex")
			if rand.Intn(2) == 0 {
				if err := rwm.RLock(); err != nil {
					t.Fatalf("could not rlock (%v)", err)
				}
				rlockedC <- rwm
			} else {
				if err := rwm.Lock(); err != nil {
					t.Fatalf("could not lock (%v)", err)
				}
				wlockedC <- rwm
			}
		}()
	}
	// unlock locked rwmutexes
	timerC := time.After(time.Duration(waiters) * time.Second)
	for i := 0; i < waiters; i++ {
		select {
		case <-timerC:
			t.Fatalf("timed out waiting for lock %d", i)
		case wl := <-wlockedC:
			select {
			case <-rlockedC:
				t.Fatalf("rlock %d readers did not wait", i)
			default:
			}
			if err := wl.Unlock(); err != nil {
				t.Fatalf("could not release lock (%v)", err)
			}
		case rl := <-rlockedC:
			select {
			case <-wlockedC:
				t.Fatalf("rlock %d writers did not wait", i)
			default:
			}
			if err := rl.RUnlock(); err != nil {
				t.Fatalf("could not release rlock (%v)", err)
			}
		}
	}
}

func makeClients(t *testing.T, clients *[]*clientv3.Client, choose func() *member) func() *clientv3.Client {
	var mu sync.Mutex
	*clients = nil
	return func() *clientv3.Client {
		cli, err := NewClientV3(choose())
		if err != nil {
			t.Fatalf("cannot create client: %v", err)
		}
		mu.Lock()
		*clients = append(*clients, cli)
		mu.Unlock()
		return cli
	}
}

func makeSingleNodeClients(t *testing.T, clus *cluster, clients *[]*clientv3.Client) func() *clientv3.Client {
	return makeClients(t, clients, func() *member {
		return clus.Members[0]
	})
}

func makeMultiNodeClients(t *testing.T, clus *cluster, clients *[]*clientv3.Client) func() *clientv3.Client {
	return makeClients(t, clients, func() *member {
		return clus.Members[rand.Intn(len(clus.Members))]
	})
}

func closeClients(t *testing.T, clients []*clientv3.Client) {
	for _, cli := range clients {
		if err := cli.Close(); err != nil {
			t.Fatal(err)
		}
	}
}
<|endoftext|>