bufferevent: Add functions to set/get max_single_read/write values.

4 files changed, 100 insertions, 7 deletions

diff --git a/bufferevent-internal.h b/bufferevent-internal.h
index 8f47979b..e0b0c0af 100644
--- a/bufferevent-internal.h
+++ b/bufferevent-internal.h
@@ -190,6 +190,14 @@ struct bufferevent_private {
 	 * If NULL, locking is disabled. */
 	void *lock;
 
+	/** No matter how big our bucket gets, don't try to read more than this
+	 * much in a single read operation. */
+	ev_ssize_t max_single_read;
+
+	/** No matter how big our bucket gets, don't try to write more than this
+	 * much in a single write operation. */
+	ev_ssize_t max_single_write;
+
 	/** Rate-limiting information for this bufferevent */
 	struct bufferevent_rate_limit *rate_limiting;
 };
@@ -403,6 +411,8 @@ int _bufferevent_decrement_read_buckets(struct bufferevent_private *bev,
 ev_ssize_t _bufferevent_get_read_max(struct bufferevent_private *bev);
 ev_ssize_t _bufferevent_get_write_max(struct bufferevent_private *bev);
 
+int _bufferevent_ratelim_init(struct bufferevent_private *bev);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/bufferevent.c b/bufferevent.c
index 3a991860..a0a107d0 100644
--- a/bufferevent.c
+++ b/bufferevent.c
@@ -298,6 +298,8 @@ bufferevent_init_common(struct bufferevent_private *bufev_private,
 
 	bufev->be_ops = ops;
 
+	_bufferevent_ratelim_init(bufev_private);
+
 	/*
 	 * Set to EV_WRITE so that using bufferevent_write is going to
 	 * trigger a callback.  Reading needs to be explicitly enabled
diff --git a/bufferevent_ratelim.c b/bufferevent_ratelim.c
index 506e82ab..cd9fcaa6 100644
--- a/bufferevent_ratelim.c
+++ b/bufferevent_ratelim.c
@@ -178,12 +178,9 @@ ev_token_bucket_cfg_free(struct ev_token_bucket_cfg *cfg)
 	mm_free(cfg);
 }
 
-/* No matter how big our bucket gets, don't try to read more than this
- * much in a single read operation. */
-#define MAX_TO_READ_EVER 16384
-/* No matter how big our bucket gets, don't try to write more than this
- * much in a single write operation. */
-#define MAX_TO_WRITE_EVER 16384
+/* Default values for max_single_read & max_single_write variables. */
+#define MAX_SINGLE_READ_DEFAULT 16384
+#define MAX_SINGLE_WRITE_DEFAULT 16384
 
 #define LOCK_GROUP(g) EVLOCK_LOCK((g)->lock, 0)
 #define UNLOCK_GROUP(g) EVLOCK_UNLOCK((g)->lock, 0)
@@ -201,7 +198,7 @@ static inline ev_ssize_t
 _bufferevent_get_rlim_max(struct bufferevent_private *bev, int is_write)
 {
 	/* needs lock on bev. */
-	ev_ssize_t max_so_far = is_write?MAX_TO_WRITE_EVER:MAX_TO_READ_EVER;
+	ev_ssize_t max_so_far = is_write?bev->max_single_write:bev->max_single_read;
 
 #define LIM(x)						\
 	(is_write ? (x).write_limit : (x).read_limit)
@@ -851,6 +848,56 @@ bufferevent_get_write_limit(struct bufferevent *bev)
 	return r;
 }
 
+int
+bufferevent_set_max_single_read(struct bufferevent *bev, size_t size)
+{
+	struct bufferevent_private *bevp;
+	BEV_LOCK(bev);
+	bevp = BEV_UPCAST(bev);
+	if (size == 0 || size > EV_SSIZE_MAX)
+		bevp->max_single_read = MAX_SINGLE_READ_DEFAULT;
+	else
+		bevp->max_single_read = size;
+	BEV_UNLOCK(bev);
+	return 0;
+}
+
+int
+bufferevent_set_max_single_write(struct bufferevent *bev, size_t size)
+{
+	struct bufferevent_private *bevp;
+	BEV_LOCK(bev);
+	bevp = BEV_UPCAST(bev);
+	if (size == 0 || size > EV_SSIZE_MAX)
+		bevp->max_single_write = MAX_SINGLE_WRITE_DEFAULT;
+	else
+		bevp->max_single_write = size;
+	BEV_UNLOCK(bev);
+	return 0;
+}
+
+ev_ssize_t
+bufferevent_get_max_single_read(struct bufferevent *bev)
+{
+	ev_ssize_t r;
+
+	BEV_LOCK(bev);
+	r = BEV_UPCAST(bev)->max_single_read;
+	BEV_UNLOCK(bev);
+	return r;
+}
+
+ev_ssize_t
+bufferevent_get_max_single_write(struct bufferevent *bev)
+{
+	ev_ssize_t r;
+
+	BEV_LOCK(bev);
+	r = BEV_UPCAST(bev)->max_single_write;
+	BEV_UNLOCK(bev);
+	return r;
+}
+
 ev_ssize_t
 bufferevent_get_max_to_read(struct bufferevent *bev)
 {
@@ -1010,3 +1057,13 @@ bufferevent_rate_limit_group_reset_totals(struct bufferevent_rate_limit_group *g
 {
 	grp->total_read = grp->total_written = 0;
 }
+
+int
+_bufferevent_ratelim_init(struct bufferevent_private *bev)
+{
+	bev->rate_limiting = NULL;
+	bev->max_single_read = MAX_SINGLE_READ_DEFAULT;
+	bev->max_single_write = MAX_SINGLE_WRITE_DEFAULT;
+
+	return 0;
+}
diff --git a/include/event2/bufferevent.h b/include/event2/bufferevent.h
index dac988e3..c9302093 100644
--- a/include/event2/bufferevent.h
+++ b/include/event2/bufferevent.h
@@ -725,6 +725,30 @@ int bufferevent_add_to_rate_limit_group(struct bufferevent *bev,
 int bufferevent_remove_from_rate_limit_group(struct bufferevent *bev);
 
 /**
+   Set the size limit for single read operation.
+
+   Set to 0 for a reasonable default.
+
+   Return 0 on success and -1 on failure.
+ */
+int bufferevent_set_max_single_read(struct bufferevent *bev, size_t size);
+
+/**
+   Set the size limit for single write operation.
+
+   Set to 0 for a reasonable default.
+
+   Return 0 on success and -1 on failure.
+ */
+int bufferevent_set_max_single_write(struct bufferevent *bev, size_t size);
+
+/** Get the current size limit for single read operation. */
+ev_ssize_t bufferevent_get_max_single_read(struct bufferevent *bev);
+
+/** Get the current size limit for single write operation. */
+ev_ssize_t bufferevent_get_max_single_write(struct bufferevent *bev);
+
+/**
    @name Rate limit inspection
 
    Return the current read or write bucket size for a bufferevent.