func NewUpsertOp(dbconnect string, tableName string, cd cube.CubeDescriber) (stream.Operator, stream.ProcessedNotifier, *Executor) { db, err := sql.Open("postgres", dbconnect) if err != nil { log.Fatal(err) } drv := db.Driver() conn, err := drv.Open(dbconnect) if err != nil { log.Fatal(err) } table := MakeTable(tableName, cd) exec := NewExecutor(table, conn) //exec.CreateBaseTable() ready := stream.NewNonBlockingProcessedNotifier(2) f := func(input stream.Object, out mapper.Outputer) { in := input.(*cube.TimeRepartitionedCube) visitor := func(part cube.Partition, c cube.Cuber) { exec.UpsertCube(part, c) } in.VisitPartitions(visitor) ready.Notify(1) } exit := func() { log.Println("Db Upser Exit: ") } name := "DbUpsert" gen := mapper.NewGenerator(f, name) gen.SingleExitCallback = exit op := mapper.NewOpFromGenerator(gen, name) op.Parallel = false return op, ready, exec }
func (src *Client) connect() error { defer func() { src.retries++ }() conn, err := net.Dial("tcp", src.addr) if err != nil { slog.Errorf("Cannot establish a connection with %s %v", src.addr, err) return err } wg_sub := &sync.WaitGroup{} defer wg_sub.Wait() rcvChData := make(chan stream.Object, 10) receiver := source.NewIOReaderSourceLengthDelim(conn) receiver.SetOut(rcvChData) rcvChCloseNotifier := make(chan bool) wg_sub.Add(1) go func() { defer wg_sub.Done() defer close(rcvChCloseNotifier) err := receiver.Run() if err != nil { slog.Errorf("Error in client reciever: %v", err) } }() //receiver will be closed by the sender after it is done sending. receiver closed via a hard stop. writeNotifier := stream.NewNonBlockingProcessedNotifier(2) sndChData := make(chan stream.Object, src.hwm) sndChCloseNotifier := make(chan bool) defer close(sndChData) sender := sink.NewMultiPartWriterSink(conn) sender.CompletedNotifier = writeNotifier sender.SetIn(sndChData) wg_sub.Add(1) go func() { defer receiver.Stop() //close receiver defer wg_sub.Done() defer close(sndChCloseNotifier) err := sender.Run() if err != nil { slog.Errorf("Error in client sender: %v", err) } }() //sender closed by closing the sndChData channel or by a hard stop if src.buf.Len() > 0 { leftover := src.buf.Reset() for i, value := range leftover { sendData(sndChData, value, i+1) } } timer := src.resetAckTimer() closing := false //defer log.Println("Exiting client loop") opName := stream.Name(src) writesNotCompleted := uint(0) for { upstreamCh := src.In() if !src.buf.CanAdd() || closing { //disable upstream listening upstreamCh = nil } if closing && src.buf.Len() == 0 { sendClose(sndChData, 100) return nil } select { case msg, ok := <-upstreamCh: if !ok { //softClose //make sure everything was sent closing = true } else { bytes := msg.([]byte) seq, err := src.buf.Add(bytes) if err != nil { slog.Fatalf("Error adding item to buffer %v", err) return err } sendData(sndChData, bytes, seq) writesNotCompleted += 1 metrics.Gm.Event(&opName) // These are batched //slog.Logf(logger.Levels.Debug, "Sent batch -- length %d seq %d", len(bytes), seq) } case cnt := <-writeNotifier.NotificationChannel(): writesNotCompleted -= cnt if timer == nil { slog.Debugf("Seting timer %v, %v", time.Now(), time.Now().UnixNano()) timer = src.resetAckTimer() } case obj, ok := <-rcvChData: slog.Debugf("in Rcv: %v", ok) if !ok { return errors.New("Connection to Server was Broken in Recieve Direction") } command, seq, _, err := parseMsg(obj.([]byte)) if err != nil { slog.Fatalf("%v", err) } if command == ACK { if src.processAck(seq) { timer = src.resetAckTimer() } } else { slog.Fatalf("Unknown Command: %v", command) } case <-rcvChCloseNotifier: //connection threw an eof to the reader? return errors.New("In Select: Recieve Closed") case <-sndChCloseNotifier: return errors.New("Connection to Server was Broken in Send Direction") case <-timer: return errors.New(fmt.Sprintf("Time Out Waiting For Ack, %d %v %v", len(rcvChData), time.Now(), time.Now().UnixNano())) case <-src.StopNotifier: sender.Stop() return nil } } }