feat: display consumer lag

- add HighWaterMark and Lag fields to TopicPartitionOffset
- add High Watermark and Lag columns to consumer group's offset lister
page

add refresh for consumer group offset page

fix tests

- fix existing tests
- update tests to include High Watermark and Lag fields

- handle case when kafka client's getOffset results in error.

- use go routines to fetch newest offset (high watermark) for a topic's partition(s)

use a mutex lock to avoid nested function, variable capturing issue, and the unnecessary for loop

make height & weight consistent for offsets page

- switch between topics table and offsets table using tab
- add a single row total table view that displays total number of partitions, sum of offsets, high watermark, and lag across partitions

fix tests

- remove the total table with total lag window title
- adjust table width

Author: offbythree

kadmin/offset_lister.go

@@ -1,15 +1,27 @@
 package kadmin
 
-import tea "github.com/charmbracelet/bubbletea"
+import (
+	"math"
+	"sync"
+
+	"github.com/IBM/sarama"
+	tea "github.com/charmbracelet/bubbletea"
+)
+
+const (
+	ErrorValue int64 = math.MinInt64
+)
 
 type OffsetLister interface {
 	ListOffsets(group string) tea.Msg
 }
 
 type TopicPartitionOffset struct {
-	Topic     string
-	Partition int32
-	Offset    int64
+	Topic         string
+	Partition     int32
+	Offset        int64
+	HighWaterMark int64
+	Lag           int64
 }
 
 type OffsetListingStartedMsg struct {
@@ -51,18 +63,46 @@ func (ka *SaramaKafkaAdmin) doListOffsets(group string, offsetsChan chan []Topic
 	listResult, err := ka.admin.ListConsumerGroupOffsets(group, nil)
 	if err != nil {
 		errChan <- err
+		return
+	}
+
+	totalPartitions := 0
+	for _, m := range listResult.Blocks {
+		totalPartitions += len(m)
 	}
 
-	var topicPartitionOffsets []TopicPartitionOffset
+	topicPartitionOffsets := make([]TopicPartitionOffset, 0, totalPartitions)
+
+	var mu sync.Mutex
+	var wg sync.WaitGroup
+
 	for t, m := range listResult.Blocks {
 		for p, block := range m {
-			topicPartitionOffsets = append(topicPartitionOffsets, TopicPartitionOffset{
-				Topic:     t,
-				Partition: p,
-				Offset:    block.Offset,
-			})
+			wg.Go(
+				func() {
+					hwm, err := ka.client.GetOffset(t, p, sarama.OffsetNewest)
+					var lag int64
+					if err != nil {
+						hwm = ErrorValue
+						lag = ErrorValue
+					} else {
+						lag = hwm - block.Offset
+					}
+					mu.Lock()
+					topicPartitionOffsets = append(topicPartitionOffsets, TopicPartitionOffset{
+						Topic:         t,
+						Partition:     p,
+						Offset:        block.Offset,
+						HighWaterMark: hwm,
+						Lag:           lag,
+					})
+					mu.Unlock()
+				},
+			)
 		}
 	}
 
+	wg.Wait()
+
 	offsetsChan <- topicPartitionOffsets
 }

kadmin/offset_lister_test.go

@@ -2,10 +2,11 @@ package kadmin
 
 import (
 	"context"
-	"github.com/IBM/sarama"
-	"github.com/stretchr/testify/assert"
 	"testing"
 	"time"
+
+	"github.com/IBM/sarama"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestConsumerGroupOffsets(t *testing.T) {
@@ -52,9 +53,11 @@ func TestConsumerGroupOffsets(t *testing.T) {
 			assert.NotNil(t, offsets)
 			assert.Len(t, offsets, 1)
 			assert.Equal(t, offsets[0], TopicPartitionOffset{
-				Topic:     topic,
-				Partition: 0,
-				Offset:    9,
+				Topic:         topic,
+				Partition:     0,
+				Offset:        9,
+				HighWaterMark: 10,
+				Lag:           1,
 			})
 		case err := <-offsetListingStartedMsg.Err:
 			t.Fatal("Error while listing offsets", err)

ui/pages/cgroups_topics_page/cgroups_parts_offsets_page.go

@@ -2,7 +2,6 @@ package cgroups_topics_page
 
 import (
 	"fmt"
-	"github.com/charmbracelet/log"
 	"ktea/kadmin"
 	"ktea/kontext"
 	"ktea/styles"
@@ -18,6 +17,8 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/charmbracelet/log"
+
 	"github.com/charmbracelet/bubbles/table"
 	tea "github.com/charmbracelet/bubbletea"
 	lg "github.com/charmbracelet/lipgloss"
@@ -29,6 +30,7 @@ type tableFocus int
 type state int
 
 const (
+	na          string     = "N/A"
 	topicFocus  tableFocus = 0
 	offsetFocus tableFocus = 1
 
@@ -38,13 +40,16 @@ const (
 )
 
 type Model struct {
+	lister            kadmin.OffsetLister
 	tableFocus        tableFocus
 	topicsTable       table.Model
 	offsetsTable      table.Model
+	totalTable        table.Model
 	offsetsBorder     *border.Model
 	topicsBorder      *border.Model
 	topicsRows        []table.Row
 	offsetRows        []table.Row
+	totalLag          int64
 	groupName         string
 	topicByPartOffset map[string][]partOffset
 	cmdBar            *CGroupCmdbar[string]
@@ -63,17 +68,24 @@ func (m *Model) View(ktx *kontext.ProgramKtx, renderer *ui.Renderer) string {
 	cmdBarView := m.cmdBar.View(ktx, renderer)
 
 	halfWidth := int(float64(ktx.WindowWidth / 2))
-	m.topicsTable.SetHeight(ktx.AvailableHeight - 4)
-	m.topicsTable.SetWidth(halfWidth - 2)
+	m.topicsTable.SetHeight(ktx.AvailableTableHeight())
+	m.topicsTable.SetWidth(int(float64(halfWidth)))
 	m.topicsTable.SetColumns([]table.Column{
-		{"Topic Name", int(float64(halfWidth - 4))},
+		{Title: "Topic Name", Width: int(float64(halfWidth - 2))},
 	})
 	m.topicsTable.SetRows(m.topicsRows)
 
-	m.offsetsTable.SetHeight(ktx.AvailableHeight - 4)
+	partitionColumnWidth := int(float64(halfWidth-4) * 0.22)
+	offsetColumnWidth := int(float64(halfWidth-4) * 0.24)
+	hwmColumnWidth := int(float64(halfWidth-4) * 0.24)
+	lagColumnWidth := int(float64(halfWidth-4) * 0.22)
+
+	m.offsetsTable.SetHeight(ktx.AvailableTableHeight())
 	m.offsetsTable.SetColumns([]table.Column{
-		{"Partition", int(float64(halfWidth-6) * 0.5)},
-		{"Offset", int(float64(halfWidth-5) * 0.5)},
+		{Title: "Partition", Width: partitionColumnWidth},
+		{Title: "Offset", Width: offsetColumnWidth},
+		{Title: "High Watermark", Width: hwmColumnWidth},
+		{Title: "Lag", Width: lagColumnWidth},
 	})
 	m.offsetsTable.SetRows(m.offsetRows)
 
@@ -106,6 +118,24 @@ func (m *Model) View(ktx *kontext.ProgramKtx, renderer *ui.Renderer) string {
 type partOffset struct {
 	partition string
 	offset    int64
+	hwm       int64
+	lag       int64
+}
+
+func (partOffset *partOffset) getHwmValue() string {
+	if partOffset.hwm == kadmin.ErrorValue {
+		return na
+	} else {
+		return humanize.Comma(partOffset.hwm)
+	}
+}
+
+func (partOffset *partOffset) getLagValue() string {
+	if partOffset.lag == kadmin.ErrorValue {
+		return na
+	} else {
+		return humanize.Comma(partOffset.lag)
+	}
 }
 
 func (m *Model) Update(msg tea.Msg) tea.Cmd {
@@ -122,6 +152,20 @@ func (m *Model) Update(msg tea.Msg) tea.Cmd {
 			if !m.cmdBar.IsFocussed() {
 				return ui.PublishMsg(nav.LoadCGroupsPageMsg{})
 			}
+		case "f5":
+			m.state = stateOffsetsLoading
+			return func() tea.Msg {
+				return m.lister.ListOffsets(m.groupName)
+			}
+		case "tab":
+			// only accept when the table is focussed
+			if !m.cmdBar.IsFocussed() {
+				if m.tableFocus == topicFocus {
+					m.tableFocus = offsetFocus
+				} else {
+					m.tableFocus = topicFocus
+				}
+			}
 		}
 	case kadmin.OffsetListingStartedMsg:
 		cmds = append(cmds, msg.AwaitCompletion)
@@ -144,8 +188,10 @@ func (m *Model) Update(msg tea.Msg) tea.Cmd {
 	if !m.cmdBar.IsFocussed() {
 		if m.tableFocus == topicFocus {
 			m.topicsTable, cmd = m.topicsTable.Update(msg)
+			m.offsetsTable.GotoTop()
 		} else {
 			m.offsetsTable, cmd = m.offsetsTable.Update(msg)
+			m.totalTable.Update(msg)
 		}
 		if cmd != nil {
 			cmds = append(cmds, cmd)
@@ -167,13 +213,18 @@ func (m *Model) recreateOffsetRows() {
 
 	selectedTopic := m.selectedRow()
 	if selectedTopic != "" {
+		totalLag := int64(0)
 		m.offsetRows = []table.Row{}
 		for _, partOffset := range m.topicByPartOffset[selectedTopic] {
+			totalLag += int64(partOffset.lag)
 			m.offsetRows = append(m.offsetRows, table.Row{
 				partOffset.partition,
 				humanize.Comma(partOffset.offset),
+				partOffset.getHwmValue(),
+				partOffset.getLagValue(),
 			})
 		}
+		m.totalLag = totalLag
 		sort.SliceStable(m.offsetRows, func(i, j int) bool {
 			a, _ := strconv.Atoi(m.offsetRows[i][0])
 			b, _ := strconv.Atoi(m.offsetRows[j][0])
@@ -201,6 +252,8 @@ func (m *Model) recreateTopicRows() {
 		partOffset := partOffset{
 			partition: strconv.FormatInt(int64(offset.Partition), 10),
 			offset:    offset.Offset,
+			hwm:       offset.HighWaterMark,
+			lag:       offset.Lag,
 		}
 		m.topicByPartOffset[offset.Topic] = append(m.topicByPartOffset[offset.Topic], partOffset)
 	}
@@ -223,9 +276,9 @@ func (m *Model) selectedRow() string {
 
 func (m *Model) Shortcuts() []statusbar.Shortcut {
 	return []statusbar.Shortcut{
-		{"Go Back", "esc"},
-		{"Search", "/"},
-		{"Refresh", "F5"},
+		{Name: "Go Back", Keybinding: "esc"},
+		{Name: "Search", Keybinding: "/"},
+		{Name: "Refresh", Keybinding: "F5"},
 	}
 }
 
@@ -268,6 +321,7 @@ func New(lister kadmin.OffsetLister, group string) (*Model, tea.Cmd) {
 	)
 
 	model := Model{
+		lister: lister,
 		cmdBar: NewCGroupCmdbar[string](
 			cmdbar.NewSearchCmdBar("Search groups by name"),
 			notifierCmdBar,
@@ -283,7 +337,10 @@ func New(lister kadmin.OffsetLister, group string) (*Model, tea.Cmd) {
 		border.WithTitleFn(func() string {
 			return border.KeyValueTitle("Total Topics", fmt.Sprintf(" %d", len(model.topicsRows)), true)
 		}))
-	model.offsetsBorder = border.New(border.WithInnerPaddingTop())
+	model.offsetsBorder = border.New(border.WithInnerPaddingTop(),
+		border.WithTitleFn(func() string {
+			return border.KeyValueTitle("Total Lag", fmt.Sprintf(" %d", model.totalLag), false)
+		}))
 	return &model, func() tea.Msg {
 		return lister.ListOffsets(group)
 	}