예제 #1
0
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)
}
예제 #2
0
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)
}