merge_write.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"time"

	log "github.com/qiniu/log"
)

type Bool struct {
	c Int64
}

func (b *Bool) Get() bool {
	return b.c.Get() != 0
}

func (b *Bool) toInt64(v bool) int64 {
	if v {
		return 1
	} else {
		return 0
	}
}

func (b *Bool) Set(v bool) {
	b.c.Set(b.toInt64(v))
}

func (b *Bool) CompareAndSwap(o, n bool) bool {
	return b.c.CompareAndSwap(b.toInt64(o), b.toInt64(n))
}

func (b *Bool) Swap(v bool) bool {
	return b.c.Swap(b.toInt64(v)) != 0
}

var bufferPool *BufferPool

func init() {
	// 4000行日志缓存
	bufferPool = NewBufferPool(4000)
}

type MergeWriter struct {
	lines  chan *bytes.Buffer
	writer io.Writer
	closed Bool
}

func NewMergeWriter(writer io.Writer) *MergeWriter {
	merger := &MergeWriter{
		lines:  make(chan *bytes.Buffer, 1000),
		writer: writer,
	}
	merger.closed.Set(false)
	merger.drainLines()
	return merger
}

func (m *MergeWriter) Close() {
	// log.Printf("Close MergeWriter")
	if m.closed.CompareAndSwap(false, true) {
		// log.Printf("Close lines chan")
		close(m.lines)
	}
}

func (m *MergeWriter) WriteStrLine(line string) {
	if m.closed.Get() {
		return
	} else {
		buffer := bufferPool.Get()
		buffer.WriteString(line)
		m.lines <- buffer
	}
}

func (m *MergeWriter) WriteLine(line *bytes.Buffer) {
	if m.closed.Get() {
		// 需要回收Buffer
		// log.Printf("Write to closed MergeWrite...")
		bufferPool.Put(line)
		return
	} else {
		m.lines <- line
	}
}

func (m *MergeWriter) drainLines() {
	go func() {
		for line := range m.lines {
			m.writer.Write(line.Bytes())
			// 回收
			bufferPool.Put(line)
		}
	}()
}

// 创建新的BufferWriter
func (m *MergeWriter) NewWriter(index int) io.Writer {
	writer := &BufferWriter{
		merge:  m,
		prefix: fmt.Sprintf(" [P%02d] ", index),
	}

	// 分配
	writer.Buffer = bufferPool.Get()
	writer.Buffer.WriteString(writer.prefix)
	return writer
}

type BufferWriter struct {
	Buffer *bytes.Buffer
	prefix string
	merge  *MergeWriter
}

func (b *BufferWriter) Write(p []byte) (n int, err error) {
	n = len(p)

	for len(p) > 0 {
		index := bytes.IndexByte(p, '\n')
		if index != -1 {
			// 写完完整的一行
			_, err = b.Buffer.Write(p[0 : index+1])
			if err != nil {
				log.Error(err, "Writer Buffer failed")
				return n, err
			}

			// 将buffer转移到merge中
			b.merge.WriteLine(b.Buffer)

			// 分配：写入新数据
			b.Buffer = bufferPool.Get()
			b.Buffer.WriteString(time.Now().Format("15:04:05") + b.prefix)
			p = p[index+1:]
		} else {
			// 剩下不足一行，一口气全部写入
			_, err = b.Buffer.Write(p)
			if err != nil {
				log.Error(err, "Writer Buffer failed")
				return n, err
			}
			break
		}
	}
	return n, nil
}