Merge branch 'TrackDiskSize' into DepthImpl

Author: kev1n80

diskqueue.go

@@ -64,6 +64,7 @@ type diskQueue struct {
 	writeFileNum  int64
 	readMessages  int64
 	writeMessages int64
+	writeBytes    int64
 	depth         int64
 
 	sync.RWMutex
@@ -107,6 +108,9 @@ type diskQueue struct {
 
 	// disk limit implementation flag
 	diskLimitFeatIsOn bool
+
+	// the size of the
+	readMsgSize int32
 }
 
 // New instantiates an instance of diskQueue, retrieving metadata
@@ -300,6 +304,7 @@ func (d *diskQueue) skipToNextRWFile() error {
 	d.depth = 0
 	d.readMessages = 0
 	d.writeMessages = 0
+	d.writeBytes = 0
 
 	return err
 }
@@ -308,7 +313,6 @@ func (d *diskQueue) skipToNextRWFile() error {
 // while advancing read positions and rolling files, if necessary
 func (d *diskQueue) readOne() ([]byte, error) {
 	var err error
-	var msgSize int32
 
 	if d.readFile == nil {
 		curFileName := d.fileName(d.readFileNum)
@@ -345,30 +349,30 @@ func (d *diskQueue) readOne() ([]byte, error) {
 		d.reader = bufio.NewReader(d.readFile)
 	}
 
-	err = binary.Read(d.reader, binary.BigEndian, &msgSize)
+	err = binary.Read(d.reader, binary.BigEndian, &d.readMsgSize)
 	if err != nil {
 		d.readFile.Close()
 		d.readFile = nil
 		return nil, err
 	}
 
-	if msgSize < d.minMsgSize || msgSize > d.maxMsgSize {
+	if d.readMsgSize < d.minMsgSize || d.readMsgSize > d.maxMsgSize {
 		// this file is corrupt and we have no reasonable guarantee on
 		// where a new message should begin
 		d.readFile.Close()
 		d.readFile = nil
-		return nil, fmt.Errorf("invalid message read size (%d)", msgSize)
+		return nil, fmt.Errorf("invalid message read size (%d)", d.readMsgSize)
 	}
 
-	readBuf := make([]byte, msgSize)
+	readBuf := make([]byte, d.readMsgSize)
 	_, err = io.ReadFull(d.reader, readBuf)
 	if err != nil {
 		d.readFile.Close()
 		d.readFile = nil
 		return nil, err
 	}
 
-	totalBytes := int64(4 + msgSize)
+	totalBytes := int64(4 + d.readMsgSize)
 
 	// we only advance next* because we have not yet sent this to consumers
 	// (where readFileNum, readPos will actually be advanced)
@@ -461,6 +465,7 @@ func (d *diskQueue) writeOne(data []byte) error {
 	if d.diskLimitFeatIsOn {
 		// save space for the number of messages in this file
 		fileSize += 8
+		d.writeBytes += totalBytes
 		d.writeMessages += 1
 	}
 
@@ -471,7 +476,12 @@ func (d *diskQueue) writeOne(data []byte) error {
 
 		d.writeFileNum++
 		d.writePos = 0
-		d.writeMessages = 0
+
+		if d.diskLimitFeatIsOn {
+			// add bytes for the number of messages in the file
+			d.writeBytes += 8
+			d.writeMessages = 0
+		}
 
 		// sync every time we start writing to a new file
 		err = d.sync()
@@ -522,10 +532,10 @@ func (d *diskQueue) retrieveMetaData() error {
 
 	// if user is using disk space limit feature
 	if d.diskLimitFeatIsOn {
-		_, err = fmt.Fscanf(f, "%d\n%d,%d,%d\n%d,%d,%d\n",
+		_, err = fmt.Fscanf(f, "%d\n%d,%d,%d\n%d,%d,%d,%d\n",
 			&d.depth,
 			&d.readFileNum, &d.readMessages, &d.readPos,
-			&d.writeFileNum, &d.writeMessages, &d.writePos)
+			&d.writeBytes, &d.writeFileNum, &d.writeMessages, &d.writePos)
 	} else {
 		_, err = fmt.Fscanf(f, "%d\n%d,%d\n%d,%d\n",
 			&d.depth,
@@ -559,10 +569,10 @@ func (d *diskQueue) persistMetaData() error {
 
 	// if user is using disk space limit feature
 	if d.diskLimitFeatIsOn {
-		_, err = fmt.Fprintf(f, "%d\n%d,%d,%d\n%d,%d,%d\n",
+		_, err = fmt.Fprintf(f, "%d\n%d,%d,%d\n%d,%d,%d,%d\n",
 			d.depth,
 			d.readFileNum, d.readMessages, d.readPos,
-			d.writeFileNum, d.writeMessages, d.writePos)
+			d.writeBytes, d.writeFileNum, d.writeMessages, d.writePos)
 	} else {
 		_, err = fmt.Fprintf(f, "%d\n%d,%d\n%d,%d\n",
 			d.depth,
@@ -629,6 +639,8 @@ func (d *diskQueue) checkTailCorruption(depth int64) {
 }
 
 func (d *diskQueue) moveForward() {
+	// add bytes for the number of messages and the size of the message
+	readFileLen := int64(d.readMsgSize) + d.readPos + 12
 	oldReadFileNum := d.readFileNum
 	d.readFileNum = d.nextReadFileNum
 	d.readPos = d.nextReadPos
@@ -647,6 +659,8 @@ func (d *diskQueue) moveForward() {
 		if err != nil {
 			d.logf(ERROR, "DISKQUEUE(%s) failed to Remove(%s) - %s", d.name, fn, err)
 		}
+
+		d.writeBytes -= readFileLen
 	}
 
 	d.checkTailCorruption(d.depth)

diskqueue_test.go

@@ -251,6 +251,7 @@ func TestDiskQueueCorruption(t *testing.T) {
 
 type md struct {
 	depth         int64
+	writeBytes    int64
 	readFileNum   int64
 	writeFileNum  int64
 	readMessages  int64
@@ -274,11 +275,12 @@ func readMetaDataFile(fileName string, retried int, diskLimitFeatIsOn bool) md {
 	defer f.Close()
 
 	var ret md
+
 	if diskLimitFeatIsOn {
-		_, err = fmt.Fscanf(f, "%d\n%d,%d,%d\n%d,%d,%d\n",
+		_, err = fmt.Fscanf(f, "%d\n%d,%d,%d\n%d,%d,%d,%d\n",
 			&ret.depth,
 			&ret.readFileNum, &ret.readMessages, &ret.readPos,
-			&ret.writeFileNum, &ret.writeMessages, &ret.writePos)
+			&ret.writeBytes, &ret.writeFileNum, &ret.writeMessages, &ret.writePos)
 	} else {
 		_, err = fmt.Fscanf(f, "%d\n%d,%d\n%d,%d\n",
 			&ret.depth,
@@ -358,6 +360,7 @@ func TestDiskQueueSyncAfterReadWithDiskSizeImplementation(t *testing.T) {
 	for i := 0; i < 10; i++ {
 		d := readMetaDataFile(dq.(*diskQueue).metaDataFileName(), 0, true)
 		if d.depth == 1 &&
+			d.writeBytes == 1004 &&
 			d.readFileNum == 0 &&
 			d.writeFileNum == 0 &&
 			d.readPos == 0 &&
@@ -378,6 +381,7 @@ next:
 	for i := 0; i < 10; i++ {
 		d := readMetaDataFile(dq.(*diskQueue).metaDataFileName(), 0, true)
 		if d.depth == 1 &&
+			d.writeBytes == 2008 &&
 			d.readFileNum == 0 &&
 			d.writeFileNum == 0 &&
 			d.readPos == 1004 &&
@@ -405,6 +409,7 @@ completeWriteFile:
 		// test the writeFileNum correctly increments
 		d := readMetaDataFile(dq.(*diskQueue).metaDataFileName(), 0, true)
 		if d.depth == 3 &&
+			d.writeBytes == 3020 &&
 			d.readFileNum == 0 &&
 			d.writeFileNum == 1 &&
 			d.readPos == 1004 &&
@@ -429,7 +434,9 @@ completeReadFile:
 		// test that read position and messages reset when a file is completely read
 		// test the readFileNum correctly increments
 		d := readMetaDataFile(dq.(*diskQueue).metaDataFileName(), 0, true)
+		t.Logf("Write bytes: %d", d.writeBytes)
 		if d.depth == 1 &&
+			d.writeBytes == 1004 &&
 			d.readFileNum == 1 &&
 			d.writeFileNum == 1 &&
 			d.readPos == 0 &&