掘金 后端 ( ) • 2024-04-20 09:52

在 Go 语言的开发过程中,我们有时需要在后台执行长时间运行的任务,例如监听或轮询某些资源。但是,如果任务执行时间过长或出现意外情况导致死循环,我们通常希望能够设置一个超时机制来中止循环。这篇文章将通过一个实例详细介绍如何为 Go 语言中的无限循环设置时间限制,保证程序的健壮性和可控性。

1_nfzvO-QW0MTyj1O7LZe9fA.jpg

问题描述

我们有一个用于检查 RabbitMQ 集群节点的 Go 函数,该函数包含一个无限循环,用于不断执行检查命令。现在的需求是,如果函数运行超过3分钟,自动终止循环。

原始代码

原始代码如下:

func checkRabbitmqClusterIfForgetNode(node string) bool {
    for {
        cmd := fmt.Sprintf("docker exec -i pam_pam-rabbitmq_1 rabbitmqctl --node %s cluster_status --formatter json", node)
        res, err := common.ExecuteCommandWithoutSpace("bash", "-c", cmd)
        if err != nil {
            log.Errorf("exec cmd %v failed, response: %v error: %v", cmd, res, err)
            continue
        }
        cluster, err := mq.ParseJSON(res)
        if err != nil {
            log.Errorf("parse json %v error: %v", res, err)
            continue
        }
        if nodes := cluster.CountDiskNodes(); nodes == 2 {
            log.Infof("rabbitmq cluster node number is %v, still not forget", nodes)
            continue
        }
        return true
    }
}

添加时间限制

要为这个无限循环设置时间限制,我们可以使用 Go 语言的 time 包。具体方法是使用 time.After 函数来创建一个超时通道,当达到指定时间后,超时通道会接收到一个时间信号。

改进后的代码如下:

func checkRabbitmqClusterIfForgetNode(node string) bool {
    timeout := time.After(3 * time.Minute) // 设置超时时间为3分钟
    for {
        select {
        case <-timeout:
            log.Info("Operation timed out")
            return false // 时间超过3分钟后退出循环
        default:
            cmd := fmt.Sprintf("docker exec -i pam_pam-rabbitmq_1 rabbitmqctl --node %s cluster_status --formatter json", node)
            res, err := common.ExecuteCommandWithoutSpace("bash", "-c", cmd)
            if err != nil {
                log.Errorf("exec cmd %v failed, response: %v error: %v", cmd, res, err)
                continue
            }
            cluster, err := mq.ParseJSON(res)
            if err != nil {
                log.Errorf("parse json %v error: %v", res, err)
                continue
            }
            if nodes := cluster.CountDiskNodes(); nodes == 2 {
                log.Infof("rabbitmq cluster node number is %v, still not forget", nodes)
                continue
            }
            return true
        }
    }
}

在这段代码中,我们使用了 select 语句来等待超时事件。如果 timeout 通道接收到了超时信号,则函数将打印超时信息并返回 false,这表明函数因为超时而终止。这种方式非常适合处理可能无限执行的循环任务,确保它们在给定时间后能够被适当中止。

结论

设置时间限制是提高长时间运行的 Go 程序健壮性的一种有效方法。通过使用 time.Afterselect 语句,我们能够控制程序在指定时间内完成任务,从而避免程序在意外情况下无限制地运行下去。这不仅保证了程序的效率,也提高了其可维护性和稳定性。