1 #ifndef IOU_REQ_REF_H
2 #define IOU_REQ_REF_H
3 
4 #include <linux/atomic.h>
5 #include <linux/io_uring_types.h>
6 
7 /*
8  * Shamelessly stolen from the mm implementation of page reference checking,
9  * see commit f958d7b528b1 for details.
10  */
11 #define req_ref_zero_or_close_to_overflow(req)	\
12 	((unsigned int) atomic_read(&(req->refs)) + 127u <= 127u)
13 
req_ref_inc_not_zero(struct io_kiocb * req)14 static inline bool req_ref_inc_not_zero(struct io_kiocb *req)
15 {
16 	WARN_ON_ONCE(!(req->flags & REQ_F_REFCOUNT));
17 	return atomic_inc_not_zero(&req->refs);
18 }
19 
req_ref_put_and_test(struct io_kiocb * req)20 static inline bool req_ref_put_and_test(struct io_kiocb *req)
21 {
22 	if (likely(!(req->flags & REQ_F_REFCOUNT)))
23 		return true;
24 
25 	WARN_ON_ONCE(req_ref_zero_or_close_to_overflow(req));
26 	return atomic_dec_and_test(&req->refs);
27 }
28 
req_ref_get(struct io_kiocb * req)29 static inline void req_ref_get(struct io_kiocb *req)
30 {
31 	WARN_ON_ONCE(!(req->flags & REQ_F_REFCOUNT));
32 	WARN_ON_ONCE(req_ref_zero_or_close_to_overflow(req));
33 	atomic_inc(&req->refs);
34 }
35 
req_ref_put(struct io_kiocb * req)36 static inline void req_ref_put(struct io_kiocb *req)
37 {
38 	WARN_ON_ONCE(!(req->flags & REQ_F_REFCOUNT));
39 	WARN_ON_ONCE(req_ref_zero_or_close_to_overflow(req));
40 	atomic_dec(&req->refs);
41 }
42 
__io_req_set_refcount(struct io_kiocb * req,int nr)43 static inline void __io_req_set_refcount(struct io_kiocb *req, int nr)
44 {
45 	if (!(req->flags & REQ_F_REFCOUNT)) {
46 		req->flags |= REQ_F_REFCOUNT;
47 		atomic_set(&req->refs, nr);
48 	}
49 }
50 
io_req_set_refcount(struct io_kiocb * req)51 static inline void io_req_set_refcount(struct io_kiocb *req)
52 {
53 	__io_req_set_refcount(req, 1);
54 }
55 #endif
56