2 * Copyright (c) 2013 Emil Mikulic <emikulic@gmail.com>
4 * Permission to use, copy, modify, and distribute this software for any
5 * purpose with or without fee is hereby granted, provided that the above
6 * copyright notice and this permission notice appear in all copies.
8 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
9 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
11 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
14 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
16 * [ http://www.openbsd.org/cgi-bin/cvsweb/src/share/misc/license.template ]
19 // buftee: like tee(1) but buffers in memory.
21 // Read from stdin into memory all the time, so that the writer doesn't
22 // block. Our memory usage is unbounded.
24 // Write to a number of file descriptors whenever we're able to.
25 // Because write() can take a long time even with O_NONBLOCK set on the
26 // fd, we have to do this from a separate kernel scheduling entity (we use
29 // We use a thread per output filedescriptor so that one slow output (e.g. disk)
30 // doesn't block another (e.g. stdout to a terminal) (or vice-versa if you use
31 // flow control on the terminal!)
33 #define _GNU_SOURCE // for clock_gettime()
35 #include <sys/queue.h>
36 #include <sys/select.h>
37 #include <sys/syscall.h>
51 #define READ_BUF_SIZE 4096
52 #define SLOW_NSEC 4000
54 // Asserted on receipt of SIGTERM, SIGINT.
55 static volatile int stopping
= 0;
57 static void sig_stopping(int _ignored_
__attribute__((__unused__
))) {
61 // Reference-counted buffer, contains data that was read in the main thread.
69 struct buf_queue_elem
{
71 STAILQ_ENTRY(buf_queue_elem
) entries
;
73 STAILQ_HEAD(buf_queue
, buf_queue_elem
); // struct buf_queue
75 // Context for a writer thread.
76 struct writer_thread
{
80 // Each writer has its own queue.
81 struct buf_queue queue
;
83 // Pointer to a shared mutex which protects all queues.
84 pthread_mutex_t
* mutex
;
86 // When a writer runs out of work to do, it sleeps on a shared condition.
89 STAILQ_ENTRY(writer_thread
) entries
;
92 // A list of writer threads.
93 STAILQ_HEAD(writer_thread_list
, writer_thread
); // struct writer_thread_list
95 static struct buf
* alloc_buf(const char* const data
, const int len
) {
97 struct buf
* buf
= malloc(sizeof(*buf
));
98 buf
->data
= malloc((size_t)len
);
99 memcpy(buf
->data
, data
, (size_t)len
);
105 static void unref_buf(struct buf
* buf
) {
106 assert(buf
->refcount
> 0);
107 if ((--buf
->refcount
) == 0) {
113 static void enqueue(struct buf_queue
* restrict queue
,
114 struct buf
* restrict buf
) {
115 struct buf_queue_elem
* elem
= malloc(sizeof(*elem
));
118 STAILQ_INSERT_TAIL(queue
, elem
, entries
);
121 static struct buf
* dequeue(struct buf_queue
* const queue
) {
122 assert(!STAILQ_EMPTY(queue
));
123 struct buf_queue_elem
* head
;
124 head
= STAILQ_FIRST(queue
);
125 STAILQ_REMOVE_HEAD(queue
, entries
);
126 struct buf
* buf
= head
->buf
;
131 static void xpthread_create(pthread_t
* thread
,
132 void* (*start_routine
)(void*),
134 int ret
= pthread_create(thread
, NULL
, start_routine
, arg
);
135 if (ret
== 0) return;
137 err(1, "pthread_create(%p) failed", thread
);
140 static void set_nonblocking(const int fd
) {
142 if ((flags
= fcntl(fd
, F_GETFL
)) == -1)
143 err(1, "fcntl(fd = %d, F_GETFL) failed", fd
);
144 if (fcntl(fd
, F_SETFL
, flags
| O_NONBLOCK
) == -1)
145 err(1, "fcntl(fd = %d, F_SETFL, O_NONBLOCK) failed", fd
);
148 static int make_file(const char* filename
) {
149 int fd
= open(filename
, O_CREAT
| O_EXCL
| O_NONBLOCK
| O_WRONLY
, 0666);
151 err(1, "failed to open(\"%s\")", filename
);
155 static void get_mono_time(struct timespec
* t
) {
156 if (clock_gettime(CLOCK_MONOTONIC
, t
) == -1)
157 err(1, "clock_gettime(CLOCK_MONOTONIC) failed");
160 static void time_diff(const struct timespec
* restrict start
,
161 const struct timespec
* restrict end
,
162 struct timespec
* restrict out
) {
163 out
->tv_sec
= end
->tv_sec
- start
->tv_sec
;
164 out
->tv_nsec
= end
->tv_nsec
- start
->tv_nsec
;
165 if (out
->tv_nsec
< 0) {
167 out
->tv_nsec
+= 1000000000;
169 assert(out
->tv_sec
>= 0);
170 assert(out
->tv_nsec
>= 0);
171 assert(out
->tv_nsec
< 1000000000);
174 static int gettid(void) {
175 return (int)syscall(SYS_gettid
);
178 static void warn_time(const char* desc
,
179 const struct timespec
* restrict start
,
180 const struct timespec
* restrict end
) {
181 struct timespec diff
;
182 time_diff(start
, end
, &diff
);
183 if (diff
.tv_sec
> 0 || diff
.tv_nsec
> SLOW_NSEC
) {
185 extern char *__progname
; // This is where glibc stashes argv[0].
186 snprintf(buf
, sizeof(buf
), "%s:tid %d: %s took %d.%09d secs\n",
187 __progname
, gettid(), desc
, (int)diff
.tv_sec
, (int)diff
.tv_nsec
);
188 // Best effort write to a non-blocking stderr.
189 (void)write(STDERR_FILENO
, buf
, strlen(buf
));
193 static int xread(const int fd
, char* const restrict buf
, const int count
) {
200 read_ret
= read(fd
, buf
, (size_t)count
);
203 warn_time("read()", &t0
, &t1
);
207 err(1, "read(fd = %d, count = %d) failed", fd
, count
);
210 assert(read_ret
>= 0);
211 return (int)read_ret
;
214 static void wait_until_writable(const int fd
) {
217 FD_SET(fd
, &write_fds
);
218 int select_ret
= select(fd
+ 1, NULL
, &write_fds
, NULL
, NULL
);
219 if (select_ret
== -1) {
220 if (errno
== EINTR
) {
221 assert(stopping
); // that should have been SIGTERM
224 err(1, "select(write fd = %d) failed", fd
);
226 if (!FD_ISSET(fd
, &write_fds
))
227 errx(1, "select() did not return writable fd = %d", fd
);
230 static int xwrite(const int fd
, struct buf
* const buf
) {
238 write_ret
= write(fd
, buf
->data
, (size_t)buf
->len
);
241 warn_time("write()", &t0
, &t1
);
244 if (write_ret
== -1) {
245 if (errno
== EAGAIN
) {
246 warn("write(fd = %d) got EAGAIN, sleeping and retrying", fd
);
247 wait_until_writable(fd
);
250 err(1, "write(fd = %d, count = %d) failed", fd
, buf
->len
);
254 assert(write_ret
>= 0);
255 if (write_ret
< buf
->len
)
256 err(1, "write(fd = %d, count = %d) stopped short (returned %d)",
257 fd
, buf
->len
, (int)write_ret
);
258 // FIXME: handle this
259 assert(write_ret
== buf
->len
);
260 return (int)write_ret
;
264 static void wait_until_readable(const int fd
) {
267 FD_SET(fd
, &read_fds
);
268 int select_ret
= select(fd
+ 1, &read_fds
, NULL
, NULL
, NULL
);
269 if (select_ret
== -1) {
270 if (errno
== EINTR
) {
271 assert(stopping
); // that should have been SIGTERM
274 err(1, "select(read fd = %d) failed", fd
);
276 if (!FD_ISSET(fd
, &read_fds
))
277 errx(1, "select() did not return readable fd = %d", fd
);
280 static void lock(pthread_mutex_t
* mutex
) {
281 int ret
= pthread_mutex_lock(mutex
);
282 if (ret
== 0) return;
284 err(1, "pthread_mutex_lock(%p) failed", mutex
);
287 static void unlock(pthread_mutex_t
* mutex
) {
288 int ret
= pthread_mutex_unlock(mutex
);
289 if (ret
== 0) return;
291 err(1, "pthread_mutex_unlock(%p) failed", mutex
);
294 static void* writer_routine(void *arg
) {
295 struct writer_thread
* my
= arg
;
296 struct buf
* buf
= NULL
;
299 while (!stopping
&& STAILQ_EMPTY(&my
->queue
)) {
301 pthread_cond_wait(my
->cond
, my
->mutex
);
303 if (!STAILQ_EMPTY(&my
->queue
))
304 buf
= dequeue(&my
->queue
);
311 int write_ret
= xwrite(my
->fd
, buf
);
312 if (write_ret
== 0) {
313 errx(1, "fd %d hit EOF", my
->fd
);
315 assert(write_ret
== buf
->len
);
317 // Unreference buffer, freeing it if we have to.
321 warnx("thread exiting cleanly");
325 static void add_writer_thread(struct writer_thread_list
* list
,
327 pthread_mutex_t
* shared_queue_mutex
,
328 pthread_cond_t
* shared_wakeup_cond
) {
329 struct writer_thread
* writer
= malloc(sizeof(*writer
));
332 writer
->mutex
= shared_queue_mutex
;
333 writer
->cond
= shared_wakeup_cond
;
334 STAILQ_INIT(&writer
->queue
);
335 STAILQ_INSERT_TAIL(list
, writer
, entries
);
336 xpthread_create(&writer
->thread
, writer_routine
, writer
);
339 static void xpthread_cond_broadcast(pthread_cond_t
* cond
) {
340 int ret
= pthread_cond_broadcast(cond
);
341 if (ret
== 0) return;
343 err(1, "pthread_cond_broadcast(%p) failed", cond
);
346 static void xpthread_join(pthread_t thread
) {
347 int ret
= pthread_join(thread
, NULL
);
348 if (ret
== 0) return;
350 err(1, "pthread_join(%lu) failed", thread
);
353 static void sig_continue(int _ignored_
__attribute__((__unused__
))) {
354 set_nonblocking(STDERR_FILENO
);
357 int main(int argc
, char **argv
) {
358 pthread_mutex_t shared_queue_mutex
= PTHREAD_MUTEX_INITIALIZER
;
359 pthread_cond_t shared_wakeup_cond
= PTHREAD_COND_INITIALIZER
;
361 struct writer_thread_list writers
;
362 STAILQ_INIT(&writers
);
364 if (signal(SIGINT
, sig_stopping
) == SIG_ERR
) err(1, "signal() failed");
365 if (signal(SIGTERM
, sig_stopping
) == SIG_ERR
) err(1, "signal() failed");
366 if (signal(SIGCONT
, sig_continue
) == SIG_ERR
) err(1, "signal() failed");
367 //if (signal(SIGPIPE, SIG_IGN) == SIG_ERR) err(1, "signal() failed");
370 // On Linux, making STDOUT non-blocking has the side-effect of
371 // also making STDIN nonblocking.
372 add_writer_thread(&writers
,
375 &shared_wakeup_cond
);
377 // Process cmdline args.
378 for (int i
= 1; i
< argc
; i
++) {
379 add_writer_thread(&writers
,
382 &shared_wakeup_cond
);
387 wait_until_readable(STDIN_FILENO
);
389 warnx("stopping after select()");
394 char data
[READ_BUF_SIZE
];
395 int read_ret
= xread(STDIN_FILENO
, data
, sizeof(data
));
397 warnx("stdin hit EOF");
400 struct buf
* buf
= alloc_buf(data
, read_ret
);
403 lock(&shared_queue_mutex
);
404 struct writer_thread
* writer
;
405 STAILQ_FOREACH(writer
, &writers
, entries
) {
406 enqueue(&(writer
->queue
), buf
);
408 xpthread_cond_broadcast(&shared_wakeup_cond
);
409 unlock(&shared_queue_mutex
);
412 // Wake and join threads.
413 lock(&shared_queue_mutex
);
415 xpthread_cond_broadcast(&shared_wakeup_cond
);
416 unlock(&shared_queue_mutex
);
418 struct writer_thread
* writer
;
419 STAILQ_FOREACH(writer
, &writers
, entries
) {
420 xpthread_join(writer
->thread
);
425 while (!STAILQ_EMPTY(&writers
)) {
426 struct writer_thread
* writer
= STAILQ_FIRST(&writers
);
427 STAILQ_REMOVE_HEAD(&writers
, entries
);
428 // FIXME: free its queue?
432 warnx("exiting cleanly");
435 // vim:set ts=2 sw=2 tw=80 et: