storj/cmd/storj-sim/prefix.go

154 lines
3.1 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
2019-01-08 15:24:15 +00:00
mu sync.Mutex
prefixlen int
dst io.Writer
}
const maxIDLength = 10
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, dst io.Writer) *PrefixWriter {
writer := &PrefixWriter{
maxline: 10000, // disable maxline cutting
dst: dst,
}
writer.root = writer.Prefixed(defaultPrefix).(*prefixWriter)
return writer
}
// prefixWriter is the implementation that handles buffering and prefixing.
type prefixWriter struct {
*PrefixWriter
prefix string
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()
2019-01-08 15:24:15 +00:00
return &prefixWriter{writer, prefix, "", 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
func (writer *prefixWriter) Write(data []byte) (int, error) {
if len(data) == 0 {
return 0, nil
}
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 {
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 = append(writer.buffer, data...) // nolint gocritic
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 := time.Now().Format("15:04:05.000")
2019-01-02 18:07:49 +00:00
for len(buffer) > 0 {
pos := bytes.IndexByte(buffer, '\n') + 1
breakline := false
if pos <= 0 {
if len(buffer) < writer.maxline {
return len(data), nil
}
}
if pos < 0 || pos > writer.maxline {
pos = writer.maxline
for p := pos; p >= writer.maxline*2/3; p-- {
if buffer[p] == ' ' {
pos = p
break
}
}
breakline = true
}
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
}
if breakline {
_, err = writer.dst.Write([]byte{'\n'})
if err != nil {
return len(data), err
}
}
prefix = ""
2019-01-08 15:24:15 +00:00
id = ""
timeText = " "
2019-01-02 18:07:49 +00:00
}
return len(data), nil
}