diff --git a/include/aws/s3/s3_client.h b/include/aws/s3/s3_client.h index 1a07f8564..93fd0d4a3 100644 --- a/include/aws/s3/s3_client.h +++ b/include/aws/s3/s3_client.h @@ -1034,6 +1034,8 @@ struct aws_s3_meta_request_options { struct aws_byte_cursor copy_source_uri; }; +typedef int(aws_s3_user_buffer_before_register_fn)(void *user_data); + /* EXPERIMENTAL - User-provided buffer options for GET/PUT operations */ struct aws_s3_user_buffer_options { /* Unified transfer buffer: PUT reads from it; GET writes into it. */ @@ -1046,6 +1048,7 @@ struct aws_s3_user_buffer_options { * is already pinned and suitable for RDMA transfers. */ bool buffer_is_rdma_registered; void *_reserved[4]; + aws_s3_user_buffer_before_register_fn *before_register_callback; }; /* Result details of a meta request. diff --git a/source/s3_rdma_buffer_manager.c b/source/s3_rdma_buffer_manager.c index b88a30018..758bf8d2c 100644 --- a/source/s3_rdma_buffer_manager.c +++ b/source/s3_rdma_buffer_manager.c @@ -82,6 +82,22 @@ static int s_default_prepare_buffer_for_rdma( AWS_LOGF_DEBUG(AWS_LS_S3_RDMA, "id=%p Buffer not suitable for RDMA", (void *)manager); return AWS_OP_SUCCESS; } + + /* Run before register callback */ + struct aws_s3_meta_request *meta_request = request->meta_request; + if (meta_request->user_buffer_options.before_register_callback) { + int before_result = meta_request->user_buffer_options.before_register_callback(meta_request->user_data); + AWS_LOGF_DEBUG(AWS_LS_S3_RDMA, "id=%p Run before register function", (void *)manager); + if (before_result != AWS_OP_SUCCESS) { + AWS_LOGF_ERROR( + AWS_LS_S3_META_REQUEST, + "id=%p Part %u: Before register function failed (result=%d)", + (void *)meta_request, + request->part_number, + before_result); + return before_result; + } + } /* Register buffer with RDMA provider */ int register_result = aws_s3_rdma_provider_register_memory(manager->rdma_provider, buffer, buffer_size);