<commit_before>#include <limits.h>

#include "aio_private.h"

#define EVEN_DISTRIBUTE

const int MAX_BUF_REQS = 1024 * 3;

/* 
 * each file gets the same number of outstanding requests.
 */
#ifdef EVEN_DISTRIBUTE
#define MAX_OUTSTANDING_NREQS (AIO_DEPTH / num_open_files())
#define ALLOW_DROP
#endif

class extended_io_request: public io_request
{
	char *orig_embedded_bufs[NUM_EMBEDDED_IOVECS];
	char **orig_bufs;
	slab_allocator *allocator;

	void assign_extended(extended_io_request &req) {
		this->allocator = req.allocator;
		if (req.orig_bufs == req.orig_embedded_bufs) {
			this->orig_bufs = this->orig_embedded_bufs;
			memcpy(this->orig_embedded_bufs, req.orig_embedded_bufs,
					sizeof(orig_embedded_bufs));
		}
		else {
			this->orig_bufs = req.orig_bufs;
			req.orig_bufs = NULL;
		}
	}
public:
	extended_io_request() {
		orig_bufs = NULL;
		allocator = NULL;
		memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
	}

	/**
	 * The buffers in the IO request are allocated in a different NUMA node
	 * than the SSDs are connected to, and allocate buffers on the local node.
	 */
	extended_io_request(io_request &req, slab_allocator *allocator) {
		init(req, allocator);
	}

	/**
	 * The buffers are allocated in the same NUMA node as the SSDs
	 * are connected to.
	 */
	extended_io_request(io_request &req): io_request(req) {
		allocator = NULL;
		orig_bufs = NULL;
		memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
	}

	extended_io_request(extended_io_request &req) {
		io_request::assign(req);
		assign_extended(req);
	}

	~extended_io_request() {
		reset();
	}

	extended_io_request &operator=(io_request &req) {
		io_request::assign(req);
		orig_bufs = NULL;
		allocator = NULL;
		memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
		return *this;
	}

	extended_io_request &operator=(extended_io_request &req) {
		io_request::assign(req);
		assign_extended(req);
		return *this;
	}

	void reset();
	void init(io_request &req, slab_allocator *allocator);

	bool is_replaced() const {
		return orig_bufs != NULL;
	}

	void use_orig_bufs();
};

void extended_io_request::use_orig_bufs()
{
	for (int i = 0; i < get_num_bufs(); i++) {
		char *buf = get_buf(i);
		// This memory copy can significantly decrease the performance.
		// But it seems there isn't a better way to avoid it.
		memcpy(orig_bufs[i], buf, get_buf_size(i));
		set_buf(i, orig_bufs[i]);
		if (this->get_buf_size(i) <= PAGE_SIZE)
			allocator->free(&buf, 1);
		else
			free(buf);
	}
	// We have to reset orig_bufs because all original buffers
	// will be destroyed when the object is destructed.
	if (orig_bufs != orig_embedded_bufs)
		delete orig_bufs;
	orig_bufs = NULL;
}

void extended_io_request::init(io_request &req, slab_allocator *allocator)
{
	io_request::assign(req);
	this->allocator = allocator;
	memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
	if (this->get_num_bufs() > NUM_EMBEDDED_IOVECS)
		orig_bufs = new char *[this->get_num_bufs()];
	else
		orig_bufs = orig_embedded_bufs;
	for (int i = 0; i < this->get_num_bufs(); i++) {
		char *remote_buf = this->get_buf(i);
		char *local_buf;
		if (this->get_buf_size(i) <= PAGE_SIZE) {
			int ret = allocator->alloc(&local_buf, 1);
			assert(ret == 1);
		}
		else
			local_buf = (char *) valloc(this->get_buf_size(i));
		this->set_buf(i, local_buf);
		orig_bufs[i] = remote_buf;
	}
}

void extended_io_request::reset()<commit_msg>Fix a bug in aio for writing data to a remote SSD.<commit_after>#include <limits.h>

#include "aio_private.h"

#define EVEN_DISTRIBUTE

const int MAX_BUF_REQS = 1024 * 3;

/* 
 * each file gets the same number of outstanding requests.
 */
#ifdef EVEN_DISTRIBUTE
#define MAX_OUTSTANDING_NREQS (AIO_DEPTH / num_open_files())
#define ALLOW_DROP
#endif

class extended_io_request: public io_request
{
	char *orig_embedded_bufs[NUM_EMBEDDED_IOVECS];
	char **orig_bufs;
	slab_allocator *allocator;

	void assign_extended(extended_io_request &req) {
		this->allocator = req.allocator;
		if (req.orig_bufs == req.orig_embedded_bufs) {
			this->orig_bufs = this->orig_embedded_bufs;
			memcpy(this->orig_embedded_bufs, req.orig_embedded_bufs,
					sizeof(orig_embedded_bufs));
		}
		else {
			this->orig_bufs = req.orig_bufs;
			req.orig_bufs = NULL;
		}
	}
public:
	extended_io_request() {
		orig_bufs = NULL;
		allocator = NULL;
		memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
	}

	/**
	 * The buffers in the IO request are allocated in a different NUMA node
	 * than the SSDs are connected to, and allocate buffers on the local node.
	 */
	extended_io_request(io_request &req, slab_allocator *allocator) {
		init(req, allocator);
	}

	/**
	 * The buffers are allocated in the same NUMA node as the SSDs
	 * are connected to.
	 */
	extended_io_request(io_request &req): io_request(req) {
		allocator = NULL;
		orig_bufs = NULL;
		memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
	}

	extended_io_request(extended_io_request &req) {
		io_request::assign(req);
		assign_extended(req);
	}

	~extended_io_request() {
		reset();
	}

	extended_io_request &operator=(io_request &req) {
		io_request::assign(req);
		orig_bufs = NULL;
		allocator = NULL;
		memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
		return *this;
	}

	extended_io_request &operator=(extended_io_request &req) {
		io_request::assign(req);
		assign_extended(req);
		return *this;
	}

	void reset();
	void init(io_request &req, slab_allocator *allocator);

	bool is_replaced() const {
		return orig_bufs != NULL;
	}

	void use_orig_bufs();
};

void extended_io_request::use_orig_bufs()
{
	for (int i = 0; i < get_num_bufs(); i++) {
		char *buf = get_buf(i);
		// This memory copy can significantly decrease the performance.
		// But it seems there isn't a better way to avoid it.
		if (this->get_access_method() == READ)
			memcpy(orig_bufs[i], buf, get_buf_size(i));
		set_buf(i, orig_bufs[i]);
		if (this->get_buf_size(i) <= PAGE_SIZE)
			allocator->free(&buf, 1);
		else
			free(buf);
	}
	// We have to reset orig_bufs because all original buffers
	// will be destroyed when the object is destructed.
	if (orig_bufs != orig_embedded_bufs)
		delete orig_bufs;
	orig_bufs = NULL;
}

void extended_io_request::init(io_request &req, slab_allocator *allocator)
{
	io_request::assign(req);
	this->allocator = allocator;
	memset(orig_embedded_bufs, 0, sizeof(orig_embedded_bufs));
	if (this->get_num_bufs() > NUM_EMBEDDED_IOVECS)
		orig_bufs = new char *[this->get_num_bufs()];
	else
		orig_bufs = orig_embedded_bufs;
	for (int i = 0; i < this->get_num_bufs(); i++) {
		char *remote_buf = this->get_buf(i);
		char *local_buf;
		if (this->get_buf_size(i) <= PAGE_SIZE) {
			int ret = allocator->alloc(&local_buf, 1);
			assert(ret == 1);
		}
		else
			local_buf = (char *) valloc(this->get_buf_size(i));
		if (this->get_access_method() == WRITE)
			memcpy(local_buf, remote_buf, this->get_buf_size(i));
		this->set_buf(i, local_buf);
		orig_bufs[i] = remote_buf;
	}
}

void extended_io_request::reset()
{'fraction_non_alphanumeric': 0.12410637689448098, 'fraction_numerical': 0.004003431512725193, 'mean_word_length': 4.700896495517522, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 42, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7411805', 'n_tokens_mistral': 2898, 'n_tokens_neox': 2635, 'n_words': 585}