From 6481cd59b07a0ac42f8277204a3afcf17d4d3f17 Mon Sep 17 00:00:00 2001
From: AndrewZuo01 <andrew.zuo@wustl.edu>
Date: Wed, 27 Dec 2023 09:58:56 +0800
Subject: [PATCH] update notification

---
 go.mod                               |  5 ++++-
 internal/rpc/friend/friend.go        |  2 ++
 pkg/rpcclient/notification/friend.go |  6 +++++-
 tools/component/component.go         | 22 +++++++++++++++-------
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 0d614c614..6efee2993 100644
--- a/go.mod
+++ b/go.mod
@@ -34,7 +34,7 @@ require github.com/google/uuid v1.3.1
 require (
 	github.com/IBM/sarama v1.41.3
 	github.com/OpenIMSDK/protocol v0.0.37
-	github.com/OpenIMSDK/tools v0.0.20
+	github.com/OpenIMSDK/tools v0.0.22
 	github.com/aliyun/aliyun-oss-go-sdk v2.2.9+incompatible
 	github.com/go-redis/redis v6.15.9+incompatible
 	github.com/redis/go-redis/v9 v9.2.1
@@ -154,3 +154,6 @@ require (
 	golang.org/x/crypto v0.14.0 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 )
+replace (
+	github.com/OpenIMSDK/protocol v0.0.37 => github.com/AndrewZuo01/protocol v0.0.0-20231226102058-3702dc94b5dd
+)
\ No newline at end of file
diff --git a/internal/rpc/friend/friend.go b/internal/rpc/friend/friend.go
index ea34ba2cd..60ded00f2 100644
--- a/internal/rpc/friend/friend.go
+++ b/internal/rpc/friend/friend.go
@@ -471,5 +471,7 @@ func (s *friendServer) UpdateFriends(
 	}
 
 	resp := &pbfriend.UpdateFriendsResp{}
+
+	s.notificationSender.FriendsInfoUpdateNotification(ctx, req.OwnerUserID)
 	return resp, nil
 }
diff --git a/pkg/rpcclient/notification/friend.go b/pkg/rpcclient/notification/friend.go
index b061a24ae..3e26b3038 100644
--- a/pkg/rpcclient/notification/friend.go
+++ b/pkg/rpcclient/notification/friend.go
@@ -196,7 +196,11 @@ func (f *FriendNotificationSender) FriendRemarkSetNotification(ctx context.Conte
 	tips.FromToUserID.ToUserID = toUserID
 	return f.Notification(ctx, fromUserID, toUserID, constant.FriendRemarkSetNotification, &tips)
 }
-
+func (f *FriendNotificationSender) FriendsInfoUpdateNotification(ctx context.Context, toUserID string) error {
+	tips := sdkws.FriendsInfoUpdateTips{}
+	tips.FromToUserID.ToUserID = toUserID
+	return f.Notification(ctx, toUserID, toUserID, constant.FriendsInfoUpdateNotification, &tips)
+}
 func (f *FriendNotificationSender) BlackAddedNotification(ctx context.Context, req *pbfriend.AddBlackReq) error {
 	tips := sdkws.BlackAddedTips{FromToUserID: &sdkws.FromToUserID{}}
 	tips.FromToUserID.FromUserID = req.OwnerUserID
diff --git a/tools/component/component.go b/tools/component/component.go
index 28ea7a2fe..616ffec2d 100644
--- a/tools/component/component.go
+++ b/tools/component/component.go
@@ -16,6 +16,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"net"
@@ -285,10 +286,23 @@ func checkZookeeper() (string, error) {
 
 	// Connect to Zookeeper
 	str := "the addr is:" + address
-	c, _, err := zk.Connect(zookeeperAddresses, time.Second) // Adjust the timeout as necessary
+	c, eventChan, err := zk.Connect(zookeeperAddresses, time.Second) // Adjust the timeout as necessary
 	if err != nil {
 		return "", errs.Wrap(errStr(err, str))
 	}
+	timeout := time.After(5 * time.Second)
+	for {
+		select {
+		case event := <-eventChan:
+			if event.State == zk.StateConnected {
+				fmt.Println("Connected to Zookeeper")
+				goto Connected
+			}
+		case <-timeout:
+			return "", errs.Wrap(errors.New("timeout waiting for Zookeeper connection"), "Zookeeper Addr: "+strings.Join(config.Config.Zookeeper.ZkAddr, " "))
+		}
+	}
+Connected:
 	defer c.Close()
 
 	// Set authentication if username and password are provided
@@ -298,12 +312,6 @@ func checkZookeeper() (string, error) {
 		}
 	}
 
-	// Check if Zookeeper is reachable
-	_, _, err = c.Get("/")
-	if err != nil {
-		return "", errs.Wrap(err, str)
-	}
-
 	return str, nil
 }