examples/demo_balancer.go

package main

import (
	"fmt"
	"jetbrains-ai-proxy/internal/balancer"
	"jetbrains-ai-proxy/internal/config"
	"log"
	"sync"
	"time"
)

func main() {
	fmt.Println("=== JWT负载均衡器演示 ===")

	// 演示轮询策略
	fmt.Println("\n1. 轮询策略演示:")
	demoRoundRobin()

	// 演示随机策略
	fmt.Println("\n2. 随机策略演示:")
	demoRandom()

	// 演示健康检查
	fmt.Println("\n3. 健康检查演示:")
	demoHealthCheck()

	// 演示并发访问
	fmt.Println("\n4. 并发访问演示:")
	demoConcurrent()

	fmt.Println("\n=== 演示完成 ===")
}

func demoRoundRobin() {
	tokens := []string{"JWT_TOKEN_1", "JWT_TOKEN_2", "JWT_TOKEN_3"}
	balancer := balancer.NewJWTBalancer(tokens, config.RoundRobin)

	fmt.Printf("配置了 %d 个JWT tokens，使用轮询策略\n", len(tokens))
	fmt.Println("获取token顺序:")

	for i := 0; i < 9; i++ {
		token, err := balancer.GetToken()
		if err != nil {
			log.Printf("错误: %v", err)
			continue
		}
		fmt.Printf("  第%d次: %s\n", i+1, token)
	}
}

func demoRandom() {
	tokens := []string{"JWT_TOKEN_A", "JWT_TOKEN_B", "JWT_TOKEN_C"}
	balancer := balancer.NewJWTBalancer(tokens, config.Random)

	fmt.Printf("配置了 %d 个JWT tokens，使用随机策略\n", len(tokens))
	fmt.Println("获取token顺序:")

	tokenCounts := make(map[string]int)
	for i := 0; i < 12; i++ {
		token, err := balancer.GetToken()
		if err != nil {
			log.Printf("错误: %v", err)
			continue
		}
		tokenCounts[token]++
		fmt.Printf("  第%d次: %s\n", i+1, token)
	}

	fmt.Println("使用统计:")
	for token, count := range tokenCounts {
		fmt.Printf("  %s: %d次\n", token, count)
	}
}

func demoHealthCheck() {
	tokens := []string{"JWT_HEALTHY_1", "JWT_HEALTHY_2", "JWT_UNHEALTHY_3"}
	balancer := balancer.NewJWTBalancer(tokens, config.RoundRobin)

	fmt.Printf("初始状态: %d/%d tokens健康\n",
		balancer.GetHealthyTokenCount(), balancer.GetTotalTokenCount())

	// 标记一个token为不健康
	fmt.Println("标记 JWT_UNHEALTHY_3 为不健康...")
	balancer.MarkTokenUnhealthy("JWT_UNHEALTHY_3")

	fmt.Printf("更新后状态: %d/%d tokens健康\n",
		balancer.GetHealthyTokenCount(), balancer.GetTotalTokenCount())

	fmt.Println("获取token（应该只返回健康的tokens）:")
	for i := 0; i < 6; i++ {
		token, err := balancer.GetToken()
		if err != nil {
			log.Printf("错误: %v", err)
			continue
		}
		fmt.Printf("  第%d次: %s\n", i+1, token)
	}

	// 恢复token健康状态
	fmt.Println("恢复 JWT_UNHEALTHY_3 为健康...")
	balancer.MarkTokenHealthy("JWT_UNHEALTHY_3")

	fmt.Printf("恢复后状态: %d/%d tokens健康\n",
		balancer.GetHealthyTokenCount(), balancer.GetTotalTokenCount())
}

func demoConcurrent() {
	tokens := []string{"JWT_CONCURRENT_1", "JWT_CONCURRENT_2", "JWT_CONCURRENT_3", "JWT_CONCURRENT_4"}
	balancer := balancer.NewJWTBalancer(tokens, config.RoundRobin)

	fmt.Printf("使用 %d 个JWT tokens进行并发测试\n", len(tokens))

	var wg sync.WaitGroup
	numGoroutines := 5
	requestsPerGoroutine := 10

	tokenCounts := make(map[string]int)
	var mutex sync.Mutex

	startTime := time.Now()

	for i := 0; i < numGoroutines; i++ {
		wg.Add(1)
		go func(goroutineID int) {
			defer wg.Done()

			for j := 0; j < requestsPerGoroutine; j++ {
				token, err := balancer.GetToken()
				if err != nil {
					log.Printf("Goroutine %d 错误: %v", goroutineID, err)
					continue
				}

				mutex.Lock()
				tokenCounts[token]++
				mutex.Unlock()

				// 模拟一些处理时间
				time.Sleep(time.Millisecond * 10)
			}
		}(i)
	}

	wg.Wait()
	duration := time.Since(startTime)

	fmt.Printf("并发测试完成，耗时: %v\n", duration)
	fmt.Printf("总请求数: %d\n", numGoroutines*requestsPerGoroutine)
	fmt.Println("Token使用分布:")

	for token, count := range tokenCounts {
		percentage := float64(count) / float64(numGoroutines*requestsPerGoroutine) * 100
		fmt.Printf("  %s: %d次 (%.1f%%)\n", token, count, percentage)
	}
}

// 演示配置加载
//func demoConfigLoading() {
//	fmt.Println("\n5. 配置加载演示:")
//
//	// 模拟环境变量配置
//	fmt.Println("模拟配置加载...")
//
//	cfg := &config.Config{}
//
//	// 设置多个JWT tokens
//	cfg.SetJetbrainsTokens("token1,token2,token3")
//	cfg.BearerToken = "bearer_token_example"
//	cfg.LoadBalanceStrategy = config.RoundRobin
//
//	fmt.Printf("JWT Tokens数量: %d\n", len(cfg.GetJetbrainsTokens()))
//	fmt.Printf("Bearer Token: %s\n", cfg.BearerToken)
//	fmt.Printf("负载均衡策略: %s\n", cfg.LoadBalanceStrategy)
//	fmt.Printf("是否有JWT Tokens: %v\n", cfg.HasJetbrainsTokens())
//
//	fmt.Println("JWT Tokens列表:")
//	for i, token := range cfg.GetJetbrainsTokens() {
//		fmt.Printf("  %d: %s\n", i+1, token)
//	}
//}