storj/cmd/storj-sim/prefix.go

178 lines
3.6 KiB
Go
Raw Normal View History

2019-01-24 20:15:10 +00:00
// Copyright (C) 2019 Storj Labs, Inc.
2019-01-02 18:07:49 +00:00
// See LICENSE for copying information.
package main
import (
"bytes"
"fmt"
"io"
"sync"
2019-01-08 15:24:15 +00:00
"time"
2019-01-02 18:07:49 +00:00
)
// PrefixWriter writes to the specified output with prefixes.
type PrefixWriter struct {
root *prefixWriter
maxline int
nowFunc func() time.Time
2019-01-02 18:07:49 +00:00
2019-01-08 15:24:15 +00:00
mu sync.Mutex
prefixlen int
dst io.Writer
}
const (
maxIDLength = 10
timeFormat = "15:04:05.000"
emptyTimeField = " "
)
2019-01-08 15:24:15 +00:00
func max(a, b int) int {
if a > b {
return a
}
return b
2019-01-02 18:07:49 +00:00
}
// NewPrefixWriter creates a writer than can prefix all lines written to it.
func NewPrefixWriter(defaultPrefix string, maxLineLen int, dst io.Writer) *PrefixWriter {
2019-01-02 18:07:49 +00:00
writer := &PrefixWriter{
maxline: maxLineLen,
2019-01-02 18:07:49 +00:00
dst: dst,
nowFunc: time.Now,
2019-01-02 18:07:49 +00:00
}
writer.root = writer.Prefixed(defaultPrefix).(*prefixWriter)
return writer
}
// prefixWriter is the implementation that handles buffering and prefixing.
type prefixWriter struct {
*PrefixWriter
prefix string
local sync.Mutex
2019-01-08 15:24:15 +00:00
id string
2019-01-02 18:07:49 +00:00
buffer []byte
}
// Prefixed returns a new writer that has writes with specified prefix.
func (writer *PrefixWriter) Prefixed(prefix string) io.Writer {
writer.mu.Lock()
2019-01-08 15:24:15 +00:00
writer.prefixlen = max(writer.prefixlen, len(prefix))
2019-01-02 18:07:49 +00:00
writer.mu.Unlock()
return &prefixWriter{
PrefixWriter: writer,
prefix: prefix,
id: "",
buffer: make([]byte, 0, writer.maxline),
}
2019-01-02 18:07:49 +00:00
}
// Write implements io.Writer that prefixes lines.
func (writer *PrefixWriter) Write(data []byte) (int, error) {
return writer.root.Write(data)
}
// Write implements io.Writer that prefixes lines.
2019-01-02 18:07:49 +00:00
func (writer *prefixWriter) Write(data []byte) (int, error) {
if len(data) == 0 {
return 0, nil
}
writer.local.Lock()
defer writer.local.Unlock()
2019-01-08 15:24:15 +00:00
var newID string
if writer.id == "" {
if start := bytes.Index(data, []byte("Node ")); start > 0 {
if end := bytes.Index(data[start:], []byte(" started")); end > 0 {
2019-01-08 15:24:15 +00:00
newID = string(data[start+5 : start+end])
if len(newID) > maxIDLength {
newID = newID[:maxIDLength]
}
}
}
}
2019-01-02 18:07:49 +00:00
buffer := data
// buffer everything that hasn't been written yet
if len(writer.buffer) > 0 {
buffer = writer.buffer
buffer = append(buffer, data...)
2019-01-02 18:07:49 +00:00
defer func() {
writer.buffer = buffer
}()
} else {
defer func() {
if len(buffer) > 0 {
writer.buffer = append(writer.buffer, buffer...)
}
}()
}
writer.mu.Lock()
defer writer.mu.Unlock()
2019-01-08 15:24:15 +00:00
if newID != "" {
writer.id = newID
}
2019-01-02 18:07:49 +00:00
prefix := writer.prefix
2019-01-08 15:24:15 +00:00
id := writer.id
timeText := writer.nowFunc().Format(timeFormat)
2019-01-02 18:07:49 +00:00
for len(buffer) > 0 {
pos := bytes.IndexByte(buffer, '\n')
insertbreak := false
// did not find a linebreak
if pos < 0 {
// wait for more data, if we haven't reached maxline
2019-01-02 18:07:49 +00:00
if len(buffer) < writer.maxline {
return len(data), nil
}
}
// try to find a nice place where to break the line
2019-01-02 18:07:49 +00:00
if pos < 0 || pos > writer.maxline {
pos = writer.maxline - 1
2019-01-02 18:07:49 +00:00
for p := pos; p >= writer.maxline*2/3; p-- {
// is there a space we can break on?
2019-01-02 18:07:49 +00:00
if buffer[p] == ' ' {
pos = p
break
}
}
insertbreak = true
2019-01-02 18:07:49 +00:00
}
2019-01-08 15:24:15 +00:00
_, err := fmt.Fprintf(writer.dst, "%-*s %-*s %s | ", writer.prefixlen, prefix, maxIDLength, id, timeText)
2019-01-02 18:07:49 +00:00
if err != nil {
return len(data), err
}
_, err = writer.dst.Write(buffer[:pos])
buffer = buffer[pos:]
if err != nil {
return len(data), err
}
_, err = writer.dst.Write([]byte{'\n'})
2019-01-02 18:07:49 +00:00
if err != nil {
return len(data), err
}
// remove the linebreak from buffer, if it's not an insert
if !insertbreak && len(buffer) > 0 {
buffer = buffer[1:]
2019-01-02 18:07:49 +00:00
}
prefix = ""
2019-01-08 15:24:15 +00:00
id = ""
timeText = emptyTimeField
2019-01-02 18:07:49 +00:00
}
return len(data), nil
}