Chunk output into files, based on number of entries and total size

Question

Below is my client code which stream all the customer URLs from a golang grpc server. It takes Request input parameter and streams customer URLs based on matching a particular clientId. In my this code, I am streaming all customer URLs for ClientId 12345 and it works fine.

I am also creating an XML file with all the URLs in it for particular clientId as shown below. I need to split an XML file into multiple XML files for same clientId to ensure it conforms to these requirements:

• A single XML file should not be more than 50MB max. It can be approximate, doesn't have to be accurate.
• A single XML file should not have more than 50K URLs max.

I know it's weird that 50k URL limit will be reached sooner than 50MB limit but this is the requirement I have. Based on above logic, I need to make multiple XML files for particular clientId. All those multiple files can be like 12345_abc_1.xml, 12345_abc_2.xml or any other better naming format.

Below is the code I got which works fine but I am opting for code review to see if anything can be improved or optimized so that it can be run efficiently.

func main() {
    // this "clientId" will be configurable in future
    clientId := 12345
    timeout := time.Duration(1000) * time.Millisecond
    ctx, _ := context.WithTimeout(context.Background(), timeout)
    conn, err := grpc.DialContext(ctx, "localhost:50005", grpc.WithInsecure())
    if err != nil {
        log.Fatalf("can not connect with server %v", err)
    }

    // create stream
    client := pb.NewCustomerServiceClient(conn)
    req := &pb.Request{ClientId: clientId}
    stream, err := client.FetchResponse(context.Background(), req)
    if err != nil {
        log.Fatalf("open stream error %v", err)
    }
    // create new object to populate all URL data in memory
    urlHolder := NewClient()
    t := time.Unix(0, 0).UTC()
    done := make(chan bool)
    go func() {
        // create new object to populate all URL data in memory
        urlHolder := NewClient()
        urlCounter := 0
        byteCounter := 0
        fileCounter := 0

        for {
            resp, err := stream.Recv()
            if err == io.EOF {
                done <- true
                file, _ := os.Create(fmt.Sprintf("%d_abc_%d.xml", clientId, fileCounter))
                urlHolder.WriteTo(file)
                return
            }
            if err != nil {
                log.Fatalf("can not receive %v", err)
            }
            log.Printf("Resp received: %s", resp.GetCustomerUrl())
            // I add the bytes of the URL here as a return
            urlBytes := urlHolder.Add(&URL{
                Loc:        resp.GetCustomerUrl(),
                LastMod:    &t,
                ChangeFreq: Daily,
                Priority:   10.2,
            })
            byteCounter += urlBytes
            urlCounter += 1
            if byteCounter > 49000000 || urlCounter >= 50000 {
                file, _ := os.Create(fmt.Sprintf("%d_abc_%d.xml", clientId, fileCounter))
                urlHolder.WriteTo(file)
                urlHolder = NewClient() // create a new object for next loop
                fileCounter += 1        // prepare fileCounter for next loop
                byteCounter = 0         // restart count variables
                urlCounter = 0
            }
        }
    }()

    <-done
    log.Printf("finished")
}

Here is my urlholder.go file:

type URL struct {
    Loc        string     `xml:"loc"`
    LastMod    *time.Time `xml:"lastmod"`
    ChangeFreq ChangeFreq `xml:"changefreq"`
    Priority   float32    `xml:"priority"`
}

type UrlMap struct {
    XMLName xml.Name `xml:"urlset"`
    Xmlns   string   `xml:"xmlns,attr"`
    URLs    []*URL   `xml:"url"`
    Minify  bool     `xml:"-"`
}

func NewClient() *UrlMap {
    return &UrlMap{
        Xmlns: "http://www.sitemaps.org/schemas/sitemap/0.9",
        URLs:  make([]*URL, 0),
    }
}

func (s *UrlMap) Add(u *URL) (int) {
    s.URLs = append(s.URLs, u)
    var urlBytes []byte
    var err error

    urlBytes, err = xml.Marshal(u) // Transform to bytes using xml.Marshal or xml.MarshalIndent
    if err != nil {
        panic(err) // or return the error
    }
    return len(urlBytes)
}

// WriteTo writes XML encoded urlMap to given io.Writer.
func (s *UrlMap) WriteTo(w io.Writer) (n int64, err error) {
    cw := NewCounterWriter(w)
    _, err = cw.Write([]byte(xml.Header))
    if err != nil {
      return cw.Count(), err
    }
    en := xml.NewEncoder(cw)
    if !s.Minify {
      en.Indent("", "  ")
    }
    err = en.Encode(s)
    cw.Write([]byte{'\n'})
    return cw.Count(), err
}

Here is my CounterWriter class -

type CounterWriter struct {
    writer io.Writer
    count  int64
}

var _ io.Writer = (*CounterWriter)(nil)

func NewCounterWriter(w io.Writer) (cw *CounterWriter) {
    return &CounterWriter{
        writer: w,
    }
}

func (cw *CounterWriter) Write(p []byte) (n int, err error) {
    n, err = cw.writer.Write(p)
    cw.count = cw.count + int64(n)
    return n, err
}

// Count returns the number of bytes written to the Writer.
func (cw *CounterWriter) Count() (n int64) {
    return cw.count
}

Problem Statement

I am looking for code review on above code:

• Is there anything that can be improved in terms of multiple XML file generation logic? I modified my Add method in urlholder.go file to return bytes (which is int) but I also have WriteTo method in same go file which also returns bytes (which is int64). Can we simplify my above multiple XML file generation logic?
• Also do we need to use interface to split out above class and make it more clean?

I am just trying to learn on how to write good go code as this code will be running in production.

Answer 1

Your code is very clear and readable in general. I came up with a very nitpicky list of style suggestions :)

A few comments on main()

• time.Duration(1000) * time.Millisecond can be written as just 1000 * time.Millisecond
• Call the cancel function from context.WithTimeout, otherwise your program has a resource leak (not that it matters from main, but still).

  ctx, cancel := context.WithTimeout(...)
  defer cancel()
  

• Use ctx consistently, there are some places in main where it refers to context.Background() instead.
• Handle the error from os.Create, or your code will be impossible to debug when it fails.
• Similarly, you probably want to print an error if urlHolder.WriteTo fails.
• I would also suggest making a helper function for creating the XML file and writing to it, in particular if you ever change the file naming pattern you would need to change it in two places.
• I think there's no need for a goroutine in main? You can just run it directly unless I'm missing something.

A few things on urlholder.go

• Time time.Time type is almost always used as a value and not as a pointer. The URL.LastMod field should be a time.Time, not a *time.Time.
• This is subjective, but UrlMap should be spelled as URLMap.
• No need to make a slice of size zero, just use nil.
• UrlMap.Add should return an error instead of panicing.
• The UrlMap.Add method uses a different method of serialization than WriteTo, so the byte count may not be correct. I would suggest calling WriteTo from inside Add, as follows

  func (m *UrlMap) Add(u *URL) int {
      s.URLs = append(s.URLs, u)
      var out strings.Builder
      m.WriteTo(&out)
      return len(out.String())
  }
  

I'm not sure what the point of CounterWriter is when the result of urlHolder.WriteTo is ignored anyway.

And since you asked about interfaces, there's no need to use an interface here. In general you should only refactor to use interfaces when you already have two implementations in mind. Using an interface preemptively for a "cleaner" design is an example of YAGNI: you ain't gonna need it.