【AI图像生成网站Golang】项目测试与优化

AI图像生成网站

目录

一、项目介绍

二、雪花算法

三、JWT认证与令牌桶算法

四、项目架构

五、图床上传与图像生成API搭建

六、项目测试与优化

---

六、项目测试与优化

在开发过程中，性能优化是保证项目可扩展性和用户体验的关键步骤。本文将详细介绍我如何使用一些主流的性能评估工具进行项目性能检测，并在此基础上进行优化。

1. 性能评估工具

我使用的性能检测工具如下：

• k6： 用于负载测试的开源工具，可以通过编写简单的javascript测试脚本，模拟高并发场景下的用户行为。
• k6 cloud： 是k6的云服务扩展，能够提供比本地运行更详细的性能监控和分析报告。
• pprof： 用于捕获CPU、内存等的使用情况进行采样分析，帮助定位性能瓶颈
• trace： 可以通过可视化的时间线查看goroutine的执行和调度情况。

在后文中，我们会使用k6运行测试文件，并根据k6 cloud上通过的请求数和RPS（Requests Per Second，每秒请求数）来简单评估性能，之后会使用pprof和trace工具查找性能瓶颈和突破口。

2. 性能评估与分析

测试流程

在进行负载测试之前，我首先进行了单用户的性能测试，用于了解每个请求的处理时长，确保基本功能的正常工作。之后，我模拟了并发场景，逐渐增加并发量测试系统的表现。

使用 k6 进行负载测试

单用户测试： 在单用户场景下，我编写了一个的k6脚本，模拟了一个用户从注册到上传作品的一套API调用流程。

import http from 'k6/http';
import { sleep, check } from 'k6';
import {generateImage, generateRandomPrompt, generateUniqueUsername,generateUniqueEmail} from "./dynamic_data.js";// 配置接口的基本 URL
const BASE_URL = 'http://localhost:8080/api/v1';// 读取 Base64 数据
const base64Image = open("image_base64.txt");// 定义一个变量来跟踪是否正在刷新 token
let refreshing = false;// k6 入口函数
export default function () {// 用户注册信息const USER_CREDENTIALS = {avatar: `data:image/png;base64,${generateImage(base64Image)}`,username: generateUniqueUsername(),password: 'password123',confirm_password: 'password123',email: generateUniqueEmail(),};// 1. 用户注册const res = http.post(`${BASE_URL}/signup`, JSON.stringify(USER_CREDENTIALS), {headers: { 'Content-Type': 'application/json' },});check(res, { 'signup succeeded': (r) => r.status === 200 });// 2. 用户登录const loginRes = http.post(`${BASE_URL}/login`, JSON.stringify(USER_CREDENTIALS), {headers: { 'Content-Type': 'application/json' },});check(loginRes, {'Login succeeded': (r) => r.status === 200,});const accessToken = loginRes.json('data.access_token');const refreshTokenValue = loginRes.json('data.refresh_token');if (!accessToken || !refreshTokenValue) {throw new Error('Login failed, tokens not received');}// 3. 创建分组const genCategory = {cover: `data:image/png;base64,${generateImage(base64Image)}`,category_name: `category_${Math.random().toString(36).substring(2, 15)}`,description: generateRandomPrompt(),};const createGroupRes = apiRequest('POST', `${BASE_URL}/createCategory`, genCategory, accessToken);check(createGroupRes, { 'Create group succeeded': (r) => r.status === 200 });// 4. 浏览分组const groupRes = apiRequest('GET', `${BASE_URL}/category`, null, accessToken);check(groupRes, { 'Browse groups succeeded': (r) => r.status === 200 });const groups = groupRes.json('data');const groupId = groups[0].category_id;// 5. 浏览作品const worksRes = apiRequest('GET', `${BASE_URL}/categoryDetail/${groupId}`, null, accessToken);check(worksRes, { 'Browse works succeeded': (r) => r.status === 200 });// 6. 提交图像处理任务const imageTaskPayload = {ori_image: `data:image/png;base64,${generateImage(base64Image)}`,category_id: groupId,prompt: generateRandomPrompt(),};const processRes = apiRequest('POST', `${BASE_URL}/processImage`, imageTaskPayload, accessToken);check(processRes, { 'Process task submitted': (r) => r.status === 200 });// 7. 上传作品const newWork = {work_image: `data:image/png;base64,${generateImage(base64Image)}`,// work_image: `data:image/png;base64,${genImage}`,category_id: groupId,prompt: generateRandomPrompt(),};const uploadRes = apiRequest('POST', `${BASE_URL}/uploadWork`, newWork, accessToken);check(uploadRes, { 'Upload work succeeded': (r) => r.status === 200 });
}// 刷新 Token 的逻辑
function refreshToken(refreshToken) {const refreshRes = http.post(`${BASE_URL}/refresh_token`, JSON.stringify({ refresh_token: refreshToken }), {headers: { 'Content-Type': 'application/json' },});check(refreshRes, {'Token refreshed': (r) => r.status === 200,});const newToken = refreshRes.json('access_token');if (!newToken) {throw new Error('Failed to refresh token');}return newToken;
}// API 请求方法
function apiRequest(method, url, body, token) {const headers = {'Content-Type': 'application/json',Authorization: `Bearer ${token}`,};let res;if (method === 'GET') {res = http.get(url, { headers });} else if (method === 'POST') {res = http.post(url, JSON.stringify(body), { headers });}// 如果遇到 401 错误，且没有正在刷新 tokenif (res.status === 401 && !refreshing) {console.log('Access token expired, refreshing...');// 设置 refreshing 为 true，避免并发请求时重复刷新refreshing = true;// 刷新 tokenconst newToken = refreshToken(refreshTokenValue);// 刷新完成后，重新发送请求res = apiRequest(method, url, body, newToken);// 刷新完成，重置 refreshing 状态refreshing = false;}return res;
}

测试结果（单用户场景）：

点击图中链接即可进入 k6 cloud 界面查看更详细的数据，可以看到单例用户的运行时长为8.3s。

进入可视化界面可以看到，单例用户的总请求为7个，平均响应时间为 4372 ms，失败率为 0% ，每秒请求数（RPS）在 1~2 之间。

在单用户场景下，系统表现良好，响应时间低，且无错误发生。接下来，我开始使用并发用户场景进行测试，模拟“高负载”下的表现（展示实验用的电脑配置低，不能模拟真正的高并发，在此仅做示例
^_^"）。

并发测试

在使用之前编写的Python API代码进行并发测试时，出现了报错： Model inference error: index 11 is out of bounds for dimension 0 with size 11 ，这时将 Python API 中这一段的 num_inference_steps=10 改小一些即可。

具体原因不在这里占用篇幅，请移步这里查看：解决并发情况下调用 Instruct-pix2pix 模型推理错误

images = pipe(prompt=request.prompt, image=image, num_inference_steps=5, image_guidance_scale=1).images

在之前的代码中加入以下部分即可进行并发测试。这里我使用了 k6 的 constant-arrival-rate 执行器模式来控制请求的速率和并发用户数。

在 constant-arrival-rate 配置下，k6 会持续不断地以固定速率发送请求，在本测试环境下，rate: 5 意味着 k6 会保持每秒发送 5 个请求，直到测试结束。这种配置适合模拟某些稳定的请求负载。

export const options = {scenarios: {my_test: {executor: 'constant-arrival-rate',rate: 5, // 每秒发送 5 个请求duration: '5m', // 持续 5 分钟preAllocatedVUs: 5, // 提前分配的虚拟用户数maxVUs: 10, // 最大虚拟用户数},},
};

当然，除了 constant-arrival-rate 之外，k6 还提供了其他几种执行器模式，适用于不同的负载场景，比如：

• virtual-users (VUs):

这是最常见的执行器模式，适用于模拟一群用户的行为，可以模拟不同用户的交替请求。

export const options = {stages: [{ duration: '10s', target: 10 }, // 10 秒内增加到 10 个虚拟用户{ duration: '1m', target: 10 },  // 保持 10 个虚拟用户 1 分钟{ duration: '10s', target: 0 },  // 10 秒内减少到 0],
};

• ramping-arrival-rate

这种模式用于模拟请求逐步增加的情况，可以通过配置startRate（起始速率）、endRate（结束速率）和 duration（持续时间）来设置请求的增长速率。

export const options = {scenarios: {my_test: {executor: 'ramping-arrival-rate',startRate: 10, // 初始请求速率每秒 10 个请求endRate: 50, // 最终请求速率每秒 50 个请求duration: '2m', // 持续 2 分钟},},
};

• shared-iterations

在这种模式下，所有虚拟用户共同执行相同数量的请求，直到所有请求完成。

export const options = {scenarios: {my_test: {executor: 'shared-iterations',iterations: 100, // 共有 100 次请求vus: 10, // 使用 10 个虚拟用户},},
};

• per-vu-iterations

这种模式下，每个虚拟用户会执行一定次数的操作。

export const options = {scenarios: {my_test: {executor: 'per-vu-iterations',vus: 10, // 使用 10 个虚拟用户iterations: 100, // 每个虚拟用户执行 100 次操作},},
};

并发测试完成后，进入 k6 cloud 界面查看测试数据。

可以看出，虽然我使用了 constant-arrival-rate 配置，但是图中的请求数量和响应数量很不稳定。

并且可以观察到，随着并发量的增加，一些需要调用API的接口响应时间明显上升，尤其是 p95 和 p99 。

初步判断是 API 调用部分出现了阻塞。

pprof 分析

使用 pprof 跟踪代码的 CPU 占用率发现，耗时最长的是 net.(*Resolver) lookupIP ，占用了 56.46% 的 CPU 时间。该函数负责 DNS 查询，用于将域名解析为IP地址，解析流程通常包括：

• 本地缓存查询
• 向 DNS 服务器发送请求
• 等待服务器返回解析结果

在项目中，如果短时间内发起大量外部请求，且每次请求都需要进行完整的 DNS 查询流程，将会导致 DNS 查询成为性能瓶颈。

此外，SetupRouter 中的日志记录（19.24%）和身份验证（12.48%）中间件是系统的次要瓶颈。

3. 优化方案

在 Go 中，默认 HTTP 客户端会在每次请求时触发新的 DNS 查询，如果复用 htttp.Client，则可以有效减少 DNS 查询次数。

同时，将外部 API 调用解耦到 Goroutine 中异步执行，可以让主线程快速处理其他任务，如日志写入等。

我们还需要引入 redis 来缓存一些会重复查询的信息，如已经验证过的用户信息等。

4. 优化实践

4.1 引入 Goroutine 解耦外部 API 调用

之前，外部 API 调用在主线程中直接执行，导致主线程需要等待外部服务返回结果。我们通过以下方式进行优化：

（1）图像上传与处理的异步化：

• 将图像上传和处理 API 任务拆分为两个 Goroutine。
• 使用 Redis 和 RabbitMQ 实现消息队列，将任务放入队列中，由独立的消费者处理，避免主线程阻塞。
• 使用 sync.WaitGroup 和 channel 进行任务同步与通信。

var wg sync.WaitGroupwg.Add(3) // 我们要等待两个消费者 goroutinego func() {defer wg.Done()controller.ConsumeUploadTasks("upload_image")}()go func() {defer wg.Done()controller.ConsumeProcessTasks("process_image")}()

（2）HTTP 客户端复用：

• 创建一个全局的 http.Client，启用连接池复用。

var ossClient *oss.Clientfunc init() {// 确保配置已经加载if err := settings.Init(); err != nil {log.Fatalf("failed to load settings: %v", err)}// 初始化 OSS 客户端配置cfg := oss.LoadDefaultConfig().WithRegion(settings.Conf.Region).WithCredentialsProvider(credentials.NewStaticCredentialsProvider(settings.Conf.AccessKeyID, settings.Conf.AccessKeySecret, ""))// 创建全局 OSS 客户端实例ossClient = oss.NewClient(cfg)
}// UploadImageToOSS 上传图片的函数
func UploadImageToOSS(base64Image string) (string, error) {if strings.HasPrefix(base64Image, "data:image") {base64Image = strings.Split(base64Image, ",")[1]}// 解码 Base64 数据imageData, err := base64.StdEncoding.DecodeString(base64Image)if err != nil {return "", fmt.Errorf("failed to decode base64 data: %v", err)}// 生成文件名objectName := GenerateUniqueFileName("png")// 创建上传对象的请求request := &oss.PutObjectRequest{Bucket: oss.Ptr(settings.Conf.Bucket), // 存储空间名称Key:    oss.Ptr(objectName),           // 对象名称Body:   bytes.NewReader(imageData),    // 图片数据}// 上传文件_, err = ossClient.PutObject(context.TODO(), request)if err != nil {return "", fmt.Errorf("failed to upload image to OSS: %v", err)}// 生成文件的 URLurl := fmt.Sprintf("%s/%s", settings.Conf.Endpoint, objectName)return url, nil
}

4.2 身份验证优化

对于身份验证逻辑，我引入了 Redis 缓存来避免频繁查询数据库。

// CacheUserInfo 缓存用户信息
func CacheUserInfo(userInfo map[string]interface{}) error {key := fmt.Sprintf("user_info:%s", userInfo["user_id"])err := client.HMSet(ctx, key, userInfo).Err()if err != nil {return fmt.Errorf("failed to HMSet: %w", err)}err = client.Expire(ctx, key, time.Duration(settings.Conf.RedisConfig.DefaultTTL)).Err()if err != nil {return fmt.Errorf("failed to set expire: %w", err)}// 如果都成功，返回 nilreturn nil}// GetCachedUserInfo 获取缓存的用户信息
func GetCachedUserInfo(userID string) (map[string]string, error) {key := fmt.Sprintf("user_info:%s", userID)return client.HGetAll(ctx, key).Result()
}

4.3 效果测试

在对项目进行优化后，通过 k6 cloud 工具对性能进行对比测试发现，系统的并发能力提升了近 7 倍，总请求数较之前增长了 2.3 倍，系统处理请求的效率得到了显著提升。

项目	优化前	优化后
总请求数	598	1.4K
请求速率峰值（reqs/s）	6.33	44.33
95% 的响应时间（ms）	32132	3.3

在 k6 可视化界面，可以看到在测试初期，请求速率（RPS）和响应时间会快速上升到一个较高的峰值。这说明系统能够快速响应并处理高并发请求。

峰值过后，请求速率和响应时间会迅速下降到一个较低的数值，并且请求速率始终保持一定的小幅波动。这是因为在高并发情况下，CPU 和内存资源的分配会随着请求数量的波动而调整，导致 RPS 的小幅变化。

在优化前，系统中耗时最多的四个请求占用大量时间，优化后有了显著的降低。

DNS 查询的 CPU 占用率从 56.46% 降低到 21.16% 。

主线程负载显著减轻，高耗时任务被转移到 Goroutine 中异步处理。

在 trace 可视化界面中，能看到Goroutines 使用数量平稳，分布较为均匀，未出现 Goroutine 阻塞的尖峰现象。同时运行的 Goroutines 数量略高，但在合理范围内。

Heap 的内存分配曲线波动较大，但回收相对及时，未出现内存泄漏或超量分配的现象。GC 活跃，Heap 内存峰值控制在可接受范围。

Threads 数量略高，可能与 HTTP 客户端连接池、Redis 以及 RabbitMQ 的 I/O 操作相关。

GC 曲线与 Heap 使用曲线基本同步，垃圾回收频繁但时间较短。未来会考虑进一步优化数据结构设计，避免频繁分配临时对象。

网络 I/O 分布均匀，RabbitMQ 和 Redis 的使用改善了对数据库的直接依赖。外部 API 调用已异步化，但响应时间稍长。

对各个 Proc，PU 核心的使用较为均匀，任务分配合理。Goroutine 与 Proc 的调度平衡较好，未观察到 Proc 长时间空闲。

在未来的工作中，将会考虑尝试将 RabbitMQ 的消费者分布式部署，以进一步增加并发量。