用go实现一个任务调度类 (泛型)

用go实现一个任务调度类 (泛型)

源码地址:

https://github.com/robinfoxnan/BirdTalkServer/blob/main/server/core/workmanager.go

1.概述

实现了一个简单的任务管理系统,允许用户定义任务和工作者,并将任务分配给工作者进行处理。这个系统旨在提供一个灵活的任务管理框架,可以根据需要动态地添加和移除工作者,以及处理任务。

2.主要功能

  1. 定义了 Task 接口和 Worker 接口,用于表示任务和工作者;
  2. 提供了基础的任务类型 BaseTask 和基础的工作者类型 BaseWorker,用户可以基于这些基础类型来实现自定义的任务和工作者。需要在 BaseTask结构上继承一个新的结构,并实现Process方法;
  3. 实现了一个泛型任务管理器 Manager,用于管理工作者并分配任务给工作者。根据最大工作者个数和任务队列长度,动态地添加工作者。提供了停止所有工作者的方法,提供了方法来等待所有工作者完成任务。

3.类型和接口

3.1Task 任务接口

go 复制代码
type Task interface {
	Process()
}

任务接口定义了一个 Process() 方法,用于执行任务的处理逻辑。

3.2Worker 接口

go 复制代码
type Worker interface {
	Init(id int64, taskChan chan Task, wg *sync.WaitGroup, f WorkerCleanF)
	Start()
	Stop()
}

工作者接口定义了三个方法:

  • Init() 方法用于初始化工作者。创建后,设置工作者ID,任务通道,同步组,以及一个析构函数类似的清理函数;
  • Start() 方法用于启动工作者协程,开始处理任务;
  • Stop() 方法用于停止工作者;(关闭通道)

3.3BaseTask 结构体

这是一个最基础的示例,后续自定义结构可以包含这个结构:

go 复制代码
type BaseTask struct {
	Id int64
}

基础任务结构体包含一个任务 ID,实现了 Task 接口的 Process() 方法,用于执行任务的处理逻辑。

3.4BaseWorker 结构体

go 复制代码
type BaseWorker struct {
	Id       int64
	waitGrp  *sync.WaitGroup
	taskChan chan Task
	cleanFun WorkerCleanF
	quitChan chan struct{}
}

基础工作者结构体包含工作者 ID、等待组、任务通道、清理函数和退出通道,实现了 Worker 接口的 Init()Start()Stop() 方法,用于初始化工作者、启动工作者和停止工作者。

4. Manager 结构体

go 复制代码
type Manager[T Task, W Worker] struct {
	workers       map[int64]W    // 使用一个map管理各个协程
	maxWorkers    int64          // 最大协程数量
	workerCounter int64          // 使用原子方式计数
	taskChan      chan Task      // 任务通道
	lock          sync.Mutex     // map用的锁
	wg            sync.WaitGroup // 同步组
	newWorkerFunc func() W       // 用于创建泛型中工作者结构的函数
	exiting       int32          // 退出状态标记,防止停止过程中加入任务
	workerIdSeq   int64          // 协程序号,可以用雪花算法代替,一般应该够用
}

任务管理器结构体包含了一个工作者映射、最大工作者数量、工作者计数器、任务通道、互斥锁、等待组、新建工作者函数、退出标志和工作者 ID 序列,提供了方法来添加任务、移除工作者、等待所有工作者完成任务和停止所有工作者。

5. 使用示例

最简单的一个测试示例

go 复制代码
    manager := NewManager[Task, *BaseWorker](20, NewBaseWorker)

	// 添加示例任务到管理器
	go func() {
		for i := 0; i < 10; i++ {
			var t = &BaseTask{Id: int64(i)}
			manager.AddTask(t)
		}
	}()

	time.Sleep(time.Minute * 1)
	manager.StopAll()
	// 等待所有工作者完成任务
	manager.Wait()

我们需要重新定义一个结构用于表示任务,通常需要更多的字段

go 复制代码
type CustomTask struct {
	BaseTask
	AdditionalInfo string
    // 这里添加更多的字段
}

// 实现 Task 接口的 Run 方法,
// 必须要实现这个函数,这是任务调度的功能入口,在协程中运行
func (t *CustomTask) Process() {
	fmt.Printf("CustomTask with additional info '%s' is running\n", t.AdditionalInfo)
	// 调用父类的 Process 方法
	//t.BaseTask.Process()
}

重写测试:

go 复制代码
func TestWorkers(t *testing.T) {
	manager := NewManager[Task, *BaseWorker](20, NewBaseWorker)

	// 添加示例工作者到管理器

	// 添加示例任务到管理器
	go func() {
		for i := 0; i < 10; i++ {
			var t = &BaseTask{Id: int64(i)}
			manager.AddTask(t)
		}
        
		for i := 10; i < 16; i++ {
			var t = &CustomTask{BaseTask: BaseTask{Id: int64(i)}, AdditionalInfo: "Custom Info"}
			manager.AddTask(t)
		}
	}()

	time.Sleep(time.Minute * 1)
	manager.StopAll()
	// 等待所有工作者完成任务
	manager.Wait()
}

结论

各个语言实现的这个轮子基本都差不多。

相关推荐
LaoZhangAI4 分钟前
沉浸式翻译API深度解析:500万用户的翻译神器如何配置[2025完整指南]
前端·后端
brzhang27 分钟前
别再梭哈 Curosr 了!这 AI 神器直接把需求、架构、任务一条龙全干了!
前端·后端·架构
安妮的心动录41 分钟前
安妮的2025 Q2 Review
后端·程序员
程序员爱钓鱼41 分钟前
Go语言数组排序(冒泡排序法)—— 用最直观的方式掌握排序算法
后端·google·go
Victor3562 小时前
MySQL(140)如何解决外键约束冲突?
后端
Victor3562 小时前
MySQL(139)如何处理MySQL字符编码问题?
后端
007php0073 小时前
服务器上PHP环境安装与更新版本和扩展(安装PHP、Nginx、Redis、Swoole和OPcache)
运维·服务器·后端·nginx·golang·测试用例·php
武子康6 小时前
Java-72 深入浅出 RPC Dubbo 上手 生产者模块详解
java·spring boot·分布式·后端·rpc·dubbo·nio
椰椰椰耶7 小时前
【Spring】拦截器详解
java·后端·spring