Support reopening container log.

Signed-off-by: Lantao Liu <lantaol@google.com>

pkg/ioutil/writer_group.go

@@ -39,16 +39,24 @@ func NewWriterGroup() *WriterGroup {
 	}
 }
 
-// Add adds a writer into the group, returns an error when writer
-// group is closed.
-func (g *WriterGroup) Add(key string, w io.WriteCloser) error {
+// Add adds a writer into the group. The writer will be closed
+// if the writer group is closed.
+func (g *WriterGroup) Add(key string, w io.WriteCloser) {
 	g.mu.Lock()
 	defer g.mu.Unlock()
 	if g.closed {
-		return errors.New("wait group closed")
+		w.Close()
+		return
 	}
 	g.writers[key] = w
-	return nil
+}
+
+// Get gets a writer from the group, returns nil if the writer
+// doesn't exist.
+func (g *WriterGroup) Get(key string) io.WriteCloser {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	return g.writers[key]
 }
 
 // Remove removes a writer from the group.
@@ -84,8 +92,8 @@ func (g *WriterGroup) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
-// Close closes the writer group. Write or Add will return error
-// after closed.
+// Close closes the writer group. Write will return error after
+// closed.
 func (g *WriterGroup) Close() {
 	g.mu.Lock()
 	defer g.mu.Unlock()

pkg/ioutil/writer_group_test.go

@@ -48,8 +48,7 @@ func TestClosedWriterGroup(t *testing.T) {
 	wc := &writeCloser{}
 	key, data := "test key", "test data"
 
-	err := wg.Add(key, wc)
-	assert.NoError(t, err)
+	wg.Add(key, wc)
 
 	n, err := wg.Write([]byte(data))
 	assert.Equal(t, len(data), n)
@@ -59,36 +58,58 @@ func TestClosedWriterGroup(t *testing.T) {
 	wg.Close()
 	assert.True(t, wc.closed)
 
-	err = wg.Add(key, &writeCloser{})
-	assert.Error(t, err)
+	newWC := &writeCloser{}
+	wg.Add(key, newWC)
+	assert.True(t, newWC.closed)
 
 	_, err = wg.Write([]byte(data))
 	assert.Error(t, err)
 }
 
-func TestAddRemoveWriter(t *testing.T) {
+func TestAddGetRemoveWriter(t *testing.T) {
 	wg := NewWriterGroup()
 	wc1, wc2 := &writeCloser{}, &writeCloser{}
 	key1, key2 := "test key 1", "test key 2"
 
-	err := wg.Add(key1, wc1)
-	assert.NoError(t, err)
-	_, err = wg.Write([]byte("test data 1"))
+	wg.Add(key1, wc1)
+	_, err := wg.Write([]byte("test data 1"))
 	assert.NoError(t, err)
 	assert.Equal(t, "test data 1", wc1.buf.String())
 
-	err = wg.Add(key2, wc2)
-	assert.NoError(t, err)
+	wg.Add(key2, wc2)
 	_, err = wg.Write([]byte("test data 2"))
 	assert.NoError(t, err)
 	assert.Equal(t, "test data 1test data 2", wc1.buf.String())
 	assert.Equal(t, "test data 2", wc2.buf.String())
 
+	assert.Equal(t, wc1, wg.Get(key1))
+
 	wg.Remove(key1)
 	_, err = wg.Write([]byte("test data 3"))
 	assert.NoError(t, err)
 	assert.Equal(t, "test data 1test data 2", wc1.buf.String())
 	assert.Equal(t, "test data 2test data 3", wc2.buf.String())
 
+	assert.Equal(t, nil, wg.Get(key1))
+
+	wg.Close()
+}
+
+func TestReplaceWriter(t *testing.T) {
+	wg := NewWriterGroup()
+	wc1, wc2 := &writeCloser{}, &writeCloser{}
+	key := "test-key"
+
+	wg.Add(key, wc1)
+	_, err := wg.Write([]byte("test data 1"))
+	assert.NoError(t, err)
+	assert.Equal(t, "test data 1", wc1.buf.String())
+
+	wg.Add(key, wc2)
+	_, err = wg.Write([]byte("test data 2"))
+	assert.NoError(t, err)
+	assert.Equal(t, "test data 1", wc1.buf.String())
+	assert.Equal(t, "test data 2", wc2.buf.String())
+
 	wg.Close()
 }