From d9f02414a18568708bd4ad05b6e9364202b6791a Mon Sep 17 00:00:00 2001
From: Zixuan Liu <nodeces@gmail.com>
Date: Mon, 1 Mar 2021 18:31:21 +0800
Subject: [PATCH] fix: add RaftConfig to Config

Signed-off-by: Zixuan Liu <nodeces@gmail.com>
---
 config.go      |  3 +++
 dispatcher.go  |  3 ++-
 store/store.go | 16 ++++++++++++++--
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/config.go b/config.go
index ff329fa..6cbb9bf 100644
--- a/config.go
+++ b/config.go
@@ -3,6 +3,7 @@ package hraftdispatcher
 import (
 	"crypto/tls"
 	"github.com/casbin/casbin/v2"
+	"github.com/hashicorp/raft"
 )
 
 // Config holds dispatcher config.
@@ -21,4 +22,6 @@ type Config struct {
 	// You have to provide a peer certificate.
 	// We recommend using cfssl tool to create this certificates.
 	TLSConfig *tls.Config
+	// RaftConfig provides any necessary configuration for the Raft server.
+	RaftConfig *raft.Config
 }
diff --git a/dispatcher.go b/dispatcher.go
index d6d682f..a544c09 100644
--- a/dispatcher.go
+++ b/dispatcher.go
@@ -80,7 +80,8 @@ func NewHRaftDispatcher(config *Config) (*HRaftDispatcher, error) {
 			MaxPool: 5,
 			Logger:  nil,
 		},
-		Enforcer: config.Enforcer,
+		Enforcer:   config.Enforcer,
+		RaftConfig: config.RaftConfig,
 	}
 	s, err := store.NewStore(storeConfig)
 	if err != nil {
diff --git a/store/store.go b/store/store.go
index ecf3c3a..1d6ec22 100644
--- a/store/store.go
+++ b/store/store.go
@@ -39,6 +39,7 @@ type Store struct {
 	serverID      string
 
 	raft                   *raft.Raft
+	raftConfig             *raft.Config
 	networkTransportConfig *raft.NetworkTransportConfig
 	transport              raft.Transport
 	snapshotStore          raft.SnapshotStore
@@ -60,6 +61,7 @@ type Config struct {
 	Dir                    string
 	NetworkTransportConfig *raft.NetworkTransportConfig
 	Enforcer               casbin.IDistributedEnforcer
+	RaftConfig             *raft.Config
 }
 
 // NewStore return a instance of Store.
@@ -70,6 +72,7 @@ func NewStore(config *Config) (*Store, error) {
 		logger:                 zap.NewExample(),
 		networkTransportConfig: config.NetworkTransportConfig,
 		enforcer:               config.Enforcer,
+		raftConfig:             config.RaftConfig,
 	}
 
 	return s, nil
@@ -77,8 +80,17 @@ func NewStore(config *Config) (*Store, error) {
 
 // Start performs initialization and runs server
 func (s *Store) Start(enableBootstrap bool) error {
-	config := raft.DefaultConfig()
-	config.LocalID = raft.ServerID(s.serverID)
+	var config *raft.Config
+	if s.raftConfig == nil {
+		config = raft.DefaultConfig()
+		s.raftConfig = config
+	} else {
+		config = s.raftConfig
+	}
+
+	if len(config.LocalID) == 0 {
+		config.LocalID = raft.ServerID(s.serverID)
+	}
 
 	var transport raft.Transport
 	if s.inMemory {