refactor load balance complete.

pkg/filters/proxies/grpcproxy/loadbalance_test.go

@@ -29,18 +29,15 @@ import (
 func TestForwardLB(t *testing.T) {
  assertions := assert.New(t)
  key := "forward-target"
- balancer := NewLoadBalancer(&LoadBalanceSpec{Policy: "forward", ForwardKey: key}, nil)
+ lb := newForwardLoadBalancer(&LoadBalanceSpec{Policy: "forward", ForwardKey: key})
 
  sm := grpcprot.NewFakeServerStream(metadata.NewIncomingContext(context.Background(), metadata.MD{}))
  req := grpcprot.NewRequestWithServerStream(sm)
 
  target := "127.0.0.1:8849"
- assertions.Nil(balancer.ChooseServer(req))
+ assertions.Nil(lb.ChooseServer(req))
 
  req.Header().Set(key, target)
 
- lb, ok := balancer.(ReusableServerLB)
- assertions.True(ok)
  assertions.Equal(target, lb.ChooseServer(req).URL)
-
 }

pkg/filters/proxies/healthcheck_test.go

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2017, MegaEase
+ * All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http:https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package proxies
+
+import (
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+)
+
+type MockHealthChecker struct {
+ result bool
+}
+
+func (c *MockHealthChecker) Check(svr *Server) bool {
+ return c.result
+}
+
+func (c *MockHealthChecker) Close() {
+}
+
+func TestHTTPHealthChecker(t *testing.T) {
+ spec := &HealthCheckSpec{}
+ c := NewHTTPHealthChecker(spec)
+ assert.NotNil(t, c)
+
+ spec = &HealthCheckSpec{Timeout: "100ms"}
+ c = NewHTTPHealthChecker(spec)
+ c.Check(&Server{URL: "https://www.megaease.com"})
+
+ c.Close()
+}

pkg/filters/proxies/httpproxy/pool_test.go

@@ -23,7 +23,6 @@ import (
  "testing"
 
  "github.com/megaease/easegress/pkg/context"
- "github.com/megaease/easegress/pkg/object/serviceregistry"
  "github.com/megaease/easegress/pkg/option"
  "github.com/megaease/easegress/pkg/protocols/httpprot"
  "github.com/megaease/easegress/pkg/resilience"
@@ -198,48 +197,6 @@ func TestCopyCORSHeaders(t *testing.T) {
  assert.Equal("dst", dst.Values("X-Dst")[0])
 }
 
-func TestUseService(t *testing.T) {
- assert := assert.New(t)
-
- yamlConfig := `spanName: test
-serverTags: [a1, a2]
-servers:
-- url: http:https://192.168.1.1
-`
-
- spec := &ServerPoolSpec{}
- err := codectool.Unmarshal([]byte(yamlConfig), spec)
- assert.NoError(err)
- assert.NoError(spec.Validate())
-
- p := &Proxy{}
- p.super = supervisor.NewMock(option.New(), nil, sync.Map{}, sync.Map{}, nil,
- nil, false, nil, nil)
- sp := NewServerPool(p, spec, "test")
- svr := sp.LoadBalancer().ChooseServer(nil)
- assert.Equal("http:https://192.168.1.1", svr.URL)
-
- sp.useService(&spec.BaseServerPoolSpec, nil)
- assert.Equal("http:https://192.168.1.1", svr.URL)
-
- sp.useService(&spec.BaseServerPoolSpec, map[string]*serviceregistry.ServiceInstanceSpec{
- "2": {
- Address: "192.168.1.2",
- Tags: []string{"a2"},
- Port: 80,
- },
- "3": {
- Address: "192.168.1.3",
- Tags: []string{"a3"},
- Port: 80,
- },
- })
- svr = sp.LoadBalancer().ChooseServer(nil)
- assert.Equal("http:https://192.168.1.2:80", svr.URL)
- svr = sp.LoadBalancer().ChooseServer(nil)
- assert.Equal("http:https://192.168.1.2:80", svr.URL)
-}
-
 func TestRemoveHopByHopHeader(t *testing.T) {
  assert := assert.New(t)
 

pkg/filters/proxies/httpproxy/wsproxy_test.go

@@ -160,13 +160,4 @@ pools:
  ctx.SetData("HTTP_RESPONSE_WRITER", httptest.NewRecorder())
  assert.Equal(resultClientError, proxy.Handle(ctx))
  }
-
- // no server
- proxy.mainPool.loadBalancer.Store(NewLoadBalancer(&LoadBalanceSpec{}, nil))
- {
- stdr, _ := http.NewRequest(http.MethodGet, "wss:https://www.megaease.com", nil)
- ctx := getCtx(stdr)
- ctx.SetData("HTTP_RESPONSE_WRITER", httptest.NewRecorder())
- assert.Equal(resultInternalError, proxy.Handle(ctx))
- }
 }

pkg/filters/proxies/loadbalance.go

@@ -49,6 +49,10 @@ type LoadBalancer interface {
 }
 
 // LoadBalanceSpec is the spec to create a load balancer.
+//
+// TODO: this spec currently include all options for all load balance policies,
+// this is not good as new policies could be added in the future, we should
+// convert it to a map later.
 type LoadBalanceSpec struct {
  Policy string `json:"policy" jsonschema:"omitempty"`
  HeaderHashKey string `json:"headerHashKey" jsonschema:"omitempty"`
@@ -91,6 +95,7 @@ func (glb *GeneralLoadBalancer) Init(
  fnNewHealthChecker func(*HealthCheckSpec) HealthChecker,
  lbp LoadBalancePolicy,
 ) {
+ // load balance policy
  if lbp == nil {
  switch glb.spec.Policy {
  case LoadBalancePolicyRoundRobin, "":
@@ -102,20 +107,22 @@ func (glb *GeneralLoadBalancer) Init(
  case LoadBalancePolicyIPHash:
  lbp = &IPHashLoadBalancePolicy{}
  case LoadBalancePolicyHeaderHash:
- lbp = &HeaderHashLoadBalancePolicy{}
+ lbp = &HeaderHashLoadBalancePolicy{spec: glb.spec}
  default:
  logger.Errorf("unsupported load balancing policy: %s", glb.spec.Policy)
  lbp = &RoundRobinLoadBalancePolicy{}
  }
  }
  glb.lbp = lbp
 
+ // sticky session
  if glb.spec.StickySession != nil {
  ss := fnNewSessionSticker(glb.spec.StickySession)
  ss.UpdateServers(glb.servers)
  glb.ss = ss
  }
 
+ // health check
  if glb.spec.HealthCheck == nil {
  return
  }
@@ -196,7 +203,7 @@ func (glb *GeneralLoadBalancer) checkServers() {
 // ChooseServer chooses a server according to the load balancing spec.
 func (glb *GeneralLoadBalancer) ChooseServer(req protocols.Request) *Server {
  sg := glb.healthyServers.Load()
- if len(sg.Servers) == 0 {
+ if sg == nil || len(sg.Servers) == 0 {
  return nil
  }