<commit_before>package main

import (
	"encoding/binary"
	"fmt"
	"encoding/hex"
	"github.com/lbarman/crypto/abstract"
	"github.com/lbarman/prifi/dcnet"
	"io"
	//"os"
	"strconv"
	"log"
	"strings"
	"net"
	"time"
	log2 "github.com/lbarman/prifi/log"
)

type Trustee struct {
	pubkey abstract.Point
}

type AnonSet struct {
	suite    abstract.Suite
	trustees []Trustee
}

// Periodic stats reporting
var begin = time.Now()
var report = begin
var numberOfReports = 0
var period, _ = time.ParseDuration("3s")
var totupcells = int64(0)
var totupbytes = int64(0)
var totdowncells = int64(0)
var totdownbytes = int64(0)

var parupcells = int64(0)
var parupbytes = int64(0)
var pardownbytes = int64(0)

func reportStatistics(payloadLength int, reportingLimit int) bool {
	now := time.Now()
	if now.After(report) {
		duration := now.Sub(begin).Seconds()

		instantUpSpeed := (float64(parupbytes)/period.Seconds())

		fmt.Printf("@ %fs; cell %f (%f) /sec, up %f (%f) B/s, down %f (%f) B/s\n",
			duration,
			 float64(totupcells)/duration, float64(parupcells)/period.Seconds(),
			 float64(totupbytes)/duration, instantUpSpeed,
			 float64(totdownbytes)/duration, float64(pardownbytes)/period.Seconds())

			// Next report time
		parupcells = 0
		parupbytes = 0
		pardownbytes = 0

		//log2.BenchmarkFloat(fmt.Sprintf("cellsize-%d-upstream-bytes", payloadLength), instantUpSpeed)

		data := struct {
		    Experiment string
		    CellSize int
		    Speed float64
		}{
		    "upstream-speed-given-cellsize",
		    payloadLength,
		    instantUpSpeed,
		}

		log2.JsonDump(data)

		report = now.Add(period)
		numberOfReports += 1

		if(reportingLimit > -1 && numberOfReports >= reportingLimit) {
			return false
		}
	}

	return true
}

type RelayCryptoParams struct {
	Name				string

	PublicKey			abstract.Point
	privateKey			abstract.Secret

	clientsConnections  []net.Conn
	trusteesConnections []net.Conn
	trusteesPublicKeys  []abstract.Point
	clientPublicKeys    []abstract.Point

	CellCoder			dcnet.CellCoder

	MessageHistory		abstract.Cipher
}

func initateRelayCrypto(nTrustees int, nClients int) *RelayCryptoParams {

	params := new(RelayCryptoParams)

	params.Name = "Relay"

	//prepare the crypto parameters
	rand 	:= suite.Cipher([]byte(params.Name))
	base	:= suite.Point().Base()

	//generate own parameters
	params.privateKey       = suite.Secret().Pick(rand)
	params.PublicKey        = suite.Point().Mul(base, params.privateKey)

	//placeholders for pubkeys and connections
	params.trusteesPublicKeys = make([]abstract.Point, nTrustees)
	params.clientPublicKeys   = make([]abstract.Point, nClients)

	params.trusteesConnections = make([]net.Conn, nTrustees)
	params.clientsConnections  = make([]net.Conn, nClients)

	//sets the cell coder, and the history
	params.CellCoder = factory()

	return params
}

var clientsConnections  []net.Conn
var trusteesConnections []net.Conn
var trusteesPublicKeys  []abstract.Point
var clientPublicKeys    []abstract.Point

func startRelay(payloadLength int, relayPort string, nClients int, nTrustees int, trusteesIp []string, reportingLimit int) {

	cryptoParams := initateRelayCrypto(nTrustees, nClients)

	//connect to the trustees
	trusteesConnections = make([]net.Conn, nTrustees)
	trusteesPublicKeys  = make([]abstract.Point, nTrustees)

	for i:= 0; i < nTrustees; i++ {
		currentTrusteeIp := strings.Replace(trusteesIp[i], "_", ":", -1) //trick for windows shell, where ":" separates args
		connectToTrustee(i, currentTrusteeIp, nClients, nTrustees, payloadLength)
	}

	//starts the client server
	lsock, err := net.Listen("tcp", relayPort)
	if err != nil {
		panic("Can't open listen socket:" + err.Error())
	}

	// Wait for all the clients to connect, and parse parameters
	clientsConnections = make([]net.Conn, nClients)
	clientPublicKeys   = make([]abstract.Point, nClients)

	for j := 0; j < nClients; j++ {
		fmt.Printf("Waiting for %d clients (on port %s)\n", nClients-j, relayPort)

		conn, err := lsock.Accept()
		if err != nil {
			panic("Listen error:" + err.Error())
		}

		buffer := make([]byte, 512)
		_, err2 := conn.Read(buffer)
		if err2 != nil {
			panic("Read error:" + err2.Error())
		}

		ver := int(binary.BigEndian.Uint32(buffer[0:4]))

		if(ver != LLD_PROTOCOL_VERSION) {
			fmt.Println(">>>> Relay client version", ver, "!= relay version", LLD_PROTOCOL_VERSION)
			panic("fatal error")
		}

		nodeId := int(binary.BigEndian.Uint32(buffer[4:8]))
		keySize := int(binary.BigEndian.Uint32(buffer[8:12]))
		keyBytes := buffer[12:(12+keySize)]

		publicKey := suite.Point()
		err3 := publicKey.UnmarshalBinary(keyBytes)

		if err3 != nil {
			panic(">>>>  Relay : can't unmarshal client key ! " + err3.Error())
		}

		if nodeId >= 0 && nodeId < nClients {
			clientsConnections[nodeId] = conn
			clientPublicKeys[nodeId] = publicKey
		} else {
			panic("illegal node number")
		}
	}
	println("All clients and trustees connected.")

	//Prepare the messages
	messageForClient   := MarshalPublicKeyArrayToByteArray(trusteesPublicKeys)
	messageForTrustees := MarshalPublicKeyArrayToByteArray(clientPublicKeys)

	//broadcast to the clients
	broadcastMessage(clientsConnections, messageForClient)
	broadcastMessage(trusteesConnections, messageForTrustees)

	println("All crypto stuff exchanged !")

	for {
		time.Sleep(5000 * time.Millisecond)
	}

	// Create ciphertext slice bufferfers for all clients and trustees
	clientPayloadLength := cryptoParams.CellCoder.ClientCellSize(payloadLength)
	clientsPayloadData  := make([][]byte, nClients)
	for i := 0; i < nClients; i++ {
		clientsPayloadData[i] = make([]byte, clientPayloadLength)
	}

	trusteePayloadLength := cryptoParams.CellCoder.TrusteeCellSize(payloadLength)
	trusteesPayloadData  := make([][]byte, nTrustees)
	for i := 0; i < nTrustees; i++ {
		trusteesPayloadData[i] = make([]byte, trusteePayloadLength)
	}

	conns := make(map[int]chan<- []byte)
	downstream := make(chan dataWithConnectionId)
	nulldown := dataWithConnectionId{} // default empty downstream cell
	window := 2           // Maximum cells in-flight
	inflight := 0         // Current cells in-flight

	for {

		//TODO: change this way of breaking the loop, it's not very elegant..
		// Show periodic reports
		if(!reportStatistics(payloadLength, reportingLimit)) {
			println("Reporting limit matched; exiting the relay")
			break;
		}

		//TODO : check if it is required to send empty cell
		// See if there's any downstream data to forward.
		var downbuffer dataWithConnectionId
		select {
			case downbuffer = <-downstream: // some data to forward downstream
				//fmt.Println("Downstream data...")
				//fmt.Printf("v %d\n", len(downbuffer)-6)
			default: // nothing at the moment to forward
				downbuffer = nulldown
		}

		downstreamDataPayloadLength := len(downbuffer.data)
		downstreamData := make([]byte, 6+downstreamDataPayloadLength)
		binary.BigEndian.PutUint32(downstreamData[0:4], uint32(downbuffer.connectionId))
		binary.BigEndian.PutUint16(downstreamData[4:6], uint16(downstreamDataPayloadLength))
		copy(downstreamData[6:], downbuffer.data)

		// Broadcast the downstream data to all clients.
		for i := 0; i < nClients; i++ {
			n, err := clientsConnections[i].Write(downstreamData)

			if n != 6+downstreamDataPayloadLength {
				panic("Relay : Write to client failed, wrote "+strconv.Itoa(downstreamDataPayloadLength+6)+" where "+strconv.Itoa(n)+" was expected : " + err.Error())
			}
		}

		totdowncells++
		totdownbytes += int64(downstreamDataPayloadLength)
		pardownbytes += int64(downstreamDataPayloadLength)

		inflight++
		if inflight < window {
			continue // Get more cells in flight
		}

		cryptoParams.CellCoder.DecodeStart(payloadLength, cryptoParams.MessageHistory)

		// Collect a cell ciphertext from each trustee
		for i := 0; i < nTrustees; i++ {

			//TODO: this looks blocking
			n, err := io.ReadFull(trusteesConnections[i], trusteesPayloadData[i])
			if n < trusteePayloadLength {
				panic("Relay : Read from trustee failed, read "+strconv.Itoa(n)+" where "+strconv.Itoa(trusteePayloadLength)+" was expected: " + err.Error())
			}

			cryptoParams.CellCoder.DecodeTrustee(trusteesPayloadData[i])
		}

		// Collect an upstream ciphertext from each client
		for i := 0; i < nClients; i++ {

			//TODO: this looks blocking
			n, err := io.ReadFull(clientsConnections[i], clientsPayloadData[i])
			if n < clientPayloadLength {
				panic("Relay : Read from client failed, read "+strconv.Itoa(n)+" where "+strconv.Itoa(clientPayloadLength)+" was expected: " + err.Error())
			}

			cryptoParams.CellCoder.DecodeClient(clientsPayloadData[i])
		}

		upstreamPlaintext := cryptoParams.CellCoder.DecodeCell()
		inflight--

		totupcells++
		totupbytes += int64(payloadLength)
		parupcells++
		parupbytes += int64(payloadLength)

		// Process the decoded cell
		if upstreamPlaintext == nil {
			continue // empty or corrupt upstream cell
		}
		if len(upstreamPlaintext) != payloadLength {
			panic("DecodeCell produced wrong-size payload")
		}

		// Decode the upstream cell header (may be empty, all zeros)
		upstreamPlainTextConnId     := int(binary.BigEndian.Uint32(upstreamPlaintext[0:4]))
		upstreamPlainTextDataLength := int(binary.BigEndian.Uint16(upstreamPlaintext[4:6]))

		if upstreamPlainTextConnId == 0 {
			continue // no upstream data
		}

		//check which connection it belongs to
		//TODO: what is that ?? this is supposed to be anonymous
		conn := conns[upstreamPlainTextConnId]

		// client initiating new connection
		if conn == nil {
			conn = relayNewConn(upstreamPlainTextConnId, downstream)
			conns[upstreamPlainTextConnId] = conn
		}

		if 6+upstreamPlainTextDataLength > payloadLength {
			log.Printf("upstream cell invalid length %d", 6+upstreamPlainTextDataLength)
			continue
		}

		conn <- upstreamPlaintext[6 : 6+upstreamPlainTextDataLength]
	}
}

func relayNewConn(connId int, downstreamData chan<- dataWithConnectionId) chan<- []byte {

	upstreamData := make(chan []byte)
	go relaySocksProxy(connId, upstreamData, downstreamData)
	return upstreamData
}

func connectToTrustee(trusteeId int, trusteeHostAddr string, nClients int, nTrustees int, payloadLength int) {
	//connect
	fmt.Println("Relay connecting to trustee", trusteeId, "on address", trusteeHostAddr)
	conn, err := net.Dial("tcp", trusteeHostAddr)
	if err != nil {
		panic("Can't connect to trustee:" + err.Error())
		//TODO : maybe code something less brutal here
	}

	//tell the trustee server our parameters
	buffer := make([]byte, 20)
	binary.BigEndian.PutUint32(buffer[0:4], uint32(LLD_PROTOCOL_VERSION))
	binary.BigEndian.PutUint32(buffer[4:8], uint32(payloadLength))
	binary.BigEndian.PutUint32(buffer[8:12], uint32(nClients))
	binary.BigEndian.PutUint32(buffer[12:16], uint32(nTrustees))
	binary.BigEndian.PutUint32(buffer[16:20], uint32(trusteeId))

	fmt.Println("Writing", LLD_PROTOCOL_VERSION, "setup is", nClients, nTrustees, "role is", trusteeId, "cellSize ", payloadLength)

	n, err := conn.Write(buffer)

	if n < 1 || err != nil {
		panic("Error writing to socket:" + err.Error())
	}

	// Now read the public key
	buffer2 := make([]byte, 1024)

	// Read the incoming connection into the buffer.
	reqLen, err := conn.Read(buffer2)
	if err != nil {
	    fmt.Println(">>>> Relay : error reading:", err.Error())
	}

	fmt.Println(">>>>  Relay : reading public key", reqLen)
	keySize := int(binary.BigEndian.Uint32(buffer2[4:8]))
	keyBytes := buffer2[8:(8+keySize)]

	fmt.Println(hex.Dump(keyBytes))

	publicKey := suite.Point()
	err2 := publicKey.UnmarshalBinary(keyBytes)

	if err2 != nil {
		panic(">>>>  Relay : can't unmarshal trustee key ! " + err2.Error())
	}

	fmt.Println("Trustee", trusteeId, "is connected.")

	//side effects
	trusteesConnections[trusteeId] = conn
	trusteesPublicKeys[trusteeId]  = publicKey
}<commit_msg>Refactors; Simplify the Statistics<commit_after>package main

import (
	"encoding/binary"
	"fmt"
	"encoding/hex"
	"github.com/lbarman/crypto/abstract"
	"github.com/lbarman/prifi/dcnet"
	"io"
	//"os"
	"strconv"
	"log"
	"strings"
	"net"
	"time"
	log2 "github.com/lbarman/prifi/log"
)

type RelayState struct {
	Name				string

	PublicKey			abstract.Point
	privateKey			abstract.Secret

	nClients			int
	nTrustees			int

	clientsConnections  []net.Conn
	trusteesConnections []net.Conn
	trusteesPublicKeys  []abstract.Point
	clientPublicKeys    []abstract.Point

	CellCoder			dcnet.CellCoder

	MessageHistory		abstract.Cipher

	PayloadLength		int
	ReportingLimit		int
}

func initiateRelayState(nTrustees int, nClients int, payloadLength int, reportingLimit int) *RelayState {

	params := new(RelayState)

	params.Name           = "Relay"
	params.PayloadLength  = payloadLength
	params.ReportingLimit = reportingLimit

	//prepare the crypto parameters
	rand 	:= suite.Cipher([]byte(params.Name))
	base	:= suite.Point().Base()

	//generate own parameters
	params.privateKey       = suite.Secret().Pick(rand)
	params.PublicKey        = suite.Point().Mul(base, params.privateKey)

	params.nClients  = nClients
	params.nTrustees = nTrustees

	//placeholders for pubkeys and connections
	params.trusteesPublicKeys = make([]abstract.Point, nTrustees)
	params.clientPublicKeys   = make([]abstract.Point, nClients)

	params.trusteesConnections = make([]net.Conn, nTrustees)
	params.clientsConnections  = make([]net.Conn, nClients)

	//sets the cell coder, and the history
	params.CellCoder = factory()

	return params
}

func startRelay(payloadLength int, relayPort string, nClients int, nTrustees int, trusteesIp []string, reportingLimit int) {

	relayState := initiateRelayState(nTrustees, nClients, payloadLength, reportingLimit)
	stats := emptyStatistics(reportingLimit)

	//connect to the trustees
	for i:= 0; i < nTrustees; i++ {
		currentTrusteeIp := strings.Replace(trusteesIp[i], "_", ":", -1) //trick for windows shell, where ":" separates args
		connectToTrustee(i, currentTrusteeIp, relayState)
	}

	//starts the client server
	lsock, err := net.Listen("tcp", relayPort)
	if err != nil {
		panic("Can't open listen socket:" + err.Error())
	}

	// Wait for all the clients to connect
	for j := 0; j < nClients; j++ {
		fmt.Printf("Waiting for %d clients (on port %s)\n", nClients-j, relayPort)
		relayAcceptOneClient(lsock, relayState)
	}
	println("All clients and trustees connected.")

	//Prepare the messages
	messageForClient   := MarshalPublicKeyArrayToByteArray(relayState.trusteesPublicKeys)
	messageForTrustees := MarshalPublicKeyArrayToByteArray(relayState.clientPublicKeys)

	//broadcast to the clients
	broadcastMessage(relayState.clientsConnections, messageForClient)
	broadcastMessage(relayState.trusteesConnections, messageForTrustees)

	println("All crypto stuff exchanged !")

	for {
		time.Sleep(5000 * time.Millisecond)
	}

	// Create ciphertext slice bufferfers for all clients and trustees
	clientPayloadLength := relayState.CellCoder.ClientCellSize(payloadLength)
	clientsPayloadData  := make([][]byte, nClients)
	for i := 0; i < nClients; i++ {
		clientsPayloadData[i] = make([]byte, clientPayloadLength)
	}

	trusteePayloadLength := relayState.CellCoder.TrusteeCellSize(payloadLength)
	trusteesPayloadData  := make([][]byte, nTrustees)
	for i := 0; i < nTrustees; i++ {
		trusteesPayloadData[i] = make([]byte, trusteePayloadLength)
	}

	conns := make(map[int]chan<- []byte)
	downstream := make(chan dataWithConnectionId)
	nulldown := dataWithConnectionId{} // default empty downstream cell
	window := 2           // Maximum cells in-flight
	inflight := 0         // Current cells in-flight

	for {

		stats.report(relayState)
		if stats.reportingDone() {
			println("Reporting limit matched; exiting the relay")
			break;
		}

		//TODO : check if it is required to send empty cell
		// See if there's any downstream data to forward.
		var downbuffer dataWithConnectionId
		select {
			case downbuffer = <-downstream: // some data to forward downstream
				//fmt.Println("Downstream data...")
				//fmt.Printf("v %d\n", len(downbuffer)-6)
			default: // nothing at the moment to forward
				downbuffer = nulldown
		}

		downstreamDataPayloadLength := len(downbuffer.data)
		downstreamData := make([]byte, 6+downstreamDataPayloadLength)
		binary.BigEndian.PutUint32(downstreamData[0:4], uint32(downbuffer.connectionId))
		binary.BigEndian.PutUint16(downstreamData[4:6], uint16(downstreamDataPayloadLength))
		copy(downstreamData[6:], downbuffer.data)

		// Broadcast the downstream data to all clients.
		for i := 0; i < nClients; i++ {
			n, err := relayState.clientsConnections[i].Write(downstreamData)

			if n != 6+downstreamDataPayloadLength {
				panic("Relay : Write to client failed, wrote "+strconv.Itoa(downstreamDataPayloadLength+6)+" where "+strconv.Itoa(n)+" was expected : " + err.Error())
			}
		}

		stats.addDownstreamCell(int64(downstreamDataPayloadLength))

		inflight++
		if inflight < window {
			continue // Get more cells in flight
		}

		relayState.CellCoder.DecodeStart(payloadLength, relayState.MessageHistory)

		// Collect a cell ciphertext from each trustee
		for i := 0; i < nTrustees; i++ {

			//TODO: this looks blocking
			n, err := io.ReadFull(relayState.trusteesConnections[i], trusteesPayloadData[i])
			if n < trusteePayloadLength {
				panic("Relay : Read from trustee failed, read "+strconv.Itoa(n)+" where "+strconv.Itoa(trusteePayloadLength)+" was expected: " + err.Error())
			}

			relayState.CellCoder.DecodeTrustee(trusteesPayloadData[i])
		}

		// Collect an upstream ciphertext from each client
		for i := 0; i < nClients; i++ {

			//TODO: this looks blocking
			n, err := io.ReadFull(relayState.clientsConnections[i], clientsPayloadData[i])
			if n < clientPayloadLength {
				panic("Relay : Read from client failed, read "+strconv.Itoa(n)+" where "+strconv.Itoa(clientPayloadLength)+" was expected: " + err.Error())
			}

			relayState.CellCoder.DecodeClient(clientsPayloadData[i])
		}

		upstreamPlaintext := relayState.CellCoder.DecodeCell()
		inflight--

		stats.addUpstreamCell(int64(payloadLength))

		// Process the decoded cell
		if upstreamPlaintext == nil {
			continue // empty or corrupt upstream cell
		}
		if len(upstreamPlaintext) != payloadLength {
			panic("DecodeCell produced wrong-size payload")
		}

		// Decode the upstream cell header (may be empty, all zeros)
		upstreamPlainTextConnId     := int(binary.BigEndian.Uint32(upstreamPlaintext[0:4]))
		upstreamPlainTextDataLength := int(binary.BigEndian.Uint16(upstreamPlaintext[4:6]))

		if upstreamPlainTextConnId == 0 {
			continue // no upstream data
		}

		//check which connection it belongs to
		//TODO: what is that ?? this is supposed to be anonymous
		conn := conns[upstreamPlainTextConnId]

		// client initiating new connection
		if conn == nil {
			conn = relayNewConn(upstreamPlainTextConnId, downstream)
			conns[upstreamPlainTextConnId] = conn
		}

		if 6+upstreamPlainTextDataLength > payloadLength {
			log.Printf("upstream cell invalid length %d", 6+upstreamPlainTextDataLength)
			continue
		}

		conn <- upstreamPlaintext[6 : 6+upstreamPlainTextDataLength]
	}
}

func relayNewConn(connId int, downstreamData chan<- dataWithConnectionId) chan<- []byte {

	upstreamData := make(chan []byte)
	go relaySocksProxy(connId, upstreamData, downstreamData)
	return upstreamData
}

func connectToTrustee(trusteeId int, trusteeHostAddr string, relayState *RelayState) {
	//connect
	fmt.Println("Relay connecting to trustee", trusteeId, "on address", trusteeHostAddr)
	conn, err := net.Dial("tcp", trusteeHostAddr)
	if err != nil {
		panic("Can't connect to trustee:" + err.Error())
		//TODO : maybe code something less brutal here
	}

	//tell the trustee server our parameters
	buffer := make([]byte, 20)
	binary.BigEndian.PutUint32(buffer[0:4], uint32(LLD_PROTOCOL_VERSION))
	binary.BigEndian.PutUint32(buffer[4:8], uint32(relayState.PayloadLength))
	binary.BigEndian.PutUint32(buffer[8:12], uint32(relayState.nClients))
	binary.BigEndian.PutUint32(buffer[12:16], uint32(relayState.nTrustees))
	binary.BigEndian.PutUint32(buffer[16:20], uint32(trusteeId))

	fmt.Println("Writing", LLD_PROTOCOL_VERSION, "setup is", relayState.nClients, relayState.nTrustees, "role is", trusteeId, "cellSize ", relayState.PayloadLength)

	n, err := conn.Write(buffer)

	if n < 1 || err != nil {
		panic("Error writing to socket:" + err.Error())
	}

	// Now read the public key
	buffer2 := make([]byte, 1024)

	// Read the incoming connection into the buffer.
	reqLen, err := conn.Read(buffer2)
	if err != nil {
	    fmt.Println(">>>> Relay : error reading:", err.Error())
	}

	fmt.Println(">>>>  Relay : reading public key", reqLen)
	keySize := int(binary.BigEndian.Uint32(buffer2[4:8]))
	keyBytes := buffer2[8:(8+keySize)]

	fmt.Println(hex.Dump(keyBytes))

	publicKey := suite.Point()
	err2 := publicKey.UnmarshalBinary(keyBytes)

	if err2 != nil {
		panic(">>>>  Relay : can't unmarshal trustee key ! " + err2.Error())
	}

	fmt.Println("Trustee", trusteeId, "is connected.")

	//side effects
	relayState.trusteesConnections[trusteeId] = conn
	relayState.trusteesPublicKeys[trusteeId]  = publicKey
}

func relayAcceptOneClient(listeningSocket net.Listener, relayState *RelayState) {
	conn, err := listeningSocket.Accept()

	if err != nil {
		panic("Listen error:" + err.Error())
	}

	buffer := make([]byte, 512)
	_, err2 := conn.Read(buffer)
	if err2 != nil {
		panic("Read error:" + err2.Error())
	}

	version := int(binary.BigEndian.Uint32(buffer[0:4]))

	if(version != LLD_PROTOCOL_VERSION) {
		fmt.Println(">>>> Relay client version", version, "!= relay version", LLD_PROTOCOL_VERSION)
		panic("fatal error")
	}

	nodeId := int(binary.BigEndian.Uint32(buffer[4:8]))
	keySize := int(binary.BigEndian.Uint32(buffer[8:12]))
	keyBytes := buffer[12:(12+keySize)]

	publicKey := suite.Point()
	err3 := publicKey.UnmarshalBinary(keyBytes)

	if err3 != nil {
		panic(">>>>  Relay : can't unmarshal client key ! " + err3.Error())
	}

	if nodeId < 0 || nodeId >= relayState.nClients {
		panic("illegal node number")
	}

	//side effect
	relayState.clientsConnections[nodeId] = conn
	relayState.clientPublicKeys[nodeId] = publicKey
}

type Statistics struct {
	begin			time.Time
	nextReport		time.Time
	nReports		int
	maxNReports		int
	period			time.Duration

	totupcells		int64
	totupbytes 		int64
	totdowncells 	int64
	totdownbytes 	int64
	parupcells		int64
	parupbytes 		int64
	pardownbytes	int64
}

func emptyStatistics(reportingLimit int) *Statistics{
	stats := Statistics{time.Now(), time.Now(), 0, reportingLimit, time.Duration(3)*time.Second, 0, 0, 0, 0, 0, 0, 0}
	return &stats
}

func (stats *Statistics) reportingDone() bool {
	return stats.nReports >= stats.maxNReports
}

func (stats *Statistics) addDownstreamCell(nBytes int64) {
	stats.totdowncells += 1
	stats.totdownbytes += nBytes
	stats.pardownbytes += nBytes
}

func (stats *Statistics) addUpstreamCell(nBytes int64) {
	stats.totupcells += 1
	stats.totupbytes += nBytes
	stats.parupcells += 1
	stats.parupbytes += nBytes
}

func (stats *Statistics) report(state *RelayState) {
	now := time.Now()
	if now.After(stats.nextReport) {
		duration := now.Sub(stats.begin).Seconds()
		instantUpSpeed := (float64(stats.parupbytes)/stats.period.Seconds())

		fmt.Printf("@ %fs; cell %f (%f) /sec, up %f (%f) B/s, down %f (%f) B/s\n",
			duration,
			 float64(stats.totupcells)/duration, float64(stats.parupcells)/stats.period.Seconds(),
			 float64(stats.totupbytes)/duration, instantUpSpeed,
			 float64(stats.totdownbytes)/duration, float64(stats.pardownbytes)/stats.period.Seconds())

		// Next report time
		stats.parupcells = 0
		stats.parupbytes = 0
		stats.pardownbytes = 0

		//log2.BenchmarkFloat(fmt.Sprintf("cellsize-%d-upstream-bytes", payloadLength), instantUpSpeed)

		//write JSON
		data := struct {
		    Experiment string
		    CellSize int
		    Speed float64
		}{
		    "upstream-speed-given-cellsize",
		    state.PayloadLength,
		    instantUpSpeed,
		}
		log2.JsonDump(data)

		stats.nextReport = now.Add(stats.period)
		stats.nReports += 1
	}
}<|endoftext|>


<commit_before>package relay

import (
	"log"
	"strings"
	"sync"
	"time"

	"github.com/armon/relay/broker"
)

// retryBroker implements a resilient AMQP broker which automatically replaces
// the underlying channel if an unexpected error is observed. Due to AMQP's
// persistent connections, this becomes necessary to survive even minor network
// hiccups which may temporarily sever the connection to the server.
//
// The retryBroker does not come without a price, however. Clients are still
// responsible for deduplicating messages or handling them idempotently, as
// duplications are a regular happening in the face of network errors due to
// message redelivery provided by the protocol. It is also important to note
// that acknowledgements - both ack and nack - cannot be retried across
// connection resets. This stems from the fact that a given consumer's channel
// cannot be revived and must be thrown away if the consumer is to be reused.
type retryBroker struct {
	broker   *relayBroker
	attempts int
	min      time.Duration
	max      time.Duration
}

// RetryBroker returns a new retrying broker with the given settings.
func (r *Relay) RetryBroker(attempts int, min, max time.Duration) *retryBroker {
	return &retryBroker{&relayBroker{r}, attempts, min, max}
}

// Publisher returns a new retrying broker.Publisher.
func (rb *retryBroker) Publisher(queue string) (broker.Publisher, error) {
	return &retryPublisher{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Consumer returns a new retrying broker.Consumer.
func (rb *retryBroker) Consumer(queue string) (broker.Consumer, error) {
	return &retryConsumer{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Close closes the broker.
func (rb *retryBroker) Close() error {
	if rb.broker != nil {
		return rb.broker.Close()
	}
	return nil
}

// retryConsumer tracks the relay broker and an associated consumer, allowing
// the consumer to be replaced transparently in the face of failure.
type retryConsumer struct {
	broker   broker.Broker
	cons     broker.Consumer
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// consumer is used to connect the consumer to the relay queue.
func (rc *retryConsumer) consumer(create bool) (broker.Consumer, error) {
	// Check for an existing consumer
	rc.l.RLock()
	cons := rc.cons
	rc.l.RUnlock()

	if cons != nil || !create {
		return cons, nil
	}

	// Make a new consumer
	rc.l.Lock()
	defer rc.l.Unlock()

	cons, err := rc.broker.Consumer(rc.queue)
	if err != nil {
		return nil, err
	}
	rc.cons = cons
	return cons, nil
}

// discard is used to remove a known-bad consumer.
func (rc *retryConsumer) discard(cons broker.Consumer) {
	cons.Close()

	rc.l.Lock()
	defer rc.l.Unlock()
	if cons == rc.cons {
		rc.cons = nil
	}
}

// Close closes the consumer.
func (rc *retryConsumer) Close() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Close()
}

// Ack marks message(s) as delivered.
func (rc *retryConsumer) Ack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Ack()
}

// Nack sends message(s) back to the queue.
func (rc *retryConsumer) Nack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Nack()
}

// ConsumeAck is used to consume with automatic ack.
func (rc *retryConsumer) ConsumeAck(out interface{}) error {
	if err := rc.Consume(out); err != nil {
		return err
	}
	if err := rc.Ack(); err != nil {
		return err
	}
	return nil
}

// Consume consumes a single message from the queue.
func (rc *retryConsumer) Consume(out interface{}) error {
	return rc.ConsumeTimeout(out, -1)
}

// ConsumeTimeout consumes a single message from the queue with an upper bound
// on the time spent waiting. This behaves slightly different in the retry
// broker, as we may encounter errors getting a valid consumer. In this case,
// the timeout may actually be longer than specified to allow reconnection
// attempts to take place.
func (rc *retryConsumer) ConsumeTimeout(out interface{}, timeout time.Duration) error {
	wait := rc.min
	for i := 0; ; i++ {
		cons, err := rc.consumer(true)
		if err != nil {
			goto RETRY
		}

		err = cons.ConsumeTimeout(out, timeout)
		if err == nil {
			break
		}

		// Respect the case where the consume legitimately times out.
		if err == TimedOut {
			return err
		}

		// Retries won't help decode errors, so just return.
		if strings.Contains(err.Error(), "Failed to decode") {
			return err
		}

	RETRY:
		log.Printf("[ERR] relay: consumer got error: %v", err)
		rc.discard(cons)
		if i == rc.attempts {
			log.Printf("[ERR] relay: consumer giving up after %d attempts", i)
			return err
		}
		log.Printf("[DEBUG] relay: consumer retrying in %s", wait)
		time.Sleep(wait)
		wait *= 2
		if wait > rc.max {
			wait = rc.max
		}
	}
	return nil
}

// retryPublisher wraps a relay broker and an associated publisher, allowing
// the underlying publisher's channel to be automatically swapped out if any
// errors are seen during a publish.
type retryPublisher struct {
	broker   broker.Broker
	pub      broker.Publisher
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// publisher is used to connect the publisher to the relay queue.
func (rp *retryPublisher) publisher(create bool) (broker.Publisher, error) {
	// Check for an existing publisher
	rp.l.RLock()
	pub := rp.pub
	rp.l.RUnlock()

	if pub != nil || !create {
		return pub, nil
	}

	// Make a new publisher
	rp.l.Lock()
	defer rp.l.Unlock()

	pub, err := rp.broker.Publisher(rp.queue)
	if err != nil {
		return nil, err
	}
	rp.pub = pub
	return pub, nil
}

// discard is used to remove a known-bad publisher.
func (rp *retryPublisher) discard(pub broker.Publisher) {
	pub.Close()

	rp.l.Lock()
	defer rp.l.Unlock()
	if pub == rp.pub {
		rp.pub = nil
	}
}

// Close closes the publisher.
func (rp *retryPublisher) Close() error {
	pub, err := rp.publisher(false)
	if err != nil || pub == nil {
		return err
	}
	return pub.Close()
}

// Publish publishes a single message to the queue. If an error is encountered,
// the broker automatically tries to replace the underlying channel.
func (rp *retryPublisher) Publish(in interface{}) error {
	wait := rp.min
	for i := 0; ; i++ {
		pub, err := rp.publisher(true)
		if err != nil {
			goto RETRY
		}

		err = pub.Publish(in)
		if err == nil {
			break
		}

	RETRY:
		log.Printf("[ERR] relay: publisher got error: %v", err)
		rp.discard(pub)
		if i == rp.attempts {
			log.Printf("[ERR] relay: publisher giving up after %d attempts", i)
			return err
		}
		log.Printf("[DEBUG] relay: publisher retrying in %s", wait)
		time.Sleep(wait)
		wait *= 2
		if wait > rp.max {
			wait = rp.max
		}
	}
	return nil
}
<commit_msg>Guard nil during discard<commit_after>package relay

import (
	"log"
	"sync"
	"time"

	"github.com/armon/relay/broker"
)

// retryBroker implements a resilient AMQP broker which automatically replaces
// the underlying channel if an unexpected error is observed. Due to AMQP's
// persistent connections, this becomes necessary to survive even minor network
// hiccups which may temporarily sever the connection to the server.
//
// The retryBroker does not come without a price, however. Clients are still
// responsible for deduplicating messages or handling them idempotently, as
// duplications are a regular happening in the face of network errors due to
// message redelivery provided by the protocol. It is also important to note
// that acknowledgements - both ack and nack - cannot be retried across
// connection resets. This stems from the fact that a given consumer's channel
// cannot be revived and must be thrown away if the consumer is to be reused.
type retryBroker struct {
	broker   *relayBroker
	attempts int
	min      time.Duration
	max      time.Duration
}

// RetryBroker returns a new retrying broker with the given settings.
func (r *Relay) RetryBroker(attempts int, min, max time.Duration) *retryBroker {
	return &retryBroker{&relayBroker{r}, attempts, min, max}
}

// Publisher returns a new retrying broker.Publisher.
func (rb *retryBroker) Publisher(queue string) (broker.Publisher, error) {
	return &retryPublisher{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Consumer returns a new retrying broker.Consumer.
func (rb *retryBroker) Consumer(queue string) (broker.Consumer, error) {
	return &retryConsumer{
		broker:   rb.broker,
		queue:    queue,
		attempts: rb.attempts,
		min:      rb.min,
		max:      rb.max,
	}, nil
}

// Close closes the broker.
func (rb *retryBroker) Close() error {
	if rb.broker != nil {
		return rb.broker.Close()
	}
	return nil
}

// retryConsumer tracks the relay broker and an associated consumer, allowing
// the consumer to be replaced transparently in the face of failure.
type retryConsumer struct {
	broker   broker.Broker
	cons     broker.Consumer
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// consumer is used to connect the consumer to the relay queue.
func (rc *retryConsumer) consumer(create bool) (broker.Consumer, error) {
	// Check for an existing consumer
	rc.l.RLock()
	cons := rc.cons
	rc.l.RUnlock()

	if cons != nil || !create {
		return cons, nil
	}

	// Make a new consumer
	rc.l.Lock()
	defer rc.l.Unlock()

	cons, err := rc.broker.Consumer(rc.queue)
	if err != nil {
		return nil, err
	}
	rc.cons = cons
	return cons, nil
}

// discard is used to remove a known-bad consumer.
func (rc *retryConsumer) discard(cons broker.Consumer) {
	if cons == nil {
		return
	}
	cons.Close()

	rc.l.Lock()
	defer rc.l.Unlock()
	if cons == rc.cons {
		rc.cons = nil
	}
}

// Close closes the consumer.
func (rc *retryConsumer) Close() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Close()
}

// Ack marks message(s) as delivered.
func (rc *retryConsumer) Ack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Ack()
}

// Nack sends message(s) back to the queue.
func (rc *retryConsumer) Nack() error {
	cons, err := rc.consumer(false)
	if err != nil || cons == nil {
		return err
	}
	return cons.Nack()
}

// ConsumeAck is used to consume with automatic ack.
func (rc *retryConsumer) ConsumeAck(out interface{}) error {
	if err := rc.Consume(out); err != nil {
		return err
	}
	if err := rc.Ack(); err != nil {
		return err
	}
	return nil
}

// Consume consumes a single message from the queue.
func (rc *retryConsumer) Consume(out interface{}) error {
	return rc.ConsumeTimeout(out, -1)
}

// ConsumeTimeout consumes a single message from the queue with an upper bound
// on the time spent waiting. This behaves slightly different in the retry
// broker, as we may encounter errors getting a valid consumer. In this case,
// the timeout may actually be longer than specified to allow reconnection
// attempts to take place.
func (rc *retryConsumer) ConsumeTimeout(out interface{}, timeout time.Duration) error {
	wait := rc.min
	for i := 0; ; i++ {
		cons, err := rc.consumer(true)
		if err != nil {
			goto RETRY
		}

		err = cons.ConsumeTimeout(out, timeout)
		if err == nil {
			break
		}

		// Respect the case where the consume legitimately times out.
		if err == TimedOut {
			return err
		}

	RETRY:
		log.Printf("[ERR] relay: consumer got error: %v", err)
		rc.discard(cons)
		if i == rc.attempts {
			log.Printf("[ERR] relay: consumer giving up after %d attempts", i)
			return err
		}
		log.Printf("[DEBUG] relay: consumer retrying in %s", wait)
		time.Sleep(wait)
		wait *= 2
		if wait > rc.max {
			wait = rc.max
		}
	}
	return nil
}

// retryPublisher wraps a relay broker and an associated publisher, allowing
// the underlying publisher's channel to be automatically swapped out if any
// errors are seen during a publish.
type retryPublisher struct {
	broker   broker.Broker
	pub      broker.Publisher
	queue    string
	attempts int
	min      time.Duration
	max      time.Duration
	l        sync.RWMutex
}

// publisher is used to connect the publisher to the relay queue.
func (rp *retryPublisher) publisher(create bool) (broker.Publisher, error) {
	// Check for an existing publisher
	rp.l.RLock()
	pub := rp.pub
	rp.l.RUnlock()

	if pub != nil || !create {
		return pub, nil
	}

	// Make a new publisher
	rp.l.Lock()
	defer rp.l.Unlock()

	pub, err := rp.broker.Publisher(rp.queue)
	if err != nil {
		return nil, err
	}
	rp.pub = pub
	return pub, nil
}

// discard is used to remove a known-bad publisher.
func (rp *retryPublisher) discard(pub broker.Publisher) {
	if pub == nil {
		return
	}
	pub.Close()

	rp.l.Lock()
	defer rp.l.Unlock()
	if pub == rp.pub {
		rp.pub = nil
	}
}

// Close closes the publisher.
func (rp *retryPublisher) Close() error {
	pub, err := rp.publisher(false)
	if err != nil || pub == nil {
		return err
	}
	return pub.Close()
}

// Publish publishes a single message to the queue. If an error is encountered,
// the broker automatically tries to replace the underlying channel.
func (rp *retryPublisher) Publish(in interface{}) error {
	wait := rp.min
	for i := 0; ; i++ {
		pub, err := rp.publisher(true)
		if err != nil {
			goto RETRY
		}

		err = pub.Publish(in)
		if err == nil {
			break
		}

	RETRY:
		log.Printf("[ERR] relay: publisher got error: %v", err)
		rp.discard(pub)
		if i == rp.attempts {
			log.Printf("[ERR] relay: publisher giving up after %d attempts", i)
			return err
		}
		log.Printf("[DEBUG] relay: publisher retrying in %s", wait)
		time.Sleep(wait)
		wait *= 2
		if wait > rp.max {
			wait = rp.max
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/rds"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsRDSClusterParameterGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsRDSClusterParameterGroupCreate,
		Read:   resourceAwsRDSClusterParameterGroupRead,
		Update: resourceAwsRDSClusterParameterGroupUpdate,
		Delete: resourceAwsRDSClusterParameterGroupDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"name": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ForceNew:      true,
				ConflictsWith: []string{"name_prefix"},
				ValidateFunc:  validateDbParamGroupName,
			},
			"name_prefix": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ForceNew:     true,
				ValidateFunc: validateDbParamGroupNamePrefix,
			},
			"family": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  "Managed by Terraform",
			},
			"parameter": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: false,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"value": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"apply_method": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "immediate",
							// this parameter is not actually state, but a
							// meta-parameter describing how the RDS API call
							// to modify the parameter group should be made.
							// Future reads of the resource from AWS don't tell
							// us what we used for apply_method previously, so
							// by squashing state to an empty string we avoid
							// needing to do an update for every future run.
							StateFunc: func(interface{}) string { return "" },
						},
					},
				},
				Set: resourceAwsDbParameterHash,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsRDSClusterParameterGroupCreate(d *schema.ResourceData, meta interface{}) error {
	rdsconn := meta.(*AWSClient).rdsconn
	tags := tagsFromMapRDS(d.Get("tags").(map[string]interface{}))

	var groupName string
	if v, ok := d.GetOk("name"); ok {
		groupName = v.(string)
	} else if v, ok := d.GetOk("name_prefix"); ok {
		groupName = resource.PrefixedUniqueId(v.(string))
	} else {
		groupName = resource.UniqueId()
	}

	createOpts := rds.CreateDBClusterParameterGroupInput{
		DBClusterParameterGroupName: aws.String(groupName),
		DBParameterGroupFamily:      aws.String(d.Get("family").(string)),
		Description:                 aws.String(d.Get("description").(string)),
		Tags:                        tags,
	}

	log.Printf("[DEBUG] Create DB Cluster Parameter Group: %#v", createOpts)
	_, err := rdsconn.CreateDBClusterParameterGroup(&createOpts)
	if err != nil {
		return fmt.Errorf("Error creating DB Cluster Parameter Group: %s", err)
	}

	d.SetId(*createOpts.DBClusterParameterGroupName)
	log.Printf("[INFO] DB Cluster Parameter Group ID: %s", d.Id())

	return resourceAwsRDSClusterParameterGroupUpdate(d, meta)
}

func resourceAwsRDSClusterParameterGroupRead(d *schema.ResourceData, meta interface{}) error {
	rdsconn := meta.(*AWSClient).rdsconn

	describeOpts := rds.DescribeDBClusterParameterGroupsInput{
		DBClusterParameterGroupName: aws.String(d.Id()),
	}

	describeResp, err := rdsconn.DescribeDBClusterParameterGroups(&describeOpts)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "DBParameterGroupNotFound" {
			log.Printf("[WARN] DB Cluster Parameter Group (%s) not found, error code (404)", d.Id())
			d.SetId("")
			return nil
		}

		return err
	}

	if len(describeResp.DBClusterParameterGroups) != 1 ||
		*describeResp.DBClusterParameterGroups[0].DBClusterParameterGroupName != d.Id() {
		return fmt.Errorf("Unable to find Cluster Parameter Group: %#v", describeResp.DBClusterParameterGroups)
	}

	d.Set("name", describeResp.DBClusterParameterGroups[0].DBClusterParameterGroupName)
	d.Set("family", describeResp.DBClusterParameterGroups[0].DBParameterGroupFamily)
	d.Set("description", describeResp.DBClusterParameterGroups[0].Description)

	// Only include user customized parameters as there's hundreds of system/default ones
	describeParametersOpts := rds.DescribeDBClusterParametersInput{
		DBClusterParameterGroupName: aws.String(d.Id()),
		Source: aws.String("user"),
	}

	describeParametersResp, err := rdsconn.DescribeDBClusterParameters(&describeParametersOpts)
	if err != nil {
		return err
	}

	d.Set("parameter", flattenParameters(describeParametersResp.Parameters))

	paramGroup := describeResp.DBClusterParameterGroups[0]
	arn, err := buildRDSCPGARN(d.Id(), meta.(*AWSClient).partition, meta.(*AWSClient).accountid, meta.(*AWSClient).region)
	if err != nil {
		name := "<empty>"
		if paramGroup.DBClusterParameterGroupName != nil && *paramGroup.DBClusterParameterGroupName != "" {
			name = *paramGroup.DBClusterParameterGroupName
		}
		log.Printf("[DEBUG] Error building ARN for DB Cluster Parameter Group, not setting Tags for Cluster Param Group %s", name)
	} else {
		d.Set("arn", arn)
		resp, err := rdsconn.ListTagsForResource(&rds.ListTagsForResourceInput{
			ResourceName: aws.String(arn),
		})

		if err != nil {
			log.Printf("[DEBUG] Error retrieving tags for ARN: %s", arn)
		}

		var dt []*rds.Tag
		if len(resp.TagList) > 0 {
			dt = resp.TagList
		}
		d.Set("tags", tagsToMapRDS(dt))
	}

	return nil
}

func resourceAwsRDSClusterParameterGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	rdsconn := meta.(*AWSClient).rdsconn

	d.Partial(true)

	if d.HasChange("parameter") {
		o, n := d.GetChange("parameter")
		if o == nil {
			o = new(schema.Set)
		}
		if n == nil {
			n = new(schema.Set)
		}

		os := o.(*schema.Set)
		ns := n.(*schema.Set)

		// Expand the "parameter" set to aws-sdk-go compat []rds.Parameter
		parameters, err := expandParameters(ns.Difference(os).List())
		if err != nil {
			return err
		}

		if len(parameters) > 0 {
			// We can only modify 20 parameters at a time, so walk them until
			// we've got them all.
			maxParams := 20
			for parameters != nil {
				paramsToModify := make([]*rds.Parameter, 0)
				if len(parameters) <= maxParams {
					paramsToModify, parameters = parameters[:], nil
				} else {
					paramsToModify, parameters = parameters[:maxParams], parameters[maxParams:]
				}
				modifyOpts := rds.ModifyDBClusterParameterGroupInput{
					DBClusterParameterGroupName: aws.String(d.Get("name").(string)),
					Parameters:                  paramsToModify,
				}

				log.Printf("[DEBUG] Modify DB Cluster Parameter Group: %s", modifyOpts)
				_, err = rdsconn.ModifyDBClusterParameterGroup(&modifyOpts)
				if err != nil {
					return fmt.Errorf("Error modifying DB Cluster Parameter Group: %s", err)
				}
			}
			d.SetPartial("parameter")
		}
	}

	if arn, err := buildRDSCPGARN(d.Id(), meta.(*AWSClient).partition, meta.(*AWSClient).accountid, meta.(*AWSClient).region); err == nil {
		if err := setTagsRDS(rdsconn, d, arn); err != nil {
			return err
		} else {
			d.SetPartial("tags")
		}
	}

	d.Partial(false)

	return resourceAwsRDSClusterParameterGroupRead(d, meta)
}

func resourceAwsRDSClusterParameterGroupDelete(d *schema.ResourceData, meta interface{}) error {
	stateConf := &resource.StateChangeConf{
		Pending:    []string{"pending"},
		Target:     []string{"destroyed"},
		Refresh:    resourceAwsRDSClusterParameterGroupDeleteRefreshFunc(d, meta),
		Timeout:    3 * time.Minute,
		MinTimeout: 1 * time.Second,
	}
	_, err := stateConf.WaitForState()
	return err
}

func resourceAwsRDSClusterParameterGroupDeleteRefreshFunc(
	d *schema.ResourceData,
	meta interface{}) resource.StateRefreshFunc {
	rdsconn := meta.(*AWSClient).rdsconn

	return func() (interface{}, string, error) {

		deleteOpts := rds.DeleteDBClusterParameterGroupInput{
			DBClusterParameterGroupName: aws.String(d.Id()),
		}

		if _, err := rdsconn.DeleteDBClusterParameterGroup(&deleteOpts); err != nil {
			rdserr, ok := err.(awserr.Error)
			if !ok {
				return d, "error", err
			}

			if rdserr.Code() != "DBParameterGroupNotFound" {
				return d, "error", err
			}
		}

		return d, "destroyed", nil
	}
}

func buildRDSCPGARN(identifier, partition, accountid, region string) (string, error) {
	if partition == "" {
		return "", fmt.Errorf("Unable to construct RDS Cluster ARN because of missing AWS partition")
	}
	if accountid == "" {
		return "", fmt.Errorf("Unable to construct RDS Cluster ARN because of missing AWS Account ID")
	}
	arn := fmt.Sprintf("arn:%s:rds:%s:%s:cluster-pg:%s", partition, region, accountid, identifier)
	return arn, nil

}
<commit_msg>:lipstick: Support > 20 parameters in rds cluster parameter groups<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/rds"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

const rdsClusterParameterGroupMaxParamsBulkEdit = 20

func resourceAwsRDSClusterParameterGroup() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsRDSClusterParameterGroupCreate,
		Read:   resourceAwsRDSClusterParameterGroupRead,
		Update: resourceAwsRDSClusterParameterGroupUpdate,
		Delete: resourceAwsRDSClusterParameterGroupDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"name": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ForceNew:      true,
				ConflictsWith: []string{"name_prefix"},
				ValidateFunc:  validateDbParamGroupName,
			},
			"name_prefix": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ForceNew:     true,
				ValidateFunc: validateDbParamGroupNamePrefix,
			},
			"family": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  "Managed by Terraform",
			},
			"parameter": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				ForceNew: false,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"value": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"apply_method": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "immediate",
							// this parameter is not actually state, but a
							// meta-parameter describing how the RDS API call
							// to modify the parameter group should be made.
							// Future reads of the resource from AWS don't tell
							// us what we used for apply_method previously, so
							// by squashing state to an empty string we avoid
							// needing to do an update for every future run.
							StateFunc: func(interface{}) string { return "" },
						},
					},
				},
				Set: resourceAwsDbParameterHash,
			},

			"tags": tagsSchema(),
		},
	}
}

func resourceAwsRDSClusterParameterGroupCreate(d *schema.ResourceData, meta interface{}) error {
	rdsconn := meta.(*AWSClient).rdsconn
	tags := tagsFromMapRDS(d.Get("tags").(map[string]interface{}))

	var groupName string
	if v, ok := d.GetOk("name"); ok {
		groupName = v.(string)
	} else if v, ok := d.GetOk("name_prefix"); ok {
		groupName = resource.PrefixedUniqueId(v.(string))
	} else {
		groupName = resource.UniqueId()
	}

	createOpts := rds.CreateDBClusterParameterGroupInput{
		DBClusterParameterGroupName: aws.String(groupName),
		DBParameterGroupFamily:      aws.String(d.Get("family").(string)),
		Description:                 aws.String(d.Get("description").(string)),
		Tags:                        tags,
	}

	log.Printf("[DEBUG] Create DB Cluster Parameter Group: %#v", createOpts)
	_, err := rdsconn.CreateDBClusterParameterGroup(&createOpts)
	if err != nil {
		return fmt.Errorf("Error creating DB Cluster Parameter Group: %s", err)
	}

	d.SetId(*createOpts.DBClusterParameterGroupName)
	log.Printf("[INFO] DB Cluster Parameter Group ID: %s", d.Id())

	return resourceAwsRDSClusterParameterGroupUpdate(d, meta)
}

func resourceAwsRDSClusterParameterGroupRead(d *schema.ResourceData, meta interface{}) error {
	rdsconn := meta.(*AWSClient).rdsconn

	describeOpts := rds.DescribeDBClusterParameterGroupsInput{
		DBClusterParameterGroupName: aws.String(d.Id()),
	}

	describeResp, err := rdsconn.DescribeDBClusterParameterGroups(&describeOpts)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "DBParameterGroupNotFound" {
			log.Printf("[WARN] DB Cluster Parameter Group (%s) not found, error code (404)", d.Id())
			d.SetId("")
			return nil
		}

		return err
	}

	if len(describeResp.DBClusterParameterGroups) != 1 ||
		*describeResp.DBClusterParameterGroups[0].DBClusterParameterGroupName != d.Id() {
		return fmt.Errorf("Unable to find Cluster Parameter Group: %#v", describeResp.DBClusterParameterGroups)
	}

	d.Set("name", describeResp.DBClusterParameterGroups[0].DBClusterParameterGroupName)
	d.Set("family", describeResp.DBClusterParameterGroups[0].DBParameterGroupFamily)
	d.Set("description", describeResp.DBClusterParameterGroups[0].Description)

	// Only include user customized parameters as there's hundreds of system/default ones
	describeParametersOpts := rds.DescribeDBClusterParametersInput{
		DBClusterParameterGroupName: aws.String(d.Id()),
		Source: aws.String("user"),
	}

	describeParametersResp, err := rdsconn.DescribeDBClusterParameters(&describeParametersOpts)
	if err != nil {
		return err
	}

	d.Set("parameter", flattenParameters(describeParametersResp.Parameters))

	paramGroup := describeResp.DBClusterParameterGroups[0]
	arn, err := buildRDSCPGARN(d.Id(), meta.(*AWSClient).partition, meta.(*AWSClient).accountid, meta.(*AWSClient).region)
	if err != nil {
		name := "<empty>"
		if paramGroup.DBClusterParameterGroupName != nil && *paramGroup.DBClusterParameterGroupName != "" {
			name = *paramGroup.DBClusterParameterGroupName
		}
		log.Printf("[DEBUG] Error building ARN for DB Cluster Parameter Group, not setting Tags for Cluster Param Group %s", name)
	} else {
		d.Set("arn", arn)
		resp, err := rdsconn.ListTagsForResource(&rds.ListTagsForResourceInput{
			ResourceName: aws.String(arn),
		})

		if err != nil {
			log.Printf("[DEBUG] Error retrieving tags for ARN: %s", arn)
		}

		var dt []*rds.Tag
		if len(resp.TagList) > 0 {
			dt = resp.TagList
		}
		d.Set("tags", tagsToMapRDS(dt))
	}

	return nil
}

func resourceAwsRDSClusterParameterGroupUpdate(d *schema.ResourceData, meta interface{}) error {
	rdsconn := meta.(*AWSClient).rdsconn

	d.Partial(true)

	if d.HasChange("parameter") {
		o, n := d.GetChange("parameter")
		if o == nil {
			o = new(schema.Set)
		}
		if n == nil {
			n = new(schema.Set)
		}

		os := o.(*schema.Set)
		ns := n.(*schema.Set)

		// Expand the "parameter" set to aws-sdk-go compat []rds.Parameter
		parameters, err := expandParameters(ns.Difference(os).List())
		if err != nil {
			return err
		}

		if len(parameters) > 0 {
			// We can only modify 20 parameters at a time, so walk them until
			// we've got them all.
			for parameters != nil {
				paramsToModify := make([]*rds.Parameter, 0)
				if len(parameters) <= rdsClusterParameterGroupMaxParamsBulkEdit {
					paramsToModify, parameters = parameters[:], nil
				} else {
					paramsToModify, parameters = parameters[:rdsClusterParameterGroupMaxParamsBulkEdit], parameters[rdsClusterParameterGroupMaxParamsBulkEdit:]
				}
				parameterGroupName := d.Get("name").(string)
				modifyOpts := rds.ModifyDBClusterParameterGroupInput{
					DBClusterParameterGroupName: aws.String(parameterGroupName),
					Parameters:                  paramsToModify,
				}

				log.Printf("[DEBUG] Modify DB Cluster Parameter Group: %s", modifyOpts)
				_, err = rdsconn.ModifyDBClusterParameterGroup(&modifyOpts)
				if err != nil {
					return fmt.Errorf("Error modifying DB Cluster Parameter Group: %s", err)
				}
			}
			d.SetPartial("parameter")
		}
	}

	if arn, err := buildRDSCPGARN(d.Id(), meta.(*AWSClient).partition, meta.(*AWSClient).accountid, meta.(*AWSClient).region); err == nil {
		if err := setTagsRDS(rdsconn, d, arn); err != nil {
			return err
		} else {
			d.SetPartial("tags")
		}
	}

	d.Partial(false)

	return resourceAwsRDSClusterParameterGroupRead(d, meta)
}

func resourceAwsRDSClusterParameterGroupDelete(d *schema.ResourceData, meta interface{}) error {
	stateConf := &resource.StateChangeConf{
		Pending:    []string{"pending"},
		Target:     []string{"destroyed"},
		Refresh:    resourceAwsRDSClusterParameterGroupDeleteRefreshFunc(d, meta),
		Timeout:    3 * time.Minute,
		MinTimeout: 1 * time.Second,
	}
	_, err := stateConf.WaitForState()
	return err
}

func resourceAwsRDSClusterParameterGroupDeleteRefreshFunc(
	d *schema.ResourceData,
	meta interface{}) resource.StateRefreshFunc {
	rdsconn := meta.(*AWSClient).rdsconn

	return func() (interface{}, string, error) {

		deleteOpts := rds.DeleteDBClusterParameterGroupInput{
			DBClusterParameterGroupName: aws.String(d.Id()),
		}

		if _, err := rdsconn.DeleteDBClusterParameterGroup(&deleteOpts); err != nil {
			rdserr, ok := err.(awserr.Error)
			if !ok {
				return d, "error", err
			}

			if rdserr.Code() != "DBParameterGroupNotFound" {
				return d, "error", err
			}
		}

		return d, "destroyed", nil
	}
}

func buildRDSCPGARN(identifier, partition, accountid, region string) (string, error) {
	if partition == "" {
		return "", fmt.Errorf("Unable to construct RDS Cluster ARN because of missing AWS partition")
	}
	if accountid == "" {
		return "", fmt.Errorf("Unable to construct RDS Cluster ARN because of missing AWS Account ID")
	}
	arn := fmt.Sprintf("arn:%s:rds:%s:%s:cluster-pg:%s", partition, region, accountid, identifier)
	return arn, nil

}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ssm"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsSsmMaintenanceWindowTask() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSsmMaintenanceWindowTaskCreate,
		Read:   resourceAwsSsmMaintenanceWindowTaskRead,
		Delete: resourceAwsSsmMaintenanceWindowTaskDelete,

		Schema: map[string]*schema.Schema{
			"window_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"max_concurrency": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"max_errors": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"task_type": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"task_arn": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"service_role_arn": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"targets": {
				Type:     schema.TypeList,
				Required: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"key": {
							Type:     schema.TypeString,
							Required: true,
							ForceNew: true,
						},
						"values": {
							Type:     schema.TypeList,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},
					},
				},
			},

			"name": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				ValidateFunc: validateAwsSSMMaintenanceWindowTaskName,
			},

			"description": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				ValidateFunc: validateAwsSSMMaintenanceWindowTaskDescription,
			},

			"priority": {
				Type:     schema.TypeInt,
				Optional: true,
				ForceNew: true,
			},

			"logging_info": {
				Type:     schema.TypeList,
				MaxItems: 1,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"s3_bucket_name": {
							Type:     schema.TypeString,
							Required: true,
						},
						"s3_region": {
							Type:     schema.TypeString,
							Required: true,
						},
						"s3_bucket_prefix": {
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},

			"task_parameters": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": {
							Type:     schema.TypeString,
							Required: true,
							ForceNew: true,
						},
						"values": {
							Type:     schema.TypeList,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},
					},
				},
			},
		},
	}
}

func expandAwsSsmMaintenanceWindowLoggingInfo(config []interface{}) *ssm.LoggingInfo {

	loggingConfig := config[0].(map[string]interface{})

	loggingInfo := &ssm.LoggingInfo{
		S3BucketName: aws.String(loggingConfig["s3_bucket_name"].(string)),
		S3Region:     aws.String(loggingConfig["s3_region"].(string)),
	}

	if s := loggingConfig["s3_bucket_prefix"].(string); s != "" {
		loggingInfo.S3KeyPrefix = aws.String(s)
	}

	return loggingInfo
}

func flattenAwsSsmMaintenanceWindowLoggingInfo(loggingInfo *ssm.LoggingInfo) []interface{} {

	result := make(map[string]interface{})
	result["s3_bucket_name"] = *loggingInfo.S3BucketName
	result["s3_region"] = *loggingInfo.S3Region

	if loggingInfo.S3KeyPrefix != nil {
		result["s3_bucket_prefix"] = *loggingInfo.S3KeyPrefix
	}

	return []interface{}{result}
}

func expandAwsSsmTaskParameters(config []interface{}) map[string]*ssm.MaintenanceWindowTaskParameterValueExpression {
	params := make(map[string]*ssm.MaintenanceWindowTaskParameterValueExpression)
	for _, v := range config {
		paramConfig := v.(map[string]interface{})
		params[paramConfig["name"].(string)] = &ssm.MaintenanceWindowTaskParameterValueExpression{
			Values: expandStringList(paramConfig["values"].([]interface{})),
		}
	}
	return params
}

func flattenAwsSsmTaskParameters(taskParameters map[string]*ssm.MaintenanceWindowTaskParameterValueExpression) []interface{} {
	result := make([]interface{}, 0, len(taskParameters))
	for k, v := range taskParameters {
		taskParam := map[string]interface{}{
			"name":   k,
			"values": flattenStringList(v.Values),
		}
		result = append(result, taskParam)
	}

	return result
}

func resourceAwsSsmMaintenanceWindowTaskCreate(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Registering SSM Maintenance Window Task")

	params := &ssm.RegisterTaskWithMaintenanceWindowInput{
		WindowId:       aws.String(d.Get("window_id").(string)),
		MaxConcurrency: aws.String(d.Get("max_concurrency").(string)),
		MaxErrors:      aws.String(d.Get("max_errors").(string)),
		TaskType:       aws.String(d.Get("task_type").(string)),
		ServiceRoleArn: aws.String(d.Get("service_role_arn").(string)),
		TaskArn:        aws.String(d.Get("task_arn").(string)),
		Name:           aws.String(d.Get("name").(string)),
		Description:    aws.String(d.Get("description").(string)),
		Targets:        expandAwsSsmTargets(d.Get("targets").([]interface{})),
	}

	if v, ok := d.GetOk("priority"); ok {
		params.Priority = aws.Int64(int64(v.(int)))
	}

	if v, ok := d.GetOk("logging_info"); ok {
		params.LoggingInfo = expandAwsSsmMaintenanceWindowLoggingInfo(v.([]interface{}))
	}

	if v, ok := d.GetOk("task_parameters"); ok {
		params.TaskParameters = expandAwsSsmTaskParameters(v.([]interface{}))
	}

	resp, err := ssmconn.RegisterTaskWithMaintenanceWindow(params)
	if err != nil {
		return err
	}

	d.SetId(*resp.WindowTaskId)

	return resourceAwsSsmMaintenanceWindowTaskRead(d, meta)
}

func resourceAwsSsmMaintenanceWindowTaskRead(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	params := &ssm.DescribeMaintenanceWindowTasksInput{
		WindowId: aws.String(d.Get("window_id").(string)),
	}

	resp, err := ssmconn.DescribeMaintenanceWindowTasks(params)
	if err != nil {
		return err
	}

	found := false
	for _, t := range resp.Tasks {
		if *t.WindowTaskId == d.Id() {
			found = true

			d.Set("window_id", t.WindowId)
			d.Set("max_concurrency", t.MaxConcurrency)
			d.Set("max_errors", t.MaxErrors)
			d.Set("task_type", t.Type)
			d.Set("service_role_arn", t.ServiceRoleArn)
			d.Set("task_arn", t.TaskArn)
			d.Set("priority", t.Priority)

			if t.Name != nil {
				d.Set("name", t.Name)
			}

			if t.Description != nil {
				d.Set("description", t.Description)
			}

			if t.LoggingInfo != nil {
				if err := d.Set("logging_info", flattenAwsSsmMaintenanceWindowLoggingInfo(t.LoggingInfo)); err != nil {
					return fmt.Errorf("[DEBUG] Error setting logging_info error: %#v", err)
				}
			}

			if t.TaskParameters != nil {
				if err := d.Set("task_parameters", flattenAwsSsmTaskParameters(t.TaskParameters)); err != nil {
					return fmt.Errorf("[DEBUG] Error setting task_parameters error: %#v", err)
				}
			}

			if err := d.Set("targets", flattenAwsSsmTargets(t.Targets)); err != nil {
				return fmt.Errorf("[DEBUG] Error setting targets error: %#v", err)
			}
		}
	}

	if !found {
		log.Printf("[INFO] Maintenance Window Target not found. Removing from state")
		d.SetId("")
		return nil
	}

	return nil
}

func resourceAwsSsmMaintenanceWindowTaskDelete(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Deregistering SSM Maintenance Window Task: %s", d.Id())

	params := &ssm.DeregisterTaskFromMaintenanceWindowInput{
		WindowId:     aws.String(d.Get("window_id").(string)),
		WindowTaskId: aws.String(d.Id()),
	}

	_, err := ssmconn.DeregisterTaskFromMaintenanceWindow(params)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>fmt checked passed<commit_after>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ssm"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsSsmMaintenanceWindowTask() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSsmMaintenanceWindowTaskCreate,
		Read:   resourceAwsSsmMaintenanceWindowTaskRead,
		Delete: resourceAwsSsmMaintenanceWindowTaskDelete,

		Schema: map[string]*schema.Schema{
			"window_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"max_concurrency": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"max_errors": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"task_type": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"task_arn": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"service_role_arn": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"targets": {
				Type:     schema.TypeList,
				Required: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"key": {
							Type:     schema.TypeString,
							Required: true,
							ForceNew: true,
						},
						"values": {
							Type:     schema.TypeList,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},
					},
				},
			},

			"name": {
				Type:         schema.TypeString,
				Optional:     true,
				ForceNew:     true,
				ValidateFunc: validateAwsSSMMaintenanceWindowTaskName,
			},

			"description": {
				Type:         schema.TypeString,
				Optional:     true,
				ForceNew:     true,
				ValidateFunc: validateAwsSSMMaintenanceWindowTaskDescription,
			},

			"priority": {
				Type:     schema.TypeInt,
				Optional: true,
				ForceNew: true,
			},

			"logging_info": {
				Type:     schema.TypeList,
				MaxItems: 1,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"s3_bucket_name": {
							Type:     schema.TypeString,
							Required: true,
						},
						"s3_region": {
							Type:     schema.TypeString,
							Required: true,
						},
						"s3_bucket_prefix": {
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},

			"task_parameters": {
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": {
							Type:     schema.TypeString,
							Required: true,
							ForceNew: true,
						},
						"values": {
							Type:     schema.TypeList,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
						},
					},
				},
			},
		},
	}
}

func expandAwsSsmMaintenanceWindowLoggingInfo(config []interface{}) *ssm.LoggingInfo {

	loggingConfig := config[0].(map[string]interface{})

	loggingInfo := &ssm.LoggingInfo{
		S3BucketName: aws.String(loggingConfig["s3_bucket_name"].(string)),
		S3Region:     aws.String(loggingConfig["s3_region"].(string)),
	}

	if s := loggingConfig["s3_bucket_prefix"].(string); s != "" {
		loggingInfo.S3KeyPrefix = aws.String(s)
	}

	return loggingInfo
}

func flattenAwsSsmMaintenanceWindowLoggingInfo(loggingInfo *ssm.LoggingInfo) []interface{} {

	result := make(map[string]interface{})
	result["s3_bucket_name"] = *loggingInfo.S3BucketName
	result["s3_region"] = *loggingInfo.S3Region

	if loggingInfo.S3KeyPrefix != nil {
		result["s3_bucket_prefix"] = *loggingInfo.S3KeyPrefix
	}

	return []interface{}{result}
}

func expandAwsSsmTaskParameters(config []interface{}) map[string]*ssm.MaintenanceWindowTaskParameterValueExpression {
	params := make(map[string]*ssm.MaintenanceWindowTaskParameterValueExpression)
	for _, v := range config {
		paramConfig := v.(map[string]interface{})
		params[paramConfig["name"].(string)] = &ssm.MaintenanceWindowTaskParameterValueExpression{
			Values: expandStringList(paramConfig["values"].([]interface{})),
		}
	}
	return params
}

func flattenAwsSsmTaskParameters(taskParameters map[string]*ssm.MaintenanceWindowTaskParameterValueExpression) []interface{} {
	result := make([]interface{}, 0, len(taskParameters))
	for k, v := range taskParameters {
		taskParam := map[string]interface{}{
			"name":   k,
			"values": flattenStringList(v.Values),
		}
		result = append(result, taskParam)
	}

	return result
}

func resourceAwsSsmMaintenanceWindowTaskCreate(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Registering SSM Maintenance Window Task")

	params := &ssm.RegisterTaskWithMaintenanceWindowInput{
		WindowId:       aws.String(d.Get("window_id").(string)),
		MaxConcurrency: aws.String(d.Get("max_concurrency").(string)),
		MaxErrors:      aws.String(d.Get("max_errors").(string)),
		TaskType:       aws.String(d.Get("task_type").(string)),
		ServiceRoleArn: aws.String(d.Get("service_role_arn").(string)),
		TaskArn:        aws.String(d.Get("task_arn").(string)),
		Name:           aws.String(d.Get("name").(string)),
		Description:    aws.String(d.Get("description").(string)),
		Targets:        expandAwsSsmTargets(d.Get("targets").([]interface{})),
	}

	if v, ok := d.GetOk("priority"); ok {
		params.Priority = aws.Int64(int64(v.(int)))
	}

	if v, ok := d.GetOk("logging_info"); ok {
		params.LoggingInfo = expandAwsSsmMaintenanceWindowLoggingInfo(v.([]interface{}))
	}

	if v, ok := d.GetOk("task_parameters"); ok {
		params.TaskParameters = expandAwsSsmTaskParameters(v.([]interface{}))
	}

	resp, err := ssmconn.RegisterTaskWithMaintenanceWindow(params)
	if err != nil {
		return err
	}

	d.SetId(*resp.WindowTaskId)

	return resourceAwsSsmMaintenanceWindowTaskRead(d, meta)
}

func resourceAwsSsmMaintenanceWindowTaskRead(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	params := &ssm.DescribeMaintenanceWindowTasksInput{
		WindowId: aws.String(d.Get("window_id").(string)),
	}

	resp, err := ssmconn.DescribeMaintenanceWindowTasks(params)
	if err != nil {
		return err
	}

	found := false
	for _, t := range resp.Tasks {
		if *t.WindowTaskId == d.Id() {
			found = true

			d.Set("window_id", t.WindowId)
			d.Set("max_concurrency", t.MaxConcurrency)
			d.Set("max_errors", t.MaxErrors)
			d.Set("task_type", t.Type)
			d.Set("service_role_arn", t.ServiceRoleArn)
			d.Set("task_arn", t.TaskArn)
			d.Set("priority", t.Priority)

			if t.Name != nil {
				d.Set("name", t.Name)
			}

			if t.Description != nil {
				d.Set("description", t.Description)
			}

			if t.LoggingInfo != nil {
				if err := d.Set("logging_info", flattenAwsSsmMaintenanceWindowLoggingInfo(t.LoggingInfo)); err != nil {
					return fmt.Errorf("[DEBUG] Error setting logging_info error: %#v", err)
				}
			}

			if t.TaskParameters != nil {
				if err := d.Set("task_parameters", flattenAwsSsmTaskParameters(t.TaskParameters)); err != nil {
					return fmt.Errorf("[DEBUG] Error setting task_parameters error: %#v", err)
				}
			}

			if err := d.Set("targets", flattenAwsSsmTargets(t.Targets)); err != nil {
				return fmt.Errorf("[DEBUG] Error setting targets error: %#v", err)
			}
		}
	}

	if !found {
		log.Printf("[INFO] Maintenance Window Target not found. Removing from state")
		d.SetId("")
		return nil
	}

	return nil
}

func resourceAwsSsmMaintenanceWindowTaskDelete(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Deregistering SSM Maintenance Window Task: %s", d.Id())

	params := &ssm.DeregisterTaskFromMaintenanceWindowInput{
		WindowId:     aws.String(d.Get("window_id").(string)),
		WindowTaskId: aws.String(d.Id()),
	}

	_, err := ssmconn.DeregisterTaskFromMaintenanceWindow(params)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package connpass

import (
	"net/url"
	"net/http"
	"io"
	"bytes"
	"encoding/json"
	"io/ioutil"
	"reflect"
	"github.com/google/go-querystring/query"
	"fmt"
)

const (
	defaultBaseURL = "https://connpass.com/api/v1/event/"
)

// レスポンスフィールド
type Connpass struct {
	ResultsReturned  *int `json:"results_returned"`
	ResultsAvailable *int `json:"results_available"`
	ResultsStart     *int `json:"results_start"`
	Events           *[]Event `json:"events"`
}

// レスポンスフィールド
type Event struct {
	EventId          *int `json:"event_id"`
	Title            *string `json:"title"`
	Catch            *string `json:"catch"`
	Description      *string `json:"description"`
	EventUrl         *string `json:"event_url"`
	HashTag          *string `json:"hash_tag"`
	StartedAt        *string `json:"started_at"`
	EndedAt          *string `json:"ended_at"`
	Limit            *int `json:"limit"`
	EventType        *string `json:"event_type"`
	Series                        `json:"series"`
	Address          *string `json:"address"`
	Place            *string `json:"place"`
	//Lat              float32 `json:"lat"`
	//Lon              float32 `json:"lon"`
	OwnerId          *int `json:"owner_id"`
	OwnerNickname    *string `json:"owner_nickname"`
	OwnerDisplayName *string `json:"owner_display_name"`
	Accepted         *int `json:"accepted"`
	Waiting          *int `json:"waiting"`
	Updated_at       *string `json:"updated_at"`
}

// レスポンスフィールド
type Series struct {
	Id    *int `json:"id"`
	Title *string `json:"title"`
	Url   *string `json:"url"`
}

// 検索クエリ
type Param struct {
	EventId       int `url:"event_id,omitempty"`
	Keyword       string `url:"keyword,omitempty"`
	KeywordOr     string `url:"keyword_or,omitempty"`
	Ym            int `url:"ym,omitempty"`
	Ymd           int `url:"ymd,omitempty"`
	Nickname      string `url:"nickname,omitempty"`
	OwnerNickname string `url:"owner_nickname,omitempty"`
	SeriesId      int `url:"series_id,omitempty"`
	Options
}

// 検索クエリオプション
type Options struct {
	Start int `url:"start,omitempty"`
	Order int `url:"order,omitempty"`
	Count int `url:"count,omitempty"`
}

type Client struct {
	client  *http.Client
	BaseURL *url.URL
}

func NewClient() *Client {

	baseURL, _ := url.Parse(defaultBaseURL)

	c := &Client{
		client:http.DefaultClient,
		BaseURL:baseURL,
	}

	return c
}

func (c *Client) Get(param *Param) (*Connpass, *http.Response, error) {

	u, err := createUrl(param)

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

	req, err := c.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	connpass := new(*Connpass)

	resp, err := c.Do(req, &connpass)
	if err != nil {
		return nil, resp, err
	}
	return *connpass, resp, err
}

func (c *Client) NewRequest(method, urlStr string, body interface{}) (*http.Request, error) {
	rel, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	u := c.BaseURL.ResolveReference(rel)

	var buf io.ReadWriter
	if body != nil {
		buf = new(bytes.Buffer)
		err := json.NewEncoder(buf).Encode(body)
		if err != nil {
			return nil, err
		}
	}

	req, err := http.NewRequest(method, u.String(), buf)
	if err != nil {
		return nil, err
	}

	return req, nil
}

func (c *Client) Do(req *http.Request, v interface{}) (*http.Response, error) {

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	defer func() {
		// Drain up to 512 bytes and close the body to let the Transport reuse the connection
		io.CopyN(ioutil.Discard, resp.Body, 512)
		resp.Body.Close()
	}()

	if v != nil {
		if w, ok := v.(io.Writer); ok {
			io.Copy(w, resp.Body)
		} else {
			err = json.NewDecoder(resp.Body).Decode(v)
			if err == io.EOF {
				err = nil // ignore EOF errors caused by empty response body
			}
		}
	}

	return resp, err
}

func createUrl(param interface{}) (string, error) {
	v := reflect.ValueOf(param)
	if v.Kind() == reflect.Ptr && v.IsNil() {
		return "", nil
	}

	qs, err := query.Values(param)
	if err != nil {
		return "", err
	}

	u := "?" + qs.Encode()

	return u, nil
}<commit_msg>update 不要なimport削除<commit_after>package connpass

import (
	"net/url"
	"net/http"
	"io"
	"bytes"
	"encoding/json"
	"io/ioutil"
	"reflect"
	"github.com/google/go-querystring/query"
)

const (
	defaultBaseURL = "https://connpass.com/api/v1/event/"
)

// レスポンスフィールド
type Connpass struct {
	ResultsReturned  *int `json:"results_returned"`
	ResultsAvailable *int `json:"results_available"`
	ResultsStart     *int `json:"results_start"`
	Events           *[]Event `json:"events"`
}

// レスポンスフィールド
type Event struct {
	EventId          *int `json:"event_id"`
	Title            *string `json:"title"`
	Catch            *string `json:"catch"`
	Description      *string `json:"description"`
	EventUrl         *string `json:"event_url"`
	HashTag          *string `json:"hash_tag"`
	StartedAt        *string `json:"started_at"`
	EndedAt          *string `json:"ended_at"`
	Limit            *int `json:"limit"`
	EventType        *string `json:"event_type"`
	Series                        `json:"series"`
	Address          *string `json:"address"`
	Place            *string `json:"place"`
	//Lat              float32 `json:"lat"`
	//Lon              float32 `json:"lon"`
	OwnerId          *int `json:"owner_id"`
	OwnerNickname    *string `json:"owner_nickname"`
	OwnerDisplayName *string `json:"owner_display_name"`
	Accepted         *int `json:"accepted"`
	Waiting          *int `json:"waiting"`
	Updated_at       *string `json:"updated_at"`
}

// レスポンスフィールド
type Series struct {
	Id    *int `json:"id"`
	Title *string `json:"title"`
	Url   *string `json:"url"`
}

// 検索クエリ
type Param struct {
	EventId       int `url:"event_id,omitempty"`
	Keyword       string `url:"keyword,omitempty"`
	KeywordOr     string `url:"keyword_or,omitempty"`
	Ym            int `url:"ym,omitempty"`
	Ymd           int `url:"ymd,omitempty"`
	Nickname      string `url:"nickname,omitempty"`
	OwnerNickname string `url:"owner_nickname,omitempty"`
	SeriesId      int `url:"series_id,omitempty"`
	Options
}

// 検索クエリオプション
type Options struct {
	Start int `url:"start,omitempty"`
	Order int `url:"order,omitempty"`
	Count int `url:"count,omitempty"`
}

type Client struct {
	client  *http.Client
	BaseURL *url.URL
}

func NewClient() *Client {

	baseURL, _ := url.Parse(defaultBaseURL)

	c := &Client{
		client:http.DefaultClient,
		BaseURL:baseURL,
	}

	return c
}

func (c *Client) Get(param *Param) (*Connpass, *http.Response, error) {

	u, err := createUrl(param)

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

	req, err := c.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	connpass := new(*Connpass)

	resp, err := c.Do(req, &connpass)
	if err != nil {
		return nil, resp, err
	}
	return *connpass, resp, err
}

func (c *Client) NewRequest(method, urlStr string, body interface{}) (*http.Request, error) {
	rel, err := url.Parse(urlStr)
	if err != nil {
		return nil, err
	}

	u := c.BaseURL.ResolveReference(rel)

	var buf io.ReadWriter
	if body != nil {
		buf = new(bytes.Buffer)
		err := json.NewEncoder(buf).Encode(body)
		if err != nil {
			return nil, err
		}
	}

	req, err := http.NewRequest(method, u.String(), buf)
	if err != nil {
		return nil, err
	}

	return req, nil
}

func (c *Client) Do(req *http.Request, v interface{}) (*http.Response, error) {

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	defer func() {
		// Drain up to 512 bytes and close the body to let the Transport reuse the connection
		io.CopyN(ioutil.Discard, resp.Body, 512)
		resp.Body.Close()
	}()

	if v != nil {
		if w, ok := v.(io.Writer); ok {
			io.Copy(w, resp.Body)
		} else {
			err = json.NewDecoder(resp.Body).Decode(v)
			if err == io.EOF {
				err = nil // ignore EOF errors caused by empty response body
			}
		}
	}

	return resp, err
}

func createUrl(param interface{}) (string, error) {
	v := reflect.ValueOf(param)
	if v.Kind() == reflect.Ptr && v.IsNil() {
		return "", nil
	}

	qs, err := query.Values(param)
	if err != nil {
		return "", err
	}

	u := "?" + qs.Encode()

	return u, nil
}<|endoftext|>


<commit_before>// Copyright 2015-2017 trivago GmbH
//
// 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 consumer

import (
	"fmt"
	"math"
	"math/rand"
	"os"
	"sync"
	"time"

	"github.com/trivago/gollum/core"
	"github.com/trivago/tgo"
)

// Profiler consumer plugin
//
// The profiler plugin generates Runs x Batches messages and send them to the
// configured streams as fast as possible. This consumer can be used to profile
// producers and/or configurations.
//
// Configuration example
//
//  - "consumer.Profile":
//    Runs: 10000
//    Batches: 10
//    TemplateCount: 10
//    Characters: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890"
//    Message: "%256s"
//	  DelayMs: 0
//    KeepRunning: false
//
// Runs defines the number of messages per batch. By default this is set to
// 10000.
//
// Batches defines the number of measurement runs to do. By default this is set
// to 10.
//
// TemplateCount defines the number of message templates to be generated.
// A random message template will be chosen when a message is sent. Templates
// are generated in advance. By default this is set to 10.
//
// Characters defines the characters to be used in generated strings. By default
// these are "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890 ".
//
// Message defines a go format string to be used to generate the message payloads.
// The length of the values generated will be deducted from the format size
// parameter. I.e. "%200d" will generate a digit between 0 and 200, "%10s" will
// generate a string with 10 characters, etc..
// By default this is set to "%256s".
//
// DelayMs defines the number of milliseconds of sleep between messages.
// By default this is set to 0.
//
// KeepRunning can be set to true to disable automatic shutdown of gollum after
// profiling is done. This can be used to e.g. read metrics after a profile run.
// By default this is set to false.
type Profiler struct {
	core.SimpleConsumer `gollumdoc:"embed_type"`
	profileRuns         int
	batches             int
	templates           [][]byte
	chars               string
	message             string
	delay               time.Duration
	keepRunning         bool
}

var profilerDefaultCharacters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890 "

func init() {
	core.TypeRegistry.Register(Profiler{})
}

// Configure initializes this consumer with values from a plugin config.
func (cons *Profiler) Configure(conf core.PluginConfigReader) error {
	cons.SimpleConsumer.Configure(conf)

	numTemplates := conf.GetInt("TemplateCount", 10)
	cons.profileRuns = conf.GetInt("Runs", 10000)
	cons.batches = conf.GetInt("Batches", 10)
	cons.chars = conf.GetString("Characters", profilerDefaultCharacters)
	cons.message = conf.GetString("Message", "%# %256s")
	cons.templates = make([][]byte, numTemplates)
	cons.keepRunning = conf.GetBool("KeepRunning", false)
	cons.delay = time.Duration(conf.GetInt("DelayMs", 0)) * time.Millisecond

	return conf.Errors.OrNil()
}

func (cons *Profiler) generateString(size int) string {
	randString := make([]byte, size)
	for i := 0; i < size; i++ {
		randString[i] = cons.chars[rand.Intn(len(cons.chars))]
	}
	return string(randString)
}

func (cons *Profiler) generateTemplate() []byte {
	var dummyValues []interface{}
	messageLen := len(cons.message)

	for idx, char := range cons.message {
		if char == '%' {
			// Get the required length
			size := 0
			searchIdx := idx + 1

			// Get format size
			for searchIdx < messageLen && cons.message[searchIdx] >= '0' && cons.message[searchIdx] <= '9' {
				size = size*10 + int(cons.message[searchIdx]-'0')
				searchIdx++
			}

			// Skip decimal places
			if cons.message[searchIdx] == '.' {
				searchIdx++
				for searchIdx < messageLen && cons.message[searchIdx] >= '0' && cons.message[searchIdx] <= '9' {
					searchIdx++
				}
			}

			// Make sure there is at least 1 number / character
			if size == 0 {
				size = 1
			}

			// Generate values
			switch cons.message[searchIdx] {
			case '%':
				// ignore

			case 'e', 'E', 'f', 'F', 'g', 'G':
				dummyValues = append(dummyValues, rand.Float64()*math.Pow10(size))

			case 'U', 'c':
				dummyValues = append(dummyValues, rune(rand.Intn(1<<16)+32))

			case 'd', 'b', 'o', 'x', 'X':
				dummyValues = append(dummyValues, rand.Intn(int(math.Pow10(size))))

			case 's', 'q', 'v', 'T':
				fallthrough
			default:
				dummyValues = append(dummyValues, cons.generateString(size))
			}
		}
	}

	return []byte(fmt.Sprintf(cons.message, dummyValues...))
}

func (cons *Profiler) profile() {
	defer cons.WorkerDone()

	for i := 0; i < len(cons.templates); i++ {
		cons.templates[i] = cons.generateTemplate()
	}

	cons.Log.Debug.Printf("Started profiling with %d byte messages", len(cons.templates[0]))

	testStart := time.Now()
	minTime := math.MaxFloat64
	maxTime := 0.0
	batchIdx := 0
	messageCount := 0

	for batchIdx = 0; batchIdx < cons.batches && cons.IsActive(); batchIdx++ {
		cons.Log.Note.Print(fmt.Sprintf("run %d/%d", batchIdx, cons.batches))
		start := time.Now()

		for i := 0; i < cons.profileRuns && cons.IsActive(); i++ {
			template := cons.templates[rand.Intn(len(cons.templates))]
			messageCount++

			cons.Enqueue(template)

			if cons.delay > 0 && cons.IsActive() {
				time.Sleep(cons.delay)
			}
		}

		runTime := time.Since(start)
		if messageCount%cons.profileRuns == 0 {
			minTime = math.Min(minTime, runTime.Seconds())
			maxTime = math.Max(maxTime, runTime.Seconds())
		}
	}

	runTime := time.Since(testStart)

	cons.Log.Note.Printf("Overview: %d messages sent in %.4f seconds",
		messageCount,
		runTime.Seconds())

	cons.Log.Note.Printf("Avg: %4.f msg/sec",
		float64(messageCount)/runTime.Seconds())

	cons.Log.Note.Printf("Best: %.4f sec = %4.f msg/sec",
		minTime,
		float64(cons.profileRuns)/minTime)

	cons.Log.Note.Printf("Worst: %.4f sec = %4.f msg/sec",
		maxTime,
		float64(cons.profileRuns)/maxTime)

	if cons.IsActive() {
		cons.Log.Debug.Print("Profiler done.")
		// Automatically shut down when done
		// TODO: Hack
		if !cons.keepRunning {
			proc, _ := os.FindProcess(os.Getpid())
			proc.Signal(os.Interrupt)
		}
	}
}

// Consume starts a profile run and exits gollum when done
func (cons *Profiler) Consume(workers *sync.WaitGroup) {
	cons.AddMainWorker(workers)

	go tgo.WithRecoverShutdown(cons.profile)
	cons.ControlLoop()
}
<commit_msg>fixed profiler consumer output<commit_after>// Copyright 2015-2017 trivago GmbH
//
// 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 consumer

import (
	"fmt"
	"math"
	"math/rand"
	"os"
	"sync"
	"time"

	"github.com/trivago/gollum/core"
	"github.com/trivago/tgo"
)

// Profiler consumer plugin
//
// The profiler plugin generates Runs x Batches messages and send them to the
// configured streams as fast as possible. This consumer can be used to profile
// producers and/or configurations.
//
// Configuration example
//
//  - "consumer.Profile":
//    Runs: 10000
//    Batches: 10
//    TemplateCount: 10
//    Characters: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890"
//    Message: "%256s"
//	  DelayMs: 0
//    KeepRunning: false
//
// Runs defines the number of messages per batch. By default this is set to
// 10000.
//
// Batches defines the number of measurement runs to do. By default this is set
// to 10.
//
// TemplateCount defines the number of message templates to be generated.
// A random message template will be chosen when a message is sent. Templates
// are generated in advance. By default this is set to 10.
//
// Characters defines the characters to be used in generated strings. By default
// these are "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890 ".
//
// Message defines a go format string to be used to generate the message payloads.
// The length of the values generated will be deducted from the format size
// parameter. I.e. "%200d" will generate a digit between 0 and 200, "%10s" will
// generate a string with 10 characters, etc..
// By default this is set to "%256s".
//
// DelayMs defines the number of milliseconds of sleep between messages.
// By default this is set to 0.
//
// KeepRunning can be set to true to disable automatic shutdown of gollum after
// profiling is done. This can be used to e.g. read metrics after a profile run.
// By default this is set to false.
type Profiler struct {
	core.SimpleConsumer `gollumdoc:"embed_type"`
	profileRuns         int
	batches             int
	templates           [][]byte
	chars               string
	message             string
	delay               time.Duration
	keepRunning         bool
}

var profilerDefaultCharacters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890 "

func init() {
	core.TypeRegistry.Register(Profiler{})
}

// Configure initializes this consumer with values from a plugin config.
func (cons *Profiler) Configure(conf core.PluginConfigReader) error {
	cons.SimpleConsumer.Configure(conf)

	numTemplates := conf.GetInt("TemplateCount", 10)
	cons.profileRuns = conf.GetInt("Runs", 10000)
	cons.batches = conf.GetInt("Batches", 10)
	cons.chars = conf.GetString("Characters", profilerDefaultCharacters)
	cons.message = conf.GetString("Message", "%# %256s")
	cons.templates = make([][]byte, numTemplates)
	cons.keepRunning = conf.GetBool("KeepRunning", false)
	cons.delay = time.Duration(conf.GetInt("DelayMs", 0)) * time.Millisecond

	return conf.Errors.OrNil()
}

func (cons *Profiler) generateString(size int) string {
	randString := make([]byte, size)
	for i := 0; i < size; i++ {
		randString[i] = cons.chars[rand.Intn(len(cons.chars))]
	}
	return string(randString)
}

func (cons *Profiler) generateTemplate() []byte {
	var dummyValues []interface{}
	messageLen := len(cons.message)

	for idx, char := range cons.message {
		if char == '%' {
			// Get the required length
			size := 0
			searchIdx := idx + 1

			// Get format size
			for searchIdx < messageLen && cons.message[searchIdx] >= '0' && cons.message[searchIdx] <= '9' {
				size = size*10 + int(cons.message[searchIdx]-'0')
				searchIdx++
			}

			// Skip decimal places
			if cons.message[searchIdx] == '.' {
				searchIdx++
				for searchIdx < messageLen && cons.message[searchIdx] >= '0' && cons.message[searchIdx] <= '9' {
					searchIdx++
				}
			}

			// Make sure there is at least 1 number / character
			if size == 0 {
				size = 1
			}

			// Generate values
			switch cons.message[searchIdx] {
			case '%':
				// ignore

			case 'e', 'E', 'f', 'F', 'g', 'G':
				dummyValues = append(dummyValues, rand.Float64()*math.Pow10(size))

			case 'U', 'c':
				dummyValues = append(dummyValues, rune(rand.Intn(1<<16)+32))

			case 'd', 'b', 'o', 'x', 'X':
				dummyValues = append(dummyValues, rand.Intn(int(math.Pow10(size))))

			case 's', 'q', 'v', 'T':
				fallthrough
			default:
				dummyValues = append(dummyValues, cons.generateString(size))
			}
		}
	}

	return []byte(fmt.Sprintf(cons.message, dummyValues...))
}

func (cons *Profiler) profile() {
	defer cons.WorkerDone()

	for i := 0; i < len(cons.templates); i++ {
		cons.templates[i] = cons.generateTemplate()
	}

	cons.Log.Debug.Printf("Started profiling with %d byte messages", len(cons.templates[0]))

	testStart := time.Now()
	minTime := math.MaxFloat64
	maxTime := 0.0
	batchIdx := 0
	messageCount := 0

	for batchIdx = 0; batchIdx < cons.batches && cons.IsActive(); batchIdx++ {
		cons.Log.Note.Print(fmt.Sprintf("batch %d/%d", batchIdx, cons.batches))
		start := time.Now()

		for i := 0; i < cons.profileRuns && cons.IsActive(); i++ {
			template := cons.templates[rand.Intn(len(cons.templates))]
			messageCount++

			cons.Enqueue(template)

			if cons.delay > 0 && cons.IsActive() {
				time.Sleep(cons.delay)
			}
		}

		runTime := time.Since(start)
		if messageCount%cons.profileRuns == 0 {
			minTime = math.Min(minTime, runTime.Seconds())
			maxTime = math.Max(maxTime, runTime.Seconds())
		}
	}

	runTime := time.Since(testStart)

	cons.Log.Note.Printf("Overview: %d messages sent in %.4f seconds",
		messageCount,
		runTime.Seconds())

	cons.Log.Note.Printf("Avg: %4.f msg/sec",
		float64(messageCount)/runTime.Seconds())

	cons.Log.Note.Printf("Best: %.4f sec = %4.f msg/sec",
		minTime,
		float64(cons.profileRuns)/minTime)

	cons.Log.Note.Printf("Worst: %.4f sec = %4.f msg/sec",
		maxTime,
		float64(cons.profileRuns)/maxTime)

	if cons.IsActive() {
		cons.Log.Debug.Print("Profiler done.")
		// Automatically shut down when done
		// TODO: Hack
		if !cons.keepRunning {
			proc, _ := os.FindProcess(os.Getpid())
			proc.Signal(os.Interrupt)
		}
	}
}

// Consume starts a profile run and exits gollum when done
func (cons *Profiler) Consume(workers *sync.WaitGroup) {
	cons.AddMainWorker(workers)

	go tgo.WithRecoverShutdown(cons.profile)
	cons.ControlLoop()
}
<|endoftext|>