From fe6674c3b7a24660f20a2a6604327d2342f9e0fd Mon Sep 17 00:00:00 2001
From: anrs <anders.hu@shopee.com>
Date: Thu, 2 Sep 2021 17:55:58 +0800
Subject: [PATCH] filter workloads in/out

---
 cmd/workload/cmd.go  |  8 +++++
 cmd/workload/list.go | 69 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/cmd/workload/cmd.go b/cmd/workload/cmd.go
index 7e48a58..c96915b 100644
--- a/cmd/workload/cmd.go
+++ b/cmd/workload/cmd.go
@@ -109,6 +109,14 @@ func Command() *cli.Command {
 						Name:  "limit",
 						Usage: "limit data size",
 					},
+					&cli.StringSliceFlag{
+						Name:  "match-ip",
+						Usage: "filter by IP",
+					},
+					&cli.StringSliceFlag{
+						Name:  "skip-ip",
+						Usage: "filter out IP",
+					},
 				},
 			},
 			{
diff --git a/cmd/workload/list.go b/cmd/workload/list.go
index 0862bd8..4c6c606 100644
--- a/cmd/workload/list.go
+++ b/cmd/workload/list.go
@@ -3,6 +3,7 @@ package workload
 import (
 	"context"
 	"io"
+	"strings"
 
 	"github.com/projecteru2/cli/cmd/utils"
 	"github.com/projecteru2/cli/describe"
@@ -19,6 +20,8 @@ type listWorkloadsOptions struct {
 	entrypoint string
 	nodename   string
 	labels     map[string]string
+	matchIPs   []string
+	skipIPs    []string
 }
 
 func (o *listWorkloadsOptions) run(ctx context.Context) error {
@@ -47,10 +50,74 @@ func (o *listWorkloadsOptions) run(ctx context.Context) error {
 		workloads = append(workloads, w)
 	}
 
+	f := filter{
+		ips:       o.matchIPs,
+		skipIPs:   o.skipIPs,
+		nodenames: []string{},
+	}
+	if len(o.nodename) > 0 {
+		f.nodenames = append(f.nodenames, o.nodename)
+	}
+
+	workloads = f.filterIn(workloads)
+
 	describe.Workloads(workloads...)
 	return nil
 }
 
+type filter struct {
+	ips       []string
+	skipIPs   []string
+	nodenames []string
+}
+
+func (wf filter) filterIn(workloads []*corepb.Workload) []*corepb.Workload {
+	ans := []*corepb.Workload{}
+	for _, workload := range workloads {
+		if !wf.skip(workload) {
+			ans = append(ans, workload)
+		}
+	}
+	return ans
+}
+
+func (wf filter) skip(workload *corepb.Workload) bool {
+	if workload == nil {
+		return true
+	}
+	if len(wf.nodenames) > 0 && !wf.hasIntersection(wf.nodenames, []string{workload.Nodename}) {
+		return true
+	}
+
+	// Don't skip any workload if there isn't Status.
+	if workload.Status == nil {
+		return false
+	}
+
+	ips := []string{}
+	for _, cidr := range workload.Status.Networks {
+		ips = append(ips, strings.Split(cidr, "/")[0])
+	}
+
+	return (len(wf.ips) > 0 && !wf.hasIntersection(wf.ips, ips)) ||
+		(len(wf.skipIPs) > 0 && wf.hasIntersection(wf.skipIPs, ips))
+}
+
+func (wf filter) hasIntersection(a, b []string) bool {
+	hash := map[string]bool{}
+	for _, v := range a {
+		hash[v] = true
+	}
+
+	for _, v := range b {
+		if _, exists := hash[v]; exists {
+			return true
+		}
+	}
+
+	return false
+}
+
 func cmdWorkloadList(c *cli.Context) error {
 	client, err := utils.NewCoreRPCClient(c)
 	if err != nil {
@@ -64,6 +131,8 @@ func cmdWorkloadList(c *cli.Context) error {
 		nodename:   c.String("nodename"),
 		labels:     utils.SplitEquality(c.StringSlice("label")),
 		limit:      c.Int64("limit"),
+		matchIPs:   c.StringSlice("match-ip"),
+		skipIPs:    c.StringSlice("skip-ip"),
 	}
 	return o.run(c.Context)
 }