File Coverage

blib/lib/Net/Async/NATS.pm

Criterion	Covered	Total	%
statement	134	245	54.6
branch	35	116	30.1
condition	19	49	38.7
subroutine	27	41	65.8
pod	15	15	100.0
total	230	466	49.3

package Net::Async::NATS;
# ABSTRACT: Async NATS client for IO::Async
our $VERSION = '0.002';
use strict;
use warnings;
use parent 'IO::Async::Notifier';

use Carp qw(croak);
use Future;
use Future::AsyncAwait;
use IO::Async::Stream;
use JSON::MaybeXS qw(encode_json decode_json);
use Scalar::Util qw(weaken);

use Net::Async::NATS::Subscription;


sub configure {
    my ($self, %params) = @_;

    for my $key (qw(host port name user pass auth_token
                     verbose pedantic reconnect max_reconnect_attempts
                     reconnect_wait on_disconnect on_error on_connect)) {
        $self->{$key} = delete $params{$key} if exists $params{$key};
    }

    # Defaults
    $self->{host}     //= 'localhost';
    $self->{port}     //= 4222;
    $self->{name}     //= 'net-async-nats-perl';
    $self->{verbose}  //= 0;
    $self->{pedantic} //= 0;
    $self->{reconnect} //= 1;
    $self->{max_reconnect_attempts} //= 10;
    $self->{reconnect_wait} //= 2;

    # Internal state
    $self->{_sid_counter}   //= 0;
    $self->{_subscriptions} //= {};
    $self->{_pending}       //= [];
    $self->{_connected}     //= 0;
    $self->{_server_info}   //= {};
    $self->{_ping_future}   //= undef;
    $self->{_connect_future} //= undef;
    $self->{_inbox_prefix}  //= '_INBOX.' . _random_id();

    $self->SUPER::configure(%params);
}


sub host     { $_[0]->{host} }
sub port     { $_[0]->{port} }
sub name     { $_[0]->{name} }
sub verbose  { $_[0]->{verbose} }
sub pedantic { $_[0]->{pedantic} }


sub server_info  { $_[0]->{_server_info} }
sub is_connected { $_[0]->{_connected} }


async sub connect {
    my ($self) = @_;

    croak "Already connected" if $self->{_connected};

    $self->{_connect_future} = $self->loop->new_future;

    my $stream = IO::Async::Stream->new(
        handle => undef,
        on_read => sub { $self->_on_read(@_) },
        on_read_eof => sub { $self->_on_disconnect('read_eof') },
        on_write_eof => sub { $self->_on_disconnect('write_eof') },
        on_read_error => sub { $self->_on_error("Read error: $_[1]") },
        on_write_error => sub { $self->_on_error("Write error: $_[1]") },
    );
    $self->{_stream} = $stream;
    $self->add_child($stream);

    # Retain the TCP connect future — if GC'd before resolution, the
    # stream never gets its handle and on_read never fires.
    $self->{_tcp_connect_future} = $stream->connect(
        host    => $self->{host},
        service => $self->{port},
    )->on_fail(sub {
        my $err = shift;
        $self->{_connect_future}->fail("Connection failed: $err")
            unless $self->{_connect_future}->is_ready;
    });

    return await $self->{_connect_future};
}


async sub publish {
    my ($self, $subject, $payload, %opts) = @_;

    croak "Not connected" unless $self->{_connected};
    croak "Subject required" unless defined $subject && length $subject;

    $payload //= '';
    my $bytes = length($payload);
    my $reply = $opts{reply_to};

    my $cmd = defined $reply
        ? "PUB $subject $reply $bytes\r\n$payload\r\n"
        : "PUB $subject $bytes\r\n$payload\r\n";

    $self->_write($cmd);
    return;
}


async sub subscribe {
    my ($self, $subject, $callback, %opts) = @_;

    croak "Not connected" unless $self->{_connected};
    croak "Subject required" unless defined $subject && length $subject;
    croak "Callback required" unless ref $callback eq 'CODE';

    my $sid   = ++$self->{_sid_counter};
    my $queue = $opts{queue};

    my $sub = Net::Async::NATS::Subscription->new(
        sid      => $sid,
        subject  => $subject,
        queue    => $queue,
        callback => $callback,
    );
    $self->{_subscriptions}{$sid} = $sub;

    my $cmd = defined $queue
        ? "SUB $subject $queue $sid\r\n"
        : "SUB $subject $sid\r\n";

    $self->_write($cmd);
    return $sub;
}


async sub unsubscribe {
    my ($self, $sub, %opts) = @_;

    croak "Not connected" unless $self->{_connected};

    my $sid = ref $sub ? $sub->sid : $sub;
    my $max = $opts{max_msgs};

    my $cmd = defined $max
        ? "UNSUB $sid $max\r\n"
        : "UNSUB $sid\r\n";

    $self->_write($cmd);
    delete $self->{_subscriptions}{$sid} unless defined $max;
    return;
}


async sub request {
    my ($self, $subject, $payload, %opts) = @_;

    croak "Not connected" unless $self->{_connected};

    my $timeout = $opts{timeout} // 30;
    my $inbox   = $self->{_inbox_prefix} . '.' . _random_id();
    my $f       = $self->loop->new_future;

    my $sub = await $self->subscribe($inbox, sub {
        my ($subj, $data, $reply) = @_;
        $f->done($data, $subj) unless $f->is_ready;
    });

    # Auto-unsubscribe after 1 message
    await $self->unsubscribe($sub, max_msgs => 1);

    # Publish with reply-to
    await $self->publish($subject, $payload, reply_to => $inbox);

    # Apply timeout
    my $timer = $self->loop->delay_future(after => $timeout);
    my $result = await Future->wait_any($f, $timer->then_fail('Request timed out'));

    return $result;
}


async sub ping {
    my ($self) = @_;
    croak "Not connected" unless $self->{_connected};

    $self->{_ping_future} = $self->loop->new_future;
    $self->_write("PING\r\n");
    return await $self->{_ping_future};
}


async sub disconnect {
    my ($self) = @_;
    return unless $self->{_connected};

    $self->{_connected}  = 0;
    $self->{reconnect}   = 0;  # suppress auto-reconnect
    $self->{_subscriptions} = {};

    if (my $stream = $self->{_stream}) {
        $stream->close_when_empty;
    }
    return;
}

# ── Wire protocol parsing ────────────────────────────────

sub _on_read {
    my ($self, $stream, $buffref, $eof) = @_;

    while ($$buffref =~ s/\A([^\r\n]*)\r\n//) {
        my $line = $1;

        # MSG   [reply-to] <#bytes>
        if ($line =~ /\AMSG\s+(\S+)\s+(\S+)\s+(?:(\S+)\s+)?(\d+)\z/i) {
            my ($subject, $sid, $reply_to, $bytes) = ($1, $2, $3, $4);

            # Need to read payload + trailing CRLF
            if (length($$buffref) < $bytes + 2) {
                # Put line back, wait for more data
                $$buffref = "MSG $subject $sid "
                    . (defined $reply_to ? "$reply_to " : '')
                    . "$bytes\r\n$$buffref";
                return 0;
            }

            my $payload = substr($$buffref, 0, $bytes, '');
            $$buffref =~ s/\A\r\n//;  # consume trailing CRLF

            $self->_dispatch_msg($subject, $sid, $reply_to, $payload);
            next;
        }

        # HMSG   [reply-to] <#header_bytes> <#total_bytes>
        if ($line =~ /\AHMSG\s+(\S+)\s+(\S+)\s+(?:(\S+)\s+)?(\d+)\s+(\d+)\z/i) {
            my ($subject, $sid, $reply_to, $hdr_bytes, $total_bytes) = ($1, $2, $3, $4, $5);

            if (length($$buffref) < $total_bytes + 2) {
                $$buffref = "HMSG $subject $sid "
                    . (defined $reply_to ? "$reply_to " : '')
                    . "$hdr_bytes $total_bytes\r\n$$buffref";
                return 0;
            }

            my $raw = substr($$buffref, 0, $total_bytes, '');
            $$buffref =~ s/\A\r\n//;

            my $payload_bytes = $total_bytes - $hdr_bytes;
            my $payload = $payload_bytes > 0
                ? substr($raw, $hdr_bytes)
                : '';

            # For now, ignore headers — deliver payload only
            $self->_dispatch_msg($subject, $sid, $reply_to, $payload);
            next;
        }

        # INFO {...}
        if ($line =~ /\AINFO\s+(\{.*\})\s*\z/i) {
            my $info = eval { decode_json($1) } // {};
            $self->{_server_info} = $info;
            $self->_handle_info($info);
            next;
        }

        # PING
        if ($line =~ /\APING\z/i) {
            $self->_write("PONG\r\n");
            next;
        }

        # PONG
        if ($line =~ /\APONG\z/i) {
            if (my $f = delete $self->{_ping_future}) {
                $f->done if !$f->is_ready;
            }
            next;
        }

        # +OK
        next if $line =~ /\A\+OK\z/i;

        # -ERR
        if ($line =~ /\A-ERR\s+'?(.+?)'?\z/i) {
            $self->_on_error($1);
            next;
        }
    }

    return 0;
}

sub _handle_info {
    my ($self, $info) = @_;

    # If we haven't sent CONNECT yet, do it now
    return if $self->{_connected};

    my %connect = (
        verbose     => $self->{verbose} ? \1 : \0,
        pedantic    => $self->{pedantic} ? \1 : \0,
        tls_required => \0,
        lang        => 'perl',
        version     => ($Net::Async::NATS::VERSION // '0.001'),
        name        => $self->{name},
        protocol    => 1,
        echo        => \1,
        headers     => \1,
        no_responders => \1,
    );

    if (defined $self->{auth_token}) {
        $connect{auth_token} = $self->{auth_token};
    }
    if (defined $self->{user}) {
        $connect{user} = $self->{user};
        $connect{pass} = $self->{pass} // '';
    }

    my $json = encode_json(\%connect);
    $self->_write("CONNECT $json\r\n");

    # Send initial PING to verify connection
    $self->_write("PING\r\n");

    $self->{_connected} = 1;
    delete $self->{_tcp_connect_future};  # no longer needed

    if (my $f = delete $self->{_connect_future}) {
        $f->done($info) unless $f->is_ready;
    }

    if (my $cb = $self->{on_connect}) {
        $cb->($self, $info);
    }
}

sub _dispatch_msg {
    my ($self, $subject, $sid, $reply_to, $payload) = @_;

    my $sub = $self->{_subscriptions}{$sid};
    return unless $sub;

    $sub->callback->($subject, $payload, $reply_to);

    if (defined $sub->max_msgs) {
        $sub->{_received}++;
        if ($sub->{_received} >= $sub->max_msgs) {
            delete $self->{_subscriptions}{$sid};
        }
    }
}

sub _write {
    my ($self, $data) = @_;
    if (my $stream = $self->{_stream}) {
        $stream->write($data);
    }
}

sub _on_disconnect {
    my ($self, $reason) = @_;
    $self->{_connected} = 0;

    if (my $cb = $self->{on_disconnect}) {
        $cb->($self, $reason);
    }

    if ($self->{reconnect}) {
        $self->_reconnect;
    }
}

sub _on_error {
    my ($self, $error) = @_;
    if (my $cb = $self->{on_error}) {
        $cb->($self, $error);
    }
}

sub _reconnect {
    my ($self) = @_;
    my $attempts = 0;
    my $max = $self->{max_reconnect_attempts};
    my $wait = $self->{reconnect_wait};

    weaken(my $weak_self = $self);

    my $try; $try = sub {
        my $self = $weak_self or return;
        return if $self->{_connected};
        return if ++$attempts > $max;

        $self->loop->delay_future(after => $wait)->on_done(sub {
            my $self = $weak_self or return;
            return if $self->{_connected};

            # Remove old stream
            if (my $old = delete $self->{_stream}) {
                $self->remove_child($old) if $old->parent;
            }

            # Save subscriptions for replay
            my %saved_subs = %{ $self->{_subscriptions} };

            eval {
                my $f = $self->connect;
                $f->on_done(sub {
                    my $self = $weak_self or return;
                    # Replay subscriptions
                    for my $sub (values %saved_subs) {
                        my $cmd = defined $sub->queue
                            ? "SUB " . $sub->subject . " " . $sub->queue . " " . $sub->sid . "\r\n"
                            : "SUB " . $sub->subject . " " . $sub->sid . "\r\n";
                        $self->_write($cmd);
                        $self->{_subscriptions}{$sub->sid} = $sub;
                    }
                });
                $f->on_fail(sub { $try->() });
            };
        });
    };

    $try->();
}

sub _random_id {
    my $bytes = '';
    if (open my $fh, '<:raw', '/dev/urandom') {
        read $fh, $bytes, 12;
        close $fh;
    } else {
        $bytes = pack('N3', rand(2**32), rand(2**32), rand(2**32));
    }
    return unpack('H*', $bytes);
}


1;

__END__