<commit_before>/*
Copyright IBM Corp. 2016 All Rights Reserved.

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 obcpbft

import (
	"fmt"
	"reflect"
	"time"

	"github.com/hyperledger/fabric/consensus"
	"github.com/hyperledger/fabric/consensus/obcpbft/events"
	pb "github.com/hyperledger/fabric/protos"

	"github.com/golang/protobuf/proto"
	"github.com/spf13/viper"
	google_protobuf "google/protobuf"
)

type obcBatch struct {
	obcGeneric
	externalEventReceiver
	pbft *pbftCore

	batchSize        int
	batchStore       []*Request
	batchTimer       events.Timer
	batchTimerActive bool
	batchTimeout     time.Duration

	manager events.Manager // TODO, remove eventually, the event manager

	incomingChan chan *batchMessage // Queues messages for processing by main thread
	idleChan     chan struct{}      // Idle channel, to be removed

	outstandingReqs map[*Request]struct{}

	persistForward
}

type batchMessage struct {
	msg    *pb.Message
	sender *pb.PeerID
}

type execInfo struct {
	seqNo uint64
	raw   []byte
}

// Event types

// batchMessageEvent is sent when a consensus messages is received to be sent to pbft
type batchMessageEvent batchMessage

// batchTimerEvent is sent when the batch timer expires
type batchTimerEvent struct{}

func newObcBatch(id uint64, config *viper.Viper, stack consensus.Stack) *obcBatch {
	var err error

	op := &obcBatch{
		obcGeneric: obcGeneric{stack: stack},
	}

	op.persistForward.persistor = stack

	logger.Debugf("Replica %d obtaining startup information", id)

	op.manager = events.NewManagerImpl() // TODO, this is hacky, eventually rip it out
	op.manager.SetReceiver(op)
	etf := events.NewTimerFactoryImpl(op.manager)
	op.pbft = newPbftCore(id, config, op, etf)
	op.manager.Start()
	op.externalEventReceiver.manager = op.manager

	op.batchSize = config.GetInt("general.batchsize")
	op.batchStore = nil
	op.batchTimeout, err = time.ParseDuration(config.GetString("general.timeout.batch"))
	if err != nil {
		panic(fmt.Errorf("Cannot parse batch timeout: %s", err))
	}
	logger.Infof("PBFT Batch size = %d", op.batchSize)
	logger.Infof("PBFT Batch timeout = %v", op.batchTimeout)

	op.incomingChan = make(chan *batchMessage)

	op.batchTimer = etf.CreateTimer()

	op.outstandingReqs = make(map[*Request]struct{})

	op.idleChan = make(chan struct{})
	close(op.idleChan) // TODO remove eventually

	return op
}

// Close tells us to release resources we are holding
func (op *obcBatch) Close() {
	op.batchTimer.Halt()
	op.pbft.close()
}

func (op *obcBatch) submitToLeader(req *Request) events.Event {
	// Broadcast the request to the network, in case we're in the wrong view
	op.broadcastMsg(&BatchMessage{&BatchMessage_Request{req}})

	// if we believe we are the leader, then process this request
	leader := op.pbft.primary(op.pbft.view)
	if leader == op.pbft.id && op.pbft.activeView {
		return op.leaderProcReq(req)
	} else {
		logger.Debugf("Replica %d add request %v to its outstanding store", op.pbft.id, req)
		op.outstandingReqs[req] = struct{}{}
		op.startTimerIfOutstandingRequests()
	}

	return nil
}

func (op *obcBatch) broadcastMsg(msg *BatchMessage) {
	msgPayload, _ := proto.Marshal(msg)
	ocMsg := &pb.Message{
		Type:    pb.Message_CONSENSUS,
		Payload: msgPayload,
	}
	op.stack.Broadcast(ocMsg, pb.PeerEndpoint_UNDEFINED)
}

// send a message to a specific replica
func (op *obcBatch) unicastMsg(msg *BatchMessage, receiverID uint64) {
	msgPayload, _ := proto.Marshal(msg)
	ocMsg := &pb.Message{
		Type:    pb.Message_CONSENSUS,
		Payload: msgPayload,
	}
	receiverHandle, err := getValidatorHandle(receiverID)
	if err != nil {
		return

	}
	op.stack.Unicast(ocMsg, receiverHandle)
}

// =============================================================================
// innerStack interface (functions called by pbft-core)
// =============================================================================

// multicast a message to all replicas
func (op *obcBatch) broadcast(msgPayload []byte) {
	op.stack.Broadcast(op.wrapMessage(msgPayload), pb.PeerEndpoint_UNDEFINED)
}

// send a message to a specific replica
func (op *obcBatch) unicast(msgPayload []byte, receiverID uint64) (err error) {
	receiverHandle, err := getValidatorHandle(receiverID)
	if err != nil {
		return
	}
	return op.stack.Unicast(op.wrapMessage(msgPayload), receiverHandle)
}

func (op *obcBatch) sign(msg []byte) ([]byte, error) {
	return op.stack.Sign(msg)
}

// verify message signature
func (op *obcBatch) verify(senderID uint64, signature []byte, message []byte) error {
	senderHandle, err := getValidatorHandle(senderID)
	if err != nil {
		return err
	}
	return op.stack.Verify(senderHandle, signature, message)
}

// validate checks whether the request is valid syntactically
// not used in obc-batch at the moment
func (op *obcBatch) validate(txRaw []byte) error {
	return nil
}

// execute an opaque request which corresponds to an OBC Transaction
func (op *obcBatch) execute(seqNo uint64, raw []byte) {
	reqs := &RequestBlock{}
	if err := proto.Unmarshal(raw, reqs); err != nil {
		logger.Warningf("Batch replica %d could not unmarshal request block: %s", op.pbft.id, err)
		return
	}

	logger.Debugf("Batch replica %d received exec for seqNo %d", op.pbft.id, seqNo)

	var txs []*pb.Transaction

	for _, req := range reqs.Requests {

		tx := &pb.Transaction{}
		if err := proto.Unmarshal(req.Payload, tx); err != nil {
			logger.Warningf("Batch replica %d could not unmarshal transaction: %s", op.pbft.id, err)
			continue
		}
		// TODO, this is a really and inefficient way to do this, but because reqs aren't comparable, they cannot be retrieved from the map directly
		for oreq := range op.outstandingReqs {
			if reflect.DeepEqual(oreq, req) {
				delete(op.outstandingReqs, oreq)
				break
			}
		}
		txs = append(txs, tx)
	}

	meta, _ := proto.Marshal(&Metadata{seqNo})

	op.stack.Execute(meta, txs) // This executes in the background, we will receive an executedEvent once it completes
}

// =============================================================================
// functions specific to batch mode
// =============================================================================

func (op *obcBatch) leaderProcReq(req *Request) events.Event {
	// XXX check req sig

	hash := hashReq(req)

	logger.Debugf("Batch primary %d queueing new request %s", op.pbft.id, hash)
	op.batchStore = append(op.batchStore, req)

	if !op.batchTimerActive {
		op.startBatchTimer()
	}

	if len(op.batchStore) >= op.batchSize {
		return op.sendBatch()
	}

	return nil
}

func (op *obcBatch) sendBatch() events.Event {
	op.stopBatchTimer()

	if len(op.batchStore) == 0 {
		logger.Error("Told to send an empty batch store for ordering, ignoring")
		return nil
	}

	earliestRequest := op.batchStore[0]

	reqBlock := &RequestBlock{op.batchStore}
	op.batchStore = nil

	reqsPacked, err := proto.Marshal(reqBlock)
	if err != nil {
		logger.Error("Unable to pack block for new batch request")
		return nil
	}

	// process internally
	logger.Infof("Creating batch with %d requests", len(reqBlock.Requests))
	return pbftMessageEvent{
		msg: &Message{&Message_Request{&Request{
			Payload:   reqsPacked,
			Timestamp: earliestRequest.Timestamp,
			ReplicaId: op.pbft.id},
		}},
		sender: op.pbft.id,
	}
}

func (op *obcBatch) txToReq(tx []byte) *Request {
	now := time.Now()
	req := &Request{
		Timestamp: &google_protobuf.Timestamp{
			Seconds: now.Unix(),
			Nanos:   int32(now.UnixNano() % 1000000000),
		},
		Payload:   tx,
		ReplicaId: op.pbft.id,
	}
	// XXX sign req
	return req
}

func (op *obcBatch) processMessage(ocMsg *pb.Message, senderHandle *pb.PeerID) events.Event {
	if ocMsg.Type == pb.Message_CHAIN_TRANSACTION {
		req := op.txToReq(ocMsg.Payload)
		return op.submitToLeader(req)
	}

	if ocMsg.Type != pb.Message_CONSENSUS {
		logger.Errorf("Unexpected message type: %s", ocMsg.Type)
		return nil
	}

	batchMsg := &BatchMessage{}
	err := proto.Unmarshal(ocMsg.Payload, batchMsg)
	if err != nil {
		logger.Errorf("Error unmarshaling message: %s", err)
		return nil
	}

	if req := batchMsg.GetRequest(); req != nil {
		if (op.pbft.primary(op.pbft.view) == op.pbft.id) && op.pbft.activeView {
			return op.leaderProcReq(req)
		}
		op.outstandingReqs[req] = struct{}{}
		return nil
	} else if pbftMsg := batchMsg.GetPbftMessage(); pbftMsg != nil {
		senderID, err := getValidatorID(senderHandle) // who sent this?
		if err != nil {
			panic("Cannot map sender's PeerID to a valid replica ID")
		}
		msg := &Message{}
		err = proto.Unmarshal(pbftMsg, msg)
		if err != nil {
			logger.Error("Error unpacking payload from message: %s", err)
			return nil
		}
		return pbftMessageEvent{
			msg:    msg,
			sender: senderID,
		}
	}

	logger.Error("Unknown request: %+v", batchMsg)

	return nil
}

// allow the primary to send a batch when the timer expires
func (op *obcBatch) ProcessEvent(event events.Event) events.Event {
	logger.Debugf("Replica %d batch main thread looping", op.pbft.id)
	switch et := event.(type) {
	case batchMessageEvent:
		ocMsg := et
		return op.processMessage(ocMsg.msg, ocMsg.sender)
	case executedEvent:
		op.stack.Commit(nil, et.tag.([]byte))
	case committedEvent:
		op.pbft.ProcessEvent(execDoneEvent{})
		op.startTimerIfOutstandingRequests()
		// If we are the primary, and know of outstanding requests, submit them for inclusion in the next batch until
		// we run out of requests, or a new batch message is triggered (this path will re-enter after execution)
		if op.pbft.primary(op.pbft.view) == op.pbft.id && op.pbft.activeView {
			for nreq := range op.outstandingReqs {
				if msg := op.leaderProcReq(nreq); msg != nil {
					return msg
				}
			}
		}
	case batchTimerEvent:
		logger.Infof("Replica %d batch timer expired", op.pbft.id)
		if op.pbft.activeView && (len(op.batchStore) > 0) {
			return op.sendBatch()
		}
	case viewChangedEvent:
		// Outstanding reqs doesn't make sense for batch, as all the requests in a batch may be processed
		// in a different batch, but PBFT core can't see through the opaque structure to see this
		// so, on view change, clear it out
		op.pbft.outstandingReqs = make(map[string]*Request)

		logger.Debugf("Replica %d batch thread recognizing new view", op.pbft.id)
		if op.batchTimerActive {
			op.stopBatchTimer()
		}

	case stateUpdatedEvent:
		// When the state is updated, clear any outstanding requests, they may have been processed while we were gone
		op.outstandingReqs = make(map[*Request]struct{})
		return op.pbft.ProcessEvent(event)
	default:
		return op.pbft.ProcessEvent(event)
	}

	return nil
}

func (op *obcBatch) startBatchTimer() {
	op.batchTimer.Reset(op.batchTimeout, batchTimerEvent{})
	logger.Debugf("Replica %d started the batch timer", op.pbft.id)
	op.batchTimerActive = true
}

func (op *obcBatch) stopBatchTimer() {
	op.batchTimer.Stop()
	logger.Debugf("Replica %d stopped the batch timer", op.pbft.id)
	op.batchTimerActive = false
}

// Wraps a payload into a batch message, packs it and wraps it into
// a Fabric message. Called by broadcast before transmission.
func (op *obcBatch) wrapMessage(msgPayload []byte) *pb.Message {
	batchMsg := &BatchMessage{&BatchMessage_PbftMessage{msgPayload}}
	packedBatchMsg, _ := proto.Marshal(batchMsg)
	ocMsg := &pb.Message{
		Type:    pb.Message_CONSENSUS,
		Payload: packedBatchMsg,
	}
	return ocMsg
}

// Retrieve the idle channel, only used for testing
func (op *obcBatch) idleChannel() <-chan struct{} {
	return op.idleChan
}

// TODO, temporary
func (op *obcBatch) getManager() events.Manager {
	return op.manager
}

func (op *obcBatch) startTimerIfOutstandingRequests() {
	if op.pbft.skipInProgress || op.pbft.currentExec != nil {
		// Do not start view change timer if some background event is in progress
		return
	}

	if len(op.outstandingReqs) == 0 {
		// Only start a timer if we are aware of outstanding requests
		return
	}
	op.pbft.softStartTimer(op.pbft.requestTimeout, "Batch outstanding requests")
}
<commit_msg>pbft/batch: goimports<commit_after>/*
Copyright IBM Corp. 2016 All Rights Reserved.

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 obcpbft

import (
	"fmt"
	"reflect"
	"time"

	"github.com/hyperledger/fabric/consensus"
	"github.com/hyperledger/fabric/consensus/obcpbft/events"
	pb "github.com/hyperledger/fabric/protos"

	google_protobuf "google/protobuf"

	"github.com/golang/protobuf/proto"
	"github.com/spf13/viper"
)

type obcBatch struct {
	obcGeneric
	externalEventReceiver
	pbft *pbftCore

	batchSize        int
	batchStore       []*Request
	batchTimer       events.Timer
	batchTimerActive bool
	batchTimeout     time.Duration

	manager events.Manager // TODO, remove eventually, the event manager

	incomingChan chan *batchMessage // Queues messages for processing by main thread
	idleChan     chan struct{}      // Idle channel, to be removed

	outstandingReqs map[*Request]struct{}

	persistForward
}

type batchMessage struct {
	msg    *pb.Message
	sender *pb.PeerID
}

type execInfo struct {
	seqNo uint64
	raw   []byte
}

// Event types

// batchMessageEvent is sent when a consensus messages is received to be sent to pbft
type batchMessageEvent batchMessage

// batchTimerEvent is sent when the batch timer expires
type batchTimerEvent struct{}

func newObcBatch(id uint64, config *viper.Viper, stack consensus.Stack) *obcBatch {
	var err error

	op := &obcBatch{
		obcGeneric: obcGeneric{stack: stack},
	}

	op.persistForward.persistor = stack

	logger.Debugf("Replica %d obtaining startup information", id)

	op.manager = events.NewManagerImpl() // TODO, this is hacky, eventually rip it out
	op.manager.SetReceiver(op)
	etf := events.NewTimerFactoryImpl(op.manager)
	op.pbft = newPbftCore(id, config, op, etf)
	op.manager.Start()
	op.externalEventReceiver.manager = op.manager

	op.batchSize = config.GetInt("general.batchsize")
	op.batchStore = nil
	op.batchTimeout, err = time.ParseDuration(config.GetString("general.timeout.batch"))
	if err != nil {
		panic(fmt.Errorf("Cannot parse batch timeout: %s", err))
	}
	logger.Infof("PBFT Batch size = %d", op.batchSize)
	logger.Infof("PBFT Batch timeout = %v", op.batchTimeout)

	op.incomingChan = make(chan *batchMessage)

	op.batchTimer = etf.CreateTimer()

	op.outstandingReqs = make(map[*Request]struct{})

	op.idleChan = make(chan struct{})
	close(op.idleChan) // TODO remove eventually

	return op
}

// Close tells us to release resources we are holding
func (op *obcBatch) Close() {
	op.batchTimer.Halt()
	op.pbft.close()
}

func (op *obcBatch) submitToLeader(req *Request) events.Event {
	// Broadcast the request to the network, in case we're in the wrong view
	op.broadcastMsg(&BatchMessage{&BatchMessage_Request{req}})

	// if we believe we are the leader, then process this request
	leader := op.pbft.primary(op.pbft.view)
	if leader == op.pbft.id && op.pbft.activeView {
		return op.leaderProcReq(req)
	} else {
		logger.Debugf("Replica %d add request %v to its outstanding store", op.pbft.id, req)
		op.outstandingReqs[req] = struct{}{}
		op.startTimerIfOutstandingRequests()
	}

	return nil
}

func (op *obcBatch) broadcastMsg(msg *BatchMessage) {
	msgPayload, _ := proto.Marshal(msg)
	ocMsg := &pb.Message{
		Type:    pb.Message_CONSENSUS,
		Payload: msgPayload,
	}
	op.stack.Broadcast(ocMsg, pb.PeerEndpoint_UNDEFINED)
}

// send a message to a specific replica
func (op *obcBatch) unicastMsg(msg *BatchMessage, receiverID uint64) {
	msgPayload, _ := proto.Marshal(msg)
	ocMsg := &pb.Message{
		Type:    pb.Message_CONSENSUS,
		Payload: msgPayload,
	}
	receiverHandle, err := getValidatorHandle(receiverID)
	if err != nil {
		return

	}
	op.stack.Unicast(ocMsg, receiverHandle)
}

// =============================================================================
// innerStack interface (functions called by pbft-core)
// =============================================================================

// multicast a message to all replicas
func (op *obcBatch) broadcast(msgPayload []byte) {
	op.stack.Broadcast(op.wrapMessage(msgPayload), pb.PeerEndpoint_UNDEFINED)
}

// send a message to a specific replica
func (op *obcBatch) unicast(msgPayload []byte, receiverID uint64) (err error) {
	receiverHandle, err := getValidatorHandle(receiverID)
	if err != nil {
		return
	}
	return op.stack.Unicast(op.wrapMessage(msgPayload), receiverHandle)
}

func (op *obcBatch) sign(msg []byte) ([]byte, error) {
	return op.stack.Sign(msg)
}

// verify message signature
func (op *obcBatch) verify(senderID uint64, signature []byte, message []byte) error {
	senderHandle, err := getValidatorHandle(senderID)
	if err != nil {
		return err
	}
	return op.stack.Verify(senderHandle, signature, message)
}

// validate checks whether the request is valid syntactically
// not used in obc-batch at the moment
func (op *obcBatch) validate(txRaw []byte) error {
	return nil
}

// execute an opaque request which corresponds to an OBC Transaction
func (op *obcBatch) execute(seqNo uint64, raw []byte) {
	reqs := &RequestBlock{}
	if err := proto.Unmarshal(raw, reqs); err != nil {
		logger.Warningf("Batch replica %d could not unmarshal request block: %s", op.pbft.id, err)
		return
	}

	logger.Debugf("Batch replica %d received exec for seqNo %d", op.pbft.id, seqNo)

	var txs []*pb.Transaction

	for _, req := range reqs.Requests {

		tx := &pb.Transaction{}
		if err := proto.Unmarshal(req.Payload, tx); err != nil {
			logger.Warningf("Batch replica %d could not unmarshal transaction: %s", op.pbft.id, err)
			continue
		}
		// TODO, this is a really and inefficient way to do this, but because reqs aren't comparable, they cannot be retrieved from the map directly
		for oreq := range op.outstandingReqs {
			if reflect.DeepEqual(oreq, req) {
				delete(op.outstandingReqs, oreq)
				break
			}
		}
		txs = append(txs, tx)
	}

	meta, _ := proto.Marshal(&Metadata{seqNo})

	op.stack.Execute(meta, txs) // This executes in the background, we will receive an executedEvent once it completes
}

// =============================================================================
// functions specific to batch mode
// =============================================================================

func (op *obcBatch) leaderProcReq(req *Request) events.Event {
	// XXX check req sig

	hash := hashReq(req)

	logger.Debugf("Batch primary %d queueing new request %s", op.pbft.id, hash)
	op.batchStore = append(op.batchStore, req)

	if !op.batchTimerActive {
		op.startBatchTimer()
	}

	if len(op.batchStore) >= op.batchSize {
		return op.sendBatch()
	}

	return nil
}

func (op *obcBatch) sendBatch() events.Event {
	op.stopBatchTimer()

	if len(op.batchStore) == 0 {
		logger.Error("Told to send an empty batch store for ordering, ignoring")
		return nil
	}

	earliestRequest := op.batchStore[0]

	reqBlock := &RequestBlock{op.batchStore}
	op.batchStore = nil

	reqsPacked, err := proto.Marshal(reqBlock)
	if err != nil {
		logger.Error("Unable to pack block for new batch request")
		return nil
	}

	// process internally
	logger.Infof("Creating batch with %d requests", len(reqBlock.Requests))
	return pbftMessageEvent{
		msg: &Message{&Message_Request{&Request{
			Payload:   reqsPacked,
			Timestamp: earliestRequest.Timestamp,
			ReplicaId: op.pbft.id},
		}},
		sender: op.pbft.id,
	}
}

func (op *obcBatch) txToReq(tx []byte) *Request {
	now := time.Now()
	req := &Request{
		Timestamp: &google_protobuf.Timestamp{
			Seconds: now.Unix(),
			Nanos:   int32(now.UnixNano() % 1000000000),
		},
		Payload:   tx,
		ReplicaId: op.pbft.id,
	}
	// XXX sign req
	return req
}

func (op *obcBatch) processMessage(ocMsg *pb.Message, senderHandle *pb.PeerID) events.Event {
	if ocMsg.Type == pb.Message_CHAIN_TRANSACTION {
		req := op.txToReq(ocMsg.Payload)
		return op.submitToLeader(req)
	}

	if ocMsg.Type != pb.Message_CONSENSUS {
		logger.Errorf("Unexpected message type: %s", ocMsg.Type)
		return nil
	}

	batchMsg := &BatchMessage{}
	err := proto.Unmarshal(ocMsg.Payload, batchMsg)
	if err != nil {
		logger.Errorf("Error unmarshaling message: %s", err)
		return nil
	}

	if req := batchMsg.GetRequest(); req != nil {
		if (op.pbft.primary(op.pbft.view) == op.pbft.id) && op.pbft.activeView {
			return op.leaderProcReq(req)
		}
		op.outstandingReqs[req] = struct{}{}
		return nil
	} else if pbftMsg := batchMsg.GetPbftMessage(); pbftMsg != nil {
		senderID, err := getValidatorID(senderHandle) // who sent this?
		if err != nil {
			panic("Cannot map sender's PeerID to a valid replica ID")
		}
		msg := &Message{}
		err = proto.Unmarshal(pbftMsg, msg)
		if err != nil {
			logger.Error("Error unpacking payload from message: %s", err)
			return nil
		}
		return pbftMessageEvent{
			msg:    msg,
			sender: senderID,
		}
	}

	logger.Error("Unknown request: %+v", batchMsg)

	return nil
}

// allow the primary to send a batch when the timer expires
func (op *obcBatch) ProcessEvent(event events.Event) events.Event {
	logger.Debugf("Replica %d batch main thread looping", op.pbft.id)
	switch et := event.(type) {
	case batchMessageEvent:
		ocMsg := et
		return op.processMessage(ocMsg.msg, ocMsg.sender)
	case executedEvent:
		op.stack.Commit(nil, et.tag.([]byte))
	case committedEvent:
		op.pbft.ProcessEvent(execDoneEvent{})
		op.startTimerIfOutstandingRequests()
		// If we are the primary, and know of outstanding requests, submit them for inclusion in the next batch until
		// we run out of requests, or a new batch message is triggered (this path will re-enter after execution)
		if op.pbft.primary(op.pbft.view) == op.pbft.id && op.pbft.activeView {
			for nreq := range op.outstandingReqs {
				if msg := op.leaderProcReq(nreq); msg != nil {
					return msg
				}
			}
		}
	case batchTimerEvent:
		logger.Infof("Replica %d batch timer expired", op.pbft.id)
		if op.pbft.activeView && (len(op.batchStore) > 0) {
			return op.sendBatch()
		}
	case viewChangedEvent:
		// Outstanding reqs doesn't make sense for batch, as all the requests in a batch may be processed
		// in a different batch, but PBFT core can't see through the opaque structure to see this
		// so, on view change, clear it out
		op.pbft.outstandingReqs = make(map[string]*Request)

		logger.Debugf("Replica %d batch thread recognizing new view", op.pbft.id)
		if op.batchTimerActive {
			op.stopBatchTimer()
		}

	case stateUpdatedEvent:
		// When the state is updated, clear any outstanding requests, they may have been processed while we were gone
		op.outstandingReqs = make(map[*Request]struct{})
		return op.pbft.ProcessEvent(event)
	default:
		return op.pbft.ProcessEvent(event)
	}

	return nil
}

func (op *obcBatch) startBatchTimer() {
	op.batchTimer.Reset(op.batchTimeout, batchTimerEvent{})
	logger.Debugf("Replica %d started the batch timer", op.pbft.id)
	op.batchTimerActive = true
}

func (op *obcBatch) stopBatchTimer() {
	op.batchTimer.Stop()
	logger.Debugf("Replica %d stopped the batch timer", op.pbft.id)
	op.batchTimerActive = false
}

// Wraps a payload into a batch message, packs it and wraps it into
// a Fabric message. Called by broadcast before transmission.
func (op *obcBatch) wrapMessage(msgPayload []byte) *pb.Message {
	batchMsg := &BatchMessage{&BatchMessage_PbftMessage{msgPayload}}
	packedBatchMsg, _ := proto.Marshal(batchMsg)
	ocMsg := &pb.Message{
		Type:    pb.Message_CONSENSUS,
		Payload: packedBatchMsg,
	}
	return ocMsg
}

// Retrieve the idle channel, only used for testing
func (op *obcBatch) idleChannel() <-chan struct{} {
	return op.idleChan
}

// TODO, temporary
func (op *obcBatch) getManager() events.Manager {
	return op.manager
}

func (op *obcBatch) startTimerIfOutstandingRequests() {
	if op.pbft.skipInProgress || op.pbft.currentExec != nil {
		// Do not start view change timer if some background event is in progress
		return
	}

	if len(op.outstandingReqs) == 0 {
		// Only start a timer if we are aware of outstanding requests
		return
	}
	op.pbft.softStartTimer(op.pbft.requestTimeout, "Batch outstanding requests")
}
<|endoftext|>


<commit_before>package structs

import (
	"reflect"
	"testing"
)

func TestEncodeDecode(t *testing.T) {
	arg := &RegisterRequest{
		Datacenter: "foo",
		Node:       "bar",
		Address:    "baz",
		Service: &NodeService{
			Service: "test",
			Address: "127.0.0.2",
		},
	}
	buf, err := Encode(RegisterRequestType, arg)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out RegisterRequest
	err = Decode(buf[1:], &out)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(arg.Service, out.Service) {
		t.Fatalf("bad: %#v %#v", arg.Service, out.Service)
	}
	if !reflect.DeepEqual(arg, &out) {
		t.Fatalf("bad: %#v %#v", arg, out)
	}
}

func TestStructs_Implements(t *testing.T) {
	var (
		_ RPCInfo          = &RegisterRequest{}
		_ RPCInfo          = &DeregisterRequest{}
		_ RPCInfo          = &DCSpecificRequest{}
		_ RPCInfo          = &ServiceSpecificRequest{}
		_ RPCInfo          = &NodeSpecificRequest{}
		_ RPCInfo          = &ChecksInStateRequest{}
		_ RPCInfo          = &KVSRequest{}
		_ RPCInfo          = &KeyRequest{}
		_ RPCInfo          = &KeyListRequest{}
		_ RPCInfo          = &SessionRequest{}
		_ RPCInfo          = &SessionSpecificRequest{}
		_ RPCInfo          = &EventFireRequest{}
		_ RPCInfo          = &ACLPolicyRequest{}
		_ RPCInfo          = &KeyringRequest{}
		_ CompoundResponse = &KeyringResponses{}
	)
}
<commit_msg>Adds unit tests for new structs clone functions.<commit_after>package structs

import (
	"reflect"
	"testing"
)

func TestEncodeDecode(t *testing.T) {
	arg := &RegisterRequest{
		Datacenter: "foo",
		Node:       "bar",
		Address:    "baz",
		Service: &NodeService{
			Service: "test",
			Address: "127.0.0.2",
		},
	}
	buf, err := Encode(RegisterRequestType, arg)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out RegisterRequest
	err = Decode(buf[1:], &out)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(arg.Service, out.Service) {
		t.Fatalf("bad: %#v %#v", arg.Service, out.Service)
	}
	if !reflect.DeepEqual(arg, &out) {
		t.Fatalf("bad: %#v %#v", arg, out)
	}
}

func TestStructs_Implements(t *testing.T) {
	var (
		_ RPCInfo          = &RegisterRequest{}
		_ RPCInfo          = &DeregisterRequest{}
		_ RPCInfo          = &DCSpecificRequest{}
		_ RPCInfo          = &ServiceSpecificRequest{}
		_ RPCInfo          = &NodeSpecificRequest{}
		_ RPCInfo          = &ChecksInStateRequest{}
		_ RPCInfo          = &KVSRequest{}
		_ RPCInfo          = &KeyRequest{}
		_ RPCInfo          = &KeyListRequest{}
		_ RPCInfo          = &SessionRequest{}
		_ RPCInfo          = &SessionSpecificRequest{}
		_ RPCInfo          = &EventFireRequest{}
		_ RPCInfo          = &ACLPolicyRequest{}
		_ RPCInfo          = &KeyringRequest{}
		_ CompoundResponse = &KeyringResponses{}
	)
}

func TestStructs_ServiceNode_Clone(t *testing.T) {
	sn := &ServiceNode{
		Node:           "node1",
		Address:        "127.0.0.1",
		ServiceID:      "service1",
		ServiceName:    "dogs",
		ServiceTags:    []string{"prod", "v1"},
		ServiceAddress: "127.0.0.2",
		ServicePort:    8080,
		RaftIndex:      RaftIndex{
			CreateIndex: 1,
			ModifyIndex: 2,
		},
	}

	clone := sn.Clone()
	if !reflect.DeepEqual(sn, clone) {
		t.Fatalf("bad: %v", clone)
	}

	sn.ServiceTags = append(sn.ServiceTags, "hello")
	if reflect.DeepEqual(sn, clone) {
		t.Fatalf("clone wasn't independent of the original")
	}
}

func TestStructs_DirEntry_Clone(t *testing.T) {
	e := &DirEntry{
		LockIndex: 5,
		Key: "hello",
		Flags: 23,
		Value: []byte("this is a test"),
		Session: "session1",
		RaftIndex:      RaftIndex{
			CreateIndex: 1,
			ModifyIndex: 2,
		},
	}

	clone := e.Clone()
	if !reflect.DeepEqual(e, clone) {
		t.Fatalf("bad: %v", clone)
	}

	e.Value = []byte("a new value")
	if reflect.DeepEqual(e, clone) {
		t.Fatalf("clone wasn't independent of the original")
	}
}
<|endoftext|>


<commit_before>package consumer

import (
	"flume-bridge/config"
	"flume-bridge/consumer/pool"
	"github.com/garyburd/redigo/redis"
	"github.com/momotech/GoRedis/libs/stdlog"
	"log"
	"os"
	"strconv"
	"sync"
	"time"
)

type poolwrapper struct {
	hostport config.HostPort

	rpool *redis.Pool

	lastValue int64

	currValue int64
}

type SourceManager struct {
	zkmanager *config.ZKManager

	sourceServers map[string]*SourceServer //业务名称和sourceserver对应

	hp2flumeClientPool map[config.HostPort]*pool.FlumePoolLink //对应的Pool

	redispool map[string][]*poolwrapper // 对应的redispool

	watcherPool map[string]*config.Watcher //watcherPool

	mutex sync.Mutex

	isRunning bool

	instancename string

	flumeLog         stdlog.Logger
	redisLog         stdlog.Logger
	watcherLog       stdlog.Logger
	flumePoolLog     stdlog.Logger
	flumeSourceLog   stdlog.Logger
	sourceManagerLog stdlog.Logger
	option           *config.Option
}

func NewSourceManager(instancename string, option *config.Option) *SourceManager {

	sourcemanager := &SourceManager{}
	sourcemanager.option = option
	sourcemanager.sourceServers = make(map[string]*SourceServer)
	sourcemanager.hp2flumeClientPool = make(map[config.HostPort]*pool.FlumePoolLink)
	sourcemanager.watcherPool = make(map[string]*config.Watcher)

	//创建使用的Logger
	basepath := option.LogPath + "/" + instancename
	sourcemanager.sourceManagerLog = buildLog(basepath, "source_manager", "source_manager.log")
	sourcemanager.flumeLog = buildLog(basepath, "flume_tps", "flume_tps.log")
	sourcemanager.flumePoolLog = buildLog(basepath, "flume_pool", "flume_pool.log")
	sourcemanager.redisLog = buildLog(basepath, "redis_tps", "redis_tps.log")
	sourcemanager.watcherLog = buildLog(basepath, "zk_watcher", "zk_watcher.log")
	sourcemanager.flumeSourceLog = buildLog(basepath, "flume_source", "flume_source.log")

	sourcemanager.redispool = initRedisQueue(option)
	//从zk中拉取flumenode的配置
	zkmanager := config.NewZKManager(option.Zkhost)
	sourcemanager.zkmanager = zkmanager
	sourcemanager.instancename = instancename

	sourcemanager.initSourceServers(option.Businesses, zkmanager)
	return sourcemanager

}

func buildLog(basepath, logname, filename string) stdlog.Logger {

	_, err := os.Stat(basepath)
	if nil != err {
		err := os.MkdirAll(basepath, os.ModePerm)
		if nil != err {
			panic(err)
		}
	}

	//创建redis的log
	f, err := os.OpenFile(basepath+"/"+filename, os.O_CREATE|os.O_WRONLY|os.O_APPEND, os.ModePerm)
	if nil != err {
		panic(err)
	}
	logger := stdlog.Log(logname)
	logger.SetOutput(f)
	logger.SetPrefix(func() string {
		now := time.Now()
		nt := now.Format("2006-01-02 15:04:05")
		return nt + "\t"
	})
	return logger
}

func initRedisQueue(option *config.Option) map[string][]*poolwrapper {
	redispool := make(map[string][]*poolwrapper, 0)

	//创建redis的消费连接
	for _, v := range option.QueueHostPorts {

		hp := v
		pool := redis.NewPool(func() (conn redis.Conn, err error) {
			conn, err = redis.DialTimeout("tcp", hp.Host+":"+strconv.Itoa(hp.Port),
				time.Duration(hp.Timeout)*time.Second,
				time.Duration(hp.Timeout)*time.Second,
				time.Duration(hp.Timeout)*time.Second)

			return
		}, hp.Maxconn/2)

		pools, ok := redispool[v.QueueName]
		if !ok {
			pools = make([]*poolwrapper, 0)
			redispool[v.QueueName] = pools
		}

		poolw := &poolwrapper{}
		poolw.hostport = v.HostPort
		poolw.rpool = pool
		redispool[v.QueueName] = append(pools, poolw)
	}

	return redispool
}

func (self *SourceManager) initSourceServers(businesses []string, zkmanager *config.ZKManager) {

	for _, business := range businesses {
		nodewatcher := newFlumeWatcher(business, self)
		flumeNode := zkmanager.GetAndWatch(business, nodewatcher)
		self.watcherPool[business] = nodewatcher
		sserver := self.initSourceServer(business, flumeNode)
		self.sourceServers[business] = sserver
	}

	//-------------------注册当前进程ID到zk
	currpid := os.Getpid()
	hostname, _ := os.Hostname()
	self.zkmanager.RegistePath(businesses, hostname+"_"+self.instancename+":"+strconv.Itoa(currpid))

}

func (self *SourceManager) initSourceServer(business string, flumenodes []config.HostPort) *SourceServer {

	//首先判断当前是否该sink支持该种business
	_, ok := self.watcherPool[business]
	if !ok {
		self.sourceManagerLog.Printf("unsupport business[%s],HostPorts:[%s]\n", business, flumenodes)
		return nil
	}

	if len(flumenodes) == 0 {
		self.sourceManagerLog.Println("no valid flume agent node for [" + business + "]")
		return nil
	}

	//新增的消费类型
	//使用的pool
	pools := self.initFlumeClientPool(business, flumenodes)

	//创建一个sourceserver
	sourceserver := newSourceServer(business, pools, self.flumeSourceLog)

	return sourceserver
}

func (self *SourceManager) initFlumeClientPool(business string, flumenodes []config.HostPort) []*pool.FlumePoolLink {

	pools := make([]*pool.FlumePoolLink, 0, 10)
	for _, hp := range flumenodes {
		poollink, ok := self.hp2flumeClientPool[hp]
		if !ok {
			err, tmppool := pool.NewFlumePoolLink(hp)
			if nil != err {
				self.sourceManagerLog.Println("SOURCE_MANGER|INIT FLUMEPOOLLINE|FAIL|%s", err)
				continue
			}
			poollink = tmppool
			self.hp2flumeClientPool[hp] = poollink
		}

		defer func() {
			if nil == poollink {
				return
			}
			if err := recover(); nil != err {
				self.sourceManagerLog.Printf("SOURCE_MANGER|CREATE FLUMECLIENT|FAIL|[%s]\n", hp)
				poollink = nil
			}
		}()

		if nil == poollink {
			continue
		}

		poollink.AttachBusiness(business)
		pools = append(pools, poollink)
	}
	self.sourceManagerLog.Printf("SOURCE_MANGER|CREATE FLUMECLIENT|SUCCESS|[%s]\n", pools)

	return pools
}

func (self *SourceManager) Start() {

	for _, v := range self.sourceServers {
		v.start()
	}
	self.isRunning = true
	go self.monitor()
	self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|[%s]|STARTED\n", self.instancename)
	self.startWorker()

}

func (self *SourceManager) startWorker() {

	for k, v := range self.redispool {
		self.sourceManagerLog.Println("LOG_SOURCE_MANGER|REDIS|[" + k + "]|START")
		for _, pool := range v {
			self.sourceManagerLog.Println("LOG_SOURCE_MANGER|REDIS|POOL|[" + pool.hostport.Host + "]|START")
			for i := 0; i < 10; i++ {
				go func(queuename string, pool *poolwrapper) {
					//批量收集数据
					conn := pool.rpool.Get()
					defer conn.Close()
					for self.isRunning {

						reply, err := conn.Do("LPOP", queuename)
						if nil != err || nil == reply {
							if nil != err {
								self.sourceManagerLog.Printf("LPOP|FAIL|%T", err)
								conn.Close()
								conn = pool.rpool.Get()
							} else {
								time.Sleep(100 * time.Millisecond)
							}
							continue
						}

						//计数器++
						pool.currValue++

						resp := reply.([]byte)

						if self.option.IsCompress {
							resp = decompress(resp)
						}
						if resp == nil {
							continue
						}
						businessName, logType, event := decodeCommand(resp)
						if nil == event {
							continue
						}

						//提交到对应business的channel中
						routeKey := businessName + logType
						defaultRoutKey := "default" + logType
						sourceServer, ok := self.sourceServers[routeKey]
						if !ok {
							//use the default channel
							sourceServer, ok := self.sourceServers[defaultRoutKey]
							if ok && nil != sourceServer && !sourceServer.isStop {
								sourceServer.buffChannel <- event
							} else {
								self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|DEFAULT SOURCE_SERVER NOT EXSIT OR STOPPED\n")
							}
						} else {
							if !sourceServer.isStop {
								sourceServer.buffChannel <- event
							} else {
								self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|SOURCE_SERVER STOPPED|%s\n", routeKey)
							}
						}
					}
					self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|REDIS-POP|EXIT|%s|%s\n", queuename, self.instancename)
				}(k, pool)
			}
		}
	}

}

func (self *SourceManager) Close() {
	self.isRunning = false

	for _, sourceserver := range self.sourceServers {
		sourceserver.stop()
	}

	for _, redispool := range self.redispool {
		for _, pool := range redispool {
			pool.rpool.Close()
		}
	}

	//关闭flumepool
	for _, flumepool := range self.hp2flumeClientPool {
		flumepool.FlumePool.Destroy()
	}

	log.Printf("LOG_SOURCE_MANGER|[%s]|STOP\n", self.instancename)
}
<commit_msg>	modified:   consumer/log_source_manager.go<commit_after>package consumer

import (
	"flume-bridge/config"
	"flume-bridge/consumer/pool"
	"github.com/garyburd/redigo/redis"
	"github.com/momotech/GoRedis/libs/stdlog"
	"log"
	"os"
	"strconv"
	"sync"
	"time"
)

type poolwrapper struct {
	hostport config.HostPort

	rpool *redis.Pool

	lastValue int64

	currValue int64
}

type SourceManager struct {
	zkmanager *config.ZKManager

	sourceServers map[string]*SourceServer //业务名称和sourceserver对应

	hp2flumeClientPool map[config.HostPort]*pool.FlumePoolLink //对应的Pool

	redispool map[string][]*poolwrapper // 对应的redispool

	watcherPool map[string]*config.Watcher //watcherPool

	mutex sync.Mutex

	isRunning bool

	instancename string

	flumeLog         stdlog.Logger
	redisLog         stdlog.Logger
	watcherLog       stdlog.Logger
	flumePoolLog     stdlog.Logger
	flumeSourceLog   stdlog.Logger
	sourceManagerLog stdlog.Logger
	option           *config.Option
}

func NewSourceManager(instancename string, option *config.Option) *SourceManager {

	sourcemanager := &SourceManager{}
	sourcemanager.option = option
	sourcemanager.sourceServers = make(map[string]*SourceServer)
	sourcemanager.hp2flumeClientPool = make(map[config.HostPort]*pool.FlumePoolLink)
	sourcemanager.watcherPool = make(map[string]*config.Watcher)

	//创建使用的Logger
	basepath := option.LogPath + "/" + instancename
	sourcemanager.sourceManagerLog = buildLog(basepath, "source_manager", "source_manager.log")
	sourcemanager.flumeLog = buildLog(basepath, "flume_tps", "flume_tps.log")
	sourcemanager.flumePoolLog = buildLog(basepath, "flume_pool", "flume_pool.log")
	sourcemanager.redisLog = buildLog(basepath, "redis_tps", "redis_tps.log")
	sourcemanager.watcherLog = buildLog(basepath, "zk_watcher", "zk_watcher.log")
	sourcemanager.flumeSourceLog = buildLog(basepath, "flume_source", "flume_source.log")

	sourcemanager.redispool = initRedisQueue(option)
	//从zk中拉取flumenode的配置
	zkmanager := config.NewZKManager(option.Zkhost)
	sourcemanager.zkmanager = zkmanager
	sourcemanager.instancename = instancename

	sourcemanager.initSourceServers(option.Businesses, zkmanager)
	sourcemanager.sourceManagerLog.Printf("SOURCE_MANGER|Init SourceManager|Succ...")
	return sourcemanager

}

func buildLog(basepath, logname, filename string) stdlog.Logger {

	_, err := os.Stat(basepath)
	if nil != err {
		err := os.MkdirAll(basepath, os.ModePerm)
		if nil != err {
			panic(err)
		}
	}

	//创建redis的log
	f, err := os.OpenFile(basepath+"/"+filename, os.O_CREATE|os.O_WRONLY|os.O_APPEND, os.ModePerm)
	if nil != err {
		panic(err)
	}
	logger := stdlog.Log(logname)
	logger.SetOutput(f)
	logger.SetPrefix(func() string {
		now := time.Now()
		nt := now.Format("2006-01-02 15:04:05")
		return nt + "\t"
	})
	return logger
}

func initRedisQueue(option *config.Option) map[string][]*poolwrapper {
	redispool := make(map[string][]*poolwrapper, 0)

	//创建redis的消费连接
	for _, v := range option.QueueHostPorts {

		hp := v
		pool := redis.NewPool(func() (conn redis.Conn, err error) {
			conn, err = redis.DialTimeout("tcp", hp.Host+":"+strconv.Itoa(hp.Port),
				time.Duration(hp.Timeout)*time.Second,
				time.Duration(hp.Timeout)*time.Second,
				time.Duration(hp.Timeout)*time.Second)

			return
		}, hp.Maxconn/2)

		pools, ok := redispool[v.QueueName]
		if !ok {
			pools = make([]*poolwrapper, 0)
			redispool[v.QueueName] = pools
		}

		poolw := &poolwrapper{}
		poolw.hostport = v.HostPort
		poolw.rpool = pool
		redispool[v.QueueName] = append(pools, poolw)
	}

	return redispool
}

func (self *SourceManager) initSourceServers(businesses []string, zkmanager *config.ZKManager) {

	for _, business := range businesses {
		nodewatcher := newFlumeWatcher(business, self)
		flumeNode := zkmanager.GetAndWatch(business, nodewatcher)
		self.watcherPool[business] = nodewatcher
		sserver := self.initSourceServer(business, flumeNode)
		self.sourceServers[business] = sserver
	}

	//-------------------注册当前进程ID到zk
	currpid := os.Getpid()
	hostname, _ := os.Hostname()
	self.zkmanager.RegistePath(businesses, hostname+"_"+self.instancename+":"+strconv.Itoa(currpid))

}

func (self *SourceManager) initSourceServer(business string, flumenodes []config.HostPort) *SourceServer {

	//首先判断当前是否该sink支持该种business
	_, ok := self.watcherPool[business]
	if !ok {
		self.sourceManagerLog.Printf("unsupport business[%s],HostPorts:[%s]\n", business, flumenodes)
		return nil
	}

	if len(flumenodes) == 0 {
		self.sourceManagerLog.Println("no valid flume agent node for [" + business + "]")
		return nil
	}

	//新增的消费类型
	//使用的pool
	pools := self.initFlumeClientPool(business, flumenodes)

	//创建一个sourceserver
	sourceserver := newSourceServer(business, pools, self.flumeSourceLog)

	return sourceserver
}

func (self *SourceManager) initFlumeClientPool(business string, flumenodes []config.HostPort) []*pool.FlumePoolLink {

	pools := make([]*pool.FlumePoolLink, 0, 10)
	for _, hp := range flumenodes {
		poollink, ok := self.hp2flumeClientPool[hp]
		if !ok {
			err, tmppool := pool.NewFlumePoolLink(hp)
			if nil != err {
				self.sourceManagerLog.Println("SOURCE_MANGER|INIT FLUMEPOOLLINE|FAIL|%s", err)
				continue
			}
			poollink = tmppool
			self.hp2flumeClientPool[hp] = poollink
		}

		defer func() {
			if nil == poollink {
				return
			}
			if err := recover(); nil != err {
				self.sourceManagerLog.Printf("SOURCE_MANGER|CREATE FLUMECLIENT|FAIL|[%s]\n", hp)
				poollink = nil
			}
		}()

		if nil == poollink {
			continue
		}

		poollink.AttachBusiness(business)
		pools = append(pools, poollink)
	}
	self.sourceManagerLog.Printf("SOURCE_MANGER|CREATE FLUMECLIENT|SUCCESS|[%s,%d]\n", business, len(flumenodes))

	return pools
}

func (self *SourceManager) Start() {

	self.sourceManagerLog.Println("SourceManager|Start|Begin...")
	for _, v := range self.sourceServers {
		v.start()
		self.sourceManagerLog.Printf("SourceManager|SourceServer|%s|started...", v.business)
	}
	self.isRunning = true
	go self.monitor()
	self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|[%s]|STARTED\n", self.instancename)
	self.startWorker()
	self.sourceManagerLog.Println("SourceManager|Start|End...")

}

func (self *SourceManager) startWorker() {

	for k, v := range self.redispool {
		self.sourceManagerLog.Println("LOG_SOURCE_MANGER|REDIS|[" + k + "]|START")
		for _, pool := range v {
			self.sourceManagerLog.Println("LOG_SOURCE_MANGER|REDIS|POOL|[" + pool.hostport.Host + "]|START")
			for i := 0; i < 10; i++ {
				go func(queuename string, pool *poolwrapper) {
					//批量收集数据
					conn := pool.rpool.Get()
					defer conn.Close()
					for self.isRunning {

						reply, err := conn.Do("LPOP", queuename)
						if nil != err || nil == reply {
							if nil != err {
								self.sourceManagerLog.Printf("LPOP|FAIL|%T", err)
								conn.Close()
								conn = pool.rpool.Get()
							} else {
								time.Sleep(100 * time.Millisecond)
							}
							continue
						}

						//计数器++
						pool.currValue++

						resp := reply.([]byte)

						if self.option.IsCompress {
							resp = decompress(resp)
						}
						if resp == nil {
							continue
						}
						businessName, logType, event := decodeCommand(resp)
						if nil == event {
							continue
						}

						//提交到对应business的channel中
						routeKey := businessName + logType
						defaultRoutKey := "default" + logType
						sourceServer, ok := self.sourceServers[routeKey]
						if !ok {
							//use the default channel
							sourceServer, ok := self.sourceServers[defaultRoutKey]
							if ok && nil != sourceServer && !sourceServer.isStop {
								sourceServer.buffChannel <- event
							} else {
								self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|DEFAULT SOURCE_SERVER NOT EXSIT OR STOPPED\n")
							}
						} else {
							if !sourceServer.isStop {
								sourceServer.buffChannel <- event
							} else {
								self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|SOURCE_SERVER STOPPED|%s\n", routeKey)
							}
						}
					}
					self.sourceManagerLog.Printf("LOG_SOURCE_MANGER|REDIS-POP|EXIT|%s|%s\n", queuename, self.instancename)
				}(k, pool)
			}
		}
	}

}

func (self *SourceManager) Close() {
	self.isRunning = false

	for _, sourceserver := range self.sourceServers {
		sourceserver.stop()
	}

	for _, redispool := range self.redispool {
		for _, pool := range redispool {
			pool.rpool.Close()
		}
	}

	//关闭flumepool
	for _, flumepool := range self.hp2flumeClientPool {
		flumepool.FlumePool.Destroy()
	}

	log.Printf("LOG_SOURCE_MANGER|[%s]|STOP\n", self.instancename)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/nbari/violetear"
)

func main() {
	d := violetear.NewDynamic()
	d.Set(":ip", `^(?:[0-9]{1,3}\.){3}[0-9]{1,3}$`)
	d.Set(":uuid", "[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}")

	uuid := "2E9C64A5-FF13-4DC5-A957-F39E39ABDC48"
	for k, v := range d {
		if v.MatchString(uuid) {
			fmt.Printf("Match ---> %v", k)
		}
		fmt.Println(k)

	}

}
<commit_msg>	modified:   dynamic.go<commit_after>package main

import (
	"fmt"
	"github.com/nbari/violetear"
)

func main() {
	d := violetear.NewDynamicSet()
	d.Set(":ip", `^(?:[0-9]{1,3}\.){3}[0-9]{1,3}$`)
	d.Set(":uuid", "[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}")

	uuid := "2E9C64A5-FF13-4DC5-A957-F39E39ABDC48"
	for k, v := range d {
		if v.MatchString(uuid) {
			fmt.Printf("Match ---> %v", k)
		}
		fmt.Println(k)

	}

}
<|endoftext|>


<commit_before>// +build e2e

/*
Copyright 2020 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package autotls

import (
	"crypto/tls"
	"crypto/x509"
	"net/http"
	"testing"

	"github.com/kelseyhightower/envconfig"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"

	"knative.dev/pkg/ptr"
	"knative.dev/pkg/test/spoof"
	"knative.dev/serving/pkg/apis/networking"
	servingv1 "knative.dev/serving/pkg/apis/serving/v1"
	routenames "knative.dev/serving/pkg/reconciler/route/resources/names"
	"knative.dev/serving/test"
	testingress "knative.dev/serving/test/conformance/ingress"
	"knative.dev/serving/test/e2e"
	v1test "knative.dev/serving/test/v1"
)

type dnsRecord struct {
	ip     string
	domain string
}

type config struct {
	// ServiceName is the name of testing Knative Service.
	// It is not required for self-signed CA or for the HTTP01 challenge when wildcard domain
	// is mapped to the Ingress IP.
	TLSServiceName string `envconfig:"tls_service_name" required: "false"`
	// AutoTLSTestName is the name of the auto tls. It is not required for local test.
	AutoTLSTestName string `envconfig:"auto_tls_test_name" required: "false" default:"TestAutoTLS"`
}

var env config

func TestAutoTLS(t *testing.T) {
	if err := envconfig.Process("", &env); err != nil {
		t.Fatalf("Failed to process environment variable: %v.", err)
	}
	t.Run(env.AutoTLSTestName, testAutoTLS)
}

func testAutoTLS(t *testing.T) {
	clients := e2e.Setup(t)

	names := test.ResourceNames{
		Service: test.ObjectNameForTest(t),
		Image:   "runtime",
	}
	if len(env.TLSServiceName) != 0 {
		names.Service = env.TLSServiceName
	}
	test.CleanupOnInterrupt(func() { test.TearDown(clients, names) })
	defer test.TearDown(clients, names)

	objects, err := v1test.CreateServiceReady(t, clients, &names)
	if err != nil {
		t.Fatalf("Failed to create initial Service: %v: %v", names.Service, err)
	}

	// The TLS info is added to the ingress after the service is created, that's
	// why we need to wait again
	err = v1test.WaitForServiceState(clients.ServingClient, names.Service, httpsReady, "HTTPSIsReady")
	if err != nil {
		t.Fatalf("Service %s did not become ready or have HTTPS URL: %v", names.Service, err)
	}

	// curl HTTPS
	certName := getCertificateName(t, clients, objects)
	rootCAs := createRootCAs(t, clients, objects.Route.Namespace, certName)
	httpsClient := createHTTPSClient(t, clients, objects, rootCAs)
	testingress.RuntimeRequest(t, httpsClient, objects.Service.Status.URL.String())

	t.Run("Tag route", func(t *testing.T) {
		// Probe main URL while we update the route
		var transportOption spoof.TransportOption = func(transport *http.Transport) *http.Transport {
			transport.TLSClientConfig = &tls.Config{RootCAs: rootCAs}
			return transport
		}
		prober := test.RunRouteProber(t.Logf, clients, objects.Service.Status.URL.URL(), transportOption)
		defer test.AssertProberDefault(t, prober)

		if _, err := v1test.UpdateServiceRouteSpec(t, clients, names, servingv1.RouteSpec{
			Traffic: []servingv1.TrafficTarget{{
				Tag:            "tag1",
				Percent:        ptr.Int64(50),
				LatestRevision: ptr.Bool(true),
			}, {
				Tag:            "tag2",
				Percent:        ptr.Int64(50),
				LatestRevision: ptr.Bool(true),
			}},
		}); err != nil {
			t.Fatalf("Failed to update Service route spec: %v", err)
		}
		if err = v1test.WaitForRouteState(clients.ServingClient, names.Route, routeTrafficHTTPS, "RouteTrafficIsHTTPS"); err != nil {
			t.Fatalf("Traffic for route: %s is not HTTPS: %v", names.Route, err)
		}

		ing, err := clients.NetworkingClient.Ingresses.Get(routenames.Ingress(objects.Route), metav1.GetOptions{})
		if err != nil {
			t.Fatalf("Failed to get ingress: %v", err)
		}
		for _, tls := range ing.Spec.TLS {
			// Each new cert has to be added to the root pool so we can make requests.
			if !rootCAs.AppendCertsFromPEM(getPEMDataFromSecret(t, clients, tls.SecretNamespace, tls.SecretName)) {
				t.Fatal("Failed to add the certificate to the root CA")
			}
		}

		route, err := clients.ServingClient.Routes.Get(objects.Route.Name, metav1.GetOptions{})
		if err != nil {
			t.Fatalf("Failed to get route: %v", err)
		}
		httpsClient := createHTTPSClient(t, clients, objects, rootCAs)
		for _, traffic := range route.Status.Traffic {
			testingress.RuntimeRequest(t, httpsClient, traffic.URL.String())
		}
	})
}

func getCertificateName(t *testing.T, clients *test.Clients, objects *v1test.ResourceObjects) string {
	t.Helper()
	ing, err := clients.NetworkingClient.Ingresses.Get(routenames.Ingress(objects.Route), metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Failed to get Ingress %s: %v", routenames.Ingress(objects.Route), err)
	}
	if len(ing.Spec.TLS) == 0 {
		t.Fatalf("IngressTLS field in Ingress %s does not exist.", ing.Name)
	}
	return ing.Spec.TLS[0].SecretName
}

func routeTrafficHTTPS(route *servingv1.Route) (bool, error) {
	for _, tt := range route.Status.Traffic {
		if tt.URL.URL().Scheme != "https" {
			return false, nil
		}
	}
	return route.Status.IsReady() && true, nil
}

func httpsReady(svc *servingv1.Service) (bool, error) {
	if ready, err := v1test.IsServiceReady(svc); err != nil {
		return ready, err
	} else if !ready {
		return false, nil
	} else {
		return svc.Status.URL.Scheme == "https", nil
	}
}

func getPEMDataFromSecret(t *testing.T, clients *test.Clients, ns, secretName string) []byte {
	t.Helper()
	secret, err := clients.KubeClient.Kube.CoreV1().Secrets(ns).Get(
		secretName, metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Failed to get Secret %s: %v", secretName, err)
	}
	return secret.Data[corev1.TLSCertKey]
}

func createRootCAs(t *testing.T, clients *test.Clients, ns, secretName string) *x509.CertPool {
	t.Helper()
	pemData := getPEMDataFromSecret(t, clients, ns, secretName)

	rootCAs, err := x509.SystemCertPool()
	if rootCAs == nil || err != nil {
		if err != nil {
			t.Logf("Failed to load cert poll from system: %v. Will create a new cert pool.", err)
		}
		rootCAs = x509.NewCertPool()
	}
	if !rootCAs.AppendCertsFromPEM(pemData) {
		t.Fatal("Failed to add the certificate to the root CA")
	}
	return rootCAs
}

func createHTTPSClient(t *testing.T, clients *test.Clients, objects *v1test.ResourceObjects, rootCAs *x509.CertPool) *http.Client {
	t.Helper()
	ing, err := clients.NetworkingClient.Ingresses.Get(routenames.Ingress(objects.Route), metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Failed to get Ingress %s: %v", routenames.Ingress(objects.Route), err)
	}
	dialer := testingress.CreateDialContext(t, ing, clients)
	tlsConfig := &tls.Config{
		RootCAs: rootCAs,
	}
	return &http.Client{
		Transport: &http.Transport{
			DialContext:     dialer,
			TLSClientConfig: tlsConfig,
		}}
}

func disableNamespaceCertWithWhiteList(t *testing.T, clients *test.Clients, whiteLists sets.String) {
	t.Helper()
	namespaces, err := clients.KubeClient.Kube.CoreV1().Namespaces().List(metav1.ListOptions{})
	if err != nil {
		t.Fatalf("Failed to list namespaces: %v", err)
	}
	for _, ns := range namespaces.Items {
		if ns.Labels == nil {
			ns.Labels = map[string]string{}
		}
		if whiteLists.Has(ns.Name) {
			delete(ns.Labels, networking.DisableWildcardCertLabelKey)
		} else {
			ns.Labels[networking.DisableWildcardCertLabelKey] = "true"
		}
		if _, err := clients.KubeClient.Kube.CoreV1().Namespaces().Update(&ns); err != nil {
			t.Errorf("Fail to disable namespace cert: %v", err)
		}
	}
}
<commit_msg>Start prober after append root CA for tag route (#7356)<commit_after>// +build e2e

/*
Copyright 2020 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package autotls

import (
	"crypto/tls"
	"crypto/x509"
	"net/http"
	"testing"

	"github.com/kelseyhightower/envconfig"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"

	"knative.dev/pkg/ptr"
	"knative.dev/pkg/test/spoof"
	"knative.dev/serving/pkg/apis/networking"
	servingv1 "knative.dev/serving/pkg/apis/serving/v1"
	routenames "knative.dev/serving/pkg/reconciler/route/resources/names"
	"knative.dev/serving/test"
	testingress "knative.dev/serving/test/conformance/ingress"
	"knative.dev/serving/test/e2e"
	v1test "knative.dev/serving/test/v1"
)

type dnsRecord struct {
	ip     string
	domain string
}

type config struct {
	// ServiceName is the name of testing Knative Service.
	// It is not required for self-signed CA or for the HTTP01 challenge when wildcard domain
	// is mapped to the Ingress IP.
	TLSServiceName string `envconfig:"tls_service_name" required: "false"`
	// AutoTLSTestName is the name of the auto tls. It is not required for local test.
	AutoTLSTestName string `envconfig:"auto_tls_test_name" required: "false" default:"TestAutoTLS"`
}

var env config

func TestAutoTLS(t *testing.T) {
	if err := envconfig.Process("", &env); err != nil {
		t.Fatalf("Failed to process environment variable: %v.", err)
	}
	t.Run(env.AutoTLSTestName, testAutoTLS)
}

func testAutoTLS(t *testing.T) {
	clients := e2e.Setup(t)

	names := test.ResourceNames{
		Service: test.ObjectNameForTest(t),
		Image:   "runtime",
	}
	if len(env.TLSServiceName) != 0 {
		names.Service = env.TLSServiceName
	}
	test.CleanupOnInterrupt(func() { test.TearDown(clients, names) })
	defer test.TearDown(clients, names)

	objects, err := v1test.CreateServiceReady(t, clients, &names)
	if err != nil {
		t.Fatalf("Failed to create initial Service: %v: %v", names.Service, err)
	}

	// The TLS info is added to the ingress after the service is created, that's
	// why we need to wait again
	err = v1test.WaitForServiceState(clients.ServingClient, names.Service, httpsReady, "HTTPSIsReady")
	if err != nil {
		t.Fatalf("Service %s did not become ready or have HTTPS URL: %v", names.Service, err)
	}

	// curl HTTPS
	certName := getCertificateName(t, clients, objects)
	rootCAs := createRootCAs(t, clients, objects.Route.Namespace, certName)
	httpsClient := createHTTPSClient(t, clients, objects, rootCAs)
	testingress.RuntimeRequest(t, httpsClient, objects.Service.Status.URL.String())

	t.Run("Tag route", func(t *testing.T) {
		// Probe main URL while we update the route
		var transportOption spoof.TransportOption = func(transport *http.Transport) *http.Transport {
			transport.TLSClientConfig = &tls.Config{RootCAs: rootCAs}
			return transport
		}
		if _, err := v1test.UpdateServiceRouteSpec(t, clients, names, servingv1.RouteSpec{
			Traffic: []servingv1.TrafficTarget{{
				Tag:            "tag1",
				Percent:        ptr.Int64(50),
				LatestRevision: ptr.Bool(true),
			}, {
				Tag:            "tag2",
				Percent:        ptr.Int64(50),
				LatestRevision: ptr.Bool(true),
			}},
		}); err != nil {
			t.Fatalf("Failed to update Service route spec: %v", err)
		}
		if err = v1test.WaitForRouteState(clients.ServingClient, names.Route, routeTrafficHTTPS, "RouteTrafficIsHTTPS"); err != nil {
			t.Fatalf("Traffic for route: %s is not HTTPS: %v", names.Route, err)
		}

		ing, err := clients.NetworkingClient.Ingresses.Get(routenames.Ingress(objects.Route), metav1.GetOptions{})
		if err != nil {
			t.Fatalf("Failed to get ingress: %v", err)
		}
		for _, tls := range ing.Spec.TLS {
			// Each new cert has to be added to the root pool so we can make requests.
			if !rootCAs.AppendCertsFromPEM(getPEMDataFromSecret(t, clients, tls.SecretNamespace, tls.SecretName)) {
				t.Fatal("Failed to add the certificate to the root CA")
			}
		}

		// Start prober after the new rootCA is added.
		prober := test.RunRouteProber(t.Logf, clients, objects.Service.Status.URL.URL(), transportOption)
		defer test.AssertProberDefault(t, prober)

		route, err := clients.ServingClient.Routes.Get(objects.Route.Name, metav1.GetOptions{})
		if err != nil {
			t.Fatalf("Failed to get route: %v", err)
		}
		httpsClient := createHTTPSClient(t, clients, objects, rootCAs)
		for _, traffic := range route.Status.Traffic {
			testingress.RuntimeRequest(t, httpsClient, traffic.URL.String())
		}
	})
}

func getCertificateName(t *testing.T, clients *test.Clients, objects *v1test.ResourceObjects) string {
	t.Helper()
	ing, err := clients.NetworkingClient.Ingresses.Get(routenames.Ingress(objects.Route), metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Failed to get Ingress %s: %v", routenames.Ingress(objects.Route), err)
	}
	if len(ing.Spec.TLS) == 0 {
		t.Fatalf("IngressTLS field in Ingress %s does not exist.", ing.Name)
	}
	return ing.Spec.TLS[0].SecretName
}

func routeTrafficHTTPS(route *servingv1.Route) (bool, error) {
	for _, tt := range route.Status.Traffic {
		if tt.URL.URL().Scheme != "https" {
			return false, nil
		}
	}
	return route.Status.IsReady() && true, nil
}

func httpsReady(svc *servingv1.Service) (bool, error) {
	if ready, err := v1test.IsServiceReady(svc); err != nil {
		return ready, err
	} else if !ready {
		return false, nil
	} else {
		return svc.Status.URL.Scheme == "https", nil
	}
}

func getPEMDataFromSecret(t *testing.T, clients *test.Clients, ns, secretName string) []byte {
	t.Helper()
	secret, err := clients.KubeClient.Kube.CoreV1().Secrets(ns).Get(
		secretName, metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Failed to get Secret %s: %v", secretName, err)
	}
	return secret.Data[corev1.TLSCertKey]
}

func createRootCAs(t *testing.T, clients *test.Clients, ns, secretName string) *x509.CertPool {
	t.Helper()
	pemData := getPEMDataFromSecret(t, clients, ns, secretName)

	rootCAs, err := x509.SystemCertPool()
	if rootCAs == nil || err != nil {
		if err != nil {
			t.Logf("Failed to load cert poll from system: %v. Will create a new cert pool.", err)
		}
		rootCAs = x509.NewCertPool()
	}
	if !rootCAs.AppendCertsFromPEM(pemData) {
		t.Fatal("Failed to add the certificate to the root CA")
	}
	return rootCAs
}

func createHTTPSClient(t *testing.T, clients *test.Clients, objects *v1test.ResourceObjects, rootCAs *x509.CertPool) *http.Client {
	t.Helper()
	ing, err := clients.NetworkingClient.Ingresses.Get(routenames.Ingress(objects.Route), metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Failed to get Ingress %s: %v", routenames.Ingress(objects.Route), err)
	}
	dialer := testingress.CreateDialContext(t, ing, clients)
	tlsConfig := &tls.Config{
		RootCAs: rootCAs,
	}
	return &http.Client{
		Transport: &http.Transport{
			DialContext:     dialer,
			TLSClientConfig: tlsConfig,
		}}
}

func disableNamespaceCertWithWhiteList(t *testing.T, clients *test.Clients, whiteLists sets.String) {
	t.Helper()
	namespaces, err := clients.KubeClient.Kube.CoreV1().Namespaces().List(metav1.ListOptions{})
	if err != nil {
		t.Fatalf("Failed to list namespaces: %v", err)
	}
	for _, ns := range namespaces.Items {
		if ns.Labels == nil {
			ns.Labels = map[string]string{}
		}
		if whiteLists.Has(ns.Name) {
			delete(ns.Labels, networking.DisableWildcardCertLabelKey)
		} else {
			ns.Labels[networking.DisableWildcardCertLabelKey] = "true"
		}
		if _, err := clients.KubeClient.Kube.CoreV1().Namespaces().Update(&ns); err != nil {
			t.Errorf("Fail to disable namespace cert: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Mirantis

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 e2e

import (
	"bytes"
	"context"
	"fmt"
	"time"

	. "github.com/onsi/gomega"

	"github.com/Mirantis/virtlet/tests/e2e/framework"
	. "github.com/Mirantis/virtlet/tests/e2e/ginkgo-ext"
)

var _ = Describe("Virtlet restart [Disruptive]", func() {
	var (
		vm *framework.VMInterface
	)

	BeforeAll(func() {
		vm = controller.VM("cirros-vm")
		vm.CreateAndWait(VMOptions{}.ApplyDefaults(), time.Minute*5, nil)
		var err error
		_, err = vm.Pod()
		Expect(err).NotTo(HaveOccurred())

		// restart virtlet before all tests
		virtletPod, err := vm.VirtletPod()
		Expect(err).NotTo(HaveOccurred())

		err = virtletPod.Delete()
		Expect(err).NotTo(HaveOccurred())

		waitVirtletPod(vm)
	})

	AfterAll(func() {
		deleteVM(vm)
	})

	It("Should allow to ssh to VM after virtlet pod restart", func() {
		waitSSH(vm)
	}, 3*60)

	It("Should keep logs from another session", func() {
		var stdout bytes.Buffer
		ctx, closeFunc := context.WithCancel(context.Background())
		defer closeFunc()
		localExecutor := framework.LocalExecutor(ctx)

		By(fmt.Sprintf("Running command: kubectl logs -n %s %s", controller.Namespace(), vm.Name))
		err := localExecutor.Run(nil, &stdout, &stdout, "kubectl", "-n", controller.Namespace(), "logs", vm.Name)
		Expect(err).NotTo(HaveOccurred())
		Expect(stdout.String()).Should(ContainSubstring("login as 'cirros' user."))

		By(fmt.Sprintf("Running command: kubectl attach -n %s -i %s", controller.Namespace(), vm.Name))
		stdin := bytes.NewBufferString("\nTESTTEXT\n\n")
		stdout.Reset()
		err = localExecutor.Run(stdin, &stdout, &stdout, "kubectl", "-n", controller.Namespace(), "attach", "-i", vm.Name)
		Expect(err).NotTo(HaveOccurred())

		By(fmt.Sprintf("Running again command: kubectl logs -n %s %s", controller.Namespace(), vm.Name))
		stdout.Reset()
		err = localExecutor.Run(nil, &stdout, &stdout, "kubectl", "-n", controller.Namespace(), "logs", vm.Name)
		Expect(err).NotTo(HaveOccurred())
		Expect(stdout.String()).Should(ContainSubstring("TESTTEXT"))
	}, 3*60)
})
<commit_msg>Remove [Disruptive] label from Virtlet restart e2e<commit_after>/*
Copyright 2018 Mirantis

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 e2e

import (
	"bytes"
	"context"
	"fmt"
	"time"

	. "github.com/onsi/gomega"

	"github.com/Mirantis/virtlet/tests/e2e/framework"
	. "github.com/Mirantis/virtlet/tests/e2e/ginkgo-ext"
)

var _ = Describe("Virtlet restart", func() {
	var (
		vm *framework.VMInterface
	)

	BeforeAll(func() {
		vm = controller.VM("cirros-vm")
		vm.CreateAndWait(VMOptions{}.ApplyDefaults(), time.Minute*5, nil)
		var err error
		_, err = vm.Pod()
		Expect(err).NotTo(HaveOccurred())

		// restart virtlet before all tests
		virtletPod, err := vm.VirtletPod()
		Expect(err).NotTo(HaveOccurred())

		err = virtletPod.Delete()
		Expect(err).NotTo(HaveOccurred())

		waitVirtletPod(vm)
	})

	AfterAll(func() {
		deleteVM(vm)
	})

	It("Should allow to ssh to VM after virtlet pod restart", func() {
		waitSSH(vm)
	}, 3*60)

	It("Should keep logs from another session", func() {
		var stdout bytes.Buffer
		ctx, closeFunc := context.WithCancel(context.Background())
		defer closeFunc()
		localExecutor := framework.LocalExecutor(ctx)

		By(fmt.Sprintf("Running command: kubectl logs -n %s %s", controller.Namespace(), vm.Name))
		err := localExecutor.Run(nil, &stdout, &stdout, "kubectl", "-n", controller.Namespace(), "logs", vm.Name)
		Expect(err).NotTo(HaveOccurred())
		Expect(stdout.String()).Should(ContainSubstring("login as 'cirros' user."))

		By(fmt.Sprintf("Running command: kubectl attach -n %s -i %s", controller.Namespace(), vm.Name))
		stdin := bytes.NewBufferString("\nTESTTEXT\n\n")
		stdout.Reset()
		err = localExecutor.Run(stdin, &stdout, &stdout, "kubectl", "-n", controller.Namespace(), "attach", "-i", vm.Name)
		Expect(err).NotTo(HaveOccurred())

		By(fmt.Sprintf("Running again command: kubectl logs -n %s %s", controller.Namespace(), vm.Name))
		stdout.Reset()
		err = localExecutor.Run(nil, &stdout, &stdout, "kubectl", "-n", controller.Namespace(), "logs", vm.Name)
		Expect(err).NotTo(HaveOccurred())
		Expect(stdout.String()).Should(ContainSubstring("TESTTEXT"))
	}, 3*60)
})
<|endoftext|>


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

package buildbucket

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"go.chromium.org/luci/buildbucket"
	bbapi "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"

	"go.chromium.org/luci/milo/api/resp"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
)

// fetchBuilds fetches builds given a criteria.
// The returned builds are sorted by build creation descending.
// count defines maximum number of builds to fetch; if <0, defaults to 100.
func fetchBuilds(c context.Context, client *bbapi.Service, bucket, builder,
	status string, limit int) ([]*bbapi.ApiCommonBuildMessage, error) {

	search := client.Search()
	search.Bucket(bucket)
	search.Status(status)
	search.Tag(strpair.Format(buildbucket.TagBuilder, builder))

	if limit < 0 {
		limit = 100
	}

	start := clock.Now(c)
	msgs, err := search.Fetch(limit, nil)
	if err != nil {
		return nil, err
	}
	logging.Infof(c, "Fetched %d %s builds in %s", len(msgs), status, clock.Since(c, start))
	return msgs, nil
}

// toMiloBuild converts a buildbucket build to a milo build.
// In case of an error, returns a build with a description of the error
// and logs the error.
func toMiloBuild(msg *bbapi.ApiCommonBuildMessage) (*resp.BuildSummary, error) {
	var b buildbucket.Build
	if err := b.ParseMessage(msg); err != nil {
		return nil, err
	}

	var params struct {
		Changes []struct {
			Author struct{ Email string }
		}
		Properties struct {
			Revision string `json:"revision"`
		}
	}
	if msg.ParametersJson != "" {
		if err := json.NewDecoder(strings.NewReader(msg.ParametersJson)).Decode(&params); err != nil {
			return nil, errors.Annotate(err, "failed to parse parameters_json of build %d", b.ID).Err()
		}
	}

	var resultDetails struct {
		Properties struct {
			GotRevision string `json:"got_revision"`
		}
		// TODO(nodir,iannucci): define a proto for build UI data
		UI struct {
			Info string
		}
	}
	if msg.ResultDetailsJson != "" {
		if err := json.NewDecoder(strings.NewReader(msg.ResultDetailsJson)).Decode(&resultDetails); err != nil {
			return nil, errors.Annotate(err, "failed to parse result_details_json of build %d", b.ID).Err()
		}
	}

	schedulingDuration, _ := b.SchedulingDuration()
	runDuration, _ := b.RunDuration()
	result := &resp.BuildSummary{
		Revision: resultDetails.Properties.GotRevision,
		Status:   parseStatus(b.Status),
		PendingTime: resp.Interval{
			Started:  b.CreationTime,
			Finished: b.StartTime,
			Duration: schedulingDuration,
		},
		ExecutionTime: resp.Interval{
			Started:  b.StartTime,
			Finished: b.CompletionTime,
			Duration: runDuration,
		},
	}
	if result.Revision == "" {
		result.Revision = params.Properties.Revision
	}
	if resultDetails.UI.Info != "" {
		result.Text = strings.Split(resultDetails.UI.Info, "\n")
	}

	for _, bs := range b.BuildSets {
		// ignore rietveld.
		cl, ok := bs.(*buildbucket.GerritChange)
		if !ok {
			continue
		}

		// support only one CL per build.
		result.Blame = []*resp.Commit{{
			Changelist: resp.NewLink(fmt.Sprintf("Gerrit CL %d", cl.Change), cl.URL(),
				fmt.Sprintf("gerrit changelist %d", cl.Change)),
			RequestRevision: resp.NewLink(params.Properties.Revision, "", fmt.Sprintf("request revision %s", params.Properties.Revision)),
		}}

		if len(params.Changes) == 1 {
			result.Blame[0].AuthorEmail = params.Changes[0].Author.Email
		}
		break
	}

	if b.Number != nil {
		numStr := strconv.Itoa(*b.Number)
		result.Link = resp.NewLink(
			numStr,
			fmt.Sprintf("/p/%s/builders/%s/%s/%s", b.Project, b.Bucket, b.Builder, numStr),
			fmt.Sprintf("build #%s", numStr))
	} else {
		idStr := strconv.FormatInt(b.ID, 10)
		result.Link = resp.NewLink(
			idStr,
			fmt.Sprintf("/p/%s/builds/b%s", b.Project, idStr),
			fmt.Sprintf("build #%s", idStr))
	}
	return result, nil
}

func getDebugBuilds(c context.Context, bucket, builder string, maxCompletedBuilds int, target *resp.Builder) error {
	// ../buildbucket below assumes that
	// - this code is not executed by tests outside of this dir
	// - this dir is a sibling of frontend dir
	resFile, err := os.Open(filepath.Join(
		"..", "buildbucket", "testdata", bucket, builder+".json"))
	if err != nil {
		return err
	}
	defer resFile.Close()

	res := &bbapi.ApiSearchResponseMessage{}
	if err := json.NewDecoder(resFile).Decode(res); err != nil {
		return err
	}

	for _, bb := range res.Builds {
		mb, err := toMiloBuild(bb)
		if err != nil {
			return err
		}
		switch mb.Status {
		case model.NotRun:
			target.PendingBuilds = append(target.PendingBuilds, mb)

		case model.Running:
			target.CurrentBuilds = append(target.CurrentBuilds, mb)

		case model.Success, model.Failure, model.InfraFailure, model.Warning:
			if len(target.FinishedBuilds) < maxCompletedBuilds {
				target.FinishedBuilds = append(target.FinishedBuilds, mb)
			}

		default:
			panic("impossible")
		}
	}
	return nil
}

func getHost(c context.Context) (string, error) {
	settings := common.GetSettings(c)
	if settings.Buildbucket == nil || settings.Buildbucket.Host == "" {
		return "", errors.New("missing buildbucket host in settings")
	}
	return settings.Buildbucket.Host, nil
}

// GetBuilder is used by buildsource.BuilderID.Get to obtain the resp.Builder.
func GetBuilder(c context.Context, bucket, builder string, limit int) (*resp.Builder, error) {
	host, err := getHost(c)
	if err != nil {
		return nil, err
	}

	if limit < 0 {
		limit = 20
	}

	result := &resp.Builder{
		Name: builder,
	}
	if host == "debug" {
		return result, getDebugBuilds(c, bucket, builder, limit, result)
	}
	client, err := newBuildbucketClient(c, host)
	if err != nil {
		return nil, err
	}

	fetch := func(statusFilter string, limit int) error {
		msgs, err := fetchBuilds(c, client, bucket, builder, statusFilter, limit)
		if err != nil {
			logging.Errorf(c, "Could not fetch %s builds: %s", statusFilter, err)
			return err
		}
		for _, m := range msgs {
			b, err := toMiloBuild(m)
			if err != nil {
				return errors.Annotate(err, "failed to convert build %d to milo build", m.Id).Err()
			}
			switch b.Status {
			case model.NotRun:
				result.PendingBuilds = append(result.PendingBuilds, b)
			case model.Running:
				result.PendingBuilds = append(result.CurrentBuilds, b)
			default:
				result.FinishedBuilds = append(result.FinishedBuilds, b)
			}
		}
		return nil
	}
	return result, parallel.FanOutIn(func(work chan<- func() error) {
		work <- func() error {
			return fetch(bbapi.StatusFilterIncomplete, -1)
		}
		work <- func() error {
			return fetch(bbapi.StatusCompleted, limit)
		}
	})
}
<commit_msg>[milo] fix current builds in buildbucket builder view<commit_after>// Copyright 2016 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package buildbucket

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"go.chromium.org/luci/buildbucket"
	bbapi "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/strpair"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"

	"go.chromium.org/luci/milo/api/resp"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
)

// fetchBuilds fetches builds given a criteria.
// The returned builds are sorted by build creation descending.
// count defines maximum number of builds to fetch; if <0, defaults to 100.
func fetchBuilds(c context.Context, client *bbapi.Service, bucket, builder,
	status string, limit int) ([]*bbapi.ApiCommonBuildMessage, error) {

	search := client.Search()
	search.Bucket(bucket)
	search.Status(status)
	search.Tag(strpair.Format(buildbucket.TagBuilder, builder))

	if limit < 0 {
		limit = 100
	}

	start := clock.Now(c)
	msgs, err := search.Fetch(limit, nil)
	if err != nil {
		return nil, err
	}
	logging.Infof(c, "Fetched %d %s builds in %s", len(msgs), status, clock.Since(c, start))
	return msgs, nil
}

// toMiloBuild converts a buildbucket build to a milo build.
// In case of an error, returns a build with a description of the error
// and logs the error.
func toMiloBuild(msg *bbapi.ApiCommonBuildMessage) (*resp.BuildSummary, error) {
	var b buildbucket.Build
	if err := b.ParseMessage(msg); err != nil {
		return nil, err
	}

	var params struct {
		Changes []struct {
			Author struct{ Email string }
		}
		Properties struct {
			Revision string `json:"revision"`
		}
	}
	if msg.ParametersJson != "" {
		if err := json.NewDecoder(strings.NewReader(msg.ParametersJson)).Decode(&params); err != nil {
			return nil, errors.Annotate(err, "failed to parse parameters_json of build %d", b.ID).Err()
		}
	}

	var resultDetails struct {
		Properties struct {
			GotRevision string `json:"got_revision"`
		}
		// TODO(nodir,iannucci): define a proto for build UI data
		UI struct {
			Info string
		}
	}
	if msg.ResultDetailsJson != "" {
		if err := json.NewDecoder(strings.NewReader(msg.ResultDetailsJson)).Decode(&resultDetails); err != nil {
			return nil, errors.Annotate(err, "failed to parse result_details_json of build %d", b.ID).Err()
		}
	}

	schedulingDuration, _ := b.SchedulingDuration()
	runDuration, _ := b.RunDuration()
	result := &resp.BuildSummary{
		Revision: resultDetails.Properties.GotRevision,
		Status:   parseStatus(b.Status),
		PendingTime: resp.Interval{
			Started:  b.CreationTime,
			Finished: b.StartTime,
			Duration: schedulingDuration,
		},
		ExecutionTime: resp.Interval{
			Started:  b.StartTime,
			Finished: b.CompletionTime,
			Duration: runDuration,
		},
	}
	if result.Revision == "" {
		result.Revision = params.Properties.Revision
	}
	if resultDetails.UI.Info != "" {
		result.Text = strings.Split(resultDetails.UI.Info, "\n")
	}

	for _, bs := range b.BuildSets {
		// ignore rietveld.
		cl, ok := bs.(*buildbucket.GerritChange)
		if !ok {
			continue
		}

		// support only one CL per build.
		result.Blame = []*resp.Commit{{
			Changelist: resp.NewLink(fmt.Sprintf("Gerrit CL %d", cl.Change), cl.URL(),
				fmt.Sprintf("gerrit changelist %d", cl.Change)),
			RequestRevision: resp.NewLink(params.Properties.Revision, "", fmt.Sprintf("request revision %s", params.Properties.Revision)),
		}}

		if len(params.Changes) == 1 {
			result.Blame[0].AuthorEmail = params.Changes[0].Author.Email
		}
		break
	}

	if b.Number != nil {
		numStr := strconv.Itoa(*b.Number)
		result.Link = resp.NewLink(
			numStr,
			fmt.Sprintf("/p/%s/builders/%s/%s/%s", b.Project, b.Bucket, b.Builder, numStr),
			fmt.Sprintf("build #%s", numStr))
	} else {
		idStr := strconv.FormatInt(b.ID, 10)
		result.Link = resp.NewLink(
			idStr,
			fmt.Sprintf("/p/%s/builds/b%s", b.Project, idStr),
			fmt.Sprintf("build #%s", idStr))
	}
	return result, nil
}

func getDebugBuilds(c context.Context, bucket, builder string, maxCompletedBuilds int, target *resp.Builder) error {
	// ../buildbucket below assumes that
	// - this code is not executed by tests outside of this dir
	// - this dir is a sibling of frontend dir
	resFile, err := os.Open(filepath.Join(
		"..", "buildbucket", "testdata", bucket, builder+".json"))
	if err != nil {
		return err
	}
	defer resFile.Close()

	res := &bbapi.ApiSearchResponseMessage{}
	if err := json.NewDecoder(resFile).Decode(res); err != nil {
		return err
	}

	for _, bb := range res.Builds {
		mb, err := toMiloBuild(bb)
		if err != nil {
			return err
		}
		switch mb.Status {
		case model.NotRun:
			target.PendingBuilds = append(target.PendingBuilds, mb)

		case model.Running:
			target.CurrentBuilds = append(target.CurrentBuilds, mb)

		case model.Success, model.Failure, model.InfraFailure, model.Warning:
			if len(target.FinishedBuilds) < maxCompletedBuilds {
				target.FinishedBuilds = append(target.FinishedBuilds, mb)
			}

		default:
			panic("impossible")
		}
	}
	return nil
}

func getHost(c context.Context) (string, error) {
	settings := common.GetSettings(c)
	if settings.Buildbucket == nil || settings.Buildbucket.Host == "" {
		return "", errors.New("missing buildbucket host in settings")
	}
	return settings.Buildbucket.Host, nil
}

// GetBuilder is used by buildsource.BuilderID.Get to obtain the resp.Builder.
func GetBuilder(c context.Context, bucket, builder string, limit int) (*resp.Builder, error) {
	host, err := getHost(c)
	if err != nil {
		return nil, err
	}

	if limit < 0 {
		limit = 20
	}

	result := &resp.Builder{
		Name: builder,
	}
	if host == "debug" {
		return result, getDebugBuilds(c, bucket, builder, limit, result)
	}
	client, err := newBuildbucketClient(c, host)
	if err != nil {
		return nil, err
	}

	fetch := func(statusFilter string, limit int) error {
		msgs, err := fetchBuilds(c, client, bucket, builder, statusFilter, limit)
		if err != nil {
			logging.Errorf(c, "Could not fetch %s builds: %s", statusFilter, err)
			return err
		}
		for _, m := range msgs {
			b, err := toMiloBuild(m)
			if err != nil {
				return errors.Annotate(err, "failed to convert build %d to milo build", m.Id).Err()
			}
			switch b.Status {
			case model.NotRun:
				result.PendingBuilds = append(result.PendingBuilds, b)
			case model.Running:
				result.CurrentBuilds = append(result.CurrentBuilds, b)
			default:
				result.FinishedBuilds = append(result.FinishedBuilds, b)
			}
		}
		return nil
	}
	return result, parallel.FanOutIn(func(work chan<- func() error) {
		work <- func() error {
			return fetch(bbapi.StatusFilterIncomplete, -1)
		}
		work <- func() error {
			return fetch(bbapi.StatusCompleted, limit)
		}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"

	"github.com/Sirupsen/logrus"
	_ "github.com/eliothedeman/bangarang/alarm/console"
	_ "github.com/eliothedeman/bangarang/alarm/email"
	_ "github.com/eliothedeman/bangarang/alarm/grafana-graphite-annotation"
	_ "github.com/eliothedeman/bangarang/alarm/pd"
	"github.com/eliothedeman/bangarang/api"
	"github.com/eliothedeman/bangarang/config"
	"github.com/eliothedeman/bangarang/pipeline"
	_ "github.com/eliothedeman/bangarang/provider/http"
	_ "github.com/eliothedeman/bangarang/provider/tcp"
)

var (
	confFile = flag.String("conf", "/etc/bangarang/conf.json", "path main config file")
	dev      = flag.Bool("dev", false, "puts bangarang in a dev testing mode")
	version  = flag.Bool("version", false, "display the version of this binary")
	confType = flag.String("conf-type", "db", `type of configuration used ["db", "json"]`)
	apiPort  = flag.Int("api-port", 8081, "port to serve the http api on")
)

const (
	versionNumber = "0.10.3"
)

func init() {
	logrus.SetLevel(logrus.InfoLevel)
	tf := &logrus.TextFormatter{}
	tf.FullTimestamp = true
	tf.ForceColors = true
	logrus.SetFormatter(tf)
}

func handleSigs() {
	stop := make(chan os.Signal)
	signal.Notify(stop, os.Kill, os.Interrupt)

	done := <-stop
	logrus.Fatal(done.String())
}

func main() {
	flag.Parse()

	// display the current version and exit
	if *version {
		fmt.Print(versionNumber)
		os.Exit(0)
	}

	// load configuration
	cp := config.GetProvider(*confType, *confFile)
	if cp == nil {
		logrus.Fatalf("Unable to load config of type %s at location %s", *confType, *confFile)
	}
	ac, err := cp.GetCurrent()
	if err != nil {
		logrus.Fatal(err)
	}

	if ac.LogLevel == "" {
		ac.LogLevel = "info"
	}

	ll, err := logrus.ParseLevel(ac.LogLevel)
	if err != nil {
		logrus.Error(err)
	} else {
		logrus.SetLevel(ll)
	}

	logrus.Infof("Starting processing pipeline with %d policie(s)", len(ac.Policies))
	// create and start up a new pipeline
	p := pipeline.NewPipeline(ac)
	p.Start()

	logrus.Infof("Serving the http api on port %d", 8081)
	// create and start a new api server
	apiServer := api.NewServer(*apiPort, p)
	go apiServer.Serve()

	handleSigs()
}
<commit_msg>version bump<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"

	"github.com/Sirupsen/logrus"
	_ "github.com/eliothedeman/bangarang/alarm/console"
	_ "github.com/eliothedeman/bangarang/alarm/email"
	_ "github.com/eliothedeman/bangarang/alarm/grafana-graphite-annotation"
	_ "github.com/eliothedeman/bangarang/alarm/pd"
	"github.com/eliothedeman/bangarang/api"
	"github.com/eliothedeman/bangarang/config"
	"github.com/eliothedeman/bangarang/pipeline"
	_ "github.com/eliothedeman/bangarang/provider/http"
	_ "github.com/eliothedeman/bangarang/provider/tcp"
)

var (
	confFile = flag.String("conf", "/etc/bangarang/conf.json", "path main config file")
	dev      = flag.Bool("dev", false, "puts bangarang in a dev testing mode")
	version  = flag.Bool("version", false, "display the version of this binary")
	confType = flag.String("conf-type", "db", `type of configuration used ["db", "json"]`)
	apiPort  = flag.Int("api-port", 8081, "port to serve the http api on")
)

const (
	versionNumber = "0.10.4"
)

func init() {
	logrus.SetLevel(logrus.InfoLevel)
	tf := &logrus.TextFormatter{}
	tf.FullTimestamp = true
	tf.ForceColors = true
	logrus.SetFormatter(tf)
}

func handleSigs() {
	stop := make(chan os.Signal)
	signal.Notify(stop, os.Kill, os.Interrupt)

	done := <-stop
	logrus.Fatal(done.String())
}

func main() {
	flag.Parse()

	// display the current version and exit
	if *version {
		fmt.Print(versionNumber)
		os.Exit(0)
	}

	// load configuration
	cp := config.GetProvider(*confType, *confFile)
	if cp == nil {
		logrus.Fatalf("Unable to load config of type %s at location %s", *confType, *confFile)
	}
	ac, err := cp.GetCurrent()
	if err != nil {
		logrus.Fatal(err)
	}

	if ac.LogLevel == "" {
		ac.LogLevel = "info"
	}

	ll, err := logrus.ParseLevel(ac.LogLevel)
	if err != nil {
		logrus.Error(err)
	} else {
		logrus.SetLevel(ll)
	}

	logrus.Infof("Starting processing pipeline with %d policie(s)", len(ac.Policies))
	// create and start up a new pipeline
	p := pipeline.NewPipeline(ac)
	p.Start()

	logrus.Infof("Serving the http api on port %d", 8081)
	// create and start a new api server
	apiServer := api.NewServer(*apiPort, p)
	go apiServer.Serve()

	handleSigs()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"strings"

	"github.com/otoolep/syslog-gollector/input"
	"github.com/otoolep/syslog-gollector/output"

	log "code.google.com/p/log4go"
)

// Program parameters
var adminIface string
var tcpIface string
var udpIface string
var kBrokers string
var kBatch int
var kTopic string
var kBufferTime int
var kBufferBytes int
var pEnabled bool
var cCapacity int

// Program servers
var tcpServer *input.TcpServer
var udpServer *input.UdpServer

// Types
const (
	adminHost        = "localhost:8080"
	connTcpHost      = "localhost:514"
	connUdpHost      = "localhost:514"
	connType         = "tcp"
	kafkaBatch       = 10
	kafkaBrokers     = "localhost:9092"
	kafkaTopic       = "logs"
	kafkaBufferTime  = 1000
	kafkaBufferBytes = 512 * 1024
	parseEnabled     = true
	chanCapacity     = 0
)

func init() {
	flag.StringVar(&adminIface, "admin", adminHost, "Admin interface")
	flag.StringVar(&tcpIface, "tcp", connTcpHost, "TCP bind interface")
	flag.StringVar(&udpIface, "udp", connUdpHost, "UDP interface")
	flag.StringVar(&kBrokers, "broker", kafkaBrokers, "comma-delimited kafka brokers")
	flag.StringVar(&kTopic, "topic", kafkaTopic, "kafka topic")
	flag.IntVar(&kBatch, "batch", kafkaBatch, "Kafka batch size")
	flag.IntVar(&kBufferTime, "maxbuff", kafkaBufferTime, "Kafka client buffer max time (ms)")
	flag.IntVar(&kBufferBytes, "maxbytes", kafkaBufferBytes, "Kafka client buffer max bytes")
	flag.BoolVar(&pEnabled, "parse", parseEnabled, "enable syslog header parsing")
	flag.IntVar(&cCapacity, "chancap", chanCapacity, "channel buffering capacity")
}

// ServeStatistics returns the statistics for the program
func ServeStatistics(w http.ResponseWriter, req *http.Request) {
	statistics := make(map[string]interface{})
	s, err := tcpServer.GetStatistics()
	if err != nil {
		log.Error("failed to get TCP stats")
		http.Error(w, "failed to get TCP stats", http.StatusInternalServerError)
		return
	}
	statistics["tcp"] = s

	s, err = udpServer.GetStatistics()
	if err != nil {
		log.Error("failed to get UDP stats")
		http.Error(w, "failed to get UDP stats", http.StatusInternalServerError)
		return
	}
	statistics["udp"] = s

	b, err := json.MarshalIndent(statistics, "", "    ")
	if err != nil {
		log.Error("failed to JSON marshal statistics map")
		http.Error(w, "failed to JSON marshal statistics map", http.StatusInternalServerError)
		return
	}
	w.Write(b)
}

func main() {
	flag.Parse()

	hostname, err := os.Hostname()
	if err != nil {
		log.Error("unable to determine hostname -- aborting")
		os.Exit(1)
	}
	log.Info("syslog server starting on %s, PID %d", hostname, os.Getpid())
	log.Info("machine has %d cores", runtime.NumCPU())

	// Log config
	log.Info("Admin server: %s", adminIface)
	log.Info("kafka brokers: %s", kBrokers)
	log.Info("kafka topic: %s", kTopic)
	log.Info("kafka batch size: %d", kBatch)
	log.Info("kafka buffer time: %dms", kBufferTime)
	log.Info("kafka buffer bytes: %d", kBufferBytes)
	log.Info("parsing enabled: %t", pEnabled)
	log.Info("channel buffering capacity: %d", cCapacity)

	// Prep the channels
	rawChan := make(chan string, cCapacity)
	prodChan := make(chan string, cCapacity)

	if pEnabled {
		// Feed the input through the Parser stage
		parser := input.NewRfc5424Parser()
		prodChan, err = parser.StreamingParse(rawChan)
	} else {
		// Pass the input directly to the output
		prodChan = rawChan
	}

	// Configure and start the Admin server
	http.HandleFunc("/statistics", ServeStatistics)
	go func() {
		err = http.ListenAndServe(adminIface, nil)
		if err != nil {
			fmt.Println("Failed to start admin server", err.Error())
			os.Exit(1)
		}
	}()
	log.Info("Admin server started")

	// Connect to Kafka
	_, err = output.NewKafkaProducer(prodChan, strings.Split(kBrokers, ","), kTopic, kBufferTime, kBufferBytes)
	if err != nil {
		fmt.Println("Failed to create Kafka producer", err.Error())
		os.Exit(1)
	}
	log.Info("connected to kafka at %s", kBrokers)

	// Start the event servers
	tcpServer = input.NewTcpServer(tcpIface)
	err = tcpServer.Start(func() chan<- string {
		return rawChan
	})
	if err != nil {
		fmt.Println("Failed to start TCP server", err.Error())
		os.Exit(1)
	}
	log.Info("listening on %s for TCP connections", tcpIface)

	udpServer = input.NewUdpServer(udpIface)
	err = udpServer.Start(func() chan<- string {
		return rawChan
	})
	if err != nil {
		fmt.Println("Failed to start UDP server", err.Error())
		os.Exit(1)
	}
	log.Info("listening on %s for UDP packets", udpIface)

	// Spin forever
	select {}
}
<commit_msg>Support pretty-printed statistics<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"strings"

	"github.com/otoolep/syslog-gollector/input"
	"github.com/otoolep/syslog-gollector/output"

	log "code.google.com/p/log4go"
)

// Program parameters
var adminIface string
var tcpIface string
var udpIface string
var kBrokers string
var kBatch int
var kTopic string
var kBufferTime int
var kBufferBytes int
var pEnabled bool
var cCapacity int

// Program servers
var tcpServer *input.TcpServer
var udpServer *input.UdpServer

// Types
const (
	adminHost        = "localhost:8080"
	connTcpHost      = "localhost:514"
	connUdpHost      = "localhost:514"
	connType         = "tcp"
	kafkaBatch       = 10
	kafkaBrokers     = "localhost:9092"
	kafkaTopic       = "logs"
	kafkaBufferTime  = 1000
	kafkaBufferBytes = 512 * 1024
	parseEnabled     = true
	chanCapacity     = 0
)

func init() {
	flag.StringVar(&adminIface, "admin", adminHost, "Admin interface")
	flag.StringVar(&tcpIface, "tcp", connTcpHost, "TCP bind interface")
	flag.StringVar(&udpIface, "udp", connUdpHost, "UDP interface")
	flag.StringVar(&kBrokers, "broker", kafkaBrokers, "comma-delimited kafka brokers")
	flag.StringVar(&kTopic, "topic", kafkaTopic, "kafka topic")
	flag.IntVar(&kBatch, "batch", kafkaBatch, "Kafka batch size")
	flag.IntVar(&kBufferTime, "maxbuff", kafkaBufferTime, "Kafka client buffer max time (ms)")
	flag.IntVar(&kBufferBytes, "maxbytes", kafkaBufferBytes, "Kafka client buffer max bytes")
	flag.BoolVar(&pEnabled, "parse", parseEnabled, "enable syslog header parsing")
	flag.IntVar(&cCapacity, "chancap", chanCapacity, "channel buffering capacity")
}

// isPretty returns whether the HTTP response body should be pretty-printed.
func isPretty(req *http.Request) (bool, error) {
	err := req.ParseForm()
	if err != nil {
		return false, err
	}
	if _, ok := req.Form["pretty"]; ok {
		return true, nil
	}
	return false, nil
}

// ServeStatistics returns the statistics for the program
func ServeStatistics(w http.ResponseWriter, req *http.Request) {
	statistics := make(map[string]interface{})
	s, err := tcpServer.GetStatistics()
	if err != nil {
		log.Error("failed to get TCP stats")
		http.Error(w, "failed to get TCP stats", http.StatusInternalServerError)
		return
	}
	statistics["tcp"] = s

	s, err = udpServer.GetStatistics()
	if err != nil {
		log.Error("failed to get UDP stats")
		http.Error(w, "failed to get UDP stats", http.StatusInternalServerError)
		return
	}
	statistics["udp"] = s

	var b []byte
	pretty, _ := isPretty(req)
	if pretty {
		b, err = json.MarshalIndent(statistics, "", "    ")
	} else {
		b, err = json.Marshal(statistics)
	}
	if err != nil {
		log.Error("failed to JSON marshal statistics map")
		http.Error(w, "failed to JSON marshal statistics map", http.StatusInternalServerError)
		return
	}
	w.Write(b)
}

func main() {
	flag.Parse()

	hostname, err := os.Hostname()
	if err != nil {
		log.Error("unable to determine hostname -- aborting")
		os.Exit(1)
	}
	log.Info("syslog server starting on %s, PID %d", hostname, os.Getpid())
	log.Info("machine has %d cores", runtime.NumCPU())

	// Log config
	log.Info("Admin server: %s", adminIface)
	log.Info("kafka brokers: %s", kBrokers)
	log.Info("kafka topic: %s", kTopic)
	log.Info("kafka batch size: %d", kBatch)
	log.Info("kafka buffer time: %dms", kBufferTime)
	log.Info("kafka buffer bytes: %d", kBufferBytes)
	log.Info("parsing enabled: %t", pEnabled)
	log.Info("channel buffering capacity: %d", cCapacity)

	// Prep the channels
	rawChan := make(chan string, cCapacity)
	prodChan := make(chan string, cCapacity)

	if pEnabled {
		// Feed the input through the Parser stage
		parser := input.NewRfc5424Parser()
		prodChan, err = parser.StreamingParse(rawChan)
	} else {
		// Pass the input directly to the output
		prodChan = rawChan
	}

	// Configure and start the Admin server
	http.HandleFunc("/statistics", ServeStatistics)
	go func() {
		err = http.ListenAndServe(adminIface, nil)
		if err != nil {
			fmt.Println("Failed to start admin server", err.Error())
			os.Exit(1)
		}
	}()
	log.Info("Admin server started")

	// Connect to Kafka
	_, err = output.NewKafkaProducer(prodChan, strings.Split(kBrokers, ","), kTopic, kBufferTime, kBufferBytes)
	if err != nil {
		fmt.Println("Failed to create Kafka producer", err.Error())
		os.Exit(1)
	}
	log.Info("connected to kafka at %s", kBrokers)

	// Start the event servers
	tcpServer = input.NewTcpServer(tcpIface)
	err = tcpServer.Start(func() chan<- string {
		return rawChan
	})
	if err != nil {
		fmt.Println("Failed to start TCP server", err.Error())
		os.Exit(1)
	}
	log.Info("listening on %s for TCP connections", tcpIface)

	udpServer = input.NewUdpServer(udpIface)
	err = udpServer.Start(func() chan<- string {
		return rawChan
	})
	if err != nil {
		fmt.Println("Failed to start UDP server", err.Error())
		os.Exit(1)
	}
	log.Info("listening on %s for UDP packets", udpIface)

	// Spin forever
	select {}
}
<|endoftext|>