init

comet/comet.go

@@ -0,0 +1,126 @@
+package comet
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"net"
+	"time"
+
+	"github.com/Terry-Mao/goim/pkg/ip"
+	"github.com/zhenjl/cityhash"
+	"gitlab.33.cn/chat/im-pkg/trace"
+	comet "gitlab.33.cn/chat/im/api/comet/grpc"
+	logic "gitlab.33.cn/chat/im/api/logic/grpc"
+	"gitlab.33.cn/chat/im/comet/conf"
+	"gitlab.33.cn/chat/im/common"
+	"gitlab.33.cn/chat/im/naming"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/resolver"
+)
+
+// Comet is comet server.
+type Comet struct {
+	c         *conf.Config
+	round     *Round
+	buckets   []*Bucket // subkey bucket
+	bucketIdx uint32
+
+	serverID       string
+	logicRPCClient logic.LogicClient
+}
+
+// NewServer returns a new Server.
+func New(c *conf.Config) *Comet {
+	s := &Comet{
+		c:              c,
+		round:          NewRound(c),
+		logicRPCClient: newLogicClient(c),
+	}
+	// init bucket
+	s.buckets = make([]*Bucket, c.Bucket.Size)
+	s.bucketIdx = uint32(c.Bucket.Size)
+	for i := 0; i < c.Bucket.Size; i++ {
+		s.buckets[i] = NewBucket(c.Bucket)
+	}
+	addr := ip.InternalIP()
+	_, port, _ := net.SplitHostPort(c.RPCServer.Addr)
+	s.serverID = "grpc://" + addr + ":" + port
+	return s
+}
+
+func newLogicClient(c *conf.Config) logic.LogicClient {
+	rb := naming.NewResolver(c.Reg.RegAddrs, c.LogicRPCClient.Schema)
+	resolver.Register(rb)
+
+	addr := fmt.Sprintf("%s:///%s", c.LogicRPCClient.Schema, c.LogicRPCClient.SrvName) // "schema://[authority]/service"
+	fmt.Println("rpc client call addr:", addr)
+
+	conn, err := common.NewGRPCConn(addr, time.Duration(c.LogicRPCClient.Dial), grpc.WithUnaryInterceptor(trace.OpentracingClientInterceptor))
+	if err != nil {
+		panic(err)
+	}
+	return logic.NewLogicClient(conn)
+}
+
+// Buckets return all buckets.
+func (s *Comet) Buckets() []*Bucket {
+	return s.buckets
+}
+
+// Bucket get the bucket by subkey.
+func (s *Comet) Bucket(subKey string) *Bucket {
+	idx := cityhash.CityHash32([]byte(subKey), uint32(len(subKey))) % s.bucketIdx
+	return s.buckets[idx]
+}
+
+// RandServerHearbeat rand server heartbeat.
+func (s *Comet) RandServerHearbeat() time.Duration {
+	return time.Duration(s.c.Protocol.MinHeartbeat) + time.Duration(rand.Int63n(int64(s.c.Protocol.MaxHeartbeat-s.c.Protocol.MinHeartbeat)))
+}
+
+// Close close the server.
+func (s *Comet) Close() (err error) {
+	return
+}
+
+// Connect connected a connection.
+func (s *Comet) Connect(c context.Context, p *comet.Proto) (key string, hb time.Duration, err error) {
+	var (
+		req   logic.ConnectReq
+		reply *logic.ConnectReply
+	)
+
+	req.Server = s.serverID
+	req.Proto = p
+	reply, err = s.logicRPCClient.Connect(c, &req)
+	if err != nil {
+		return
+	}
+
+	return reply.Key, time.Duration(reply.Heartbeat), nil
+}
+
+// Disconnect disconnected a connection.
+func (s *Comet) Disconnect(c context.Context, key string) (err error) {
+	_, err = s.logicRPCClient.Disconnect(context.Background(), &logic.DisconnectReq{
+		Server: s.serverID,
+		Key:    key,
+	})
+	return
+}
+
+// Heartbeat heartbeat a connection session.
+func (s *Comet) Heartbeat(ctx context.Context, key string) (err error) {
+	_, err = s.logicRPCClient.Heartbeat(ctx, &logic.HeartbeatReq{
+		Server: s.serverID,
+		Key:    key,
+	})
+	return
+}
+
+// Receive receive a message.
+func (s *Comet) Receive(ctx context.Context, key string, p *comet.Proto) (err error) {
+	_, err = s.logicRPCClient.Receive(ctx, &logic.ReceiveReq{Key: key, Proto: p})
+	return
+}