Browse Source

Fix a panic that could happen when using the TokenAwareHostPolicy

This was due to that the fallbackHost could be nil but was still checked
accessed to check if the host was the same as the filtered.
Max Ekman 10 năm trước cách đây
mục cha
commit
6eee003f22
2 tập tin đã thay đổi với 39 bổ sung1 xóa
  1. 1 1
      policies.go
  2. 38 0
      policies_test.go

+ 1 - 1
policies.go

@@ -370,7 +370,7 @@ func (t *tokenAwareHostPolicy) Pick(qry *Query) NextHost {
 		fallbackHost := fallbackIter()
 
 		// filter the token aware selected hosts from the fallback hosts
-		if fallbackHost.Info() == host {
+		if fallbackHost != nil && fallbackHost.Info() == host {
 			fallbackHost = fallbackIter()
 		}
 

+ 38 - 0
policies_test.go

@@ -192,6 +192,44 @@ func TestRoundRobinNilHostInfo(t *testing.T) {
 	}
 }
 
+func TestTokenAwareNilHostInfo(t *testing.T) {
+	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())
+
+	hosts := []*HostInfo{
+		{peer: "0", tokens: []string{"00"}},
+		{peer: "1", tokens: []string{"25"}},
+		{peer: "2", tokens: []string{"50"}},
+		{peer: "3", tokens: []string{"75"}},
+	}
+	policy.SetHosts(hosts)
+	policy.SetPartitioner("OrderedPartitioner")
+
+	query := &Query{}
+	query.RoutingKey([]byte("20"))
+
+	iter := policy.Pick(query)
+	next := iter()
+	if next == nil {
+		t.Fatal("got nil host")
+	} else if v := next.Info(); v == nil {
+		t.Fatal("got nil HostInfo")
+	} else if v.Peer() != "1" {
+		t.Fatalf("expected peer 1 got %v", v.Peer())
+	}
+
+	// Empty the hosts to trigger the panic when using the fallback.
+	hosts = []*HostInfo{}
+	policy.SetHosts(hosts)
+
+	next = iter()
+	if next != nil {
+		t.Errorf("expected to get nil host got %+v", next)
+		if next.Info() == nil {
+			t.Fatalf("HostInfo is nil")
+		}
+	}
+}
+
 func TestCOWList_Add(t *testing.T) {
 	var cow cowHostList