akkoma/lib/load_testing/generator.ex

defmodule Pleroma.LoadTesting.Generator do
  use Pleroma.LoadTesting.Helper
  alias Pleroma.Web.CommonAPI

  def generate_users(opts) do
    IO.puts("Starting generating #{opts[:users_max]} users...")
    {time, _} = :timer.tc(fn -> do_generate_users(opts) end)

    IO.puts("Inserting users take #{to_sec(time)} sec.\n")
  end

  defp do_generate_users(opts) do
    max = Keyword.get(opts, :users_max)

    Task.async_stream(
      1..max,
      &generate_user_data(&1),
      max_concurrency: 10,
      timeout: 30_000
    )
    |> Enum.to_list()
  end

  defp generate_user_data(i) do
    remote = Enum.random([true, false])

    user = %User{
      name: "Test テスト User #{i}",
      email: "user#{i}@example.com",
      nickname: "nick#{i}",
      password_hash: Comeonin.Pbkdf2.hashpwsalt("test"),
      bio: "Tester Number #{i}",
      info: %{},
      local: remote
    }

    user_urls =
      if remote do
        base_url =
          Enum.random(["https://domain1.com", "https://domain2.com", "https://domain3.com"])

        ap_id = "#{base_url}/users/#{user.nickname}"

        %{
          ap_id: ap_id,
          follower_address: ap_id <> "/followers",
          following_address: ap_id <> "/following",
          following: [ap_id]
        }
      else
        %{
          ap_id: User.ap_id(user),
          follower_address: User.ap_followers(user),
          following_address: User.ap_following(user),
          following: [User.ap_id(user)]
        }
      end

    user = Map.merge(user, user_urls)

    Repo.insert!(user)
  end

  def generate_activities(user, users) do
    do_generate_activities(user, users)
  end

  defp do_generate_activities(user, users) do
    IO.puts("Starting generating 20000 common activities...")

    {time, _} =
      :timer.tc(fn ->
        Task.async_stream(
          1..20_000,
          fn _ ->
            do_generate_activity([user | users])
          end,
          max_concurrency: 10,
          timeout: 30_000
        )
        |> Stream.run()
      end)

    IO.puts("Inserting common activities take #{to_sec(time)} sec.\n")

    IO.puts("Starting generating 20000 activities with mentions...")

    {time, _} =
      :timer.tc(fn ->
        Task.async_stream(
          1..20_000,
          fn _ ->
            do_generate_activity_with_mention(user, users)
          end,
          max_concurrency: 10,
          timeout: 30_000
        )
        |> Stream.run()
      end)

    IO.puts("Inserting activities with menthions take #{to_sec(time)} sec.\n")

    IO.puts("Starting generating 10000 activities with threads...")

    {time, _} =
      :timer.tc(fn ->
        Task.async_stream(
          1..10_000,
          fn _ ->
            do_generate_threads([user | users])
          end,
          max_concurrency: 10,
          timeout: 30_000
        )
        |> Stream.run()
      end)

    IO.puts("Inserting activities with threads take #{to_sec(time)} sec.\n")
  end

  defp do_generate_activity(users) do
    post = %{
      "status" => "Some status without mention with random user"
    }

    CommonAPI.post(Enum.random(users), post)
  end

  defp do_generate_activity_with_mention(user, users) do
    mentions_cnt = Enum.random([2, 3, 4, 5])
    with_user = Enum.random([true, false])
    users = Enum.shuffle(users)
    mentions_users = Enum.take(users, mentions_cnt)
    mentions_users = if with_user, do: [user | mentions_users], else: mentions_users

    mentions_str =
      Enum.map(mentions_users, fn user -> "@" <> user.nickname end) |> Enum.join(", ")

    post = %{
      "status" => mentions_str <> "some status with mentions random users"
    }

    CommonAPI.post(Enum.random(users), post)
  end

  defp do_generate_threads(users) do
    thread_length = Enum.random([2, 3, 4, 5])
    actor = Enum.random(users)

    post = %{
      "status" => "Start of the thread"
    }

    {:ok, activity} = CommonAPI.post(actor, post)

    Enum.each(1..thread_length, fn _ ->
      user = Enum.random(users)

      post = %{
        "status" => "@#{actor.nickname} reply to thread",
        "in_reply_to_status_id" => activity.id
      }

      CommonAPI.post(user, post)
    end)
  end

  def generate_dms(user, users, opts) do
    IO.puts("Starting generating #{opts[:dms_max]} DMs")
    {time, _} = :timer.tc(fn -> do_generate_dms(user, users, opts) end)
    IO.puts("Inserting dms take #{to_sec(time)} sec.\n")
  end

  defp do_generate_dms(user, users, opts) do
    Task.async_stream(
      1..opts[:dms_max],
      fn _ ->
        do_generate_dm(user, users)
      end,
      max_concurrency: 10,
      timeout: 30_000
    )
    |> Stream.run()
  end

  defp do_generate_dm(user, users) do
    post = %{
      "status" => "@#{user.nickname} some direct message",
      "visibility" => "direct"
    }

    CommonAPI.post(Enum.random(users), post)
  end

  def generate_long_thread(user, users, opts) do
    IO.puts("Starting generating long thread with #{opts[:thread_length]} replies")
    {time, activity} = :timer.tc(fn -> do_generate_long_thread(user, users, opts) end)
    IO.puts("Inserting long thread replies take #{to_sec(time)} sec.\n")
    {:ok, activity}
  end

  defp do_generate_long_thread(user, users, opts) do
    {:ok, %{id: id} = activity} = CommonAPI.post(user, %{"status" => "Start of long thread"})

    Task.async_stream(
      1..opts[:thread_length],
      fn _ -> do_generate_thread(users, id) end,
      max_concurrency: 10,
      timeout: 30_000
    )
    |> Stream.run()

    activity
  end

  defp do_generate_thread(users, activity_id) do
    CommonAPI.post(Enum.random(users), %{
      "status" => "reply to main post",
      "in_reply_to_status_id" => activity_id
    })
  end

  def generate_non_visible_message(user, users) do
    IO.puts("Starting generating 1000 non visible posts")

    {time, _} =
      :timer.tc(fn ->
        do_generate_non_visible_posts(user, users)
      end)

    IO.puts("Inserting non visible posts take #{to_sec(time)} sec.\n")
  end

  defp do_generate_non_visible_posts(user, users) do
    [not_friend | users] = users

    make_friends(user, users)

    Task.async_stream(1..1000, fn _ -> do_generate_non_visible_post(not_friend, users) end,
      max_concurrency: 10,
      timeout: 30_000
    )
    |> Stream.run()
  end

  defp make_friends(_user, []), do: nil

  defp make_friends(user, [friend | users]) do
    {:ok, _} = User.follow(user, friend)
    {:ok, _} = User.follow(friend, user)
    make_friends(user, users)
  end

  defp do_generate_non_visible_post(not_friend, users) do
    post = %{
      "status" => "some non visible post",
      "visibility" => "private"
    }

    {:ok, activity} = CommonAPI.post(not_friend, post)

    thread_length = Enum.random([2, 3, 4, 5])

    Enum.each(1..thread_length, fn _ ->
      user = Enum.random(users)

      post = %{
        "status" => "@#{not_friend.nickname} reply to non visible post",
        "in_reply_to_status_id" => activity.id,
        "visibility" => "private"
      }

      CommonAPI.post(user, post)
    end)
  end
end