将一个切片拆分为 N 个切片

Question

我正在尝试实现一个将 TCP 端口切片分成 x 个其他切片的功能。这些切片将发送给将扫描这些端口的工作人员，因此 x 由工作人员的数量设置。

代码如下：

// createJobs split portsToScan from a specified protocol into an equal number
// of jobs that will be returned.
func (t *Target) createJobs(proto string) ([]jobMsg, error) {
    // init jobs slice
    jobs := []jobMsg{}

    // check protocol accordance
    if _, ok := t.portsToScan[proto]; !ok {
        return nil, fmt.Errorf("no such protocol %q in current protocol list", proto)
    }

    // if proto is ICMP, we do not need to scan ports
    if proto == "icmp" {
        return []jobMsg{
            jobMsg{ip: t.ip, protocol: proto},
        }, nil
    }

    step := (len(t.portsToScan[proto]) + t.workers - 1) / t.workers

    for i := 0; i < len(t.portsToScan[proto]); i += step {
        batch := t.portsToScan[proto][i:min(i+step, len(t.portsToScan[proto]))]

        jobs = append(jobs, jobMsg{
            ip:       t.ip,
            protocol: proto,
            ports:    batch,
        })
    }

    return jobs, nil
}

这里是对应的单元测试：

func TestTarget_createJobs(t *testing.T) {
    tests := []struct {
        name         string
        pts          map[string][]string
        workersCount int
        wantErr      bool
    }{
        {
            name:         "5-1",
            pts:          map[string][]string{"tcp": []string{"1", "2", "3", "4", "5"}},
            workersCount: 1,
        },
        {
            name:         "5-2",
            pts:          map[string][]string{"tcp": []string{"1", "2", "3", "4", "5"}},
            workersCount: 2,
        },
        {
            name:         "5-3",
            pts:          map[string][]string{"tcp": []string{"1", "2", "3", "4", "5"}},
            workersCount: 3,
        },
        {
            name:         "5-4",
            pts:          map[string][]string{"tcp": []string{"1", "2", "3", "4", "5"}},
            workersCount: 4,
        },
        {
            name:         "5-5",
            pts:          map[string][]string{"tcp": []string{"1", "2", "3", "4", "5"}},
            workersCount: 5,
        },
    }
    for _, tt := range tests {
        t.Run(tt.name, func(t *testing.T) {
            tg := &Target{
                portsToScan: tt.pts,
                workers:     tt.workersCount,
            }
            got, err := tg.createJobs("tcp")
            if (err != nil) != tt.wantErr {
                t.Errorf("Target.createJobs() error = %v, wantErr %v", err, tt.wantErr)
                return
            }
            if len(got) != tt.workersCount {
                t.Errorf("Target.createJobs() = %d, wanted %d jobs; joblist %v", len(got), tt.workersCount, got)
            }
        })
    }
}


func min(a, b int) int {
    if a <= b {
        return a
    }
    return b
}

测试的输出给了我这个结果：

--- FAIL: TestTarget_createJobs/5-4 (0.00s)
scan_test.go:309: Target.createJobs() = 3, wanted 4 jobs; joblist [{ 0  tcp [1 2]} { 0  tcp [3 4]} { 0  tcp [5]}]

初始端口列表保存在t.portsToScan[proto] 中，工作人员的数量（因此我要创建的切片数量）由t.workers 设置。

最后，len(jobs) 必须等于 t.workers，但我不知道该怎么做。

Answer 1

您的算法使用step 作为批次大小：

step := (len(t.portsToScan[proto]) + t.workers - 1) / t.workers

这不是最佳尺寸。例如，如果您有 4 个要扫描的端口和 3 个工作人员，这将导致 step = 2，这意味着您将只有 2 个工作 (2+2=4)。但最好有 3 个批次（大小为2+1+1=4）。

所以批次的大小应该是

defSize := len(t.portsToScan[proto]) / t.workers

这样做的问题是，如果长度不是t.workers 的倍数，则最后一些元素（端口）将不会分配给任何作业。对所有作业都使用defSize+1 太多了。

所以最佳解决方案是在“中间”：一些作业将有defSize 端口要扫描，而有些则有defSize+1。有多少必须有defSize+1？如果所有人都拥有defSize，则与失踪一样多：

numBigger := len(t.portsToScan[proto]) - defSize*t.workers

请注意，如果要扫描的端口数少于工作人员，则上述计算得出defSize=0，因此一些工作人员将得到0 端口进行扫描，而一些工作人员将得到1。没关系，但您不应添加带有 0 端口的作业进行扫描。

使用这个分布：

defSize := len(t.portsToScan[proto]) / t.workers
numBigger := len(t.portsToScan[proto]) - defSize*t.workers

size := defSize+1
for i, idx := 0, 0; i < t.workers; i++ {
    if i == numBigger {
        size--
        if size == 0 {
            break // 0 ports left to scan
        }
    }
    jobs = append(jobs, jobMsg{
        ip:       t.ip,
        protocol: proto,
        ports:    t.portsToScan[proto][idx : idx+size],
    })
    idx += size
}