func main() {
flag.StringVar(&myId, "id", "", "cluster id, typically a consuming service name")
flag.StringVar(&topic, "topic", "", "topic to consume")
flag.Parse()
if *cpuprofile != "" {
f, err := os.Create(*cpuprofile)
if err != nil {
log.Printf("Unable to create cpuprofile file: %v, moving on\n", err)
} else {
pprof.StartCPUProfile(f)
defer dumpCPUProfile()
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt)
go func() {
<-c
dumpCPUProfile()
os.Exit(0)
}()
}
}
log.V(2).Printf("Identified brokers: %v\n", kafkas)
server, err := comm.Listen(":4987")
if err != nil {
log.Printf("Unable to listen on port :4987: %v", err)
os.Exit(1)
}
_, consumers := newConsumer()
for _, consumer := range consumers {
go handleConsumer(consumer)
}
go func() {
for message := range blackHole.inbox {
log.V(2).Printf("Got message in blackhole: %v\n", message)
go func(message messages.Message) {
log.V(2).Printf("Gonna re-send message from blackhole: %v\n", message)
nextRoundRobinClient().inbox <- message
}(message)
}
}()
log.Println("Listening on port :4987")
for {
conn, err := server.Accept()
if err != nil {
log.Printf("Unable to accept connection: %v", err)
continue
}
go handleClient(conn)
}
}
func (c *Consumer) listen() error {
go func() {
var err error
c.stream, err = comm.Dial(c.fanouts[0])
if err != nil {
log.Printf("Unable to obtain stream: %v\n", err)
return
}
defer c.stream.Close()
err = c.stream.WriteLine(c.instanceId)
if err != nil {
log.Printf("Unable to send own instance id: %v\n", err)
return
}
for {
select {
case _, ok := <-c.closeChan:
if ok {
return
}
default:
}
var message messages.Message
err := c.stream.ReadMessage(&message)
if err != nil {
log.Printf("Unable to unmarshal message: %v\n", err)
continue
}
go func(message messages.Message) {
if c.SendAcks {
ack := &messages.MessageAck{
Partition: message.Partition,
Offset: message.Offset,
}
err := c.stream.WriteMessage(ack)
if err != nil {
log.Printf("Unable to send message ack: %v\n", err)
}
log.V(2).Printf("Sent ack from %s: %v\n", c.instanceId, *ack)
}
}(message)
if !message.IsAckRequest {
c.messages <- message
}
}
}()
return nil
}
func Ping(fanouts []string) error {
var response string
stream, err := comm.Dial(fanouts[0])
if err != nil {
log.Printf("Unable to create stream: %v\n", err)
return err
}
defer stream.Close()
err = stream.WriteLine("-PING")
if err != nil {
return err
}
response, err = stream.ReadLine()
if err != nil {
return err
}
if response != "+PONG\n" {
return NoPongFromServer
}
return nil
}
func sendMessageAndExpectAnAck(stream *comm.Stream, dead chan bool, message *messages.Message) {
err := stream.WriteMessage(message)
if err != nil {
log.Printf("Unable to send message to client: %v\n", err)
dead <- true
return
}
expectedAck := messages.MessageAck{
Partition: message.Partition,
Offset: message.Offset,
}
acks[expectedAck] = make(chan bool)
go func() {
select {
case _ = <-acks[expectedAck]:
break
case <-time.After(50 * time.Millisecond):
if !message.IsAckRequest {
blackHole.inbox <- *message
}
dead <- true
}
delete(acks, expectedAck)
}()
}
func assertReceived(t *testing.T, inbox chan *userneed.UserNeed, expected userNeedSet) *userneed.UserNeed {
select {
case actual := <-inbox:
log.Printf("Got something: %v :)\n", *actual)
if present, found := expected[*actual]; !found || !present {
t.Errorf("Expected %v to be in %v", *actual, expected)
return nil
}
expected[*actual] = false
return actual
case <-time.After(2000 * time.Millisecond):
log.Printf("Timed out :(\n")
t.Errorf("Expected to receive one of %v, got nothing", expected)
}
return nil
}
func retryCustom(times int, interval time.Duration, fn func() error) (err error) {
for i := 0; i < times; i++ {
err = fn()
if err == nil {
return
}
log.Printf("Got error: %v, retrying..\n", err)
time.Sleep(interval)
}
log.Println("Gave up :(")
return
}
func subscriptionInbox(t *testing.T, instanceId string) (inbox chan *userneed.UserNeed, consumer *fanout.Consumer) {
inbox = make(chan *userneed.UserNeed)
consumer = newConsumer(instanceId)
consumer.Subscribe(func(raw messages.Message) {
message := &userneed.UserNeed{}
log.Printf("Got raw message: %v with id=%d:%d\n", raw.Value, raw.Partition, raw.Offset)
err := proto.Unmarshal(raw.Value, message)
if err != nil {
t.Errorf("Unmarshal error: %v", err)
}
inbox <- message
})
return
}
func TestOneConsumer(t *testing.T) {
log.Printf("===== TestOneConsumer =====")
inbox, consumer := subscriptionInbox(t, "INSTANCE_0")
producer := newProducer()
defer consumer.Close()
producer.Publish(messageA)
producer.Publish(messageB)
producer.Publish(messageC)
expected := userNeedSet{
*messageA: true,
*messageB: true,
*messageC: true,
}
assertReceived(t, inbox, expected)
assertReceived(t, inbox, expected)
assertReceived(t, inbox, expected)
}
func handleClient(stream *comm.Stream) {
defer stream.Close()
var instanceId string
instanceId, err := stream.ReadLine()
if err != nil {
log.Printf("Unable to identify client: %v\n", err)
return
}
if instanceId == "-PING\n" {
log.V(1).Printf("got ping: -PING; answering with: +PONG\n")
err = stream.WriteLine("+PONG")
if err != nil {
log.Printf("Unable to answer with +PONG: %v\n", err)
}
return
}
log.V(1).Printf("got client: %s\n", instanceId)
inbox := make(chan messages.Message, CHANNEL_BUFFER_SIZE)
client := clientInbox{
inbox: inbox,
count: 1,
}
addClient(instanceId, client)
ackErrors := make(chan error)
go func() {
for {
var ack messages.MessageAck
err = stream.ReadMessage(&ack)
if err != nil {
ackErrors <- err
continue
}
if _, found := acks[ack]; found {
log.V(2).Printf("Got ack: %v\n", ack)
acks[ack] <- true
}
}
}()
dead := make(chan bool)
for {
select {
case message := <-inbox:
go sendMessageAndExpectAnAck(stream, dead, &message)
case <-time.Tick(75 * time.Millisecond):
go sendMessageAndExpectAnAck(stream, dead, ackRequest)
case err := <-ackErrors:
log.V(3).Printf("Got ack error: %v\n", err)
case <-dead:
log.V(2).Printf("Client %v disconnected\n", instanceId)
clientDead(instanceId)
}
}
}