(OraklNode) Track peers manually

node/pkg/raft/accessors.go

@@ -1,7 +1,5 @@
 package raft
 
-import "github.com/libp2p/go-libp2p/core/peer"
-
 func (r *Raft) IncreaseTerm() {
 	r.Mutex.Lock()
 	defer r.Mutex.Unlock()
@@ -75,11 +73,7 @@ func (r *Raft) UpdateVotedFor(votedFor string) {
 }
 
 func (r *Raft) SubscribersCount() int {
-	return len(r.Subscribers())
-}
-
-func (r *Raft) Subscribers() []peer.ID {
-	return r.Ps.ListPeers(r.Topic.String())
+	return r.Peers.Size()
 }
 
 func (r *Raft) GetHostId() string {

node/pkg/raft/raft.go

@@ -10,6 +10,7 @@ import (
 	"github.com/rs/zerolog/log"
 
 	errorSentinel "bisonai.com/orakl/node/pkg/error"
+	"bisonai.com/orakl/node/pkg/utils/set"
 	pubsub "github.com/libp2p/go-libp2p-pubsub"
 	"github.com/libp2p/go-libp2p/core/host"
 )
@@ -40,6 +41,9 @@ func NewRaftNode(
 		HeartbeatTimeout: HEARTBEAT_TIMEOUT,
 
 		LeaderJobTimeout: leaderJobTimeout,
+
+		PrevPeers: *set.NewSet[string](),
+		Peers:     *set.NewSet[string](),
 	}
 	return r
 }
@@ -104,12 +108,17 @@ func (r *Raft) handleMessage(ctx context.Context, msg Message) error {
 		return r.handleRequestVote(msg)
 	case ReplyVote:
 		return r.handleReplyVote(ctx, msg)
+	case ReplyHeartbeat:
+		return r.handleReplyHeartbeat(msg)
 	default:
 		return r.HandleCustomMessage(ctx, msg)
 	}
 }
 
 func (r *Raft) handleHeartbeat(msg Message) error {
+	r.Peers = r.PrevPeers
+	r.PrevPeers = *set.NewSet[string]()
+
 	if msg.SentFrom == r.GetHostId() {
 		return nil
 	}
@@ -151,7 +160,7 @@ func (r *Raft) handleHeartbeat(msg Message) error {
 		r.UpdateLeader(heartbeatMessage.LeaderID)
 	}
 
-	return nil
+	return r.sendReplyHeartbeat()
 }
 
 func (r *Raft) handleRequestVote(msg Message) error {
@@ -216,6 +225,17 @@ func (r *Raft) handleReplyVote(ctx context.Context, msg Message) error {
 	return nil
 }
 
+func (r *Raft) handleReplyHeartbeat(msg Message) error {
+	var replyHeartbeatMessage ReplyHeartbeatMessage
+	err := json.Unmarshal(msg.Data, &replyHeartbeatMessage)
+	if err != nil {
+		return err
+	}
+
+	r.PrevPeers.Add(msg.SentFrom)
+	return nil
+}
+
 // publishing messages
 
 func (r *Raft) PublishMessage(msg Message) error {
@@ -227,6 +247,7 @@ func (r *Raft) PublishMessage(msg Message) error {
 }
 
 func (r *Raft) sendHeartbeat() error {
+
 	heartbeatMessage := HeartbeatMessage{
 		LeaderID: r.GetHostId(),
 		Term:     r.GetCurrentTerm(),
@@ -250,6 +271,24 @@ func (r *Raft) sendHeartbeat() error {
 	return nil
 }
 
+func (r *Raft) sendReplyHeartbeat() error {
+	replyHeartbeatMessage := ReplyHeartbeatMessage{}
+	marshalledReplyHeartbeatMsg, err := json.Marshal(replyHeartbeatMessage)
+	if err != nil {
+		return err
+	}
+	message := Message{
+		Type:     ReplyHeartbeat,
+		SentFrom: r.GetHostId(),
+		Data:     json.RawMessage(marshalledReplyHeartbeatMsg),
+	}
+	err = r.PublishMessage(message)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (r *Raft) sendReplyVote(to string, voteGranted bool) error {
 	replyVoteMessage := ReplyRequestVoteMessage{
 		VoteGranted: voteGranted,

node/pkg/raft/types.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 	"time"
 
+	"bisonai.com/orakl/node/pkg/utils/set"
 	pubsub "github.com/libp2p/go-libp2p-pubsub"
 	"github.com/libp2p/go-libp2p/core/host"
 )
@@ -20,6 +21,7 @@ const (
 	ReplyVote          MessageType = "replyVote"
 	AppendEntries      MessageType = "appendEntries"
 	ReplyAppendEntries MessageType = "replyAppendEntries"
+	ReplyHeartbeat     MessageType = "replyHeartbeat"
 
 	Leader    RoleType = "leader"
 	Candidate RoleType = "candidate"
@@ -41,6 +43,8 @@ type HeartbeatMessage struct {
 	Term     int    `json:"term"`
 }
 
+type ReplyHeartbeatMessage struct{}
+
 type ReplyRequestVoteMessage struct {
 	VoteGranted bool   `json:"voteGranted"`
 	LeaderID    string `json:"leaderID"`
@@ -68,4 +72,7 @@ type Raft struct {
 	LeaderJobTicker     *time.Ticker
 	HandleCustomMessage func(context.Context, Message) error
 	LeaderJob           func() error
+
+	PrevPeers set.Set[string]
+	Peers     set.Set[string]
 }

node/pkg/utils/set/set.go

@@ -0,0 +1,26 @@
+package set
+
+type Set[T comparable] struct {
+	data map[T]struct{}
+}
+
+func NewSet[T comparable]() *Set[T] {
+	return &Set[T]{data: make(map[T]struct{})}
+}
+
+func (s *Set[T]) Add(element T) {
+	s.data[element] = struct{}{}
+}
+
+func (s *Set[T]) Remove(element T) {
+	delete(s.data, element)
+}
+
+func (s *Set[T]) Contains(element T) bool {
+	_, exists := s.data[element]
+	return exists
+}
+
+func (s *Set[T]) Size() int {
+	return len(s.data)
+}

node/pkg/utils/tests/set_test.go

@@ -0,0 +1,37 @@
+package tests
+
+import (
+	"testing"
+
+	"bisonai.com/orakl/node/pkg/utils/set"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSetBasicOperations(t *testing.T) {
+	s := set.NewSet[int]()
+	assert.Equal(t, 0, s.Size())
+
+	s.Add(1)
+	assert.Equal(t, 1, s.Size())
+	assert.True(t, s.Contains(1))
+	assert.False(t, s.Contains(2))
+
+	s.Remove(1)
+	assert.Equal(t, 0, s.Size())
+	assert.False(t, s.Contains(1))
+}
+
+func TestSetAddDuplicates(t *testing.T) {
+	s := set.NewSet[int]()
+	s.Add(1)
+	s.Add(1) // Attempt to add duplicate
+	assert.Equal(t, 1, s.Size(), "Set should not allow duplicates")
+	assert.True(t, s.Contains(1))
+}
+
+func TestSetOperationsOnEmpty(t *testing.T) {
+	s := set.NewSet[int]()
+	assert.False(t, s.Contains(1), "Empty set should not contain any element")
+	s.Remove(1) // Attempt to remove from empty set
+	assert.Equal(t, 0, s.Size(), "Size should remain 0 after remove operation on empty set")
+}