github.com/AntonOrnatskyi/goproxy@v0.0.0-20190205095733-4526a9fa18b4/utils/lb/lb.go (about) 1 package lb 2 3 import ( 4 "crypto/md5" 5 "log" 6 "net" 7 "sync" 8 9 "github.com/AntonOrnatskyi/goproxy/utils/dnsx" 10 ) 11 12 const ( 13 SELECT_ROUNDROBIN = iota 14 SELECT_LEASTCONN 15 SELECT_HASH 16 SELECT_WEITHT 17 SELECT_LEASTTIME 18 ) 19 20 type Selector interface { 21 Select(srcAddr string) (addr string) 22 SelectBackend(srcAddr string) (b *Backend) 23 IncreasConns(addr string) 24 DecreaseConns(addr string) 25 Stop() 26 Reset(configs BackendsConfig, dr *dnsx.DomainResolver, log *log.Logger) 27 IsActive() bool 28 ActiveCount() (count int) 29 Backends() (bs []*Backend) 30 } 31 32 type Group struct { 33 selector *Selector 34 log *log.Logger 35 dr *dnsx.DomainResolver 36 lock *sync.Mutex 37 last *Backend 38 debug bool 39 bks []*Backend 40 } 41 42 func NewGroup(selectType int, configs BackendsConfig, dr *dnsx.DomainResolver, log *log.Logger, debug bool) Group { 43 bks := []*Backend{} 44 for _, c := range configs { 45 b, _ := NewBackend(*c, dr, log) 46 bks = append(bks, b) 47 } 48 if len(bks) > 1 { 49 for _, b := range bks { 50 b.StartHeartCheck() 51 } 52 } 53 var s Selector 54 switch selectType { 55 case SELECT_ROUNDROBIN: 56 s = NewRoundRobin(bks, log, debug) 57 case SELECT_LEASTCONN: 58 s = NewLeastConn(bks, log, debug) 59 case SELECT_HASH: 60 s = NewHash(bks, log, debug) 61 case SELECT_WEITHT: 62 s = NewWeight(bks, log, debug) 63 case SELECT_LEASTTIME: 64 s = NewLeastTime(bks, log, debug) 65 } 66 return Group{ 67 selector: &s, 68 log: log, 69 dr: dr, 70 lock: &sync.Mutex{}, 71 debug: debug, 72 bks: bks, 73 } 74 } 75 76 func (g *Group) Select(srcAddr string, onlyHa bool) (addr string) { 77 _, addr = g.Select2(srcAddr, onlyHa) 78 return 79 } 80 func (g *Group) Select2(srcAddr string, onlyHa bool) (isEmpty bool, addr string) { 81 addr = "" 82 if len(g.bks) == 1 { 83 return false, g.bks[0].Address 84 } 85 if onlyHa { 86 g.lock.Lock() 87 defer g.lock.Unlock() 88 if g.last != nil && (g.last.Active || g.last.ConnectUsedMillisecond == 0) { 89 if g.debug { 90 g.log.Printf("############ choosed %s from lastest ############", g.last.Address) 91 printDebug(true, g.log, nil, srcAddr, (*g.selector).Backends()) 92 } 93 return false, g.last.Address 94 } 95 g.last = (*g.selector).SelectBackend(srcAddr) 96 if !g.last.Active && g.last.ConnectUsedMillisecond > 0 { 97 g.log.Printf("###warn### lb selected empty , return default , for : %s", srcAddr) 98 } 99 return true, g.last.Address 100 } 101 b := (*g.selector).SelectBackend(srcAddr) 102 if !b.Active && b.ConnectUsedMillisecond > 0 { 103 g.log.Printf("###warn### lb selected empty , return default , for : %s", srcAddr) 104 return true, b.Address 105 } 106 return false, b.Address 107 108 } 109 func (g *Group) IncreasConns(addr string) { 110 (*g.selector).IncreasConns(addr) 111 } 112 func (g *Group) DecreaseConns(addr string) { 113 (*g.selector).DecreaseConns(addr) 114 } 115 func (g *Group) Stop() { 116 if g.selector != nil { 117 (*g.selector).Stop() 118 } 119 } 120 func (g *Group) IsActive() bool { 121 return (*g.selector).IsActive() 122 } 123 func (g *Group) ActiveCount() (count int) { 124 return (*g.selector).ActiveCount() 125 } 126 func (g *Group) Reset(addrs []string) { 127 bks := (*g.selector).Backends() 128 if len(bks) == 0 { 129 return 130 } 131 cfg := bks[0].BackendConfig 132 configs := BackendsConfig{} 133 for _, addr := range addrs { 134 c := cfg 135 c.Address = addr 136 configs = append(configs, &c) 137 } 138 (*g.selector).Reset(configs, g.dr, g.log) 139 g.bks = (*g.selector).Backends() 140 } 141 func (g *Group) Backends() []*Backend { 142 return (*g.selector).Backends() 143 } 144 145 //########################RoundRobin########################## 146 type RoundRobin struct { 147 sync.Mutex 148 backendIndex int 149 backends Backends 150 log *log.Logger 151 debug bool 152 } 153 154 func NewRoundRobin(backends Backends, log *log.Logger, debug bool) Selector { 155 return &RoundRobin{ 156 backends: backends, 157 log: log, 158 debug: debug, 159 } 160 161 } 162 func (r *RoundRobin) Select(srcAddr string) (addr string) { 163 return r.SelectBackend(srcAddr).Address 164 } 165 func (r *RoundRobin) SelectBackend(srcAddr string) (b *Backend) { 166 r.Lock() 167 defer r.Unlock() 168 defer func() { 169 printDebug(r.debug, r.log, b, srcAddr, r.backends) 170 }() 171 if len(r.backends) == 0 { 172 return 173 } 174 if len(r.backends) == 1 { 175 return r.backends[0] 176 } 177 RETRY: 178 found := false 179 for _, b := range r.backends { 180 if b.Active { 181 found = true 182 break 183 } 184 } 185 if !found { 186 return r.backends[0] 187 } 188 r.backendIndex++ 189 if r.backendIndex > len(r.backends)-1 { 190 r.backendIndex = 0 191 } 192 if !r.backends[r.backendIndex].Active { 193 goto RETRY 194 } 195 return r.backends[r.backendIndex] 196 } 197 func (r *RoundRobin) IncreasConns(addr string) { 198 199 } 200 func (r *RoundRobin) DecreaseConns(addr string) { 201 202 } 203 func (r *RoundRobin) Stop() { 204 for _, b := range r.backends { 205 b.StopHeartCheck() 206 } 207 } 208 func (r *RoundRobin) Backends() []*Backend { 209 return r.backends 210 } 211 func (r *RoundRobin) IsActive() bool { 212 for _, b := range r.backends { 213 if b.Active { 214 return true 215 } 216 } 217 return false 218 } 219 func (r *RoundRobin) ActiveCount() (count int) { 220 for _, b := range r.backends { 221 if b.Active { 222 count++ 223 } 224 } 225 return 226 } 227 func (r *RoundRobin) Reset(configs BackendsConfig, dr *dnsx.DomainResolver, log *log.Logger) { 228 r.Lock() 229 defer r.Unlock() 230 r.Stop() 231 bks := []*Backend{} 232 for _, c := range configs { 233 b, _ := NewBackend(*c, dr, log) 234 bks = append(bks, b) 235 } 236 if len(bks) > 1 { 237 for _, b := range bks { 238 b.StartHeartCheck() 239 } 240 } 241 r.backends = bks 242 } 243 244 //########################LeastConn########################## 245 246 type LeastConn struct { 247 sync.Mutex 248 backends Backends 249 log *log.Logger 250 debug bool 251 } 252 253 func NewLeastConn(backends []*Backend, log *log.Logger, debug bool) Selector { 254 lc := LeastConn{ 255 backends: backends, 256 log: log, 257 debug: debug, 258 } 259 return &lc 260 } 261 262 func (lc *LeastConn) Select(srcAddr string) (addr string) { 263 return lc.SelectBackend(srcAddr).Address 264 } 265 func (lc *LeastConn) SelectBackend(srcAddr string) (b *Backend) { 266 lc.Lock() 267 defer lc.Unlock() 268 defer func() { 269 printDebug(lc.debug, lc.log, b, srcAddr, lc.backends) 270 }() 271 if len(lc.backends) == 0 { 272 return 273 } 274 if len(lc.backends) == 1 { 275 return lc.backends[0] 276 } 277 found := false 278 for _, b := range lc.backends { 279 if b.Active { 280 found = true 281 break 282 } 283 } 284 if !found { 285 return lc.backends[0] 286 } 287 min := lc.backends[0].Connections 288 index := 0 289 for i, b := range lc.backends { 290 if b.Active { 291 min = b.Connections 292 index = i 293 break 294 } 295 } 296 for i, b := range lc.backends { 297 if b.Active && b.Connections <= min { 298 min = b.Connections 299 index = i 300 } 301 } 302 return lc.backends[index] 303 } 304 func (lc *LeastConn) IncreasConns(addr string) { 305 for _, a := range lc.backends { 306 if a.Address == addr { 307 a.IncreasConns() 308 return 309 } 310 } 311 } 312 func (lc *LeastConn) DecreaseConns(addr string) { 313 for _, a := range lc.backends { 314 if a.Address == addr { 315 a.DecreaseConns() 316 return 317 } 318 } 319 } 320 func (lc *LeastConn) Stop() { 321 for _, b := range lc.backends { 322 b.StopHeartCheck() 323 } 324 } 325 func (lc *LeastConn) IsActive() bool { 326 for _, b := range lc.backends { 327 if b.Active { 328 return true 329 } 330 } 331 return false 332 } 333 func (lc *LeastConn) ActiveCount() (count int) { 334 for _, b := range lc.backends { 335 if b.Active { 336 count++ 337 } 338 } 339 return 340 } 341 func (lc *LeastConn) Reset(configs BackendsConfig, dr *dnsx.DomainResolver, log *log.Logger) { 342 lc.Lock() 343 defer lc.Unlock() 344 lc.Stop() 345 bks := []*Backend{} 346 for _, c := range configs { 347 b, _ := NewBackend(*c, dr, log) 348 bks = append(bks, b) 349 } 350 if len(bks) > 1 { 351 for _, b := range bks { 352 b.StartHeartCheck() 353 } 354 } 355 lc.backends = bks 356 } 357 func (lc *LeastConn) Backends() []*Backend { 358 return lc.backends 359 } 360 361 //########################Hash########################## 362 type Hash struct { 363 sync.Mutex 364 backends Backends 365 log *log.Logger 366 debug bool 367 } 368 369 func NewHash(backends Backends, log *log.Logger, debug bool) Selector { 370 return &Hash{ 371 backends: backends, 372 log: log, 373 debug: debug, 374 } 375 } 376 func (h *Hash) Select(srcAddr string) (addr string) { 377 return h.SelectBackend(srcAddr).Address 378 } 379 func (h *Hash) SelectBackend(srcAddr string) (b *Backend) { 380 h.Lock() 381 defer h.Unlock() 382 defer func() { 383 printDebug(h.debug, h.log, b, srcAddr, h.backends) 384 }() 385 if len(h.backends) == 0 { 386 return 387 } 388 if len(h.backends) == 1 { 389 return h.backends[0] 390 } 391 i := 0 392 host, _, err := net.SplitHostPort(srcAddr) 393 if err != nil { 394 return 395 } 396 //porti, _ := strconv.Atoi(port) 397 //i += porti 398 for _, b := range md5.Sum([]byte(host)) { 399 i += int(b) 400 } 401 RETRY: 402 found := false 403 for _, b := range h.backends { 404 if b.Active { 405 found = true 406 break 407 } 408 } 409 if !found { 410 return h.backends[0] 411 } 412 k := i % len(h.backends) 413 if !h.backends[k].Active { 414 i++ 415 goto RETRY 416 } 417 return h.backends[k] 418 } 419 func (h *Hash) IncreasConns(addr string) { 420 421 } 422 func (h *Hash) DecreaseConns(addr string) { 423 424 } 425 func (h *Hash) Stop() { 426 for _, b := range h.backends { 427 b.StopHeartCheck() 428 } 429 } 430 func (h *Hash) IsActive() bool { 431 for _, b := range h.backends { 432 if b.Active { 433 return true 434 } 435 } 436 return false 437 } 438 func (h *Hash) ActiveCount() (count int) { 439 for _, b := range h.backends { 440 if b.Active { 441 count++ 442 } 443 } 444 return 445 } 446 func (h *Hash) Reset(configs BackendsConfig, dr *dnsx.DomainResolver, log *log.Logger) { 447 h.Lock() 448 defer h.Unlock() 449 h.Stop() 450 bks := []*Backend{} 451 for _, c := range configs { 452 b, _ := NewBackend(*c, dr, log) 453 bks = append(bks, b) 454 } 455 if len(bks) > 1 { 456 for _, b := range bks { 457 b.StartHeartCheck() 458 } 459 } 460 h.backends = bks 461 } 462 func (h *Hash) Backends() []*Backend { 463 return h.backends 464 } 465 466 //########################Weight########################## 467 type Weight struct { 468 sync.Mutex 469 backends Backends 470 log *log.Logger 471 debug bool 472 } 473 474 func NewWeight(backends Backends, log *log.Logger, debug bool) Selector { 475 return &Weight{ 476 backends: backends, 477 log: log, 478 debug: debug, 479 } 480 } 481 func (w *Weight) Select(srcAddr string) (addr string) { 482 return w.SelectBackend(srcAddr).Address 483 } 484 func (w *Weight) SelectBackend(srcAddr string) (b *Backend) { 485 w.Lock() 486 defer w.Unlock() 487 defer func() { 488 printDebug(w.debug, w.log, b, srcAddr, w.backends) 489 }() 490 if len(w.backends) == 0 { 491 return 492 } 493 if len(w.backends) == 1 { 494 return w.backends[0] 495 } 496 497 found := false 498 for _, b := range w.backends { 499 if b.Active { 500 found = true 501 break 502 } 503 } 504 if !found { 505 return w.backends[0] 506 } 507 508 min := w.backends[0].Connections / w.backends[0].Weight 509 index := 0 510 for i, b := range w.backends { 511 if b.Active { 512 min = b.Connections / b.Weight 513 index = i 514 break 515 } 516 } 517 for i, b := range w.backends { 518 if b.Active && b.Connections/b.Weight <= min { 519 min = b.Connections 520 index = i 521 } 522 } 523 return w.backends[index] 524 } 525 func (w *Weight) IncreasConns(addr string) { 526 w.Lock() 527 defer w.Unlock() 528 for _, a := range w.backends { 529 if a.Address == addr { 530 a.IncreasConns() 531 return 532 } 533 } 534 } 535 func (w *Weight) DecreaseConns(addr string) { 536 w.Lock() 537 defer w.Unlock() 538 for _, a := range w.backends { 539 if a.Address == addr { 540 a.DecreaseConns() 541 return 542 } 543 } 544 } 545 func (w *Weight) Stop() { 546 for _, b := range w.backends { 547 b.StopHeartCheck() 548 } 549 } 550 func (w *Weight) IsActive() bool { 551 for _, b := range w.backends { 552 if b.Active { 553 return true 554 } 555 } 556 return false 557 } 558 func (w *Weight) ActiveCount() (count int) { 559 for _, b := range w.backends { 560 if b.Active { 561 count++ 562 } 563 } 564 return 565 } 566 func (w *Weight) Reset(configs BackendsConfig, dr *dnsx.DomainResolver, log *log.Logger) { 567 w.Lock() 568 defer w.Unlock() 569 w.Stop() 570 bks := []*Backend{} 571 for _, c := range configs { 572 b, _ := NewBackend(*c, dr, log) 573 bks = append(bks, b) 574 } 575 if len(bks) > 1 { 576 for _, b := range bks { 577 b.StartHeartCheck() 578 } 579 } 580 w.backends = bks 581 } 582 func (w *Weight) Backends() []*Backend { 583 return w.backends 584 } 585 586 //########################LeastTime########################## 587 588 type LeastTime struct { 589 sync.Mutex 590 backends Backends 591 log *log.Logger 592 debug bool 593 } 594 595 func NewLeastTime(backends []*Backend, log *log.Logger, debug bool) Selector { 596 lt := LeastTime{ 597 backends: backends, 598 log: log, 599 debug: debug, 600 } 601 return < 602 } 603 604 func (lt *LeastTime) Select(srcAddr string) (addr string) { 605 return lt.SelectBackend(srcAddr).Address 606 } 607 func (lt *LeastTime) SelectBackend(srcAddr string) (b *Backend) { 608 lt.Lock() 609 defer lt.Unlock() 610 defer func() { 611 printDebug(lt.debug, lt.log, b, srcAddr, lt.backends) 612 }() 613 if len(lt.backends) == 0 { 614 return 615 } 616 if len(lt.backends) == 1 { 617 return lt.backends[0] 618 } 619 found := false 620 for _, b := range lt.backends { 621 if b.Active { 622 found = true 623 break 624 } 625 } 626 if !found { 627 return lt.backends[0] 628 } 629 min := lt.backends[0].ConnectUsedMillisecond 630 index := 0 631 for i, b := range lt.backends { 632 if b.Active { 633 min = b.ConnectUsedMillisecond 634 index = i 635 break 636 } 637 } 638 for i, b := range lt.backends { 639 if b.Active && b.ConnectUsedMillisecond > 0 && b.ConnectUsedMillisecond <= min { 640 min = b.ConnectUsedMillisecond 641 index = i 642 } 643 } 644 return lt.backends[index] 645 } 646 func (lt *LeastTime) IncreasConns(addr string) { 647 648 } 649 func (lt *LeastTime) DecreaseConns(addr string) { 650 651 } 652 func (lt *LeastTime) Stop() { 653 for _, b := range lt.backends { 654 b.StopHeartCheck() 655 } 656 } 657 func (lt *LeastTime) IsActive() bool { 658 for _, b := range lt.backends { 659 if b.Active { 660 return true 661 } 662 } 663 return false 664 } 665 func (lt *LeastTime) ActiveCount() (count int) { 666 for _, b := range lt.backends { 667 if b.Active { 668 count++ 669 } 670 } 671 return 672 } 673 func (lt *LeastTime) Reset(configs BackendsConfig, dr *dnsx.DomainResolver, log *log.Logger) { 674 lt.Lock() 675 defer lt.Unlock() 676 lt.Stop() 677 bks := []*Backend{} 678 for _, c := range configs { 679 b, _ := NewBackend(*c, dr, log) 680 bks = append(bks, b) 681 } 682 if len(bks) > 1 { 683 for _, b := range bks { 684 b.StartHeartCheck() 685 } 686 } 687 lt.backends = bks 688 } 689 func (lt *LeastTime) Backends() []*Backend { 690 return lt.backends 691 } 692 func printDebug(isDebug bool, log *log.Logger, selected *Backend, srcAddr string, backends []*Backend) { 693 if isDebug { 694 log.Printf("############ LB start ############\n") 695 if selected != nil { 696 log.Printf("choosed %s for %s\n", selected.Address, srcAddr) 697 } 698 for _, v := range backends { 699 log.Printf("addr:%s,conns:%d,time:%d,weight:%d,active:%v\n", v.Address, v.Connections, v.ConnectUsedMillisecond, v.Weight, v.Active) 700 } 701 log.Printf("############ LB end ############\n") 702 } 703 }