package proxy import ( "context" "errors" "log" "github.com/zilliztech/milvus-distributed/internal/msgstream" "github.com/zilliztech/milvus-distributed/internal/proto/commonpb" "github.com/zilliztech/milvus-distributed/internal/proto/internalpb" "github.com/zilliztech/milvus-distributed/internal/proto/masterpb" "github.com/zilliztech/milvus-distributed/internal/proto/servicepb" ) type task interface { Id() UniqueID // return ReqId Type() internalpb.MsgType BeginTs() Timestamp EndTs() Timestamp SetTs(ts Timestamp) PreExecute() error Execute() error PostExecute() error WaitToFinish() error Notify(err error) } type baseInsertTask = msgstream.InsertMsg type InsertTask struct { baseInsertTask ts Timestamp done chan error resultChan chan *servicepb.IntegerRangeResponse manipulationMsgStream *msgstream.PulsarMsgStream ctx context.Context cancel context.CancelFunc } func (it *InsertTask) SetTs(ts Timestamp) { it.ts = ts } func (it *InsertTask) BeginTs() Timestamp { return it.ts } func (it *InsertTask) EndTs() Timestamp { return it.ts } func (it *InsertTask) Id() UniqueID { return it.ReqId } func (it *InsertTask) Type() internalpb.MsgType { return it.MsgType } func (it *InsertTask) PreExecute() error { return nil } func (it *InsertTask) Execute() error { var tsMsg msgstream.TsMsg = it msgPack := &msgstream.MsgPack{ BeginTs: it.BeginTs(), EndTs: it.EndTs(), Msgs: make([]*msgstream.TsMsg, 1), } msgPack.Msgs[0] = &tsMsg it.manipulationMsgStream.Produce(msgPack) return nil } func (it *InsertTask) PostExecute() error { return nil } func (it *InsertTask) WaitToFinish() error { defer it.cancel() for { select { case err := <-it.done: return err case <-it.ctx.Done(): log.Print("wait to finish failed, timeout!") return errors.New("wait to finish failed, timeout!") } } } func (it *InsertTask) Notify(err error) { it.done <- err } type CreateCollectionTask struct { internalpb.CreateCollectionRequest masterClient masterpb.MasterClient done chan error resultChan chan *commonpb.Status ctx context.Context cancel context.CancelFunc } func (cct *CreateCollectionTask) Id() UniqueID { return cct.ReqId } func (cct *CreateCollectionTask) Type() internalpb.MsgType { return cct.MsgType } func (cct *CreateCollectionTask) BeginTs() Timestamp { return cct.Timestamp } func (cct *CreateCollectionTask) EndTs() Timestamp { return cct.Timestamp } func (cct *CreateCollectionTask) SetTs(ts Timestamp) { cct.Timestamp = ts } func (cct *CreateCollectionTask) PreExecute() error { return nil } func (cct *CreateCollectionTask) Execute() error { resp, err := cct.masterClient.CreateCollection(cct.ctx, &cct.CreateCollectionRequest) if err != nil { log.Printf("create collection failed, error= %v", err) cct.resultChan <- &commonpb.Status{ ErrorCode: commonpb.ErrorCode_UNEXPECTED_ERROR, Reason: err.Error(), } } else { cct.resultChan <- resp } return err } func (cct *CreateCollectionTask) PostExecute() error { return nil } func (cct *CreateCollectionTask) WaitToFinish() error { defer cct.cancel() for { select { case err := <-cct.done: return err case <-cct.ctx.Done(): log.Print("wait to finish failed, timeout!") return errors.New("wait to finish failed, timeout!") } } } func (cct *CreateCollectionTask) Notify(err error) { cct.done <- err }