milvus/internal/rootcoord/proxy_client_manager.go

// 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 rootcoord

import (
	"context"
	"fmt"
	"sync"

	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	"github.com/milvus-io/milvus/internal/log"
	"github.com/milvus-io/milvus/internal/proto/commonpb"
	"github.com/milvus-io/milvus/internal/proto/internalpb"
	"github.com/milvus-io/milvus/internal/proto/proxypb"
	"github.com/milvus-io/milvus/internal/types"
	"github.com/milvus-io/milvus/internal/util/sessionutil"
)

type proxyClientManager struct {
	core        *Core
	lock        sync.RWMutex
	proxyClient map[int64]types.Proxy
	helper      proxyClientManagerHelper
}

type proxyClientManagerHelper struct {
	afterConnect func()
}

var defaultClientManagerHelper = proxyClientManagerHelper{
	afterConnect: func() {},
}

func newProxyClientManager(c *Core) *proxyClientManager {
	return &proxyClientManager{
		core:        c,
		proxyClient: make(map[int64]types.Proxy),
		helper:      defaultClientManagerHelper,
	}
}

func (p *proxyClientManager) GetProxyClients(sessions []*sessionutil.Session) {
	for _, session := range sessions {
		p.AddProxyClient(session)
	}
}

func (p *proxyClientManager) AddProxyClient(session *sessionutil.Session) {
	p.lock.RLock()
	_, ok := p.proxyClient[session.ServerID]
	p.lock.RUnlock()
	if ok {
		return
	}

	go p.connect(session)
}

func (p *proxyClientManager) connect(session *sessionutil.Session) {
	pc, err := p.core.NewProxyClient(session)
	if err != nil {
		log.Warn("failed to create proxy client", zap.String("address", session.Address), zap.Int64("serverID", session.ServerID), zap.Error(err))
		return
	}

	p.lock.Lock()
	defer p.lock.Unlock()

	_, ok := p.proxyClient[session.ServerID]
	if ok {
		pc.Stop()
		return
	}
	p.proxyClient[session.ServerID] = pc
	log.Debug("succeed to create proxy client", zap.String("address", session.Address), zap.Int64("serverID", session.ServerID))
	p.helper.afterConnect()
}

func (p *proxyClientManager) DelProxyClient(s *sessionutil.Session) {
	p.lock.Lock()
	defer p.lock.Unlock()

	cli, ok := p.proxyClient[s.ServerID]
	if ok {
		cli.Stop()
	}

	delete(p.proxyClient, s.ServerID)
	log.Debug("remove proxy client", zap.String("proxy address", s.Address), zap.Int64("proxy id", s.ServerID))
}

func (p *proxyClientManager) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) error {
	p.lock.Lock()
	defer p.lock.Unlock()

	if len(p.proxyClient) == 0 {
		log.Warn("proxy client is empty, InvalidateCollectionMetaCache will not send to any client")
		return nil
	}

	group := &errgroup.Group{}
	for k, v := range p.proxyClient {
		k, v := k, v
		group.Go(func() error {
			sta, err := v.InvalidateCollectionMetaCache(ctx, request)
			if err != nil {
				return fmt.Errorf("InvalidateCollectionMetaCache failed, proxyID = %d, err = %s", k, err)
			}
			if sta.ErrorCode != commonpb.ErrorCode_Success {
				return fmt.Errorf("InvalidateCollectionMetaCache failed, proxyID = %d, err = %s", k, sta.Reason)
			}
			return nil
		})
	}
	return group.Wait()
}

func (p *proxyClientManager) ReleaseDQLMessageStream(ctx context.Context, in *proxypb.ReleaseDQLMessageStreamRequest) error {
	p.lock.Lock()
	defer p.lock.Unlock()

	if len(p.proxyClient) == 0 {
		log.Warn("proxy client is empty, ReleaseDQLMessageStream will not send to any client")
		return nil
	}

	group := &errgroup.Group{}
	for k, v := range p.proxyClient {
		k, v := k, v
		group.Go(func() error {
			sta, err := v.ReleaseDQLMessageStream(ctx, in)
			if err != nil {
				return fmt.Errorf("ReleaseDQLMessageStream failed, proxyID = %d, err = %s", k, err)
			}
			if sta.ErrorCode != commonpb.ErrorCode_Success {
				return fmt.Errorf("ReleaseDQLMessageStream failed, proxyID = %d, err = %s", k, sta.Reason)
			}
			return nil
		})
	}
	return group.Wait()
}

func (p *proxyClientManager) InvalidateCredentialCache(ctx context.Context, request *proxypb.InvalidateCredCacheRequest) error {
	p.lock.Lock()
	defer p.lock.Unlock()

	if len(p.proxyClient) == 0 {
		log.Warn("proxy client is empty, InvalidateCredentialCache will not send to any client")
		return nil
	}

	group := &errgroup.Group{}
	for k, v := range p.proxyClient {
		k, v := k, v
		group.Go(func() error {
			sta, err := v.InvalidateCredentialCache(ctx, request)
			if err != nil {
				return fmt.Errorf("InvalidateCredentialCache failed, proxyID = %d, err = %s", k, err)
			}
			if sta.ErrorCode != commonpb.ErrorCode_Success {
				return fmt.Errorf("InvalidateCredentialCache failed, proxyID = %d, err = %s", k, sta.Reason)
			}
			return nil
		})
	}
	return group.Wait()
}

func (p *proxyClientManager) UpdateCredentialCache(ctx context.Context, request *proxypb.UpdateCredCacheRequest) error {
	p.lock.Lock()
	defer p.lock.Unlock()

	if len(p.proxyClient) == 0 {
		log.Warn("proxy client is empty, UpdateCredentialCache will not send to any client")
		return nil
	}

	group := &errgroup.Group{}
	for k, v := range p.proxyClient {
		k, v := k, v
		group.Go(func() error {
			sta, err := v.UpdateCredentialCache(ctx, request)
			if err != nil {
				return fmt.Errorf("UpdateCredentialCache failed, proxyID = %d, err = %s", k, err)
			}
			if sta.ErrorCode != commonpb.ErrorCode_Success {
				return fmt.Errorf("UpdateCredentialCache failed, proxyID = %d, err = %s", k, sta.Reason)
			}
			return nil
		})
	}
	return group.Wait()
}

func (p *proxyClientManager) ClearCredUsersCache(ctx context.Context, request *internalpb.ClearCredUsersCacheRequest) error {
	p.lock.Lock()
	defer p.lock.Unlock()

	if len(p.proxyClient) == 0 {
		log.Warn("proxy client is empty, ClearCredUsersCache will not send to any client")
		return nil
	}

	group := &errgroup.Group{}
	for k, v := range p.proxyClient {
		k, v := k, v
		group.Go(func() error {
			sta, err := v.ClearCredUsersCache(ctx, request)
			if err != nil {
				return fmt.Errorf("ClearCredUsersCache failed, proxyID = %d, err = %s", k, err)
			}
			if sta.ErrorCode != commonpb.ErrorCode_Success {
				return fmt.Errorf("ClearCredUsersCache failed, proxyID = %d, err = %s", k, sta.Reason)
			}
			return nil
		})
	}
	return group.Wait()
}