c9b8bd8de2
Queue output items and reuse them. Remove the unneeded type system in sql and just use the Go type system. In best case this is more than an order of magnitude speedup: ``` BenchmarkSelectAll_1M-12 1 1841049400 ns/op 274299728 B/op 4198522 allocs/op BenchmarkSelectAll_1M-12 14 84833400 ns/op 169228346 B/op 3146541 allocs/op ```
427 lines
15 KiB
Go
427 lines
15 KiB
Go
/*
|
|
* MinIO Cloud Storage, (C) 2019 MinIO, Inc.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
package s3select
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"hash/crc32"
|
|
"net/http"
|
|
"strconv"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// A message is in the format specified in
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/images/s3select-frame-diagram-frame-overview.png
|
|
// hence the calculation is made accordingly.
|
|
func totalByteLength(headerLength, payloadLength int) int {
|
|
return 4 + 4 + 4 + headerLength + payloadLength + 4
|
|
}
|
|
|
|
func genMessage(header, payload []byte) []byte {
|
|
headerLength := len(header)
|
|
payloadLength := len(payload)
|
|
totalLength := totalByteLength(headerLength, payloadLength)
|
|
|
|
buf := new(bytes.Buffer)
|
|
binary.Write(buf, binary.BigEndian, uint32(totalLength))
|
|
binary.Write(buf, binary.BigEndian, uint32(headerLength))
|
|
prelude := buf.Bytes()
|
|
binary.Write(buf, binary.BigEndian, crc32.ChecksumIEEE(prelude))
|
|
buf.Write(header)
|
|
if payload != nil {
|
|
buf.Write(payload)
|
|
}
|
|
message := buf.Bytes()
|
|
binary.Write(buf, binary.BigEndian, crc32.ChecksumIEEE(message))
|
|
|
|
return buf.Bytes()
|
|
}
|
|
|
|
// Refer genRecordsHeader().
|
|
var recordsHeader = []byte{
|
|
13, ':', 'm', 'e', 's', 's', 'a', 'g', 'e', '-', 't', 'y', 'p', 'e', 7, 0, 5, 'e', 'v', 'e', 'n', 't',
|
|
13, ':', 'c', 'o', 'n', 't', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 24, 'a', 'p', 'p', 'l', 'i', 'c', 'a', 't', 'i', 'o', 'n', '/', 'o', 'c', 't', 'e', 't', '-', 's', 't', 'r', 'e', 'a', 'm',
|
|
11, ':', 'e', 'v', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 7, 'R', 'e', 'c', 'o', 'r', 'd', 's',
|
|
}
|
|
|
|
const (
|
|
maxRecordMessageLength = 128 * 1024 // Chosen for compatibility with AWS JAVA SDK
|
|
)
|
|
|
|
var (
|
|
bufLength = payloadLenForMsgLen(maxRecordMessageLength)
|
|
)
|
|
|
|
// newRecordsMessage - creates new Records Message which can contain a single record, partial records,
|
|
// or multiple records. Depending on the size of the result, a response can contain one or more of these messages.
|
|
//
|
|
// Header specification
|
|
// Records messages contain three headers, as follows:
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/images/s3select-frame-diagram-record.png
|
|
//
|
|
// Payload specification
|
|
// Records message payloads can contain a single record, partial records, or multiple records.
|
|
func newRecordsMessage(payload []byte) []byte {
|
|
return genMessage(recordsHeader, payload)
|
|
}
|
|
|
|
// payloadLenForMsgLen computes the length of the payload in a record
|
|
// message given the length of the message.
|
|
func payloadLenForMsgLen(messageLength int) int {
|
|
headerLength := len(recordsHeader)
|
|
payloadLength := messageLength - 4 - 4 - 4 - headerLength - 4
|
|
return payloadLength
|
|
}
|
|
|
|
// continuationMessage - S3 periodically sends this message to keep the TCP connection open.
|
|
// These messages appear in responses at random. The client must detect the message type and process accordingly.
|
|
//
|
|
// Header specification:
|
|
// Continuation messages contain two headers, as follows:
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/images/s3select-frame-diagram-cont.png
|
|
//
|
|
// Payload specification:
|
|
// Continuation messages have no payload.
|
|
var continuationMessage = []byte{
|
|
0, 0, 0, 57, // total byte-length.
|
|
0, 0, 0, 41, // headers byte-length.
|
|
139, 161, 157, 242, // prelude crc.
|
|
13, ':', 'm', 'e', 's', 's', 'a', 'g', 'e', '-', 't', 'y', 'p', 'e', 7, 0, 5, 'e', 'v', 'e', 'n', 't', // headers.
|
|
11, ':', 'e', 'v', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 4, 'C', 'o', 'n', 't', // headers.
|
|
156, 134, 74, 13, // message crc.
|
|
}
|
|
|
|
// Refer genProgressHeader().
|
|
var progressHeader = []byte{
|
|
13, ':', 'm', 'e', 's', 's', 'a', 'g', 'e', '-', 't', 'y', 'p', 'e', 7, 0, 5, 'e', 'v', 'e', 'n', 't',
|
|
13, ':', 'c', 'o', 'n', 't', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 8, 't', 'e', 'x', 't', '/', 'x', 'm', 'l',
|
|
11, ':', 'e', 'v', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 8, 'P', 'r', 'o', 'g', 'r', 'e', 's', 's',
|
|
}
|
|
|
|
// newProgressMessage - creates new Progress Message. S3 periodically sends this message, if requested.
|
|
// It contains information about the progress of a query that has started but has not yet completed.
|
|
//
|
|
// Header specification:
|
|
// Progress messages contain three headers, as follows:
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/images/s3select-frame-diagram-progress.png
|
|
//
|
|
// Payload specification:
|
|
// Progress message payload is an XML document containing information about the progress of a request.
|
|
// * BytesScanned => Number of bytes that have been processed before being uncompressed (if the file is compressed).
|
|
// * BytesProcessed => Number of bytes that have been processed after being uncompressed (if the file is compressed).
|
|
// * BytesReturned => Current number of bytes of records payload data returned by S3.
|
|
//
|
|
// For uncompressed files, BytesScanned and BytesProcessed are equal.
|
|
//
|
|
// Example:
|
|
//
|
|
// <?xml version="1.0" encoding="UTF-8"?>
|
|
// <Progress>
|
|
// <BytesScanned>512</BytesScanned>
|
|
// <BytesProcessed>1024</BytesProcessed>
|
|
// <BytesReturned>1024</BytesReturned>
|
|
// </Progress>
|
|
//
|
|
func newProgressMessage(bytesScanned, bytesProcessed, bytesReturned int64) []byte {
|
|
payload := []byte(`<?xml version="1.0" encoding="UTF-8"?><Progress><BytesScanned>` +
|
|
strconv.FormatInt(bytesScanned, 10) + `</BytesScanned><BytesProcessed>` +
|
|
strconv.FormatInt(bytesProcessed, 10) + `</BytesProcessed><BytesReturned>` +
|
|
strconv.FormatInt(bytesReturned, 10) + `</BytesReturned></Stats>`)
|
|
return genMessage(progressHeader, payload)
|
|
}
|
|
|
|
// Refer genStatsHeader().
|
|
var statsHeader = []byte{
|
|
13, ':', 'm', 'e', 's', 's', 'a', 'g', 'e', '-', 't', 'y', 'p', 'e', 7, 0, 5, 'e', 'v', 'e', 'n', 't',
|
|
13, ':', 'c', 'o', 'n', 't', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 8, 't', 'e', 'x', 't', '/', 'x', 'm', 'l',
|
|
11, ':', 'e', 'v', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 5, 'S', 't', 'a', 't', 's',
|
|
}
|
|
|
|
// newStatsMessage - creates new Stats Message. S3 sends this message at the end of the request.
|
|
// It contains statistics about the query.
|
|
//
|
|
// Header specification:
|
|
// Stats messages contain three headers, as follows:
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/images/s3select-frame-diagram-stats.png
|
|
//
|
|
// Payload specification:
|
|
// Stats message payload is an XML document containing information about a request's stats when processing is complete.
|
|
// * BytesScanned => Number of bytes that have been processed before being uncompressed (if the file is compressed).
|
|
// * BytesProcessed => Number of bytes that have been processed after being uncompressed (if the file is compressed).
|
|
// * BytesReturned => Total number of bytes of records payload data returned by S3.
|
|
//
|
|
// For uncompressed files, BytesScanned and BytesProcessed are equal.
|
|
//
|
|
// Example:
|
|
//
|
|
// <?xml version="1.0" encoding="UTF-8"?>
|
|
// <Stats>
|
|
// <BytesScanned>512</BytesScanned>
|
|
// <BytesProcessed>1024</BytesProcessed>
|
|
// <BytesReturned>1024</BytesReturned>
|
|
// </Stats>
|
|
func newStatsMessage(bytesScanned, bytesProcessed, bytesReturned int64) []byte {
|
|
payload := []byte(`<?xml version="1.0" encoding="UTF-8"?><Stats><BytesScanned>` +
|
|
strconv.FormatInt(bytesScanned, 10) + `</BytesScanned><BytesProcessed>` +
|
|
strconv.FormatInt(bytesProcessed, 10) + `</BytesProcessed><BytesReturned>` +
|
|
strconv.FormatInt(bytesReturned, 10) + `</BytesReturned></Stats>`)
|
|
return genMessage(statsHeader, payload)
|
|
}
|
|
|
|
// endMessage - indicates that the request is complete, and no more messages will be sent.
|
|
// You should not assume that the request is complete until the client receives an End message.
|
|
//
|
|
// Header specification:
|
|
// End messages contain two headers, as follows:
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/images/s3select-frame-diagram-end.png
|
|
//
|
|
// Payload specification:
|
|
// End messages have no payload.
|
|
var endMessage = []byte{
|
|
0, 0, 0, 56, // total byte-length.
|
|
0, 0, 0, 40, // headers byte-length.
|
|
193, 198, 132, 212, // prelude crc.
|
|
13, ':', 'm', 'e', 's', 's', 'a', 'g', 'e', '-', 't', 'y', 'p', 'e', 7, 0, 5, 'e', 'v', 'e', 'n', 't', // headers.
|
|
11, ':', 'e', 'v', 'e', 'n', 't', '-', 't', 'y', 'p', 'e', 7, 0, 3, 'E', 'n', 'd', // headers.
|
|
207, 151, 211, 146, // message crc.
|
|
}
|
|
|
|
// newErrorMessage - creates new Request Level Error Message. S3 sends this message if the request failed for any reason.
|
|
// It contains the error code and error message for the failure. If S3 sends a RequestLevelError message,
|
|
// it doesn't send an End message.
|
|
//
|
|
// Header specification:
|
|
// Request-level error messages contain three headers, as follows:
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/images/s3select-frame-diagram-error.png
|
|
//
|
|
// Payload specification:
|
|
// Request-level error messages have no payload.
|
|
func newErrorMessage(errorCode, errorMessage []byte) []byte {
|
|
buf := new(bytes.Buffer)
|
|
|
|
buf.Write([]byte{13, ':', 'm', 'e', 's', 's', 'a', 'g', 'e', '-', 't', 'y', 'p', 'e', 7, 0, 5, 'e', 'r', 'r', 'o', 'r'})
|
|
|
|
buf.Write([]byte{14, ':', 'e', 'r', 'r', 'o', 'r', '-', 'm', 'e', 's', 's', 'a', 'g', 'e', 7})
|
|
binary.Write(buf, binary.BigEndian, uint16(len(errorMessage)))
|
|
buf.Write(errorMessage)
|
|
|
|
buf.Write([]byte{11, ':', 'e', 'r', 'r', 'o', 'r', '-', 'c', 'o', 'd', 'e', 7})
|
|
binary.Write(buf, binary.BigEndian, uint16(len(errorCode)))
|
|
buf.Write(errorCode)
|
|
|
|
return genMessage(buf.Bytes(), nil)
|
|
}
|
|
|
|
// NewErrorMessage - creates new Request Level Error Message specified in
|
|
// https://docs.aws.amazon.com/AmazonS3/latest/API/RESTObjectSELECTContent.html.
|
|
func NewErrorMessage(errorCode, errorMessage string) []byte {
|
|
return newErrorMessage([]byte(errorCode), []byte(errorMessage))
|
|
}
|
|
|
|
type messageWriter struct {
|
|
writer http.ResponseWriter
|
|
getProgressFunc func() (int64, int64)
|
|
bytesReturned int64
|
|
|
|
payloadBuffer []byte
|
|
payloadBufferIndex int
|
|
payloadCh chan *bytes.Buffer
|
|
|
|
finBytesScanned, finBytesProcessed int64
|
|
|
|
errCh chan []byte
|
|
doneCh chan struct{}
|
|
}
|
|
|
|
func (writer *messageWriter) write(data []byte) bool {
|
|
if _, err := writer.writer.Write(data); err != nil {
|
|
return false
|
|
}
|
|
|
|
writer.writer.(http.Flusher).Flush()
|
|
return true
|
|
}
|
|
|
|
func (writer *messageWriter) start() {
|
|
keepAliveTicker := time.NewTicker(1 * time.Second)
|
|
var progressTicker *time.Ticker
|
|
var progressTickerC <-chan time.Time
|
|
if writer.getProgressFunc != nil {
|
|
progressTicker = time.NewTicker(1 * time.Minute)
|
|
progressTickerC = progressTicker.C
|
|
}
|
|
recordStagingTicker := time.NewTicker(500 * time.Millisecond)
|
|
|
|
// Exit conditions:
|
|
//
|
|
// 1. If a writer.write() returns false, select loop below exits and
|
|
// closes `doneCh` to indicate to caller to also exit.
|
|
//
|
|
// 2. If caller (Evaluate()) has an error, it sends an error
|
|
// message and waits for this go-routine to quit in
|
|
// FinishWithError()
|
|
//
|
|
// 3. If caller is done, it waits for this go-routine to exit
|
|
// in Finish()
|
|
|
|
quitFlag := false
|
|
for !quitFlag {
|
|
select {
|
|
case data := <-writer.errCh:
|
|
quitFlag = true
|
|
// Flush collected records before sending error message
|
|
if !writer.flushRecords() {
|
|
break
|
|
}
|
|
writer.write(data)
|
|
|
|
case payload, ok := <-writer.payloadCh:
|
|
if !ok {
|
|
// payloadCh is closed by caller to
|
|
// indicate finish with success
|
|
quitFlag = true
|
|
|
|
if !writer.flushRecords() {
|
|
break
|
|
}
|
|
// Write Stats message, then End message
|
|
bytesReturned := atomic.LoadInt64(&writer.bytesReturned)
|
|
if !writer.write(newStatsMessage(writer.finBytesScanned, writer.finBytesProcessed, bytesReturned)) {
|
|
break
|
|
}
|
|
writer.write(endMessage)
|
|
} else {
|
|
for payload.Len() > 0 {
|
|
copiedLen := copy(writer.payloadBuffer[writer.payloadBufferIndex:], payload.Bytes())
|
|
writer.payloadBufferIndex += copiedLen
|
|
payload.Next(copiedLen)
|
|
|
|
// If buffer is filled, flush it now!
|
|
freeSpace := bufLength - writer.payloadBufferIndex
|
|
if freeSpace == 0 {
|
|
if !writer.flushRecords() {
|
|
quitFlag = true
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
bufPool.Put(payload)
|
|
}
|
|
|
|
case <-recordStagingTicker.C:
|
|
if !writer.flushRecords() {
|
|
quitFlag = true
|
|
}
|
|
|
|
case <-keepAliveTicker.C:
|
|
if !writer.write(continuationMessage) {
|
|
quitFlag = true
|
|
}
|
|
|
|
case <-progressTickerC:
|
|
bytesScanned, bytesProcessed := writer.getProgressFunc()
|
|
bytesReturned := atomic.LoadInt64(&writer.bytesReturned)
|
|
if !writer.write(newProgressMessage(bytesScanned, bytesProcessed, bytesReturned)) {
|
|
quitFlag = true
|
|
}
|
|
}
|
|
}
|
|
close(writer.doneCh)
|
|
|
|
recordStagingTicker.Stop()
|
|
keepAliveTicker.Stop()
|
|
if progressTicker != nil {
|
|
progressTicker.Stop()
|
|
}
|
|
|
|
// Whatever drain the payloadCh to prevent from memory leaking.
|
|
for len(writer.payloadCh) > 0 {
|
|
payload := <-writer.payloadCh
|
|
bufPool.Put(payload)
|
|
}
|
|
}
|
|
|
|
// Sends a single whole record.
|
|
func (writer *messageWriter) SendRecord(payload *bytes.Buffer) error {
|
|
select {
|
|
case writer.payloadCh <- payload:
|
|
return nil
|
|
case <-writer.doneCh:
|
|
return fmt.Errorf("messageWriter is done")
|
|
}
|
|
}
|
|
|
|
func (writer *messageWriter) flushRecords() bool {
|
|
if writer.payloadBufferIndex == 0 {
|
|
return true
|
|
}
|
|
result := writer.write(newRecordsMessage(writer.payloadBuffer[0:writer.payloadBufferIndex]))
|
|
if result {
|
|
atomic.AddInt64(&writer.bytesReturned, int64(writer.payloadBufferIndex))
|
|
writer.payloadBufferIndex = 0
|
|
}
|
|
return result
|
|
}
|
|
|
|
// Finish is the last call to the message writer - it sends any
|
|
// remaining record payload, then sends statistics and finally the end
|
|
// message.
|
|
func (writer *messageWriter) Finish(bytesScanned, bytesProcessed int64) error {
|
|
select {
|
|
case <-writer.doneCh:
|
|
return fmt.Errorf("messageWriter is done")
|
|
default:
|
|
writer.finBytesScanned = bytesScanned
|
|
writer.finBytesProcessed = bytesProcessed
|
|
close(writer.payloadCh)
|
|
// Wait until the `start` go-routine is done.
|
|
<-writer.doneCh
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (writer *messageWriter) FinishWithError(errorCode, errorMessage string) error {
|
|
select {
|
|
case <-writer.doneCh:
|
|
return fmt.Errorf("messageWriter is done")
|
|
case writer.errCh <- newErrorMessage([]byte(errorCode), []byte(errorMessage)):
|
|
// Wait until the `start` go-routine is done.
|
|
<-writer.doneCh
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// newMessageWriter creates a message writer that writes to the HTTP
|
|
// response writer
|
|
func newMessageWriter(w http.ResponseWriter, getProgressFunc func() (bytesScanned, bytesProcessed int64)) *messageWriter {
|
|
writer := &messageWriter{
|
|
writer: w,
|
|
getProgressFunc: getProgressFunc,
|
|
|
|
payloadBuffer: make([]byte, bufLength),
|
|
payloadCh: make(chan *bytes.Buffer, 1),
|
|
|
|
errCh: make(chan []byte),
|
|
doneCh: make(chan struct{}),
|
|
}
|
|
go writer.start()
|
|
return writer
|
|
}
|