Skip to content

Commit

Permalink
Ability to build logger with custom Sink (#572)
Browse files Browse the repository at this point in the history
Using the same global-registry pattern that we're using for encoders, let users register new sinks (like the built-in `stdout` and `stderr`).
  • Loading branch information
dimroc authored and akshayjshah committed Jun 27, 2018
1 parent f4243df commit 7e7e266
Show file tree
Hide file tree
Showing 4 changed files with 240 additions and 36 deletions.
94 changes: 94 additions & 0 deletions sink.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zap

import (
"errors"
"fmt"
"io"
"os"
"sync"

"go.uber.org/zap/zapcore"
)

var (
_sinkMutex sync.RWMutex
_sinkFactories map[string]func() (Sink, error)
)

func init() {
resetSinkRegistry()
}

func resetSinkRegistry() {
_sinkMutex.Lock()
defer _sinkMutex.Unlock()
_sinkFactories = map[string]func() (Sink, error){
"stdout": func() (Sink, error) { return nopCloserSink{os.Stdout}, nil },
"stderr": func() (Sink, error) { return nopCloserSink{os.Stderr}, nil },
}
}

type errSinkNotFound struct {
key string
}

func (e *errSinkNotFound) Error() string {
return fmt.Sprintf("no sink found for %q", e.key)
}

// Sink defines the interface to write to and close logger destinations.
type Sink interface {
zapcore.WriteSyncer
io.Closer
}

// RegisterSink adds a Sink at the given key so it can be referenced
// in config OutputPaths.
func RegisterSink(key string, sinkFactory func() (Sink, error)) error {
_sinkMutex.Lock()
defer _sinkMutex.Unlock()
if key == "" {
return errors.New("sink key cannot be blank")
}
if _, ok := _sinkFactories[key]; ok {
return fmt.Errorf("sink already registered for key %q", key)
}
_sinkFactories[key] = sinkFactory
return nil
}

// newSink invokes the registered sink factory to create and return the
// sink for the given key. Returns errSinkNotFound if the key cannot be found.
func newSink(key string) (Sink, error) {
_sinkMutex.RLock()
defer _sinkMutex.RUnlock()
sinkFactory, ok := _sinkFactories[key]
if !ok {
return nil, &errSinkNotFound{key}
}
return sinkFactory()
}

type nopCloserSink struct{ zapcore.WriteSyncer }

func (nopCloserSink) Close() error { return nil }
76 changes: 76 additions & 0 deletions sink_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package zap

import (
"errors"
"os"
"testing"

"github.com/stretchr/testify/assert"
)

func TestRegisterSink(t *testing.T) {
tests := []struct {
name string
key string
factory func() (Sink, error)
wantError bool
}{
{"valid", "valid", func() (Sink, error) { return nopCloserSink{os.Stdout}, nil }, false},
{"empty", "", func() (Sink, error) { return nopCloserSink{os.Stdout}, nil }, true},
{"stdout", "stdout", func() (Sink, error) { return nopCloserSink{os.Stdout}, nil }, true},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := RegisterSink(tt.key, tt.factory)
if tt.wantError {
assert.NotNil(t, err)
} else {
assert.Nil(t, err)
assert.NotNil(t, _sinkFactories[tt.key], "expected the factory to be present")
}
})
}
}

func TestNewSink(t *testing.T) {
defer resetSinkRegistry()
errTestSink := errors.New("test erroring")
err := RegisterSink("errors", func() (Sink, error) { return nil, errTestSink })
assert.Nil(t, err)
tests := []struct {
key string
err error
}{
{"stdout", nil},
{"errors", errTestSink},
{"nonexistent", &errSinkNotFound{"nonexistent"}},
}

for _, tt := range tests {
t.Run(tt.key, func(t *testing.T) {
_, err := newSink(tt.key)
assert.Equal(t, tt.err, err)
})
}
}
34 changes: 19 additions & 15 deletions writer.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
package zap

import (
"io"
"io/ioutil"
"os"

Expand Down Expand Up @@ -49,29 +50,32 @@ func Open(paths ...string) (zapcore.WriteSyncer, func(), error) {
func open(paths []string) ([]zapcore.WriteSyncer, func(), error) {
var openErr error
writers := make([]zapcore.WriteSyncer, 0, len(paths))
files := make([]*os.File, 0, len(paths))
closers := make([]io.Closer, 0, len(paths))
close := func() {
for _, f := range files {
f.Close()
for _, c := range closers {
c.Close()
}
}
for _, path := range paths {
switch path {
case "stdout":
writers = append(writers, os.Stdout)
// Don't close standard out.
sink, err := newSink(path)
if err == nil {
// Using a registered sink constructor.
writers = append(writers, sink)
closers = append(closers, sink)
continue
case "stderr":
writers = append(writers, os.Stderr)
// Don't close standard error.
}
if _, ok := err.(*errSinkNotFound); ok {
// No named sink constructor, use key as path to log file.
f, e := os.OpenFile(path, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0644)
openErr = multierr.Append(openErr, e)
if e == nil {
writers = append(writers, f)
closers = append(closers, f)
}
continue
}
f, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0644)
// Sink constructor failed.
openErr = multierr.Append(openErr, err)
if err == nil {
writers = append(writers, f)
files = append(files, f)
}
}

if openErr != nil {
Expand Down
72 changes: 51 additions & 21 deletions writer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,12 @@
package zap

import (
"encoding/hex"
"errors"
"io/ioutil"
"math/rand"
"os"
"path/filepath"
"testing"

"github.com/stretchr/testify/assert"
Expand All @@ -44,28 +48,25 @@ func TestOpenNoPaths(t *testing.T) {
}

func TestOpen(t *testing.T) {
temp, err := ioutil.TempFile("", "zap-open-test")
require.NoError(t, err, "Couldn't create a temporary file for test.")
defer os.Remove(temp.Name())
tempName := tempFileName("", "zap-open-test")
assert.False(t, fileExists(tempName))

tests := []struct {
paths []string
filenames []string
error string
paths []string
error string
}{
{[]string{"stdout"}, []string{os.Stdout.Name()}, ""},
{[]string{"stderr"}, []string{os.Stderr.Name()}, ""},
{[]string{temp.Name()}, []string{temp.Name()}, ""},
{[]string{"/foo/bar/baz"}, []string{}, "open /foo/bar/baz: no such file or directory"},
{[]string{"stdout"}, ""},
{[]string{"stderr"}, ""},
{[]string{tempName}, ""},
{[]string{"/foo/bar/baz"}, "open /foo/bar/baz: no such file or directory"},
{
paths: []string{"stdout", "/foo/bar/baz", temp.Name(), "/baz/quux"},
filenames: []string{os.Stdout.Name(), temp.Name()},
error: "open /foo/bar/baz: no such file or directory; open /baz/quux: no such file or directory",
paths: []string{"stdout", "/foo/bar/baz", tempName, "/baz/quux"},
error: "open /foo/bar/baz: no such file or directory; open /baz/quux: no such file or directory",
},
}

for _, tt := range tests {
wss, cleanup, err := open(tt.paths)
_, cleanup, err := Open(tt.paths...)
if err == nil {
defer cleanup()
}
Expand All @@ -75,14 +76,10 @@ func TestOpen(t *testing.T) {
} else {
assert.Equal(t, tt.error, err.Error(), "Unexpected error opening paths %v.", tt.paths)
}
names := make([]string, len(wss))
for i, ws := range wss {
f, ok := ws.(*os.File)
require.True(t, ok, "Expected all WriteSyncers returned from open() to be files.")
names[i] = f.Name()
}
assert.Equal(t, tt.filenames, names, "Opened unexpected files given paths %v.", tt.paths)
}

assert.True(t, fileExists(tempName))
os.Remove(tempName)
}

func TestOpenFails(t *testing.T) {
Expand Down Expand Up @@ -118,8 +115,41 @@ func (w *testWriter) Sync() error {
return nil
}

func TestOpenWithCustomSink(t *testing.T) {
defer resetSinkRegistry()
tw := &testWriter{"test", t}
ctr := func() (Sink, error) { return nopCloserSink{tw}, nil }
assert.Nil(t, RegisterSink("TestOpenWithCustomSink", ctr))
w, cleanup, err := Open("TestOpenWithCustomSink")
assert.Nil(t, err)
defer cleanup()
w.Write([]byte("test"))
}

func TestOpenWithErroringSinkFactory(t *testing.T) {
defer resetSinkRegistry()
expectedErr := errors.New("expected factory error")
ctr := func() (Sink, error) { return nil, expectedErr }
assert.Nil(t, RegisterSink("TestOpenWithErroringSinkFactory", ctr))
_, _, err := Open("TestOpenWithErroringSinkFactory")
assert.Equal(t, expectedErr, err)
}

func TestCombineWriteSyncers(t *testing.T) {
tw := &testWriter{"test", t}
w := CombineWriteSyncers(tw)
w.Write([]byte("test"))
}

func tempFileName(prefix, suffix string) string {
randBytes := make([]byte, 16)
rand.Read(randBytes)
return filepath.Join(os.TempDir(), prefix+hex.EncodeToString(randBytes)+suffix)
}

func fileExists(name string) bool {
if _, err := os.Stat(name); os.IsNotExist(err) {
return false
}
return true
}

0 comments on commit 7e7e266

Please sign in to comment.