storage/mem: implement size enforcer for #88

pkg/storage/mem/maxsize.go

@@ -0,0 +1,73 @@
+package mem
+
+import "container/list"
+
+type msgDone struct {
+	msg  *Message
+	done chan struct{}
+}
+
+// enforceMaxSize will delete the oldest message until the entire mail store is equal to or less
+// than Store.maxSize bytes.
+func (s *Store) maxSizeEnforcer(maxSize int64) {
+	all := &list.List{}
+	curSize := int64(0)
+	for {
+		select {
+		case md, ok := <-s.incoming:
+			if !ok {
+				return
+			}
+			// Add message to all.
+			m := md.msg
+			el := all.PushBack(m)
+			m.el = el
+			curSize += int64(m.Size())
+			for curSize > maxSize {
+				// Remove oldest message.
+				el := all.Front()
+				all.Remove(el)
+				m := el.Value.(*Message)
+				if s.removeMessage(m.mailbox, m.id) != nil {
+					curSize -= int64(m.Size())
+				}
+			}
+			close(md.done)
+		case md, ok := <-s.remove:
+			if !ok {
+				return
+			}
+			// Remove message from all.
+			m := md.msg
+			el := all.Remove(m.el)
+			if el != nil {
+				curSize -= int64(m.Size())
+			}
+			close(md.done)
+		}
+	}
+}
+
+// enforcerDeliver sends delivery to enforcer if configured, and waits for completion.
+func (s *Store) enforcerDeliver(m *Message) {
+	if s.incoming != nil {
+		md := &msgDone{
+			msg:  m,
+			done: make(chan struct{}),
+		}
+		s.incoming <- md
+		<-md.done
+	}
+}
+
+// enforcerRemove sends removal to enforcer if configured, and waits for completion.
+func (s *Store) enforcerRemove(m *Message) {
+	if s.remove != nil {
+		md := &msgDone{
+			msg:  m,
+			done: make(chan struct{}),
+		}
+		s.remove <- md
+		<-md.done
+	}
+}

pkg/storage/mem/message.go

@@ -2,6 +2,7 @@ package mem
 
 import (
 	"bytes"
+	"container/list"
 	"io"
 	"io/ioutil"
 	"net/mail"
@@ -20,6 +21,7 @@ type Message struct {
 	date    time.Time
 	subject string
 	source  []byte
+	el      *list.Element // This message in Store.messages
 }
 
 var _ storage.Message = &Message{}

pkg/storage/mem/store.go

@@ -1,6 +1,7 @@
 package mem
 
 import (
+	"fmt"
 	"io/ioutil"
 	"sort"
 	"strconv"
@@ -13,8 +14,10 @@ import (
 // Store implements an in-memory message store.
 type Store struct {
 	sync.Mutex
-	boxes map[string]*mbox
-	cap   int
+	boxes    map[string]*mbox
+	cap      int           // Per-mailbox message cap.
+	incoming chan *msgDone // New messages for size enforcer.
+	remove   chan *msgDone // Remove deleted messages from size enforcer.
 }
 
 type mbox struct {
@@ -29,38 +32,51 @@ var _ storage.Store = &Store{}
 
 // New returns an emtpy memory store.
 func New(cfg config.Storage) (storage.Store, error) {
-	return &Store{
+	s := &Store{
 		boxes: make(map[string]*mbox),
 		cap:   cfg.MailboxMsgCap,
-	}, nil
+	}
+	if str, ok := cfg.Params["maxkb"]; ok {
+		maxKB, err := strconv.ParseInt(str, 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse maxkb: %v", err)
+		}
+		if maxKB > 0 {
+			// Setup enforcer.
+			s.incoming = make(chan *msgDone)
+			s.remove = make(chan *msgDone)
+			go s.maxSizeEnforcer(maxKB * 1024)
+		}
+	}
+	return s, nil
 }
 
 // AddMessage stores the message, message ID and Size will be ignored.
 func (s *Store) AddMessage(message storage.Message) (id string, err error) {
+	r, ierr := message.Source()
+	if ierr != nil {
+		err = ierr
+		return
+	}
+	source, ierr := ioutil.ReadAll(r)
+	if ierr != nil {
+		err = ierr
+		return
+	}
+	m := &Message{
+		mailbox: message.Mailbox(),
+		from:    message.From(),
+		to:      message.To(),
+		date:    message.Date(),
+		subject: message.Subject(),
+	}
 	s.withMailbox(message.Mailbox(), true, func(mb *mbox) {
-		r, ierr := message.Source()
-		if ierr != nil {
-			err = ierr
-			return
-		}
-		source, ierr := ioutil.ReadAll(r)
-		if ierr != nil {
-			err = ierr
-			return
-		}
 		// Generate message ID.
 		mb.last++
+		m.index = mb.last
 		id = strconv.Itoa(mb.last)
-		m := &Message{
-			index:   mb.last,
-			mailbox: message.Mailbox(),
-			id:      id,
-			from:    message.From(),
-			to:      message.To(),
-			date:    message.Date(),
-			subject: message.Subject(),
-			source:  source,
-		}
+		m.id = id
+		m.source = source
 		mb.messages[id] = m
 		if s.cap > 0 {
 			// Enforce cap.
@@ -70,6 +86,7 @@ func (s *Store) AddMessage(message storage.Message) (id string, err error) {
 			}
 		}
 	})
+	s.enforcerDeliver(m)
 	return id, err
 }
 
@@ -97,17 +114,38 @@ func (s *Store) GetMessages(mailbox string) (ms []storage.Message, err error) {
 
 // PurgeMessages deletes the contents of a mailbox.
 func (s *Store) PurgeMessages(mailbox string) error {
+	var messages map[string]*Message
 	s.withMailbox(mailbox, true, func(mb *mbox) {
+		messages = mb.messages
 		mb.messages = make(map[string]*Message)
 	})
+	if len(messages) > 0 && s.remove != nil {
+		for _, m := range messages {
+			s.enforcerRemove(m)
+		}
+	}
 	return nil
 }
 
+// removeMessage deletes a single message without notifying the size enforcer.  Returns the message
+// that was removed.
+func (s *Store) removeMessage(mailbox, id string) *Message {
+	var m *Message
+	s.withMailbox(mailbox, true, func(mb *mbox) {
+		m = mb.messages[id]
+		if m != nil {
+			delete(mb.messages, id)
+		}
+	})
+	return m
+}
+
 // RemoveMessage deletes a single message.
 func (s *Store) RemoveMessage(mailbox, id string) error {
-	s.withMailbox(mailbox, true, func(mb *mbox) {
-		delete(mb.messages, id)
-	})
+	m := s.removeMessage(mailbox, id)
+	if m != nil {
+		s.enforcerRemove(m)
+	}
 	return nil
 }
 

pkg/storage/mem/store_test.go

@@ -1,7 +1,9 @@
 package mem
 
 import (
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/jhillyerd/inbucket/pkg/config"
 	"github.com/jhillyerd/inbucket/pkg/storage"
@@ -16,3 +18,65 @@ func TestSuite(t *testing.T) {
 		return s, destroy, nil
 	})
 }
+
+// TestMessageList verifies the operation of the global message list: mem.Store.messages.
+func TestMaxSize(t *testing.T) {
+	maxSize := int64(2048)
+	s, _ := New(config.Storage{Params: map[string]string{"maxkb": "2"}})
+	boxes := []string{"alpha", "beta", "whiskey", "tango", "foxtrot"}
+	n := 10
+	// total := 50
+	sizeChan := make(chan int64, len(boxes))
+	// Populate mailboxes concurrently.
+	for _, mailbox := range boxes {
+		go func(mailbox string) {
+			size := int64(0)
+			for i := 0; i < n; i++ {
+				_, nbytes := test.DeliverToStore(t, s, mailbox, "subject", time.Now())
+				size += nbytes
+			}
+			sizeChan <- size
+		}(mailbox)
+	}
+	// Wait for sizes.
+	sentBytesTotal := int64(0)
+	for range boxes {
+		sentBytesTotal += <-sizeChan
+	}
+	// Calculate actual size.
+	gotSize := int64(0)
+	s.VisitMailboxes(func(messages []storage.Message) bool {
+		for _, m := range messages {
+			gotSize += m.Size()
+		}
+		return true
+	})
+	// Verify state. Messages are ~75 bytes each.
+	if gotSize < 2048-75 {
+		t.Errorf("Got total size %v, want greater than: %v", gotSize, 2048-75)
+	}
+	if gotSize > maxSize {
+		t.Errorf("Got total size %v, want less than: %v", gotSize, maxSize)
+	}
+	// Purge all messages concurrently, testing for deadlocks.
+	wg := &sync.WaitGroup{}
+	wg.Add(len(boxes))
+	for _, mailbox := range boxes {
+		go func(mailbox string) {
+			err := s.PurgeMessages(mailbox)
+			if err != nil {
+				t.Fatal(err)
+			}
+			wg.Done()
+		}(mailbox)
+	}
+	wg.Wait()
+	count := 0
+	s.VisitMailboxes(func(messages []storage.Message) bool {
+		count += len(messages)
+		return true
+	})
+	if count != 0 {
+		t.Errorf("Got %v total messages, want: %v", count, 0)
+	}
+}