source/transfcount/transfcount.c

/**
 * Measure data transferred by a process.
 *
 * Used with LD_PRELOAD.
 */

#define _GNU_SOURCE
#include <dlfcn.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>

static unsigned long long total_sent = 0;
static unsigned long long total_recv = 0;
static time_t start_time;

static void transfcount_init(void);
static void transfcount_atexit(void);
static void transfcount_print_bytes(unsigned long long n);

typedef ssize_t (*send_type)(int sockfd, const void *buf, size_t len,
                int flags);
typedef int (*sendmmsg_type)(int sockfd, struct mmsghdr *msgvec,
                unsigned int vlen, int flags);
typedef ssize_t (*recv_type)(int sockfd, void *buf, size_t len, int flags);
typedef ssize_t (*read_type)(int fd, void *buf, size_t count);
typedef ssize_t (*write_type)(int fd, const void *buf, size_t count);

static send_type orig_send;
static sendmmsg_type orig_sendmmsg;
static recv_type orig_recv;
static read_type orig_read;
static write_type orig_write;


/**
 * send - send a message on a socket.
 */
ssize_t send(int sockfd, const void *buf, size_t len, int flags)
{
        ssize_t bytes;

        transfcount_init();

        /*printf("######## send(sockfd=%i, buf=%p, len=%zi, flags=%x) -> ",
                        sockfd, buf, len, flags);*/

        bytes = orig_send(sockfd, buf, len, flags);
        total_sent += bytes;
        /*printf("%lu\n", bytes);*/

        return bytes;
}


/**
 * sendmmsg - send multiple messages on a socket.
 */
int sendmmsg(int sockfd, struct mmsghdr *msgvec, unsigned int vlen,
                    int flags)
{
        int sent, i;

        transfcount_init();

        /*printf("######## sendmmsg(sockfd=%i, msgvec=%p, vlen=%u, flags=%x) -> ",
                        sockfd, msgvec, vlen, flags);*/

        sent = orig_sendmmsg(sockfd, msgvec, vlen, flags);
        /*printf("%i\n", sent);*/
        for (i = 0; i < sent; i++)
                total_sent += msgvec[i].msg_len;

        return sent;
}


/**
 * recv - receive a message from a socket.
 */
ssize_t recv(int sockfd, void *buf, size_t len, int flags)
{
        ssize_t bytes;

        transfcount_init();

        /*printf("######## recv(sockfd=%i, buf=%p, len=%zi, flags=%x) -> ",
                        sockfd, buf, len, flags);*/

        bytes = orig_recv(sockfd, buf, len, flags);
        total_recv += bytes;
        /*printf("%lu\n", bytes);*/

        return bytes;
}


/**
 * read - read from a file descriptor.
 */
ssize_t read(int fd, void *buf, size_t count)
{
        struct stat statbuf;
        ssize_t bytes;

        transfcount_init();

        fstat(fd, &statbuf);
        if (!S_ISSOCK(statbuf.st_mode))
                return orig_read(fd, buf, count);

        /*printf("######## read(fd=%i, buf=%p, count=%zi) -> ",
                        fd, buf, count);*/

        bytes = orig_read(fd, buf, count);
        if (bytes != -1)
                total_recv += bytes;
        /*printf("%lu\n", bytes);*/

        return bytes;
}


/**
 * write - write to a file descriptor.
 */
ssize_t write(int fd, const void *buf, size_t count)
{
        struct stat statbuf;
        ssize_t bytes;

        transfcount_init();

        fstat(fd, &statbuf);
        if (!S_ISSOCK(statbuf.st_mode))
                return orig_write(fd, buf, count);

        /*printf("######## write(fd=%i, buf=%p, count=%zi) -> ",
                        fd, buf, count);*/

        bytes = orig_write(fd, buf, count);
        if (bytes != -1)
                total_sent += bytes;
        /*printf("%lu\n", bytes);*/

        return bytes;
}


/**
 * Set up library.
 */
void transfcount_init(void)
{
        static bool init_done = false;

        if (init_done)
                return;

        orig_send = (send_type)dlsym(RTLD_NEXT, "send");
        orig_sendmmsg = (sendmmsg_type)dlsym(RTLD_NEXT, "sendmmsg");
        orig_recv = (recv_type)dlsym(RTLD_NEXT, "recv");
        orig_read = (read_type)dlsym(RTLD_NEXT, "read");
        orig_write = (write_type)dlsym(RTLD_NEXT, "write");

        atexit(transfcount_atexit);

        start_time = time(NULL);

        init_done = true;
}


/**
 * Output transfer statistics.
 */
void transfcount_atexit(void)
{
        time_t end_time = time(NULL);
        time_t duration = end_time - start_time;

        if (total_sent == 0 && total_recv == 0)
                return;

        if (duration == 0)
                duration = 1;

        fprintf(stderr, "[>>> Sent ");
        transfcount_print_bytes(total_sent);
        fprintf(stderr, "B (");
        transfcount_print_bytes(total_sent / duration);
        fprintf(stderr, "B/s) >>>] [<<< Received ");
        transfcount_print_bytes(total_recv);
        fprintf(stderr, "B (");
        transfcount_print_bytes(total_recv / duration);
        fprintf(stderr, "B/s) <<<]\n");
}


/**
 * Print a size as human readable.
 */
void transfcount_print_bytes(unsigned long long n)
{
        if (1024 * 1024 * 1024 <= n) {
                fprintf(stderr, "%.1fGi", (double) n / 1024 / 1024 / 1024);
        } else if (1024 * 1024 <= n) {
                fprintf(stderr, "%.1fMi", (double) n / 1024 / 1024);
        } else if (1024 <= n) {
                fprintf(stderr, "%.1fKi", (double) n / 1024);
        } else {
                fprintf(stderr, "%llu", n);
        }
}

1	james	88	/**
2			* Measure data transferred by a process.
3			*
4			* Used with LD_PRELOAD.
5			*/
6
7			#define _GNU_SOURCE
8			#include <dlfcn.h>
9			#include <stdbool.h>
10			#include <stdio.h>
11			#include <stdlib.h>
12			#include <time.h>
13			#include <unistd.h>
14			#include <sys/types.h>
15			#include <sys/socket.h>
16			#include <sys/stat.h>
17
18			static unsigned long long total_sent = 0;
19			static unsigned long long total_recv = 0;
20			static time_t start_time;
21
22			static void transfcount_init(void);
23			static void transfcount_atexit(void);
24			static void transfcount_print_bytes(unsigned long long n);
25
26			typedef ssize_t (send_type)(int sockfd, const void buf, size_t len,
27			int flags);
28			typedef int (sendmmsg_type)(int sockfd, struct mmsghdr msgvec,
29			unsigned int vlen, int flags);
30			typedef ssize_t (recv_type)(int sockfd, void buf, size_t len, int flags);
31			typedef ssize_t (read_type)(int fd, void buf, size_t count);
32			typedef ssize_t (write_type)(int fd, const void buf, size_t count);
33
34			static send_type orig_send;
35			static sendmmsg_type orig_sendmmsg;
36			static recv_type orig_recv;
37			static read_type orig_read;
38			static write_type orig_write;
39
40
41			/**
42			* send - send a message on a socket.
43			*/
44			ssize_t send(int sockfd, const void *buf, size_t len, int flags)
45			{
46			ssize_t bytes;
47
48			transfcount_init();
49
50			/*printf("######## send(sockfd=%i, buf=%p, len=%zi, flags=%x) -> ",
51			sockfd, buf, len, flags);*/
52
53			bytes = orig_send(sockfd, buf, len, flags);
54			total_sent += bytes;
55			/printf("%lu\n", bytes);/
56
57			return bytes;
58			}
59
60
61			/**
62			* sendmmsg - send multiple messages on a socket.
63			*/
64			int sendmmsg(int sockfd, struct mmsghdr *msgvec, unsigned int vlen,
65			int flags)
66			{
67			int sent, i;
68
69			transfcount_init();
70
71			/*printf("######## sendmmsg(sockfd=%i, msgvec=%p, vlen=%u, flags=%x) -> ",
72			sockfd, msgvec, vlen, flags);*/
73
74			sent = orig_sendmmsg(sockfd, msgvec, vlen, flags);
75			/printf("%i\n", sent);/
76			for (i = 0; i < sent; i++)
77			total_sent += msgvec[i].msg_len;
78
79			return sent;
80			}
81
82
83			/**
84			* recv - receive a message from a socket.
85			*/
86			ssize_t recv(int sockfd, void *buf, size_t len, int flags)
87			{
88			ssize_t bytes;
89
90			transfcount_init();
91
92			/*printf("######## recv(sockfd=%i, buf=%p, len=%zi, flags=%x) -> ",
93			sockfd, buf, len, flags);*/
94
95			bytes = orig_recv(sockfd, buf, len, flags);
96			total_recv += bytes;
97			/printf("%lu\n", bytes);/
98
99			return bytes;
100			}
101
102
103			/**
104			* read - read from a file descriptor.
105			*/
106			ssize_t read(int fd, void *buf, size_t count)
107			{
108			struct stat statbuf;
109			ssize_t bytes;
110
111			transfcount_init();
112
113			fstat(fd, &statbuf);
114			if (!S_ISSOCK(statbuf.st_mode))
115			return orig_read(fd, buf, count);
116
117			/*printf("######## read(fd=%i, buf=%p, count=%zi) -> ",
118			fd, buf, count);*/
119
120			bytes = orig_read(fd, buf, count);
121			if (bytes != -1)
122			total_recv += bytes;
123			/printf("%lu\n", bytes);/
124
125			return bytes;
126			}
127
128
129			/**
130			* write - write to a file descriptor.
131			*/
132			ssize_t write(int fd, const void *buf, size_t count)
133			{
134			struct stat statbuf;
135			ssize_t bytes;
136
137			transfcount_init();
138
139			fstat(fd, &statbuf);
140			if (!S_ISSOCK(statbuf.st_mode))
141			return orig_write(fd, buf, count);
142
143			/*printf("######## write(fd=%i, buf=%p, count=%zi) -> ",
144			fd, buf, count);*/
145
146			bytes = orig_write(fd, buf, count);
147			if (bytes != -1)
148			total_sent += bytes;
149			/printf("%lu\n", bytes);/
150
151			return bytes;
152			}
153
154
155			/**
156			* Set up library.
157			*/
158			void transfcount_init(void)
159			{
160			static bool init_done = false;
161
162			if (init_done)
163			return;
164
165			orig_send = (send_type)dlsym(RTLD_NEXT, "send");
166			orig_sendmmsg = (sendmmsg_type)dlsym(RTLD_NEXT, "sendmmsg");
167			orig_recv = (recv_type)dlsym(RTLD_NEXT, "recv");
168			orig_read = (read_type)dlsym(RTLD_NEXT, "read");
169			orig_write = (write_type)dlsym(RTLD_NEXT, "write");
170
171			atexit(transfcount_atexit);
172
173			start_time = time(NULL);
174
175			init_done = true;
176			}
177
178
179			/**
180			* Output transfer statistics.
181			*/
182			void transfcount_atexit(void)
183			{
184			time_t end_time = time(NULL);
185			time_t duration = end_time - start_time;
186
187			if (total_sent == 0 && total_recv == 0)
188			return;
189
190			if (duration == 0)
191			duration = 1;
192
193			fprintf(stderr, "[>>> Sent ");
194			transfcount_print_bytes(total_sent);
195			fprintf(stderr, "B (");
196			transfcount_print_bytes(total_sent / duration);
197			fprintf(stderr, "B/s) >>>] [<<< Received ");
198			transfcount_print_bytes(total_recv);
199			fprintf(stderr, "B (");
200			transfcount_print_bytes(total_recv / duration);
201			fprintf(stderr, "B/s) <<<]\n");
202			}
203
204
205			/**
206			* Print a size as human readable.
207			*/
208			void transfcount_print_bytes(unsigned long long n)
209			{
210			if (1024 * 1024 * 1024 <= n) {
211			fprintf(stderr, "%.1fGi", (double) n / 1024 / 1024 / 1024);
212			} else if (1024 * 1024 <= n) {
213			fprintf(stderr, "%.1fMi", (double) n / 1024 / 1024);
214			} else if (1024 <= n) {
215			fprintf(stderr, "%.1fKi", (double) n / 1024);
216			} else {
217			fprintf(stderr, "%llu", n);
218			}
219			}
220