reconnect clients (#16)

* publisher now attempts to reconnect if the server is down

* consumer now attempts to reconnect and resubscribe

client/subscriber.go

···
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"log/slog"
 	"net"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/willdot/messagebroker/internal/server"
···
 
 // Subscriber allows subscriptions to a server and the consumption of messages
 type Subscriber struct {
-	conn   net.Conn
-	connMu sync.Mutex
+	conn             net.Conn
+	connMu           sync.Mutex
+	subscribedTopics []string
+	addr             string
 }
 
 // NewSubscriber will connect to the server at the given address
···
 
 	return &Subscriber{
 		conn: conn,
+		addr: addr,
 	}, nil
 }
 
+func (s *Subscriber) reconnect() error {
+	conn, err := net.Dial("tcp", s.addr)
+	if err != nil {
+		return fmt.Errorf("failed to dial: %w", err)
+	}
+
+	s.conn = conn
+	return nil
+}
+
 // Close cleanly shuts down the subscriber
 func (s *Subscriber) Close() error {
 	return s.conn.Close()
···
 		return subscribeToTopics(conn, topicNames, startAtType, startAtIndex)
 	}
 
-	return s.connOperation(op)
+	err := s.connOperation(op)
+	if err != nil {
+		return fmt.Errorf("failed to subscribe to topics: %w", err)
+	}
+
+	s.addToSubscribedTopics(topicNames)
+
+	return nil
+}
+
+func (s *Subscriber) addToSubscribedTopics(topics []string) {
+	existingSubs := make(map[string]struct{})
+	for _, topic := range s.subscribedTopics {
+		existingSubs[topic] = struct{}{}
+	}
+
+	for _, topic := range topics {
+		existingSubs[topic] = struct{}{}
+	}
+
+	subs := make([]string, 0, len(existingSubs))
+	for topic := range existingSubs {
+		subs = append(subs, topic)
+	}
+
+	s.subscribedTopics = subs
+}
+
+func (s *Subscriber) removeTopicsFromSubscription(topics []string) {
+	existingSubs := make(map[string]struct{})
+	for _, topic := range s.subscribedTopics {
+		existingSubs[topic] = struct{}{}
+	}
+
+	for _, topic := range topics {
+		delete(existingSubs, topic)
+	}
+
+	subs := make([]string, 0, len(existingSubs))
+	for topic := range existingSubs {
+		subs = append(subs, topic)
+	}
+
+	s.subscribedTopics = subs
 }
 
 // UnsubscribeToTopics will unsubscribe to the provided topics
···
 		return unsubscribeToTopics(conn, topicNames)
 	}
 
-	return s.connOperation(op)
+	err := s.connOperation(op)
+	if err != nil {
+		return fmt.Errorf("failed to unsubscribe to topics: %w", err)
+	}
+
+	s.removeTopicsFromSubscription(topicNames)
+
+	return nil
 }
 
 func subscribeToTopics(conn net.Conn, topicNames []string, startAtType server.StartAtType, startAtIndex int) error {
···
 		}
 
 		err := s.readMessage(ctx, consumer.msgs)
-		if err != nil {
-			// TODO: if broken pipe, we need to somehow reconnect and subscribe again....YIKES
+		if err == nil {
+			continue
+		}
+
+		// if we couldn't connect to the server, attempt to reconnect
+		if !errors.Is(err, syscall.EPIPE) && !errors.Is(err, io.EOF) {
+			slog.Error("failed to read message", "error", err)
 			consumer.Err = err
 			return
 		}
+
+		slog.Info("attempting to reconnect")
+
+		for i := 0; i < 5; i++ {
+			time.Sleep(time.Millisecond * 500)
+			err = s.reconnect()
+			if err == nil {
+				break
+			}
+
+			slog.Error("Failed to reconnect", "error", err, "attempt", i)
+		}
+
+		slog.Info("attempting to resubscribe")
+
+		err = s.SubscribeToTopics(s.subscribedTopics, server.Current, 0)
+		if err != nil {
+			consumer.Err = fmt.Errorf("failed to subscribe to topics after reconnecting: %w", err)
+			return
+		}
+
 	}
 }
 

example/main.go

···
 	"github.com/willdot/messagebroker/internal/server"
 )
 
-var publish *bool
+// var publish *bool
+// var consume *bool
 var consumeFrom *int
+var clientType *string
 
 const (
 	topic = "topic-a"
 )
 
 func main() {
-	publish = flag.Bool("publish", false, "will also publish messages every 500ms until client is stopped")
+	clientType = flag.String("client-type", "consume", "consume or publish (default consume)")
+	// publish = flag.Bool("publish", false, "will publish messages every 500ms until client is stopped")
+	// consume = flag.Bool("consume", false, "will consume messages until client is stopped")
 	consumeFrom = flag.Int("consume-from", -1, "index of message to start consuming from. If not set it will consume from the most recent")
 	flag.Parse()
 
-	if *publish {
-		go sendMessages()
+	switch *clientType {
+	case "consume":
+		consume()
+	case "publish":
+		sendMessages()
+	default:
+		fmt.Println("unknown client type")
 	}
+}
 
+func consume() {
 	sub, err := client.NewSubscriber(":3000")
 	if err != nil {
 		panic(err)
···
 		slog.Info("received message", "message", string(msg.Data))
 		msg.Ack(true)
 	}
-
-	time.Sleep(time.Second * 30)
-
 }
 
 func sendMessages() {