老司机们坐稳了 – 将Amazon EC2到Amazon S3的数据传输推向100Gbps线速

#include <fcntl.h>
#include <unistd.h>
#include <aws/auth/credentials.h>
#include <aws/common/condition_variable.h>
#include <aws/common/mutex.h>
#include <aws/common/zero.h>
#include <aws/io/channel_bootstrap.h>
#include <aws/io/event_loop.h>
#include <aws/io/logging.h>
#include <aws/http/request_response.h>
#include <aws/s3/s3.h>
#include <aws/s3/s3_client.h>
#include <aws/common/private/thread_shared.h>
#include <aws/common/clock.h>

#define TEST_REGION "ap-northeast-1"
#define TEST_S3_EP "testbucket.s3.ap-northeast-1.amazonaws.com"
#define TEST_MAIN_EP "s3.ap-northeast-1.amazonaws.com"
#define TEST_FILE "200G_test.file"
#define TEST_RESULT_FILE "test_stats.result"
#define TEST_RANGE_SIZE 8 * 1024 * 1024

struct perf_item {
    uint64_t ns;
    uint64_t len;
} __attribute__((packed, aligned(8)));

static const struct aws_byte_cursor g_host_header_name = AWS_BYTE_CUR_INIT_FROM_STRING_LITERAL("Host");

struct app_ctx {
    struct aws_allocator *allocator;
    struct aws_s3_client *client;
    struct aws_credentials_provider *credentials_provider;
    struct aws_client_bootstrap *client_bootstrap;
    struct aws_logger logger;
    struct aws_mutex mutex;
    struct aws_condition_variable c_var;
    bool execution_completed;
    struct aws_signing_config_aws signing_config;
    const char *region;
    enum aws_log_level log_level;
    bool help_requested;
    void *sub_command_data;
    size_t expected_transfers;
    size_t completed_transfers;
};

struct transfer_ctx {
    struct aws_s3_meta_request *meta_request;
    struct app_ctx *app_ctx;
    struct aws_atomic_var bytes_done;
    struct aws_atomic_var index;
    struct perf_item *stats;
};

void s_get_request_finished(
    struct aws_s3_meta_request *meta_request,
    const struct aws_s3_meta_request_result *meta_request_result,
    void *user_data) {

    struct transfer_ctx *transfer_ctx = user_data;
    struct perf_item *items = transfer_ctx->stats;
    int i = 0;

    uint32_t x = aws_atomic_load_int(&transfer_ctx->index);

    uint64_t total_ns = last - first;
    double total_ms = total_ns / 1000.0 / 1000.0;
    printf("first: %ld, last: %ld, total ms: %.2lf\n", first, last, total_ms);

    ssize_t total_bytes = aws_atomic_load_int(&transfer_ctx->bytes_done);

    double gbps = ((total_bytes / total_ms) * 1000) * 8 / 1000 / 1000 / 1000;
    printf("total bytes downloaded: %ld, line speed: %.2lf Gbps\n", total_bytes, gbps);

    /* write test stats */
    int fd;
    fd = open(TEST_RESULT_FILE, O_CREAT | O_TRUNC | O_WRONLY);
    write(fd, items, sizeof(struct perf_item) * x);
    close(fd);

    transfer_ctx->app_ctx->completed_transfers++;

    aws_condition_variable_notify_one(&transfer_ctx->app_ctx->c_var);
    aws_s3_meta_request_release(transfer_ctx->meta_request);
    aws_mem_release(transfer_ctx->app_ctx->allocator, transfer_ctx->stats);
    aws_mem_release(transfer_ctx->app_ctx->allocator, transfer_ctx);

    return;
}

int s_get_body_callback(
    struct aws_s3_meta_request *meta_request,
    const struct aws_byte_cursor *body,
    uint64_t range_start,
    void *user_data) {

    struct transfer_ctx *transfer_ctx = user_data;

    uint64_t now;
    size_t index;

    aws_high_res_clock_get_ticks(&now);
    index = aws_atomic_fetch_add(&transfer_ctx->index, 1);
    (transfer_ctx->stats)[index].ns = now;
    (transfer_ctx->stats)[index].len = body->len;

    aws_atomic_fetch_add(&transfer_ctx->bytes_done, body->len);

    return AWS_OP_SUCCESS;
}

bool s_are_all_transfers_done(void *arg) {

    struct app_ctx *app_ctx = arg;

    return app_ctx->expected_transfers == app_ctx->completed_transfers;
}

int s3_get(struct app_ctx *app_ctx, struct aws_s3_client *client) {
    char source_endpoint[1024];
    AWS_ZERO_ARRAY(source_endpoint);
    sprintf(source_endpoint, TEST_S3_EP);

    char main_endpoint[1024];
    AWS_ZERO_ARRAY(main_endpoint);
    sprintf(main_endpoint, TEST_MAIN_EP);

    struct aws_byte_cursor slash_cur = aws_byte_cursor_from_c_str("/");
    struct aws_byte_cursor keyname = aws_byte_cursor_from_c_str(TEST_FILE);
    struct aws_byte_cursor *key = &keyname;

    struct transfer_ctx *xfer_ctx = aws_mem_calloc(app_ctx->allocator, 1, sizeof(struct transfer_ctx));
    xfer_ctx->app_ctx = app_ctx;

    xfer_ctx->stats = aws_mem_calloc(app_ctx->allocator, 1000000, sizeof(struct perf_item));
    aws_atomic_init_int(&xfer_ctx->bytes_done, 0);
    aws_atomic_init_int(&xfer_ctx->index, 0);
    
    struct aws_s3_meta_request_options request_options = {
        .user_data = xfer_ctx,
        .signing_config = &app_ctx->signing_config,
        .type = AWS_S3_META_REQUEST_TYPE_GET_OBJECT,
        .finish_callback = s_get_request_finished,
        .body_callback = s_get_body_callback,
        .headers_callback = NULL,
        .shutdown_callback = NULL,
        .progress_callback = NULL,
    };

    struct aws_http_header host_header = {
        .name = g_host_header_name,
        .value = aws_byte_cursor_from_c_str(source_endpoint),
    };

    struct aws_http_header accept_header = {
        .name = aws_byte_cursor_from_c_str("accept"),
        .value = aws_byte_cursor_from_c_str("*/*"),
    };

    struct aws_http_header user_agent_header = {
        .name = aws_byte_cursor_from_c_str("user-agent"),
        .value = aws_byte_cursor_from_c_str("AWS common runtime command-line client"),
    };

    request_options.message = aws_http_message_new_request(app_ctx->allocator);
    aws_http_message_add_header(request_options.message, host_header);
    aws_http_message_add_header(request_options.message, accept_header);
    aws_http_message_add_header(request_options.message, user_agent_header);
    aws_http_message_set_request_method(request_options.message, aws_http_method_get);

    struct aws_byte_buf path_buf;
    aws_byte_buf_init(&path_buf, app_ctx->allocator, key->len + 1);
    aws_byte_buf_append_dynamic(&path_buf, &slash_cur);
    aws_byte_buf_append_dynamic(&path_buf, key);
    struct aws_byte_cursor path_cur = aws_byte_cursor_from_buf(&path_buf);
    aws_http_message_set_request_path(request_options.message, path_cur);
    aws_byte_buf_clean_up(&path_buf);

    struct aws_s3_meta_request *meta_request = aws_s3_client_make_meta_request(client, &request_options);
    xfer_ctx->meta_request = meta_request;

    aws_mutex_lock(&app_ctx->mutex);
    app_ctx->expected_transfers++;
    aws_mutex_unlock(&app_ctx->mutex);

    aws_mutex_lock(&app_ctx->mutex);
    aws_condition_variable_wait_pred(&app_ctx->c_var, &app_ctx->mutex, s_are_all_transfers_done, app_ctx);
    aws_mutex_unlock(&app_ctx->mutex);

    return 0;
}

int main() {
    struct aws_allocator *allocator = aws_default_allocator();
    aws_s3_library_init(allocator);

    struct app_ctx app_ctx;
    AWS_ZERO_STRUCT(app_ctx);
    app_ctx.allocator = allocator;
    app_ctx.c_var = (struct aws_condition_variable)AWS_CONDITION_VARIABLE_INIT;
    aws_mutex_init(&app_ctx.mutex);
    app_ctx.expected_transfers = 0;
    app_ctx.completed_transfers = 0;

    app_ctx.log_level = AWS_LOG_LEVEL_NONE;

    struct aws_logger_standard_options logger_options = {
        .level = app_ctx.log_level,
        .file = stderr,
    };

    aws_logger_init_standard(&app_ctx.logger, app_ctx.allocator, &logger_options);
    aws_logger_set(&app_ctx.logger);

    /* event loop */
    struct aws_event_loop_group *event_loop_group = aws_event_loop_group_new_default(allocator, 0, NULL);

    /* resolver */
    struct aws_host_resolver_default_options resolver_options = {
        .el_group = event_loop_group,
        .max_entries = 128,
    };
    struct aws_host_resolver *resolver = aws_host_resolver_new_default(allocator, &resolver_options);

    /* client bootstrap */
    struct aws_client_bootstrap_options bootstrap_options = {
        .event_loop_group = event_loop_group,
        .host_resolver = resolver,
    };
    app_ctx.client_bootstrap = aws_client_bootstrap_new(allocator, &bootstrap_options);
    if (app_ctx.client_bootstrap == NULL) {
        printf("ERROR initializing client bootstrap\n");
        return -1;
    }

    /* credentials */
    struct aws_credentials_provider_chain_default_options credentials_provider_options;
    AWS_ZERO_STRUCT(credentials_provider_options);
    credentials_provider_options.bootstrap = app_ctx.client_bootstrap;
    app_ctx.credentials_provider = aws_credentials_provider_new_chain_default(allocator, &credentials_provider_options);

    /* signing config */
    aws_s3_init_default_signing_config(
        &app_ctx.signing_config, aws_byte_cursor_from_c_str(TEST_REGION), app_ctx.credentials_provider);
    app_ctx.signing_config.flags.use_double_uri_encode = false;

    /* s3 client */
    struct aws_s3_client_config client_config;
    AWS_ZERO_STRUCT(client_config);
    client_config.client_bootstrap = app_ctx.client_bootstrap;
    client_config.region = aws_byte_cursor_from_c_str(TEST_REGION);
    client_config.signing_config = &app_ctx.signing_config;
    client_config.max_active_connections_override = 0;
    client_config.throughput_target_gbps = 100;
    client_config.part_size = TEST_RANGE_SIZE;

    struct aws_s3_client *client = aws_s3_client_new(app_ctx.allocator, &client_config);
    app_ctx.client = client;
    s3_get(&app_ctx, client);

    /* release resources */
    aws_s3_client_release(app_ctx.client);
    aws_credentials_provider_release(app_ctx.credentials_provider);
    aws_client_bootstrap_release(app_ctx.client_bootstrap);
    aws_host_resolver_release(resolver);
    aws_event_loop_group_release(event_loop_group);
    aws_mutex_clean_up(&app_ctx.mutex);

    aws_s3_library_clean_up();

    return 0;
}

[ec2-user@ip-172-31-32-214 clang]$ ./s3bench
first: 30633404368026, last: 30653346831918, total ms: 19942.46
total bytes downloaded: 214748364800, line speed: 86.15 Gbps

[ec2-user@ip-172-31-40-21 ~]$ aws s3 cp s3://testbucket/200_test.file  /mnt/200_test.file
Completed 29.6 GiB/200.0 GiB (2.0 GiB/s) with 1 file(s) remaining

(gdb) info thread
  Id   Target Id         Frame
* 1    Thread 0x7f467761e740 (LWP 18862) 0x00007f467645ea46 in do_futex_wait.constprop ()
   from /lib64/libpthread.so.0
  2    Thread 0x7f45d8dfa700 (LWP 18904) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  3    Thread 0x7f45661fc700 (LWP 18926) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  4    Thread 0x7f45da1fc700 (LWP 18902) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  5    Thread 0x7f46561fc700 (LWP 18873) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  6    Thread 0x7f4657fff700 (LWP 18870) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  7    Thread 0x7f45d97fb700 (LWP 18903) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  8    Thread 0x7f4664e86700 (LWP 18868) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  9    Thread 0x7f46161fc700 (LWP 18890) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  10   Thread 0x7f45975fe700 (LWP 18912) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  11   Thread 0x7f4596bfd700 (LWP 18913) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  12   Thread 0x7f45db5fe700 (LWP 18900) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  13   Thread 0x7f44c21fc700 (LWP 18962) 0x00007f467645ea46 in do_futex_wait.constprop ()
   from /lib64/libpthread.so.0
  14   Thread 0x7f451d7fb700 (LWP 18945) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  15   Thread 0x7f461ffff700 (LWP 18882) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  16   Thread 0x7f453d7fb700 (LWP 18939) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  17   Thread 0x7f457d7fb700 (LWP 18921) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  18   Thread 0x7f464f5fe700 (LWP 18876) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  19   Thread 0x7f45957fb700 (LWP 18915) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  20   Thread 0x7f4614dfa700 (LWP 18892) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6
  21   Thread 0x7f453ebfd700 (LWP 18937) 0x00007f4676d3088c in epoll_wait () from /lib64/libc.so.6

(gdb) thread 1
[Switching to thread 1 (Thread 0x7f467761e740 (LWP 18862))]
#0  0x00007f467645ea46 in do_futex_wait.constprop () from /lib64/libpthread.so.0
(gdb) bt
#0  0x00007f467645ea46 in do_futex_wait.constprop () from /lib64/libpthread.so.0
#1  0x00007f467645eb22 in __new_sem_wait_slow.constprop.0 () from /lib64/libpthread.so.0
#2  0x00007f46768812ff in PyThread_acquire_lock_timed (lock=lock@entry=0x561a67fcf390,
    microseconds=microseconds@entry=-1000000, intr_flag=intr_flag@entry=1) at Python/thread_pthread.h:483
#3  0x00007f46768dcaf4 in acquire_timed (timeout=-1000000000, lock=0x561a67fcf390)
    at ./Modules/_threadmodule.c:63
#4  lock_PyThread_acquire_lock (self=0x7f4664407ea0, args=<optimized out>, kwds=<optimized out>)
    at ./Modules/_threadmodule.c:146
#5  0x00007f46767368a2 in method_vectorcall_VARARGS_KEYWORDS (func=0x7f46775c92c0, args=0x561a67ee7f48,
    nargsf=<optimized out>, kwnames=<optimized out>) at Objects/descrobject.c:348
#6  0x00007f46766d80a8 in _PyObject_VectorcallTstate (kwnames=<optimized out>, nargsf=<optimized out>,
    args=<optimized out>, callable=<optimized out>, tstate=<optimized out>)
    at ./Include/cpython/abstract.h:118
#7  PyObject_Vectorcall (kwnames=<optimized out>, nargsf=<optimized out>, args=<optimized out>,
    callable=<optimized out>) at ./Include/cpython/abstract.h:127
#8  trace_call_function (kwnames=<optimized out>, nargs=<optimized out>, args=<optimized out>,
    func=<optimized out>, tstate=<optimized out>) at Python/ceval.c:5058
#9  call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x561a674afe40)
    at Python/ceval.c:5074
#10 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=<optimized out>, throwflag=<optimized out>)
    at Python/ceval.c:3506
#11 0x00007f467682aec1 in _PyEval_EvalFrame (throwflag=0, f=0x561a67ee7db0, tstate=0x561a674afe40)
    at ./Include/internal/pycore_ceval.h:40

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_s3::{Client, Region, PKG_VERSION};
use aws_sdk_s3::Endpoint;
use http::Uri;
use tokio::runtime::{Builder, Runtime};
use tokio::fs::File;
use tokio::io::AsyncWriteExt;
use tokio::io::AsyncSeekExt;
use tokio::io::SeekFrom;
use tokio::fs::OpenOptions;
use tokio::time::{Duration, Instant};
use std::error::Error;
use indicatif::{ProgressBar, ProgressStyle};

macro_rules! thread_tasks {
    ($tasks:expr, $manager:expr) => {{
        let mut _tasks = Vec::new();
        for i in 0..$manager.threads {
            let mut thread_tasks = Vec::new();
            _tasks.push(thread_tasks);
        }
        let mut pos = 0;
        for task in $tasks {
            let i = pos % $manager.threads;
            _tasks[i].push(task);
            pos += 1;
        }
        _tasks
    }}
}

type Range = (isize, isize);

type GetTask = Range;

#[derive(Clone)]
struct TransferManager {
    s3: Client,
    threads: usize,
    chunks: isize,
    bucket: String,
    key: String,
    filename: String,
    objsz: isize,
    forcewrite: bool,
    bar: ProgressBar,
}

impl TransferManager {

    pub async fn new(is_get: bool, forcewrite: bool, filename: String, bucket: String, key: String, thread_cnt: usize, chunks: isize, endpoint: Option<String>) -> Self {

        let region_provider = RegionProviderChain::default_provider().or_else("ap-northeast-1");
        let shared_config = aws_config::from_env().region(region_provider).load().await;

        let client = if endpoint.is_some() {
            let s3_config = aws_sdk_s3::config::Builder::from(&shared_config)
                .endpoint_resolver(
                    Endpoint::immutable(Uri::try_from(endpoint.unwrap()).expect("error: bad endpoint")),
                )
                .build();
            Client::from_conf(s3_config)
        } else {
            Client::new(&shared_config)
        };

        let mut manager = Self {
            s3: client,
            threads: thread_cnt,
            chunks: chunks,
            bucket: bucket.clone(),
            key: key.clone(),
            filename: filename.clone(),
            objsz: 0,
            forcewrite: forcewrite,
            bar: ProgressBar::new(0),
        };

        let mut objsz = 0;
        if is_get {
            objsz = manager.head_object().await;
            if forcewrite {
                let mut f = File::create(manager.filename.clone()).await.expect("failed to create file");
                f.set_len(objsz as u64).await.expect("failed to set local file len");
            }
        } else {
            objsz = manager.get_file_size().await;
        }

        let bar = ProgressBar::new(objsz as u64);
        bar.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta}) ({bytes_per_sec})")
            .unwrap()
            .progress_chars("#>-"));

        manager.objsz = objsz;
        manager.bar = bar;
        manager
    }

    pub fn alloc_get_tasks(&self, tasks: &mut Vec<GetTask>) {

        let mut remains = self.objsz;
        let mut offset = 0;
        while remains > 0 {
            let mut range: Range = (0, 0);
            if remains > self.chunks {
                range = (offset, offset + self.chunks - 1);
            } else {
                range = (offset, offset + remains - 1);
            }
            let task = range;
            tasks.push(task);
            offset += self.chunks;
            remains -= self.chunks;
        }
    }
    
    pub async fn do_get_tasks(&self) {

        let mut tasks = Vec::new();

        self.alloc_get_tasks(&mut tasks);
        let mut thr_tasks: Vec<Vec<GetTask>> = thread_tasks!(tasks, self);

        let mut joins = Vec::new();
        while let Some(mut thr_taskq) = thr_tasks.pop() {
            let mut m = self.clone();
            let bar = self.bar.clone();
            let join = tokio::spawn(async move {
                while let Some(task) = thr_taskq.pop() {
                    let res = m.download_object(task).await;
                    if res.is_err() {
                        println!("download error: {:?}", res);
                    }
                    bar.inc(res.unwrap() as u64);
                }
                0
            });
            joins.push(join);
        }

        let instant = Instant::now();
        self.bar.reset();
        for join in joins {
            join.await;
        }

        let diff = instant.elapsed();
        let bw = ((self.objsz as u128 / diff.as_millis()) as f64 / (1024.0 * 1024.0 * 1024.0)) * 1000.0;
        self.bar.finish();
    }

    pub async fn download_object(&mut self, range: GetTask) -> Result<isize, Box<dyn Error>> {
        let resp = self.s3
            .get_object()
            .bucket(&self.bucket)
            .key(&self.key)
            .range(format!("bytes={}-{}", range.0, range.1))
            .send()
            .await?;
        let data = resp.body.collect().await.expect("error reading data");
        let mut bytes = data.into_bytes();

        if !self.forcewrite {
            return Ok(bytes.len() as isize);
        }

        let mut file = OpenOptions::new()
            .write(true)
            .open(self.filename.clone())
            .await?;

        file.seek(SeekFrom::Start(range.0 as u64)).await?;
        file.write_all(&bytes).await?;

        Ok(bytes.len() as isize)
    }

    pub async fn head_object(&self) -> isize {
        let resp = self.s3
            .head_object()
            .bucket(&self.bucket)
            .key(&self.key)
            .send()
            .await
            .expect("failed to get object size from S3");
        resp.content_length() as isize
    }

    pub async fn get_file_size(&self) -> isize {
        let attr = tokio::fs::metadata(self.filename.clone()).await.expect("failed to get file size");
        attr.len() as isize
    }
}

fn main() -> Result<(), Box<dyn Error>> {

    let is_get = true;
    let mut region = "ap-northeast-1".to_string();
    let mut bucket = "testbucket".to_string();
    let mut object = "200G_test.file".to_string();
    let mut filename = "200G_test.file".to_string();
    let mut forcewrite = false;
    let mut threadpool = 30;
    let mut parallel_tasks = 300;
    let mut chunksize = 8 * 1024 * 1024;
    let mut endpoint = None;

    let rt = Builder::new_multi_thread()
        .worker_threads(threadpool)
        .enable_all()
        .build()
        .unwrap();

    rt.block_on(async {
        let mut tm = TransferManager::new(is_get, forcewrite, filename, bucket, object, parallel_tasks, chunksize, endpoint).await;

        tm.do_get_tasks().await;
    });
    Ok(())
}

  [ec2-user@ip-172-31-32-214 src]$./s3perf
  [00:00:24] [#######################################################################] 200.00 GiB/200.00 GiB (0s) (8.24 GiB/s)
  [ec2-user@ip-172-31-32-214 src]$