-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
6e02254
commit dc1be16
Showing
6 changed files
with
255 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,187 @@ | ||
| // TODO | ||
| // context | ||
| // wait for confirms before closing the connection | ||
| // handle errors | ||
| package amqp091 | ||
|
|
||
| import ( | ||
| "context" | ||
| "math/rand" | ||
| "sync" | ||
| "sync/atomic" | ||
| "time" | ||
|
|
||
| amqp091 "github.com/rabbitmq/amqp091-go" | ||
| "github.com/rabbitmq/omq/pkg/config" | ||
| "github.com/rabbitmq/omq/pkg/log" | ||
| "github.com/rabbitmq/omq/pkg/metrics" | ||
| "github.com/rabbitmq/omq/pkg/utils" | ||
| ) | ||
|
|
||
| type Amqp091Publisher struct { | ||
| Id int | ||
| Connection *amqp091.Connection | ||
| Channel *amqp091.Channel | ||
| confirms chan amqp091.Confirmation | ||
| publishTimes sync.Map | ||
| Terminus string | ||
| Config config.Config | ||
| msg []byte | ||
| whichUri int | ||
| ctx context.Context | ||
| } | ||
|
|
||
| func NewPublisher(ctx context.Context, cfg config.Config, id int) *Amqp091Publisher { | ||
| publisher := &Amqp091Publisher{ | ||
| Id: id, | ||
| Connection: nil, | ||
| Config: cfg, | ||
| Terminus: utils.InjectId(cfg.PublishTo, id), | ||
| whichUri: 0, | ||
| ctx: ctx, | ||
| } | ||
|
|
||
| if cfg.SpreadConnections { | ||
| publisher.whichUri = (id - 1) % len(cfg.PublisherUri) | ||
| } | ||
|
|
||
| publisher.Connect() | ||
|
|
||
| return publisher | ||
| } | ||
|
|
||
| func (p *Amqp091Publisher) Connect() { | ||
| var conn *amqp091.Connection | ||
| var err error | ||
|
|
||
| if p.Connection != nil { | ||
| _ = p.Connection.Close() | ||
| } | ||
| p.Connection = nil | ||
|
|
||
| for p.Connection == nil { | ||
| if p.whichUri >= len(p.Config.PublisherUri) { | ||
| p.whichUri = 0 | ||
| } | ||
| uri := p.Config.PublisherUri[p.whichUri] | ||
| p.whichUri++ | ||
| // hostname, vhost := hostAndVHost(uri) | ||
| conn, err = amqp091.Dial(uri) | ||
| if err != nil { | ||
| log.Error("connection failed", "id", p.Id, "error", err.Error()) | ||
| select { | ||
| case <-time.After(1 * time.Second): | ||
| continue | ||
| case <-p.ctx.Done(): | ||
| return | ||
| } | ||
| } else { | ||
| log.Debug("publisher connected", "id", p.Id, "uri", uri) | ||
| p.Connection = conn | ||
| } | ||
| } | ||
| p.Channel, err = p.Connection.Channel() // TODO | ||
| if err != nil { | ||
| log.Error("channel creation failed", "id", p.Id, "error", err.Error()) | ||
| } | ||
| p.confirms = make(chan amqp091.Confirmation) // p.Config.MaxInFlight ? | ||
| _ = p.Channel.Confirm(false) | ||
| p.Channel.NotifyPublish(p.confirms) | ||
| } | ||
|
|
||
| func (p *Amqp091Publisher) Start(ctx context.Context, publisherReady chan bool, startPublishing chan bool) { | ||
| p.msg = utils.MessageBody(p.Config.Size) | ||
|
|
||
| close(publisherReady) | ||
|
|
||
| select { | ||
| case <-ctx.Done(): | ||
| return | ||
| case <-startPublishing: | ||
| // short random delay to avoid all publishers publishing at the same time | ||
| time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond) | ||
| } | ||
|
|
||
| log.Info("publisher started", "id", p.Id, "rate", utils.Rate(p.Config.Rate), "destination", p.Terminus) | ||
| var farewell string | ||
| if p.Config.Rate == 0 { | ||
| // idle connection | ||
| <-ctx.Done() | ||
| farewell = "context cancelled" | ||
| } else { | ||
| farewell = p.StartPublishing(ctx) | ||
| } | ||
| p.Stop(farewell) | ||
| } | ||
|
|
||
| func (p *Amqp091Publisher) StartPublishing(ctx context.Context) string { | ||
| limiter := utils.RateLimiter(p.Config.Rate) | ||
|
|
||
| go p.handleConfirms() | ||
|
|
||
| var msgSent atomic.Int64 | ||
| for { | ||
| select { | ||
| case <-ctx.Done(): | ||
| return "context cancelled" | ||
| default: | ||
| n := uint64(msgSent.Add(1)) | ||
| if n > uint64(p.Config.PublishCount) { | ||
| return "--pmessages value reached" | ||
| } | ||
| if p.Config.Rate > 0 { | ||
| _ = limiter.Wait(ctx) | ||
| } | ||
| err := p.SendAsync(ctx, n) | ||
| if err != nil { | ||
| p.Connect() | ||
| } else { | ||
| log.Debug("message sent", "id", p.Id, "deliveryTag", n) | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (p *Amqp091Publisher) SendAsync(ctx context.Context, n uint64) error { | ||
| msg := p.prepareMessage() | ||
|
|
||
| p.publishTimes.Store(n, time.Now()) | ||
| err := p.Channel.Publish("", p.Config.PublishTo, false, false, msg) | ||
| return err | ||
| } | ||
|
|
||
| func (p *Amqp091Publisher) handleConfirms() { | ||
| for confirm := range p.confirms { | ||
| if confirm.Ack { | ||
| pubTime, _ := p.publishTimes.LoadAndDelete(confirm.DeliveryTag) | ||
| latency := time.Since(pubTime.(time.Time)) | ||
| metrics.PublishingLatency.Update(latency.Seconds()) | ||
| metrics.MessagesPublished.Inc() | ||
| log.Debug("message confirmed", "id", p.Id, "delivery_tag", confirm.DeliveryTag, "latency", latency) | ||
| } else { | ||
| p.publishTimes.Delete(confirm.DeliveryTag) | ||
| log.Debug("message not confirmed by the broker", "id", p.Id) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (p *Amqp091Publisher) Stop(reason string) { | ||
| // for p.publishTimes. > 0 { | ||
| // time.Sleep(100 * time.Millisecond) | ||
| // } | ||
| log.Debug("closing publisher connection", "id", p.Id, "reason", reason) | ||
| if p.Channel != nil { | ||
| _ = p.Channel.Close() | ||
| } | ||
| if p.Connection != nil { | ||
| _ = p.Connection.Close() | ||
| } | ||
| } | ||
|
|
||
| func (p *Amqp091Publisher) prepareMessage() amqp091.Publishing { | ||
| utils.UpdatePayload(p.Config.UseMillis, &p.msg) | ||
| return amqp091.Publishing{ | ||
| DeliveryMode: amqp091.Persistent, | ||
| Body: p.msg, | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -10,6 +10,7 @@ type Protocol int | |
|
|
||
| const ( | ||
| AMQP Protocol = iota | ||
| AMQP091 | ||
| STOMP | ||
| MQTT | ||
| MQTT5 | ||
|
|
||