Go writer with buffer and timeout

I would like to send requests to AWS from SQS using Go. There are examples of doing this in single event mode, using sqs.SendMessageInput from the v2 SDK, but I would like to batch send instead. I have created an interface to abstract away the implementation details from the rest of my code, along the lines of this:

type UserRepository interface {
   Save(context.Context, User) error
}

As you can see, the interface is not SQS specific and could easily be replaced by e.g. a Postgres implementation. I would really like to keep the interface as clean as possible.

Using SQS there are some considerations that I can think of and probably also some that I didn't.
A batch needs to be sent:

1. after every nth Save
2. after a configurable timeout
3. after context cancellation
4. after the calling site is done with saving all User objects

Note that instead of sending to SQS this should also work similarly for writing to console, which is what I have created in my example.

Is this design possible or will I always have to have a Close function in my interface?

The problem with the code is that the last 4 (nr_of_items % batch_size) will not be "saved."

package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"time"
)

type LoggingBufferedUserRepository struct {
	buffer        []string
	bufferSize    int
	bufferTimeout time.Duration
	mutex         sync.Mutex
	closeChan     chan struct{}
}

func NewLoggingBufferedUserRepository(
	ctx context.Context, bufferSize int, bufferTimeout time.Duration,
) *LoggingBufferedUserRepository {
	client := &LoggingBufferedUserRepository{
		bufferSize:    bufferSize,
		bufferTimeout: bufferTimeout,
		closeChan:     make(chan struct{}),
	}

	go client.bufferMonitor(ctx)
	return client
}

func (c *LoggingBufferedUserRepository) SendMessage(ctx context.Context, input string) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.buffer = append(c.buffer, input)
	if len(c.buffer) >= c.bufferSize {
		go c.flush(ctx, c.buffer)
		c.buffer = []string{}
	}
	return
}

func (c *LoggingBufferedUserRepository) flush(ctx context.Context, buffer []string) {
	if len(buffer) == 0 {
		return
	}

	// This is the actual batch 'save':
	fmt.Printf("flushing buffer, size=%d, cid=%s, buffer=$%v\n", len(buffer), ctx.Value("cid"), buffer)
}

func (c *LoggingBufferedUserRepository) bufferMonitor(ctx context.Context) {
	timeout := time.NewTimer(c.bufferTimeout)
	for {
		select {
		case <-timeout.C:
			c.flush(ctx, c.buffer)
			c.buffer = []string{}
		}

		timeout.Reset(c.bufferTimeout)
	}
}

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
	defer stop()
	g := NewLoggingBufferedUserRepository(ctx, 10, 1*time.Second)

	wg := &sync.WaitGroup{}
	for i := 0; i < 15; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			fmt.Printf("sending message %d\n", i)
			g.SendMessage(ctx, fmt.Sprintf("a%d", i))
		}(i)
		time.Sleep(100 * time.Millisecond)
	}
	wg.Wait()
	fmt.Println("done")
}

There are multiple ways of doing this. In general, it is not advisable to flush such buffered implementations based on timeout, because you have no control over what will happen if things fail.

One way is to make an explicit batch operation:

type UserRepository interface {
// Save single user
Save(context.Context, User) error
// Save batch
SaveBatch(context.Context) UserBatch
}

where

type UserBatch interface {
Save(context.Context,User) error
Flush(context.Context) error
Close(context.Context) error
}

> Can UserRepository interface support the batch operations without
> needing a close function?

Yes, one solution is to have a different Flush function in your interface that the caller can use to explicitly flush any remaining buffered items. Your interface:

type UserRepository interface {
    Save(context.Context, User) error
    Flush(context.Context) error
    SaveAll(context.Context) error
}

One example to implement this UserRepository for your likings, could be something like playground:

// for SO testing, by JS

package main

import (
	"context"
	"fmt"
	"sync"
	"time"
)

type User struct {
	ID   int
	Name string
}
type UserRepository interface {
	Save(context.Context, User) error
	Flush(context.Context) error
	SaveAll(context.Context) error
}

type LoggingBufferedUserRepository struct {
	buffer        []string
	bufferSize    int
	bufferTimeout time.Duration
	mutex         sync.Mutex
	// we don't need close here now
}

// add the item to the buffer and flushes the buffer if it reaches the configured size
func (c *LoggingBufferedUserRepository) Save(ctx context.Context, user User) error {
	// implement save functionality
	fmt.Printf("Save called %v: %v\n", user.Name, user.ID)
	return nil
}

// can - Flush function immediately flushes any buffered items
func (c *LoggingBufferedUserRepository) Flush(ctx context.Context) error {
	// implement Flush functionality
	fmt.Printf("Flush called\n")
	return nil
}

// waits for any pending items to be saved before returning, and flushes any remaining buffered items if there are any
func (c *LoggingBufferedUserRepository) SaveAll(ctx context.Context) error {
	// implement SaveAll functionality
	fmt.Printf("SaveAll called\n")
	return nil
}

func (c *LoggingBufferedUserRepository) flush(ctx context.Context) error {
	if len(c.buffer) == 0 {
		return nil
	}

	// This is your actual batch 'save':
	fmt.Printf("flushing buffer, size=%d, cid=%s, buffer=$%v\n", len(c.buffer), ctx.Value("cid"), c.buffer)
	c.buffer = []string{}

	return nil
}

func main() {
	repo := &LoggingBufferedUserRepository{
		bufferSize:    5,
		bufferTimeout: 1 * time.Second,
	}

	users := []User{
		{ID: 1, Name: "Alice"},
		{ID: 2, Name: "Bob"},
		{ID: 3, Name: "Charlie"},
		{ID: 4, Name: "Dave"},
		{ID: 5, Name: "Eve"},
	}

	for _, user := range users {
		if err := repo.Save(context.Background(), user); err != nil {
			fmt.Printf("error saving user %v: %v\n", user, err)
		}
	}

	fmt.Printf("Waiting for flushing... (5 seconds)\n")
	time.Sleep(5 * time.Second)

	// Save remaining items
	if err := repo.SaveAll(context.Background()); err != nil {
		fmt.Printf("error saving remaining items: %v\n", err)
	}

	// Flush remaining items
	if err := repo.Flush(context.Background()); err != nil {
		fmt.Printf("error flushing remaining items: %v\n", err)
	}
}

The output of this should be:

Save called Alice: 1
Save called Bob: 2
Save called Charlie: 3
Save called Dave: 4
Save called Eve: 5
Waiting for flushing... (5 seconds)
SaveAll called
Flush called
Program exited.