mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-06 17:18:35 +08:00
e6da4a64b5
Pre-check import message to prevent pipeline block indefinitely. issue: https://github.com/milvus-io/milvus/issues/42414 --------- Signed-off-by: bigsheeper <yihao.dai@zilliz.com> Co-authored-by: chyezh <chyezh@outlook.com>
59 lines
1.8 KiB
Go
59 lines
1.8 KiB
Go
package registry
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
|
|
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
|
|
)
|
|
|
|
// init the message ack callbacks
|
|
func init() {
|
|
resetMessageAckCallbacks()
|
|
resetMessageCheckCallbacks()
|
|
}
|
|
|
|
// resetMessageAckCallbacks resets the message ack callbacks.
|
|
func resetMessageAckCallbacks() {
|
|
messageAckCallbacks = map[message.MessageType]*syncutil.Future[MessageAckCallback]{
|
|
message.MessageTypeDropPartition: syncutil.NewFuture[MessageAckCallback](),
|
|
message.MessageTypeImport: syncutil.NewFuture[MessageAckCallback](),
|
|
}
|
|
}
|
|
|
|
// MessageAckCallback is the callback function for the message type.
|
|
type MessageAckCallback = func(ctx context.Context, msg message.MutableMessage) error
|
|
|
|
// messageAckCallbacks is the map of message type to the callback function.
|
|
var messageAckCallbacks map[message.MessageType]*syncutil.Future[MessageAckCallback]
|
|
|
|
// RegisterMessageAckCallback registers the callback function for the message type.
|
|
func RegisterMessageAckCallback(typ message.MessageType, callback MessageAckCallback) {
|
|
future, ok := messageAckCallbacks[typ]
|
|
if !ok {
|
|
panic(fmt.Sprintf("the future of message callback for type %s is not registered", typ))
|
|
}
|
|
if future.Ready() {
|
|
// only for test, the register callback should be called once and only once
|
|
return
|
|
}
|
|
future.Set(callback)
|
|
}
|
|
|
|
// CallMessageAckCallback calls the callback function for the message type.
|
|
func CallMessageAckCallback(ctx context.Context, msg message.MutableMessage) error {
|
|
callbackFuture, ok := messageAckCallbacks[msg.MessageType()]
|
|
if !ok {
|
|
// No callback need tobe called, return nil
|
|
return nil
|
|
}
|
|
callback, err := callbackFuture.GetWithContext(ctx)
|
|
if err != nil {
|
|
return errors.Wrap(err, "when waiting callback registered")
|
|
}
|
|
return callback(ctx, msg)
|
|
}
|