ngx_http_status_code_counter_module.c


/*
 * Copyright (C) Kennon Ballou
 */

#include <sys/ipc.h>
#include <sys/shm.h>
#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>

#define NGX_HTTP_LAST_LEVEL_500  508
#define NGX_HTTP_NUM_STATUS_CODES (NGX_HTTP_LAST_LEVEL_500 - NGX_HTTP_OK)
#define SHMKEY 22223
#define SHMFLAG IPC_CREAT | (S_IRUSR|S_IWUSR|S_IRGRP|S_IWGRP|S_IROTH|S_IWOTH)

/* 
  This is really wasteful to store such a sparse array of status code counts,
  but this will allow new status codes to be counted without having any hardcoding 
*/
ngx_atomic_t *ngx_http_status_code_counts;

typedef struct {
    ngx_int_t   startup;
} ngx_http_status_code_counter_conf_t;

static char *ngx_http_set_status_code_counter(ngx_conf_t *cf, ngx_command_t *cmd,
                                 void *conf);
static ngx_int_t init_worker_sharedmemory(ngx_cycle_t *cycle);
static ngx_int_t ngx_http_status_code_counter_init(ngx_conf_t *cf);
static void * ngx_http_status_code_counter_create_conf(ngx_conf_t *cf);
static ngx_command_t  ngx_http_status_code_counter_commands[] = {
    { ngx_string("show_status_code_count"),
      NGX_HTTP_SRV_CONF|NGX_HTTP_LOC_CONF|NGX_CONF_FLAG,
      ngx_http_set_status_code_counter,
      0,
      0,
      NULL },
      ngx_null_command
};


static ngx_http_module_t  ngx_http_status_code_counter_module_ctx = {
    NULL,                                  /* preconfiguration */
    ngx_http_status_code_counter_init,     /* postconfiguration */
    ngx_http_status_code_counter_create_conf, /* create main configuration */
    NULL,                                  /* init main configuration */
    NULL,                                  /* create server configuration */
    NULL,                                  /* merge server configuration */
    NULL,                                  /* create location configuration */
    NULL                                   /* merge location configuration */
};


ngx_module_t  ngx_http_status_code_counter_module = {
    NGX_MODULE_V1,
    &ngx_http_status_code_counter_module_ctx,      /* module context */
    ngx_http_status_code_counter_commands,              /* module directives */
    NGX_HTTP_MODULE,                       /* module type */
    NULL,                                  /* init master */
    NULL,                                  /* init module */
    init_worker_sharedmemory,              /* init process */
    NULL,                                  /* init thread */
    NULL,                                  /* exit thread */
    NULL,                                  /* exit process */
    NULL,                                  /* exit master */
    NGX_MODULE_V1_PADDING
};

static void *
ngx_http_status_code_counter_create_conf(ngx_conf_t *cf)
{
    ngx_http_status_code_counter_conf_t    *conf;
    ngx_time_t  *time;

    conf = ngx_pcalloc(cf->pool, sizeof(ngx_http_status_code_counter_conf_t));
    if (NULL == conf) {
        return NULL;
    }

    time = ngx_timeofday();
    conf->startup = time->sec;

    return conf;
}


static ngx_int_t ngx_http_status_code_counter_handler(ngx_http_request_t *r)
{
    size_t             size;
    ngx_int_t          rc, i, j;
    ngx_buf_t         *b;
    ngx_chain_t        out;
    ngx_time_t        *time;
    ngx_http_status_code_counter_conf_t *cf;
        
    if (r->method != NGX_HTTP_GET && r->method != NGX_HTTP_HEAD) {
        return NGX_HTTP_NOT_ALLOWED;
    }

    rc = ngx_http_discard_request_body(r);

    if (rc != NGX_OK) {
        return rc;
    }

    cf = ngx_http_get_module_main_conf(r, ngx_http_status_code_counter_module);
    if (NULL == cf) {
        return NGX_ERROR;
    }

    ngx_str_set(&r->headers_out.content_type, "application/json");

    if (r->method == NGX_HTTP_HEAD) {
        r->headers_out.status = NGX_HTTP_OK;

        rc = ngx_http_send_header(r);

        if (rc == NGX_ERROR || rc > NGX_OK || r->header_only) {
            return rc;
        }
    }

    /* count number of seen status codes to determine how much buffer we need */
    j = 0;
    for(i = 0; i < NGX_HTTP_NUM_STATUS_CODES; i++ )
    {
      if(ngx_http_status_code_counts[i] > 0)
      {
        j++;
      }
    }

    size = sizeof("{}") + j * (sizeof("\"XXX\":,") + NGX_ATOMIC_T_LEN);

    b = ngx_create_temp_buf(r->pool, size);
    if (b == NULL) {
        return NGX_HTTP_INTERNAL_SERVER_ERROR;
    }

    out.buf = b;
    out.next = NULL;

    time = ngx_timeofday();

    b->last = ngx_sprintf(b->last, "{\"uptime\":%uA", time->sec - cf->startup);
    for(i = 0; i < NGX_HTTP_NUM_STATUS_CODES; i++ )
    {
      if(ngx_http_status_code_counts[i] > 0)
      {
        b->last = ngx_sprintf(b->last, ",\"%uA\":%uA", i+NGX_HTTP_OK, ngx_http_status_code_counts[i]);
      }
    }

    b->last = ngx_sprintf(b->last, "}");
    
    r->headers_out.status = NGX_HTTP_OK;
    r->headers_out.content_length_n = b->last - b->pos;

    b->last_buf = 1;

    rc = ngx_http_send_header(r);

    if (rc == NGX_ERROR || rc > NGX_OK || r->header_only) {
        return rc;
    }

    return ngx_http_output_filter(r, &out);
}


static char *ngx_http_set_status_code_counter(ngx_conf_t *cf, ngx_command_t *cmd, void *conf)
{
    ngx_http_core_loc_conf_t  *clcf;

    clcf = ngx_http_conf_get_module_loc_conf(cf, ngx_http_core_module);
    clcf->handler = ngx_http_status_code_counter_handler;

    return NGX_CONF_OK;
}

ngx_int_t
ngx_http_status_code_count_handler(ngx_http_request_t *r)
{
  if( r->headers_out.status >= NGX_HTTP_OK && r->headers_out.status < NGX_HTTP_LAST_LEVEL_500)
  {
    if (NULL != (void *)ngx_http_status_code_counts) {
      ngx_http_status_code_counts[r->headers_out.status - NGX_HTTP_OK]++;
    } 
  }

  return NGX_OK;
}

static ngx_int_t
ngx_http_status_code_counter_init(ngx_conf_t *cf)
{
    ngx_http_handler_pt        *h;
    ngx_http_core_main_conf_t  *cmcf;
    ngx_int_t i;
    int shmid = -1;
    
    cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);

    h = ngx_array_push(&cmcf->phases[NGX_HTTP_LOG_PHASE].handlers);
    if (h == NULL) {
        return NGX_ERROR;
    }

    *h = ngx_http_status_code_count_handler;

    shmid = shmget(SHMKEY, sizeof(ngx_atomic_t) * NGX_HTTP_NUM_STATUS_CODES, SHMFLAG);
    if (shmid > 0) {
      ngx_http_status_code_counts = (void *)shmat(shmid, NULL, 0);
      if (NULL != (void *)ngx_http_status_code_counts) {
        for(i = 0; i < NGX_HTTP_NUM_STATUS_CODES; i++) {
          ngx_http_status_code_counts[i] = 0;
        }
      } else {
        perror("ngx_http_status_code_counter_init:shmat");
      }
    } else { 
      perror("ngx_http_status_code_counter_init:shmget");
    }
    
    return NGX_OK;
}

static ngx_int_t init_worker_sharedmemory(ngx_cycle_t *cycle)
{
  int shmid = -1;
  
  shmid = shmget(SHMKEY, sizeof(ngx_atomic_t) * NGX_HTTP_NUM_STATUS_CODES, SHMFLAG);
  if (shmid > 0) {
    ngx_http_status_code_counts = (void *)shmat(shmid, NULL, 0);
  } else { 
    perror("ngx_http_status_code_counter_init");
  }
    
  return NGX_OK;
}