diff --git a/client/subscriber.go b/client/subscriber.go --- a/client/subscriber.go +++ b/client/subscriber.go @@ -6,8 +6,11 @@ "encoding/binary" "encoding/json" "errors" "fmt" + "io" + "log/slog" "net" "sync" + "syscall" "time" "github.com/willdot/messagebroker/internal/server" @@ -17,8 +20,10 @@ type connOpp func(conn net.Conn) error // 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 @@ -30,9 +35,20 @@ } 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() @@ -44,7 +60,50 @@ op := func(conn net.Conn) error { 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 @@ -53,7 +112,14 @@ op := func(conn net.Conn) error { 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 { @@ -189,11 +255,37 @@ return } 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 + } + } } diff --git a/example/main.go b/example/main.go --- a/example/main.go +++ b/example/main.go @@ -11,22 +11,33 @@ "github.com/willdot/messagebroker/client" "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) @@ -56,9 +67,6 @@ for msg := range consumer.Messages() { slog.Info("received message", "message", string(msg.Data)) msg.Ack(true) } - - time.Sleep(time.Second * 30) - } func sendMessages() {