Skip to content

Commit

Permalink
Feature: optional rate limiter on concurrent sitemap functions execut…
Browse files Browse the repository at this point in the history
…ed (#62)
  • Loading branch information
AjaniBilby authored Nov 17, 2023
1 parent f0a6e9a commit aacec50
Show file tree
Hide file tree
Showing 5 changed files with 190 additions and 29 deletions.
24 changes: 16 additions & 8 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,19 +25,26 @@

## 📚 Table Of Contents

- [Getting Started](#-getting-started)
- [✨ Features](#-features)
- [📚 Table Of Contents](#-table-of-contents)
- [🚀 Getting Started](#-getting-started)
- [Installation](#installation)
- [Usage](#usage)
- [Runtime Generation](#runtime-generation)
- [Build time Generation](#build-time-generation)
- [Guides](#-guides)
- [Generate Sitemap for Dynamic Routes](#generate-sitemap-for-dynamic-routes)
- [Exclude Route](#exclude-route)
- [Google: News, Image and Video](#google-news-image-and-video)
- [Splitting Sitemaps](#splitting-sitemaps)
- [Caching](#caching)
- [API Reference](#-api-reference)
- [📝 Guides](#-guides)
- [Generate Sitemap for Dynamic Routes](#generate-sitemap-for-dynamic-routes)
- [Exclude Route](#exclude-route)
- [Google: News, Image and Video](#google-news-image-and-video)
- [Splitting Sitemaps](#splitting-sitemaps)
- [Caching](#caching)
- [📖 API Reference](#-api-reference)
- [Config](#config)
- [RobotsTxtOptions](#robotstxtoptions)
- [👤 Author](#-author)
- [🤝 Contributing](#-contributing)
- [Show your support](#show-your-support)
- [Acknowledgements](#acknowledgements)


## 🚀 Getting Started
Expand Down Expand Up @@ -254,6 +261,7 @@ createSitemapGenerator({
- `alternateRefs = []`: (*optional*) default multi language support by unique url for all entries.
- `generateRobotsTxt = false`: (*optional*) generate `robots.txt` file.
- `robotsTxtOptions`: (*optional*) options for generating `robots.txt` [details](#RobotsTxtOptions).
- `rateLimit`: (*optional*) limits the number of `sitemap` functions that can execute at once.

**Runtime only**
- `headers = {}`: (*optional*) headers for the sitemap and robots.txt response.
Expand Down
49 changes: 28 additions & 21 deletions src/builders/sitemap.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { XMLBuilder } from 'fast-xml-parser';
import { cleanDoubleSlashes } from 'ufo';
import { XMLBuilder } from 'fast-xml-parser';
import type {
SitemapEntry,
AlternateRef,
Expand All @@ -10,6 +10,7 @@ import type {
EntryContext
} from '../lib/types';
import { getBooleanValue, getOptionalValue } from '../utils/xml';
import { RateLimiter } from '../utils/rate-limiter';
import { getEntry } from '../utils/entries';
import { truthy } from '../utils/truthy';
import { chunk } from '../utils/chunk';
Expand Down Expand Up @@ -138,37 +139,43 @@ export type GetSitemapParams = {
request: Request;
};

export async function buildSitemap(params: GetSitemapParams): Promise<string> {
async function IngestRoutes(params: GetSitemapParams) {
const { config, context, request } = params;

const routes = Object.keys(context.manifest.routes);

const entriesPromise = routes.map(route =>
getEntry({ route, config, context, request })
);
if (config.rateLimit) {
const limiter = new RateLimiter(config.rateLimit);

const entriesPromise = routes.map(async route => {
await limiter.allocate();
const out = await getEntry({ route, config, context, request });
limiter.free();
return out;
});

const entries = (await Promise.all(entriesPromise)).flat().filter(truthy);
return (await Promise.all(entriesPromise)).flat().filter(truthy);
} else {
const entriesPromise = routes.map(route =>
getEntry({ route, config, context, request })
);

return buildSitemapXml(entries, config.format);
return (await Promise.all(entriesPromise)).flat().filter(truthy);
}
}

export async function buildSitemaps(params: GetSitemapParams) {
const { config, context, request } = params;
export async function buildSitemap(params: GetSitemapParams): Promise<string> {
const entries = await IngestRoutes(params);
return buildSitemapXml(entries, params.config.format);
}

if (!config.size)
export async function buildSitemaps(params: GetSitemapParams) {
if (!params.config.size)
throw new Error('You must specify a size for sitemap splitting');

const routes = Object.keys(context.manifest.routes);

const entriesPromise = routes.map(route =>
getEntry({ route, config, context, request })
);

const entries = (await Promise.all(entriesPromise)).flat().filter(truthy);

const sitemaps = chunk(entries, config.size);
const entries = await IngestRoutes(params);

return sitemaps.map(urls => buildSitemapXml(urls, config.format));
const sitemaps = chunk(entries, params.config.size);
return sitemaps.map(urls => buildSitemapXml(urls, params.config.format));
}

export function buildSitemapIndex(sitemaps: string[], config: Config): string {
Expand Down
5 changes: 5 additions & 0 deletions src/lib/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,11 @@ export interface RemixSitemapConfig {
*/
headers?: HeadersInit;

/**
* Limit the number of routes processed at a time
*/
rateLimit?: number;

/**
* The cache to use.
*/
Expand Down
91 changes: 91 additions & 0 deletions src/utils/__tests__/rate-limiter.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
import { RateLimiter } from '../rate-limiter';

describe('RateLimiter', () => {
it('should not allow more than the rate limit to run concurrently', async () => {
const limit = new RateLimiter(3);
let activeTasks = 0;

const increaseActive = () => {
activeTasks++;
};

const decreaseActive = () => {
activeTasks--;
limit.free();
};

const tasks = Array(5)
.fill(null)
.map(async () => {
await limit.allocate();

increaseActive();
expect(activeTasks).toBeLessThanOrEqual(3);

// Mock task duration (randomness to affect completion order)
await new Promise(resolve =>
setTimeout(resolve, 500 + 500 * Math.random())
);
decreaseActive();
});
await Promise.all(tasks);
});

it('should queue tasks and execute them in order', async () => {
const limit = new RateLimiter(1);
const taskOrder: number[] = [];
const mockTask = async (id: number) => {
await limit.allocate();

// Make each task successively longer to ensure execution completion order
await new Promise(resolve => setTimeout(resolve, 100));

taskOrder.push(id);
limit.free();
};

const tasks = [1, 2, 3, 4].map(id => mockTask(id));
await Promise.all(tasks);
expect(taskOrder).toEqual([1, 2, 3, 4]); // Ensure tasks are executed in the expected order
expect(limit.getProcessing()).toEqual(0);
expect(limit.getWaiting()).toEqual(0);
});

it('ensure too many tasks never run at once', async () => {
const maxConcurrent = 20;
const limit = new RateLimiter(maxConcurrent);
const mockTask = async () => {
await limit.allocate();

expect(limit.getProcessing()).toBeLessThanOrEqual(maxConcurrent);

// Mock task duration (randomness to affect completion order)
await new Promise(resolve => setTimeout(resolve, 500 * Math.random()));

expect(limit.getProcessing()).toBeLessThanOrEqual(maxConcurrent);

limit.free();
};
await Promise.all(
Array(100)
.fill(0)
.map(() => mockTask())
);

expect(limit.getProcessing()).toEqual(0);
expect(limit.getWaiting()).toEqual(0);
});

it('should handle rapid calls to free correctly', async () => {
const limit = new RateLimiter(1);
const mockTask = async () => {
await limit.allocate();
limit.free();
limit.free(); // Intentional rapid call to free
};

await Promise.all([mockTask(), mockTask()]);
expect(limit.getProcessing()).toEqual(0);
expect(limit.getWaiting()).toEqual(0);
});
});
50 changes: 50 additions & 0 deletions src/utils/rate-limiter.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
/**
* Usage:
* const limit = new RateLimiter(3);
* Promise.all(arr.map(async x => {
* await limit.allocate(); // waits until there are less than three active
* await longTask(x);
* limit.free(); // mark the task as done, allowing another to start
* }))
*/

export class RateLimiter {
rate: number;
queue: Array<() => void>;
active: number;

constructor(rate: number) {
this.rate = rate;
this.queue = [];
this.active = 0;
}

allocate(): Promise<void> {
return new Promise(res => {
if (this.active < this.rate) {
this.active++;
res();
} else {
this.queue.push(res);
}
});
}

free(): void {
const first = this.queue.shift();
if (first) {
// Don't change activity count
// just substitute the finished work with new work
first();
} else {
this.active = Math.max(0, this.active - 1);
}
}

getProcessing(): number {
return this.active;
}
getWaiting(): number {
return this.queue.length;
}
}

0 comments on commit aacec50

Please sign in to comment.