Implementing Request Rate Limiting in Moco with a Custom LimitHandler

Overview

The built‑in documentation of the Moco API does not provide a way to limit the number of requests that can be made within a short time window. A custom LimitHandle response handler is introduced to enforce a per‑request‑signature rate limit, independent of sessions or users. It is useful for protecting endpoints such as lottery draws, order submissions, or any operation that should not be invoked repeatedly in rapid succession during testing.

Implementation Details

The handler records the timestamp of the last request for each unique request signature (the concatenation of the request URI and its body). A ConcurrentHashMap<String, Long> stores these timestamps, guaranteeing thread‑safe updates when Moco processes concurrent requests.

package com.fun.moco.support;

import com.fun.utils.Time;
import com.github.dreamhead.moco.HttpRequest;
import com.github.dreamhead.moco.MocoConfig;
import com.github.dreamhead.moco.ResponseHandler;
import com.github.dreamhead.moco.handler.AbstractResponseHandler;
import com.github.dreamhead.moco.internal.SessionContext;
import com.github.dreamhead.moco.model.MessageContent;
import com.google.common.base.Function;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * ResponseHandler that limits request frequency.
 */
@SuppressWarnings("all")
public class LimitHandle extends AbstractResponseHandler {
    private final ResponseHandler limit;      // response when request is limited
    private final ResponseHandler unlimit;    // response when request is allowed
    private final Map<String, Long> timestamps = new ConcurrentHashMap<>();
    private final int interval;               // minimum interval in milliseconds

    private LimitHandle(ResponseHandler limit, ResponseHandler unlimit, int interval) {
        this.limit = limit;
        this.unlimit = unlimit;
        this.interval = interval;
    }

    public static ResponseHandler newSeq(ResponseHandler limit, ResponseHandler unlimit, int interval) {
        return new LimitHandle(limit, unlimit, interval);
    }

    @Override
    public void writeToResponse(SessionContext context) {
        HttpRequest request = (HttpRequest) context.getRequest();
        String uri = request.getUri();
        MessageContent content = request.getContent();
        String key = uri + content;                     // unique signature of the request
        (limited(key) ? limit : unlimit).writeToResponse(context);
    }

    @Override
    public ResponseHandler apply(MocoConfig config) {
        // Preserve normal Moco configuration handling for response IDs
        if (config.isFor(MocoConfig.RESPONSE_ID)) {
            return super.apply(config);
        }
        return new LimitHandle(limit, unlimit, interval);
    }

    /**
     * Determines whether the request identified by {@code info} should be limited.
     * The method stores the current timestamp and returns {@code true} when the
     * elapsed time since the previous request exceeds {@code interval}.
     */
    public boolean limited(String info) {
        long now = Time.getTimeStamp();
        long previous = timestamps.getOrDefault(info, 0L);
        timestamps.put(info, now);
        return now - previous > interval;
    }
}

Key points:

Signature : The map key is {@code uri + content}, ensuring that identical requests are throttled together.

Interval logic : If the time difference between the current request and the stored timestamp is greater than the configured {@code interval} (in milliseconds), the request is considered “unlimited” and the normal response is returned; otherwise the “limited” response is sent.

Thread safety : {@code ConcurrentHashMap} guarantees safe updates when multiple threads invoke the handler concurrently.

Helper Factory Methods

Static convenience methods create ResponseHandler instances with default or custom intervals. The default interval is 1000 ms.

/**
 * Create a limiter with the default 1000 ms interval.
 */
static ResponseHandler limit(ResponseHandler limited, ResponseHandler unlimited) {
    return limit(limited, unlimited, 1000);
}

/**
 * Create a limiter with a custom interval (milliseconds).
 */
static ResponseHandler limit(ResponseHandler limited, ResponseHandler unlimited, int interval) {
    return LimitHandle.newSeq(limited, unlimited, interval);
}

/**
 * Overloads that accept raw string bodies or JSON objects.
 */
static ResponseHandler limit(String limitedBody, String unlimitedBody) {
    return limit(contentResponse(limitedBody), contentResponse(unlimitedBody));
}

static ResponseHandler limit(JSONObject limitedJson, JSONObject unlimitedJson) {
    return limit(jsonResponse(limitedJson), jsonResponse(unlimitedJson));
}

static ResponseHandler limit(String limitedBody, String unlimitedBody, int interval) {
    return limit(contentResponse(limitedBody), contentResponse(unlimitedBody), interval);
}

static ResponseHandler limit(JSONObject limitedJson, JSONObject unlimitedJson, int interval) {
    return limit(limitedJson.toString(), unlimitedJson.toString(), interval);
}

Integration Example

In a Moco configuration file, the limiter can be attached to any endpoint:

server = Moco.httpServer(8080);

ResponseHandler success = Moco.text("OK");
ResponseHandler tooFast = Moco.text("Too many requests");

// Apply a 2‑second limit to /order endpoint
server.request(Moco.by(uri("/order")))
      .response(LimitHandle.limit(success, tooFast, 2000));

This configuration returns “OK” when the interval between two identical requests exceeds 2000 ms; otherwise it returns “Too many requests”. The approach works for any HTTP method supported by Moco.