static int
stats_dostuff(TSCont contp, TSEvent event, void *edata)
{
stats_state *my_state = TSContDataGet(contp);
if (event == TS_EVENT_NET_ACCEPT) {
my_state->net_vc = (TSVConn) edata;
stats_process_accept(contp, my_state);
} else if (edata == my_state->read_vio) {
stats_process_read(contp, event, my_state);
} else if (edata == my_state->write_vio) {
stats_process_write(contp, event, my_state);
} else {
TSReleaseAssert(!"Unexpected Event");
}
return 0;
}
static int
thread_plugin(TSCont contp ATS_UNUSED, TSEvent event, void *edata)
{
switch (event) {
case TS_EVENT_HTTP_OS_DNS:
/**
* Check if the thread has been created successfully or not.
* If the thread has not been created successfully, assert.
*/
if (!TSThreadCreate(reenable_txn, edata)) {
TSReleaseAssert(!"Failure in thread creation");
}
return 0;
default:
break;
}
return 0;
}
/* Implement the server intercept */
static int
acme_intercept(TSCont contp, TSEvent event, void *edata)
{
AcmeState *my_state = TSContDataGet(contp);
if (event == TS_EVENT_NET_ACCEPT) {
my_state->net_vc = (TSVConn)edata;
acme_process_accept(contp, my_state);
} else if (edata == my_state->read_vio) { /* All read events */
acme_process_read(contp, event, my_state);
} else if (edata == my_state->write_vio) { /* All write events */
acme_process_write(contp, event, my_state);
} else {
TSReleaseAssert(!"Unexpected Event");
}
return 0;
}
static void
stats_process_write(TSCont contp, TSEvent event, stats_state * my_state)
{
if (event == TS_EVENT_VCONN_WRITE_READY) {
if (my_state->body_written == 0) {
TSDebug("istats", "plugin adding response body");
my_state->body_written = 1;
json_out_stats(my_state);
TSVIONBytesSet(my_state->write_vio, my_state->output_bytes);
}
TSVIOReenable(my_state->write_vio);
} else if (TS_EVENT_VCONN_WRITE_COMPLETE) {
stats_cleanup(contp, my_state);
} else if (event == TS_EVENT_ERROR) {
TSError("stats_process_write: Received TS_EVENT_ERROR\n");
} else {
TSReleaseAssert(!"Unexpected Event");
}
}
static void
stats_process_read(TSCont contp, TSEvent event, stats_state * my_state)
{
TSDebug("istats", "stats_process_read(%d)", event);
if (event == TS_EVENT_VCONN_READ_READY) {
my_state->output_bytes = stats_add_resp_header(my_state);
TSVConnShutdown(my_state->net_vc, 1, 0);
my_state->write_vio = TSVConnWrite(my_state->net_vc, contp, my_state->resp_reader, INT64_MAX);
} else if (event == TS_EVENT_ERROR) {
TSError("stats_process_read: Received TS_EVENT_ERROR\n");
} else if (event == TS_EVENT_VCONN_EOS) {
/* client may end the connection, simply return */
return;
} else if (event == TS_EVENT_NET_ACCEPT_FAILED) {
TSError("stats_process_read: Received TS_EVENT_NET_ACCEPT_FAILED\n");
} else {
printf("Unexpected Event %d\n", event);
TSReleaseAssert(!"Unexpected Event");
}
}
/* Process a write event from the SM */
static void
acme_process_write(TSCont contp, TSEvent event, AcmeState *my_state)
{
if (event == TS_EVENT_VCONN_WRITE_READY) {
char buf[64]; /* Plenty of space for CL: header */
int len;
len = snprintf(buf, sizeof(buf) - 1, "Content-Length: %zd\r\n\r\n", my_state->stat_buf.st_size);
my_state->output_bytes += add_data_to_resp(buf, len, my_state);
my_state->output_bytes += add_file_to_resp(my_state);
TSVIONBytesSet(my_state->write_vio, my_state->output_bytes);
TSVIOReenable(my_state->write_vio);
} else if (TS_EVENT_VCONN_WRITE_COMPLETE) {
cleanup(contp, my_state);
} else if (event == TS_EVENT_ERROR) {
TSError("[%s] acme_process_write: Received TS_EVENT_ERROR", PLUGIN_NAME);
} else {
TSReleaseAssert(!"Unexpected Event");
}
}
/* Process a read event from the SM */
static void
acme_process_read(TSCont contp, TSEvent event, AcmeState *my_state)
{
if (event == TS_EVENT_VCONN_READ_READY) {
if (-1 == my_state->fd) {
my_state->output_bytes = add_data_to_resp(ACME_DENIED_RESP, strlen(ACME_DENIED_RESP), my_state);
} else {
my_state->output_bytes = add_data_to_resp(ACME_OK_RESP, strlen(ACME_OK_RESP), my_state);
}
TSVConnShutdown(my_state->net_vc, 1, 0);
my_state->write_vio = TSVConnWrite(my_state->net_vc, contp, my_state->resp_reader, INT64_MAX);
} else if (event == TS_EVENT_ERROR) {
TSError("[%s] acme_process_read: Received TS_EVENT_ERROR", PLUGIN_NAME);
} else if (event == TS_EVENT_VCONN_EOS) {
/* client may end the connection, simply return */
return;
} else if (event == TS_EVENT_NET_ACCEPT_FAILED) {
TSError("[%s] acme_process_read: Received TS_EVENT_NET_ACCEPT_FAILED", PLUGIN_NAME);
} else {
TSReleaseAssert(!"Unexpected Event");
}
}
/* Process a write event from the SM */
static void
hc_process_write(TSCont contp, TSEvent event, HCState *my_state)
{
if (event == TS_EVENT_VCONN_WRITE_READY) {
char buf[48];
int len;
len = snprintf(buf, sizeof(buf)-1, "Content-Length: %d\r\n\r\n", my_state->data->b_len);
my_state->output_bytes += add_data_to_resp(buf, len, my_state);
if (my_state->data->b_len > 0)
my_state->output_bytes += add_data_to_resp(my_state->data->body, my_state->data->b_len, my_state);
else
my_state->output_bytes += add_data_to_resp("\r\n", 2, my_state);
TSVIONBytesSet(my_state->write_vio, my_state->output_bytes);
TSVIOReenable(my_state->write_vio);
} else if (TS_EVENT_VCONN_WRITE_COMPLETE) {
cleanup(contp, my_state);
} else if (event == TS_EVENT_ERROR) {
TSError("hc_process_write: Received TS_EVENT_ERROR\n");
} else {
TSReleaseAssert(!"Unexpected Event");
}
}
/* Process a read event from the SM */
static void
hc_process_read(TSCont contp, TSEvent event, HCState *my_state)
{
if (event == TS_EVENT_VCONN_READ_READY) {
if (my_state->data->exists) {
TSDebug(PLUGIN_NAME, "Setting OK response header");
my_state->output_bytes = add_data_to_resp(my_state->info->ok, my_state->info->o_len, my_state);
} else {
TSDebug(PLUGIN_NAME, "Setting MISS response header");
my_state->output_bytes = add_data_to_resp(my_state->info->miss, my_state->info->m_len, my_state);
}
TSVConnShutdown(my_state->net_vc, 1, 0);
my_state->write_vio = TSVConnWrite(my_state->net_vc, contp, my_state->resp_reader, INT64_MAX);
} else if (event == TS_EVENT_ERROR) {
TSError("hc_process_read: Received TS_EVENT_ERROR\n");
} else if (event == TS_EVENT_VCONN_EOS) {
/* client may end the connection, simply return */
return;
} else if (event == TS_EVENT_NET_ACCEPT_FAILED) {
TSError("hc_process_read: Received TS_EVENT_NET_ACCEPT_FAILED\n");
} else {
TSReleaseAssert(!"Unexpected Event");
}
}
static int
handle_buffering(TSCont contp, MyData * data)
{
TSVIO write_vio;
int towrite;
int avail;
/* Get the write VIO for the write operation that was performed on
ourself. This VIO contains the buffer that we are to read from
as well as the continuation we are to call when the buffer is
empty. */
write_vio = TSVConnWriteVIOGet(contp);
/* Create the output buffer and its associated reader */
if (!data->output_buffer) {
data->output_buffer = TSIOBufferCreate();
TSAssert(data->output_buffer);
data->output_reader = TSIOBufferReaderAlloc(data->output_buffer);
TSAssert(data->output_reader);
}
/* We also check to see if the write VIO's buffer is non-NULL. A
NULL buffer indicates that the write operation has been
shutdown and that the continuation does not want us to send any
more WRITE_READY or WRITE_COMPLETE events. For this buffered
transformation that means we're done buffering data. */
if (!TSVIOBufferGet(write_vio)) {
data->state = STATE_OUTPUT_DATA;
return 0;
}
/* Determine how much data we have left to read. For this bnull
transform plugin this is also the amount of data we have left
to write to the output connection. */
towrite = TSVIONTodoGet(write_vio);
if (towrite > 0) {
/* The amount of data left to read needs to be truncated by
the amount of data actually in the read buffer. */
avail = TSIOBufferReaderAvail(TSVIOReaderGet(write_vio));
if (towrite > avail) {
towrite = avail;
}
if (towrite > 0) {
/* Copy the data from the read buffer to the input buffer. */
TSIOBufferCopy(data->output_buffer, TSVIOReaderGet(write_vio), towrite, 0);
/* Tell the read buffer that we have read the data and are no
longer interested in it. */
TSIOBufferReaderConsume(TSVIOReaderGet(write_vio), towrite);
/* Modify the write VIO to reflect how much data we've
completed. */
TSVIONDoneSet(write_vio, TSVIONDoneGet(write_vio) + towrite);
}
}
/* Now we check the write VIO to see if there is data left to read. */
if (TSVIONTodoGet(write_vio) > 0) {
if (towrite > 0) {
/* Call back the write VIO continuation to let it know that we
are ready for more data. */
TSContCall(TSVIOContGet(write_vio), TS_EVENT_VCONN_WRITE_READY, write_vio);
}
} else {
data->state = STATE_OUTPUT_DATA;
/* Call back the write VIO continuation to let it know that we
have completed the write operation. */
TSContCall(TSVIOContGet(write_vio), TS_EVENT_VCONN_WRITE_COMPLETE, write_vio);
}
return 1;
/* If we are in this code path then something is seriously wrong. */
TSError("[bnull-transform] Fatal error in plugin");
TSReleaseAssert(!"[bnull-transform] Fatal error in plugin\n");
return 0;
}
uint64_t RateLimiter::GetMaxUnits(int counter_index, const char * key, const timeval& time, uint64_t amount) {
LimiterEntry * limiter_entry = counters_[counter_index];
TSReleaseAssert(!pthread_rwlock_rdlock(&rwlock_keymap_));
std::map<const char *,LimiterState *>::iterator it = keymap_.find(key);
LimiterState * state = NULL;
TSReleaseAssert(!pthread_rwlock_unlock(&rwlock_keymap_));
if ( it == keymap_.end() ) {
char * key_copy = (char *)malloc( strlen(key) + 1 );
strcpy(key_copy, key);
//NOTE: consider letting only limiterentry construct new states
state = new LimiterState(GetCounterArray(), GetTimevalArray(time));
TSReleaseAssert(!pthread_rwlock_wrlock(&rwlock_keymap_));
keymap_.insert( std::pair<const char *, LimiterState *> ( key_copy, state ));
TSReleaseAssert(!pthread_rwlock_unlock(&rwlock_keymap_));
} else {
state = it->second;
}
//******** fixme -> this should have a critical section per key
TSMutexLock(update_mutex_);
timeval elapsed;
timeval stime = state->time(counter_index);
timersub(&time, &stime, &elapsed);
//FIXME: tv_usec seems to differ on different compilers/platforms :S
float elapsed_ms = (elapsed.tv_sec * 1000.0f) + ( elapsed.tv_usec/1000.0f );
float rate_timeslice = 1.0f - (limiter_entry->milliseconds() - elapsed_ms) / limiter_entry->milliseconds();
//FIXME: -> this is very very ugly
if (rate_timeslice<0) rate_timeslice=0;
//FIXME: this one sometimes gets slightly negative ... shudder
TSReleaseAssert(rate_timeslice >= 0.0f);
float replenishment = rate_timeslice * limiter_entry->max_rate();
float newallowance = state->allowance(counter_index) + replenishment;
//clip the new allowance value at the max rate for this limiter
newallowance = newallowance > limiter_entry->max_rate() ? limiter_entry->max_rate() : newallowance;
TSReleaseAssert(newallowance >= 0.0f);
int rv = amount;
if (amount > newallowance) {
amount = rv = newallowance;
}
TSReleaseAssert(rv >= 0);
//now substract the specified amount
newallowance -= amount;
if (newallowance >= 0.0f ) {
//update the state
state->set_allowance(counter_index, newallowance);
state->set_time(counter_index, time);
}
TSMutexUnlock(update_mutex_);
//******** end fixme
return rv;
}
//TODO: thread safety!! this assumes the stl map is implemented using a r/b tree that does not change
// during lookups. e.g. concurrent reads are allowed.
//TODO: a hash table would be more appropriate then a std::map
//this function updates the specified counter for the specified key, substring amount units at the specified time.
//the return value is 0 when no throttling is needed. it returns the number of ms that needs te be waited for to allow the
//specified amount to be substracted from the remaining allowance
uint64_t RateLimiter::Register(int counter_index, const char * key, const timeval& time, uint64_t amount) {
LimiterEntry * limiter_entry = counters_[counter_index];
TSReleaseAssert(!pthread_rwlock_rdlock(&rwlock_keymap_));
std::map<const char *,LimiterState *>::iterator it = keymap_.find(key);
LimiterState * state = NULL;
TSReleaseAssert(!pthread_rwlock_unlock(&rwlock_keymap_));
if ( it == keymap_.end() ) {
char * key_copy = (char *)malloc( strlen(key) + 1 );
strcpy(key_copy, key);
//NOTE: consider letting only limiterentry construct new states
state = new LimiterState(GetCounterArray(), GetTimevalArray(time));
TSReleaseAssert(!pthread_rwlock_wrlock(&rwlock_keymap_));
keymap_.insert( std::pair<const char *, LimiterState *> ( key_copy, state ));
TSReleaseAssert(!pthread_rwlock_unlock(&rwlock_keymap_));
} else {
state = it->second;
}
//******** fixme -> this could have a critical section per key,
//but that would generate a lot of lock structures
TSMutexLock(update_mutex_);
timeval elapsed;
timeval stime = state->time(counter_index);
timersub(&time, &stime, &elapsed);
//FIXME: tv_usec seems to differ on different compilers/platforms :S
float elapsed_ms = (elapsed.tv_sec * 1000.0f) + ( elapsed.tv_usec/1000.0f );
float rate_timeslice = 1.0f - (limiter_entry->milliseconds() - elapsed_ms) / limiter_entry->milliseconds();
//FIXME: -> this is very very ugly
if (rate_timeslice<0) rate_timeslice=0;
//FIXME: this one sometimes gets slightly negative ... shudder
TSReleaseAssert(rate_timeslice >= 0.0f);
float replenishment = rate_timeslice * limiter_entry->max_rate();
float newallowance = state->allowance(counter_index) + replenishment;
//clip the new allowance value at the max rate for this limiter
newallowance = newallowance > limiter_entry->max_rate() ? limiter_entry->max_rate() : newallowance;
//now substract the specified amount
newallowance -= amount;
if (newallowance >= 0.0f ) {
//update the state
state->set_allowance(counter_index, newallowance);
state->set_time(counter_index, time);
}
state->set_taken(counter_index, state->taken(counter_index) + amount);
dbg("added amount, currently taken %f", state->taken(counter_index));
TSMutexUnlock(update_mutex_);
//******** end fixme
if (newallowance >= 0.0f) {
//dbg("OK : rate timeslice: %f (elapsed ms: %f), replenish %f units, substract %ld, new allowance %f\n", rate_timeslice, elapsed_ms, replenishment, amount, newallowance);
return 0;
} else {
float debt_factor = 1.0f - ( (limiter_entry->max_rate() + newallowance) / limiter_entry->max_rate() );
uint64_t wait_ms = ceil((float)limiter_entry->milliseconds() * debt_factor);
//clog("RL!!: rate timeslice: %f, replenish %f units, substract %ld, if allowed, debt would be %f (%ld ms)\n", rate_timeslice, replenishment, amount, newallowance, wait_ms);
return wait_ms;
}
}