milvus/pkg/eventlog/grpc.go

173 lines
3.6 KiB
Go
Raw Normal View History

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package eventlog
import (
"net"
"sync"
"time"
"github.com/milvus-io/milvus/pkg/util/conc"
"github.com/milvus-io/milvus/pkg/util/funcutil"
"github.com/milvus-io/milvus/pkg/util/typeutil"
"go.uber.org/atomic"
"google.golang.org/grpc"
)
var (
grpcLog atomic.Pointer[grpcLogger]
sf conc.Singleflight[*grpcLogger]
)
// grpcLogger is a Logger with dispatches streaming Evt to client listeners.
type grpcLogger struct {
level atomic.Int32
lis net.Listener
port int
clients *typeutil.ConcurrentMap[string, *listenerClient]
}
func (l *grpcLogger) SetLevel(lvl Level) {
l.level.Store(int32(lvl))
}
func (l *grpcLogger) GetLevel() Level {
return Level(l.level.Load())
}
func (l *grpcLogger) Record(evt Evt) {
if evt.Level() < Level(l.level.Load()) {
return
}
l.clients.Range(func(key string, client *listenerClient) bool {
client.Notify(evt)
return true
})
}
func (l *grpcLogger) RecordFunc(lvl Level, fn func() Evt) {
if lvl < l.GetLevel() {
return
}
l.Record(fn())
}
func (l *grpcLogger) Flush() error {
return nil
}
func (l *grpcLogger) Listen(req *ListenRequest, svr EventLogService_ListenServer) error {
client := newListenerClient()
key := funcutil.RandomString(8)
l.clients.Insert(key, client)
defer func() {
client, ok := l.clients.GetAndRemove(key)
if ok {
client.Stop()
}
}()
for {
select {
case evt := <-client.ch:
err := svr.Send(&Event{
Level: evt.Level(),
Type: evt.Type(),
Data: evt.Raw(),
Ts: time.Now().UnixNano(),
})
if err != nil {
return nil
}
case <-svr.Context().Done():
return nil
case <-client.closed:
}
}
}
func (l *grpcLogger) Close() {
if l.lis != nil {
l.lis.Close()
}
}
// getGrpcLogger starts or returns the singleton grpcLogger listening port.
func getGrpcLogger() (int, error) {
if l := grpcLog.Load(); l != nil {
return l.port, nil
}
l, err, _ := sf.Do("grpc_evt_log", func() (*grpcLogger, error) {
if grpcLog.Load() != nil {
return grpcLog.Load(), nil
}
lis, err := net.Listen("tcp", ":0")
if err != nil {
return nil, err
}
port := lis.Addr().(*net.TCPAddr).Port
svr := grpc.NewServer()
l := &grpcLogger{
lis: lis,
port: port,
clients: typeutil.NewConcurrentMap[string, *listenerClient](),
}
l.SetLevel(Level_Debug)
RegisterEventLogServiceServer(svr, l)
go svr.Serve(lis)
grpcLog.Store(l)
getGlobalLogger().Register("grpc_logger", l)
return l, nil
})
if err != nil {
return -1, err
}
return l.port, nil
}
type listenerClient struct {
ch chan Evt
closed chan struct{}
once sync.Once
}
func newListenerClient() *listenerClient {
return &listenerClient{
ch: make(chan Evt, 100),
closed: make(chan struct{}),
}
}
func (c *listenerClient) Notify(l Evt) {
select {
case c.ch <- l:
default:
}
}
func (c *listenerClient) Stop() {
c.once.Do(func() {
close(c.ch)
})
}