summaryrefslogtreecommitdiff
path: root/src/mongo/gotools/src/github.com/mongodb/mongo-tools/mongostat/stat_consumer/stat_consumer.go
blob: eb3fecba6ca110d66ea05d9c3d97b1a12ea1c396 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
// Copyright (C) MongoDB, Inc. 2014-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0

package stat_consumer

import (
	"fmt"
	"io"
	"os"

	"github.com/mongodb/mongo-tools-common/util"
	"github.com/mongodb/mongo-tools/mongostat/stat_consumer/line"
	"github.com/mongodb/mongo-tools/mongostat/status"
)

// StatConsumer maintains the current set of headers and the most recent
// ServerStatus for each host. It creates a StatLine when passed a new record
// and can format and write groups of StatLines.
type StatConsumer struct {
	formatter              LineFormatter
	readerConfig           *status.ReaderConfig
	oldStats               map[string]*status.ServerStatus
	headers, customHeaders []string
	keyNames               map[string]string
	writer                 io.Writer
	flags                  int
}

// NewStatConsumer creates a new StatConsumer with no previous records
func NewStatConsumer(flags int, customHeaders []string, keyNames map[string]string, readerConfig *status.ReaderConfig, formatter LineFormatter, writer io.Writer) (sc *StatConsumer) {
	sc = &StatConsumer{
		formatter:     formatter,
		readerConfig:  readerConfig,
		oldStats:      make(map[string]*status.ServerStatus),
		customHeaders: customHeaders,
		keyNames:      keyNames,
		writer:        writer,
		flags:         flags,
	}
	if flags == 0 {
		sc.headers = customHeaders
	}
	return sc
}

// Update takes in a ServerStatus and returns a StatLine if it has a previous record
func (sc *StatConsumer) Update(newStat *status.ServerStatus) (l *line.StatLine, seen bool) {
	oldStat, seen := sc.oldStats[newStat.Host]
	sc.oldStats[newStat.Host] = newStat
	if seen {
		l = line.NewStatLine(oldStat, newStat, sc.headers, sc.readerConfig)
		return
	}

	if sc.flags != 0 {
		if status.IsMMAP(newStat) {
			sc.flags |= line.FlagMMAP
		} else if status.IsWT(newStat) {
			sc.flags |= line.FlagWT
		}
		if status.IsReplSet(newStat) {
			sc.flags |= line.FlagRepl
		}
		if status.HasLocks(newStat) {
			sc.flags |= line.FlagLocks
		}

		// Modify headers
		sc.headers = []string{}
		for _, desc := range line.CondHeaders {
			if desc.Flag&sc.flags == desc.Flag {
				sc.headers = append(sc.headers, desc.Key)
			}
		}
		sc.headers = append(sc.headers, sc.customHeaders...)
	}
	return
}

// FormatLines consumes StatLines, formats them, and sends them to its writer
// It returns true if the formatter should no longer receive data
func (sc *StatConsumer) FormatLines(lines []*line.StatLine) bool {
	str := sc.formatter.FormatLines(lines, sc.headers, sc.keyNames)
	_, err := fmt.Fprintf(sc.writer, "%s", str)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error writing formatted output: %v", err)
		os.Exit(util.ExitFailure)
	}
	return sc.formatter.IsFinished()
}