func (h Handler) Set(cmd common.SetRequest, src *bufio.Reader) error {
// TODO: should there be a unique flags value for regular data?
// Write command header
if err := binprot.WriteSetCmd(h.rw.Writer, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data))); err != nil {
return err
}
// Write value
h.rw.Write(cmd.Data)
if err := h.rw.Flush(); err != nil {
return err
}
// Read server's response
resHeader, err := readResponseHeader(h.rw.Reader)
if err != nil {
// Discard request body
if _, ioerr := src.Discard(len(cmd.Data)); ioerr != nil {
return ioerr
}
// Discard response body
if _, ioerr := h.rw.Discard(int(resHeader.TotalBodyLength)); ioerr != nil {
return ioerr
}
return err
}
return nil
}
func (h Handler) realHandleSet(cmd common.SetRequest, reqType common.RequestType) error {
// TODO: should there be a unique flags value for regular data?
// Write command header
switch reqType {
case common.RequestSet:
if err := binprot.WriteSetCmd(h.rw.Writer, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data))); err != nil {
return err
}
case common.RequestAdd:
if err := binprot.WriteAddCmd(h.rw.Writer, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data))); err != nil {
return err
}
case common.RequestReplace:
if err := binprot.WriteReplaceCmd(h.rw.Writer, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data))); err != nil {
return err
}
default:
// I know. It's all wrong. By rights we shouldn't even be here. But we are.
panic("Unrecognized request type in realHandleSet!")
}
// Write value
h.rw.Write(cmd.Data)
metrics.IncCounterBy(common.MetricBytesWrittenLocal, uint64(len(cmd.Data)))
if err := h.rw.Flush(); err != nil {
return err
}
// Read server's response
resHeader, err := readResponseHeader(h.rw.Reader)
if err != nil {
// Discard response body
n, ioerr := h.rw.Discard(int(resHeader.TotalBodyLength))
metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
if ioerr != nil {
return ioerr
}
// For Add and Replace, the error here will be common.ErrKeyExists or common.ErrKeyNotFound
// respectively. For each, this is the right response to send to the requestor. The error
// here is overloaded because it would signal a true error for sets, but a normal "error"
// response for Add and Replace.
return err
}
return nil
}
func (h Handler) realHandleSet(cmd common.SetRequest, reqType common.RequestType) error {
// Specialized chunk reader to make the code here much simpler
dataSize, fullSize := chunkSize(len(cmd.Key))
limChunkReader := newChunkLimitedReader(bytes.NewBuffer(cmd.Data), int64(dataSize), int64(len(cmd.Data)))
numChunks := int(math.Ceil(float64(len(cmd.Data)) / float64(dataSize)))
token := <-tokens
metaKey := metaKey(cmd.Key)
metaData := metadata{
Length: uint32(len(cmd.Data)),
OrigFlags: cmd.Flags,
NumChunks: uint32(numChunks),
ChunkSize: dataSize,
Token: token,
}
// Write metadata key
// TODO: should there be a unique flags value for chunked data?
switch reqType {
case common.RequestSet:
if err := binprot.WriteSetCmd(h.rw.Writer, metaKey, cmd.Flags, cmd.Exptime, metadataSize); err != nil {
return err
}
case common.RequestAdd:
if err := binprot.WriteAddCmd(h.rw.Writer, metaKey, cmd.Flags, cmd.Exptime, metadataSize); err != nil {
return err
}
case common.RequestReplace:
if err := binprot.WriteReplaceCmd(h.rw.Writer, metaKey, cmd.Flags, cmd.Exptime, metadataSize); err != nil {
return err
}
default:
// I know. It's all wrong. By rights we shouldn't even be here. But we are.
panic("Unrecognized request type in realHandleSet!")
}
// Write value
writeMetadata(h.rw, metaData)
if err := h.rw.Flush(); err != nil {
return err
}
// Read server's response
resHeader, err := readResponseHeader(h.rw.Reader)
if err != nil {
// Discard response body
n, ioerr := h.rw.Discard(int(resHeader.TotalBodyLength))
metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
if ioerr != nil {
return ioerr
}
// For Add and Replace, the error here will be common.ErrKeyExists or common.ErrKeyNotFound
// respectively. For each, this is the right response to send to the requestor. The error
// here is overloaded because it would signal a true error for sets, but a normal "error"
// response for Add and Replace.
return err
}
// Write all the data chunks
// TODO: Clean up if a data chunk write fails
// Failure can mean the write failing at the I/O level
// or at the memcached level, e.g. response == ERROR
chunkNum := 0
for limChunkReader.More() {
// Build this chunk's key
key := chunkKey(cmd.Key, chunkNum)
// Write the key
if err := binprot.WriteSetCmd(h.rw.Writer, key, cmd.Flags, cmd.Exptime, fullSize); err != nil {
return err
}
// Write token
n, err := h.rw.Write(token[:])
metrics.IncCounterBy(common.MetricBytesWrittenLocal, uint64(n))
if err != nil {
return err
}
// Write value
n2, err := io.Copy(h.rw.Writer, limChunkReader)
metrics.IncCounterBy(common.MetricBytesWrittenLocal, uint64(n2))
if err != nil {
return err
}
// There's some additional overhead here calling Flush() because it causes a write() syscall
// The set case is already a slow path and is async from the client perspective for our use
// case so this is not a problem.
if err := h.rw.Flush(); err != nil {
return err
}
// Read server's response
resHeader, err = readResponseHeader(h.rw.Reader)
if err != nil {
if err == common.ErrNoMem {
metrics.IncCounter(MetricCmdSetErrorsOOM)
}
// Reset the ReadWriter to prevent sending garbage to memcached
// otherwise we get disconnected. This is last-ditch and probably won't help. We should
// probably just disconnect and reconnect to clear OS buffers
h.reset()
// Discard repsonse body
n, ioerr := h.rw.Discard(int(resHeader.TotalBodyLength))
metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
if ioerr != nil {
return ioerr
}
return err
}
// Reset for next iteration
limChunkReader.NextChunk()
chunkNum++
}
return nil
}
func (h Handler) Set(cmd common.SetRequest) error {
if err := binprot.WriteSetCmd(h.rw.Writer, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data))); err != nil {
return err
}
return h.handleSetCommon(cmd)
}
func (h Handler) Set(cmd common.SetRequest, src *bufio.Reader) error {
// For writing chunks, the specialized chunked reader is appropriate.
// for unchunked, a limited reader will be needed since the text protocol
// includes a /r/n at the end and there's no EOF to be had with a long-lived
// connection.
limChunkReader := newChunkLimitedReader(bytes.NewBuffer(cmd.Data), int64(chunkSize), int64(len(cmd.Data)))
numChunks := int(math.Ceil(float64(len(cmd.Data)) / float64(chunkSize)))
token := <-tokens
metaKey := metaKey(cmd.Key)
metaData := metadata{
Length: uint32(len(cmd.Data)),
OrigFlags: cmd.Flags,
NumChunks: uint32(numChunks),
ChunkSize: chunkSize,
Token: token,
}
metaDataBuf := new(bytes.Buffer)
binary.Write(metaDataBuf, binary.BigEndian, metaData)
// Write metadata key
// TODO: should there be a unique flags value for chunked data?
if err := binprot.WriteSetCmd(h.rw.Writer, metaKey, cmd.Flags, cmd.Exptime, metadataSize); err != nil {
return err
}
// Write value
if _, err := io.Copy(h.rw.Writer, metaDataBuf); err != nil {
return err
}
if err := h.rw.Flush(); err != nil {
return err
}
// Read server's response
resHeader, err := readResponseHeader(h.rw.Reader)
if err != nil {
// Discard request body
if _, ioerr := src.Discard(len(cmd.Data)); ioerr != nil {
return ioerr
}
// Discard response body
if _, ioerr := h.rw.Discard(int(resHeader.TotalBodyLength)); ioerr != nil {
return ioerr
}
return err
}
// Write all the data chunks
// TODO: Clean up if a data chunk write fails
// Failure can mean the write failing at the I/O level
// or at the memcached level, e.g. response == ERROR
chunkNum := 0
for limChunkReader.More() {
// Build this chunk's key
key := chunkKey(cmd.Key, chunkNum)
// Write the key
if err := binprot.WriteSetCmd(h.rw.Writer, key, cmd.Flags, cmd.Exptime, fullDataSize); err != nil {
return err
}
// Write token
if _, err := h.rw.Write(token[:]); err != nil {
return err
}
// Write value
if _, err := io.Copy(h.rw.Writer, limChunkReader); err != nil {
return err
}
// There's some additional overhead here calling Flush() because it causes a write() syscall
// The set case is already a slow path and is async from the client perspective for our use
// case so this is not a problem.
if err := h.rw.Flush(); err != nil {
return err
}
// Read server's response
resHeader, err = readResponseHeader(h.rw.Reader)
if err != nil {
// Reset the ReadWriter to prevent sending garbage to memcached
// otherwise we get disconnected
h.reset()
// Discard request body
// This is more complicated code but more straightforward than attempting to get at
// the underlying reader and discard directly, since we don't exactly know how many
// bytes were sent already
for limChunkReader.More() {
if _, ioerr := io.Copy(ioutil.Discard, limChunkReader); ioerr != nil {
return ioerr
}
limChunkReader.NextChunk()
}
// Discard repsonse body
if _, ioerr := h.rw.Discard(int(resHeader.TotalBodyLength)); ioerr != nil {
return ioerr
}
return err
}
// Reset for next iteration
limChunkReader.NextChunk()
chunkNum++
}
return nil
}