plugin/transport: Prototype QUIC transport

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/juju/ratelimit v1.0.2-0.20191002062651-f60b32039441
 	github.com/lni/goutils v1.3.1-0.20210207055804-2f3468487e42
 	github.com/lni/vfs v0.2.0
+	github.com/lucas-clemente/quic-go v0.20.1
 	github.com/stretchr/testify v1.7.0
 )
 

plugin/plugin_test.go

@@ -16,23 +16,28 @@ package plugin
 
 import (
 	"os"
+	"path"
 	"testing"
 
 	"github.com/lni/dragonboat/v3"
 	"github.com/lni/dragonboat/v3/config"
 	"github.com/lni/dragonboat/v3/internal/vfs"
 	"github.com/lni/dragonboat/v3/plugin/rocksdb"
+	"github.com/lni/dragonboat/v3/plugin/transport/quic"
 )
 
-var (
+const (
 	singleNodeHostTestDir = "plugin_test_dir_safe_to_delete"
+	caFile                = "../internal/transport/tests/test-root-ca.crt"
+	certFile              = "../internal/transport/tests/localhost.crt"
+	keyFile               = "../internal/transport/tests/localhost.key"
 )
 
 func testLogDBPluginCanBeUsed(t *testing.T, f config.LogDBFactory) {
-	os.RemoveAll(singleNodeHostTestDir)
-	defer os.RemoveAll(singleNodeHostTestDir)
+	testDir := path.Join(t.TempDir(), singleNodeHostTestDir)
+	defer os.RemoveAll(testDir)
 	nhc := config.NodeHostConfig{
-		NodeHostDir:    singleNodeHostTestDir,
+		NodeHostDir:    testDir,
 		RTTMillisecond: 20,
 		RaftAddress:    "localhost:26000",
 		Expert:         config.ExpertConfig{FS: vfs.DefaultFS, LogDBFactory: f},
@@ -44,6 +49,30 @@ func testLogDBPluginCanBeUsed(t *testing.T, f config.LogDBFactory) {
 	defer nh.Close()
 }
 
+func testTransportPluginCanBeUsed(t *testing.T, f config.TransportFactory) {
+	testDir := path.Join(t.TempDir(), singleNodeHostTestDir)
+	defer os.RemoveAll(testDir)
+	nhc := config.NodeHostConfig{
+		NodeHostDir:    testDir,
+		RTTMillisecond: 20,
+		RaftAddress:    "localhost:26000",
+		Expert:         config.ExpertConfig{FS: vfs.DefaultFS, TransportFactory: f},
+		MutualTLS:      true,
+		CAFile:         caFile,
+		CertFile:       certFile,
+		KeyFile:        keyFile,
+	}
+	nh, err := dragonboat.NewNodeHost(nhc)
+	if err != nil {
+		t.Fatalf("failed to create nodehost %v", err)
+	}
+	defer nh.Close()
+}
+
 func TestLogDBPluginsCanBeUsed(t *testing.T) {
 	testLogDBPluginCanBeUsed(t, &rocksdb.Factory{})
 }
+
+func TestTransportPluginsCanBeUsed(t *testing.T) {
+	testTransportPluginCanBeUsed(t, &quic.TransportFactory{})
+}

plugin/transport/quic/factory.go

@@ -0,0 +1,42 @@
+// Copyright 2017-2021 Lei Ni ([email protected]) and other contributors.
+//
+// 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://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 quic
+
+import (
+	"github.com/lni/dragonboat/v3/config"
+	"github.com/lni/dragonboat/v3/raftio"
+	"github.com/lni/goutils/stringutil"
+	"github.com/lni/goutils/syncutil"
+)
+
+// TransportFactory is a prototype QUIC transport factory, use at your own risk.
+type TransportFactory struct {
+}
+
+// Create a new raftio.ITransport over QUIC protocol.
+func (t *TransportFactory) Create(config config.NodeHostConfig, mHandler raftio.MessageHandler, chHandler raftio.ChunkHandler) raftio.ITransport {
+	return &quicTransport{
+		nhConfig:       config,
+		messageHandler: mHandler,
+		chunkHandler:   chHandler,
+		stopper:        syncutil.NewStopper(),
+		connStopper:    syncutil.NewStopper(),
+	}
+}
+
+// Validate address.
+func (t *TransportFactory) Validate(addr string) bool {
+	return stringutil.IsValidAddress(addr)
+}

plugin/transport/quic/header.go

@@ -0,0 +1,49 @@
+// Copyright 2017-2021 Lei Ni ([email protected]) and other contributors.
+//
+// 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://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 quic
+
+import (
+	"encoding/binary"
+)
+
+const requestHeaderSize = 10
+
+type requestHeader struct {
+	method uint16
+	size   uint64
+}
+
+func (h *requestHeader) encode(buf []byte) []byte {
+	if len(buf) < requestHeaderSize {
+		plog.Panicf("requestHeader input buf too small")
+	}
+	binary.BigEndian.PutUint16(buf, h.method)
+	binary.BigEndian.PutUint64(buf[2:], h.size)
+	return buf[:requestHeaderSize]
+}
+
+func (h *requestHeader) decode(buf []byte) bool {
+	if len(buf) < requestHeaderSize {
+		return false
+	}
+	method := binary.BigEndian.Uint16(buf)
+	if method != raftType && method != snapshotType {
+		plog.Errorf("invalid method type")
+		return false
+	}
+	h.method = method
+	h.size = binary.BigEndian.Uint64(buf[2:])
+	return true
+}

plugin/transport/quic/transport_client.go

@@ -0,0 +1,172 @@
+// Copyright 2017-2021 Lei Ni ([email protected]) and other contributors.
+//
+// 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://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 quic
+
+import (
+	"context"
+	"time"
+
+	"github.com/lni/dragonboat/v3/raftio"
+	"github.com/lni/dragonboat/v3/raftpb"
+	"github.com/lucas-clemente/quic-go"
+)
+
+func (q *quicTransport) GetConnection(ctx context.Context, target string) (raftio.IConnection, error) {
+	str, err := q.openSessionTo(ctx, target)
+	if err != nil {
+		return nil, err
+	}
+	return NewMessageConnection(str), nil
+}
+
+func (q *quicTransport) GetSnapshotConnection(ctx context.Context, target string) (raftio.ISnapshotConnection, error) {
+	str, err := q.openSessionTo(ctx, target)
+	if err != nil {
+		return nil, err
+	}
+	return NewSnapshotConnection(str), nil
+}
+
+func (q *quicTransport) openSessionTo(ctx context.Context, target string) (quic.Session, error) {
+	tlsConf, err := q.nhConfig.GetClientTLSConfig(target)
+	if err != nil {
+		return nil, err
+	}
+	tlsConf = tlsConf.Clone()
+	tlsConf.NextProtos = []string{TransportName}
+	session, err := quic.DialAddrContext(ctx, target, tlsConf, nil)
+	if err != nil {
+		return nil, err
+	}
+	return session, nil
+}
+
+type quicConnWriter struct {
+	header  []byte
+	payload []byte
+	session quic.Session
+}
+
+func NewMessageConnection(session quic.Session) raftio.IConnection {
+	return &quicMessageConnection{
+		quicConnWriter: &quicConnWriter{
+			session: session,
+			header:  make([]byte, requestHeaderSize),
+			payload: make([]byte, perConnBufSize),
+		},
+	}
+}
+
+type quicMessageConnection struct {
+	*quicConnWriter
+}
+
+func (q *quicMessageConnection) SendMessageBatch(batch raftpb.MessageBatch) error {
+	header := requestHeader{method: raftType}
+	sz := batch.SizeUpperLimit()
+	var buf []byte
+	if len(q.payload) < sz {
+		buf = make([]byte, sz)
+	} else {
+		buf = q.payload
+	}
+	n, err := batch.MarshalTo(buf)
+	if err != nil {
+		panic(err)
+	}
+	return q.writeMessage(header, buf[:n])
+}
+
+func (q *quicMessageConnection) Close() {
+	_ = q.session.CloseWithError(0, "")
+}
+
+func NewSnapshotConnection(session quic.Session) raftio.ISnapshotConnection {
+	return &quicSnapshotConnection{
+		quicConnWriter: &quicConnWriter{
+			session: session,
+			header:  make([]byte, requestHeaderSize),
+			payload: make([]byte, perConnBufSize),
+		},
+	}
+}
+
+type quicSnapshotConnection struct {
+	*quicConnWriter
+}
+
+func (q *quicSnapshotConnection) SendChunk(chunk raftpb.Chunk) error {
+	header := requestHeader{method: snapshotType}
+	sz := chunk.Size()
+	buf := make([]byte, sz)
+	n, err := chunk.MarshalTo(buf)
+	if err != nil {
+		panic(err)
+	}
+	return q.writeMessage(header, buf[:n])
+}
+
+func (q *quicSnapshotConnection) Close() {
+	str, err := q.session.OpenStreamSync(context.Background())
+	if err != nil {
+		return
+	}
+	defer func() {
+		_ = str.Close()
+	}()
+	if err := sendPoison(str, poisonNumber[:]); err != nil {
+		return
+	}
+	waitPoisonAck(q.session)
+	_ = q.session.CloseWithError(0, "")
+}
+
+func (q *quicConnWriter) writeMessage(header requestHeader, buf []byte) error {
+	str, err := q.session.OpenUniStreamSync(context.Background())
+	if err != nil {
+		return err
+	}
+	header.size = uint64(len(buf))
+	q.header = header.encode(q.header)
+	tt := time.Now().Add(magicNumberDuration).Add(headerDuration)
+	if err := str.SetWriteDeadline(tt); err != nil {
+		return err
+	}
+	if _, err := str.Write(magicNumber[:]); err != nil {
+		return err
+	}
+	if _, err := str.Write(q.header); err != nil {
+		return err
+	}
+	sent := 0
+	bufSize := int(recvBufSize)
+	for sent < len(buf) {
+		if sent+bufSize > len(buf) {
+			bufSize = len(buf) - sent
+		}
+		tt = time.Now().Add(writeDuration)
+		if err := str.SetWriteDeadline(tt); err != nil {
+			return err
+		}
+		if _, err := str.Write(buf[sent : sent+bufSize]); err != nil {
+			return err
+		}
+		sent += bufSize
+	}
+	if sent != len(buf) {
+		plog.Panicf("sent %d, buf len %d", sent, len(buf))
+	}
+	return nil
+}