diff --git a/f3read.c b/f3read.c index b060662..12686f9 100644 --- a/f3read.c +++ b/f3read.c @@ -191,15 +191,16 @@ static ssize_t read_all(int fd, char *buf, size_t count) return done; } -static ssize_t check_chunk(int fd, uint64_t *p_expected_offset, - uint64_t chunk_size, struct file_stats *stats) +static ssize_t check_chunk(struct dynamic_buffer *dbuf, int fd, + uint64_t *p_expected_offset, uint64_t chunk_size, + struct file_stats *stats) { - char buf[MAX_BUFFER_SIZE]; + char *buf = dbuf_get_buf(dbuf, chunk_size); + size_t len = dbuf_get_len(dbuf); ssize_t tot_bytes_read = 0; while (chunk_size > 0) { - size_t turn_size = chunk_size <= MAX_BUFFER_SIZE - ? chunk_size : MAX_BUFFER_SIZE; + size_t turn_size = chunk_size <= len ? chunk_size : len; ssize_t bytes_read = read_all(fd, buf, turn_size); if (bytes_read < 0) { @@ -235,6 +236,7 @@ static void validate_file(const char *path, int number, struct flow *fw, int fd, saved_errno; ssize_t bytes_read; uint64_t expected_offset; + struct dynamic_buffer dbuf; zero_fstats(stats); @@ -263,11 +265,12 @@ static void validate_file(const char *path, int number, struct flow *fw, /* Help the kernel to help us. */ assert(!posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL)); + dbuf_init(&dbuf); saved_errno = 0; expected_offset = (uint64_t)number * GIGABYTES; start_measurement(fw); while (true) { - bytes_read = check_chunk(fd, &expected_offset, + bytes_read = check_chunk(&dbuf, fd, &expected_offset, get_rem_chunk_size(fw), stats); if (bytes_read == 0) break; @@ -297,6 +300,7 @@ static void validate_file(const char *path, int number, struct flow *fw, } printf("\n"); + dbuf_free(&dbuf); close(fd); free(full_fn); } diff --git a/f3write.c b/f3write.c index 8942dc5..112e7c6 100644 --- a/f3write.c +++ b/f3write.c @@ -150,13 +150,14 @@ static int write_all(int fd, const char *buf, size_t count) return 0; } -static int write_chunk(int fd, size_t chunk_size, uint64_t *poffset) +static int write_chunk(struct dynamic_buffer *dbuf, int fd, size_t chunk_size, + uint64_t *poffset) { - char buf[MAX_BUFFER_SIZE]; + char *buf = dbuf_get_buf(dbuf, chunk_size); + size_t len = dbuf_get_len(dbuf); while (chunk_size > 0) { - size_t turn_size = chunk_size <= MAX_BUFFER_SIZE - ? chunk_size : MAX_BUFFER_SIZE; + size_t turn_size = chunk_size <= len ? chunk_size : len; int ret; chunk_size -= turn_size; *poffset = fill_buffer(buf, turn_size, *poffset); @@ -177,6 +178,7 @@ static int create_and_fill_file(const char *path, long number, size_t size, int fd, saved_errno; size_t remaining; uint64_t offset; + struct dynamic_buffer dbuf; assert(size > 0); assert(size % fw->block_size == 0); @@ -198,6 +200,7 @@ static int create_and_fill_file(const char *path, long number, size_t size, assert(fd >= 0); /* Write content. */ + dbuf_init(&dbuf); saved_errno = 0; offset = (uint64_t)number * GIGABYTES; remaining = size; @@ -206,7 +209,7 @@ static int create_and_fill_file(const char *path, long number, size_t size, uint64_t write_size = get_rem_chunk_size(fw); if (write_size > remaining) write_size = remaining; - saved_errno = write_chunk(fd, write_size, &offset); + saved_errno = write_chunk(&dbuf, fd, write_size, &offset); if (saved_errno) break; remaining -= write_size; @@ -220,6 +223,7 @@ static int create_and_fill_file(const char *path, long number, size_t size, if (!saved_errno) saved_errno = errno; } + dbuf_free(&dbuf); close(fd); free(full_fn); diff --git a/libflow.c b/libflow.c index 9ddc81a..52efb53 100644 --- a/libflow.c +++ b/libflow.c @@ -366,3 +366,45 @@ int end_measurement(int fd, struct flow *fw) } return ret; } + +static inline void __dbuf_free(struct dynamic_buffer *dbuf) +{ + if (dbuf->buf != dbuf->backup_buf) + free(dbuf->buf); +} + +void dbuf_free(struct dynamic_buffer *dbuf) +{ + __dbuf_free(dbuf); + dbuf->buf = NULL; + dbuf->len = 0; + dbuf->max_buf = true; +} + +char *dbuf_get_buf(struct dynamic_buffer *dbuf, size_t size) +{ + /* If enough buffer, or it's already the largest buffer, return it. */ + if (size <= dbuf->len || dbuf->max_buf) + return dbuf->buf; + + /* + * Allocate a new buffer. + */ + + __dbuf_free(dbuf); + do { + dbuf->buf = malloc(size); + if (dbuf->buf != NULL) { + dbuf->len = size; + return dbuf->buf; + } else { + dbuf->max_buf = true; + } + size /= 2; + } while (size > sizeof(dbuf->backup_buf)); + + /* A larger buffer is not available; failsafe. */ + dbuf->buf = dbuf->backup_buf; + dbuf->len = sizeof(dbuf->backup_buf); + return dbuf->buf; +} diff --git a/libflow.h b/libflow.h index bce2f28..272a5de 100644 --- a/libflow.h +++ b/libflow.h @@ -2,6 +2,7 @@ #define HEADER_LIBFLOW_H #include +#include struct flow; @@ -90,6 +91,31 @@ static inline uint64_t get_rem_chunk_size(struct flow *fw) return (fw->blocks_per_delay - fw->processed_blocks) * fw->block_size; } -#define MAX_BUFFER_SIZE (1<<21) /* 2MB */ +struct dynamic_buffer { + char *buf; + size_t len; + bool max_buf; + char backup_buf[1 << 21]; /* 2MB */ +}; + +static inline void dbuf_init(struct dynamic_buffer *dbuf) +{ + dbuf->buf = dbuf->backup_buf; + dbuf->len = sizeof(dbuf->backup_buf); + dbuf->max_buf = false; +} + +void dbuf_free(struct dynamic_buffer *dbuf); + +/* + * Although the returned buffer may be smaller than @size, + * this function never returns NULL. + */ +char *dbuf_get_buf(struct dynamic_buffer *dbuf, size_t size); + +static inline size_t dbuf_get_len(const struct dynamic_buffer *dbuf) +{ + return dbuf->len; +} #endif /* HEADER_LIBFLOW_H */