Fixed issue with peekChan

Author: Mikhail Faraponov

diskqueue.go

@@ -152,6 +152,7 @@ func NewWithDiskSpace(name string, dataPath string,
 		minMsgSize:           minMsgSize,
 		maxMsgSize:           maxMsgSize,
 		readChan:             make(chan []byte),
+		peekChan:             make(chan []byte),
 		depthChan:            make(chan int64),
 		writeChan:            make(chan []byte),
 		writeResponseChan:    make(chan error),
@@ -648,23 +649,6 @@ func (d *diskQueue) writeOne(data []byte) error {
 		return fmt.Errorf("invalid message write size (%d) minMsgSize=%d maxMsgSize=%d", dataLen, d.minMsgSize, d.maxMsgSize)
 	}
 
-	if d.enableDiskLimitation {
-		expectedBytesIncrease := totalBytes
-		// check if we will reach or surpass file size limit
-		if d.writePos+totalBytes+numFileMsgBytes >= d.maxBytesPerFile {
-			reachedFileSizeLimit = true
-			expectedBytesIncrease += numFileMsgBytes
-		}
-
-		// free disk space if needed
-		err = d.checkDiskSpace(expectedBytesIncrease)
-		if err != nil {
-			return err
-		}
-	} else if d.writePos+totalBytes >= d.maxBytesPerFile {
-		reachedFileSizeLimit = true
-	}
-
 	// will not wrap-around if maxBytesPerFile + maxMsgSize < Int64Max
 	if d.writePos > 0 && d.writePos+totalBytes > d.maxBytesPerFile {
 		if d.readFileNum == d.writeFileNum {
@@ -685,6 +669,7 @@ func (d *diskQueue) writeOne(data []byte) error {
 			d.writeFile = nil
 		}
 	}
+
 	if d.writeFile == nil {
 		curFileName := d.fileName(d.writeFileNum)
 		d.writeFile, err = os.OpenFile(curFileName, os.O_RDWR|os.O_CREATE, 0600)
@@ -704,6 +689,23 @@ func (d *diskQueue) writeOne(data []byte) error {
 		}
 	}
 
+	if d.enableDiskLimitation {
+		expectedBytesIncrease := totalBytes
+		// check if we will reach or surpass file size limit
+		if d.writePos+totalBytes+numFileMsgBytes >= d.maxBytesPerFile {
+			reachedFileSizeLimit = true
+			expectedBytesIncrease += numFileMsgBytes
+		}
+
+		// free disk space if needed
+		err = d.checkDiskSpace(expectedBytesIncrease)
+		if err != nil {
+			return err
+		}
+	} else if d.writePos+totalBytes >= d.maxBytesPerFile {
+		reachedFileSizeLimit = true
+	}
+
 	d.writeBuf.Reset()
 	err = binary.Write(&d.writeBuf, binary.BigEndian, dataLen)
 	if err != nil {

diskqueue_test.go

@@ -149,6 +149,7 @@ func TestDiskQueuePeek(t *testing.T) {
 	t.Run("roll", func(t *testing.T) {
 		for i := 0; i < 10; i++ {
 			err := dq.Put(msg)
+
 			Nil(t, err)
 			Equal(t, int64(i+1), dq.Depth())
 		}
@@ -327,7 +328,7 @@ func TestDiskQueueCorruption(t *testing.T) {
 	Equal(t, msg, <-dq.ReadChan())
 	badFilesCount = numberOfBadFiles(dqName, tmpDir)
 	if badFilesCount != 2 {
-		panic("fail")
+		panic(badFilesCount)
 	}
 
 	// write a corrupt (len 0) message at the 5th (current) file
@@ -339,12 +340,12 @@ func TestDiskQueueCorruption(t *testing.T) {
 
 	Equal(t, msg, <-dq.ReadChan())
 
-  // conflict was here
+	// conflict was here
 	badFilesCount = numberOfBadFiles(dqName, tmpDir)
 	if badFilesCount != 3 {
 		panic("fail")
 	}
-  //
+	//
 
 	dq.Put(msg)
 	dq.Put(msg)