Implement Koa request queue middleware based on Promise.resolve

The author of this article is the front-end engineer of 360 Qiwu Troupe

foreword

AIGCI am working on a project recently , and the backend is based on Koa2the implementation. One of the requirements is to call the brother business line server AIGCcapability to generate pictures. However, since AIGCthe project of the brother business line is still in the testing stage, the server resources that can be provided are limited. When the concurrent request resources cannot be satisfied, it will respond with [Server is busy], which is very Cunfriendly to us displayed on the terminal. Based on the current predicament, the first solution that comes to mind is Kafkaor RabbitMQ, but in fact, for our current user volume, it is simply overkill. So I changed my mind, can I use js to simulate the queue to solve the problem? The answer is: yes, the Promisequeue Resolve!

analyze

Resolveunderstanding

PromiseThe core usage of is to use Resolvethe function to do chain transfer. For example:

new Promise(resolve => {
  resolve('ok')
}).then(res => {
  console.log(res)
})
// 输出结果:ok

Through the above example, we can understand that Resolvechanging Promisethe state of the object to is pendingcalled fullfilledwhen the asynchronous operation is successful, and the result of the asynchronous operation is passed as a parameter.

Core point: asynchronous

At this point, a question is raised : If I resolveput all the callback functions into a queue, Promisewill it always be in pendingthe state? pendingThe state means that the then function is always in the state, and the function cannot be executed until the function waittingin the queue is executed?resolvethen

Promisefunction that creates blocking

const queue = []
new Promise(resolve => {
  queue.push(resolve)
}).then(res => {
  console.log(res)
})
// 输出结果:Promise {<pending>}

queue[0]('ok')
// 输出结果:ok

For proof, directly map:

6f6ba4680028921155a502a9f4ab0f83.png
image.png

asynchronous to synchronous

Koa2It belongs to the onion model. When the request comes, we need to call nextthe function to continue to penetrate, and our requirement is to limit the flow, which means that we have to block the request. At this moment, awaitI raised my hands and blocked this kind of shameless thing. ah!

const queue = []
const fn = async = () => {
  await new Promise(resolve => {
    queue.push(resolve)
  })
  // ...一大波操作
}
// queue[0]()

If queue[0]not executed, the code will always be blocked. Then we can use await to write a middleware to achieve blocking certain apirequirements.

// 阻塞所有请求,知道queue中的resolve函数被执行才会执行next
const queue = []
module.exports = function () {
  return async function (ctx, next) {
    await new Promise(resolve => {
      queue.push(resolve)
    })
    await next();
  };
};

Implement middleware

The principles and ideas are straightened out, so let’s get started. Without further ado, post the code:

const resolveMap = {};

/**
 * 请求队列
 * @param {*} ctx
 * @param {*} ifer 是否是图生图
 * @param {*} maxReqNumber 最大请求数量
 * @returns
 * @description
 * 使用promise解决请求队列问题
 * 1. 用于限制aicg的并发请求
 * 2. 当文生图是,根据风格分类存储resolve,当前请求响应完成时,触发消费队列中下一个请求
 * 3. 当图生图是,直接存储resolve到image风格,当前请求响应完成时,触发消费队列中下一个请求
 * 4. 同时处理的请求数量不超过maxReqNumber个,否则加入队列等待。
 */
function requestQueue(ctx, maxReqNumber) {
  const params = ctx.request.body ?? ctx.request.query ?? ctx.request.params ?? {};
  const style = params.style ?? 'pruned_cgfull';

  resolveMap[style] = resolveMap[style] || { list: [], processNumber: 0 };
  const currentResolve = resolveMap[style];

  ((currentResolve) => {
    ctx.res.on('close', () => {
      saveNumberMinus(currentResolve);
      // 当前请求响应完成时,触发消费队列中下一个请求
      if (currentResolve.list.length !== 0) {
        const node = currentResolve.list.shift();
        node.resolve();
        currentResolve.processNumber++;
      }
      currentResolve = null;
    });
  })(currentResolve);
  // 当前请求正在处理中,将resolve存储到队列中
  if (currentResolve.processNumber + 1 > maxReqNumber) {
    // 利用promise阻塞请求
    return new Promise((resolve, reject) => {
      // 当前请求正在处理中,将resolve存储到队列中
      currentResolve.list.push({ resolve, reject, timeStamp: Date.now(), params });
    });
  } else {
    currentResolve.processNumber++;
    return Promise.resolve();
  }
}

module.exports = function (options = {}) {
  const { maxReqNumber = 2, apis = [] } = options;
  return async function (ctx, next) {
    const url = ctx.url;
    if (apis.includes(url)) {
      try {
        await requestQueue(ctx, maxReqNumber);
      } catch (error) {
        console.log(error);
        ctx.body = {
          code: 0,
          msg: error,
        };
        return;
      }
    }
    await next();
  };
};

const fiveMinutes = 5 * 60 * 1000;
setInterval(() => {
  Object.values(resolveMap).forEach((item) => {
    const { timeStamp, resolve } = item;
    if (Date.now() - timeStamp > fiveMinutes) {
      resolve(); // 执行并释放请求,防止用户请求因异常积压导致一直挂起
      saveNumberMinus(item);
    }
  });
}, 5 * 60 * 1000);

Here we want to emphasize one point, the use of closures. The reason why closures are used is to ensure that the object closecan be used when the event of the current request is triggered currentResolve. Because the current request is placed in the array corresponding to its own style, closethe next waiting request must be consumed, and don't forget to release resources manually.

app.jslogic part

const requsetQueue = require('./app/middleware/request-queue');
const app = new Koa();
app.use(
  requsetQueue({
    maxReqNumber: 1,
    apis: ['/api/aigc/image', '/api/aigc/textToImage', '/api/aigc/img2img'],
  })
);
app.listen(process.env.NODE_ENV === 'development' ? '9527' : '3000');

Summarize

In fact, based on Promisethe Resolvequeue, we can also implement some other functions, such as: collect some requests in the unlogged state in the front-end code, and send the request after the login is successful . I also hope that everyone can explore and discuss the implementation of Promise's other resolution capabilities.


- END -

About Qi Wu Troupe

Qi Wu Troupe is the largest front-end team of 360 Group, and participates in the work of W3C and ECMA members (TC39) on behalf of the group. Qi Wu Troupe attaches great importance to talent training, and has various development directions such as engineers, lecturers, translators, business interface people, and team leaders for employees to choose from, and provides corresponding technical, professional, general, and leadership training course. Qi Dance Troupe welcomes all kinds of outstanding talents to pay attention to and join Qi Dance Troupe with an open and talent-seeking attitude.

0c94266f4114bc66f653a577a9889d38.png

Guess you like

Origin blog.csdn.net/qiwoo_weekly/article/details/132200340
Recommended