1// Copyright 2015 The etcd Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package rafthttp
16
17import (
18	"context"
19	"net/http/httptest"
20	"reflect"
21	"testing"
22	"time"
23
24	stats "go.etcd.io/etcd/etcdserver/api/v2stats"
25	"go.etcd.io/etcd/pkg/types"
26	"go.etcd.io/etcd/raft"
27	"go.etcd.io/etcd/raft/raftpb"
28)
29
30func TestSendMessage(t *testing.T) {
31	// member 1
32	tr := &Transport{
33		ID:          types.ID(1),
34		ClusterID:   types.ID(1),
35		Raft:        &fakeRaft{},
36		ServerStats: newServerStats(),
37		LeaderStats: stats.NewLeaderStats("1"),
38	}
39	tr.Start()
40	srv := httptest.NewServer(tr.Handler())
41	defer srv.Close()
42
43	// member 2
44	recvc := make(chan raftpb.Message, 1)
45	p := &fakeRaft{recvc: recvc}
46	tr2 := &Transport{
47		ID:          types.ID(2),
48		ClusterID:   types.ID(1),
49		Raft:        p,
50		ServerStats: newServerStats(),
51		LeaderStats: stats.NewLeaderStats("2"),
52	}
53	tr2.Start()
54	srv2 := httptest.NewServer(tr2.Handler())
55	defer srv2.Close()
56
57	tr.AddPeer(types.ID(2), []string{srv2.URL})
58	defer tr.Stop()
59	tr2.AddPeer(types.ID(1), []string{srv.URL})
60	defer tr2.Stop()
61	if !waitStreamWorking(tr.Get(types.ID(2)).(*peer)) {
62		t.Fatalf("stream from 1 to 2 is not in work as expected")
63	}
64
65	data := []byte("some data")
66	tests := []raftpb.Message{
67		// these messages are set to send to itself, which facilitates testing.
68		{Type: raftpb.MsgProp, From: 1, To: 2, Entries: []raftpb.Entry{{Data: data}}},
69		{Type: raftpb.MsgApp, From: 1, To: 2, Term: 1, Index: 3, LogTerm: 0, Entries: []raftpb.Entry{{Index: 4, Term: 1, Data: data}}, Commit: 3},
70		{Type: raftpb.MsgAppResp, From: 1, To: 2, Term: 1, Index: 3},
71		{Type: raftpb.MsgVote, From: 1, To: 2, Term: 1, Index: 3, LogTerm: 0},
72		{Type: raftpb.MsgVoteResp, From: 1, To: 2, Term: 1},
73		{Type: raftpb.MsgSnap, From: 1, To: 2, Term: 1, Snapshot: raftpb.Snapshot{Metadata: raftpb.SnapshotMetadata{Index: 1000, Term: 1}, Data: data}},
74		{Type: raftpb.MsgHeartbeat, From: 1, To: 2, Term: 1, Commit: 3},
75		{Type: raftpb.MsgHeartbeatResp, From: 1, To: 2, Term: 1},
76	}
77	for i, tt := range tests {
78		tr.Send([]raftpb.Message{tt})
79		msg := <-recvc
80		if !reflect.DeepEqual(msg, tt) {
81			t.Errorf("#%d: msg = %+v, want %+v", i, msg, tt)
82		}
83	}
84}
85
86// TestSendMessageWhenStreamIsBroken tests that message can be sent to the
87// remote in a limited time when all underlying connections are broken.
88func TestSendMessageWhenStreamIsBroken(t *testing.T) {
89	// member 1
90	tr := &Transport{
91		ID:          types.ID(1),
92		ClusterID:   types.ID(1),
93		Raft:        &fakeRaft{},
94		ServerStats: newServerStats(),
95		LeaderStats: stats.NewLeaderStats("1"),
96	}
97	tr.Start()
98	srv := httptest.NewServer(tr.Handler())
99	defer srv.Close()
100
101	// member 2
102	recvc := make(chan raftpb.Message, 1)
103	p := &fakeRaft{recvc: recvc}
104	tr2 := &Transport{
105		ID:          types.ID(2),
106		ClusterID:   types.ID(1),
107		Raft:        p,
108		ServerStats: newServerStats(),
109		LeaderStats: stats.NewLeaderStats("2"),
110	}
111	tr2.Start()
112	srv2 := httptest.NewServer(tr2.Handler())
113	defer srv2.Close()
114
115	tr.AddPeer(types.ID(2), []string{srv2.URL})
116	defer tr.Stop()
117	tr2.AddPeer(types.ID(1), []string{srv.URL})
118	defer tr2.Stop()
119	if !waitStreamWorking(tr.Get(types.ID(2)).(*peer)) {
120		t.Fatalf("stream from 1 to 2 is not in work as expected")
121	}
122
123	// break the stream
124	srv.CloseClientConnections()
125	srv2.CloseClientConnections()
126	var n int
127	for {
128		select {
129		// TODO: remove this resend logic when we add retry logic into the code
130		case <-time.After(time.Millisecond):
131			n++
132			tr.Send([]raftpb.Message{{Type: raftpb.MsgHeartbeat, From: 1, To: 2, Term: 1, Commit: 3}})
133		case <-recvc:
134			if n > 50 {
135				t.Errorf("disconnection time = %dms, want < 50ms", n)
136			}
137			return
138		}
139	}
140}
141
142func newServerStats() *stats.ServerStats {
143	return stats.NewServerStats("", "")
144}
145
146func waitStreamWorking(p *peer) bool {
147	for i := 0; i < 1000; i++ {
148		time.Sleep(time.Millisecond)
149		if _, ok := p.msgAppV2Writer.writec(); !ok {
150			continue
151		}
152		if _, ok := p.writer.writec(); !ok {
153			continue
154		}
155		return true
156	}
157	return false
158}
159
160type fakeRaft struct {
161	recvc     chan<- raftpb.Message
162	err       error
163	removedID uint64
164}
165
166func (p *fakeRaft) Process(ctx context.Context, m raftpb.Message) error {
167	select {
168	case p.recvc <- m:
169	default:
170	}
171	return p.err
172}
173
174func (p *fakeRaft) IsIDRemoved(id uint64) bool { return id == p.removedID }
175
176func (p *fakeRaft) ReportUnreachable(id uint64) {}
177
178func (p *fakeRaft) ReportSnapshot(id uint64, status raft.SnapshotStatus) {}
179