/* GStreamer plugin for forward error correction

 * Copyright (C) 2017 Pexip

 *

 * This library is free software; you can redistribute it and/or

 * modify it under the terms of the GNU Lesser General Public

 * License as published by the Free Software Foundation; either

 * version 2.1 of the License, or (at your option) any later version.

 *

 * This library is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

 * Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public

 * License along with this library; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

 *

 * Author: Mikhail Fludkov <misha@pexip.com>

 */

#include <gst/rtp/gstrtpbuffer.h>

#include <gst/check/gstcheck.h>

#include <gst/check/gstharness.h>

#include "../../../gst/rtp/rtpstorage.h"

#include "../../../gst/rtp/rtpstoragestream.c"

#include "../../../gst/rtp/rtpstorage.c"

#define RTP_CLOCK_RATE (90000)

#define RTP_FRAME_DUR (RTP_CLOCK_RATE / 30)

#define RTP_TSTAMP_BASE (0x11111111)

#define GST_TSTAMP_BASE (0x22222222)

#define RTP_TSTAMP(i) (RTP_FRAME_DUR * (i) + RTP_TSTAMP_BASE)

#define GST_TSTAMP(i) (RTP_PACKET_DUR * (i) + GST_TSTAMP_BASE)

#define RTP_PACKET_DUR (10 * GST_MSECOND)

static GstBufferList *

get_packets_for_recovery (GstHarness * h, gint fec_pt, guint32 ssrc,

    guint16 lost_seq)

{

  GstBufferList *res;

  RtpStorage *internal_storage;

  g_object_get (h->element, "internal-storage", &internal_storage, NULL);

  res =

      rtp_storage_get_packets_for_recovery (internal_storage, fec_pt, ssrc,

      lost_seq);

  g_object_unref (internal_storage);

  return res;

}

static void

put_recovered_packet (GstHarness * h, GstBuffer * buffer, guint8 pt,

    guint32 ssrc, guint16 seq)

{

  RtpStorage *internal_storage;

  g_object_get (h->element, "internal-storage", &internal_storage, NULL);

  rtp_storage_do_put_recovered_packet (internal_storage, buffer, pt, ssrc, seq);

  g_object_unref (internal_storage);

}

static GstBuffer *

create_rtp_packet (guint8 pt, guint32 ssrc, guint32 timestamp, guint16 seq)

{

  GstBuffer *buf = gst_rtp_buffer_new_allocate (0, 0, 0);

  GstRTPBuffer rtp = GST_RTP_BUFFER_INIT;

  fail_unless (gst_rtp_buffer_map (buf, GST_MAP_WRITE, &rtp));

  gst_rtp_buffer_set_ssrc (&rtp, ssrc);

  gst_rtp_buffer_set_payload_type (&rtp, pt);

  gst_rtp_buffer_set_timestamp (&rtp, timestamp);

  gst_rtp_buffer_set_seq (&rtp, seq);

  GST_BUFFER_DTS (buf) = GST_TSTAMP (seq);

  gst_rtp_buffer_unmap (&rtp;;

  return buf;

}

GST_START_TEST (rtpstorage_up_and_down)

{

  GstHarness *h = gst_harness_new ("rtpstorage");

  gst_harness_set_src_caps_str (h, "application/x-rtp");

  gst_harness_teardown (h);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_resize)

{

  guint i, j;

  GstBuffer *bufin, *bufout, *bufs[10];

  GstHarness *h = gst_harness_new ("rtpstorage");

  gst_harness_set_src_caps_str (h, "application/x-rtp");

  g_object_set (h->element, "size-time", 0, NULL);

  bufin = create_rtp_packet (96, 0xabe2b0b, 0x111111, 0);

  bufout = gst_harness_push_and_pull (h, bufin);

  fail_unless (bufin == bufout);

  fail_unless (gst_buffer_is_writable (bufout));

  g_object_set (h->element, "size-time",

      (G_N_ELEMENTS (bufs) - 1) * RTP_PACKET_DUR, NULL);

  // Pushing 10 buffers all of them should have ref. count =2

  for (i = 0; i < G_N_ELEMENTS (bufs); ++i) {

    bufs[i] =

        gst_harness_push_and_pull (h, create_rtp_packet (96, 0xabe2b0b,

            0x111111, i));

    for (j = 0; j <= i; ++j)

      fail_unless (!gst_buffer_is_writable (bufs[j]));

  }

  // The next 10 buffers should expel the first 10

  for (i = 0; i < G_N_ELEMENTS (bufs); ++i) {

    gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

                0xabe2b0b, 0x111111, G_N_ELEMENTS (bufs) + i)));

    for (j = 0; j <= i; ++j)

      fail_unless (gst_buffer_is_writable (bufs[j]));

  }

  for (i = 0; i < G_N_ELEMENTS (bufs); ++i)

    gst_buffer_unref (bufs[i]);

  gst_buffer_unref (bufout);

  gst_harness_teardown (h);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_stop_redundant_packets)

{

  GstHarness *h = gst_harness_new ("rtpstorage");

  GstBuffer *bufinp;

  g_object_set (h->element, "size-time", 2 * RTP_PACKET_DUR, NULL);

  gst_harness_set_src_caps_str (h, "application/x-rtp");

  bufinp = create_rtp_packet (96, 0xabe2b0b, 0x111111, 0);

  GST_BUFFER_FLAG_SET (bufinp, GST_RTP_BUFFER_FLAG_REDUNDANT);

  gst_harness_push (h, bufinp);

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

              0xabe2b0b, 0x111111, 1)));

  fail_unless_equals_int (gst_harness_buffers_received (h), 1);

  gst_harness_teardown (h);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_unknown_ssrc)

{

  GstBufferList *bufs_out;

  GstHarness *h = gst_harness_new ("rtpstorage");

  g_object_set (h->element, "size-time", RTP_PACKET_DUR, NULL);

  gst_harness_set_src_caps_str (h, "application/x-rtp");

  /* No packets has been pushed through yet */

  bufs_out = get_packets_for_recovery (h, 100, 0xabe2b0b, 0);

  fail_unless (NULL == bufs_out);

  /* 1 packet with ssrc=0xabe2bob pushed. Asking for ssrc=0xdeadbeef */

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

              0xabe2b0b, 0x111111, 0)));

  bufs_out = get_packets_for_recovery (h, 100, 0xdeadbeef, 0);

  fail_unless (NULL == bufs_out);

  gst_harness_teardown (h);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_packet_not_lost)

{

  GstBuffer *buf;

  GstBufferList *bufs_out;

  GstHarness *h = gst_harness_new ("rtpstorage");

  g_object_set (h->element, "size-time", 10 * RTP_PACKET_DUR, NULL);

  gst_harness_set_src_caps_str (h, "application/x-rtp");

  /* Pushing through 2 frames + 2 FEC */

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

              0xabe2b0b, RTP_TSTAMP (0), 0)));

  gst_buffer_unref (gst_harness_push_and_pull (h, (buf =

              create_rtp_packet (96, 0xabe2b0b, RTP_TSTAMP (1), 1))));

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

              0xabe2b0b, RTP_TSTAMP (1), 2)));

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

              0xabe2b0b, RTP_TSTAMP (1), 3)));

  /* Asking for a packet which was pushed before */

  bufs_out = get_packets_for_recovery (h, 100, 0xabe2b0b, 1);

  fail_unless (NULL != bufs_out);

  fail_unless_equals_int (1, gst_buffer_list_length (bufs_out));

  fail_unless (gst_buffer_list_get (bufs_out, 0) == buf);

  gst_buffer_list_unref (bufs_out);

  gst_harness_teardown (h);

}

GST_END_TEST;

GST_START_TEST (test_rtpstorage_put_recovered_packet)

{

  GstBuffer *bufs_in[4];

  GstBufferList *bufs_out;

  GstHarness *h = gst_harness_new ("rtpstorage");

  g_object_set (h->element, "size-time", 10 * RTP_PACKET_DUR, NULL);

  gst_harness_set_src_caps_str (h, "application/x-rtp");

  /* Pushing through 2 frames + 2 FEC

   * Packets with sequence numbers 1 and 2 are lost */

  bufs_in[0] = create_rtp_packet (96, 0xabe2b0b, RTP_TSTAMP (0), 0);

  bufs_in[1] = NULL;

  bufs_in[2] = NULL;

  bufs_in[3] = create_rtp_packet (100, 0xabe2b0b, RTP_TSTAMP (1), 3);

  gst_buffer_unref (gst_harness_push_and_pull (h, bufs_in[0]));

  gst_buffer_unref (gst_harness_push_and_pull (h, bufs_in[3]));

  /* 1 more frame + 1 FEC */

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

              0xabe2b0b, RTP_TSTAMP (2), 4)));

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (100,

              0xabe2b0b, RTP_TSTAMP (2), 5)));

  /* Asking for the lost packet seq=1 */

  bufs_out = get_packets_for_recovery (h, 100, 0xabe2b0b, 1);

  fail_unless (NULL != bufs_out);

  fail_unless_equals_int (2, gst_buffer_list_length (bufs_out));

  fail_unless (gst_buffer_list_get (bufs_out, 0) == bufs_in[0]);

  fail_unless (gst_buffer_list_get (bufs_out, 1) == bufs_in[3]);

  gst_buffer_list_unref (bufs_out);

  /* During recovery the packet of a new frame has arrived */

  gst_buffer_unref (gst_harness_push_and_pull (h, create_rtp_packet (96,

              0xabe2b0b, RTP_TSTAMP (3), 6)));

  /* Say we recovered packet with seq=1 and put it back in the storage */

  bufs_in[1] = create_rtp_packet (96, 0xabe2b0b, RTP_TSTAMP (1), 1);

  put_recovered_packet (h, bufs_in[1], 96, 0xabe2b0b, 1);

  /* Asking for the lost packet seq=2 */

  bufs_out = get_packets_for_recovery (h, 100, 0xabe2b0b, 2);

  fail_unless (NULL != bufs_out);

  fail_unless_equals_int (3, gst_buffer_list_length (bufs_out));

  fail_unless (gst_buffer_list_get (bufs_out, 0) == bufs_in[0]);

  fail_unless (gst_buffer_list_get (bufs_out, 1) == bufs_in[1]);

  fail_unless (gst_buffer_list_get (bufs_out, 2) == bufs_in[3]);

  gst_buffer_list_unref (bufs_out);

  gst_harness_teardown (h);

}

GST_END_TEST;

static void

_single_ssrc_test (GstHarness * h, guint32 ssrc,

    guint16 seq_start, guint16 nth_to_loose,

    gsize expected_buf_size, gsize expected_first_buffer_idx)

{

  guint i;

  GPtrArray *bufs_in =

      g_ptr_array_new_with_free_func ((GDestroyNotify) gst_buffer_unref);

  GstBufferList *bufs_out;

  /* 2 frames + 2 FEC */

  g_ptr_array_add (bufs_in, create_rtp_packet (96, ssrc, RTP_TSTAMP (0),

          seq_start + 0));

  g_ptr_array_add (bufs_in, create_rtp_packet (96, ssrc, RTP_TSTAMP (1),

          seq_start + 1));

  g_ptr_array_add (bufs_in, create_rtp_packet (100, ssrc, RTP_TSTAMP (1),

          seq_start + 2));

  g_ptr_array_add (bufs_in, create_rtp_packet (100, ssrc, RTP_TSTAMP (1),

          seq_start + 3));

  /* 3 frames + 2 FEC */

  g_ptr_array_add (bufs_in, create_rtp_packet (96, ssrc, RTP_TSTAMP (2),

          seq_start + 4));

  g_ptr_array_add (bufs_in, create_rtp_packet (96, ssrc, RTP_TSTAMP (3),

          seq_start + 5));

  g_ptr_array_add (bufs_in, create_rtp_packet (96, ssrc, RTP_TSTAMP (4),

          seq_start + 6));

  g_ptr_array_add (bufs_in, create_rtp_packet (100, ssrc, RTP_TSTAMP (4),

          seq_start + 7));

  g_ptr_array_add (bufs_in, create_rtp_packet (100, ssrc, RTP_TSTAMP (4),

          seq_start + 8));

  g_ptr_array_add (bufs_in, create_rtp_packet (100, ssrc, RTP_TSTAMP (4),

          seq_start + 9));

  /* 2 frames + no FEC */

  g_ptr_array_add (bufs_in, create_rtp_packet (96, ssrc, RTP_TSTAMP (5),

          seq_start + 10));

  g_ptr_array_add (bufs_in, create_rtp_packet (96, ssrc, RTP_TSTAMP (6),

          seq_start + 11));

  /* Loosing one */

  g_ptr_array_remove_index (bufs_in, nth_to_loose);

  /* Push all of them through */

  for (i = 0; i < bufs_in->len; ++i)

    gst_buffer_unref (gst_harness_push_and_pull (h,

            gst_buffer_ref (g_ptr_array_index (bufs_in, i))));

  bufs_out =

      get_packets_for_recovery (h, 100, ssrc,

      (guint16) (seq_start + nth_to_loose));

  if (0 == expected_buf_size) {

    fail_unless (NULL == bufs_out);

  } else {

    fail_unless (NULL != bufs_out);

    fail_unless_equals_int (expected_buf_size,

        gst_buffer_list_length (bufs_out));

    for (i = 0; i < gst_buffer_list_length (bufs_out); ++i)

      fail_unless (gst_buffer_list_get (bufs_out, i) ==

          g_ptr_array_index (bufs_in, expected_first_buffer_idx + i));

    gst_buffer_list_unref (bufs_out);

  }

  g_ptr_array_unref (bufs_in);

}

static void

_multiple_ssrcs_test (guint16 nth_to_loose,

    gsize expected_buf_size, gsize expected_first_buffer_idx)

{

  guint16 stream0_seq_start = 200;

  guint16 stream1_seq_start = 65529;

  GstHarness *h = gst_harness_new ("rtpstorage");

  g_object_set (h->element, "size-time", 12 * RTP_PACKET_DUR, NULL);

  gst_harness_set_src_caps_str (h, "application/x-rtp");

  _single_ssrc_test (h, 0x0abe2b0b, stream0_seq_start,

      nth_to_loose, expected_buf_size, expected_first_buffer_idx);

  _single_ssrc_test (h, 0xdeadbeef, stream1_seq_start,

      nth_to_loose, expected_buf_size, expected_first_buffer_idx);

  gst_harness_teardown (h);

}

GST_START_TEST (rtpstorage_loss_pattern0)

{

  _multiple_ssrcs_test (1, 3, 0);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern1)

{

  _multiple_ssrcs_test (2, 3, 0);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern2)

{

  _multiple_ssrcs_test (3, 6, 3);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern3)

{

  _multiple_ssrcs_test (4, 5, 4);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern4)

{

  _multiple_ssrcs_test (5, 5, 4);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern5)

{

  _multiple_ssrcs_test (6, 5, 4);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern6)

{

  _multiple_ssrcs_test (7, 5, 4);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern7)

{

  _multiple_ssrcs_test (8, 5, 4);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern8)

{

  _multiple_ssrcs_test (9, 0, 0);

}

GST_END_TEST;

GST_START_TEST (rtpstorage_loss_pattern9)

{

  _multiple_ssrcs_test (10, 0, 0);

}

GST_END_TEST;

#define STRESS_TEST_SSRCS (8)

#define STRESS_TEST_STORAGE_DEPTH (50)

typedef struct _StressTestData StressTestData;

struct _StressTestData

{

  guint16 seq[STRESS_TEST_SSRCS];

  guint32 ssrc[STRESS_TEST_SSRCS];

  gsize count[STRESS_TEST_SSRCS];

  GRand *rnd;

};

static GstBuffer *

rtpstorage_stress_prepare_buffer (GstHarness * h, gpointer data)

{

  static const guint8 fec_pt = 100;

  static const guint8 media_pt = 96;

  StressTestData *test_data = data;

  gsize ssrc_idx = g_rand_int_range (test_data->rnd, 0, STRESS_TEST_SSRCS);

  guint16 seq = test_data->seq[ssrc_idx];

  guint32 ssrc = test_data->ssrc[ssrc_idx];

  gboolean is_fec = test_data->count[ssrc_idx] > 0 && (seq % 5 == 0

      || seq % 5 == 1);

  guint8 pt = is_fec ? fec_pt : media_pt;

  GstBuffer *buf = create_rtp_packet (pt, ssrc, RTP_TSTAMP (0), seq);

  ++test_data->seq[ssrc_idx];

  ++test_data->count[ssrc_idx];

  return buf;

}

GST_START_TEST (rtpstorage_stress)

{

  GRand *rnd;

  GTimer *timer;

  GstCaps *caps;

  GstSegment segment;

  GstHarnessThread *ht;

  StressTestData test_data;

  guint seed, i, total, requested;

  GstHarness *h = gst_harness_new ("rtpstorage");

  g_object_set (h->element, "size-time",

      STRESS_TEST_STORAGE_DEPTH * RTP_PACKET_DUR, NULL);

  /* The stress test pushes buffers with STRESS_TEST_SSRCS different

   * ssrcs from one thread and requests packets for FEC recovery from

   * another thread.

   * */

  memset (&test_data, 0, sizeof (test_data));

  seed = g_random_int ();

  test_data.rnd = g_rand_new_with_seed (seed);

  for (i = 0; i < STRESS_TEST_SSRCS; ++i) {

    test_data.ssrc[i] = 0x00112233 + i * 0x01000000;

    test_data.seq[i] = g_rand_int_range (test_data.rnd, 0, 0x10000);

  }

  gst_segment_init (&segment, GST_FORMAT_TIME);

  caps = gst_caps_from_string ("application/x-rtp");

  rnd = g_rand_copy (test_data.rnd);

  GST_INFO ("%u seed", seed);

  ht = gst_harness_stress_push_buffer_with_cb_start (h, caps, &segment,

      rtpstorage_stress_prepare_buffer, &test_data, NULL);

  requested = 0;

  timer = g_timer_new ();

  while (g_timer_elapsed (timer, NULL) < 2) {

    gsize ssrc_idx = g_rand_int_range (rnd, 0, STRESS_TEST_SSRCS);

    /* The following if statement is simply keeping the log

     * clean from ERROR messages */

    if (*((volatile gsize *) &test_data.count[ssrc_idx]) > 1) {

      guint16 lost_seq = *((volatile guint16 *) &test_data.seq[ssrc_idx]) - 5;

      GstBufferList *bufs_out = get_packets_for_recovery (h, 100,

          test_data.ssrc[ssrc_idx], lost_seq);

      if (bufs_out) {

        requested += gst_buffer_list_length (bufs_out);

        gst_buffer_list_unref (bufs_out);

      }

    }

    /* Having sleep here makes it hard to detect the race, but we need it to

     * allow another thread to push more buffers when running under valgrind */

    g_usleep (G_USEC_PER_SEC / 10000);

  }

  gst_harness_stress_thread_stop (ht);

  for (i = 0, total = 0; i < STRESS_TEST_SSRCS; ++i) {

    GST_INFO ("SSRC 0x%08x: %u packets", test_data.ssrc[i],

        (guint32) test_data.count[i]);

    total += test_data.count[i];

  }

  GST_INFO ("%u packets pushed through, %u requested", total, requested);

  g_rand_free (rnd);

  g_rand_free (test_data.rnd);

  gst_caps_unref (caps);

  g_timer_destroy (timer);

  gst_harness_teardown (h);

}

GST_END_TEST;

static Suite *

rtpstorage_suite (void)

{

  Suite *s = suite_create ("rtpstorage");

  TCase *tc_chain = tcase_create ("general");

  suite_add_tcase (s, tc_chain);

  tcase_add_test (tc_chain, rtpstorage_up_and_down);

  tcase_add_test (tc_chain, rtpstorage_resize);

  tcase_add_test (tc_chain, rtpstorage_stop_redundant_packets);

  tcase_add_test (tc_chain, rtpstorage_unknown_ssrc);

  tcase_add_test (tc_chain, rtpstorage_packet_not_lost);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern0);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern1);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern2);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern3);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern4);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern5);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern6);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern7);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern8);

  tcase_add_test (tc_chain, rtpstorage_loss_pattern9);

  tcase_add_test (tc_chain, test_rtpstorage_put_recovered_packet);

  tcase_add_test (tc_chain, rtpstorage_stress);

  return s;

}

GST_CHECK_MAIN (rtpstorage)