1#include <iostream>
2#include <chrono>
3#include <numeric>
4#include <sys/types.h>
5#include <sys/stat.h>
6#include <fcntl.h>
7#include <linux/fs.h>
8#include <unistd.h>
9#include <sys/swap.h>
10
11using namespace std;
12
13const char zram_blkdev_path[] = "/dev/block/zram0";
14const size_t sector_size = 512;
15const size_t page_size = 4096;
16
17void fillPageRand(uint32_t *page) {
18    int start = rand();
19    for (int i = 0; i < page_size / sizeof(int); i++) {
20        page[i] = start+i;
21    }
22}
23void fillPageCompressible(uint32_t *page) {
24    int val = rand() & 0xfff;
25    for (int i = 0; i < page_size / sizeof(int); i++) {
26        page[i] = val;
27    }
28}
29
30class AlignedAlloc {
31    void *m_ptr;
32public:
33    AlignedAlloc(size_t size, size_t align) {
34        posix_memalign(&m_ptr, align, size);
35    }
36    ~AlignedAlloc() {
37        free(m_ptr);
38    }
39    void *ptr() {
40        return m_ptr;
41    }
42};
43
44class BlockFd {
45    int m_fd = -1;
46public:
47    BlockFd(const char *path, bool direct) {
48        m_fd = open(path, O_RDWR | (direct ? O_DIRECT : 0));
49    }
50    size_t getSize() {
51        size_t blockSize = 0;
52        int result = ioctl(m_fd, BLKGETSIZE, &blockSize);
53        if (result < 0) {
54            cout << "ioctl failed" << endl;
55        }
56        return blockSize * sector_size;
57    }
58    ~BlockFd() {
59        if (m_fd >= 0) {
60            close(m_fd);
61        }
62    }
63    void fillWithCompressible() {
64        size_t devSize = getSize();
65        AlignedAlloc page(page_size, page_size);
66        for (uint64_t offset = 0; offset < devSize; offset += page_size) {
67            fillPageCompressible((uint32_t*)page.ptr());
68            ssize_t ret = write(m_fd, page.ptr(), page_size);
69            if (ret != page_size) {
70                cout << "write() failed" << endl;
71            }
72        }
73    }
74    void benchSequentialRead() {
75        chrono::time_point<chrono::high_resolution_clock> start, end;
76        size_t devSize = getSize();
77        size_t passes = 4;
78        AlignedAlloc page(page_size, page_size);
79
80        start = chrono::high_resolution_clock::now();
81        for (int i = 0; i < passes; i++) {
82            for (uint64_t offset = 0; offset < devSize; offset += page_size) {
83                if (offset == 0)
84                    lseek(m_fd, offset, SEEK_SET);
85                ssize_t ret = read(m_fd, page.ptr(), page_size);
86                if (ret != page_size) {
87                    cout << "read() failed" << endl;
88                }
89            }
90        }
91        end = chrono::high_resolution_clock::now();
92        size_t duration = chrono::duration_cast<chrono::microseconds>(end - start).count();
93        cout << "read: " << (double)devSize * passes / 1024.0 / 1024.0 / (duration / 1000.0 / 1000.0) << "MB/s" << endl;
94    }
95    void benchSequentialWrite() {
96        chrono::time_point<chrono::high_resolution_clock> start, end;
97        size_t devSize = getSize();
98        size_t passes = 4;
99        AlignedAlloc page(page_size, page_size);
100
101        start = chrono::high_resolution_clock::now();
102        for (int i = 0; i < passes; i++) {
103            for (uint64_t offset = 0; offset < devSize; offset += page_size) {
104                fillPageCompressible((uint32_t*)page.ptr());
105                if (offset == 0)
106                    lseek(m_fd, offset, SEEK_SET);
107                ssize_t ret = write(m_fd, page.ptr(), page_size);
108                if (ret != page_size) {
109                    cout << "write() failed" << endl;
110                }
111            }
112        }
113        end = chrono::high_resolution_clock::now();
114        size_t duration = chrono::duration_cast<chrono::microseconds>(end - start).count();
115        cout << "write: " << (double)devSize * passes / 1024.0 / 1024.0 / (duration / 1000.0 / 1000.0) << "MB/s" << endl;
116
117    }
118};
119
120int bench(bool direct)
121{
122    BlockFd zramDev{zram_blkdev_path, direct};
123
124    zramDev.fillWithCompressible();
125    zramDev.benchSequentialRead();
126    zramDev.benchSequentialWrite();
127    return 0;
128}
129
130int main(int argc, char *argv[])
131{
132    int result = swapoff(zram_blkdev_path);
133    if (result < 0) {
134        cout << "swapoff failed: " << strerror(errno) << endl;
135    }
136
137    bench(1);
138
139    result = system((string("mkswap ") + string(zram_blkdev_path)).c_str());
140    if (result < 0) {
141        cout << "mkswap failed: " <<  strerror(errno) << endl;
142        return -1;
143    }
144
145    result = swapon(zram_blkdev_path, 0);
146    if (result < 0) {
147        cout << "swapon failed: " <<  strerror(errno) << endl;
148        return -1;
149    }
150    return 0;
151}
152