[AI生成] k8s ThreadSafeStore使用案例

package main

import (
	"fmt"
	"sync"

	"k8s.io/client-go/tools/cache"
)

// 1. 定义一个模拟的 K8s 资源结构体（如 Pod/Service）
type MyPod struct {
	Name      string
	Namespace string
	NodeName  string // 我们用它做自定义索引
	Status    string
}

// 2. 自定义 Key 函数，直接处理我们的 *MyPod 类型
func podKeyFunc(obj interface{}) (string, error) {
	pod, ok := obj.(*MyPod)
	if !ok {
		return "", fmt.Errorf("object is not a *MyPod")
	}
	if pod.Namespace == "" {
		return pod.Name, nil
	}
	return pod.Namespace + "/" + pod.Name, nil
}

func main() {
	// ==========================
	// 步骤1：创建 ThreadSafeStore
	// ==========================
	// 先定义索引：按 Namespace + 按 NodeName
	indexers := cache.Indexers{
		"namespace": func(obj interface{}) ([]string, error) {
			pod, ok := obj.(*MyPod)
			if !ok {
				return nil, fmt.Errorf("object is not a *MyPod")
			}
			return []string{pod.Namespace}, nil
		},
		"node": func(obj interface{}) ([]string, error) {
			pod, ok := obj.(*MyPod)
			if !ok {
				return nil, fmt.Errorf("object is not a *MyPod")
			}
			return []string{pod.NodeName}, nil
		},
	}

	// 创建 store - 使用 NewIndexer 和我们自定义的 key 函数
	store := cache.NewIndexer(
		podKeyFunc, // 使用我们自定义的 key 函数
		indexers,
	)

	// ==========================
	// 步骤2：添加数据
	// ==========================
	pod1 := &MyPod{Name: "pod-1", Namespace: "default", NodeName: "node-1", Status: "Running"}
	pod2 := &MyPod{Name: "pod-2", Namespace: "default", NodeName: "node-1", Status: "Running"}
	pod3 := &MyPod{Name: "pod-3", Namespace: "kube-system", NodeName: "node-2", Status: "Pending"}

	// Add 方法直接接受对象
	if err := store.Add(pod1); err != nil {
		fmt.Println("Error adding pod1:", err)
	}
	if err := store.Add(pod2); err != nil {
		fmt.Println("Error adding pod2:", err)
	}
	if err := store.Add(pod3); err != nil {
		fmt.Println("Error adding pod3:", err)
	}

	fmt.Println("=== 全部存储的资源 ===")
	for _, obj := range store.List() {
		pod := obj.(*MyPod)
		fmt.Printf("Pod: %s\tNS: %s\tNode: %s\n", pod.Name, pod.Namespace, pod.NodeName)
	}

	// ==========================
	// 步骤3：按索引查询（核心能力）
	// ==========================
	fmt.Println("\n=== 查询 namespace=default 的所有 Pod ==")
	// 创建一个临时对象，它的 namespace 是 default，然后用 Index() 查询
	tempPod1 := &MyPod{Namespace: "default"}
	objs1, err := store.Index("namespace", tempPod1)
	if err != nil {
		fmt.Println("Error getting index:", err)
	} else {
		fmt.Printf("找到 %d 个 Pod\n", len(objs1))
		for _, obj := range objs1 {
			pod := obj.(*MyPod)
			fmt.Printf("- %s (namespace: %s)\n", pod.Name, pod.Namespace)
		}
	}

	fmt.Println("\n=== 查询运行在 node-1 上的所有 Pod ==")
	// 同样，我们需要传入一个具有 NodeName="node-1" 的临时对象
	tempPod2 := &MyPod{NodeName: "node-1"}
	objs2, err := store.Index("node", tempPod2)
	if err != nil {
		fmt.Println("Error getting index:", err)
	} else {
		fmt.Printf("找到 %d 个 Pod\n", len(objs2))
		for _, obj := range objs2 {
			pod := obj.(*MyPod)
			fmt.Printf("- %s (node: %s)\n", pod.Name, pod.NodeName)
		}
	}

	// ==========================
	// 步骤4：Get / Update / Delete
	// ==========================
	obj, exists, _ := store.Get(pod1) // 直接用对象获取
	if exists {
		fmt.Println("\n=== 获取到 Pod:", obj.(*MyPod).Name)
	}

	// 更新
	pod1.Status = "Failed"
	if err := store.Update(pod1); err != nil {
		fmt.Println("Error updating pod:", err)
	}

	// 删除
	if err := store.Delete(pod2); err != nil {
		fmt.Println("Error deleting pod:", err)
	}

	fmt.Println("\n=== 更新 & 删除后的列表 ===")
	for _, obj := range store.List() {
		pod := obj.(*MyPod)
		fmt.Printf("%s 状态: %s\n", pod.Name, pod.Status)
	}

	// ==========================
	// 步骤5：演示并发安全（多 goroutine 读写）
	// ==========================
	fmt.Println("\n=== 启动 7 个 goroutine 并发读写 ===")
	var wg sync.WaitGroup
	for i := 4; i <= 10; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			pod := &MyPod{
				Name:      fmt.Sprintf("pod-%d", i),
				Namespace: "default",
				NodeName:  "node-3",
				Status:    "Running",
			}
			if err := store.Add(pod); err != nil {
				fmt.Printf("Error adding pod-%d: %v\n", i, err)
			}
		}(i)
	}
	wg.Wait()

	fmt.Println("最终资源数量:", len(store.List()))
}

=== 全部存储的资源 ===
Pod: pod-1	NS: default	Node: node-1
Pod: pod-2	NS: default	Node: node-1
Pod: pod-3	NS: kube-system	Node: node-2

=== 查询 namespace=default 的所有 Pod ==
找到 2 个 Pod
- pod-2 (namespace: default)
- pod-1 (namespace: default)

=== 查询运行在 node-1 上的所有 Pod ==
找到 2 个 Pod
- pod-1 (node: node-1)
- pod-2 (node: node-1)

=== 获取到 Pod: pod-1

=== 更新 & 删除后的列表 ===
pod-1 状态: Failed
pod-3 状态: Pending

=== 启动 7 个 goroutine 并发读写 ===
最终资源数量: 9