func sendProgressUpdate(conn net.Conn, connMu *sync.Mutex, filesProcessed, filesTotal int) (int64, error) { seg := capn.NewBuffer(nil) z := proto.NewRootZ(seg) p := proto.NewProgressUpdate(seg) p.SetFilesprocessed(uint64(filesProcessed)) p.SetFilestotal(uint64(filesTotal)) z.SetProgressupdate(p) connMu.Lock() defer connMu.Unlock() return seg.WriteToPacked(conn) }
func queryBackend(queryid string, backend string, backendidx int, sourceQuery []byte) { // When exiting this function, check that all results were processed. If // not, the backend query must have failed for some reason. Send a progress // update to prevent the query from running forever. defer func() { filesTotal := state[queryid].filesTotal[backendidx] if state[queryid].filesProcessed[backendidx] == filesTotal { return } if filesTotal == -1 { filesTotal = 0 } seg := capn.NewBuffer(nil) p := proto.NewProgressUpdate(seg) p.SetFilesprocessed(uint64(filesTotal)) p.SetFilestotal(uint64(filesTotal)) storeProgress(queryid, backendidx, p) addEventMarshal(queryid, &Error{ Type: "error", ErrorType: "backendunavailable", }) }() // TODO: switch in the config log.Printf("[%s] [src:%s] connecting...\n", queryid, backend) conn, err := net.DialTimeout("tcp", strings.Replace(backend, "28082", "26082", -1), 5*time.Second) if err != nil { log.Printf("[%s] [src:%s] Connection failed: %v\n", queryid, backend, err) return } defer conn.Close() if _, err := conn.Write(sourceQuery); err != nil { log.Printf("[%s] [src:%s] could not send query: %v\n", queryid, backend, err) return } bufferedReader := bufio.NewReaderSize(conn, 65536) bstate := state[queryid].perBackend[backendidx] tempFileWriter := bstate.tempFileWriter var capnbuf bytes.Buffer var written countingWriter for !state[queryid].done { conn.SetReadDeadline(time.Now().Add(10 * time.Second)) written = 0 tee := io.TeeReader(bufferedReader, io.MultiWriter( tempFileWriter, &written)) seg, err := capn.ReadFromPackedStream(tee, &capnbuf) if err != nil { if err == io.EOF { log.Printf("[%s] [src:%s] EOF\n", queryid, backend) return } else { log.Printf("[%s] [src:%s] Error decoding result stream: %v\n", queryid, backend, err) return } } z := proto.ReadRootZ(seg) if z.Which() == proto.Z_PROGRESSUPDATE { storeProgress(queryid, backendidx, z.Progressupdate()) } else { storeResult(queryid, backendidx, z.Match()) } bstate.tempFileOffset += int64(written) } log.Printf("[%s] [src:%s] query done, disconnecting\n", queryid, backend) }