func (foRunner *foRunner) Encode(pack *PipelinePack) (output []byte, err error) { var encoded []byte if encoded, err = foRunner.encoder.Encode(pack); err != nil || encoded == nil { return } if foRunner.useFraming { client.CreateHekaStream(encoded, &output, nil) } else { output = encoded } return }
func (b *BufferedOutput) QueueRecord(pack *PipelinePack) (err error) { var msgBytes []byte if msgBytes, err = b.or.Encode(pack); err != nil { return } // If framing isn't already in place then we need to add it. if b.or.UsesFraming() { b.outBytes = msgBytes } else { if err = client.CreateHekaStream(msgBytes, &b.outBytes, nil); err != nil { return } } if _, err = b.writeFile.Write(b.outBytes); err != nil { return fmt.Errorf("writing to %s: %s", getQueueFilename(b.queue, b.writeId), err) } return nil }
// QueueRecord adds a new record to the end of the current queue buffer. Note // that QueueRecord is *not* thread safe, it should only ever be called by one // goroutine at a time. func (bf *BufferFeeder) QueueRecord(pack *PipelinePack) error { maxQueueSize := bf.Config.MaxBufferSize if maxQueueSize > 0 && (bf.queueSize.Get()+uint64(len(pack.MsgBytes)) > maxQueueSize) { return QueueIsFull } maxQueueFileSize := bf.Config.MaxFileSize if bf.writeFileSize+uint64(len(pack.MsgBytes)) > maxQueueFileSize { if err := bf.RollQueue(); err != nil { return fmt.Errorf("queue file rotation error: %s", err) } } var outBytes []byte err := client.CreateHekaStream(pack.MsgBytes, &outBytes, nil) if err != nil { return fmt.Errorf("message framing error: %s", err) } n, err := bf.writeFile.Write(outBytes) if err != nil { if n > 0 { // If we wrote some data but there was an error, that data is // suspect so we need to seek back to before the bogus write // happened. ret, e := bf.writeFile.Seek(int64(-n), 2) if e != nil { return QueueCorrupt } if e = bf.writeFile.Truncate(ret); e != nil { return QueueCorrupt } } return fmt.Errorf("can't write to queue: %s", err) } bf.queueSize.Add(uint64(n)) bf.writeFileSize += uint64(n) return nil }
func BufferedOutputSpec(c gs.Context) { tmpDir, tmpErr := ioutil.TempDir("", "bufferedout-tests") defer func() { tmpErr = os.RemoveAll(tmpDir) c.Expect(tmpErr, gs.Equals, nil) }() t := &ts.SimpleT{} ctrl := gomock.NewController(t) defer ctrl.Finish() c.Specify("BufferedOutput Internals", func() { encoder := new(ProtobufEncoder) encoder.sample = false encoder.sampleDenominator = 1000 or := NewMockOutputRunner(ctrl) bufferedOutput, err := NewBufferedOutput(tmpDir, "test", or) c.Expect(err, gs.IsNil) msg := ts.GetTestMessage() c.Specify("fileExists", func() { c.Expect(fileExists(tmpDir), gs.IsTrue) c.Expect(fileExists(filepath.Join(tmpDir, "test.log")), gs.IsFalse) }) c.Specify("extractBufferId", func() { id, err := extractBufferId("555.log") c.Expect(err, gs.IsNil) c.Expect(id, gs.Equals, uint(555)) id, err = extractBufferId("") c.Expect(err, gs.Not(gs.IsNil)) id, err = extractBufferId("a.log") c.Expect(err, gs.Not(gs.IsNil)) }) c.Specify("findBufferId", func() { c.Expect(findBufferId(tmpDir, true), gs.Equals, uint(0)) c.Expect(findBufferId(tmpDir, false), gs.Equals, uint(0)) fd, err := os.Create(filepath.Join(tmpDir, "4.log")) c.Expect(err, gs.IsNil) fd.Close() fd, err = os.Create(filepath.Join(tmpDir, "5.log")) c.Expect(err, gs.IsNil) fd.Close() fd, err = os.Create(filepath.Join(tmpDir, "6a.log")) c.Expect(err, gs.IsNil) fd.Close() c.Expect(findBufferId(tmpDir, false), gs.Equals, uint(4)) c.Expect(findBufferId(tmpDir, true), gs.Equals, uint(5)) }) c.Specify("writeCheckpoint", func() { bufferedOutput.checkpointFilename = filepath.Join(tmpDir, "cp.txt") err := bufferedOutput.writeCheckpoint(43, 99999) c.Expect(err, gs.IsNil) c.Expect(fileExists(bufferedOutput.checkpointFilename), gs.IsTrue) id, offset, err := readCheckpoint(bufferedOutput.checkpointFilename) c.Expect(err, gs.IsNil) c.Expect(id, gs.Equals, uint(43)) c.Expect(offset, gs.Equals, int64(99999)) err = bufferedOutput.writeCheckpoint(43, 1) c.Expect(err, gs.IsNil) id, offset, err = readCheckpoint(bufferedOutput.checkpointFilename) c.Expect(err, gs.IsNil) c.Expect(id, gs.Equals, uint(43)) c.Expect(offset, gs.Equals, int64(1)) bufferedOutput.checkpointFile.Close() }) c.Specify("readCheckpoint", func() { cp := filepath.Join(tmpDir, "cp.txt") file, err := os.OpenFile(cp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644) c.Expect(err, gs.IsNil) id, offset, err := readCheckpoint(cp) c.Expect(err, gs.Not(gs.IsNil)) file.WriteString("22") id, offset, err = readCheckpoint(cp) c.Expect(err.Error(), gs.Equals, "invalid checkpoint format") file.Seek(0, 0) file.WriteString("aa 22") id, offset, err = readCheckpoint(cp) c.Expect(err.Error(), gs.Equals, "invalid checkpoint id") file.Seek(0, 0) file.WriteString("43 aa") id, offset, err = readCheckpoint(cp) c.Expect(err.Error(), gs.Equals, "invalid checkpoint offset") file.Seek(0, 0) file.WriteString("43 22") file.Close() id, offset, err = readCheckpoint(cp) c.Expect(err, gs.IsNil) c.Expect(id, gs.Equals, uint(43)) c.Expect(offset, gs.Equals, int64(22)) }) c.Specify("RollQueue", func() { bufferedOutput.checkpointFilename = filepath.Join(tmpDir, "cp.txt") bufferedOutput.queue = tmpDir err := bufferedOutput.RollQueue() c.Expect(err, gs.IsNil) c.Expect(fileExists(getQueueFilename(bufferedOutput.queue, bufferedOutput.writeId)), gs.IsTrue) bufferedOutput.writeFile.WriteString("this is a test item") bufferedOutput.writeFile.Close() bufferedOutput.writeCheckpoint(bufferedOutput.writeId, 10) bufferedOutput.checkpointFile.Close() err = bufferedOutput.readFromNextFile() buf := make([]byte, 4) n, err := bufferedOutput.readFile.Read(buf) c.Expect(n, gs.Equals, 4) c.Expect("test", gs.Equals, string(buf)) bufferedOutput.writeFile.Close() bufferedOutput.readFile.Close() }) c.Specify("QueueRecord", func() { bufferedOutput.checkpointFilename = filepath.Join(tmpDir, "cp.txt") bufferedOutput.queue = tmpDir newpack := NewPipelinePack(nil) newpack.Message = msg newpack.Decoded = true payload := "Write me out to the network" newpack.Message.SetPayload(payload) protoBytes, err := encoder.Encode(newpack) expectedLen := 115 c.Specify("adds framing when necessary", func() { or.EXPECT().Encode(newpack).Return(protoBytes, err) or.EXPECT().UsesFraming().Return(false) err = bufferedOutput.RollQueue() c.Expect(err, gs.IsNil) err = bufferedOutput.QueueRecord(newpack) fName := getQueueFilename(bufferedOutput.queue, bufferedOutput.writeId) c.Expect(fileExists(fName), gs.IsTrue) c.Expect(err, gs.IsNil) bufferedOutput.writeFile.Close() f, err := os.Open(fName) c.Expect(err, gs.IsNil) n, record, err := bufferedOutput.parser.Parse(f) f.Close() c.Expect(n, gs.Equals, expectedLen) c.Expect(err, gs.IsNil) headerLen := int(record[1]) + message.HEADER_FRAMING_SIZE record = record[headerLen:] outMsg := new(message.Message) proto.Unmarshal(record, outMsg) c.Expect(outMsg.GetPayload(), gs.Equals, payload) }) c.Specify("doesn't add framing if it's already there", func() { var framed []byte client.CreateHekaStream(protoBytes, &framed, nil) or.EXPECT().Encode(newpack).Return(framed, err) or.EXPECT().UsesFraming().Return(true) err = bufferedOutput.RollQueue() c.Expect(err, gs.IsNil) err = bufferedOutput.QueueRecord(newpack) fName := getQueueFilename(bufferedOutput.queue, bufferedOutput.writeId) c.Expect(fileExists(fName), gs.IsTrue) c.Expect(err, gs.IsNil) bufferedOutput.writeFile.Close() f, err := os.Open(fName) c.Expect(err, gs.IsNil) n, record, err := bufferedOutput.parser.Parse(f) f.Close() c.Expect(n, gs.Equals, expectedLen) c.Expect(err, gs.IsNil) headerLen := int(record[1]) + message.HEADER_FRAMING_SIZE record = record[headerLen:] outMsg := new(message.Message) proto.Unmarshal(record, outMsg) c.Expect(outMsg.GetPayload(), gs.Equals, payload) }) }) }) }