// Test to make sure that we can send and async receive messages on
// different subjects within a callback.
func TestAsyncSubscriberStarvation(t *testing.T) {
s := RunDefaultServer()
defer s.Shutdown()
nc := NewDefaultConnection(t)
defer nc.Close()
// Helper
nc.Subscribe("helper", func(m *nats.Msg) {
nc.Publish(m.Reply, []byte("Hello"))
})
ch := make(chan bool)
// Kickoff
nc.Subscribe("start", func(m *nats.Msg) {
// Helper Response
response := nats.NewInbox()
nc.Subscribe(response, func(_ *nats.Msg) {
ch <- true
})
nc.PublishRequest("helper", response, []byte("Help Me!"))
})
nc.Publish("start", []byte("Begin"))
nc.Flush()
if e := Wait(ch); e != nil {
t.Fatal("Was stalled inside of callback waiting on another callback")
}
}
func (n *natsServer) AcceptLoop() error {
sub, err := n.conn.QueueSubscribe(n.subject, queue, func(msg *nats.Msg) {
if msg.Reply != "" {
var (
heartbeat = nats.NewInbox()
listenTo = nats.NewInbox()
client, err = n.accept(listenTo, msg.Reply, heartbeat)
)
if err != nil {
log.Println("thrift_nats: error accepting client transport:", err)
return
}
if n.isHeartbeating() {
n.mu.Lock()
n.clients[heartbeat] = client
n.mu.Unlock()
}
connectMsg := heartbeat + " " + strconv.FormatInt(int64(n.heartbeatDeadline.Seconds())*1000, 10)
if err := n.conn.PublishRequest(msg.Reply, listenTo, []byte(connectMsg)); err != nil {
log.Println("thrift_nats: error publishing transport inbox:", err)
if n.isHeartbeating() {
n.remove(heartbeat)
}
} else if n.isHeartbeating() {
go n.acceptHeartbeat(heartbeat)
}
} else {
log.Printf("thrift_nats: discarding invalid connect message %+v", msg)
}
})
if err != nil {
return err
}
n.conn.Flush()
n.mu.Lock()
n.closed = false
n.mu.Unlock()
log.Println("thrift_nats: server running...")
<-n.quit
return sub.Unsubscribe()
}
func connect(conn *nats.Conn, subj string, timeout time.Duration) (*nats.Msg, string, error) {
inbox := nats.NewInbox()
s, err := conn.Subscribe(inbox, nil)
if err != nil {
return nil, "", err
}
s.AutoUnsubscribe(1)
err = conn.PublishRequest(subj, inbox, nil)
if err != nil {
return nil, "", err
}
msg, err := s.NextMsg(timeout)
if err != nil {
return nil, "", err
}
return msg, inbox, nil
}
func (n *ntport) Dial(addr string, dialOpts ...transport.DialOption) (transport.Client, error) {
dopts := transport.DialOptions{
Timeout: transport.DefaultDialTimeout,
}
for _, o := range dialOpts {
o(&dopts)
}
opts := nats.DefaultOptions
opts.Servers = n.addrs
opts.Secure = n.opts.Secure
opts.TLSConfig = n.opts.TLSConfig
opts.Timeout = dopts.Timeout
// secure might not be set
if n.opts.TLSConfig != nil {
opts.Secure = true
}
c, err := opts.Connect()
if err != nil {
return nil, err
}
id := nats.NewInbox()
sub, err := c.SubscribeSync(id)
if err != nil {
return nil, err
}
return &ntportClient{
conn: c,
addr: addr,
id: id,
sub: sub,
opts: n.opts,
}, nil
}
func (n *ntport) Listen(addr string, listenOpts ...transport.ListenOption) (transport.Listener, error) {
opts := nats.DefaultOptions
opts.Servers = n.addrs
opts.Secure = n.opts.Secure
opts.TLSConfig = n.opts.TLSConfig
// secure might not be set
if n.opts.TLSConfig != nil {
opts.Secure = true
}
c, err := opts.Connect()
if err != nil {
return nil, err
}
return &ntportListener{
addr: nats.NewInbox(),
conn: c,
exit: make(chan bool, 1),
so: make(map[string]*ntportSocket),
opts: n.opts,
}, nil
}
func TestAsyncPubSubWithReply(t *testing.T) {
// Run a NATS Streaming server
s := RunServer(clusterName)
defer s.Shutdown()
sc := NewDefaultConnection(t)
defer sc.Close()
ch := make(chan bool)
hw := []byte("Hello World")
inbox := nats.NewInbox()
sub, err := sc.Subscribe("foo", func(m *Msg) {
if m.Subject != "foo" {
t.Fatalf("Expected subject of 'foo', got '%s'\n", m.Subject)
}
if !bytes.Equal(m.Data, hw) {
t.Fatalf("Wrong payload, got %q\n", m.Data)
}
if m.Reply != inbox {
t.Fatalf("Expected reply subject of '%s', got '%s'\n", inbox, m.Reply)
}
ch <- true
})
if err != nil {
t.Fatalf("Unexpected error on Subscribe, got %v", err)
}
defer sub.Unsubscribe()
sc.PublishAsyncWithReply("foo", inbox, hw, nil)
if err := WaitTime(ch, 1*time.Second); err != nil {
t.Fatal("Did not receive our messages")
}
}
func BenchmarkInboxCreation(b *testing.B) {
for i := 0; i < b.N; i++ {
nats.NewInbox()
}
}
func TestInbox(t *testing.T) {
inbox := nats.NewInbox()
if matched, _ := regexp.Match(`_INBOX.\S`, []byte(inbox)); !matched {
t.Fatal("Bad INBOX format")
}
}
func (n *natsRegistry) query(s string, quorum int) ([]*registry.Service, error) {
conn, err := n.getConn()
if err != nil {
return nil, err
}
var action string
var service *registry.Service
if len(s) > 0 {
action = "get"
service = ®istry.Service{Name: s}
} else {
action = "list"
}
inbox := nats.NewInbox()
response := make(chan *registry.Service, 10)
sub, err := conn.Subscribe(inbox, func(m *nats.Msg) {
var service *registry.Service
if err := json.Unmarshal(m.Data, &service); err != nil {
return
}
select {
case response <- service:
case <-time.After(DefaultTimeout):
}
})
if err != nil {
return nil, err
}
defer sub.Unsubscribe()
b, err := json.Marshal(®istry.Result{Action: action, Service: service})
if err != nil {
return nil, err
}
if err := conn.PublishMsg(&nats.Msg{
Subject: QueryTopic,
Reply: inbox,
Data: b,
}); err != nil {
return nil, err
}
timeoutChan := time.After(n.opts.Timeout)
serviceMap := make(map[string]*registry.Service)
loop:
for {
select {
case service := <-response:
key := service.Name + "-" + service.Version
srv, ok := serviceMap[key]
if ok {
srv.Nodes = append(srv.Nodes, service.Nodes...)
serviceMap[key] = srv
} else {
serviceMap[key] = service
}
if quorum > 0 && len(serviceMap[key].Nodes) >= quorum {
break loop
}
case <-timeoutChan:
break loop
}
}
var services []*registry.Service
for _, service := range serviceMap {
services = append(services, service)
}
return services, nil
}
// subscribe will perform a subscription with the given options to the NATS Streaming cluster.
func (sc *conn) subscribe(subject, qgroup string, cb MsgHandler, options ...SubscriptionOption) (Subscription, error) {
sub := &subscription{subject: subject, qgroup: qgroup, inbox: nats.NewInbox(), cb: cb, sc: sc, opts: DefaultSubscriptionOptions}
for _, opt := range options {
if err := opt(&sub.opts); err != nil {
return nil, err
}
}
sc.Lock()
if sc.nc == nil {
sc.Unlock()
return nil, ErrConnectionClosed
}
// Register subscription.
sc.subMap[sub.inbox] = sub
nc := sc.nc
sc.Unlock()
// Hold lock throughout.
sub.Lock()
defer sub.Unlock()
// Listen for actual messages.
nsub, err := nc.Subscribe(sub.inbox, sc.processMsg)
if err != nil {
return nil, err
}
sub.inboxSub = nsub
// Create a subscription request
// FIXME(dlc) add others.
sr := &pb.SubscriptionRequest{
ClientID: sc.clientID,
Subject: subject,
QGroup: qgroup,
Inbox: sub.inbox,
MaxInFlight: int32(sub.opts.MaxInflight),
AckWaitInSecs: int32(sub.opts.AckWait / time.Second),
StartPosition: sub.opts.StartAt,
DurableName: sub.opts.DurableName,
}
// Conditionals
switch sr.StartPosition {
case pb.StartPosition_TimeDeltaStart:
sr.StartTimeDelta = time.Now().UnixNano() - sub.opts.StartTime.UnixNano()
case pb.StartPosition_SequenceStart:
sr.StartSequence = sub.opts.StartSequence
}
b, _ := sr.Marshal()
reply, err := sc.nc.Request(sc.subRequests, b, 2*time.Second)
if err != nil {
sub.inboxSub.Unsubscribe()
return nil, err
}
r := &pb.SubscriptionResponse{}
if err := r.Unmarshal(reply.Data); err != nil {
sub.inboxSub.Unsubscribe()
return nil, err
}
if r.Error != "" {
sub.inboxSub.Unsubscribe()
return nil, errors.New(r.Error)
}
sub.ackInbox = r.AckInbox
return sub, nil
}
// Connect will form a connection to the NATS Streaming subsystem.
func Connect(stanClusterID, clientID string, options ...Option) (Conn, error) {
// Process Options
c := conn{clientID: clientID, opts: DefaultOptions}
for _, opt := range options {
if err := opt(&c.opts); err != nil {
return nil, err
}
}
// Check if the user has provided a connection as an option
c.nc = c.opts.NatsConn
// Create a NATS connection if it doesn't exist.
if c.nc == nil {
nc, err := nats.Connect(c.opts.NatsURL)
if err != nil {
return nil, err
}
c.nc = nc
c.ncOwned = true
} else if !c.nc.IsConnected() {
// Bail if the custom NATS connection is disconnected
return nil, ErrBadConnection
}
// Create a heartbeat inbox
hbInbox := nats.NewInbox()
var err error
if c.hbSubscription, err = c.nc.Subscribe(hbInbox, c.processHeartBeat); err != nil {
c.Close()
return nil, err
}
// Send Request to discover the cluster
discoverSubject := c.opts.DiscoverPrefix + "." + stanClusterID
req := &pb.ConnectRequest{ClientID: clientID, HeartbeatInbox: hbInbox}
b, _ := req.Marshal()
reply, err := c.nc.Request(discoverSubject, b, c.opts.ConnectTimeout)
if err != nil {
c.Close()
if err == nats.ErrTimeout {
return nil, ErrConnectReqTimeout
}
return nil, err
}
// Process the response, grab server pubPrefix
cr := &pb.ConnectResponse{}
err = cr.Unmarshal(reply.Data)
if err != nil {
c.Close()
return nil, err
}
if cr.Error != "" {
c.Close()
return nil, errors.New(cr.Error)
}
// Capture cluster configuration endpoints to publish and subscribe/unsubscribe.
c.pubPrefix = cr.PubPrefix
c.subRequests = cr.SubRequests
c.unsubRequests = cr.UnsubRequests
c.closeRequests = cr.CloseRequests
// Setup the ACK subscription
c.ackSubject = DefaultACKPrefix + "." + nuid.Next()
if c.ackSubscription, err = c.nc.Subscribe(c.ackSubject, c.processAck); err != nil {
c.Close()
return nil, err
}
c.ackSubscription.SetPendingLimits(1024*1024, 32*1024*1024)
c.pubAckMap = make(map[string]*ack)
// Create Subscription map
c.subMap = make(map[string]*subscription)
c.pubAckChan = make(chan struct{}, c.opts.MaxPubAcksInflight)
// Attach a finalizer
runtime.SetFinalizer(&c, func(sc *conn) { sc.Close() })
return &c, nil
}