sparrow/services/emqx-agent/main.go

package main

import (
	"context"
	"log"
	"math"
	"sparrow/pkg/rpcs"
	"sparrow/pkg/server"
	"sync"
	"time"

	"google.golang.org/grpc"
	pb "sparrow/services/emqx-agent/protobuf"
)

type Counter struct {
	count    int64
	start    int64
	duration int64
	maxCount int64
	lock     sync.Mutex
}

func NewCounter(duration int64, maxCount int64) *Counter {
	counter := &Counter{}
	counter.count = 0
	counter.start = time.Now().UnixNano() / 1e6 // 当前毫秒时间戳
	if duration == 0 {
		counter.duration = math.MaxInt64 // duration传0表示没有时间间隔限制，计数器不刷新
	} else {
		counter.duration = duration
	}
	counter.maxCount = maxCount
	return counter
}

// Count 计数器计数
// n: 计数值
// refersh: 计数器是否刷新
// limit: 是否达到计数最大值
// num: 计数后计数器的值
func (counter *Counter) Count(n int64) (refresh bool, limit bool, num int64) {
	now := time.Now().UnixNano() / 1e6
	counter.lock.Lock()
	defer counter.lock.Unlock()

	if now-counter.start < counter.duration {
		counter.count += n
		num = counter.count
		limit = num > counter.maxCount
	} else {
		// num = counter.count  // 刷新前的最大计数
		counter.start = now
		counter.count = 0
		refresh = true
	}
	return
}

func (counter *Counter) GetCount() (num int64) {
	counter.lock.Lock()
	defer counter.lock.Unlock()
	return counter.count
}

const (
	port = ":9000"
)

var cnter *Counter = NewCounter(0, 100)

// emqttServer is used to implement emqx_exhook_v1.s *emqttServer
type emqttServer struct {
	pb.UnimplementedHookProviderServer
}

// HookProviderServer callbacks

func (s *emqttServer) OnProviderLoaded(ctx context.Context, in *pb.ProviderLoadedRequest) (*pb.LoadedResponse, error) {
	cnter.Count(1)
	hooks := []*pb.HookSpec{
		&pb.HookSpec{Name: "client.connect"},
		&pb.HookSpec{Name: "client.connack"},
		&pb.HookSpec{Name: "client.connected"},
		&pb.HookSpec{Name: "client.disconnected"},
		&pb.HookSpec{Name: "client.authenticate"},
		&pb.HookSpec{Name: "client.authorize"},
		&pb.HookSpec{Name: "client.subscribe"},
		&pb.HookSpec{Name: "client.unsubscribe"},
		&pb.HookSpec{Name: "session.created"},
		&pb.HookSpec{Name: "session.subscribed"},
		&pb.HookSpec{Name: "session.unsubscribed"},
		&pb.HookSpec{Name: "session.resumed"},
		&pb.HookSpec{Name: "session.discarded"},
		&pb.HookSpec{Name: "session.takenover"},
		&pb.HookSpec{Name: "session.terminated"},
		&pb.HookSpec{Name: "message.publish"},
		&pb.HookSpec{Name: "message.delivered"},
		&pb.HookSpec{Name: "message.acked"},
		&pb.HookSpec{Name: "message.dropped"},
	}
	return &pb.LoadedResponse{Hooks: hooks}, nil
}

func (s *emqttServer) OnProviderUnloaded(ctx context.Context, in *pb.ProviderUnloadedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnClientConnect(ctx context.Context, in *pb.ClientConnectRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnClientConnack(ctx context.Context, in *pb.ClientConnackRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnClientConnected(ctx context.Context, in *pb.ClientConnectedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnClientDisconnected(ctx context.Context, in *pb.ClientDisconnectedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnClientAuthenticate(ctx context.Context, in *pb.ClientAuthenticateRequest) (*pb.ValuedResponse, error) {
	cnter.Count(1)
	reply := &pb.ValuedResponse{}
	reply.Type = pb.ValuedResponse_STOP_AND_RETURN
	reply.Value = &pb.ValuedResponse_BoolResult{BoolResult: true}
	return reply, nil
}

func (s *emqttServer) OnClientAuthorize(ctx context.Context, in *pb.ClientAuthorizeRequest) (*pb.ValuedResponse, error) {
	cnter.Count(1)
	reply := &pb.ValuedResponse{}
	reply.Type = pb.ValuedResponse_STOP_AND_RETURN
	reply.Value = &pb.ValuedResponse_BoolResult{BoolResult: true}
	return reply, nil
}

func (s *emqttServer) OnClientSubscribe(ctx context.Context, in *pb.ClientSubscribeRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnClientUnsubscribe(ctx context.Context, in *pb.ClientUnsubscribeRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnSessionCreated(ctx context.Context, in *pb.SessionCreatedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}
func (s *emqttServer) OnSessionSubscribed(ctx context.Context, in *pb.SessionSubscribedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnSessionUnsubscribed(ctx context.Context, in *pb.SessionUnsubscribedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnSessionResumed(ctx context.Context, in *pb.SessionResumedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnSessionDiscarded(ctx context.Context, in *pb.SessionDiscardedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnSessionTakenover(ctx context.Context, in *pb.SessionTakenoverRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnSessionTerminated(ctx context.Context, in *pb.SessionTerminatedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnMessagePublish(ctx context.Context, in *pb.MessagePublishRequest) (*pb.ValuedResponse, error) {
	cnter.Count(1)
	in.Message.Payload = []byte("hardcode payload by exhook-svr-go :)")
	reply := &pb.ValuedResponse{}
	reply.Type = pb.ValuedResponse_STOP_AND_RETURN
	reply.Value = &pb.ValuedResponse_Message{Message: in.Message}
	return reply, nil
}

//case2: stop publish the `t/d` messages
//func (s *emqttServer) OnMessagePublish(ctx context.Context, in *pb.MessagePublishRequest) (*pb.ValuedResponse, error) {
//	cnter.Count(1)
//    if in.Message.Topic == "t/d" {
//        in.Message.Headers["allow_publish"] = "false"
//        in.Message.Payload = []byte("")
//    }
//	reply := &pb.ValuedResponse{}
//	reply.Type = pb.ValuedResponse_STOP_AND_RETURN
//	reply.Value = &pb.ValuedResponse_Message{Message: in.Message}
//	return reply, nil
//}

func (s *emqttServer) OnMessageDelivered(ctx context.Context, in *pb.MessageDeliveredRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnMessageDropped(ctx context.Context, in *pb.MessageDroppedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func (s *emqttServer) OnMessageAcked(ctx context.Context, in *pb.MessageAckedRequest) (*pb.EmptySuccess, error) {
	cnter.Count(1)
	return &pb.EmptySuccess{}, nil
}

func main() {

	// init emqttServer
	err := server.Init(rpcs.EmqxAgentServiceName)
	if err != nil {
		server.Log.Fatal(err)
		return
	}
	s := grpc.NewServer()
	pb.RegisterHookProviderServer(s, &emqttServer{})
	log.Println("Started gRPC emqttServer on ::9000")
	// register a http handler
	err = server.RegisterHTTPHandler(s)
	if err != nil {
		server.Log.Errorf("RegisterHTTPHandler Error: %s", err)
		return
	}
}